Skip to content

[NFC] Refactor code annotations to make using them easier #7599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

[NFC] Refactor code annotations to make using them easier #7599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
122ef79
start
kripken May 15, 2025
8e9035a
format
kripken May 15, 2025
2997134
finish
kripken May 15, 2025
80ee9ba
format
kripken May 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions src/wasm-ir-builder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,21 +116,22 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
// signatures ensure that there are no missing fields.
Result<> makeNop();
Result<> makeBlock(Name label, Signature sig);
Result<>
makeIf(Name label, Signature sig, std::optional<bool> likely = std::nullopt);
Result<> makeIf(Name label,
Signature sig,
CodeAnnotation::BranchLikely likely = std::nullopt);
Result<> makeLoop(Name label, Signature sig);
Result<> makeBreak(Index label,
bool isConditional,
std::optional<bool> likely = std::nullopt);
CodeAnnotation::BranchLikely likely = std::nullopt);
Result<> makeSwitch(const std::vector<Index>& labels, Index defaultLabel);
// Unlike Builder::makeCall, this assumes the function already exists.
Result<> makeCall(Name func,
bool isReturn,
std::optional<std::uint8_t> inline_ = std::nullopt);
CodeAnnotation::Inline inline_ = std::nullopt);
Result<> makeCallIndirect(Name table,
HeapType type,
bool isReturn,
std::optional<std::uint8_t> inline_ = std::nullopt);
CodeAnnotation::Inline inline_ = std::nullopt);
Result<> makeLocalGet(Index local);
Result<> makeLocalSet(Index local);
Result<> makeLocalTee(Index local);
Expand Down Expand Up @@ -206,14 +207,14 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
Result<> makeI31Get(bool signed_);
Result<> makeCallRef(HeapType type,
bool isReturn,
std::optional<std::uint8_t> inline_ = std::nullopt);
CodeAnnotation::Inline inline_ = std::nullopt);
Result<> makeRefTest(Type type);
Result<> makeRefCast(Type type);
Result<> makeBrOn(Index label,
BrOnOp op,
Type in = Type::none,
Type out = Type::none,
std::optional<bool> likely = std::nullopt);
CodeAnnotation::BranchLikely likely = std::nullopt);
Result<> makeStructNew(HeapType type);
Result<> makeStructNewDefault(HeapType type);
Result<>
Expand Down Expand Up @@ -706,10 +707,10 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
void fixLoopWithInput(Loop* loop, Type inputType, Index scratch);

// Add a branch hint, if |likely| is present.
void addBranchHint(Expression* expr, std::optional<bool> likely);
void addBranchHint(Expression* expr, CodeAnnotation::BranchLikely likely);

// Add an inlining hint, if |inline_| is present.
void addInlineHint(Expression* expr, std::optional<std::uint8_t> inline_);
void addInlineHint(Expression* expr, CodeAnnotation::Inline inline_);

void dump();
};
Expand Down
28 changes: 15 additions & 13 deletions src/wasm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2137,6 +2137,21 @@ struct BinaryLocations {
std::unordered_map<Function*, FunctionLocations> functions;
};

// Code annotations for VMs. As with debug info, we do not store these on
// Expressions 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.
using BranchLikely = std::optional<bool>;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think these type aliases are beneficial. Users of BranchLikely still need to know that it is actually a std::optional<bool>, so hiding that fact from them just introduces more friction. If we do want to introduce a new type name here, using a three-state enum is probably the way to go. Same for Inline below.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree users need to know likely is a bool and inline is an int. But I am trying to help in this way:

// ugly "true" - what is it?
makeIf(condition, arm, true);

// comment... slightly better, but easy to forget
makeIf(condition, arm, true /* likely */);

// type - much nicer
makeIf(condition, arm, BranchLikely(true));

(Maybe an enforced type would be even better?)

If we do want to introduce a new type name here, using a three-state enum is probably the way to go. Same for Inline below

For Inline, a 129-state enum..? 😉

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about using enum BranchLikeliness { Unlikely, Likely }; and a wrapper type with an explicit constructor:

struct InlineHint {
  uint8_t value;
  explicit Inline(uint8_t value) : value(value) {}
);

Then for the parameters in IRBuilder use std::optional<BranchLikeliness> and std::optional<InlineHint>. That puts the optionality where it belongs and makes call sites more readable.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The optionality is also needed in CodeAnnotation itself, not just as call params, since each hint is optional when stored (as appearing in the code before this PR). So we could have BranchLikeliness and InlineHint as you suggested - those look good - but we should also have a nice name for the optional of each of those, given the multiple uses in multiple places... and I don't have a good name idea for them?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would just continue using std::optional explicitly wherever they need to be stored optionally. I generally don't think introducing a new name to avoid writing std::optional is a good trade off.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair enough. In that case, I'm not sure this PR is needed. I'll close it for now - maybe when we start using annotations in more places in the code, we'll have further thoughts on how to improve.

BranchLikely branchLikely;

// Compilation Hints proposal.
static const uint8_t NeverInline = 0;
static const uint8_t AlwaysInline = 127;
using Inline = std::optional<uint8_t>;
Inline inline_;
};

// Forward declaration for FuncEffectsMap.
class EffectAnalyzer;

Expand Down Expand Up @@ -2191,19 +2206,6 @@ class Function : public Importable {
delimiterLocations;
BinaryLocations::FunctionLocations funcLocation;

// Code annotations for VMs. As with debug info, we do not store these on
// Expressions 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<bool> branchLikely;

// Compilation Hints proposal.
static const uint8_t NeverInline = 0;
static const uint8_t AlwaysInline = 127;
std::optional<uint8_t> inline_;
};

// Function-level annotations are implemented with a key of nullptr, matching
// the 0 byte offset in the spec.
std::unordered_map<Expression*, CodeAnnotation> codeAnnotations;
Expand Down
70 changes: 31 additions & 39 deletions src/wasm/wasm-binary.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1576,7 +1576,7 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeExpressionHints(
Expression* expr;
// The offset we will write in the custom section.
BinaryLocation offset;
Function::CodeAnnotation* hint;
CodeAnnotation* hint;
};

struct FuncHints {
Expand Down Expand Up @@ -1668,11 +1668,8 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::writeExpressionHints(
std::optional<BufferWithRandomAccess> WasmBinaryWriter::getBranchHintsBuffer() {
return writeExpressionHints(
Annotations::BranchHint,
[](const Function::CodeAnnotation& annotation) {
return annotation.branchLikely;
},
[](const Function::CodeAnnotation& annotation,
BufferWithRandomAccess& buffer) {
[](const CodeAnnotation& annotation) { return annotation.branchLikely; },
[](const CodeAnnotation& annotation, BufferWithRandomAccess& buffer) {
// Hint size, always 1 for now.
buffer << U32LEB(1);

Expand All @@ -1687,11 +1684,8 @@ std::optional<BufferWithRandomAccess> WasmBinaryWriter::getBranchHintsBuffer() {
std::optional<BufferWithRandomAccess> WasmBinaryWriter::getInlineHintsBuffer() {
return writeExpressionHints(
Annotations::InlineHint,
[](const Function::CodeAnnotation& annotation) {
return annotation.inline_;
},
[](const Function::CodeAnnotation& annotation,
BufferWithRandomAccess& buffer) {
[](const CodeAnnotation& annotation) { return annotation.inline_; },
[](const CodeAnnotation& annotation, BufferWithRandomAccess& buffer) {
// Hint size, always 1 for now.
buffer << U32LEB(1);

Expand Down Expand Up @@ -5323,39 +5317,37 @@ void WasmBinaryReader::readExpressionHints(Name sectionName,
}

void WasmBinaryReader::readBranchHints(size_t payloadLen) {
readExpressionHints(Annotations::BranchHint,
payloadLen,
[&](Function::CodeAnnotation& annotation) {
auto size = getU32LEB();
if (size != 1) {
throwError("bad BranchHint size");
}
readExpressionHints(
Annotations::BranchHint, payloadLen, [&](CodeAnnotation& annotation) {
auto size = getU32LEB();
if (size != 1) {
throwError("bad BranchHint size");
}

auto likely = getU32LEB();
if (likely != 0 && likely != 1) {
throwError("bad BranchHint value");
}
auto likely = getU32LEB();
if (likely != 0 && likely != 1) {
throwError("bad BranchHint value");
}

annotation.branchLikely = likely;
});
annotation.branchLikely = likely;
});
}

void WasmBinaryReader::readInlineHints(size_t payloadLen) {
readExpressionHints(Annotations::InlineHint,
payloadLen,
[&](Function::CodeAnnotation& annotation) {
auto size = getU32LEB();
if (size != 1) {
throwError("bad InlineHint size");
}

uint8_t inline_ = getInt8();
if (inline_ > 127) {
throwError("bad InlineHint value");
}

annotation.inline_ = inline_;
});
readExpressionHints(
Annotations::InlineHint, payloadLen, [&](CodeAnnotation& annotation) {
auto size = getU32LEB();
if (size != 1) {
throwError("bad InlineHint size");
}

uint8_t inline_ = getInt8();
if (inline_ > 127) {
throwError("bad InlineHint value");
}

annotation.inline_ = inline_;
});
}

Index WasmBinaryReader::readMemoryAccess(Address& alignment, Address& offset) {
Expand Down
28 changes: 16 additions & 12 deletions src/wasm/wasm-ir-builder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1380,8 +1380,9 @@ Result<> IRBuilder::makeBlock(Name label, Signature sig) {
return visitBlockStart(block, sig.params);
}

Result<>
IRBuilder::makeIf(Name label, Signature sig, std::optional<bool> likely) {
Result<> IRBuilder::makeIf(Name label,
Signature sig,
CodeAnnotation::BranchLikely likely) {
auto* iff = wasm.allocator.alloc<If>();
iff->type = sig.results;
addBranchHint(iff, likely);
Expand All @@ -1397,7 +1398,7 @@ Result<> IRBuilder::makeLoop(Name label, Signature sig) {

Result<> IRBuilder::makeBreak(Index label,
bool isConditional,
std::optional<bool> likely) {
CodeAnnotation::BranchLikely likely) {
auto name = getLabelName(label);
CHECK_ERR(name);
auto labelType = getLabelType(label);
Expand Down Expand Up @@ -1441,9 +1442,8 @@ Result<> IRBuilder::makeSwitch(const std::vector<Index>& labels,
return Ok{};
}

Result<> IRBuilder::makeCall(Name func,
bool isReturn,
std::optional<std::uint8_t> inline_) {
Result<>
IRBuilder::makeCall(Name func, bool isReturn, CodeAnnotation::Inline inline_) {
auto sig = wasm.getFunction(func)->getSig();
Call curr(wasm.allocator);
curr.target = func;
Expand All @@ -1459,7 +1459,7 @@ Result<> IRBuilder::makeCall(Name func,
Result<> IRBuilder::makeCallIndirect(Name table,
HeapType type,
bool isReturn,
std::optional<std::uint8_t> inline_) {
CodeAnnotation::Inline inline_) {
CallIndirect curr(wasm.allocator);
curr.heapType = type;
curr.operands.resize(type.getSignature().params.size());
Expand Down Expand Up @@ -1959,7 +1959,7 @@ Result<> IRBuilder::makeI31Get(bool signed_) {

Result<> IRBuilder::makeCallRef(HeapType type,
bool isReturn,
std::optional<std::uint8_t> inline_) {
CodeAnnotation::Inline inline_) {
CallRef curr(wasm.allocator);
if (!type.isSignature()) {
return Err{"expected function type"};
Expand Down Expand Up @@ -1991,8 +1991,11 @@ Result<> IRBuilder::makeRefCast(Type type) {
return Ok{};
}

Result<> IRBuilder::makeBrOn(
Index label, BrOnOp op, Type in, Type out, std::optional<bool> likely) {
Result<> IRBuilder::makeBrOn(Index label,
BrOnOp op,
Type in,
Type out,
CodeAnnotation::BranchLikely likely) {
BrOn curr;
curr.op = op;
curr.castType = out;
Expand Down Expand Up @@ -2548,7 +2551,8 @@ Result<> IRBuilder::makeStackSwitch(HeapType ct, Name tag) {
return Ok{};
}

void IRBuilder::addBranchHint(Expression* expr, std::optional<bool> likely) {
void IRBuilder::addBranchHint(Expression* expr,
CodeAnnotation::BranchLikely likely) {
if (likely) {
// Branches are only possible inside functions.
assert(func);
Expand All @@ -2557,7 +2561,7 @@ void IRBuilder::addBranchHint(Expression* expr, std::optional<bool> likely) {
}

void IRBuilder::addInlineHint(Expression* expr,
std::optional<uint8_t> inline_) {
CodeAnnotation::Inline inline_) {
if (inline_) {
// Branches are only possible inside functions.
assert(func);
Expand Down
Loading