Skip to content

Commit 65bbb80

Browse files
committed
Enhance layout inference and copy operations with 1D TMA support
- Updated `CopyNode` to introduce separate handling for 1D bulk load/store operations, including new methods for checking and lowering these operations. - Modified `InferLayout` and `GetCopyInst` to accommodate additional parameters for layout maps and analyzers. - Enhanced `AtomicAddNode` and `FillNode` to utilize the updated layout inference logic. - Improved buffer out-of-bounds checks during layout inference to ensure safe memory access. This update improves the efficiency and correctness of memory operations in the TileLang framework.
1 parent f07f31c commit 65bbb80

File tree

7 files changed

+287
-271
lines changed

7 files changed

+287
-271
lines changed

src/op/atomic_add.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,7 @@ Stmt AtomicAddNode::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
361361
InferLevel::kFree};
362362
for (auto level : levels) {
363363
(par_op)->InferLayout(
364-
{T.target, T.thread_bounds, T.layout_map, T.buffer_remap}, level);
364+
{T.target, T.thread_bounds, T.layout_map, analyzer, false, T.buffer_remap}, level);
365365
}
366366
auto loop_layout = par_op->GetLoopLayout();
367367
Var thread_var = T.thread_var;

0 commit comments

Comments
 (0)