This is an automated email from the ASF dual-hosted git repository.
syfeng 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 fa4d73ac51 [TIR] Handle DeclBuffer in IRSubstitute (#15054)
fa4d73ac51 is described below
commit fa4d73ac51b03555e40b57e8f3ca17c250024bd6
Author: Eric Lunderberg <[email protected]>
AuthorDate: Thu Jun 8 02:15:36 2023 -0400
[TIR] Handle DeclBuffer in IRSubstitute (#15054)
Previously, occurrences of `DeclBuffer` were not handled by
`IRSubstitute`, and the buffer object was left unhandled. This commit
updates `IRSubstitute` to modify the buffer object in `DeclBuffer`,
similar to how it is already modified in `BufferLoad` and
`BufferStore` instances.
This is a subset of changes being split out from
https://github.com/apache/tvm/pull/14778 into independent portions.
---
src/tir/ir/stmt_functor.cc | 26 ++++++-
.../unittest/test_tir_stmt_functor_substitute.py | 83 ++++++++++++++++++++++
2 files changed, 106 insertions(+), 3 deletions(-)
diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc
index ae797c0d79..1c15f95826 100644
--- a/src/tir/ir/stmt_functor.cc
+++ b/src/tir/ir/stmt_functor.cc
@@ -675,6 +675,11 @@ class IRSubstitute : public StmtExprMutator {
return VisitBufferAccess(std::move(node));
}
+ Stmt VisitStmt_(const DeclBufferNode* op) final {
+ auto node = Downcast<DeclBuffer>(StmtExprMutator::VisitStmt_(op));
+ return VisitBufferAccess(std::move(node));
+ }
+
template <typename Node>
Node VisitBufferAccess(Node node) {
Buffer new_buf = GetRemappedBuffer(node->buffer);
@@ -694,10 +699,25 @@ class IRSubstitute : public StmtExprMutator {
return it->second;
}
- auto new_buffer_var = vmap_(buf->data);
- if (new_buffer_var.defined() &&
!new_buffer_var.value().same_as(buf->data)) {
+ PrimExpr new_buffer_var_expr = VisitExpr(buf->data);
+ CHECK(new_buffer_var_expr->IsInstance<VarNode>())
+ << "Buffer " << buf << " uses backing allocation " << buf->data
+ << ", which was substituted into the expression " <<
new_buffer_var_expr << ". "
+ << "However, this expression is of type " <<
new_buffer_var_expr->GetTypeKey()
+ << " and the backing allocation must be a tir::Var";
+
+ Var buffer_var = Downcast<Var>(new_buffer_var_expr);
+ auto elem_offset = VisitExpr(buf->elem_offset);
+ auto shape = buf->shape.Map([this](const auto& expr) { return
VisitExpr(expr); });
+ auto strides = buf->strides.Map([this](const auto& expr) { return
VisitExpr(expr); });
+
+ if (!buffer_var.same_as(buf->data) ||
!elem_offset.same_as(buf->elem_offset) ||
+ !shape.same_as(buf->shape) || !strides.same_as(buf->strides)) {
auto writer = buf.CopyOnWrite();
- writer->data = Downcast<Var>(new_buffer_var);
+ writer->data = buffer_var;
+ writer->elem_offset = elem_offset;
+ writer->shape = shape;
+ writer->strides = strides;
}
buf_remap_[key] = buf;
diff --git a/tests/python/unittest/test_tir_stmt_functor_substitute.py
b/tests/python/unittest/test_tir_stmt_functor_substitute.py
new file mode 100644
index 0000000000..dffcb04da8
--- /dev/null
+++ b/tests/python/unittest/test_tir_stmt_functor_substitute.py
@@ -0,0 +1,83 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import tvm
+import tvm.testing
+
+from tvm.script import tir as T
+from tvm.tir.stmt_functor import substitute
+
+
+class BaseCompare(tvm.testing.CompareBeforeAfter):
+ def transform(self):
+ def inner(mod):
+ func = mod["main"]
+ vmap = {func.params[0]: 16}
+ new_func = tvm.tir.PrimFunc(params=[], body=substitute(func.body,
vmap))
+ return tvm.IRModule.from_expr(new_func)
+
+ return inner
+
+
+class TestBasicSubstitute(BaseCompare):
+ def before(n: T.int32):
+ for i in range(n):
+ T.evaluate(i)
+
+ def expected():
+ for i in range(16):
+ T.evaluate(i)
+
+
+class TestSubstituteAllocate(BaseCompare):
+ def before(n: T.int32):
+ A_data = T.allocate([n], "float32")
+ T.evaluate(A_data)
+
+ def expected():
+ A_data = T.allocate([16], "float32")
+ T.evaluate(A_data)
+
+
+class TestSubstituteBufferLoad(BaseCompare):
+ def before(n: T.int32):
+ A_data = T.allocate([n], "float32")
+ A = T.Buffer(n, "float32", data=A_data)
+ for i in range(n):
+ T.evaluate(A[i])
+
+ def expected():
+ A_data = T.allocate([16], "float32")
+ A = T.Buffer(16, "float32", data=A_data)
+ for i in range(16):
+ T.evaluate(A[i])
+
+
+class TestSubstituteDeclBuffer(BaseCompare):
+ def before(n: T.int32):
+ A_data = T.allocate([n], "float32")
+ A = T.decl_buffer(n, "float32", data=A_data)
+ T.evaluate(A.data)
+
+ def expected():
+ A_data = T.allocate([16], "float32")
+ A = T.decl_buffer(16, "float32", data=A_data)
+ T.evaluate(A.data)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
