[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:
parent
2e81492b8a
commit
ef19256d33
2 changed files with 30 additions and 2 deletions
|
|
@ -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();
|
||||
}
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
Loading…
Reference in a new issue