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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
