[EagerJIT] Fix bug in handling of positional arguments

[kurisu6912]

This pr fix some wrong code in jit argument parsing, which results in a bug that positional arguments are ignored.

import tilelang
import tilelang.language as T


@tilelang.jit
def vectorize_broadcast_int8(vec_num):
    with T.Kernel(1, threads=128):
        a = T.alloc_local((64,), 'int8')
        b = T.alloc_var('int8')
        for i in T.vectorized(vec_num):
            a[i] = b


kernel = vectorize_broadcast_int8.compile(128)
print(kernel)

Summary by CodeRabbit

• Style

  • Adjusted formatting in code examples with blank line and quote standardization.

• Refactor

  • Improved internal argument handling logic for better consistency across function calls.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[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

The pull request includes formatting adjustments, quote style consistency updates, and argument parsing refactoring across example files and the builder module. Multi-line method calls are collapsed to single-line equivalents, exception print statements are updated to use double-quotes, and argument forwarding in the builder's get_tir method is refactored to use keyword-only propagation.

Changes

Cohort / File(s)	Summary
Example code formatting examples/autodd/tilelang_buggy.py, examples/autodd/tilelang_minimized_expected.py	Collapsed multi-line method calls to single-line format; removed extra blank lines; updated exception print statements from single to double-quote style for consistency
Builder argument refactoring tilelang/language/eager/builder.py	Refactored get_tir to use _parse_phase1_key for argument parsing; changed template construction to use keyword-only argument propagation (**kwargs) instead of positional and keyword arguments (*args, **kwargs)

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [AutoDD] Add Tilelang AutoDD to Reduce Buggy Program #1639 — Modifies the same example files and MatmulConfig methods with identical function refactoring patterns

Suggested labels

enhancement

Suggested reviewers

• LeiWang1999
• tzj-fxz

Poem

🐰 Hoppy changes hopping through the code,
Lines collapsed on the developer's road,
Quotes aligned in double-quote grace,
Arguments flow in their proper place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.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 accurately describes the main fix: addressing a bug in positional argument handling in the EagerJIT system, which is reflected in the key changes to argument parsing and forwarding in builder.py.

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

✨ Finishing touches

• 📝 Generate docstrings

---

---

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2d8d367 and 65b68f2.

📒 Files selected for processing (3)

• examples/autodd/tilelang_buggy.py
• examples/autodd/tilelang_minimized_expected.py
• tilelang/language/eager/builder.py

🧰 Additional context used

🧬 Code graph analysis (1)

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

tilelang/jit/__init__.py (1)

• get_tir (290-302)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
• GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (5)

examples/autodd/tilelang_minimized_expected.py (1)

46-46: LGTM!

Minor formatting change to use double quotes for the f-string. No functional impact.

examples/autodd/tilelang_buggy.py (2)

75-79: LGTM!

Formatting improvements collapsing multi-line calls to single lines. No functional changes.

---

105-106: LGTM!

Whitespace formatting adjustment.

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

1062-1062: LGTM!

Correctly passes only non-tensor keyword arguments to _build_tir_template after tensor args have been extracted by _parse_phase1_key.

---

1067-1074: LGTM! Proper fix for positional argument handling.

The refactored get_tir method correctly:

1. Uses _parse_phase1_key to convert positional args to kwargs and extract tensor_args separately
2. Passes only non-tensor **kwargs to _build_tir_template
3. Passes **tensor_args, **kwargs to TirTemplate.get_tir, which expects keyword-only arguments

Since tensor_args are popped from kwargs in _parse_phase1_key (line 995), there's no risk of key collision when merging them back.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

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.}