Lunderberg commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r805050148
##########
File path: include/tvm/tir/function.h
##########
@@ -91,11 +91,23 @@ class PrimFuncNode : public BaseFuncNode {
*/
Map<tir::Var, Buffer> buffer_map;
+ /*! \brief The buffer map prior to flattening.
+ *
+ * This contains the buffers as they exists prior to flattening, and
+ * is used for validating an input tensor passed into the packed
+ * API. Any buffer that is present in `buffer_map` but not present
+ * in `preflattened_buffer_map` is assumed to be the same before
+ * and after flattening (e.g. a 1-d tensor that is backed by 1-d
+ * flat memory).
+ */
+ Map<tir::Var, Buffer> preflattened_buffer_map;
Review comment:
That's correct, it's for the validation of user-provided tensors. The
buffer_map defines the buffers that will be used in the body of the function,
and is updated with the flattened buffer definitions either in `StorageFlatten`
for TE-based schedules or `FlattenBuffer` for TIR-based schedules. However,
the unpacking of arguments and insertion of checks to validate user-provided
input isn't added until much later, during the `MakePackedAPI` pass. Unless
the buffer's shape prior to flattening is available, the check would be
generated based on the post-flattening shape, which would be a major
user-facing change. (e.g. Needing to pass a numpy array of shape `[6]` to
match the flattened shape of a `[2,3]` tensor, rather than the original `[2,3]`
shape.)
That said, I'd like to it in future PRs, by keeping the buffer_map
untouched, but defining a flattened buffer that aliases the original buffer.
This will remove the need for the buffer_map, but would make it trickier for
passes to manage which buffer objects have replaced which other buffer objects,
and what relationships need to be maintained. It's not hard to do for any one
pass, but is tedious, and so I want to first improve the ergonomics of
`StmtExprMutator`'s handling of buffers before removing the
preflattened_buffer_map.
--
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]
