This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new c1bf9ce656 [Refactor][Arith] Move conditional bounds into S-TIR
(#20355)
c1bf9ce656 is described below
commit c1bf9ce6561929b40c12058931c43c47a893024b
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 16 07:23:01 2026 -0400
[Refactor][Arith] Move conditional bounds into S-TIR (#20355)
S-TIR buffer compaction and block-access analysis are the remaining
compiler users of the integer-inequality solver. Move their scoped
conditional bounds and the required solver implementation into private
S-TIR analysis files, keeping the general arithmetic analyses in arith.
Remove the retired equation and deskew solvers, their
constraint-transform machinery, and the old public C++/Python/FFI
interfaces. Retain the useful conditional-bound examples through
block-access analysis.
---
include/tvm/arith/int_solver.h | 328 -------------
python/tvm/arith/__init__.py | 1 -
python/tvm/arith/int_solver.py | 183 -------
python/tvm/backend/cuda/cpp/asm.py | 24 +-
python/tvm/backend/cuda/op.py | 7 +-
python/tvm/testing/utils.py | 57 ---
src/arith/int_constraints.cc | 300 ------------
src/arith/presburger_set.cc | 1 -
src/arith/solve_linear_equation.cc | 489 -------------------
.../analysis/conditional_bounds.cc} | 540 +++++++++++++--------
src/s_tir/analysis/conditional_bounds.h | 82 ++++
.../analysis/sblock_access_region_detector.cc | 14 +-
src/s_tir/transform/compact_buffer_region.cc | 2 +-
src/tirx/script/printer/stmt.cc | 2 +
src/tirx/transform/ir_utils.cc | 145 ------
src/tirx/transform/ir_utils.h | 43 --
.../arith/test_arith_solve_linear_equations.py | 186 -------
.../arith/test_arith_solve_linear_inequality.py | 223 ---------
.../s_tir/analysis/test_sblock_access_region.py | 67 +++
tests/python/tirx/codegen/test_cuda_wait_until.py | 8 +-
20 files changed, 511 insertions(+), 2191 deletions(-)
diff --git a/include/tvm/arith/int_solver.h b/include/tvm/arith/int_solver.h
deleted file mode 100644
index ccca3dcbcb..0000000000
--- a/include/tvm/arith/int_solver.h
+++ /dev/null
@@ -1,328 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/arith/int_solver.h
- * \brief integer constraints data structures and solvers
- */
-#ifndef TVM_ARITH_INT_SOLVER_H_
-#define TVM_ARITH_INT_SOLVER_H_
-
-#include <tvm/ir/expr.h>
-#include <tvm/ir/prim/expr.h>
-
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-#include "analyzer.h"
-
-namespace tvm {
-namespace arith {
-
-// According to experiments two best simplifications orders were can->rw and
rw->can->rw,
-// but rw->can->rw is better for a couple of cases.
-// Also we should end with rw because it factors multipliers out.
-constexpr int kSimplifyRewriteCanonicalRewrite = 3;
-
-/*!
- * \brief Represent integer grouped bounds which are classified into
- * lower bounds (inclusive), upper bounds (inclusive) and equalities.
- * It also contains coefficient as a multiplier for the bounds, i.e.,
- * coef * var >= lower
- * coef * var == equal
- * coef * var <= upper
- * \sa IntGroupBounds
- */
-class IntGroupBoundsNode : public ffi::Object {
- public:
- PrimExpr coef;
- ffi::Array<PrimExpr> lower;
- ffi::Array<PrimExpr> equal;
- ffi::Array<PrimExpr> upper;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IntGroupBoundsNode>()
- .def_ro("coef", &IntGroupBoundsNode::coef)
- .def_ro("lower", &IntGroupBoundsNode::lower)
- .def_ro("equal", &IntGroupBoundsNode::equal)
- .def_ro("upper", &IntGroupBoundsNode::upper);
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntGroupBounds",
IntGroupBoundsNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to IntGroupBoundsNode.
- * \sa IntGroupBoundsNode
- */
-class IntGroupBounds : public ffi::ObjectRef {
- public:
- /*!
- * \brief Constructor by fields
- * \param coef The coefficient. Must be integer.
- * coef * var >= lower
- * coef * var == equal
- * coef * var >= upper
- * \param lower the lower bounds (include)
- * \param equal equalities
- * \param upper the upper bounds (include)
- */
- TVM_DLL IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower,
ffi::Array<PrimExpr> equal,
- ffi::Array<PrimExpr> upper);
-
- /*!
- * \brief Construct bounds from a range.
- * \param r The range
- * \return constructed bounds.
- */
- static IntGroupBounds FromRange(const Range& r);
-
- /*!
- * \brief Find the best range from the grouped bounds.
- * \param vranges_addl additional variable ranges that help infer the best
range.
- * \return The best range (has the least difference between the lower bound
and upper bound).
- * undefined if (-inf, +inf).
- */
- Range FindBestRange(const ffi::Map<Var, Range>& vranges_addl = {}) const;
-
- /*!
- * \brief Combine the bounds with another range.
- * \param r range to be combined.
- * \return combined bounds.
- */
- IntGroupBounds operator+(const Range& r);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntGroupBounds, ffi::ObjectRef,
IntGroupBoundsNode);
-};
-
-/*!
- * \brief Represent integer constrains including (integer) variables, their
ranges and
- * the relations between them (either equations or inequalities).
- * \sa LinearSystem
- */
-class IntConstraintsNode : public ffi::Object {
- public:
- // e.g., \alpha, \beta, must be integers
- ffi::Array<PrimVar> variables;
- // e.g., 1 <= \alpha <= N, etc.
- // it is absolutely ok to include ranges for parameters
- // (variables that are not in this->variables) in this map
- ffi::Map<Var, Range> ranges;
- // linear equalities or inequalities
- // e.g., A \alpha = \beta or A \alpha <= \beta
- ffi::Array<PrimExpr> relations;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IntConstraintsNode>()
- .def_ro("variables", &IntConstraintsNode::variables)
- .def_ro("ranges", &IntConstraintsNode::ranges)
- .def_ro("relations", &IntConstraintsNode::relations);
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraints",
IntConstraintsNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to IntConstraintsNode.
- * \sa IntConstraintsNode
- */
-class IntConstraints : public ffi::ObjectRef {
- public:
- /*!
- * \brief Constructor by fields
- * \param variables The variables in the constraints, must be integers.
- * \param ranges The ranges of the variables.
- * \param relations The linear relations between the variables
- * (either equations or inequalities)
- */
- TVM_DLL IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var, Range>
ranges,
- ffi::Array<PrimExpr> relations);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraints, ffi::ObjectRef,
IntConstraintsNode);
-};
-
-/*!
- * \brief We can have different set of variables to represent the same
constraints.
- * For example, the following two systems 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 linear systems.
- * 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}
- * \sa IntConstraintsTransform
- */
-class IntConstraintsTransformNode : public ffi::Object {
- public:
- IntConstraints src;
- IntConstraints dst;
- ffi::Map<Var, PrimExpr> src_to_dst;
- ffi::Map<Var, PrimExpr> dst_to_src;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IntConstraintsTransformNode>()
- .def_ro("src", &IntConstraintsTransformNode::src)
- .def_ro("dst", &IntConstraintsTransformNode::dst)
- .def_ro("src_to_dst", &IntConstraintsTransformNode::src_to_dst)
- .def_ro("dst_to_src", &IntConstraintsTransformNode::dst_to_src);
- }
-
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraintsTransform",
IntConstraintsTransformNode,
- ffi::Object);
-};
-
-/*!
- * \brief Managed reference to IntConstraintsTransformNode.
- * \sa IntConstraintsTransformNode
- */
-class IntConstraintsTransform : public ffi::ObjectRef {
- public:
- /*!
- * \brief Constructor by fields
- * \param src source integer constraints, e.g., {a + b = 0 | a >= 0,
b >= 0}
- * \param dst integer constraints equivalent to the source,
- * e.g., {m - n = 0 | m >= 0, n <= 0}
- * \param src_to_dst mapping from variables in the \p src to the variables
in the \p dst,
- * e.g., {a -> m, b -> -n}
- * \param dst_to_src mapping from variables in the \p dst to the variables
in the \p src,
- * e.g., {m -> a, n -> -b}
- */
- TVM_DLL IntConstraintsTransform(IntConstraints src, IntConstraints dst,
- ffi::Map<Var, PrimExpr> src_to_dst,
- ffi::Map<Var, PrimExpr> dst_to_src);
-
- /*!
- * \brief Chain-compose two IntConstraintsTransform together.
- * this->dst must be the same as other->src.
- * @param other another IntConstraintsTransform whose src is same as
this->dst.
- * @return composed IntConstraintsTransform(this->src, other->dst)
- * with its variables and ranges are properly modified.
- */
- IntConstraintsTransform operator+(const IntConstraintsTransform& other)
const;
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraintsTransform,
ffi::ObjectRef,
- IntConstraintsTransformNode);
-};
-
-typedef std::pair<ffi::Map<Var, IntGroupBounds>, ffi::Array<PrimExpr>>
PartialSolvedInequalities;
-
-/*!
- * \brief Obtain Smith Normal Form of linear equation A x = y.
- * Smith Normal Form of matrix A_{mxn} is S_{mxn} = U_{mxm} A_{mxn}
V_{nxn},
- * in which S_{mxn} is diag(s1, s2, ..., sr, 0, ..., 0) and r is the
rank of A.
- * NOTE: Although in standard Smith Normal Form the diagonal elements
satisfy
- * s_i | s_{i+1} (| means divides), the implement here does not
guarantee it.
- * TODO(yzhliu): From sergei-grechanik:
- * computing the proper Smith normal form may improve stability of
automatic
- * differentiation (generating the same gradient code for slightly different
but equivalent input
- * code U_{mxm} and V_{nxn} are invertible matrices. This function modifies \p
S to be S_{mxn}, \p V
- * to be V_{nxn}, \p y to be U_{mxm} y_{mx1} and \p x to be V^{-1} x. \param S
the original
- * A_{mxn}, it will be modified to S_{mxn} \param V an identity matrix, it
will be modified to
- * V_{nxn} \param x the x in A x = y. it will be modified to V^{-1}_{nxn}
x_{nx1} \param y the y
- * in A x = y. it will be modified to U_{mxm} y_{mx1}
- */
-void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S,
std::vector<std::vector<int64_t>>* V,
- std::vector<PrimExpr>* x, std::vector<PrimExpr>* y);
-
-/*!
- * \brief Solve linear equations.
- * \param system_to_solve the variables to solve, their ranges, and a list of
equations.
- * \return A new linear system, with less variables (if \p system_to_solve is
NOT of full rank),
- * or no variable (if \p system_to_solve is of full rank),
- * or an empty linear system (if \p system_to_solve is unsolvable).
- * It also provides the ranges of the variables in the new system,
- * as well as inequalities inferred from the \p system_to_solve.
- * You can get the mapping from the original variables to the
solution via ret->src_to_dst.
- */
-IntConstraintsTransform SolveLinearEquations(const IntConstraints&
system_to_solve);
-
-/*!
- * \brief Solve linear inequalities.
- * \param system_to_solve the variables to solve, their ranges, and a list of
inequalities.
- * The inequalities are rewritten using Fourier-Motzkin elimination.
- * This function takes an array of (in)equalities and an array of
variables, and essentially
- * rewrites the (in)equalities into an array of (in)equalities of the
following form,
- *
- * x0 >= f0(x1, x2, ..., xn)
- * x0 <= g0(x1, x2, ..., xn)
- * x1 >= f1(x2, ..., xn)
- * x1 <= g1(x2, ..., xn)
- * ...
- * xn >= fn() // just a constant
- * xn <= gn() // just a constant
- *
- * \return A map of variables and their solved bounds,
- * and constrains that cannot be solved to bounds.
- */
-PartialSolvedInequalities SolveLinearInequalities(const IntConstraints&
system_to_solve);
-
-/*!
- * \brief Combine the information into an array of (in)equalities.
- * \param variables The variables in \p bounds.
- * It is used to determine the iteration order to avoid indeterministic
results.
- * \param bounds grouped boundary of the variables.
- * \param relations other relations.
- */
-ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables,
- const ffi::Map<Var, IntGroupBounds>& bounds,
- const ffi::Array<PrimExpr>& relations);
-
-/*!
- * \brief Solve linear inequalities and infer the range of each variable.
- * \param system_to_solve the variables to solve, their ranges, and a list of
inequalities.
- * \return The result ranges for each variables.
- * The returned IntConstraints(variables, ranges, relations) contains,
- * 1. variables - the variables that have been solved.
- * 2. ranges - the best range of each variable.
- * 3. relations - constraints that cannot be transformed to
- * Range will be stored in relations.
- */
-IntConstraints SolveInequalitiesToRange(const IntConstraints& system_to_solve);
-
-/*!
- * \brief Solve linear inequalities and deskew the ranges towards zero.
- * \param system_to_solve the variables to solve, their ranges, and a list of
inequalities.
- * \return A transform (src IntConstraints -> dst IntConstraints)
- * from original variables to a set of new variables.
- * The ranges of new variables always start from zero,
- * their extents are solved from \p system_to_solve.
- * src IntConstraints is the same as \p system_to_solve.
- * dst IntConstraints(variables, ranges, relations) contains,
- * 1. variables - the variables that have been solved.
- * 2. ranges - the best range (start from zero) of each variable.
- * 3. relations - constraints that cannot be transformed to
- * Range will be stored in relations.
- * Variable mapping can be obtained from
- * IntConstraintsTransform.src_to_dst and
IntConstraintsTransform.dst_to_src.
- */
-IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints&
system_to_solve);
-
-} // namespace arith
-} // namespace tvm
-#endif // TVM_ARITH_INT_SOLVER_H_
diff --git a/python/tvm/arith/__init__.py b/python/tvm/arith/__init__.py
index e2a4bf664b..b646a6bf69 100644
--- a/python/tvm/arith/__init__.py
+++ b/python/tvm/arith/__init__.py
@@ -37,7 +37,6 @@ from .analyzer import (
)
from .bound import deduce_bound
from .pattern import detect_linear_equation, detect_clip_bound
-from .int_solver import solve_linear_equations, solve_linear_inequalities
from .iter_affine_map import IterMapExpr, IterMark, IterSplitExpr, IterSumExpr
from .iter_affine_map import (
detect_iter_map,
diff --git a/python/tvm/arith/int_solver.py b/python/tvm/arith/int_solver.py
deleted file mode 100644
index a50d6818fd..0000000000
--- a/python/tvm/arith/int_solver.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# 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)
diff --git a/python/tvm/backend/cuda/cpp/asm.py
b/python/tvm/backend/cuda/cpp/asm.py
index b2c94b47f5..ccbe995be1 100644
--- a/python/tvm/backend/cuda/cpp/asm.py
+++ b/python/tvm/backend/cuda/cpp/asm.py
@@ -35,7 +35,6 @@ from ..codegen.schema import device_intrinsic
from ..codegen.types import PTXDataType
from ..codegen.utils import parse_str
-
# =============================================================================
# Declared synchronization words. The four direct forms emit exactly what their
# raw PTX spellings do; only the operation's identity differs, which is what
@@ -109,9 +108,10 @@ def _wait_until_word_suffix(ptr, requested, what):
# A 128-bit word spans two 64-bit elements, so a pointer into the pair is
# how a kernel names it; asking for `.b128` there is not a respelling of
# the pointee but the statement that the word is the wider one.
- if requested == "b128" and _WAIT_UNTIL_PTX_WIDTH.get(
- _WAIT_UNTIL_SCALARS.get(pointee, "")
- ) in (32, 64):
+ if requested == "b128" and
_WAIT_UNTIL_PTX_WIDTH.get(_WAIT_UNTIL_SCALARS.get(pointee, "")) in (
+ 32,
+ 64,
+ ):
return requested
if pointee in _WAIT_UNTIL_SCALARS:
return _wait_until_scalar_suffix(pointee, requested)
@@ -191,26 +191,20 @@ def cuda_wait_until(dst, ptr, condition, scope, space,
ptx_type, backoff_ns):
# ISA 8.4.2 puts it and `.relaxed` in one class, and measured they are
# within 0.1% everywhere -- but `.relaxed.<scope>` says the scope out loud
# instead of resting on `volatile` meaning `.sys`.
- load_call, tags = _wait_until_forward(
- f"ld.relaxed.{scope}.{space}.{suffix}", dst, ptr
- )
+ load_call, tags =
_wait_until_forward(f"ld.relaxed.{scope}.{space}.{suffix}", dst, ptr)
# The helper is named after the poll, which is the load that repeats; the
# closing acquire hangs off that name with an `_acquire` suffix.
load_name = parse_str(load_call.args[0])
name = load_name.replace("ptx_ld_", "cuda_wait_until_")
source = load_call.args[-1].value.replace(load_name, name + "_load")
- acquire_call, _ = _wait_until_forward(
- f"ld.acquire.{scope}.{space}.{suffix}", dst, ptr
- )
+ acquire_call, _ =
_wait_until_forward(f"ld.acquire.{scope}.{space}.{suffix}", dst, ptr)
acquire_name = parse_str(acquire_call.args[0])
acquire_source = acquire_call.args[-1].value.replace(acquire_name, name +
"_acquire")
source += "\n" + acquire_source
# NVRTC has no `__typeof__`, and `decltype` on the destination yields a
# reference that cannot be declared uninitialized, so the scratch takes
# the C type the generated helper already spells in its signature.
- scratch_type = re.search(
- rf"void\s+{re.escape(name)}_acquire\(\s*([\w:]+)\s*&", acquire_source
- )
+ scratch_type =
re.search(rf"void\s+{re.escape(name)}_acquire\(\s*([\w:]+)\s*&", acquire_source)
if scratch_type is None: # pragma: no cover - the helper shape is fixed
raise RuntimeError(f"cannot read the destination type of
{name}_acquire")
closing = (
@@ -247,7 +241,7 @@ def cuda_wait_until(dst, ptr, condition, scope, space,
ptx_type, backoff_ns):
source += (
f"\n#define {name}(dst, ptr, predicate) "
f"do {{ {name}_load((dst), (ptr)); if (!(predicate)) {{ "
- f"_Pragma(\"unroll 1\") "
+ f'_Pragma("unroll 1") '
f"do {{ {name}_load((dst), (ptr)); }} while (!(predicate)); }}"
f"{closing} }} while (0)\n"
)
@@ -260,7 +254,7 @@ def cuda_wait_until(dst, ptr, condition, scope, space,
ptx_type, backoff_ns):
source += (
f"\n#define {name}(dst, ptr, predicate, backoff_ns) "
f"do {{ {name}_load((dst), (ptr)); if (!(predicate)) {{ "
- f"_Pragma(\"unroll 1\") "
+ f'_Pragma("unroll 1") '
f"while (1) {{ __nanosleep(backoff_ns); {name}_load((dst), (ptr));
"
f"if (predicate) break; }} }}"
f"{closing} }} while (0)\n"
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 0d6ea8806c..8c6a6a02ef 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -421,13 +421,10 @@ def _validate_wait_until_attrs(scope, space,
ptx_type=None):
if space == "shared"
else ""
)
- raise ValueError(
- f"invalid space={space!r}; expected one of
{_WAIT_UNTIL_SPACE}{detail}"
- )
+ raise ValueError(f"invalid space={space!r}; expected one of
{_WAIT_UNTIL_SPACE}{detail}")
if ptx_type is not None and ptx_type not in _WAIT_UNTIL_PTX_TYPES:
raise ValueError(
- f"invalid ptx_type={ptx_type!r}; expected one of "
- f"{tuple(sorted(_WAIT_UNTIL_PTX_TYPES))}"
+ f"invalid ptx_type={ptx_type!r}; expected one of
{tuple(sorted(_WAIT_UNTIL_PTX_TYPES))}"
)
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index 076a12c361..05f1e97c24 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -337,63 +337,6 @@ def check_bool_expr_is_true(bool_expr, vranges, cond=None):
)
-def check_int_constraints_trans_consistency(constraints_trans, vranges=None):
- """Check IntConstraintsTransform is a bijective transformation.
-
- Parameters
- ----------
- constraints_trans : arith.IntConstraintsTransform
- Integer constraints transformation
- vranges: Dict[tvm.tirx.Var, tvm.ir.Range]
- Free variables and their ranges
- """
- if vranges is None:
- vranges = {}
-
- def _check_forward(constraints1, constraints2, varmap, backvarmap):
- ana = tvm.arith.Analyzer()
- all_vranges = vranges.copy()
- all_vranges.update({v: r for v, r in constraints1.ranges.items()})
-
- # Check that the transformation is injective
- cond_on_vars = tvm.tirx.const(1, "bool")
- for v in constraints1.variables:
- if v in varmap:
- # variable mapping is consistent
- v_back = ana.simplify(_substitute_tir_vars(varmap[v],
backvarmap))
- cond_on_vars = tvm.te.all(cond_on_vars, v == v_back)
- # Also we have to check that the new relations are true when old
relations are true
- cond_subst = _substitute_tir_vars(
- tvm.te.all(tvm.tirx.const(1, "bool"), *constraints2.relations),
backvarmap
- )
- # We have to include relations from vranges too
- for v in constraints2.variables:
- if v in constraints2.ranges:
- r = constraints2.ranges[v]
- range_cond = tvm.te.all(v >= r.min, v < r.min + r.extent)
- range_cond = _substitute_tir_vars(range_cond, backvarmap)
- cond_subst = tvm.te.all(cond_subst, range_cond)
- cond_subst = ana.simplify(cond_subst)
- check_bool_expr_is_true(
- tvm.te.all(cond_subst, cond_on_vars),
- all_vranges,
- cond=tvm.te.all(tvm.tirx.const(1, "bool"),
*constraints1.relations),
- )
-
- _check_forward(
- constraints_trans.src,
- constraints_trans.dst,
- constraints_trans.src_to_dst,
- constraints_trans.dst_to_src,
- )
- _check_forward(
- constraints_trans.dst,
- constraints_trans.src,
- constraints_trans.dst_to_src,
- constraints_trans.src_to_dst,
- )
-
-
def _get_targets(target_names=None):
if target_names is None:
target_names = _tvm_test_targets()
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
deleted file mode 100644
index 1d7bdf90e0..0000000000
--- a/src/arith/int_constraints.cc
+++ /dev/null
@@ -1,300 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file int_constraints.cc
- * \brief The integer constraints data structures.
- */
-#include <tvm/arith/analyzer.h>
-#include <tvm/arith/int_solver.h>
-#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/tirx/op.h>
-
-#include <algorithm>
-#include <unordered_map>
-#include <utility>
-
-namespace tvm {
-namespace arith {
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- IntGroupBoundsNode::RegisterReflection();
- IntConstraintsNode::RegisterReflection();
- IntConstraintsTransformNode::RegisterReflection();
-}
-
-ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables,
- const ffi::Map<Var, IntGroupBounds>& bounds,
- const ffi::Array<PrimExpr>& relations) {
- ffi::Array<PrimExpr> res;
- // use variables to keep the order of iteration
- // so as to get rid of any non-determinism.
- TVM_FFI_ICHECK_EQ(variables.size(), bounds.size());
- for (const auto v : variables) {
- TVM_FFI_ICHECK(bounds.count(v));
- const auto& bnds = bounds[v];
- PrimExpr lhs = bnds->coef * v.as_or_throw<PrimExpr>();
- for (const PrimExpr& rhs : bnds->equal) {
- res.push_back(lhs == rhs);
- }
- for (const PrimExpr& rhs : bnds->lower) {
- res.push_back(lhs >= rhs);
- }
- for (const PrimExpr& rhs : bnds->upper) {
- res.push_back(lhs <= rhs);
- }
- }
- for (const PrimExpr& e : relations) {
- res.push_back(e);
- }
- return res;
-}
-
-IntGroupBounds::IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower,
- ffi::Array<PrimExpr> equal,
ffi::Array<PrimExpr> upper) {
- PrimType coef_ty = coef.ty();
- TVM_FFI_ICHECK(coef_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
- << "Coefficient in IntGroupBounds must be integers";
- ffi::ObjectPtr<IntGroupBoundsNode> node =
ffi::make_object<IntGroupBoundsNode>();
- node->coef = std::move(coef);
- node->lower = std::move(lower);
- node->equal = std::move(equal);
- node->upper = std::move(upper);
- data_ = std::move(node);
-}
-
-IntGroupBounds IntGroupBounds::FromRange(const Range& r) {
- Analyzer analyzer;
- PrimExpr coef = tirx::MakeConst(r->min.ty(), 1);
- ffi::Array<PrimExpr> equal;
- ffi::Array<PrimExpr> lower;
- ffi::Array<PrimExpr> upper;
- if (tirx::is_one(r->extent)) {
- equal.push_back(r->min);
- } else {
- lower.push_back(r->min);
- upper.push_back(analyzer->Simplify(r->min + r->extent - 1));
- }
- return IntGroupBounds(coef, lower, equal, upper);
-}
-
-IntGroupBounds IntGroupBounds::operator+(const Range& r) {
- Analyzer analyzer;
- ffi::Array<PrimExpr> equal;
- ffi::Array<PrimExpr> lower;
- ffi::Array<PrimExpr> upper;
- const PrimExpr& coef = operator->()->coef;
- if (tirx::is_one(r->extent)) {
- equal.push_back(analyzer->Simplify(r->min * coef));
- } else {
- lower.push_back(analyzer->Simplify(r->min * coef));
- upper.push_back(analyzer->Simplify((r->min + r->extent - 1) * coef));
- }
- for (const auto& eq : operator->()->equal) equal.push_back(eq);
- for (const auto& lb : operator->()->lower) lower.push_back(lb);
- for (const auto& ub : operator->()->upper) upper.push_back(ub);
- return IntGroupBounds(coef, lower, equal, upper);
-}
-
-Range IntGroupBounds::FindBestRange(const ffi::Map<Var, Range>& vranges_addl)
const {
- Analyzer analyzer;
- analyzer->Bind(vranges_addl);
-
- std::unordered_map<const VarNode*, IntSet> var_intsets;
- for (auto kv : vranges_addl) {
- var_intsets[kv.first.get()] = IntSet::FromRange(kv.second);
- }
-
- const ffi::Array<PrimExpr>& equal = operator->()->equal;
- const PrimExpr& coef = operator->()->coef;
-
- std::vector<PrimExpr> lowers(equal.begin(), equal.end());
- std::vector<PrimExpr> uppers(equal.begin(), equal.end());
- for (const auto& expr : operator->()->lower) {
- lowers.push_back(expr);
- }
- for (const auto& expr : operator->()->upper) {
- uppers.push_back(expr);
- }
-
- if (lowers.size() == 1 && uppers.size() == 1 && tirx::is_one(coef)) {
- return Range(analyzer->Simplify(lowers[0]), analyzer->Simplify(uppers[0] +
1));
- }
-
- // Here we will try all pairs of lower and upper bounds and find the best
pair, that is, the
- // pair with the minimal difference between the upper and the lower.
- // Note that the bounds are for v, not for v*coef
-
- // The lower bound of the best pair so far
- PrimExpr best_lower;
- // The difference between the upper and the lower of the best pair, maybe
overapproximation
- PrimExpr best_diff_over;
-
- for (const PrimExpr& low : lowers) {
- for (const PrimExpr& upp : uppers) {
- // Since diff may depend on some other variables, we compute its
overapproximation
- ffi::Optional<PrimExpr> diff_over;
- PrimExpr diff_1 = analyzer->Simplify(floordiv(upp - low, coef), 3);
- IntSet diff_set1 = EvalSet(diff_1, var_intsets);
- if (diff_set1.HasUpperBound()) {
- diff_over = analyzer->Simplify(diff_set1.max(), 3);
- }
-
- // low is the lower bound for v*coef, but we need the lower bound for v.
- // We use rounding-up division to compute it. Since we want to use a
single formula
- PrimExpr low_divided = analyzer->Simplify(floordiv(low + coef - 1,
coef), 3);
-
- // Compute another difference which may be more precise (or not).
- PrimExpr diff_2 = analyzer->Simplify(floordiv(upp, coef) - low_divided,
3);
- IntSet diff_set2 = EvalSet(diff_2, var_intsets);
- if (diff_set2.HasUpperBound()) {
- PrimExpr diff_over_2 = analyzer->Simplify(diff_set2.max(), 3);
- diff_over = diff_over.has_value() ? (analyzer->CanProve(diff_over_2 -
diff_over.value() < 0)
- ? diff_over_2
- : diff_over.value())
- : diff_over_2;
- }
-
- // If it is provable that the new one is strictly better than the
current best one,
- // then replace it. Note that we are biased towards earlier pairs which
should be simpler.
- if (diff_over.has_value() && (!best_diff_over.defined() ||
- analyzer->CanProve(diff_over.value() -
best_diff_over < 0))) {
- best_lower = low_divided;
- best_diff_over = diff_over.value();
- }
- }
- }
-
- if (!best_lower.defined()) {
- TVM_FFI_ICHECK(!best_diff_over.defined());
- return Range();
- }
- return Range::FromMinExtent(best_lower, analyzer->Simplify(best_diff_over +
1));
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("arith.IntGroupBounds",
- [](PrimExpr coef, ffi::Array<PrimExpr> lower, ffi::Array<PrimExpr>
equal,
- ffi::Array<PrimExpr> upper) { return IntGroupBounds(coef, lower,
equal, upper); })
- .def("arith.IntGroupBounds_from_range", IntGroupBounds::FromRange)
- .def_packed("arith.IntGroupBounds_FindBestRange", [](ffi::PackedArgs
args, ffi::Any* ret) {
- TVM_FFI_ICHECK(args.size() == 1 || args.size() == 2);
- auto bounds = args[0].cast<IntGroupBounds>();
- if (args.size() == 1) {
- *ret = bounds.FindBestRange();
- } else if (args.size() == 2) {
- *ret = bounds.FindBestRange(args[1].cast<ffi::Map<Var, Range>>());
- }
- });
-}
-
-// Pattern A (RM): auto-default repr from reflection.
-
-IntConstraints::IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var,
Range> ranges,
- ffi::Array<PrimExpr> relations) {
- ffi::ObjectPtr<IntConstraintsNode> node =
ffi::make_object<IntConstraintsNode>();
- if (!variables.defined()) {
- variables = ffi::Array<PrimVar>();
- }
- if (!ranges.defined()) {
- ranges = ffi::Map<Var, Range>();
- }
- TVM_FFI_ICHECK(relations.defined());
- for (const PrimVar& var : variables) {
- TVM_FFI_CHECK(var.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt), TypeError)
- << "Variables in IntConstraints must be integers";
- }
- node->variables = std::move(variables);
- node->ranges = std::move(ranges);
- node->relations = std::move(relations);
- data_ = std::move(node);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def(
- "arith.IntConstraints",
- [](ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
- ffi::Array<PrimExpr> relations) { return IntConstraints(variables,
ranges, relations); });
-}
-
-// Pattern A (RM): auto-default repr from reflection.
-
-IntConstraintsTransform::IntConstraintsTransform(IntConstraints src,
IntConstraints dst,
- ffi::Map<Var, PrimExpr>
src_to_dst,
- ffi::Map<Var, PrimExpr>
dst_to_src) {
- ffi::ObjectPtr<IntConstraintsTransformNode> node =
- ffi::make_object<IntConstraintsTransformNode>();
- node->src = std::move(src);
- node->dst = std::move(dst);
- node->src_to_dst = std::move(src_to_dst);
- node->dst_to_src = std::move(dst_to_src);
- data_ = std::move(node);
-}
-
-IntConstraintsTransform IntConstraintsTransform::operator+(
- const IntConstraintsTransform& other) const {
- TVM_FFI_ICHECK(other->src.same_as(operator->()->dst));
- ffi::Map<Var, PrimExpr> dst_to_src;
- ffi::Map<Var, PrimExpr> src_to_dst;
-
- Analyzer ana_first;
- ana_first->Bind(operator->()->src->ranges);
- auto f_dst_to_src = [this](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto repl = operator->()->dst_to_src.Get(var)) return *std::move(repl);
- return ffi::Unchanged();
- };
- for (auto p : other->dst_to_src) {
- dst_to_src.Set(p.first,
ana_first->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
- p.second, f_dst_to_src)
- .as_or_throw<PrimExpr>()));
- }
-
- Analyzer ana_second;
- ana_second->Bind(other->dst->ranges);
- auto f_src_to_dst = [&other](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto repl = other->src_to_dst.Get(var)) return *std::move(repl);
- return ffi::Unchanged();
- };
- for (auto p : operator->()->src_to_dst) {
- src_to_dst.Set(p.first,
ana_second->Simplify(ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
- p.second, f_src_to_dst)
-
.as_or_throw<PrimExpr>()));
- }
- return IntConstraintsTransform(operator->()->src, other->dst, src_to_dst,
dst_to_src);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("arith.IntConstraintsTransform",
- [](IntConstraints src, IntConstraints dst,
- ffi::Map<Var, PrimExpr> src_to_dst, ffi::Map<Var,
PrimExpr> dst_to_src) {
- return IntConstraintsTransform(src, dst, src_to_dst,
dst_to_src);
- });
-}
-
-// Pattern A (RM): auto-default repr from reflection.
-
-} // namespace arith
-} // namespace tvm
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index 2f38200d61..be79bd3d10 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -24,7 +24,6 @@
#include "presburger_set.h"
#include <tvm/arith/int_set.h>
-#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
diff --git a/src/arith/solve_linear_equation.cc
b/src/arith/solve_linear_equation.cc
deleted file mode 100644
index 40e5e0883e..0000000000
--- a/src/arith/solve_linear_equation.cc
+++ /dev/null
@@ -1,489 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/arith/solve_linear_equation.cc
- * \brief Solve linear equations.
- */
-#include <tvm/arith/analyzer.h>
-#include <tvm/arith/int_solver.h>
-#include <tvm/arith/pattern.h>
-#include <tvm/ffi/dtype.h>
-#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/tirx/op.h>
-
-#include <unordered_set>
-#include <utility>
-
-#include "int_operator.h"
-
-namespace tvm {
-namespace arith {
-
-using namespace tvm::runtime;
-
-void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S,
std::vector<std::vector<int64_t>>* V,
- std::vector<PrimExpr>* x, std::vector<PrimExpr>* y) {
- if (S->empty() || V->empty()) return;
- size_t m = S->size();
- size_t n = (*S)[0].size(); // n is # of variables
- TVM_FFI_ICHECK_EQ(V->size(), n);
- TVM_FFI_ICHECK_EQ((*V)[0].size(), n);
-
- for (size_t index = 0; index < std::min(m, n); ++index) {
- // Here A is partially diagonalized, that is A[i, j] is zero for all i, j
- // such that (i < index) or (j < index), unless (i == j).
- // That is, now we are diagonalizing the submatrix with i >= index and j
>= index
-
- // Find a row with a nonzero element in the index-th column
- // (We also prefer rows where this element has minimal abs value)
- size_t best_i = index;
- for (size_t i = best_i; i < m; ++i) {
- int64_t s_old = (*S)[best_i][index];
- int64_t s_new = (*S)[i][index];
- if (s_new != 0) {
- if (s_old == 0 || std::abs(s_new) < std::abs(s_old)) {
- best_i = i;
- }
- }
- }
- // Move the row we found to the index-th position
- std::swap((*S)[index], (*S)[best_i]);
- std::swap((*y)[index], (*y)[best_i]);
-
- // If the index-th diagonal element is still zero, try to find a column
with nonzero index-th
- // element and move it to the index-th position
- if ((*S)[index][index] == 0) {
- for (size_t j = index + 1; j < n; ++j) {
- if ((*S)[index][j] != 0) {
- for (size_t i = index; i < m; ++i) {
- std::swap((*S)[i][index], (*S)[i][j]);
- }
- // swapping columns corresponds to swapping the corresponding x
- std::swap((*x)[index], (*x)[j]);
- for (size_t i = 0; i < n; ++i) {
- std::swap((*V)[i][index], (*V)[i][j]);
- }
- break;
- }
- }
- }
-
- // If the index-th diagonal element is still zero, then both the index-th
row and the index-th
- // column are completely zero, and we don't need to do anything; just go
to the next index
- if ((*S)[index][index] == 0) {
- continue;
- }
-
- // Now the index-th diagonal element is non-zero and we can zero all the
index-th column
- // below it by subtracting rows from each other
- for (auto i = index + 1; i < m; ++i) {
- if ((*S)[i][index] != 0) {
- int64_t g, a, b;
- // g = a*matrix[index][index] + b*matrix[i][index]
- if ((*S)[i][index] % (*S)[index][index] != 0) {
- g = ExtendedEuclidean((*S)[index][index], (*S)[i][index], &a, &b);
- } else {
- // Explicitly avoid changing the index-th row. This is important to
avoid infinite loop.
- g = (*S)[index][index];
- a = 1;
- b = 0;
- }
-
- // Let m = S[index][index], n = S[i][index], then the following is
true:
- //
- // [ a n/g ][ m/g n/g ] = [ 1 0 ]
- // [ b -m/g ][ b -a ] = [ 0 1 ]
- //
- // Note that the two matrices are integer (since g = gcd(m, n)).
- // We will essentially multiply our matrix on the left by a dilated
and transposed version
- // of the first of these two matrices. The second matrix is not needed
here, however we will
- // use it while zeroing the index-th row.
-
- int64_t m_g = (*S)[index][index] / g;
- int64_t n_g = (*S)[i][index] / g;
-
- // Note that j is the index of the column, not the row
- for (size_t j = index; j < (*S)[i].size(); ++j) {
- // Multiply index-th row by a and add the i-th row multiplied by b
- // This will make the index-th diagonal element equal to the gcd
- int64_t new_index_j = a * (*S)[index][j] + b * (*S)[i][j];
- // This transformation performs zeroing of matrix[i][index]
- int64_t new_i_j = n_g * (*S)[index][j] - m_g * (*S)[i][j];
- (*S)[index][j] = new_index_j;
- (*S)[i][j] = new_i_j;
- }
- // We have to do the same with rhs
- PrimExpr ea = tirx::MakeConst((*y)[index].ty(), a);
- PrimExpr eb = tirx::MakeConst((*y)[i].ty(), b);
- PrimExpr e_m_g = tirx::MakeConst((*y)[i].ty(), m_g);
- PrimExpr e_n_g = tirx::MakeConst((*y)[index].ty(), n_g);
- PrimExpr new_index_rhs = ea * (*y)[index] + eb * (*y)[i];
- PrimExpr new_i_rhs = e_n_g * (*y)[index] - e_m_g * (*y)[i];
- (*y)[index] = new_index_rhs;
- (*y)[i] = new_i_rhs;
- }
- }
-
- bool changed = false;
-
- // Now we have to zero the elements of the index-th row by manipulating
columns.
- // This is more difficult because column manipulation corresponds to
variable manipulation,
- // but the algorithm is essentially the same as before.
- for (size_t j = index + 1; j < n; ++j) {
- if ((*S)[index][j] != 0) {
- int64_t g, a, b;
- // g = a*matrix[index][index] + b*matrix[index][j]
- if ((*S)[index][j] % (*S)[index][index] != 0) {
- g = ExtendedEuclidean((*S)[index][index], (*S)[index][j], &a, &b);
- // During this phase we may disrupt the zeroness of the index-th
column, so we will
- // have to take some action if this might have happened.
- changed = true;
- } else {
- // Explicitly avoid changing the index-th column. This is important
to avoid infinite
- // loop. Note that here we don't have to set `changed` to true since
we don't change the
- // index-th column.
- g = (*S)[index][index];
- a = 1;
- b = 0;
- }
-
- // Let m = S[index][index], n = S[index][j], then the following is
true:
- //
- // [ a n/g ][ m/g n/g ] = [ 1 0 ]
- // [ b -m/g ][ b -a ] = [ 0 1 ]
- //
- // Now we are going to multiply our matrix on the right (to manipulate
columns instead of
- // rows), we will also transform the old_to_new matrix the same way,
and we will use the
- // second matrix to transform new_to_old.
-
- int64_t m_g = (*S)[index][index] / g;
- int64_t n_g = (*S)[index][j] / g;
-
- for (size_t i = index; i < m; ++i) {
- int64_t new_i_index = a * (*S)[i][index] + b * (*S)[i][j];
- int64_t new_i_j = n_g * (*S)[i][index] - m_g * (*S)[i][j];
- (*S)[i][index] = new_i_index;
- (*S)[i][j] = new_i_j;
- }
- // We do exactly the same transformations with V
- for (size_t i = 0; i < n; ++i) {
- int64_t new_i_index = a * (*V)[i][index] + b * (*V)[i][j];
- int64_t new_i_j = n_g * (*V)[i][index] - m_g * (*V)[i][j];
- (*V)[i][index] = new_i_index;
- (*V)[i][j] = new_i_j;
- }
- // And apply reverse transformations to new_to_old.
- PrimExpr ea = tirx::MakeConst((*x)[j].ty(), a);
- PrimExpr eb = tirx::MakeConst((*x)[index].ty(), b);
- PrimExpr e_m_g = tirx::MakeConst((*x)[index].ty(), m_g);
- PrimExpr e_n_g = tirx::MakeConst((*x)[j].ty(), n_g);
- PrimExpr new_index = e_m_g * (*x)[index] + e_n_g * (*x)[j];
- PrimExpr new_j = eb * (*x)[index] - ea * (*x)[j];
- (*x)[index] = new_index;
- (*x)[j] = new_j;
- }
- }
-
- if (changed) {
- // We might have changed the first column, so we have to zero it once
more
- // (or at least check if it's zero), so just perform this iteration once
more.
- index -= 1;
- }
- }
-}
-
-ffi::Map<Var, Range> InferRange(const ffi::Map<Var, PrimExpr>& vars_to_infer,
- const ffi::Array<PrimVar>& ori_vars,
- const ffi::Map<Var, Range>& ori_ranges) {
- // The resulting ranges
- ffi::Map<Var, Range> new_ranges;
-
- std::unordered_set<const VarNode*> ori_vset;
- for (const PrimVar& v : ori_vars) {
- ori_vset.insert(v.get());
- }
-
- std::unordered_map<const VarNode*, IntSet> var_intsets;
- for (const auto& p : ori_ranges) {
- if (!ori_vset.count(p.first.get())) {
- // First of all, fill the new ranges with outer variable ranges
- new_ranges.Set(p.first, p.second);
- }
- // Convert original ranges to IntSets
- var_intsets[p.first.get()] = IntSet::FromRange(p.second);
- }
-
- // Infer ranges for the new variables and add them to the resulting ranges
- for (const auto& p : vars_to_infer) {
- const auto& var = p.first;
- const auto& expr = p.second;
- Range range = EvalSet(expr, var_intsets).CoverRange(Range());
- if (range.defined()) {
- new_ranges.Set(var, range);
- }
- }
- return new_ranges;
-}
-
-// pretty print matrix equation
-void DebugPrint(const std::vector<std::vector<int64_t>>& S,
- const std::vector<std::vector<int64_t>>& V, const
std::vector<PrimExpr>& V_inv_x,
- const std::vector<PrimExpr>& rhs) {
- std::cout << "S:\n";
- for (size_t i = 0; i < S.size(); ++i) {
- for (auto e : S[i]) {
- std::cout << e << "\t";
- }
- std::cout << "\t->\t" << rhs[i];
- std::cout << "\n";
- }
- std::cout << "V:\n";
- for (const auto& r : V) {
- for (auto e : r) {
- std::cout << e << "\t";
- }
- std::cout << "\n";
- }
- std::cout << "V_inv x:\n" << ffi::Array<PrimExpr>(V_inv_x);
- std::cout << "\n" << std::endl;
-}
-
-IntConstraintsTransform SolveLinearEquations(const IntConstraints&
system_to_solve) {
- // m: # of equations
- // n: # of variables
- // we first construct A_{mxn} x_{nx1} = y_{mx1}
- // then get Smith normal form of matrix A,
- // S_{mxn} = U_{mxm} A_{mxn} V_{nxn}
- // => U^{-1} S V^{-1} x = y
- // S V^{-1} x = U y
- std::vector<PrimExpr> Uy; // mx1
- std::vector<std::vector<int64_t>> S; // mxn
- std::vector<std::vector<int64_t>> V; // nxn
- std::vector<PrimExpr> V_inv_x; // V^{-1} x, nx1
- // Conditions we don't know what to do with
- std::vector<PrimExpr> rest;
-
- Analyzer analyzer_problem;
- analyzer_problem->Bind(system_to_solve->ranges);
-
- size_t num_vars = system_to_solve->variables.size();
-
- // initialize V_{nxn} with identity matrix,
- // initialize V^{-1} x as x
- for (size_t i = 0; i < num_vars; ++i) {
- V.emplace_back(num_vars);
- V.back()[i] = 1;
- V_inv_x.push_back(system_to_solve->variables[i].as_or_throw<PrimExpr>());
- }
-
- // Transform formulas into rows of the matrix
- // S_{mxn} V^{-1}_{nxn} x_{nx1} = U y, in which n is # of variables
- // here we initialize S_{mxn} to be A, U to be identity matrix.
- for (const PrimExpr& equation : system_to_solve->relations) {
- if (const prim::EQNode* eq = equation.as<prim::EQNode>()) {
- // a-b = sum_{i=0}^{n-1} variables[i] * coeff[i] + coeff[n]
- ffi::Array<PrimExpr> coeffs = arith::DetectLinearEquation(
- analyzer_problem->Simplify(eq->a - eq->b),
system_to_solve->variables);
- if (!coeffs.empty()) {
- std::vector<int64_t> row;
- for (size_t j = 0; j < coeffs.size() - 1; ++j) {
- PrimExpr c = coeffs[j];
- if (const IntImmNode* ic = c.as<IntImmNode>()) {
- row.push_back(ic->value);
- } else {
- // elements in matrix S V must be integers
- // ignore equations that we cannot deal with.
- LOG(WARNING) << "Cannot deal with non-integer coefficients, ignore
equation "
- << equation;
- row.clear();
- break;
- }
- }
-
- if (!row.empty()) {
- // S V^{-1} (a-b) = Uy
- // V is identity for now
- S.push_back(row);
- Uy.push_back(-coeffs[coeffs.size() - 1]);
- continue;
- }
- }
- }
-
- // otherwise
- rest.push_back(equation);
- }
-
- // After diagonalizing, we have
- // S_{mxn} is the Smith normal form (diagonal matrix)
- // V_{nxn} is invertible
- // V_inv_x is V^{-1} \times x
- // Uy is U \times y
- SmithNormalFormDiag(&S, &V, &V_inv_x, &Uy);
-
- ffi::Array<PrimVar> new_vars;
- ffi::Array<PrimExpr> new_relations;
- ffi::Map<Var, PrimExpr> new_to_old_map;
- ffi::Map<Var, PrimExpr> old_to_new_map;
-
- // Simplify right hand sides
- for (PrimExpr r : Uy) {
- r = analyzer_problem->Simplify(r);
- }
-
- // Create the relations of the existence of a solution
- for (size_t j = 0; j < S.size(); ++j) {
- PrimExpr new_relation;
- if (j >= num_vars || S[j][j] == 0) {
- // The row of matrix is zero. A solution exists only if the Ub[j] is
also zero
- new_relation = (Uy[j] == 0);
- } else {
- // The diagonal element is non-zero. A solution exists only if the
diagonal element
- // is a divisor of the Ub[j]
- new_relation = (floormod(Uy[j], std::abs(S[j][j])) == 0);
- }
- new_relation = analyzer_problem->Simplify(new_relation);
- if (tirx::is_const_int(new_relation, 0)) {
- // unable to solve the system.
- return IntConstraintsTransform(system_to_solve,
- IntConstraints(
- /*variables=*/{},
- /*ranges=*/{},
- /*relations=*/{IntImm::Bool(false)}),
- {}, {});
- } else if (!tirx::is_const_int(new_relation, 1)) {
- new_relations.push_back(new_relation);
- }
- }
-
- ffi::Array<PrimExpr> solution_for_V_inv_x;
- // Now create new variables or directly solve the equations
- // suppose the rank of A is r, aka r = # of non-zeros in S
- // the solution of S_{mxn} V^{-1}_{nxn} x_{nx1} = U b
- // is
- // x = (pseudo-inverse of A) b + K_{(n)x(n-r)} z_{n-r}
- // = V_{nxn} S^{-1}_{nxm} (Ub)_{mxn} + K_{(n)x(n-r)} z_{n-r}
- // in which K is the right n-r columns of V, z is variable vector
- // thus,
- // V^{-1} x = S^{-1}_{nxm} (Ub)_{mxn} +
- // [[0, ... 0]_{n-r}, ... [0, ..., 0], diag(1, ...,
1)_{(n-r)x(n-r)}] z_{n-r}
- for (size_t j = 0; j < num_vars; ++j) {
- if (j >= S.size() || S[j][j] == 0) {
- // The j-th variable can take any integer value, create a tvm variable
for it
- PrimExpr to_old = analyzer_problem->Simplify(V_inv_x[j]);
- std::string name_hint = "n" + std::to_string(new_vars.size());
- if (auto old_var = to_old.as<PrimVar>()) {
- name_hint += "_" + (*old_var)->name;
- }
- PrimVar v(name_hint, V_inv_x[j].ty());
- solution_for_V_inv_x.push_back(v);
- new_vars.push_back(v);
- new_to_old_map.Set(v, to_old);
- } else {
- // The j-th variable is just a single value, don't create a tvm variable
- // S^{-1}_{nxm} Uy_{mxn}
- if (S[j][j] >= 0) {
- PrimExpr a = IntImm(Uy[j].ty(), S[j][j]);
-
solution_for_V_inv_x.push_back(analyzer_problem->Simplify(floordiv(Uy[j], a)));
- } else {
- // This is required because some simplifiers
- // have problems with dividing by negative numbers
- PrimExpr a = IntImm(Uy[j].ty(), -S[j][j]);
-
solution_for_V_inv_x.push_back(analyzer_problem->Simplify(floordiv(-Uy[j], a)));
- }
- }
- }
-
- // V V^{-1} x = x
- for (size_t i = 0; i < num_vars; ++i) {
- PrimExpr e =
IntImm(system_to_solve->variables[i]->ty.as_or_throw<PrimType>(), 0);
- for (size_t j = 0; j < num_vars; ++j) {
- e = e + tirx::MakeConst(e.ty(), V[i][j]) * solution_for_V_inv_x[j];
- }
- e = analyzer_problem->Simplify(e);
- old_to_new_map.Set(system_to_solve->variables[i], e);
- }
-
- // The resulting ranges
- ffi::Map<Var, Range> new_ranges =
- InferRange(new_to_old_map, system_to_solve->variables,
system_to_solve->ranges);
- Analyzer analyzer_solution;
- analyzer_solution->Bind(new_ranges);
-
- // We have to transform ranges of the old variables into relations over new
variables because
- // new ranges are not enough usually.
- for (const auto& old_var : system_to_solve->variables) {
- if (system_to_solve->ranges.find(old_var) !=
system_to_solve->ranges.end()) {
- const Range& old_range = system_to_solve->ranges.at(old_var);
- PrimExpr express_by_new_vars = old_to_new_map.at(old_var);
- PrimExpr lower_cond = analyzer_solution->Simplify(old_range->min <=
express_by_new_vars);
- PrimExpr upper_cond =
- analyzer_solution->Simplify(express_by_new_vars < old_range->min +
old_range->extent);
- if (!tirx::is_const_int(lower_cond, 1)) {
- new_relations.push_back(lower_cond);
- }
- if (!tirx::is_const_int(upper_cond, 1)) {
- new_relations.push_back(upper_cond);
- }
- }
- }
-
- // Add the rest conditions
- auto f_subst = [&old_to_new_map](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto repl = old_to_new_map.Get(var)) return *std::move(repl);
- return ffi::Unchanged();
- };
- for (const PrimExpr& cond : rest) {
- new_relations.push_back(
- ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(cond,
f_subst).as_or_throw<PrimExpr>());
- }
-
- IntConstraints solution(new_vars, new_ranges, new_relations);
- IntConstraintsTransform transform(system_to_solve, solution, old_to_new_map,
new_to_old_map);
-
- return transform;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_packed(
- "arith.SolveLinearEquations", [](ffi::PackedArgs args, ffi::Any* ret) {
- if (args.size() == 1) {
- *ret = SolveLinearEquations(args[0].cast<IntConstraints>());
- } else if (args.size() == 3) {
- auto opt_vars = args[0].cast<ffi::Optional<ffi::Array<PrimVar>>>();
- auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var, Range>>>();
- auto opt_relations =
args[2].cast<ffi::Optional<ffi::Array<PrimExpr>>>();
- IntConstraints problem(opt_vars.value_or({}), opt_map.value_or({}),
- opt_relations.value_or({}));
- *ret = SolveLinearEquations(problem);
- } else {
- TVM_FFI_THROW(InternalError)
- << "arith.SolveLinearEquations expects 1 or 3 arguments, gets "
<< args.size();
- }
- });
-}
-
-} // namespace arith
-} // namespace tvm
diff --git a/src/arith/solve_linear_inequality.cc
b/src/s_tir/analysis/conditional_bounds.cc
similarity index 53%
rename from src/arith/solve_linear_inequality.cc
rename to src/s_tir/analysis/conditional_bounds.cc
index 4bc103708a..2576fb62db 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/s_tir/analysis/conditional_bounds.cc
@@ -18,66 +18,219 @@
*/
/*!
- * \file tvm/arith/solve_linear_inequality.cc
- * \brief Solve linear inequalities.
+ * \file s_tir/analysis/conditional_bounds.cc
+ * \brief Scoped conditional bounds and private inequality support for S-TIR.
*/
+#include "conditional_bounds.h"
+
#include <tvm/arith/analyzer.h>
-#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
-#include <tvm/ffi/dtype.h>
#include <tvm/ffi/extra/structural_mutate.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr_functor.h>
-#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/builtin.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
+#include <algorithm>
+#include <functional>
#include <utility>
-#include "int_operator.h"
+#include "../../arith/int_operator.h"
namespace tvm {
-namespace arith {
+namespace s_tir {
-using namespace tvm::runtime;
using namespace tvm::tirx;
-
-struct ExprLess {
- bool operator()(const PrimExpr& l, const PrimExpr& r) const {
- return CalculateExprComplexity(l) < CalculateExprComplexity(r);
+using arith::Analyzer;
+using arith::AnalyzerObj;
+using arith::EvalSet;
+using arith::IntSet;
+
+namespace {
+using arith::ExtendedEuclidean;
+using arith::LeastCommonMultiple;
+
+// The solver's intermediate representations remain local to this analysis.
+struct IntGroupBounds {
+ PrimExpr coef;
+ ffi::Array<PrimExpr> lower;
+ ffi::Array<PrimExpr> equal;
+ ffi::Array<PrimExpr> upper;
+
+ IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower,
ffi::Array<PrimExpr> equal,
+ ffi::Array<PrimExpr> upper)
+ : coef(std::move(coef)),
+ lower(std::move(lower)),
+ equal(std::move(equal)),
+ upper(std::move(upper)) {
+ TVM_FFI_ICHECK(this->coef.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
+ << "Coefficient in IntGroupBounds must be integers";
}
+
+ Range FindBestRange(const ffi::Map<Var, Range>& vranges_addl) const;
+ IntGroupBounds operator+(const Range& range);
};
-void DebugPrint(const std::vector<PrimExpr>& current_ineq_set,
- const std::vector<PrimExpr>& next_ineq_set, const
std::vector<PrimExpr>& rest,
- const std::vector<std::pair<int64_t, PrimExpr>>& coef_pos,
- const std::vector<std::pair<int64_t, PrimExpr>>& coef_neg) {
- std::cout << "Current ineq set:\n[";
- for (auto& ineq : current_ineq_set) {
- std::cout << ineq << ", ";
+struct IntConstraints {
+ ffi::Array<PrimVar> variables;
+ ffi::Map<Var, Range> ranges;
+ ffi::Array<PrimExpr> relations;
+
+ IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
+ ffi::Array<PrimExpr> relations);
+};
+
+using GroupedBounds =
+ std::unordered_map<Var, IntGroupBounds, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>;
+using PartialSolvedInequalities = std::pair<GroupedBounds,
ffi::Array<PrimExpr>>;
+
+// Rewrite, canonicalize, then rewrite to factor multipliers out.
+constexpr int kSimplifyRewriteCanonicalRewrite = 3;
+
+IntConstraints::IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var,
Range> ranges,
+ ffi::Array<PrimExpr> relations) {
+ if (!variables.defined()) {
+ variables = ffi::Array<PrimVar>();
+ }
+ if (!ranges.defined()) {
+ ranges = ffi::Map<Var, Range>();
+ }
+ TVM_FFI_ICHECK(relations.defined());
+ for (const PrimVar& var : variables) {
+ TVM_FFI_CHECK(var.ty().MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt), TypeError)
+ << "Variables in IntConstraints must be integers";
+ }
+ this->variables = std::move(variables);
+ this->ranges = std::move(ranges);
+ this->relations = std::move(relations);
+}
+
+ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables, const
GroupedBounds& bounds,
+ const ffi::Array<PrimExpr>& relations) {
+ ffi::Array<PrimExpr> res;
+ // use variables to keep the order of iteration
+ // so as to get rid of any non-determinism.
+ TVM_FFI_ICHECK_EQ(variables.size(), bounds.size());
+ for (const auto v : variables) {
+ TVM_FFI_ICHECK(bounds.count(v));
+ const auto& bnds = bounds.at(v);
+ PrimExpr lhs = bnds.coef * v.as_or_throw<PrimExpr>();
+ for (const PrimExpr& rhs : bnds.equal) {
+ res.push_back(lhs == rhs);
+ }
+ for (const PrimExpr& rhs : bnds.lower) {
+ res.push_back(lhs >= rhs);
+ }
+ for (const PrimExpr& rhs : bnds.upper) {
+ res.push_back(lhs <= rhs);
+ }
+ }
+ for (const PrimExpr& e : relations) {
+ res.push_back(e);
+ }
+ return res;
+}
+
+IntGroupBounds IntGroupBounds::operator+(const Range& r) {
+ Analyzer analyzer;
+ ffi::Array<PrimExpr> equal;
+ ffi::Array<PrimExpr> lower;
+ ffi::Array<PrimExpr> upper;
+ const PrimExpr& coef = this->coef;
+ if (tirx::is_one(r->extent)) {
+ equal.push_back(analyzer->Simplify(r->min * coef));
+ } else {
+ lower.push_back(analyzer->Simplify(r->min * coef));
+ upper.push_back(analyzer->Simplify((r->min + r->extent - 1) * coef));
+ }
+ for (const auto& eq : this->equal) equal.push_back(eq);
+ for (const auto& lb : this->lower) lower.push_back(lb);
+ for (const auto& ub : this->upper) upper.push_back(ub);
+ return IntGroupBounds(coef, lower, equal, upper);
+}
+
+Range IntGroupBounds::FindBestRange(const ffi::Map<Var, Range>& vranges_addl)
const {
+ Analyzer analyzer;
+ analyzer->Bind(vranges_addl);
+
+ std::unordered_map<const VarNode*, IntSet> var_intsets;
+ for (auto kv : vranges_addl) {
+ var_intsets[kv.first.get()] = IntSet::FromRange(kv.second);
+ }
+
+ const ffi::Array<PrimExpr>& equal = this->equal;
+ const PrimExpr& coef = this->coef;
+
+ std::vector<PrimExpr> lowers(equal.begin(), equal.end());
+ std::vector<PrimExpr> uppers(equal.begin(), equal.end());
+ for (const auto& expr : this->lower) {
+ lowers.push_back(expr);
+ }
+ for (const auto& expr : this->upper) {
+ uppers.push_back(expr);
}
- std::cout << "]\n";
- std::cout << "Next ineq set:\n[";
- for (auto& ineq : next_ineq_set) {
- std::cout << ineq << ", ";
+ if (lowers.size() == 1 && uppers.size() == 1 && tirx::is_one(coef)) {
+ return Range(analyzer->Simplify(lowers[0]), analyzer->Simplify(uppers[0] +
1));
}
- std::cout << "]\n";
- std::cout << "coef_pos:\n[";
- for (auto& coef : coef_pos) {
- std::cout << "(" << coef.first << ", " << coef.second << "), ";
+ // Here we will try all pairs of lower and upper bounds and find the best
pair, that is, the
+ // pair with the minimal difference between the upper and the lower.
+ // Note that the bounds are for v, not for v*coef
+
+ // The lower bound of the best pair so far
+ PrimExpr best_lower;
+ // The difference between the upper and the lower of the best pair, maybe
overapproximation
+ PrimExpr best_diff_over;
+
+ for (const PrimExpr& low : lowers) {
+ for (const PrimExpr& upp : uppers) {
+ // Since diff may depend on some other variables, we compute its
overapproximation
+ ffi::Optional<PrimExpr> diff_over;
+ PrimExpr diff_1 = analyzer->Simplify(floordiv(upp - low, coef), 3);
+ IntSet diff_set1 = EvalSet(diff_1, var_intsets);
+ if (diff_set1.HasUpperBound()) {
+ diff_over = analyzer->Simplify(diff_set1.max(), 3);
+ }
+
+ // low is the lower bound for v*coef, but we need the lower bound for v.
+ // We use rounding-up division to compute it. Since we want to use a
single formula
+ PrimExpr low_divided = analyzer->Simplify(floordiv(low + coef - 1,
coef), 3);
+
+ // Compute another difference which may be more precise (or not).
+ PrimExpr diff_2 = analyzer->Simplify(floordiv(upp, coef) - low_divided,
3);
+ IntSet diff_set2 = EvalSet(diff_2, var_intsets);
+ if (diff_set2.HasUpperBound()) {
+ PrimExpr diff_over_2 = analyzer->Simplify(diff_set2.max(), 3);
+ diff_over = diff_over.has_value() ? (analyzer->CanProve(diff_over_2 -
diff_over.value() < 0)
+ ? diff_over_2
+ : diff_over.value())
+ : diff_over_2;
+ }
+
+ // If it is provable that the new one is strictly better than the
current best one,
+ // then replace it. Note that we are biased towards earlier pairs which
should be simpler.
+ if (diff_over.has_value() && (!best_diff_over.defined() ||
+ analyzer->CanProve(diff_over.value() -
best_diff_over < 0))) {
+ best_lower = low_divided;
+ best_diff_over = diff_over.value();
+ }
+ }
}
- std::cout << "]\n";
- std::cout << "coef_neg:\n[";
- for (auto& coef : coef_neg) {
- std::cout << "(" << coef.first << ", " << coef.second << "), ";
+ if (!best_lower.defined()) {
+ TVM_FFI_ICHECK(!best_diff_over.defined());
+ return Range();
}
- std::cout << "]\n";
+ return Range::FromMinExtent(best_lower, analyzer->Simplify(best_diff_over +
1));
}
+struct ExprLess {
+ bool operator()(const PrimExpr& l, const PrimExpr& r) const {
+ return CalculateExprComplexity(l) < CalculateExprComplexity(r);
+ }
+};
+
/*!
* \brief normalize to the form `expr <= 0`
*/
@@ -206,7 +359,7 @@ void MoveEquality(std::vector<PrimExpr>* upper_bounds,
std::vector<PrimExpr>* lo
PartialSolvedInequalities SolveLinearInequalities(const IntConstraints&
system_to_solve) {
arith::Analyzer analyzer;
- analyzer->Bind(system_to_solve->ranges);
+ analyzer->Bind(system_to_solve.ranges);
// The algorithm consists in doing the following things for each variable v
// - Take formulas from `current_ineq_set_to_solve` and
@@ -231,14 +384,14 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
// Simplify each inequality into the form `expr <= 0` and add to current
formulas
auto normalizer = ffi::make_object<NormalizeComparisons>();
- for (const PrimExpr& ineq : system_to_solve->relations) {
+ for (const PrimExpr& ineq : system_to_solve.relations) {
PrimExpr simplified = analyzer->Simplify(ineq,
kSimplifyRewriteCanonicalRewrite);
PrimExpr normalized =
normalizer->Mutate(simplified).ValueOrUnchanged(simplified);
AddInequality(¤t_ineq_set_to_solve, normalized, analyzer.get());
}
- ffi::Map<Var, IntGroupBounds> res_bounds;
- for (const PrimVar& v : system_to_solve->variables) {
+ GroupedBounds res_bounds;
+ for (const PrimVar& v : system_to_solve.variables) {
TVM_FFI_ICHECK(!res_bounds.count(v))
<< "Variable " << v
<< " appears more than one time in the `variables` which might be a
bug";
@@ -248,8 +401,8 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
coef_neg.clear();
// Add bounds from vranges
- if (system_to_solve->ranges.count(v)) {
- const Range& range = system_to_solve->ranges[v];
+ if (system_to_solve.ranges.count(v)) {
+ const Range& range = system_to_solve.ranges[v];
PrimExpr range_lbound = analyzer->Simplify(range->min,
kSimplifyRewriteCanonicalRewrite);
PrimExpr range_ubound =
analyzer->Simplify(range->min + range->extent - 1,
kSimplifyRewriteCanonicalRewrite);
@@ -352,7 +505,7 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
ffi::Array<PrimExpr>(lower_bounds.begin(),
lower_bounds.end()),
ffi::Array<PrimExpr>(equal_list.begin(),
equal_list.end()),
ffi::Array<PrimExpr>(upper_bounds.begin(),
upper_bounds.end()));
- res_bounds.Set(v, bnds);
+ res_bounds.emplace(v, bnds);
std::swap(current_ineq_set_to_solve, next_ineq_set_to_solve);
}
@@ -384,40 +537,39 @@ PartialSolvedInequalities SolveLinearInequalities(const
IntConstraints& system_t
#endif
IntConstraints SolveInequalitiesToRange(const IntConstraints& inequalities) {
// Resulting ranges will contain ranges for the new variables and for the
variables that are
- // not in the inequalities->variables but are in inequalities->ranges
- // It will be useful when solving Jacobian axes jac_xxx)
+ // not in the inequalities.variables but are in inequalities.ranges
ffi::Map<Var, Range> res_ranges;
// we get a set of equality, lower, upper bound of each variable.
auto solved_system = SolveLinearInequalities(inequalities);
- ffi::Map<Var, IntGroupBounds> solved_bounds = solved_system.first;
+ GroupedBounds solved_bounds = solved_system.first;
ffi::Array<PrimExpr> solved_other_relations = solved_system.second;
ffi::Array<PrimExpr> res_relations;
// this keeps being updated during determining the range of each variable.
ffi::Map<Var, Range> vranges;
- for (std::pair<Var, Range> vr : inequalities->ranges) {
+ for (std::pair<Var, Range> vr : inequalities.ranges) {
vranges.Set(vr.first, vr.second);
}
// We process variables in the reverse direction to start with the most
independent one.
// This order is needed to compute new ranges.
- for (auto it = inequalities->variables.rbegin(); it !=
inequalities->variables.rend(); ++it) {
+ for (auto it = inequalities.variables.rbegin(); it !=
inequalities.variables.rend(); ++it) {
arith::Analyzer analyzer;
analyzer->Bind(vranges);
const PrimVar& var = *it;
TVM_FFI_ICHECK(solved_bounds.count(var));
- auto bnd = solved_bounds[var];
- if (is_one(bnd->coef) && !bnd->equal.empty()) {
+ auto bnd = solved_bounds.at(var);
+ if (is_one(bnd.coef) && !bnd.equal.empty()) {
// There is an equation of the form `v == expr`, so this variable can be
completely removed.
// Note that we use the 0-th expression because they are ordered by
complexity,
// so it must be the simplest one.
- // The MSVC compiler optimization must be disabled for the expression
`bnd->equal[0]` which
+ // The MSVC compiler optimization must be disabled for the expression
`bnd.equal[0]` which
// triggers an internal compiler error.
- Range best_range(bnd->equal[0],
- analyzer->Simplify(bnd->equal[0] + 1,
kSimplifyRewriteCanonicalRewrite));
+ Range best_range(bnd.equal[0],
+ analyzer->Simplify(bnd.equal[0] + 1,
kSimplifyRewriteCanonicalRewrite));
res_ranges.Set(var, best_range);
vranges.Set(var, best_range);
} else {
@@ -443,14 +595,14 @@ IntConstraints SolveInequalitiesToRange(const
IntConstraints& inequalities) {
arith::Analyzer analyzer;
analyzer->Bind(vranges);
for (const PrimExpr& old_cond :
- AsConditions(inequalities->variables, solved_bounds,
solved_other_relations)) {
+ AsConditions(inequalities.variables, solved_bounds,
solved_other_relations)) {
if (!analyzer->CanProve(old_cond)) {
// those not represented in vranges (res_ranges)
res_relations.push_back(old_cond);
}
}
- IntConstraints system(inequalities->variables, res_ranges, res_relations);
+ IntConstraints system(inequalities.variables, res_ranges, res_relations);
return system;
}
@@ -458,171 +610,151 @@ IntConstraints SolveInequalitiesToRange(const
IntConstraints& inequalities) {
#pragma optimize("g", on)
#endif
-IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints&
inequalities) {
- // Resulting ranges will contain ranges for the new variables and for the
variables that are
- // not in the inequalities->variables but are in inequalities->ranges
(jac_xxx)
- ffi::Map<Var, Range> res_ranges;
- // we get a set of equality, lower, upper bound of each variable.
- auto solved_system = SolveLinearInequalities(inequalities);
- ffi::Map<Var, IntGroupBounds> solved_bounds = solved_system.first;
- ffi::Array<PrimExpr> solved_other_relations = solved_system.second;
+} // namespace
+ffi::Optional<ffi::Map<Var, Range>>
ConditionalBoundsContext::TrySolveCondition() {
+ // extract equations and related vars from condition expression.
+ // currently only extract simple integral equations which could be solvable.
arith::Analyzer analyzer;
-
- ffi::Map<Var, PrimExpr> res_src_to_dst;
- ffi::Map<Var, PrimExpr> res_dst_to_src;
- ffi::Array<PrimVar> res_variables;
- ffi::Array<PrimExpr> res_relations;
-
- // this keeps being updated during determining the range of each variable.
- ffi::Map<Var, Range> vranges;
- for (std::pair<Var, Range> vr : inequalities->ranges) {
- vranges.Set(vr.first, vr.second);
+ PrimExpr condition = analyzer->Simplify(condition_);
+ if (is_const_int(condition)) {
+ return std::nullopt;
}
- analyzer->Bind(vranges);
-
- auto subst = [&res_src_to_dst](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto repl = res_src_to_dst.Get(var)) return *std::move(repl);
- return ffi::Unchanged();
- };
- auto f_dst_to_src =
- [&res_dst_to_src](const Var& var) ->
ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
- if (auto repl = res_dst_to_src.Get(var)) return *std::move(repl);
- return ffi::Unchanged();
+ ffi::Array<PrimExpr> equations;
+ ffi::Array<PrimVar> vars;
+ std::function<void(const PrimExpr&)> fvisit = [&equations, &vars,
&fvisit](const PrimExpr& e) {
+ if (e->IsInstance<prim::GENode>() || e->IsInstance<prim::GTNode>() ||
+ e->IsInstance<prim::LENode>() || e->IsInstance<prim::LTNode>() ||
+ e->IsInstance<prim::EQNode>() || e->IsInstance<prim::NENode>()) {
+ bool is_simple = true;
+ std::vector<PrimVar> cand_vars;
+ auto walk_fn = [&cand_vars, &is_simple,
+ &e](const PrimExpr& obj) ->
ffi::Expected<ffi::WalkResult> {
+ if (obj.same_as(e)) {
+ return ffi::WalkResult::Advance();
+ } else if (const VarNode* var = obj.as<VarNode>()) {
+ PrimType var_ty = var->ty.as_or_throw<PrimType>();
+ if (var_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
+ cand_vars.push_back(ffi::GetRef<Var>(var).as_or_throw<PrimVar>());
+ }
+ } else {
+ is_simple &= obj->IsInstance<prim::AddNode>() ||
obj->IsInstance<prim::SubNode>() ||
+ obj->IsInstance<prim::MulNode>() ||
obj->IsInstance<prim::FloorDivNode>() ||
+ obj->IsInstance<prim::FloorModNode>() ||
obj->IsInstance<IntImmNode>();
+ }
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(e, walk_fn);
+ if (is_simple && !cand_vars.empty()) {
+ for (const PrimVar& new_var : cand_vars) {
+ if (!std::any_of(vars.begin(), vars.end(),
+ [&new_var](const PrimVar& v) { return
v.same_as(new_var); })) {
+ vars.push_back(new_var);
+ }
+ }
+ equations.push_back(e.as_or_throw<PrimExpr>());
+ }
+ } else if (e->IsInstance<prim::AndNode>()) {
+ prim::And op = e.as_or_throw<prim::And>();
+ fvisit(op->a);
+ fvisit(op->b);
+ } else if (e->IsInstance<CallNode>()) {
+ Call op = e.as_or_throw<Call>();
+ if (op->op.same_as(prim::builtin::likely())) {
+ fvisit(op->args[0].as_or_throw<PrimExpr>());
+ }
+ }
};
-
- // We process variables in the reverse direction to start with the most
independent one.
- // This order is needed to compute new ranges.
- for (auto it = inequalities->variables.rbegin(); it !=
inequalities->variables.rend(); ++it) {
- const PrimVar& var = *it;
- auto bnd = solved_bounds[var];
- // Note that we replace old vars with new ones
- bnd = ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(bnd,
subst).as_or_throw<IntGroupBounds>();
-
- if (is_one(bnd->coef) && !bnd->equal.empty()) {
- // There is an equation of the form `v == expr`,
- // so this variable can be completely removed.
- // Note that we use the 0-th expression because they are ordered by
complexity,
- // so it must be the simplest one.
- res_src_to_dst.Set(var, bnd->equal[0]);
+ fvisit(condition);
+ if (equations.empty() || vars.empty()) {
+ return std::nullopt;
+ }
+ // build dom ranges for related vars
+ ffi::Map<Var, Range> ranges;
+ for (const Var& v : vars) {
+ arith::IntSet dom;
+ auto relax_it = relax_map_->find(v.get());
+ if (relax_it != relax_map_->end()) {
+ dom = relax_it->second;
} else {
- if (vranges.count(var) > 0) {
- bnd = bnd + vranges[var];
+ auto hint_it = hint_map_->find(v.get());
+ if (hint_it != hint_map_->end()) {
+ dom = hint_it->second;
}
+ }
+ if (dom.defined()) {
+ ranges.Set(v, Range::FromMinExtent(dom.min(),
analyzer->Simplify(dom.max() - dom.min() + 1)));
+ }
+ }
+ // solve constraints
+ IntConstraints constraint(vars, ranges, equations);
+ IntConstraints result = SolveInequalitiesToRange(constraint);
+ if (!result.relations.empty()) {
+ return std::nullopt;
+ }
+ return result.ranges;
+}
- auto best_range = bnd.FindBestRange(vranges);
-
- PrimVar new_var = var.CopyWithSuffix(".shifted");
- if (!best_range.defined()) {
- res_src_to_dst.Set(var, var.as_or_throw<PrimExpr>());
- res_dst_to_src.Set(var, var.as_or_throw<PrimExpr>());
- res_variables.push_back(var);
- } else if (is_const_int(best_range->extent, 1)) {
- // Don't create an itervar, just replace it everywhere with its min
- res_src_to_dst.Set(var, best_range->min);
- } else if (analyzer->CanProveGreaterEqual(-best_range->extent, 0)) {
- // range.extent <= 0 implies the input inequality system is unsolvable
- return IntConstraintsTransform(inequalities,
- IntConstraints(
- /*variables=*/{},
- /*ranges=*/{},
-
/*relations=*/{IntImm::Bool(false)}),
- {}, {});
+ConditionalBoundsContext::ConditionalBoundsContext(
+ const PrimExpr& condition, std::unordered_map<const VarNode*,
arith::IntSet>* relax_map,
+ std::unordered_map<const VarNode*, arith::IntSet>* hint_map,
+ std::vector<PrimExpr>* pending_conditions)
+ : condition_(condition),
+ relax_map_(relax_map),
+ hint_map_(hint_map),
+ pending_conditions_(pending_conditions),
+ origin_pending_conditions_num_(pending_conditions->size()) {}
+
+void ConditionalBoundsContext::EnterWithScope() {
+ ffi::Optional<ffi::Map<Var, Range>> constraints = TrySolveCondition();
+ if (!constraints.has_value()) {
+ // fail to process the condition, add to unresolved
+ pending_conditions_->push_back(condition_);
+ return;
+ }
+ // update solved var ranges
+ for (const auto& kv : constraints.value()) {
+ const VarNode* var = kv.first.get();
+ arith::IntSet new_dom = arith::IntSet::FromRange(kv.second);
+ auto relax_it = relax_map_->find(var);
+ if (relax_it != relax_map_->end()) {
+ // this is a bound for relaxed var
+ origin_map_.emplace(var, relax_it->second);
+ relax_it->second = arith::Intersect({relax_it->second, new_dom});
+ } else {
+ // this is a bound for free var
+ auto hint_it = hint_map_->find(var);
+ if (hint_it != hint_map_->end()) {
+ origin_map_.emplace(var, hint_it->second);
+ hint_it->second = arith::Intersect({hint_it->second, new_dom});
} else {
- // created new_var starts from 0
- res_src_to_dst.Set(var, new_var.as_or_throw<PrimExpr>() +
best_range->min);
- // Note that we are substituting old with new, so best_range contains
new var,
- // that is we have to substitute new with old in best_range here
- res_dst_to_src.Set(new_var,
- analyzer->Simplify(var.as_or_throw<PrimExpr>() -
-
ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(
- best_range->min,
f_dst_to_src)
- .as_or_throw<PrimExpr>()));
-
- // Add the new var to the resulting axis
- auto range = Range(IntImm(new_var->ty.as_or_throw<PrimType>(), 0),
best_range->extent);
- res_variables.push_back(new_var);
- res_ranges.Set(new_var, range);
-
- vranges.Set(new_var, range);
- analyzer->Bind(new_var, range);
+ origin_map_.emplace(var, arith::IntSet::Nothing());
+ hint_map_->insert(hint_it, {var, new_dom});
}
}
}
+}
- // Add the original conditions (with variables substituted) to the resulting
conditions
- for (const PrimExpr& old_cond :
- AsConditions(inequalities->variables, solved_bounds,
solved_other_relations)) {
- PrimExpr new_cond = analyzer->Simplify(
- ffi::StructuralMap<ffi::WalkOrder::kPreOrder>(old_cond,
subst).as_or_throw<PrimExpr>());
- if (!is_const_int(new_cond, 1)) {
- // those not represented in vranges (res_ranges)
- res_relations.push_back(new_cond);
+void ConditionalBoundsContext::ExitWithScope() {
+ pending_conditions_->resize(origin_pending_conditions_num_);
+ for (const auto& p : origin_map_) {
+ const auto* var = p.first;
+ auto relax_it = relax_map_->find(var);
+ if (relax_it != relax_map_->end()) {
+ // recover bound for relaxed var
+ relax_it->second = p.second;
+ } else {
+ // recover bound for free var
+ auto hint_it = hint_map_->find(var);
+ TVM_FFI_ICHECK(hint_it != hint_map_->end());
+ if (p.second.IsNothing()) {
+ hint_map_->erase(hint_it);
+ } else {
+ hint_it->second = p.second;
+ }
}
}
-
- // Reverse the axis so that it matches the order of the original variables
- res_variables = ffi::Array<PrimVar>(res_variables.rbegin(),
res_variables.rend());
-
- IntConstraints new_inequalities(res_variables, res_ranges, res_relations);
- IntConstraintsTransform transform(inequalities, new_inequalities,
res_src_to_dst, res_dst_to_src);
-
- return transform;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def_packed("arith.SolveInequalitiesAsCondition",
- [](ffi::PackedArgs args, ffi::Any* ret) {
- IntConstraints problem;
- PartialSolvedInequalities ret_ineq;
- if (args.size() == 1) {
- problem = args[0].cast<IntConstraints>();
- ret_ineq = SolveLinearInequalities(problem);
- } else if (args.size() == 3) {
- problem =
IntConstraints(args[0].cast<ffi::Array<PrimVar>>(),
- args[1].cast<ffi::Map<Var,
Range>>(),
-
args[2].cast<ffi::Array<PrimExpr>>());
- ret_ineq = SolveLinearInequalities(problem);
- } else {
- TVM_FFI_THROW(InternalError)
- << "arith.SolveInequalitiesAsCondition expects 1 or
3 arguments, gets "
- << args.size();
- }
- *ret = AsConditions(problem->variables, ret_ineq.first,
ret_ineq.second);
- })
- .def_packed("arith.SolveInequalitiesToRange",
- [](ffi::PackedArgs args, ffi::Any* ret) {
- if (args.size() == 1) {
- *ret =
SolveInequalitiesToRange(args[0].cast<IntConstraints>());
- } else if (args.size() == 3) {
- auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var,
Range>>>();
- IntConstraints
problem(args[0].cast<ffi::Array<PrimVar>>(),
- opt_map.value_or({}),
-
args[2].cast<ffi::Array<PrimExpr>>());
- *ret = SolveInequalitiesToRange(problem);
- } else {
- TVM_FFI_THROW(InternalError)
- << "arith.SolveInequalitiesToRange expects 1 or 3
arguments, gets "
- << args.size();
- }
- })
- .def_packed("arith.SolveInequalitiesDeskewRange", [](ffi::PackedArgs
args, ffi::Any* ret) {
- if (args.size() == 1) {
- *ret = SolveInequalitiesDeskewRange(args[0].cast<IntConstraints>());
- } else if (args.size() == 3) {
- auto opt_map = args[1].cast<ffi::Optional<ffi::Map<Var, Range>>>();
- IntConstraints problem(args[0].cast<ffi::Array<PrimVar>>(),
opt_map.value_or({}),
- args[2].cast<ffi::Array<PrimExpr>>());
- *ret = SolveInequalitiesDeskewRange(problem);
- } else {
- TVM_FFI_THROW(InternalError)
- << "arith.SolveInequalitiesDeskewRange expects 1 or 3 arguments,
gets "
- << args.size();
- }
- });
}
-} // namespace arith
+} // namespace s_tir
} // namespace tvm
diff --git a/src/s_tir/analysis/conditional_bounds.h
b/src/s_tir/analysis/conditional_bounds.h
new file mode 100644
index 0000000000..779e89bb5a
--- /dev/null
+++ b/src/s_tir/analysis/conditional_bounds.h
@@ -0,0 +1,82 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file s_tir/analysis/conditional_bounds.h
+ * \brief Scoped conditional bounds for S-TIR buffer analysis.
+ */
+#ifndef TVM_S_TIR_ANALYSIS_CONDITIONAL_BOUNDS_H_
+#define TVM_S_TIR_ANALYSIS_CONDITIONAL_BOUNDS_H_
+
+#include <tvm/arith/int_set.h>
+#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/with_context.h>
+
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace s_tir {
+
+/*!
+ * \brief Context helper to update domain map within conditional scope.
+ * Assume the condition is `0 <= i && i < 9` and domain of i is [0, 20], Then
+ * `With<ConditionalBoundsContext> ctx(condition, &relax_map, &hint_map,
&constraints)`
+ * step into scope where dom_map[i] is [0, 8]; and
+ * `With<ConditionalBoundsContext> ctx(!condition, &relax_map, &hint_map,
&constraints)`
+ * step into scope where dom_map[i] is [9, 20]
+ */
+class ConditionalBoundsContext {
+ private:
+ friend class With<ConditionalBoundsContext>;
+ /*!
+ * \brief Construct a condition bounds context.
+ * \param condition The condition holds on true branch.
+ * \param relax_map The domain map for relaxed vars to update.
+ * \param hint_map The domain map for free vars to update.
+ * \param pending_conditions The stack of unresolved constraints.
+ */
+ ConditionalBoundsContext(const PrimExpr& condition,
+ std::unordered_map<const VarNode*, arith::IntSet>*
relax_map,
+ std::unordered_map<const VarNode*, arith::IntSet>*
hint_map,
+ std::vector<PrimExpr>* pending_constraints);
+ void EnterWithScope();
+ void ExitWithScope();
+
+ /*! \brief Helper to solve related variable's bound within conditional
scope.*/
+ ffi::Optional<ffi::Map<Var, Range>> TrySolveCondition();
+
+ /*! \brief the condition holds on true branch. */
+ const PrimExpr& condition_;
+ /*! \brief domain map for relaxed vars to update */
+ std::unordered_map<const VarNode*, arith::IntSet>* relax_map_;
+ /*! \brief domain map for free vars to update */
+ std::unordered_map<const VarNode*, arith::IntSet>* hint_map_;
+ /*! \brief unresolved condition stack */
+ std::vector<PrimExpr>* pending_conditions_;
+ /*! \brief used to record and restore original var bounds */
+ std::unordered_map<const VarNode*, arith::IntSet> origin_map_;
+ /*! \brief used to record unresolved conditions num. */
+ size_t origin_pending_conditions_num_;
+};
+
+} // namespace s_tir
+} // namespace tvm
+
+#endif // TVM_S_TIR_ANALYSIS_CONDITIONAL_BOUNDS_H_
diff --git a/src/s_tir/analysis/sblock_access_region_detector.cc
b/src/s_tir/analysis/sblock_access_region_detector.cc
index 9b13d6c511..b8699b1526 100644
--- a/src/s_tir/analysis/sblock_access_region_detector.cc
+++ b/src/s_tir/analysis/sblock_access_region_detector.cc
@@ -33,6 +33,8 @@
#include <unordered_set>
#include "../../tirx/transform/ir_utils.h"
+#include "conditional_bounds.h"
+
namespace tvm {
namespace tirx {
@@ -215,12 +217,14 @@ ffi::Optional<VisitInterrupt>
BlockReadWriteDetector::Visit_(const IfThenElseNod
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->condition));
{
// Visit then branch
- With<ConditionalBoundsContext> ctx(op->condition, &dom_map_, &hint_map_,
&pending_conditions_);
+ With<s_tir::ConditionalBoundsContext> ctx(op->condition, &dom_map_,
&hint_map_,
+ &pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->then_case));
}
if (op->else_case) {
// Visit else branch
- With<ConditionalBoundsContext> ctx(!op->condition, &dom_map_, &hint_map_,
&pending_conditions_);
+ With<s_tir::ConditionalBoundsContext> ctx(!op->condition, &dom_map_,
&hint_map_,
+ &pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit(op->else_case.value()));
}
return std::nullopt;
@@ -316,13 +320,15 @@ ffi::Optional<VisitInterrupt>
BlockReadWriteDetector::Visit_(const CallNode* op)
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(condition));
{
// Visit then branch
- With<ConditionalBoundsContext> ctx(condition, &dom_map_, &hint_map_,
&pending_conditions_);
+ With<s_tir::ConditionalBoundsContext> ctx(condition, &dom_map_,
&hint_map_,
+ &pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
StmtExprVisitor::Visit(op->args[1].as_or_throw<PrimExpr>()));
}
{
// Visit else branch
- With<ConditionalBoundsContext> ctx(!condition, &dom_map_, &hint_map_,
&pending_conditions_);
+ With<s_tir::ConditionalBoundsContext> ctx(!condition, &dom_map_,
&hint_map_,
+ &pending_conditions_);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(
StmtExprVisitor::Visit(op->args[2].as_or_throw<PrimExpr>()));
}
diff --git a/src/s_tir/transform/compact_buffer_region.cc
b/src/s_tir/transform/compact_buffer_region.cc
index 3c81566cf6..5cacf41bad 100644
--- a/src/s_tir/transform/compact_buffer_region.cc
+++ b/src/s_tir/transform/compact_buffer_region.cc
@@ -23,7 +23,6 @@
*/
#include <tvm/arith/int_set.h>
-#include <tvm/arith/int_solver.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
@@ -38,6 +37,7 @@
#include "../../support/arena.h"
#include "../../support/utils.h"
#include "../../tirx/transform/ir_utils.h"
+#include "../analysis/conditional_bounds.h"
#include "../schedule/utils.h"
#include "../support/nd_int_set.h"
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 864f781c06..2901c01c85 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -16,6 +16,8 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/arith/analyzer.h>
+
#include <algorithm>
#include "../../../tirx/transform/ir_utils.h" // For `GetPtrStorageScope`
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index 653e61a6a7..c629688ee0 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -24,7 +24,6 @@
#include "ir_utils.h"
#include <tvm/arith/analyzer.h>
-#include <tvm/arith/int_solver.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
@@ -785,150 +784,6 @@ Region ConvertRegion(const MatchBufferRegion&
match_buffer, const Region& region
return result;
}
-ffi::Optional<arith::IntConstraints>
ConditionalBoundsContext::TrySolveCondition() {
- // extract equations and related vars from condition expression.
- // currently only extract simple integral equations which could be solvable.
- arith::Analyzer analyzer;
- PrimExpr condition = analyzer->Simplify(condition_);
- if (is_const_int(condition)) {
- return std::nullopt;
- }
- ffi::Array<PrimExpr> equations;
- ffi::Array<PrimVar> vars;
- std::function<void(const PrimExpr&)> fvisit = [&equations, &vars,
&fvisit](const PrimExpr& e) {
- if (e->IsInstance<prim::GENode>() || e->IsInstance<prim::GTNode>() ||
- e->IsInstance<prim::LENode>() || e->IsInstance<prim::LTNode>() ||
- e->IsInstance<prim::EQNode>() || e->IsInstance<prim::NENode>()) {
- bool is_simple = true;
- std::vector<PrimVar> cand_vars;
- auto walk_fn = [&cand_vars, &is_simple,
- &e](const PrimExpr& obj) ->
ffi::Expected<ffi::WalkResult> {
- if (obj.same_as(e)) {
- return ffi::WalkResult::Advance();
- } else if (const VarNode* var = obj.as<VarNode>()) {
- PrimType var_ty = var->ty.as_or_throw<PrimType>();
- if (var_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
- cand_vars.push_back(ffi::GetRef<Var>(var).as_or_throw<PrimVar>());
- }
- } else {
- is_simple &= obj->IsInstance<prim::AddNode>() ||
obj->IsInstance<prim::SubNode>() ||
- obj->IsInstance<prim::MulNode>() ||
obj->IsInstance<prim::FloorDivNode>() ||
- obj->IsInstance<prim::FloorModNode>() ||
obj->IsInstance<IntImmNode>();
- }
- return ffi::WalkResult::Advance();
- };
- ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(e, walk_fn);
- if (is_simple && !cand_vars.empty()) {
- for (const PrimVar& new_var : cand_vars) {
- if (!std::any_of(vars.begin(), vars.end(),
- [&new_var](const PrimVar& v) { return
v.same_as(new_var); })) {
- vars.push_back(new_var);
- }
- }
- equations.push_back(e.as_or_throw<PrimExpr>());
- }
- } else if (e->IsInstance<prim::AndNode>()) {
- prim::And op = e.as_or_throw<prim::And>();
- fvisit(op->a);
- fvisit(op->b);
- } else if (e->IsInstance<CallNode>()) {
- Call op = e.as_or_throw<Call>();
- if (op->op.same_as(prim::builtin::likely())) {
- fvisit(op->args[0].as_or_throw<PrimExpr>());
- }
- }
- };
- fvisit(condition);
- if (equations.empty() || vars.empty()) {
- return std::nullopt;
- }
- // build dom ranges for related vars
- ffi::Map<Var, Range> ranges;
- for (const Var& v : vars) {
- arith::IntSet dom;
- auto relax_it = relax_map_->find(v.get());
- if (relax_it != relax_map_->end()) {
- dom = relax_it->second;
- } else {
- auto hint_it = hint_map_->find(v.get());
- if (hint_it != hint_map_->end()) {
- dom = hint_it->second;
- }
- }
- if (dom.defined()) {
- ranges.Set(v, Range::FromMinExtent(dom.min(),
analyzer->Simplify(dom.max() - dom.min() + 1)));
- }
- }
- // solve constraints
- arith::IntConstraints constraint(vars, ranges, equations);
- arith::IntConstraints result = arith::SolveInequalitiesToRange(constraint);
- if (!result->relations.empty()) {
- return std::nullopt;
- }
- return result;
-}
-
-ConditionalBoundsContext::ConditionalBoundsContext(
- const PrimExpr& condition, std::unordered_map<const VarNode*,
arith::IntSet>* relax_map,
- std::unordered_map<const VarNode*, arith::IntSet>* hint_map,
- std::vector<PrimExpr>* pending_conditions)
- : condition_(condition),
- relax_map_(relax_map),
- hint_map_(hint_map),
- pending_conditions_(pending_conditions),
- origin_pending_conditions_num_(pending_conditions->size()) {}
-
-void ConditionalBoundsContext::EnterWithScope() {
- ffi::Optional<arith::IntConstraints> constraints = TrySolveCondition();
- if (!constraints.has_value()) {
- // fail to process the condition, add to unresolved
- pending_conditions_->push_back(condition_);
- return;
- }
- // update solved var ranges
- for (const auto& kv : constraints.value()->ranges) {
- const VarNode* var = kv.first.get();
- arith::IntSet new_dom = arith::IntSet::FromRange(kv.second);
- auto relax_it = relax_map_->find(var);
- if (relax_it != relax_map_->end()) {
- // this is a bound for relaxed var
- origin_map_.emplace(var, relax_it->second);
- relax_it->second = arith::Intersect({relax_it->second, new_dom});
- } else {
- // this is a bound for free var
- auto hint_it = hint_map_->find(var);
- if (hint_it != hint_map_->end()) {
- origin_map_.emplace(var, hint_it->second);
- hint_it->second = arith::Intersect({hint_it->second, new_dom});
- } else {
- origin_map_.emplace(var, arith::IntSet::Nothing());
- hint_map_->insert(hint_it, {var, new_dom});
- }
- }
- }
-}
-
-void ConditionalBoundsContext::ExitWithScope() {
- pending_conditions_->resize(origin_pending_conditions_num_);
- for (const auto& p : origin_map_) {
- const auto* var = p.first;
- auto relax_it = relax_map_->find(var);
- if (relax_it != relax_map_->end()) {
- // recover bound for relaxed var
- relax_it->second = p.second;
- } else {
- // recover bound for free var
- auto hint_it = hint_map_->find(var);
- TVM_FFI_ICHECK(hint_it != hint_map_->end());
- if (p.second.IsNothing()) {
- hint_map_->erase(hint_it);
- } else {
- hint_it->second = p.second;
- }
- }
- }
-}
-
std::pair<PrimExpr, PrimExpr> GetAsyncWaitAttributes(const AttrStmtNode* op) {
TVM_FFI_ICHECK(op && op->attr_key == s_tir::attr::async_wait_queue_scope);
auto inner = op->body.as<AttrStmtNode>();
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index 9ad1099a91..1b24e462ee 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -25,7 +25,6 @@
#define TVM_TIR_TRANSFORM_IR_UTILS_H_
#include <tvm/arith/int_set.h>
-#include <tvm/arith/int_solver.h>
#include <tvm/ffi/container/tuple.h>
#include <tvm/ir/prim/builtin.h>
#include <tvm/ir/prim/expr.h>
@@ -255,48 +254,6 @@ Region ConvertRegion(const MatchBufferRegion&
match_buffer, const Region& region
*/
ffi::Array<PrimExpr> GetBufferAllocationShape(const BufferVar& buffer);
-/*!
- * \brief Context helper to update domain map within conditional scope.
- * Assume the condition is `0 <= i && i < 9` and domain of i is [0, 20], Then
- * `With<ConditionalBoundsContext> ctx(condition, &relax_map, &hint_map,
&constraints)`
- * step into scope where dom_map[i] is [0, 8]; and
- * `With<ConditionalBoundsContext> ctx(!condition, &relax_map, &hint_map,
&constraints)`
- * step into scope where dom_map[i] is [9, 20]
- */
-class ConditionalBoundsContext {
- private:
- friend class With<ConditionalBoundsContext>;
- /*!
- * \brief Construct a condition bounds context.
- * \param condition The condition holds on true branch.
- * \param relax_map The domain map for relaxed vars to update.
- * \param hint_map The domain map for free vars to update.
- * \param pending_conditions The stack of unresolved constraints.
- */
- ConditionalBoundsContext(const PrimExpr& condition,
- std::unordered_map<const VarNode*, arith::IntSet>*
relax_map,
- std::unordered_map<const VarNode*, arith::IntSet>*
hint_map,
- std::vector<PrimExpr>* pending_constraints);
- void EnterWithScope();
- void ExitWithScope();
-
- /*! \brief Helper to solve related variable's bound within conditional
scope.*/
- ffi::Optional<arith::IntConstraints> TrySolveCondition();
-
- /*! \brief the condition holds on true branch. */
- const PrimExpr& condition_;
- /*! \brief domain map for relaxed vars to update */
- std::unordered_map<const VarNode*, arith::IntSet>* relax_map_;
- /*! \brief domain map for free vars to update */
- std::unordered_map<const VarNode*, arith::IntSet>* hint_map_;
- /*! \brief unresolved condition stack */
- std::vector<PrimExpr>* pending_conditions_;
- /*! \brief used to record and restore original var bounds */
- std::unordered_map<const VarNode*, arith::IntSet> origin_map_;
- /*! \brief used to record unresolved conditions num. */
- size_t origin_pending_conditions_num_;
-};
-
// Information of tensor core fragment.
struct FragmentInfo {
// fragment shape
diff --git a/tests/python/arith/test_arith_solve_linear_equations.py
b/tests/python/arith/test_arith_solve_linear_equations.py
deleted file mode 100644
index b9550d570d..0000000000
--- a/tests/python/arith/test_arith_solve_linear_equations.py
+++ /dev/null
@@ -1,186 +0,0 @@
-# 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.
-# ruff: noqa: F401
-import random
-import sys
-
-import pytest
-import tvm_ffi
-
-import tvm
-from tvm import arith, ir, testing, tirx
-from tvm.script import tirx as T
-
-
-def test_solution_consistency():
- seed = random.randrange(sys.maxsize)
- print(
- "\nThis test is intentionally non-deterministic, "
- f"if it fails please report it in GitHub issue together with this seed
{seed}\n"
- )
- random.seed(seed)
-
- def _check(num_vars, num_formulas, coef=(-5, 5), bounds=(-20, 20)):
- variables = [tvm.tirx.Var("x" + str(i), "int32") for i in
range(num_vars)]
-
- relations = []
- for i in range(num_formulas):
- s1 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
- s1 += random.randint(coef[0], coef[1])
- s2 = sum([v * random.randint(coef[0], coef[1]) for v in variables])
- s2 += random.randint(coef[0], coef[1])
- if random.random() < 0.7:
- op = tvm.tirx.EQ
- else:
- # we also make sure it can correctly handle inequalities
- op = random.choice([tvm.tirx.LE, tvm.tirx.LT, tvm.tirx.GE,
tvm.tirx.GT])
- relations.append(op(s1, s2))
-
- vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in
variables}
- solution = arith.solve_linear_equations(relations, variables, vranges)
-
- testing.check_int_constraints_trans_consistency(solution)
-
- # leaving some variables as parameters should also be ok
- for k in [1, 2]:
- if len(variables) > k:
- solution = arith.solve_linear_equations(relations,
variables[:-k], vranges)
- param_ranges = {v: vranges[v] for v in variables[-k:]}
- testing.check_int_constraints_trans_consistency(solution,
param_ranges)
-
- for i in range(2):
- _check(num_vars=1, num_formulas=1)
- for i in range(2):
- _check(num_vars=1, num_formulas=2)
-
- for i in range(2):
- _check(num_vars=2, num_formulas=1)
- for i in range(2):
- _check(num_vars=2, num_formulas=2)
- for i in range(2):
- _check(num_vars=2, num_formulas=3)
-
- for i in range(3):
- _check(num_vars=3, num_formulas=3, coef=(-2, 2))
- for i in range(3):
- _check(num_vars=3, num_formulas=4, coef=(-2, 2))
-
- for i in range(3):
- _check(num_vars=4, num_formulas=3, coef=(-1, 1))
-
- for i in range(3):
- _check(num_vars=10, num_formulas=2, coef=(-1, 1), bounds=(0, 4))
- for i in range(3):
- _check(num_vars=10, num_formulas=3, coef=(0, 1), bounds=(0, 4))
-
-
-def test_empty_var_to_solve():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
- equations = [
- tvm.tirx.EQ(x + y, 20),
- tvm.tirx.EQ(x - y, 10),
- ]
- solution = arith.solve_linear_equations(equations)
- assert len(solution.src_to_dst) == 0
- assert len(solution.dst_to_src) == 0
- assert len(solution.src.variables) == 0
- assert len(solution.src.ranges) == 0
- assert tvm_ffi.structural_equal(solution.src.relations, equations)
- assert tvm_ffi.structural_equal(solution.src, solution.dst)
-
-
-def test_unique_solution():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
-
- solution = arith.solve_linear_equations(
- [
- tvm.tirx.EQ(x + y, 20),
- tvm.tirx.EQ(x - y, 10),
- ],
- [x, y],
- )
- assert list(solution.dst.variables) == []
- assert tvm_ffi.structural_equal(solution.src_to_dst[x], T.int32(15))
- assert tvm_ffi.structural_equal(solution.src_to_dst[y], T.int32(5))
-
-
-def test_low_rank():
- x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"),
tvm.tirx.Var("z", "int32")
- ranges = {}
-
- solution = arith.solve_linear_equations(
- [
- tvm.tirx.EQ(x + y + z, 15),
- tvm.tirx.EQ(x + y, 10),
- ],
- [x, y, z],
- ranges,
- )
- [n0] = solution.dst.variables
- assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0 + 10)
- assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
- assert tvm_ffi.structural_equal(solution.src_to_dst[z], T.int32(5))
-
-
-def test_infer_range():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
- ranges = {
- x: tvm.ir.Range.from_min_extent(-5, 10),
- y: tvm.ir.Range.from_min_extent(0, 10),
- }
-
- solution = arith.solve_linear_equations(
- [
- tvm.tirx.EQ(x + y, 0),
- ],
- [x, y],
- ranges,
- )
- [n0] = solution.dst.variables
- assert tvm_ffi.structural_equal(solution.src_to_dst[x], n0)
- assert tvm_ffi.structural_equal(solution.src_to_dst[y], -n0)
- # inferred from y's range
- assert tvm_ffi.structural_equal(solution.dst.ranges[n0].min, T.int32(-9))
- assert tvm_ffi.structural_equal(solution.dst.ranges[n0].extent,
T.int32(10))
- # additional inequality is added into the system for x
- [ineq] = solution.dst.relations
- assert isinstance(ineq, tvm.tirx.LE)
- assert tvm_ffi.structural_equal(ineq.a, T.int32(-5))
- assert tvm_ffi.structural_equal(ineq.b, n0)
-
-
-def test_ill_formed():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
-
- solution = arith.solve_linear_equations(
- [
- tvm.tirx.EQ(x + y, 0),
- tvm.tirx.EQ(x - y, 0),
- tvm.tirx.EQ(x, 5),
- ],
- [x, y],
- {},
- )
- assert list(solution.dst.variables) == []
- [rel] = solution.dst.relations
- ir.assert_structural_equal(rel, tirx.const(False))
- assert len(solution.src_to_dst) == 0
- assert len(solution.dst_to_src) == 0
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/arith/test_arith_solve_linear_inequality.py
b/tests/python/arith/test_arith_solve_linear_inequality.py
deleted file mode 100644
index 04b109fef3..0000000000
--- a/tests/python/arith/test_arith_solve_linear_inequality.py
+++ /dev/null
@@ -1,223 +0,0 @@
-# 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.
-import random
-import sys
-
-import pytest
-import tvm_ffi
-
-import tvm
-from tvm import arith, ir, testing, tirx
-from tvm.script import tirx as T
-
-
[email protected](reason="See https://github.com/apache/tvm/issues/11458";)
-def test_solution_consistency():
- seed = random.randrange(sys.maxsize)
- print(
- "\nThis test is intentionally non-deterministic, "
- f"if it fails please report it in GitHub issue together with this seed
{seed}\n"
- )
- random.seed(seed)
-
- def _check(variables, formulas, coef=(-5, 5), bounds=(-20, 20)):
- vs = [tvm.tirx.Var("x" + str(i), "int32") for i in range(variables)]
-
- fs = []
- for i in range(formulas):
- s1 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
- s1 += random.randint(coef[0], coef[1])
- s2 = sum([v * random.randint(coef[0], coef[1]) for v in vs])
- s2 += random.randint(coef[0], coef[1])
- op = random.choice(
- [tirx.expr.EQ, tirx.expr.LE, tirx.expr.LT, tirx.expr.GE,
tirx.expr.GT]
- )
- fs.append(op(s1, s2))
-
- vranges = {v: tvm.ir.expr.Range(bounds[0], bounds[1] + 1) for v in vs}
- before = tvm.tirx.all(tirx.const(1, "bool"), *fs)
- after = arith._ffi_api.SolveInequalitiesAsCondition(vs, vranges, fs)
- after = tvm.tirx.all(tirx.const(1, "bool"), *after)
- testing.check_bool_expr_is_true(before == after, vranges)
-
- solution = arith.solve_linear_inequalities(fs, vs, vranges,
deskew_range=True)
- testing.check_int_constraints_trans_consistency(solution)
-
- for i in range(3):
- _check(1, 1)
- for i in range(3):
- _check(1, 2)
-
- for i in range(3):
- _check(2, 1)
- for i in range(3):
- _check(2, 2)
- for i in range(3):
- _check(2, 3)
-
- # Somewhere here coefficients in the results become too large, leading to
overflow,
- # so we use smaller initial coefficients
- for i in range(5):
- _check(3, 3, coef=(-2, 2))
- for i in range(5):
- _check(3, 4, coef=(-2, 2))
-
- for i in range(5):
- _check(4, 3, coef=(-1, 1))
-
- for i in range(5):
- _check(10, 2, coef=(-1, 1), bounds=(0, 4))
- for i in range(5):
- _check(10, 3, coef=(0, 1), bounds=(0, 4))
-
-
-def test_dual_variable():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
-
- variables = [x, y]
- ranges = {
- x: tvm.ir.Range(-100, 100),
- y: tvm.ir.Range(0, 10),
- }
- problem = [
- tvm.tirx.LE(x + y, 20),
- tvm.tirx.GE(x - y, 10),
- ]
-
- # solution as conditions
- solution = arith._ffi_api.SolveInequalitiesAsCondition(variables, ranges,
problem)
- assert tvm_ffi.structural_equal(solution[0], x >= (y + 10))
- assert tvm_ffi.structural_equal(solution[1], x <= (20 - y))
- assert tvm_ffi.structural_equal(solution[2], y >= 0)
- assert tvm_ffi.structural_equal(solution[3], y <= 5)
-
- # solve and get the ranges
- solution = arith.solve_linear_inequalities(problem, variables, ranges)
- # 0 <= y <=5
- assert solution.ranges[y].min == 0
- assert solution.ranges[y].extent == 6
- # y + 10 <= x <= 20 - y
- assert tvm_ffi.structural_equal(solution.ranges[x].min, y + 10)
- assert solution.ranges[x].extent == 11 # max(10 - 2y)
-
- # deskew the solved ranges to be starting from zero
- solution = arith.solve_linear_inequalities(problem, variables, ranges,
deskew_range=True)
- [x_new, y_new] = solution.dst.variables
- [rel] = solution.dst.relations
- assert tvm_ffi.structural_equal(rel, (y_new * 2) + x_new <= 10)
- assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].min, T.int32(0))
- assert tvm_ffi.structural_equal(solution.dst.ranges[x_new].extent,
T.int32(11))
- assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].min, T.int32(0))
- assert tvm_ffi.structural_equal(solution.dst.ranges[y_new].extent,
T.int32(6))
- assert tvm_ffi.structural_equal(solution.src_to_dst[x], x_new + (y_new +
10))
- assert tvm_ffi.structural_equal(solution.src_to_dst[y], y_new)
- assert tvm_ffi.structural_equal(solution.dst_to_src[x_new], x - y - 10)
- assert tvm_ffi.structural_equal(solution.dst_to_src[y_new], y)
-
-
-def test_equal():
- x, y = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32")
- problem = [
- tvm.tirx.GE(x + y, 10),
- tvm.tirx.GE(x - y, 2),
- tvm.tirx.LE(x, 6),
- ]
-
- solution = arith.solve_linear_inequalities(problem, [x, y])
- assert solution.ranges[x].min == 6
- assert solution.ranges[x].extent == 1
- assert solution.ranges[y].min == 4
- assert solution.ranges[y].extent == 1
-
- solution = arith.solve_linear_inequalities(problem, [x, y],
deskew_range=True)
- assert len(solution.dst.variables) == 0
- assert len(solution.dst.ranges) == 0
- assert len(solution.dst.relations) == 0
- assert solution.src_to_dst[x] == 6
- assert solution.src_to_dst[y] == 4
-
-
-def test_multi_equal():
- x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"),
tvm.tirx.Var("z", "int32")
- problem = [
- tvm.tirx.LE(x, 6),
- tvm.tirx.GE(x, 6),
- tvm.tirx.GE(x - z * y, 0),
- tvm.tirx.LE(x - z * y, 0),
- ]
-
- solution = arith.solve_linear_inequalities(problem, [x, y, z])
- assert solution.ranges[x].min == 6
- assert solution.ranges[x].extent == 1
- assert len(solution.relations) == 3
- assert tvm_ffi.structural_equal(solution.relations[0], x == z * y)
-
- assert isinstance(solution.relations[1], tvm.tirx.LE)
- assert solution.relations[1].b == 0
- assert isinstance(solution.relations[2], tvm.tirx.LE)
- assert solution.relations[2].b == 0
- # (z*y - 6) <= 0 && (6 - z*y) <= 0
- ana = tvm.arith.Analyzer()
- assert ana.simplify(solution.relations[1].a + solution.relations[2].a) == 0
- assert tvm_ffi.structural_equal(
- solution.relations[1].a, (z * y - 6)
- ) or tvm_ffi.structural_equal(solution.relations[2].a, (z * y - 6))
-
- solution = arith.solve_linear_inequalities(problem, [x, y, z],
deskew_range=True)
- assert solution.src_to_dst[y] == y
- assert solution.src_to_dst[z] == z
- assert solution.src_to_dst[x] == 6
-
-
-def test_no_solution():
- x = tvm.tirx.Var("x0", "int32")
- vranges = {x: tvm.ir.Range.from_min_extent(-20, 41)}
- problem = [-x - 4 <= -5 * x + 2, x * 4 + 5 <= x * 5]
-
- solution = arith.solve_linear_inequalities(problem, [x], vranges,
deskew_range=True)
- assert list(solution.dst.variables) == []
- [rel] = solution.dst.relations
- ir.assert_structural_equal(rel, tirx.const(False))
- assert len(solution.src_to_dst) == 0
- assert len(solution.dst_to_src) == 0
-
- solution = arith.solve_linear_inequalities(problem, [x], vranges)
- assert len(solution.variables) == 0
- assert len(solution.ranges) == 0
- [rel] = solution.relations
- assert not rel
-
-
-def test_unbound_var_range():
- x = tvm.tirx.Var("x0", "int32")
- free_var = tvm.tirx.Var("fv", "int32")
- vranges = {
- x: tvm.ir.Range.from_min_extent(0, tvm.tirx.Cast("int32", 1 +
tvm.tirx.log(free_var)))
- }
- problem = [x > 3]
- solution = arith.solve_linear_inequalities(
- problem,
- [x],
- vranges,
- )
- assert len(solution.variables) == 1
- assert len(solution.ranges) == 0
- assert len(solution.relations) == 3
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/s_tir/analysis/test_sblock_access_region.py
b/tests/python/s_tir/analysis/test_sblock_access_region.py
index b9266be09a..7218fad53f 100644
--- a/tests/python/s_tir/analysis/test_sblock_access_region.py
+++ b/tests/python/s_tir/analysis/test_sblock_access_region.py
@@ -457,5 +457,72 @@ def test_buffer_access_with_nested_let_binding():
tvm.ir.assert_structural_equal(block.writes, ret[1])
[email protected]("case", ["coupled", "equal", "nonlinear", "empty",
"unbounded", "rounded"])
+def test_conditional_inequality_access_regions(case):
+ # Retain the live cases from the former arith inequality solver tests
through
+ # the block-access consumer, including its conservative unresolved
fallback.
+ tirx = tvm.tirx
+ x, y, z = [tirx.Var(name, "int32") for name in ("x", "y", "z")]
+ free_var = tirx.Var("free_var", "int32")
+ unbounded_extent = tirx.Cast("int32", 1 + tirx.log(free_var))
+ cases = {
+ "coupled": (
+ [x, y],
+ [(-100, 200), (0, 10)],
+ tirx.all(x + y <= 20, x - y >= 10),
+ [(-100, 200), (0, 10)],
+ ),
+ "equal": (
+ [x, y],
+ [(-100, 200), (-100, 200)],
+ tirx.all(x + y >= 10, x - y >= 2, x <= 6),
+ [(6, 1), (4, 1)],
+ ),
+ "nonlinear": (
+ [x, y, z],
+ [(-100, 200)] * 3,
+ tirx.all(x <= 6, x >= 6, x - z * y >= 0, x - z * y <= 0),
+ [(-100, 200)] * 3,
+ ),
+ "empty": (
+ [x],
+ [(-20, 41)],
+ tirx.all(-x - 4 <= -5 * x + 2, x * 4 + 5 <= x * 5),
+ [(-20, 41)],
+ ),
+ "unbounded": ([x], [(0, unbounded_extent)], x > 3, [(0, 256)]),
+ "rounded": ([x], [(-20, 41)], tirx.all(x * 3 >= -7, x * 2 <= 9), [(-2,
7)]),
+ }
+ variables, domains, condition, expected = cases[case]
+ inside = tirx.decl_buffer([256] * len(variables), name="inside")
+ outside = tirx.decl_buffer([256] * len(variables), name="outside")
+ body = tirx.SeqStmt(
+ [
+ tirx.IfThenElse(condition,
tirx.Evaluate(inside[tuple(variables)]), None),
+ tirx.Evaluate(outside[tuple(variables)]),
+ ]
+ )
+ for var, (minimum, extent) in reversed(list(zip(variables, domains))):
+ body = tirx.For(var, minimum, extent, tirx.ForKind.SERIAL, body)
+ block = tirx.SBlock([], [], [], "conditional", body)
+ # Unbounded access sets conservatively cover the whole buffer.
+ outside_expected = [(0, 256)] if case == "unbounded" else domains
+ reads, writes, opaque = s_tir.analysis.get_sblock_access_region(
+ block, {inside: inside, outside: outside}
+ )
+ tvm.ir.assert_structural_equal(
+ reads,
+ [
+ tirx.BufferRegion(inside, [Range.from_min_extent(*bounds) for
bounds in expected]),
+ # Leaving the conditional scope must restore the original domains.
+ tirx.BufferRegion(
+ outside, [Range.from_min_extent(*bounds) for bounds in
outside_expected]
+ ),
+ ],
+ )
+ assert not writes
+ assert not opaque
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx/codegen/test_cuda_wait_until.py
b/tests/python/tirx/codegen/test_cuda_wait_until.py
index 26fa522626..20be1e8da9 100644
--- a/tests/python/tirx/codegen/test_cuda_wait_until.py
+++ b/tests/python/tirx/codegen/test_cuda_wait_until.py
@@ -95,9 +95,7 @@ def packed_contribution():
def _wait_macro(source):
return next(
- line
- for line in source.splitlines()
- if line.startswith("#define") and "wait_until" in line
+ line for line in source.splitlines() if line.startswith("#define") and
"wait_until" in line
)
@@ -327,9 +325,7 @@ def
test_a_declared_word_is_global_and_says_where_a_shared_wait_belongs():
"""
with pytest.raises(ValueError, match="mbarrier"):
- T.cuda.wait_until(
- None, None, None, scope="cta", space="shared", ptx_type="b32"
- )
+ T.cuda.wait_until(None, None, None, scope="cta", space="shared",
ptx_type="b32")
def test_the_four_direct_forms_are_gone():
