caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][spirv][cf] legalize block arguments when convert cf to spirv (#71288)

Browse Source 
When converting branches, legalize target block arguments first.

Fixes llvm#70813
This commit is contained in:
Xiang Li
2023-11-04 23:48:59 -04:00
committed by GitHub
parent 437c62178c
commit c8bc72dc79
2 changed files with 36 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp
View File

@@ -28,18 +28,33 @@
using namespace mlir;
/// Checks that the target block arguments are legal.
static LogicalResult checkBlockArguments(Block &block, Operation *op,
PatternRewriter &rewriter,
const TypeConverter &converter) {
for (BlockArgument arg : block.getArguments()) {
if (!converter.isLegal(arg.getType())) {
/// Legailze target block arguments.
static LogicalResult legalizeBlockArguments(Block &block, Operation *op,
PatternRewriter &rewriter,
const TypeConverter &converter) {
auto builder = OpBuilder::atBlockBegin(&block);
for (unsigned i = 0; i < block.getNumArguments(); ++i) {
const auto arg = block.getArgument(i);
if (converter.isLegal(arg.getType()))
continue;
Type ty = arg.getType();
Type newTy = converter.convertType(ty);
if (!newTy) {
return rewriter.notifyMatchFailure(
op,
llvm::formatv(
"failed to match, destination argument not legalized (found {0})",
arg));
op, llvm::formatv("failed to legalize type for argument {0})", arg));
}
unsigned argNum = arg.getArgNumber();
Location loc = arg.getLoc();
Value newArg = block.insertArgument(argNum, newTy, loc);
Value convertedValue = converter.materializeSourceConversion(
builder, op->getLoc(), ty, newArg);
if (!convertedValue) {
return rewriter.notifyMatchFailure(
op, llvm::formatv("failed to cast new argument {0} to type {1})",
newArg, ty));
}
arg.replaceAllUsesWith(convertedValue);
block.eraseArgument(argNum + 1);
}
return success();
}
@@ -56,8 +71,8 @@ struct BranchOpPattern final : OpConversionPattern<cf::BranchOp> {
LogicalResult
matchAndRewrite(cf::BranchOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (failed(checkBlockArguments(*op.getDest(), op, rewriter,
*getTypeConverter())))
if (failed(legalizeBlockArguments(*op.getDest(), op, rewriter,
*getTypeConverter())))
return failure();
rewriter.replaceOpWithNewOp<spirv::BranchOp>(op, op.getDest(),
@@ -73,12 +88,12 @@ struct CondBranchOpPattern final : OpConversionPattern<cf::CondBranchOp> {
LogicalResult
matchAndRewrite(cf::CondBranchOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
if (failed(checkBlockArguments(*op.getTrueDest(), op, rewriter,
*getTypeConverter())))
if (failed(legalizeBlockArguments(*op.getTrueDest(), op, rewriter,
*getTypeConverter())))
return failure();
if (failed(checkBlockArguments(*op.getFalseDest(), op, rewriter,
*getTypeConverter())))
if (failed(legalizeBlockArguments(*op.getFalseDest(), op, rewriter,
*getTypeConverter())))
return failure();
rewriter.replaceOpWithNewOp<spirv::BranchConditionalOp>(

6
mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir
View File

@@ -39,16 +39,20 @@ func.func @simple_loop(%begin: i32, %end: i32, %step: i32) {
// -----
// TODO: We should handle blocks whose arguments require type conversion.
// Handle blocks whose arguments require type conversion.
// CHECK-LABEL: func.func @main_graph
func.func @main_graph(%arg0: index) {
%c3 = arith.constant 1 : index
// CHECK: spirv.Branch ^bb1({{.*}} : i32)
cf.br ^bb1(%arg0 : index)
// CHECK: ^bb1({{.*}}: i32): // 2 preds: ^bb0, ^bb2
^bb1(%0: index): // 2 preds: ^bb0, ^bb2
%1 = arith.cmpi slt, %0, %c3 : index
// CHECK: spirv.BranchConditional {{.*}}, ^bb2, ^bb3
cf.cond_br %1, ^bb2, ^bb3
^bb2: // pred: ^bb1
// CHECK: spirv.Branch ^bb1({{.*}} : i32)
cf.br ^bb1(%c3 : index)
^bb3: // pred: ^bb1
return