gpu

package
v0.1.1 Latest Latest
Warning

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

 Go to latest Published: Jun 11, 2026 License: MIT Imports: 3 Imported by: 0

README

gpu

Package gpu provides a compute backend abstraction for accelerated factor operations. The default implementation uses pure Go on the CPU. Future CGO backends (CUDA, OpenCL) can be plugged in via the Backend interface.

import "github.com/asymmetric-effort/datascience/lib/gpu"

Backend Interface

All compute backends implement the Backend interface:

type Backend interface {
    Name() string
    IsAvailable() bool
    MatMul(a, b []float64, m, k, n int) []float64
    ElementWiseMul(a, b []float64) []float64
    Sum(a []float64) float64
    Normalize(a []float64) []float64
    FactorProduct(aValues []float64, aShape []int, bValues []float64, bShape []int, resultShape []int) []float64
    Marginalize(values []float64, shape []int, axis int) ([]float64, []int)
    Close() error
}
Method Description
Name Backend identifier (e.g. "cpu", "cuda")
IsAvailable Whether the backend can run on the current system
MatMul Matrix multiplication of (m x k) and (k x n) matrices in row-major order
ElementWiseMul Element-wise product of two equal-length slices
Sum Sum of all elements
Normalize Divide each element by the total sum (produces a distribution summing to 1)
FactorProduct Outer product of two discrete factors given their shapes
Marginalize Sum out one axis from a multi-dimensional tensor, returning reduced values and new shape
Close Release resources held by the backend

CPUBackend (Default)

The CPUBackend is a pure-Go implementation that is always available.

backend := gpu.NewCPUBackend()

backend.Name()        // "cpu"
backend.IsAvailable() // true

// Matrix multiplication: (2x3) * (3x2) -> (2x2)
a := []float64{1, 2, 3, 4, 5, 6}
b := []float64{7, 8, 9, 10, 11, 12}
c := backend.MatMul(a, b, 2, 3, 2) // row-major result

// Element-wise multiplication
result := backend.ElementWiseMul([]float64{1, 2, 3}, []float64{4, 5, 6})

// Sum
total := backend.Sum([]float64{1, 2, 3}) // 6.0

// Normalize to a probability distribution
dist := backend.Normalize([]float64{2, 3, 5}) // [0.2, 0.3, 0.5]

// Factor product (outer product of two discrete factors)
aVals := []float64{0.3, 0.7}
bVals := []float64{0.4, 0.6}
product := backend.FactorProduct(aVals, []int{2}, bVals, []int{2}, []int{2, 2})
// [0.12, 0.18, 0.28, 0.42]

// Marginalize: sum out axis 1 from a 2x3 tensor
vals := []float64{1, 2, 3, 4, 5, 6}
reduced, newShape := backend.Marginalize(vals, []int{2, 3}, 1)
// reduced = [6, 15], newShape = [2]

backend.Close() // no-op for CPU

Global Backend Management

The package maintains a global default backend used by other datascience components.

// Get the current backend (CPUBackend by default)
b := gpu.GetBackend()

// Set a custom backend
gpu.SetBackend(myCustomBackend)

// Reset to the default CPUBackend
gpu.ResetBackend()

Implementing a Custom Backend

To add GPU acceleration, implement the Backend interface and register it:

type CUDABackend struct {
    // device handle, context, etc.
}

func (c *CUDABackend) Name() string        { return "cuda" }
func (c *CUDABackend) IsAvailable() bool   { /* check for CUDA device */ }
func (c *CUDABackend) MatMul(a, b []float64, m, k, n int) []float64 { /* CUDA kernel */ }
func (c *CUDABackend) ElementWiseMul(a, b []float64) []float64      { /* CUDA kernel */ }
func (c *CUDABackend) Sum(a []float64) float64                      { /* CUDA reduction */ }
func (c *CUDABackend) Normalize(a []float64) []float64              { /* CUDA kernel */ }
func (c *CUDABackend) FactorProduct(aValues []float64, aShape []int, bValues []float64, bShape []int, resultShape []int) []float64 { /* ... */ }
func (c *CUDABackend) Marginalize(values []float64, shape []int, axis int) ([]float64, []int) { /* ... */ }
func (c *CUDABackend) Close() error { /* release CUDA resources */ }

// Register the backend
gpu.SetBackend(&CUDABackend{})

API Summary

Category Types / Functions
Interface Backend
CPU Backend CPUBackend, NewCPUBackend
Global Management GetBackend, SetBackend, ResetBackend

Documentation

Overview

Package gpu provides a compute backend abstraction for accelerated factor operations. The default implementation uses pure Go on CPU. Future CGO backends (CUDA, OpenCL) can be plugged in via the Backend interface.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func AcceleratedBP 

func AcceleratedBP(backend Backend, cliques [][]string, cliqueFactors []FactorData) error

AcceleratedBP runs loopy belief propagation on a cluster graph using the provided GPU backend.

cliques lists the variable names in each clique. cliqueFactors provides the initial factor (potential) for each clique. Messages are passed between neighboring cliques (those sharing at least one variable) until convergence or the iteration limit is reached.

On return, cliqueFactors[i].Values contains the updated (unnormalized) belief for clique i.

func AcceleratedSample 

func AcceleratedSample(backend Backend, factors []FactorData, topoOrder []string, n int) [][]int

AcceleratedSample generates n samples via forward sampling using the GPU backend for probability computations. topoOrder gives the topological ordering of variables. factors must include a factor for each variable.

Returns an n x len(topoOrder) matrix where result[i][j] is the sampled value of the j-th variable in sample i.

func AcceleratedVE 

func AcceleratedVE(backend Backend, factors []FactorData, queryVars, elimVars []string, evidence map[string]int) ([]float64, error)

AcceleratedVE runs variable elimination using the provided GPU backend.

factors is the set of factors in the model. queryVars are the variables whose joint distribution is desired. elimVars lists variables to eliminate (in the order they should be eliminated). evidence maps variable names to observed values (0-indexed).

The returned slice is the normalized distribution over queryVars.

func ResetBackend 

func ResetBackend()

ResetBackend restores the default CPU backend.

func SetBackend 

func SetBackend(b Backend)

SetBackend sets the active compute backend.

Types

type Backend 

type Backend interface {
	// Name returns the backend identifier (e.g. "cpu", "cuda", "opencl").
	Name() string

	// IsAvailable reports whether this backend can run on the current system.
	IsAvailable() bool

	// MatMul performs matrix multiplication of a (m x k) and b (k x n),
	// returning the result as a flat slice of length m*n in row-major order.
	MatMul(a, b []float64, m, k, n int) []float64

	// ElementWiseMul returns the element-wise product of a and b.
	// The slices must have equal length.
	ElementWiseMul(a, b []float64) []float64

	// Sum returns the sum of all elements in a.
	Sum(a []float64) float64

	// Normalize divides each element by the sum of all elements.
	// Returns a new slice; does not modify the input.
	Normalize(a []float64) []float64

	// FactorProduct computes the product of two discrete factors.
	// aValues/bValues are the flat value arrays; aShape/bShape describe their
	// dimensionality; resultShape is the shape of the output factor.
	FactorProduct(aValues []float64, aShape []int, bValues []float64, bShape []int, resultShape []int) []float64

	// Marginalize sums out the given axis from a multi-dimensional array,
	// returning the reduced values and the new shape with that axis removed.
	Marginalize(values []float64, shape []int, axis int) ([]float64, []int)

	// Close releases any resources held by the backend.
	Close() error

	// ElementWiseAdd returns the element-wise sum of a and b.
	ElementWiseAdd(a, b []float64) []float64

	// ElementWiseSub returns the element-wise difference a - b.
	ElementWiseSub(a, b []float64) []float64

	// ElementWiseDiv returns the element-wise quotient a / b.
	ElementWiseDiv(a, b []float64) []float64

	// ScalarMul multiplies every element of a by the scalar s.
	ScalarMul(a []float64, s float64) []float64

	// ScalarAdd adds the scalar s to every element of a.
	ScalarAdd(a []float64, s float64) []float64

	// Exp returns the element-wise exponential of a.
	Exp(a []float64) []float64

	// Log returns the element-wise natural logarithm of a.
	Log(a []float64) []float64

	// Sqrt returns the element-wise square root of a.
	Sqrt(a []float64) []float64

	// Abs returns the element-wise absolute value of a.
	Abs(a []float64) []float64

	// Max returns the maximum element in a.
	Max(a []float64) float64

	// Min returns the minimum element in a.
	Min(a []float64) float64

	// ArgMax returns the index of the maximum element in a.
	ArgMax(a []float64) int

	// ArgMin returns the index of the minimum element in a.
	ArgMin(a []float64) int

	// Dot returns the dot product of a and b.
	Dot(a, b []float64) float64

	// FactorReduce fixes the variable at the given axis to the specified index,
	// returning the reduced values and the new shape with that axis removed.
	FactorReduce(values []float64, shape []int, axis int, index int) ([]float64, []int)

	// FactorMaximize takes the maximum over the given axis,
	// returning the reduced values and the new shape with that axis removed.
	FactorMaximize(values []float64, shape []int, axis int) ([]float64, []int)

	// LogSumExp computes log(sum(exp(a))) in a numerically stable way.
	LogSumExp(a []float64) float64

	// Softmax returns the softmax of a: exp(a_i) / sum(exp(a)).
	Softmax(a []float64) []float64

	// BatchMatMul performs batched matrix multiplication. a and b contain
	// batchSize matrices of dimensions (m x k) and (k x n) respectively,
	// packed contiguously. Returns batchSize result matrices of (m x n).
	BatchMatMul(a, b []float64, batchSize, m, k, n int) []float64

	// BatchNormalize normalizes batchSize vectors of length n packed
	// contiguously in a, so each sub-vector sums to 1.
	BatchNormalize(a []float64, batchSize, n int) []float64

	// Alloc allocates a zeroed slice of the given size on the device.
	Alloc(size int) []float64

	// Free releases device memory. For CPU this is a no-op.
	Free(data []float64)

	// CopyToDevice copies data from host to device memory.
	CopyToDevice(data []float64) []float64

	// CopyFromDevice copies data from device to host memory.
	CopyFromDevice(data []float64) []float64

	// DeviceCount returns the number of available compute devices.
	DeviceCount() int

	// DeviceName returns the name of the device at the given index.
	DeviceName(index int) string

	// MemoryUsed returns the approximate device memory in use (bytes).
	MemoryUsed() int64

	// MemoryTotal returns the total device memory capacity (bytes).
	MemoryTotal() int64
}

Backend defines the compute backend interface for accelerated factor operations. Implementations may use CPU, CUDA, OpenCL, or other compute backends.

func GetBackend 

func GetBackend() Backend

GetBackend returns the active compute backend.

func SelectBackend 

func SelectBackend() Backend

SelectBackend returns a Backend for the best available device. It prefers CUDA > OpenCL > MPS > CPU. Currently only CPU is implemented.

type CPUBackend 

type CPUBackend struct{}

CPUBackend is a pure-Go compute backend that runs all operations on the CPU.

func NewCPUBackend 

func NewCPUBackend() *CPUBackend

NewCPUBackend creates a new CPU-based compute backend.

func (*CPUBackend) Abs 

func (c *CPUBackend) Abs(a []float64) []float64

Abs returns the element-wise absolute value of a.

func (*CPUBackend) Alloc 

func (c *CPUBackend) Alloc(size int) []float64

Alloc allocates a zeroed slice of the given size. On CPU this is make.

func (*CPUBackend) ArgMax 

func (c *CPUBackend) ArgMax(a []float64) int

ArgMax returns the index of the maximum element in a. Panics if a is empty.

func (*CPUBackend) ArgMin 

func (c *CPUBackend) ArgMin(a []float64) int

ArgMin returns the index of the minimum element in a. Panics if a is empty.

func (*CPUBackend) BatchMatMul 

func (c *CPUBackend) BatchMatMul(a, b []float64, batchSize, m, k, n int) []float64

BatchMatMul performs batched matrix multiplication. a and b contain batchSize matrices packed contiguously.

func (*CPUBackend) BatchNormalize 

func (c *CPUBackend) BatchNormalize(a []float64, batchSize, n int) []float64

BatchNormalize normalizes batchSize vectors of length n packed contiguously.

func (*CPUBackend) Close 

func (c *CPUBackend) Close() error

Close is a no-op for the CPU backend.

func (*CPUBackend) CopyFromDevice 

func (c *CPUBackend) CopyFromDevice(data []float64) []float64

CopyFromDevice copies data from device memory. On CPU this copies the slice.

func (*CPUBackend) CopyToDevice 

func (c *CPUBackend) CopyToDevice(data []float64) []float64

CopyToDevice copies data to device memory. On CPU this copies the slice.

func (*CPUBackend) DeviceCount 

func (c *CPUBackend) DeviceCount() int

DeviceCount returns 1 for the CPU backend.

func (*CPUBackend) DeviceName 

func (c *CPUBackend) DeviceName(index int) string

DeviceName returns "cpu" for index 0.

func (*CPUBackend) Dot 

func (c *CPUBackend) Dot(a, b []float64) float64

Dot returns the dot product of a and b.

func (*CPUBackend) ElementWiseAdd 

func (c *CPUBackend) ElementWiseAdd(a, b []float64) []float64

ElementWiseAdd returns the element-wise sum of a and b.

func (*CPUBackend) ElementWiseDiv 

func (c *CPUBackend) ElementWiseDiv(a, b []float64) []float64

ElementWiseDiv returns the element-wise quotient a / b.

func (*CPUBackend) ElementWiseMul 

func (c *CPUBackend) ElementWiseMul(a, b []float64) []float64

ElementWiseMul returns the element-wise product of a and b.

func (*CPUBackend) ElementWiseSub 

func (c *CPUBackend) ElementWiseSub(a, b []float64) []float64

ElementWiseSub returns the element-wise difference a - b.

func (*CPUBackend) Exp 

func (c *CPUBackend) Exp(a []float64) []float64

Exp returns the element-wise exponential of a.

func (*CPUBackend) FactorMaximize 

func (c *CPUBackend) FactorMaximize(values []float64, shape []int, axis int) ([]float64, []int)

FactorMaximize takes the maximum over the given axis, returning the reduced values and the new shape with that axis removed.

func (*CPUBackend) FactorProduct 

func (c *CPUBackend) FactorProduct(aValues []float64, aShape []int, bValues []float64, bShape []int, resultShape []int) []float64

FactorProduct computes the product of two discrete factors.

Each factor is represented as a flat value array with an associated shape. The result shape must be provided and defines the dimensionality of the output.

The algorithm treats the first factor's dimensions as the leading axes and the second factor's dimensions as the trailing axes of the result tensor, computing the outer product and storing it in row-major order.

func (*CPUBackend) FactorReduce 

func (c *CPUBackend) FactorReduce(values []float64, shape []int, axis int, index int) ([]float64, []int)

FactorReduce fixes the variable at the given axis to the specified index, returning the reduced values and the new shape with that axis removed.

func (*CPUBackend) Free 

func (c *CPUBackend) Free(data []float64)

Free releases device memory. On CPU this is a no-op; the GC handles it.

func (*CPUBackend) IsAvailable 

func (c *CPUBackend) IsAvailable() bool

IsAvailable always returns true for the CPU backend.

func (*CPUBackend) Log 

func (c *CPUBackend) Log(a []float64) []float64

Log returns the element-wise natural logarithm of a.

func (*CPUBackend) LogSumExp 

func (c *CPUBackend) LogSumExp(a []float64) float64

LogSumExp computes log(sum(exp(a))) in a numerically stable way.

func (*CPUBackend) Marginalize 

func (c *CPUBackend) Marginalize(values []float64, shape []int, axis int) ([]float64, []int)

Marginalize sums out the given axis from a multi-dimensional array. Returns the reduced values and the new shape with the axis removed.

The values slice is interpreted as a row-major tensor with the given shape. The axis parameter specifies which dimension to sum over (0-indexed).

func (*CPUBackend) MatMul 

func (c *CPUBackend) MatMul(a, b []float64, m, k, n int) []float64

MatMul performs matrix multiplication of a (m x k) and b (k x n). Both a and b are in row-major order. Returns a flat slice of length m*n.

func (*CPUBackend) Max 

func (c *CPUBackend) Max(a []float64) float64

Max returns the maximum element in a. Panics if a is empty.

func (*CPUBackend) MemoryTotal 

func (c *CPUBackend) MemoryTotal() int64

MemoryTotal returns 0 for the CPU backend (system RAM is not tracked).

func (*CPUBackend) MemoryUsed 

func (c *CPUBackend) MemoryUsed() int64

MemoryUsed returns the approximate heap memory in use (bytes).

func (*CPUBackend) Min 

func (c *CPUBackend) Min(a []float64) float64

Min returns the minimum element in a. Panics if a is empty.

func (*CPUBackend) Name 

func (c *CPUBackend) Name() string

Name returns "cpu".

func (*CPUBackend) Normalize 

func (c *CPUBackend) Normalize(a []float64) []float64

Normalize divides each element by the total sum, producing a distribution that sums to 1. Returns a new slice.

func (*CPUBackend) ScalarAdd 

func (c *CPUBackend) ScalarAdd(a []float64, s float64) []float64

ScalarAdd adds the scalar s to every element of a.

func (*CPUBackend) ScalarMul 

func (c *CPUBackend) ScalarMul(a []float64, s float64) []float64

ScalarMul multiplies every element of a by the scalar s.

func (*CPUBackend) Softmax 

func (c *CPUBackend) Softmax(a []float64) []float64

Softmax returns the softmax of a: exp(a_i - max(a)) / sum(exp(a - max(a))).

func (*CPUBackend) Sqrt 

func (c *CPUBackend) Sqrt(a []float64) []float64

Sqrt returns the element-wise square root of a.

func (*CPUBackend) Sum 

func (c *CPUBackend) Sum(a []float64) float64

Sum returns the sum of all elements.

type DeviceInfo 

type DeviceInfo struct {
	// Type is the kind of device (CPU, CUDA, OpenCL, MPS).
	Type DeviceType
	// Name is the human-readable device name.
	Name string
	// Index is the device index within its type.
	Index int
	// Memory is the total device memory in bytes (0 if unknown).
	Memory int64
	// Available indicates whether the device can be used.
	Available bool
}

DeviceInfo describes a single compute device.

func DetectDevices 

func DetectDevices() []DeviceInfo

DetectDevices returns a list of available compute devices. Currently only the CPU device is detected. Future CGO backends will add CUDA, OpenCL, and MPS detection here.

type DeviceType 

type DeviceType int

DeviceType represents the type of compute device. 

const (
	// CPU represents a CPU compute device.
	CPU DeviceType = iota
	// CUDA represents an NVIDIA CUDA GPU device.
	CUDA
	// OpenCL represents an OpenCL compute device.
	OpenCL
	// MPS represents an Apple Metal Performance Shaders device.
	MPS
)

func (DeviceType) String 

func (d DeviceType) String() string

String returns the human-readable name of the device type.

type FactorData 

type FactorData struct {
	Variables []string
	Shape     []int
	Values    []float64
}

FactorData represents a discrete factor (conditional probability table) with named variables, a shape describing each variable's cardinality, and flat values in row-major order.

type Tensor 

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

Tensor is a multi-dimensional array backed by a compute Backend. It stores data as a flat float64 slice in row-major order along with the shape describing its dimensionality.

func NewTensor 

func NewTensor(shape []int, data []float64, backend Backend) *Tensor

NewTensor creates a Tensor with the given shape and data on the specified backend. If data is nil, the tensor is zero-initialized. If data is non-nil its length must equal the product of shape dimensions.

func (*Tensor) Add 

func (t *Tensor) Add(other *Tensor) *Tensor

Add performs element-wise addition with another tensor. Both tensors must have the same shape.

func (*Tensor) Clone 

func (t *Tensor) Clone() *Tensor

Clone returns a deep copy of the tensor on the same device.

func (*Tensor) Data 

func (t *Tensor) Data() []float64

Data returns a copy of the tensor's underlying data on the host.

func (*Tensor) Exp 

func (t *Tensor) Exp() *Tensor

Exp returns the element-wise exponential.

func (*Tensor) Log 

func (t *Tensor) Log() *Tensor

Log returns the element-wise natural logarithm.

func (*Tensor) MatMul 

func (t *Tensor) MatMul(other *Tensor) *Tensor

MatMul performs matrix multiplication (t @ other). t must be 2-D (m x k), other must be 2-D (k x n).

func (*Tensor) Max 

func (t *Tensor) Max(axis int) *Tensor

Max reduces the tensor along the given axis by taking the maximum. axis must be in [0, ndim). Returns a tensor with that axis removed.

func (*Tensor) Mul 

func (t *Tensor) Mul(other *Tensor) *Tensor

Mul performs element-wise multiplication with another tensor.

func (*Tensor) Normalize 

func (t *Tensor) Normalize() *Tensor

Normalize returns a new tensor whose elements sum to 1.

func (*Tensor) Reshape 

func (t *Tensor) Reshape(shape []int) *Tensor

Reshape returns a new tensor with the same data but a different shape. The total number of elements must remain the same.

func (*Tensor) ScalarMul 

func (t *Tensor) ScalarMul(s float64) *Tensor

ScalarMul multiplies every element by a scalar.

func (*Tensor) Shape 

func (t *Tensor) Shape() []int

Shape returns the dimensions of the tensor.

func (*Tensor) Size 

func (t *Tensor) Size() int

Size returns the total number of elements in the tensor.

func (*Tensor) Sub 

func (t *Tensor) Sub(other *Tensor) *Tensor

Sub performs element-wise subtraction (t - other).

func (*Tensor) Sum 

func (t *Tensor) Sum(axis int) *Tensor

Sum reduces the tensor along the given axis by summation. axis must be in [0, ndim). Returns a tensor with that axis removed.

func (*Tensor) ToDevice 

func (t *Tensor) ToDevice(backend Backend) *Tensor

ToDevice copies this tensor to the target backend, returning a new Tensor.

Source Files

View all Source files

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.