From ef19256d338802a6bb4a8b3761d675dfd65e8dec Mon Sep 17 00:00:00 2001 From: crZhao Date: Fri, 31 Jul 2026 19:24:21 -0700 Subject: [PATCH] [mlir][Shard] Propagate failures in ConvertShardToMPI (#211713) This PR fixes a crash in ConvertShardToMPI caused by ignoring a failed partial conversion and continuing to apply folding patterns to partially converted IR. Assisted-by: codex --------- Co-authored-by: Maksim Levental --- mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp | 8 +++++-- .../convert-shard-to-mpi-invalid.mlir | 24 +++++++++++++++++++ 2 files changed, 30 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 ea2825bff6bf..6ad80619a664 100644 --- a/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp +++ b/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp @@ -1212,13 +1212,17 @@ struct ConvertShardToMPIPass populateCallOpTypeConversionPattern(patterns, typeConverter); populateReturnOpTypeConversionPattern(patterns, typeConverter); - (void)applyPartialConversion(getOperation(), target, std::move(patterns)); + if (failed(applyPartialConversion(getOperation(), target, + std::move(patterns)))) { + return signalPassFailure(); + } // 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)))) + return 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 000000000000..4d93e131e14b --- /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 , %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@+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) +}