[Backend] Add unified backend resolution policy#2318
Open
SiriusNEO wants to merge 2 commits into
Open
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
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (11)
💤 Files with no reviewable changes (11)
📝 WalkthroughWalkthroughThis PR centralizes backend resolution through new backend descriptors and a registry, updates engine lowering and JIT execution-backend selection to use them, and adds backend-owned implementations for CPU, CUDA, ROCm, Metal, and WebGPU. It also expands registry tests and updates the backend README. ChangesCentralized Backend Registry and Descriptor System
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
10 tasks
![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
Comment on lines
+15
to
+16
| def webgpu_device_codegen_without_compile(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: | ||
| return build_device_with_global_func(device_mod, target, "target.build.webgpu") |
Contributor
There was a problem hiding this comment.
The default WebGPU backend advertises “no device compile” but still builds the device module.
Line 29 sets enable_device_compile=False, yet webgpu_device_codegen_without_compile() still calls build_device_with_global_func(..., "target.build.webgpu"). So the default tvm_ffi policy claims to skip device compilation while the actual hook still performs it. Please either switch this to a real no-compile codegen path or set enable_device_compile=True so the descriptor matches reality.
Also applies to: 19-30
🤖 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/common.py` around lines 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.
Comment on lines
+12
to
+13
| def metal_device_codegen(device_mod: tvm.IRModule, target: Target) -> tvm.IRModule: | ||
| return build_device_with_global_func(device_mod, target, "target.build.tilelang_metal") |
Contributor
There was a problem hiding this comment.
device_codegen_without_compile still goes through the compile hook.
metal_device_codegen() calls build_device_with_global_func(..., "target.build.tilelang_metal"), and Line 26 reuses that same function for the "without compile" path. That means any execution backend that is supposed to skip device compilation still requires the Metal build step, which breaks the new execution-backend policy contract. Either add a true source-only/no-compile hook here, or mark the affected execution backend as enable_device_compile=True.
Also applies to: 20-35
🤖 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/metal/backend.py` around lines 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.
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.
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
1 participant
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 by CodeRabbit
New Features
Backenddescriptor + backend registry as the single source of truth for target matching, pass/pipeline wiring, codegen hooks, callback registration, feature checks, and JIT execution-backend policy.Refactor
resolve_backend(target)(instead of resolving pipelines/FFI paths in disparate places).Bug Fixes / Hardening
override=True) to correctly prune old target mappings and keep indices consistent.nvrtconly allowed for CUDA targets, plus CPU/LLVM validation).Documentation
Tests
features/execution_backends, correct pipeline/codegen behavior, override/pruning semantics, and JIT kernel source helper adapter-key behavior.