[Misc] add env for default target/backend/verbose#1512
[Misc] add env for default target/backend/verbose#1512LeiWang1999 merged 8 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! 🚀 |
📝 WalkthroughWalkthroughEnvironment-driven defaults were added for compilation parameters. New TILELANG_* env vars and Environment accessors were introduced. Autotuner, JIT, and cache APIs now accept None for target/execution_backend/verbose and resolve defaults from the environment when parameters are omitted. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant API as "Autotuner/JIT/Cache API"
participant Env as "tilelang.env"
participant Resolver as "Target/Backend Resolver"
participant Backend as "Execution Backend"
Caller->>API: call(func, target=None, execution_backend=None, verbose=None)
API->>Env: get_default_target()
API->>Env: get_default_execution_backend()
API->>Env: get_default_verbose()
Env-->>API: return defaults (e.g., "auto", "auto", False)
API->>Resolver: normalize target/backend (may remain "auto")
Resolver-->>Backend: select/initialize backend
Backend-->>API: backend handle
API-->>Caller: return compiled kernel / JIT wrapper
Note over API,Env: Env-driven defaults applied before normalization
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
⏰ 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)
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: 2
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tilelang/autotuner/tuner.pytilelang/env.pytilelang/jit/__init__.py
🧰 Additional context used
🧬 Code graph analysis (2)
tilelang/autotuner/tuner.py (1)
tilelang/env.py (3)
get_default_target(298-300)get_default_execution_backend(302-304)get_default_verbose(306-308)
tilelang/jit/__init__.py (1)
tilelang/env.py (3)
get_default_target(298-300)get_default_execution_backend(302-304)get_default_verbose(306-308)
🪛 Ruff (0.14.10)
tilelang/jit/__init__.py
148-148: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
⏰ 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 Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
🔇 Additional comments (9)
tilelang/env.py (2)
248-252: LGTM!The new environment variable declarations follow the existing patterns in the codebase with sensible defaults. The naming convention using
TILELANG_prefix for the environment variable keys is consistent.
298-308: LGTM!The accessor methods are clean and follow existing patterns (e.g.,
is_cache_globally_disabled). Theget_default_verbosemethod correctly parses truthy string values to boolean.tilelang/autotuner/tuner.py (3)
12-12: LGTM!The import is correctly placed with other tilelang imports.
172-178: LGTM!The environment variable defaults are applied correctly before the target resolution logic. The order ensures that explicit parameters take precedence over environment variables.
164-167: LGTM!Good documentation of the environment variables. This helps users understand how to configure defaults.
tilelang/jit/__init__.py (4)
27-27: LGTM!The import is correctly placed with other tilelang imports.
94-100: LGTM!The environment variable defaulting logic is correctly placed before the assertion and target resolution. The pattern is consistent with other functions in the PR.
184-191: LGTM!Consistent implementation with the
compilefunction. The defaults are applied before passing to the thread pool executor.
546-552: LGTM!The
jitdecorator correctly applies environment defaults before creating theJITImpl. This ensures the decorator can be used without arguments while still respecting environment configuration.
LeiWang1999
left a comment
LeiWang1999
left a comment
There was a problem hiding this comment.
Overall LGTM, but left some performance concerns.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tilelang/cache/kernel_cache.py (1)
128-128: Consider explicit Optional type annotation.The type annotation
dict = Noneshould bedict | None = Nonefor PEP 484 compliance. This is pre-existing code but can be improved for consistency with the newly updated parameters above.🔎 Proposed fix for explicit Optional
- pass_configs: dict = None, + pass_configs: dict | None = None,Similarly, check line 129 for
compile_flags:- compile_flags: list[str] | str | None = None, + compile_flags: list[str] | str | None = None,(Line 129 already appears correct)
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tilelang/cache/__init__.pytilelang/cache/kernel_cache.pytilelang/jit/__init__.py
🧰 Additional context used
🧬 Code graph analysis (3)
tilelang/cache/kernel_cache.py (1)
tilelang/env.py (3)
get_default_target(298-300)get_default_execution_backend(302-304)get_default_verbose(306-308)
tilelang/cache/__init__.py (1)
tilelang/language/ast/ir.py (1)
target(1677-1707)
tilelang/jit/__init__.py (3)
tilelang/language/ast/ir.py (1)
target(1677-1707)tilelang/jit/kernel.py (1)
JITKernel(38-767)tilelang/env.py (1)
get(178-181)
🪛 Ruff (0.14.10)
tilelang/cache/kernel_cache.py
128-128: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
tilelang/jit/__init__.py
120-120: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
🔇 Additional comments (7)
tilelang/jit/__init__.py (4)
48-109: Well-structured environment variable integration.The function signature correctly accepts None for compilation parameters, and the docstring comprehensively documents the environment variable fallback behavior. The delegation of None resolution to
cached()maintains clean separation of concerns.
436-534: Comprehensive environment variable integration for jit decorator.The overloads, implementation, and documentation are well-aligned. The docstring clearly explains the environment variable fallback mechanism, making the API behavior transparent to users.
536-593: Clean implementation with appropriate documentation reference.The
lazy_jitdecorator correctly mirrors thejitpattern for None-accepting parameters. The concise docstring that referencesjitfor detailed parameter and environment variable documentation follows good DRY principles.
99-108: The downstream None handling forexecution_backend,target, andverboseis correctly implemented inKernelCache.cached()(tilelang/cache/kernel_cache.py:119-160). These parameters are properly resolved using environment variable accessors (env.get_default_target(),env.get_default_execution_backend(),env.get_default_verbose()) when None is passed, before being forwarded to theJITKernelconstructor. No changes needed.tilelang/cache/kernel_cache.py (3)
124-127: LGTM! Signature correctly updated for environment-driven defaults.The parameter types now properly accept None, enabling callers to opt into environment-variable resolution. The union type syntax is clean and follows modern Python typing conventions.
134-157: LGTM! Documentation clearly explains the environment-driven defaults.The documentation properly describes the centralized environment variable processing and the behavior when parameters are None. The emphasis on this being the "ONLY place" for such processing establishes a clear architectural pattern.
158-166: LGTM! Environment variable processing is clean and correctly placed.The environment variable resolution happens at the right point in the flow—before target normalization and backend resolution. The logic correctly preserves explicitly-provided values while filling in defaults only when parameters are None.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tilelang/jit/__init__.py (1)
421-421: Cleaner cache lookup implementation.The simplification from
self._kernel_cache.get(key, None)toself._kernel_cache.get(key)is a good stylistic improvement sincedict.get()returnsNoneby default when the key is not found.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
examples/flash_decoding/example_gqa_decode_varlen_logits.pytilelang/cache/__init__.pytilelang/jit/__init__.py
🧰 Additional context used
🧬 Code graph analysis (3)
tilelang/cache/__init__.py (1)
tilelang/language/ast/ir.py (1)
target(1677-1707)
examples/flash_decoding/example_gqa_decode_varlen_logits.py (1)
tilelang/autotuner/tuner.py (1)
autotune(691-786)
tilelang/jit/__init__.py (2)
tilelang/jit/kernel.py (1)
JITKernel(38-767)tilelang/env.py (1)
get(178-181)
⏰ 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 (8)
examples/flash_decoding/example_gqa_decode_varlen_logits.py (1)
200-200: LGTM! Explicit decorator path improves consistency.The change from
@autotuneto@tilelang.autotunemakes the decorator path explicit and aligns with the PR's introduction of environment-driven defaults. This is consistent with the existing usage of@tilelang.jiton Line 201 andtilelang.disable_cache()on Line 10.tilelang/cache/__init__.py (1)
20-23: LGTM! Type annotations and defaults are now consistent.The function signature correctly accepts
Nonefortarget,target_host,execution_backend, andverbose, with proper type hints including| Nonefor all parameters. The implementation appropriately delegates environment variable resolution to the kernel cache layer.tilelang/jit/__init__.py (6)
48-108: Well-structured environment-driven defaults implementation.The
compilefunction correctly acceptsNoneforexecution_backend,target, andverbose, and the docstring clearly documents how environment variables (TILELANG_TARGET,TILELANG_EXECUTION_BACKEND,TILELANG_VERBOSE) provide defaults when parameters are omitted. The implementation properly delegates to the cached layer.
111-190: Parallel compilation correctly inherits environment variable support.The
par_compilefunction signature and docstring are consistent with thecompilefunction. Type annotations properly include| Nonefor all optional parameters, including the previously fixednum_workers. The implementation correctly delegates tocompile()which handles environment variable resolution.
265-268: Dataclass fields correctly support optional environment-driven defaults.The
JITImpldataclass properly declaresexecution_backend,target,target_host, andverbosewith| Nonetype hints, allowing environment variable defaults to be resolved at compilation time.
436-533: JIT decorator properly supports environment-driven configuration.The
jitfunction and its overloads correctly acceptNonefortarget,target_host,execution_backend, andverbose. The comprehensive docstring clearly explains how environment variables provide defaults, making it easy for users to configure batch operations viaTILELANG_TARGET,TILELANG_EXECUTION_BACKEND, andTILELANG_VERBOSE.
536-593: Lazy JIT decorator aligns with environment variable pattern.The
lazy_jitfunction correctly mirrors thejitfunction's support for environment-driven defaults. The concise docstring appropriately references thejitdocumentation for parameter details and environment variables, avoiding duplication while maintaining clarity.
48-593: Environment variable resolution is correctly implemented.Verification confirms:
TILELANG_TARGETdefaults to"auto"(set in tilelang/env.py line 250)TILELANG_EXECUTION_BACKENDdefaults to"auto"(set in tilelang/env.py line 251)TILELANG_VERBOSEdefaults toFalse(set to"0"in tilelang/env.py line 252, parsed in get_default_verbose())The
compile()function correctly delegates tocached()in tilelang/cache/kernel_cache.py, which applies these environment variable defaults when parameters areNone(lines 159-162 of kernel_cache.py).
Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
…1483) * refactor kernel_cache Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * set key to auto when execution_backend is None Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * fix type error on _load_kernel_source override Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * put kernel cache subclasses in adapter Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * remove auto in execution_backend Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * remove _save_kernel_source_code_to_disk for NVRTCKernelCache Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * refactor all cutedsl if statements Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * remove unused kernel_cubin_path Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * add func arguments Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * remove ctypes Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * update ci.yaml to test kernel cache Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * fix bugs after #1512 merged Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> * test: fix kernel cache test * fix: ruff * fix: ruff * fix: coderabbit * fix: coderabbit * fix: ruff * fix: MMA * fix: coderabbit * fix: codex * fix: coderabbit * fix: coderabbit --------- Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn> Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
2 participants
With these env vars, we can batch run examples for a specific target/backend.
Summary by CodeRabbit
New Features
Documentation
✏️ Tip: You can customize this high-level summary in your review settings.