AMD ROCm & Efficient GPU Training

10.1 ROCm Concepts

Why ROCm over CUDA

• AMD GPUs are more cost-effective
• Open-source ecosystem
• HIP (Heterogeneous-compute Interface for Portability) is similar to CUDA

Key Concepts

• Kernels — functions executed on the GPU across many parallel threads
• Waves/Workgroups — groups of threads executing in lockstep (analogous to CUDA warps/blocks)
• Local/Global Memory — per-workgroup cache vs. global GPU VRAM
• hipBLAS — AMD's BLAS library for GPU matrix operations (use this, don't implement your own gemm)

10.2 Setting Up ROCm

wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/ubuntu focal main' \
  | sudo tee /etc/apt/sources.list.d/rocm.list

sudo apt update
sudo apt install -y rocm-dkms rocm-libs rocm-hipclang rocm-cmake

sudo usermod -a -G video $USER  # Logout and back in after this

# Verify
rocm-smi
hipcc --version

See Resources & Setup for the full installation walkthrough and troubleshooting.

10.3 Matrix Operations in Rust/ROCm

This is where you switch from Haskell to Rust for performance-critical code. hipBLAS handles the matrix multiplications on the GPU.

use std::ffi::c_void;

pub struct HipMatrix {
    data: *mut f32,
    rows: usize,
    cols: usize,
    handle: hipblasHandle_t,
}

impl HipMatrix {
    pub fn new(rows: usize, cols: usize) -> Self {
        let size = rows * cols;
        let mut data: *mut f32 = std::ptr::null_mut();

        unsafe {
            hipMalloc(
                &mut data as *mut _ as *mut *mut c_void,
                size * std::mem::size_of::<f32>()
            );
        }

        let mut handle: hipblasHandle_t = std::ptr::null_mut();
        unsafe { hipblasCreate(&mut handle); }

        HipMatrix { data, rows, cols, handle }
    }

    /// Matrix multiply: C = A @ B  using hipBLAS SGEMM
    pub fn matmul(handle: hipblasHandle_t, a: &HipMatrix, b: &HipMatrix, c: &mut HipMatrix) {
        unsafe {
            hipblasSgemm(
                handle,
                HIPBLAS_OP_N,   // A not transposed
                HIPBLAS_OP_N,   // B not transposed
                a.rows as i32,
                b.cols as i32,
                a.cols as i32,
                &1.0f32,        // alpha
                a.data, a.rows as i32,
                b.data, b.rows as i32,
                &0.0f32,        // beta
                c.data, c.rows as i32,
            );
        }
    }
}

impl Drop for HipMatrix {
    fn drop(&mut self) {
        unsafe {
            hipFree(self.data as *mut c_void);
            hipblasDestroy(self.handle);
        }
    }
}

[package]
name    = "llm-hip-kernels"
version = "0.1.0"
edition = "2021"

[dependencies]
# HIP bindings — adjust crate/version as ecosystem matures

[lib]
crate-type = ["cdylib"]  # Produce .so for Haskell FFI

10.4 Haskell–Rust FFI Bridge

{-# LANGUAGE ForeignFunctionInterface #-}
import Foreign.Ptr   (Ptr)
import Foreign.C.Types

-- Declare the Rust function exported with #[no_mangle] as a C symbol
foreign import ccall "hip_matmul" c_hip_matmul
  :: Ptr CFloat -> CInt -> CInt     -- Matrix A: data, rows, cols
  -> Ptr CFloat -> CInt -> CInt     -- Matrix B: data, rows, cols
  -> Ptr CFloat -> CInt -> CInt     -- Matrix C (output)
  -> IO ()

-- Haskell wrapper: converts DenseMatrix to pinned CFloat arrays and calls kernel
hipMatmul :: DenseMatrix -> DenseMatrix -> IO DenseMatrix
hipMatmul a b = do
  -- Allocate pinned memory, copy data, call c_hip_matmul, read back result
  undefined

10.5 Profiling & Optimisation

# Kernel statistics
rocprof --stats ./your_program

# Full execution trace
rocprof --trace ./your_program

# GPU utilisation monitor
watch -n 1 rocm-smi

Optimisation Checklist

• Use hipBLAS for matrix operations — never roll your own GEMM
• Fuse operations where possible (e.g., A@B + bias in one kernel)
• Use FP16 for forward/backward passes (2× speedup)
• Minimise CPU↔GPU synchronisation
• Profile before optimising — find actual bottlenecks first
• Batch small operations to amortise kernel launch overhead