Lunderberg commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r805056854
##########
File path: include/tvm/tir/buffer.h
##########
@@ -55,8 +55,22 @@ class BufferNode : public Object {
Var data;
/*! \brief data type in the content of the tensor */
DataType dtype;
- /*! \brief The shape of the buffer */
+ /*! \brief The type of the buffer prior to flattening
+ *
+ * This contains the shape as it is accessed by
+ * BufferLoad/BufferStore nodes, and used by the low-level code
+ * generators.
+ */
Array<PrimExpr> shape;
+ /*!
+ * \brief Separators between input axes when generating flattened output axes
+ *
+ * For buffers representing flat 1-d memory (e.g. any buffer in
+ * RAM), this should be an empty array. For buffers representing
+ * non-flat memory, each entry in axis_separators should be the
+ * first input axis that is part of a new flattened axis.
+ */
+ Array<IntImm> axis_separators;
Review comment:
There are two separate steps to the transformation, one which re-orders
the existing indices and a later transformation that produces flat memory
buffers. The former is either on-demand (TIR schedules), or during the
building of a PrimFunc (SchedulePostProcToPrimFunc for TE-based schedules).
The latter is applied during the lowering steps, either in StorageFlatten
(TE-based schedules) or FlattenBuffer (TIR-based schedules). We had some
discussion ([link to RFC section with some
details](https://github.com/Lunderberg/tvm-rfcs/blob/data_layout/rfcs/0039-buffer-physical-layout.md#unresolved-questions),
[link to comments on the
RFC](https://github.com/apache/tvm-rfcs/pull/0039#issuecomment-961358086)) on
having the flattening to 1-d be a special case of layout transformations that
are applied on-demand. However, @vinx13 brought up that some of the TIR
lowering steps would need to apply after axis transformation makes the shape of
the transformed buffer available, but before bu
ffer flattening removes that shape information.
Maintaining the `axis_separators` through to the buffer flattening allows
the later transformation to produce buffers representing non-flat memory spaces
(e.g. texture memory on GPUs).
--
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]
