[mlir][llvm] Add branch weights to call and invoke

[Dinistro]

Summary

This commit introduces branch weight attributes to the LLVM::CallOp and
LLVM::InvokeOp and adds both import and export of them.

[gysit]

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

120–127

nit: do we want to use a TypeSwitch now that the list of operation starts to get long?

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

326

nit: add a brief doc comment

330

nit: I believe the llvm prefix is not needed for SmallVector and APInt?

354–357

I think "If the called function has a result, remap the corresponding value. Note that LLVM IR dialect CallOp has either 0 or 1 result." should move down to where it was before?

if (opInst.getNumResults() != 0) {
  moduleTranslation.mapValue(opInst.getResult(0), call);
  return success();
}

mlir/test/Target/LLVMIR/llvmir.mlir

1638

nit: indent by 2 spaces.

1653

nit: indent by 2 spaces.

[Dinistro]

address review comments

[gysit]

LGTM!

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

 // Call-related operations.
 def LLVM_InvokeOp : LLVM_Op<"invoke", [
                       AttrSizedOperandSegments,
                       DeclareOpInterfaceMethods<BranchOpInterface>,
                       DeclareOpInterfaceMethods<CallOpInterface>, Terminator]> {
   let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                    Variadic<LLVM_Type>:$callee_operands,
                    Variadic<LLVM_Type>:$normalDestOperands,
-                   Variadic<LLVM_Type>:$unwindDestOperands);
+                   Variadic<LLVM_Type>:$unwindDestOperands,
+                   OptionalAttr<ElementsAttr>:$branch_weights);
   let results = (outs Variadic<LLVM_Type>);
   let successors = (successor AnySuccessor:$normalDest,
                               AnySuccessor:$unwindDest);
 
   let builders = [
     OpBuilder<(ins "TypeRange":$tys, "FlatSymbolRefAttr":$callee,
       "ValueRange":$ops, "Block*":$normal, "ValueRange":$normalOps,
       "Block*":$unwind, "ValueRange":$unwindOps),
     [{
       $_state.addAttribute("callee", callee);
       build($_builder, $_state, tys, ops, normal, normalOps, unwind, unwindOps);
     }]>,
     OpBuilder<(ins "TypeRange":$tys, "ValueRange":$ops, "Block*":$normal,
       "ValueRange":$normalOps, "Block*":$unwind, "ValueRange":$unwindOps),
     [{
       build($_builder, $_state, tys, /*callee=*/FlatSymbolRefAttr(), ops, normalOps,
-            unwindOps, normal, unwind);
+            unwindOps, nullptr, normal, unwind);
     }]>];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_LandingpadOp : LLVM_Op<"landingpad"> {
   let arguments = (ins UnitAttr:$cleanup, Variadic<LLVM_Type>);
   let results = (outs LLVM_Type:$res);
let description = [{
     // Indirect call with an argument and without a result.
     llvm.call %1(%0) : (f32) -> ()
     ```
   }];
 
   let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                        Variadic<LLVM_Type>,
                        DefaultValuedAttr<LLVM_FastmathFlagsAttr,
-                                         "{}">:$fastmathFlags);
+                                         "{}">:$fastmathFlags,
+                       OptionalAttr<ElementsAttr>:$branch_weights);
   let results = (outs Optional<LLVM_Type>:$result);
 
   let builders = [
     OpBuilder<(ins "LLVMFuncOp":$func, "ValueRange":$args)>,
     OpBuilder<(ins "TypeRange":$results, "StringAttr":$callee,
                    CArg<"ValueRange", "{}">:$args)>,
+    OpBuilder<(ins "TypeRange":$results, "FlatSymbolRefAttr":$callee,
+                   CArg<"ValueRange", "{}">:$args)>,
     OpBuilder<(ins "TypeRange":$results, "StringRef":$callee,
                    CArg<"ValueRange", "{}">:$args)>
   ];
 
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

 
 void CallOp::build(OpBuilder &builder, OperationState &state, TypeRange results,
                    StringRef callee, ValueRange args) {
   build(builder, state, results, builder.getStringAttr(callee), args);
 }
 
 void CallOp::build(OpBuilder &builder, OperationState &state, TypeRange results,
                    StringAttr callee, ValueRange args) {
-  build(builder, state, results, SymbolRefAttr::get(callee), args, nullptr);
+  build(builder, state, results, SymbolRefAttr::get(callee), args, nullptr,
+        nullptr);
+}
+
+void CallOp::build(OpBuilder &builder, OperationState &state, TypeRange results,
+                   FlatSymbolRefAttr callee, ValueRange args) {
+  build(builder, state, results, callee, args, nullptr, nullptr);
 }
 
 void CallOp::build(OpBuilder &builder, OperationState &state, LLVMFuncOp func,
                    ValueRange args) {
   SmallVector<Type> results;
   Type resultType = func.getFunctionType().getReturnType();
   if (!resultType.isa<LLVM::LLVMVoidType>())
     results.push_back(resultType);
-  build(builder, state, results, SymbolRefAttr::get(func), args, nullptr);
+  build(builder, state, results, SymbolRefAttr::get(func), args, nullptr,
+        nullptr);
 }
 
 CallInterfaceCallable CallOp::getCallableForCallee() {
   // Direct call.
   if (FlatSymbolRefAttr calleeAttr = getCalleeAttr())
     return calleeAttr;
   // Indirect call, callee Value is the first operand.
   return getOperand(0);

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/ModuleImport.h"
 
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
static LogicalResult setProfilingAttrs(OpBuilder &builder, llvm::MDNode *node,
   branchWeights.reserve(node->getNumOperands() - 1);
   for (unsigned i = 1, e = node->getNumOperands(); i != e; ++i) {
     llvm::ConstantInt *branchWeight =
         llvm::mdconst::extract<llvm::ConstantInt>(node->getOperand(i));
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 
   // Attach the branch weights to the operations that support it.
-  if (auto condBrOp = dyn_cast<CondBrOp>(op)) {
-    condBrOp.setBranchWeightsAttr(builder.getI32VectorAttr(branchWeights));
+  return llvm::TypeSwitch<Operation *, LogicalResult>(op)
+      .Case<CondBrOp, SwitchOp, CallOp, InvokeOp>([&](auto branchWeightOp) {
+        branchWeightOp.setBranchWeightsAttr(
+            builder.getI32VectorAttr(branchWeights));
         return success();
-  }
-  if (auto switchOp = dyn_cast<SwitchOp>(op)) {
-    switchOp.setBranchWeightsAttr(builder.getI32VectorAttr(branchWeights));
-    return success();
-  }
-  return failure();
+      })
+      .Default([](auto) { return failure(); });
 }
 
 /// Attaches the given TBAA metadata `node` to the imported operation.
 /// Returns success, if the metadata has been converted and the attachment
 /// succeeds, failure - otherwise.
 static LogicalResult setTBAAAttrs(const llvm::MDNode *node, Operation *op,
                                   LLVM::ModuleImport &moduleImport) {
   SymbolRefAttr tbaaTagSym = moduleImport.lookupTBAAAttr(node);
   if (!tbaaTagSym)

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

convertCallLLVMIntrinsicOp(CallIntrinsicOp &op, llvm::IRBuilderBase &builder,
 
   auto *inst =
       builder.CreateCall(fn, moduleTranslation.lookupValues(op.getOperands()));
   if (op.getNumResults() == 1)
     moduleTranslation.mapValue(op->getResults().front()) = inst;
   return success();
 }
 
+/// Constructs branch weights metadata if the provided `weights` hold a value,
+/// otherwise returns nullptr.
+static llvm::MDNode *
+convertBranchWeights(std::optional<ElementsAttr> weights,
+                     LLVM::ModuleTranslation &moduleTranslation) {
+  if (!weights)
+    return nullptr;
+  SmallVector<uint32_t> weightValues;
+  weightValues.reserve(weights->size());
+  for (APInt weight : weights->cast<DenseIntElementsAttr>())
+    weightValues.push_back(weight.getLimitedValue());
+  return llvm::MDBuilder(moduleTranslation.getLLVMContext())
+      .createBranchWeights(weightValues);
+}
+
 static LogicalResult
 convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
                      LLVM::ModuleTranslation &moduleTranslation) {
 
   llvm::IRBuilder<>::FastMathFlagGuard fmfGuard(builder);
   if (auto fmf = dyn_cast<FastmathFlagsInterface>(opInst))
     builder.setFastMathFlags(getFastmathFlags(fmf));
 
 #include "mlir/Dialect/LLVMIR/LLVMConversions.inc"
 #include "mlir/Dialect/LLVMIR/LLVMIntrinsicConversions.inc"
 
   // Emit function calls.  If the "callee" attribute is present, this is a
   // direct function call and we also need to look up the remapped function
   // itself.  Otherwise, this is an indirect call and the callee is the first
-  // operand, look it up as a normal value.  Return the llvm::Value
-  // representing the function result, which may be of llvm::VoidTy type.
-  auto convertCall = [&](Operation &op) -> llvm::Value * {
-    auto operands = moduleTranslation.lookupValues(op.getOperands());
+  // operand, look it up as a normal value.
+  if (auto callOp = dyn_cast<LLVM::CallOp>(opInst)) {
+    auto operands = moduleTranslation.lookupValues(callOp.getOperands());
     ArrayRef<llvm::Value *> operandsRef(operands);
-    if (auto attr = op.getAttrOfType<FlatSymbolRefAttr>("callee"))
-      return builder.CreateCall(
+    llvm::CallInst *call;
+    if (auto attr = callOp.getCalleeAttr()) {
+      call = builder.CreateCall(
           moduleTranslation.lookupFunction(attr.getValue()), operandsRef);
+    } else {
       auto calleeType =
-        op.getOperands().front().getType().cast<LLVMPointerType>();
+          callOp->getOperands().front().getType().cast<LLVMPointerType>();
       auto *calleeFunctionType = cast<llvm::FunctionType>(
           moduleTranslation.convertType(calleeType.getElementType()));
-    return builder.CreateCall(calleeFunctionType, operandsRef.front(),
+      call = builder.CreateCall(calleeFunctionType, operandsRef.front(),
                                 operandsRef.drop_front());
-  };
-
-  // Emit calls.  If the called function has a result, remap the corresponding
-  // value.  Note that LLVM IR dialect CallOp has either 0 or 1 result.
-  if (isa<LLVM::CallOp>(opInst)) {
-    llvm::Value *result = convertCall(opInst);
+    }
+    llvm::MDNode *branchWeights =
+        convertBranchWeights(callOp.getBranchWeights(), moduleTranslation);
+    if (branchWeights)
+      call->setMetadata(llvm::LLVMContext::MD_prof, branchWeights);
+    // If the called function has a result, remap the corresponding value.  Note
+    // that LLVM IR dialect CallOp has either 0 or 1 result.
     if (opInst.getNumResults() != 0) {
-      moduleTranslation.mapValue(opInst.getResult(0), result);
+      moduleTranslation.mapValue(opInst.getResult(0), call);
       return success();
     }
     // Check that LLVM call returns void for 0-result functions.
-    return success(result->getType()->isVoidTy());
+    return success(call->getType()->isVoidTy());
   }
 
   if (auto inlineAsmOp = dyn_cast<LLVM::InlineAsmOp>(opInst)) {
     // TODO: refactor function type creation which usually occurs in std-LLVM
     // conversion.
     SmallVector<Type, 8> operandTypes;
     llvm::append_range(operandTypes, inlineAsmOp.getOperands().getTypes());
 
if (auto attr = opInst.getAttrOfType<FlatSymbolRefAttr>("callee")) {
       auto *calleeFunctionType = cast<llvm::FunctionType>(
           moduleTranslation.convertType(calleeType.getElementType()));
       result = builder.CreateInvoke(
           calleeFunctionType, operandsRef.front(),
           moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
           moduleTranslation.lookupBlock(invOp.getSuccessor(1)),
           operandsRef.drop_front());
     }
+    llvm::MDNode *branchWeights =
+        convertBranchWeights(invOp.getBranchWeights(), moduleTranslation);
+    if (branchWeights)
+      result->setMetadata(llvm::LLVMContext::MD_prof, branchWeights);
     moduleTranslation.mapBranch(invOp, result);
     // InvokeOp can only have 0 or 1 result
     if (invOp->getNumResults() != 0) {
       moduleTranslation.mapValue(opInst.getResult(0), result);
       return success();
     }
     return success(result->getType()->isVoidTy());
   }
#include "mlir/Dialect/LLVMIR/LLVMIntrinsicConversions.inc"
   if (auto brOp = dyn_cast<LLVM::BrOp>(opInst)) {
     llvm::BranchInst *branch =
         builder.CreateBr(moduleTranslation.lookupBlock(brOp.getSuccessor()));
     moduleTranslation.mapBranch(&opInst, branch);
     setLoopMetadata(opInst, *branch, builder, moduleTranslation);
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
-    llvm::MDNode *branchWeights = nullptr;
-    if (auto weights = condbrOp.getBranchWeights()) {
-      // Map weight attributes to LLVM metadata.
-      auto weightValues = weights->getValues<APInt>();
-      auto trueWeight = weightValues[0].getSExtValue();
-      auto falseWeight = weightValues[1].getSExtValue();
-      branchWeights =
-          llvm::MDBuilder(moduleTranslation.getLLVMContext())
-              .createBranchWeights(static_cast<uint32_t>(trueWeight),
-                                   static_cast<uint32_t>(falseWeight));
-    }
+    llvm::MDNode *branchWeights =
+        convertBranchWeights(condbrOp.getBranchWeights(), moduleTranslation);
     llvm::BranchInst *branch = builder.CreateCondBr(
         moduleTranslation.lookupValue(condbrOp.getOperand(0)),
         moduleTranslation.lookupBlock(condbrOp.getSuccessor(0)),
         moduleTranslation.lookupBlock(condbrOp.getSuccessor(1)), branchWeights);
     moduleTranslation.mapBranch(&opInst, branch);
     setLoopMetadata(opInst, *branch, builder, moduleTranslation);
     return success();
   }
   if (auto switchOp = dyn_cast<LLVM::SwitchOp>(opInst)) {
-    llvm::MDNode *branchWeights = nullptr;
-    if (auto weights = switchOp.getBranchWeights()) {
-      llvm::SmallVector<uint32_t> weightValues;
-      weightValues.reserve(weights->size());
-      for (llvm::APInt weight : weights->cast<DenseIntElementsAttr>())
-        weightValues.push_back(weight.getLimitedValue());
-      branchWeights = llvm::MDBuilder(moduleTranslation.getLLVMContext())
-                          .createBranchWeights(weightValues);
-    }
-
+    llvm::MDNode *branchWeights =
+        convertBranchWeights(switchOp.getBranchWeights(), moduleTranslation);
     llvm::SwitchInst *switchInst = builder.CreateSwitch(
         moduleTranslation.lookupValue(switchOp.getValue()),
         moduleTranslation.lookupBlock(switchOp.getDefaultDestination()),
         switchOp.getCaseDestinations().size(), branchWeights);
 
     auto *ty = llvm::cast<llvm::IntegerType>(
         moduleTranslation.convertType(switchOp.getValue().getType()));
     for (auto i :

mlir/test/Target/LLVMIR/Import/profiling-metadata.ll

bb1:
   ret i32 %arg1
 bb2:
   ret i32 %arg1
 bbd:
   ret i32 %arg1
 }
 
 !0 = !{!"branch_weights", i32 42, i32 3, i32 5}
+
+; // -----
+
+; CHECK: llvm.func @fn()
+declare void @fn()
+
+; CHECK-LABEL: @call_branch_weights
+define void @call_branch_weights() {
+  ; CHECK:  llvm.call @fn() {branch_weights = dense<42> : vector<1xi32>}
+  call void @fn(), !prof !0
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 42}
+
+; // -----
+
+declare void @foo()
+declare i32 @__gxx_personality_v0(...)
+
+; CHECK-LABEL: @invoke_branch_weights
+define i32 @invoke_branch_weights() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+  ; CHECK: llvm.invoke @foo() to ^bb2 unwind ^bb1 {branch_weights = dense<[42, 99]> : vector<2xi32>} : () -> ()
+  invoke void @foo() to label %bb2 unwind label %bb1, !prof !0
+bb1:
+  %1 = landingpad { i8*, i32 } cleanup
+  br label %bb2
+bb2:
+  ret i32 1
+
+}
+
+!0 = !{!"branch_weights", i32 42, i32 99}

mlir/test/Target/LLVMIR/llvmir.mlir

 ^bb2:  // pred: ^bb0
   llvm.return %arg1 : i32
 }
 
 // CHECK: ![[NODE]] = !{!"branch_weights", i32 5, i32 10}
 
 // -----
 
+llvm.func @fn()
+
+// CHECK-LABEL: @call_branch_weights
+llvm.func @call_branch_weights() {
+  // CHECK: !prof ![[NODE:[0-9]+]]
+  llvm.call @fn() {branch_weights = dense<42> : vector<1xi32>} : () -> ()
+  llvm.return
+}
+
+// CHECK: ![[NODE]] = !{!"branch_weights", i32 42}
+
+// -----
+
+llvm.func @foo()
+llvm.func @__gxx_personality_v0(...) -> i32
+
+// CHECK-LABEL: @invoke_branch_weights
+llvm.func @invoke_branch_weights() -> i32 attributes {personality = @__gxx_personality_v0} {
+  %0 = llvm.mlir.constant(1 : i32) : i32
+  // CHECK: !prof ![[NODE:[0-9]+]]
+  llvm.invoke @foo() to ^bb2 unwind ^bb1 {branch_weights = dense<[42, 99]> : vector<2xi32>} : () -> ()
+^bb1:  // pred: ^bb0
+  %1 = llvm.landingpad cleanup : !llvm.struct<(ptr<i8>, i32)>
+  llvm.br ^bb2
+^bb2:  // 2 preds: ^bb0, ^bb1
+  llvm.return %0 : i32
+}
+
+// CHECK: ![[NODE]] = !{!"branch_weights", i32 42, i32 99}
+
+// -----
+
 llvm.func @volatile_store_and_load() {
   %val = llvm.mlir.constant(5 : i32) : i32
   %size = llvm.mlir.constant(1 : i64) : i64
   %0 = llvm.alloca %size x i32 : (i64) -> (!llvm.ptr<i32>)
   // CHECK: store volatile i32 5, ptr %{{.*}}
   llvm.store volatile %val, %0 : !llvm.ptr<i32>
   // CHECK: %{{.*}} = load volatile i32, ptr %{{.*}}
   %1 = llvm.load volatile %0: !llvm.ptr<i32>