[BUG]: tileiras SIGSEGV when occupancy=2 is requested for a 32-wide fused tile kernel #96

import torch
import cuda.tile as ct


ConstInt = ct.Constant[int]
ZERO = ct.PaddingMode.ZERO


def factor(a, block: ConstInt):
    cols = ct.arange(block, dtype=ct.int32)[None, :]
    for p in range(block):
        pivot = ct.extract(a, (p, p), shape=(1, 1))
        column = ct.extract(a, (0, p), shape=(block, 1)) / pivot
        outer = column * column.transpose(0, 1)
        a = ct.where(cols > p, a - outer, a)
    return a


def solve(panel, diagonal, block: ConstInt):
    cols = ct.arange(block, dtype=ct.int32)[None, :]
    for p in range(block):
        pivot = ct.extract(diagonal, (p, p), shape=(1, 1))
        solved = ct.extract(panel, (0, p), shape=(block, 1)) / pivot
        column = ct.extract(diagonal, (0, p), shape=(block, 1))
        panel = ct.where(
            cols > p,
            panel - solved * column.transpose(0, 1),
            panel,
        )
    return panel


@ct.kernel(opt_level=2, occupancy=2)
def kernel(a, step, block: ConstInt):
    rows = block
    width = 2 * block
    next_step = step + 1

    work = ct.load(a, (0, 0, 0), shape=(1, rows, width), padding_mode=ZERO)
    work = work.reshape((rows, width))
    cross = ct.load(a, (0, 1, 0), shape=(1, block, block), padding_mode=ZERO)
    cross = cross.reshape((block, block))
    diagonal = ct.load(a, (0, 1, 1), shape=(1, block, block), padding_mode=ZERO)
    diagonal = diagonal.reshape((block, block))

    for prior in range(step):
        left = ct.load(a, (0, 0, prior), shape=(1, rows, width), padding_mode=ZERO)
        right = ct.load(a, (0, 0, prior), shape=(1, width, width), padding_mode=ZERO)
        left = left.reshape((rows, width))
        right = right.reshape((width, width))
        work = ct.mma(
            left.astype(ct.tfloat32),
            (-right.transpose(0, 1)).astype(ct.tfloat32),
            work,
        )
        diagonal = ct.mma(
            left.astype(ct.tfloat32),
            (-left.transpose(0, 1)).astype(ct.tfloat32),
            diagonal,
        )

    first_panel = ct.extract(work, (0, 0), shape=(rows, block))
    second_panel = ct.mma(
        first_panel.astype(ct.tfloat32),
        (-cross.transpose(0, 1)).astype(ct.tfloat32),
        first_panel,
    )
    second_panel = solve(second_panel, diagonal, block)

    ct.store(a, (0, 0, 0), first_panel.reshape((1, rows, block)))
    if ct.bid(1) == 0:
        ct.store(a, (0, 1, 1), diagonal.reshape((1, block, block)))

    if ct.bid(1) == next_step:
        next_diagonal = ct.load(
            a,
            (0, next_step, next_step),
            shape=(1, block, block),
            padding_mode=ZERO,
        ).reshape((block, block))
        next_diagonal = ct.mma(
            second_panel.astype(ct.tfloat32),
            (-second_panel.transpose(0, 1)).astype(ct.tfloat32),
            next_diagonal,
        )
        next_diagonal = factor(next_diagonal, block)
        ct.store(
            a,
            (0, next_step, next_step),
            next_diagonal.reshape((1, block, block)),
        )


a = torch.empty((1, 4, 4), device="cuda", dtype=torch.float32)
ct.launch(torch.cuda.current_stream(), (1, 1), kernel, (a, 0, 32))