nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/addr.py

@@ -0,0 +1,68 @@
+"""Addr (outer product add) — pure @triton.jit kernel.
+
+Computes: output = beta * input + alpha * (vec1 outer vec2)
+
+Where vec1 is [M], vec2 is [N], input is [M, N] (broadcast-compatible),
+and output is [M, N].
+
+Extracted from FlagGems reference (addr.py).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit(do_not_specialize=['beta', 'alpha'])
+def addr_kernel(
+    input_ptr,
+    vec1_ptr,
+    vec2_ptr,
+    output_ptr,
+    beta,
+    alpha,
+    M, N,
+    stride_input_m,
+    stride_input_n,
+    stride_vec1,
+    stride_vec2,
+    stride_output_m,
+    stride_output_n,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Outer product add: out = beta * input + alpha * (vec1 outer vec2).
+
+    Grid: (cdiv(M, BLOCK_M), cdiv(N, BLOCK_N))
+    """
+    pid_m = tl.program_id(0).to(tl.int64)
+    pid_n = tl.program_id(1).to(tl.int64)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    mask_m = offs_m < M
+    mask_n = offs_n < N
+    mask_2d = mask_m[:, None] & mask_n[None, :]
+
+    # Load vectors
+    v1 = tl.load(vec1_ptr + offs_m * stride_vec1, mask=mask_m, other=0.0).to(tl.float32)
+    v2 = tl.load(vec2_ptr + offs_n * stride_vec2, mask=mask_n, other=0.0).to(tl.float32)
+
+    # Load input matrix
+    input_ptrs = (
+        input_ptr
+        + offs_m[:, None] * stride_input_m
+        + offs_n[None, :] * stride_input_n
+    )
+    inp = tl.load(input_ptrs, mask=mask_2d, other=0.0).to(tl.float32)
+
+    # Compute: beta * input + alpha * outer(vec1, vec2)
+    result = beta * inp + alpha * (v1[:, None] * v2[None, :])
+
+    # Store
+    output_ptrs = (
+        output_ptr
+        + offs_m[:, None] * stride_output_m
+        + offs_n[None, :] * stride_output_n
+    )
+    tl.store(output_ptrs, result, mask=mask_2d)

nbxtensor/ops/all_reduce.py

@@ -0,0 +1,83 @@
+"""All reduction (logical AND) — pure @triton.jit kernels.
+
+Tests if ALL elements evaluate to True (non-zero).
+Global two-pass reduction and dim-specific kernel.
+Extracted from FlagGems (Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def reduce_all(a, b):
+    return a & b
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def all_kernel_mid(
+    inp,
+    mid,
+    n_elements,
+    mid_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pass 1: check if all elements in this block are non-zero."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    inp_ptrs = inp + offset
+    mask = offset < n_elements
+    inp_val = tl.load(inp_ptrs, mask=mask, other=1.0)
+    all_val = tl.reduce(inp_val != 0, axis=0, combine_fn=reduce_all)
+    mid_ptr = mid + pid
+    tl.store(mid_ptr, all_val)
+
+
+@triton.jit
+def all_kernel_result(mid, out, MID_SIZE, BLOCK_MID: tl.constexpr):
+    """Pass 2: combine partial results into scalar bool."""
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid + offset
+    mask = offset < MID_SIZE
+    mid_val = tl.load(mid_ptrs, mask=mask, other=1).to(tl.int1)
+    all_val = tl.reduce(mid_val, axis=0, combine_fn=reduce_all)
+    tl.store(out, all_val)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: all along one axis
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def all_kernel_dim(
+    inp,
+    out,
+    M,
+    N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Check all non-zero along compressed inner dim N.
+
+    inp: [M, N] (dim-compressed) → out: [M] (bool)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    inp = inp + rows * N
+    out = out + rows
+    row_mask = rows < M
+
+    _all = tl.full([BLOCK_M, BLOCK_N], value=1, dtype=tl.int1)
+    for off in range(0, N, BLOCK_N):
+        cols = off + tl.arange(0, BLOCK_N)[None, :]
+        col_mask = cols < N
+        mask = row_mask & col_mask
+
+        a = tl.load(inp + cols, mask, other=1.0)
+        _all = _all & (a != 0)
+    result = tl.reduce(_all, axis=1, combine_fn=reduce_all)
+    tl.store(out, result[:, None], row_mask)

nbxtensor/ops/amax.py

@@ -0,0 +1,40 @@
+"""Max reduction — pure @triton.jit kernel, the row walked in tiles of the certified width."""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import nbx_autotune, row_reduce_configs, prune_row_reduce
+
+
+@nbx_autotune(configs=row_reduce_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET'],
+              prune_configs_by={'early_config_prune': prune_row_reduce}, cache_results=True)
+@triton.jit
+def amax_forward_kernel(
+    input_ptr, output_ptr,
+    batch_dim, feat_dim, BATCH_BUCKET, FEAT_BUCKET,
+    input_batch_stride, input_feat_stride,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+):
+    """Max reduction over the last dimension: input [batch_dim, feat_dim] -> output [batch_dim].
+
+    The first tile is reduced alone and every further tile combined into it, so a row that fits one
+    tile is reduced by exactly the instructions it was before the loop existed.
+    """
+    batch_pid = tl.program_id(0).to(tl.int64)
+    batch_offset = batch_pid * BLOCK_SIZE_BATCH + tl.arange(0, BLOCK_SIZE_BATCH)
+    feat_offset = tl.arange(0, BLOCK_SIZE_FEAT)
+    batch_mask = batch_offset < batch_dim
+    row_ptr = input_ptr + input_batch_stride * batch_offset[:, None]
+
+    feat_mask = feat_offset < feat_dim
+    inp = tl.load(row_ptr + input_feat_stride * feat_offset[None, :],
+                  mask=batch_mask[:, None] & feat_mask[None, :], other=-float('inf'))
+    result = tl.max(inp, axis=1)
+    for start in range(BLOCK_SIZE_FEAT, feat_dim, BLOCK_SIZE_FEAT):
+        offs = start + feat_offset
+        feat_mask = offs < feat_dim
+        inp = tl.load(row_ptr + input_feat_stride * offs[None, :].to(tl.int64),
+                      mask=batch_mask[:, None] & feat_mask[None, :], other=-float('inf'))
+        result = tl.maximum(result, tl.max(inp, axis=1))
+    tl.store(output_ptr + batch_offset, result, mask=batch_mask)

nbxtensor/ops/any_reduce.py

@@ -0,0 +1,83 @@
+"""Any reduction (logical OR) — pure @triton.jit kernels.
+
+Tests if ANY element evaluates to True (non-zero).
+Global two-pass reduction and dim-specific kernel.
+Extracted from FlagGems (Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def reduce_any(a, b):
+    return a | b
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def any_kernel_mid(
+    inp,
+    mid,
+    n_elements,
+    mid_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pass 1: check if any element in this block is non-zero."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    inp_ptrs = inp + offset
+    mask = offset < n_elements
+    inp_val = tl.load(inp_ptrs, mask=mask, other=0.0)
+    any_val = tl.reduce(inp_val != 0, axis=0, combine_fn=reduce_any)
+    mid_ptr = mid + pid
+    tl.store(mid_ptr, any_val)
+
+
+@triton.jit
+def any_kernel_result(mid, out, MID_SIZE, BLOCK_MID: tl.constexpr):
+    """Pass 2: combine partial results into scalar bool."""
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid + offset
+    mask = offset < MID_SIZE
+    mid_val = tl.load(mid_ptrs, mask=mask, other=0).to(tl.int1)
+    any_val = tl.reduce(mid_val, axis=0, combine_fn=reduce_any)
+    tl.store(out, any_val)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: any along one axis
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def any_kernel_dim(
+    inp,
+    out,
+    M,
+    N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Check any non-zero along compressed inner dim N.
+
+    inp: [M, N] (dim-compressed) → out: [M] (bool)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    rows = pid * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    inp = inp + rows * N
+    out = out + rows
+    row_mask = rows < M
+
+    _any = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.int1)
+    for off in range(0, N, BLOCK_N):
+        cols = off + tl.arange(0, BLOCK_N)[None, :]
+        col_mask = cols < N
+        mask = row_mask & col_mask
+
+        a = tl.load(inp + cols, mask, other=0.0)
+        _any = _any | (a != 0)
+    result = tl.reduce(_any, axis=1, combine_fn=reduce_any)
+    tl.store(out, result[:, None], row_mask)

nbxtensor/ops/arange_op.py

@@ -0,0 +1,27 @@
+"""Arange — pure @triton.jit kernel.
+
+Ported from FlagGems arange_func (Apache-2.0 license).
+Generates sequential values: start, start+step, start+2*step, ...
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def arange_kernel(
+    output_ptr,
+    start,
+    step,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Generate sequential values: output[i] = start + i * step."""
+    pid = tl.program_id(0).to(tl.int64)
+    block_start = pid * BLOCK_SIZE
+    offsets = tl.arange(0, BLOCK_SIZE)
+    idx = block_start + offsets
+    mask = idx < n_elements
+
+    value = start + idx * step
+    tl.store(output_ptr + idx, value, mask=mask)

nbxtensor/ops/argmax.py

@@ -0,0 +1,107 @@
+"""Argmax — pure @triton.jit kernels. Extracted from FlagGems.
+
+Three-kernel approach:
+  argmax_kernel_1: per-block max reduction (flat tensor, dim=None)
+  argmax_kernel_2: final reduction across blocks (flat tensor, dim=None)
+  argmax_kernel_inner: along last dim (dim=-1)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def argmax_kernel_1(
+    inp,
+    mid_value,
+    mid_index,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    inp_ptrs = inp + offset
+    mask = offset < M
+    inp_val = tl.load(inp_ptrs, mask=mask, other=float('-inf'))
+    max_val, max_index = tl.max(inp_val, axis=0, return_indices=True)
+    max_index = max_index + pid * BLOCK_SIZE
+    mid_value_ptr = mid_value + pid
+    max_index_ptr = mid_index + pid
+    tl.store(mid_value_ptr, max_val)
+    tl.store(max_index_ptr, max_index)
+
+
+@triton.jit
+def argmax_kernel_2(mid_value, mid_index, out, mid_size, BLOCK_MID: tl.constexpr):
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid_value + offset
+    mask = offset < mid_size
+    mid_val = tl.load(mid_ptrs, mask=mask, other=float('-inf'))
+    index_val = tl.argmax(mid_val, axis=0)
+    mid_index_ptrs = mid_index + index_val
+    out_val = tl.load(mid_index_ptrs)
+    tl.store(out, out_val)
+
+
+@triton.jit
+def argmax_kernel_inner(
+    inp,
+    out_index,
+    M,
+    N,
+    TILE_N: tl.constexpr,
+    ONE_TILE_PER_CTA: tl.constexpr,
+):
+    """Argmax along the innermost (last) dimension."""
+    pid_m = tl.program_id(0).to(tl.int64)
+
+    min_value = float('-inf')
+
+    if ONE_TILE_PER_CTA:
+        n_offset = tl.arange(0, TILE_N)
+        offset = pid_m * N + n_offset
+        mask = n_offset < N
+        inp_ptrs = inp + offset
+        inp_vals = tl.load(inp_ptrs, mask=mask, other=min_value)
+        local_max, local_argmax = tl.max(
+            inp_vals, 0, return_indices=True, return_indices_tie_break_left=True
+        )
+        out_index_ptrs = out_index + pid_m
+        tl.store(out_index_ptrs, local_argmax)
+    else:
+        max_values = min_value
+        argmax_values = 0
+
+        loop_time = N // TILE_N
+        remainder = N % TILE_N
+        for start_n in range(0, loop_time):
+            n_offset = start_n * TILE_N + tl.arange(0, TILE_N)
+            offset = pid_m * N + n_offset
+            inp_ptrs = inp + offset
+            inp_vals = tl.load(inp_ptrs)
+            local_max, local_argmax = tl.max(
+                inp_vals, 0, return_indices=True, return_indices_tie_break_left=True
+            )
+            update = local_max > max_values
+            max_values = tl.where(update, local_max, max_values)
+            argmax_values = tl.where(
+                update, start_n * TILE_N + local_argmax, argmax_values
+            )
+
+        if remainder:
+            n_offset = loop_time * TILE_N + tl.arange(0, TILE_N)
+            offset = pid_m * N + n_offset
+            mask = n_offset < N
+            inp_ptrs = inp + offset
+            inp_vals = tl.load(inp_ptrs, mask=mask, other=min_value)
+            local_max, local_argmax = tl.max(
+                inp_vals, 0, return_indices=True, return_indices_tie_break_left=True
+            )
+            update = local_max > max_values
+            max_values = tl.where(update, local_max, max_values)
+            argmax_values = tl.where(
+                update, loop_time * TILE_N + local_argmax, argmax_values
+            )
+
+        out_index_ptrs = out_index + pid_m
+        tl.store(out_index_ptrs, argmax_values)

nbxtensor/ops/argmin.py

@@ -0,0 +1,93 @@
+"""Argmin — pure @triton.jit kernels. Extracted from FlagGems.
+
+Three-kernel approach:
+  argmin_kernel_1: per-block min reduction (flat tensor, dim=None)
+  argmin_kernel_2: final reduction across blocks (flat tensor, dim=None)
+  argmin_kernel: along arbitrary dim (MxNxK layout)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def argmin_kernel_1(
+    inp,
+    mid_value,
+    mid_index,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    inp_ptrs = inp + offset
+    mask = offset < M
+
+    inp_val = tl.load(inp_ptrs, mask=mask, other=float('inf'))
+    min_val, min_index = tl.min(inp_val, axis=0, return_indices=True)
+    min_index = min_index + pid * BLOCK_SIZE
+    mid_value_ptr = mid_value + pid
+    min_index_ptr = mid_index + pid
+    tl.store(mid_value_ptr, min_val)
+    tl.store(min_index_ptr, min_index)
+
+
+@triton.jit
+def argmin_kernel_2(
+    mid_value,
+    mid_index,
+    out,
+    mid_size,
+    BLOCK_MID: tl.constexpr,
+):
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid_value + offset
+    mask = offset < mid_size
+    mid_val = tl.load(mid_ptrs, mask=mask, other=float('inf'))
+    index_val = tl.argmin(mid_val, axis=0)
+    mid_index_ptrs = mid_index + index_val
+    out_val = tl.load(mid_index_ptrs)
+    tl.store(out, out_val)
+
+
+@triton.jit
+def argmin_kernel(
+    inp,
+    out_index,
+    M,
+    N,
+    K,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Argmin along an arbitrary dimension.
+
+    Input is viewed as [M, N, K] where N is the reduction dim.
+    Output shape is [M, K].
+    """
+    pid_m = tl.program_id(0).to(tl.int64)
+    pid_k = tl.program_id(1).to(tl.int64)
+    m_offset = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    dtype = inp.type.element_ty
+    acc_type = tl.float32 if dtype is tl.bfloat16 else dtype
+    max_value = float('inf')
+    min_values = tl.full([BLOCK_M], dtype=acc_type, value=max_value)
+    argmin_values = tl.full([BLOCK_M], dtype=tl.int64, value=0)
+    for start_n in range(0, N, BLOCK_N):
+        n_offset = start_n + tl.arange(0, BLOCK_N)
+        offset = m_offset[:, None] * N * K + n_offset[None, :] * K + pid_k
+        mask = (m_offset[:, None] < M) & (n_offset[None, :] < N)
+        inp_ptrs = inp + offset
+        inp_vals = tl.load(inp_ptrs, mask=mask, other=max_value)
+        local_min, local_argmin = tl.min(
+            inp_vals, 1, return_indices=True, return_indices_tie_break_left=True
+        )
+        update = local_min < min_values
+        min_values = tl.where(update, local_min, min_values)
+        argmin_values = tl.where(update, start_n + local_argmin, argmin_values)
+
+    offset_index = m_offset * K + pid_k
+    out_index_ptrs = out_index + offset_index
+    mask1 = m_offset < M
+    tl.store(out_index_ptrs, argmin_values, mask=mask1)

nbxtensor/ops/atan2.py

@@ -0,0 +1,71 @@
+"""atan2(y, x) — pure @triton.jit kernel (element-wise, fp32).
+
+`tl.extra.CUDA.libdevice.atan2` stood here. That namespace is NVIDIA's device
+library and does not lower on any other backend: Kokoro-82M died at
+`aten.angle::0` (1, 11, 15361) on triton-ext with "PassManager::run failed",
+`angle` being atan2(imag, real). Triton 3.8.0 exposes no portable `atan` either
+(`tl.math` has abs, ceil, cos, div_rn, erf, exp, exp2, fdiv, floor, fma, log,
+log2, rsqrt, sin, sqrt, sqrt_rn, umulhi), so the function is built here.
+"""
+
+import triton
+import triton.language as tl
+
+# Triton only lets a @jit kernel read a global that is a constexpr.
+_PI = tl.constexpr(3.14159265358979323846)
+_PI_2 = tl.constexpr(1.57079632679489661923)
+
+
+@triton.jit
+def _atan_unit(z):
+    """atan(z) for |z| <= 1, odd minimax polynomial in z^2.
+
+    Degree-17 odd fit (the classic SunPro/Cephes coefficients). A degree-9 fit
+    was tried first and measured 1.167e-05 max absolute error against numpy over
+    all four quadrants — too coarse for a phase feeding audio synthesis, which
+    is what `aten::angle` is here. This one is checked in the probe beside this
+    file, not asserted in the kernel.
+    """
+    z2 = z * z
+    p = 0.0028662257
+    p = -0.0161657367 + z2 * p
+    p = 0.0429096138 + z2 * p
+    p = -0.0752896400 + z2 * p
+    p = 0.1065626393 + z2 * p
+    p = -0.1420889944 + z2 * p
+    p = 0.1999355085 + z2 * p
+    p = -0.3333314528 + z2 * p
+    return z + z * z2 * p
+
+
+@triton.jit
+def _atan(z):
+    """atan(z) over the whole line: fold |z| > 1 onto the unit interval."""
+    a = tl.abs(z)
+    big = a > 1.0
+    # for |z| > 1, atan(z) = pi/2 - atan(1/z), with the sign carried outside
+    u = tl.where(big, 1.0 / a, a)
+    r = _atan_unit(u)
+    r = tl.where(big, _PI_2 - r, r)
+    return tl.where(z < 0, -r, r)
+
+
+@triton.jit
+def atan2_kernel(y_ptr, x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0).to(tl.int64)
+    offs = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offs < n_elements
+    y = tl.load(y_ptr + offs, mask=mask).to(tl.float32)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+
+    # Quadrant reconstruction, matching numpy/libm atan2:
+    #   x > 0            atan(y/x)
+    #   x < 0, y >= 0    atan(y/x) + pi
+    #   x < 0, y <  0    atan(y/x) - pi
+    #   x = 0, y > 0     +pi/2        x = 0, y < 0   -pi/2      x = y = 0   0
+    safe_x = tl.where(x == 0, 1.0, x)
+    base = _atan(y / safe_x)
+    shifted = base + tl.where(y >= 0, _PI, -_PI)
+    on_axis = tl.where(y > 0, _PI_2, tl.where(y < 0, -_PI_2, 0.0))
+    out = tl.where(x > 0, base, tl.where(x < 0, shifted, on_axis))
+    tl.store(out_ptr + offs, out, mask=mask)