[MLIR][SPIRVToLLVM] Branch weights support for BranchConditional conversion

[georgemitenkov]

Summary

Conversion of spv.BranchConditional now supports branch weights
that are mapped to weights vector in llvm.cond_br.

[rriddle]

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

465

You could construct this with DenseElementsAttr directly(it accepts ArrayRef<Attribute>):

ShapedType weightType = VectorType::get(2, rewriter.getI32Type());
branchWeights = DenseElementsAttr::get(weightType, weights.getValue());

[georgemitenkov]

Addressed comments on constructing weights vector.

mlir/lib/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVM.cpp

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.
-    if (auto weights = op.branch_weights())
-      return failure();
+    // If branch weights exist, map them to 32-bit integer vector.
+    ElementsAttr branchWeights = nullptr;
+    if (auto weights = op.branch_weights()) {
+      VectorType weightType = VectorType::get(2, rewriter.getI32Type());
+      branchWeights =
+          DenseElementsAttr::get(weightType, weights.getValue().getValue());
+    }
 
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
-        op, op.condition(), op.getTrueBlock(), op.getTrueBlockArguments(),
-        op.getFalseBlock(), op.getFalseBlockArguments());
+        op, op.condition(), op.getTrueBlockArguments(),
+        op.getFalseBlockArguments(), branchWeights, op.getTrueBlock(),
+        op.getFalseBlock());
     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> {

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

^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}}
+    // CHECK: llvm.cond_br %{{.*}} weights(dense<[1, 2]> : vector<2xi32>), ^bb1, ^bb2
     spv.BranchConditional %cond [1, 2], ^true, ^false
+  // CHECK: ^bb1:
   ^true:
     spv.Return
+  // CHECK: ^bb2:
   ^false:
     spv.Return
   }
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//