Replaced LLVM dialect's BranchWeightOpInterface with WeightedBranchOpInterface.

Author: vzakhari

flang/test/Fir/invalid.fir

@@ -1397,7 +1397,7 @@ fir.local {type = local_init} @x.localizer : f32 init {
 // -----
 
 func.func @wrong_weights_number_in_if_then(%cond: i1) {
-// expected-error @below {{number of weights (1) does not match the number of regions (2)}}
+// expected-error @below {{expects number of region weights to match number of regions: 1 vs 2}}
   fir.if %cond weights([50]) {
   }
   return
@@ -1406,7 +1406,7 @@ func.func @wrong_weights_number_in_if_then(%cond: i1) {
 // -----
 
 func.func @wrong_weights_number_in_if_then_else(%cond: i1) {
-// expected-error @below {{number of weights (3) does not match the number of regions (2)}}
+// expected-error @below {{expects number of region weights to match number of regions: 3 vs 2}}
   fir.if %cond weights([50, 40, 10]) {
   } else {
   }
@@ -1421,12 +1421,3 @@ func.func @negative_weight_in_if_then(%cond: i1) {
   }
   return
 }
-
-// -----
-
-func.func @wrong_total_weight_in_if_then(%cond: i1) {
-// expected-error @below {{total weight 101 is not 100}}
-  fir.if %cond weights([1, 100]) {
-  }
-  return
-}

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

@@ -168,42 +168,6 @@ def NonNegFlagInterface : OpInterface<"NonNegFlagInterface"> {
   ];
 }
 
-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=*/  "::mlir::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 "::mlir::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.

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

@@ -660,12 +660,12 @@ def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "FPTrunc",
                                  LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
 
 // Call-related operations.
-def LLVM_InvokeOp : LLVM_Op<"invoke", [
-                      AttrSizedOperandSegments,
-                      DeclareOpInterfaceMethods<BranchOpInterface>,
-                      DeclareOpInterfaceMethods<CallOpInterface>,
-                      DeclareOpInterfaceMethods<BranchWeightOpInterface>,
-                      Terminator]> {
+def LLVM_InvokeOp
+    : LLVM_Op<"invoke", [AttrSizedOperandSegments,
+                         DeclareOpInterfaceMethods<BranchOpInterface>,
+                         DeclareOpInterfaceMethods<CallOpInterface>,
+                         DeclareOpInterfaceMethods<WeightedBranchOpInterface>,
+                         Terminator]> {
   let arguments = (ins
                    OptionalAttr<TypeAttrOf<LLVM_FunctionType>>:$var_callee_type,
                    OptionalAttr<FlatSymbolRefAttr>:$callee,
@@ -734,12 +734,13 @@ def LLVM_VaArgOp : LLVM_Op<"va_arg"> {
 // CallOp
 //===----------------------------------------------------------------------===//
 
-def LLVM_CallOp : LLVM_MemAccessOpBase<"call",
-                    [AttrSizedOperandSegments,
-                     DeclareOpInterfaceMethods<FastmathFlagsInterface>,
-                     DeclareOpInterfaceMethods<CallOpInterface>,
-                     DeclareOpInterfaceMethods<SymbolUserOpInterface>,
-                     DeclareOpInterfaceMethods<BranchWeightOpInterface>]> {
+def LLVM_CallOp
+    : LLVM_MemAccessOpBase<
+          "call", [AttrSizedOperandSegments,
+                   DeclareOpInterfaceMethods<FastmathFlagsInterface>,
+                   DeclareOpInterfaceMethods<CallOpInterface>,
+                   DeclareOpInterfaceMethods<SymbolUserOpInterface>,
+                   DeclareOpInterfaceMethods<WeightedBranchOpInterface>]> {
   let summary = "Call to an LLVM function.";
   let description = [{
     In LLVM IR, functions may return either 0 or 1 value. LLVM IR dialect
@@ -1047,11 +1048,12 @@ def LLVM_BrOp : LLVM_TerminatorOp<"br",
     LLVM_TerminatorPassthroughOpBuilder
   ];
 }
-def LLVM_CondBrOp : LLVM_TerminatorOp<"cond_br",
-    [AttrSizedOperandSegments,
-     DeclareOpInterfaceMethods<BranchOpInterface>,
-     DeclareOpInterfaceMethods<BranchWeightOpInterface>,
-     Pure]> {
+def LLVM_CondBrOp
+    : LLVM_TerminatorOp<
+          "cond_br", [AttrSizedOperandSegments,
+                      DeclareOpInterfaceMethods<BranchOpInterface>,
+                      DeclareOpInterfaceMethods<WeightedBranchOpInterface>,
+                      Pure]> {
   let arguments = (ins I1:$condition,
                    Variadic<LLVM_Type>:$trueDestOperands,
                    Variadic<LLVM_Type>:$falseDestOperands,
@@ -1136,11 +1138,12 @@ def LLVM_UnreachableOp : LLVM_TerminatorOp<"unreachable"> {
   }];
 }
 
-def LLVM_SwitchOp : LLVM_TerminatorOp<"switch",
-    [AttrSizedOperandSegments,
-     DeclareOpInterfaceMethods<BranchOpInterface>,
-     DeclareOpInterfaceMethods<BranchWeightOpInterface>,
-     Pure]> {
+def LLVM_SwitchOp
+    : LLVM_TerminatorOp<
+          "switch", [AttrSizedOperandSegments,
+                     DeclareOpInterfaceMethods<BranchOpInterface>,
+                     DeclareOpInterfaceMethods<WeightedBranchOpInterface>,
+                     Pure]> {
   let arguments = (ins
     AnySignlessInteger:$value,
     Variadic<AnyType>:$defaultOperands,

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

@@ -385,12 +385,12 @@ def WeightedBranchOpInterface : OpInterface<"WeightedBranchOpInterface"> {
     operations, i.e. terminator operations with successors.
 
     This interface provides methods for getting/setting integer non-negative
-    weight of each branch in the range from 0 to 100. The sum of weights
-    must be 100. The number of weights must match the number of successors
-    of the operation.
-
-    The weights specify the probability (in percents) of taking
-    a particular branch.
+    weight of each branch. The probability of executing a branch
+    is computed as the ratio between the branch's weight and the total
+    sum of the weights.
+    The number of weights must match the number of successors of the operation,
+    with one exception for CallOpInterface operations, which may only
+    have on weight when they do not have any successors.
 
     The default implementations of the methods expect the operation
     to have an attribute of type DenseI32ArrayAttr named branch_weights.
@@ -440,15 +440,16 @@ def WeightedRegionBranchOpInterface
     that exhibit branching behavior between held regions.
 
     This interface provides methods for getting/setting integer non-negative
-    weight of each branch in the range from 0 to 100. The sum of weights
-    must be 100. The number of weights must match the number of regions
+    weight of each branch. The probability of executing a region is computed
+    as the ratio between the region branch's weight and the total sum
+    of the weights.
+    The number of weights must match the number of regions
     held by the operation (including empty regions).
 
-    The weights specify the probability (in percents) of branching
-    to a particular region when first executing the operation.
+    The weights specify the probability of branching to a particular
+    region when first executing the operation.
     For example, for loop-like operations with a single region
     the weight specifies the probability of entering the loop.
-    In this case, the weight must be either 0 or 100.
 
     The default implementations of the methods expect the operation
     to have an attribute of type DenseI32ArrayAttr named branch_weights.

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

@@ -189,7 +189,7 @@ class ModuleTranslation {
                                   llvm::Instruction *inst);
 
   /// Sets LLVM profiling metadata for operations that have branch weights.
-  void setBranchWeightsMetadata(BranchWeightOpInterface op);
+  void setBranchWeightsMetadata(WeightedBranchOpInterface op);
 
   /// Sets LLVM loop metadata for branch operations that have a loop annotation
   /// attribute.

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

@@ -9,6 +9,7 @@
 #include <utility>
 
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
@@ -84,24 +85,33 @@ detail::verifyBranchSuccessorOperands(Operation *op, unsigned succNo,
 // WeightedBranchOpInterface
 //===----------------------------------------------------------------------===//
 
-LogicalResult detail::verifyBranchWeights(Operation *op) {
-  auto weights = cast<WeightedBranchOpInterface>(op).getBranchWeightsOrNull();
+static LogicalResult verifyWeights(Operation *op, DenseI32ArrayAttr weights,
+                                   int64_t weightsNum,
+                                   llvm::StringRef weightAnchorName,
+                                   llvm::StringRef weightRefName) {
   if (weights) {
-    if (weights.size() != op->getNumSuccessors())
-      return op->emitError() << "number of weights (" << weights.size()
-                             << ") does not match the number of successors ("
-                             << op->getNumSuccessors() << ")";
-    int32_t total = 0;
-    for (auto weight : llvm::enumerate(weights.asArrayRef())) {
+    if (weights.size() != weightsNum)
+      return op->emitError() << "expects number of " << weightAnchorName
+                             << " weights to match number of " << weightRefName
+                             << ": " << weights.size() << " vs " << weightsNum;
+
+    for (auto weight : llvm::enumerate(weights.asArrayRef()))
       if (weight.value() < 0)
         return op->emitError()
                << "weight #" << weight.index() << " must be non-negative";
-      total += weight.value();
-    }
-    if (total != 100)
-      return op->emitError() << "total weight " << total << " is not 100";
   }
-  return mlir::success();
+  return success();
+}
+
+LogicalResult detail::verifyBranchWeights(Operation *op) {
+  auto weights = cast<WeightedBranchOpInterface>(op).getBranchWeightsOrNull();
+  unsigned successorsNum = op->getNumSuccessors();
+  // CallOpInterface operations without successors may only have
+  // one weight, though it seems to be redundant and indicate
+  // 100% probability of calling the callee(s).
+  int64_t weightsNum =
+      (successorsNum == 0 && isa<CallOpInterface>(op)) ? 1 : successorsNum;
+  return verifyWeights(op, weights, weightsNum, "branch", "successors");
 }
 
 //===----------------------------------------------------------------------===//
@@ -111,22 +121,7 @@ LogicalResult detail::verifyBranchWeights(Operation *op) {
 LogicalResult detail::verifyRegionBranchWeights(Operation *op) {
   auto weights =
       cast<WeightedRegionBranchOpInterface>(op).getRegionWeightsOrNull();
-  if (weights) {
-    if (weights.size() != op->getNumRegions())
-      return op->emitError() << "number of weights (" << weights.size()
-                             << ") does not match the number of regions ("
-                             << op->getNumRegions() << ")";
-    int32_t total = 0;
-    for (auto weight : llvm::enumerate(weights.asArrayRef())) {
-      if (weight.value() < 0)
-        return op->emitError()
-               << "weight #" << weight.index() << " must be non-negative";
-      total += weight.value();
-    }
-    if (total != 100)
-      return op->emitError() << "total weight " << total << " is not 100";
-  }
-  return mlir::success();
+  return verifyWeights(op, weights, op->getNumRegions(), "region", "regions");
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -146,7 +146,7 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 
-  if (auto iface = dyn_cast<BranchWeightOpInterface>(op)) {
+  if (auto iface = dyn_cast<WeightedBranchOpInterface>(op)) {
     iface.setBranchWeights(builder.getDenseI32ArrayAttr(branchWeights));
     return success();
   }

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

@@ -1055,7 +1055,7 @@ LogicalResult ModuleTranslation::convertBlockImpl(Block &bb,
       return failure();
 
     // Set the branch weight metadata on the translated instruction.
-    if (auto iface = dyn_cast<BranchWeightOpInterface>(op))
+    if (auto iface = dyn_cast<WeightedBranchOpInterface>(op))
       setBranchWeightsMetadata(iface);
   }
 
@@ -2026,7 +2026,7 @@ void ModuleTranslation::setDereferenceableMetadata(
   inst->setMetadata(kindId, derefSizeNode);
 }
 
-void ModuleTranslation::setBranchWeightsMetadata(BranchWeightOpInterface op) {
+void ModuleTranslation::setBranchWeightsMetadata(WeightedBranchOpInterface op) {
   DenseI32ArrayAttr weightsAttr = op.getBranchWeightsOrNull();
   if (!weightsAttr)
     return;

mlir/test/Dialect/ControlFlow/invalid.mlir

@@ -72,7 +72,7 @@ func.func @switch_missing_default(%flag : i32, %caseOperand : i32) {
 
 // CHECK-LABEL: func @wrong_weights_number
 func.func @wrong_weights_number(%cond: i1) {
-  // expected-error@+1 {{number of weights (1) does not match the number of successors (2)}}
+  // expected-error@+1 {{expects number of branch weights to match number of successors: 1 vs 2}}
   cf.cond_br %cond weights([100]), ^bb1, ^bb2
   ^bb1:
     return
@@ -91,15 +91,3 @@ func.func @wrong_total_weight(%cond: i1) {
   ^bb2:
     return
 }
-
-// -----
-
-// CHECK-LABEL: func @wrong_total_weight
-func.func @wrong_total_weight(%cond: i1) {
-  // expected-error@+1 {{total weight 101 is not 100}}
-  cf.cond_br %cond weights([100, 1]), ^bb1, ^bb2
-  ^bb1:
-    return
-  ^bb2:
-    return
-}