---
title: "GPU / Vulkan Backend"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{06. GPU / Vulkan Backend}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include=FALSE}
# Vignette code is executed locally (NOT_CRAN=true) but not on CRAN, where
# the CPU fallback would multi-thread and trip the "CPU time > elapsed" NOTE.
knitr::opts_chunk$set(eval = identical(Sys.getenv("NOT_CRAN"), "true"))
```

ggmlR is GPU-first: Vulkan is auto-detected at install time and used by default
when available. All APIs (sequential, functional, autograd, ONNX) fall back to
CPU transparently when no GPU is present.

```{r}
library(ggmlR)
```

---

## 1. Check availability

```{r}
if (ggml_vulkan_available()) {
  cat("Vulkan is available\n")
  ggml_vulkan_status()              # print device list and properties
} else {
  cat("No Vulkan GPU — running on CPU\n")
}

n <- ggml_vulkan_device_count()
cat("Vulkan device count:", n, "\n")
```

---

## 2. Device enumeration

```{r}
# Low-level device registry (all backends including CPU)
ggml_backend_load_all()

n_dev <- ggml_backend_dev_count()
for (i in seq_len(n_dev)) {
  dev  <- ggml_backend_dev_get(i - 1L)   # 0-based
  name <- ggml_backend_dev_name(dev)
  desc <- ggml_backend_dev_description(dev)
  mem  <- ggml_backend_dev_memory(dev)
  cat(sprintf("[%d] %s — %s\n", i, name, desc))
  cat(sprintf("    %.1f GB free / %.1f GB total\n",
              mem["free"] / 1e9, mem["total"] / 1e9))
}
```

Full example: `inst/examples/device_discovery.R`

`ggml_vulkan_status()` prints a summary and also *returns* it invisibly, as
`list(available, n_devices, devices)` — so a script can branch on the result
without re-querying with `ggml_vulkan_available()` and
`ggml_vulkan_device_count()`.

```{r}
st <- ggml_vulkan_status()
if (st$available && st$n_devices > 0) {
  # Each entry is a list: index, name, free_memory, total_memory
  nm <- vapply(st$devices, function(d) d$name, character(1))
  cat("Vulkan devices:", paste(nm, collapse = ", "), "\n")
} else {
  cat("No Vulkan device — CPU only\n")
}
```

---

## 3. Autograd device selection

```{r}
# Select GPU (falls back to CPU if unavailable)
device <- tryCatch({
  ag_device("gpu")
  "gpu"
}, error = function(e) {
  message("GPU not available, using CPU")
  "cpu"
})

cat("Active device:", device, "\n")
```

After `ag_device("gpu")`, all subsequent `ag_param` and `ag_tensor` calls
allocate on the GPU.  Switch back with `ag_device("cpu")`.

### 3.1 What stays on the GPU during training

Weights, Adam moments, gradients and forward activations stay in device
buffers across a training step, so a step makes 4 host/device crossings
instead of 10, and the Adam update does not touch the host at all. This is the
default and changes nothing numerically. Transformers train on the GPU too:
`ag_flash_attention()` computes all heads in one fused op and its gradient in
one more, on both CPU and Vulkan. See `vignette("autograd-engine")`,
sections 3.1 and 11.

---

## 4. Mixed precision (f16 / bf16)

```{r}
if (device == "gpu") {
  ag_dtype("f16")     # half-precision on Vulkan GPU
  # ag_dtype("bf16") # bfloat16 — falls back to f16 on Vulkan automatically
} else {
  ag_dtype("f32")     # full precision on CPU
}

cat("Active dtype:", ag_default_dtype(), "\n")
```

`bf16 → f16` fallback is automatic on Vulkan because bf16 is not natively
supported in GLSL shaders.

---

## 5. Vulkan memory info

```{r}
if (ggml_vulkan_available()) {
  mem <- ggml_vulkan_device_memory(0L)
  cat(sprintf("GPU memory: %.1f MB free / %.1f MB total\n",
              mem$free / 1e6, mem$total / 1e6))
}
```

---

## 6. Multi-GPU

ggmlR supports multiple Vulkan devices via the backend scheduler.
`dp_train()` distributes data across replicas automatically:

```{r}
n_gpu <- ggml_vulkan_device_count()
cat(sprintf("Using %d GPU(s)\n", n_gpu))

# dp_train handles multi-GPU internally — see vignette("data-parallel-training")
```

For low-level multi-GPU scheduler usage see `inst/examples/multi_gpu_example.R`.

---

## 7. Sequential / functional API on GPU

The high-level `ggml_fit()` API picks up the Vulkan backend automatically —
no extra configuration needed:

```{r}
data(iris)
x_train <- scale(as.matrix(iris[, 1:4]))
y_train <- model.matrix(~ Species - 1, iris)

model <- ggml_model_sequential() |>
  ggml_layer_dense(64L, activation = "relu", input_shape = 4L) |>
  ggml_layer_dense(3L,  activation = "softmax") |>
  ggml_compile(optimizer = "adam", loss = "categorical_crossentropy")

# Training runs on GPU if Vulkan is available
model <- ggml_fit(model, x_train, y_train, epochs = 50L,
                  batch_size = 32L, verbose = 0L)
```

### 7.1 Precision of predictions on Vulkan

On a Vulkan device with fp16 support, matrix multiplies keep their tiles in
f16 by default. A single row goes through an exact f32 kernel, larger batches
through the tiled one, so predictions carry a relative error around 1e-3 that
depends on the batch size. Training is not hurt by noise of that size;
comparisons against a reference, tests and reproducible evaluation are.

`predict_precision = "f32"` in `ggml_compile()` applies to every prediction of
the model; `matmul_precision = "f32"` applies to one `ggml_predict()` (or a
trainer's `$predict()`) call and wins over the model's setting. The result then
matches the CPU to about 1e-6. Training graphs are not affected. The f32 tiles
cost throughput — about a seventh of the f16 speed for large multiplies on an
RX 9070 under RADV; `inst/examples/benchmark_coopmat.R` measures it on your
device. The CPU backend is exact either way.

```{r}
p_def <- ggml_predict(model, x_train, batch_size = 32L)
p_f32 <- ggml_predict(model, x_train, batch_size = 32L,
                      matmul_precision = "f32")
cat(sprintf("max |default - f32|: %.2e\n", max(abs(p_def - p_f32))))
```

---

## 8. State-space models on the GPU

`ggml_ssm_conv()` and `ggml_ssm_scan()` have Vulkan backward shaders, so a
Mamba block trains entirely on the GPU — for `d_state` 128 or 256 and Mamba-2
shapes; other shapes compute the backward on the CPU. The backward kernels of
the RWKV ops and `ggml_gated_linear_attn()` are CPU-only.
`inst/examples/mamba_train_demo.R` trains a block end to end; the Vulkan
backward shaders brought it from 21.7 s to 7.7 s.

---

## 9. ONNX inference on GPU

```{r, eval = FALSE}
# Weights loaded to GPU once at load time
model_onnx <- onnx_load("model.onnx", device = "vulkan")

# Repeated inference — no weight re-transfer
for (i in seq_len(100L)) {
  out <- onnx_run(model_onnx, list(input = batch[[i]]))
}
```

Against ONNX Runtime on the 15 reference ONNX Model Zoo models, the quantised
MaskRCNN-12-int8 matches bit for bit (`max|d| = 0`) on Vulkan as on the CPU.
The others agree within 1e-3 on Vulkan, except five transformer models (XCiT,
CaiT, BERT, RoBERTa, GPT-NeoX) at 1e-3 to 7e-3, from long float reductions
summed in a different order. See
`vignette("onnx-import")` for which operators run on the GPU.

---

## 10. Build configuration

Vulkan support is compiled in automatically when `libvulkan-dev` and `glslc`
are detected at install time.  SIMD (AVX2/AVX512) for the CPU fallback
requires an explicit flag:

```bash
# Default install (Vulkan auto-detected, CPU fallback without SIMD)
R CMD INSTALL .

# With CPU SIMD acceleration
R CMD INSTALL . --configure-args="--with-simd"

# Also enable the hard-exit path for multi-GPU scripts (see below)
R CMD INSTALL . --configure-args="--with-simd --enable-hard-exit"
```

On Windows R ignores `configure.args`, so set the equivalent environment
variables before installing from source:
`Sys.setenv(GGML_USE_SIMD = "1")`, `Sys.setenv(GGML_VK_HARD_EXIT = "1")`.

`--enable-hard-exit` is **off by default** and is only relevant to standalone
**multi-GPU** scripts. It controls whether `ggml_vulkan_shutdown(hard = TRUE)`
may call `_exit()` to skip the exit-time Vulkan loader teardown, which can
otherwise flakily segfault *after* such a script has printed its results. The
released package omits it because CRAN policy forbids a package from
terminating the user's R session; in a default build `hard = TRUE` tears Vulkan
down normally and emits a warning rather than silently doing nothing. See
`vignette("multi-gpu")` for the full story.

To check what was compiled in:

```{r}
cat(ggml_version(), "\n")
ggml_vulkan_status()   # shows "Vulkan not available" if not compiled in
ggml_vulkan_hard_exit_available()   # TRUE only with --enable-hard-exit
```

---

## Restoring the default device

`ag_device()` and `ag_dtype()` set package-level state that outlives this
document: anything running later in the same R session keeps the GPU and the
f16 default unless it is put back. Restore it once the GPU section is done.

```{r}
ag_device("cpu")
ag_dtype("f32")
cat("Restored:", ag_default_dtype(), "on cpu\n")
```