Lunderberg commented on a change in pull request #9091:
URL: https://github.com/apache/tvm/pull/9091#discussion_r716931686
##########
File path: src/tir/transforms/storage_flatten.cc
##########
@@ -507,8 +1340,20 @@ PrimFunc StorageFlatten(PrimFunc func, int
cache_line_size, bool create_bound_at
IRVisitorWithAnalyzer bound_analyzer;
bound_analyzer(fptr->body);
+
+ fptr->body = BufferShapeLegalize(fptr->buffer_map,
&bound_analyzer)(std::move(fptr->body));
+
+ auto stride_legalize = BufferStrideLegalize(fptr->buffer_map,
&bound_analyzer);
+ fptr->body = stride_legalize(std::move(fptr->body));
+ fptr->buffer_map = stride_legalize.UpdatedExternBufferMap();
+
+ fptr->body = ThreadScopePropagate(fptr->buffer_map)(std::move(fptr->body));
+
+ fptr->body = BufferBindUnwrapper(fptr->buffer_map,
&bound_analyzer)(std::move(fptr->body));
Review comment:
I agree that I'd want to unwrap the binds earlier, if I could, to
prevent the amount of rewriting needed to pass the updated buffers along. The
main issue I ran into was for IR definitions that directly reference
`buf.elem_offset`
([example](https://github.com/apache/tvm/blob/main/tests/python/unittest/test_te_schedule_tensor_core.py#L54)).
In order to determine the offset of the bufffer view relative to the `data`
pointer of the parent buffer, the shape and strides of the parent buffer need
to be determined first.
I have two ideas for making the implementation be cleaner and more readable.
One is to change how data are packed in an `AttrStmtNode` for
`buffer_bind_scope`, to use [the `MatchBufferRegion`
class](https://github.com/apache/tvm/blob/main/include/tvm/tir/stmt.h#L1059).
The other is to extend `StmtExprMutator` to act on `BufferNode`, so that the
buffer replacements only need to be done in a single location for each pass
through. As it is, rewriting the `BufferStoreNode`, `BufferLoadNode`,
`AttrStmtNode`, and `CallNode` must be done each time the buffer gets modified,
even if it's just to use the modified Buffer object.
--
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]
