Bug 235581 – Add support for WASM branch hinting proposal

• Status: RESOLVED
• Resolution: FIXED
• Component: JavaScriptCore

[Tom Tartarin]

See the proposal for a more detailed description: https://github.com/WebAssembly/branch-hinting.

This allows parsing a "code_annotation.branch_hint" custom section, as per the code annotation proposal: https://github.com/WebAssembly/annotations.
This section provides per function information about how likely a branch at a
given offset is to be taken. It is similar to branch weight metadata in LLVM.
For B3 and AIR this implies generating "hinted" blocks, whose frequencies are already set.

The provided test is inspired by how the name section is tested: a wasm file containing the encoded custom section is read, compiled and run. It only tests the integration of the WasmBranchHintsSectionParser in the parsing pipeline, since more accurate testing would require inspecting the internal state of the VM.

[Tom Tartarin]

Created attachment 449921 [details]
Patch

[Tom Tartarin]

Created attachment 449923 [details]
bug-235581: add support for WASM branch hinting proposal

[Yusuke Suzuki]

Comment on attachment 449921 [details]
Patch

View in context: https://bugs.webkit.org/attachment.cgi?id=449921&action=review

Nice!

> Source/JavaScriptCore/b3/B3EstimateStaticExecutionCounts.cpp:45
>          constexpr double base = 10.0;
> +        if (block->isHinted())
> +            continue;
>          block->setFrequency(pow(base, naturalLoops.loopDepth(block)));

Should we treat the hint as a coefficient to this calculated count?
Let's consider that here is a block in the deep loop (4-depth). Then, the frequency becomes 10e4.
It is much higher than the hint's value (0 or 1.0). So, currently, if we use hint, it lowers the value compared to the value offered by static-execution-counts.

Background: In Wasm B3 IR generator, except for FrequencyClass::Rare, we are not attaching frequency and calculating it in this pass.
In JS, we are using execution counter to profile block's frequency and reflect it at runtime.

[Yusuke Suzuki]

Comment on attachment 449921 [details]
Patch

View in context: https://bugs.webkit.org/attachment.cgi?id=449921&action=review

>> Source/JavaScriptCore/b3/B3EstimateStaticExecutionCounts.cpp:45
>>          block->setFrequency(pow(base, naturalLoops.loopDepth(block)));
> 
> Should we treat the hint as a coefficient to this calculated count?
> Let's consider that here is a block in the deep loop (4-depth). Then, the frequency becomes 10e4.
> It is much higher than the hint's value (0 or 1.0). So, currently, if we use hint, it lowers the value compared to the value offered by static-execution-counts.
> 
> Background: In Wasm B3 IR generator, except for FrequencyClass::Rare, we are not attaching frequency and calculating it in this pass.
> In JS, we are using execution counter to profile block's frequency and reflect it at runtime.

Ah, let me wait. I'm currently discussing with Saam about the way to use branch hinting.
It would be much simpler.

[Saam Barati]

Comment on attachment 449921 [details]
Patch

View in context: https://bugs.webkit.org/attachment.cgi?id=449921&action=review

Your implementation in B3/Air is more complicated than it needs to be. All you need to do is set the FrequencyClass based on the hint. If the hint is Likely, set the taken side  "Normal" and not taken side "Rare". If the hint is Unlikely, set it the taken side to "Rare" and the not taken sideT "Normal". That'll get you the behavior you want. No need to change static execution count estimate, no need to add the "isHinted" stuff to any of the basic block classes.

> Source/JavaScriptCore/b3/B3BasicBlock.h:150
> +    void setHinted(bool hinted) { m_hinted = hinted; }

none of this is necessary. See my suggestion on how to implement this patch.

> Source/JavaScriptCore/b3/air/AirBasicBlock.h:137
> +    void setHinted(bool hinted) { m_hinted = hinted; }
> +    bool isHinted() const { return m_hinted; }

none of this is necessary. See my suggestion on how to implement this patch.

> Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:2498
> +auto B3IRGenerator::addIfWithHint(ExpressionType condition, BlockSignature signature, Stack& enclosingStack, ControlType& result, Stack& newStack, const BranchHint& hint) -> PartialResult

ditto to what I wrote about the AirIRGenerator having two versions of this function.

> Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:2755
> +    if (Options::useWebAssemblyBranchHints()) {
> +        const BranchHint& hint = m_info.getBranchHint(m_functionIndex, m_parser->currentOpcodeStartingOffset());
> +        if (hint != BranchHint::Invalid)
> +            return addBranchWithHint(data, condition, returnValues, hint);
> +    }

Instead of the second "addBranchWithHint" function, just define default values for FrequencyClass for taken/not taken. Then if we find a branch hint inside the map, you can change the values of those frequency classes.

Also, no need for "const BranchHint&". Just "BranchHint hint = m_info.get(...);"

> Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:2775
> +auto B3IRGenerator::addBranchWithHint(ControlData& data, ExpressionType condition, const Stack& returnValues, const BranchHint& hint) -> PartialResult

You shouldn't use two versions of essentially an identical function. Just have default values for the FrequencyClass used and augment them when we have hint data.

> Source/JavaScriptCore/wasm/WasmBranchHints.h:59
> +    StdUnorderedMap<uint32_t, BranchHint> m_map;

you should use HashMap here (and with UnsignedWithZeroKeyHashTraits  if the branch offset can be zero)

> Source/JavaScriptCore/wasm/WasmBranchHints.h:62
> +

all functions below here can be constexpr.

Also, no need to really take in this enum as "const BranchHint&"

> Source/JavaScriptCore/wasm/WasmBranchHints.h:70
> +inline BranchHint branchHintToInverse(const BranchHint& hint)

branchHintToInverse => "invertBranchHint"

> Source/JavaScriptCore/wasm/WasmBranchHints.h:84
> +        break;

maybe just return false here?

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:39
> +    int64_t lastFunctionIndex = -1;

name nit: "lastFunctionIndex" => "previousFunctionIndex"

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:46
> +        WASM_PARSER_FAIL_IF(int64_t(functionIndex) <= lastFunctionIndex, "invalid function index ", functionIndex, " for function ", i);

webkit style guide uses static_cast<int64_t>, not C style casts.

Also, is this really "<="? Why not strictly "<"? The spec says we can keep having functions of the same index? That has consequences on your algorithm below using "emplace"

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:60
> +            WASM_PARSER_FAIL_IF(int64_t(branchOffset) <= lastBranchOffset, "invalid branch offset ", branchOffset, " for hint ", j);

same nit regarding c style casts


Also, when using HashMap, we should fail if the branchOffset is either the empty or deleted value in the hash map. See the hash trait I was talking about:

template<typename T> struct UnsignedWithZeroKeyHashTraits : GenericHashTraits<T> {
    static constexpr bool emptyValueIsZero = false;
    static T emptyValue() { return std::numeric_limits<T>::max(); }
    static void constructDeletedValue(T& slot) { slot = std::numeric_limits<T>::max() - 1; }
    static bool isDeletedValue(T value) { return value == std::numeric_limits<T>::max() - 1; }
};


Which in general should be fine. Nobody should be using such a large offset anyways.

Also, same question regarding "<=" instead of "<"

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:71
> +            BranchHint branchHint = static_cast<BranchHint>(int32_t(parsedBranchHint));

don't use C style cast here. But also, I don't see why it's necessary at all. Shouldn't static_cast<BranchHint>(parsedBranchHint) work?

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:72
> +            WASM_PARSER_FAIL_IF(!isValidBranchHint(branchHint), "invalid branch hint value ", parsedBranchHint, " for hint ", j);

how is this possible if we're parsing a VarUint1? The value has to be zero or one, right? Should this be an assert?

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:74
> +            branchHintsForFunction.insert(branchOffset, branchHint);

related question regarding the "<=" above. If functionIndex already exists, what semantics do we want? Merge? Overwrite? Skip?

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:76
> +        m_info->branchHints.emplace(functionIndex, WTFMove(branchHintsForFunction));

related question regarding the "<=" above. If functionIndex already exists, what semantics do we want? Merge? Overwrite? Skip?

> Source/JavaScriptCore/wasm/WasmModuleInformation.h:40
> +    using BranchHints = StdUnorderedMap<uint32_t, BranchHintMap>;

you should use HashMap here with UnsignedWithZeroKeyHashTraits

[Tom Tartarin]

Thank you for the detailed review. I will rework the patch and submit soon.

[Tom Tartarin]

Created attachment 450087 [details]
Updated patch after review

[Yusuke Suzuki]

Comment on attachment 450087 [details]
Updated patch after review

View in context: https://bugs.webkit.org/attachment.cgi?id=450087&action=review

r=me

> Source/JavaScriptCore/wasm/WasmAirIRGenerator.cpp:3186
> +        if (hint == BranchHint::Unlikely)
> +            takenFrequency = B3::FrequencyClass::Rare;
> +        else if (hint == BranchHint::Likely)
> +            notTakenFrequency = B3::FrequencyClass::Rare;

Let's use switch.

> Source/JavaScriptCore/wasm/WasmAirIRGenerator.cpp:3439
> +        if (hint == BranchHint::Unlikely)
> +            targetFrequency = B3::FrequencyClass::Rare;
> +        else if (hint == BranchHint::Likely)
> +            continuationFrequency = B3::FrequencyClass::Rare;

Ditto.

> Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:2486
> +        if (hint == BranchHint::Unlikely)
> +            takenFrequency = FrequencyClass::Rare;
> +        else if (hint == BranchHint::Likely)
> +            notTakenFrequency = FrequencyClass::Rare;

Let's use switch, which can ensure that all enums are handled.

switch (hint) {
case BranchHint::Unlikely:
    takenFrequency = FrequencyClass::Rare;
    break;
case BranchHint::Likely:
    notTakenFrequency = FrequencyClass::Rare;
    break;
}

> Source/JavaScriptCore/wasm/WasmB3IRGenerator.cpp:2742
> +        if (hint == BranchHint::Unlikely)
> +            targetFrequency = FrequencyClass::Rare;
> +        else if (hint == BranchHint::Likely)
> +            continuationFrequency = FrequencyClass::Rare;

Ditto.

> Source/JavaScriptCore/wasm/WasmBranchHintsSectionParser.cpp:70
> +            BranchHint branchHint = static_cast<BranchHint>(parsedBranchHint);

Let's add

ASSERT(isValidBranchHint(branchHint));

It should be due to parseVarUInt1. But such an ASSERT is helpful if this branch hint is extended to have more values (then, assert hits and we can fix the issue :)).

[Tom Tartarin]

Created attachment 450230 [details]
Use switch instead of if else. Add ASSERT for branch hint validation

[Tom Tartarin]

Thanks for the review!

[EWS]

Committed r288758 (246546@main): <https://commits.webkit.org/246546@main>

All reviewed patches have been landed. Closing bug and clearing flags on attachment 450230 [details].

[Radar WebKit Bug Importer]

<rdar://problem/88199746>