Skip to content

[BugFix] Fix JITKernel export_library bug #1699

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 1 commit into from
Jan 20, 2026
Merged

[BugFix] Fix JITKernel export_library bug #1699

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
13 changes: 11 additions & 2 deletions tilelang/jit/kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -637,8 +637,17 @@ def export_library(self, kernel_file: str) -> None:
# rt_module: use export_library to export
# rt_params: use cloudpickle to serialize

# Export the compiled kernel function to a shared library file.
self.rt_module.export_library(kernel_file)
if self.artifact is None or self.artifact.rt_mod is None:
raise AttributeError(
'Runtime module is not available. Please compile the kernel with `execution_backend="tvm_ffi"` before exporting.'
)

dir_path = os.path.dirname(kernel_file)
if dir_path:
os.makedirs(dir_path, exist_ok=True)

self.artifact.rt_mod.export_library(kernel_file)
Comment on lines +645 to +649
Copy link
Contributor

@coderabbitai coderabbitai bot Jan 20, 2026

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟡 Minor

Validate kernel_file before path ops

export_ptx/export_sass reject empty paths; export_library should do the same to avoid confusing export errors with an empty string.

✅ Suggested guard 
     def export_library(self, kernel_file: str) -> None:
@@
-        dir_path = os.path.dirname(kernel_file)
+        if not kernel_file:
+            raise ValueError("kernel_file must be provided to export library")
+
+        dir_path = os.path.dirname(kernel_file)
📝 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
dir_path = os.path.dirname(kernel_file)
if dir_path:
os.makedirs(dir_path, exist_ok=True)
self.artifact.rt_mod.export_library(kernel_file)
if not kernel_file:
raise ValueError("kernel_file must be provided to export library")
dir_path = os.path.dirname(kernel_file)
if dir_path:
os.makedirs(dir_path, exist_ok=True)
self.artifact.rt_mod.export_library(kernel_file)
🤖 Prompt for AI Agents 
In `@tilelang/jit/kernel.py` around lines 645 - 649, The code should validate the
kernel_file before performing path operations and calling
self.artifact.rt_mod.export_library(kernel_file); add a guard in the function
containing kernel_file (check that kernel_file is a non-empty, non-whitespace
string) and raise a clear ValueError (consistent with export_ptx/export_sass
behavior) if invalid, then proceed with os.path.dirname/kernel dir creation and
the call to export_library; reference kernel_file and
self.artifact.rt_mod.export_library to locate the change.

logger.info(f"Kernel library exported to {os.path.abspath(kernel_file)}")

def _get_ptx(self, verbose: bool | None = None) -> str:
"""
Expand Down
Loading