yzh119 commented on code in PR #14224:
URL: https://github.com/apache/tvm/pull/14224#discussion_r1127986799
##########
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:
Why changed to `StmtMutator`?
##########
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:
Should we also consider warp-level memory (identical to local for CUDA).
--
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]
