submission 920766

import torch
import triton
import triton.language as tl

from task import input_t, output_t


@triton.jit
def _chol_small(
    a,
    l,
    matrix_stride: tl.constexpr,
    N: tl.constexpr,
):
    batch_idx = tl.program_id(0)

    idx = tl.arange(0, N)
    row = idx[:, None]
    col = idx[None, :]
    offsets = batch_idx * matrix_stride + row * N + col

    # Keeping the complete lower tile in registers avoids global-memory
    # round-trips and CTA synchronization for the heavily batched cases.
    values = tl.load(a + offsets, mask=row >= col, other=0.0)

    for k in tl.static_range(0, N):
        pivot_row = tl.sum(tl.where(row == k, values, 0.0), axis=0)
        diagonal = tl.sum(tl.where(idx == k, pivot_row, 0.0), axis=0)
        diagonal -= tl.sum(
            tl.where(idx < k, pivot_row * pivot_row, 0.0), axis=0
        )
        diagonal = tl.sqrt(diagonal)

        pivot_col = tl.sum(tl.where(col == k, values, 0.0), axis=1)
        correction = tl.where(col < k, values * pivot_row[None, :], 0.0)
        pivot_col = (pivot_col - tl.sum(correction, axis=1)) / diagonal

        values = tl.where((row == k) & (col == k), diagonal, values)
        values = tl.where((row > k) & (col == k), pivot_col[:, None], values)

    tl.store(l + offsets, values)


def custom_kernel(data: input_t) -> output_t:
    batch, n, _ = data.shape

    if n in (32, 64, 128):
        out = torch.empty_like(data)
        _chol_small[(batch,)](
            data,
            out,
            n * n,
            N=n,
            num_warps=1 if n == 32 else (2 if n == 64 else 8),
            num_stages=1,
        )
        return out

    return torch.linalg.cholesky_ex(data, check_errors=False).L