[Mlir-commits] [mlir] c8bc72d - [mlir][spirv][cf] legalize block arguments when convert cf to spirv (#71288)

[llvmlistbot at llvm.org]

Author: Xiang Li
Date: 2023-11-04T23:48:59-04:00
New Revision: c8bc72dc79eec511d76c741fe6b6334355903dbe

URL: https://github.com/llvm/llvm-project/commit/c8bc72dc79eec511d76c741fe6b6334355903dbe
DIFF: https://github.com/llvm/llvm-project/commit/c8bc72dc79eec511d76c741fe6b6334355903dbe.diff

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

When converting branches, legalize target block arguments first.

Fixes llvm#70813

Added: 
    

Modified: 
    mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp
    mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp b/mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp
index 5dba79016120b38..8907756d33845a0 100644
--- a/mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp
+++ b/mlir/lib/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.cpp
@@ -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>(

diff  --git a/mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir b/mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir
index a10a1468e582499..4e2f6a12c8e7ee4 100644
--- a/mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/ControlFlowToSPIRV/cf-ops-to-spirv.mlir
@@ -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