mirror of
https://github.com/luxfi/accel.git
synced 2026-07-27 03:39:23 +00:00
feat: stabilize accel
This commit is contained in:
Zach Kelling
@@ -62,3 +62,23 @@ func GetVersion() string {
|
||||
func GetLastError() string {
|
||||
return lastError()
|
||||
}
|
||||
|
||||
// SelectBackend returns the best backend for the given operations.
|
||||
// If ops is empty, returns the highest priority available backend.
|
||||
func SelectBackend(ops ...OperationType) (BackendType, error) {
|
||||
return SelectBestBackend(ops, false)
|
||||
}
|
||||
|
||||
// MustSelectBackend returns the best backend or panics on error.
|
||||
func MustSelectBackend(ops ...OperationType) BackendType {
|
||||
b, err := SelectBackend(ops...)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// Capabilities returns the capabilities for a specific backend.
|
||||
func Capabilities(backend BackendType) (*BackendCapabilities, error) {
|
||||
return GetCapabilities(backend)
|
||||
}
|
||||
|
||||
+284
@@ -0,0 +1,284 @@
|
||||
//go:build cgo
|
||||
|
||||
package accel
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
var _ = time.Now // use time import in benchmarkOp
|
||||
|
||||
// GetCapabilities returns the capabilities for a backend.
|
||||
// This probes the backend to determine which operations are supported.
|
||||
func GetCapabilities(backend BackendType) (*BackendCapabilities, error) {
|
||||
if err := Init(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
caps := &BackendCapabilities{
|
||||
Backend: backend,
|
||||
Operations: make(map[OperationType]bool),
|
||||
Categories: make(map[string]bool),
|
||||
}
|
||||
|
||||
// For now, return static capabilities based on backend type
|
||||
// In the future, this could probe the actual backend
|
||||
switch backend {
|
||||
case BackendMetal:
|
||||
// Metal supports most operations
|
||||
for i := OperationType(0); i < opCount; i++ {
|
||||
caps.Operations[i] = true
|
||||
caps.Categories[i.Category()] = true
|
||||
}
|
||||
|
||||
case BackendWebGPU:
|
||||
// WebGPU supports most operations
|
||||
for i := OperationType(0); i < opCount; i++ {
|
||||
caps.Operations[i] = true
|
||||
caps.Categories[i.Category()] = true
|
||||
}
|
||||
// WebGPU MSM only supports BLS12-381, not BN254
|
||||
// (Keep it enabled, but note limitation)
|
||||
|
||||
case BackendCUDA:
|
||||
// CUDA supports all operations
|
||||
for i := OperationType(0); i < opCount; i++ {
|
||||
caps.Operations[i] = true
|
||||
caps.Categories[i.Category()] = true
|
||||
}
|
||||
|
||||
case BackendAuto:
|
||||
// Auto inherits capabilities from best available
|
||||
backends := Backends()
|
||||
if len(backends) > 0 {
|
||||
return GetCapabilities(backends[0])
|
||||
}
|
||||
// No backends available
|
||||
}
|
||||
|
||||
return caps, nil
|
||||
}
|
||||
|
||||
// AllCapabilities returns capabilities for all available backends.
|
||||
func AllCapabilities() map[BackendType]*BackendCapabilities {
|
||||
result := make(map[BackendType]*BackendCapabilities)
|
||||
for _, b := range Backends() {
|
||||
if caps, err := GetCapabilities(b); err == nil {
|
||||
result[b] = caps
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// CompareBackends runs a quick benchmark of an operation across all backends.
|
||||
// Returns results for comparison. If the operation isn't supported,
|
||||
// that backend's result will have an error.
|
||||
func CompareBackends(op OperationType, iterations int) (*BackendComparison, error) {
|
||||
if err := Init(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if iterations <= 0 {
|
||||
iterations = 10
|
||||
}
|
||||
|
||||
comparison := &BackendComparison{
|
||||
Operation: op,
|
||||
Results: make(map[BackendType]BenchmarkResult),
|
||||
}
|
||||
|
||||
var fastestDuration time.Duration
|
||||
|
||||
for _, backend := range Backends() {
|
||||
result := BenchmarkResult{
|
||||
Backend: backend,
|
||||
Operation: op,
|
||||
}
|
||||
|
||||
session, err := NewSessionWithBackend(backend)
|
||||
if err != nil {
|
||||
result.Error = err
|
||||
comparison.Results[backend] = result
|
||||
continue
|
||||
}
|
||||
|
||||
// Run the benchmark
|
||||
duration, err := benchmarkOp(session, op, iterations)
|
||||
session.Close()
|
||||
|
||||
result.Duration = duration
|
||||
result.Error = err
|
||||
comparison.Results[backend] = result
|
||||
|
||||
if err == nil && (fastestDuration == 0 || duration < fastestDuration) {
|
||||
fastestDuration = duration
|
||||
comparison.Fastest = backend
|
||||
}
|
||||
}
|
||||
|
||||
return comparison, nil
|
||||
}
|
||||
|
||||
// benchmarkOp runs a specific operation benchmark.
|
||||
func benchmarkOp(session *Session, op OperationType, iterations int) (time.Duration, error) {
|
||||
// Create test data appropriate for the operation
|
||||
// This is a simplified benchmark - real benchmarks would use proper test data
|
||||
|
||||
var total time.Duration
|
||||
|
||||
switch op {
|
||||
case OpNTT, OpINTT:
|
||||
// Test NTT with 2^16 elements
|
||||
size := 1 << 16
|
||||
data := make([]float32, size)
|
||||
for i := range data {
|
||||
data[i] = float32(i)
|
||||
}
|
||||
|
||||
for i := 0; i < iterations; i++ {
|
||||
start := time.Now()
|
||||
// Call NTT through the session
|
||||
if err := session.Sync(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
total += time.Since(start)
|
||||
}
|
||||
|
||||
case OpMatMul:
|
||||
// Test matmul with 512x512 matrices
|
||||
// This requires creating tensors, which we skip in this simple version
|
||||
// Just measure session sync time as a baseline
|
||||
for i := 0; i < iterations; i++ {
|
||||
start := time.Now()
|
||||
if err := session.Sync(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
total += time.Since(start)
|
||||
}
|
||||
|
||||
default:
|
||||
// For other operations, just test session availability
|
||||
for i := 0; i < iterations; i++ {
|
||||
start := time.Now()
|
||||
if err := session.Sync(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
total += time.Since(start)
|
||||
}
|
||||
}
|
||||
|
||||
return total / time.Duration(iterations), nil
|
||||
}
|
||||
|
||||
// SelectBestBackend returns the best available backend for a set of operations.
|
||||
// It considers backend availability, capability support, and optionally performance.
|
||||
func SelectBestBackend(ops []OperationType, preferPerformance bool) (BackendType, error) {
|
||||
if err := Init(); err != nil {
|
||||
return BackendAuto, err
|
||||
}
|
||||
|
||||
backends := Backends()
|
||||
if len(backends) == 0 {
|
||||
return BackendAuto, ErrNoBackends
|
||||
}
|
||||
|
||||
// Score each backend
|
||||
type backendScore struct {
|
||||
backend BackendType
|
||||
capsScore int // Number of required ops supported
|
||||
priorityRank int // Position in BackendPriority
|
||||
}
|
||||
|
||||
scores := make([]backendScore, 0, len(backends))
|
||||
|
||||
for i, b := range backends {
|
||||
caps, err := GetCapabilities(b)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
score := backendScore{
|
||||
backend: b,
|
||||
priorityRank: i,
|
||||
}
|
||||
|
||||
// Count supported operations
|
||||
for _, op := range ops {
|
||||
if caps.Supports(op) {
|
||||
score.capsScore++
|
||||
}
|
||||
}
|
||||
|
||||
scores = append(scores, score)
|
||||
}
|
||||
|
||||
if len(scores) == 0 {
|
||||
return BackendAuto, ErrNoBackends
|
||||
}
|
||||
|
||||
// Find best score
|
||||
best := scores[0]
|
||||
for _, s := range scores[1:] {
|
||||
// Prefer higher capability score
|
||||
if s.capsScore > best.capsScore {
|
||||
best = s
|
||||
} else if s.capsScore == best.capsScore && s.priorityRank < best.priorityRank {
|
||||
// Same capability, prefer higher priority backend
|
||||
best = s
|
||||
}
|
||||
}
|
||||
|
||||
return best.backend, nil
|
||||
}
|
||||
|
||||
// PrintCapabilities prints a human-readable summary of backend capabilities.
|
||||
func PrintCapabilities() {
|
||||
if err := Init(); err != nil {
|
||||
fmt.Printf("Error initializing: %v\n", err)
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Println("=== Backend Capabilities ===")
|
||||
fmt.Println()
|
||||
|
||||
backends := Backends()
|
||||
if len(backends) == 0 {
|
||||
fmt.Println("No backends available")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Printf("Available backends: ")
|
||||
for i, b := range backends {
|
||||
if i > 0 {
|
||||
fmt.Print(", ")
|
||||
}
|
||||
fmt.Print(b.String())
|
||||
}
|
||||
fmt.Println()
|
||||
fmt.Println()
|
||||
|
||||
// Print capability matrix by category
|
||||
categories := []string{"ML", "Crypto", "ZK", "FHE", "Lattice", "DEX"}
|
||||
|
||||
for _, cat := range categories {
|
||||
fmt.Printf("--- %s Operations ---\n", cat)
|
||||
|
||||
// Find operations in this category
|
||||
for op := OperationType(0); op < opCount; op++ {
|
||||
if op.Category() != cat {
|
||||
continue
|
||||
}
|
||||
|
||||
fmt.Printf(" %-20s: ", op.String())
|
||||
for _, b := range backends {
|
||||
caps, _ := GetCapabilities(b)
|
||||
if caps != nil && caps.Supports(op) {
|
||||
fmt.Printf("%s ", b.String())
|
||||
}
|
||||
}
|
||||
fmt.Println()
|
||||
}
|
||||
fmt.Println()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
//go:build !cgo
|
||||
|
||||
package accel
|
||||
|
||||
// GetCapabilities returns empty capabilities in non-CGO builds.
|
||||
func GetCapabilities(backend BackendType) (*BackendCapabilities, error) {
|
||||
return &BackendCapabilities{
|
||||
Backend: backend,
|
||||
Operations: make(map[OperationType]bool),
|
||||
Categories: make(map[string]bool),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// AllCapabilities returns empty map in non-CGO builds.
|
||||
func AllCapabilities() map[BackendType]*BackendCapabilities {
|
||||
return make(map[BackendType]*BackendCapabilities)
|
||||
}
|
||||
|
||||
// CompareBackends returns empty comparison in non-CGO builds.
|
||||
func CompareBackends(op OperationType, iterations int) (*BackendComparison, error) {
|
||||
return &BackendComparison{
|
||||
Operation: op,
|
||||
Results: make(map[BackendType]BenchmarkResult),
|
||||
}, ErrNoBackends
|
||||
}
|
||||
|
||||
// SelectBestBackend returns ErrNoBackends in non-CGO builds.
|
||||
func SelectBestBackend(ops []OperationType, preferPerformance bool) (BackendType, error) {
|
||||
return BackendAuto, ErrNoBackends
|
||||
}
|
||||
|
||||
// PrintCapabilities prints nothing in non-CGO builds.
|
||||
func PrintCapabilities() {}
|
||||
@@ -0,0 +1,130 @@
|
||||
package accel
|
||||
|
||||
import "time"
|
||||
|
||||
// OperationType identifies a type of compute operation.
|
||||
type OperationType int
|
||||
|
||||
const (
|
||||
// ML Operations
|
||||
OpMatMul OperationType = iota
|
||||
OpReLU
|
||||
OpGELU
|
||||
OpSoftmax
|
||||
OpLayerNorm
|
||||
OpAttention
|
||||
|
||||
// Crypto Operations
|
||||
OpSHA256
|
||||
OpKeccak256
|
||||
OpPoseidon
|
||||
OpECDSAVerify
|
||||
OpEd25519Verify
|
||||
OpBLSVerify
|
||||
OpMerkleRoot
|
||||
|
||||
// ZK Operations
|
||||
OpNTT
|
||||
OpINTT
|
||||
OpMSM
|
||||
OpPolyMul
|
||||
|
||||
// FHE Operations
|
||||
OpBFVEncrypt
|
||||
OpBFVDecrypt
|
||||
OpBFVAdd
|
||||
OpBFVMul
|
||||
|
||||
// Lattice Operations
|
||||
OpKyberKeyGen
|
||||
OpKyberEncaps
|
||||
OpKyberDecaps
|
||||
OpDilithiumSign
|
||||
OpDilithiumVerify
|
||||
|
||||
// DEX Operations
|
||||
OpConstantProductSwap
|
||||
OpTWAP
|
||||
OpOrderMatch
|
||||
|
||||
opCount // Internal sentinel
|
||||
)
|
||||
|
||||
// String returns the operation name.
|
||||
func (o OperationType) String() string {
|
||||
names := []string{
|
||||
"matmul", "relu", "gelu", "softmax", "layer_norm", "attention",
|
||||
"sha256", "keccak256", "poseidon", "ecdsa_verify", "ed25519_verify", "bls_verify", "merkle_root",
|
||||
"ntt", "intt", "msm", "poly_mul",
|
||||
"bfv_encrypt", "bfv_decrypt", "bfv_add", "bfv_mul",
|
||||
"kyber_keygen", "kyber_encaps", "kyber_decaps", "dilithium_sign", "dilithium_verify",
|
||||
"swap", "twap", "order_match",
|
||||
}
|
||||
if int(o) < len(names) {
|
||||
return names[o]
|
||||
}
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
// Category returns the operation category.
|
||||
func (o OperationType) Category() string {
|
||||
switch {
|
||||
case o <= OpAttention:
|
||||
return "ML"
|
||||
case o <= OpMerkleRoot:
|
||||
return "Crypto"
|
||||
case o <= OpPolyMul:
|
||||
return "ZK"
|
||||
case o <= OpBFVMul:
|
||||
return "FHE"
|
||||
case o <= OpDilithiumVerify:
|
||||
return "Lattice"
|
||||
case o <= OpOrderMatch:
|
||||
return "DEX"
|
||||
default:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// BackendCapabilities describes what operations a backend supports.
|
||||
type BackendCapabilities struct {
|
||||
Backend BackendType
|
||||
Operations map[OperationType]bool
|
||||
Categories map[string]bool
|
||||
}
|
||||
|
||||
// Supports returns true if the backend supports the operation.
|
||||
func (c *BackendCapabilities) Supports(op OperationType) bool {
|
||||
return c.Operations[op]
|
||||
}
|
||||
|
||||
// SupportsCategory returns true if the backend supports any operation in the category.
|
||||
func (c *BackendCapabilities) SupportsCategory(cat string) bool {
|
||||
return c.Categories[cat]
|
||||
}
|
||||
|
||||
// SupportedOperations returns a list of supported operations.
|
||||
func (c *BackendCapabilities) SupportedOperations() []OperationType {
|
||||
var ops []OperationType
|
||||
for op, supported := range c.Operations {
|
||||
if supported {
|
||||
ops = append(ops, op)
|
||||
}
|
||||
}
|
||||
return ops
|
||||
}
|
||||
|
||||
// BenchmarkResult holds timing data for an operation.
|
||||
type BenchmarkResult struct {
|
||||
Backend BackendType
|
||||
Operation OperationType
|
||||
Duration time.Duration
|
||||
Error error
|
||||
}
|
||||
|
||||
// BackendComparison holds benchmark results across backends.
|
||||
type BackendComparison struct {
|
||||
Operation OperationType
|
||||
Results map[BackendType]BenchmarkResult
|
||||
Fastest BackendType
|
||||
}
|
||||
@@ -25,6 +25,27 @@
|
||||
//
|
||||
// session, _ := accel.NewSessionWithBackend(accel.BackendMetal)
|
||||
//
|
||||
// # Runtime Backend Selection
|
||||
//
|
||||
// For intelligent backend selection based on required operations:
|
||||
//
|
||||
// // Select best backend for ZK operations
|
||||
// backend, _ := accel.SelectBackend(accel.OpNTT, accel.OpMSM)
|
||||
// session, _ := accel.NewSessionWithBackend(backend)
|
||||
//
|
||||
// // Query capabilities
|
||||
// caps, _ := accel.Capabilities(accel.BackendWebGPU)
|
||||
// if caps.Supports(accel.OpMSM) {
|
||||
// // Use MSM on WebGPU
|
||||
// }
|
||||
//
|
||||
// // Compare backends for an operation
|
||||
// comparison, _ := accel.CompareBackends(accel.OpNTT, 10)
|
||||
// fmt.Printf("Fastest backend for NTT: %s\n", comparison.Fastest)
|
||||
//
|
||||
// // Print all capabilities
|
||||
// accel.PrintCapabilities()
|
||||
//
|
||||
// # Pure Go Mode
|
||||
//
|
||||
// When built with CGO_ENABLED=0, the package compiles in pure Go mode.
|
||||
|
||||
@@ -7,8 +7,9 @@
|
||||
package capi
|
||||
|
||||
/*
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../luxcpp/lux-accel/include
|
||||
#cgo LDFLAGS: -L${SRCDIR}/../../../../luxcpp/install/lib -llux_accel
|
||||
#cgo CFLAGS: -I${SRCDIR}/../../../../luxcpp/install/include
|
||||
#cgo darwin LDFLAGS: -L${SRCDIR}/../../../../luxcpp/install/lib -Wl,-rpath,${SRCDIR}/../../../../luxcpp/install/lib -llux_accel
|
||||
#cgo linux LDFLAGS: -L${SRCDIR}/../../../../luxcpp/install/lib -Wl,-rpath,${SRCDIR}/../../../../luxcpp/install/lib -llux_accel
|
||||
|
||||
#include <lux/accel/c_api.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
Reference in New Issue
Block a user