Slice layout wtf

triton-lang · Nov 20, 2024 · 3841712 · 3841712
1 parent 5b7dc8b
commit 3841712
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -206,6 +206,11 @@ SmallVector<unsigned> getShapePerCTATile(Attribute layout) {
           mlir::dyn_cast<DistributedEncodingTrait>(layout)) {
     auto sizePerThread = distributedLayout.getSizePerThread();
     auto threadsPerWarp = distributedLayout.getThreadsPerWarp();
+    // ThreadsPerWarp does not align with this function for slice layout
+    if (auto sliceLayout = mlir::dyn_cast<SliceEncodingAttr>(layout)) {
+      threadsPerWarp = getThreadsPerWarp(sliceLayout.getParent());
+      threadsPerWarp.erase(threadsPerWarp.begin() + sliceLayout.getDim());
+    }
     auto warpsPerCTA = distributedLayout.getWarpsPerCTA();
     assert(sizePerThread.size() == threadsPerWarp.size() &&
            sizePerThread.size() == warpsPerCTA.size());