jroesch commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r776900143
##########
File path: src/tir/transforms/inject_virtual_thread.cc
##########
@@ -354,46 +391,44 @@ class VTInjector : public StmtExprMutator {
}
// Allocate
Stmt VisitStmt_(const AllocateNode* op) final {
+ Allocate node = GetRef<Allocate>(op);
+
PrimExpr condition = this->VisitExpr(op->condition);
+
+ Array<PrimExpr> extents = op->extents;
+ extents.MutateByApply([this](const PrimExpr& extent) { return
this->VisitExpr(extent); });
+
if (visit_touched_var_ && !vt_loop_injected_) {
return InjectVTLoop(GetRef<Stmt>(op), true);
}
- bool changed = false;
- Array<PrimExpr> extents;
- for (size_t i = 0; i < op->extents.size(); i++) {
- PrimExpr new_ext = this->VisitExpr(op->extents[i]);
- if (visit_touched_var_ && !vt_loop_injected_) {
- return InjectVTLoop(GetRef<Stmt>(op), true);
- }
- if (!new_ext.same_as(op->extents[i])) changed = true;
- extents.push_back(new_ext);
- }
visit_touched_var_ = false;
- Stmt body;
- // always rewrite if not allow sharing.
+ // Rewrite the buffer if its shape or any value stored in it
+ // depends on the virtual thread var. If `allow_share_` is false,
+ // then the buffer is always rewritten, even if separate virtual
+ // threads only read from the buffer.
if (touched_var_.count(op->buffer_var.get()) || !allow_share_) {
// place v on highest dimension.
- PrimExpr stride = foldl([](PrimExpr a, PrimExpr b, Span span) { return
mul(a, b, span); },
- make_const(DataType::Int(32), 1), op->extents) *
- op->dtype.lanes();
- Array<PrimExpr> other;
- other.push_back(make_const(op->extents[0].dtype(), num_threads_));
- for (PrimExpr e : extents) {
- other.push_back(e);
- }
- extents = other;
- changed = true;
- // mark this buffer get touched.
+
+ // TODO(Lunderberg): Move pass to apply before
Review comment:
For these TODOs can we make sure we have some kind of tracking
issue/roadmap issue cc @denise-k on this one?
--
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]
