use sizevar when convert any to tir (apache#8555)

ylc · Jan 13, 2022 · 722b6ae · 722b6ae
1 parent 0d086c8
commit 722b6ae
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/include/tvm/tir/expr.h b/include/tvm/tir/expr.h
@@ -1126,6 +1126,9 @@ class AnyNode : public PrimExprNode {
   /*! \brief Convert to var. */
   Var ToVar() const { return Var("any_dim", DataType::Int(32)); }
 
+  /*! \brief Convert to SizeVar. */
+  SizeVar ToSizeVar() const { return SizeVar("any_dim", DataType::Int(32)); }
+
   static constexpr const char* _type_key = "tir.Any";
   TVM_DECLARE_FINAL_OBJECT_INFO(AnyNode, PrimExprNode);
 };

diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc
@@ -98,7 +98,8 @@ Array<IndexExpr> GetShape(const Array<IndexExpr>& shape) {
       res.push_back(val);
 #endif  // TVM_INDEX_DEFAULT_I64
     } else if (val->IsInstance<tir::AnyNode>()) {
-      res.push_back(val.as<tir::AnyNode>()->ToVar());
+      // currently all 'any' we meet in shape function are non-negative.
+      res.push_back(val.as<tir::AnyNode>()->ToSizeVar());
     } else {
       res.push_back(val);
     }