Strengthened constraint on the number of weights for calls.

Author: vzakhari

mlir/include/mlir/Interfaces/ControlFlowInterfaces.td

@@ -387,11 +387,9 @@ def WeightedBranchOpInterface : OpInterface<"WeightedBranchOpInterface"> {
     This interface provides methods for getting/setting integer non-negative
     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.
+    sum of the weights (which cannot be zero).
     The weights are optional. If they are provided, then their number
-    must match the number of successors of the operation,
-    with one exception for CallOpInterface operations, which may only
-    have one weight when they do not have any successors.
+    must match the number of successors of the operation.
 
     The default implementations of the methods expect the operation
     to have an attribute of type DenseI32ArrayAttr named branch_weights.
@@ -445,7 +443,7 @@ def WeightedRegionBranchOpInterface
     This interface provides methods for getting/setting integer non-negative
     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.
+    of the weights (which cannot be zero).
     The weights are optional. If they are provided, then their number
     must match the number of regions held by the operation
     (including empty regions).

mlir/lib/Interfaces/ControlFlowInterfaces.cpp

@@ -103,20 +103,17 @@ static LogicalResult verifyWeights(Operation *op,
     if (weight < 0)
       return op->emitError() << "weight #" << index << " must be non-negative";
 
+  if (llvm::all_of(weights, [](int32_t value) { return value == 0; }))
+    return op->emitError() << "branch weights cannot all be zero";
+
   return success();
 }
 
 LogicalResult detail::verifyBranchWeights(Operation *op) {
   llvm::ArrayRef<int32_t> weights =
       cast<WeightedBranchOpInterface>(op).getWeights();
-  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).
-  // TODO: maybe we should disallow weights for calls without successors.
-  std::size_t weightsNum =
-      (successorsNum == 0 && isa<CallOpInterface>(op)) ? 1 : successorsNum;
-  return verifyWeights(op, weights, weightsNum, "branch", "successors");
+  return verifyWeights(op, weights, op->getNumSuccessors(), "branch",
+                       "successors");
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -147,7 +147,14 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
   }
 
   if (auto iface = dyn_cast<WeightedBranchOpInterface>(op)) {
-    iface.setWeights(branchWeights);
+    // LLVM allows attaching a single weight to call instructions.
+    // This is used for carrying the execution count information
+    // in PGO modes. MLIR WeightedBranchOpInterface does not allow this,
+    // so we drop the metadata in this case.
+    // LLVM should probably use the VP form of MD_prof metadata
+    // for such cases.
+    if (op->getNumSuccessors() != 0)
+      iface.setWeights(branchWeights);
     return success();
   }
   return failure();

mlir/test/Dialect/ControlFlow/invalid.mlir

@@ -91,3 +91,15 @@ func.func @wrong_total_weight(%cond: i1) {
   ^bb2:
     return
 }
+
+// -----
+
+// CHECK-LABEL: func @zero_weights
+func.func @wrong_total_weight(%cond: i1) {
+  // expected-error@+1 {{branch weights cannot all be zero}}
+  cf.cond_br %cond weights([0, 0]), ^bb1, ^bb2
+  ^bb1:
+    return
+  ^bb2:
+    return
+}

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

@@ -36,14 +36,17 @@ bbd:
 
 ; // -----
 
+; Verify that a single weight attached to a call is not translated.
+; The MLIR WeightedBranchOpInterface does not support this case.
+
 ; CHECK: llvm.func @fn()
-declare void @fn()
+declare i32 @fn()
 
 ; CHECK-LABEL: @call_branch_weights
-define void @call_branch_weights() {
-  ; CHECK:  llvm.call @fn() {branch_weights = array<i32: 42>}
-  call void @fn(), !prof !0
-  ret void
+define i32 @call_branch_weights() {
+  ; CHECK:  llvm.call @fn() : () -> i32
+  %1 = call i32 @fn(), !prof !0
+  ret i32 %1
 }
 
 !0 = !{!"branch_weights", i32 42}

mlir/test/Target/LLVMIR/llvmir-invalid.mlir

@@ -448,3 +448,23 @@ llvm.mlir.global external constant @const() {addr_space = 0 : i32, dso_local} :
 }
 
 llvm.func extern_weak @extern_func()
+
+// -----
+
+llvm.func @fn()
+
+llvm.func @call_branch_weights() {
+  // expected-error @below{{expects number of branch weights to match number of successors: 1 vs 0}}
+  llvm.call @fn() {branch_weights = array<i32 : 42>} : () -> ()
+  llvm.return
+}
+
+// -----
+
+llvm.func @fn() -> i32
+
+llvm.func @call_branch_weights() {
+  // expected-error @below{{expects number of branch weights to match number of successors: 1 vs 0}}
+  %res = llvm.call @fn() {branch_weights = array<i32 : 42>} : () -> i32
+  llvm.return
+}

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -1906,32 +1906,6 @@ llvm.func @cond_br_weights(%cond : i1, %arg0 : i32,  %arg1 : i32) -> i32 {
 
 // -----
 
-llvm.func @fn()
-
-// CHECK-LABEL: @call_branch_weights
-llvm.func @call_branch_weights() {
-  // CHECK: !prof ![[NODE:[0-9]+]]
-  llvm.call @fn() {branch_weights = array<i32 : 42>} : () -> ()
-  llvm.return
-}
-
-// CHECK: ![[NODE]] = !{!"branch_weights", i32 42}
-
-// -----
-
-llvm.func @fn() -> i32
-
-// CHECK-LABEL: @call_branch_weights
-llvm.func @call_branch_weights() {
-  // CHECK: !prof ![[NODE:[0-9]+]]
-  %res = llvm.call @fn() {branch_weights = array<i32 : 42>} : () -> i32
-  llvm.return
-}
-
-// CHECK: ![[NODE]] = !{!"branch_weights", i32 42}
-
-// -----
-
 llvm.func @foo()
 llvm.func @__gxx_personality_v0(...) -> i32