[Host] Provide post process to customize host code and enhance nullable check#1562
[Host] Provide post process to customize host code and enhance nullable check#1562LeiWang1999 merged 2 commits intotile-ai:mainfrom
Conversation
…e buffer handling - Introduced a new callback for C host code generation that allows for post-processing of generated code before it is wrapped into a CSourceModule. - Enhanced the ArgBinder to log shape variable sources and ensure safe binding of symbolic shape variables, preventing potential segmentation faults when dealing with nullable buffers. - Added a regression test to verify that a single buffer with a symbolic shape variable must be non-null, ensuring robustness against null inputs.
… generation and ArgBinder - Adjusted indentation in the C host code generation to enhance clarity. - Reformatted comments and code structure in ArgBinder for better readability, ensuring consistent style throughout the file. - Minor whitespace adjustments to maintain code consistency.
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughThe PR introduces a C host post-processing callback mechanism and adds runtime null-safety guards for shape variable binding. It exposes a new public API for registering C post-processing functions and includes a test validating null-safe shape binding for single-buffer nullable scenarios. Changes
Sequence Diagram(s)sequenceDiagram
participant CodeGen as BuildTileLangCHost
participant Registry as TVM Global<br/>Function Registry
participant Callback as User Callback<br/>Function
CodeGen->>CodeGen: Generate C code
CodeGen->>Registry: Lookup tilelang_callback_c_host_postproc
alt Callback registered
Registry-->>CodeGen: Return callback function
CodeGen->>Callback: Invoke(generated_code, target)
Callback-->>CodeGen: Return modified code
CodeGen->>CodeGen: Use post-processed code
else Callback not found
Registry-->>CodeGen: Return null/not found
CodeGen->>CodeGen: Use original code (fallback)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tilelang/engine/callback.py (1)
110-140: Decorator implementation is functionally correct.The decorator properly handles all three invocation patterns and mirrors the existing CUDA/HIP decorator implementations, ensuring API consistency.
📝 Optional type hint improvement
For stricter PEP 484 compliance, consider using
Optionalexplicitly:-def register_c_postproc_callback(func: Callable | bool = None, override: bool = True): +def register_c_postproc_callback(func: Optional[Callable] | bool = None, override: bool = True):Add the import at the top:
from __future__ import annotations -from typing import Callable +from typing import Callable, OptionalThis makes the nullable parameter more explicit, though the current implementation works correctly.
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
src/target/codegen_c_host.ccsrc/transform/arg_binder.cctesting/python/transform/test_nullable_buffer_params.pytilelang/__init__.pytilelang/engine/__init__.pytilelang/engine/callback.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
src/target/codegen_c_host.cctilelang/engine/__init__.pytilelang/__init__.pytilelang/engine/callback.py
🧬 Code graph analysis (3)
testing/python/transform/test_nullable_buffer_params.py (1)
tilelang/language/symbolics.py (1)
dynamic(10-21)
tilelang/engine/__init__.py (1)
tilelang/engine/callback.py (3)
register_cuda_postproc(7-15)register_hip_postproc(18-26)register_c_postproc(29-41)
tilelang/__init__.py (1)
tilelang/engine/callback.py (3)
register_cuda_postproc(7-15)register_hip_postproc(18-26)register_c_postproc(29-41)
🪛 Ruff (0.14.10)
testing/python/transform/test_nullable_buffer_params.py
84-84: Unused function argument: x
(ARG001)
tilelang/engine/__init__.py
4-4: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
5-5: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
6-6: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
tilelang/__init__.py
148-148: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
tilelang/engine/callback.py
110-110: PEP 484 prohibits implicit Optional
Convert to Optional[T]
(RUF013)
140-140: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
- GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
🔇 Additional comments (8)
tilelang/__init__.py (1)
148-148: LGTM! Consistent API extension for C host post-processing.The addition of
register_c_postprocfollows the established pattern for CUDA and HIP post-processing hooks, providing a uniform interface for code generation customization across different backends.Note: The static analysis warning about the
noqadirective is a false positive—F401suppression is necessary for re-exported symbols.tilelang/engine/__init__.py (1)
3-7: LGTM! Proper module-level API exposure.The multi-line import correctly groups all post-processing registration functions together and exposes
register_c_postprocalongside the existing CUDA and HIP variants.Note: The static analysis warnings about
noqadirectives are false positives—these symbols are intentionally re-exported.src/transform/arg_binder.cc (3)
602-609: Excellent null-safety improvement for shape loading.The guarded load prevents segmentation faults by checking
is_nullbefore dereferencing the buffer handle. When the buffer is null, the shape dimension safely defaults to zero, allowing downstream code to handle the null case gracefully.
664-696: Well-designed cascaded binding for multi-source symbolic shapes.The logic correctly builds a cascaded
if_then_elseexpression in reverse order, binding the symbolic variable to the shape from the first non-null buffer. The handling distinguishes between:
- Used buffers (non-nullable): directly use their raw shape value
- Nullable buffers: guard with
is_nullchecks and cascade fallbackThis ensures correctness when multiple buffers share a symbolic dimension, relying on the earlier assertion (lines 644-645) that at least one source is non-null.
707-708: Clear documentation of single-source nullable behavior.The updated comment accurately reflects the safer binding semantics: when the only source buffer is null, the symbolic shape variable binds to zero, preventing segfaults while allowing the kernel to execute (likely as a no-op for zero-sized dimensions).
tilelang/engine/callback.py (1)
29-41: LGTM! Consistent implementation of C host post-processing registration.The function follows the established pattern for CUDA and HIP post-processing, registering the global function
tilelang_callback_c_host_postprocthat will be invoked by the code generator.src/target/codegen_c_host.cc (1)
498-501: Perfect integration point for post-processing hook.The hook is invoked at the optimal stage—after code generation but before wrapping into
CSourceModule. The optional pattern (only invoking if registered) ensures backward compatibility, and the implementation correctly passes both the generated code and target to the callback.testing/python/transform/test_nullable_buffer_params.py (1)
72-101: Excellent regression test for nullable single-source shape binding.The test correctly validates the null-safety improvements in
arg_binder.cc:
- First invocation with a valid tensor confirms normal operation
- Second invocation with
Noneensures no segfault occurs and that the symbolic variablemsafely binds to 0 (as implemented in the guarded shape loading)Note: The static analysis warning about unused parameter
xat line 84 is expected—the parameter must be declared to test the nullable buffer path, even though the kernel body doesn't use it.
1 participant
This pull request introduces enhanced support for post-processing C host code generation in TileLang, improves the safety of buffer shape binding for nullable tensors, and adds corresponding tests. The most important changes are grouped below:
C Host Code Generation Post-Processing
register_c_postproc, and its decorator variant,register_c_postproc_callback, to allow users to intercept and modify C host code emitted by TileLang before it is wrapped into aCSourceModule. This includes updates totilelang/engine/callback.py,tilelang/engine/__init__.py, andtilelang/__init__.pyto expose the new API. [1] [2] [3] [4]BuildTileLangCHost, the callback is invoked if registered, allowing custom modification of the emitted code.Buffer Shape Binding Safety
ArgBinder::BindDLTensorsto safely handle nullable buffers: shape loads are now guarded to avoid dereferencing null handles, and symbolic shape variables bind to zero if the source buffer is null, preventing segfaults. [1] [2] [3] [4]Testing
test_nullable_single_source_shape, to ensure that binding a symbolic shape variable from a nullable buffer does not cause segfaults and safely binds to zero when the buffer isNone.Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.