cannon: Add support for --type mt to enable multithreaded cannon (#11411)

* cannon: Add support for --type mt to enable multithreaded cannon

* cannon: Rename vm type to cannon-mt

Makefile

@@ -141,6 +141,12 @@ cannon-prestate: op-program cannon ## Generates prestate using cannon and op-pro
 	mv op-program/bin/0.json op-program/bin/prestate-proof.json
 .PHONY: cannon-prestate
 
+cannon-prestate-mt: op-program cannon ## Generates prestate using cannon and op-program in the multithreaded cannon format
+	./cannon/bin/cannon load-elf --type mt --path op-program/bin/op-program-client.elf --out op-program/bin/prestate-mt.json --meta op-program/bin/meta-mt.json
+	./cannon/bin/cannon run --type mt --proof-at '=0' --stop-at '=1' --input op-program/bin/prestate-mt.json --meta op-program/bin/meta-mt.json --proof-fmt 'op-program/bin/%d-mt.json' --output ""
+	mv op-program/bin/0-mt.json op-program/bin/prestate-proof-mt.json
+.PHONY: cannon-prestate
+
 mod-tidy: ## Cleans up unused dependencies in Go modules
 	# Below GOPRIVATE line allows mod-tidy to be run immediately after
 	# releasing new versions. This bypasses the Go modules proxy, which

cannon/cmd/load_elf.go

@@ -4,6 +4,8 @@ import (
 	"debug/elf"
 	"fmt"
 
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm"
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm/multithreaded"
 	"github.com/urfave/cli/v2"
 
 	"github.com/ethereum-optimism/optimism/cannon/mipsevm/program"
@@ -39,6 +41,28 @@ var (
 )
 
 func LoadELF(ctx *cli.Context) error {
+	var createInitialState func(f *elf.File) (mipsevm.FPVMState, error)
+	var writeState func(path string, state mipsevm.FPVMState) error
+
+	if vmType, err := vmTypeFromString(ctx); err != nil {
+		return err
+	} else if vmType == cannonVMType {
+		createInitialState = func(f *elf.File) (mipsevm.FPVMState, error) {
+			return program.LoadELF(f, singlethreaded.CreateInitialState)
+		}
+		writeState = func(path string, state mipsevm.FPVMState) error {
+			return jsonutil.WriteJSON[*singlethreaded.State](path, state.(*singlethreaded.State), OutFilePerm)
+		}
+	} else if vmType == mtVMType {
+		createInitialState = func(f *elf.File) (mipsevm.FPVMState, error) {
+			return program.LoadELF(f, multithreaded.CreateInitialState)
+		}
+		writeState = func(path string, state mipsevm.FPVMState) error {
+			return jsonutil.WriteJSON[*multithreaded.State](path, state.(*multithreaded.State), OutFilePerm)
+		}
+	} else {
+		return fmt.Errorf("invalid VM type: %q", vmType)
+	}
 	elfPath := ctx.Path(LoadELFPathFlag.Name)
 	elfProgram, err := elf.Open(elfPath)
 	if err != nil {
@@ -47,7 +71,7 @@ func LoadELF(ctx *cli.Context) error {
 	if elfProgram.Machine != elf.EM_MIPS {
 		return fmt.Errorf("ELF is not big-endian MIPS R3000, but got %q", elfProgram.Machine.String())
 	}
-	state, err := program.LoadELF(elfProgram, singlethreaded.CreateInitialState)
+	state, err := createInitialState(elfProgram)
 	if err != nil {
 		return fmt.Errorf("failed to load ELF data into VM state: %w", err)
 	}
@@ -71,7 +95,7 @@ func LoadELF(ctx *cli.Context) error {
 	if err := jsonutil.WriteJSON[*program.Metadata](ctx.Path(LoadELFMetaFlag.Name), meta, OutFilePerm); err != nil {
 		return fmt.Errorf("failed to output metadata: %w", err)
 	}
-	return jsonutil.WriteJSON[*singlethreaded.State](ctx.Path(LoadELFOutFlag.Name), state, OutFilePerm)
+	return writeState(ctx.Path(LoadELFOutFlag.Name), state)
 }
 
 var LoadELFCommand = &cli.Command{
@@ -80,6 +104,7 @@ var LoadELFCommand = &cli.Command{
 	Description: "Load ELF file into Cannon JSON state, optionally patch out functions",
 	Action:      LoadELF,
 	Flags: []cli.Flag{
+		VMTypeFlag,
 		LoadELFPathFlag,
 		LoadELFPatchFlag,
 		LoadELFOutFlag,

cannon/cmd/run.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm/multithreaded"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/log"
@@ -24,13 +25,6 @@ import (
 )
 
 var (
-	RunType = &cli.StringFlag{
-		Name:  "type",
-		Usage: "VM type to run. Options are 'cannon' (default)",
-		Value: "cannon",
-		// TODO(client-pod#903): This should be required once we have additional vm types
-		Required: false,
-	}
 	RunInputFlag = &cli.PathFlag{
 		Name:      "input",
 		Usage:     "path of input JSON state. Stdin if left empty.",
@@ -258,20 +252,14 @@ func Guard(proc *os.ProcessState, fn StepFn) StepFn {
 
 var _ mipsevm.PreimageOracle = (*ProcessPreimageOracle)(nil)
 
-type VMType string
-
-var cannonVMType VMType = "cannon"
-
 func Run(ctx *cli.Context) error {
 	if ctx.Bool(RunPProfCPU.Name) {
 		defer profile.Start(profile.NoShutdownHook, profile.ProfilePath("."), profile.CPUProfile).Stop()
 	}
 
-	var vmType VMType
-	if vmTypeStr := ctx.String(RunType.Name); vmTypeStr == string(cannonVMType) {
-		vmType = cannonVMType
-	} else {
-		return fmt.Errorf("unknown VM type %q", vmType)
+	vmType, err := vmTypeFromString(ctx)
+	if err != nil {
+		return err
 	}
 
 	guestLogger := Logger(os.Stderr, log.LevelInfo)
@@ -366,6 +354,7 @@ func Run(ctx *cli.Context) error {
 	var vm mipsevm.FPVM
 	var debugProgram bool
 	if vmType == cannonVMType {
+		l.Info("Using cannon VM")
 		cannon, err := singlethreaded.NewInstrumentedStateFromFile(ctx.Path(RunInputFlag.Name), po, outLog, errLog, meta)
 		if err != nil {
 			return err
@@ -380,6 +369,22 @@ func Run(ctx *cli.Context) error {
 			}
 		}
 		vm = cannon
+	} else if vmType == mtVMType {
+		l.Info("Using cannon multithreaded VM")
+		cannon, err := multithreaded.NewInstrumentedStateFromFile(ctx.Path(RunInputFlag.Name), po, outLog, errLog, l)
+		if err != nil {
+			return err
+		}
+		debugProgram = ctx.Bool(RunDebugFlag.Name)
+		if debugProgram {
+			if metaPath := ctx.Path(RunMetaFlag.Name); metaPath == "" {
+				return fmt.Errorf("cannot enable debug mode without a metadata file")
+			}
+			if err := cannon.InitDebug(meta); err != nil {
+				return fmt.Errorf("failed to initialize debug mode: %w", err)
+			}
+		}
+		vm = cannon
 	} else {
 		return fmt.Errorf("unknown VM type %q", vmType)
 	}
@@ -503,7 +508,7 @@ var RunCommand = &cli.Command{
 	Description: "Run VM step(s) and generate proof data to replicate onchain. See flags to match when to output a proof, a snapshot, or to stop early.",
 	Action:      Run,
 	Flags: []cli.Flag{
-		RunType,
+		VMTypeFlag,
 		RunInputFlag,
 		RunOutputFlag,
 		RunProofAtFlag,

cannon/cmd/vmtype.go

+package cmd
+
+import (
+	"fmt"
+
+	"github.com/urfave/cli/v2"
+)
+
+type VMType string
+
+var cannonVMType VMType = "cannon"
+var mtVMType VMType = "cannon-mt"
+
+var VMTypeFlag = &cli.StringFlag{
+	Name:     "type",
+	Usage:    "VM type to create state for. Options are 'cannon' (default), 'cannon-mt'",
+	Value:    "cannon",
+	Required: false,
+}
+
+func vmTypeFromString(ctx *cli.Context) (VMType, error) {
+	if vmTypeStr := ctx.String(VMTypeFlag.Name); vmTypeStr == string(cannonVMType) {
+		return cannonVMType, nil
+	} else if vmTypeStr == string(mtVMType) {
+		return mtVMType, nil
+	} else {
+		return "", fmt.Errorf("unknown VM type %q", vmTypeStr)
+	}
+}

cannon/cmd/witness.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm"
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm/multithreaded"
 	"github.com/urfave/cli/v2"
 
 	"github.com/ethereum-optimism/optimism/cannon/mipsevm/singlethreaded"
@@ -27,10 +29,23 @@ var (
 func Witness(ctx *cli.Context) error {
 	input := ctx.Path(WitnessInputFlag.Name)
 	output := ctx.Path(WitnessOutputFlag.Name)
-	state, err := jsonutil.LoadJSON[singlethreaded.State](input)
+	var state mipsevm.FPVMState
+	if vmType, err := vmTypeFromString(ctx); err != nil {
+		return err
+	} else if vmType == cannonVMType {
+		state, err = jsonutil.LoadJSON[singlethreaded.State](input)
 		if err != nil {
 			return fmt.Errorf("invalid input state (%v): %w", input, err)
 		}
+	} else if vmType == mtVMType {
+		state, err = jsonutil.LoadJSON[multithreaded.State](input)
+		if err != nil {
+			return fmt.Errorf("invalid input state (%v): %w", input, err)
+		}
+	} else {
+		return fmt.Errorf("invalid VM type: %q", vmType)
+	}
+
 	witness, h := state.EncodeWitness()
 	if output != "" {
 		if err := os.WriteFile(output, witness, 0755); err != nil {
@@ -47,6 +62,7 @@ var WitnessCommand = &cli.Command{
 	Description: "Convert a Cannon JSON state into a binary witness. The hash of the witness is written to stdout",
 	Action:      Witness,
 	Flags: []cli.Flag{
+		VMTypeFlag,
 		WitnessInputFlag,
 		WitnessOutputFlag,
 	},