colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[MLIR][KernelGen] Generate assertion message in `transform_unranked_hlo` pass 

Use constant to generate the correct assertion message. This avoids
confusion when lowering the max rank specialization for debugging.

PiperOrigin-RevId: 344769021
This commit is contained in:
A. Unique TensorFlower 2020-11-30 01:40:08 -08:00 committed by TensorFlow MLIR Team
parent 1fe04c0aff
commit dd15c6cd84
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
View File

@ -400,7 +400,7 @@ struct ConvertUnrankedDynamicBroadcastBinaryOp
rewriter.create<SelectOp>(loc, greater_rank_lhs, lhs_rank, rhs_rank); rewriter.create<SelectOp>(loc, greater_rank_lhs, lhs_rank, rhs_rank);
// Generate a list of nested if/else statements to handle rank // Generate a list of nested if/else statements to handle rank
// specializations from 1-6. // specializations from 1 to `kMaxRankSpecialization`.
scf::IfOp if_op = createIfOpForRankSpecializedBroadcastAndOp( scf::IfOp if_op = createIfOpForRankSpecializedBroadcastAndOp(
rewriter, op, greater_rank, 1); rewriter, op, greater_rank, 1);
OpBuilder if_builder = if_op.getThenBodyBuilder(rewriter.getListener()); OpBuilder if_builder = if_op.getThenBodyBuilder(rewriter.getListener());
@ -419,13 +419,13 @@ struct ConvertUnrankedDynamicBroadcastBinaryOp
else_builder = inner_if.getElseBodyBuilder(rewriter.getListener()); else_builder = inner_if.getElseBodyBuilder(rewriter.getListener());
} }
// Fire an assertion if none of the rank specializations applied (one of // Fire an assertion if none of the rank specializations applied (one of
// the ranks was greater than 6). // the ranks was greater than `kMaxRankSpecialization`).
else_builder.create<AssertOp>( else_builder.create<AssertOp>(
loc, loc,
GreaterRankIsN(else_builder, op.getLoc(), greater_rank, GreaterRankIsN(else_builder, op.getLoc(), greater_rank,
kMaxRankSpecialization), kMaxRankSpecialization),
"Input for dynamic binary op lowering was of a rank greater than " "Input for dynamic binary op lowering was of a rank greater than " +
"6"); std::to_string(kMaxRankSpecialization));
// Add the rank 6 specialization to the innermost else block. // Add the rank 6 specialization to the innermost else block.
createRankSpecializedBroadcastAndOp(else_builder, op, lhs, rhs, createRankSpecializedBroadcastAndOp(else_builder, op, lhs, rhs,
kMaxRankSpecialization); kMaxRankSpecialization);