fpga/versal.go

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

//go:build fpga
// +build fpga

package fpga

import (
	"encoding/binary"
	"errors"
	"fmt"
	"sync"
	"time"
	"unsafe"
)

func init() {
	RegisterBackend(BackendAMDVersal, func() Accelerator {
		return NewVersalAccelerator()
	})
}

// VersalAccelerator implements FPGAAccelerator for AMD Versal AI Edge/Premium series
// Supports VE2802, VPK180, VPK280 with AI Engines and DSP slices
type VersalAccelerator struct {
	config Config
	stats  Stats

	// Device info
	deviceID    string
	model       DeviceType
	pciSlot     string
	temperature float64
	powerUsage  float64

	// AI Engine array
	aiEngines    int // 400 for VE2802, 304 for VPK180
	aiEngineFreq int // MHz (1-1.25 GHz typical)

	// DSP resources
	dspSlices int // 1968 for VPK180
	dspFreq   int // MHz

	// Memory hierarchy
	ddrChannels  int // Number of DDR4/5 channels
	ddrBandwidth int // GB/s per channel

	// Network on Chip (NoC)
	nocBandwidth int // TB/s

	// DMA engines
	dmaEngines []*versalDMAEngine

	// Kernels
	kernels map[KernelType]*Kernel

	// ZK-specific kernels
	nttKernel      *aiEngineKernel
	msmKernel      *aiEngineKernel
	poseidonKernel *aiEngineKernel
	fheKernel      *aiEngineKernel

	// DEX-specific kernels
	matchingKernel *aiEngineKernel
	riskKernel     *aiEngineKernel

	// Runtime state
	initialized bool
	mu          sync.RWMutex
}

type versalDMAEngine struct {
	id           int
	channelCount int
	maxBandwidth int // GB/s
	bufferSize   int // MB
	isActive     bool
}

type aiEngineKernel struct {
	name          string
	tileStart     int
	tileEnd       int
	inputBuffers  []unsafe.Pointer
	outputBuffers []unsafe.Pointer
	isRunning     bool
}

// NewVersalAccelerator creates a new AMD Versal accelerator
func NewVersalAccelerator() *VersalAccelerator {
	return &VersalAccelerator{
		model:        DeviceVersalVPK180,
		aiEngines:    304,
		aiEngineFreq: 1250,
		dspSlices:    1968,
		ddrChannels:  4,
		ddrBandwidth: 25,
		nocBandwidth: 5000,
		kernels:      make(map[KernelType]*Kernel),
		dmaEngines:   make([]*versalDMAEngine, 0),
	}
}

func (v *VersalAccelerator) Initialize(config Config) error {
	v.mu.Lock()
	defer v.mu.Unlock()

	if v.initialized {
		return errors.New("already initialized")
	}

	v.config = config
	v.deviceID = config.DeviceID
	v.pciSlot = config.PCIeSlot

	if config.AIEngines > 0 {
		v.aiEngines = config.AIEngines
	}
	if config.DSPSlices > 0 {
		v.dspSlices = config.DSPSlices
	}

	// Initialize AI Engine array
	if err := v.initializeAIEngines(); err != nil {
		return fmt.Errorf("failed to initialize AI engines: %w", err)
	}

	// Initialize DMA engines
	if err := v.initializeDMAEngines(config.DMAChannels); err != nil {
		return fmt.Errorf("failed to initialize DMA engines: %w", err)
	}

	// Load ZK kernels if enabled
	if config.EnableZKKernels {
		if err := v.loadZKKernels(); err != nil {
			return fmt.Errorf("failed to load ZK kernels: %w", err)
		}
	}

	// Load DEX kernels if enabled
	if config.EnableDEXKernels {
		if err := v.loadDEXKernels(); err != nil {
			return fmt.Errorf("failed to load DEX kernels: %w", err)
		}
	}

	v.initialized = true
	return nil
}

func (v *VersalAccelerator) Shutdown() error {
	v.mu.Lock()
	defer v.mu.Unlock()

	if !v.initialized {
		return nil
	}

	v.stopAIEngines()

	for _, dma := range v.dmaEngines {
		dma.isActive = false
	}

	v.initialized = false
	return nil
}

func (v *VersalAccelerator) Reset() error {
	if err := v.Shutdown(); err != nil {
		return err
	}
	v.clearAIEngineMemory()
	v.stats = Stats{}
	return nil
}

func (v *VersalAccelerator) IsHealthy() bool {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return v.initialized
}

func (v *VersalAccelerator) GetStats() *Stats {
	v.mu.RLock()
	defer v.mu.RUnlock()
	stats := v.stats
	stats.LastUpdate = time.Now()
	return &stats
}

func (v *VersalAccelerator) GetTemperature() float64 {
	return v.temperature
}

func (v *VersalAccelerator) GetPowerUsage() float64 {
	return v.powerUsage
}

func (v *VersalAccelerator) Backend() Backend {
	return BackendAMDVersal
}

func (v *VersalAccelerator) DeviceType() DeviceType {
	return v.model
}

func (v *VersalAccelerator) DeviceID() string {
	return v.deviceID
}

func (v *VersalAccelerator) Capabilities() Capabilities {
	return Capabilities{
		MaxKernelClockMHz: 1333,
		MaxMemoryClockMHz: 2400,
		MaxDMAChannels:    8,
		MaxDMABufferSize:  1024 * 1024 * 1024,
		LogicCells:        1000000,
		DSPSlices:         v.dspSlices,
		BRAMBlocks:        2000,
		AIEngines:         v.aiEngines,
		DDRChannels:       v.ddrChannels,
		DDRSizeGB:         64,
		PCIeGen:           5,
		PCIeLanes:         16,
		PCIeGbps:          64.0,
		Has100GEthernet:   true,
		HasRDMA:           true,
		SupportedKernels: []KernelType{
			KernelNTT, KernelINTT, KernelMSM, KernelPoseidon, KernelPoseidon2,
			KernelFHEAdd, KernelFHEMul, KernelFHEBootstrap,
			KernelOrderMatch, KernelRiskCheck, KernelOrderBook,
		},
	}
}

func (v *VersalAccelerator) LoadKernel(kernelType KernelType, bitstream []byte) error {
	v.mu.Lock()
	defer v.mu.Unlock()

	v.kernels[kernelType] = &Kernel{
		Type:     kernelType,
		Name:     kernelTypeName(kernelType),
		IsLoaded: true,
		ClockMHz: v.aiEngineFreq,
	}

	return nil
}

func (v *VersalAccelerator) UnloadKernel(kernelType KernelType) error {
	v.mu.Lock()
	defer v.mu.Unlock()

	delete(v.kernels, kernelType)
	return nil
}

func (v *VersalAccelerator) IsKernelLoaded(kernelType KernelType) bool {
	v.mu.RLock()
	defer v.mu.RUnlock()

	kernel, ok := v.kernels[kernelType]
	return ok && kernel.IsLoaded
}

func (v *VersalAccelerator) AllocateDMABuffer(size int64, direction DMADirection) (*DMABuffer, error) {
	// Allocate contiguous memory for DMA
	data := make([]byte, size)
	buffer := &DMABuffer{
		VirtualAddr:  unsafe.Pointer(&data[0]),
		PhysicalAddr: 0, // Would be mapped by kernel driver
		Size:         size,
		Direction:    direction,
		IsMapped:     true,
	}
	return buffer, nil
}

func (v *VersalAccelerator) FreeDMABuffer(buffer *DMABuffer) error {
	buffer.IsMapped = false
	return nil
}

func (v *VersalAccelerator) DMATransfer(src, dst *DMABuffer, size int64) error {
	// Simulate DMA via NoC
	time.Sleep(time.Microsecond)
	v.stats.OperationsProcessed++
	return nil
}

func (v *VersalAccelerator) DMATransferAsync(src, dst *DMABuffer, size int64) (<-chan error, error) {
	ch := make(chan error, 1)
	go func() {
		ch <- v.DMATransfer(src, dst, size)
		close(ch)
	}()
	return ch, nil
}

func (v *VersalAccelerator) Execute(kernel KernelType, input, output []byte) error {
	if !v.IsKernelLoaded(kernel) {
		return ErrKernelNotLoaded
	}

	// Execute on AI engines
	result := v.processInAIEngines(input)
	copy(output, result)

	v.stats.OperationsProcessed++
	return nil
}

func (v *VersalAccelerator) ExecuteAsync(kernel KernelType, input []byte) (<-chan []byte, error) {
	if !v.IsKernelLoaded(kernel) {
		return nil, ErrKernelNotLoaded
	}

	ch := make(chan []byte, 1)
	go func() {
		output := make([]byte, len(input))
		_ = v.Execute(kernel, input, output)
		ch <- output
		close(ch)
	}()
	return ch, nil
}

func (v *VersalAccelerator) SetKernelClock(mhz int) error {
	if mhz < 1000 || mhz > 1333 {
		return fmt.Errorf("frequency %d MHz out of range (1000-1333)", mhz)
	}
	v.aiEngineFreq = mhz
	return nil
}

func (v *VersalAccelerator) GetKernelClock() int {
	return v.aiEngineFreq
}

// ZK operations

func (v *VersalAccelerator) NTT(input []uint64, logN int, forward bool) ([]uint64, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	// Encode for AI engine processing
	data := make([]byte, len(input)*8)
	for i, val := range input {
		binary.LittleEndian.PutUint64(data[i*8:], val)
	}

	// Process through AI engines
	result := v.processInAIEngines(data)

	// Decode result
	output := make([]uint64, len(input))
	for i := range output {
		output[i] = binary.LittleEndian.Uint64(result[i*8:])
	}

	v.stats.OperationsProcessed++
	return output, nil
}

func (v *VersalAccelerator) NTTBatch(inputs [][]uint64, logN int, forward bool) ([][]uint64, error) {
	outputs := make([][]uint64, len(inputs))
	for i, input := range inputs {
		out, err := v.NTT(input, logN, forward)
		if err != nil {
			return nil, err
		}
		outputs[i] = out
	}
	return outputs, nil
}

func (v *VersalAccelerator) MSM(pointsX, pointsY, scalars []uint64, count int) ([]uint64, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	// Prepare data for AI engines
	dataSize := (len(pointsX) + len(pointsY) + len(scalars)) * 8
	data := make([]byte, dataSize)
	offset := 0

	for _, val := range pointsX {
		binary.LittleEndian.PutUint64(data[offset:], val)
		offset += 8
	}
	for _, val := range pointsY {
		binary.LittleEndian.PutUint64(data[offset:], val)
		offset += 8
	}
	for _, val := range scalars {
		binary.LittleEndian.PutUint64(data[offset:], val)
		offset += 8
	}

	// Process on AI engines
	result := v.processInAIEngines(data)

	// Decode result (X and Y coordinates, 4 limbs each)
	output := make([]uint64, 8)
	for i := range output {
		if i*8 < len(result) {
			output[i] = binary.LittleEndian.Uint64(result[i*8:])
		}
	}

	v.stats.OperationsProcessed++
	return output, nil
}

func (v *VersalAccelerator) MSMBatch(pointsX, pointsY, scalars [][]uint64) ([][]uint64, error) {
	outputs := make([][]uint64, len(pointsX))
	for i := range pointsX {
		out, err := v.MSM(pointsX[i], pointsY[i], scalars[i], len(pointsX[i])/4)
		if err != nil {
			return nil, err
		}
		outputs[i] = out
	}
	return outputs, nil
}

func (v *VersalAccelerator) PoseidonHash(input []uint64, rate int) ([]uint64, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	// Encode input
	data := make([]byte, len(input)*8)
	for i, val := range input {
		binary.LittleEndian.PutUint64(data[i*8:], val)
	}

	// Process
	result := v.processInAIEngines(data)

	// Decode (256-bit hash = 4 uint64)
	output := make([]uint64, 4)
	for i := range output {
		if i*8 < len(result) {
			output[i] = binary.LittleEndian.Uint64(result[i*8:])
		}
	}

	v.stats.OperationsProcessed++
	return output, nil
}

func (v *VersalAccelerator) PoseidonHashBatch(inputs [][]uint64, rate int) ([][]uint64, error) {
	outputs := make([][]uint64, len(inputs))
	for i, input := range inputs {
		out, err := v.PoseidonHash(input, rate)
		if err != nil {
			return nil, err
		}
		outputs[i] = out
	}
	return outputs, nil
}

func (v *VersalAccelerator) FHEAdd(ctA, ctB []byte) ([]byte, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	combined := append(ctA, ctB...)
	result := v.processInAIEngines(combined)

	v.stats.OperationsProcessed++
	return result[:len(ctA)], nil
}

func (v *VersalAccelerator) FHEMul(ctA, ctB []byte) ([]byte, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	combined := append(ctA, ctB...)
	result := v.processInAIEngines(combined)

	// FHE mul increases ciphertext size
	v.stats.OperationsProcessed++
	return result, nil
}

func (v *VersalAccelerator) FHEBootstrap(ct []byte) ([]byte, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	// Bootstrapping involves many NTTs
	result := v.processInAIEngines(ct)

	v.stats.OperationsProcessed++
	return result, nil
}

// DEX operations

func (v *VersalAccelerator) ProcessOrder(order *Order) (*OrderResult, error) {
	if !v.initialized {
		return nil, ErrNotInitialized
	}

	start := time.Now()

	// Encode order
	data := v.encodeOrder(order)

	// Process through AI engines
	result := v.processInAIEngines(data)

	// Decode result
	orderResult := v.decodeOrderResult(result)

	latency := time.Since(start)
	v.stats.OperationsProcessed++
	v.stats.AverageLatencyNs = uint64(latency.Nanoseconds())

	return orderResult, nil
}

func (v *VersalAccelerator) ProcessOrderBatch(orders []*Order) ([]*OrderResult, error) {
	results := make([]*OrderResult, len(orders))
	for i, order := range orders {
		result, err := v.ProcessOrder(order)
		if err != nil {
			return nil, err
		}
		results[i] = result
	}
	return results, nil
}

func (v *VersalAccelerator) CancelOrder(orderID uint64) error {
	v.stats.OperationsProcessed++
	return nil
}

func (v *VersalAccelerator) UpdateOrderBook(symbol string, bids, asks []PriceLevel) error {
	data := v.encodeOrderBook(symbol, bids, asks)
	_ = v.broadcastToAIEngines(data)
	v.stats.OperationsProcessed++
	return nil
}

func (v *VersalAccelerator) GetOrderBook(symbol string) (*OrderBookSnapshot, error) {
	return &OrderBookSnapshot{
		Symbol:    symbol,
		Bids:      []PriceLevel{},
		Asks:      []PriceLevel{},
		Timestamp: uint64(time.Now().UnixNano()),
	}, nil
}

// Internal methods

func (v *VersalAccelerator) initializeAIEngines() error {
	halfEngines := v.aiEngines / 2

	// ZK kernels
	v.nttKernel = &aiEngineKernel{name: "ntt", tileStart: 0, tileEnd: halfEngines / 4}
	v.msmKernel = &aiEngineKernel{name: "msm", tileStart: halfEngines / 4, tileEnd: halfEngines / 2}
	v.poseidonKernel = &aiEngineKernel{name: "poseidon", tileStart: halfEngines / 2, tileEnd: halfEngines * 3 / 4}
	v.fheKernel = &aiEngineKernel{name: "fhe", tileStart: halfEngines * 3 / 4, tileEnd: halfEngines}

	// DEX kernels
	v.matchingKernel = &aiEngineKernel{name: "matching", tileStart: halfEngines, tileEnd: halfEngines + halfEngines/2}
	v.riskKernel = &aiEngineKernel{name: "risk", tileStart: halfEngines + halfEngines/2, tileEnd: v.aiEngines}

	return nil
}

func (v *VersalAccelerator) initializeDMAEngines(count int) error {
	if count == 0 {
		count = 4
	}
	for i := 0; i < count; i++ {
		dma := &versalDMAEngine{
			id:           i,
			channelCount: 4,
			maxBandwidth: 25,
			bufferSize:   64,
			isActive:     true,
		}
		v.dmaEngines = append(v.dmaEngines, dma)
	}
	return nil
}

func (v *VersalAccelerator) loadZKKernels() error {
	v.nttKernel.isRunning = true
	v.msmKernel.isRunning = true
	v.poseidonKernel.isRunning = true
	v.fheKernel.isRunning = true

	v.kernels[KernelNTT] = &Kernel{Type: KernelNTT, IsLoaded: true}
	v.kernels[KernelINTT] = &Kernel{Type: KernelINTT, IsLoaded: true}
	v.kernels[KernelMSM] = &Kernel{Type: KernelMSM, IsLoaded: true}
	v.kernels[KernelPoseidon] = &Kernel{Type: KernelPoseidon, IsLoaded: true}
	v.kernels[KernelFHEAdd] = &Kernel{Type: KernelFHEAdd, IsLoaded: true}
	v.kernels[KernelFHEMul] = &Kernel{Type: KernelFHEMul, IsLoaded: true}
	v.kernels[KernelFHEBootstrap] = &Kernel{Type: KernelFHEBootstrap, IsLoaded: true}

	return nil
}

func (v *VersalAccelerator) loadDEXKernels() error {
	v.matchingKernel.isRunning = true
	v.riskKernel.isRunning = true

	v.kernels[KernelOrderMatch] = &Kernel{Type: KernelOrderMatch, IsLoaded: true}
	v.kernels[KernelRiskCheck] = &Kernel{Type: KernelRiskCheck, IsLoaded: true}
	v.kernels[KernelOrderBook] = &Kernel{Type: KernelOrderBook, IsLoaded: true}

	return nil
}

func (v *VersalAccelerator) stopAIEngines() {
	if v.nttKernel != nil {
		v.nttKernel.isRunning = false
	}
	if v.msmKernel != nil {
		v.msmKernel.isRunning = false
	}
	if v.poseidonKernel != nil {
		v.poseidonKernel.isRunning = false
	}
	if v.fheKernel != nil {
		v.fheKernel.isRunning = false
	}
	if v.matchingKernel != nil {
		v.matchingKernel.isRunning = false
	}
	if v.riskKernel != nil {
		v.riskKernel.isRunning = false
	}
}

func (v *VersalAccelerator) clearAIEngineMemory() {
	// Clear local and program memory
}

func (v *VersalAccelerator) processInAIEngines(data []byte) []byte {
	// Simulate AI engine processing
	result := make([]byte, len(data))
	copy(result, data)
	return result
}

func (v *VersalAccelerator) broadcastToAIEngines(data []byte) error {
	return nil
}

func (v *VersalAccelerator) encodeOrder(order *Order) []byte {
	buf := make([]byte, 64)
	binary.LittleEndian.PutUint64(buf[0:], order.OrderID)
	binary.LittleEndian.PutUint32(buf[8:], order.Symbol)
	buf[12] = order.Side
	buf[13] = order.Type
	binary.LittleEndian.PutUint64(buf[16:], order.Price)
	binary.LittleEndian.PutUint64(buf[24:], order.Quantity)
	return buf
}

func (v *VersalAccelerator) decodeOrderResult(data []byte) *OrderResult {
	return &OrderResult{
		OrderID:       binary.LittleEndian.Uint64(data[0:]),
		Status:        data[8],
		ExecutedQty:   binary.LittleEndian.Uint64(data[16:]),
		ExecutedPrice: binary.LittleEndian.Uint64(data[24:]),
		Timestamp:     uint64(time.Now().UnixNano()),
	}
}

func (v *VersalAccelerator) encodeOrderBook(symbol string, bids, asks []PriceLevel) []byte {
	size := 16 + len(bids)*24 + len(asks)*24
	buf := make([]byte, size)
	binary.LittleEndian.PutUint64(buf[0:], uint64(len(symbol)))
	binary.LittleEndian.PutUint32(buf[8:], uint32(len(bids)))
	binary.LittleEndian.PutUint32(buf[12:], uint32(len(asks)))

	offset := 16
	for _, bid := range bids {
		binary.LittleEndian.PutUint64(buf[offset:], bid.Price)
		binary.LittleEndian.PutUint64(buf[offset+8:], bid.Quantity)
		binary.LittleEndian.PutUint32(buf[offset+16:], bid.Orders)
		offset += 24
	}
	for _, ask := range asks {
		binary.LittleEndian.PutUint64(buf[offset:], ask.Price)
		binary.LittleEndian.PutUint64(buf[offset+8:], ask.Quantity)
		binary.LittleEndian.PutUint32(buf[offset+16:], ask.Orders)
		offset += 24
	}

	return buf
}

// Verify VersalAccelerator implements all interfaces
var _ Accelerator = (*VersalAccelerator)(nil)
var _ ZKAccelerator = (*VersalAccelerator)(nil)
var _ DEXAccelerator = (*VersalAccelerator)(nil)