[Mlir-commits] [mlir] [mlir][shard] Fix FoldDuplicateShardOp assertion on duplicates in different blocks (PR #232953)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Stanislav Bardyuk (kodlan)

<details>
<summary>Changes</summary>

`FoldDuplicateShardOp` looks for another `shard.shard` of the same value and checks `otherOp->isBeforeInBlock(op)` to decide which one is the duplicate. `isBeforeInBlock` asserts when the two ops are not in the same block, so a value that is annotated once in the function body and once inside a region (for example a `linalg.generic` or `scf.if` body) crashes `--canonicalize`:

```
Operation.cpp:382: bool mlir::Operation::isBeforeInBlock(mlir::Operation *): Assertion `other && other->block == block && "Expected other operation to have the same parent block."' failed.
```

The pattern was written for the same-block duplicates that sharding propagation produces, so just skip the fold when the other op lives in a different block. Both annotations stay in place, which is always valid.

Adds a canonicalization test with a duplicate inside an `scf.if` region.

Fixes #<!-- -->223332

---
Full diff: https://github.com/llvm/llvm-project/pull/232953.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Shard/IR/ShardOps.cpp (+3-1) 
- (modified) mlir/test/Dialect/Shard/canonicalization.mlir (+25) 


``````````diff
diff --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index a4206d6717687c..1bc1fae851738c 100644
--- a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
+++ b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
@@ -870,7 +870,9 @@ class FoldDuplicateShardOp final : public OpRewritePattern<ShardOp> {
     for (auto &use : value.getUses()) {
       if (use.getOwner() != op.getOperation()) {
         auto otherOp = dyn_cast<ShardOp>(use.getOwner());
-        if (!otherOp || !otherOp->isBeforeInBlock(op)) {
+        // Only fold duplicates in the same block; isBeforeInBlock requires it.
+        if (!otherOp || otherOp->getBlock() != op->getBlock() ||
+            !otherOp->isBeforeInBlock(op)) {
           return failure();
         }
         // Create a Sharding object for the current and the other ShardOp
diff --git a/mlir/test/Dialect/Shard/canonicalization.mlir b/mlir/test/Dialect/Shard/canonicalization.mlir
index a3a9c592ff0aec..e16b84dfb23764 100644
--- a/mlir/test/Dialect/Shard/canonicalization.mlir
+++ b/mlir/test/Dialect/Shard/canonicalization.mlir
@@ -246,3 +246,28 @@ func.func @test_duplicate_shardops_diff() -> (tensor<1024x1024xf32>, tensor<1024
   // CHECK-NEXT: return [[vsharded_1]], [[vsharded]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>
   return %sharded_1, %sharded_2 : tensor<1024x1024xf32>, tensor<1024x1024xf32>
 }
+
+// CHECK-LABEL: func @test_duplicate_shardops_different_blocks
+// CHECK-SAME: [[varg0:%.*]]: tensor<1024x1024xf32>, [[varg1:%.*]]: tensor<1024x1024xf32>, [[varg2:%.*]]: i1
+func.func @test_duplicate_shardops_different_blocks(%arg0: tensor<1024x1024xf32>, %arg1: tensor<1024x1024xf32>, %arg2: i1) -> (tensor<1024x1024xf32>, tensor<1024x1024xf32>) {
+  // CHECK-NEXT: [[vsharding:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}0]] : !shard.sharding
+  %sharding_1 = shard.sharding @grid4x4 split_axes = [[0]] : !shard.sharding
+  // CHECK-NEXT: [[vsharded:%.*]] = shard.shard [[varg0]] to [[vsharding]] : tensor<1024x1024xf32>
+  %sharded_1 = shard.shard %arg0 to %sharding_1 : tensor<1024x1024xf32>
+  // CHECK-NEXT: [[vif:%.*]] = scf.if [[varg2]] -> (tensor<1024x1024xf32>) {
+  %0 = scf.if %arg2 -> (tensor<1024x1024xf32>) {
+    // CHECK-NEXT: [[vsharding_0:%.*]] = shard.sharding @grid4x4 split_axes = {{\[\[}}1]] : !shard.sharding
+    %sharding_2 = shard.sharding @grid4x4 split_axes = [[1]] : !shard.sharding
+    // CHECK-NEXT: [[vsharded_0:%.*]] = shard.shard [[varg0]] to [[vsharding_0]] : tensor<1024x1024xf32>
+    %sharded_2 = shard.shard %arg0 to %sharding_2 : tensor<1024x1024xf32>
+    // CHECK-NEXT: scf.yield [[vsharded_0]] : tensor<1024x1024xf32>
+    scf.yield %sharded_2 : tensor<1024x1024xf32>
+  // CHECK-NEXT: } else {
+  } else {
+    // CHECK-NEXT: scf.yield [[varg1]] : tensor<1024x1024xf32>
+    scf.yield %arg1 : tensor<1024x1024xf32>
+  // CHECK-NEXT: }
+  }
+  // CHECK-NEXT: return [[vsharded]], [[vif]] : tensor<1024x1024xf32>, tensor<1024x1024xf32>
+  return %sharded_1, %0 : tensor<1024x1024xf32>, tensor<1024x1024xf32>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/232953