fix(language): coerce coalesced_width to IntImm in T.Parallel - #3374

Problem & User Impact

When constructing nested parallel loops via T.Parallel(*extents, coalesced_width=<int>) or using loop annotations T.Parallel(*extents, annotations={"coalesced_width": <int>}), compilation fails during LayoutInference with a fatal TVM error:

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

Even though the parameter coalesced_width is explicitly typed as int | None and documented as Optional[int], passing a Python int caused an immediate compiler crash. This prevented users from controlling memory coalescing widths on parallel loops unless they manually wrapped the argument in T.int32(...).

Root Cause

In tilelang/language/loop.py, Parallel attached the raw argument directly into merged_annotations["coalesced_width"] = coalesced_width. When converted to TVM's Map<String, Any> through TVM FFI, a Python int is not boxed as a tvm::tir::IntImmNode. Later in lowering, ParallelOpNode::ComputePlanCandidate (src/op/parallel.cc) accesses root_->annotations.Get(attr::kCoalescedWidth) and inspects coalesced_width->as<IntImmNode>(). Because the annotation was not an IntImmNode, it triggered LOG(FATAL) << "coalesced_width should be an IntImmNode.";. In contrast, sibling operations like T.copy route annotations through call_intrin, which normalizes integer entries to IntImm.

Solution

In tilelang/language/loop.py:

• In Parallel(...), if "coalesced_width" is present in merged_annotations and is not already a tirx.PrimExpr, wrap it using IntImm("int32", int(cw)).
• This ensures that both direct keyword arguments (coalesced_width=<int>) and dictionary annotations (annotations={"coalesced_width": <int>}) as well as numpy integers are coerced into an IntImmNode before passing to C++ lowering, while preserving existing IntImm or PrimExpr values.

Tests Performed

1. Verified reproduction on main before fix: confirmed that lowering a kernel with T.Parallel(..., coalesced_width=4) crashes with tvm.error.InternalError: coalesced_width should be an IntImmNode..
2. Verified fix: confirmed that tl.lower(..., enable_device_compile=False) compiles cleanly and lowers the loop according to the geometry-supported vector width.
3. Added automated regression tests to testing/python/transform/test_tilelang_transform_coalesced_width.py:
   • test_parallel_coalesced_width_integer: tests T.Parallel with integer coalesced_width values 2 and 4.
   • test_parallel_coalesced_width_annotation_dict: tests T.Parallel with integer coalesced_width in the annotations dictionary.
4. Ran the entire test module pytest testing/python/transform/test_tilelang_transform_coalesced_width.py -v: all tests passed (3 passed, 6 skipped hardware-dependent).
5. Ran pytest testing/python/transform/test_tilelang_transform_verify_parallel_loop.py -v: all 7 passed.
6. Ran ruff check on modified files: all checks passed.

References

Resolves #3013

Summary

• T.Parallel now converts non-PrimExpr coalesced_width values to IntImm("int32", int(value)) before passing annotations to the FFI builder. Existing PrimExpr values remain unchanged.
• Added CUDA lowering regression tests for integer widths passed directly and through the annotations dictionary.

Validation

• Coalesced-width tests: 3 passed and 6 hardware-dependent tests skipped.
• Parallel-loop verification tests: 7 passed.
• ruff check on the modified files passed.