[TIR][Runtime] Enforce host-evaluable assumptions at runtime#2655
Open
LeiWang1999 wants to merge 2 commits into
Open
[TIR][Runtime] Enforce host-evaluable assumptions at runtime#2655LeiWang1999 wants to merge 2 commits into
LeiWang1999 wants to merge 2 commits into
Conversation
|
👋 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! 🚀 |
Contributor
📝 WalkthroughWalkthroughChangesAssumption runtime validation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AssumptionInjector
participant SplitDeviceFunc
participant HostLaunch
participant MakePackedAPI
participant RuntimeAssertion
AssumptionInjector->>SplitDeviceFunc: provide marked assumption attributes
SplitDeviceFunc->>HostLaunch: wrap launch with host-evaluable checks
HostLaunch->>MakePackedAPI: pass packed host body
MakePackedAPI->>RuntimeAssertion: lower marked checks to AssertStmt
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
Member
Author
|
@regression-perf |
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/transform/split_host_device.cc (1)
135-141: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
HostEvaluableConditionChecker::Checkshould resetis_host_evaluable_before each visit.The field is initialized to
truein the constructor but never reset inCheck. If the same checker instance were reused across multiple conditions, a priorfalseresult would persist and cause valid conditions to be incorrectly rejected. Currently masked because a new checker is created per attribute (line 182), but this is a latent bug that could surface during refactoring.🛡️ Proposed fix
bool Check(const PrimExpr &condition) { + is_host_evaluable_ = true; if (!condition->dtype.is_bool() || condition->dtype.lanes() != 1) { return false; } VisitExpr(condition); return is_host_evaluable_; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/transform/split_host_device.cc` around lines 135 - 141, Reset is_host_evaluable_ to true at the start of HostEvaluableConditionChecker::Check before validating or visiting each condition, so reused checker instances evaluate every condition independently. Preserve the existing boolean/scalar validation and VisitExpr flow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/transform/split_host_device.cc`:
- Around line 135-141: Reset is_host_evaluable_ to true at the start of
HostEvaluableConditionChecker::Check before validating or visiting each
condition, so reused checker instances evaluate every condition independently.
Preserve the existing boolean/scalar validation and VisitExpr flow.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: c1fe1d41-244b-4d01-87e1-a573840207c6
📒 Files selected for processing (6)
src/transform/common/attr.hsrc/transform/inject_assumes.ccsrc/transform/make_packed_api.ccsrc/transform/split_host_device.cctesting/python/transform/test_tilelang_transform_make_packed_api.pytesting/python/transform/test_tilelang_transform_split_host_device.py
🚧 Files skipped from review as they are similar to previous changes (1)
- testing/python/transform/test_tilelang_transform_split_host_device.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
T.assumeconditions in the host launcher before a device kernel is launched.Changes
InjectAssumeswhile preserving the pairedtl.assumeoptimizer fact.SplitHostDevice, and remove runtime-check markers from device functions.MakePackedAPIand emitAssertStmtnodes outside optimizer assumptions so later simplification cannot erase validation.Validation
./format.shcmake --build build -jTILELANG_DISABLE_CACHE=1 TVM_TEST_TARGETS='cuda;llvm' python -m pytest testing/python/transform/test_tilelang_transform_make_packed_api.py testing/python/transform/test_tilelang_transform_split_host_device.py testing/python/language/test_tilelang_language_assume.py testing/python/transform/test_nullable_buffer_params.py -q(21 passed, 4 skipped)Risk
Summary
tl.assumechecks that are statically host-evaluable to the host launcher, while keeping the original device-side assumptions for optimization.MakePackedAPIto extracttl::attr::kAssumeRuntimeCheckattributes and lower them to host-side runtimeAssertStmtchecks (preserving condition, message, and source span), while leaving device-only/unsupported assumptions unchanged.SplitHostDeviceto conservatively lift eligible runtime-check assumptions to the host-side kernel-launch code, strip the marker attributes from the replaced regions, and wrap the host call so the runtime checks execute.inject_assumesso generated user-authored assumptions include the runtime-check marker (with intentional exclusions for buffer shape/stride assumptions).Validation
C++ style / lint notes
docs/developer_guide/cpp_style.md.