Vulkan: JSON backend placement

github.com/openfluke/welvet/nativevulkanlinux amd64 routing

Why it exists

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

What it is

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

Explicit placement

The example embeds placement.json; use -config to load another file. nativevulkan.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.

Requirements: chapter 103.

Go example

examples/107-vulkan-json-routing/main.go

Run:cd welvet/examples/107-vulkan-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/nativevulkan"
	"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: Vulkan 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 := nativevulkan.EnableRouting()
	if e != nil {
		fmt.Printf("SKIP: Vulkan 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: Vulkan 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 Vulkan backward/SGD")
	}
	fmt.Printf("loss=%.8g vulkan 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=vulkan
layer=2 backend=vulkan
loss=0.055711714 vulkan routing={Mounts:2 Forwards:4 Backwards:2 Updates:2 TrainingSteps:0}
Routed updates are already in Go stores; Close releases cached mounts.

Authorization required, but no authorization protocol specified

Authorization required, but no authorization protocol specified

Validate (harness)

cd welvet/examples/107-vulkan-json-routing && source ../env.sh && go run .