Add cannon memory benchmarks (ethereum-optimism#10942)

* Add cannon memory benchmarks

* fix comment

* go mod tidy alloc program

cannon/cmd/run.go

@@ -103,6 +103,12 @@ var (
 		Name:  "debug",
 		Usage: "enable debug mode, which includes stack traces and other debug info in the output. Requires --meta.",
 	}
+	RunDebugInfoFlag = &cli.PathFlag{
+		Name:      "debug-info",
+		Usage:     "path to write debug info to",
+		TakesFile: true,
+		Required:  false,
+	}
 
 	OutFilePerm = os.FileMode(0o755)
 )
@@ -466,6 +472,11 @@ func Run(ctx *cli.Context) error {
 	if err := jsonutil.WriteJSON(ctx.Path(RunOutputFlag.Name), state, OutFilePerm); err != nil {
 		return fmt.Errorf("failed to write state output: %w", err)
 	}
+	if debugInfoFile := ctx.Path(RunDebugInfoFlag.Name); debugInfoFile != "" {
+		if err := jsonutil.WriteJSON(debugInfoFile, us.GetDebugInfo(), OutFilePerm); err != nil {
+			return fmt.Errorf("failed to write benchmark data: %w", err)
+		}
+	}
 	return nil
 }
 
@@ -489,5 +500,6 @@ var RunCommand = &cli.Command{
 		RunInfoAtFlag,
 		RunPProfCPU,
 		RunDebugFlag,
+		RunDebugInfoFlag,
 	},
 }

cannon/example/alloc/go.mod

@@ -0,0 +1,14 @@
+module alloc
+
+go 1.21
+
+toolchain go1.21.1
+
+require github.com/ethereum-optimism/optimism v0.0.0
+
+require (
+	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/sys v0.21.0 // indirect
+)
+
+replace github.com/ethereum-optimism/optimism v0.0.0 => ../../..

cannon/example/alloc/go.sum

@@ -0,0 +1,12 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
+golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
+golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
+golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

cannon/example/alloc/main.go

@@ -0,0 +1,31 @@
+package main
+
+import (
+	"encoding/binary"
+	"fmt"
+	"runtime"
+
+	preimage "github.com/ethereum-optimism/optimism/op-preimage"
+)
+
+func main() {
+	var mem []byte
+	po := preimage.NewOracleClient(preimage.ClientPreimageChannel())
+	numAllocs := binary.LittleEndian.Uint64(po.Get(preimage.LocalIndexKey(0)))
+
+	fmt.Printf("alloc program. numAllocs=%d\n", numAllocs)
+	var alloc int
+	for i := 0; i < int(numAllocs); i++ {
+		mem = make([]byte, 32*1024*1024)
+		alloc += len(mem)
+		// touch a couple pages to prevent the runtime from overcommitting memory
+		for j := 0; j < len(mem); j += 1024 {
+			mem[j] = 0xFF
+		}
+		fmt.Printf("allocated %d bytes\n", alloc)
+	}
+
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	fmt.Printf("alloc program exit. memstats: heap_alloc=%d frees=%d mallocs=%d\n", m.HeapAlloc, m.Frees, m.Mallocs)
+}

cannon/mipsevm/instrumented.go

@@ -26,7 +26,7 @@ type InstrumentedState struct {
 	memProofEnabled bool
 	memProof        [28 * 32]byte
 
-	preimageOracle PreimageOracle
+	preimageOracle *trackingOracle
 
 	// cached pre-image data, including 8 byte length prefix
 	lastPreimage []byte
@@ -59,7 +59,7 @@ func NewInstrumentedState(state *State, po PreimageOracle, stdOut, stdErr io.Wri
 		state:          state,
 		stdOut:         stdOut,
 		stdErr:         stdErr,
-		preimageOracle: po,
+		preimageOracle: &trackingOracle{po: po},
 	}
 }
 
@@ -103,3 +103,34 @@ func (m *InstrumentedState) Step(proof bool) (wit *StepWitness, err error) {
 func (m *InstrumentedState) LastPreimage() ([32]byte, []byte, uint32) {
 	return m.lastPreimageKey, m.lastPreimage, m.lastPreimageOffset
 }
+
+func (d *InstrumentedState) GetDebugInfo() *DebugInfo {
+	return &DebugInfo{
+		Pages:               d.state.Memory.PageCount(),
+		NumPreimageRequests: d.preimageOracle.numPreimageRequests,
+		TotalPreimageSize:   d.preimageOracle.totalPreimageSize,
+	}
+}
+
+type DebugInfo struct {
+	Pages               int `json:"pages"`
+	NumPreimageRequests int `json:"num_preimage_requests"`
+	TotalPreimageSize   int `json:"total_preimage_size"`
+}
+
+type trackingOracle struct {
+	po                  PreimageOracle
+	totalPreimageSize   int
+	numPreimageRequests int
+}
+
+func (d *trackingOracle) Hint(v []byte) {
+	d.po.Hint(v)
+}
+
+func (d *trackingOracle) GetPreimage(k [32]byte) []byte {
+	d.numPreimageRequests++
+	preimage := d.po.GetPreimage(k)
+	d.totalPreimageSize += len(preimage)
+	return preimage
+}

cannon/mipsevm/state_test.go

@@ -127,15 +127,7 @@ func TestStateHash(t *testing.T) {
 }
 
 func TestHello(t *testing.T) {
-	elfProgram, err := elf.Open("../example/bin/hello.elf")
-	require.NoError(t, err, "open ELF file")
-
-	state, err := LoadELF(elfProgram)
-	require.NoError(t, err, "load ELF into state")
-
-	err = PatchGo(elfProgram, state)
-	require.NoError(t, err, "apply Go runtime patches")
-	require.NoError(t, PatchStack(state), "add initial stack")
+	state := loadELFProgram(t, "../example/bin/hello.elf")
 
 	var stdOutBuf, stdErrBuf bytes.Buffer
 	us := NewInstrumentedState(state, nil, io.MultiWriter(&stdOutBuf, os.Stdout), io.MultiWriter(&stdErrBuf, os.Stderr))
@@ -225,15 +217,7 @@ func claimTestOracle(t *testing.T) (po PreimageOracle, stdOut string, stdErr str
 }
 
 func TestClaim(t *testing.T) {
-	elfProgram, err := elf.Open("../example/bin/claim.elf")
-	require.NoError(t, err, "open ELF file")
-
-	state, err := LoadELF(elfProgram)
-	require.NoError(t, err, "load ELF into state")
-
-	err = PatchGo(elfProgram, state)
-	require.NoError(t, err, "apply Go runtime patches")
-	require.NoError(t, PatchStack(state), "add initial stack")
+	state := loadELFProgram(t, "../example/bin/claim.elf")
 
 	oracle, expectedStdOut, expectedStdErr := claimTestOracle(t)
 
@@ -255,6 +239,44 @@ func TestClaim(t *testing.T) {
 	require.Equal(t, expectedStdErr, stdErrBuf.String(), "stderr")
 }
 
+func TestAlloc(t *testing.T) {
+	t.Skip("TODO(client-pod#906): Currently fails on Single threaded Cannon. Re-enable for the MT FPVM")
+
+	state := loadELFProgram(t, "../example/bin/alloc.elf")
+	const numAllocs = 100 // where each alloc is a 32 MiB chunk
+	oracle := allocOracle(t, numAllocs)
+
+	// completes in ~870 M steps
+	us := NewInstrumentedState(state, oracle, os.Stdout, os.Stderr)
+	for i := 0; i < 20_000_000_000; i++ {
+		if us.state.Exited {
+			break
+		}
+		_, err := us.Step(false)
+		require.NoError(t, err)
+		if state.Step%10_000_000 == 0 {
+			t.Logf("Completed %d steps", state.Step)
+		}
+	}
+	t.Logf("Completed in %d steps", state.Step)
+	require.True(t, state.Exited, "must complete program")
+	require.Equal(t, uint8(0), state.ExitCode, "exit with 0")
+	require.Less(t, state.Memory.PageCount()*PageSize, 1*1024*1024*1024, "must not allocate more than 1 GiB")
+}
+
+func loadELFProgram(t *testing.T, name string) *State {
+	elfProgram, err := elf.Open(name)
+	require.NoError(t, err, "open ELF file")
+
+	state, err := LoadELF(elfProgram)
+	require.NoError(t, err, "load ELF into state")
+
+	err = PatchGo(elfProgram, state)
+	require.NoError(t, err, "apply Go runtime patches")
+	require.NoError(t, PatchStack(state), "add initial stack")
+	return state
+}
+
 func staticOracle(t *testing.T, preimageData []byte) *testOracle {
 	return &testOracle{
 		hint: func(v []byte) {},
@@ -289,6 +311,18 @@ func staticPrecompileOracle(t *testing.T, precompile common.Address, input []byt
 	}
 }
 
+func allocOracle(t *testing.T, numAllocs int) *testOracle {
+	return &testOracle{
+		hint: func(v []byte) {},
+		getPreimage: func(k [32]byte) []byte {
+			if k != preimage.LocalIndexKey(0).PreimageKey() {
+				t.Fatalf("invalid preimage request for %x", k)
+			}
+			return binary.LittleEndian.AppendUint64(nil, uint64(numAllocs))
+		},
+	}
+}
+
 func selectOracleFixture(t *testing.T, programName string) PreimageOracle {
 	if strings.HasPrefix(programName, "oracle_kzg") {
 		precompile := common.BytesToAddress([]byte{0xa})