Fix atomic_load access_ptr lowering for dynamic indices#2157
Conversation
Add the issue 2123 regression test and keep the access_ptr base BufferLoad intact through safe-memory legalization and lowering. Closes tile-ai#2123
Move shared tl.access_ptr BufferLoad base visitation into a common transform helper so LowerAccessPtr and safe-memory legalization stay aligned. Extend the issue 2123 regression with a direct LowerAccessPtr pass check while keeping the LowerAndLegalize pipeline coverage.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a reusable helper to visit/rebuild a ChangesAccess Pointer Lowering and Legalization Fix
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
|
👋 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! 🚀 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/transform/common/access_ptr_utils.h (1)
7-9: ⚡ Quick winAvoid
using namespacedirectives in headers.
using namespace tir;at namespace scope in a header injects the entiretirnamespace intotvm::tlfor every translation unit that includes this file. This can cause name-lookup surprises and ODR clashes as the codebase grows. Prefer fully-qualifying names (tir::BufferLoad,tir::PrimExpr, etc.) within the header.♻️ Suggested change
namespace tvm { namespace tl { -using namespace tir; namespace detail { template <typename VisitExprFn> -BufferLoad VisitAccessPtrBase(const PrimExpr &expr, VisitExprFn &&visit_expr) { - const auto *base_load_node = expr.as<BufferLoadNode>(); +tir::BufferLoad VisitAccessPtrBase(const tir::PrimExpr &expr, VisitExprFn &&visit_expr) { + const auto *base_load_node = expr.as<tir::BufferLoadNode>(); ICHECK(base_load_node) << "tl.access_ptr base must be BufferLoad, but got " << expr; - BufferLoad base_load = tvm::ffi::GetRef<BufferLoad>(base_load_node); + tir::BufferLoad base_load = tvm::ffi::GetRef<tir::BufferLoad>(base_load_node); - Array<PrimExpr> indices; + tvm::Array<tir::PrimExpr> indices; bool changed = false; - for (const PrimExpr &index : base_load->indices) { + for (const tir::PrimExpr &index : base_load->indices) { ... } - Optional<PrimExpr> predicate = base_load->predicate; + tvm::Optional<tir::PrimExpr> predicate = base_load->predicate; ... - return BufferLoad(base_load->buffer, indices, predicate, base_load->span); + return tir::BufferLoad(base_load->buffer, indices, predicate, base_load->span); }🤖 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/common/access_ptr_utils.h` around lines 7 - 9, Remove the header-level "using namespace tir;" from src/transform/common/access_ptr_utils.h and update all references that rely on it to fully-qualified names (e.g., tir::BufferLoad, tir::PrimExpr, tir::Var, etc.) inside the tvm::tl namespace and in functions/classes such as any occurrences in AccessPtr utilities; this keeps the header from polluting translation units and prevents lookup/ODR issues while preserving the existing API and symbols.
🤖 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.
Nitpick comments:
In `@src/transform/common/access_ptr_utils.h`:
- Around line 7-9: Remove the header-level "using namespace tir;" from
src/transform/common/access_ptr_utils.h and update all references that rely on
it to fully-qualified names (e.g., tir::BufferLoad, tir::PrimExpr, tir::Var,
etc.) inside the tvm::tl namespace and in functions/classes such as any
occurrences in AccessPtr utilities; this keeps the header from polluting
translation units and prevents lookup/ODR issues while preserving the existing
API and symbols.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 13264419-1cb3-4c81-92f3-d99c3796f3ca
📒 Files selected for processing (1)
src/transform/common/access_ptr_utils.h
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/transform/common/access_ptr_utils.h (1)
8-9: 💤 Low valueRemove unnecessary include
<tvm/ffi/reflection/registry.h>.This header is only used for FFI registration macros, not for
tvm::ffi::GetRef. The codebase showsGetRefworks without it (e.g.,loop_vectorization_utils.husesGetRefextensively without this header). The required declarations come from<tvm/tir/expr.h>or its transitive includes.🤖 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/common/access_ptr_utils.h` around lines 8 - 9, Remove the unnecessary include of <tvm/ffi/reflection/registry.h> from the include block — it is only for FFI registration macros and not required for tvm::ffi::GetRef or the code in this file; delete that include line in access_ptr_utils.h and keep the existing <tvm/tir/expr.h> (and other necessary includes) so compilation uses the transitive declarations instead.
🤖 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.
Nitpick comments:
In `@src/transform/common/access_ptr_utils.h`:
- Around line 8-9: Remove the unnecessary include of
<tvm/ffi/reflection/registry.h> from the include block — it is only for FFI
registration macros and not required for tvm::ffi::GetRef or the code in this
file; delete that include line in access_ptr_utils.h and keep the existing
<tvm/tir/expr.h> (and other necessary includes) so compilation uses the
transitive declarations instead.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: dbd08278-cfbd-468d-b7d2-82a67e67b051
📒 Files selected for processing (3)
src/transform/common/access_ptr_utils.hsrc/transform/legalize_safe_memory_access.ccsrc/transform/lower_access_ptr.cc
…ssue-2123-access-ptr-helper
|
Thanks for tracking this down and putting together a fix. However, I do not think accepting a non-BufferLoad as the first argument of In other words, we should avoid producing IR like this: tl.access_ptr(
tirx.if_then_else(
in_bounds(look),
BufferLoad(status, [look]),
safe_value,
),
1,
read_mask,
)and instead preserve the pointer shape and put the safety guard around the operation that consumes the pointer, for example: tirx.if_then_else(
in_bounds(look),
tl.atomic_load_elem_op(
tl.access_ptr(
BufferLoad(status, [look]),
1,
read_mask,
),
...
),
safe_value,
)I will make some follow-up improvements in LegalizeSafeMemoryAccess so it treats |
…ssue-2123-access-ptr-helper
…s.cc Updated the return structure of AccessPtrInfo to improve readability and maintainability. The logic for handling the tvm_access_ptr operation has been streamlined, ensuring clearer checks and consistent formatting of returned values.
|
@regression-perf |
Performance Regression Test ReportTriggered by: @LeiWang1999 Results
Artifacts
|
2 participants
Summary
T.atomic_loadfails inLowerAccessPtrwhen the address uses a block-derived loop variable #2123, Fixtl.access_ptrlowering for dynamic indices and share the base-load helper between lowering and legalization.Validation
./format.shpassed pre-commit hooks; the script reported unstaged changes before commit.Summary by CodeRabbit