[TIR] Improve well-formed check's handling of match buffer #16655
Conversation
|
Hi @Lunderberg , I tested my cases with this pr, and they are all working properly now. Thanks for your help! |
- The `T.match_buffer` at the start of a function may contain repeated use of the same data var. For example, a function that must accept two `DLTensor` objects with the same backing allocation. - The `"buffer_bind_scope"` is an older style of match buffer, and may be the point of definition for variables.
Lunderberg
force-pushed
the
tir_well_formed_bugfixes
branch
from
848d539 to
8ffcead
Compare
| if (auto iter_var = op->node.as<IterVar>(); | ||
| iter_var && (op->attr_key == attr::thread_extent || op->attr_key == attr::virtual_thread)) { | ||
| // Some attributes serve as a source of definition for the | ||
| // tir::Var they annotate. | ||
| context = WithDef(iter_var.value(), path->Attr("node")); | ||
| context.push_back(WithDef(iter_var.value(), path->Attr("node"))); | ||
| } else if (op->attr_key == attr::buffer_bind_scope) { |
There was a problem hiding this comment.
Probably worth commenting that this acts as an older form of MatchBuffer, per the PR description.
There was a problem hiding this comment.
Thank you, and added a comment with description.
| @@ -199,12 +174,23 @@ void TIRVisitorWithPath::VisitStmt_(const LetStmtNode* op, ObjectPath path) { | |||
| void TIRVisitorWithPath::VisitStmt_(const AttrStmtNode* op, ObjectPath path) { | |||
| Visit(op->value, path->Attr("value")); | |||
|
|
|||
| std::optional<DefContext<IterVar>> context = std::nullopt; | |||
| std::vector<std::variant<DefContext<IterVar>, DefContext<Var>>> context; | |||
There was a problem hiding this comment.
Is this ever used? I don't see any reads from it.
There was a problem hiding this comment.
There aren't any reads from it, as it holds a scoped context manager. On destruction, the DefContext<T> object removes items from TIRVisitorWithPath::in_scope_definitions_, and calls the ExitDef handler of the child class.
Also, thank you for pointing this one out. When switching from std::optional to std::vector, I forgot to add a while(context.size()) context.pop_back(); loop in case child classes rely on ExitDef being called in the reverse order from EnterDef.
| @@ -250,7 +236,8 @@ void TIRVisitorWithPath::VisitStmt_(const BufferStoreNode* op, ObjectPath path) | |||
| void TIRVisitorWithPath::VisitStmt_(const BufferRealizeNode* op, ObjectPath path) { | |||
| Visit(op->condition, path->Attr("condition")); | |||
| Visit(op->bounds, path->Attr("bounds")); | |||
| auto context = WithDef(op->buffer, path->Attr("buffer")); | |||
| auto context = WithDefIfUndefined(op->buffer->data, path->Attr("buffer")->Attr("data")); | |||
There was a problem hiding this comment.
I imagine this accounts for the case where a BufferRealize can act as a point of definition?
There was a problem hiding this comment.
That's correct. In cases where the buffer's backing allocation is defined externally, the BufferRealize is an annotation of the bounds where the external buffer is accessed. Otherwise, BufferRealize is an allocation. Prior to this commit, only the external backing allocation was handled.
| @@ -195,6 +204,10 @@ void TIRVisitorWithPath::VisitStmt_(const AttrStmtNode* op, ObjectPath path) { | |||
| Visit(expr.value(), path->Attr("node")); | |||
| } | |||
| Visit(op->body, path->Attr("body")); | |||
|
|
|||
| while (context.size()) { | |||
There was a problem hiding this comment.
Thanks for the change. It's probably worth also putting in a comment that this is to ensure that the defs expire, otherwise it seems like spooky action at a distance. Not 100% sure, as the name "DefContext" does imply that it's an RAII sort of thing.
There was a problem hiding this comment.
Yeah, I tried to follow the existing FooContext naming structure (e.g. arith::ConstraintContext).
) * [TIR] Improve well-formed check's handling of match buffer - The `T.match_buffer` at the start of a function may contain repeated use of the same data var. For example, a function that must accept two `DLTensor` objects with the same backing allocation. - The `"buffer_bind_scope"` is an older style of match buffer, and may be the point of definition for variables. * Improved comment, added context.pop_back()
The
T.match_bufferat the start of a function may contain repeated use of the same data var. For example, a function that must accept twoDLTensorobjects with the same backing allocation.The
"buffer_bind_scope"is an older style of match buffer, and may be the point of definition for variables.A
BufferRealizenode may act either as an annotation of an external buffer (which indices are used), or as a point of definition for a local buffer (allocation extents of that buffer).