cmd/internal/obj/riscv: Reuse REG_TMP for consecutive memory access with large offset

[ctk-1998]

Related Issue(s) & Descriptions

During testing golang/benchmarks, three benchmarks in the wazero library performed relatively poorly. When capturing hot functions in the SG server environment, it was found that the hot function in all cases was callNativeFunc, with CPU time reaching 97% in each. This function is a pure Go function implemented by wazero. Static code analysis revealed that it generates an excessively large number of memory-access-related instructions. Compared to the 31,901 lines shown in the ARM disassembly, the RISC-V instruction count reached 49,317 lines. The reasons for this instruction bloat are as follows:

• RISC-V lacks instructions similar to ARM's LDP and STP, which can load or store two registers at once.
• RISC-V's immediate values are only 12 bits wide, limiting the size of offset addresses.

Currently, the lack of instructions cannot be supplemented. However, when calculating memory access addresses, if the offset exceeds the 12-bit immediate limit, two additional instructions—LUI and ADD—are generated to compute a new base address. At present, this handling is applied uniformly in such scenarios, but many of these LUI and ADD calculations are redundant and could be reused from previous results.

Solution

During preprocess, memory-access prog nodes eligible for address-base reuse are identified and marked. In instructionsForLoad/Store, a marked node emits only the memory operation with REG_TMP+low, omitting LUI+ADD.

Reuse is allowed only if the current and previous instructions are both optimizable TYPE_MEM accesses, share the same base register, have identical immediate high parts, and the offset delta is < 2048; additionally, the previous instruction must not clobber the base register or overwrite REG_TMP.

Control-flow safety is enforced: neither the current node nor any node in the prev->p interval may be a branch/jump landing target, preventing reuse on paths where the prior address computation is not guaranteed to execute.

For store operations, consecutive reuse length is capped to avoid regressions from excessively long optimized store chains.

Performance data

Tested on SG2044, taskset to an idle core:

goos: linux
goarch: riscv64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
                                                   │ sg_old.txt  │            sg_new.txt             │
                                                   │   sec/op    │   sec/op     vs base              │
Invocation/interpreter/fib_for_20                    9.194m ± 2%   8.515m ± 1%  -7.38% (p=0.002 n=6)
Invocation/interpreter/string_manipulation_size_50   3.916m ± 7%   3.638m ± 6%  -7.09% (p=0.004 n=6)
Invocation/interpreter/random_mat_mul_size_20        24.75m ± 1%   23.72m ± 0%  -4.17% (p=0.002 n=6)
geomean                                              9.623m        9.023m       -6.23%

Todo

The existing implementation is limited to reuse scenarios involving consecutive memory access operations, with no other instructions interspersed, which restricts its practicality. Work is already underway to support interleaving other instructions, but an issue has currently been encountered where certain incorrect reuse optimizations result in the use of a modified REG_TMP, causing a segmentation violation. Once this is resolved, a new PR will be submitted as a supplement.

Checklist

• Tests were added or are not required
• Documentation was added or is not required

[wangpc-pp]

IIRC, there is a related issue in upstream Golang repo?

---

Edit: I think I found it: golang#77541

[wangpc-pp]

I don't know if we can do this in a SSA pass, instead of optimizing it when emitting machine instructions.

[ctk-1998]

I don't know if we can do this in a SSA pass, instead of optimizing it when emitting machine instructions.

I considered implementing this in the SSA phase. My approach is to compute baseHi = base + (high << 12) once at an appropriate position within the block, then reuse it for subsequent accesses, retaining only the low‑order offset. In the SSA, this is expressed only as “a baseHi value is needed”; the actual LUI+ADD pattern is determined by the riscv64 lowering pass. Whether a subsequent load/store “only needs a MOV” depends on whether the lowering can encode baseHi + low into an addressing mode with rs1 = baseHi, imm = low. However, in the case where ssa.OpLoadReg / ssa.OpStoreReg is originally used, ssagen.AddrAuto fixes the address to a.Reg = REGSP along this path, so I cannot simply change the offset while converting the base to baseHi. Therefore, I will not perform the optimization for such cases.

I previously created an implementation by adding a new pass in compile.go, inserted after the late lower stage, and performed the optimization in that new pass. The code is as follows:

func riscv64BaseHiReuse(f *Func) {
	if f == nil || f.Config == nil || f.Config.arch != "riscv64" {
		return
	}

	type key struct {
		base *Value
		high int64
	}

	isCallLike := func(op Op) bool {
		switch op {
		case OpRISCV64CALLstatic, OpRISCV64CALLclosure, OpRISCV64CALLinter, OpRISCV64CALLtail:
			return true
		case OpRISCV64LoweredWB, OpRISCV64LoweredNilCheck:
			return true
		case OpRISCV64LoweredPanicBoundsA, OpRISCV64LoweredPanicBoundsB, OpRISCV64LoweredPanicBoundsC:
			return true
		default:
			return false
		}
	}

	isAtomicLike := func(op Op) bool {
		switch op {
		case OpRISCV64LoweredAtomicLoad8,
			OpRISCV64LoweredAtomicLoad32, OpRISCV64LoweredAtomicLoad64,
			OpRISCV64LoweredAtomicStore8,
			OpRISCV64LoweredAtomicStore32, OpRISCV64LoweredAtomicStore64,
			OpRISCV64LoweredAtomicAdd32, OpRISCV64LoweredAtomicAdd64,
			OpRISCV64LoweredAtomicExchange8,
			OpRISCV64LoweredAtomicExchange32, OpRISCV64LoweredAtomicExchange64,
			OpRISCV64LoweredAtomicCas32, OpRISCV64LoweredAtomicCas64,
			OpRISCV64LoweredAtomicAnd8, OpRISCV64LoweredAtomicOr8,
			OpRISCV64LoweredAtomicAnd32, OpRISCV64LoweredAtomicOr32:
			return true
		default:
			return false
		}
	}

	isMemOp := func(op Op) bool {
		switch op {
		case OpRISCV64MOVBload, OpRISCV64MOVHload, OpRISCV64MOVWload, OpRISCV64MOVDload,
			OpRISCV64MOVBUload, OpRISCV64MOVHUload, OpRISCV64MOVWUload,
			OpRISCV64FMOVWload, OpRISCV64FMOVDload,
			OpRISCV64MOVBstore, OpRISCV64MOVHstore, OpRISCV64MOVWstore, OpRISCV64MOVDstore,
			OpRISCV64FMOVWstore, OpRISCV64FMOVDstore:
			return true
		default:
			return false
		}
	}

	// splitRISCVImm32 splits an address offset into low (signed 12-bit) and high (remaining, in units of 4KiB),
	// matching the backend convention: offset == (high<<12) + low, with low in [-2048, 2047].
	splitRISCVImm32 := func(off int64) (low, high int64, ok bool) {
		// We only care about values that fit in signed 32-bit (matches common backend constraints).
		if off < -0x80000000 || off > 0x7fffffff {
			return 0, 0, false
		}
		// Compute low as signed 12-bit part.
		low = off & 0xfff
		if low >= 0x800 {
			low -= 0x1000
		}
		high = (off - low) >> 12
		return low, high, true
	}

	for _, b := range f.Blocks {
		// Skip blocks with call-like/atomic-like ops to keep this pass conservative.
		skip := false
		for _, v := range b.Values {
			if isCallLike(v.Op) || isAtomicLike(v.Op) {
				skip = true
				break
			}
		}
		if skip {
			continue
		}

		groups := make(map[key][]*Value)
		for _, v := range b.Values {
			if !isMemOp(v.Op) {
				continue
			}
			// Only rewrite pure base+const offsets without symbol addends.
			if v.Aux != nil {
				continue
			}
			if len(v.Args) < 1 {
				continue
			}
			base := v.Args[0]
			off := v.AuxInt
			low, high, ok := splitRISCVImm32(off)
			if !ok || high == 0 {
				continue
			}
			groups[key{base: base, high: high}] = append(groups[key{base: base, high: high}], v)
			_ = low // recomputed when rewriting
		}

		for k, vs := range groups {
			if len(vs) < 2 {
				continue
			}

			// Materialize baseHi = base + (high<<12).
			hiOff := k.high << 12
			// Use the base's type to keep pointer/uintptr typing consistent.
			hiConst := b.NewValue0(vs[0].Pos, OpRISCV64MOVDconst, k.base.Type)
			hiConst.AuxInt = int64ToAuxInt(hiOff)
			baseHi := b.NewValue0(vs[0].Pos, OpRISCV64ADD, k.base.Type)
			baseHi.AddArg2(k.base, hiConst)

			// Rewrite each mem op to use baseHi with low 12-bit offset.
			for _, v := range vs {
				low, _, ok := splitRISCVImm32(v.AuxInt)
				if !ok {
					continue
				}
				v.SetArg(0, baseHi)
				v.AuxInt = int64ToAuxInt(low)
			}
		}
	}
}

After that, I tested it in the same environment and obtained the following results:

goos: linux
goarch: riscv64
pkg: github.com/tetratelabs/wazero/internal/integration_test/bench
                                                   │ sg_old.txt  │            sg_ssa.txt             │
                                                   │   sec/op    │   sec/op     vs base              │
Invocation/interpreter/fib_for_20                    9.194m ± 2%   8.948m ± 0%  -2.67% (p=0.002 n=6)
Invocation/interpreter/string_manipulation_size_50   3.916m ± 7%   3.851m ± 5%       ~ (p=0.310 n=6)
Invocation/interpreter/random_mat_mul_size_20        24.75m ± 1%   24.49m ± 1%  -1.04% (p=0.002 n=6)
geomean                                              9.623m        9.450m       -1.79%

The results are disappointing.

[BoyaoWang430]

Why 3 here？ Any benchmark result to justify it?

[ctk-1998]

Actually, I think the value should be related to the number of pipelines in the Load/Store execution units and cycle counts, but I don't know how to configure it dynamically.

[BoyaoWang430]

No need to use a sync.Map here, we can use 1 bit in Mark of obj.Prog

[BoyaoWang430]

We can use a linked list here. The slice here may be very large, and we only need prev.