cuda_runtime.go

//go:build linux

package main

/*
#cgo LDFLAGS: -ldl
#include <dlfcn.h>
#include <stddef.h>
#include <string.h>

// CUDA types (matching cuda_runtime_api.h without requiring the header)
typedef int cudaError_t;
typedef void* cudaEvent_t;
typedef void* cudaStream_t;

// cudaMemcpyKind enum values
enum {
	CUDA_MEMCPY_HOST_TO_DEVICE = 1,
	CUDA_MEMCPY_DEVICE_TO_HOST = 2
};

// Function pointer types
typedef cudaError_t (*fn_cudaSetDevice)(int);
typedef cudaError_t (*fn_cudaDeviceSynchronize)(void);
typedef cudaError_t (*fn_cudaMemGetInfo)(size_t*, size_t*);
typedef cudaError_t (*fn_cudaMalloc)(void**, size_t);
typedef cudaError_t (*fn_cudaFree)(void*);
typedef cudaError_t (*fn_cudaMemcpy)(void*, const void*, size_t, int);
typedef cudaError_t (*fn_cudaMallocHost)(void**, size_t);
typedef cudaError_t (*fn_cudaFreeHost)(void*);
typedef cudaError_t (*fn_cudaMemset)(void*, int, size_t);
typedef cudaError_t (*fn_cudaEventCreate)(cudaEvent_t*);
typedef cudaError_t (*fn_cudaEventRecord)(cudaEvent_t, cudaStream_t);
typedef cudaError_t (*fn_cudaEventSynchronize)(cudaEvent_t);
typedef cudaError_t (*fn_cudaEventElapsedTime)(float*, cudaEvent_t, cudaEvent_t);
typedef cudaError_t (*fn_cudaEventDestroy)(cudaEvent_t);

static void* cudart_handle = NULL;

// Resolved function pointers (cached after first dlsym)
static fn_cudaSetDevice           p_cudaSetDevice = NULL;
static fn_cudaDeviceSynchronize   p_cudaDeviceSynchronize = NULL;
static fn_cudaMemGetInfo          p_cudaMemGetInfo = NULL;
static fn_cudaMalloc              p_cudaMalloc = NULL;
static fn_cudaFree                p_cudaFree = NULL;
static fn_cudaMemcpy              p_cudaMemcpy = NULL;
static fn_cudaMallocHost          p_cudaMallocHost = NULL;
static fn_cudaFreeHost            p_cudaFreeHost = NULL;
static fn_cudaMemset              p_cudaMemset = NULL;
static fn_cudaEventCreate         p_cudaEventCreate = NULL;
static fn_cudaEventRecord         p_cudaEventRecord = NULL;
static fn_cudaEventSynchronize    p_cudaEventSynchronize = NULL;
static fn_cudaEventElapsedTime    p_cudaEventElapsedTime = NULL;
static fn_cudaEventDestroy        p_cudaEventDestroy = NULL;

static int open_cudart(void) {
	// Try versioned names first, then unversioned
	const char* names[] = {
		"libcudart.so.12",
		"libcudart.so.11",
		"libcudart.so",
		NULL
	};
	for (int i = 0; names[i] != NULL; i++) {
		cudart_handle = dlopen(names[i], RTLD_LAZY | RTLD_GLOBAL);
		if (cudart_handle) return 0;
	}
	return -1;
}

static void close_cudart(void) {
	if (cudart_handle) {
		dlclose(cudart_handle);
		cudart_handle = NULL;
	}
	// Reset all pointers
	p_cudaSetDevice = NULL;
	p_cudaDeviceSynchronize = NULL;
	p_cudaMemGetInfo = NULL;
	p_cudaMalloc = NULL;
	p_cudaFree = NULL;
	p_cudaMemcpy = NULL;
	p_cudaMallocHost = NULL;
	p_cudaFreeHost = NULL;
	p_cudaMemset = NULL;
	p_cudaEventCreate = NULL;
	p_cudaEventRecord = NULL;
	p_cudaEventSynchronize = NULL;
	p_cudaEventElapsedTime = NULL;
	p_cudaEventDestroy = NULL;
}

// Macro to resolve and cache a function pointer
#define RESOLVE(name) \
	if (!p_##name) { \
		p_##name = (fn_##name)dlsym(cudart_handle, #name); \
		if (!p_##name) return -1; \
	}

static cudaError_t call_cudaSetDevice(int dev) {
	RESOLVE(cudaSetDevice);
	return p_cudaSetDevice(dev);
}

static cudaError_t call_cudaDeviceSynchronize(void) {
	RESOLVE(cudaDeviceSynchronize);
	return p_cudaDeviceSynchronize();
}

static cudaError_t call_cudaMemGetInfo(size_t* free_bytes, size_t* total_bytes) {
	RESOLVE(cudaMemGetInfo);
	return p_cudaMemGetInfo(free_bytes, total_bytes);
}

static cudaError_t call_cudaMalloc(void** devPtr, size_t size) {
	RESOLVE(cudaMalloc);
	return p_cudaMalloc(devPtr, size);
}

static cudaError_t call_cudaFree(void* devPtr) {
	RESOLVE(cudaFree);
	return p_cudaFree(devPtr);
}

static cudaError_t call_cudaMemcpy(void* dst, const void* src, size_t count, int kind) {
	RESOLVE(cudaMemcpy);
	return p_cudaMemcpy(dst, src, count, kind);
}

static cudaError_t call_cudaMallocHost(void** ptr, size_t size) {
	RESOLVE(cudaMallocHost);
	return p_cudaMallocHost(ptr, size);
}

static cudaError_t call_cudaFreeHost(void* ptr) {
	RESOLVE(cudaFreeHost);
	return p_cudaFreeHost(ptr);
}

static cudaError_t call_cudaMemset(void* devPtr, int value, size_t count) {
	RESOLVE(cudaMemset);
	return p_cudaMemset(devPtr, value, count);
}

static cudaError_t call_cudaEventCreate(cudaEvent_t* event) {
	RESOLVE(cudaEventCreate);
	return p_cudaEventCreate(event);
}

static cudaError_t call_cudaEventRecord(cudaEvent_t event, cudaStream_t stream) {
	RESOLVE(cudaEventRecord);
	return p_cudaEventRecord(event, stream);
}

static cudaError_t call_cudaEventSynchronize(cudaEvent_t event) {
	RESOLVE(cudaEventSynchronize);
	return p_cudaEventSynchronize(event);
}

static cudaError_t call_cudaEventElapsedTime(float* ms, cudaEvent_t start, cudaEvent_t end) {
	RESOLVE(cudaEventElapsedTime);
	return p_cudaEventElapsedTime(ms, start, end);
}

static cudaError_t call_cudaEventDestroy(cudaEvent_t event) {
	RESOLVE(cudaEventDestroy);
	return p_cudaEventDestroy(event);
}
*/
import "C"

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

// CUDARuntime provides access to the CUDA Runtime API via dlopen.
// No compile-time CUDA dependency — libcudart.so is loaded at runtime.
type CUDARuntime struct {
	available bool
}

// NewCUDARuntime attempts to load libcudart.so via dlopen.
// Returns an error if the library is not available.
func NewCUDARuntime() (*CUDARuntime, error) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()

	if C.open_cudart() != 0 {
		return nil, fmt.Errorf("libcudart.so not found (CUDA toolkit not installed?)")
	}
	return &CUDARuntime{available: true}, nil
}

// Close releases the dlopen handle.
func (c *CUDARuntime) Close() {
	if c.available {
		C.close_cudart()
		c.available = false
	}
}

// Available returns whether the CUDA runtime was successfully loaded.
func (c *CUDARuntime) Available() bool {
	return c.available
}

// SetDevice sets the active CUDA device.
func (c *CUDARuntime) SetDevice(index int) error {
	ret := C.call_cudaSetDevice(C.int(index))
	if ret != 0 {
		return fmt.Errorf("cudaSetDevice(%d) failed: error %d", index, int(ret))
	}
	return nil
}

// DeviceSynchronize blocks until the device has completed all preceding tasks.
func (c *CUDARuntime) DeviceSynchronize() error {
	ret := C.call_cudaDeviceSynchronize()
	if ret != 0 {
		return fmt.Errorf("cudaDeviceSynchronize failed: error %d", int(ret))
	}
	return nil
}

// MemGetInfo returns the free and total device memory in bytes.
func (c *CUDARuntime) MemGetInfo() (free, total uint64, err error) {
	var cFree, cTotal C.size_t
	ret := C.call_cudaMemGetInfo(&cFree, &cTotal)
	if ret != 0 {
		return 0, 0, fmt.Errorf("cudaMemGetInfo failed: error %d", int(ret))
	}
	return uint64(cFree), uint64(cTotal), nil
}

// malloc allocates device memory.
func (c *CUDARuntime) malloc(size uint64) (unsafe.Pointer, error) {
	var devPtr unsafe.Pointer
	ret := C.call_cudaMalloc(&devPtr, C.size_t(size))
	if ret != 0 {
		return nil, fmt.Errorf("cudaMalloc(%d) failed: error %d", size, int(ret))
	}
	return devPtr, nil
}

// free releases device memory.
func (c *CUDARuntime) free(devPtr unsafe.Pointer) {
	C.call_cudaFree(devPtr)
}

// mallocHost allocates pinned (page-locked) host memory.
func (c *CUDARuntime) mallocHost(size uint64) (unsafe.Pointer, error) {
	var ptr unsafe.Pointer
	ret := C.call_cudaMallocHost(&ptr, C.size_t(size))
	if ret != 0 {
		return nil, fmt.Errorf("cudaMallocHost(%d) failed: error %d", size, int(ret))
	}
	return ptr, nil
}

// freeHost releases pinned host memory.
func (c *CUDARuntime) freeHost(ptr unsafe.Pointer) {
	C.call_cudaFreeHost(ptr)
}

// memset fills device memory with a byte value.
func (c *CUDARuntime) memset(devPtr unsafe.Pointer, value int, size uint64) error {
	ret := C.call_cudaMemset(devPtr, C.int(value), C.size_t(size))
	if ret != 0 {
		return fmt.Errorf("cudaMemset failed: error %d", int(ret))
	}
	return nil
}

// memcpyH2D copies from host to device.
func (c *CUDARuntime) memcpyH2D(dst, src unsafe.Pointer, size uint64) error {
	ret := C.call_cudaMemcpy(dst, src, C.size_t(size), C.CUDA_MEMCPY_HOST_TO_DEVICE)
	if ret != 0 {
		return fmt.Errorf("cudaMemcpy H2D failed: error %d", int(ret))
	}
	return nil
}

// memcpyD2H copies from device to host.
func (c *CUDARuntime) memcpyD2H(dst, src unsafe.Pointer, size uint64) error {
	ret := C.call_cudaMemcpy(dst, src, C.size_t(size), C.CUDA_MEMCPY_DEVICE_TO_HOST)
	if ret != 0 {
		return fmt.Errorf("cudaMemcpy D2H failed: error %d", int(ret))
	}
	return nil
}

// timedOp runs a function and returns the elapsed time in milliseconds
// using CUDA events for precise GPU-side timing.
func (c *CUDARuntime) timedOp(fn func() error) (float64, error) {
	var start, stop C.cudaEvent_t

	if ret := C.call_cudaEventCreate(&start); ret != 0 {
		return 0, fmt.Errorf("cudaEventCreate(start) failed: error %d", int(ret))
	}
	defer C.call_cudaEventDestroy(start)

	if ret := C.call_cudaEventCreate(&stop); ret != 0 {
		return 0, fmt.Errorf("cudaEventCreate(stop) failed: error %d", int(ret))
	}
	defer C.call_cudaEventDestroy(stop)

	C.call_cudaEventRecord(start, nil)

	if err := fn(); err != nil {
		return 0, err
	}

	C.call_cudaEventRecord(stop, nil)
	C.call_cudaEventSynchronize(stop)

	var ms C.float
	if ret := C.call_cudaEventElapsedTime(&ms, start, stop); ret != 0 {
		return 0, fmt.Errorf("cudaEventElapsedTime failed: error %d", int(ret))
	}

	return float64(ms), nil
}

// BenchmarkAlloc measures CUDA memory allocation time for the given size in bytes.
// Returns elapsed time in milliseconds.
func (c *CUDARuntime) BenchmarkAlloc(sizeBytes uint64) (float64, error) {
	var devPtr unsafe.Pointer
	ms, err := c.timedOp(func() error {
		var err error
		devPtr, err = c.malloc(sizeBytes)
		return err
	})
	if devPtr != nil {
		c.free(devPtr)
	}
	return ms, err
}

// BenchmarkH2D measures host-to-device bandwidth using pinned memory.
// Returns bandwidth in GB/s.
func (c *CUDARuntime) BenchmarkH2D(sizeBytes uint64) (float64, error) {
	hostPtr, err := c.mallocHost(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate pinned host memory: %w", err)
	}
	defer c.freeHost(hostPtr)

	devPtr, err := c.malloc(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate device memory: %w", err)
	}
	defer c.free(devPtr)

	// Fill host memory so we're not copying zeros
	C.memset(hostPtr, 0xAB, C.size_t(sizeBytes))

	// Warmup
	c.memcpyH2D(devPtr, hostPtr, sizeBytes)
	c.DeviceSynchronize()

	ms, err := c.timedOp(func() error {
		return c.memcpyH2D(devPtr, hostPtr, sizeBytes)
	})
	if err != nil {
		return 0, err
	}
	if ms <= 0 {
		return 0, fmt.Errorf("invalid elapsed time: %.4f ms", ms)
	}

	gbps := float64(sizeBytes) / (ms / 1000.0) / 1e9
	return gbps, nil
}

// BenchmarkDiskToGPU measures the end-to-end disk→RAM→GPU pipeline bandwidth.
// Creates a temp file, reads it into pinned host memory, then copies to GPU.
// Returns bandwidth in GB/s.
func (c *CUDARuntime) BenchmarkDiskToGPU(dir string, sizeBytes uint64) (float64, error) {
	// Create temp file
	path, err := createTempFile(dir, sizeBytes)
	if err != nil {
		return 0, err
	}
	defer os.Remove(path)

	// Allocate pinned host memory and device memory
	hostPtr, err := c.mallocHost(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate pinned host memory: %w", err)
	}
	defer c.freeHost(hostPtr)

	devPtr, err := c.malloc(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate device memory: %w", err)
	}
	defer c.free(devPtr)

	// Open file for reading
	f, err := os.Open(path)
	if err != nil {
		return 0, fmt.Errorf("failed to open temp file: %w", err)
	}
	defer f.Close()

	// Create a Go slice backed by the pinned host memory
	hostSlice := sliceFromPtr(hostPtr, sizeBytes)

	// Time the full pipeline: disk read + H2D memcpy
	start := time.Now()

	// Read file into pinned host memory
	var totalRead uint64
	for totalRead < sizeBytes {
		n, err := f.Read(hostSlice[totalRead:])
		if n > 0 {
			totalRead += uint64(n)
		}
		if err != nil {
			break
		}
	}

	// Copy from pinned host to device
	if err := c.memcpyH2D(devPtr, hostPtr, sizeBytes); err != nil {
		return 0, err
	}
	c.DeviceSynchronize()

	elapsed := time.Since(start)

	if elapsed.Seconds() <= 0 {
		return 0, fmt.Errorf("invalid elapsed time")
	}
	return float64(sizeBytes) / elapsed.Seconds() / 1e9, nil
}

// BenchmarkD2H measures device-to-host bandwidth using pinned memory.
// Returns bandwidth in GB/s.
func (c *CUDARuntime) BenchmarkD2H(sizeBytes uint64) (float64, error) {
	hostPtr, err := c.mallocHost(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate pinned host memory: %w", err)
	}
	defer c.freeHost(hostPtr)

	devPtr, err := c.malloc(sizeBytes)
	if err != nil {
		return 0, fmt.Errorf("failed to allocate device memory: %w", err)
	}
	defer c.free(devPtr)

	// Fill device memory
	c.memset(devPtr, 0xCD, sizeBytes)

	// Warmup
	c.memcpyD2H(hostPtr, devPtr, sizeBytes)
	c.DeviceSynchronize()

	ms, err := c.timedOp(func() error {
		return c.memcpyD2H(hostPtr, devPtr, sizeBytes)
	})
	if err != nil {
		return 0, err
	}
	if ms <= 0 {
		return 0, fmt.Errorf("invalid elapsed time: %.4f ms", ms)
	}

	gbps := float64(sizeBytes) / (ms / 1000.0) / 1e9
	return gbps, nil
}