feat: allow T.const() variables used only in grid dims or computations

[NolanHo]

Problem

In eager (@tilelang.jit) mode, T.const() is the only way to declare a
compile-time symbolic variable resolved from tensor arguments at runtime.
TirTemplate.create() requires every such variable to appear directly
in some buffer's shape or stride, and raises otherwise:

RuntimeError: Constexpr variable `X` is not used in any buffer shape or stride.

This rules out two common patterns where a compile-time value is known at
call time but does not itself index a buffer — a grid-only variable (number
of CTAs) and a computation-only variable (e.g. a scale factor).

Closes #2760.

Solution

The phase-2 value-resolution path (_parse_phase2_key) already supports
resolving a constexpr via an explicit keyword argument (if name in kwargs),
and the JIT argument binder (#2357) already routes extra kwargs into
compile_kwargs. The only thing blocking this path was create() raising
before the template was cached.

This PR replaces the raise with a sentinel entry so grid-only /
computation-only constexpr vars flow through the existing explicit-kwarg
path — no binder changes, no new code paths:

# before: raise RuntimeError(...)
# after:
matcher[s] = (s.name, "__kwarg__", -1, s.name)

Value flow

kernel(a, b, num_blocks=4)
  → binder routes num_blocks into compile_kwargs (extra_kwargs, #2357)
  → _parse_phase2_key: name "num_blocks" in kwargs → returns 4
  → subs = {orig_name: 4}
  → phase2 T.const("num_blocks") → eager_jit_subs["num_blocks"] → 4

Verification

Tested on 0.1.12 (wheel + source patch on the installed builder.py,
which is byte-identical to main):

python3 -m pytest testing/python/issue/test_tilelang_issue_2760.py -v
test_const_grid_only             PASSED
test_const_computation_only      PASSED
test_const_mixed                 PASSED
test_const_missing_kwarg_errors  PASSED
4 passed in 3.66s

• test_const_grid_only — T.const("num_blocks") used only in T.Kernel(...), resolved via num_blocks=4 kwarg.
• test_const_computation_only — T.const("scale") used only in kernel body, resolved via scale=3.
• test_const_mixed — mix of shape-derived (N) and explicit-kwarg (num_blocks, scale).
• test_const_missing_kwarg_errors — omitting the kwarg raises ValueError with a clear message.

Note: I don't have a main source build with CUDA locally, so verification
was done by applying the source patch to the 0.1.12 wheel installation
(builder.py is identical between 0.1.12 and main — 0 lines diff).

Summary

• Allow T.const() constexpr variables used only for grid/kernel extent and/or inside kernel computations (i.e., not referenced by any buffer shape or stride) to be resolved via explicit call-time keyword arguments instead of being rejected during TirTemplate.create().
• Update TirTemplate.create() to record such “missing from buffer shape/stride” constexpr variables into the template matcher using a "__kwarg__" sentinel (keyword-only substitution), so phase-2 matching/JIT argument binding can pull their values from provided **kwargs.
• Add regression coverage for issue #2760, including grid-only, computation-only, and mixed usage; cache isolation across different constexpr-kwarg configurations; float keyword handling; and a clear error when a required constexpr keyword is omitted (“Cannot find value for constexpr variable”).

Testing

• Added testing/python/issue/test_tilelang_issue_2760.py with CUDA-gated tests:
  • test_const_grid_only
  • test_const_computation_only
  • test_const_mixed
  • test_const_pipelined_gemm
  • test_const_missing_kwarg_errors
  • test_const_cache_isolation
  • test_const_float_value

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The eager template matcher now retains T.const() variables absent from buffer shapes or strides for explicit keyword resolution. CUDA regression tests cover grid-only, computation-only, mixed, pipelined GEMM, cache isolation, float values, and missing-keyword errors.

Changes

Constexpr keyword support

Layer / File(s)	Summary
Matcher support for explicit constexpr values tilelang/language/eager/builder.py	TirTemplate.create records unmatched constexpr variables with a __kwarg__ sentinel instead of raising during template creation.
Regression coverage for constexpr patterns and caching testing/python/issue/test_tilelang_issue_2760.py	Tests validate grid-only, computation-only, mixed, pipelined GEMM, cache-isolated, and float T.const() values supplied by keyword arguments, along with the expected error when a required keyword is omitted.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant TestKernelCall
  participant TirTemplate
  participant Phase2Matcher
  participant CUDAKernel
  TestKernelCall->>TirTemplate: create JIT kernel with T.const variables
  TirTemplate->>Phase2Matcher: retain unmatched constexpr as __kwarg__
  TestKernelCall->>Phase2Matcher: pass num_blocks and scale kwargs
  Phase2Matcher->>CUDAKernel: resolve constexpr values and launch cached variant
  CUDAKernel-->>TestKernelCall: produce output for assertions

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: allowing T.const() values used only in grid dimensions or computations.
Linked Issues check	✅ Passed	The code and tests address #2760 by permitting constexprs resolved from explicit kwargs and covering the required usage patterns.
Out of Scope Changes check	✅ Passed	The added tests and builder change stay within the issue scope and support the new explicit-kwarg constexpr resolution path.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[LeiWang1999]

@NolanHo Thanks for your contributions, but two comment left

1. error message

import torch

import tilelang
import tilelang.language as T


@tilelang.jit
def copy_kernel(A, B):
    N = T.const("N")
    num_blocks = T.const("num_blocks")

    A: T.Tensor((N,), T.float32)
    B: T.Tensor((N,), T.float32)

    with T.Kernel(num_blocks, threads=128) as bx:
        B[bx] = A[bx]


def main() -> None:
    A = torch.empty(4, dtype=torch.float32)
    B = torch.empty_like(A)

    try:
        # Deliberately omit num_blocks. It is a kwarg-only constexpr and has no
        # corresponding tensor argument from which TileLang could infer it.
        copy_kernel(A, B)
    except ValueError as error:
        message = str(error)
        print(message)
        assert "Or provide the corresponding tensor argument `num_blocks`." in message
        return

    raise AssertionError("Expected a missing-constexpr ValueError")


if __name__ == "__main__":
    main()

the output is

Cannot find value for constexpr variable `num_blocks`
  Please provide it as a keyword argument, e.g. `num_blocks=<value>`
  Or provide the corresponding tensor argument `num_blocks`.

the Or provide the corresponding tensor argument num_blocks. is not precise.

2. doc update.