AMD ROCm: JSON backend placement

github.com/openfluke/welvet/nativerocmlinux amd64 routing

Why it exists

Choose which layers run on AMD GPU while retaining ordinary Go training entry points.

What it is

Enable the ROCm router and apply JSON placement to a Stack: layer 0 on Go CPU, layers 1-2 on ROCm.

Explicit placement

The example embeds placement.json; use -config to load another file. nativerocm.EnableRouting() installs the router for its owner lifetime (Linux amd64). placement.Stack applies immediate-child numeric layer keys.

Routed execution maintains native sessions and returns updated parameters to Go automatically. Changing source weights remounts a session and resets transient native state. Automatic hardware/performance planning remains separate from explicit placement.

Windows amd64 supports explicit Mount (chapters 96-99); automatic JSON routing is Linux-only in this release.

Requirements: chapter 96.

Go example

examples/100-amd-json-routing/main.go

Run:cd welvet/examples/100-amd-json-routing && source ../env.sh && go run .

// JSON selects backends; the existing Go Stack scheduler moves activations/gradients.
package main

import (
	_ "embed"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"

	"github.com/openfluke/welvet/core"
	"github.com/openfluke/welvet/layers/dense"
	"github.com/openfluke/welvet/layers/parallel"
	"github.com/openfluke/welvet/nativecpp"
	"github.com/openfluke/welvet/nativerocm"
	"github.com/openfluke/welvet/quant"
	"github.com/openfluke/welvet/runtime/placement"
)

//go:embed placement.json
var defaultPlan []byte

func main() {
	if e := run(); e != nil {
		log.Fatal(e)
	}
}

func run() error {
	path := flag.String("config", "", "optional replacement execution JSON")
	flag.Parse()
	if runtime.GOOS != "linux" || runtime.GOARCH != "amd64" {
		fmt.Println("SKIP: ROCm JSON routing needs Linux amd64 and cgo")
		return nil
	}
	raw := defaultPlan
	if *path != "" {
		var e error
		raw, e = os.ReadFile(*path)
		if e != nil {
			return e
		}
	}
	cfg, e := nativecpp.ParseExecConfigJSON(raw)
	if e != nil {
		return e
	}
	router, e := nativerocm.EnableRouting()
	if e != nil {
		fmt.Printf("SKIP: ROCm routing unavailable (%v)\n", e)
		return nil
	}
	defer router.Close()
	var ops []any
	for i := 0; i < 3; i++ {
		l, e := dense.NewConfigured(2, 2, core.ActivationTanh, core.DTypeFloat32, quant.FormatNone, []float32{.8, .1, -.1, .9})
		if e != nil {
			return e
		}
		ops = append(ops, l)
	}
	model, e := parallel.NewStack(ops...)
	if e != nil {
		return e
	}
	if e = placement.Stack(model, cfg, "main", "0,0,0,0"); e != nil {
		return e
	}
	for i, op := range ops {
		fmt.Printf("layer=%d backend=%s\n", i, op.(*dense.Layer).Exec.Backend)
	}
	x := &core.Tensor[float32]{Shape: []int{1, 2}, Data: []float32{.5, -.3}}
	target := core.NewTensor[float32](1, 2)
	loss, e := parallel.TrainStackMSE(model, x, target, parallel.ModeNormalBP, .03)
	if e != nil {
		fmt.Printf("SKIP: ROCm train unavailable (%v)\n", e)
		return nil
	}
	stats := router.Stats()
	if stats.Backwards == 0 || stats.Updates == 0 {
		return fmt.Errorf("plan did not execute native ROCm backward/SGD")
	}
	fmt.Printf("loss=%.8g rocm routing=%+v\nRouted updates are already in Go stores; Close releases cached mounts.\n", loss, stats)
	return nil
}

Output

exit 0 · last run via go run .

layer=0 backend=cpu-tiled
layer=1 backend=rocm
layer=2 backend=rocm
SKIP: ROCm train unavailable (stack fwd child 1: nativerocm: load: libamdhip64.so.7: cannot open shared object file: No such file or directory)

Validate (harness)

cd welvet/examples/100-amd-json-routing && source ../env.sh && go run .