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

[llvmlistbot at llvm.org]

https://github.com/cr-zhao created https://github.com/llvm/llvm-project/pull/211713

### 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

>From bb851d774439f78fb0600b7df557df96960b7606 Mon Sep 17 00:00:00 2001
From: real-cpu <zhaochenrui757 at gmail.com>
Date: Thu, 23 Jul 2026 19:45:49 -0700
Subject: [PATCH] [mlir][Shard] Propagate failures in ConvertShardToMPI

---
 mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp |  9 +++++--
 .../convert-shard-to-mpi-invalid.mlir         | 24 +++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi-invalid.mlir

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)
+}