Repository navigation
[Backend] Add unified backend resolution policy - #2318
Merged
Merged
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! 🚀 |
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR replaces separate pipeline, device-codegen, host-codegen, and execution-backend registries with unified ChangesCentralized backend architecture
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Target
participant Registry
participant Backend
participant JITKernel
participant Cache
Target->>Registry: resolve target backend
Registry->>Backend: validate availability and callbacks
Backend-->>Target: lower and generate host/device code
JITKernel->>Backend: resolve execution-backend specification
JITKernel->>Cache: dispatch using resolved adapter
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
13 tasks done
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
tilelang/engine/lower.py (1)
132-139:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInclude LLVM in the CPU device split.
Line 132 still derives the host/device filters from
is_cpu_device_backend(target), but that helper only returnsTruefor"c". The new registry path also routes"llvm"through the CPU backend, so LLVM CPU kernels will fall back to the genericDEVICE_KERNEL_LAUNCHcheck and can be filtered intohost_modinstead ofdevice_mod. This also affectstilelang/jit/adapter/utils.py, which imports the same helper.def is_cpu_device_backend(target: Target): return target.kind.name in {"c", "llvm"}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/engine/lower.py` around lines 132 - 139, The host/device split is using is_cpu_device_backend(target) but that helper currently only flags "c" and must also include "llvm" so LLVM CPU kernels are treated as host-side; update the is_cpu_device_backend implementation (referenced by get_host_call and get_device_call invocations in lower.py) to return True for target.kind.name in {"c", "llvm"}, and propagate the same change to any other callers (e.g., tilelang/jit/adapter/utils.py) so backend.resolve_backend and backend.lower continue to route LLVM through the CPU path and the host_mod/device_mod filtering behaves correctly.tilelang/jit/kernel.py (2)
451-460:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUse the resolved adapter key in source helpers, not the execution-backend name.
These branches still key off
self.execution_backend, but the new abstraction explicitly allowsExecutionBackendSpec.name != ExecutionBackendSpec.adapter(the new registry tests already exercise that shape). In that case, a backend likeExecutionBackendSpec("fast", adapter="tvm_ffi")will compile correctly and then take the wrong branch here when fetching kernel/host source. Compare againstself.execution_backend_spec.adapter(or a capability on the spec) instead.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/jit/kernel.py` around lines 451 - 460, The source helper methods get_kernel_source and get_host_source currently branch on self.execution_backend but must use the resolved adapter key from the backend spec; update the conditional checks to compare against self.execution_backend_spec.adapter (or a capability on self.execution_backend_spec) instead of self.execution_backend so backends whose .name differs from .adapter (e.g. ExecutionBackendSpec(name="fast", adapter="tvm_ffi")) correctly route to self.adapter.get_kernel_source/get_host_source; leave the fallback to self.artifact.kernel_source unchanged.
350-403:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHandle the
torchadapter in database reconstruction.Per this PR, Metal can now resolve to a
torchexecution backend, but_create_adapter_from_database()still rejectsadapter_name == "torch"outright.JITKernel.from_database(..., execution_backend="torch", target="metal")will now deterministically fail with the fallbackValueError, even though the normal compile path supports that adapter. Either add theMetalKernelAdapterreconstruction path here or reject that combination earlier before caching/loading it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/jit/kernel.py` around lines 350 - 403, The _create_adapter_from_database branch rejects adapter_name=="torch" but Metal-backed JITs resolve to the "torch" adapter; update _create_adapter_from_database to handle adapter_name == "torch" by constructing the adapter via MetalKernelAdapter.from_database with the same params/result_idx/target/func_or_mod/host_kernel_source/device_kernel_source/kernel_lib_path/pass_configs/compile_flags signature used for other adapters (import MetalKernelAdapter from the appropriate module where adapters live), so JITKernel.from_database(..., execution_backend="torch", target="metal") can reconstruct the adapter from the DB; alternatively, if you prefer rejecting earlier, add a clear early check in JITKernel.from_database to raise a descriptive error before caching/loading instead of falling through to the ValueError.tilelang/backend/README.md (1)
152-173:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winInclude
backend.pyin the backend package layout example.This tree omits the new backend-owned
backend.pymodules even though the rest of the document describes them as the entry point for descriptor registration. That makes the package layout section misleading for contributors using it as the template for new backends.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/backend/README.md` around lines 152 - 173, The README backend package layout examples omit the new backend-owned backend.py entry points; update each example tree to include backend.py (e.g., add backend.py under tilelang/cuda/, tilelang/rocm/, tilelang/cpu/, and tilelang/metal/) so the documented package layout matches the described entry-point used for descriptor registration (refer to the backend.py module name when editing the examples).
🧹 Nitpick comments (2)
tilelang/cuda/backend.py (1)
152-159: ⚡ Quick winReuse the shared cutedsl target predicate here.
These hooks reimplement target classification with raw key checks instead of
_is_cutedsl_target(). Reusing the helper keeps codegen hook selection aligned with the backend’s centralized target policy.Suggested refactor
def cuda_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + global_func = "target.build.tilelang_" + ("cutedsl" if _is_cutedsl_target(target) else "cuda") return build_device_with_global_func(device_mod, target, global_func) def cuda_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: - global_func = "target.build.tilelang_" + ("cutedsl" if "cutedsl" in target.keys else "cuda") + "_without_compile" + global_func = ( + "target.build.tilelang_" + + ("cutedsl" if _is_cutedsl_target(target) else "cuda") + + "_without_compile" + ) return build_device_with_global_func(device_mod, target, global_func)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/cuda/backend.py` around lines 152 - 159, The two hooks cuda_device_codegen and cuda_device_codegen_without_compile re-check target type by inspecting target.keys; replace that raw check with the centralized predicate _is_cutedsl_target(target) so both functions pick the same cutedsl vs cuda suffix as other backends; construct the global_func string the same way but use _is_cutedsl_target(target) to decide between "cutedsl" and "cuda" before passing it to build_device_with_global_func.tilelang/backend/registry.py (1)
73-76: ⚡ Quick winMake callback registration one-shot.
resolve_backend()is now used during lowering, device codegen, and host preprocessing, so the same backend can hitensure_callbacks_registered()several times in one compile. If a registrar does non-idempotent TVM global registration, this will double-register callbacks or repeat expensive setup. Cache successful registrations by backend name.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/backend/registry.py` around lines 73 - 76, The call site in resolve_backend() calls backend.ensure_callbacks_registered() every time which can re-run non-idempotent registrar work; make callback registration one-shot by caching successful registrations keyed by backend name (or a per-backend flag). Modify resolve_backend() (and/or ensure_callbacks_registered()) to check a module-level set (e.g., registered_backends) or an attribute on the backend object before invoking ensure_callbacks_registered(), add the backend.name to that cache after successful registration, and skip calling ensure_callbacks_registered() on subsequent resolves for the same backend.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tilelang/backend/backend.py`:
- Line 64: Backend is frozen but its nested mappings remain mutable; wrap both
the features and execution_backends mappings in an immutable mapping before
storing so callers can't mutate backend.features or backend.execution_backends.
Import and use types.MappingProxyType (or another immutable mapping wrapper) for
the dataclass default_factory for features and when constructing Backend
instances (the code paths around the Backend dataclass definition and the
registry creation code that sets execution_backends at the block around the
creation lines referenced). Ensure default_factory returns an immutable mapping
(e.g., MappingProxyType({})) and any dict passed into the Backend constructor is
wrapped with MappingProxyType(...) so get_backend() and list_backends() return
truly immutable descriptors.
In `@tilelang/backend/common.py`:
- Around line 15-16: The function webgpu_device_codegen_without_compile
currently calls build_device_with_global_func(device_mod, target,
"target.build.webgpu") which performs device compilation despite the default
policy advertising enable_device_compile=False; either change this function to a
true no-compile codegen path (e.g., return the device module unchanged or call
the appropriate non-compiling helper) or update the default policy to set
enable_device_compile=True so the descriptor matches behavior; update the same
logic in the other affected functions/lines (the block around 19-30) and ensure
references to build_device_with_global_func and enable_device_compile are
consistent.
In `@tilelang/backend/registry.py`:
- Around line 20-26: When overriding a backend in registry.py the code removes
the backend name from _TARGET_INDEX but leaves empty lists in the dict, causing
_ensure_loaded_for_kind to treat the kind as already loaded and skip lazy
imports; update the removal logic that iterates old.target_kinds (and the
analogous block around lines 49-54) to delete the key from _TARGET_INDEX if
names becomes empty (i.e., if not names: del _TARGET_INDEX[kind]) so that
_ensure_loaded_for_kind will still attempt lazy import for that kind.
In `@tilelang/cuda/backend.py`:
- Around line 34-37: The helper _is_nvrtc_available currently returns
bool(is_nvrtc_available) which truth-tests the imported function object instead
of executing the probe; change it to call the probe and return its boolean
result by invoking is_nvrtc_available() and casting/returning that value. Update
the function _is_nvrtc_available to import is_nvrtc_available and return the
result of is_nvrtc_available() (or bool(is_nvrtc_available()) for explicitness)
so the NVRTC backend is only advertised when the probe reports availability.
In `@tilelang/metal/backend.py`:
- Around line 12-13: metal_device_codegen currently calls
build_device_with_global_func(device_mod, target, "target.build.tilelang_metal")
which is also reused for the "without compile" path, causing backends that
should skip device compilation to still invoke the Metal build; either implement
a true no-compile path (e.g., add a device_codegen_without_compile variant that
returns source-only IRModule or a thin wrapper that avoids calling
build_device_with_global_func) or explicitly mark this execution backend as
enable_device_compile=True so the contract is honored; update
metal_device_codegen (and the code path that reuses it for the without-compile
flow) to use the new no-compile hook or set the backend flag accordingly.
In `@tilelang/rocm/backend.py`:
- Around line 9-37: The eager top-level import of tilelang.contrib.hipcc causes
its module-level callback to register on import; to keep registration lazy,
remove the module-level import and perform a local import of hipcc where needed
(either inside _tilelang_callback_hip_compile or at the start of
register_rocm_callbacks) so the hipcc module is only imported when the ROCm
backend is actually resolved; ensure _tilelang_callback_hip_compile still calls
hipcc.compile_hip and register_rocm_callbacks still calls
tvm_ffi.register_global_func("tilelang_callback_hip_compile",
f=_tilelang_callback_hip_compile, override=True).
---
Outside diff comments:
In `@tilelang/backend/README.md`:
- Around line 152-173: The README backend package layout examples omit the new
backend-owned backend.py entry points; update each example tree to include
backend.py (e.g., add backend.py under tilelang/cuda/, tilelang/rocm/,
tilelang/cpu/, and tilelang/metal/) so the documented package layout matches the
described entry-point used for descriptor registration (refer to the backend.py
module name when editing the examples).
In `@tilelang/engine/lower.py`:
- Around line 132-139: The host/device split is using
is_cpu_device_backend(target) but that helper currently only flags "c" and must
also include "llvm" so LLVM CPU kernels are treated as host-side; update the
is_cpu_device_backend implementation (referenced by get_host_call and
get_device_call invocations in lower.py) to return True for target.kind.name in
{"c", "llvm"}, and propagate the same change to any other callers (e.g.,
tilelang/jit/adapter/utils.py) so backend.resolve_backend and backend.lower
continue to route LLVM through the CPU path and the host_mod/device_mod
filtering behaves correctly.
In `@tilelang/jit/kernel.py`:
- Around line 451-460: The source helper methods get_kernel_source and
get_host_source currently branch on self.execution_backend but must use the
resolved adapter key from the backend spec; update the conditional checks to
compare against self.execution_backend_spec.adapter (or a capability on
self.execution_backend_spec) instead of self.execution_backend so backends whose
.name differs from .adapter (e.g. ExecutionBackendSpec(name="fast",
adapter="tvm_ffi")) correctly route to
self.adapter.get_kernel_source/get_host_source; leave the fallback to
self.artifact.kernel_source unchanged.
- Around line 350-403: The _create_adapter_from_database branch rejects
adapter_name=="torch" but Metal-backed JITs resolve to the "torch" adapter;
update _create_adapter_from_database to handle adapter_name == "torch" by
constructing the adapter via MetalKernelAdapter.from_database with the same
params/result_idx/target/func_or_mod/host_kernel_source/device_kernel_source/kernel_lib_path/pass_configs/compile_flags
signature used for other adapters (import MetalKernelAdapter from the
appropriate module where adapters live), so JITKernel.from_database(...,
execution_backend="torch", target="metal") can reconstruct the adapter from the
DB; alternatively, if you prefer rejecting earlier, add a clear early check in
JITKernel.from_database to raise a descriptive error before caching/loading
instead of falling through to the ValueError.
---
Nitpick comments:
In `@tilelang/backend/registry.py`:
- Around line 73-76: The call site in resolve_backend() calls
backend.ensure_callbacks_registered() every time which can re-run non-idempotent
registrar work; make callback registration one-shot by caching successful
registrations keyed by backend name (or a per-backend flag). Modify
resolve_backend() (and/or ensure_callbacks_registered()) to check a module-level
set (e.g., registered_backends) or an attribute on the backend object before
invoking ensure_callbacks_registered(), add the backend.name to that cache after
successful registration, and skip calling ensure_callbacks_registered() on
subsequent resolves for the same backend.
In `@tilelang/cuda/backend.py`:
- Around line 152-159: The two hooks cuda_device_codegen and
cuda_device_codegen_without_compile re-check target type by inspecting
target.keys; replace that raw check with the centralized predicate
_is_cutedsl_target(target) so both functions pick the same cutedsl vs cuda
suffix as other backends; construct the global_func string the same way but use
_is_cutedsl_target(target) to decide between "cutedsl" and "cuda" before passing
it to build_device_with_global_func.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 376e089b-d168-482e-9efe-5f00e6b9cab6
📒 Files selected for processing (19)
testing/python/backend/test_tilelang_backend_registry.pytilelang/backend/README.mdtilelang/backend/__init__.pytilelang/backend/backend.pytilelang/backend/codegen.pytilelang/backend/common.pytilelang/backend/registry.pytilelang/cpu/__init__.pytilelang/cpu/backend.pytilelang/cuda/__init__.pytilelang/cuda/backend.pytilelang/engine/lower.pytilelang/jit/adapter/utils.pytilelang/jit/execution_backend.pytilelang/jit/kernel.pytilelang/metal/__init__.pytilelang/metal/backend.pytilelang/rocm/__init__.pytilelang/rocm/backend.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@testing/python/backend/test_tilelang_backend_registry.py`:
- Around line 50-61: The test only checks that views are read-only but not that
Backend makes defensive copies of input mappings; update the
test_backend_descriptor_freezes_nested_mappings to keep references to the
original dicts passed into Backend (e.g., features_input and
execution_backends_input), mutate those originals after constructing Backend
(add/remove keys or replace values), and assert backend.features and
backend.execution_backends do not reflect those mutations; if the Backend
implementation is missing defensive copies, modify Backend's constructor to copy
incoming mappings (e.g., copy/deepcopy features and execution_backends) so
external changes to the original dicts cannot mutate the backend descriptor.
- Around line 146-161: The test bypasses the new adapter-key resolution by
directly assigning kernel.adapter and leaving kernel.artifact None; change it to
not inject kernel.adapter and instead set kernel.artifact to contain sources
keyed by the resolved adapter name ("tvm_ffi") so get_kernel_source and
get_host_source exercise the "fast" -> "tvm_ffi" resolution path: leave
execution_backend_spec = ExecutionBackendSpec("fast", adapter="tvm_ffi"), remove
or avoid setting kernel.adapter, set kernel.artifact to a structure that
provides kernel and host sources under the "tvm_ffi" key, then assert
kernel.get_kernel_source(kernel_only=True) and kernel.get_host_source() still
return the expected values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 90e63e1f-9701-47e0-88b9-23e79dc16621
📒 Files selected for processing (8)
testing/python/backend/test_tilelang_backend_registry.pytilelang/backend/README.mdtilelang/backend/backend.pytilelang/backend/registry.pytilelang/cuda/backend.pytilelang/engine/lower.pytilelang/jit/kernel.pytilelang/rocm/backend.py
🚧 Files skipped from review as they are similar to previous changes (7)
- tilelang/rocm/backend.py
- tilelang/backend/README.md
- tilelang/engine/lower.py
- tilelang/jit/kernel.py
- tilelang/backend/backend.py
- tilelang/cuda/backend.py
- tilelang/backend/registry.py
Comment on lines
+50
to
+61
| def test_backend_descriptor_freezes_nested_mappings(): | ||
| backend = Backend( | ||
| "unit-test", | ||
| ("c",), | ||
| features={"feature": lambda target: True}, | ||
| execution_backends={"fast": ExecutionBackendSpec("fast")}, | ||
| ) | ||
|
|
||
| with pytest.raises(TypeError): | ||
| backend.features["other"] = lambda target: False | ||
| with pytest.raises(TypeError): | ||
| backend.execution_backends["slow"] = ExecutionBackendSpec("slow") |
Contributor
There was a problem hiding this comment.
Verify defensive copies, not just read-only views.
This only proves writes through backend.features / backend.execution_backends fail. If Backend keeps a proxy over the caller-owned dict, mutating the original mapping after construction still changes the descriptor and this test stays green. Keep handles to the input dicts and assert post-construction mutations do not leak into backend.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/backend/test_tilelang_backend_registry.py` around lines 50 -
61, The test only checks that views are read-only but not that Backend makes
defensive copies of input mappings; update the
test_backend_descriptor_freezes_nested_mappings to keep references to the
original dicts passed into Backend (e.g., features_input and
execution_backends_input), mutate those originals after constructing Backend
(add/remove keys or replace values), and assert backend.features and
backend.execution_backends do not reflect those mutations; if the Backend
implementation is missing defensive copies, modify Backend's constructor to copy
incoming mappings (e.g., copy/deepcopy features and execution_backends) so
external changes to the original dicts cannot mutate the backend descriptor.
Comment on lines
+146
to
+161
| def test_jit_source_helpers_use_resolved_adapter_key(): | ||
| class FakeAdapter: | ||
| def get_kernel_source(self, kernel_only=True): | ||
| return f"kernel:{kernel_only}" | ||
|
|
||
| def get_host_source(self): | ||
| return "host" | ||
|
|
||
| kernel = JITKernel.__new__(JITKernel) | ||
| kernel.execution_backend = "fast" | ||
| kernel.execution_backend_spec = ExecutionBackendSpec("fast", adapter="tvm_ffi") | ||
| kernel.adapter = FakeAdapter() | ||
| kernel.artifact = None | ||
|
|
||
| assert kernel.get_kernel_source(kernel_only=True) == "kernel:True" | ||
| assert kernel.get_host_source() == "host" |
Contributor
There was a problem hiding this comment.
This bypasses the adapter-key lookup path the PR is changing.
Because kernel.adapter is injected directly and kernel.artifact is None, these assertions still pass if the helpers ignore execution_backend_spec.adapter entirely. To catch regressions in the new "fast" -> "tvm_ffi" resolution path, populate kernel.artifact with sources keyed only by the resolved adapter name and assert get_kernel_source() / get_host_source() still succeed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/backend/test_tilelang_backend_registry.py` around lines 146 -
161, The test bypasses the new adapter-key resolution by directly assigning
kernel.adapter and leaving kernel.artifact None; change it to not inject
kernel.adapter and instead set kernel.artifact to contain sources keyed by the
resolved adapter name ("tvm_ffi") so get_kernel_source and get_host_source
exercise the "fast" -> "tvm_ffi" resolution path: leave execution_backend_spec =
ExecutionBackendSpec("fast", adapter="tvm_ffi"), remove or avoid setting
kernel.adapter, set kernel.artifact to a structure that provides kernel and host
sources under the "tvm_ffi" key, then assert
kernel.get_kernel_source(kernel_only=True) and kernel.get_host_source() still
return the expected values.
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tilelang/backend/README.md (1)
79-85: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the WebGPU descriptor location.
The layout shows
tilelang/<backend>/backend.pyfor every backend. This PR defineswebgpu_backendintilelang/backend/common.py. Add a WebGPU exception to the layout, or move the descriptor totilelang/webgpu/backend.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tilelang/backend/README.md` around lines 79 - 85, Update the backend layout documentation to explicitly identify the WebGPU descriptor exception in tilelang/backend/common.py, or relocate webgpu_backend to tilelang/webgpu/backend.py and update the layout accordingly; ensure the documented location matches the implementation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@testing/python/jit/test_tilelang_jit_diagnostics.py`:
- Around line 197-200: Disable CUDA binary caching around the
_tilelang_callback_cuda_compile invocation in this test so CUDABinaryCache.load
cannot bypass fake_compile_cuda. Preserve the existing fake compiler assertion
and target setup, ensuring captured["target_format"] is populated
deterministically.
In `@tilelang/backend/pipeline.py`:
- Around line 95-96: Update the error message string near the format validation
in the pipeline code so there is a space before the example parenthesis and the
period follows the closing parenthesis, producing correct punctuation without
changing the message content.
---
Nitpick comments:
In `@tilelang/backend/README.md`:
- Around line 79-85: Update the backend layout documentation to explicitly
identify the WebGPU descriptor exception in tilelang/backend/common.py, or
relocate webgpu_backend to tilelang/webgpu/backend.py and update the layout
accordingly; ensure the documented location matches the implementation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: d29be022-fe62-4586-ac34-839f33745e1d
📒 Files selected for processing (47)
testing/python/backend/test_tilelang_backend_registry.pytesting/python/cache/test_tilelang_cuda_binary_cache.pytesting/python/debug/test_lower_trace.pytesting/python/jit/test_tilelang_jit_diagnostics.pytesting/python/llvm/test_tilelang_llvm_while_repro.pytesting/python/target/test_tilelang_target.pytesting/python/transform/test_tilelang_transform_verify_parallel_loop.pytesting/python/webgpu/test_webgpu_codegen.pytilelang/backend/README.mdtilelang/backend/__init__.pytilelang/backend/backend.pytilelang/backend/codegen.pytilelang/backend/common.pytilelang/backend/device_codegen.pytilelang/backend/execution_backend.pytilelang/backend/host_codegen.pytilelang/backend/pass_pipeline/__init__.pytilelang/backend/pass_pipeline/pipeline.pytilelang/backend/pipeline.pytilelang/cache/__init__.pytilelang/contrib/hipcc.pytilelang/cpu/__init__.pytilelang/cpu/backend.pytilelang/cpu/codegen.pytilelang/cpu/execution_backend.pytilelang/cpu/pipeline.pytilelang/cuda/__init__.pytilelang/cuda/backend.pytilelang/cuda/codegen.pytilelang/cuda/execution_backend.pytilelang/cuda/pipeline.pytilelang/engine/lower.pytilelang/jit/kernel.pytilelang/metal/__init__.pytilelang/metal/backend.pytilelang/metal/codegen.pytilelang/metal/execution_backend.pytilelang/metal/pipeline.pytilelang/rocm/__init__.pytilelang/rocm/backend.pytilelang/rocm/codegen.pytilelang/rocm/execution_backend.pytilelang/rocm/pipeline.pytilelang/tools/lower_trace/core.pytilelang/webgpu/__init__.pytilelang/webgpu/codegen.pytilelang/webgpu/pipeline.py
💤 Files with no reviewable changes (19)
- tilelang/webgpu/codegen.py
- tilelang/metal/execution_backend.py
- tilelang/backend/pass_pipeline/init.py
- tilelang/rocm/codegen.py
- tilelang/cuda/codegen.py
- tilelang/rocm/execution_backend.py
- tilelang/backend/host_codegen.py
- tilelang/rocm/init.py
- tilelang/cpu/execution_backend.py
- tilelang/metal/init.py
- tilelang/contrib/hipcc.py
- tilelang/backend/pass_pipeline/pipeline.py
- tilelang/cuda/init.py
- tilelang/metal/codegen.py
- tilelang/cuda/execution_backend.py
- tilelang/backend/init.py
- tilelang/backend/device_codegen.py
- tilelang/cpu/codegen.py
- tilelang/cpu/init.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tilelang/rocm/backend.py
- tilelang/cuda/backend.py
- tilelang/engine/lower.py
- tilelang/jit/kernel.py
Comment on lines
+197
to
+200
| monkeypatch.setattr(cuda_backend.nvcc, "compile_cuda", fake_compile_cuda) | ||
|
|
||
| target = determine_target({"kind": "cuda", "arch": "sm_100f", "code": ["sm_100a", "sm_103a"]}, return_object=True) | ||
| lower.tilelang_callback_cuda_compile("__global__ void kernel() {}", target) | ||
| cuda_backend._tilelang_callback_cuda_compile("__global__ void kernel() {}", target) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline testing/python/jit/test_tilelang_jit_diagnostics.py --items all
rg -n -C 4 'CUDABinaryCache|TILELANG_DISABLE_CACHE|enable_cache|disable_cache|test_cuda_compile_callback_uses_fatbin' \
testing/python/jit/test_tilelang_jit_diagnostics.py tilelang/cacheRepository: tile-ai/tilelang
Length of output: 6138
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== test snippet ==="
sed -n '161,206p' testing/python/jit/test_tilelang_jit_diagnostics.py
echo "=== cuda_binary_cache relevant methods ==="
sed -n '1,180p' tilelang/cache/cuda_binary_cache.py
echo "=== tilelang cache disable/enable symbols and usages ==="
rg -n 'disable_cache|enable_cache|CUDABinaryCache|_tilelang_callback_cuda_compile|compile_cuda' tilelang testing/python/jit/test_tilelang_jit_diagnostics.pyRepository: tile-ai/tilelang
Length of output: 11356
Disable CUDA binary caching for this fake-NVCC assertion.
_tilelang_callback_cuda_compile() calls CUDABinaryCache.load() before running the patched compiler. If a matching cache file exists, fake_compile_cuda() is skipped and captured["target_format"] raises KeyError; disable the cache around this callback call unless this test intentionally covers the cache path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@testing/python/jit/test_tilelang_jit_diagnostics.py` around lines 197 - 200,
Disable CUDA binary caching around the _tilelang_callback_cuda_compile
invocation in this test so CUDABinaryCache.load cannot bypass fake_compile_cuda.
Preserve the existing fake compiler assertion and target setup, ensuring
captured["target_format"] is populated deterministically.
Comment on lines
+95
to
+96
| "You can choose one of the valid formats or a comma-separated list of " | ||
| "formats.(e.g. 'txt,png,pdf')" |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the error-message punctuation.
The message renders formats.(e.g. 'txt,png,pdf'). Add a space before the parenthesis and place the period after the example.
Proposed fix
- "formats.(e.g. 'txt,png,pdf')"
+ "formats (e.g. 'txt,png,pdf')."📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| "You can choose one of the valid formats or a comma-separated list of " | |
| "formats.(e.g. 'txt,png,pdf')" | |
| "You can choose one of the valid formats or a comma-separated list of " | |
| "formats (e.g. 'txt,png,pdf')." |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tilelang/backend/pipeline.py` around lines 95 - 96, Update the error message
string near the format validation in the pipeline code so there is a space
before the example parenthesis and the period follows the closing parenthesis,
producing correct punctuation without changing the message content.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Tracking Issue: #2115
This PR centralizes backend resolution around a
Backenddescriptor so target-level capability, codegen, pass pipeline, and JIT execution policy are resolved from one backend-owned entry point.Summary
Backenddescriptor and registry as the source of truth for target matching, pass pipelines, codegen hooks, callbacks, feature checks, and execution-backend policy.PassPipelineentry points while exposing them throughBackend.pipeline.Summary
Backenddescriptors and a global registry.Backend.pipeline.