CUTLASS 3.0 Hopper GEMMs are GETTs in disguise (#897)

2023-03-29 10:42:40 -04:00
parent 1eef5c3cf1
commit 15d9d31f1f
10 changed files with 1231 additions and 71 deletions
@@ -762,7 +762,17 @@ make_tma_copy(CopyOp,
  print("layout_tv     :  "); print(layout_tv); print("\n");
 #endif

-  return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_tv), decltype(cta_tile)>{tma_desc, gmem_stride_bases};
+  // If CTA_Tile and SLayout are incompatible, product_each makes sure
+  // that the TiledCopy generates consistent accesses.
+  auto cta_tile_tiled = [&]() {
+    if constexpr (compatible(shape(CTA_Tile{}), shape(SLayout{}))) {
+      return cta_tile;
+    } else {
+      return product_each(cta_tile);
+    }
+  }();
+
+  return TiledCopy<Copy_Atom<Traits,T>, decltype(layout_tv), decltype(cta_tile_tiled)>{tma_desc, gmem_stride_bases};
 }

 // Explicit defaulting