Matmul Reference and cuBLAS Routines

Low-level matrix-multiplication routines used to compare the CPU reference implementation with the explicit FP32 and FP64 cuBLAS paths. User-facing execution selects kernels through the runtime device and backend profile; this module does not define a separate execution mode.

inductive Runtime.Autograd.FastKernels.GpuMatmulPrecision : Type

Precision selector for GPU-backed fast matmul over Lean Float tensors.

• .fp32 routes through Cuda.Buffer and cuBLAS SGEMM, matching the precision used by the eager CUDA tensor-buffer path.
• .fp64 routes through the host FloatArray DGEMM bridge and cuBLAS DGEMM, preserving Lean Float precision for matmul-only research paths.

• fp32 : GpuMatmulPrecision
• fp64 : GpuMatmulPrecision

def Runtime.Autograd.FastKernels.instReprGpuMatmulPrecision.repr : GpuMatmulPrecision → ℕ → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.FastKernels.instReprGpuMatmulPrecision : Repr GpuMatmulPrecision

@[implicit_reducible]

instance Runtime.Autograd.FastKernels.instDecidableEqGpuMatmulPrecision : DecidableEq GpuMatmulPrecision

def Runtime.Autograd.FastKernels.matToRows {α : Type} {m n : ℕ} : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)) → Array (Array α)

Convert an (m×n) matrix tensor into an array-of-rows representation.

This is purely a representation change to make runtime loops faster/easier to write.

def Runtime.Autograd.FastKernels.matmulForward {α : Type} [Context α] {m n p : ℕ} (a : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Spec.Tensor α (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))

Fast (runtime-only) 2D matmul kernel.

This is a tight-loop kernel (array-of-rows representation) intended to avoid the overhead of the spec-layer definitions when running eager autograd.

def Runtime.Autograd.FastKernels.Cuda.natToU32! (n : ℕ) : UInt32

Convert an FFI dimension to UInt32, failing before the native call on overflow.

def Runtime.Autograd.FastKernels.Cuda.matmulForwardcuBLAS64 {m n p : ℕ} (a : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Spec.Tensor Float (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))

2D matmul forward via cuBLAS DGEMM (torchlean_dgemm_cuda / Cuda.torchleanDgemmCuda).

def Runtime.Autograd.FastKernels.Cuda.matmulForwardcuBLAS32 {m n p : ℕ} (a : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Spec.Tensor Float (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))

2D matmul forward via the float32 CUDA buffer path.

This path uploads Lean Float values to Cuda.Buffer (rounding to float32), calls the existing Buffer.bmm SGEMM implementation with batch = 1, then downloads the float32 result back to Lean Float.

def Runtime.Autograd.FastKernels.Cuda.matmulForwardcuBLASWith (precision : GpuMatmulPrecision) {m n p : ℕ} (a : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Spec.Tensor Float (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : Spec.Tensor Float (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar))

Dispatch to the requested GPU matmul precision.