[BUG]: tileiras 13.4 ignores the occupancy hint on sm_120 when ct.gather is inside a runtime-trip-count loop #117

"""tileiras ignores the occupancy hint for a W4A16 GEMM whose activation gather sits inside a runtime loop.

One kernel: FP4 (E2M1) weights with E4M3 block scales per 16, BF16 activations, tile 64x128x128. The activation
is read with a tile load or a row gather (GATHER), and the M-block loop runs a runtime trip count or once (LOOP).
For each case and occupancy hint the script prints the compiled CTA size, registers and shared memory, and the
kernel time.

usage: CUDA_TILE_TEMP_DIR=$(mktemp -d) python occupancy_repro.py   (cuobjdump on PATH)
"""
import glob
import os
import re
import statistics
import subprocess

import cuda.tile as ct
import torch

ConstInt = ct.Constant[int]
ConstBool = ct.Constant[bool]
TM, TN, TK, SB = 64, 128, 128, 16


@ct.kernel
def w4a16_gemm(X, ROWS, W, S, OUT, num_iterations, grid_m, K: ConstInt, GATHER: ConstBool, LOOP: ConstBool):
    nb = ct.bid(1)
    k_off = ct.arange(TK, dtype=ct.int32)
    for it in range(num_iterations if LOOP else 1):
        mb = ct.bid(0) + it * grid_m
        rows = ct.gather(ROWS, (mb * TM + ct.arange(TM, dtype=ct.int32),), padding_value=0)
        acc = ct.zeros((TM, TN), dtype=ct.float32)
        for kt in range(K // TK):
            if GATHER:
                idx = ct.reshape(rows, (TM, 1)) * K + ct.reshape(kt * TK + k_off, (1, TK))
                a = ct.gather(X, (idx,), padding_value=0, check_bounds=True, latency=3)
            else:
                a = ct.load(X, index=(mb, kt), shape=(TM, TK), latency=3)
            wb = ct.load(W, index=(nb, kt), shape=(TN, TK // 2), latency=3)
            w = ct.reshape(ct.unpack_from_bytes(ct.reshape(wb, (-1,)), ct.float4_e2m1fn), (TN, TK))
            s = ct.load(S, index=(nb, kt), shape=(TN, TK // SB), latency=3)
            wd = ct.reshape(ct.astype(w, ct.bfloat16), (TN, TK // SB, SB)) * ct.expand_dims(
                ct.astype(s, ct.bfloat16), axis=2)
            acc = ct.mma(a, ct.permute(ct.reshape(wd, (TN, TK)), (1, 0)), acc)
        ct.store(OUT, index=(mb, nb), tile=ct.astype(acc, ct.bfloat16))


def cubin_resources(path):
    elf = subprocess.run(["cuobjdump", "-elf", path], capture_output=True, text=True).stdout
    ru = subprocess.run(["cuobjdump", "-res-usage", path], capture_output=True, text=True).stdout
    nt = re.search(r"EIATTR_REQNTID\s.*?Value:\s*(0x[0-9a-f]+)", elf, re.S)
    r = re.search(r"REG:(\d+) STACK:\d+ SHARED:(\d+)", ru)
    threads = int(nt.group(1), 16) if nt else "?"
    return f"threads {threads:>3}  regs {r.group(1):>3}  smem {int(r.group(2)):>6} B" if r else "?"


M, N, K = 16384, 6144, 6144
dev = torch.device("cuda")
g = torch.Generator(device=dev).manual_seed(0)
X = torch.randn(M, K, device=dev, generator=g).to(torch.bfloat16)
W = torch.randint(0, 256, (N, K // 2), device=dev, generator=g, dtype=torch.uint8)
S = (torch.rand(N, K // SB, device=dev, generator=g) * 0.01 + 0.001).to(torch.float8_e4m3fn)
E2M1 = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], device=dev)
Wf = torch.stack((E2M1[(W & 15).long()], E2M1[(W >> 4).long()]), -1).reshape(N, K)
Wf = Wf * S.float().repeat_interleave(SB, dim=1)
ROWS = torch.randperm(M, device=dev, generator=g).to(torch.int32)
ref_gather = X[ROWS.long()].float() @ Wf.T
ref_load = X.float() @ Wf.T
tmp = os.environ["CUDA_TILE_TEMP_DIR"]
num_mb = M // TM
from cuda.tile import _compile  # noqa: E402

print(f"tileiras {_compile._find_compiler_bin().path}")
print(f"GPU {torch.cuda.get_device_name(dev)} sm_{''.join(map(str, torch.cuda.get_device_capability(dev)))}  "
      f"cuda-tile {ct.__version__ if hasattr(ct, '__version__') else '?'}  torch {torch.__version__}")
stream = torch.cuda.current_stream()
for gather in (False, True):
    for loop in (False, True):
        for occ in (1, 2, 4):
            kern = w4a16_gemm.replace_hints(occupancy=occ)
            out = torch.empty(M, N, device=dev, dtype=torch.bfloat16)
            before = set(glob.glob(os.path.join(tmp, "w4a16_gemm*.cubin")))

            def run():
                # one CTA per M block in both cases; LOOP only changes how the kernel is written
                ct.launch(stream, (num_mb, N // TN), kern,
                          (X.reshape(-1) if gather else X, ROWS, W, S, out, 1, num_mb, K, gather, loop))

            run()
            torch.cuda.synchronize()
            ref = ref_gather if gather else ref_load
            err = float((out.float() - ref).norm() / ref.norm())
            ts = []
            for _ in range(10):
                a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
                a.record()
                run()
                b.record()
                b.synchronize()
                ts.append(a.elapsed_time(b))
            new = sorted(set(glob.glob(os.path.join(tmp, "w4a16_gemm*.cubin"))) - before)
            res = cubin_resources(new[-1]) if new else "(no new cubin)"
            print(f"gather={gather!s:5} runtime_loop={loop!s:5} occupancy={occ}:  {res}  "
                  f"{statistics.median(ts):7.2f} ms  rel err {err:.1e}", flush=True)