nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/neg.py

@@ -0,0 +1,17 @@
+"""Negation — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def neg_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    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)
+    mask = offset < n_elements
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, -x, mask=mask)

nbxtensor/ops/nllloss.py

@@ -0,0 +1,124 @@
+"""NLL Loss forward kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+
+NLL Loss: l_n = -w_{y_n} * x_{n, y_n}
+
+Reduction modes:
+  0 = none: per-sample loss
+  1 = mean: weighted mean
+  2 = sum:  weighted sum
+
+The kernel handles ignore_index to skip certain targets.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def nll_loss_forward_kernel(
+    inp_ptr,
+    tgt_ptr,
+    wgt_ptr,
+    out_ptr,
+    ignore_index,
+    N,
+    C,
+    reduction: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """NLL Loss forward for 1D/2D input [N, C].
+
+    Args:
+        inp_ptr: Log-probability input [N, C].
+        tgt_ptr: Target class indices [N].
+        wgt_ptr: Per-class weights [C] or None.
+        out_ptr: Output. Shape depends on reduction mode.
+        ignore_index: Target index to ignore.
+        N: Batch size.
+        C: Number of classes.
+        reduction: 0=none, 1=mean, 2=sum.
+        BLOCK_N: Block size for batch dimension.
+    """
+    pid_n = tl.program_id(0).to(tl.int64)
+    offsets_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offsets_n < N
+
+    tgt = tl.load(tgt_ptr + offsets_n, mask=mask_n, other=0)
+    ignore_mask = (tgt != ignore_index) & mask_n
+
+    # Load per-class weight for target, or 1.0 if no weights
+    if wgt_ptr is None:
+        wgt_tgt = ignore_mask.to(tl.float32)
+    else:
+        wgt_tgt = tl.load(wgt_ptr + tgt, mask=ignore_mask, other=0).to(tl.float32)
+
+    # Gather input at target index: inp[n, tgt[n]]
+    inp_tgt_ptrs = inp_ptr + offsets_n * C + tgt
+    inp_tgt = tl.load(inp_tgt_ptrs, mask=ignore_mask, other=0).to(tl.float32)
+    out = inp_tgt * wgt_tgt * -1
+
+    # none
+    if reduction == 0:
+        tl.store(out_ptr + offsets_n, out, mask=mask_n)
+    # mean
+    elif reduction == 1:
+        total_out = tl.sum(out)
+        total_wgt = tl.sum(wgt_tgt)
+        tl.atomic_add(out_ptr, total_out, sem="relaxed")
+        tl.atomic_add(out_ptr + 1, total_wgt, sem="relaxed")
+        tl.atomic_add(out_ptr + 2, 1, sem="release")
+        counter = tl.load(out_ptr + 2)
+        if counter == tl.num_programs(0):
+            total_out = tl.load(out_ptr)
+            total_wgt = tl.load(out_ptr + 1)
+            tl.store(out_ptr + 3, total_out / total_wgt)
+    # sum
+    else:
+        total_out = tl.sum(out)
+        tl.atomic_add(out_ptr, total_out, sem="relaxed")
+
+
+@triton.jit
+def nll_loss_backward_kernel(
+    out_grad_ptr,
+    tgt_ptr,
+    wgt_ptr,
+    inp_grad_ptr,
+    ignore_index,
+    total_weight,
+    N,
+    C,
+    reduction: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """NLL Loss backward for 1D/2D input [N, C].
+
+    Computes gradient w.r.t. input: grad_input[n, tgt[n]] = -w * grad_output / total_weight.
+    """
+    pid_n = tl.program_id(0).to(tl.int64)
+    offsets_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask_n = offsets_n < N
+
+    tgt = tl.load(tgt_ptr + offsets_n, mask=mask_n, other=0)
+    ignore_mask = (tgt != ignore_index) & mask_n
+
+    if wgt_ptr is None:
+        wgt_tgt = ignore_mask.to(tl.float32)
+    else:
+        wgt_tgt = tl.load(wgt_ptr + tgt, mask=ignore_mask, other=0).to(tl.float32)
+
+    if reduction == 0:
+        out_grad = tl.load(out_grad_ptr + offsets_n, mask=mask_n, other=0).to(tl.float32)
+    else:
+        out_grad = tl.load(out_grad_ptr).to(tl.float32)
+
+    if reduction == 1:
+        total_w = tl.load(total_weight).to(tl.float32)
+    else:
+        total_w = 1
+
+    inp_grad = tl.where(ignore_mask, -1 * out_grad * wgt_tgt / total_w, 0)
+    inp_grad_ptrs = inp_grad_ptr + offsets_n * C + tgt
+    tl.store(inp_grad_ptrs, inp_grad, mask=mask_n)

nbxtensor/ops/nonzero.py

@@ -0,0 +1,33 @@
+"""1-D boolean-mask compaction (nonzero) — pure @triton.jit kernel.
+
+Deterministic by construction: a SINGLE program scans the mask in
+BLOCK-sized chunks with an in-block tl.cumsum and a running cursor, so
+the compacted indices are emitted in ascending order without atomics.
+Mask sizes on this path are decode-control scale (visual/attention
+masks, thousands of elements) — the serial-block scan costs
+microseconds there and keeps the op bit-reproducible across runs.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def nonzero_1d_kernel(
+    mask_ptr, out_ptr, count_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    cursor = tl.zeros((1,), dtype=tl.int32)
+    for start in range(0, n_elements, BLOCK_SIZE):
+        offs = start + tl.arange(0, BLOCK_SIZE)
+        in_range = offs < n_elements
+        m = tl.load(mask_ptr + offs, mask=in_range, other=0)
+        pred = (m != 0) & in_range
+        p32 = pred.to(tl.int32)
+        excl = tl.cumsum(p32, axis=0) - p32
+        base = tl.sum(cursor, axis=0)
+        tl.store(out_ptr + (base + excl).to(tl.int64), offs.to(tl.int64),
+                 mask=pred)
+        cursor += tl.sum(p32, axis=0)
+    tl.store(count_ptr, tl.sum(cursor, axis=0))

nbxtensor/ops/nonzero_op.py

@@ -0,0 +1,62 @@
+"""Nonzero — pure @triton.jit kernel.
+
+Extracts indices of non-zero elements from a flattened boolean tensor.
+Requires a pre-computed prefix sum (cumsum of bool input) to determine
+the output position of each non-zero element.
+
+The Python wrapper must:
+1. Flatten and convert to bool: inp_bool = (inp != 0)
+2. Compute prefix sum: prefix_sum = inp_bool.cumsum(0)
+3. Allocate output: out = empty(numel, ndim, dtype=int64)
+4. Launch this kernel
+5. Slice to actual count: out = out[:prefix_sum[-1]]
+
+Extracted from FlagGems (Apache 2.0).
+
+Grid: (ceil(n_elements / BLOCK_SIZE),)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def nonzero_kernel(
+    inp,
+    prefix_sum,
+    out,
+    n_elements,
+    shape,
+    ndim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Write indices of non-zero elements into output tensor.
+
+    For each non-zero element at flat index i, decomposes i into
+    multi-dimensional indices using the original shape, then writes
+    each dimension's index into out[prefix_sum[i]-1, dim].
+
+    Args:
+        inp: flattened boolean input (1 = nonzero)
+        prefix_sum: cumulative sum of inp (1-based output position)
+        out: output tensor of shape (n_elements, ndim), int64
+        shape: original tensor shape stored as int32 on device
+        ndim: number of dimensions (compile-time constant)
+    """
+    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
+
+    inp_vals = tl.load(inp + offset, mask=mask).to(tl.int1)
+    out_offset = tl.load(prefix_sum + offset, mask=mask) - 1
+
+    nonzero_mask = mask & inp_vals
+
+    # Decompose flat index into multi-dimensional indices
+    idx_flat = offset
+    for dim in range(ndim - 1, -1, -1):
+        dim_size = tl.load(shape + dim)
+        remainder = idx_flat % dim_size
+        idx_flat = idx_flat // dim_size
+        tl.store(out + out_offset * ndim + dim, remainder, mask=nonzero_mask)