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

Show file tree
Hide file tree
Changes from 5 commits
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
3 changes: 2 additions & 1 deletion benchmark/matmul/benchmark_matmul.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,4 +243,5 @@ def main(
print(f"Best TFlops: {total_flops / best_latency * 1e-9:.3f}")
print(f"Best config: {best_config}")

print(f"Reference TFlops: {total_flops / ref_latency * 1e-9:.3f}")
if ref_latency is not None:
print(f"Reference TFlops: {total_flops / ref_latency * 1e-9:.3f}")
15 changes: 8 additions & 7 deletions src/target/codegen_cuda.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -325,16 +325,12 @@ void CodeGenTileLangCUDA::PrintType(DataType t, std::ostream &os) { // NOLINT(*)
enable_fp6_ = true;
if (t.lanes() <= 4) {
os << GetFP6Type(t);
} else {
fail = true;
}
return;
} else if (t.is_float4()) {
enable_fp4_ = true;
if (t.lanes() <= 4) {
os << GetFP4Type(t);
} else {
fail = true;
}
return;
} else if (t == DataType::Bool()) {
Expand Down Expand Up @@ -1960,13 +1956,17 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os,
// Type code is kBFloat
if (op->dtype.is_bfloat16()) {
os << "bfloat16_t";
os << '(' << std::scientific << op->value << 'f' << ')';
os << '(' << std::hexfloat << op->value << 'f';
os << "/*" << std::scientific << op->value << "*/";
os << ')';
Comment on lines +1959 to +1961

@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.

💡 Verification agent

🧩 Analysis chain

Restore stream format flags after hex/scientific emission

std::hexfloat and std::scientific persist on the stream and can leak into later emissions. Also add a separating space before the comment for readability.

-    os << '(' << std::hexfloat << op->value << 'f';
-    os << "/*" << std::scientific << op->value << "*/";
+    os << '(' << std::hexfloat << op->value << 'f' << std::defaultfloat;
+    os << " /*" << std::scientific << op->value << std::defaultfloat << "*/";

Repeat the same change in the float8/float4 branch.

Also applies to: 1967-1969

🏁 Script executed:

#!/bin/bash
set -e
# Locate the CUDA code generator file
FILE=$(fd --full-path codegen_cuda.cc)
echo "Inspecting $FILE"
# Display lines around the reported snippet (1959–1961)
sed -n '1950,1970p' "$FILE"
# Search for any existing stream-flag restorations (defaultfloat) and hexfloat usages
rg -nP '\b(hexfloat|scientific|defaultfloat)\b' -C3 "$FILE" || true

Length of output: 1897

Restore default float formatting after hex/scientific and add space before comment

std::hexfloat and std::scientific persist on the stream and will affect later emissions—append std::defaultfloat to reset formatting and insert a space before the comment. Apply in both the bfloat16 (≈1959–1961) and float8/float4 (≈1967–1969) branches:

-    os << '(' << std::hexfloat << op->value << 'f';
-    os << "/*" << std::scientific << op->value << "*/";
+    os << '(' << std::hexfloat << op->value << 'f' << std::defaultfloat;
+    os << " /*" << std::scientific << op->value << std::defaultfloat << "*/";
     os << ')';
🤖 Prompt for AI Agents 
In src/target/codegen_cuda.cc around lines 1959–1961 (bfloat16 branch) and
1967–1969 (float8/float4 branch), the stream manipulators std::hexfloat and
std::scientific are left set which will affect subsequent output; update each
sequence so that after emitting the hexfloat literal and the scientific comment
you emit std::defaultfloat to restore the stream formatting, and insert a space
before the comment (i.e. " /*...*/" instead of "/*...*/") so the output is
spaced cleanly. Ensure both branches are changed consistently.

return;
}
// Type code is kFloat8_e5m2 or kE4M4Float
if (op->dtype.is_float8() || op->dtype.is_float4()) {
p->PrintType(op->dtype, os);
os << '(' << std::scientific << op->value << 'f' << ')';
os << '(' << std::hexfloat << op->value << 'f';
os << "/*" << std::scientific << op->value << "*/";
os << ')';
return;
}
// Type code is kFloat
Expand All @@ -1984,9 +1984,10 @@ inline void PrintConst(const FloatImmNode *op, std::ostream &os,
temp << ((op->dtype.bits() == 32) ? "CUDART_NAN_F" : "CUDART_NAN");
p->need_math_constants_h_ = true;
} else {
temp << std::scientific << op->value;
temp << std::hexfloat << op->value;
if (op->dtype.bits() == 32)
temp << 'f';
temp << "/*" << std::scientific << op->value << "*/";
}
p->MarkConst(temp.str());
os << temp.str();
Expand Down
2 changes: 1 addition & 1 deletion tilelang/autotuner/param.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ def __hash__(self):
"execution_backend":
self.execution_backend,
"target":
self.target,
str(self.target),
"target_host":
str(self.target_host) if self.target_host else None,
"verbose":
Expand Down
3 changes: 2 additions & 1 deletion tilelang/autotuner/tuner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
from tilelang import env
from tilelang.autotuner.param import CompileArgs, ProfileArgs, AutotuneResult
from tilelang.autotuner.capture import get_autotune_inputs
from tilelang.utils.target import determine_target
from tilelang.jit.param import _P, _RProg
from tilelang.version import __version__

Expand Down Expand Up @@ -150,7 +151,7 @@ def set_compile_args(self,
"""
self.compile_args = CompileArgs(
out_idx=out_idx,
target=target,
target=Target(determine_target(target)),
execution_backend=execution_backend,
target_host=target_host,
verbose=verbose,
Expand Down
7 changes: 6 additions & 1 deletion tilelang/jit/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from tvm.target import Target

from tilelang.jit.kernel import JITKernel
from tilelang.utils.target import determine_target
from tilelang.cache import cached
from os import path, makedirs
from logging import getLogger
Expand All @@ -34,7 +35,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 +70,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(determine_target(target))

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