diff --git a/src/te/schedule/schedule_postproc_to_primfunc.cc b/src/te/schedule/schedule_postproc_to_primfunc.cc index c7d5d7a8dafa..7204ee059455 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 tuple = Downcast>(op->node); + Buffer view_buffer = Downcast(tuple[0]); + Buffer source_buffer = Downcast(tuple[1]); + return AttrStmt( + Array{GetRemappedBuffer(view_buffer), GetRemappedBuffer(source_buffer)}, + op->attr_key, op->value, op->body); } else { return ret; }