bug fix

bytedance · Jul 2, 2024 · 9924fa8 · 9924fa8
1 parent b5f4c1a
commit 9924fa8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/compiler/lib/Dialect/GPU/Transforms/GPUBlockSwizzle.cpp b/compiler/lib/Dialect/GPU/Transforms/GPUBlockSwizzle.cpp
@@ -131,6 +131,8 @@ static LogicalResult reorderForallOpMappedToBlock(scf::ForallOp forallOp,
   IRMapping bvm;
   bvm.map(originWorkgroupIds[0], swizzledIdX);
   bvm.map(originWorkgroupIds[1], swizzledIdY);
+  if (mapping.size() == 3)
+    bvm.map(originWorkgroupIds[2], workgroupIds[2]);
   for (auto &op : forallOp.getBody()->getOperations()) {
     b.clone(op, bvm);
   }