dex/pkg/lxgpu/orderbook_bridge.go

// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

//go:build cgo

// dexvm_gpu.go — direct dlopen + dlsym to the gpu-kernels DEX plugin.
//
// Why dlopen, not pkg-config: the OrderBook host launchers live in a
// per-backend plugin (libluxgpu_backend_{cuda,hip,metal,vulkan,webgpu})
// whose presence is platform-dependent — Apple has only Metal+Vulkan
// (via MoltenVK), Linux has CUDA+Vulkan+WebGPU, ROCm only on a small
// set of AMD hosts. pkg-config would force a build-time decision; we
// want a runtime probe so the same go binary picks up whichever
// plugin is on disk. This mirrors the dlopen pattern luxcpp/gpu uses
// for its own backend loader (luxcpp/gpu/test/test_plugin_loader.cpp).
//
// Host launcher ABI — exposed by every plugin under the per-backend
// symbol prefix `lux_<backend>_*`. The C signatures come straight from
// the GPU plugin install tree backends/{metal/src/dex_launchers.mm,
// vulkan/src/dex_launchers.cpp}:
//
//   int dex_orderbook_match(void* arena, const uint8_t* calldata,
//                      uint8_t* out, uint32_t* num_fills);
//   int dex_orderbook_arena_create(void** out_arena);
//   int dex_orderbook_arena_destroy(void* arena);
//
// All inputs are HOST pointers (no separate device address space is
// exposed across the ABI — the launcher does H2D/D2H copies
// internally). `arena` for OrderBook is opaque — created/destroyed via the symmetric
// pair, persisted across calls so the BookArena lives device-resident
// for the lifetime of the book.

package lxgpu

/*
#cgo LDFLAGS: -ldl

#include <dlfcn.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

// Typed function-pointer typedefs for the launcher symbols. cgo can't
// cast `unsafe.Pointer` directly into a typed C function pointer, so we
// route every call through a tiny trampoline that takes the dlsym'd
// pointer and the typed arguments.

typedef int (*lux_dex_orderbook_match_fn)(
    void*          arena,
    const uint8_t* calldata,
    uint8_t*       out,
    uint32_t*      num_fills);

typedef int (*lux_dex_orderbook_arena_create_fn)(void** out_arena);
typedef int (*lux_dex_orderbook_arena_destroy_fn)(void* arena);

static int call_dex_orderbook_match(
    void* fn,
    void* arena,
    const uint8_t* calldata,
    uint8_t* out,
    uint32_t* num_fills)
{
    return ((lux_dex_orderbook_match_fn)fn)(arena, calldata, out, num_fills);
}

static int call_dex_orderbook_arena_create(void* fn, void** out_arena) {
    return ((lux_dex_orderbook_arena_create_fn)fn)(out_arena);
}

static int call_dex_orderbook_arena_destroy(void* fn, void* arena) {
    return ((lux_dex_orderbook_arena_destroy_fn)fn)(arena);
}

// dlopen / dlsym wrappers. dlerror() returns thread-local NULL when
// the last call succeeded; we call it after every failed lookup to
// produce a meaningful error string back in Go.
static void* lux_dlopen(const char* path) {
    return dlopen(path, RTLD_NOW | RTLD_LOCAL);
}

static void lux_dlclose(void* handle) {
    if (handle) dlclose(handle);
}

static void* lux_dlsym(void* handle, const char* name) {
    return dlsym(handle, name);
}

static const char* lux_dlerror(void) {
    return dlerror();
}
*/
import "C"

import (
	"fmt"
	"os"
	"runtime"
	"sync"
	"unsafe"
)

// gpuHandle bundles the dlopen handle and the four resolved symbol
// pointers for one backend. Created by init() exactly once per
// process; safe for concurrent use from multiple goroutines because
// the underlying plugin uses thread-safe per-backend launchers (Metal
// caches an MTLPipelineState behind a once-flag; Vulkan keeps a
// process-global VkDevice + pool; CUDA/HIP launchers are reentrant).
type gpuHandle struct {
	backend GPUBackend
	libPath string

	hLib unsafe.Pointer // dlopen handle

	fnOrderBookMatch unsafe.Pointer // lux_<backend>_dex_orderbook_match
	fnArenaCreate    unsafe.Pointer // lux_<backend>_dex_orderbook_arena_create
	fnArenaDestroy   unsafe.Pointer // lux_<backend>_dex_orderbook_arena_destroy
}

// globalGPU holds the result of the init() probe. nil after a clean
// init when no plugin was found on disk — every method returns
// ErrGPUNotAvailable in that case.
var (
	globalGPU   *gpuHandle
	globalGPUMu sync.RWMutex
)

// init runs the dlopen probe at package load time. Order is
// cuda → hip → metal → vulkan → webgpu, matching the GPUBackend
// numeric ordering. The first plugin whose dylib loads AND exposes
// all four required symbols wins; later candidates are skipped.
//
// Search path order, per candidate:
//  1. ${LUX_GPU_PLUGIN_DIR}/lib<name>.<ext>  (operator-set override)
//  2. lib<name>.<ext>                        (system dlopen lookup —
//     LD_LIBRARY_PATH, DYLD_LIBRARY_PATH, rpath, default loader paths)
func init() {
	for _, b := range probeOrder() {
		if h := tryLoadBackend(b); h != nil {
			globalGPUMu.Lock()
			globalGPU = h
			globalGPUMu.Unlock()
			return
		}
	}
	// No plugin available — every method returns ErrGPUNotAvailable.
	// Intentionally NOT a panic: a node on a CPU-only host is fine,
	// it just falls back to the in-Go orderbook path.
}

// probeOrder returns the canonical dlopen probe order. CUDA first so
// NVIDIA hosts pick up the native ICD; HIP next for AMD ROCm; Metal
// is the Apple Silicon native path; Vulkan covers the rest (cross-
// vendor portable, MoltenVK on Apple, native ICDs on Linux/Windows);
// WebGPU (wgpu / Dawn) is the universal last-resort.
func probeOrder() []GPUBackend {
	return []GPUBackend{
		GPUBackendCUDA,
		GPUBackendHIP,
		GPUBackendMetal,
		GPUBackendVulkan,
		GPUBackendWebGPU,
	}
}

// dylibExt returns the platform-specific shared-library extension. We
// search both the platform-native form (.dylib on darwin, .so on
// linux, .dll on windows-via-cgo) and the unix-style .so as a
// fallback — the Vulkan loader, for example, sometimes lives under
// libluxgpu_backend_vulkan.so even on darwin (when built by CMake's
// MODULE_LIBRARY default).
func dylibExt() []string {
	switch runtime.GOOS {
	case "darwin":
		return []string{".dylib", ".so"}
	case "windows":
		return []string{".dll"}
	default:
		return []string{".so"}
	}
}

// pluginCandidatePaths returns the dlopen candidates for one backend,
// in lookup order: the operator override directory first (when set),
// then the bare library name so the system loader's standard search
// path takes over (LD_LIBRARY_PATH, DYLD_LIBRARY_PATH, rpath, default
// loader paths).
func pluginCandidatePaths(b GPUBackend) []string {
	base := "libluxgpu_backend_" + b.String()
	exts := dylibExt()

	var out []string
	if dir := os.Getenv("LUX_GPU_PLUGIN_DIR"); dir != "" {
		for _, ext := range exts {
			out = append(out, dir+"/"+base+ext)
		}
	}
	for _, ext := range exts {
		out = append(out, base+ext)
	}
	return out
}

// tryLoadBackend dlopen's the candidate paths for one backend, and
// resolves the three required OrderBook symbols. Returns nil if either
// the dylib won't load OR any required symbol is missing. Symbol misses
// are treated as fatal-for-this-backend (we don't accept a partial
// plugin — that would mask a build skew and produce a "GPU loaded
// but OrderBook returns garbage" failure mode that's worse than CPU).
func tryLoadBackend(b GPUBackend) *gpuHandle {
	candidates := pluginCandidatePaths(b)
	for _, path := range candidates {
		// Clear dlerror() before dlopen so a previous miss doesn't
		// confuse the post-failure diagnostic.
		C.lux_dlerror()
		cpath := C.CString(path)
		hLib := C.lux_dlopen(cpath)
		C.free(unsafe.Pointer(cpath))
		if hLib == nil {
			continue
		}

		prefix := "lux_" + b.String() + "_"
		var (
			fnOrderBook = dlsymOrNil(hLib, prefix+"dex_orderbook_match")
			fnCreate    = dlsymOrNil(hLib, prefix+"dex_orderbook_arena_create")
			fnDestr     = dlsymOrNil(hLib, prefix+"dex_orderbook_arena_destroy")
		)
		if fnOrderBook == nil || fnCreate == nil || fnDestr == nil {
			// Partial export set — plausible cause is a stale plugin
			// build that pre-dates the DEX launcher landing. Close
			// the handle so we don't keep a dangling mapping, and
			// keep probing.
			C.lux_dlclose(hLib)
			continue
		}
		return &gpuHandle{
			backend:          b,
			libPath:          path,
			hLib:             hLib,
			fnOrderBookMatch: fnOrderBook,
			fnArenaCreate:    fnCreate,
			fnArenaDestroy:   fnDestr,
		}
	}
	return nil
}

// dlsymOrNil wraps dlsym() — returns nil on lookup failure so the
// caller can fail uniformly without parsing dlerror() strings.
func dlsymOrNil(h unsafe.Pointer, name string) unsafe.Pointer {
	cname := C.CString(name)
	defer C.free(unsafe.Pointer(cname))
	C.lux_dlerror()
	p := C.lux_dlsym(h, cname)
	if p == nil {
		return nil
	}
	return p
}

// AutoBackend returns the GPU backend bound at init(), or
// GPUBackendNone if no plugin loaded. Callers branch on this to
// decide between the GPU path and the CPU fallback.
func AutoBackend() GPUBackend {
	globalGPUMu.RLock()
	defer globalGPUMu.RUnlock()
	if globalGPU == nil {
		return GPUBackendNone
	}
	return globalGPU.backend
}

// GPUPluginPath returns the dylib path the active backend was loaded
// from, useful for diagnostic logging at node startup. Empty string
// when no plugin loaded.
func GPUPluginPath() string {
	globalGPUMu.RLock()
	defer globalGPUMu.RUnlock()
	if globalGPU == nil {
		return ""
	}
	return globalGPU.libPath
}

// activeGPU returns the loaded handle or ErrGPUNotAvailable. Every
// public method funnels through this so the nil-check is in one place.
func activeGPU() (*gpuHandle, error) {
	globalGPUMu.RLock()
	defer globalGPUMu.RUnlock()
	if globalGPU == nil {
		return nil, ErrGPUNotAvailable
	}
	return globalGPU, nil
}

// OrderBookArena is the opaque handle to a device-resident BookArena. The
// pointer crosses the cgo boundary as void* — Go never dereferences
// it, only passes it to ArenaDestroy or OrderBookMatch.
//
// ArenaCreate hands ownership to the Go caller; ArenaDestroy releases
// it. The arena lives across many OrderBookMatch calls — that's the
// entire reason it's device-resident: the BookArena holds the
// resting bids/asks for a single book, and successive OrderBookMatch
// calls on the same book MUST share state.
type OrderBookArena struct {
	ptr unsafe.Pointer
}

// ArenaCreate allocates a fresh device-resident BookArena (zero
// initialized). The returned handle is owned by the caller; call
// ArenaDestroy to free it. Concurrent ArenaCreate is safe.
func ArenaCreate() (*OrderBookArena, error) {
	h, err := activeGPU()
	if err != nil {
		return nil, err
	}
	var raw unsafe.Pointer
	rc := C.call_dex_orderbook_arena_create(h.fnArenaCreate, &raw)
	if rc != 0 {
		return nil, fmt.Errorf("dexvm.ArenaCreate: %s plugin returned rc=%d",
			h.backend, int(rc))
	}
	if raw == nil {
		return nil, fmt.Errorf("dexvm.ArenaCreate: %s plugin returned NULL arena",
			h.backend)
	}
	return &OrderBookArena{ptr: raw}, nil
}

// ArenaDestroy releases the device-resident BookArena. Calling it on
// a nil arena is a no-op (matches the C-side symmetry). Calling it
// twice on the same arena is UB at the plugin level — don't.
func ArenaDestroy(a *OrderBookArena) error {
	if a == nil || a.ptr == nil {
		return nil
	}
	h, err := activeGPU()
	if err != nil {
		return err
	}
	rc := C.call_dex_orderbook_arena_destroy(h.fnArenaDestroy, a.ptr)
	a.ptr = nil
	if rc != 0 {
		return fmt.Errorf("dexvm.ArenaDestroy: %s plugin returned rc=%d",
			h.backend, int(rc))
	}
	return nil
}

// OrderBookMatch runs one matcher step against the arena. calldata is the
// 117-byte EVM precompile 0x100 input — see backend.go for the byte
// layout. Output is a fixed-size 68-byte buffer (filled + avg_price +
// num_fills) plus a num_fills uint32 mirror for callers that want it
// without re-parsing the 68 bytes.
//
// Single matcher pass per call; concurrency is across book_ids — each
// book gets its own arena, and independent arenas can be matched in
// parallel on different goroutines without serialization.
func OrderBookMatch(a *OrderBookArena, calldata []byte) (out [OrderBookOutLen]byte, numFills uint32, err error) {
	h, gerr := activeGPU()
	if gerr != nil {
		return out, 0, gerr
	}
	if a == nil || a.ptr == nil {
		return out, 0, fmt.Errorf("dexvm.OrderBookMatch: arena is nil")
	}
	if len(calldata) != OrderBookCalldataLen {
		return out, 0, fmt.Errorf("dexvm.OrderBookMatch: calldata len=%d, want %d",
			len(calldata), OrderBookCalldataLen)
	}

	var pinner runtime.Pinner
	defer pinner.Unpin()
	pinner.Pin(&calldata[0])
	pinner.Pin(&out[0])

	var nf C.uint32_t
	rc := C.call_dex_orderbook_match(
		h.fnOrderBookMatch,
		a.ptr,
		(*C.uint8_t)(unsafe.Pointer(&calldata[0])),
		(*C.uint8_t)(unsafe.Pointer(&out[0])),
		&nf,
	)
	runtime.KeepAlive(calldata)
	runtime.KeepAlive(out)
	if rc != 0 {
		return out, 0, fmt.Errorf("dexvm.OrderBookMatch: %s plugin returned rc=%d",
			h.backend, int(rc))
	}
	return out, uint32(nf), nil
}