tqchen commented on code in PR #14224:
URL: https://github.com/apache/tvm/pull/14224#discussion_r1128107313
##########
src/tir/transforms/unroll_loop.cc:
##########
@@ -137,14 +165,36 @@ class LoopUnroller : public StmtExprMutator {
LOG(FATAL) << "Unexpected use of deprecated StoreNode. Please use
BufferStoreNode instead.";
}
+ PrimExpr VisitExpr_(const BufferLoadNode* op) final {
+ if (unroll_local_access_) {
+ auto storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data));
+ if (storage_scope.rank == runtime::StorageRank::kLocal) {
+ VarLocalAccessMarker marker(&var_touched_local_);
+ for (PrimExpr e : op->indices) {
+ marker(e);
+ }
+ }
+ }
+ return GetRef<PrimExpr>(op);
+ }
+
Stmt VisitStmt_(const BufferStoreNode* op) final {
++step_count_;
- return StmtExprMutator::VisitStmt_(op);
+ if (unroll_local_access_) {
+ auto storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data));
+ if (storage_scope.rank == runtime::StorageRank::kLocal) {
+ VarLocalAccessMarker marker(&var_touched_local_);
+ for (PrimExpr e : op->indices) {
+ marker(e);
+ }
+ }
+ }
+ return StmtMutator::VisitStmt_(op);
}
Stmt VisitStmt_(const EvaluateNode* op) final {
++step_count_;
- return StmtExprMutator::VisitStmt_(op);
+ return StmtMutator::VisitStmt_(op);
Review Comment:
good catch updated
##########
src/tir/transforms/unroll_loop.cc:
##########
@@ -137,14 +165,36 @@ class LoopUnroller : public StmtExprMutator {
LOG(FATAL) << "Unexpected use of deprecated StoreNode. Please use
BufferStoreNode instead.";
}
+ PrimExpr VisitExpr_(const BufferLoadNode* op) final {
+ if (unroll_local_access_) {
+ auto storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(op->buffer->data));
+ if (storage_scope.rank == runtime::StorageRank::kLocal) {
Review Comment:
added
--
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]
