Lunderberg commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r810384766
##########
File path: src/tir/transforms/vectorize_loop.cc
##########
@@ -62,30 +62,89 @@ class VecAllocAccess : public StmtExprMutator {
public:
VecAllocAccess(const VarNode* buf, Var var, int var_lanes)
: buf_(buf), var_(var), var_lanes_(var_lanes) {}
- // Load
+
PrimExpr VisitExpr_(const LoadNode* op) final {
- PrimExpr expr = StmtExprMutator::VisitExpr_(op);
- op = expr.as<LoadNode>();
- if (op->buffer_var.get() == buf_) {
- return Load(op->dtype, op->buffer_var, op->index * var_lanes_ + var_,
op->predicate);
- } else {
- return expr;
- }
+ LOG(FATAL) << "Unexpected use of deprecated LoadNode. Please use
BufferLoadNode instead.";
+ return PrimExpr();
}
- // Store
+
Stmt VisitStmt_(const StoreNode* op) final {
- Stmt stmt = StmtExprMutator::VisitStmt_(op);
- op = stmt.as<StoreNode>();
- if (op->buffer_var.get() == buf_) {
- return Store(op->buffer_var, op->value, op->index * var_lanes_ + var_,
op->predicate);
+ LOG(FATAL) << "Unexpected use of deprecated StoreNode. Please use
BufferStoreNode instead.";
+ return Stmt();
+ }
+
+ PrimExpr VisitExpr_(const BufferLoadNode* op) final {
+ auto load = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));
+ return UpdateBufferAccess(load);
+ }
+
+ Stmt VisitStmt_(const BufferStoreNode* op) final {
+ auto store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
+ return UpdateBufferAccess(store);
+ }
+
+ private:
+ template <typename Node>
+ Node UpdateBufferAccess(Node node) {
+ // Only update the buffer that's being replaced.
+ if (node->buffer->data.get() != buf_) {
+ return node;
+ }
+
+ arith::Analyzer analyzer;
Review comment:
Sounds good, and done.
--
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]
