# 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 import tvm_ffi from tvm import ir, tirx from tvm.arith import IntSet from tvm.runtime import Object from . import _ffi_api class ProofStrength(enum.IntEnum): """Proof strength of the analysis""" DEFAULT = 0 SYMBOLIC_BOUND = 1 class CompareResult(enum.IntEnum): """Result of a transitive comparison. Values must match the C++ ``arith::CompareResult`` enum. """ INCONSISTENT = 0 EQ = 1 LT = 2 LE = 3 GT = 4 GE = 5 NE = 6 UNKNOWN = 7 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() @tvm_ffi.register_object("arith.Analyzer") class Analyzer(Object): """Integer arithmetic analyzer This is a stateful analyzer class that can be used to perform various symbolic integer analysis. The same analyzer instance can be passed to FFI APIs to share accumulated facts across calls. """ def __init__(self): self.__init_handle_by_constructor__(_ffi_api.Analyzer) @property def is_z3_enabled(self) -> bool: """Whether this build includes the Z3 backend (``USE_Z3=ON``). The Z3-specific methods (:py:meth:`get_smtlib2`, :py:meth:`get_z3_stats`, :py:meth:`set_z3_timeout_ms`, :py:meth:`set_z3_rlimit`) only work when this is ``True``. """ return bool(_ffi_api.AnalyzerIsZ3Enabled(self)) def _check_z3_enabled(self) -> None: if not self.is_z3_enabled: raise RuntimeError( "The Z3 backend is not available in this build. " "Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs." ) def get_smtlib2(self, expr: tirx.Expr | None = None) -> str: """Get the current Z3 problem in SMT-LIB2 format. Raises ------ RuntimeError If TVM was built without Z3 (``USE_Z3=OFF``), since there is no solver state to export. Use :py:attr:`is_z3_enabled` to check first. Parameters ---------- expr : Optional[Expr] The expression to prove. If provided, its negation is added to the problem. """ self._check_z3_enabled() return _ffi_api.AnalyzerGetSMTLIB2(self, expr) def set_z3_timeout_ms(self, timeout_ms: int) -> None: """Set Z3 timeout in milliseconds. Raises ------ RuntimeError If TVM was built without Z3 (``USE_Z3=OFF``). Parameters ---------- timeout_ms : int The timeout in milliseconds. """ self._check_z3_enabled() _ffi_api.AnalyzerSetZ3TimeoutMs(self, timeout_ms) def set_z3_rlimit(self, rlimit: int) -> None: """Set Z3 resource limit. The resource limit gives deterministic solver budgeting (unlike a wall clock timeout). A value of ``0`` disables the limit. Raises ------ RuntimeError If TVM was built without Z3 (``USE_Z3=OFF``). Parameters ---------- rlimit : int The resource limit. """ self._check_z3_enabled() _ffi_api.AnalyzerSetZ3RLimit(self, rlimit) def get_z3_stats(self) -> str: """Get Z3 solver statistics. Raises ------ RuntimeError If TVM was built without Z3 (``USE_Z3=OFF``). Returns ------- stats : str The Z3 statistics. """ self._check_z3_enabled() return _ffi_api.AnalyzerGetZ3Stats(self) def const_int_bound(self, expr: tirx.Expr) -> ConstIntBound: """Find constant integer bound for expr. Parameters ---------- expr : Expr The expression. Returns ------- bound : ConstIntBound The result bound """ return _ffi_api.AnalyzerConstIntBound(self, expr) def const_int_bound_is_bound(self, var: tirx.Var) -> bool: """Check if a variable is bound to a range. Parameters ---------- var : tvm.tirx.Var The variable. Returns ------- result : bool Whether the variable is bound to a range. """ return _ffi_api.AnalyzerConstIntBoundIsBound(self, var) def modular_set(self, expr: tirx.Expr) -> ModularSet: """Find a modular set that expr belongs to. Parameters ---------- expr : Expr The expression. Returns ------- result : ModularSet The result. """ return _ffi_api.AnalyzerModularSet(self, expr) def simplify(self, expr: tirx.Expr, steps: int = 2) -> tirx.Expr: """Simplify expression via both rewrite and canonicalization. Parameters ---------- expr : Expr 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 _ffi_api.AnalyzerSimplify(self, expr, steps) def clone(self) -> "Analyzer": """Return a deep copy of this analyzer with independent state. The returned analyzer carries the same accumulated facts (variable bounds, modular sets, bindings, integer-set domains, literal constraints and transitive comparisons) as this one, but owns its own state: binding or simplifying on either analyzer afterwards does not affect the other. Unlike copying the handle, this is a true deep copy. Do not call this while a constraint scope is active on this analyzer. Returns ------- result : Analyzer A new analyzer holding an independent copy of the facts. """ return _ffi_api.AnalyzerClone(self) def rewrite_simplify(self, expr: tirx.Expr) -> tirx.Expr: """Simplify expression via rewriting rules. Parameters ---------- expr : Expr The expression. Returns ------- result : Expr The result. """ return _ffi_api.AnalyzerRewriteSimplify(self, expr) @property def rewrite_simplify_stats(self): return _ffi_api.AnalyzerGetRewriteSimplifyStats(self) def reset_rewrite_simplify_stats(self): _ffi_api.AnalyzerResetRewriteSimplifyStats(self) def canonical_simplify(self, expr: tirx.Expr) -> tirx.Expr: """Simplify expression via canonicalization. Parameters ---------- expr : Expr The expression. Returns ------- result : Expr The result. """ return _ffi_api.AnalyzerCanonicalSimplify(self, expr) def int_set(self, expr: tirx.Expr, dom_map: dict[tirx.Var, IntSet] | None = None) -> IntSet: """Compute a symbolic IntSet that covers expr for all values in dom_map. Parameters ---------- expr : Expr The expression. dom_map : Optional[Dict[tvm.tirx.Var, tvm.arith.IntSet]] The domain for variables to be relaxed. When omitted, the analyzer uses the domains of the variables already bound to it. Returns ------- result : IntSet The result. """ return _ffi_api.AnalyzerIntSet(self, expr, dom_map) def can_prove(self, expr: tirx.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool: """Check whether we can prove expr to be true. Parameters ---------- expr : Expr The expression. strength: ProofStrength The proof strength. When TVM is built with Z3 (``USE_Z3=ON``), the optional Z3 fallback is only consulted at ``SYMBOLIC_BOUND`` or higher, after the native analyzers fail to prove the predicate. Returns ------- result : Expr The result. """ return _ffi_api.AnalyzerCanProve(self, expr, strength) def set_maximum_rewrite_steps(self, maximum: int) -> None: """Set the maximum allowed number of rewrite-simplify steps. When a positive limit is set, the simplifier raises an exception once it exceeds that number of rewrite steps. This is useful for guarding against performance regressions in tests. Parameters ---------- maximum : int The maximum number of rewrite steps, or a non-positive value to allow an unlimited number of steps. """ _ffi_api.AnalyzerSetMaximumRewriteSteps(self, maximum) def bind( self, var: tirx.Var, expr: tirx.Expr | ir.Range, allow_override: bool = False, ) -> None: """Bind a variable to the expression. Parameters ---------- var : tvm.tirx.Var The variable. expr : Union[tirx.Expr, ir.Range] The expression or the range to bind to. allow_override : bool Whether to allow overriding an existing binding for the variable. """ return _ffi_api.AnalyzerBind(self, var, expr, allow_override) def constraint_scope(self, constraint: tirx.Expr) -> ConstraintScope: """Create a constraint scope. Parameters ---------- constraint : Expr The constraint expression. returns ------- scope : ConstraintScope The constraint scope Examples -------- .. code-block:: python x = te.var("x") analyzer = tvm.arith.Analyzer() with analyzer.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 _ffi_api.AnalyzerEnterConstraintContext(self, constraint) return ConstraintScope(_fenter) def update( self, var: tirx.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False ) -> None: """Update information about var. Parameters ---------- var : tvm.tirx.Var The variable. info : Union[ConstIntBound, ModularSet, IntSet] Related information. A ``ConstIntBound`` updates the constant integer bound, a ``ModularSet`` updates the modular set, and an ``IntSet`` updates the integer-set domain of ``var``. override : bool Whether allow override. """ if isinstance(info, ConstIntBound): _ffi_api.AnalyzerConstIntBoundUpdate(self, var, info, override) elif isinstance(info, ModularSet): _ffi_api.AnalyzerModularSetUpdate(self, var, info, override) elif isinstance(info, IntSet): _ffi_api.AnalyzerIntSetUpdate(self, var, info, override) else: raise TypeError(f"Do not know how to handle type {type(info)}") def can_prove_equal(self, lhs: tirx.Expr, rhs: tirx.Expr) -> bool: """Whether we can prove that lhs == rhs Parameters ---------- lhs: Expr The left-hand side of the comparison rhs: Expr The right-hand side of the comparison Returns ------- result: bool Whether we can prove that lhs == rhs """ return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs) def try_compare( self, lhs: tirx.Expr, rhs: tirx.Expr, propagate_inequalities: bool = True ) -> CompareResult: """Compare lhs and rhs using previously provided known comparisons. Parameters ---------- lhs : Expr The left-hand side of the comparison. rhs : Expr The right-hand side of the comparison. propagate_inequalities : bool If true, attempt to find a sequence of transitive inequalities that allow lhs and rhs to be compared. Returns ------- result : CompareResult The most specific result that can be proven about the comparison. Returns ``CompareResult.UNKNOWN`` when nothing can be proven. """ return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs, propagate_inequalities)) @property def enabled_extensions(self) -> Extension: """Return the currently enabled extensions""" value = _ffi_api.AnalyzerGetEnabledExtensions(self) return Extension(value) @enabled_extensions.setter def enabled_extensions(self, flags: int | Extension): """Enable extensions for the analyzer Parameters ---------- flags: Union[int,Extension] The extensions to enable. """ flags = Extension(flags).value _ffi_api.AnalyzerSetEnabledExtensions(self, flags)