diff --git a/src/parser/contexts.h b/src/parser/contexts.h index fd835134eb5..b4d27133dcb 100644 --- a/src/parser/contexts.h +++ b/src/parser/contexts.h @@ -23,6 +23,7 @@ #include "support/name.h" #include "support/result.h" #include "support/string.h" +#include "wasm-annotations.h" #include "wasm-builder.h" #include "wasm-ir-builder.h" #include "wasm.h" @@ -2339,11 +2340,47 @@ struct ParseDefsCtx : TypeParserCtx { return withLoc(pos, irBuilder.makeCallIndirect(*t, type, isReturn)); } + // Return the branch hint for a branching instruction, if there is one. + std::optional + getBranchHint(const std::vector& annotations) { + // Find and apply (the last) branch hint. + const Annotation* hint = nullptr; + for (auto& a : annotations) { + if (a.kind == Annotations::BranchHint) { + hint = &a; + } + } + if (!hint) { + return std::nullopt; + } + + Lexer lexer(hint->contents); + if (lexer.empty()) { + std::cerr << "warning: empty BranchHint\n"; + return std::nullopt; + } + + auto str = lexer.takeString(); + if (!str || str->size() != 1) { + std::cerr << "warning: invalid BranchHint string\n"; + return std::nullopt; + } + + auto value = (*str)[0]; + if (value != 0 && value != 1) { + std::cerr << "warning: invalid BranchHint value\n"; + return std::nullopt; + } + + return bool(value); + } + Result<> makeBreak(Index pos, const std::vector& annotations, Index label, bool isConditional) { - return withLoc(pos, irBuilder.makeBreak(label, isConditional)); + auto likely = getBranchHint(annotations); + return withLoc(pos, irBuilder.makeBreak(label, isConditional, likely)); } Result<> makeSwitch(Index pos, diff --git a/src/passes/Print.cpp b/src/passes/Print.cpp index 35fa57ea51a..435b5061dad 100644 --- a/src/passes/Print.cpp +++ b/src/passes/Print.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -267,15 +268,17 @@ struct PrintSExpression : public UnifiedExpressionVisitor { void printDebugLocation(const std::optional& location); - void printDebugLocation(Expression* curr); // Prints debug info for a delimiter in an expression. void printDebugDelimiterLocation(Expression* curr, Index i); + // Prints debug info and code annotations. + void printMetadata(Expression* curr); + void printExpressionContents(Expression* curr); void visit(Expression* curr) { - printDebugLocation(curr); + printMetadata(curr); UnifiedExpressionVisitor::visit(curr); } @@ -2677,20 +2680,20 @@ void PrintSExpression::printDebugLocation( doIndent(o, indent); } -void PrintSExpression::printDebugLocation(Expression* curr) { +void PrintSExpression::printMetadata(Expression* curr) { if (currFunction) { - // show an annotation, if there is one - auto& debugLocations = currFunction->debugLocations; - auto iter = debugLocations.find(curr); - if (iter != debugLocations.end()) { + // Show a debug location, if there is one. + if (auto iter = currFunction->debugLocations.find(curr); + iter != currFunction->debugLocations.end()) { printDebugLocation(iter->second); } else { printDebugLocation(std::nullopt); } - // show a binary position, if there is one + + // Show a binary position, if there is one. if (debugInfo) { - auto iter = currFunction->expressionLocations.find(curr); - if (iter != currFunction->expressionLocations.end()) { + if (auto iter = currFunction->expressionLocations.find(curr); + iter != currFunction->expressionLocations.end()) { Colors::grey(o); o << ";; code offset: 0x" << std::hex << iter->second.start << std::dec << '\n'; @@ -2698,6 +2701,19 @@ void PrintSExpression::printDebugLocation(Expression* curr) { doIndent(o, indent); } } + + // Show a code annotation, if there is one. + if (auto iter = currFunction->codeAnnotations.find(curr); + iter != currFunction->codeAnnotations.end()) { + auto& annotation = iter->second; + if (annotation.branchLikely) { + Colors::grey(o); + o << "(@" << Annotations::BranchHint << " \"\\0" + << (*annotation.branchLikely ? "1" : "0") << "\")\n"; + restoreNormalColor(o); + doIndent(o, indent); + } + } } } @@ -2781,7 +2797,7 @@ void PrintSExpression::visitBlock(Block* curr) { while (1) { if (stack.size() > 0) { doIndent(o, indent); - printDebugLocation(curr); + printMetadata(curr); } stack.push_back(curr); o << '('; diff --git a/src/wasm-annotations.h b/src/wasm-annotations.h new file mode 100644 index 00000000000..42f0e56d732 --- /dev/null +++ b/src/wasm-annotations.h @@ -0,0 +1,32 @@ +/* + * Copyright 2025 WebAssembly Community Group participants + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// +// Support for code annotations. +// + +#ifndef wasm_annotations_h +#define wasm_annotations_h + +#include "support/name.h" + +namespace wasm::Annotations { + +extern const Name BranchHint; + +} // namespace wasm::Annotations + +#endif // wasm_annotations_h diff --git a/src/wasm-ir-builder.h b/src/wasm-ir-builder.h index 7fbe089c127..d979fb52597 100644 --- a/src/wasm-ir-builder.h +++ b/src/wasm-ir-builder.h @@ -72,8 +72,7 @@ class IRBuilder : public UnifiedExpressionVisitor> { this->codeSectionOffset = codeSectionOffset; } - // Set the function used to add scratch locals when constructing an isolated - // sequence of IR. + // The current function, used to create scratch locals, add annotations, etc. void setFunction(Function* func) { this->func = func; } // Handle the boundaries of control flow structures. Users may choose to use @@ -119,7 +118,9 @@ class IRBuilder : public UnifiedExpressionVisitor> { Result<> makeBlock(Name label, Signature sig); Result<> makeIf(Name label, Signature sig); Result<> makeLoop(Name label, Signature sig); - Result<> makeBreak(Index label, bool isConditional); + Result<> makeBreak(Index label, + bool isConditional, + std::optional likely = std::nullopt); Result<> makeSwitch(const std::vector& labels, Index defaultLabel); // Unlike Builder::makeCall, this assumes the function already exists. Result<> makeCall(Name func, bool isReturn); diff --git a/src/wasm.h b/src/wasm.h index b6541fc6cf5..60efbfd65b8 100644 --- a/src/wasm.h +++ b/src/wasm.h @@ -2191,6 +2191,16 @@ class Function : public Importable { delimiterLocations; BinaryLocations::FunctionLocations funcLocation; + // Code annotations. As with debug info, we do not store these on Expressions + // themselves, as we assume most instances are unannotated, and do not want to + // add constant memory overhead. + struct CodeAnnotation { + // Branch hinting proposal: Whether the branch is likely, or unlikely. + std::optional branchLikely; + }; + + std::unordered_map codeAnnotations; + // The effects for this function, if they have been computed. We use a shared // ptr here to avoid compilation errors with the forward-declared // EffectAnalyzer. diff --git a/src/wasm/wasm-ir-builder.cpp b/src/wasm/wasm-ir-builder.cpp index 0928c99c927..269682c5fa0 100644 --- a/src/wasm/wasm-ir-builder.cpp +++ b/src/wasm/wasm-ir-builder.cpp @@ -1393,7 +1393,9 @@ Result<> IRBuilder::makeLoop(Name label, Signature sig) { return visitLoopStart(loop, sig.params); } -Result<> IRBuilder::makeBreak(Index label, bool isConditional) { +Result<> IRBuilder::makeBreak(Index label, + bool isConditional, + std::optional likely) { auto name = getLabelName(label); CHECK_ERR(name); auto labelType = getLabelType(label); @@ -1404,7 +1406,15 @@ Result<> IRBuilder::makeBreak(Index label, bool isConditional) { // Use a dummy condition value if we need to pop a condition. curr.condition = isConditional ? &curr : nullptr; CHECK_ERR(ChildPopper{*this}.visitBreak(&curr, *labelType)); - push(builder.makeBreak(curr.name, curr.value, curr.condition)); + auto* br = builder.makeBreak(curr.name, curr.value, curr.condition); + push(br); + + if (likely) { + // Branches are only possible inside functions. + assert(func); + func->codeAnnotations[br].branchLikely = likely; + } + return Ok{}; } diff --git a/src/wasm/wasm.cpp b/src/wasm/wasm.cpp index 3f885cbccdf..b66a2a7c6d3 100644 --- a/src/wasm/wasm.cpp +++ b/src/wasm/wasm.cpp @@ -16,6 +16,7 @@ #include "wasm.h" #include "ir/branch-utils.h" +#include "wasm-annotations.h" #include "wasm-traversal.h" #include "wasm-type.h" @@ -63,6 +64,12 @@ const char* CustomDescriptorsFeature = "custom-descriptors"; } // namespace BinaryConsts::CustomSections +namespace Annotations { + +const Name BranchHint = "metadata.code.branch_hint"; + +} // namespace Annotations + Name STACK_POINTER("__stack_pointer"); Name MODULE("module"); Name START("start"); diff --git a/test/lit/wat-annotations.wast b/test/lit/wat-annotations.wast new file mode 100644 index 00000000000..6d5be1f0a9d --- /dev/null +++ b/test/lit/wat-annotations.wast @@ -0,0 +1,43 @@ +;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited. + +;; RUN: wasm-opt -all %s -S -o - | filecheck %s + +(module + ;; CHECK: (type $0 (func (param i32))) + + ;; CHECK: (func $branch-hints-br_if (type $0) (param $x i32) + ;; CHECK-NEXT: (block $out + ;; CHECK-NEXT: (@metadata.code.branch_hint "\00") + ;; CHECK-NEXT: (br_if $out + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (@metadata.code.branch_hint "\01") + ;; CHECK-NEXT: (br_if $out + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: (@metadata.code.branch_hint "\00") + ;; CHECK-NEXT: (br_if $out + ;; CHECK-NEXT: (local.get $x) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + ;; CHECK-NEXT: ) + (func $branch-hints-br_if (param $x i32) + (block $out + ;; A branch annotated as unlikely, and one as likely. + (@metadata.code.branch_hint "\00") + (br_if $out + (local.get $x) + ) + (@metadata.code.branch_hint "\01") + (br_if $out + (local.get $x) + ) + ;; The last one wins. + (@metadata.code.branch_hint "\01") + (@metadata.code.branch_hint "\00") + (br_if $out + (local.get $x) + ) + ) + ) +)