[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 <maksim.levental@gmail.com>
This commit is contained in:
crZhao 2026-07-31 19:24:21 -07:00 committed by GitHub
parent 2e81492b8a
commit ef19256d33
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 30 additions and 2 deletions

View file

@ -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();
}
};

View file

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