Optimize/streamline fill operations (wazero#2395)

Author: ncruces

internal/engine/interpreter/interpreter.go

@@ -1885,12 +1885,17 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance
 			if fillSize+offset > uint64(len(memoryInst.Buffer)) {
 				panic(wasmruntime.ErrRuntimeOutOfBoundsMemoryAccess)
 			} else if fillSize != 0 {
-				// Uses the copy trick for faster filling buffer.
-				// https://gist.github.com/taylorza/df2f89d5f9ab3ffd06865062a4cf015d
+				// Uses the copy trick for faster filling the buffer with the value.
+				// https://github.com/golang/go/blob/go1.24.0/src/bytes/bytes.go#L664-L673
 				buf := memoryInst.Buffer[offset : offset+fillSize]
-				buf[0] = value
-				for i := 1; i < len(buf); i *= 2 {
-					copy(buf[i:], buf[:i])
+				if value == 0 {
+					clear(buf)
+				} else {
+					buf[0] = value
+					for i := 1; i < len(buf); {
+						chunk := min(i, 8192)
+						i += copy(buf[i:], buf[:chunk])
+					}
 				}
 			}
 			frame.pc++
@@ -1964,7 +1969,7 @@ func (ce *callEngine) callNativeFunc(ctx context.Context, m *wasm.ModuleInstance
 				panic(wasmruntime.ErrRuntimeInvalidTableAccess)
 			} else if num > 0 {
 				// Uses the copy trick for faster filling the region with the value.
-				// https://gist.github.com/taylorza/df2f89d5f9ab3ffd06865062a4cf015d
+				// https://github.com/golang/go/blob/go1.24.0/src/slices/slices.go#L514-L517
 				targetRegion := table.References[offset : offset+num]
 				targetRegion[0] = ref
 				for i := 1; i < len(targetRegion); i *= 2 {

internal/engine/wazevo/frontend/lower.go

@@ -665,19 +665,22 @@ func (c *Compiler) lowerCurrentOpcode() {
 			tableBaseAddr := c.loadTableBaseAddr(tableInstancePtr)
 			addr := builder.AllocateInstruction().AsIadd(tableBaseAddr, offsetInBytes).Insert(builder).Return()
 
-			// Prepare the loop and following block.
-			beforeLoop := builder.AllocateBasicBlock()
-			loopBlk := builder.AllocateBasicBlock()
-			loopVar := loopBlk.AddParam(builder, ssa.TypeI64)
-			followingBlk := builder.AllocateBasicBlock()
-
 			// Uses the copy trick for faster filling buffer like memory.fill, but in this case we copy 8 bytes at a time.
+			// Tables are rarely huge, so ignore the 8KB maximum.
+			// https://github.com/golang/go/blob/go1.24.0/src/slices/slices.go#L514-L517
+			//
 			// 	buf := memoryInst.Buffer[offset : offset+fillSize]
 			// 	buf[0:8] = value
 			// 	for i := 8; i < fillSize; i *= 2 { Begin with 8 bytes.
 			// 		copy(buf[i:], buf[:i])
 			// 	}
 
+			// Prepare the loop and following block.
+			beforeLoop := builder.AllocateBasicBlock()
+			loopBlk := builder.AllocateBasicBlock()
+			loopVar := loopBlk.AddParam(builder, ssa.TypeI64)
+			followingBlk := builder.AllocateBasicBlock()
+
 			// Insert the jump to the beforeLoop block; If the fillSize is zero, then jump to the following block to skip entire logics.
 			zero := builder.AllocateInstruction().AsIconst64(0).Insert(builder).Return()
 			ifFillSizeZero := builder.AllocateInstruction().AsIcmp(fillSizeExt, zero, ssa.IntegerCmpCondEqual).
@@ -688,32 +691,24 @@ func (c *Compiler) lowerCurrentOpcode() {
 			// buf[0:8] = value
 			builder.SetCurrentBlock(beforeLoop)
 			builder.AllocateInstruction().AsStore(ssa.OpcodeStore, value, addr, 0).Insert(builder)
-			initValue := builder.AllocateInstruction().AsIconst64(8).Insert(builder).Return()
-			c.insertJumpToBlock(c.allocateVarLengthValues(1, initValue), loopBlk)
+			eight := builder.AllocateInstruction().AsIconst64(8).Insert(builder).Return()
+			c.insertJumpToBlock(c.allocateVarLengthValues(1, eight), loopBlk)
 
 			builder.SetCurrentBlock(loopBlk)
 			dstAddr := builder.AllocateInstruction().AsIadd(addr, loopVar).Insert(builder).Return()
 
-			// If loopVar*2 > fillSizeInBytes, then count must be fillSizeInBytes-loopVar.
-			var count ssa.Value
-			{
-				loopVarDoubled := builder.AllocateInstruction().AsIadd(loopVar, loopVar).Insert(builder).Return()
-				loopVarDoubledLargerThanFillSize := builder.
-					AllocateInstruction().AsIcmp(loopVarDoubled, fillSizeInBytes, ssa.IntegerCmpCondUnsignedGreaterThanOrEqual).
-					Insert(builder).Return()
-				diff := builder.AllocateInstruction().AsIsub(fillSizeInBytes, loopVar).Insert(builder).Return()
-				count = builder.AllocateInstruction().AsSelect(loopVarDoubledLargerThanFillSize, diff, loopVar).Insert(builder).Return()
-			}
+			newLoopVar := builder.AllocateInstruction().AsIadd(loopVar, loopVar).Insert(builder).Return()
+			newLoopVarLessThanFillSize := builder.AllocateInstruction().
+				AsIcmp(newLoopVar, fillSizeInBytes, ssa.IntegerCmpCondUnsignedLessThan).Insert(builder).Return()
 
-			c.callMemmove(dstAddr, addr, count)
+			// On the last iteration, count must be fillSizeInBytes-loopVar.
+			diff := builder.AllocateInstruction().AsIsub(fillSizeInBytes, loopVar).Insert(builder).Return()
+			count := builder.AllocateInstruction().AsSelect(newLoopVarLessThanFillSize, loopVar, diff).Insert(builder).Return()
 
-			shiftAmount := builder.AllocateInstruction().AsIconst64(1).Insert(builder).Return()
-			newLoopVar := builder.AllocateInstruction().AsIshl(loopVar, shiftAmount).Insert(builder).Return()
-			loopVarLessThanFillSize := builder.AllocateInstruction().
-				AsIcmp(newLoopVar, fillSizeInBytes, ssa.IntegerCmpCondUnsignedLessThan).Insert(builder).Return()
+			c.callMemmove(dstAddr, addr, count)
 
 			builder.AllocateInstruction().
-				AsBrnz(loopVarLessThanFillSize, c.allocateVarLengthValues(1, newLoopVar), loopBlk).
+				AsBrnz(newLoopVarLessThanFillSize, c.allocateVarLengthValues(1, newLoopVar), loopBlk).
 				Insert(builder)
 
 			c.insertJumpToBlock(ssa.ValuesNil, followingBlk)
@@ -741,11 +736,15 @@ func (c *Compiler) lowerCurrentOpcode() {
 			// Calculate the base address:
 			addr := builder.AllocateInstruction().AsIadd(c.getMemoryBaseValue(false), offset).Insert(builder).Return()
 
-			// Uses the copy trick for faster filling buffer: https://gist.github.com/taylorza/df2f89d5f9ab3ffd06865062a4cf015d
+			// Uses the copy trick for faster filling buffer, with a maximum chunk size of 8KB.
+			// https://github.com/golang/go/blob/go1.24.0/src/bytes/bytes.go#L664-L673
+			//
 			// 	buf := memoryInst.Buffer[offset : offset+fillSize]
 			// 	buf[0] = value
-			// 	for i := 1; i < fillSize; i *= 2 {
-			// 		copy(buf[i:], buf[:i])
+			// 	for i := 1; i < fillSize; {
+			// 		chunk := ((i - 1) & 8191) + 1
+			// 		copy(buf[i:], buf[:chunk])
+			// 		i += chunk
 			// 	}
 
 			// Prepare the loop and following block.
@@ -764,32 +763,31 @@ func (c *Compiler) lowerCurrentOpcode() {
 			// buf[0] = value
 			builder.SetCurrentBlock(beforeLoop)
 			builder.AllocateInstruction().AsStore(ssa.OpcodeIstore8, value, addr, 0).Insert(builder)
-			initValue := builder.AllocateInstruction().AsIconst64(1).Insert(builder).Return()
-			c.insertJumpToBlock(c.allocateVarLengthValues(1, initValue), loopBlk)
+			one := builder.AllocateInstruction().AsIconst64(1).Insert(builder).Return()
+			c.insertJumpToBlock(c.allocateVarLengthValues(1, one), loopBlk)
 
 			builder.SetCurrentBlock(loopBlk)
 			dstAddr := builder.AllocateInstruction().AsIadd(addr, loopVar).Insert(builder).Return()
 
-			// If loopVar*2 > fillSizeExt, then count must be fillSizeExt-loopVar.
-			var count ssa.Value
-			{
-				loopVarDoubled := builder.AllocateInstruction().AsIadd(loopVar, loopVar).Insert(builder).Return()
-				loopVarDoubledLargerThanFillSize := builder.
-					AllocateInstruction().AsIcmp(loopVarDoubled, fillSize, ssa.IntegerCmpCondUnsignedGreaterThanOrEqual).
-					Insert(builder).Return()
-				diff := builder.AllocateInstruction().AsIsub(fillSize, loopVar).Insert(builder).Return()
-				count = builder.AllocateInstruction().AsSelect(loopVarDoubledLargerThanFillSize, diff, loopVar).Insert(builder).Return()
-			}
-
-			c.callMemmove(dstAddr, addr, count)
+			// chunk := ((i - 1) & 8191) + 1
+			mask := builder.AllocateInstruction().AsIconst64(16383).Insert(builder).Return()
+			tmp1 := builder.AllocateInstruction().AsIsub(loopVar, one).Insert(builder).Return()
+			tmp2 := builder.AllocateInstruction().AsBand(tmp1, mask).Insert(builder).Return()
+			chunk := builder.AllocateInstruction().AsIadd(tmp2, one).Insert(builder).Return()
 
-			shiftAmount := builder.AllocateInstruction().AsIconst64(1).Insert(builder).Return()
-			newLoopVar := builder.AllocateInstruction().AsIshl(loopVar, shiftAmount).Insert(builder).Return()
-			loopVarLessThanFillSize := builder.AllocateInstruction().
+			// i += chunk
+			newLoopVar := builder.AllocateInstruction().AsIadd(loopVar, chunk).Insert(builder).Return()
+			newLoopVarLessThanFillSize := builder.AllocateInstruction().
 				AsIcmp(newLoopVar, fillSize, ssa.IntegerCmpCondUnsignedLessThan).Insert(builder).Return()
 
+			// count = min(chunk, fillSize-loopVar)
+			diff := builder.AllocateInstruction().AsIsub(fillSize, loopVar).Insert(builder).Return()
+			count := builder.AllocateInstruction().AsSelect(newLoopVarLessThanFillSize, chunk, diff).Insert(builder).Return()
+
+			c.callMemmove(dstAddr, addr, count)
+
 			builder.AllocateInstruction().
-				AsBrnz(loopVarLessThanFillSize, c.allocateVarLengthValues(1, newLoopVar), loopBlk).
+				AsBrnz(newLoopVarLessThanFillSize, c.allocateVarLengthValues(1, newLoopVar), loopBlk).
 				Insert(builder)
 
 			c.insertJumpToBlock(ssa.ValuesNil, followingBlk)

internal/wasm/table.go

@@ -326,10 +326,14 @@ func (t *TableInstance) Grow(delta uint32, initialRef Reference) (currentLen uin
 	newLen >= math.MaxUint32 || (t.Max != nil && newLen > int64(*t.Max)) {
 		return 0xffffffff // = -1 in signed 32-bit integer.
 	}
+
 	t.References = append(t.References, make([]uintptr, delta)...)
+	if initialRef == 0 {
+		return
+	}
 
 	// Uses the copy trick for faster filling the new region with the initial value.
-	// https://gist.github.com/taylorza/df2f89d5f9ab3ffd06865062a4cf015d
+	// https://github.com/golang/go/blob/go1.24.0/src/slices/slices.go#L514-L517
 	newRegion := t.References[currentLen:]
 	newRegion[0] = initialRef
 	for i := 1; i < len(newRegion); i *= 2 {