Relax AST

CMakeLists.txt

@@ -249,6 +249,7 @@ file(GLOB_RECURSE COMPILER_SRCS
     src/parser/*.cc
     src/printer/*.cc
     src/support/*.cc
+    src/relax/*.cc
     )
 
 file(GLOB CODEGEN_SRCS
@@ -276,6 +277,7 @@ file(GLOB_RECURSE RELAY_IR_SRCS
 file(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})

include/tvm/ir/expr.h

@@ -130,6 +130,7 @@ class PrimExpr : public BaseExpr {
   TVM_DLL static PrimExpr FromObject_(ObjectRef ref);
 };
 
+class RelayExpr;
 /*!
  * \brief Base node of all non-primitive expressions.
  *
@@ -148,10 +149,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.
+  */
+  mutable 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.

include/tvm/ir/op.h

@@ -146,7 +146,7 @@ class OpNode : public RelayExprNode {
   // Internal function to compute if it is primitive op
   bool IsPrimitiveOp_() const {
     const auto& fn_ty = this->op_type;
-    ICHECK(fn_ty.get() != nullptr) << "op_type of " << this->name << "is not registered";
+    ICHECK(fn_ty.get() != nullptr) << "op_type of " << this->name << " is not registered";
     if (fn_ty->type_constraints.size() != 1) return false;
     const TypeRelationNode* rel = fn_ty->type_constraints[0].as<TypeRelationNode>();
     if (rel == nullptr) return false;

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_
@@ -173,7 +173,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); }
   /*!