# 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.
# pylint: disable=invalid-name
"""Arithmetic data structure and utility"""
import enum
from typing import Union
import tvm._ffi
from tvm import tir, ir
from tvm.runtime import Object
from . import _ffi_api
class ProofStrength(enum.IntEnum):
"""Proof strength of the analysis"""
DEFAULT = 0
SYMBOLIC_BOUND = 1
class Extension(enum.Flag):
"""Extensions enabled for RewriteSimplifier
Values should match `RewriteSimplifier::Extensions`
"""
NoExtensions = 0
TransitivelyProveInequalities = 1 << 0
ConvertBooleanToAndOfOrs = 1 << 1
ApplyConstraintsToBooleanBranches = 1 << 2
ComparisonOfProductAndSum = 1 << 3
@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__(_ffi_api.ModularSet, coeff, base)
@tvm._ffi.register_object("arith.ConstIntBound")
class ConstIntBound(Object):
"""Represent constant integer bound
Parameters
----------
min_value : int
The minimum value of the bound.
max_value : int
The maximum value of the bound.
"""
POS_INF = (1 << 63) - 1
NEG_INF = -POS_INF
def __init__(self, min_value, max_value):
self.__init_handle_by_constructor__(_ffi_api.ConstIntBound, min_value, max_value)
class ConstraintScope:
"""Constraint scope.
Parameters
----------
fenter : function
A function that will be called to create an enter context.
Note
----
Do not create object directly, use Analyzer.constraint_scope
"""
def __init__(self, fenter):
self._fenter = fenter
self._fexit = None
def __enter__(self):
self._fexit = self._fenter()
def __exit__(self, ptype, value, trace):
self._fexit()
[文档]
class Analyzer:
"""Integer arithmetic analyzer
This is a stateful analyzer class that can
be used to perform various symbolic integer analysis.
"""
def __init__(self):
_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._get_rewrite_simplify_stats = _mod("get_rewrite_simplify_stats")
self._reset_rewrite_simplify_stats = _mod("reset_rewrite_simplify_stats")
self._canonical_simplify = _mod("canonical_simplify")
self._int_set = _mod("int_set")
self._enter_constraint_context = _mod("enter_constraint_context")
self._can_prove_equal = _mod("can_prove_equal")
self._can_prove = _mod("can_prove")
self._get_enabled_extensions = _mod("get_enabled_extensions")
self._set_enabled_extensions = _mod("set_enabled_extensions")
[文档]
def const_int_bound(self, expr):
"""Find constant integer bound for expr.
Parameters
----------
expr : PrimExpr
The expression.
Returns
-------
bound : ConstIntBound
The result bound
"""
return self._const_int_bound(expr)
[文档]
def modular_set(self, expr):
"""Find a modular set that expr belongs to.
Parameters
----------
expr : PrimExpr
The expression.
Returns
-------
result : ModularSet
The result.
"""
return self._modular_set(expr)
[文档]
def simplify(self, expr, steps=2):
"""Simplify expression via both rewrite and canonicalization.
Parameters
----------
expr : PrimExpr
The expression.
steps : The simplification runs in the order of
rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
param steps controls how many steps to run.
Default is 2, i.e., rewrite_simplify + canonical_simplify.
Returns
-------
result : Expr
The result.
"""
return self._simplify(expr, steps)
[文档]
def rewrite_simplify(self, expr):
"""Simplify expression via rewriting rules.
Parameters
----------
expr : PrimExpr
The expression.
Returns
-------
result : Expr
The result.
"""
return self._rewrite_simplify(expr)
@property
def rewrite_simplify_stats(self):
return self._get_rewrite_simplify_stats()
def reset_rewrite_simplify_stats(self):
self._reset_rewrite_simplify_stats()
[文档]
def canonical_simplify(self, expr):
"""Simplify expression via canonicalization.
Parameters
----------
expr : PrimExpr
The expression.
Returns
-------
result : Expr
The result.
"""
return self._canonical_simplify(expr)
[文档]
def int_set(self, expr, dom_map):
"""Compute a symbolic IntSet that covers expr for all values in dom_map.
Parameters
----------
expr : PrimExpr
The expression.
dom_map : Dict[tvm.tir.Var, tvm.arith.IntSet]
The domain for variables to be relaxed.
Returns
-------
result : IntSet
The result.
"""
return self._int_set(expr, dom_map)
[文档]
def can_prove(self, expr, strength=ProofStrength.DEFAULT):
"""Check whether we can prove expr to be true.
Parameters
----------
expr : PrimExpr
The expression.
strength: ProofStrength
The proof strength
Returns
-------
result : Expr
The result.
"""
return self._can_prove(expr, strength)
[文档]
def bind(self, var: tir.Var, expr: Union[tir.PrimExpr, ir.Range]):
"""Bind a variable to the expression.
Parameters
----------
var : tvm.tir.Var
The variable.
expr : Union[tir.PrimExpr, ir.Range]
The expression or the range to bind to.
"""
return self._bind(var, expr)
[文档]
def constraint_scope(self, constraint):
"""Create a constraint scope.
Parameters
----------
constraint : PrimExpr
The constraint expression.
returns
-------
scope : ConstraintScope
The constraint scope
Examples
--------
.. code-block:: python
x = te.var("x")
analyzer = tvm.arith.Analyzer()
with analzyer.constraint_scope(x % 3 == 0):
# constraint in effect
assert analyzer.modular_set(x).coeff == 3
# constraint no longer in effect
assert analyzer.modular_set(x).coeff != 3
"""
def _fenter():
return self._enter_constraint_context(constraint)
return ConstraintScope(_fenter)
[文档]
def update(self, var, info, override=False):
"""Update infomation about var
Parameters
----------
var : tvm.tir.Var
The variable.
info : tvm.Object
Related information.
override : bool
Whether allow override.
"""
if isinstance(info, ConstIntBound):
self._const_int_bound_update(var, info, override)
else:
raise TypeError("Do not know how to handle type {}".format(type(info)))
[文档]
def can_prove_equal(self, lhs: "PrimExpr", rhs: "PrimExpr"):
"""Whether we can prove that lhs == rhs
Parameters
----------
lhs: PrimExpr
The left-hand side of the comparison
rhs: PrimExpr
The right-hand side of the comparison
Returns
-------
result: bool
Whether we can prove that lhs == rhs
"""
return self._can_prove_equal(lhs, rhs)
@property
def enabled_extensions(self) -> Extension:
"""Return the currently enabled extensions"""
value = self._get_enabled_extensions()
return Extension(value)
@enabled_extensions.setter
def enabled_extensions(self, flags: Union[int, Extension]):
"""Enable extensions for the analyzer
Parameters
----------
flags: Union[int,Extension]
The extensions to enable.
"""
flags = Extension(flags).value
self._set_enabled_extensions(flags)
