[Typing] Update Kernel signature and add type hints for buffer operations

[clouds56]

I know it's not perfect, but better than nothing

def __getitem__(self, idx) -> Buffer

This is not breaking, I think T.Kernel is intended to call with T.Kernel(64, 128) but not T.Kernel([64], [128])

Summary by CodeRabbit

• New Features

  • Added new utility operations for expression handling.
  • Typing support added for indexing on buffer-like objects.

• API Changes

  • Simplified kernel-like call signature to accept block extents directly (no per-block list wrapping required).

_{✏️ 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

Kernel signature changed to accept variadic tir.PrimExpr for blocks (no per-block lists); Buffer gained __getitem__/__setitem__ declarations in a TYPE_CHECKING block; Cast, max, and min function stubs added to TIR type stubs.

Changes

Cohort / File(s)	Summary
Kernel API Update tilelang/language/kernel.py	Signature changed from def Kernel(*blocks: list[tir.PrimExpr], ...) to def Kernel(*blocks: tir.PrimExpr, ...); docstring/type hint updated from blocks : List[int] → blocks : int.
Buffer Type Hints tilelang/language/v2/annot.py	Added __getitem__(self, idx) -> Buffer and __setitem__(self, idx, val) declarations inside the TYPE_CHECKING block (typing/IDE-only).
TIR Function Stubs tilelang/language/tir/ir.pyi	Added public stubs: `def Cast(dtype, value: _T, span: Span

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Poem

🐰 I hopped through code with nimble feet,
Blocks unwrapped made my rhythm sweet,
Cast, max, and min now join the crew,
Buffers index with hints anew,
A tiny rabbit's dev-time tune — hooray for you!

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.56% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ 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 summarizes the main changes: updating the Kernel signature (from list-wrapped to variadic PrimExpr) and adding type hints for buffer operations (getitem, setitem, and new functions in the type stub file).

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

[clouds56]

@coderabbitai generate docstrings

[coderabbitai]

✅ Actions performed

Initiated docstring generation; will generate only if new commits exist.

[coderabbitai]

Note

Docstrings generation - SUCCESS
Generated docstrings for this pull request at #1546

[coderabbitai]

Actionable comments posted: 3

📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0c44d59 and e5d95aa.

📒 Files selected for processing (1)

• tilelang/language/tir/ir.pyi

[kurisu6912]

LGTM!