[GitHub] [tvm] Hzfengsy commented on a diff in pull request #14021: [TIR] More flexible buffer compaction

["Hzfengsy (via GitHub)" <gi...@apache.org>]

Hzfengsy commented on code in PR #14021:
URL: https://github.com/apache/tvm/pull/14021#discussion_r1121240663


##########
python/tvm/tir/transform/transform.py:
##########
@@ -780,7 +780,7 @@ def ConvertBlocksToOpaque():
     return _ffi_api.ConvertBlocksToOpaque()  # type: ignore
 
 
-def CompactBufferAllocation():
+def CompactBufferAllocation(is_strict=True):

Review Comment:
   ```suggestion
   def CompactBufferAllocation(is_strict: bool = True):
   ```



##########
src/tir/transforms/ir_utils.cc:
##########
@@ -455,5 +489,52 @@ std::pair<PrimExpr, PrimExpr> GetAsyncWaitAttributes(const AttrStmtNode* op) {
   return std::make_pair(op->value, inner->value);
 }
 
+/*! \brief Collect storage alignment information from annotations. */
+class StorageAlignCollector : public StmtVisitor {
+ private:
+  friend std::unordered_map<Var, StorageAlignAnnotation, ObjectPtrHash, ObjectPtrEqual>
+  CollectStorageAlignAnnotation(const Stmt& body);
+
+  /*! \brief For s-stir, the alignment annotations reside in block annotations. */
+  void VisitStmt_(const BlockNode* op) final {
+    auto it = op->annotations.find(attr::buffer_dim_align);
+    if (it != op->annotations.end()) {
+      auto storage_align_annotation = Downcast<StorageAlignAnnotation>((*it).second);
+      for (const auto& storage_align_tuple : storage_align_annotation) {
+        int buffer_index = storage_align_tuple[0]->value;
+        const Buffer& buffer = op->writes[buffer_index]->buffer;
+        storage_align_[buffer->data].push_back(storage_align_tuple);
+      }
+    }
+    StmtVisitor::VisitStmt_(op);
+  }
+
+  /*! \brief For lowerred tir, the alignment annotations reside in allocate annotations. */

Review Comment:
   ```suggestion
     /*! \brief For lowered tir, the alignment annotations reside in allocate annotations. */
   ```



-- 
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