Skip to content

Optimize tuple.extract of gets in BinaryInstWriter #5941

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

Merged
merged 2 commits into from
Sep 14, 2023
Merged

Optimize tuple.extract of gets in BinaryInstWriter #5941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
418379f
Optimize tuple.extract of gets in BinaryInstWriter
tlively Sep 14, 2023
20d855e
address comments
tlively Sep 14, 2023
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
5 changes: 5 additions & 0 deletions src/wasm-stack.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,11 @@ class BinaryInstWriter : public OverriddenVisitor<BinaryInstWriter> {
InsertOrderedMap<Type, Index> scratchLocals;
void countScratchLocals();
void setScratchLocals();

// local.get, local.tee, and glboal.get expressions that will be followed by
// tuple.extracts. We can optimize these by getting only the local for the
// extracted index.
std::unordered_map<Expression*, Index> extractedGets;
};

// Takes binaryen IR and converts it to something else (binary or stack IR)
Expand Down
43 changes: 42 additions & 1 deletion src/wasm/wasm-stack.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,13 @@ void BinaryInstWriter::visitCallIndirect(CallIndirect* curr) {
}

void BinaryInstWriter::visitLocalGet(LocalGet* curr) {
if (auto it = extractedGets.find(curr); it != extractedGets.end()) {
// We have a tuple of locals to get, but we will only end up using one of
// them, so we can just emit that one.
o << int8_t(BinaryConsts::LocalGet)
<< U32LEB(mappedLocals[std::make_pair(curr->index, it->second)]);
return;
}
size_t numValues = func->getLocalType(curr->index).size();
for (Index i = 0; i < numValues; ++i) {
o << int8_t(BinaryConsts::LocalGet)
Expand All @@ -96,14 +103,28 @@ void BinaryInstWriter::visitLocalGet(LocalGet* curr) {

void BinaryInstWriter::visitLocalSet(LocalSet* curr) {
size_t numValues = func->getLocalType(curr->index).size();
// If this is a tuple, set all the elements with nonzero index.
for (Index i = numValues - 1; i >= 1; --i) {
o << int8_t(BinaryConsts::LocalSet)
<< U32LEB(mappedLocals[std::make_pair(curr->index, i)]);
}
if (!curr->isTee()) {
// This is not a tee, so just finish setting the values.
o << int8_t(BinaryConsts::LocalSet)
<< U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
} else if (auto it = extractedGets.find(curr); it != extractedGets.end()) {
// We only need to get the single extracted value.
if (it->second == 0) {
o << int8_t(BinaryConsts::LocalTee)
<< U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
} else {
o << int8_t(BinaryConsts::LocalSet)
<< U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
o << int8_t(BinaryConsts::LocalGet)
<< U32LEB(mappedLocals[std::make_pair(curr->index, it->second)]);
}
} else {
// We need to get all the values.
o << int8_t(BinaryConsts::LocalTee)
<< U32LEB(mappedLocals[std::make_pair(curr->index, 0)]);
for (Index i = 1; i < numValues; ++i) {
Expand All @@ -114,8 +135,14 @@ void BinaryInstWriter::visitLocalSet(LocalSet* curr) {
}

void BinaryInstWriter::visitGlobalGet(GlobalGet* curr) {
// Emit a global.get for each element if this is a tuple global
Index index = parent.getGlobalIndex(curr->name);
if (auto it = extractedGets.find(curr); it != extractedGets.end()) {
// We have a tuple of globals to get, but we will only end up using one of
// them, so we can just emit that one.
o << int8_t(BinaryConsts::GlobalGet) << U32LEB(index + it->second);
return;
}
// Emit a global.get for each element if this is a tuple global
Copy link
Member

Choose a reason for hiding this comment

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

Wait, we allow tuple globals? Does the wasm spec?

Copy link
Member Author

Choose a reason for hiding this comment

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

Yes, we weren't originally going to support tuple globals, but when we first implemented all this tuple business we found that some passes needed them. I don't remember which ones those were, though.

The Wasm spec supports tuple globals only as much as it supports tuple locals, i.e. not at all. But that's ok!

Copy link
Member

Choose a reason for hiding this comment

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

Hmm... a little odd, but ok I guess 😄 I suppose the passes would need to create more globals or such, otherwise.

size_t numValues = curr->type.size();
for (Index i = 0; i < numValues; ++i) {
o << int8_t(BinaryConsts::GlobalGet) << U32LEB(index + i);
Expand Down Expand Up @@ -1970,6 +1997,10 @@ void BinaryInstWriter::visitTupleMake(TupleMake* curr) {
}

void BinaryInstWriter::visitTupleExtract(TupleExtract* curr) {
if (extractedGets.count(curr->tuple)) {
// We already have just the extracted value on the stack.
return;
}
size_t numVals = curr->tuple->type.size();
// Drop all values after the one we want
for (size_t i = curr->index + 1; i < numVals; ++i) {
Expand Down Expand Up @@ -2506,6 +2537,7 @@ void BinaryInstWriter::mapLocalsAndEmitHeader() {
}
}
setScratchLocals();

o << U32LEB(numLocalsByType.size());
for (auto& localType : localTypes) {
o << U32LEB(numLocalsByType.at(localType));
Expand All @@ -2532,6 +2564,15 @@ void BinaryInstWriter::countScratchLocals() {
for (auto& [type, _] : scratchLocals) {
noteLocalType(type);
}
// While we have all the tuple.extracts, also find extracts of local.gets,
// local.tees, and global.gets that we can optimize.
for (auto* extract : extracts.list) {
auto* tuple = extract->tuple;
if (tuple->is<LocalGet>() || tuple->is<LocalSet>() ||
tuple->is<GlobalGet>()) {
extractedGets.insert({tuple, extract->index});
}
}
}

void BinaryInstWriter::setScratchLocals() {
Expand Down
11 changes: 1 addition & 10 deletions test/exception-handling.wast.fromBinary
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
(local $1 i64)
(local $2 (i32 i64))
(local $3 i32)
(local $4 i32)
(try $label$3
(do
(throw $e-i32
Expand Down Expand Up @@ -60,15 +59,7 @@
)
)
(drop
(block (result i32)
(local.set $4
(local.get $x)
)
(drop
(local.get $1)
)
(local.get $4)
)
(local.get $x)
)
)
)
Expand Down
11 changes: 1 addition & 10 deletions test/exception-handling.wast.fromBinary.noDebugInfo
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
(local $1 i64)
(local $2 (i32 i64))
(local $3 i32)
(local $4 i32)
(try $label$3
(do
(throw $tag$0
Expand Down Expand Up @@ -60,15 +59,7 @@
)
)
(drop
(block (result i32)
(local.set $4
(local.get $0)
)
(drop
(local.get $1)
)
(local.get $4)
)
(local.get $0)
)
)
)
Expand Down
86 changes: 4 additions & 82 deletions test/lit/multivalue.wast
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,13 +155,6 @@
;; CHECK-NEXT: (local $5 (i32 i64 f32))
;; CHECK-NEXT: (local $6 i64)
;; CHECK-NEXT: (local $7 i32)
;; CHECK-NEXT: (local $8 i64)
;; CHECK-NEXT: (local $9 i32)
;; CHECK-NEXT: (local $10 i64)
;; CHECK-NEXT: (local $11 i32)
;; CHECK-NEXT: (local $12 i64)
;; CHECK-NEXT: (local $13 i32)
;; CHECK-NEXT: (local $14 f32)
;; CHECK-NEXT: (local.set $5
;; CHECK-NEXT: (call $triple)
;; CHECK-NEXT: )
Expand Down Expand Up @@ -190,69 +183,10 @@
;; CHECK-NEXT: (local.get $7)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (local.set $9
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i64)
;; CHECK-NEXT: (local.set $8
;; CHECK-NEXT: (local.get $1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $4
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $8)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $9)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (tuple.make
;; CHECK-NEXT: (block (result f32)
;; CHECK-NEXT: (local.set $14
;; CHECK-NEXT: (local.get $4)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (local.set $11
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (block (result i64)
;; CHECK-NEXT: (local.set $10
;; CHECK-NEXT: (local.get $1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $10)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $11)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $14)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (local.set $13
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i64)
;; CHECK-NEXT: (local.set $12
;; CHECK-NEXT: (local.get $1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $12)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $13)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $3)
;; CHECK-NEXT: (local.get $1)
;; CHECK-NEXT: (local.get $x)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
(func $reverse (result f32 i64 i32)
Expand Down Expand Up @@ -296,7 +230,6 @@
;; CHECK: (func $global (type $0) (result i32 i64)
;; CHECK-NEXT: (local $0 i64)
;; CHECK-NEXT: (local $1 i32)
;; CHECK-NEXT: (local $2 i32)
;; CHECK-NEXT: (global.set $g1
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (local.set $1
Expand All @@ -309,18 +242,7 @@
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (block (result i32)
;; CHECK-NEXT: (local.set $2
;; CHECK-NEXT: (global.get $g1)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.set $0
;; CHECK-NEXT: (global.get $g2)
;; CHECK-NEXT: )
;; CHECK-NEXT: (local.get $2)
;; CHECK-NEXT: )
;; CHECK-NEXT: )
;; CHECK-NEXT: (drop
;; CHECK-NEXT: (local.get $0)
;; CHECK-NEXT: (global.get $g2)
;; CHECK-NEXT: )
;; CHECK-NEXT: (tuple.make
;; CHECK-NEXT: (global.get $global$2)
Expand Down