# Stencils as Small Matrix Multiplications

## How to leverage low-precision compute?

LIBXSTREAM stencil sample

---

## Abstract

<span style="font-size: 75%;">
High-order FD stencils for seismic wave propagation are bandwidth-bound on modern GPUs.
We reformulate the 3D isotropic Laplacian as three small dense matrix multiplications per axis,
mapping the banded Toeplitz operator to hardware matrix engines (Intel DPAS) that would otherwise sit idle.
To preserve FP32 accuracy from an INT8 datapath, we apply Ozaki-1 slicing
(1-3 operator digits × 1-3 adaptive wavefield digits with a carried-forward exponent).
No prior work combines Ozaki digit splitting with hardware matrix engines
for finite-difference stencil operators.
</span>

---

## Outline

- Seismic stencils and their long spatial legs.
- Mapping stencil operators to small dense GEMMs.
- INT8 DPAS preserving FP32 accuracy.
- Compact operators as alternative paths.
- Implementation and performance.

---

## Timeline

| Time      | Topic                                   |
|-----------|-----------------------------------------|
| 0-5 min   | Application: RTM and wave propagation   |
| 5-10 min  | Stencil math and long legs              |
| 10-18 min | Mapping stencils to GEMM/DPAS           |
| 18-24 min | INT8-DPAS Ozaki-1                       |
| 24-30 min | Implementation, performance, discussion |

---

## Why Seismic Stencils?

Reverse Time Migration (RTM) and Full Waveform Inversion (FWI)  
solve wave equations over 3D grids.

$$p_{next} = 2 \cdot p_{now} - p_{prev} + \Delta t^2 \cdot v^2 \cdot \mathcal{L}(p_{now})$$

- $p$: pressure wavefield
- $v$: velocity model
- $\mathcal{L}$: spatial differential operator

The hot loop is repeated stencil evaluation over (many) grid points.

---

## Implemented in the Sample

The current sample is a GPU stencil benchmark and integration example.

| Mode                          | CLI                      | Implemented path                              |
|-------------------------------|--------------------------|-----------------------------------------------|
| FP32 stencil (default)        |                          | SLM-tiled banded FMA, XYZ and ZYX layouts     |
| Isotropic RTM-style Laplacian | `-d 3`                   | fused 3-axis DPAS apply                       |
| TTI-style anisotropic terms   | `-d 9`                   | pure terms plus cross-derivative DPAS phases  |
| Direct high-order stencil     | `-m 0`                   | radius-4 per axis                             |
| Compact variants              | `-m 1`, `-m 2`           | radius-1/radius-2 compact runtime paths       |
| Compact dispersion-fit        | `-m 3`                   | minimax-fitted coefficients (PPW=8 default)   |
| INT8-DPAS Ozaki-1 (Intel)     | `STENCIL_INT8=1`         | signed 8-bit slicing with carried exponent    |
| INT8 dp4a Ozaki-1 (NV>=2)     | `STENCIL_INT8=1`         | Ozaki-1 slicing, PTX dp4a on NVIDIA SM>=7.5   |
| INT8 scalar fallback          | `STENCIL_INT8=2`         | Ozaki-1 slicing, scalar multiply-add          |
| INT8 implicit operator (def.) | `STENCIL_I8_OP=implicit` | dense compact `A⁻¹ B`, radius-2 gather        |
| INT8 banded operator          | `STENCIL_I8_OP=banded`   | explicit radius-4 FD weights                  |

Note: the INT8 path defaults to the implicit-dense compact operator (see below).
INT8 supports XYZ and ZYX; ZYX no longer forces an FP32 fallback.

---

## Block View

The sample updates one `32 × 32 × 32` output cube per block.

```text
BLK       = 32
RADIUS    = 4          direct 8th-order FD
K_BASE    = BLK + 2*RADIUS = 40
K_PAD_I8  = 64                    (INT8, k=32 DPAS alignment)
XMX tile  = 8 x 16
```

For one axis, each block becomes a matrix multiplication.

```text
INT8: D[32 x 64]  * P[64 x 16]    -> Y[32 x 16]  (per strip)
```

---

## Direct Long-Leg Stencil

An 8th-order second derivative has a radius-4 stencil.

$$u''(i) \approx c_0 \cdot u(i) + \sum_{k=1}^{4} c_k \bigl[u(i-k) + u(i+k)\bigr]$$

Each output point reads nine positions along one axis.  
In 3D isotropic mode this is applied along $x$, $y$, and $z$.

---

## Long Legs as a Matrix

The 1D stencil is a banded operator matrix.

$$D = \begin{bmatrix}
c_4 & c_3 & c_2 & c_1 & c_0 & c_1 & c_2 & c_3 & c_4 & 0 & \cdots \\\\
0 & c_4 & c_3 & c_2 & c_1 & c_0 & c_1 & c_2 & c_3 & c_4 & \cdots \\\\
& & & & \ddots & & & & & &
\end{bmatrix}$$

$$Y = D \cdot P$$

The sample stores $D$ as a dense surface because DPAS wants regular tiles.  
The zeros are structural convenience.

---

## Three Isotropic GEMMs

The isotropic Laplacian separates into three 1D operators.

$$\mathcal{L}(p) = D_x \cdot p + D_y \cdot p + D_z \cdot p$$

For each axis, the kernel gathers a `K_PAD x XMX_N` panel into SLM
and applies the same DPAS micro-kernel.

```text
for dim in x, y, z:
    gather haloed lines into SLM
    slice wavefield into INT8 digits
    accumulate D_dim * P_dim into INT32
```

---

## Low Precision Compute

Use matrix compute units but without giving up accuracy.

<span style="opacity: 0.4; font-size: 50%;">
[ConvStencil 2024] stencil-as-matmul on NVIDIA TC,
[Ichimura+ 2025] INT8 TC elastic wave FE</span>

This work represents FP32 values are as digit sums (Ozaki):

$$D \cdot P \approx \sum_i \sum_j D_i \cdot P_j \qquad \text{each } D_i \cdot P_j \text{ is one DPAS call}$$

| Datatype | Digit width | D slices | P slices | Products/dim |
|----------|-------------|----------|----------|--------------|
| INT8     | 7 signed    | 1-3      | 1-3      | 1-9          |

The number of D-slices is configurable via `STENCIL_NDIGITS_A` (default 1).

---

## INT8 Ozaki-1

FD operator weights fit in a single 7-bit signed digit (NSLICES_D=1).
Only the wavefield (P-side) needs multi-slice representation.

$$D \cdot P \approx D_0 \sum_j P_j$$

Advantage: half the operator storage, fewer DPAS products.
The number of P slices adapts at runtime to the local exponent range.

```text
nslices_eff = 1  if assumed_exp <= 7
            = 2  if assumed_exp <= 14
            = 3  otherwise
```

---

## INT8 Implicit Operator (default)

Lele 6th-order tridiagonal $A y = B p$ inverted offline into a dense  
translation-invariant operator that fits the DPAS micro-kernel.

$$\mathcal{L}(p) \approx D' \cdot p, \qquad D' = A^{-1} B$$

- Same DPAS kernel and single gather as the banded path.
- `STENCIL_R_GATHER=2` truncates the tail; quantization-limited, so free.
- `STENCIL_NDIGITS_X=2` shrinks wavefield SLM, raises occupancy.
- `STENCIL_I8_OP=banded` selects the classic radius-4 operator.

Perf. 10% over banded at N=800 (PVC-1T), accuracy
Linf_rel=1.28e-2 vs FP32 reference.

Note: Row weights are re-normalized to sum to zero after truncation,
preserving the constant-annihilation property that keeps the leapfrog
stable.

---

## Carried-Forward Exponent

INT8 slicing needs a shared exponent per spatial strip.

```text
step N:   read exp_buf[old]  →  slice P  →  DPAS  →  write p_new
          scan p_new exponents  →  write exp_buf[new]

step N+1: read exp_buf[new]  →  ...         (buffers flip)
```

- Margin of +1: covers one-step neighbor growth propagation lag
- Output-based: tracks what was written, not what was read
- Double-buffered: no read/write race between neighbors

---

## DPAS Work Count

INT8 path: $N_A \times$ P-slices$_{eff}$ DPAS products per axis (default $N_A = 1$).

| Operator family   | INT8 work/block (NDA=1) |
|-------------------|-------------------------|
| Isotropic, direct | 3 axes × 1-3 = 3-9      |
| Compact           | same DPAS, fewer K      |

The shape is always small and regular: `8 × 16` DPAS tiles over the K dimension.

---

## Hardware Mapping

The kernel uses Intel GPU matrix and block I/O features.

```text
A-side: 8 rows x K_PAD    operator D (2D block read)
B-side: K_PAD x 16 cols   wavefield (SLM block read)
C:      8 x 16            FP32 accumulator
```

|      | A load         | B load          | DPAS         |
|------|----------------|-----------------|--------------|
| INT8 | 2D 8b 8r32x1c  | block_read8 SLM | i8 mad k32   |

---

## TTI: Why It Is Different

Tilted Transverse Isotropy introduces mixed derivatives.

$$\mathcal{L}_\text{TTI}(p) = \text{pure terms} + \text{cross terms}$$

$$\text{cross term: } D_i\bigl(c_{ij} \cdot D_j \cdot p\bigr)$$

Not a wider 1D stencil — a composition of two directional  
derivatives with a pointwise anisotropy field in between.

---

## TTI as Two GEMM Phases

Each cross term is a two-phase DPAS pipeline:

$$T = D_j \cdot P \qquad T = c_{ij} \cdot T \qquad Y \mathrel{+}= D_i \cdot T$$

- `x_slm`: gathered wavefield digits
- `t_slm`: re-split intermediate after pointwise scaling
- Pure terms reuse the isotropic `stencil_apply` path

---

## Stencil Kinds and GEMM Shapes

| Stencil kind        | Math                            | GEMM form                             |
|---------------------|---------------------------------|---------------------------------------|
| 1D direct FD        | $D_i \cdot P$                   | $32 \times 48$ by $48 \times 1024$    |
| 3D isotropic        | $D_x P + D_y P + D_z P$         | three independent 1D GEMMs            |
| VTI-like pure terms | scaled pure axes                | three GEMMs plus coefficients         |
| TTI cross term      | $D_i(c_{ij} \cdot D_j \cdot P)$ | GEMM, scale, GEMM                     |
| Compact             | repeated compact evolution      | smaller-radius $D_r \cdot P$ per step |

The key design choice is to make the stencil look like many dense,  
small, predictable GEMMs.

---

## Long-Leg Motivation

High-order FD stencils use long spatial legs to reduce dispersion error.

| Benefit                          | Cost on large 3D grids       |
|----------------------------------|------------------------------|
| better wave propagation accuracy | wider block halos            |
| fewer time-step artifacts        | more distant memory accesses |
| familiar RTM/TTI formulation     | more L2/TLB pressure         |

Can time evolution provide the effective reach while each update  
touches only a compact neighborhood?
<span style="opacity: 0.3;">Yes, approximately.</span>

Note: Repeated compact updates compose into a wider domain of dependence.
The hard part is fitting the compact coefficients so the composed symbol
matches the long-leg stencil's dispersion behavior.

---

## Compact Operator Idea

Instead of applying the long-leg radius-4 operator directly,  
use compact operators over time.

<span style="opacity: 0.4; font-size: 50%;">
[Dablain 1986] cascaded Laplacians,
[Liu&Sen 2009] time-space optimized explicit,
[Zhang+ 2016] operator splitting for compact FD</span>

```text
direct:       one radius-4 operator

compact-r1:   radius-1 operator, K=4
compact-r2:   radius-2 operator, K=2
compact-fit:  radius-2/3 with dispersion-optimized coefficients
```

Compact-fit uses Golden Section Search by default to minimize  
the worst-case dispersion error over the band $[0, 2\pi/\text{PPW}]$.

Effective reach arises from repeated time updates,  
not from loading the long halo every step.

Note: PPW = Points Per Wavelength — the number of grid points
that resolve one shortest wavelength of interest. Higher PPW
means the fitting band covers only well-resolved frequencies.

---

## What Is Implemented Today

| Method     | CLI              | Radius | Status                              |
|------------|------------------|--------|-------------------------------------|
| Direct     | `-m 0`           | `r=4`  | baseline high-order path (default)  |
| Compact r1 | `-m 1`           | `r=1`  | compact isotropic path              |
| Compact r2 | `-m 2`           | `r=2`  | compact isotropic path              |
| Compact fit| `-m 3`           | `r=2/3`| dispersion-fitted (minimax, PPW=8)  |
| TTI        | `-d 9`           | direct | cross terms implemented             |
| INT8       | `STENCIL_INT8=1` | all    | Ozaki-1 with exp_buf, XYZ and ZYX   |

Environment variables `STENCIL_PPW`, `STENCIL_FIT`, `STENCIL_RADIUS_FIT`
control the target wavelength, fitting method, and fit radius.
`STENCIL_NDIGITS_A` controls operator digit count (1-3, default 1).
`STENCIL_I8_OP`, `STENCIL_R_GATHER`, `STENCIL_NDIGITS_X` control the INT8
operator family, gather radius, and wavefield digit count.

---

## Kernel Structure

```text
host:  precompute D (INT8+scale), JIT kernel

INT8:  gather+slice → DPAS → update → scan exp → exp_buf_out
TTI:   GEMM → scale → re-split → GEMM
```

All paths: one dispatch per time step, 3-dim loop inside kernel.

---

## Runtime Controls

Kernel path selection (default: FP32):

```text
STENCIL_INT8=1            INT8 Ozaki-1 (Intel DPAS or NVIDIA dp4a)
STENCIL_INT8=2            INT8 structure, scalar fallback (any device)
```

Tuning and accuracy controls:

```text
STENCIL_NDIGITS_A=N       operator digits: 1-3 (default 1, use 2 for zero-sum)
STENCIL_I8_OP=X           INT8 operator: implicit (default) or banded
STENCIL_R_GATHER=N        INT8 gather radius: 1..RADIUS (implicit default 2)
STENCIL_NDIGITS_X=N       INT8 wavefield digits: 1-3 (implicit default 2)
STENCIL_STRIPS_PER_WG=2   default, best measured grouping
STENCIL_TRIM=N            accuracy/performance tradeoff (INT8)
STENCIL_GRF256=1          tested slower on target system
STENCIL_LAYOUT=N          0=XYZ (default), 1=blocked, 2=ZYX
STENCIL_HALO=N            halo padding size per axis
STENCIL_PML=1             enable PML absorbing boundaries
STENCIL_PML_ETA=X         PML damping peak (default calibrated 0.0199)
STENCIL_PPW=8             points-per-wavelength for compact-fit
STENCIL_FIT=N             fit method: 0=L2, 1=Ricker, 2=minimax
STENCIL_RADIUS_FIT=N      fit radius (2 or 3, default 3)
```

Note: PML is stable and production-ready on the FP32 kernel (XYZ and ZYX).
The INT8 low-precision kernel with PML is under validation: the PML
`phi` feedback in the absorbing layer amplifies a small boundary asymmetry in
the matrix-engine Laplacian and can drift over long runs. The interior
(non-PML) INT8 path is stable.

---

## Demo Script

Build the code:

```bash
git clone https://github.com/hfp/libxs.git
git clone https://github.com/hfp/libxstream.git
cd libxstream/samples/stencil
echo "Make OpenCL runtime available"
make GNU=1
```

Run the code:

```text
./stencil.x -n 800 -d 3 -m 0

-n N    grid dimension (NxNxN)
-d 3    isotropic pure terms
-d 9    TTI-style pure plus cross terms
-m 0    direct radius-4
-m 1    compact-r1
-m 2    compact-r2
-m 3    compact-fit (dispersion-optimized)
```

---

## GPoints/s @ N=800

Uses XYZ layout and individual coefficients per axis.

| Path         | B580 | B70  | PVC-1T | PVC-2T\* | H100 |
|--------------|------|------|--------|----------|------|
| FP32 direct  | 26.6 | 36.5 | 40.4   | 80.8     | 77.7 |
| FP32 compact | 27.4 | 37.1 | 45.3   | 90.6     | 82.1 |
| INT8 direct  | 10.8 | 14.3 | 16.3   | 32.6     | -    |
| INT8 compact | 10.7 | -    | 16.1   | 32.2     | -    |

- **B580**: Intel® Arc™ B580 Graphics
- **B70**: Intel® Arc™ Pro B70 Graphics
- **PVC**: Intel® Data Center GPU Max 1550 (450W TDP)\*
- **H100**: NVIDIA H100 80GB HBM3

\* Two tiles (2T) that can be utilized with MPI (shown result was trivially doubled).

</div>

Note: The "compact" case uses a radius-3 stencil with fitted coefficients.

---

## GPoints/s @ N=800 (Minimod)

Uses ZYX layout and individual coefficients per axis and Perfectly Matched Layer.  
There is a point-source injection every time step (one grid point).

Starting point of shown iterations is either zero- or random-initialization.

| GPU      | Zero-1k | Zero-2k | Zero-4k | Rand-1k | Rand-2k | Rand-4k |
|----------|---------|---------|---------|---------|---------|---------|
| H100     | 58.9    | 59.0    | 58.7    | 58.9    | 59.0    | 58.7    |
| PVC-1T\* | 40.0    | 39.2    | 37.7    | 36.4    | 36.4    | 36.4    |
| B70      | 60.8    | 56.7    | 45.6    | 30.8    | 30.8    | 30.8    |

\* Two tiles (2T) that can be utilized with MPI (shown result was trivially doubled).

</div>

---

## Collect Results

Write CSV-files:

```bash
STENCIL_CHECK=0 ./stencil.py --kernel int8 --sizes 50:850:50
STENCIL_CHECK=0 ./stencil.py --kernel fp32 --sizes 50:850:50 --pml --layout 2
STENCIL_CHECK=0 ./stencil.py --kernel fp32 --sizes 50:850:50
```

Plot graph:

```bash
./stencil.py --input stencil-fp32.csv --peak-bandwidth-gbs 540 --dark
```

Note: Results on B580 and B70 may have been collected using LIBXSTREAM_USM=0.

---

## Intel® Arc™ Pro B70 (INT8, XYT)

Note: `TRIM` drops least-significant digit products (accuracy tradeoff).
INT8 is bound by gather and slicing overhead.

---

## Intel® Arc™ Pro B70 (FP32+PML, ZYX)

Note: Performance is fully bandwidth-bound.

---

## Intel® Arc™ Pro B70 (FP32, XYZ)

Note: Performance is fully bandwidth-bound.

---

## Accuracy vs. Performance (FP32)

Compact operators trade spatial accuracy for throughput. Measured on PVC-1T,  
N=800, random init, no PML, `STENCIL_CHECK=1` with a CPU based radius-4 direct  
stencil as reference, and an error measure using the relative Infinity Norm.

| Method\*     | Radius | K | GPoints/s | t=10   | t=100  |
|--------------|--------|---|-----------|--------|--------|
| Direct       | r=4    | 1 | 42.6      | 1.8e-6 | 1.9e-5 |
| Compact-fit  | r=3    | 2 | 44.6      | 2.1e-2 | 1.0    |
| Compact-fit  | r=2    | 2 | 51.0      | 7.3e-2 | 5.8    |
| Compact r2   | r=2    | 2 | 50.9      | 8.3e-2 | 5.8    |
| Compact r1   | r=1    | 4 | 53.4      | 2.1e-1 | 11     |

\* Compact operators use different dispersion curves
with "compact-fit" using minimax fitting (PPW=8).

</div>

The direct path on GPU matches the CPU to single-precision
noise (even at t=100). Compact vs direct shows cumulative
phase drift — still physically valid wave propagation but
diverged on a per-point basis over time.

---

## Takeaway

Seismic stencils as dense small matrix multiplications.

- FP32 banded-FMA: fastest on BW-bound hardware
- INT8-DPAS via Ozaki-1: 1-3 D digits, 1-3 adaptive P digits
- INT8 default: implicit `A⁻¹ B`, +10% over banded
- Compact paths: long-leg reach, short-leg cost
- RTM isotropic: three directional GEMMs per step

→ Expressing stencil structure so matrix engines can execute it.

Note: Can LP exceed FP32 (assuming FP32 remains native)? Only by moving
less data. STENCIL_BF16S selects the wavefield storage limbs: =2 keeps two
BF16 limbs per point (4 bytes, FP32-equivalent bandwidth and accuracy),
while =1 keeps a single limb (2 bytes) and halves wavefield traffic at
reduced accuracy. Low-precision compute alone does not help a
bandwidth-bound kernel.

---

## LIBXSTREAM

Minimal compiler requirements (C90), e.g., GNU\* Compiler.

- API to ease buffer management and carrying kernel code
- Abstracts memory model (pointer arithmetic, USM, etc.)
- Based on LIBXS, can make use of powerful primitives
  - For example, predicting tuning parameters

Leverage runtime code generation to specialize kernels (JIT).

---

## OpenCL

OpenCL is interoperable with respective vendor model, e.g., SYCL.

- Intel: Driver and SYCL already deliver OpenCL, otherwise
  - Install opencl-c-headers, ocl-icd-libopencl1, ocl-icd-opencl-dev
  - Install https://github.com/intel/compute-runtime
- Nvidia: Driver and CUDA already deliver OpenCL