[Mlir-commits] [mlir] [mlir][Shard] Propagate failures in ConvertShardToMPI (PR #211713)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: crZhao (cr-zhao)

<details>
<summary>Changes</summary>

### Summary

- Propagate failures from partial Shard-to-MPI conversion and stop the pass before running folding patterns on partially converted IR.
- Propagate failures from greedy folding pattern application.
- Add a regression test for the assertion crash reported in #<!-- -->208331.

Fixes #<!-- -->208331

### Testing

- ShardToMPI invalid conversion regression test
- Existing ShardToMPI conversion test

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


2 Files Affected:

- (modified) mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp (+7-2) 
- (added) mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi-invalid.mlir (+24) 


``````````diff
diff --git a/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp b/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp
index ea2825bff6bf7..94eb9a3e0f031 100644
--- a/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp
+++ b/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp
@@ -1212,13 +1212,18 @@ struct ConvertShardToMPIPass
     populateCallOpTypeConversionPattern(patterns, typeConverter);
     populateReturnOpTypeConversionPattern(patterns, typeConverter);
 
-    (void)applyPartialConversion(getOperation(), target, std::move(patterns));
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns)))) {
+      signalPassFailure();
+      return;
+    }
 
     // Folding patterns cannot be mixed with conversion patterns -> extra pass.
     patterns.clear();
     SymbolTableCollection symbolTableCollection;
     mlir::shard::populateFoldingPatterns(patterns, symbolTableCollection);
-    (void)applyPatternsGreedily(getOperation(), std::move(patterns));
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+      signalPassFailure();
   }
 };
 
diff --git a/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi-invalid.mlir b/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi-invalid.mlir
new file mode 100644
index 0000000000000..4d93e131e14be
--- /dev/null
+++ b/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi-invalid.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s -convert-shard-to-mpi -verify-diagnostics
+
+module {
+  func.func @vector_ops(%arg0: memref<4xf32>) {
+    %mask = vector.constant_mask [8] : vector<8xi1>
+    %zero = arith.constant 0.000000e+00 : f32
+    %broadcast = vector.broadcast %zero : f32 to vector<8xf32>
+    %reduction = vector.mask %mask {
+      vector.multi_reduction <add>, %broadcast, %zero [0]
+        : vector<8xf32> to f32
+    } : vector<8xi1> -> f32
+    %result = arith.addf %reduction, %zero : f32
+    return
+  }
+
+  func.func @unsupported_send(%arg0: tensor<4xf32>) {
+    // expected-error at +1 {{failed to legalize operation 'shard.send' that was explicitly marked illegal}}
+    %send = shard.send %arg0 on @grid destination = []
+      : (tensor<4xf32>) -> tensor<4xf32>
+    return
+  }
+
+  shard.grid @grid(shape = 2)
+}

``````````

</details>


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