introduce `@branchWeight` builtin with numerical branch weights

[andrewrk]

Originally specified at #489 (comment).

This proposal supplants:

• @expect() hint to optimizer #489
• replace @setCold with @cold #5177

Summary:

• No "expect" builtin or similar
• Each branch gets a default weight of 10
• Error branches (weight error branches against errors #84) get a default weight of 1
• Branch weight can be overridden with @branchWeight builtin.
• Branch weights are of type u32

This generalizes better to switch expressions.

This allows source code to be annotated with PGO data.

[mlugg]

The only issue I see here is that this doesn't work so well with #8220, because @branchWeight tells us nothing about the source of the branch; perhaps it is unlikely for a given prong to be reached from the initial switch, but very likely to be reached from a labeled continue. You could also argue that this applies to loops: maybe a loop body is unlikely to be initially entered, but highly likely to be repeated once entered (although this case can't really be nicely modeled with @expect either). However, on the whole I agree that this is a better model than @expect.

[rohlem]

@mlugg If I understand your comment correctly, the actual distinction here would be to mark the branches = code paths rather than the target block they lead to.
I think this could be modeled by introducing additional blocks along the way, and only specifying @branchWeight for these intermediate blocks.

Current syntax might make this more verbose than would be ideal, but here's some examples anyway:

if (iterator.hasNext()) {
  @branchWeight(2); //unlikely to enter
  while(true) {
    const next = iterator.next().?;

    // actual logic - no @branchWeight directly here

    if (!iterator.hasNext()) {
      @branchWeight(1); //very unlikely to exit
      break;
    }
    @branchWeight(20); //very likely to continue looping
  }
}

switch(nextToken()) {
  .end => {},
  else => |tok| {
    switch (tok) {
      .a => @branchWeight(1), //unlikely to be .a initially
    }
    inner: switch(tok) {
      .a => {

        // actual logic - no @branchWeight directly here

        const next_token = nextToken();
        switch(next_token) {
          .a => @branchWeight(20), //very likely to be .a as a followup
          // other cases...
        }
        continue :inner next_token; //or alternatively inline this continue into the switch above with `@branchWeight` to make it easier for the optimizer to understand this?
      },
      // other cases...
    }
  }
}

[kprotty]

The idea of a builtin on the block instead of the condition makes sense to support cases where the condition isn't visible (or would be annoying to make visible) in the source code like captures (while (x) |y|, if (x) |y|) and monadic control flow (orelse {}, catch {}).

One thing that still feels fuzzy is the idea of having arbitrary weight values. The original mentions that probability values are problematic when keeping consistent across code edits. But more concretely, there's the argument for switch cases and profiling:

Switch cases in reality have little control on how they're dispatched: If a case is marked as "likely/unlikely" but a jump table would be ideal, what should codegen do here? If cases are sparse (jump table unavailable), what is the codegen for two equally "likely/unlikely" cases? Branch hints primarily affect which path falls through a conditional and which is jumped to. Default Switch statements don't execute comparisons in order (you can have else => first) so there's no concept of fall-through. Labeled continue may allow expressing this, but branching within them is still a matter of likeliness (binary), not probability/weight.

Profiling (Profile Guided Optimization) tracks which branches are taken to form a probability model, but how this affects codegen is another story: PGO can either 1) reorganize branching to make most/least probable ones fall-through/jumped-to or 2) eliminate the (conditional) branch entirely. The latter is aggressive and only correct when the optimized program will only run again with similar inputs to when profiled. The former is still a decision that's binary, not probability/weighted.

Note

It's also unclear how PGO could take advantage of "source code annotations" here since it bases optimizations on what's executed not speculated. Maybe it could report if an unlikely-marked branch was profiled to be more likely? Either ways, no codegen changes that this builtin wouldn't already do.

---

I propose the builtin be renamed to something like @{set}branchHint(x) which takes in a comptime enum with variant names like .likely/.unlikely or .fall_through/.jumped_to. The semantic edge cases could then be defined:

• Q: How to simulate old @setCold(x) behavior?
  A: The old "coldness" is a more a property of inlining (a type of branching, with more implications). At the function level, use noinline or callconv(if (x) .NoInline else .Auto) if x is comptime-computed (requires a new CallingConvention.NoInline). At the scope level, use this branching hint. We could also instead define branchHint at the function level to be similar to inline from C and __attribute__((cold)) from GCC-likes respectively.

• Q: What happens if two branchHint()s are called in the same scope?
  A: We should mimic whatever @setRuntimeSafety does here. Would propose to just ignore any subsequent others that aren't the first.

• Q: What happens if branchHint occurs at function level of inline fn?
  A: Inline function bodies IIRC are semantically copied to the caller's invocation-site. It should then either be cancelled out (as noted above) or affect the branch hint for the caller's body which contains the invocation.

• Q: What to do if branching to the same location must be jumped-to in one path and fall-through in another?
  A: Borrowing from @rohlem, add the respective branch hints to the conditional block at each path before branching to the same location. Example:

  while (true) {
      if (x) break; // unlikely
      if (y) break; // likely
  }
  // branchHint(???)
  sameLocation();
  
  // should become
  
  while (true) {
      if (x) { branchHint(.unlikely); break; }
      if (y) { branchHint(.likely); break; }
  }
  sameLocation();

[MatthiasPortzel]

I’m not a fan of “weight” values with undefined semantics (in general), because it leads to paranoid developers doing things like @branchWeight(999999999) or @branchWeight(4_294_967_295).

(It may be too early, but) it would be nice if this had a documentation that;

1. the compiler would use these values only relative to each other, never implement thresholds where weights over a magic number trigger different code-gen and
2. these values would be compared only locally, and if I call a function that uses large branch-weights, I don’t have to increase the branch weights in my code.

These are just the concerns that pop into my head, I don’t know how it would actually be implemented.

[Snektron]

Each branch gets a default weight of 10

This seems a little bit arbitrary. Why not use floats and specify 1.0 as the default weight?

[andrewrk]

LLVM's branch weight metadata encodes probability. To those who suggest to toss out this data, can you explain why LLVM devs have decided to include this information, and why they have made a mistake by doing so?

[Rexicon226]

LLVM optimizes branches using weights to be compatible with PGO:

Branch weights might be fetch from the profiling file, or generated based on __builtin_expect

The weights are also influenced by variables like what's inside the branch. LLVM internally warns against using arbitrary weights given the clash with implementation thresholds. We should take their advice and prefer the binary likely/unlikely API

[dweiller]

The only issue I see here is that this doesn't work so well with #8220, because @branchWeight tells us nothing about the source of the branch; perhaps it is unlikely for a given prong to be reached from the initial switch, but very likely to be reached from a labeled continue.

Is there any kind of control flow that would want information about the source of a branch other than replicated dispatch switch loops? In a replicated dispatch switch loop you probably would want to have more control than simply saying which is the mostly likely target anyway, so I don't think @expect would be the right thing. If there aren't other control flow constructs that benefit from the source information and we want to give the programmer enough control to specify per-branch dispatch weightings/expectations, I think that @branchWeight could generalise to this situation by allowing it to annotate the continue expression in a labelled switch loop:

sw_loop: switch (instructions[i].tag) {
    .add => {
        i += 1;
        continue :sw_loop @branchWeight(.{ .add 4 }, .{ .mul, 2 }, .{ .div, 1 }) intructions[i].tag;
    },
    .mul => {
        i += 1;
        continue :sw_loop @branchWeight(.{ .add, 2 }, .{ .mul, 3 },}) intructions[i].tag;
    },
    .div => {
        i += 1;
        continue :sw_loop intructions[i].tag;
    },
}

That looks a bit unweildy, but maybe it could take an anonymous struct instead allowing the weight map to be specified elsewhere and leaving continue: sw_loop @branchWeight(add_weights) instructions[i].tag;. There might still be a readability issue as it's not obvious whether branch weights provided to a continue statement merely override @branchWeight statements in the toplevel body of switch cases or if @branchWeight in a branch only applies to the initial switch and any continues without a @branchWeight annotation.

You could also argue that this applies to loops: maybe a loop body is unlikely to be initially entered, but highly likely to be repeated once entered (although this case can't really be nicely modeled with @expect either). However, on the whole I agree that this is a better model than @expect.

This could be modelled with either @branchWeight or @expect by manully doing loop inversion, though readability suffers:

if (condition) {
    @branchWeight(1);
    while (true) {
        @branchWeight(2);
        // loop code
        if (!condition) { 
            @branchWeight(1);
            break;
        }
    }
} else {
    @branchWeight(2); // shouldn't be needed given default weight of 10
}

It may be sensible to just make the above the interpretation of an 'unexpected' loop body so it could be simply written as

while (condition) {
    @branchWeight(1);
    // loop code
} else {
    @branchWeight(2); //shouldn't be needed given default weight of 10
}

as I'm not sure it makes sense (and as far as I know it's not possible) to have a loop continually predicted to exit.

Switch cases in reality have little control on how they're dispatched: If a case is marked as "likely/unlikely" but a jump table would be ideal, what should codegen do here? If cases are sparse (jump table unavailable), what is the codegen for two equally "likely/unlikely" cases? Branch hints primarily affect which path falls through a conditional and which is jumped to. Default Switch statements don't execute comparisons in order (you can have else => first) so there's no concept of fall-through. Labeled continue may allow expressing this, but branching within them is still a matter of likeliness (binary), not probability/weight.

I would expect the branch weightings to just be additional input to assist codegen heuristics - there doesn't have to be a order-preserving map between the weightings and the branch's preference in the generated code, rather I would expect them to be used to try and optimise something like branch-weight-weighted average performance. If a jump table would be ideal just ignore them and generate a jump table. For sparse cases there could be a whole rabbit hole of techniques here, but the weights provide the information needed to make decisions. It should be relatively straightforward to map branch weights to a binary decision tree that optimises fall-through/jump behaviour relative to the provided weights; while any specific branch in that tree is binary the best way to arrange the whole tree requires the per-branch weighting information precisely because there is no comparison order information in a switch statement.

LLVM optimizes branches using weights to be compatible with PGO:

Branch weights might be fetch from the profiling file, or generated based on __builtin_expect

The weights are also influenced by variables like what's inside the branch. LLVM internally warns against using arbitrary weights given the clash with implementation thresholds. We should take their advice and prefer the binary likely/unlikely API

Even if we use the binary API, we can still use something like @branchWeight to decide when a branch is likely/unlikely - the weight doesn't necessarily need to be mapped directly to LLVM's branch weight metadata probability.

[kprotty]

In a replicated dispatch switch loop you probably would want to have more control than simply saying which is the mostly likely target anyway

An efficient dispatch switch should have dense comparison targets so that it generates a jump-table. If the cases are sparse, then dispatch is no longer efficient and hints wont help: It then requires at most log(N) comparisons and choosing which to fall-through into with a continue :sw weight can be cancelled out by that target choosing another, even if Zig added a guaranteed comparison order.

the best way to arrange the whole tree requires the per-branch weighting information precisely because there is no comparison order information in a switch statement.

A weight can dictate the level in the comparison tree that a conditional appears (and maybe the order in that level), but not which side or the depth of the tree. So in these tree-based sparse comparisons, one still can't accurately control branching codegen.

As a practical example, something like this could be written using a switch statement. But unsure how weights at each case could achieve the desired codegen compared to its explicit use of likely ifs / writing out the desired comparison-tree.

as I'm not sure it makes sense (and as far as I know it's not possible) to have a loop continually predicted to exit.

I think so too. A while loop at codegen is two conditionals: 1) condition to enter the loop and 2) condition to keep going. For a loop to jump back up, 2) must conditionally branch so its fundamentally unlikely. The branch hint would then apply to 1) for whether it falls-through to loop body or jumps-away to it.

[rohlem]

A while loop at codegen is two conditionals: 1) condition to enter the loop and 2) condition to keep going. For a loop to jump back up, 2) must conditionally branch so its fundamentally unlikely.

Technically unrolling the loop is a mechanism to model 2) as the likely branch via inlining, although the resulting code size limits the scalability of this.
(No idea how small a loop would have to be to make the likely-unlikely tradeoff realistically worth the size increase - I imagine pretty small.)

[dweiller]

the best way to arrange the whole tree requires the per-branch weighting information precisely because there is no comparison order information in a switch statement.

A weight can dictate the level in the comparison tree that a conditional appears (and maybe the order in that level), but not which side or the depth of the tree. So in these tree-based sparse comparisons, one still can't accurately control branching codegen.

I agree that the programmer is not going to have precise control over the branching and I don't think they ever will with a sparse switch; what I'm getting at is that I would expect a 'sparse switch optimizer' to be minimizing some (potentially target specific) objective function of branch weights.

In a replicated dispatch switch loop you probably would want to have more control than simply saying which is the mostly likely target anyway

An efficient dispatch switch should have dense comparison targets so that it generates a jump-table. If the cases are sparse, then dispatch is no longer efficient and hints wont help: It then requires at most log(N) comparisons and choosing which to fall-through into with a continue :sw weight can be cancelled out by that target choosing another, even if Zig added a guaranteed comparison order.

I think I agree with you, though I might have to think about that last sentence a bit more (not 100% sure at the moment what you mean with the cancelling out). I wasn't really trying to suggest the branch weights in switch continue loops is going to be very useful because, as you say, when you care you'll make jump tables. I was just saying that I think if someone did end up wanting source and target branch information for some kind of optimization pass here I think a special form of @branchWeight could be used to do it.

[andrewrk]

@jacobly0 I would value your input here if you have the energy to spare. Otherwise I'll go ahead and make an executive decision.