Repository: tetratelabs/wazero
Pull request #1407 (2023-04-26 15:01:16)

api: adds CallWithStack to avoid allocations
This implements the API discussed in #1332.

It still needs unit tests, and redoing the benchmarks to confirm the improvements. 

ncruces: Re-ran the benchmarks from #1332 ([`github.com/maragudk/sqlite-benchmark`](https://github.com/maragudk/sqlite-benchmark) adapted to use my [driver](https://pkg.go.dev/github.com/ncruces/go-sqlite3/driver) rather than [mattn](https://pkg.go.dev/github.com/mattn/go-sqlite3)).

Again this benchmark is I/O bound, so improvements to CPU time are negligible, and there's plenty of non-determinism. OTOH, it's not a synthetic made to make this look great.

The reduction in `allocs/op` is stark (this is a rough proxy for the amount of GC work, not necessarily peak memory usage).

Before:
```
goos: darwin
goarch: amd64
pkg: github.com/maragudk/sqlite-benchmark
cpu: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
BenchmarkSelect1-12                                         2000              8643 ns/op           28301 B/op         13 allocs/op
BenchmarkDB_ReadPost/mutex_false-12                         2000             30699 ns/op           26302 B/op         76 allocs/op
BenchmarkDB_ReadPost/mutex_true-12                          2000             29405 ns/op           28953 B/op         77 allocs/op
BenchmarkDB_WritePost/mutex_false-12                        2000            982031 ns/op           20309 B/op         37 allocs/op
BenchmarkDB_WritePost/mutex_true-12                         2000            767473 ns/op            7634 B/op         36 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_0.01-12                   2000             43738 ns/op           24638 B/op        218 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_0.1-12                    2000            195674 ns/op           31664 B/op       1476 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_1-12                      2000          11684738 ns/op          230575 B/op      15605 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_0.01-12                    2000             62366 ns/op           96816 B/op        220 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_0.1-12                     2000            301926 ns/op          627289 B/op       1633 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_1-12                       2000           2598808 ns/op          225410 B/op      15603 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_0.01-12                       2000             42212 ns/op           36322 B/op        183 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_0.1-12                        2000            202609 ns/op          102164 B/op       1727 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_1-12                          2000           2937426 ns/op          225967 B/op      15606 allocs/op
PASS
ok      github.com/maragudk/sqlite-benchmark    40.663s
```

After:
```
goos: darwin
goarch: amd64
pkg: github.com/maragudk/sqlite-benchmark
cpu: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
BenchmarkSelect1-12                                         2000             10291 ns/op           23294 B/op          8 allocs/op
BenchmarkDB_ReadPost/mutex_false-12                         2000             29984 ns/op           26373 B/op         47 allocs/op
BenchmarkDB_ReadPost/mutex_true-12                          2000             33873 ns/op           32017 B/op         49 allocs/op
BenchmarkDB_WritePost/mutex_false-12                        2000            973562 ns/op           24745 B/op         24 allocs/op
BenchmarkDB_WritePost/mutex_true-12                         2000            794107 ns/op            7523 B/op         23 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_0.01-12                   2000             38152 ns/op           29614 B/op         74 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_0.1-12                    2000            320714 ns/op           22022 B/op        509 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_false/comment_rate_1-12                      2000          10405274 ns/op          159338 B/op       6665 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_0.01-12                    2000             68912 ns/op          100010 B/op        117 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_0.1-12                     2000            315064 ns/op          617469 B/op        758 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteComment/mutex_true/comment_rate_1-12                       2000           2373769 ns/op          154330 B/op       6664 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_0.01-12                       2000             53201 ns/op           46634 B/op        123 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_0.1-12                        2000            192710 ns/op           96767 B/op        662 allocs/op
BenchmarkDB_ReadPostAndMaybeWriteCommentWithPool/comment_rate_1-12                          2000           2552369 ns/op          155078 B/op       6666 allocs/op
PASS
ok      github.com/maragudk/sqlite-benchmark    37.593s
```
ncruces: @inkeliz, since you wanted this, I'd really appreciate your feedback on whether this solves your issue, and what's the more ergonomic API.

I don't want to get too much caught up into my own solution.

Thanks!
ncruces: I think in #1332 and here we've established the usefulness of this, but can't settle for an API design (which in this case, I find, is understandably hard).

I'm happy to keep this updated (tests passing) until more people are ready to share instances where they'd use this, so we can improve ergonomics and settle on a design.

Or someone makes the call, because the perfect is the enemy of the good.

I'd use this in exactly one call site, so I'm probably not the right one to make the call, any design works for me.
codefromthecrypt: My summary comment is that the api surface of a codebase should have coherence with itself as a feature, even if it wasn't also optimizable due to this.

It would almost be self-deprecating humor to have to tell people the secret way to make our API whose goal is to reduce allocations, do less via things like this: `CallWithStack(ctx, params, params[:0])` when we already have an API design the same people will have to understand to reduce allocations host side.

When designing something explicitly for advanced users, we shouldn't fear them being able to do min(len(params), len(results)) if trying to size things. The exact concept is required for the host side, even if they aren't the ones sizing the slice.

All this said, If others feel like they really want the third API style and that's better than two, go for it. I've said my piece and feel folks here can grok where I'm coming from.

codefromthecrypt: If we do go for api style 3, please don't link it to GoModuleFunc, as it no longer has any relationship to it.
anuraaga: I happened to find this and #1332 - as a user with go-re2, something like this would definitely be useful. For example, with WAF, hundreds of regex are executed per request, which is hundreds of wasm calls. Reducing allocations would be helpful if possible, though given everything else most logic that uses big regex does, it wouldn't be a greatly sad event if nothing was added as it's still within margin of error in the big scheme of things I guess.

I see mention of "api style 3" but wasn't actually able to find 3 options so not sure which it is, but personally I find adapting `GoModuleFunc` to be very natural, I assume that option looks like `func (f Module) CallWithStack(ctx context.Context, stack []uint64) error`. If so, I would use it in go-re2 like this

https://github.com/wasilibs/go-re2/blob/main/internal/re2_wazero.go#L212

becomes

```go
func match(re *Regexp, s cString, matchesPtr uintptr, nMatches uint32) bool {
	ctx := context.Background()
        stack := [8]uint64{uint64(re.ptr), uint64(s.ptr), uint64(s.length), 0, uint64(s.length), 0, uint64(matchesPtr), uint64(nMatches)}
	if err := re.abi.cre2Match.CallWithStack(ctx, stack[:]); err != nil {
		panic(err)
	}

	return stack[0] == 1
}
```

For a function that takes no parameters and returns a result, `var stack [1]uint64`.

I really enjoy the fact that the result slice isn't returned as it makes error handling simpler - for me, `max(len(params), len(results))` to size things would not be a problem at all especially if it gives me this code benefit. It's a very minor point though.

I'm sure any of the other options would work fine too, to quote @ncruces "any design works for me". Thought it might be helpful to have another code example of where this API would end up being used.
codefromthecrypt: @anuraaga thanks for the feedback! I should have linked api style 3 is from this thread, and worth a look in case changes your opinion https://github.com/tetratelabs/wazero/pull/1407#discussion_r1178069476
anuraaga: @codefromthecrypt Ah thanks for the link, so that's style 3. I'd still use it, but it would be a bit annoying to have to add `var out [1]uint64` on every invocation. I suppose lines-of-code could be returned by inlining the array declaration `[8]uint64{a,b,c,d,e...}[:]` but I'd probably avoid it since declaring an array and slicing it in the same code looks confusing to me. So for go-re2, it's probably not optimal, but I do see mention of some use cases like invoking in a loop where it may be so understand the motivation. go-re2 is just one example so meant as a data point more than anything else.
inkeliz: > but it would be a bit annoying to have to add var out [1]uint64 on every invocation

I think that is the "pain-point" of the 3rd design. You don't _need_ `var out [1]uint64`, but that is not clear, but could be documented. I guess "Style 3" can be used in three (or more) ways:

Option A:

```go
var out [1]uint64
var param [8]uint64{a,b,c,d,e...}

CallWithStack(ctx, param[:], out[:])
```

Option B:

```go
var param [9]uint64{a,b,c,d,e...}
CallWithStack(ctx, param[:8], out[8:])

out := param[8] // output is the last item, not overwrites
```

Option C:

```go
var param [8]uint64{a,b,c,d,e...}

CallWithStack(ctx, param[:], param[:0])

out := param[0] // output is the first, and overwrite the first param
```

In all cases, we are potentially reducing allocations. Currently, this PR forces you to use the "Option C" (which overwrites), while the old CallTo (or the "Style 3") makes it more flexible. Is that flexibility worth it? I'm not sure. _The only specific case is when you do multiple calls with the same arguments_.

The `[:0]` is not so weird, that is used in some std-lib packages, such as `Seal` on AEAD (https://pkg.go.dev/crypto/cipher#AEAD).
anuraaga: Ah got it - yeah I agree option B there would have similar writability, it didn't come to mind from the API definition but with examples it would make sense. Thanks for clarifying it 🙏
ncruces: Thanks, @inkeliz, @anuraaga, this discussion is super helpful. As I expected I'm pigeonholed into my own usage, where (for other reasons) I concentrated all my `api.Function.Call`s into a wrapper.

I agree that an API where each call site can use stack allocated arrays (this requires constant size) to make the call is probably significantly more generally useful, _if these don't escape to heap_, in other words, _if we can make it work_.

Unfortunately, I don't think it's that easy: `CallWithStack`[^1] is an interface call, so its arguments will escape to heap, regardless. Right?

[^1]: whatever the name, with this API `CallTo` is more appropriate
anuraaga: I'm not particularly authoritative but my understanding is while an interfere implies any struct implementing it that is used as an interface has it's contents on the heap, function calls themselves are not related. Definitely would expect being able to keep the arguments on the stack for this API so hope it's possible!
achille-roussel: @ncruces is right that since it's an interface method, Go cannot currently determine that the values do not escape and has to put them on the heap.

I believe the goal of this PR is to offer a way to reuse the slices, not to allow stack-allocation; for this, we would need to use a concrete type allowing the compiler to track the lifetime of method arguments.
achille-roussel: @ncruces @codefromthecrypt I buy the argument of API consistency, I didn't see the relation with `api.GoModuleFunction` at first but keeping it coherent makes most sense 👍 Thanks for spending the time to share your perspectives!

Let's go with this?
```go
CallWithStack(ctx context.Context, stack []uint64) error
```
ncruces: @anuraaga please see [this](https://go.dev/play/p/sQEJqww59ug).

Building it with go 1.20.3:
```sh
go build -gcflags=-m buh.go   
# command-line-arguments
./main.go:9:6: can inline (*impl).Call
./main.go:13:6: can inline directCall
./main.go:16:8: inlining call to (*impl).Call
./main.go:19:6: can inline ifaceCall
./main.go:27:12: inlining call to directCall
./main.go:28:11: inlining call to ifaceCall
./main.go:27:12: inlining call to (*impl).Call
./main.go:28:11: devirtualizing v.Call to *impl
./main.go:9:19: in does not escape
./main.go:9:32: out does not escape
./main.go:13:17: v does not escape
./main.go:19:16: leaking param: v
./main.go:20:6: moved to heap: in
./main.go:21:6: moved to heap: out
```

On an interface call all arguments escape to heap. I'd really like to support this, but can't see how.

The idea of this PR is not to prevent them from escaping, they always do. Rather it is to allow you to reuse an heap allocated slice you have lying around, somewhere. If concurrent calls into a module are not OK, you can keep a single slice per module for all calls, otherwise, probably a `sync.Pool` is better.
ncruces: @achille-roussel, thanks! Slices are subtle, though.

At entry of `CallWithStack` in https://github.com/tetratelabs/wazero/pull/1407#issuecomment-1526003485, do we assume:
1. `len(stack) == len(args) && cap(stack) >= len(res)`
2. `len(stack) == max(len(args), len(res))`
3. `len(stack) >= max(len(args), len(res))`
4. …?

I get the parallel with `GoModuleFunction`, and agree consistency reduces cognitive load, so thought experiment…

Would it be fair to consider a `GoModuleFunc` that uses `CallWithStack` to forward a "host call" to a "guest call" to either the same or another module?

If so, if that's the "consistency" we want, according to the [docs](https://pkg.go.dev/github.com/tetratelabs/wazero/api#GoModuleFunction), we want option **2** (or maybe **3**) above.

If everyone approves of this I'll make the change.
cc: @codefromthecrypt
codefromthecrypt: I think we can change the comment on GoModuleFunc so that it says "at least" not "equals" max blah blah https://github.com/tetratelabs/wazero/pull/1407#issuecomment-1526124384. Then we can be consistent on this point between the two sides of the API.


anuraaga: > On an interface call all arguments escape to heap. I'd really like to support this, but can't see how.

Thanks for sharing this! Learned something new. It's too bad that stack allocation can't work but it is what it is.

Luckily, currently wasm calls are guarded by a lock so it'll be easy to reuse a heap slice. Eventually when atomics are supported I'm looking forward to removing that lock and per-regex module instantiation and then I'd probably have to see if sync.Pool performs better than the small allocation. But that'll be a long ways away.
codefromthecrypt: > Not necessarily the best example. A v_v function can use Call without allocations: the params slice is nil, when varargs are empty; we take care never to allocate for empty results (and this is actually tested).

good point and this is why you should think about my suggestions before implementing them ;)
ncruces: > Can you parameterize BenchmarkHostFunctionCall so we can see any overall diff before/after?

```sh
% go test -bench=BenchmarkHostFunctionCall -benchmem
goos: darwin
goarch: amd64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
cpu: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
BenchmarkHostFunctionCall/call_go_reflect_host-12                         515949              2197 ns/op             160 B/op          8 allocs/op
BenchmarkHostFunctionCall/call_go_reflect_host_with_stack-12              529288              2200 ns/op             144 B/op          6 allocs/op
BenchmarkHostFunctionCall/call_go_host-12                                9936541               120.9 ns/op            16 B/op          2 allocs/op
BenchmarkHostFunctionCall/call_go_host_with_stack-12                    13865826                85.96 ns/op            0 B/op          0 allocs/op
PASS
ok      github.com/tetratelabs/wazero/internal/integration_test/bench   5.955s
```

Patch:
diff --git a/api/wasm.go b/api/wasm.go
index ffab8928..38728b42 100644
--- a/api/wasm.go
+++ b/api/wasm.go
@@ -366,6 +366,34 @@ type Function interface {
 	// the end-to-end demonstrations of how these terminations can be performed.
 	Call(ctx context.Context, params ...uint64) ([]uint64, error)
 
+	// CallWithStack is an optimized variation of Call that saves memory
+	// allocations when the stack slice is reused across calls.
+	//
+	// Stack length must be at least the max of parameter or result length.
+	// The caller adds parameters in order to the stack, and reads any results
+	// in order from the stack, except in the error case.
+	//
+	// For example, the following reuses the same stack slice to call searchFn
+	// repeatedly saving one allocation per iteration:
+	//
+	//	stack := make([]uint64, 4)
+	//	for i, search := range searchParams {
+	//		// copy the next params to the stack
+	//		copy(stack, search)
+	//		if err := searchFn.CallWithStack(ctx, stack); err != nil {
+	//			return err
+	//		} else if stack[0] == 1 { // found
+	//			return i // searchParams[i] matched!
+	//		}
+	//	}
+	//
+	// # Notes
+	//
+	//   - This is similar to GoModuleFunction, except for using calling functions
+	//     instead of implementing them. Moreover, this is used regardless of
+	//     whether the callee is a host or wasm defined function.
+	CallWithStack(ctx context.Context, stack []uint64) error
+
 	internalapi.WazeroOnly
 }
 
diff --git a/internal/emscripten/emscripten.go b/internal/emscripten/emscripten.go
index b014d09c..6cda874c 100644
--- a/internal/emscripten/emscripten.go
+++ b/internal/emscripten/emscripten.go
@@ -76,8 +76,11 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 		panic(err)
 	}
 
+	// This needs copy (not reslice) because the stack is reused for results.
+	// Consider invoke_i (zero arguments, one result): index zero (tableOffset)
+	// is needed to store the result.
 	tableOffset := wasm.Index(stack[0]) // position in the module's only table.
-	params := stack[1:]                 // parameters to the dynamic function being called
+	copy(stack, stack[1:])              // pop the tableOffset.
 
 	// Lookup the table index we will call.
 	t := m.Tables[0] // Note: Emscripten doesn't use multiple tables
@@ -86,10 +89,8 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 		panic(err)
 	}
 
-	ret, err := m.Engine.NewFunction(idx).Call(ctx, params...)
+	err = m.Engine.NewFunction(idx).CallWithStack(ctx, stack)
 	if err != nil {
 		panic(err)
 	}
-	// if there are any results, copy them back to the stack
-	copy(stack, ret)
 }
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 24da7889..b6f1104a 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -705,6 +705,24 @@ func (ce *callEngine) Definition() api.FunctionDefinition {
 
 // Call implements the same method as documented on wasm.ModuleEngine.
 func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uint64, err error) {
+	ft := ce.initialFn.funcType
+	if n := ft.ParamNumInUint64; n != len(params) {
+		return nil, fmt.Errorf("expected %d params, but passed %d", n, len(params))
+	}
+	return ce.call(ctx, params, nil)
+}
+
+// CallWithStack implements the same method as documented on wasm.ModuleEngine.
+func (ce *callEngine) CallWithStack(ctx context.Context, stack []uint64) error {
+	params, results, err := wasm.SplitCallStack(ce.initialFn.funcType, stack)
+	if err != nil {
+		return err
+	}
+	_, err = ce.call(ctx, params, results)
+	return err
+}
+
+func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []uint64, err error) {
 	m := ce.initialFn.moduleInstance
 	if ce.ensureTermination {
 		select {
@@ -717,13 +735,6 @@ func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uin
 		}
 	}
 
-	tp := ce.initialFn.funcType
-
-	paramCount := len(params)
-	if tp.ParamNumInUint64 != paramCount {
-		return nil, fmt.Errorf("expected %d params, but passed %d", ce.initialFn.funcType.ParamNumInUint64, paramCount)
-	}
-
 	// We ensure that this Call method never panics as
 	// this Call method is indirectly invoked by embedders via store.CallFunction,
 	// and we have to make sure that all the runtime errors, including the one happening inside
@@ -736,7 +747,8 @@ func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uin
 		}
 	}()
 
-	ce.initializeStack(tp, params)
+	ft := ce.initialFn.funcType
+	ce.initializeStack(ft, params)
 
 	if ce.ensureTermination {
 		done := m.CloseModuleOnCanceledOrTimeout(ctx)
@@ -747,12 +759,12 @@ func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uin
 
 	// This returns a safe copy of the results, instead of a slice view. If we
 	// returned a re-slice, the caller could accidentally or purposefully
-	// corrupt the stack of subsequent calls
-	if resultCount := tp.ResultNumInUint64; resultCount > 0 {
-		results = make([]uint64, resultCount)
-		copy(results, ce.stack[:resultCount])
+	// corrupt the stack of subsequent calls.
+	if results == nil && ft.ResultNumInUint64 > 0 {
+		results = make([]uint64, ft.ResultNumInUint64)
 	}
-	return
+	copy(results, ce.stack)
+	return results, nil
 }
 
 // initializeStack initializes callEngine.stack before entering native code.
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index ac8e0553..d10f2571 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -116,6 +116,10 @@ func (ce *callEngine) pushValue(v uint64) {
 	ce.stack = append(ce.stack, v)
 }
 
+func (ce *callEngine) pushValues(v []uint64) {
+	ce.stack = append(ce.stack, v...)
+}
+
 func (ce *callEngine) popValue() (v uint64) {
 	// No need to check stack bound
 	// as we can assume that all the operations
@@ -129,6 +133,12 @@ func (ce *callEngine) popValue() (v uint64) {
 	return
 }
 
+func (ce *callEngine) popValues(v []uint64) {
+	stackTopIndex := len(ce.stack) - len(v)
+	copy(v, ce.stack[stackTopIndex:])
+	ce.stack = ce.stack[:stackTopIndex]
+}
+
 // peekValues peeks api.ValueType values from the stack and returns them.
 func (ce *callEngine) peekValues(count int) []uint64 {
 	if count == 0 {
@@ -445,10 +455,24 @@ func (ce *callEngine) Definition() api.FunctionDefinition {
 
 // Call implements the same method as documented on api.Function.
 func (ce *callEngine) Call(ctx context.Context, params ...uint64) (results []uint64, err error) {
-	return ce.call(ctx, ce.compiled, params)
+	ft := ce.compiled.funcType
+	if n := ft.ParamNumInUint64; n != len(params) {
+		return nil, fmt.Errorf("expected %d params, but passed %d", n, len(params))
+	}
+	return ce.call(ctx, params, nil)
+}
+
+// CallWithStack implements the same method as documented on api.Function.
+func (ce *callEngine) CallWithStack(ctx context.Context, stack []uint64) error {
+	params, results, err := wasm.SplitCallStack(ce.compiled.funcType, stack)
+	if err != nil {
+		return err
+	}
+	_, err = ce.call(ctx, params, results)
+	return err
 }
 
-func (ce *callEngine) call(ctx context.Context, tf *function, params []uint64) (results []uint64, err error) {
+func (ce *callEngine) call(ctx context.Context, params, results []uint64) (_ []uint64, err error) {
 	m := ce.compiled.moduleInstance
 	if ce.compiled.parent.ensureTermination {
 		select {
@@ -461,13 +485,6 @@ func (ce *callEngine) call(ctx context.Context, tf *function, params []uint64) (
 		}
 	}
 
-	ft := tf.funcType
-	paramSignature := ft.ParamNumInUint64
-	paramCount := len(params)
-	if paramSignature != paramCount {
-		return nil, fmt.Errorf("expected %d params, but passed %d", paramSignature, paramCount)
-	}
-
 	defer func() {
 		// If the module closed during the call, and the call didn't err for another reason, set an ExitError.
 		if err == nil {
@@ -480,22 +497,24 @@ func (ce *callEngine) call(ctx context.Context, tf *function, params []uint64) (
 		}
 	}()
 
-	for _, param := range params {
-		ce.pushValue(param)
-	}
+	ce.pushValues(params)
 
 	if ce.compiled.parent.ensureTermination {
 		done := m.CloseModuleOnCanceledOrTimeout(ctx)
 		defer done()
 	}
 
-	ce.callFunction(ctx, m, tf)
+	ce.callFunction(ctx, m, ce.compiled)
 
 	// This returns a safe copy of the results, instead of a slice view. If we
 	// returned a re-slice, the caller could accidentally or purposefully
 	// corrupt the stack of subsequent calls.
-	results = wasm.PopValues(ft.ResultNumInUint64, ce.popValue)
-	return
+	ft := ce.compiled.funcType
+	if results == nil && ft.ResultNumInUint64 > 0 {
+		results = make([]uint64, ft.ResultNumInUint64)
+	}
+	ce.popValues(results)
+	return results, nil
 }
 
 // recoverOnCall takes the recovered value `recoverOnCall`, and wraps it
diff --git a/internal/wasm/gofunc.go b/internal/wasm/gofunc.go
index 6f77a831..9510c258 100644
--- a/internal/wasm/gofunc.go
+++ b/internal/wasm/gofunc.go
@@ -79,26 +79,6 @@ func (f *reflectGoFunction) Call(ctx context.Context, stack []uint64) {
 	callGoFunc(ctx, nil, f.fn, stack)
 }
 
-// PopValues pops the specified number of api.ValueType parameters off the
-// stack into a parameter slice for use in api.GoFunction or api.GoModuleFunction.
-//
-// For example, if the host function F requires the (x1 uint32, x2 float32)
-// parameters, and the stack is [..., A, B], then the function is called as
-// F(A, B) where A and B are interpreted as uint32 and float32 respectively.
-//
-// Note: the popper intentionally doesn't return bool or error because the
-// caller's stack depth is trusted.
-func PopValues(count int, popper func() uint64) []uint64 {
-	if count == 0 {
-		return nil
-	}
-	params := make([]uint64, count)
-	for i := count - 1; i >= 0; i-- {
-		params[i] = popper()
-	}
-	return params
-}
-
 // callGoFunc executes the reflective function by converting params to Go
 // types. The results of the function call are converted back to api.ValueType.
 func callGoFunc(ctx context.Context, mod api.Module, fn *reflect.Value, stack []uint64) {


Test patch:
diff --git a/internal/engine/compiler/engine_test.go b/internal/engine/compiler/engine_test.go
index b61442b0..e624ebe9 100644
--- a/internal/engine/compiler/engine_test.go
+++ b/internal/engine/compiler/engine_test.go
@@ -67,6 +67,16 @@ func TestCompiler_ModuleEngine_Call(t *testing.T) {
 `, "\n"+functionLog.String())
 }
 
+func TestCompiler_ModuleEngine_CallWithStack(t *testing.T) {
+	defer functionLog.Reset()
+	requireSupportedOSArch(t)
+	enginetest.RunTestModuleEngineCallWithStack(t, et)
+	require.Equal(t, `
+--> .$0(1,2)
+<-- (1,2)
+`, "\n"+functionLog.String())
+}
+
 func TestCompiler_ModuleEngine_Call_HostFn(t *testing.T) {
 	defer functionLog.Reset()
 	requireSupportedOSArch(t)
diff --git a/internal/engine/interpreter/interpreter_test.go b/internal/engine/interpreter/interpreter_test.go
index 65723e64..7f03b9e0 100644
--- a/internal/engine/interpreter/interpreter_test.go
+++ b/internal/engine/interpreter/interpreter_test.go
@@ -105,6 +105,15 @@ func TestInterpreter_ModuleEngine_Call(t *testing.T) {
 `, "\n"+functionLog.String())
 }
 
+func TestCompiler_ModuleEngine_CallWithStack(t *testing.T) {
+	defer functionLog.Reset()
+	enginetest.RunTestModuleEngineCallWithStack(t, et)
+	require.Equal(t, `
+--> .$0(1,2)
+<-- (1,2)
+`, "\n"+functionLog.String())
+}
+
 func TestInterpreter_ModuleEngine_Call_HostFn(t *testing.T) {
 	defer functionLog.Reset()
 	enginetest.RunTestModuleEngineCallHostFn(t, et)
diff --git a/internal/integration_test/bench/hostfunc_bench_test.go b/internal/integration_test/bench/hostfunc_bench_test.go
index fcd5f795..e69c4bad 100644
--- a/internal/integration_test/bench/hostfunc_bench_test.go
+++ b/internal/integration_test/bench/hostfunc_bench_test.go
@@ -58,6 +58,23 @@ func BenchmarkHostFunctionCall(b *testing.B) {
 				}
 			}
 		})
+
+		b.Run(fn+"_with_stack", func(b *testing.B) {
+			ce := getCallEngine(m, fn)
+
+			b.ResetTimer()
+			stack := make([]uint64, 1)
+			for i := 0; i < b.N; i++ {
+				stack[0] = offset
+				err := ce.CallWithStack(testCtx, stack)
+				if err != nil {
+					b.Fatal(err)
+				}
+				if uint32(stack[0]) != math.Float32bits(val) {
+					b.Fail()
+				}
+			}
+		})
 	}
 }
 
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index b3300bf6..d443bb37 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -191,6 +191,58 @@ func RunTestModuleEngineCall(t *testing.T, et EngineTester) {
 	})
 }
 
+func RunTestModuleEngineCallWithStack(t *testing.T, et EngineTester) {
+	e := et.NewEngine(api.CoreFeaturesV2)
+
+	// Define a basic function which defines two parameters and two results.
+	// This is used to test results when incorrect arity is used.
+	m := &wasm.Module{
+		TypeSection: []wasm.FunctionType{
+			{
+				Params:            []wasm.ValueType{i64, i64},
+				Results:           []wasm.ValueType{i64, i64},
+				ParamNumInUint64:  2,
+				ResultNumInUint64: 2,
+			},
+		},
+		FunctionSection: []wasm.Index{0},
+		CodeSection: []wasm.Code{
+			{Body: []byte{wasm.OpcodeLocalGet, 0, wasm.OpcodeLocalGet, 1, wasm.OpcodeEnd}},
+		},
+	}
+
+	m.BuildFunctionDefinitions()
+	listeners := buildListeners(et.ListenerFactory(), m)
+	err := e.CompileModule(testCtx, m, listeners, false)
+	require.NoError(t, err)
+
+	// To use the function, we first need to add it to a module.
+	module := &wasm.ModuleInstance{
+		ModuleName: t.Name(), TypeIDs: []wasm.FunctionTypeID{0},
+		Definitions: m.FunctionDefinitionSection,
+	}
+
+	// Compile the module
+	me, err := e.NewModuleEngine(m, module)
+	require.NoError(t, err)
+	linkModuleToEngine(module, me)
+
+	// Ensure the base case doesn't fail: A single parameter should work as that matches the function signature.
+	const funcIndex = 0
+	ce := me.NewFunction(funcIndex)
+
+	stack := []uint64{1, 2}
+	err = ce.CallWithStack(testCtx, stack)
+	require.NoError(t, err)
+	require.Equal(t, []uint64{1, 2}, stack)
+
+	t.Run("errs when not enough parameters", func(t *testing.T) {
+		ce := me.NewFunction(funcIndex)
+		err = ce.CallWithStack(testCtx, nil)
+		require.EqualError(t, err, "need 2 params, but stack size is 0")
+	})
+}
+
 func RunTestModuleEngineLookupFunction(t *testing.T, et EngineTester) {
 	e := et.NewEngine(api.CoreFeaturesV1)
 
diff --git a/internal/wasm/gofunc_test.go b/internal/wasm/gofunc_test.go
index 64269bf8..8b2e8b8c 100644
--- a/internal/wasm/gofunc_test.go
+++ b/internal/wasm/gofunc_test.go
@@ -124,55 +124,6 @@ func Test_parseGoFunc_Errors(t *testing.T) {
 	}
 }
 
-// stack simulates the value stack in a way easy to be tested.
-type stack struct {
-	vals []uint64
-}
-
-func (s *stack) pop() (result uint64) {
-	stackTopIndex := len(s.vals) - 1
-	result = s.vals[stackTopIndex]
-	s.vals = s.vals[:stackTopIndex]
-	return
-}
-
-func TestPopValues(t *testing.T) {
-	stackVals := []uint64{1, 2, 3, 4, 5, 6, 7}
-	tests := []struct {
-		name     string
-		count    int
-		expected []uint64
-	}{
-		{
-			name: "pop zero doesn't allocate a slice ",
-		},
-		{
-			name:     "pop 1",
-			count:    1,
-			expected: []uint64{7},
-		},
-		{
-			name:     "pop 2",
-			count:    2,
-			expected: []uint64{6, 7},
-		},
-		{
-			name:     "pop 3",
-			count:    3,
-			expected: []uint64{5, 6, 7},
-		},
-	}
-
-	for _, tt := range tests {
-		tc := tt
-
-		t.Run(tc.name, func(t *testing.T) {
-			vals := PopValues(tc.count, (&stack{stackVals}).pop)
-			require.Equal(t, tc.expected, vals)
-		})
-	}
-}
-
 func Test_callGoFunc(t *testing.T) {
 	tPtr := uintptr(unsafe.Pointer(t))
 	inst := &ModuleInstance{}
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 6d736b40..135f688c 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -488,12 +488,16 @@ func (e *mockModuleEngine) Close(context.Context) {
 func (ce *mockCallEngine) Definition() api.FunctionDefinition { return nil }
 
 // Call implements the same method as documented on api.Function.
-func (ce *mockCallEngine) Call(_ context.Context, _ ...uint64) (results []uint64, err error) {
+func (ce *mockCallEngine) Call(ctx context.Context, _ ...uint64) (results []uint64, err error) {
+	return nil, ce.CallWithStack(ctx, nil)
+}
+
+// CallWithStack implements the same method as documented on api.Function.
+func (ce *mockCallEngine) CallWithStack(_ context.Context, _ []uint64) error {
 	if ce.callFailIndex >= 0 && ce.index == Index(ce.callFailIndex) {
-		err = errors.New("call failed")
-		return
+		return errors.New("call failed")
 	}
-	return
+	return nil
 }
 
 func TestStore_getFunctionTypeID(t *testing.T) {