[WIP] implement all benchmarks on triton

environment.yml

@@ -13,5 +13,6 @@ dependencies:
   - numba
   - pythran
   - cupy
+  - triton
   # pygount has to be installed via pip
 

framework_info/triton.json

@@ -0,0 +1,10 @@
+{
+    "framework": {
+        "simple_name": "triton",
+        "full_name": "Triton",
+        "prefix": "tr",
+        "postfix": "triton",
+        "class": "TritonFramework",
+        "arch": "gpu"
+    }
+}

npbench/benchmarks/azimint_hist/azimint_hist.py

@@ -1,8 +1,9 @@
 # Copyright 2021 ETH Zurich and the NPBench authors. All rights reserved.
+import numpy as np
 
 
-def initialize(N):
+def initialize(N, datatype=np.float32):
     from numpy.random import default_rng
     rng = default_rng(42)
-    data, radius = rng.random((N, )), rng.random((N, ))
+    data, radius = rng.random((N, ), dtype=datatype), rng.random((N, ), dtype=datatype)
     return data, radius

npbench/benchmarks/azimint_hist/azimint_hist_dace.py

@@ -36,12 +36,13 @@
 
 import numpy as np
 import dace as dc
+from npbench.infrastructure.dace_framework import dc_float
 
 N, bins, npt = (dc.symbol(s, dtype=dc.int64) for s in ('N', 'bins', 'npt'))
 
 
 @dc.program
-def get_bin_edges(a: dc.float64[N], bin_edges: dc.float64[bins + 1]):
+def get_bin_edges(a: dc_float[N], bin_edges: dc_float[bins + 1]):
     a_min = np.amin(a)
     a_max = np.amax(a)
     delta = (a_max - a_min) / bins
@@ -52,15 +53,15 @@ def get_bin_edges(a: dc.float64[N], bin_edges: dc.float64[bins + 1]):
 
 
 @dc.program
-def compute_bin(x: dc.float64, bin_edges: dc.float64[bins + 1]):
+def compute_bin(x: dc_float, bin_edges: dc_float[bins + 1]):
     # assuming uniform bins for now
     a_min = bin_edges[0]
     a_max = bin_edges[bins]
     return dc.int64(bins * (x - a_min) / (a_max - a_min))
 
 
 @dc.program
-def histogram(a: dc.float64[N], bin_edges: dc.float64[bins + 1]):
+def histogram(a: dc_float[N], bin_edges: dc_float[bins + 1]):
     hist = np.ndarray((bins, ), dtype=np.int64)
     hist[:] = 0
     get_bin_edges(a, bin_edges)
@@ -73,8 +74,8 @@ def histogram(a: dc.float64[N], bin_edges: dc.float64[bins + 1]):
 
 
 @dc.program
-def histogram_weights(a: dc.float64[N], bin_edges: dc.float64[bins + 1],
-                      weights: dc.float64[N]):
+def histogram_weights(a: dc_float[N], bin_edges: dc_float[bins + 1],
+                      weights: dc_float[N]):
     hist = np.ndarray((bins, ), dtype=weights.dtype)
     hist[:] = 0
     get_bin_edges(a, bin_edges)
@@ -87,11 +88,11 @@ def histogram_weights(a: dc.float64[N], bin_edges: dc.float64[bins + 1],
 
 
 @dc.program
-def azimint_hist(data: dc.float64[N], radius: dc.float64[N]):
+def azimint_hist(data: dc_float[N], radius: dc_float[N]):
     # histu = np.histogram(radius, npt)[0]
-    bin_edges_u = np.ndarray((npt + 1, ), dtype=np.float64)
+    bin_edges_u = np.ndarray((npt + 1, ), dtype=dc_float)
     histu = histogram(radius, bin_edges_u)
     # histw = np.histogram(radius, npt, weights=data)[0]
-    bin_edges_w = np.ndarray((npt + 1, ), dtype=np.float64)
+    bin_edges_w = np.ndarray((npt + 1, ), dtype=dc_float)
     histw = histogram_weights(radius, bin_edges_w, data)
     return histw / histu

npbench/benchmarks/azimint_hist/azimint_hist_triton.py

@@ -0,0 +1,110 @@
+import torch
+import triton
+import triton.language as tl
+import itertools
+
+
+def get_configs():
+    return [
+        triton.Config({"BLOCK_SIZE": block_size}, num_warps=num_warps)
+        for block_size, num_warps in itertools.product(
+            [32, 64, 128, 256, 512, 1024], [1, 2, 4, 8]
+        )
+    ]
+
+
+@triton.autotune(
+    configs=get_configs(),
+    key=["N", "npt"],
+    cache_results=True,
+)
+@triton.jit
+def azimint_hist_kernel(
+    data_ptr,
+    radius_ptr,
+    histw_ptr,
+    histu_ptr,
+    N,
+    npt,
+    rmin,
+    rmax,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Kernel 1: Computes weighted and unweighted histograms for azimuthal integration.
+    Equivalent to
+    histu = np.histogram(radius, npt)[0]
+    histw = np.histogram(radius, npt, weights=data)[0]
+    """
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < N
+    r = tl.load(radius_ptr + offsets, mask=mask)
+    d = tl.load(data_ptr + offsets, mask=mask)
+
+    rmin = tl.load(rmin)
+    rmax = tl.load(rmax)
+
+    # TODO: avoid division by zero
+    normalized_r = npt * (r - rmin) / (rmax - rmin)
+    bin_idx = tl.floor(normalized_r)
+    bin_idx = tl.clamp(bin_idx, 0, npt - 1).to(tl.int32)
+
+    histw_offsets = histw_ptr + bin_idx
+    histu_offsets = histu_ptr + bin_idx
+
+    tl.atomic_add(histw_offsets, d, mask=mask)
+    tl.atomic_add(histu_offsets, 1.0, mask=mask)
+
+
+@triton.autotune(
+    configs=get_configs(),
+    key=["npt"],
+    cache_results=True,
+)
+@triton.jit
+def azimint_div_kernel(
+    histw_ptr,
+    histu_ptr,
+    result_ptr,
+    npt,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    Kernel 2: Computes the final azimuthal integration result by dividing
+    the weighted histogram by the unweighted histogram.
+    Equivalent to
+    return histw / histu
+    """
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    mask = offsets < npt
+    histw = tl.load(histw_ptr + offsets, mask=mask)
+    histu = tl.load(histu_ptr + offsets, mask=mask)
+
+    tl.store(result_ptr + offsets, histw / histu, mask=mask)
+
+
+def azimint_hist(data: torch.Tensor, radius: torch.Tensor, npt: int):
+    """
+    histu = np.histogram(radius, npt)[0]
+    histw = np.histogram(radius, npt, weights=data)[0]
+    return histw / histu
+    """
+    rmin = radius.min().to(data.dtype)
+    rmax = radius.max().to(data.dtype)
+
+    histw = torch.zeros(npt, dtype=data.dtype, device=data.device)
+    histu = torch.zeros(npt, dtype=data.dtype, device=data.device)
+
+    N = data.shape[0]
+
+    grid_hist = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),)
+    azimint_hist_kernel[grid_hist](data, radius, histw, histu, N, npt, rmin, rmax)
+    result = torch.zeros(npt, dtype=data.dtype, device=data.device)
+    grid_div = lambda meta: (triton.cdiv(npt, meta["BLOCK_SIZE"]),)
+    azimint_div_kernel[grid_div](histw, histu, result, npt)
+
+    return result

npbench/benchmarks/azimint_naive/azimint_naive.py

@@ -1,8 +1,8 @@
 # Copyright 2021 ETH Zurich and the NPBench authors. All rights reserved.
 
 
-def initialize(N):
+def initialize(N, datatype):
     from numpy.random import default_rng
     rng = default_rng(42)
-    data, radius = rng.random((N, )), rng.random((N, ))
+    data, radius = rng.random((N, ), dtype=datatype), rng.random((N, ), dtype=datatype)
     return data, radius

npbench/benchmarks/azimint_naive/azimint_naive_dace.py

@@ -10,15 +10,16 @@
 
 import numpy as np
 import dace as dc
+from npbench.infrastructure.dace_framework import dc_float
 
 N, npt = (dc.symbol(s, dtype=dc.int64) for s in ('N', 'npt'))
 
 
 @dc.program
-def azimint_naive(data: dc.float64[N], radius: dc.float64[N]):
+def azimint_naive(data: dc_float[N], radius: dc_float[N]):
     # rmax = radius.max()
     rmax = np.amax(radius)
-    res = np.zeros((npt, ), dtype=np.float64)  # Fix in np.full
+    res = np.zeros((npt, ), dtype=dc_float)  # Fix in np.full
     for i in range(npt):
         # for i in dc.map[0:npt]:  # Optimization
         r1 = rmax * i / npt
@@ -27,7 +28,7 @@ def azimint_naive(data: dc.float64[N], radius: dc.float64[N]):
         # values_r12 = data[mask_r12]
         # res[i] = np.mean(values_r12)
         on_values = 0
-        tmp = np.float64(0)
+        tmp = dc_float(0)
         for j in dc.map[0:N]:
             if mask_r12[j]:
                 tmp += data[j]

npbench/benchmarks/azimint_naive/azimint_naive_triton.py

@@ -0,0 +1,115 @@
+import itertools
+import torch
+import triton
+import triton.language as tl
+
+def generate_config():
+    return [
+        triton.Config(kwargs={"BLOCK_SIZE": m}, num_warps=w)
+        for m, w in itertools.product(
+            [256, 512, 1024], [1, 2, 4, 8]
+        )
+    ]
+
+
+def generate_config_npt():
+    return [
+        triton.Config(kwargs={"BLOCK_SIZE_NPT": m}, num_warps=w)
+        for m, w in itertools.product(
+            [8, 16, 32, 64, 128], [1, 2, 4, 8]
+        )
+    ]
+
+@triton.autotune(configs=generate_config(), key=["N"], cache_results=True)
+@triton.jit
+def _kernel_max(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
+    pid  = tl.program_id(axis=0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < N
+    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
+    m = tl.max(x, axis=0)
+    tl.store(out_ptr + pid, m)
+
+
+def triton_max(x: torch.Tensor):
+    cur = x
+    n = cur.numel()
+
+    MIN_BLOCK = 8
+    while n > 1:
+        grid_size = triton.cdiv(n, MIN_BLOCK)
+        out = torch.empty(grid_size, dtype=cur.dtype)
+        _kernel_max[(grid_size,)](cur, out, n)
+        cur = out
+        n = cur.numel()
+    return cur[0]
+
+@triton.autotune(configs=generate_config(), key=["N"], cache_results=True)
+@triton.jit
+def _accumulate_bins_kernel(data_ptr, radius_ptr,
+                            sums_ptr, counts_ptr,
+                            N, n_bins, rmax: tl.float64,
+                            BLOCK_SIZE: tl.constexpr):
+    # axis 0 = bin index; axis 1 = block id over the data
+    bin_idx = tl.program_id(axis=0)
+    pid = tl.program_id(axis=1)
+
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < N
+
+    r = tl.load(radius_ptr + offs, mask=mask, other=0.0)
+    rmax64 = tl.full((), rmax, tl.float64)
+    n_bins64 = tl.full((), n_bins, tl.float64)
+    bin64 = bin_idx.to(tl.float64)
+
+    r1 = rmax64 * bin64 / n_bins64
+    r2 = rmax64 * (bin64 + 1.0) / n_bins64
+
+    # faster version but worse error:
+    # r1 = rmax * bin_idx / n_bins
+    # r2 = rmax * (bin_idx + 1.0) / n_bins
+
+    in_bin = (r1 <= r) & (r < r2)
+
+    v = tl.load(data_ptr + offs, mask=mask & in_bin, other=0.0)
+    value = tl.sum(v, axis=0)
+
+    counter = tl.sum((in_bin & mask), axis=0)
+
+    tl.atomic_add(sums_ptr + bin_idx, value)
+    tl.atomic_add(counts_ptr + bin_idx, counter)
+
+@triton.autotune(configs=generate_config_npt(), key=["n_bins"], cache_results=True)
+@triton.jit
+def _finalize_means_kernel(sums_ptr, counts_ptr, means_ptr, n_bins,
+                           BLOCK_SIZE_NPT: tl.constexpr):
+    pid  = tl.program_id(axis=0)
+    offs = pid * BLOCK_SIZE_NPT + tl.arange(0, BLOCK_SIZE_NPT)
+    mask = offs < n_bins
+
+    s = tl.load(sums_ptr   + offs, mask=mask, other=0.0)
+    c = tl.load(counts_ptr + offs, mask=mask, other=0)
+
+    mean = tl.where(c > 0, s / c, 0.0)
+    tl.store(means_ptr + offs, mean, mask=mask)
+
+
+def azimint_naive(data: torch.Tensor, radius: torch.Tensor, npt: int):
+    N = data.numel()
+
+    rmax = triton_max(radius).item()
+
+    sums   = torch.zeros(npt, dtype=torch.float64)
+    counts = torch.zeros(npt, dtype=torch.int32)
+    means  = torch.empty(npt, dtype=torch.float64)
+
+    grid = lambda meta: (npt, triton.cdiv(N, meta["BLOCK_SIZE"]),)
+    _accumulate_bins_kernel[grid](
+        data, radius, sums, counts, N, npt, rmax,
+    )
+
+    grid = lambda meta: (triton.cdiv(npt, meta["BLOCK_SIZE_NPT"]),)
+    _finalize_means_kernel[grid](
+        sums, counts, means, npt,
+    )
+    return means

npbench/benchmarks/cavity_flow/cavity_flow.py

@@ -3,10 +3,10 @@
 import numpy as np
 
 
-def initialize(ny, nx):
-    u = np.zeros((ny, nx), dtype=np.float64)
-    v = np.zeros((ny, nx), dtype=np.float64)
-    p = np.zeros((ny, nx), dtype=np.float64)
+def initialize(ny, nx, datatype=np.float32):
+    u = np.zeros((ny, nx), dtype=datatype)
+    v = np.zeros((ny, nx), dtype=datatype)
+    p = np.zeros((ny, nx), dtype=datatype)
     dx = 2 / (nx - 1)
     dy = 2 / (ny - 1)
     dt = .1 / ((nx - 1) * (ny - 1))