[Language] Make TL scripts friendly to Python syntax highlights#1466
[Language] Make TL scripts friendly to Python syntax highlights#1466LeiWang1999 merged 3 commits intotile-ai:mainfrom
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
WalkthroughReorganizes tilelang.language public surface: moves several imports to new *_op modules, expands reduce exports, replaces Serial/Unroll aliases with wrapper functions, and updates many module docstrings and example file headers. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tilelang/language/loop.py (1)
178-186: Consider adding explicit type signatures to improve IDE support.The wrapper functions use
*args, **kwargs, which loses the detailed type information from the underlyingserialandunrollfunctions. This actually hinders IDE autocomplete and type checking, which seems contrary to making the code "friendly to Python syntax highlights" (per the PR title).🔎 Apply this diff to add explicit type signatures matching the underlying functions:
# "Serial" and "Unroll" are aliases of "T.serial" and "T.unroll". We use uppercase to emphasize that they are tile-level loops. -def Serial(*args, **kwargs): - return serial(*args, **kwargs) +def Serial( + start: tir.PrimExpr, + stop: tir.PrimExpr | None = None, + step: tir.PrimExpr | None = None, + *, + annotations: dict[str, Any] | None = None, +) -> frame.ForFrame: + """Wrapper for serial() - uppercase emphasizes tile-level loop. + + See serial() for detailed parameter documentation. + """ + return serial(start, stop, step, annotations=annotations) -def Unroll(*args, **kwargs): - return unroll(*args, **kwargs) +def Unroll( + start: tir.PrimExpr, + stop: tir.PrimExpr | None = None, + step: tir.PrimExpr | None = None, + *, + explicit: bool = False, + unroll_factor: int | None = None, + annotations: dict[str, Any] | None = None, +) -> frame.ForFrame: + """Wrapper for unroll() - uppercase emphasizes tile-level loop. + + See unroll() for detailed parameter documentation. + """ + return unroll(start, stop, step, explicit=explicit, unroll_factor=unroll_factor, annotations=annotations)tilelang/language/__init__.py (1)
63-84: Verify module file renames are complete.The import paths have been updated from module names that shadow built-ins (e.g.,
_opsuffixed names. All new module files (gemm_op.py,fill_op.py,print_op.py,reduce_op.py) exist, old module files have been removed, and all imported functions are properly defined.Note: The
# noqa: F401directives on lines 63, 65, and 84 are unused.The project's ruff configuration ignores F401 globally, making these directives unnecessary. Remove them to keep the code clean:
-from .gemm_op import gemm, gemm_v1, gemm_v2 # noqa: F401 +from .gemm_op import gemm, gemm_v1, gemm_v2 from .experimental.gemm_sp import gemm_sp, gemm_sp_v2 # noqa: F401 -from .fill_op import fill, clear # noqa: F401 +from .fill_op import fill, clear from .reduce_op import ( reduce, # noqa: F401 reduce_max, # noqa: F401 reduce_min, # noqa: F401 reduce_sum, # noqa: F401 reduce_abssum, # noqa: F401 reduce_absmax, # noqa: F401 reduce_bitand, # noqa: F401 reduce_bitor, # noqa: F401 reduce_bitxor, # noqa: F401 cumsum, # noqa: F401 finalize_reducer, # noqa: F401 warp_reduce_sum, # noqa: F401 warp_reduce_max, # noqa: F401 warp_reduce_min, # noqa: F401 warp_reduce_bitand, # noqa: F401 warp_reduce_bitor, # noqa: F401 ) -from .print_op import print, device_assert # noqa: F401 +from .print_op import print, device_assert
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
3rdparty/tvm(1 hunks)tilelang/language/__init__.py(2 hunks)tilelang/language/loop.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/language/loop.py (1)
tilelang/language/tir/ir.py (2)
serial(9-28)unroll(75-101)
tilelang/language/__init__.py (4)
tilelang/language/gemm_op.py (2)
gemm_v1(129-154)gemm_v2(158-183)tilelang/language/experimental/gemm_sp.py (2)
gemm_sp(19-97)gemm_sp_v2(101-223)tilelang/language/fill_op.py (2)
fill(9-36)clear(39-62)tilelang/language/print_op.py (1)
device_assert(127-138)
🪛 Ruff (0.14.8)
tilelang/language/__init__.py
63-63: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
64-64: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
65-65: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
84-84: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
⏰ 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). (3)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- 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 (1)
3rdparty/tvm (1)
1-1: Invalid commit hash reference for TVM submodule.The commit
2b1ead1a375704c75af563cc800aa9347583ba2breferenced in the review does not exist. However, the TVM submodule is already at commitea89d1f6a7b6a91d6314de67aedd031060f4b6ef, which implements "[Language] Make TL scripts friendly to Python syntax highlights"—exactly matching the PR objective. Verify that the submodule pointer is correctly set to the intended commit.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tilelang/language/math_intrinsics.py (1)
133-147: Fix docstring:__expcomputes ex, not 2x.The docstring at line 134 incorrectly states "Calculate 2x with fast math". The function
__expcomputes the natural exponential (e^x) using thetl.__expintrinsic, consistent with standard mathematical conventions and the Triton library (which provides__logfor natural log,__log2for log base 2, and__exp10for base-10 exponential). Update the docstring to "Calculate ex with fast math".
🧹 Nitpick comments (2)
tilelang/language/loop.py (1)
178-186: Consider adding type hints and docstrings to wrappers.The wrapper functions lose the type hints and documentation from
serial/unroll. For better IDE support and discoverability, consider forwarding the type signature:🔎 Suggested improvement
-def Serial(*args, **kwargs): - return serial(*args, **kwargs) +def Serial( + start: tir.PrimExpr, stop: tir.PrimExpr | None = None, step: tir.PrimExpr | None = None, *, annotations: dict[str, Any] | None = None +) -> frame.ForFrame: + """Alias for serial. See serial for documentation.""" + return serial(start, stop, step, annotations=annotations) -def Unroll(*args, **kwargs): - return unroll(*args, **kwargs) +def Unroll( + start: tir.PrimExpr, + stop: tir.PrimExpr | None = None, + step: tir.PrimExpr | None = None, + *, + explicit: bool = False, + unroll_factor: int | None = None, + annotations: dict[str, Any] | None = None, +) -> frame.ForFrame: + """Alias for unroll. See unroll for documentation.""" + return unroll(start, stop, step, explicit=explicit, unroll_factor=unroll_factor, annotations=annotations)tilelang/language/__init__.py (1)
61-65: LGTM on module reorganization.The
_opsuffix convention properly distinguishes these modules from Python built-ins like# noqa: F401directives as unused on these lines—consider removing them if F401 is not enabled in your linter configuration.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
examples/gemm_sp/example_custom_compress.py(0 hunks)examples/gemm_sp/example_gemm_sp.py(0 hunks)tilelang/language/__init__.py(2 hunks)tilelang/language/atomic.py(1 hunks)tilelang/language/builtin.py(1 hunks)tilelang/language/copy_op.py(1 hunks)tilelang/language/customize.py(1 hunks)tilelang/language/fastmath.py(1 hunks)tilelang/language/fill_op.py(1 hunks)tilelang/language/gemm_op.py(1 hunks)tilelang/language/kernel.py(1 hunks)tilelang/language/logical.py(1 hunks)tilelang/language/loop.py(2 hunks)tilelang/language/math_intrinsics.py(1 hunks)tilelang/language/proxy.py(1 hunks)tilelang/language/reduce_op.py(1 hunks)tilelang/quantize/quantization.py(1 hunks)
💤 Files with no reviewable changes (2)
- examples/gemm_sp/example_custom_compress.py
- examples/gemm_sp/example_gemm_sp.py
✅ Files skipped from review due to trivial changes (11)
- tilelang/language/proxy.py
- tilelang/language/kernel.py
- tilelang/language/builtin.py
- tilelang/language/reduce_op.py
- tilelang/language/gemm_op.py
- tilelang/language/fill_op.py
- tilelang/language/copy_op.py
- tilelang/language/atomic.py
- tilelang/language/logical.py
- tilelang/quantize/quantization.py
- tilelang/language/fastmath.py
🧰 Additional context used
🧬 Code graph analysis (1)
tilelang/language/loop.py (1)
tilelang/language/tir/ir.py (2)
serial(9-28)unroll(75-101)
🪛 Ruff (0.14.8)
tilelang/language/__init__.py
61-61: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
62-62: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
63-63: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
64-64: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
65-65: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
84-84: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
⏰ 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). (1)
- GitHub Check: Quick Lint
🔇 Additional comments (5)
tilelang/language/math_intrinsics.py (1)
1-2: LGTM!The module-level docstring addition aligns well with the PR objective of improving Python syntax highlighting friendliness.
tilelang/language/customize.py (1)
1-1: LGTM!The updated module docstring clearly describes the purpose of this module.
tilelang/language/loop.py (1)
1-1: LGTM!The updated docstring accurately reflects the module's purpose.
tilelang/language/__init__.py (2)
66-83: LGTM!The expanded reduce-related exports provide a more complete public API surface for reduction operations.
84-84: Good renaming to avoid built-in shadowing.Importing from
print_opinstead of
SiriusNEO
changed the title
2 participants
Summary by CodeRabbit
New Features
Refactor
Documentation
Chores
✏️ Tip: You can customize this high-level summary in your review settings.