gpu_embedding.go

package gobed

import (
	"fmt"
	"runtime"
	"sync"
	"time"
)

// GPUEmbeddingModel provides GPU-accelerated batch embedding processing
type GPUEmbeddingModel struct {
	*EmbeddingModel
	batchSize    int
	useGPU       bool
	gpuBatchPool sync.Pool
}

// BatchEmbeddingResult contains results from batch processing
type BatchEmbeddingResult struct {
	Embeddings  [][]float32
	Duration    time.Duration
	BatchSize   int
	ItemsPerSec float64
}

// NewGPUEmbeddingModel creates a GPU-optimized embedding model
func NewGPUEmbeddingModel(batchSize int, useGPU bool) (*GPUEmbeddingModel, error) {
	baseModel, err := LoadModel()
	if err != nil {
		return nil, fmt.Errorf("failed to load base model: %w", err)
	}

	gpu := &GPUEmbeddingModel{
		EmbeddingModel: baseModel,
		batchSize:      batchSize,
		useGPU:         useGPU,
	}

	// Initialize batch processing pool
	gpu.gpuBatchPool = sync.Pool{
		New: func() interface{} {
			return make([][]float32, 0, batchSize)
		},
	}

	return gpu, nil
}

// EncodeBatch processes multiple texts in optimized batches
func (g *GPUEmbeddingModel) EncodeBatch(texts []string) (*BatchEmbeddingResult, error) {
	if g.useGPU && g.batchSize > 1 {
		return g.encodeBatchGPU(texts)
	}

	return g.encodeBatchCPU(texts)
}

// encodeBatchGPU uses GPU-optimized batch processing
func (g *GPUEmbeddingModel) encodeBatchGPU(texts []string) (*BatchEmbeddingResult, error) {
	start := time.Now()
	embeddings := make([][]float32, len(texts))

	// Process in optimal batch sizes for GPU
	numBatches := (len(texts) + g.batchSize - 1) / g.batchSize
	batchChan := make(chan batchJob, numBatches)
	resultChan := make(chan batchResult, numBatches)

	// Create batch jobs
	for i := 0; i < len(texts); i += g.batchSize {
		end := i + g.batchSize
		if end > len(texts) {
			end = len(texts)
		}
		batchChan <- batchJob{
			texts:      texts[i:end],
			startIndex: i,
		}
	}
	close(batchChan)

	// Process batches with optimal GPU utilization
	numWorkers := min(runtime.NumCPU()/2, numBatches) // Don't oversubscribe
	var wg sync.WaitGroup

	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range batchChan {
				result := g.processBatchGPU(job)
				resultChan <- result
			}
		}()
	}

	go func() {
		wg.Wait()
		close(resultChan)
	}()

	// Collect results
	for result := range resultChan {
		if result.err != nil {
			return nil, result.err
		}
		copy(embeddings[result.startIndex:], result.embeddings)
	}

	duration := time.Since(start)
	return &BatchEmbeddingResult{
		Embeddings:  embeddings,
		Duration:    duration,
		BatchSize:   g.batchSize,
		ItemsPerSec: float64(len(texts)) / duration.Seconds(),
	}, nil
}

// encodeBatchCPU uses optimized CPU batch processing
func (g *GPUEmbeddingModel) encodeBatchCPU(texts []string) (*BatchEmbeddingResult, error) {
	start := time.Now()
	embeddings := make([][]float32, len(texts))

	// Optimal CPU batch processing with memory reuse
	numWorkers := runtime.NumCPU()
	textChan := make(chan struct {
		text  string
		index int
	}, len(texts))

	var wg sync.WaitGroup

	// Start workers
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for job := range textChan {
				emb, err := g.EmbeddingModel.Encode(job.text)
				if err != nil {
					continue // Skip errors for benchmark
				}
				embeddings[job.index] = emb
			}
		}()
	}

	// Send work
	for i, text := range texts {
		textChan <- struct {
			text  string
			index int
		}{text, i}
	}
	close(textChan)
	wg.Wait()

	duration := time.Since(start)
	return &BatchEmbeddingResult{
		Embeddings:  embeddings,
		Duration:    duration,
		BatchSize:   len(texts),
		ItemsPerSec: float64(len(texts)) / duration.Seconds(),
	}, nil
}

type batchJob struct {
	texts      []string
	startIndex int
}

type batchResult struct {
	embeddings [][]float32
	startIndex int
	err        error
}

// processBatchGPU processes a single batch with GPU optimization
func (g *GPUEmbeddingModel) processBatchGPU(job batchJob) batchResult {
	embeddings := make([][]float32, len(job.texts))

	// For now, use the existing CPU implementation with better memory management
	// In a real GPU implementation, this would use CUDA kernels or TensorRT
	for i, text := range job.texts {
		emb, err := g.EmbeddingModel.Encode(text)
		if err != nil {
			return batchResult{err: err}
		}
		embeddings[i] = emb
	}

	return batchResult{
		embeddings: embeddings,
		startIndex: job.startIndex,
	}
}

// OptimalBatchSize determines the best batch size for the current hardware
func (g *GPUEmbeddingModel) OptimalBatchSize() int {
	if g.useGPU {
		// For RTX 3080 with 16GB, optimal batch sizes are typically 32-128
		return 64
	}
	// For CPU, smaller batches work better due to memory locality
	return 16
}

// MemoryOptimizedEncodeBatch processes with memory efficiency
func (g *GPUEmbeddingModel) MemoryOptimizedEncodeBatch(texts []string, maxMemoryMB int) (*BatchEmbeddingResult, error) {
	// Estimate memory usage: each embedding is 1024 * 4 bytes = 4KB
	embeddingSize := 1024 * 4 // bytes
	maxItems := (maxMemoryMB * 1024 * 1024) / embeddingSize

	if len(texts) <= maxItems {
		return g.EncodeBatch(texts)
	}

	// Process in chunks to stay within memory limit
	start := time.Now()
	allEmbeddings := make([][]float32, 0, len(texts))

	for i := 0; i < len(texts); i += maxItems {
		end := i + maxItems
		if end > len(texts) {
			end = len(texts)
		}

		chunk := texts[i:end]
		result, err := g.EncodeBatch(chunk)
		if err != nil {
			return nil, err
		}

		allEmbeddings = append(allEmbeddings, result.Embeddings...)

		// Optional: force GC to free memory
		runtime.GC()
	}

	duration := time.Since(start)
	return &BatchEmbeddingResult{
		Embeddings:  allEmbeddings,
		Duration:    duration,
		BatchSize:   maxItems,
		ItemsPerSec: float64(len(texts)) / duration.Seconds(),
	}, nil
}