[BUG][Fuzzer][ice-on-valid-code] T.Parallel(coalesced_width=<int>) aborts compilation instead of accepting the documented int #3013

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker and Discussions that this hasn't already been reported.

What version of TileLang are you using?

0.1.13

System information

TileLang 0.1.13 / CUDA 12.8 / PyTorch 2.8 ; run on NVIDIA L40S (sm_89). The failure is in LayoutInference (host-side layout planning), before any device code is generated, so it is architecture-independent.

Problem description

Passing an int to T.Parallel's documented coalesced_width keyword aborts compilation:

tvm.error.InternalError: coalesced_width should be an IntImmNode.

coalesced_width is a public keyword of T.Parallel, typed int | None in the signature (loop.py#L15) and documented as Optional[int] (loop.py#L28), so a plain Python int is exactly what the signature promises. It never reaches codegen — the crash fires during LayoutInference. The same int works everywhere else it is accepted: passing it via T.copy(..., coalesced_width=<int>) compiles fine, and wrapping the int in tir.IntImm before handing it to T.Parallel also compiles fine (both verified below), so the value itself is legal.

Not a regression — see Provenance.

Reproducible example code

import tilelang
import tilelang.language as T
import torch
# 0.1.13: IntImm moved from tvm.tir to tvm.tirx; tilelang re-exports it as T.IntImm.

N, blk = 4096, 256
A = torch.arange(N, device="cuda", dtype=torch.float32)

# --- FAIL: documented int coalesced_width aborts compilation ---
@T.prim_func
def bug(A: T.Tensor((N,), "float32"), B: T.Tensor((N,), "float32")):
    with T.Kernel(N // blk, threads=64) as bx:
        for i in T.Parallel(blk, coalesced_width=4):   # int, per the signature
            B[bx * blk + i] = A[bx * blk + i]

tilelang.compile(bug, out_idx=[1], target="cuda")
# -> tvm.error.InternalError: coalesced_width should be an IntImmNode.

# --- CONTROL: identical kernel, coalesced_width wrapped in IntImm -> compiles & correct ---
@T.prim_func
def ok(A: T.Tensor((N,), "float32"), B: T.Tensor((N,), "float32")):
    with T.Kernel(N // blk, threads=64) as bx:
        for i in T.Parallel(blk, coalesced_width=T.IntImm("int32", 4)):
            B[bx * blk + i] = A[bx * blk + i]

k = tilelang.compile(ok, out_idx=[1], target="cuda")
print("control match:", torch.equal(k(A), A))   # -> True

Traceback

File ".../tilelang/cuda/pipeline.py", line 117, in CUDAPassPipelineBodyPrologue
    mod = tilelang.transform.LayoutInference()(mod)
  ...
  tvm::tl::ParallelOpNode::ComputePlanCandidate(...) const
tvm.error.InternalError: coalesced_width should be an IntImmNode.

Expected behavior

T.Parallel(..., coalesced_width=4) should compile and use the requested coalesced width, matching its own int | None signature and the sibling T.copy(..., coalesced_width=<int>) path, which accepts a bare int over the same value. At minimum a documented, correctly-typed argument should not abort compilation.

Additional context

Root cause. T.Parallel attaches the coalesced_width argument to the loop-annotation map as a raw Python int, but the C++ layout planner reads it back as an IntImmNode and hard-fails when it is not one. In loop.py#L79-L87 the value is placed in merged_annotations unconverted (merged_annotations["coalesced_width"] = coalesced_width) and handed straight to _ffi_api.Parallel(extents, merged_annotations), which attaches the dict to the For node's annotation Map; the value arrives as a runtime int, not a tir.IntImm. Then in parallel.cc#L798-L809, ParallelOpNode::ComputePlanCandidate does coalesced_width->as<IntImmNode>(), which returns null, and takes the LOG(FATAL) branch (parallel.cc:808).

The T.copy family stores the argument the same raw way (copy_op.py#L124-L125), and its lowering (copy.cc#L538-L540) does re-attach kCoalescedWidth onto a ForKind::kParallel loop that reaches the same ComputePlanCandidate / as<IntImmNode>() read. The reason copy does not crash is not that it skips that read — it is that copy passes its annotations through call_intrin(..., annotations=ann), which builds a TVM Map and coerces the Python int to a tir.IntImm (a proper IntImmNode) as part of node construction, whereas _ffi_api.Parallel attaches the raw Python dict directly and no such coercion happens. So the divergence is a coercion gap at the FFI boundary of T.Parallel, not a difference in what the C++ planner reads. (Verified this session: T.copy(..., coalesced_width=4) compiles and returns the correct result on CUDA; see Reach.)

Suggested fix. Wrapping the value at the Python boundary — e.g. merged_annotations["coalesced_width"] = IntImm("int32", coalesced_width) in T.Parallel — makes the annotation an IntImmNode and lets the existing parallel.cc branch run; the CONTROL above (which does exactly this wrap at the call site) compiles and returns the correct result, so this is the minimal change. Alternatively ComputePlanCandidate could coerce the annotation via Downcast/arith before the as<IntImmNode>() read.

Provenance. The coalesced_width keyword on T.Parallel was added by #1887 (merged 2026-03-06), which introduced the raw-int passthrough at loop.py:80; that is when the crash first became reachable. The as<IntImmNode>() FATAL in parallel.cc predates it (present since the codebase migration #10, 2025-01-11), so the two halves never agreed on the type — the argument crashes on every release that exposes the keyword, not a regression.

Dedup. Searched the open and closed tracker (coalesced_width, IntImmNode, should be an IntImmNode) and found no existing report of this defect.

Reach. The trigger is documented and typed: coalesced_width is declared int | None and documented Optional[int], so an int is the type the API asks for. coalesced_width appears in shipped code (git grep coalesced_width examples/ testing/ at v0.1.13): every site routes it through T.copy / T.async_copy, e.g. examples/amd/example_amd_flash_attn_bwd.py:163,172-173 and examples/amd/example_amd_flash_attn_fwd.py:183,192-193 (coalesced_width=vec_size), and testing/python/amd/test_tilelang_gfx950_copy_async.py. No example or test passes it to T.Parallel (git grep 'Parallel(' examples/ testing/ | grep coalesced_width at v0.1.13 is empty), so CI never exercises the crashing path. These shipped sites are all AMD gfx950-targeted and were not runnable on this session's hardware (L40S, sm_89), but the load-bearing fact they establish — that a bare int through the copy family is a legal coalesced_width — was verified directly on CUDA this session: T.copy(dst, src, coalesced_width=4) compiled and returned torch.equal(...) == True. Given the keyword on T.Parallel, any value trips the crash — the failure is unconditional on that int path (the CONTROL only differs by wrapping the same value in IntImm), and it also fires when the int is supplied via the generic annotations={"coalesced_width": <int>} dict, so it is the FFI-attach path of T.Parallel, not the specific keyword.

Impact. The narrow ingredient is passing coalesced_width as a bare int through T.Parallel specifically (the sibling T.copy path and the IntImm-wrapped form both compile). When it fires it is a loud compile-time abort during LayoutInference — no device code is produced, so nothing is silently miscompiled and no output is corrupted; the failure is deterministic and caught immediately at build time. Fixing it removes a type disagreement between the Python signature (int | None) and the C++ as<IntImmNode>() read, letting a documented, correctly-typed argument build the kernel it should.