[REFACTOR][PY] Establish tvm.arith (#4904)

python/tvm/arith/__init__.py

@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Integer bound analysis, simplification and pattern detection."""
+
+from .int_set import IntSet, IntervalSet
+from .analyzer import ModularSet, ConstIntBound, Analyzer
+from .bound import deduce_bound
+from .pattern import detect_linear_equation, detect_clip_bound

python/tvm/arith/_ffi_api.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""FFI APIs for tvm.arith"""
+import tvm._ffi
+
+
+tvm._ffi._init_api("arith", __name__)

python/tvm/arith.py → python/tvm/arith/analyzer.py

@@ -17,42 +17,15 @@
 """Arithmetic data structure and utility"""
 import tvm._ffi
 from tvm.runtime import Object
-
-
-class IntSet(Object):
-    """Represent a set of integer in one dimension."""
-    def is_nothing(self):
-        """Whether the set represent nothing"""
-        return _IntSetIsNothing(self)
-
-    def is_everything(self):
-        """Whether the set represent everything"""
-        return _IntSetIsEverything(self)
-
-
-@tvm._ffi.register_object("arith.IntervalSet")
-class IntervalSet(IntSet):
-    """Represent set of continuous interval [min_value, max_value]
-
-    Parameters
-    ----------
-    min_value : Expr
-        The minimum value in the interval.
-
-    max_value : Expr
-        The maximum value in the interval.
-    """
-    def __init__(self, min_value, max_value):
-        self.__init_handle_by_constructor__(
-            _make_IntervalSet, min_value, max_value)
+from . import _ffi_api
 
 
 @tvm._ffi.register_object("arith.ModularSet")
 class ModularSet(Object):
     """Represent range of (coeff * x + base) for x in Z """
     def __init__(self, coeff, base):
         self.__init_handle_by_constructor__(
-            _make_ModularSet, coeff, base)
+            _ffi_api.ModularSet, coeff, base)
 
 
 @tvm._ffi.register_object("arith.ConstIntBound")
@@ -72,7 +45,7 @@ class ConstIntBound(Object):
 
     def __init__(self, min_value, max_value):
         self.__init_handle_by_constructor__(
-            _make_ConstIntBound, min_value, max_value)
+            _ffi_api.ConstIntBound, min_value, max_value)
 
 
 class ConstraintScope:
@@ -105,11 +78,12 @@ class Analyzer:
     be used to perform various symbolic integer analysis.
     """
     def __init__(self):
-        _mod = _CreateAnalyzer()
+        _mod = _ffi_api.CreateAnalyzer()
         self._const_int_bound = _mod("const_int_bound")
         self._const_int_bound_update = _mod("const_int_bound_update")
         self._bind = _mod("bind")
         self._modular_set = _mod("modular_set")
+        self._simplify = _mod("Simplify")
         self._rewrite_simplify = _mod("rewrite_simplify")
         self._canonical_simplify = _mod("canonical_simplify")
         self._int_set = _mod("int_set")
@@ -120,7 +94,7 @@ def const_int_bound(self, expr):
 
         Parameters
         ----------
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
 
         Returns
@@ -135,7 +109,7 @@ def modular_set(self, expr):
 
         Parameters
         ----------
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
 
         Returns
@@ -145,12 +119,27 @@ def modular_set(self, expr):
         """
         return self._modular_set(expr)
 
+    def simplify(self, expr):
+        """Simplify expression via both rewrite and canonicalization.
+
+        Parameters
+        ----------
+        expr : PrimExpr
+            The expression.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return self._simplify(expr)
+
     def rewrite_simplify(self, expr):
         """Simplify expression via rewriting rules.
 
         Parameters
         ----------
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
 
         Returns
@@ -165,7 +154,7 @@ def canonical_simplify(self, expr):
 
         Parameters
         ----------
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
 
         Returns
@@ -180,7 +169,7 @@ def int_set(self, expr, dom_map):
 
         Parameters
         ----------
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
 
         dom_map : Dict[Var, tvm.arith.IntSet]
@@ -198,10 +187,10 @@ def bind(self, var, expr):
 
         Parameters
         ----------
-        var : tvm.Var
+        var : tvm.tir.Var
             The variable.
 
-        expr : tvm.Expr
+        expr : PrimExpr
             The expression.
         """
         return self._bind(var, expr)
@@ -211,7 +200,7 @@ def constraint_scope(self, constraint):
 
         Parameters
         ----------
-        constraint : tvm.Expr
+        constraint : PrimExpr
             The constraint expression.
 
         returns
@@ -240,7 +229,7 @@ def update(self, var, info, override=False):
 
         Parameters
         ----------
-        var : tvm.Var
+        var : tvm.tir.Var
             The variable.
 
         info : tvm.Object
@@ -254,6 +243,3 @@ def update(self, var, info, override=False):
         else:
             raise TypeError(
                 "Do not know how to handle type {}".format(type(info)))
-
-
-tvm._ffi._init_api("tvm.arith")

python/tvm/arith/bound.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Bound deduction."""
+from . import _ffi_api
+
+
+def deduce_bound(var, cond, hint_map, relax_map):
+    """Deduce the bound of the target variable in the cond.
+
+    Parameters
+    ----------
+    var : Var
+        The target variable to be deduced.
+
+    cond : PrimExpr
+        The condition
+
+    hint_map : Map[Var, IntSet]
+        Domain of variables used to help deduction.
+
+    relax_map : Map[Var, IntSet]
+        The fomain of the variables to be relaxed
+        using the provided domain.
+    """
+    return _ffi_api.DeduceBound(var, cond, hint_map, relax_map)

python/tvm/arith/int_set.py

@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Integer set."""
+import tvm._ffi
+from tvm.runtime import Object
+from . import _ffi_api
+
+
+class IntSet(Object):
+    """Represent a set of integer in one dimension."""
+    def is_nothing(self):
+        """Whether the set represent nothing"""
+        return _ffi_api.IntSetIsNothing(self)
+
+    def is_everything(self):
+        """Whether the set represent everything"""
+        return _ffi_api.IntSetIsEverything(self)
+
+    @staticmethod
+    def vector(vec):
+        """Construct an integer set that covers the vector expr
+
+        Parameters
+        ----------
+        vec : PrimExpr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_vector(vec)
+
+    @staticmethod
+    def single_point(point):
+        """Construct a point set.
+
+        Parameters
+        ----------
+        point : PrimExpr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_single_point(point)
+
+
+@tvm._ffi.register_object("arith.IntervalSet")
+class IntervalSet(IntSet):
+    """Represent set of continuous interval [min_value, max_value]
+
+    Parameters
+    ----------
+    min_value : PrimExpr
+        The minimum value in the interval.
+
+    max_value : PrimExpr
+        The maximum value in the interval.
+    """
+    def __init__(self, min_value, max_value):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntervalSet, min_value, max_value)

python/tvm/arith/pattern.py

@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Detect common patterns."""
+from . import _ffi_api
+
+
+def detect_linear_equation(expr, var_list):
+    """Match `expr = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]`
+
+    Where coeff[i] and base are invariant of var[j] for all i and j.
+
+    Parameters
+    ----------
+    expr : PrimExpr
+        The expression to be matched.
+
+    var_list : List[tvm.tir.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[PrimExpr]
+        A list of co-efficients if the match is successful.
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectLinearEquation(expr, var_list)
+
+
+def detect_clip_bound(expr, var_list):
+    """ Detect if expression corresponds to clip bound of the vars
+
+    Parameters
+    ----------
+    expr : PrimExpr
+        The expression to be matched.
+
+    var_list : List[tvm.tir.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[PrimExpr]
+        `concat([min_value[i], max_value[i]] for i, v in enumerate(var_list))`
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectClipBound(expr, var_list)