[EagerJIT] Add Support for Parameter Only Kernel Compilation

[kurisu6912]

This pr add support for the grammar to compile kernel with only constexpr parameters.

@tilelang.jit
def transpose(X, Y, block_M, block_N):
    M, N = T.const("M N")
    X: T.Tensor[[M, N], T.float32]
    Y: T.Tensor[[N, M], T.float32]
    ...

transpose.compile(M=1024, N=1024, block_M=64, block_N=64)
transpose.par_compile(
    [
        {'M': M, 'N': N, 'block_M': 128, 'block_N': 128}
        for M, N in [(256, 256), (512, 512), (1024, 1024)]
    ]
)

[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]

📝 Walkthrough

Walkthrough

Changes add support for compile-time constants in JIT-compiled functions, expand AST argument type annotation parsing to handle Call and Subscript patterns with explicit kwargs parameter representation, and refactor the constraint matcher from 3-tuple to 4-tuple structure to include buffer element names for phase-2 key resolution.

Changes

Cohort / File(s)	Summary
Test Additions testing/python/language/test_tilelang_language_eager_jit.py	New test function test_jit2_compile_with_consts that exercises JIT compilation with constant values for M, N, block_M, and block_N parameters.
Test Refactoring testing/python/layout/test_tilelang_annotate_loop_layout.py	Updated test cases to pass dimension integers (M, N) directly to compile() instead of tensor arguments; adjusted expected indexing expressions to reflect dimension-based kernel generation.
Core Language Changes tilelang/language/eager/ast.py	Expanded argument type annotation parsing to handle Call and Subscript patterns with inner Attribute references to TensorProxy, StridedTensorProxy, or ptr; added explicit __kwargs parameter to function AST.
Constraint Matcher Refactoring tilelang/language/eager/builder.py	Extended matcher tuples from 3-element to 4-element structure (adding buffer element name); updated phase-2 key parsing to perform name-based keyword lookup before positional fallback; enhanced exception handling to catch TypeError.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Poem

🐰 Constants compile with grace so true,
Matchers now hold names in their crew,
Annotations expand in patterns bright,
AST's kwargs shine with new light!
Dimensions flow where tensors once flew. ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding support for parameter-only kernel compilation in the EagerJIT system, which is reflected in the core changes to ast.py, builder.py, and test files.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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 and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

testing/python/language/test_tilelang_language_eager_jit.py (1)

228-244: Test compiles but doesn't verify execution correctness.

This test validates that parameter-only compilation works, but unlike other tests in this file (e.g., test_jit2_gemm), it doesn't execute the kernel or verify results. Consider adding execution and assertion to catch potential issues in the transpose logic.

♻️ Suggested enhancement to verify execution

     transpose.compile(M=1024, N=1024, block_M=64, block_N=64)
+
+    # Verify correctness
+    X = torch.randn(1024, 1024, dtype=torch.float32, device="cuda")
+    Y = torch.empty(1024, 1024, dtype=torch.float32, device="cuda")
+    transpose(X, Y)
+    torch.testing.assert_close(Y, X.T)

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 4fb96a0 and b26e090.

📒 Files selected for processing (4)

• testing/python/language/test_tilelang_language_eager_jit.py
• testing/python/layout/test_tilelang_annotate_loop_layout.py
• tilelang/language/eager/ast.py
• tilelang/language/eager/builder.py

🧰 Additional context used

🧠 Learnings (3)

📚 Learning: 2026-01-06T05:20:51.649Z

Learnt from: silentCoder-dev
Repo: tile-ai/tilelang PR: 1606
File: testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py:30-30
Timestamp: 2026-01-06T05:20:51.649Z
Learning: In `testing/python/transform/test_tilelang_transform_hoist_broadcast_values.py`, the test validates that the `hoist_broadcast_values` transformation pass correctly identifies and hoists broadcast operations by checking for patterns in the generated kernel source code. The specific literal values used (e.g., 430) are not important for the test's purpose, as it does not validate numerical precision or actual stored tensor values.

Applied to files:

• testing/python/layout/test_tilelang_annotate_loop_layout.py
• testing/python/language/test_tilelang_language_eager_jit.py

📚 Learning: 2025-12-18T04:50:00.512Z

Learnt from: silentCoder-dev
Repo: tile-ai/tilelang PR: 1464
File: testing/python/language/test_tilelang_language_rand.py:14-14
Timestamp: 2025-12-18T04:50:00.512Z
Learning: In `testing/python/language/test_tilelang_language_rand.py`, the TileLang kernel uses `blk_M = M` (single block) and calls `rng_rand()` four times per element to align results with the Triton implementation, which uses `blk_M = 128` (multiple blocks) and calls the RNG once per element. These differences compensate for internal RNG behavior differences between TileLang and Triton.

Applied to files:

• testing/python/layout/test_tilelang_annotate_loop_layout.py

📚 Learning: 2025-11-14T07:56:11.098Z

Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.

Applied to files:

• testing/python/language/test_tilelang_language_eager_jit.py

🧬 Code graph analysis (4)

testing/python/layout/test_tilelang_annotate_loop_layout.py (2)

tilelang/jit/__init__.py (2)

• compile (47-107)
• compile (372-398)

tilelang/jit/adapter/tvm_ffi.py (1)

• get_kernel_source (311-316)

tilelang/language/eager/builder.py (1)

tilelang/jit/exceptions.py (2)

• JITNoBuilderError (4-13)
• EagerJITBuildError (16-24)

tilelang/language/eager/ast.py (4)

tilelang/language/eager/builder.py (1)

• arg (670-674)

tilelang/language/dtypes.py (1)

• dtype (14-19)

tilelang/jit/adapter/tvm_ffi.py (1)

• func (206-260)

tilelang/language/proxy.py (3)

• TensorProxy (137-155)
• StridedTensorProxy (158-167)
• ptr (256-275)

testing/python/language/test_tilelang_language_eager_jit.py (2)

tilelang/language/eager/builder.py (1)

• const (851-879)

tilelang/language/allocate.py (1)

• alloc_shared (39-54)

🪛 Ruff (0.14.10)

tilelang/language/eager/builder.py

942-946: Avoid specifying long messages outside the exception class

(TRY003)

🔇 Additional comments (10)

testing/python/layout/test_tilelang_annotate_loop_layout.py (4)

28-32: LGTM! Consistent update to dimension-based compilation.

The change from tensor-based to keyword-argument-based compilation (M=M, N=N) aligns with the new parameter-only kernel compilation feature. The assertion correctly validates the expected vectorized memory access pattern.

---

44-47: LGTM!

Same consistent pattern as above for the identity layout test.

---

70-76: LGTM!

Updated to dimension-based compilation with the reformatted assertion checking the expected bitwise masking pattern (& 63).

---

102-108: LGTM!

Consistent with other test updates in this file.

tilelang/language/eager/ast.py (2)

475-475: LGTM! Enables kwargs passthrough for parameter-only compilation.

Adding __kwargs to the function arguments allows keyword arguments to flow through the AST transformation, supporting the new dimension-based compilation API.

---

514-527: LGTM! Expanded type hint parsing for tensor annotations.

The new case handling correctly identifies T.Tensor[...], T.StridedTensor[...], and T.ptr(...) patterns by checking both ast.Call and ast.Subscript node types. Storing ptr for all tensor-like types allows the eager JIT builder to recognize these as buffer arguments requiring constexpr substitution.

tilelang/language/eager/builder.py (4)

893-893: LGTM! Type annotation updated for 4-tuple matcher.

The matcher now stores (buffer_name, type, index, constexpr_name) to support both tensor-based inference and direct keyword argument lookup.

---

900-905: LGTM! Matcher construction captures constexpr variable names.

Storing s.name (the constexpr variable name) alongside the buffer metadata enables the phase-2 parser to resolve constants via direct keyword arguments like M=1024.

---

929-946: LGTM! Phase-2 parsing supports both named constants and tensor inference.

The updated logic correctly prioritizes direct keyword lookup (kwargs.get(name)) before falling back to tensor-based shape/stride extraction. The error message is clear and actionable, guiding users to provide either the constant value or the corresponding tensor.

---

1027-1033: LGTM! TypeError handling improves lazy/eager detection robustness.

Adding TypeError to the caught exceptions handles cases where calling the original function fails due to argument mismatches, correctly identifying such functions as eager-style.