nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/vector_norm.py

@@ -0,0 +1,216 @@
+"""Vector norm — pure @triton.jit kernels.
+
+Computes L2, L1, L0 (count non-zero), Linf (max abs), L-inf (min abs),
+and general Lp norms.
+
+Two kernel variants:
+  - Row-wise: input is [M, N], computes norm over N for each of M rows
+  - Flat (two-pass): input is flat [M], first pass reduces blocks, second
+    pass reduces block results to scalar
+
+Extracted from FlagGems reference (vector_norm.py).
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Row-wise kernels: input [M, N] -> output [M]
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def l2_norm_kernel(
+    X, Out,
+    M, N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """L2 norm per row: out[m] = sqrt(sum(x[m,:]^2))."""
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Out = Out + pid
+    row_mask = pid < M
+
+    _sum = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    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(X + cols, mask=mask, other=0.0).to(tl.float32)
+        _sum += a * a
+
+    s = tl.sum(_sum, axis=1)
+    out = tl.sqrt(s)[:, None]
+    tl.store(Out, out, mask=row_mask)
+
+
+@triton.jit
+def l1_norm_kernel(
+    X, Out,
+    M, N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """L1 norm per row: out[m] = sum(|x[m,:]|)."""
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Out = Out + pid
+    row_mask = pid < M
+
+    _sum = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    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(X + cols, mask=mask, other=0.0).to(tl.float32)
+        _sum += tl.abs(a)
+
+    s = tl.sum(_sum, axis=1)
+    out = s[:, None]
+    tl.store(Out, out, mask=row_mask)
+
+
+@triton.jit
+def linf_norm_kernel(
+    X, Out,
+    M, N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """L-infinity norm per row: out[m] = max(|x[m,:]|)."""
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Out = Out + pid
+    row_mask = pid < M
+
+    _max = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    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(X + cols, mask=mask, other=0.0).to(tl.float32)
+        _max = tl.maximum(tl.abs(a), _max)
+
+    m = tl.max(_max, axis=1)
+    out = m[:, None]
+    tl.store(Out, out, mask=row_mask)
+
+
+@triton.jit
+def l0_norm_kernel(
+    X, Out,
+    M, N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """L0 norm per row: out[m] = count(x[m,:] != 0)."""
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Out = Out + pid
+    row_mask = pid < M
+
+    _sum = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    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(X + cols, mask=mask, other=0.0).to(tl.float32)
+        _sum += tl.where(a != 0, 1.0, 0.0)
+
+    s = tl.sum(_sum, axis=1)
+    out = s[:, None]
+    tl.store(Out, out, mask=row_mask)
+
+
+@triton.jit(do_not_specialize=['ord'])
+def lp_norm_kernel(
+    X, Out,
+    M, N,
+    ord,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """General Lp norm per row: out[m] = sum(|x[m,:]|^p)^(1/p)."""
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Out = Out + pid
+    row_mask = pid < M
+
+    _sum = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    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(X + cols, mask=mask, other=0.0).to(tl.float32)
+        # |a|**ord for ord > 0 and a >= 0: exp/log is exact enough and portable.
+        # `libdevice.pow` is NVIDIA-only and fails to lower elsewhere.
+        _sum += tl.exp(ord * tl.log(tl.abs(a) + 1e-38))
+
+    s = tl.sum(_sum, axis=1)
+    out = tl.exp((1.0 / ord) * tl.log(s + 1e-38))[:, None]
+    tl.store(Out, out, mask=row_mask)
+
+
+# ---------------------------------------------------------------------------
+# Flat two-pass kernels: input [M] -> scalar
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def l2_norm_pass1_kernel(
+    X, Mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """First pass of flat L2 norm: each block computes partial sum of squares."""
+    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)
+    mask = offset < M
+
+    x = tl.load(X + offset, mask=mask, other=0.0).to(tl.float32)
+    mid = tl.sum(x * x)
+    tl.store(Mid + pid, mid)
+
+
+@triton.jit
+def l2_norm_pass2_kernel(
+    Mid, Out,
+    MID_SIZE,
+    BLOCK_MID: tl.constexpr,
+):
+    """Second pass: reduce partial sums and take sqrt."""
+    offset = tl.arange(0, BLOCK_MID)
+    mask = offset < MID_SIZE
+    mid = tl.load(Mid + offset, mask=mask, other=0.0).to(tl.float32)
+    out = tl.sqrt(tl.sum(mid))
+    tl.store(Out, out)
+
+
+@triton.jit
+def linf_norm_pass1_kernel(
+    X, Mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """First pass of flat Linf norm: each block computes partial max abs."""
+    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)
+    mask = offset < M
+
+    x = tl.load(X + offset, mask=mask, other=0.0).to(tl.float32)
+    mid = tl.max(tl.abs(x))
+    tl.store(Mid + pid, mid)
+
+
+@triton.jit
+def linf_norm_pass2_kernel(
+    Mid, Out,
+    MID_SIZE,
+    BLOCK_MID: tl.constexpr,
+):
+    """Second pass: reduce partial maxes."""
+    offset = tl.arange(0, BLOCK_MID)
+    mask = offset < MID_SIZE
+    mid = tl.load(Mid + offset, mask=mask, other=0.0).to(tl.float32)
+    out = tl.max(mid)
+    tl.store(Out, out)

nbxtensor/ops/weight_norm.py

@@ -0,0 +1,104 @@
+"""Weight normalization (_weight_norm_interface) — pure @triton.jit kernels.
+
+Computes w = g * v / ||v|| where g is a per-channel gain scalar.
+Two variants: norm along dim=0 ("first") and dim=last ("last").
+
+Ported from FlagGems weightnorm.py. Stripped FlagGems decorators.
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---- Forward: dim = 0 (norm along first dimension, one gain per row) ----
+
+@triton.jit
+def weight_norm_kernel_first(
+    output_ptr,
+    norm_ptr,
+    v_ptr,
+    g_ptr,
+    M,
+    N,
+    eps,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Weight norm forward, normalizing along dim=0 (rows).
+
+    v: [M, N], g: [M], norm: [M], output: [M, N]
+    Each program handles BLOCK_M rows.
+    """
+    ty = tl.arange(0, BLOCK_M)[:, None]
+    by = tl.program_id(0).to(tl.int64) * BLOCK_M
+    row_offset = by + ty
+    row_mask = row_offset < M
+
+    tx = tl.arange(0, BLOCK_N)[None, :]
+
+    # Pass 1: compute ||v|| per row
+    v_sq = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    for base in range(0, N, BLOCK_N):
+        col_offset = base + tx
+        mask = (col_offset < N) & row_mask
+        v_val = tl.load(v_ptr + row_offset * N + col_offset, mask=mask).to(tl.float32)
+        v_sq += v_val * v_val
+
+    norm_val = tl.sqrt(tl.sum(v_sq, axis=1) + eps)
+    tl.store(norm_ptr + row_offset, norm_val[:, None], mask=row_mask)
+    g_val = tl.load(g_ptr + row_offset, mask=row_mask).to(tl.float32)
+
+    # Pass 2: output = g * v / ||v||
+    for base in range(0, N, BLOCK_N):
+        col_offset = base + tx
+        mask = (col_offset < N) & row_mask
+        v_val = tl.load(v_ptr + row_offset * N + col_offset, mask=mask).to(tl.float32)
+        out = (v_val / norm_val[:, None]) * g_val
+        tl.store(output_ptr + row_offset * N + col_offset, out, mask=mask)
+
+
+# ---- Forward: dim = last (norm along last dimension, one gain per column) ----
+
+@triton.jit
+def weight_norm_kernel_last(
+    output_ptr,
+    norm_ptr,
+    v_ptr,
+    g_ptr,
+    M,
+    N,
+    eps,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Weight norm forward, normalizing along last dim (columns).
+
+    v: [M, N], g: [N], norm: [N], output: [M, N]
+    Each program handles BLOCK_N columns.
+    """
+    tx = tl.arange(0, BLOCK_N)[:, None]
+    bx = tl.program_id(0).to(tl.int64) * BLOCK_N
+    col_offset = bx + tx
+    col_mask = col_offset < N
+
+    ty = tl.arange(0, BLOCK_M)[None, :]
+
+    # Pass 1: compute ||v|| per column
+    v_sq = tl.zeros([BLOCK_N, BLOCK_M], dtype=tl.float32)
+    for base in range(0, M, BLOCK_M):
+        row_offset = base + ty
+        mask = (row_offset < M) & col_mask
+        v_val = tl.load(v_ptr + row_offset * N + col_offset, mask=mask).to(tl.float32)
+        v_sq += v_val * v_val
+
+    norm_val = tl.sqrt(tl.sum(v_sq, axis=1) + eps)
+    tl.store(norm_ptr + col_offset, norm_val[:, None], mask=col_mask)
+    g_val = tl.load(g_ptr + col_offset, mask=col_mask).to(tl.float32)
+
+    # Pass 2: output = g * v / ||v||
+    for base in range(0, M, BLOCK_M):
+        row_offset = base + ty
+        mask = (row_offset < M) & col_mask
+        v_val = tl.load(v_ptr + row_offset * N + col_offset, mask=mask).to(tl.float32)
+        out = (v_val / norm_val[:, None]) * g_val
+        tl.store(output_ptr + row_offset * N + col_offset, out, mask=mask)

nbxtensor/ops/where.py

@@ -0,0 +1,39 @@
+"""Where — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def where_forward_kernel(
+    cond_ptr, x_ptr, y_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = where(cond, x, y)"""
+    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)
+    mask = offset < n_elements
+
+    cond = tl.load(cond_ptr + offset, mask=mask)
+    x = tl.load(x_ptr + offset, mask=mask)
+    y = tl.load(y_ptr + offset, mask=mask)
+    # `.to(tl.int1)` and not `cond` alone: a boolean crosses the NBX boundary in
+    # a uint8 container (that is how a bool tensor is stored), so the loaded
+    # value is an INTEGER. `tl.where` accepts that today with a deprecation
+    # warning and will raise on it in a future version — one of the two
+    # scheduled breakages a user's A40 report surfaced, four warnings per run.
+    #
+    # This is a forward-compatibility edit and it changes NO instruction: the
+    # cast `tl.where` already performed internally on a non-boolean condition is
+    # `not_equal(x, null_of(x.dtype))` (semantic.py, the is_bool() branch of
+    # cast), which is exactly what `.to(tl.int1)` emits. Proven, not asserted —
+    # the TTIR is identical to the pre-edit form
+    # (`tools/kernel_boolean_ir_equality.py`).
+    #
+    # Why NOT `cond != 0`, which was the first form written here: the Python
+    # literal `0` promotes to i32, so the comparison gains an `arith.extsi`
+    # i8 -> i32 on every element before it. Same boolean for every input —
+    # extsi maps 0 to 0 and is injective — but one extra elementwise op, on the
+    # kernel that 39 of the 56 local containers reach. "Costs nothing" was a
+    # claim nobody had checked; the instrument checked it.
+    tl.store(output_ptr + offset, tl.where(cond.to(tl.int1), x, y), mask=mask)