Relax AST (apache#2)

Co-authored-by: ZihengJiang <ziheng@apache.org>
Lunderberg · Feb 9, 2023 · 68e15a8 · 68e15a8
1 parent 3bc1ce1
commit 68e15a8
Show file tree

Hide file tree

Showing 19 changed files with 1,181 additions and 8 deletions.
diff --git a/3rdparty/cutlass b/3rdparty/cutlass
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -322,6 +322,7 @@ tvm_file_glob(GLOB_RECURSE RELAY_IR_SRCS
 tvm_file_glob(GLOB_RECURSE RELAY_QNN_SRCS
     src/relay/qnn/*.cc
 )
+
 list(APPEND COMPILER_SRCS ${RELAY_OP_SRCS})
 list(APPEND COMPILER_SRCS ${RELAY_PASS_SRCS})
 list(APPEND COMPILER_SRCS ${RELAY_BACKEND_SRCS})

diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
@@ -135,6 +135,7 @@ class PrimExpr : public BaseExpr {
   TVM_DLL static PrimExpr FromObject_(ObjectRef ref);
 };
 
+class RelayExpr;
 /*!
  * \brief add operator
  *
@@ -367,10 +368,27 @@ class RelayExprNode : public BaseExprNode {
    *       This value is discarded during serialization.
    */
   mutable Type checked_type_ = Type(nullptr);
+
+  /*!
+  * \brief Stores the result of static shape analysis.
+  *
+  * \note The value will be optional if a static shape can not be inferred.
+  * use .shape() instead to acesss an always defined shape expression.
+  */
+  Optional<Array<PrimExpr>> shape_ = Optional<Array<PrimExpr>>();
+
   /*!
    * \return The checked_type
    */
   inline const Type& checked_type() const;
+
+  /*!
+   * \return An expression which corresponds to the shape of the expression.
+   *
+   * Only valid when the expression's type is a Tensor.
+   */
+  inline RelayExpr shape() const;
+
   /*!
    * \brief Check if the inferred(checked) type of the Expr
    *  is backed by a TTypeNode and return it.

diff --git a/include/tvm/node/structural_hash.h b/include/tvm/node/structural_hash.h
@@ -17,7 +17,7 @@
  * under the License.
  */
 /*!
- * \file tvm/node/structural_equal.h
+ * \file tvm/node/structural_hash.h
  * \brief Structural hash class.
  */
 #ifndef TVM_NODE_STRUCTURAL_HASH_H_
@@ -174,7 +174,7 @@ class SHashReducer {
   /*!
    * \brief Push hash of key to the current sequence of hash values.
    * \param key The key to be hashed.
-   * \note This function indicate key could contain var defintions.
+   * \note This function indicates key could contain variable defintions.
    */
   void DefHash(const ObjectRef& key) const { return handler_->SHashReduce(key, true); }
   /*!