86. Apple Metal: JSON backend placement

Part: XII · Apple acceleration (welvet.metal)
Package: nativemetal + runtime/placement
Status: ok: verified macOS/arm64 on Apple M5

When

Assign Go and Metal operations using the shared execution JSON.

Where

From example/welvet/86-apple-json-routing:

go run .
go run . -config placement.json

Why

Use ordinary Forward/TrainStackMSE calls after installing the routing owner.

What

Embeds placement.json: layer 0 runs Go CPU, layers 1-2 run Metal. EnableRouting retains native mounts and returns updated parameters to Go automatically. Close the router after the workload.

Requirements

Apple silicon macOS, Go with cgo, and the installed independent Metal dylib. Set WELVET_METAL_LIBRARY or install it from the Metal checkout:

cd /Users/openfluke/gs/welvet.metal
./scripts/macos_install_go.sh

Use the Go v1.3.0 checkout containing nativemetal. Examples 82-86 can also use go run -tags welvet_metal . to exclude the independent C++ binding. Example 87 needs the normal build to keep both native backends. Unsupported operating systems print SKIP. Missing libraries on a supported host report an error.

These are usage examples with small shapes. For speed measurements, use ../apple_stress_test. The results are numerical CPU/GPU comparisons, not model prediction accuracy.

Backend labels in routed examples: cpu-tiled is Go CPU, native_cpu is welvet.cpp, and metal is the independent Apple GPU backend.

Captured sample

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.

main.go

Download source ↓

// 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
}