Relax AST (#2)

Co-authored-by: ZihengJiang <[email protected]>
tqchen · Jan 14, 2023 · 03fd11e · 03fd11e
1 parent 4e29a2f
commit 03fd11e
Show file tree

Hide file tree

Showing 18 changed files with 1,183 additions and 7 deletions.
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
@@ -133,6 +133,7 @@ class PrimExpr : public BaseExpr {
   TVM_DLL static PrimExpr FromObject_(ObjectRef ref);
 };
 
+class RelayExpr;
 /*!
  * \brief add operator
  *
@@ -365,10 +366,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); }
   /*!