[TVMScript][Relax] Use tir.SizeVar for shape variables

As reported in apache#16877, shape inference performed during a Relax
transformation may produce different results than shape inference
performed during TVMScript parsing.  While Relax transformations
call `Analyzer::MarkGlobalNonNegValue` for each shape expression, this
is not provided during TVMScript parsing.  As a result, output shapes
that are conditional on the sign of a variable may produce different
results when simplifying.

This commit provides a partial resolution for this issue.  Where prior
to this commit, the TVMScript parser generated a `tir.Var` for each
symbolic variable, the TVMScript parser now provides a `tir.SizeVar`
for symbolic variables that are defined in contexts that require a
non-negative size.

Author: Lunderberg

python/tvm/script/parser/relax/entry.py

@@ -154,6 +154,9 @@ def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> Struc
     def get_symbolic_vars(self) -> Set[str]:
         return {}
 
+    def get_symbolic_size_vars(self) -> Set[str]:
+        return self.get_symbolic_vars()
+
     def asobject(self):
         return self.as_struct_info(None)
 
@@ -176,9 +179,6 @@ class ObjectProxy(StructInfoProxy):
     def __init__(self) -> None:
         pass
 
-    def get_symbolic_vars(self) -> Set[str]:
-        return set()
-
     def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> ShapeStructInfo:
         return ObjectStructInfo()
 
@@ -341,6 +341,12 @@ def get_symbolic_vars(self) -> Set[str]:
         else:
             return set().union(*[p.get_symbolic_vars() for p in self.params])
 
+    def get_symbolic_size_vars(self) -> Set[str]:
+        if self.params is None:
+            return set()
+        else:
+            return set().union(*[p.get_symbolic_size_vars() for p in self.params])
+
     def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> FuncStructInfo:
         if self.ret is None:
             ret = None
@@ -394,6 +400,9 @@ def __init__(
     def get_symbolic_vars(self) -> Set[str]:
         return set().union(*[f.get_symbolic_vars() for f in self.fields])
 
+    def get_symbolic_size_vars(self) -> Set[str]:
+        return set().union(*[f.get_symbolic_size_vars() for f in self.fields])
+
     def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> TupleStructInfo:
         fields = [field.as_struct_info(dict_globals) for field in self.fields]
         return TupleStructInfo(fields)
@@ -480,6 +489,13 @@ def get_symbolic_vars(self) -> Set[str]:
         else:
             return set()
 
+    def get_symbolic_size_vars(self) -> Set[str]:
+        # While variables defined by R.Shape and R.Tensor arguments
+        # are known to be non-negative, R.Prim arguments may be
+        # negative.  Overriding the default implementation of
+        # `get_symbolic_size_vars()`
+        return set()
+
     def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> ShapeStructInfo:
         if self.value is None:
             return PrimStructInfo(dtype=self.dtype)