[MLIR][LLVMDialect] Added branch weights attribute to CondBrOp

[georgemitenkov]

Summary

This patch introduces branch weights metadata to llvm.cond_br op in LLVM Dialect. It is modelled as optional ElementsAttr, for example:

llvm.cond_br %cond weights(dense<[1, 3]> : vector<2xi32>), ^bb1, ^bb2

When exporting to proper LLVM, this attribute is transformed into metadata node. The test for metadata creation is added to ../Target/llvmir.mlir.

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

180

Can you please avoid writing C++ parsers if you don't have to? Or file bugs for things that are missing? Seems like the only thing missing here is parsing optional attributes, as it looks like you are rewriting the parser for ArrayAttr. I would expect https://reviews.llvm.org/D83712 to be all you need to make this declarative.

[georgemitenkov]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

180

Thanks for the comment! I have followed a similar approach taken for some ops in the codebase where optional parsing is done manually, but I guess the codebase can be not very consistent then.

Using assembly format is much nicer, I agree. I will make changes to use the patch you mentioned.

[mehdi_amini]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

180

Many ops in the codebase also predates the declarative assembly feature, any migration of existing ops is welcome (either River or I would be happy to review patches to migrate to it)

[rriddle]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

180

Submitted the patch I referred to previously.

No worries George, it's hard to keep track of what can and can't be done sometimes in the codebase especially when not everything is using the desired thing. As Mehdi mentions, a lot of operations still haven't been updated to use the declarative form as new feature have been added. I try to do sweeps occasionally to port new ones over, but MLIR is getting larger and it can be hard to keep track of everything. If you see some in the future could you file bugs, or send a patch to fix if you can(that would be awesome). (Same applies to any other area of the codebase that is a bit inconsistent with the current recommended practices)

Thanks

[georgemitenkov]

Removed custom parser and printer for cond_br, and used assembly format for parsing optional attributes instead.

[georgemitenkov]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

180

Great, thanks!

[ftynse]

Thanks! A couple of small comments.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

Can this use DenseElementsAttr instead of ArrayAttr? It is more efficient and likely easier to work with in this use case.

483

Please format

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

447

Can we declare this above the if, initialize to nullptr and remove the duplication of calls to builder.CreateCondBr?

[georgemitenkov]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

If using DenseElementsAttr, then the op would be printed as

llvm.cond_br %cond dense<[5, 10]> : vector<2xi32>, ^bb1, ^bb2

which is not as concise and clear as if it was using ArrayAttr?

[georgemitenkov]

Addressed comments on branch weights attributes:

• improved format
• removed the duplication of calls to builder.CreateCondBr()

TODO: decide on whether to use DenseElementsAttr or ArrayAttr

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

I wouldn't motivate modeling choices by pretty-printing. It should be possible to wrap it with an extra keyword for better readability:

llvm.cond_br %cond weights(dense<[5, 10]> : vector<2xi32>), ^bb1, ^bb2

[georgemitenkov]

Better formatting in builder.

[georgemitenkov]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

I see, thanks!

[georgemitenkov]

Changed branch_weights from ArrayAttr to ElementsAttr.

[ftynse]

Thanks!

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

If you could constrain the type of the attribute in ODS to one with a builder, it should be elided from the printed format as well.

[mehdi_amini]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

The pretty-printing would be nice if we could customize it even more, e.g. something like:

llvm.cond_br %cond ^bb1, ^bb2  weights[5, 10]

[georgemitenkov]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

474

@rriddle @mehdi_amini I am not sure we can do this at the moment with the current infrastructure unless using a custom parser? Especially if using ElementsAttr.

branch_weights are not heavily used explicitly so I guess this printing will suffice for now?

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_BrOp : LLVM_TerminatorOp<"br",
   }];
   let builders = [LLVM_TerminatorPassthroughOpBuilder];
 }
 def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br",
     [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
      NoSideEffect]> {
   let arguments = (ins LLVMI1:$condition,
                    Variadic<LLVM_Type>:$trueDestOperands,
-                   Variadic<LLVM_Type>:$falseDestOperands);
+                   Variadic<LLVM_Type>:$falseDestOperands,
+                   OptionalAttr<I32ArrayAttr>:$branch_weights);
   let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
-  let assemblyFormat = [{
-    $condition `,`
-    $trueDest (`(` $trueDestOperands^ `:` type($trueDestOperands) `)`)? `,`
-    $falseDest (`(` $falseDestOperands^ `:` type($falseDestOperands) `)`)?
-    attr-dict
-  }];
 
   let builders = [OpBuilder<
      "OpBuilder &builder, OperationState &result, Value condition,"
      "Block *trueDest, ValueRange trueOperands,"
-    "Block *falseDest, ValueRange falseOperands", [{
-      build(builder, result, condition, trueOperands, falseOperands, trueDest,
-            falseDest);
+     "Block *falseDest, ValueRange falseOperands,"
+     "Optional<std::pair<uint32_t, uint32_t>> weights = {}", [{
+       ArrayAttr weightsAttr;
+          if (weights) {
+          weightsAttr =
+              builder.getI32ArrayAttr({static_cast<int32_t>(weights->first),
+                                       static_cast<int32_t>(weights->second)});
+          }
+        build(builder, result, condition, trueOperands, falseOperands, weightsAttr,
+            trueDest, falseDest);
   }]>, OpBuilder<
     "OpBuilder &builder, OperationState &result, Value condition,"
     "Block *trueDest, Block *falseDest, ValueRange falseOperands = {}", [{
       build(builder, result, condition, trueDest, ValueRange(), falseDest,
             falseOperands);
   }]>, LLVM_TerminatorPassthroughOpBuilder];
+
+  let parser = [{ return parseCondBrOp(parser, result); }];
+  let printer = [{ printCondBrOp(p, *this); }];
 }
 def LLVM_ReturnOp : LLVM_TerminatorOp<"return", [NoSideEffect]>,
                     Arguments<(ins Variadic<LLVM_Type>:$args)> {
   string llvmBuilder = [{
     if ($_numOperands != 0)
       builder.CreateRet($args[0]);
     else
       builder.CreateRetVoid();

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

 //===----------------------------------------------------------------------===//
 
 Optional<MutableOperandRange>
 CondBrOp::getMutableSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return index == 0 ? trueDestOperandsMutable() : falseDestOperandsMutable();
 }
 
+// <operation> ::= `llvm.cond_br` ssa-use
+//                 (`[` true-weight `:` type `,` false-weight `:` type `]`)? `,`
+//                 [<true-successor>(<true-operands>),
+//                  <false-successor>(<false-operands>)]
+static ParseResult parseCondBrOp(OpAsmParser &parser, OperationState &result) {
+  OpAsmParser::OperandType condition;
+  auto &builder = parser.getBuilder();
+  Type trueWeightType, falseWeightType;
+  Block *dest;
+
+  // Parse the condition.
+  auto int1Ty = LLVMType::getInt1Ty(
+      builder.getContext()->getRegisteredDialect<LLVM::LLVMDialect>());
+  if (parser.parseOperand(condition) ||
+      parser.resolveOperand(condition, int1Ty, result.operands))
+    return failure();
+
+  // Parse optional branch weights.
+  if (succeeded(parser.parseOptionalLSquare())) {
+    IntegerAttr trueWeight, falseWeight;
+    NamedAttrList weights;
+
+    auto i32Type = builder.getIntegerType(32);
+    if (parser.parseAttribute(trueWeight, i32Type, "true_branch_weight",
+                              weights) ||
+        parser.parseColon() || parser.parseType(trueWeightType) ||
+        !trueWeightType.isInteger(32) || parser.parseComma() ||
+        parser.parseAttribute(falseWeight, i32Type, "false_branch_weight",
+                              weights) ||
+        parser.parseColon() || parser.parseType(falseWeightType) ||
+        !falseWeightType.isInteger(32) || parser.parseRSquare())
+      return failure();
+
+    result.addAttribute("branch_weights",
+                        builder.getArrayAttr({trueWeight, falseWeight}));
+  }
+
+  // Parse the true branch.
+  SmallVector<Value, 4> trueOperands;
+  if (parser.parseComma() ||
+      parser.parseSuccessorAndUseList(dest, trueOperands))
+    return failure();
+  result.addSuccessors(dest);
+  result.addOperands(trueOperands);
+
+  // Parse the false branch.
+  SmallVector<Value, 4> falseOperands;
+  if (parser.parseComma() ||
+      parser.parseSuccessorAndUseList(dest, falseOperands))
+    return failure();
+  result.addSuccessors(dest);
+  result.addOperands(falseOperands);
+
+  result.addAttribute(
+      CondBrOp::getOperandSegmentSizeAttr(),
+      builder.getI32VectorAttr({1, static_cast<int32_t>(trueOperands.size()),
+                                static_cast<int32_t>(falseOperands.size())}));
+  return success();
+}
+
+static void printCondBrOp(OpAsmPrinter &p, CondBrOp &op) {
+  p << op.getOperationName() << ' ' << op.condition();
+  if (auto weights = op.branch_weights()) {
+    p << ' ';
+    p.printAttribute(weights.getValue());
+  }
+  p << ", ";
+  p.printSuccessorAndUseList(op.trueDest(), op.trueDestOperands());
+  p << ", ";
+  p.printSuccessorAndUseList(op.falseDest(), op.falseDestOperands());
+}
+
 //===----------------------------------------------------------------------===//
 // Printing/parsing for LLVM::LoadOp.
 //===----------------------------------------------------------------------===//
 
 static void printLoadOp(OpAsmPrinter &p, LoadOp &op) {
   p << op.getOperationName() << ' ' << op.addr();
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.addr().getType();

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
 
 #include "mlir/Dialect/LLVMIR/LLVMConversionEnumsToLLVM.inc"
#include "mlir/Dialect/LLVMIR/LLVMConversions.inc"
 
   // Emit branches.  We need to look up the remapped blocks and ignore the block
   // arguments that were transformed into PHI nodes.
   if (auto brOp = dyn_cast<LLVM::BrOp>(opInst)) {
     builder.CreateBr(blockMapping[brOp.getSuccessor()]);
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
+    auto weights = condbrOp.branch_weights();
+    if (weights) {
+      // Map weight attributes to LLVM metadata.
+      auto weightsVector = llvm::to_vector<2>(
+          llvm::map_range(weights.getValue(), [&](Attribute attr) {
+            return static_cast<uint32_t>(attr.cast<IntegerAttr>().getInt());
+          }));
+      llvm::MDNode *branchWeights = llvm::MDBuilder(llvmModule->getContext())
+                                        .createBranchWeights(weightsVector);
+      builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
+                           blockMapping[condbrOp.getSuccessor(0)],
+                           blockMapping[condbrOp.getSuccessor(1)],
+                           branchWeights);
+      return success();
+    }
     builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
                          blockMapping[condbrOp.getSuccessor(0)],
                          blockMapping[condbrOp.getSuccessor(1)]);
     return success();
   }
 
   // Emit addressof.  We need to look up the global value referenced by the
   // operation and store it in the MLIR-to-LLVM value mapping.  This does not

mlir/test/Target/llvmir.mlir

 llvm.func @address_taken() {
   llvm.return
 }
 
 llvm.mlir.global internal constant @taker_of_address() : !llvm<"void()*"> {
   %0 = llvm.mlir.addressof @address_taken : !llvm<"void()*">
   llvm.return %0 : !llvm<"void()*">
 }
+
+// -----
+
+// Check that branch weight attributes are exported properly as metadata.
+llvm.func @cond_br_weights(%cond : !llvm.i1, %arg0 : !llvm.i32,  %arg1 : !llvm.i32) -> !llvm.i32 {
+  // CHECK: !prof ![[NODE:[0-9]+]]
+  llvm.cond_br %cond [5 : i32, 10 : i32], ^bb1, ^bb2
+^bb1:  // pred: ^bb0
+  llvm.return %arg0 : !llvm.i32
+^bb2:  // pred: ^bb0
+  llvm.return %arg1 : !llvm.i32
+}
+
+// CHECK: ![[NODE]] = !{!"branch_weights", i32 5, i32 10}