# 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 constraints data structures and solvers""" import tvm_ffi from tvm.runtime import Object from . import _ffi_api @tvm_ffi.register_object("arith.IntGroupBounds") class IntGroupBounds(Object): """Represent integer grouped bounds which are classified into lower bounds (include), upper bounds (include) and equalities. Parameters ---------- coef : tvm.ir.Expr The coefficient. Must be integer type. coef * var >= lower coef * var == equal coef * var >= upper lower : List[tvm.ir.Expr] the lower bounds (include) equal : List[tvm.ir.Expr] equalities upper : List[tvm.ir.Expr] the upper bounds (include) """ def __init__(self, coef, lower, equal, upper): self.__init_handle_by_constructor__(_ffi_api.IntGroupBounds, coef, lower, equal, upper) @staticmethod def from_range(rng): """Construct a IntGroupedBounds by Range. Parameters ---------- rng : tvm.ir.Range Returns ------- ret : Range The constructed range. """ return _ffi_api.IntGroupBounds_from_range(rng) def find_best_range(self): """Return the best range from the grouped bounds. None if (-inf, +inf). """ return _ffi_api.IntGroupBounds_FindBestRange(self) @tvm_ffi.register_object("arith.IntConstraints") class IntConstraints(Object): """Represent a set of integer constraints including variables, their ranges and the relations between them (either equations or inequalities) Parameters ---------- variables : List[tvm.tirx.Var] The variables in the constraints. Must be integers ranges : Map[tvm.tirx.Var, tvm.ir.Range] The ranges of the variables. relations : List[tvm.ir.Expr] The relations between the variables (either equations or inequalities) """ def __init__(self, variables, ranges, relations): self.__init_handle_by_constructor__(_ffi_api.IntConstraints, variables, ranges, relations) @tvm_ffi.register_object("arith.IntConstraintsTransform") class IntConstraintsTransform(Object): """We can have different set of variables to represent the same integer constraints. For example, the following two constrains are equivalent, {a + b = 0 | a >= 0, b >= 0} and {m - n = 0 | m >= 0, n <= 0} This data structure represents the transformation between two equivalent integer constraints. In the above example, src : {a + b = 0 | a >= 0, b >= 0} dst : {m - n = 0 | m >= 0, n <= 0} src_to_dst : {a -> m, b -> -n} dst_to_src : {m -> a, n -> -b} Parameters ---------- src : arith.IntConstraints source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0} dst : arith.IntConstraints integer constraints equivalent to the source, e.g., {m - n = 0 | m >= 0, n <= 0} src_to_dst : Map[tvm.tirx.Var, tvm.ir.Expr] mapping from variables in the src to the variables in the dst, e.g., {a -> m, b -> -n} dst_to_src : Map[tvm.tirx.Var, tvm.ir.Expr] mapping from variables in the dst to the variables in the src, e.g., {m -> a, n -> -b} """ def __init__(self, src, dst, src_to_dst, dst_to_src): self.__init_handle_by_constructor__( _ffi_api.IntConstraintsTransform, src, dst, src_to_dst, dst_to_src ) def solve_linear_equations(equations, variables=None, ranges=None): """Solve linear equations. Parameters ---------- equations: List[tvm.ir.Expr] or IntConstraints The equations of the variables variables : Optional[List[tvm.tirx.Var]] The variables in the system. ranges : Optional[Map[tvm.tirx.Var, tvm.ir.Range]] The ranges of the variables. Returns ------- int_constraints_transform : IntConstraintsTransform New integer constraints, with less variables (if the problem is NOT of full rank), or no variable (if the problem is of full rank), or an empty integer constraints (if the problem is unsolvable). It also provides the ranges of the variables in the new system, as well as inequalities inferred from the problem. You can get the mapping from the original variables to the solution via int_constraints_transform.src_to_dst. """ if isinstance(equations, IntConstraints): return _ffi_api.SolveLinearEquations(equations) return _ffi_api.SolveLinearEquations(variables, ranges, equations) def solve_linear_inequalities(equations, variables=None, ranges=None, deskew_range=False): """Solve linear inequalities. Parameters ---------- equations : List[tvm.ir.Expr] or IntConstraints The inequalities of the variables variables : Optional[List[tvm.tirx.Var]] The variables in the system. ranges : Optional[Map[tvm.tirx.Var, tvm.ir.Range]] The ranges of the variables. deskew_range: Optional[bool] Whether deskew the result ranges to be started from zero. Default false. Returns ------- ret_ranges: IntConstraints or IntConstraintsTransform The result ranges for each variables. Constrains that cannot be transformed to Range will be stored in IntConstraints.relations. If deskew_range is set (=True), the result ranges will be deskewed to be started from zero. New variables are created accordingly therefore IntConstraintsTransform is returned. """ solver = ( _ffi_api.SolveInequalitiesDeskewRange if deskew_range else _ffi_api.SolveInequalitiesToRange ) if isinstance(equations, IntConstraints): assert variables is None assert ranges is None return solver(equations) return solver(variables, ranges, equations)