wrongtest-intellif commented on code in PR #14021:
URL: https://github.com/apache/tvm/pull/14021#discussion_r1169679801
##########
src/tir/transforms/compact_buffer_region.cc:
##########
@@ -103,17 +59,62 @@ NDIntSet NDIntSetEval(Region region, PrimExpr predicate,
return support::NDIntSetEval(support::NDIntSetFromRegion(region), dom_map);
}
+/*!
+ * \brief Collect buffer aliasing information.
+ */
+class Var2BufferCollector : public StmtExprVisitor {
+ public:
+ /*! \brief Map the buffer var to all aliased buffers. */
+ std::unordered_map<Var, std::unordered_set<Buffer, ObjectPtrHash,
ObjectPtrEqual>, ObjectPtrHash,
+ ObjectPtrEqual>
+ var2buffer_;
+
+ private:
+ void VisitStmt_(const BufferStoreNode* op) final {
+ var2buffer_[op->buffer->data].insert(op->buffer);
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitExpr_(const BufferLoadNode* op) final {
+ var2buffer_[op->buffer->data].insert(op->buffer);
+ StmtExprVisitor::VisitExpr_(op);
+ }
+
+ void VisitStmt_(const BlockNode* op) final {
+ for (const Buffer& buffer : op->alloc_buffers) {
+ var2buffer_[buffer->data].insert(buffer);
+ }
+ for (const MatchBufferRegion& region : op->match_buffers) {
+ var2buffer_[region->buffer->data].insert(region->buffer);
+ var2buffer_[region->source->buffer->data].insert(region->source->buffer);
+ }
+ StmtExprVisitor::VisitStmt_(op);
+ }
+
+ void VisitStmt_(const DeclBufferNode* op) final {
+ var2buffer_[op->buffer->data].insert(op->buffer);
+ StmtExprVisitor::VisitStmt_(op);
+ }
+};
+
/*!
* \brief Collect the access region of each buffer.
* \note The param buffer regions will not be collected.
*/
class BufferAccessRegionCollector : public StmtExprVisitor {
public:
static std::unordered_map<Buffer, Region, ObjectPtrHash, ObjectPtrEqual>
Collect(
- const PrimFunc& f) {
- BufferAccessRegionCollector collector;
- collector(f->body);
- return std::move(collector.buffer_access_region_);
+ const PrimFunc& f, bool collect_inbound) {
Review Comment:
Hi~ @Lunderberg I totally agree the `DomainTouched` utility is of the
similar purpose . Here are some differences I think:
- `BufferAccessRegionCollector` collect region wrt the allocation point,
that is,
```python
for i in range(10):
a = T.allocate([10])
A = T.decl_buffer( [10, 10], data=a)
for j in range(10):
# use A[i, j]
```
would give region `A[i, 0:10]`, which is the domain "touched" under each
allocation scope.
- `BufferAccessRegionCollector` take some special considerations on thread
bindings.
- `BufferAccessRegionCollector` try use more arith utilities to improve
intset analysis.
It would be great to share the same implementation (maybe at the cost of
adding cost to where we use `DomainTouched`?). I'd like to try it in a
standalone MR.
--
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]
