[BugFix] Fix JITKernel export_library bug

[chengyupku]

This pull request improves the robustness and usability of the kernel export functionality in tilelang/jit/kernel.py. The main change is to add error handling and ensure the target directory exists before exporting the compiled kernel library.

Kernel export improvements:

• Added a check to raise an AttributeError if self.artifact or self.artifact.rt_mod is not available, with a clear error message instructing the user to compile the kernel with execution_backend="tvm_ffi" before exporting.
• Ensured the target directory for the kernel file exists by creating it if necessary before exporting the library.
• Added a log message to confirm successful export and show the absolute path of the exported kernel library.

Summary by CodeRabbit

• Bug Fixes
  • Improved kernel library export robustness with enhanced validation.
  • Automatically creates required output directories during export.
  • Enhanced logging for export operations with file path information.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds runtime module export validation and initialization to the kernel library export process. It verifies the compiled runtime module exists, creates the target directory recursively, exports the module artifact, and logs the export path.

Changes

Cohort / File(s)	Summary
Runtime Module Export tilelang/jit/kernel.py	Added validation to ensure compiled runtime module exists before export; creates target directory recursively and exports module artifact with informational logging.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• [Bugfix] Fix saving kernel source code where JITKernel.artifact is None #921: Handles JITKernel artifacts in the same file, addressing kernel source saving when artifact is missing—complements this PR's validation approach.

Suggested reviewers

• LeiWang1999

Poem

🐰 A kernel exported with validation so bright,
Directory paths created, everything right!
Runtime modules bundled with logs to compile,
The export enchantment completed in style! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately reflects the main change: fixing a bug in the JITKernel export_library function by adding runtime validation, directory creation, and logging.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@tilelang/jit/kernel.py`:
- Around line 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.

🧹 Nitpick comments (1)

tilelang/jit/kernel.py (1)

640-643: Optional: encapsulate the export precondition error

Ruff (TRY003) flags long inline messages; consider a small custom exception (or at least a module-level constant) so the message is centralized and reusable.

♻️ Possible refactor

+class RuntimeModuleUnavailableError(AttributeError):
+    def __init__(self) -> None:
+        super().__init__(
+            'Runtime module is not available. Please compile the kernel with `execution_backend="tvm_ffi"` before exporting.'
+        )
+
@@
-        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.'
-            )
+        if self.artifact is None or self.artifact.rt_mod is None:
+            raise RuntimeModuleUnavailableError()

[coderabbitai]

⚠️ 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.