[tvm] branch main updated: [TE] Correctly generate buffer binds with axis separators (#10819)

[masahi]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new b5cad84  [TE] Correctly generate buffer binds with axis separators 
(#10819)
b5cad84 is described below

commit b5cad84f795a63ee8cec9fa28a6678d55495f283
Author: Eric Lunderberg <lunderb...@users.noreply.github.com>
AuthorDate: Wed Mar 30 14:44:26 2022 -0500

    [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`.
---
 src/te/schedule/schedule_postproc_to_primfunc.cc | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/te/schedule/schedule_postproc_to_primfunc.cc 
b/src/te/schedule/schedule_postproc_to_primfunc.cc
index c7d5d7a..7204ee0 100644
--- a/src/te/schedule/schedule_postproc_to_primfunc.cc
+++ b/src/te/schedule/schedule_postproc_to_primfunc.cc
@@ -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;
     }