Skip to content

Commit

Permalink
[TE] Correctly generate buffer binds with axis separators (apache#10819)
Browse files Browse the repository at this point in the history
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`.
  • Loading branch information
Lunderberg authored and mehrdadh committed Apr 11, 2022
1 parent c800862 commit ff2c0ab
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions src/te/schedule/schedule_postproc_to_primfunc.cc
Original file line number Diff line number Diff line change
Expand Up @@ -289,6 +289,13 @@ class AxisSeparatorsAttrUnwrapper : StmtExprMutator {

if (op->attr_key == tir::attr::axis_separators) {
return op->body;
} else if (op->attr_key == tir::attr::buffer_bind_scope) {
Array<ObjectRef> tuple = Downcast<Array<ObjectRef>>(op->node);
Buffer view_buffer = Downcast<Buffer>(tuple[0]);
Buffer source_buffer = Downcast<Buffer>(tuple[1]);
return AttrStmt(
Array<ObjectRef>{GetRemappedBuffer(view_buffer), GetRemappedBuffer(source_buffer)},
op->attr_key, op->value, op->body);
} else {
return ret;
}
Expand Down

0 comments on commit ff2c0ab

Please sign in to comment.