Skip to content

Commit b5cad84

Browse files
authored
[TE] Correctly generate buffer binds with axis separators (#10819)
In SchedulePostProcToPrimfunc, when the axis separator attribute is moved to the buffer properties, it doesn't update buffers that are in the buffer bind scope. This occurs if `Stage.tensorize` is called for a stage whose layout transformation includes `te.AXIS_SEPARATOR`.
1 parent de93d81 commit b5cad84

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

src/te/schedule/schedule_postproc_to_primfunc.cc

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,13 @@ class AxisSeparatorsAttrUnwrapper : StmtExprMutator {
289289

290290
if (op->attr_key == tir::attr::axis_separators) {
291291
return op->body;
292+
} else if (op->attr_key == tir::attr::buffer_bind_scope) {
293+
Array<ObjectRef> tuple = Downcast<Array<ObjectRef>>(op->node);
294+
Buffer view_buffer = Downcast<Buffer>(tuple[0]);
295+
Buffer source_buffer = Downcast<Buffer>(tuple[1]);
296+
return AttrStmt(
297+
Array<ObjectRef>{GetRemappedBuffer(view_buffer), GetRemappedBuffer(source_buffer)},
298+
op->attr_key, op->value, op->body);
292299
} else {
293300
return ret;
294301
}

0 commit comments

Comments
 (0)