[MLIR][SPIRVToLLVM] Conversion of SPIR-V branch ops

[georgemitenkov]

Summary

This patch introduces conversion for spv.Branch and spv.BranchConditional ops. Branch weigths for spv.BranchConditional are not supported at the moment, and conversion in this case fails.

This revision relies on https://reviews.llvm.org/D83786.

[antiagainst]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

374

You can just write this as if (auto weights = ...) return failure.

[georgemitenkov]

Addressed comments on style.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

matchAndRewrite(spirv::BitFieldUExtractOp op, ArrayRef<Value> operands,
     // Shift `Base` by `Offset` and apply the mask on it.
     Value shiftedBase =
         rewriter.create<LLVM::LShrOp>(loc, dstType, op.base(), offset);
     rewriter.replaceOpWithNewOp<LLVM::AndOp>(op, dstType, shiftedBase, mask);
     return success();
   }
 };
 
+class BranchConversionPattern : public SPIRVToLLVMConversion<spirv::BranchOp> {
+public:
+  using SPIRVToLLVMConversion<spirv::BranchOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::BranchOp branchOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<LLVM::BrOp>(branchOp, operands,
+                                            branchOp.getTarget());
+    return success();
+  }
+};
+
+class BranchConditionalConversionPattern
+    : public SPIRVToLLVMConversion<spirv::BranchConditionalOp> {
+public:
+  using SPIRVToLLVMConversion<
+      spirv::BranchConditionalOp>::SPIRVToLLVMConversion;
+
+  LogicalResult
+  matchAndRewrite(spirv::BranchConditionalOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // There is no support of branch weights in LLVM dialect at the moment.
+    auto weights = op.branch_weights();
+    if (weights)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
+        op, op.condition(), op.getTrueBlock(), op.getTrueBlockArguments(),
+        op.getFalseBlock(), op.getFalseBlockArguments());
+    return success();
+  }
+};
+
 /// Converts SPIR-V operations that have straightforward LLVM equivalent
 /// into LLVM dialect operations.
 template <typename SPIRVOp, typename LLVMOp>
 class DirectConversionPattern : public SPIRVToLLVMConversion<SPIRVOp> {
 public:
   using SPIRVToLLVMConversion<SPIRVOp>::SPIRVToLLVMConversion;
 
   LogicalResult
patterns.insert<
       IComparePattern<spirv::UGreaterThanOp, LLVM::ICmpPredicate::ugt>,
       IComparePattern<spirv::UGreaterThanEqualOp, LLVM::ICmpPredicate::uge>,
       IComparePattern<spirv::ULessThanEqualOp, LLVM::ICmpPredicate::ule>,
       IComparePattern<spirv::ULessThanOp, LLVM::ICmpPredicate::ult>,
 
       // Constant op
       ConstantScalarAndVectorPattern,
 
+      // Control Flow ops
+      BranchConversionPattern, BranchConditionalConversionPattern,
+
       // Function Call op
       FunctionCallPattern,
 
       // Logical ops
       DirectConversionPattern<spirv::LogicalAndOp, LLVM::AndOp>,
       DirectConversionPattern<spirv::LogicalOrOp, LLVM::OrOp>,
       IComparePattern<spirv::LogicalEqualOp, LLVM::ICmpPredicate::eq>,
       IComparePattern<spirv::LogicalNotEqualOp, LLVM::ICmpPredicate::ne>,

mlir/test/Conversion/SPIRVToLLVM/control-flow-ops-to-llvm.mlir

+// RUN: mlir-opt -convert-spirv-to-llvm -split-input-file -verify-diagnostics %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.Branch
+//===----------------------------------------------------------------------===//
+
+spv.module Logical GLSL450 {
+  spv.func @branch_without_arguments() -> () "None" {
+    // CHECK: llvm.br ^bb1
+    spv.Branch ^label
+  // CHECK: ^bb1
+  ^label:
+    spv.Return
+  }
+
+  spv.func @branch_with_arguments() -> () "None" {
+    %0 = spv.constant 0 : i32
+    %1 = spv.constant true
+    // CHECK: llvm.br ^bb1(%{{.*}}, %{{.*}} : !llvm.i32, !llvm.i1)
+    spv.Branch ^label(%0, %1: i32, i1)
+  // CHECK: ^bb1(%{{.*}}: !llvm.i32, %{{.*}}: !llvm.i1)
+  ^label(%arg0: i32, %arg1: i1):
+    spv.Return
+  }
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.BranchConditional
+//===----------------------------------------------------------------------===//
+
+spv.module Logical GLSL450 {
+  spv.func @cond_branch_without_arguments() -> () "None" {
+    // CHECK: %[[COND:.*]] = llvm.mlir.constant(true) : !llvm.i1
+    %cond = spv.constant true
+    // CHECK: lvm.cond_br %[[COND]], ^bb1, ^bb2
+    spv.BranchConditional %cond, ^true, ^false
+    // CHECK: ^bb1:
+  ^true:
+    spv.Return
+    // CHECK: ^bb2:
+  ^false:
+    spv.Return
+  }
+
+  spv.func @cond_branch_with_arguments_nested() -> () "None" {
+    // CHECK: %[[COND1:.*]] = llvm.mlir.constant(true) : !llvm.i1
+    %cond = spv.constant true
+    %0 = spv.constant 0 : i32
+    // CHECK: %[[COND2:.*]] = llvm.mlir.constant(false) : !llvm.i1
+    %false = spv.constant false
+    // CHECK: llvm.cond_br %[[COND1]], ^bb1(%{{.*}}, %[[COND2]] : !llvm.i32, !llvm.i1), ^bb2
+    spv.BranchConditional %cond, ^outer_true(%0, %false: i32, i1), ^outer_false
+  // CHECK: ^bb1(%{{.*}}: !llvm.i32, %[[COND:.*]]: !llvm.i1):
+  ^outer_true(%arg0: i32, %arg1: i1):
+    // CHECK: llvm.cond_br %[[COND]], ^bb3, ^bb4(%{{.*}}, %{{.*}} : !llvm.i32, !llvm.i32)
+    spv.BranchConditional %arg1, ^inner_true, ^inner_false(%arg0, %arg0: i32, i32)
+  // CHECK: ^bb2:
+  ^outer_false:
+    spv.Return
+  // CHECK: ^bb3:
+  ^inner_true:
+    spv.Return
+  // CHECK: ^bb4(%{{.*}}: !llvm.i32, %{{.*}}: !llvm.i32):
+  ^inner_false(%arg3: i32, %arg4: i32):
+    spv.Return
+  }
+}
+
+// -----
+
+spv.module Logical GLSL450 {
+  spv.func @cond_branch_with_weights(%cond: i1) -> () "None" {
+    // expected-error@+1 {{failed to legalize operation 'spv.BranchConditional' that was explicitly marked illegal}}
+    spv.BranchConditional %cond [1, 2], ^true, ^false
+  ^true:
+    spv.Return
+  ^false:
+    spv.Return
+  }
+}