[GitHub] [tvm] Lunderberg commented on a change in pull request #9091: [TIR] tir.transform.StorageFlatten refactor

[GitBox <gi...@apache.org>]

Lunderberg commented on a change in pull request #9091:
URL: https://github.com/apache/tvm/pull/9091#discussion_r715831607



##########
File path: src/tir/transforms/storage_flatten.cc
##########
@@ -50,6 +50,913 @@ using runtime::StorageRank;
 using runtime::StorageScope;
 using runtime::ThreadScope;
 
+/* Make buffer realize extents and buffer shapes consistent
+ *
+ * For external buffers, verify that the extents of BufferRealize
+ * nodes match the shape of the external buffer.  For internal
+ * buffers, rewrite the shape of the Buffer objects to match the
+ * extent of the BufferRealize, and rewrite indices of
+ * BufferLoad/BufferStore nodes to match.
+ */
+class BufferShapeLegalize : public StmtExprMutator {
+ public:
+  explicit BufferShapeLegalize(const Map<Var, Buffer>& extern_buffer_map,
+                               IRVisitorWithAnalyzer* bound_analyzer)
+      : bound_analyzer_(bound_analyzer) {
+    for (auto kv : extern_buffer_map) {
+      extern_buffers_.insert(kv.second);
+    }
+  }
+
+  Stmt VisitStmt_(const BufferRealizeNode* op) final {
+    // External buffers should not be changed.
+    if (extern_buffers_.count(op->buffer)) {
+      ICHECK_EQ(op->buffer->shape.size(), op->bounds.size())
+          << "External buffer realize has mismatched dimension";

Review comment:
       Makes sense, and changed.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org