wagogpu

README

wago-gpu

An optional Go plugin for Wago. It translates selected WebAssembly instructions into WGSL and runs elementwise kernels on a GPU. Wago is an unchanged dependency.

The buffer API has 34 imports, CPU fallback, checked memory32/memory64/GC transfers, typed storage, and explicit dispatch. Kernels can share buffers without an intermediate guest transfer. The earlier scalar wago_gpu.run and run_batch imports remain available.

v0.0.0 is an experimental release. The ABI remains provisional. See the specification, current results, and canonical ABI table.

Build and test

Go 1.25 or later is required. The tested Go version is 1.27.1. The module path is github.com/jtenner/wago-gpu.

git clone https://github.com/jtenner/wago-gpu.git
cd wago-gpu

Run the commands below from this checkout. GPU builds need the local dependency repairs selected by this repository's go.mod. Go does not apply a dependency's replace directives to a consuming module. A GPU host in another module must apply both replacements to this checkout as well; the unpatched upstream binding is not supported. For a host module beside the checkout, use:

go get github.com/jtenner/wago-gpu@v0.0.0
go mod edit -replace github.com/oliverbestmann/webgpu=../wago-gpu/third_party/webgpu
go mod edit -replace github.com/oliverbestmann/webgpu/libs-linux=../wago-gpu/third_party/webgpu/libs-linux

Build the native library in the checkout with ./native/build.sh before building that host with -tags webgpu. The CPU path needs no native replacement build.

# No GPU, native library, or C compiler is needed.
CGO_ENABLED=0 go test ./...
go run ./examples/buffers -first-cpu -second-cpu
go run ./examples/tinygo -cpu

# Complete generated-Wasm tests also need wasm-tools on PATH.
go test -race ./...

# GPU build: Linux, Vulkan driver, C compiler, Rust/Cargo, and Git.
# This builds the pinned native dependency with the checked-error patch.
./native/build.sh
go run -tags webgpu ./examples/buffers -require-gpu
go run -tags webgpu ./examples/tinygo -require-gpu
WAGO_GPU_TEST=1 go test -race -tags webgpu ./...

# Benchmark both math kernels at all four required array sizes.
go run -tags webgpu ./cmd/buffer-bench > buffer-results.jsonl

-require-gpu and WAGO_GPU_TEST=1 fail when a hardware GPU cannot run. Software adapters are rejected. Without webgpu, or without CGO, the plugin uses CPU fallback. GPU support is currently enabled only on Linux. Other platforms use the CPU path.

The GPU dependency has local callback and error-reporting repairs. Read the patch notes before changing it. These repairs include a small native Rust patch and C binding changes. No GPU code or patch is inside Wago. Native driver calls cannot be forcibly interrupted by a Go context.

Complete examples

The WAT start module creates two F32 buffers, computes B = 2*A, then A = B+1, and copies [3, 5, 7, 9] to guest memory. Its Go host checks the result after instantiation. Each dispatch can fall back independently:

go run -tags webgpu ./examples/buffers -first-cpu
go run -tags webgpu ./examples/buffers -second-cpu

The complete TinyGo guest uses the single-result packed allocation import. It includes allocation, transfer, dispatch, CPU fallback, result checks, and cleanup. Rebuild its checked-in Wasm with:

tinygo build -target=wasm-unknown -scheduler=none -panic=trap \
  -gc=leaking -opt=2 -no-debug -o examples/tinygo/guest.wasm \
  examples/tinygo/guest.go

TinyGo 0.42.0 was tested. The tiny guest uses fixed arrays. -gc=leaking is a build setting for this example, not a requirement for plugin-owned buffers.

Typed storage examples are also available:

go run ./examples/storage -kind f16 -cpu
go run ./examples/storage -kind memory64 -cpu
go run ./examples/storage -kind gc -cpu
go run -tags webgpu ./examples/storage -kind f16 -require-gpu

The storage WAT files show F16 upcasting and explicit transfers. The memory64 and GC examples test transfers; they do not dispatch GPU work. Regeneration needs wasm-tools.

Host contract

Use NewHost for the controlled compilation path:

host, err := wagogpu.NewHost(ctx, wagogpu.Config{
    Kernels: []wagogpu.KernelConfig{{
        ID: 1,
        Export: "wago_gpu.kernel.double",
        CPUExport: "wago_gpu.cpu.double",
        RelaxedFloat: true,
        Bindings: []wagogpu.BindingConfig{
            {Slot: 0, Type: wagogpu.TypeF32, Access: wagogpu.AccessRead},
            {Slot: 1, Type: wagogpu.TypeF32, Access: wagogpu.AccessWrite},
        },
    }},
})

Then call host.Compile(wasm) and host.Instantiate(ctx, module). Close the instance, module, and host with an independent cleanup context. See the TinyGo host for a complete sequence that returns cleanup errors.

Host also offers source preparations with Compile and Close. It exposes no artifact-adoption or import-override route. With the lower-level PluginSet, the host must enforce these restrictions itself:

• Use Runtime.Compile or PreparedCompile.Compile.
• Do not adopt an unrelated compiled artifact. A source digest alone cannot prove that an adopted artifact matches inspected source.
• Do not override imports that the shader compiler treats as intrinsics.
• Consume or close prepared compilations. Pending plugin records are bounded; an evicted preparation that finishes later is unverified.

Missing exports, wrong signatures, and established contract violations return INVALID_KERNEL. Unsupported GPU instructions use fallback only after current CPU data is available. Unknown final source metadata is never trusted.

Guest contract

A kernel exports (i32 index) -> (). Its CPU runner exports (i32 count) -> (). Names use wago_gpu.kernel.<name> and wago_gpu.cpu.<name>. The host selects kernel IDs, names, slot types, access permissions, and float consent.

Buffers belong to one instance. Handles are opaque nonzero i32 bit patterns. The guest uses createBufferPacked, bindBuffer, setBuffer32/64/GC, dispatch, copyBuffer32/64/GC, and freeBuffer. A kernel uses getBuffer and typed readBuffer* / writeBuffer* intrinsics.

Dispatch status 0 means success. Status 1 means the guest must run its original CPU runner. Other statuses are errors; do not treat them as fallback. BufferSnapshot().Last preserves the dispatch record through the CPU loop. Scalar intrinsics do not replace it.

GPU outputs remain private until all checks and completion succeed. A failed operation does not publish partial output. Device loss or uncertain completion quarantines the runtime's GPU. Current CPU copies remain usable. GPU-only contents become lost; a full successful set operation can replace them.

A later trap or cancellation of the enclosing Wasm call does not undo an import that already committed. Do not automatically retry a chain on that assumption.

Supported compiler subset

• local.get, local.set, local.tee for i32/f32 values.
• i32.const, f32.const, f32.add, f32.sub, f32.mul.
• Exact buffer intrinsic imports and a structural final end; restricted void return directly before that end.
• One independent invocation per element. Access uses the original index, including unchanged copies through locals.

GPU types: I8, U8, I16, U16, I32, U32, F16, F32. Narrow elements use separate 32-bit GPU words. I64, U64, and F64 have CPU imports and require fallback. There is no general integer arithmetic, gather/scatter, reduction, matrix compiler, arbitrary control flow, or automatic function replacement.

CompileBufferWGSL returns inspectable WGSL. Its *CompileError category is available through errors.As, including after wrapping. Export names are not inserted into shader identifiers.

F16 conversion and float copies use integer bits. Scalar F16 conversions use round-to-nearest, ties-to-even and canonical NaNs. Bulk copies preserve bits. GPU float arithmetic can fuse operations and change subnormal or non-finite behavior. RelaxedFloat is required per float kernel. Exact transfer is not a promise of exact arithmetic.

Limits, measurements, and maintenance

Defaults are 64 buffers per instance, 64 selected kernels, 8 slots per kernel, 10 million elements per buffer, 256 MiB tracked instance storage, and 512 MiB tracked runtime storage. Hosts can lower them. One idle scratch buffer per instance can be reused; idle capacity is reclaimed before rejecting a budget reservation. Driver overhead is outside these byte counts.

Snapshots distinguish submission from confirmed execution, guest copies from host/device transfers, and device-local seed copies. GPU timestamps are not available: TimestampsValid is false. Profile fields are host wall times and include queue waits. ProfileStages adds waits, so benchmark results state whether it is enabled.

Source-transform hooks disable Wago prepared-artifact cache eligibility. Pipelines are cached while their modules or instances own them. No Wago public API change is required under the documented host restrictions.

go generate ./...                 # Requires WABT and wasm-tools for fixtures.
python3 spec/check_spec.py --tinygo
go test -run '^$' -fuzz FuzzBufferCompiler -fuzztime 10s .

The module path and plugin provenance identify jtenner/wago-gpu. The ABI remains provisional in v0.0.0. Historical scalar measurements remain in REPORT.md and IMPROVEMENTS.md; they are not buffer benchmark results.

Documentation

Overview

Code generated by go generate; DO NOT EDIT.

Package wagogpu provides an optional Wago plugin for small f32 GPU kernels.

Index

• Constants
• func CompileBufferWGSL(source []byte, kernel KernelConfig) (string, error)
• func CompileWGSL(source []byte, export string) (string, error)
• type BindingConfig
• type BufferAccess
• type BufferOperation
• type BufferSnapshot
• type BuildTiming
• type CompileError
  • func (e *CompileError) Error() string
• type CompileErrorKind
• type Config
• type ElementType
• type Host
  • func NewHost(ctx context.Context, config Config) (*Host, error)
  • func (h *Host) BufferSnapshot() BufferSnapshot
  • func (h *Host) Close(ctx context.Context) error
  • func (h *Host) Compile(source []byte) (*Module, error)
  • func (h *Host) Instantiate(ctx context.Context, module *Module) (*wago.Instance, error)
  • func (h *Host) Prepare(source []byte) (*Preparation, error)
  • func (h *Host) Snapshot() Snapshot
• type InstanceDiagnostic
• type KernelConfig
• type KernelDiagnostic
• type Module
  • func (m *Module) Close() error
• type Plugin
  • func New(config Config) (*Plugin, error)
  • func (p *Plugin) BufferSnapshot() BufferSnapshot
  • func (p *Plugin) PluginSet() (wago.PluginSet, error)
  • func (p *Plugin) Register(reg *wago.Registrar) error
  • func (p *Plugin) Snapshot() Snapshot
• type Preparation
  • func (p *Preparation) Close() error
  • func (p *Preparation) Compile() (*Module, error)
• type Snapshot
• type Timing
• type TransferCounters

Constants

const (
	TypeI8               ElementType = 1
	TypeU8               ElementType = 2
	TypeI16              ElementType = 3
	TypeU16              ElementType = 4
	TypeI32              ElementType = 5
	TypeU32              ElementType = 6
	TypeF16              ElementType = 7
	TypeF32              ElementType = 8
	TypeI64              ElementType = 9
	TypeU64              ElementType = 10
	TypeF64              ElementType = 11
	V1OK                 int32       = 0
	V1CPUFallback        int32       = 1
	V1InvalidHandle      int32       = 2
	V1InvalidRange       int32       = 3
	V1TypeMismatch       int32       = 4
	V1InvalidBinding     int32       = 5
	V1LimitExceeded      int32       = 6
	V1DeviceError        int32       = 7
	V1ContentsLost       int32       = 8
	V1Cancelled          int32       = 9
	V1InvalidKernel      int32       = 10
	V1UnsupportedStorage int32       = 11
	V1InvalidState       int32       = 12
)

const (
	Success int32 = iota
	Fallback
	InvalidRange
	InvalidBatch
	DefaultMaxElements uint32 = 10_000_000
	MaxBatchPasses     uint32 = 64
)

Functions

func CompileBufferWGSL

func CompileBufferWGSL(source []byte, kernel KernelConfig) (string, error)

CompileBufferWGSL translates the selected kernel, without creating a device. Every failure is a *CompileError. This decoder is bounded; Wago remains the validator for the whole module before any runtime pipeline can be attached.

func CompileWGSL

func CompileWGSL(source []byte, export string) (string, error)

CompileWGSL translates one explicitly named, defined Wasm function export. This is a bounded subset decoder, not a general Wasm validator. Wago validates the entire module before the plugin creates a GPU pipeline.

Types

type BindingConfig

type BindingConfig struct {
	Slot   uint32
	Type   ElementType
	Access BufferAccess
}

type BufferAccess

type BufferAccess uint8

const (
	AccessRead BufferAccess = iota + 1
	AccessWrite
	AccessReadWrite
)

type BufferOperation

type BufferOperation struct {
	Sequence                                                                                   uint64
	Instance                                                                                   wago.InstanceIdentity
	KernelID                                                                                   uint32
	Import                                                                                     string
	Status                                                                                     int32
	ReasonCode, Reason, Outcome                                                                string
	Count                                                                                      uint32
	HardwareSubmitted, TransferCountsComplete, ProfileValid, TimestampsValid                   bool
	QueueWait, Total, Upload, DeviceCopy, Compute, Download, Conversion, Commit, DeviceCompute time.Duration
	TransferCounters
}

type BufferSnapshot

type BufferSnapshot struct {
	PeakRuntimeBufferBytes                                     uint64
	State, Reason, DeviceState                                 string
	Last                                                       BufferOperation
	Totals                                                     TransferCounters
	TotalsComplete                                             bool
	Kernels                                                    []KernelDiagnostic
	Instances                                                  []InstanceDiagnostic
	PipelineHits, PipelineMisses                               uint64
	PendingCompilations, PendingContracts, PendingTranslations uint32
	PendingContractBytes, PendingWGSLBytes                     uint64
	ActiveCompilations                                         uint32
	RuntimeBufferBytes, RuntimePoolBytes, LegacyBufferBytes    uint64
}

type BuildTiming

type BuildTiming struct{ Translate, Pipeline time.Duration }

type CompileError

type CompileError struct {
	Kind    CompileErrorKind
	Message string
}

CompileError is inspectable with errors.As, including after wrapping.

func (*CompileError) Error

func (e *CompileError) Error() string

type CompileErrorKind

type CompileErrorKind uint8

const (
	CompileInvalidContract CompileErrorKind = iota + 1
	CompileUnsupported
	CompileLimit
)

type Config

type Config struct {
	Kernels                                                          []KernelConfig
	MaxBuffersPerInstance, MaxKernelsPerModule, MaxBindingsPerKernel uint32
	MaxInstanceBufferBytes, MaxRuntimeBufferBytes                    uint64
	KernelExport                                                     string
	RelaxedFloat                                                     bool
	Disabled                                                         bool
	// ProfileStages adds queue waits to measure each stage. Leave false for use.
	ProfileStages bool
	MaxElements   uint32
	// MinElements selects CPU below this size. Zero always attempts GPU execution.
	// The limit applies to a complete batch, not to each pass separately.
	MinElements uint32
	// RunTimeout bounds result polling. Native driver entry points may still block.
	RunTimeout time.Duration
}

Config selects one export per module. RelaxedFloat must be an explicit opt-in: GPU f32 can fuse operations and flush subnormals (see README).

type ElementType

type ElementType uint32

type Host

type Host struct {
	// contains filtered or unexported fields
}

Host is the controlled v1 integration. It exposes source compilation only and supplies no per-instance import overrides. It never adopts compiled artifacts. Use this type when the host cannot otherwise enforce the source-identity rules.

func NewHost

func NewHost(ctx context.Context, config Config) (*Host, error)

func (*Host) BufferSnapshot

func (h *Host) BufferSnapshot() BufferSnapshot

func (*Host) Close

func (h *Host) Close(ctx context.Context) error

func (*Host) Compile

func (h *Host) Compile(source []byte) (*Module, error)

func (*Host) Instantiate

func (h *Host) Instantiate(ctx context.Context, module *Module) (*wago.Instance, error)

func (*Host) Prepare

func (h *Host) Prepare(source []byte) (*Preparation, error)

func (*Host) Snapshot

func (h *Host) Snapshot() Snapshot

type InstanceDiagnostic

type InstanceDiagnostic struct {
	PeakBytes                                                                                      uint64
	Instance                                                                                       wago.InstanceIdentity
	Buffers, CPUCurrent, GPUCurrent, BothCurrent, ContentsLost                                     uint32
	LogicalBytes, CPUBytes, GPUBytes, StagingBytes, ScratchBytes, RetainedPoolBytes, ReservedBytes uint64
}

type KernelConfig

type KernelConfig struct {
	ID                uint32
	Export, CPUExport string
	Bindings          []BindingConfig
	RelaxedFloat      bool
	MinElements       uint32
}

type KernelDiagnostic

type KernelDiagnostic struct {
	Translate, Pipeline time.Duration
	Module              wago.ModuleIdentity
	KernelID            uint32
	State, Reason       string
	ContractVerified    bool
}

type Module

type Module struct {
	// contains filtered or unexported fields
}

func (*Module) Close

func (m *Module) Close() error

type Plugin

type Plugin struct {
	// contains filtered or unexported fields
}

Plugin owns all GPU state. Make a new Plugin for each Wago Runtime. Different instances may call it concurrently; GPU work is serialized so that cached buffers can be reused. Wago controls calls within a single instance.

func New

func New(config Config) (*Plugin, error)

func (*Plugin) BufferSnapshot

func (p *Plugin) BufferSnapshot() BufferSnapshot

func (*Plugin) PluginSet

func (p *Plugin) PluginSet() (wago.PluginSet, error)

PluginSet grants only this experiment's declared authorities. The host opts into these grants by passing this set to Runtime.LoadPlugins.

func (*Plugin) Register

func (p *Plugin) Register(reg *wago.Registrar) error

func (*Plugin) Snapshot

func (p *Plugin) Snapshot() Snapshot

type Preparation

type Preparation struct {
	// contains filtered or unexported fields
}

func (*Preparation) Close

func (p *Preparation) Close() error

func (*Preparation) Compile

func (p *Preparation) Compile() (*Module, error)

type Snapshot

type Snapshot struct {
	Device, Reason                         string
	GPUFailed                              bool
	DeviceInit                             time.Duration
	Build                                  BuildTiming
	Last                                   Timing
	Successes, Fallbacks                   uint64
	Modules, Instances, Pipelines, Pending int
	BufferBytes                            uint64
}

type Timing

type Timing struct {
	Upload, Compute, Download, Commit, Total time.Duration
	DeviceCompute                            time.Duration
	TimestampsValid                          bool
	UploadBytes, DownloadBytes               uint64
	Passes                                   uint32
}

Timing contains host wall times. DeviceCompute and TimestampsValid remain zero and false until a supported backend exposes reliable GPU timestamps. Stage times are present only with ProfileStages. Total includes guest commit.

type TransferCounters

type TransferCounters struct {
	GuestSetBytes, GuestCopyBytes, GuestSetCount, GuestCopyCount                                           uint64
	LogicalUploadBytes, LogicalDownloadBytes, GPUUploadBytes, GPUDownloadBytes, UploadCount, DownloadCount uint64
	DeviceCopyBytes, DeviceCopyCount, ParameterUploadBytes, ParameterUploadCount                           uint64
}