[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 ↗

(On Diff #277337)

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 ↗

(On Diff #277337)

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 ↗

(On Diff #277337)

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 ↗

(On Diff #277337)

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 ↗

(On Diff #277337)

Great, thanks!

[ftynse]

Thanks! A couple of small comments.

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

518

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

533

Please format

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

605

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

518

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

518

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

518

I see, thanks!

[georgemitenkov]

Changed branch_weights from ArrayAttr to ElementsAttr.

[ftynse]

Thanks!

[rriddle]

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

518

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

518

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

518

@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<ElementsAttr>:$branch_weights);
   let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
   let assemblyFormat = [{
-    $condition `,`
+    $condition ( `weights` `(` $branch_weights^ `)` )? `,`
     $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 = {}", [{
+        ElementsAttr weightsAttr;
+        if (weights) {
+          weightsAttr =
+              builder.getI32VectorAttr({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];
 }
 def LLVM_ReturnOp : LLVM_TerminatorOp<"return", [NoSideEffect]>,

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/CFG.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/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
 
#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();
+    llvm::MDNode *branchWeights = nullptr;
+    if (weights) {
+      // Map weight attributes to LLVM metadata.
+      auto trueWeight =
+          weights.getValue().getValue(0).cast<IntegerAttr>().getInt();
+      auto falseWeight =
+          weights.getValue().getValue(1).cast<IntegerAttr>().getInt();
+      branchWeights =
+          llvm::MDBuilder(llvmModule->getContext())
+              .createBranchWeights(static_cast<uint32_t>(trueWeight),
+                                   static_cast<uint32_t>(falseWeight));
+    }
     builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
                          blockMapping[condbrOp.getSuccessor(0)],
-                         blockMapping[condbrOp.getSuccessor(1)]);
+                         blockMapping[condbrOp.getSuccessor(1)], branchWeights);
     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
   // emit any LLVM instruction.
   if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
     LLVM::GlobalOp global = addressOfOp.getGlobal();

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 weights(dense<[5, 10]> : vector<2xi32>), ^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}