Skip to content

Commit 688cd37

Browse files
tlivelyradekdoulik
tlively
authored and
radekdoulik
committed
Preserve multivalue drops in IRBuilder (WebAssembly#6150)
In Binaryen IR, we allow single `Drop` expressions to drop multiple values packaged up as a tuple. When using IRBuilder to rebuild IR containing such a drop, it previously treated the drop as a normal WebAssembly drop that dropped only a single value, producing invalid IR that had extra, undropped values. Fix the problem by preserving the arity of `Drop` inputs in IRBuilder. To avoid bloating the IR, thread the size of the desired value through IRBuilder's pop implementation so that tuple values do not need to be split up and recombined.
1 parent 6842a7f commit 688cd37

File tree

3 files changed

+72
-36
lines changed

3 files changed

+72
-36
lines changed

‎src/wasm-ir-builder.h‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,8 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
200200

201201
// Private functions that must be public for technical reasons.
202202
[[nodiscard]] Result<> visitExpression(Expression*);
203+
[[nodiscard]] Result<>
204+
visitDrop(Drop*, std::optional<uint32_t> arity = std::nullopt);
203205
[[nodiscard]] Result<> visitIf(If*);
204206
[[nodiscard]] Result<> visitReturn(Return*);
205207
[[nodiscard]] Result<> visitStructNew(StructNew*);
@@ -463,7 +465,7 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
463465
[[nodiscard]] Result<Name> getLabelName(Index label);
464466
[[nodiscard]] Result<Name> getDelegateLabelName(Index label);
465467
[[nodiscard]] Result<Index> addScratchLocal(Type);
466-
[[nodiscard]] Result<Expression*> pop();
468+
[[nodiscard]] Result<Expression*> pop(size_t size = 1);
467469

468470
struct HoistedVal {
469471
// The index in the stack of the original value-producing expression.
@@ -478,8 +480,12 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
478480
// Transform the stack as necessary such that the original producer of the
479481
// hoisted value will be popped along with the final expression that produces
480482
// the value, if they are different. May only be called directly after
481-
// hoistLastValue().
482-
[[nodiscard]] Result<> packageHoistedValue(const HoistedVal&);
483+
// hoistLastValue(). `sizeHint` is the size of the type we ultimately want to
484+
// consume, so if the hoisted value has `sizeHint` elements, it is left intact
485+
// even if it is a tuple. Otherwise, hoisted tuple values will be broken into
486+
// pieces.
487+
[[nodiscard]] Result<> packageHoistedValue(const HoistedVal&,
488+
size_t sizeHint = 1);
483489

484490
[[nodiscard]] Result<Expression*> getBranchValue(Name labelName,
485491
std::optional<Index> label);

‎src/wasm/wasm-ir-builder.cpp‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,8 @@ MaybeResult<IRBuilder::HoistedVal> IRBuilder::hoistLastValue() {
9090
return HoistedVal{Index(index), get};
9191
}
9292

93-
Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted) {
93+
Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted,
94+
size_t sizeHint) {
9495
auto& scope = getScope();
9596
assert(!scope.exprStack.empty());
9697

@@ -106,7 +107,7 @@ Result<> IRBuilder::packageHoistedValue(const HoistedVal& hoisted) {
106107

107108
auto type = scope.exprStack.back()->type;
108109

109-
if (!type.isTuple()) {
110+
if (type.size() == sizeHint) {
110111
if (hoisted.get) {
111112
packageAsBlock(type);
112113
}
@@ -158,7 +159,8 @@ void IRBuilder::push(Expression* expr) {
158159
DBG(dump());
159160
}
160161

161-
Result<Expression*> IRBuilder::pop() {
162+
Result<Expression*> IRBuilder::pop(size_t size) {
163+
assert(size >= 1);
162164
auto& scope = getScope();
163165

164166
// Find the suffix of expressions that do not produce values.
@@ -173,11 +175,25 @@ Result<Expression*> IRBuilder::pop() {
173175
return Err{"popping from empty stack"};
174176
}
175177

176-
CHECK_ERR(packageHoistedValue(*hoisted));
178+
CHECK_ERR(packageHoistedValue(*hoisted, size));
177179

178180
auto* ret = scope.exprStack.back();
179-
scope.exprStack.pop_back();
180-
return ret;
181+
if (ret->type.size() == size) {
182+
scope.exprStack.pop_back();
183+
return ret;
184+
}
185+
186+
// The last value-producing expression did not produce exactly the right
187+
// number of values, so we need to construct a tuple piecewise instead.
188+
assert(size > 1);
189+
std::vector<Expression*> elems;
190+
elems.resize(size);
191+
for (int i = size - 1; i >= 0; --i) {
192+
auto elem = pop();
193+
CHECK_ERR(elem);
194+
elems[i] = *elem;
195+
}
196+
return builder.makeTupleMake(elems);
181197
}
182198

183199
Result<Expression*> IRBuilder::build() {
@@ -319,6 +335,20 @@ Result<> IRBuilder::visitExpression(Expression* curr) {
319335
return Ok{};
320336
}
321337

338+
Result<> IRBuilder::visitDrop(Drop* curr, std::optional<uint32_t> arity) {
339+
// Multivalue drops must remain multivalue drops.
340+
if (!arity) {
341+
arity = curr->value->type.size();
342+
}
343+
if (*arity >= 2) {
344+
auto val = pop(*arity);
345+
CHECK_ERR(val);
346+
curr->value = *val;
347+
return Ok{};
348+
}
349+
return visitExpression(curr);
350+
}
351+
322352
Result<> IRBuilder::visitIf(If* curr) {
323353
// Only the condition is popped from the stack. The ifTrue and ifFalse are
324354
// self-contained so we do not modify them.
@@ -1183,7 +1213,7 @@ Result<> IRBuilder::makeSelect(std::optional<Type> type) {
11831213

11841214
Result<> IRBuilder::makeDrop() {
11851215
Drop curr;
1186-
CHECK_ERR(visitDrop(&curr));
1216+
CHECK_ERR(visitDrop(&curr, 1));
11871217
push(builder.makeDrop(curr.value));
11881218
return Ok{};
11891219
}

‎test/lit/passes/outlining.wast‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -729,20 +729,16 @@
729729
)
730730

731731
;; Test outlining works with call_indirect
732-
;; 0 results, 2 params, 3 operands
732+
;; 0 results, 0 params, 1 operand
733733
(module
734734
(table funcref)
735735
;; CHECK: (type $0 (func))
736736

737-
;; CHECK: (type $1 (func (param i32 i32)))
738-
739737
;; CHECK: (table $0 0 funcref)
740738

741739
;; CHECK: (func $outline$ (type $0)
742-
;; CHECK-NEXT: (call_indirect $0 (type $1)
740+
;; CHECK-NEXT: (call_indirect $0 (type $0)
743741
;; CHECK-NEXT: (i32.const 0)
744-
;; CHECK-NEXT: (i32.const 1)
745-
;; CHECK-NEXT: (i32.const 2)
746742
;; CHECK-NEXT: )
747743
;; CHECK-NEXT: )
748744

@@ -752,31 +748,29 @@
752748
;; CHECK-NEXT: )
753749
(func $a
754750
(call_indirect
755-
(param i32 i32)
756751
(i32.const 0)
757-
(i32.const 1)
758-
(i32.const 2)
759752
)
760753
(call_indirect
761-
(param i32 i32)
762754
(i32.const 0)
763-
(i32.const 1)
764-
(i32.const 2)
765755
)
766756
)
767757
)
768758

769759
;; Test outlining works with call_indirect
770-
;; 0 results, 0 params, 1 operand
760+
;; 1 result, 0 params, 1 operand
771761
(module
772762
(table funcref)
773763
;; CHECK: (type $0 (func))
774764

765+
;; CHECK: (type $1 (func (result i32)))
766+
775767
;; CHECK: (table $0 0 funcref)
776768

777769
;; CHECK: (func $outline$ (type $0)
778-
;; CHECK-NEXT: (call_indirect $0 (type $0)
779-
;; CHECK-NEXT: (i32.const 0)
770+
;; CHECK-NEXT: (drop
771+
;; CHECK-NEXT: (call_indirect $0 (type $1)
772+
;; CHECK-NEXT: (i32.const 0)
773+
;; CHECK-NEXT: )
780774
;; CHECK-NEXT: )
781775
;; CHECK-NEXT: )
782776

@@ -785,27 +779,33 @@
785779
;; CHECK-NEXT: (call $outline$)
786780
;; CHECK-NEXT: )
787781
(func $a
788-
(call_indirect
789-
(i32.const 0)
782+
(drop
783+
(call_indirect
784+
(result i32)
785+
(i32.const 0)
786+
)
790787
)
791-
(call_indirect
792-
(i32.const 0)
788+
(drop
789+
(call_indirect
790+
(result i32)
791+
(i32.const 0)
792+
)
793793
)
794794
)
795795
)
796796

797797
;; Test outlining works with call_indirect
798-
;; 1 result, 0 params, 1 operand
798+
;; 2 results, 0 params, 1 operand
799799
(module
800800
(table funcref)
801801
;; CHECK: (type $0 (func))
802802

803-
;; CHECK: (type $1 (func (result i32)))
803+
;; CHECK: (type $1 (func (result i32 i32)))
804804

805805
;; CHECK: (table $0 0 funcref)
806806

807807
;; CHECK: (func $outline$ (type $0)
808-
;; CHECK-NEXT: (drop
808+
;; CHECK-NEXT: (tuple.drop 2
809809
;; CHECK-NEXT: (call_indirect $0 (type $1)
810810
;; CHECK-NEXT: (i32.const 0)
811811
;; CHECK-NEXT: )
@@ -817,15 +817,15 @@
817817
;; CHECK-NEXT: (call $outline$)
818818
;; CHECK-NEXT: )
819819
(func $a
820-
(drop
820+
(tuple.drop 2
821821
(call_indirect
822-
(result i32)
822+
(result i32 i32)
823823
(i32.const 0)
824824
)
825825
)
826-
(drop
826+
(tuple.drop 2
827827
(call_indirect
828-
(result i32)
828+
(result i32 i32)
829829
(i32.const 0)
830830
)
831831
)

0 commit comments

Comments
 (0)