gpu

package
v1.0.44 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: May 2, 2026 License: MIT Imports: 17 Imported by: 0

Documentation

Overview

Package gpu provides GPU acceleration for NornicDB vector operations. This file provides the high-level accelerator that integrates with EmbeddingIndex.

Package gpu provides optional GPU acceleration for NornicDB vector operations.

This package implements GPU-accelerated vector similarity search using OpenCL, CUDA, Metal, and Vulkan compute backends. The design is optimized for the common case: fast vector search with minimal memory overhead.

Architecture (Simplified & Focused):

  • GPU VRAM stores ONLY embeddings as contiguous float32 arrays
  • CPU RAM stores nodeID mappings and all other graph data
  • Vector queries are offloaded to GPU for parallel computation
  • Results (nodeID indices) are returned to CPU for graph operations
  • No complex graph algorithms on GPU (CPU is better for traversal)

Performance Benefits:

  • 10-100x speedup for vector similarity search
  • Parallel cosine similarity computation
  • Efficient batch operations
  • Reduced CPU load for embedding-heavy workloads

Memory Usage (1024-dim float32 embeddings):

  • 100K nodes ≈ 400MB VRAM (100K × 1024 × 4 bytes)
  • 500K nodes ≈ 2GB VRAM
  • 1M nodes ≈ 4GB VRAM
  • 10M nodes ≈ 40GB VRAM (requires high-end GPU)

Example Usage: 

// Initialize GPU manager
config := gpu.DefaultConfig()
config.Enabled = true
config.PreferredBackend = gpu.BackendOpenCL
config.MaxMemoryMB = 8192 // 8GB limit

manager, err := gpu.NewManager(config)
if err != nil {
	log.Printf("GPU not available: %v", err)
	// Fall back to CPU-only mode
}

// Create embedding index
indexConfig := gpu.DefaultEmbeddingIndexConfig(1024) // 1024 dimensions
index := gpu.NewEmbeddingIndex(manager, indexConfig)

// Add embeddings
embedding := make([]float32, 1024)
// ... populate embedding ...
index.Add("node-123", embedding)

// Batch add for efficiency
nodeIDs := []string{"node-1", "node-2", "node-3"}
embeddings := [][]float32{emb1, emb2, emb3}
index.AddBatch(nodeIDs, embeddings)

// Sync to GPU for acceleration
if err := index.SyncToGPU(); err != nil {
	log.Printf("GPU sync failed: %v", err)
}

// Perform similarity search
query := make([]float32, 1024)
// ... populate query embedding ...
results, err := index.Search(query, 10) // Top 10 similar
if err != nil {
	log.Fatal(err)
}

for _, result := range results {
	fmt.Printf("Node %s: similarity %.3f\n", result.ID, result.Score)
}

// Check performance stats
stats := index.Stats()
fmt.Printf("GPU searches: %d, CPU fallbacks: %d\n",
	stats.SearchesGPU, stats.SearchesCPU)

Supported Backends:

1. **OpenCL** (Cross-platform):

  • Works with NVIDIA, AMD, Intel GPUs
  • Best compatibility across hardware
  • Good performance for most workloads

2. **CUDA** (NVIDIA only):

  • Highest performance on NVIDIA GPUs
  • Requires CUDA toolkit installation
  • Best for production NVIDIA deployments

3. **Metal** (Apple Silicon):

  • Native acceleration on M1/M2/M3 Macs
  • Excellent performance and power efficiency
  • Automatic on macOS with Apple Silicon

4. **Vulkan** (Cross-platform):

  • Modern compute API
  • Good performance across vendors
  • Future-proof choice

Performance Characteristics:

Vector Search (1024-dim, cosine similarity):

  • CPU (single-thread): ~1K vectors/sec
  • CPU (multi-thread): ~10K vectors/sec
  • GPU (mid-range): ~100K-1M vectors/sec
  • GPU (high-end): ~1M-10M vectors/sec

Memory Bandwidth:

  • System RAM: ~50-100 GB/s
  • GPU VRAM: ~500-1000 GB/s (10x faster)
  • This is why GPU excels at vector operations

When to Use GPU: 

✅ Large embedding collections (>10K vectors)
✅ Frequent similarity searches
✅ Batch processing workloads
✅ Real-time recommendation systems
❌ Small datasets (<1K vectors)
❌ Infrequent searches
❌ Memory-constrained environments

ELI12 (Explain Like I'm 12):

Think of your computer like a kitchen:

  1. **CPU = Chef**: Really smart, can do complex recipes (graph traversal), but can only work on one thing at a time.

  2. **GPU = Assembly line**: Not as smart as the chef, but can do simple tasks (vector math) REALLY fast with hundreds of workers in parallel.

  3. **Vector search**: Like comparing the "taste" of 1 million dishes to find the 10 most similar. The chef would take forever doing this one by one, but the assembly line can compare them all at the same time!

  4. **Memory**: The assembly line has its own super-fast ingredients storage (VRAM) that's much faster than the main kitchen storage (RAM).

So we use the assembly line for the repetitive math work, then send the results back to the chef for the complex decision-making!

Package gpu provides GPU-accelerated k-means clustering for NornicDB.

This file implements ClusterIndex, which extends EmbeddingIndex with k-means clustering capabilities for faster semantic search.

Architecture: 

ClusterIndex
    ├── EmbeddingIndex (inherited)      <- GPU vector search
    ├── centroids [][]float32           <- cluster centers
    ├── assignments []int               <- embedding→cluster mapping
    └── clusterMap map[int][]int        <- cluster→embeddings lookup

Performance (M1/M2/M3 GPU):

  • 10K embeddings, 100 clusters: ~50-100ms
  • 100K embeddings, 500 clusters: ~500ms-1s
  • Search speedup: 10-50x vs brute-force

Usage: 

index := gpu.NewClusterIndex(manager, nil, nil)
for _, emb := range embeddings {
    index.Add(nodeID, emb)
}
index.Cluster()  // Run k-means
results := index.SearchWithClusters(query, 10, 3)  // Search 3 clusters

Index

Constants

This section is empty.

Variables

View Source 
var (
	ErrGPUNotAvailable   = errors.New("gpu: no compatible GPU found")
	ErrGPUDisabled       = errors.New("gpu: acceleration disabled")
	ErrOutOfMemory       = errors.New("gpu: out of GPU memory")
	ErrKernelFailed      = errors.New("gpu: kernel execution failed")
	ErrDataTooLarge      = errors.New("gpu: data exceeds GPU memory")
	ErrInvalidDimensions = errors.New("gpu: vector dimension mismatch")
)

Errors

View Source 
var (
	ErrNotClustered     = errors.New("gpu: clustering not yet performed")
	ErrTooFewEmbeddings = errors.New("gpu: too few embeddings for requested clusters")
	ErrInvalidK         = errors.New("gpu: invalid number of clusters")
)

Errors for k-means clustering

Functions

This section is empty.

Types

type Accelerator 

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

Accelerator provides GPU-accelerated vector operations. It automatically selects the best available backend (Metal on macOS, OpenCL/CUDA on other platforms).

Usage: 

accel, err := gpu.NewAccelerator(nil)
if err != nil {
	// Fall back to CPU
}
defer accel.Release()

// Create GPU-backed embedding index
index := accel.NewEmbeddingIndex(1024)
index.Add("doc-1", embedding1)
index.SyncToGPU()
results, _ := index.Search(query, 10)

func NewAccelerator 

func NewAccelerator(config *Config) (*Accelerator, error)

NewAccelerator creates a new GPU accelerator with auto-detection.

The accelerator automatically detects and initializes the best available GPU backend for the current platform:

  • macOS: Metal (Apple Silicon optimized)
  • Linux/Windows: OpenCL or CUDA (when implemented)

If no GPU is available and config.FallbackOnError is true (default), the accelerator runs in CPU-only mode.

func (*Accelerator) Backend 

func (a *Accelerator) Backend() Backend

Backend returns the active GPU backend.

func (*Accelerator) DeviceMemoryMB 

func (a *Accelerator) DeviceMemoryMB() int

DeviceMemoryMB returns the GPU memory in megabytes.

func (*Accelerator) DeviceName 

func (a *Accelerator) DeviceName() string

DeviceName returns the GPU device name.

func (*Accelerator) IsEnabled 

func (a *Accelerator) IsEnabled() bool

IsEnabled returns whether GPU acceleration is active.

func (*Accelerator) NewGPUEmbeddingIndex 

func (a *Accelerator) NewGPUEmbeddingIndex(dimensions int) *GPUEmbeddingIndex

NewGPUEmbeddingIndex creates a new GPU-accelerated embedding index.

func (*Accelerator) Release 

func (a *Accelerator) Release()

Release frees all GPU resources.

func (*Accelerator) Stats 

func (a *Accelerator) Stats() AcceleratorStats

Stats returns GPU usage statistics.

type AcceleratorStats 

type AcceleratorStats struct {
	SearchesGPU      int64
	SearchesCPU      int64
	BytesUploaded    int64
	BytesDownloaded  int64
	KernelExecutions int64
}

AcceleratorStats tracks GPU usage statistics.

type Backend 

type Backend string

Backend represents the GPU compute backend. 

const (
	BackendNone   Backend = "none"   // CPU fallback
	BackendOpenCL Backend = "opencl" // Cross-platform (AMD + NVIDIA)
	BackendCUDA   Backend = "cuda"   // NVIDIA only
	BackendMetal  Backend = "metal"  // Apple Silicon
	BackendVulkan Backend = "vulkan" // Cross-platform compute
)

type BenchmarkResult 

type BenchmarkResult struct {
	DeviceID          int
	VectorOpsPerSec   int64
	MemoryBandwidthGB float64
	LatencyUs         int64
}

BenchmarkResult holds GPU benchmark results.

func BenchmarkDevice 

func BenchmarkDevice(deviceID int) (*BenchmarkResult, error)

BenchmarkDevice runs a simple benchmark on a GPU.

type ClusterIndex 

type ClusterIndex struct {
	*EmbeddingIndex
	// contains filtered or unexported fields
}

ClusterIndex extends EmbeddingIndex with k-means clustering.

Architecture: 

┌─────────────────────────────────────────────────────┐
│                  ClusterIndex                        │
├─────────────────────────────────────────────────────┤
│  EmbeddingIndex (embedded)                          │
│    ├── cpuVectors []float32   <- all embeddings     │
│    ├── nodeIDs []string       <- ID mapping         │
│    └── GPU buffers            <- Metal/CUDA         │
├─────────────────────────────────────────────────────┤
│  Clustering State                                    │
│    ├── centroids [][]float32  <- K cluster centers  │
│    ├── assignments []int      <- embedding→cluster  │
│    └── clusterMap map[int][]int <- cluster→indices  │
└─────────────────────────────────────────────────────┘

Usage: 

index := gpu.NewClusterIndex(manager, embConfig, kmeansConfig)

// Add embeddings
for i, emb := range embeddings {
    index.Add(nodeIDs[i], emb)
}

// Run clustering
if err := index.Cluster(); err != nil {
    log.Fatal(err)
}

// Fast cluster-based search
results, _ := index.SearchWithClusters(query, 10, 3)

Thread Safety: All methods are thread-safe.

func NewClusterIndex 

func NewClusterIndex(manager *Manager, embConfig *EmbeddingIndexConfig, kmeansConfig *KMeansConfig) *ClusterIndex

NewClusterIndex creates a clusterable embedding index.

Parameters:

  • manager: GPU manager (can be nil for CPU-only mode)
  • embConfig: Embedding index config (nil uses defaults)
  • kmeansConfig: K-means config (nil uses defaults)

Example: 

// Default configuration
index := gpu.NewClusterIndex(manager, nil, nil)

// Custom configuration
embConfig := &gpu.EmbeddingIndexConfig{
    Dimensions: 1024,
    InitialCap: 100000,
}
kmeansConfig := &gpu.KMeansConfig{
    NumClusters:   500,
    MaxIterations: 50,
}
index = gpu.NewClusterIndex(manager, embConfig, kmeansConfig)

func (*ClusterIndex) Clear 

func (ci *ClusterIndex) Clear()

Clear removes all embeddings and cluster state from the index. This overrides EmbeddingIndex.Clear() to also reset clustering state.

func (*ClusterIndex) Cluster 

func (ci *ClusterIndex) Cluster() error

Cluster performs k-means clustering on current embeddings.

This method:

  1. Determines optimal K (if AutoK enabled)
  2. Initializes centroids (k-means++ or random)
  3. Iterates assignment/update steps until convergence
  4. Builds cluster membership map for fast lookup

Returns error if too few embeddings or invalid configuration.

Example: 

if err := index.Cluster(); err != nil {
    log.Printf("Clustering failed: %v", err)
}

stats := index.ClusterStats()
fmt.Printf("Created %d clusters in %v\n",
    stats.NumClusters, stats.LastClusterTime)

Cluster runs k-means clustering. For cancellable clustering (e.g. on shutdown), use ClusterWithContext.

func (*ClusterIndex) ClusterStats 

func (ci *ClusterIndex) ClusterStats() ClusterStats

ClusterStats returns clustering statistics.

func (*ClusterIndex) ClusterWithContext 

func (ci *ClusterIndex) ClusterWithContext(ctx context.Context) error

ClusterWithContext runs k-means clustering and stops promptly if ctx is cancelled (e.g. process shutdown). It copies embedding data out, runs the iteration without holding clusterMu, then applies the result. This allows search to continue using the previous clustering state while k-means runs in the background.

func (*ClusterIndex) Config 

func (ci *ClusterIndex) Config() KMeansConfig

Config returns a copy of the k-means configuration.

func (*ClusterIndex) Dimensions 

func (ci *ClusterIndex) Dimensions() int

Dimensions returns the embedding dimensions.

func (*ClusterIndex) FindNearestCentroid 

func (ci *ClusterIndex) FindNearestCentroid(embedding []float32) int

FindNearestCentroid finds the cluster ID nearest to the given embedding.

func (*ClusterIndex) FindNearestClusters 

func (ci *ClusterIndex) FindNearestClusters(embedding []float32, k int) []int

FindNearestClusters finds the k nearest cluster IDs to the given embedding.

func (*ClusterIndex) GetCentroids 

func (ci *ClusterIndex) GetCentroids() [][]float32

GetCentroids returns a copy of the centroid vectors for persistence. Returns nil if not clustered.

func (*ClusterIndex) GetClusterMemberIDs 

func (ci *ClusterIndex) GetClusterMemberIDs(clusterIDs []int) []string

GetClusterMemberIDs returns node IDs belonging to the given clusters.

This is the stable, package-level API for consuming cluster membership outside the gpu package. It intentionally copies the IDs to avoid exposing internal slices that may be mutated during re-clustering.

func (*ClusterIndex) GetClusterMemberIDsForCluster 

func (ci *ClusterIndex) GetClusterMemberIDsForCluster(clusterID int) []string

GetClusterMemberIDsForCluster returns node IDs for a single cluster ID.

func (*ClusterIndex) GetClusterMembers 

func (ci *ClusterIndex) GetClusterMembers(clusterIDs []int) []int

GetClusterMembers returns the embedding indices belonging to the given clusters.

func (*ClusterIndex) GetConfig 

func (ci *ClusterIndex) GetConfig() *KMeansConfig

GetConfig returns the k-means configuration.

func (*ClusterIndex) GetIndicesForNodeIDs 

func (ci *ClusterIndex) GetIndicesForNodeIDs(nodeIDs []string) []int

GetIndicesForNodeIDs returns embedding indices for the provided node IDs.

func (*ClusterIndex) IsClustered 

func (ci *ClusterIndex) IsClustered() bool

IsClustered returns true if clustering has been performed.

func (*ClusterIndex) NumClusters 

func (ci *ClusterIndex) NumClusters() int

NumClusters returns the number of clusters.

func (*ClusterIndex) OnNodeUpdate 

func (ci *ClusterIndex) OnNodeUpdate(nodeID string, embedding []float32) error

OnNodeUpdate handles real-time embedding changes (Tier 1).

This method:

  1. Adds/updates the embedding in the index
  2. If clustered, reassigns to nearest centroid
  3. Tracks update for potential batch centroid recalculation

Example: 

// Called when a node's embedding changes
if err := index.OnNodeUpdate("node-123", newEmbedding); err != nil {
    log.Printf("Update failed: %v", err)
}

func (*ClusterIndex) RestoreClusteringState 

func (ci *ClusterIndex) RestoreClusteringState(centroids [][]float32, idToCluster map[string]int) error

RestoreClusteringState sets clustering state from persisted centroids and id->cluster map so that k-means can be skipped on load. Call after AddBatch has populated the index. Any node ID not in idToCluster is assigned to cluster 0.

func (*ClusterIndex) SearchCandidates 

func (ci *ClusterIndex) SearchCandidates(ctx context.Context, query []float32, candidateIndices []int, topK int) ([]SearchResult, error)

SearchCandidates performs similarity search on a subset of embeddings.

func (*ClusterIndex) SearchWithClusters 

func (ci *ClusterIndex) SearchWithClusters(query []float32, topK, numClusters int) ([]SearchResult, error)

SearchWithClusters performs cluster-accelerated similarity search.

This method:

  1. Finds the k nearest clusters to the query
  2. Gets all embeddings from those clusters as candidates
  3. Performs exact similarity search on candidates only

Parameters:

  • query: Query embedding vector
  • topK: Number of results to return
  • numClusters: Number of clusters to search (expansion factor)

Returns: SearchResult slice sorted by similarity (descending)

Example: 

// Search 3 nearest clusters for top 10 results
results, err := index.SearchWithClusters(query, 10, 3)

func (*ClusterIndex) SetPreferredSeedIndices 

func (ci *ClusterIndex) SetPreferredSeedIndices(indices []int)

SetPreferredSeedIndices sets optional preferred indices for the next Cluster/ClusterWithContext call. Indices outside the current embedding range are ignored during clustering.

func (*ClusterIndex) ShouldRecluster 

func (ci *ClusterIndex) ShouldRecluster() bool

ShouldRecluster checks if re-clustering is needed based on thresholds.

func (*ClusterIndex) UpdateCentroidsBatch 

func (ci *ClusterIndex) UpdateCentroidsBatch()

UpdateCentroidsBatch recomputes centroids for affected clusters (Tier 2). Call periodically to keep centroids accurate after node updates.

type ClusterStats 

type ClusterStats struct {
	EmbeddingCount  int
	NumClusters     int
	AvgClusterSize  float64
	MinClusterSize  int
	MaxClusterSize  int
	Iterations      int
	LastClusterTime time.Duration
	CentroidDrift   float32
	Clustered       bool
}

ClusterStats holds clustering statistics.

type Config 

type Config struct {
	// Enabled toggles GPU acceleration on/off
	Enabled bool

	// PreferredBackend selects compute backend (auto-detected if empty)
	PreferredBackend Backend

	// MaxMemoryMB limits GPU memory usage (0 = use 80% of available)
	MaxMemoryMB int

	// BatchSize for bulk operations
	BatchSize int

	// SyncInterval for async GPU->CPU sync
	SyncInterval time.Duration

	// FallbackOnError falls back to CPU on GPU errors
	FallbackOnError bool

	// DeviceID selects specific GPU (for multi-GPU systems)
	DeviceID int
}

Config holds GPU acceleration configuration options.

The configuration allows fine-tuning of GPU usage, memory limits, and fallback behavior. All settings have sensible defaults.

Example: 

// Production configuration
config := &gpu.Config{
	Enabled:          true,
	PreferredBackend: gpu.BackendOpenCL,
	MaxMemoryMB:      8192, // 8GB limit
	BatchSize:        50000, // Larger batches for throughput
	SyncInterval:     50 * time.Millisecond, // Faster sync
	FallbackOnError:  true, // Always fall back to CPU
	DeviceID:         0, // Use first GPU
}

// Development configuration
config = gpu.DefaultConfig()
config.Enabled = false // Disable for development

func DefaultConfig 

func DefaultConfig() *Config

DefaultConfig returns sensible defaults for GPU acceleration.

The defaults are conservative and prioritize stability over performance:

  • GPU disabled by default (must opt-in)
  • Automatic backend detection
  • 80% of available GPU memory
  • Medium batch sizes
  • CPU fallback enabled

Example: 

config := gpu.DefaultConfig()
config.Enabled = true // Enable GPU acceleration
manager, err := gpu.NewManager(config)

type DeviceInfo 

type DeviceInfo struct {
	ID           int
	Name         string
	Vendor       string
	Backend      Backend
	MemoryMB     int
	ComputeUnits int
	MaxWorkGroup int
	Available    bool
}

DeviceInfo contains information about a GPU device.

func ListDevices 

func ListDevices() ([]DeviceInfo, error)

ListDevices returns all available GPU devices.

type EmbeddingIndex 

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

EmbeddingIndex provides GPU-accelerated vector similarity search.

This is the core GPU acceleration feature. It stores embeddings in GPU VRAM as contiguous float32 arrays for optimal parallel processing, while keeping nodeID mappings and metadata on CPU.

Memory Layout (Optimized for GPU):

GPU VRAM (contiguous float32 array): 

[vec0[0], vec0[1], ..., vec0[D-1], vec1[0], vec1[1], ..., vec1[D-1], ...]
Pure numerical data, perfect for SIMD/parallel computation

CPU RAM (nodeID mapping): 

nodeIDs[0] = "node-123"  -> corresponds to vec0 in GPU
nodeIDs[1] = "node-456"  -> corresponds to vec1 in GPU
idToIndex["node-123"] = 0  -> fast lookup

Search Flow:

  1. Upload query vector to GPU (single float32 array)
  2. GPU computes cosine similarity for ALL embeddings in parallel
  3. GPU performs parallel reduction to find top-k indices
  4. CPU maps indices back to nodeIDs: [5, 12, 3] -> ["node-456", "node-789", "node-234"]

Performance:

  • CPU (1M vectors): ~1-10 seconds
  • GPU (1M vectors): ~10-100 milliseconds (10-100x speedup)
  • Memory bandwidth: GPU VRAM ~10x faster than system RAM

Example: 

// Create index
config := gpu.DefaultEmbeddingIndexConfig(1024)
index := gpu.NewEmbeddingIndex(manager, config)

// Add embeddings (CPU side)
for i, nodeID := range nodeIDs {
	index.Add(nodeID, embeddings[i])
}

// Sync to GPU for acceleration
if err := index.SyncToGPU(); err != nil {
	log.Printf("GPU sync failed, using CPU: %v", err)
}

// Fast similarity search
results, err := index.Search(queryEmbedding, 10)
if err != nil {
	log.Fatal(err)
}

// Results are automatically sorted by similarity (descending)
for i, result := range results {
	fmt.Printf("%d. %s (%.3f similarity)\n",
		i+1, result.ID, result.Score)
}

Memory Efficiency:

  • Only embeddings stored in GPU (no strings, metadata, properties)
  • Contiguous layout maximizes memory bandwidth
  • CPU overhead: ~32 bytes per nodeID (string + map entry)
  • GPU overhead: dimensions × 4 bytes per embedding

Thread Safety: 

All methods are thread-safe. Concurrent searches are supported.

func NewEmbeddingIndex 

func NewEmbeddingIndex(manager *Manager, config *EmbeddingIndexConfig) *EmbeddingIndex

NewEmbeddingIndex creates a new GPU-accelerated embedding index.

The index is created in CPU memory initially. Call SyncToGPU() to upload embeddings to GPU for acceleration. The index gracefully falls back to CPU computation when GPU is unavailable.

Parameters:

  • manager: GPU manager (can be nil for CPU-only mode)
  • config: Index configuration (uses defaults if nil)

Returns:

  • EmbeddingIndex ready for use

Example: 

// Create with custom config
config := &gpu.EmbeddingIndexConfig{
	Dimensions:     1024,
	InitialCap:     100000, // Pre-allocate for 100K embeddings
	GPUEnabled:     true,
	AutoSync:       false,  // Manual sync control
	BatchThreshold: 5000,   // Sync every 5K additions
}
index := gpu.NewEmbeddingIndex(manager, config)

// Or use defaults
index = gpu.NewEmbeddingIndex(manager, nil)

Memory Pre-allocation: 

Setting InitialCap avoids repeated memory allocations during bulk loading.

func (*EmbeddingIndex) Add 

func (ei *EmbeddingIndex) Add(nodeID string, embedding []float32) error

Add inserts or updates an embedding for a node.

The embedding is stored in CPU memory and the GPU sync flag is cleared. Call SyncToGPU() to upload changes to GPU for acceleration.

Parameters:

  • nodeID: Unique identifier for the node
  • embedding: Vector embedding (must match index dimensions)

Returns:

  • ErrInvalidDimensions if embedding size doesn't match

Example: 

// Add single embedding
embedding := make([]float32, 1024)
// ... populate embedding from model ...
err := index.Add("user-123", embedding)
if err != nil {
	log.Fatal(err)
}

// Update existing embedding
newEmbedding := make([]float32, 1024)
// ... compute updated embedding ...
index.Add("user-123", newEmbedding) // Overwrites previous

Performance:

  • O(1) for new insertions
  • O(1) for updates (overwrites in-place)
  • Thread-safe (uses mutex)

Memory:

  • Embedding is copied (safe to modify original after Add)
  • GPU sync is deferred until SyncToGPU() is called

func (*EmbeddingIndex) AddBatch 

func (ei *EmbeddingIndex) AddBatch(nodeIDs []string, embeddings [][]float32) error

AddBatch inserts multiple embeddings efficiently.

func (*EmbeddingIndex) Clear 

func (ei *EmbeddingIndex) Clear()

Clear removes all embeddings from the index.

func (*EmbeddingIndex) Count 

func (ei *EmbeddingIndex) Count() int

Count returns the number of embeddings in the index.

func (*EmbeddingIndex) Deserialize 

func (ei *EmbeddingIndex) Deserialize(data []byte) error

Deserialize loads the index from bytes.

func (*EmbeddingIndex) GPUMemoryUsageMB 

func (ei *EmbeddingIndex) GPUMemoryUsageMB() float64

GPUMemoryUsageMB returns GPU memory usage.

func (*EmbeddingIndex) Get 

func (ei *EmbeddingIndex) Get(nodeID string) ([]float32, bool)

Get retrieves the embedding for a nodeID.

func (*EmbeddingIndex) Has 

func (ei *EmbeddingIndex) Has(nodeID string) bool

Has checks if a nodeID exists in the index.

func (*EmbeddingIndex) MemoryUsageMB 

func (ei *EmbeddingIndex) MemoryUsageMB() float64

MemoryUsageMB returns estimated memory usage.

func (*EmbeddingIndex) Release 

func (ei *EmbeddingIndex) Release()

Release frees all GPU resources associated with this index. Call this when the index is no longer needed to free GPU memory.

func (*EmbeddingIndex) Remove 

func (ei *EmbeddingIndex) Remove(nodeID string) bool

Remove deletes an embedding from the index.

func (*EmbeddingIndex) ScoreSubset 

func (ei *EmbeddingIndex) ScoreSubset(query []float32, ids []string) ([]SearchResult, error)

ScoreSubset computes similarity scores for a specific subset of node IDs. Missing IDs are ignored. Results are sorted by score descending.

func (*EmbeddingIndex) Search 

func (ei *EmbeddingIndex) Search(query []float32, k int) ([]SearchResult, error)

Search finds the k most similar embeddings to the query vector.

The search automatically uses GPU acceleration if available and synced, otherwise falls back to optimized CPU computation. Results are sorted by similarity score in descending order.

Parameters:

  • query: Query embedding vector (must match index dimensions)
  • k: Number of most similar results to return

Returns:

  • SearchResult slice with nodeIDs and similarity scores
  • ErrInvalidDimensions if query size doesn't match

Example: 

// Search for similar items
queryEmbedding := getEmbedding("search query")
results, err := index.Search(queryEmbedding, 10)
if err != nil {
	log.Fatal(err)
}

// Process results (sorted by similarity)
for i, result := range results {
	fmt.Printf("%d. %s (similarity: %.3f, distance: %.3f)\n",
		i+1, result.ID, result.Score, result.Distance)
}

// Check if GPU was used
stats := index.Stats()
if stats.SearchesGPU > stats.SearchesCPU {
	fmt.Println("GPU acceleration is working!")
}

Performance:

  • CPU: O(n×d) where n=embeddings, d=dimensions
  • GPU: O(d) with massive parallelization
  • Typical speedup: 10-100x for large datasets

Similarity Metric: 

Uses cosine similarity: score = dot(a,b) / (||a|| × ||b||)
Range: [-1, 1] where 1 = identical, 0 = orthogonal, -1 = opposite

func (*EmbeddingIndex) Serialize 

func (ei *EmbeddingIndex) Serialize() ([]byte, error)

Serialize exports the index to bytes for persistence.

func (*EmbeddingIndex) Stats 

func (ei *EmbeddingIndex) Stats() EmbeddingIndexStats

Stats returns index statistics.

func (*EmbeddingIndex) SyncToGPU 

func (ei *EmbeddingIndex) SyncToGPU() error

SyncToGPU uploads the current embeddings to GPU memory.

type EmbeddingIndexConfig 

type EmbeddingIndexConfig struct {
	Dimensions     int  // Embedding dimensions (e.g., 1024)
	InitialCap     int  // Initial capacity (number of embeddings)
	GPUEnabled     bool // Use GPU if available
	AutoSync       bool // Auto-sync to GPU on Add
	BatchThreshold int  // Batch size before GPU sync
}

EmbeddingIndexConfig configures the embedding index.

func DefaultEmbeddingIndexConfig 

func DefaultEmbeddingIndexConfig(dimensions int) *EmbeddingIndexConfig

DefaultEmbeddingIndexConfig returns sensible defaults.

type EmbeddingIndexStats 

type EmbeddingIndexStats struct {
	Count        int
	Dimensions   int
	GPUSynced    bool
	SearchesGPU  int64
	SearchesCPU  int64
	UploadsCount int64
	UploadBytes  int64
}

EmbeddingIndexStats holds embedding index statistics.

type GPUEmbeddingIndex 

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

GPUEmbeddingIndex provides GPU-accelerated embedding storage and search.

func (*GPUEmbeddingIndex) Add 

func (idx *GPUEmbeddingIndex) Add(nodeID string, embedding []float32) error

Add inserts or updates an embedding.

func (*GPUEmbeddingIndex) AddBatch 

func (idx *GPUEmbeddingIndex) AddBatch(nodeIDs []string, embeddings [][]float32) error

AddBatch adds multiple embeddings efficiently.

func (*GPUEmbeddingIndex) Count 

func (idx *GPUEmbeddingIndex) Count() int

Count returns the number of embeddings.

func (*GPUEmbeddingIndex) IsGPUSynced 

func (idx *GPUEmbeddingIndex) IsGPUSynced() bool

IsGPUSynced returns whether GPU buffer is up-to-date.

func (*GPUEmbeddingIndex) Release 

func (idx *GPUEmbeddingIndex) Release()

Release frees GPU resources.

func (*GPUEmbeddingIndex) Remove 

func (idx *GPUEmbeddingIndex) Remove(nodeID string) bool

Remove deletes an embedding by nodeID.

func (*GPUEmbeddingIndex) Search 

func (idx *GPUEmbeddingIndex) Search(query []float32, k int) ([]SearchResult, error)

Search finds the k most similar embeddings.

func (*GPUEmbeddingIndex) Stats 

func (idx *GPUEmbeddingIndex) Stats() GPUEmbeddingIndexStats

Stats returns index statistics.

func (*GPUEmbeddingIndex) SyncToGPU 

func (idx *GPUEmbeddingIndex) SyncToGPU() error

SyncToGPU uploads embeddings to GPU memory.

type GPUEmbeddingIndexStats 

type GPUEmbeddingIndexStats struct {
	Count       int
	Dimensions  int
	GPUSynced   bool
	SearchesGPU int64
	SearchesCPU int64
	MemoryMB    float64
}

GPUEmbeddingIndexStats holds index statistics.

type KMeansConfig 

type KMeansConfig struct {
	// NumClusters is the K value. If 0 and AutoK=true, auto-detected.
	NumClusters int

	// MaxIterations limits convergence iterations (default: 15)
	MaxIterations int

	// Tolerance is the convergence threshold (default: 0.0001)
	// Clustering stops when centroid drift < tolerance
	Tolerance float32

	// InitMethod: "kmeans++" (better) or "random" (faster)
	InitMethod string

	// AutoK enables automatic cluster count selection
	AutoK bool

	// DriftThreshold triggers re-clustering when centroids drift > this (default: 0.1)
	DriftThreshold float32

	// MinClusterSize is the minimum embeddings per cluster (default: 10)
	MinClusterSize int
}

KMeansConfig configures k-means clustering behavior.

Example: 

config := &gpu.KMeansConfig{
    NumClusters:    100,       // Fixed K
    MaxIterations:  50,        // Converge faster
    Tolerance:      0.001,     // Stricter convergence
    InitMethod:     "kmeans++",
    DriftThreshold: 0.05,      // Recluster on 5% drift
}

func DefaultKMeansConfig 

func DefaultKMeansConfig() *KMeansConfig

DefaultKMeansConfig returns sensible defaults. Dimensions are auto-detected from the first embedding added.

type Manager 

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

Manager handles GPU resources and operations for vector acceleration.

The Manager provides a simplified interface focused on vector similarity search. It handles device detection, memory management, and fallback to CPU when GPU is unavailable or encounters errors.

Key responsibilities:

  • GPU device detection and initialization
  • Memory allocation and tracking
  • Performance statistics
  • Graceful fallback to CPU operations

Example: 

config := gpu.DefaultConfig()
config.Enabled = true

manager, err := gpu.NewManager(config)
if err != nil {
	log.Printf("GPU unavailable: %v", err)
	return // Use CPU-only mode
}

if manager.IsEnabled() {
	device := manager.Device()
	fmt.Printf("Using GPU: %s (%s)\n", device.Name, device.Backend)
	fmt.Printf("Memory: %d MB\n", device.MemoryMB)
}

// Check usage periodically
stats := manager.Stats()
fmt.Printf("GPU operations: %d, CPU fallbacks: %d\n",
	stats.OperationsGPU, stats.FallbackCount)

Thread Safety: 

All methods are thread-safe and can be called concurrently.

func NewManager 

func NewManager(config *Config) (*Manager, error)

NewManager creates a new GPU manager with the given configuration.

The manager attempts to detect and initialize a compatible GPU device. If GPU is disabled in config or no compatible device is found, the manager operates in CPU-only mode.

Parameters:

  • config: GPU configuration (uses DefaultConfig() if nil)

Returns:

  • Manager instance (always succeeds if FallbackOnError=true)
  • Error if GPU required but unavailable

Example: 

// Try to use GPU, fall back to CPU
config := gpu.DefaultConfig()
config.Enabled = true
config.FallbackOnError = true

manager, err := gpu.NewManager(config)
if err != nil {
	log.Fatal(err) // Should not happen with fallback enabled
}

if manager.IsEnabled() {
	fmt.Println("GPU acceleration active")
} else {
	fmt.Println("Using CPU-only mode")
}

Device Detection: 

The manager tries backends in order: Preferred -> OpenCL -> CUDA -> Vulkan -> Metal

func (*Manager) AllocatedMemoryMB 

func (m *Manager) AllocatedMemoryMB() int

AllocatedMemoryMB returns current GPU memory usage.

func (*Manager) Device 

func (m *Manager) Device() *DeviceInfo

Device returns current GPU device info.

func (*Manager) Disable 

func (m *Manager) Disable()

Disable deactivates GPU acceleration.

func (*Manager) Enable 

func (m *Manager) Enable() error

Enable activates GPU acceleration.

func (*Manager) IsEnabled 

func (m *Manager) IsEnabled() bool

IsEnabled returns whether GPU acceleration is active.

func (*Manager) Stats 

func (m *Manager) Stats() Stats

Stats returns GPU usage statistics.

type SearchResult 

type SearchResult struct {
	ID       string
	Score    float32
	Distance float32
}

SearchResult holds a search result.

type Stats 

type Stats struct {
	OperationsGPU       int64
	OperationsCPU       int64
	BytesTransferred    int64
	KernelExecutions    int64
	FallbackCount       int64
	AverageKernelTimeNs int64
}

Stats tracks GPU usage statistics.

type VectorIndex 

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

VectorIndex provides GPU-accelerated vector operations. Legacy implementation - use EmbeddingIndex for production.

func NewVectorIndex 

func NewVectorIndex(manager *Manager, dimensions int) *VectorIndex

NewVectorIndex creates a GPU-accelerated vector index.

func (*VectorIndex) Add 

func (vi *VectorIndex) Add(id string, vector []float32) error

Add inserts a vector into the index.

func (*VectorIndex) Search 

func (vi *VectorIndex) Search(query []float32, k int) ([]SearchResult, error)

Search finds the k nearest neighbors.

Source Files

View all Source files

Directories

Path Synopsis
Package cuda provides NVIDIA GPU acceleration using CUDA.
 Package cuda provides NVIDIA GPU acceleration using CUDA.
Package metal provides Metal GPU acceleration for macOS and Apple Silicon.
 Package metal provides Metal GPU acceleration for macOS and Apple Silicon.
Package opencl provides cross-platform GPU acceleration using OpenCL.
 Package opencl provides cross-platform GPU acceleration using OpenCL.
Package vulkan provides cross-platform GPU acceleration using Vulkan Compute.
 Package vulkan provides cross-platform GPU acceleration using Vulkan Compute.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.