[Enhancement] Improve equality checks in layout nodes and fragment validation#1573
[Enhancement] Improve equality checks in layout nodes and fragment validation#1573LeiWang1999 merged 2 commits intomainfrom
Conversation
…lidation * Enhanced the IsEqual method in LayoutNode and FragmentNode to include detailed comparisons of forward mappings, ensuring accurate equality checks. * Introduced a new parameter in ProveFragmentContains to validate physical indices when checking fragment containment, improving correctness in layout validation. * Removed obsolete test file related to layout inference. This update aims to strengthen the integrity of layout comparisons and fragment validations in the system.
|
👋 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! 🚀 |
📝 WalkthroughWalkthroughThis change enhances Layout and Fragment equality comparison by comparing forward-mapped synthetic variables instead of raw indices, and introduces forward-index validation in fragment containment checks via a new parameter in ProveFragmentContains. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
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 |
|
@codex review |
![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)
testing/python/layout/test_tilelang_layout_equal.py (1)
8-8: Consider removingset_random_seed()for this test module.This test file contains only deterministic equality comparisons with no random number generation. The
set_random_seed()call is harmless but unnecessary, adding a slight inconsistency since the tests don't rely on randomness.🔎 Proposed fix
-tilelang.testing.set_random_seed() - - + class TestLayoutEqual:
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
src/layout/layout.ccsrc/op/parallel.ccsrc/op/parallel.htesting/python/layout/test_tilelang_layout_equal.pytesting/python/layout/test_tilelang_layout_inference.py
🧰 Additional context used
🧬 Code graph analysis (1)
testing/python/layout/test_tilelang_layout_equal.py (3)
src/layout/layout.cc (4)
Layout(60-72)Layout(74-77)Fragment(518-540)Fragment(542-552)tilelang/testing/__init__.py (1)
set_random_seed(32-37)tilelang/layout/fragment.py (1)
replicate(138-152)
🪛 Ruff (0.14.10)
testing/python/layout/test_tilelang_layout_equal.py
86-86: Unused lambda argument: j
(ARG005)
87-87: Unused lambda argument: i
(ARG005)
98-98: Unused lambda argument: j
(ARG005)
99-99: Unused lambda argument: j
(ARG005)
104-104: Unused lambda argument: j
(ARG005)
105-105: Unused lambda argument: j
(ARG005)
164-164: Unused lambda argument: i
(ARG005)
164-164: Unused lambda argument: j
(ARG005)
165-165: Unused lambda argument: i
(ARG005)
165-165: Unused lambda argument: j
(ARG005)
172-172: Unused lambda argument: i
(ARG005)
172-172: Unused lambda argument: j
(ARG005)
173-173: Unused lambda argument: j
(ARG005)
⏰ 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 Nightly-ROCm-7.1 (on self-hosted-amd)
- 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)
🔇 Additional comments (8)
src/op/parallel.h (1)
27-31: LGTM!The API extension is well-designed with a default parameter value of
false, ensuring backward compatibility with existing call sites while enabling the new forward-index validation when explicitly requested.src/layout/layout.cc (2)
695-718: LGTM!The fix correctly addresses the equality comparison issue by constructing fresh common variables and comparing the actual forward mapping results rather than the raw AST structure. The early return pattern improves readability.
720-764: LGTM!The implementation correctly handles the broadcast case and uses the fresh common variable approach for both forward index and forward thread comparisons. The addition of
ForwardThreadcomparison ensures complete fragment equality validation.src/op/parallel.cc (2)
44-70: LGTM!The forward index validation logic is well-implemented. The dimension check followed by component-wise equality verification using
analyzer_.Simplifyis the correct approach for comparing symbolic expressions.
621-637: LGTM!The call site correctly enables forward index validation during candidate validation. The comment clearly explains the rationale: ensuring physical indices match for correct code generation. Other call sites appropriately use the default
falsesince they're checking containment relationships rather than strict validation.testing/python/layout/test_tilelang_layout_equal.py (3)
11-67: LGTM!Comprehensive test coverage for layout equality including shape mismatches, coefficient differences, variable name independence, 2D outputs, and the notable edge case of commutative expressions where the test correctly validates behavior without assuming AST-level equivalence.
69-125: LGTM!Good coverage of fragment equality scenarios including thread mapping, forward index, replicate factors, and the
forward_fnAPI. The unused lambda arguments flagged by static analysis are intentional—the lambdas must match the expected arity even when only a subset of parameters is used in the expression.
127-175: LGTM!Excellent edge case coverage including single/multi-dimensional layouts, empty forward index, and the important distinction between constant and variable mappings using
tvm.tir.const.
|
@regression-perf |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| auto small_physical = small_frag->Forward(small_frag_indices); | ||
| auto large_physical = large_frag->Forward(large_frag_indices); | ||
|
|
||
| // Dimension mismatch means they are not equal. | ||
| if (small_physical.size() != large_physical.size()) { |
There was a problem hiding this comment.
When check_forward_index is enabled, any size mismatch between small_frag->Forward(...) and large_frag->Forward(...) is treated as incompatible. This causes plan candidates (which often have a non-empty forward_index, e.g., from PlanLoopPartition) to be rejected against buffer fragments that intentionally use an empty forward_index (see CompleteBufferFragment/ComputeLoopLayoutFromBuffer constructing Fragment(..., {})). The result is that ValidateCandidateAgainstFragments can mark the plan candidate invalid and ChooseBestCandidate falls back to the more replicated buffer-based candidate, introducing a performance regression even though no physical index compatibility can be checked. Consider only enforcing the forward-index equality when both fragments define a non-empty forward index (or treating empty as “don’t care”).
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Thanks for the review! However, this concern doesn't apply here.
When a Fragment is constructed with an empty forward_index (e.g., Fragment(..., {}, ...)), the constructor automatically infers and fills it via infer_fragment_index():
// layout.cc:510-513
if (forward_index.empty()) {
forward_index = {
infer_fragment_index(getVarMap(), forward_thread_, &analyzer)};
}
So after construction, all Fragments have non-empty forward_index_. The "empty forward_index" scenario you described doesn't exist at runtime - it's only a construction-time shorthand that gets filled automatically.
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
1 participant
Summary
This PR improves the correctness of layout and fragment equality checks, and adds validation for physical index compatibility in
ProveFragmentContains.Problem
LayoutNode::IsEqualandFragmentNode::IsEqualhad incorrect comparison logic: The original implementation usedStructuralEqualdirectly onforward_index_, which could incorrectly compareInputPlaceholdervariables. For example, two layouts with different mappings likei * 16 + jandj * 16 + imight be incorrectly considered equal if theirInputPlaceholderindices happened to match structurally.ProveFragmentContainsonly checked thread containment, not physical index compatibility: When validating a loop layout candidate against a buffer fragment, the function only verified that threads accessing small fragment elements are a subset of threads accessing large fragment elements. However, it didn't verify that the physical indices match, which could lead to incorrect code generation.Example of incorrect code generation:
Without physical index validation, code generation would produce:
Solution
Fix
IsEqualby usingForwardwith common variables:Add
check_forward_indexparameter toProveFragmentContains:When
check_forward_index=true, validates physical index equality:Changes
src/op/parallel.h: Addcheck_forward_indexparameter toProveFragmentContainssrc/op/parallel.cc: Implement physical index validation; update call sites inValidateCandidateAgainstFragmentsandInferLayoutsrc/layout/layout.cc: FixLayoutNode::IsEqualandFragmentNode::IsEqualto useForwardwith common variablestesting/python/layout/test_tilelang_layout_equal.py: Add comprehensive tests for layout/fragment equalitytest_tilelang_issue_layout.pytotesting/python/layout/test_tilelang_layout_inference.pyTest plan
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.