[GitHub] [tvm] manupa-arm commented on a change in pull request #9727: [TE][TIR] Implement layout transformations, non-flat memory buffers

[GitBox <gi...@apache.org>]

manupa-arm commented on a change in pull request #9727:
URL: https://github.com/apache/tvm/pull/9727#discussion_r802967813



##########
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:
       What is the reasoning behind preserving this?
   
   The comment sounds like it is for validation. Do we need to maintain a member just for validation or is there more uses than that ?
   

##########
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:
       Maybe this was discussed elsewhere (in which case please point me at that), what is the reasoning behind maintaining this in the IR. I would naively assume shape to be transformed on-demand basis ?

##########
File path: include/tvm/te/schedule.h
##########
@@ -771,6 +835,36 @@ class Singleton : public IterVarRelation {
   TVM_DEFINE_OBJECT_REF_METHODS(Singleton, IterVarRelation, SingletonNode);
 };
 
+/*!
+ * \brief Transform iterator according to some arbitrary expression.
+ */
+class TransformNode : public IterVarRelationNode {
+ public:
+  Array<IterVar> original_variables;

Review comment:
       Should we have docs for these variables ?




-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org