High-order PDE solvers often spend most of their time applying structured differential operators. The arithmetic is local, but boundary closures, variable coefficients, and memory layout make the GPU implementation less straightforward than a textbook stencil.
This experiment ports a two-dimensional, variable-coefficient summation-by-parts/simultaneous-approximation-term (SBP-SAT) operator from CUDA.jl-style kernels to Triton. The implementation also includes the diagonal SBP norm and its inverse, along with prolongation and restriction kernels for a multilevel solver.
The complete experiment is available in the Triton_SBP_SAT.ipynb notebook.
The discrete operator
On a mapped two-dimensional grid, the second-order variable-coefficient operator has the schematic form
\[\mathcal{L}u = \partial_r\!\left(c_{rr}\,\partial_r u + c_{rs}\,\partial_s u\right) + \partial_s\!\left(c_{rs}\,\partial_r u + c_{ss}\,\partial_s u\right).\]The arrays crr, css, and crs hold the metric-dependent coefficients. In the interior, each output uses the center point, its four axial neighbors, and four diagonal neighbors. The implementation groups those contributions as
where the first two terms approximate the pure derivatives and the final two terms approximate the mixed derivatives.
SBP operators are designed so that their discrete derivative and norm matrices reproduce integration by parts. In one dimension this defining relation is
\[H D + D^T H = B,\]where $H$ is a positive-definite discrete quadrature norm and $B$ contains the boundary contribution. SAT terms then impose boundary or interface conditions weakly while preserving an energy estimate. That split is reflected directly in the GPU implementation: one kernel handles regular interior points and four specialized kernels handle the faces and corners.
Preserving the storage contract
The source arrays use Julia’s column-major linearization. If $(i,j)$ denotes a one-based grid coordinate and Nr1 is the leading dimension, the Julia index is
Triton and PyTorch use zero-based offsets, so the corresponding address is
\[g = (i-1)N_r + (j-1).\]Keeping this contract explicit was essential. A kernel can execute successfully while silently transposing its numerical interpretation if the flat storage order is changed during a port.
The interior launch maps a two-dimensional Triton program grid directly to the physical grid:
def launch_x_interior(hr, hs, x, Nr1, Ns1, crr, css, crs, out):
grid = (max(Ns1 - 2, 0), max(Nr1 - 2, 0))
x_interior_kernel[grid](
hr, hs, x, Nr1, Ns1, crr, css, crs, out
)
return out
Each program computes one interior point. This is deliberately simple: the first goal was a transparent translation whose indexing could be checked against the original formulas.
Making the boundary update race-free
The original monolithic operator allowed boundary threads to write the boundary value and neighboring output entries. Translating that literally would make ownership difficult to reason about and could introduce write races.
The Triton version instead decomposes the application into five launches:
x_interior_kernelwrites regular interior points.x_f1_kernelandx_f2_kernelapply the left and right SAT closures.x_f3_kernelandx_f4_kernelapply the bottom and top closures.- Explicit corner branches apply the half-weighted two-dimensional closures.
The face kernels retain the one-sided derivative stencil used by the SAT term. For example, the normal derivative at a boundary uses
\[\partial_n u \approx \frac{1}{h} \left(\frac{3}{2}u_0 - 2u_1 + \frac{1}{2}u_2\right).\]Separating the faces costs additional launches, but it makes output ownership visible and keeps the irregular boundary logic out of the interior kernel. That is a useful trade while establishing correctness; fusion can be revisited after reference tests are in place.
The SBP norm and transfer operators
For this second-order tensor-product SBP discretization, the diagonal norm weights are $1$ in the interior, $1/2$ on an edge, and $1/4$ at a corner. The Triton kernel applies $H$ without assembling a matrix:
boundary_x = (i == 1) | (i == ns1)
boundary_y = (j == 1) | (j == nr1)
corner = boundary_x & boundary_y
edge = boundary_x | boundary_y
scale = hr * hs
out = tl.where(
corner,
0.25 * scale * value,
tl.where(edge, 0.5 * scale * value, scale * value),
)
The inverse kernel uses the reciprocal weights. The notebook also implements bilinear prolongation and full-weighting restriction, including separate formulas for edges and corners. Together, these kernels cover the main building blocks needed to use the SBP-SAT operator inside a geometric multigrid iteration.
Benchmark setup
The saved run used an NVIDIA Tesla T4 with single-precision PyTorch tensors. Square grids ranged from $32^2$ to $512^2$. Each measurement used 10 warm-up launches followed by 100 timed launches, with CUDA synchronization before and after the timed region.
Performance results and analysis
Three representative kernels were measured:
| Kernel | Work per launch | Observed behavior |
|---|---|---|
| Interior operator | $O(N^2)$ | Runtime grows strongly once the grid is large enough to dominate launch overhead. |
| Left-face SAT closure | $O(N)$ | Runtime remains nearly flat over the measured range. |
| Diagonal $H$ application | $O(N^2)$ | Runtime grows with the number of grid points, below the more arithmetic-heavy interior operator. |
The interior kernel rises from roughly 35 microseconds at $64^2$ to about 1.25 milliseconds at $512^2$. Over that interval the number of grid points grows by $64\times$, while runtime grows by about $36\times$. The sublinear increase relative to point count is consistent with fixed launch overhead and low GPU occupancy dominating the smaller cases. It should not be interpreted as a better-than-linear algorithm: each interior point still performs a fixed stencil, so the asymptotic work is $O(N^2)$.
The diagonal $H$ kernel follows the same area-dependent trend but remains substantially faster, reaching roughly 450 microseconds at $512^2$. This is expected because each point needs only one load, a boundary-dependent scale, and one store. The interior operator loads several coefficient and neighboring solution values and evaluates pure and mixed derivative terms, so it has much greater memory traffic and arithmetic work per point.
The left-face SAT kernel stays in a narrow range of roughly 30–45 microseconds after the smallest case. Its work grows only with one grid dimension, from 32 to 512 boundary points, and that amount of parallelism is too small to amortize launch overhead effectively on a T4. The slight non-monotonicity is therefore best treated as timing noise and launch-level effects, not as meaningful scaling.
At $512^2$, the interior launch is about $35\times$ slower than one face launch. This confirms that optimizing the full operator should focus first on the two-dimensional interior and on reducing total launch and memory traffic across all five operator kernels. Possible next experiments include fusing compatible face work, processing multiple points per Triton program, and reporting effective memory bandwidth. Those changes should come only after comparison against a numerical reference.
These timings are best read as a scaling sanity check, not a performance comparison. The notebook does not retain the raw timing arrays, and it does not benchmark the original CUDA.jl implementation or a tuned PyTorch baseline.
What the experiment establishes
The notebook’s smoke test successfully launches the interior, all four face closures, $H$, $H^{-1}$, prolongation, and restriction on the Tesla T4. The benchmark also separates the expected $O(N)$ boundary work from the $O(N^2)$ interior work.
There is an important distinction between runs successfully and is numerically verified. Before this port is used in a solver, the next validation steps should be:
- Compare every kernel against a CPU or CUDA.jl reference on small deterministic arrays.
- Verify the discrete SBP identity and an energy estimate, including corner terms.
- Use a manufactured solution to measure convergence under grid refinement.
- Check prolongation and restriction for constant preservation and adjoint consistency.
- Benchmark end-to-end operator applications after correctness is established.
The main implementation lesson is that boundary closures should remain first-class computational regions. Triton makes the interior stencil concise, but the reliability of an SBP-SAT port still depends on explicit indexing, unambiguous write ownership, and tests derived from the numerical method rather than from kernel execution alone.
Full code
Show the complete Triton implementation and benchmark
# Check that a Tesla T4 is connected
!nvidia-smi
# Install or upgrade Triton and PyTorch
!pip install --upgrade torch triton
import torch
from scipy.sparse import random
import numpy as np
# Import the Triton components we defined earlier
import triton
import triton.language as tl
# Notebook cell 1
"""
Triton translations of the CUDA.jl kernels in
AlexandreChern/IJHPCA_revision/CUDA_kernels_second_new.jl.
IMPORTANT:
- All PDE vectors/coefficient arrays are 1-D CUDA tensors containing the
Julia column-major storage, i.e. linear index k in Julia maps to k-1 here.
- Floating point arrays should normally be torch.float32.
- Nr1/Ns1 are the Julia dimensions.
- The four face kernels are intentionally separate, matching the Julia code.
- The large `cuda_knl_2_x` Julia kernel writes neighboring output entries
from boundary threads. Here it is represented by the race-free decomposition
x_interior + x_f1 + x_f2 + x_f3 + x_f4.
"""
import torch
import triton
import triton.language as tl
# ---------------------------------------------------------------------------
# Utility: 2-D interior variable-coefficient operator
# ---------------------------------------------------------------------------
@triton.jit
def x_interior_kernel(
hr, hs, x_ptr, nr1, ns1,
crr_ptr, css_ptr, crs_ptr, out_ptr,
):
# Julia:
# i = global x coordinate, j = global y coordinate
# global_index = (j-1)*Nr1 + i
pid_i = tl.program_id(0)
pid_j = tl.program_id(1)
# Julia source: i runs over Ns1 (slow/column dimension), j over Nr1.
i = pid_i + 2 # Julia i: 2..Ns1-1
j = pid_j + 2 # Julia j: 2..Nr1-1
if i <= ns1 - 1 and j <= nr1 - 1:
g = (i - 1) * nr1 + (j - 1) # 0-based flat offset
crr_m = tl.load(crr_ptr + g - 1)
crr_0 = tl.load(crr_ptr + g)
crr_p = tl.load(crr_ptr + g + 1)
css_m = tl.load(css_ptr + g - nr1)
css_0 = tl.load(css_ptr + g)
css_p = tl.load(css_ptr + g + nr1)
crs_m = tl.load(crs_ptr + g - 1)
crs_p = tl.load(crs_ptr + g + 1)
crs_d = tl.load(crs_ptr + g - nr1)
crs_u = tl.load(crs_ptr + g + nr1)
xm = tl.load(x_ptr + g - 1)
x0 = tl.load(x_ptr + g)
xp = tl.load(x_ptr + g + 1)
xd = tl.load(x_ptr + g - nr1)
xu = tl.load(x_ptr + g + nr1)
xdm = tl.load(x_ptr + g - nr1 - 1)
xdp = tl.load(x_ptr + g - nr1 + 1)
xum = tl.load(x_ptr + g + nr1 - 1)
xup = tl.load(x_ptr + g + nr1 + 1)
arr = (
(-0.5 * crr_m - 0.5 * crr_0) * xm
+ (0.5 * crr_m + crr_0 + 0.5 * crr_p) * x0
+ (-0.5 * crr_0 - 0.5 * crr_p) * xp
)
ass = (
(-0.5 * css_m - 0.5 * css_0) * xd
+ (0.5 * css_m + css_0 + 0.5 * css_p) * x0
+ (-0.5 * css_0 - 0.5 * css_p) * xu
)
ars = (
0.5 * crs_m * (-0.5 * xdm + 0.5 * xum)
- 0.5 * crs_p * (-0.5 * xdp + 0.5 * xup)
)
asr = (
0.5 * crs_d * (-0.5 * xdm + 0.5 * xdp)
- 0.5 * crs_u * (-0.5 * xum + 0.5 * xup)
)
y = hs * (1.0 / hr) * (arr + ass + ars + asr)
tl.store(out_ptr + g, y)
# ---------------------------------------------------------------------------
# Face 1: left boundary (Julia cuda_knl_2_x_f1)
# ---------------------------------------------------------------------------
@triton.jit
def x_f1_kernel(
hr, hs, x_ptr, nr1, ns1,
crr_ptr, css_ptr, crs_ptr, psi1_ptr, psi2_ptr,
out_ptr,
):
i0 = tl.program_id(0) + 1 # Julia i = 1..Ns1
if i0 <= ns1:
g = (i0 - 1) * nr1 # j=1 in Julia, 0-based
if i0 >= 2 and i0 <= ns1 - 1:
crr0 = tl.load(crr_ptr + g)
crr1 = tl.load(crr_ptr + g + 1)
cssm = tl.load(css_ptr + g - nr1)
css0 = tl.load(css_ptr + g)
cssp = tl.load(css_ptr + g + nr1)
crs0 = tl.load(crs_ptr + g)
crs1 = tl.load(crs_ptr + g + 1)
crsd = tl.load(crs_ptr + g - nr1)
crsu = tl.load(crs_ptr + g + nr1)
x0 = tl.load(x_ptr + g)
x1 = tl.load(x_ptr + g + 1)
xd = tl.load(x_ptr + g - nr1)
xu = tl.load(x_ptr + g + nr1)
xu1 = tl.load(x_ptr + g + nr1 + 1)
x2 = tl.load(x_ptr + g + 2 * 0 + 2)
arr = (0.5 * crr0 + 0.5 * crr1) * x0 + (-0.5 * crr0 - 0.5 * crr1) * x1
ass = (
(-0.5 * cssm - 0.5 * css0) * xd
+ (0.5 * cssm + css0 + 0.5 * cssp) * x0
+ (-0.5 * css0 - 0.5 * cssp) * xu
)
ars = (
-0.5 * crs0 * (-0.5 * xd + 0.5 * xu)
-0.5 * crs1 * (-0.5 * xu1 + 0.5 * tl.load(x_ptr + g + 1 - nr1))
)
asr = (
0.5 * crsd * (-0.5 * xd + 0.5 * x1)
-0.5 * crsu * (-0.5 * xu + 0.5 * xu1)
)
mod = (
-hs * (1.0 / hr) * crr0 * (1.5 * x0 - 2.0 * x1 + 0.5 * x2)
+ crs0 * (-0.5 * xd + 0.5 * xu)
+ crr0 * (4.0 + crr0 / tl.load(psi1_ptr + i0 - 1)) * x0
- hs * crr0 * ((1.0 / hr) * 1.5 * x0)
+ 0.5 * crsd * xd
- 0.5 * crsu * xu
)
y = hs * (1.0 / hr) * (arr + ass * 0.5 * (hr / hs) * (hs / hr) + ars + asr) + mod
# The expression above deliberately keeps the Julia factors explicit.
# Reconstruct exact Julia scaling for the first three terms:
y = (
hs * (1.0 / hr) * arr
+ hr * 0.5 * (1.0 / hs) * ass
+ ars + asr + mod
)
tl.store(out_ptr + g, y)
# The Julia f1 kernel also writes the two neighboring entries.
tl.store(
out_ptr + g + 1,
-hs * (crr0 * (1.0 / hr) * (-2.0) * x0)
)
tl.store(
out_ptr + g + 2,
-hs * (crr0 * (1.0 / hr) * 0.5 * x0)
)
# Julia special corner i==1
if i0 == 1:
crr0 = tl.load(crr_ptr + g)
crr1 = tl.load(crr_ptr + g + 1)
crs0 = tl.load(crs_ptr + g)
crs1 = tl.load(crs_ptr + g + 1)
crsu = tl.load(crs_ptr + g + nr1)
css0 = tl.load(css_ptr + g)
cssu = tl.load(css_ptr + g + nr1)
x0 = tl.load(x_ptr + g)
x1 = tl.load(x_ptr + g + 1)
xu = tl.load(x_ptr + g + nr1)
xu1 = tl.load(x_ptr + g + nr1 + 1)
x2 = tl.load(x_ptr + g + 2)
y = (
hs * 0.5 * (1.0 / hr) *
((0.5 * crr0 + 0.5 * crr1) * x0
+ (-0.5 * crr0 - 0.5 * crr1) * x1)
+ (-0.5 * crs0 * (-0.5 * x0 + 0.5 * xu)
- 0.5 * crs1 * (-0.5 * x1 + 0.5 * xu1))
- hs * 0.5 * (1.0 / hr) * crr0 * (1.5 * x0 - 2.0 * x1 + 0.5 * x2)
+ 0.5 * crs0 * (-x0 + xu)
+ 0.5 * crr0 * (4.0 + crr0 / tl.load(psi1_ptr)) * x0
- hs * 0.5 * crr0 * ((1.0 / hr) * 1.5 * x0)
- 0.5 * crsu * xu
- 0.5 * crs0 * x0
+ hr * 0.5 * (1.0 / hs) *
((0.5 * css0 + 0.5 * cssu) * x0
+ (-0.5 * css0 - 0.5 * cssu) * xu)
- 0.5 * (crs0 * (-0.5 * x0 + 0.5 * x1)
+ crsu * (-0.5 * xu + 0.5 * xu1))
)
tl.store(out_ptr + g, y)
tl.store(out_ptr + g + 1, -hs * 0.5 * crr0 * (1.0 / hr) * (-2.0) * x0)
tl.store(out_ptr + g + 2, -hs * 0.5 * crr0 * (1.0 / hr) * 0.5 * x0)
# Julia special corner i==Ns1
if i0 == ns1:
g = (ns1 - 1) * nr1
crr0 = tl.load(crr_ptr + g)
crr1 = tl.load(crr_ptr + g + 1)
crs0 = tl.load(crs_ptr + g)
crs1 = tl.load(crs_ptr + g + 1)
crsd = tl.load(crs_ptr + g - nr1)
cssd = tl.load(css_ptr + g - nr1)
css0 = tl.load(css_ptr + g)
x0 = tl.load(x_ptr + g)
x1 = tl.load(x_ptr + g + 1)
xd = tl.load(x_ptr + g - nr1)
xd1 = tl.load(x_ptr + g - nr1 + 1)
x2 = tl.load(x_ptr + g + 2)
y = (
hs * 0.5 * (1.0 / hr) *
((0.5 * crr0 + 0.5 * crr1) * x0
+ (-0.5 * crr0 - 0.5 * crr1) * x1)
+ (-0.5 * crs0 * (-0.5 * xd + 0.5 * x0)
- 0.5 * crs1 * (-0.5 * xd1 + 0.5 * x1))
- hs * 0.5 * (1.0 / hr) * crr0 * (1.5 * x0 - 2.0 * x1 + 0.5 * x2)
+ 0.5 * crs0 * (-xd + x0)
+ 0.5 * crr0 * (4.0 + crr0 / tl.load(psi1_ptr + ns1 - 1)) * x0
- hs * 0.5 * crr0 * ((1.0 / hr) * 1.5 * x0)
+ 0.5 * crsd * xd
+ 0.5 * crs0 * x0
+ hr * 0.5 * (1.0 / hs) *
((-0.5 * cssd - 0.5 * css0) * xd
+ (0.5 * cssd + 0.5 * css0) * x0)
)
tl.store(out_ptr + g, y)
tl.store(out_ptr + g + 1, -hs * 0.5 * crr0 * (1.0 / hr) * (-2.0) * x0)
tl.store(out_ptr + g + 2, -hs * 0.5 * crr0 * (1.0 / hr) * 0.5 * x0)
# ---------------------------------------------------------------------------
# Face 2 (right), face 3 (bottom), face 4 (top)
#
# These use the same mathematical expressions as the Julia source, with
# one Triton program per face coordinate.
# ---------------------------------------------------------------------------
@triton.jit
def x_f2_kernel(
hr, hs, x_ptr, nr1, ns1,
crr_ptr, css_ptr, crs_ptr, psi1_ptr, psi2_ptr,
out_ptr,
):
j = tl.program_id(0) + 1
if j <= ns1:
g = j * nr1 - 1
if j >= 2 and j <= ns1 - 1:
crrm = tl.load(crr_ptr + g - 1)
crr0 = tl.load(crr_ptr + g)
cssm = tl.load(css_ptr + g - nr1)
css0 = tl.load(css_ptr + g)
cssp = tl.load(css_ptr + g + nr1)
crsm = tl.load(crs_ptr + g - 1)
crs0 = tl.load(crs_ptr + g)
crsd = tl.load(crs_ptr + g - nr1)
crsu = tl.load(crs_ptr + g + nr1)
x_m = tl.load(x_ptr + g - 1)
x0 = tl.load(x_ptr + g)
x_d = tl.load(x_ptr + g - nr1)
x_u = tl.load(x_ptr + g + nr1)
x_dm = tl.load(x_ptr + g - nr1 - 1)
x_um = tl.load(x_ptr + g + nr1 - 1)
y = (
hs * (1.0 / hr) * ((-0.5*crrm-0.5*crr0)*x_m + (0.5*crrm+0.5*crr0)*x0)
+ hr * 0.5 * (1.0 / hs) *
((-0.5*cssm-0.5*css0)*x_d
+ (0.5*cssm+css0+0.5*cssp)*x0
+ (-0.5*css0-0.5*cssp)*x_u)
+ 0.5*crsm*(-0.5*x_dm+0.5*x_um)
+ 0.5*crs0*(-0.5*x_d+0.5*x_u)
- hs*(1.0/hr)*crr0*(1.5*x0-2.0*x_m+0.5*tl.load(x_ptr+g-2))
- crs0*(-0.5*tl.load(x_ptr+g-nr1)+0.5*tl.load(x_ptr+g+nr1))
+ crr0*(4.0+crr0/tl.load(psi2_ptr+j-1))*x0
- hs*crr0*((1.0/hr)*1.5*x0)
- 0.5*crsd* x_d + 0.5*crsu*x_u
)
tl.store(out_ptr+g, y)
tl.store(out_ptr+g-1, -hs*crr0*(1.0/hr)*(-2.0)*x0)
tl.store(out_ptr+g-2, -hs*crr0*(1.0/hr)*0.5*x0)
if j == 1:
g = nr1 - 1
# exact right/bottom corner from f2
crra = tl.load(crr_ptr+g-1); crr0 = tl.load(crr_ptr+g)
crsa = tl.load(crs_ptr+g-1); crs0 = tl.load(crs_ptr+g)
crsd = tl.load(crs_ptr+g-nr1); crsu = tl.load(crs_ptr+g+nr1)
cssd = tl.load(css_ptr+g-nr1); css0 = tl.load(css_ptr+g)
x0 = tl.load(x_ptr+g); xm = tl.load(x_ptr+g-1)
xd = tl.load(x_ptr+g-nr1); xd_m = tl.load(x_ptr+g-nr1-1)
x2m = tl.load(x_ptr+g-2)
y = (
hs*0.5*(1.0/hr)*((-0.5*crra-0.5*crr0)*xm+(0.5*crra+0.5*crr0)*x0)
+ 0.5*crsa*(-0.5*xm+0.5*x0)
+ 0.5*crs0*(-0.5*xd+0.5*x0)
- hs*0.5*(1.0/hr)*crr0*(1.5*x0-2*xm+0.5*x2m)
- 0.5*crs0*(-xd+x0)
+ 0.5*crr0*(4.0+crr0/tl.load(psi2_ptr))*x0
- hs*0.5*crr0*((1.0/hr)*1.5*x0)
+ 0.5*crsu*xd
+ 0.5*crs0*x0
)
tl.store(out_ptr+g, y)
tl.store(out_ptr+g-1, -hs*0.5*crr0*(1.0/hr)*(-2)*x0)
tl.store(out_ptr+g-2, -hs*0.5*crr0*(1.0/hr)*0.5*x0)
if j == ns1:
g = ns1*nr1 - 1
crrm=tl.load(crr_ptr+g-1); crr0=tl.load(crr_ptr+g)
crsm=tl.load(crs_ptr+g-1); crs0=tl.load(crs_ptr+g)
crsd=tl.load(crs_ptr+g-nr1)
cssd=tl.load(css_ptr+g-nr1); css0=tl.load(css_ptr+g)
x0=tl.load(x_ptr+g); xm=tl.load(x_ptr+g-1)
xd=tl.load(x_ptr+g-nr1); xdm=tl.load(x_ptr+g-nr1-1)
x2m=tl.load(x_ptr+g-2)
y=(
hs*0.5*(1.0/hr)*((-0.5*crrm-0.5*crr0)*xm+(0.5*crrm+0.5*crr0)*x0)
+ 0.5*crsm*(-0.5*xdm+0.5*xm)
+ 0.5*crs0*(-0.5*xd+0.5*x0)
- hs*0.5*(1.0/hr)*crr0*(1.5*x0-2*xm+0.5*x2m)
- 0.5*crs0*(-xd+x0)
+ 0.5*crr0*(4.0+crr0/tl.load(psi2_ptr+ns1-1))*x0
- hs*0.5*crr0*((1.0/hr)*1.5*x0)
- 0.5*crsd*xd
- 0.5*crs0*x0
+ hr*0.5*(1.0/hs)*((-0.5*cssd-0.5*css0)*xd+(0.5*cssd+0.5*css0)*x0)
)
tl.store(out_ptr+g,y)
tl.store(out_ptr+g-1,-hs*0.5*crr0*(1.0/hr)*(-2)*x0)
tl.store(out_ptr+g-2,-hs*0.5*crr0*(1.0/hr)*0.5*x0)
@triton.jit
def x_f3_kernel(hr, hs, x_ptr, nr1, ns1, crr_ptr, css_ptr, crs_ptr,
psi1_ptr, psi2_ptr, out_ptr):
i = tl.program_id(0) + 1
if i <= nr1:
g = i - 1
if i >= 2 and i <= nr1 - 1:
cm=tl.load(crr_ptr+g-1); c0=tl.load(crr_ptr+g); cp=tl.load(crr_ptr+g+1)
s0=tl.load(css_ptr+g); su=tl.load(css_ptr+g+nr1)
rm=tl.load(crs_ptr+g-1); rp=tl.load(crs_ptr+g+1); r0=tl.load(crs_ptr+g); ru=tl.load(crs_ptr+g+nr1)
xm=tl.load(x_ptr+g-1); x0=tl.load(x_ptr+g); xp=tl.load(x_ptr+g+1)
xu=tl.load(x_ptr+g+nr1); xum=tl.load(x_ptr+g+nr1-1); xup=tl.load(x_ptr+g+nr1+1)
y=(
hs*0.5*(1/hr)*((-0.5*cm-0.5*c0)*xm+(0.5*cm+c0+0.5*cp)*x0+(-0.5*c0-0.5*cp)*xp)
+ hr*(1/hs)*((0.5*s0+0.5*su)*x0+(-0.5*s0-0.5*su)*xu)
+ 0.5*rm*(-0.5*xm+0.5*xum)
-0.5*rp*(-0.5*xp+0.5*xup)
-0.5*r0*(-0.5*xm+0.5*xp)
-0.5*ru*(-0.5*xum+0.5*xup)
)
tl.store(out_ptr+g,y)
if i == 1:
s0=tl.load(css_ptr+g); su=tl.load(css_ptr+g+nr1)
r0=tl.load(crs_ptr+g); ru=tl.load(crs_ptr+g+nr1)
x0=tl.load(x_ptr+g); x1=tl.load(x_ptr+g+1); xu=tl.load(x_ptr+g+nr1); xu1=tl.load(x_ptr+g+nr1+1)
y=hr*0.5*(1/hs)*((0.5*s0+0.5*su)*x0+(-0.5*s0-0.5*su)*xu) \
-0.5*(r0*(-0.5*x0+0.5*x1)+ru*(-0.5*xu+0.5*xu1))
tl.store(out_ptr+g,y)
if i == nr1:
s0=tl.load(css_ptr+g); su=tl.load(css_ptr+g+nr1)
r0=tl.load(crs_ptr+g); rm=tl.load(crs_ptr+g-1); ru=tl.load(crs_ptr+g+nr1)
x0=tl.load(x_ptr+g); xm=tl.load(x_ptr+g-1); xu=tl.load(x_ptr+g+nr1); xum=tl.load(x_ptr+g+nr1-1)
y=hr*0.5*(1/hs)*((0.5*s0+0.5*su)*x0+(-0.5*s0-0.5*su)*xu) \
-0.5*(r0*(-0.5*xm+0.5*x0)+ru*(-0.5*xum+0.5*xu))
tl.store(out_ptr+g,y)
@triton.jit
def x_f4_kernel(hr, hs, x_ptr, nr1, ns1, crr_ptr, css_ptr, crs_ptr,
psi1_ptr, psi2_ptr, out_ptr):
i = tl.program_id(0) + 1
if i <= nr1:
g = (ns1-1)*nr1 + i - 1
if i >= 2 and i <= nr1 - 1:
cm=tl.load(crr_ptr+g-1); c0=tl.load(crr_ptr+g); cp=tl.load(crr_ptr+g+1)
sm=tl.load(css_ptr+g-nr1); s0=tl.load(css_ptr+g)
rm=tl.load(crs_ptr+g-1); rp=tl.load(crs_ptr+g+1); rd=tl.load(crs_ptr+g-nr1); r0=tl.load(crs_ptr+g)
xm=tl.load(x_ptr+g-1); x0=tl.load(x_ptr+g); xp=tl.load(x_ptr+g+1)
xd=tl.load(x_ptr+g-nr1); xdm=tl.load(x_ptr+g-nr1-1); xdp=tl.load(x_ptr+g-nr1+1)
y=(
hs*0.5*(1/hr)*((-0.5*cm-0.5*c0)*xm+(0.5*cm+c0+0.5*cp)*x0+(-0.5*c0-0.5*cp)*xp)
+ hr*(1/hs)*((-0.5*sm-0.5*s0)*xd+(0.5*sm+0.5*s0)*x0)
+ 0.5*rm*(-0.5*xdm+0.5*xm)
-0.5*rp*(-0.5*xdp+0.5*xp)
+0.5*rd*(-0.5*xm+0.5*xp)
+0.5*r0*(-0.5*xm+0.5*xp)
)
tl.store(out_ptr+g,y)
if i == 1:
sm=tl.load(css_ptr+g-nr1); s0=tl.load(css_ptr+g)
rd=tl.load(crs_ptr+g-nr1); r0=tl.load(crs_ptr+g)
x0=tl.load(x_ptr+g); xp=tl.load(x_ptr+g+1); xd=tl.load(x_ptr+g-nr1); xdp=tl.load(x_ptr+g+1-nr1)
y=hr*0.5*(1/hs)*((-0.5*sm-0.5*s0)*xd+(0.5*sm+0.5*s0)*x0) \
+0.5*(rd*(-0.5*xd+0.5*xdp)+r0*(-0.5*x0+0.5*xp))
tl.store(out_ptr+g,y)
if i == nr1:
sm=tl.load(css_ptr+g-nr1); s0=tl.load(css_ptr+g)
rd=tl.load(crs_ptr+g-nr1); r0=tl.load(crs_ptr+g)
x0=tl.load(x_ptr+g); xm=tl.load(x_ptr+g-1); xd=tl.load(x_ptr+g-nr1)
xdm=tl.load(x_ptr+g-1-nr1)
y=hr*0.5*(1/hs)*((-0.5*sm-0.5*s0)*xd+(0.5*sm+0.5*s0)*x0) \
+0.5*(rd*(-0.5*xdm+0.5*xd)+r0*(-0.5*xm+0.5*x0))
tl.store(out_ptr+g,y)
# ---------------------------------------------------------------------------
# H and H^{-1}
# ---------------------------------------------------------------------------
@triton.jit
def H_kernel(hr, hs, x_ptr, nr1, ns1, out_ptr, inverse: tl.constexpr):
pid_x = tl.program_id(0)
pid_y = tl.program_id(1)
i = pid_x + 1
j = pid_y + 1
if i <= ns1 and j <= nr1:
g = (i - 1) * nr1 + (j - 1)
boundary_x = (i == 1) | (i == ns1)
boundary_y = (j == 1) | (j == nr1)
corner = boundary_x & boundary_y
edge = boundary_x | boundary_y
val = tl.load(x_ptr + g)
if inverse:
scale = 1.0 / (hr * hs)
out = tl.where(corner, 4.0 * scale * val,
tl.where(edge, 2.0 * scale * val, scale * val))
else:
scale = hr * hs
out = tl.where(corner, 0.25 * scale * val,
tl.where(edge, 0.5 * scale * val, scale * val))
tl.store(out_ptr + g, out)
# ---------------------------------------------------------------------------
# Prolongation / restriction
# ---------------------------------------------------------------------------
@triton.jit
def prolongation_2d_kernel(idata_ptr, odata_ptr, nx, ny):
pid_x = tl.program_id(0)
pid_y = tl.program_id(1)
i = pid_x + 1
j = pid_y + 1
# Julia source uses 1-based column-major indexing.
if i <= nx - 1 and j <= ny - 1:
g = (i - 1) * nx + (j - 1)
o_stride = 2 * nx - 1
o00 = (2*i - 2) * (2*nx - 1) + (2*j - 2)
o01 = o00 + 1
o10 = o00 + (2*ny - 1)
o11 = o10 + 1
a = tl.load(idata_ptr + g)
b = tl.load(idata_ptr + g + 1)
c = tl.load(idata_ptr + ny)
d = tl.load(idata_ptr + ny + 1)
tl.store(odata_ptr + o00, a)
tl.store(odata_ptr + o01, 0.5*(a+b))
tl.store(odata_ptr + o10, 0.5*(a+c))
tl.store(odata_ptr + o11, 0.25*(a+c+b+d))
if j <= ny - 1 and i == nx:
g = (i - 1) * nx + (j - 1)
o_stride = 2 * nx - 1
o00 = (2*i - 2) * o_stride + (2*j - 2)
a = tl.load(idata_ptr + g)
b = tl.load(idata_ptr + g + 1)
tl.store(odata_ptr + o00, a)
tl.store(odata_ptr + o00 + 1, 0.5*(a+b))
if i <= nx - 1 and j == ny:
g = (i - 1) * nx + (j - 1)
o_stride = 2 * nx - 1
o00 = (2*i - 2) * o_stride + (2*j - 2)
a = tl.load(idata_ptr + g)
c = tl.load(idata_ptr + ny)
tl.store(odata_ptr + o00, a)
tl.store(odata_ptr + o00 + o_stride, 0.5*(a+c))
if i == nx and j == ny:
g = (i - 1) * nx + (j - 1)
o_stride = 2 * nx - 1
o00 = (2*i - 2) * o_stride + (2*j - 2)
tl.store(odata_ptr + o00, tl.load(idata_ptr + g))
@triton.jit
def restriction_2d_kernel(idata_ptr, odata_ptr, nx, ny):
# This follows the exact indexing formulas in the Julia restriction kernel.
pid_x = tl.program_id(0)
pid_y = tl.program_id(1)
i = pid_x + 1
j = pid_y + 1
nox = (nx + 1) // 2
noy = (ny + 1) // 2
if i <= nox and j <= noy:
# Interior
if i >= 2 and i <= nox-1 and j >= 2 and j <= noy-1:
a = (2*i - 2) * nx + (2*j - 2)
v = (
4.0 * tl.load(idata_ptr + a)
+ 2.0 * (
tl.load(idata_ptr + a + nx)
+ tl.load(idata_ptr + a - nx)
+ tl.load(idata_ptr + a + 1)
+ tl.load(idata_ptr + a - 1)
)
+ tl.load(idata_ptr + a - nx - 1)
+ tl.load(idata_ptr + a + nx + 1)
+ tl.load(idata_ptr + a - nx + 1)
+ tl.load(idata_ptr + a + nx - 1)
) / 16.0
tl.store(odata_ptr + (i-1)*noy + (j-1), v)
# Four corners
if i == 1 and j == 1:
v = (tl.load(idata_ptr) + tl.load(idata_ptr+nx)
+ tl.load(idata_ptr+1) + tl.load(idata_ptr+nx+1)) / 4.0
tl.store(odata_ptr, v)
if i == nox and j == 1:
a = (2*i-2)*ny
v = (tl.load(idata_ptr+a) + tl.load(idata_ptr+a-nx)
+ tl.load(idata_ptr+a+1) + tl.load(idata_ptr+a-nx+1)) / 4.0
tl.store(odata_ptr + (i-1)*noy, v)
if i == 1 and j == noy:
a = 2*j-2
v = (tl.load(idata_ptr+a) + tl.load(idata_ptr+a+nx)
+ tl.load(idata_ptr+a-1) + tl.load(idata_ptr+a+nx-1)) / 4.0
tl.store(odata_ptr + (j-1), v)
if i == nox and j == noy:
a = (2*i-2)*ny + (2*j-2)
v = (tl.load(idata_ptr+a) + tl.load(idata_ptr+a-nx)
+ tl.load(idata_ptr+a-1) + tl.load(idata_ptr+a-nx-1)) / 4.0
tl.store(odata_ptr + (i-1)*noy + (j-1), v)
# Bottom/top edges
if i >= 2 and i <= nox-1 and j == 1:
a = (2*i-2)*ny
v = (
2*tl.load(idata_ptr+a)
+ tl.load(idata_ptr+a-nx)
+ tl.load(idata_ptr+a+nx)
+ 2*tl.load(idata_ptr+a+1)
+ tl.load(idata_ptr+a-nx+1)
+ tl.load(idata_ptr+a+nx+1)
) / 8.0
tl.store(odata_ptr+(i-1)*noy, v)
if i >= 2 and i <= nox-1 and j == noy:
a = (2*i-2)*ny + (2*j-2)
v = (
2*tl.load(idata_ptr+a)
+ tl.load(idata_ptr+a-nx)
+ tl.load(idata_ptr+a+nx)
+ 2*tl.load(idata_ptr+a-1)
+ tl.load(idata_ptr+a-nx-1)
+ tl.load(idata_ptr+a+nx-1)
) / 8.0
tl.store(odata_ptr+(i-1)*noy+(j-1), v)
# Left/right edges
if i == 1 and j >= 2 and j <= noy-1:
a = 2*j-2
v = (
2*tl.load(idata_ptr+a)
+ tl.load(idata_ptr+a-1)
+ tl.load(idata_ptr+a+1)
+ 2*tl.load(idata_ptr+a+nx)
+ tl.load(idata_ptr+a+nx-1)
+ tl.load(idata_ptr+a+nx+1)
) / 8.0
tl.store(odata_ptr+(j-1), v)
if i == nox and j >= 2 and j <= noy-1:
a = (2*i-2)*ny + (2*j-2)
v = (
2*tl.load(idata_ptr+a)
+ tl.load(idata_ptr+a-1)
+ tl.load(idata_ptr+a+1)
+ 2*tl.load(idata_ptr+a-nx)
+ tl.load(idata_ptr+a-nx-1)
+ tl.load(idata_ptr+a-nx+1)
) / 8.0
tl.store(odata_ptr+(i-1)*noy+(j-1), v)
# ---------------------------------------------------------------------------
# Python launch wrappers
# ---------------------------------------------------------------------------
def launch_x_interior(hr, hs, x, Nr1, Ns1, crr, css, crs, out):
grid = (max(Ns1-2, 0), max(Nr1-2, 0))
x_interior_kernel[grid](hr, hs, x, Nr1, Ns1, crr, css, crs, out)
return out
def launch_x_f1(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out):
x_f1_kernel[(Ns1,)](hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
return out
def launch_x_f2(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out):
x_f2_kernel[(Ns1,)](hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
return out
def launch_x_f3(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out):
x_f3_kernel[(Nr1,)](hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
return out
def launch_x_f4(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out):
x_f4_kernel[(Nr1,)](hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
return out
def launch_H(hr, hs, x, Nr1, Ns1, out):
H_kernel[(Ns1, Nr1)](hr, hs, x, Nr1, Ns1, out, inverse=False)
return out
def launch_H_inverse(hr, hs, x, Nr1, Ns1, out):
H_kernel[(Ns1, Nr1)](hr, hs, x, Nr1, Ns1, out, inverse=True)
return out
def launch_prolongation(x, out, Nx, Ny):
grid = ((Nx + 15)//16, (Ny + 15)//16)
prolongation_2d_kernel[grid](x, out, Nx, Ny)
return out
def launch_restriction(x, out, Nx, Ny):
nox = (Nx + 1)//2
noy = (Ny + 1)//2
grid = ((nox + 15)//16, (noy + 15)//16)
restriction_2d_kernel[grid](x, out, Nx, Ny)
return out
# ---------------------------------------------------------------------------
# Smoke test
# ---------------------------------------------------------------------------
def smoke_test():
assert torch.cuda.is_available(), "CUDA GPU is required"
device = "cuda"
Nr1, Ns1 = 64, 48
n = Nr1 * Ns1
hr = 0.1
hs = 0.2
x = torch.randn(n, device=device, dtype=torch.float32)
crr = torch.rand(n, device=device, dtype=torch.float32) + 1.0
css = torch.rand(n, device=device, dtype=torch.float32) + 1.0
crs = torch.rand(n, device=device, dtype=torch.float32) * 0.1
psi1 = torch.rand(Ns1, device=device, dtype=torch.float32) + 1.0
psi2 = torch.rand(Ns1, device=device, dtype=torch.float32) + 1.0
out = torch.empty_like(x)
launch_x_interior(hr, hs, x, Nr1, Ns1, crr, css, crs, out)
launch_x_f1(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
launch_x_f2(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
launch_x_f3(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
launch_x_f4(hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
h_out = torch.empty_like(x)
launch_H(hr, hs, x, Nr1, Ns1, h_out)
launch_H_inverse(hr, hs, x, Nr1, Ns1, h_out)
# For transfer kernels use dimensions for which the formulas are valid.
nx, ny = 17, 17
fine_x = 2*nx - 1
fine_y = 2*ny - 1
coarse = torch.randn(nx*ny, device=device)
fine = torch.empty(fine_x*fine_y, device=device)
launch_prolongation(coarse, fine, nx, ny)
restricted = torch.empty(((fine_x+1)//2)*((fine_y+1)//2),
device=device)
launch_restriction(fine, restricted, fine_x, fine_y)
torch.cuda.synchronize()
print("Triton smoke test completed.")
if __name__ == "__main__":
smoke_test()
# Notebook cell 2
import time
import matplotlib.pyplot as plt
def benchmark_kernel(kernel_func, *args, repeats=100):
# Warmup
for _ in range(10):
kernel_func(*args)
torch.cuda.synchronize()
start_time = time.perf_counter()
for _ in range(repeats):
kernel_func(*args)
torch.cuda.synchronize()
end_time = time.perf_counter()
return (end_time - start_time) / repeats * 1e6 # Time in microseconds
def run_benchmarks():
sizes = [32, 64, 128, 256, 512]
interior_times = []
f1_times = []
h_times = []
device = "cuda"
for N in sizes:
Nr1, Ns1 = N, N
n = Nr1 * Ns1
hr, hs = 0.1, 0.2
x = torch.randn(n, device=device, dtype=torch.float32)
crr = torch.rand(n, device=device, dtype=torch.float32) + 1.0
css = torch.rand(n, device=device, dtype=torch.float32) + 1.0
crs = torch.rand(n, device=device, dtype=torch.float32) * 0.1
psi1 = torch.rand(Ns1, device=device, dtype=torch.float32) + 1.0
psi2 = torch.rand(Ns1, device=device, dtype=torch.float32) + 1.0
out = torch.empty_like(x)
t_int = benchmark_kernel(launch_x_interior, hr, hs, x, Nr1, Ns1, crr, css, crs, out)
t_f1 = benchmark_kernel(launch_x_f1, hr, hs, x, Nr1, Ns1, crr, css, crs, psi1, psi2, out)
t_h = benchmark_kernel(launch_H, hr, hs, x, Nr1, Ns1, out)
interior_times.append(t_int)
f1_times.append(t_f1)
h_times.append(t_h)
# Plotting
plt.figure(figsize=(10, 6))
plt.plot(sizes, interior_times, marker='o', label='Interior Kernel')
plt.plot(sizes, f1_times, marker='s', label='Face 1 Kernel')
plt.plot(sizes, h_times, marker='^', label='H Matrix Kernel')
plt.title('Triton Kernel Execution Time vs. Grid Dimension')
plt.xlabel('Grid Dimension (N x N)')
plt.ylabel('Execution Time (microseconds)')
plt.yscale('log')
plt.grid(True, which="both", ls="--")
plt.legend()
plt.show()
run_benchmarks()
# Notebook cell 4
!jupyter nbconvert --to markdown Triton_SBP_SAT.ipynb
Discussion