Apple Metal: JSON backend placement

github.com/openfluke/welvet/nativemetalverified Apple M5

Why it exists

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

What it is

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

Explicit placement

The example embeds placement.json; use -config to load another file. nativemetal.EnableRouting() installs the router for its owner lifetime. placement.Stack applies immediate-child numeric layer keys. Grid placement also uses sector keys such as z,y,x,0 and local 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.

Requirements and library installation: chapter 82.

Go example

examples/86-apple-json-routing/main.go

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

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

import (
	_ "embed"
	"flag"
	"fmt"
	"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/nativemetal"
	"github.com/openfluke/welvet/quant"
	"github.com/openfluke/welvet/runtime/placement"
	"log"
	"os"
	"runtime"
)

//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 != "darwin" || runtime.GOARCH != "arm64" {
		fmt.Println("SKIP: Apple Metal needs Apple silicon macOS 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 := nativemetal.EnableRouting()
	if e != nil {
		return e
	}
	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 {
		return e
	}
	stats := router.Stats()
	if stats.Backwards == 0 || stats.Updates == 0 {
		return fmt.Errorf("plan did not execute native Metal backward/SGD")
	}
	fmt.Printf("loss=%.8g metal 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=metal
layer=2 backend=metal
loss=0.055711714 metal routing={Mounts:2 Forwards:4 Backwards:2 Updates:2 TrainingSteps:0}
Routed updates are already in Go stores; Close releases cached mounts.

Validate (harness)

cd welvet/examples/86-apple-json-routing && source ../env.sh && go run .