[mlir][llvm] Add branch weight op interface

[gysit]

Summary

This revision adds a branch weight op interface for the call / branch
operations that support branch weights. It can be used in the LLVM IR
import and export to simplify the branch weight conversion. An
additional mapping between call operations and instructions ensures
the actual conversion can be done in the module translation itself,
rather than in the dialect translation interface. It also has the
benefit that downstream users can amend custom metadata to the call
operation during the export to LLVM IR.

[zero9178]

LGTM % nits

mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td

69

I know this hasn't been done consistently in this file but you could use auto here and below to avoid repeating the type twice.

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

122

[gysit]

Address comments.

[definelicht]

LGTM!

mlir/include/mlir/Dialect/LLVMIR/LLVMInterfaces.td

def FastmathFlagsInterface : OpInterface<"FastmathFlagsInterface"> {
   let methods = [
     InterfaceMethod<
       /*desc=*/        "Returns a FastmathFlagsAttr attribute for the operation",
       /*returnType=*/  "FastmathFlagsAttr",
       /*methodName=*/  "getFastmathAttr",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         return op.getFastmathFlagsAttr();
       }]
       >,
     StaticInterfaceMethod<
       /*desc=*/        [{Returns the name of the FastmathFlagsAttr attribute
                          for the operation}],
       /*returnType=*/  "StringRef",
       /*methodName=*/  "getFastmathAttrName",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
         return "fastmathFlags";
       }]
       >
   ];
 }
 
+def BranchWeightOpInterface : OpInterface<"BranchWeightOpInterface"> {
+  let description = [{
+    An interface for operations that can carry branch weights metadata. It
+    provides setters and getters for the operation's branch weights attribute.
+    The default implementation of the interface methods expect the operation to
+    have an attribute of type DenseI32ArrayAttr named branch_weights.
+  }];
+
+  let cppNamespace = "::mlir::LLVM";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/        "Returns the branch weights attribute or nullptr",
+      /*returnType=*/  "DenseI32ArrayAttr",
+      /*methodName=*/  "getBranchWeightsOrNull",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        auto op = cast<ConcreteOp>(this->getOperation());
+        return op.getBranchWeightsAttr();
+      }]
+      >,
+    InterfaceMethod<
+      /*desc=*/        "Sets the branch weights attribute",
+      /*returnType=*/  "void",
+      /*methodName=*/  "setBranchWeights",
+      /*args=*/        (ins "DenseI32ArrayAttr":$attr),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        auto op = cast<ConcreteOp>(this->getOperation());
+        op.setBranchWeightsAttr(attr);
+      }]
+      >
+  ];
+}
+
 def AccessGroupOpInterface : OpInterface<"AccessGroupOpInterface"> {
   let description = [{
     An interface for memory operations that can carry access groups metadata.
     It provides setters and getters for the operation's access groups attribute.
     The default implementations of the interface methods expect the operation
     to have an attribute of type ArrayAttr named access_groups.
   }];
 
   let cppNamespace = "::mlir::LLVM";
   let verify = [{ return detail::verifyAccessGroupOpInterface($_op); }];
 
   let methods = [
     InterfaceMethod<
       /*desc=*/        "Returns the access groups attribute or nullptr",
       /*returnType=*/  "ArrayAttr",
       /*methodName=*/  "getAccessGroupsOrNull",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         return op.getAccessGroupsAttr();
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Sets the access groups attribute",
       /*returnType=*/  "void",
       /*methodName=*/  "setAccessGroups",
       /*args=*/        (ins "const ArrayAttr":$attr),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         op.setAccessGroupsAttr(attr);
       }]
       >
   ];
 }
 
 def AliasAnalysisOpInterface : OpInterface<"AliasAnalysisOpInterface"> {
   let description = [{
def AliasAnalysisOpInterface : OpInterface<"AliasAnalysisOpInterface"> {
   let methods = [
     InterfaceMethod<
       /*desc=*/        "Returns the alias scopes attribute or nullptr",
       /*returnType=*/  "ArrayAttr",
       /*methodName=*/  "getAliasScopesOrNull",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         return op.getAliasScopesAttr();
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Sets the alias scopes attribute",
       /*returnType=*/  "void",
       /*methodName=*/  "setAliasScopes",
       /*args=*/        (ins "const ArrayAttr":$attr),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         op.setAliasScopesAttr(attr);
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Returns the noalias scopes attribute or nullptr",
       /*returnType=*/  "ArrayAttr",
       /*methodName=*/  "getNoAliasScopesOrNull",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         return op.getNoaliasScopesAttr();
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Sets the noalias scopes attribute",
       /*returnType=*/  "void",
       /*methodName=*/  "setNoAliasScopes",
       /*args=*/        (ins "const ArrayAttr":$attr),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         op.setNoaliasScopesAttr(attr);
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Returns the tbaa attribute or nullptr",
       /*returnType=*/  "ArrayAttr",
       /*methodName=*/  "getTBAATagsOrNull",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         return op.getTbaaAttr();
       }]
       >,
     InterfaceMethod<
       /*desc=*/        "Sets the tbaa attribute",
       /*returnType=*/  "void",
       /*methodName=*/  "setTBAATags",
       /*args=*/        (ins "const ArrayAttr":$attr),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
-        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        auto op = cast<ConcreteOp>(this->getOperation());
         op.setTbaaAttr(attr);
       }]
       >
   ];
 }
 
 def GetResultPtrElementType : OpInterface<"GetResultPtrElementType"> {
   let description = [{

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

 def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "FPTrunc",
                                  LLVM_ScalarOrVectorOf<LLVM_AnyFloat>,
                                  LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
 
 // Call-related operations.
 def LLVM_InvokeOp : LLVM_Op<"invoke", [
                       AttrSizedOperandSegments,
                       DeclareOpInterfaceMethods<BranchOpInterface>,
-                      DeclareOpInterfaceMethods<CallOpInterface>, Terminator]> {
+                      DeclareOpInterfaceMethods<CallOpInterface>,
+                      DeclareOpInterfaceMethods<BranchWeightOpInterface>,
+                      Terminator]> {
   let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                    Variadic<LLVM_Type>:$callee_operands,
                    Variadic<LLVM_Type>:$normalDestOperands,
                    Variadic<LLVM_Type>:$unwindDestOperands,
-                   OptionalAttr<ElementsAttr>:$branch_weights);
+                   OptionalAttr<DenseI32ArrayAttr>:$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),

 //===----------------------------------------------------------------------===//
 
 // FIXME: Add a type attribute that carries the LLVM function type to support
 // indirect calls to variadic functions. The type attribute is necessary to
 // distinguish normal and variadic arguments.
 def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
                     [DeclareOpInterfaceMethods<FastmathFlagsInterface>,
                      DeclareOpInterfaceMethods<CallOpInterface>,
-                     DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+                     DeclareOpInterfaceMethods<SymbolUserOpInterface>,
+                     DeclareOpInterfaceMethods<BranchWeightOpInterface>]> {
   let summary = "Call to an LLVM function.";
   let description = [{
     In LLVM IR, functions may return either 0 or 1 value. LLVM IR dialect
     implements this behavior by providing a variadic `call` operation for 0- and
     1-result functions. Even though MLIR supports multi-result functions, LLVM
     IR dialect disallows them.
 
     The `call` instruction supports both direct and indirect calls. Direct calls
let description = [{
     llvm.call %1(%0) : !llvm.ptr, (f32) -> ()
     ```
   }];
 
   dag args = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                   Variadic<LLVM_Type>,
                   DefaultValuedAttr<LLVM_FastmathFlagsAttr,
                                    "{}">:$fastmathFlags,
-                  OptionalAttr<ElementsAttr>:$branch_weights);
+                  OptionalAttr<DenseI32ArrayAttr>:$branch_weights);
   // Append the aliasing related attributes defined in LLVM_MemAccessOpBase.
   let arguments = !con(args, aliasAttrs);
   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,
let builders = [
     }]>,
     OpBuilder<(ins "ValueRange":$operands, "Block *":$dest), [{
       build($_builder, $_state, operands, /*loop_annotation=*/{}, dest);
     }]>,
     LLVM_TerminatorPassthroughOpBuilder
   ];
 }
 def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br",
-    [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<BranchOpInterface>,
+     DeclareOpInterfaceMethods<BranchWeightOpInterface>,
      Pure]> {
   let arguments = (ins I1:$condition,
                    Variadic<LLVM_Type>:$trueDestOperands,
                    Variadic<LLVM_Type>:$falseDestOperands,
-                   OptionalAttr<ElementsAttr>:$branch_weights,
+                   OptionalAttr<DenseI32ArrayAttr>:$branch_weights,
                    OptionalAttr<LoopAnnotationAttr>:$loop_annotation);
   let successors = (successor AnySuccessor:$trueDest, AnySuccessor:$falseDest);
   let assemblyFormat = [{
     $condition ( `weights` `(` $branch_weights^ `)` )? `,`
     $trueDest (`(` $trueDestOperands^ `:` type($trueDestOperands) `)`)? `,`
     $falseDest (`(` $falseDestOperands^ `:` type($falseDestOperands) `)`)?
     attr-dict
   }];
 
   let builders = [
     OpBuilder<(ins "Value":$condition, "Block *":$trueDest,
       "ValueRange":$trueOperands, "Block *":$falseDest,
       "ValueRange":$falseOperands,
       CArg<"std::optional<std::pair<uint32_t, uint32_t>>", "{}">:$weights)>,
   OpBuilder<(ins "Value":$condition, "Block *":$trueDest,
     "Block *":$falseDest, CArg<"ValueRange", "{}">:$falseOperands),
   [{
       build($_builder, $_state, condition, trueDest, ValueRange(), falseDest,
             falseOperands);
   }]>,
   OpBuilder<(ins "Value":$condition, "ValueRange":$trueOperands, "ValueRange":$falseOperands,
-    "ElementsAttr":$branchWeights, "Block *":$trueDest, "Block *":$falseDest),
+    "DenseI32ArrayAttr":$branchWeights, "Block *":$trueDest, "Block *":$falseDest),
   [{
       build($_builder, $_state, condition, trueOperands, falseOperands, branchWeights,
       {}, trueDest, falseDest);
   }]>, LLVM_TerminatorPassthroughOpBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
def LLVM_UnreachableOp : LLVM_TerminatorOp<"unreachable"> {
   string llvmInstName = "Unreachable";
   string llvmBuilder = [{ builder.CreateUnreachable(); }];
   string mlirBuilder = [{
     $_op = $_builder.create<LLVM::UnreachableOp>($_location);
   }];
 }
 
 def LLVM_SwitchOp : LLVM_TerminatorOp<"switch",
-    [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<BranchOpInterface>,
+     DeclareOpInterfaceMethods<BranchWeightOpInterface>,
      Pure]> {
   let arguments = (ins
     AnyInteger:$value,
     Variadic<AnyType>:$defaultOperands,
     VariadicOfVariadic<AnyType, "case_operand_segments">:$caseOperands,
     OptionalAttr<AnyIntElementsAttr>:$case_values,
     DenseI32ArrayAttr:$case_operand_segments,
-    OptionalAttr<ElementsAttr>:$branch_weights
+    OptionalAttr<DenseI32ArrayAttr>:$branch_weights
   );
   let successors = (successor
     AnySuccessor:$defaultDestination,
     VariadicSuccessor<AnySuccessor>:$caseDestinations
   );
 
   let assemblyFormat = [{
     $value `:` type($value) `,`

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

public:
   }
 
   /// Finds an LLVM IR instruction that corresponds to the given MLIR operation
   /// with successors.
   llvm::Instruction *lookupBranch(Operation *op) const {
     return branchMapping.lookup(op);
   }
 
+  /// Stores a mapping between an MLIR call operation and a corresponding LLVM
+  /// call instruction.
+  void mapCall(Operation *mlir, llvm::CallInst *llvm) {
+    auto result = callMapping.try_emplace(mlir, llvm);
+    (void)result;
+    assert(result.second && "attempting to map a call that is already mapped");
+  }
+
+  /// Finds an LLVM call instruction that corresponds to the given MLIR call
+  /// operation.
+  llvm::CallInst *lookupCall(Operation *op) const {
+    return callMapping.lookup(op);
+  }
+
   /// Removes the mapping for blocks contained in the region and values defined
   /// in these blocks.
   void forgetMapping(Region &region);
 
   /// Returns the LLVM metadata corresponding to a mlir LLVM dialect alias scope
   /// attribute.
   llvm::MDNode *getAliasScope(AliasScopeAttr aliasScopeAttr) const;
 
   /// Returns the LLVM metadata corresponding to an array of mlir LLVM dialect
   /// alias scope attributes.
   llvm::MDNode *getAliasScopes(ArrayRef<AliasScopeAttr> aliasScopeAttrs) const;
 
   // Sets LLVM metadata for memory operations that are in a parallel loop.
   void setAccessGroupsMetadata(AccessGroupOpInterface op,
                                llvm::Instruction *inst);
 
   // Sets LLVM metadata for memory operations that have alias scope information.
   void setAliasScopeMetadata(AliasAnalysisOpInterface op,
                              llvm::Instruction *inst);
 
   /// Sets LLVM TBAA metadata for memory operations that have TBAA attributes.
   void setTBAAMetadata(AliasAnalysisOpInterface op, llvm::Instruction *inst);
 
+  /// Sets LLVM profiling metadata for operations that have branch weights.
+  void setBranchWeightsMetadata(BranchWeightOpInterface op);
+
   /// Sets LLVM loop metadata for branch operations that have a loop annotation
   /// attribute.
   void setLoopMetadata(Operation *op, llvm::Instruction *inst);
 
   /// Converts the type from MLIR LLVM dialect to LLVM.
   llvm::Type *convertType(Type type);
 
   /// Returns the MLIR context of the module being translated.
private:
   DenseMap<Value, llvm::Value *> valueMapping;
   DenseMap<Block *, llvm::BasicBlock *> blockMapping;
 
   /// A mapping between MLIR LLVM dialect terminators and LLVM IR terminators
   /// they are converted to. This allows for connecting PHI nodes to the source
   /// values after all operations are converted.
   DenseMap<Operation *, llvm::Instruction *> branchMapping;
 
+  /// A mapping between MLIR LLVM dialect call operations and LLVM IR call
+  /// instructions. This allows for adding branch weights after the operations
+  /// have been converted.
+  DenseMap<Operation *, llvm::CallInst *> callMapping;
+
   /// Mapping from an alias scope metadata operation to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<Attribute, llvm::MDNode *> aliasScopeMetadataMapping;
 
   /// Mapping from a tbaa attribute to its LLVM metadata.
   /// This map is populated on module entry.
   DenseMap<Attribute, llvm::MDNode *> tbaaMetadataMapping;
 

mlir/lib/Conversion/SPIRVToLLVM/SPIRVToLLVM.cpp

 public:
   using SPIRVToLLVMConversion<
       spirv::BranchConditionalOp>::SPIRVToLLVMConversion;
 
   LogicalResult
   matchAndRewrite(spirv::BranchConditionalOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // If branch weights exist, map them to 32-bit integer vector.
-    ElementsAttr branchWeights = nullptr;
+    DenseI32ArrayAttr branchWeights = nullptr;
     if (auto weights = op.getBranchWeights()) {
-      VectorType weightType = VectorType::get(2, rewriter.getI32Type());
-      branchWeights = DenseElementsAttr::get(weightType, weights->getValue());
+      SmallVector<int32_t> weightValues;
+      for (auto weight : weights->getAsRange<IntegerAttr>())
+        weightValues.push_back(weight.getInt());
+      branchWeights = DenseI32ArrayAttr::get(getContext(), weightValues);
     }
 
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
         op, op.getCondition(), op.getTrueBlockArguments(),
         op.getFalseBlockArguments(), branchWeights, op.getTrueBlock(),
         op.getFalseBlock());
     return success();
   }

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

SuccessorOperands CondBrOp::getSuccessorOperands(unsigned index) {
   return SuccessorOperands(index == 0 ? getTrueDestOperandsMutable()
                                       : getFalseDestOperandsMutable());
 }
 
 void CondBrOp::build(OpBuilder &builder, OperationState &result,
                      Value condition, Block *trueDest, ValueRange trueOperands,
                      Block *falseDest, ValueRange falseOperands,
                      std::optional<std::pair<uint32_t, uint32_t>> weights) {
-  ElementsAttr weightsAttr;
+  DenseI32ArrayAttr weightsAttr;
   if (weights)
     weightsAttr =
-        builder.getI32VectorAttr({static_cast<int32_t>(weights->first),
+        builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights->first),
                                       static_cast<int32_t>(weights->second)});
 
   build(builder, result, condition, trueOperands, falseOperands, weightsAttr,
         /*loop_annotation=*/{}, trueDest, falseDest);
 }
 
 //===----------------------------------------------------------------------===//
 // LLVM::SwitchOp
 //===----------------------------------------------------------------------===//
 
 void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
                      Block *defaultDestination, ValueRange defaultOperands,
                      DenseIntElementsAttr caseValues,
                      BlockRange caseDestinations,
                      ArrayRef<ValueRange> caseOperands,
                      ArrayRef<int32_t> branchWeights) {
-  ElementsAttr weightsAttr;
+  DenseI32ArrayAttr weightsAttr;
   if (!branchWeights.empty())
-    weightsAttr = builder.getI32VectorAttr(llvm::to_vector<4>(branchWeights));
+    weightsAttr = builder.getDenseI32ArrayAttr(branchWeights);
 
   build(builder, result, value, defaultOperands, caseOperands, caseValues,
         weightsAttr, defaultDestination, caseDestinations);
 }
 
 void SwitchOp::build(OpBuilder &builder, OperationState &result, Value value,
                      Block *defaultDestination, ValueRange defaultOperands,
                      ArrayRef<APInt> caseValues, BlockRange caseDestinations,

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

static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
   for (unsigned i = 1, e = node->getNumOperands(); i != e; ++i) {
     llvm::ConstantInt *branchWeight =
         llvm::mdconst::dyn_extract<llvm::ConstantInt>(node->getOperand(i));
     if (!branchWeight)
       return failure();
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 
-  return TypeSwitch<Operation *, LogicalResult>(op)
-      .Case<CondBrOp, SwitchOp, CallOp, InvokeOp>([&](auto branchWeightOp) {
-        branchWeightOp.setBranchWeightsAttr(
-            builder.getI32VectorAttr(branchWeights));
-        return success();
-      })
-      .Default([](auto) { return failure(); });
+  if (auto iface = dyn_cast<BranchWeightOpInterface>(op)) {
+    iface.setBranchWeights(builder.getDenseI32ArrayAttr(branchWeights));
+  }
+  return failure();
 }
 
 /// Searches for the attribute that maps to the given TBAA metadata `node` and
 /// attaches it to the imported operation if the lookup succeeds. Returns
 /// failure otherwise.
 static LogicalResult setTBAAAttr(const llvm::MDNode *node, Operation *op,
                                  LLVM::ModuleImport &moduleImport) {
   Attribute tbaaTagSym = moduleImport.lookupTBAAAttr(node);

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 : llvm::cast<DenseIntElementsAttr>(*weights))
-    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));
 
if (auto callOp = dyn_cast<LLVM::CallOp>(opInst)) {
     if (auto attr = callOp.getCalleeAttr()) {
       call = builder.CreateCall(
           moduleTranslation.lookupFunction(attr.getValue()), operandsRef);
     } else {
       call = builder.CreateCall(getCalleeFunctionType(callOp.getResultTypes(),
                                                       callOp.getArgOperands()),
                                 operandsRef.front(), operandsRef.drop_front());
     }
-    llvm::MDNode *branchWeights =
-        convertBranchWeights(callOp.getBranchWeights(), moduleTranslation);
-    if (branchWeights)
-      call->setMetadata(llvm::LLVMContext::MD_prof, branchWeights);
     moduleTranslation.setAccessGroupsMetadata(callOp, call);
     moduleTranslation.setAliasScopeMetadata(callOp, call);
     moduleTranslation.setTBAAMetadata(callOp, call);
     // 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), call);
       return success();
     }
     // Check that LLVM call returns void for 0-result functions.
-    return success(call->getType()->isVoidTy());
+    if (!call->getType()->isVoidTy())
+      return failure();
+    moduleTranslation.mapCall(callOp, call);
+    return success();
   }
 
   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 invOp = dyn_cast<LLVM::InvokeOp>(opInst)) {
     } else {
       result = builder.CreateInvoke(
           getCalleeFunctionType(invOp.getResultTypes(), invOp.getArgOperands()),
           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);
     moduleTranslation.setLoopMetadata(&opInst, branch);
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
-    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.lookupBlock(condbrOp.getSuccessor(1)));
     moduleTranslation.mapBranch(&opInst, branch);
     moduleTranslation.setLoopMetadata(&opInst, branch);
     return success();
   }
   if (auto switchOp = dyn_cast<LLVM::SwitchOp>(opInst)) {
-    llvm::MDNode *branchWeights =
-        convertBranchWeights(switchOp.getBranchWeights(), moduleTranslation);
     llvm::SwitchInst *switchInst = builder.CreateSwitch(
         moduleTranslation.lookupValue(switchOp.getValue()),
         moduleTranslation.lookupBlock(switchOp.getDefaultDestination()),
-        switchOp.getCaseDestinations().size(), branchWeights);
+        switchOp.getCaseDestinations().size());
 
     auto *ty = llvm::cast<llvm::IntegerType>(
         moduleTranslation.convertType(switchOp.getValue().getType()));
     for (auto i :
          llvm::zip(llvm::cast<DenseIntElementsAttr>(*switchOp.getCaseValues()),
                    switchOp.getCaseDestinations()))
       switchInst->addCase(
           llvm::ConstantInt::get(ty, std::get<0>(i).getLimitedValue()),

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments,
   // Traverse operations.
   for (auto &op : bb) {
     // Set the current debug location within the builder.
     builder.SetCurrentDebugLocation(
         debugTranslation->translateLoc(op.getLoc(), subprogram));
 
     if (failed(convertOperation(op, builder)))
       return failure();
+
+    // Set the branch weight metadata on the translated instruction.
+    if (auto iface = dyn_cast<BranchWeightOpInterface>(op))
+      setBranchWeightsMetadata(iface);
   }
 
   return success();
 }
 
 /// A helper method to get the single Block in an operation honoring LLVM's
 /// module requirements.
 static Block &getModuleBody(Operation *module) {
op.emitWarning() << "TBAA access tags were not translated, because LLVM "
                         "IR only supports a single tag per instruction";
     return;
   }
 
   llvm::MDNode *node = getTBAANode(cast<TBAATagAttr>(tagRefs[0]));
   inst->setMetadata(llvm::LLVMContext::MD_tbaa, node);
 }
 
+void ModuleTranslation::setBranchWeightsMetadata(BranchWeightOpInterface op) {
+  DenseI32ArrayAttr weightsAttr = op.getBranchWeightsOrNull();
+  if (!weightsAttr)
+    return;
+
+  llvm::Instruction *inst = isa<CallOp>(op) ? lookupCall(op) : lookupBranch(op);
+  assert(inst && "expected the operation to have a mapping to an instruction");
+  SmallVector<uint32_t> weights(weightsAttr.asArrayRef());
+  inst->setMetadata(
+      llvm::LLVMContext::MD_prof,
+      llvm::MDBuilder(getLLVMContext()).createBranchWeights(weights));
+}
+
 LogicalResult ModuleTranslation::createTBAAMetadata() {
   llvm::LLVMContext &ctx = llvmModule->getContext();
   llvm::IntegerType *offsetTy = llvm::IntegerType::get(ctx, 64);
 
   // Walk the entire module and create all metadata nodes for the TBAA
   // attributes. The code below relies on two invariants of the
   // `AttrTypeWalker`:
   // 1. Attributes are visited in post-order: Since the attributes create a DAG,

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

spirv.module Logical GLSL450 {
   ^inner_true:
     spirv.Return
   // CHECK: ^bb4(%{{.*}}: i32, %{{.*}}: i32):
   ^inner_false(%arg3: i32, %arg4: i32):
     spirv.Return
   }
 
   spirv.func @cond_branch_with_weights(%cond: i1) -> () "None" {
-    // CHECK: llvm.cond_br %{{.*}} weights(dense<[1, 2]> : vector<2xi32>), ^bb1, ^bb2
+    // CHECK: llvm.cond_br %{{.*}} weights([1, 2]), ^bb1, ^bb2
     spirv.BranchConditional %cond [1, 2], ^true, ^false
   // CHECK: ^bb1:
   ^true:
     spirv.Return
   // CHECK: ^bb2:
   ^false:
     spirv.Return
   }

mlir/test/Dialect/LLVMIR/invalid.mlir

 }
 
 // -----
 
 func.func @switch_wrong_number_of_weights(%arg0 : i32) {
   // expected-error@+1 {{expects number of branch weights to match number of successors: 3 vs 2}}
   llvm.switch %arg0 : i32, ^bb1 [
     42: ^bb2(%arg0, %arg0 : i32, i32)
-  ] {branch_weights = dense<[13, 17, 19]> : vector<3xi32>}
+  ] {branch_weights = array<i32: 13, 17, 19>}
 
 ^bb1: // pred: ^bb0
   llvm.return
 
 ^bb2(%1: i32, %2: i32): // pred: ^bb0
   llvm.return
 }
 

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

 ; RUN: mlir-translate -import-llvm -split-input-file %s | FileCheck %s
 
 ; CHECK-LABEL: @cond_br
 define i64 @cond_br(i1 %arg1, i64 %arg2) {
 entry:
   ; CHECK: llvm.cond_br
-  ; CHECK-SAME: weights(dense<[0, 3]> : vector<2xi32>)
+  ; CHECK-SAME: weights([0, 3])
   br i1 %arg1, label %bb1, label %bb2, !prof !0
 bb1:
   ret i64 %arg2
 bb2:
   ret i64 %arg2
 }
 
 !0 = !{!"branch_weights", i32 0, i32 3}
 
 ; // -----
 
 ; CHECK-LABEL: @simple_switch(
 define i32 @simple_switch(i32 %arg1) {
   ; CHECK: llvm.switch
-  ; CHECK: {branch_weights = dense<[42, 3, 5]> : vector<3xi32>}
+  ; CHECK: {branch_weights = array<i32: 42, 3, 5>}
   switch i32 %arg1, label %bbd [
     i32 0, label %bb1
     i32 9, label %bb2
   ], !prof !0
 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>}
+  ; CHECK:  llvm.call @fn() {branch_weights = array<i32: 42>}
   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 ptr @__gxx_personality_v0 {
-  ; CHECK: llvm.invoke @foo() to ^bb2 unwind ^bb1 {branch_weights = dense<[42, 99]> : vector<2xi32>} : () -> ()
+  ; CHECK: llvm.invoke @foo() to ^bb2 unwind ^bb1 {branch_weights = array<i32: 42, 99>} : () -> ()
   invoke void @foo() to label %bb2 unwind label %bb1, !prof !0
 bb1:
   %1 = landingpad { ptr, i32 } cleanup
   br label %bb2
 bb2:
   ret i32 1
 
 }
 
 !0 = !{!"branch_weights", i32 42, i32 99}

mlir/test/Target/LLVMIR/llvmir.mlir

llvm.func @foo() {
   llvm.return
 }
 
 // -----
 
 // Check that branch weight attributes are exported properly as metadata.
 llvm.func @cond_br_weights(%cond : i1, %arg0 : i32,  %arg1 : i32) -> i32 {
   // CHECK: !prof ![[NODE:[0-9]+]]
-  llvm.cond_br %cond weights(dense<[5, 10]> : vector<2xi32>), ^bb1, ^bb2
+  llvm.cond_br %cond weights([5, 10]), ^bb1, ^bb2
 ^bb1:  // pred: ^bb0
   llvm.return %arg0 : i32
 ^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.call @fn() {branch_weights = array<i32 : 42>} : () -> ()
   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>} : () -> ()
+  llvm.invoke @foo() to ^bb2 unwind ^bb1 {branch_weights = array<i32 : 42, 99>} : () -> ()
 ^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 @switch_weights(%arg0: i32) -> i32 {
   %0 = llvm.mlir.constant(19 : i32) : i32
   %1 = llvm.mlir.constant(23 : i32) : i32
   %2 = llvm.mlir.constant(29 : i32) : i32
   // CHECK: !prof ![[SWITCH_WEIGHT_NODE:[0-9]+]]
   llvm.switch %arg0 : i32, ^bb1(%0 : i32) [
     9: ^bb2(%1, %2 : i32, i32),
     99: ^bb3
-  ] {branch_weights = dense<[13, 17, 19]> : vector<3xi32>}
+  ] {branch_weights = array<i32 : 13, 17, 19>}
 
 ^bb1(%3: i32):  // pred: ^bb0
   llvm.return %3 : i32
 
 ^bb2(%4: i32, %5: i32): // pred: ^bb0
   llvm.return %5 : i32
 
 ^bb3: // pred: ^bb0