Skip to content

[Cache] Introduce detailed target information for the disk kernel cache #780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
LeiWang1999 merged 6 commits into from
Sep 2, 2025
Merged

[Cache] Introduce detailed target information for the disk kernel cache #780

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
287bf3a
Fix type hint for target_host parameter in compile function to allow …
LeiWang1999 Sep 2, 2025
5f44aef
Refactor target handling in compile function to utilize determine_tar…
LeiWang1999 Sep 2, 2025
b2e85c0
Update PrintConst function in codegen_cuda.cc to use hexfloat format …
LeiWang1999 Sep 2, 2025
aeff94d
Refactor PrintType function in codegen_cuda.cc to remove unnecessary …
LeiWang1999 Sep 2, 2025
200e366
Enhance benchmark_matmul.py to conditionally print Reference TFlops o…
LeiWang1999 Sep 2, 2025
0b4dc07
Remove automatic commit and push step from AMD and NVIDIA CI workflow…
LeiWang1999 Sep 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion tilelang/jit/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def compile(
out_idx: Union[List[int], int, None] = None,
execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
target: Union[str, Target] = "auto",
target_host: Union[str, Target] = None,
target_host: Union[str, Target, None] = None,
verbose: bool = False,
pass_configs: Optional[Dict[str, Any]] = None,
compile_flags: Optional[Union[List[str], str]] = None,
Expand Down Expand Up @@ -69,6 +69,10 @@ def compile(
assert isinstance(func, PrimFunc), f"target function must be a PrimFunc but got {type(func)}"
if isinstance(compile_flags, str):
compile_flags = [compile_flags]

# This path is not a performance critical path, so we can afford to convert the target.
target, target_host = Target(target), Target(target_host) if target_host else None

@gemini-code-assist gemini-code-assist Bot Sep 2, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

This line will cause a runtime error when target has its default value of "auto". The tvm.target.Target constructor does not support the "auto" keyword and will raise an exception. The logic to resolve "auto" into a concrete target (e.g., "cuda" or "hip") needs to be executed before creating the Target object. This logic appears to exist in the determine_target utility function, which should be used here. A similar issue exists for target_host if it is set to "auto".

Note: You will need to add from tilelang.utils.target import determine_target to the file's imports for the suggestion to work.

Suggested change
target, target_host = Target(target), Target(target_host) if target_host else None
target, target_host = determine_target(target, return_object=True), determine_target(target_host, return_object=True) if target_host else None


Comment on lines +72 to +76

@coderabbitai coderabbitai Bot Sep 2, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Do not convert the 'auto' sentinel into a TVM Target; this will fail.

Target("auto") is not a valid TVM target string and will raise at runtime. Gate the conversion and only normalize when the input is not the 'auto' sentinel and not already a Target. Also, independently gate target_host.

Apply:

-    # This path is not a performance critical path, so we can afford to convert the target.
-    target, target_host = Target(target), Target(target_host) if target_host else None
+    # Normalize targets for caching while preserving the 'auto' sentinel.
+    if not isinstance(target, Target) and not (isinstance(target, str) and target == "auto"):
+        target = Target(target)
+    if target_host is not None and not isinstance(target_host, Target):
+        target_host = Target(target_host)

Optionally, to centralize validation/coercion, consider using tilelang.language.ast.ir.target(target, host=target_host) and skipping it when target == "auto".

📝 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
# This path is not a performance critical path, so we can afford to convert the target.
target, target_host = Target(target), Target(target_host) if target_host else None
# Normalize targets for caching while preserving the 'auto' sentinel.
if not isinstance(target, Target) and not (isinstance(target, str) and target == "auto"):
target = Target(target)
if target_host is not None and not isinstance(target_host, Target):
target_host = Target(target_host)
🤖 Prompt for AI Agents 
In tilelang/jit/__init__.py around lines 72 to 75, the code converts both target
and target_host unconditionally which will raise for the sentinel "auto"
(Target("auto") is invalid); change the logic to only call Target(...) when the
value is not the string "auto" and is not already a Target instance, and apply
the same guarded conversion separately for target_host; optionally, replace the
ad-hoc conversion with a call to tilelang.language.ast.ir.target(target,
host=target_host) but skip calling it when target == "auto".

return cached(
func=func,
out_idx=out_idx,
Expand Down
Loading