kernels

Documentation

Overview

Package kernels provides Go wrappers for custom HIP kernels via purego dlopen. Build libhipkernels.so first: cd internal/hip/kernels && make. No build tags required; use kernels.Available() to check runtime availability.

Index

• func Add(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func AddScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Available() bool
• func Div(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func DivScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Exp(a, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Fill(data unsafe.Pointer, value float32, n int, s unsafe.Pointer) error
• func FlashAttentionForward(Q, K, V, O unsafe.Pointer, batch, heads, seqLen, headDim int, causal bool, ...) error
• func Log(a, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Mul(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func MulScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Pow(base, exp, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Rsqrt(a, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Softmax(input, output unsafe.Pointer, outer, inner, axisSize int, s unsafe.Pointer) error
• func Sqrt(a, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func Sub(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func SumAxis(input, output unsafe.Pointer, outer, inner, axisSize int, s unsafe.Pointer) error
• func Tanh(a, c unsafe.Pointer, n int, s unsafe.Pointer) error
• func TanhPrime(a, upstream, c unsafe.Pointer, n int, s unsafe.Pointer) error
• type KernelLib

Functions

func Add

func Add(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error

Add launches the elementwise add kernel: c = a + b.

func AddScalar

func AddScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error

AddScalar launches the scalar add kernel: c = a + scalar.

func Available

func Available() bool

Available returns true if the HIP kernel library is loadable.

func Div

func Div(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error

Div launches the elementwise divide kernel: c = a / b.

func DivScalar

func DivScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error

DivScalar launches the scalar divide kernel: c = a / scalar.

func Exp

func Exp(a, c unsafe.Pointer, n int, s unsafe.Pointer) error

Exp launches the elementwise exp kernel: c = exp(a).

func Fill

func Fill(data unsafe.Pointer, value float32, n int, s unsafe.Pointer) error

Fill launches the fill kernel: sets all elements to value.

func FlashAttentionForward

func FlashAttentionForward(
	Q, K, V, O unsafe.Pointer,
	batch, heads, seqLen, headDim int,
	causal bool,
	stream unsafe.Pointer,
) error

FlashAttentionForward computes scaled dot-product attention using a fused tiled kernel. All tensors are in [batch, heads, seq_len, head_dim] layout. When causal is true, an upper-triangular mask is applied.

func Log

func Log(a, c unsafe.Pointer, n int, s unsafe.Pointer) error

Log launches the elementwise log kernel: c = log(a).

func Mul

func Mul(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error

Mul launches the elementwise multiply kernel: c = a * b.

func MulScalar

func MulScalar(a unsafe.Pointer, scalar float32, c unsafe.Pointer, n int, s unsafe.Pointer) error

MulScalar launches the scalar multiply kernel: c = a * scalar.

func Pow

func Pow(base, exp, c unsafe.Pointer, n int, s unsafe.Pointer) error

Pow launches the elementwise power kernel: c = base ^ exp.

func Rsqrt

func Rsqrt(a, c unsafe.Pointer, n int, s unsafe.Pointer) error

Rsqrt launches the elementwise rsqrt kernel: c = 1/sqrt(a).

func Softmax

func Softmax(input, output unsafe.Pointer, outer, inner, axisSize int, s unsafe.Pointer) error

Softmax launches the softmax kernel along an axis.

func Sqrt

func Sqrt(a, c unsafe.Pointer, n int, s unsafe.Pointer) error

Sqrt launches the elementwise sqrt kernel: c = sqrt(a).

func Sub

func Sub(a, b, c unsafe.Pointer, n int, s unsafe.Pointer) error

Sub launches the elementwise subtract kernel: c = a - b.

func SumAxis

func SumAxis(input, output unsafe.Pointer, outer, inner, axisSize int, s unsafe.Pointer) error

SumAxis launches the sum-reduction kernel along an axis.

func Tanh

func Tanh(a, c unsafe.Pointer, n int, s unsafe.Pointer) error

Tanh launches the elementwise tanh kernel: c = tanh(a).

func TanhPrime

func TanhPrime(a, upstream, c unsafe.Pointer, n int, s unsafe.Pointer) error

TanhPrime launches the tanh derivative kernel: c = (1 - tanh(a)^2) * upstream.

Types

type KernelLib

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

KernelLib holds dlopen'd function pointers for custom HIP kernels compiled into libhipkernels.so.