mirrors/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/arith/analyzer.h
+993
View File
@@ -0,0 +1,993 @@
/*
* 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/analyzer.h
* \brief Algebra expression simplifications.
*/
#ifndef TVM_ARITH_ANALYZER_H_
#define TVM_ARITH_ANALYZER_H_
#include <tvm/arith/int_set.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/with_context.h>
#include <limits>
#include <memory>
#include <unordered_map>
#include <utility>
#include <vector>
namespace tvm {
/*! \brief namespace of arithmetic analysis. */
namespace arith {
//-------------------------------------------------------
// Base integer analysis API.
//
// We have multiple type of analyzers to do relaxed
// integer set analysis(bound analysis, modulo) and
// equivalence checking and simplification.
//
// Importantly, each analyzer may need result from
// another analyzer.
//-------------------------------------------------------
// Forward declare the analyzer object and its reference handle.
class AnalyzerObj;
class Analyzer;
class ConstraintContext;
using tirx::Var;
enum DivMode {
/*! \brief Truncated division. */
kTruncDiv,
/*! \brief Floor division. */
kFloorDiv
};
/*!
* \brief The strength used in top-level condition proves
* \note The higher, the more time consuming it can be.
*
* Do not use level beyond kDefault in internal recursive rewriting in arith
* analysis and only use it at top-level simplification to avoid speed issues.
*/
enum class ProofStrength : int {
/*! \brief default strength, can be used in. */
kDefault = 0,
/*!
* \brief Prove using symbolic bound analysis
*/
kSymbolicBound = 1,
};
/*!
* \brief Constant integer up and lower bound(inclusive).
* Useful for value bound analysis.
*
* set = [min_value, max_value]
*/
class ConstIntBoundNode : public ffi::Object {
public:
int64_t min_value;
int64_t max_value;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ConstIntBoundNode>()
.def_ro("min_value", &ConstIntBoundNode::min_value)
.def_ro("max_value", &ConstIntBoundNode::max_value);
}
/*! \brief Number to represent +inf */
static const constexpr int64_t kPosInf = std::numeric_limits<int64_t>::max();
/*!
* \brief Number to represent -inf
* \note We can make use the of fact that -kPosInf == kNegInf in the project.
*/
static const constexpr int64_t kNegInf = -kPosInf;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ConstIntBound", ConstIntBoundNode, ffi::Object);
};
/*!
* \brief reference class to ConstIntBoundNode
* \sa ConstIntBoundNode
*/
class ConstIntBound : public ffi::ObjectRef {
public:
/*!
* \brief constructor by fields.
* \param min_value The mininum value.
* \param max_value The maximum value.
*/
TVM_DLL ConstIntBound(int64_t min_value, int64_t max_value);
static const constexpr int64_t kPosInf = ConstIntBoundNode::kPosInf;
static const constexpr int64_t kNegInf = ConstIntBoundNode::kNegInf;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ConstIntBound, ffi::ObjectRef, ConstIntBoundNode);
};
/*!
* \brief Analyzer to get constant integer bound over expression.
*/
class ConstIntBoundAnalyzer {
public:
using BoundMapType =
std::unordered_map<PrimExpr, ConstIntBound, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL ConstIntBound operator()(const PrimExpr& expr) const;
/*!
* \brief analyze the expr with the intermediate memorized to avoid redundant computation
* \param expr The expression of interest.
* \param bound The lookup table to store the intermediate results
* \return the result of the analysis.
*/
TVM_DLL ConstIntBound operator()(const PrimExpr& expr, BoundMapType* bound);
/*!
* \brief Update constant int bound information of var.
*
* \param var The variable of interest.
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const ConstIntBound& info, bool allow_override = false);
/*!
* \brief Bind variable to a range.
*
* \param var The variable.
* \param range The range we bind to.
* \param allow_override Whether we allow overriding an existing var's range.
*/
TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Check if a variable is bound to a range.
* \param var The variable.
* \return Whether the variable is bound to a range.
*/
TVM_DLL bool IsBound(const Var& var) const;
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ConstIntBoundAnalyzer(AnalyzerObj* parent);
TVM_DLL ~ConstIntBoundAnalyzer();
void CopyFrom(const ConstIntBoundAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
struct Entry;
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Range of a linear integer function.
* Use to do specify the possible index values.
*
* set = { coeff * x + base | x in Z }
*
* When coeff != 0, it can also be written as
* set = { n | n % coeff == base }
*
* This is useful to decide if the index is dividable by certain value.
* For example, if index = 0 + 4 x, then we know it can be divided by 4.
*/
class ModularSetNode : public ffi::Object {
public:
/*! \brief linear co-efficient */
int64_t coeff;
/*! \brief The base */
int64_t base;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ModularSetNode>()
.def_ro("coeff", &ModularSetNode::coeff)
.def_ro("base", &ModularSetNode::base);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.ModularSet", ModularSetNode, ffi::Object);
};
/*!
* \brief reference of ModularSetNode
* \sa ModularSetNode
*/
class ModularSet : public ffi::ObjectRef {
public:
TVM_DLL ModularSet(int64_t coeff, int64_t base);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ModularSet, ffi::ObjectRef, ModularSetNode);
};
/*!
* \brief Analyzer to get modular information over expression.
*/
class ModularSetAnalyzer {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL ModularSet operator()(const PrimExpr& expr);
/*!
* \brief Update constant int bound information of var.
*
* \param var The variable of interest.
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const ModularSet& info, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ModularSetAnalyzer(AnalyzerObj* parent);
TVM_DLL ~ModularSetAnalyzer();
void CopyFrom(const ModularSetAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
struct Entry;
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Rewrite-rule based simplifier.
*/
class RewriteSimplifier {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_expr
* \param allow_override Whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
/*! \brief Flags to enable more computationally-intensive simplifications
*
* These simplifications may be required for specific schedules, but
* would impose too high a compile-time cost to enable by default.
* They can be enabled on an as-needed basis by calling
* `RewriteSimplifier::SetEnabledExtensions` prior to using
* `RewriteSimplifier::operator()`.
*
* Flags are defined as powers of two to allow future expansion. To
* enable multiple extensions, a user should pass a bitwise OR of the
* flags for each desired extension.
*/
enum Extension {
// No extensions enabled
kNone = 0,
/* When simplifying an inequality, attempt to use scope-based knowns.
*
* Example:
* if_then_else(i<j && j<k, i<k, false) => if_then_else(i<j && j<k, true, false)
*/
kTransitivelyProveInequalities = (1 << 0),
/* When simplifying a boolean expression, convert to an AND of ORs
* (conjunctive normal form).
*
* Example:
* (a && b) || c => (a || c) && (b || c)
*/
kConvertBooleanToAndOfOrs = (1 << 1),
/* When simplifying a boolean AND or a boolean OR, simplify each
* branch under the assumption that the other branch does not
* already dominate the result. That is, simplify each branch of
* (A && B) under the assumption that the other branch is true,
* and simplify each branch of (A || B) under the assumption that
* the other branch is false.
*
* Example:
* (n < 10) && (n < 5) => (n < 10)
* (n < 10) || (n < 5) => (n < 5)
*/
kApplyConstraintsToBooleanBranches = (1 << 2),
/* Special handling for expressions `(A+B)*C < (A*B)*D`
*
* Expressions of the form `(A+B)*C < (A*B)*D` can occur occur
* when comparing the number of operations required for two
* different orderings in which matrix multiplications can be
* performed. Proving or disproving this conditional allows an
* optimal order of execution to be selected, even for dynamic
* argument shapes.
*
* The default behavior of `ConstIntBounds` assumes that each term
* in an expression is independent, and is insufficient to prove
* these inequalities. For example, the maximum value of `(A+B)*C
* - (A*B)*D` is determined by taking the maximum value of
* `(A+B)*C` and subtracting the minimum value of `(A*B)*D`.
* While this algorithm can be applied in all cases, the bound it
* provides is looser than strictly required.
*
* This extension adds a check for this case. When `A`, `B`, `C`,
* and `D` are all positive values, as is the case for tensor
* shapes, the inequality can be written as `1/A + 1/B < D/C`. If
* this inequality holds for the minimum values of `A`, `B`, and
* `D`, along with the maximum value of `C`, then the inequality
* holds for all values.
*
* This extension requires little to no performance overhead, and
* may be enabled by default in future releases.
*/
kComparisonOfProductAndSum = (1 << 3),
};
/*! \brief Enable an optional extension or extensions
*
* \param flags A bitwise OR of all optional extensions that should
* be enabled.
*/
TVM_DLL void SetEnabledExtensions(Extension flags);
/*! \brief Return the currently enabled extensions */
TVM_DLL Extension GetEnabledExtensions() const;
/*! \brief Return the statistics counters */
TVM_DLL ffi::ObjectRef GetStatsCounters() const;
/*! \brief Reset the statistics counters */
TVM_DLL void ResetStatsCounters();
/*! \brief Set the maximum allowed number of rewrite steps
*
* By default, the simplifier may perform as many steps as are
* required. If a positive limit is set, then the simplifier will
* throw an exception when exceeding that number of rewrite steps.
* This allows tests to guard against performance regressions.
*
* Note: To maintain accurate usage counters, `Analyzer` instances
* should be re-used wherever possible. For example, TIR
* transformations should declare a single `Analyzer` that is used
* throughout the pass. Internal helper functions that only borrow
* the analyzer temporarily may receive the underlying `AnalyzerObj*`
* from their calling scope.
*/
TVM_DLL void SetMaximumRewriteSteps(int64_t maximum);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
friend class CanonicalSimplifier;
explicit RewriteSimplifier(AnalyzerObj* parent);
TVM_DLL ~RewriteSimplifier();
void CopyFrom(const RewriteSimplifier& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*!
* \brief Canonical-form based simplifier.
*/
class CanonicalSimplifier {
public:
/*!
* \brief analyze the expr
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_expr
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit CanonicalSimplifier(AnalyzerObj* parent);
TVM_DLL ~CanonicalSimplifier();
void CopyFrom(const CanonicalSimplifier& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
/*! \brief Structure for representing result of known
*
* Values are assigned to allow these flags to be used in bitwise
* operations.
*/
enum class CompareResult : int {
kInconsistent = 0,
kEQ = 1,
kLT = 2,
kLE = 3,
kGT = 4,
kGE = 5,
kNE = 6,
kUnknown = 7
};
inline constexpr CompareResult operator&(CompareResult lhs, CompareResult rhs) {
return CompareResult(static_cast<int>(lhs) & static_cast<int>(rhs));
}
inline constexpr CompareResult operator|(CompareResult lhs, CompareResult rhs) {
return CompareResult(static_cast<int>(lhs) | static_cast<int>(rhs));
}
/*!
* \brief Using previously specified knowns, compare the expressions provided
*
* Given known expressions [(a OP b), (b OP c), ..., (y OP z)], search
* for a known result for `(a OP z)`.
*/
class TransitiveComparisonAnalyzer {
public:
/* \brief Using previously specified knowns, compare the expressions provided
*
* \param lhs The left-hand side of the comparison
*
* \param rhs The right-hand side of the comparison
*
* \param propagate_inequalities If true, attempt to find a sequence
* of transitive inequalities that allow the lhs and rhs to be
* compared. If false, only use the known comparison that have been
* directly provided. Using `propagate_inequalities = false` is
* roughly equivalent to comparing against all known inequality
* expressions using `ExprDeepEqual`, but also allows for constant
* offsets on either side of the inequality.
*
* \return The most specific result that can be proven about the
* comparison. If nothing can be proven, returns kUnknown.
*/
TVM_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs,
bool propagate_inequalities = true);
/*! \brief Bind a variable as being equal to a known expression
*
* \param var The variable of interest.
* \param expr The bound expression
* \param allow_override Whether to allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*! \brief Bind a variable as being within a specified range
*
* \param var The variable of interest.
* \param range The known range
* \param allow_override Whether to allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
* \param constraint A constraint expression.
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
TransitiveComparisonAnalyzer();
TVM_DLL ~TransitiveComparisonAnalyzer();
void CopyFrom(const TransitiveComparisonAnalyzer& other);
class Impl;
/*! \brief Internal impl */
std::unique_ptr<Impl> impl_;
};
/*!
* \brief Integer set analyzer.
*/
class IntSetAnalyzer {
public:
/*!
* \brief Find a symbolic integer set that contains all possible values of
* expr given the domain of each variables.
*
* \param expr The expression of interest.
* \param dom_map The domain map to indicate which variable to relax.
* \return the result of the analysis.
*/
TVM_DLL IntSet operator()(const PrimExpr& expr, const ffi::Map<Var, IntSet>& dom_map);
/*!
* \brief Find a symbolic integer set that contains all possible
* values of expr given the domain of each variables, using
* the domain map defined by bound variables.
*
* \param expr The expression of interest.
* \return the result of the analysis.
*/
TVM_DLL IntSet operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_interval_set The set of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Update(const Var& var, const IntSet& new_interval_set, bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
explicit IntSetAnalyzer(AnalyzerObj* parent);
TVM_DLL ~IntSetAnalyzer();
void CopyFrom(const IntSetAnalyzer& other);
class Impl;
/*! \brief Internal impl */
Impl* impl_;
};
class Z3Prover {
public:
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
*
* \param var The variable of interest.
* \param expr The bound expression.
* \param allow_override whether we allow override of existing information.
*/
TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*!
* \brief Whether the Z3 backend is compiled into this build (USE_Z3=ON).
*
* \return true if the real Z3 prover is available, false for the stub.
*/
TVM_DLL bool IsEnabled() const;
/*!
* \brief Whether can we prove expr is always true.
*
* \param expr The expression.
* \return Whether we can prove it.
*/
TVM_DLL bool CanProve(const PrimExpr& expr);
/*!
* \brief Update the internal state to enter constraint.
*
* \param constraint A constraint expression.
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
std::function<void()> EnterConstraint(const PrimExpr& constraint);
/*!
* \brief Get the SMTLIB2 representation of the current context.
*
* \param expr The optional expression to check.
* \return The SMTLIB2 string.
*/
ffi::String GetSMTLIB2(const ffi::Optional<PrimExpr> expr);
/*!
* \brief Get statistics about Z3 prover.
*
* \return The statistics string.
*/
ffi::String GetStats();
/*!
* \brief Set timeout in milliseconds for Z3 prover.
*
* \param timeout_ms The timeout in milliseconds.
*/
void SetTimeoutMs(unsigned timeout_ms);
/*!
* \brief Set resource limitation for Z3 prover.
*
* \param rlimit the resource limitation.
*/
void SetRLimit(unsigned rlimit);
/*!
* \brief Get the Z3 model for the given expression if satisfiable.
*
* \param expr The expression to get the model for.
* \return The model as a string.
*/
ffi::String GetModel(const PrimExpr& expr);
/*!
* \brief Count the number of integer values that satisfy the current constraints.
*
* This method uses Z3's model enumeration to count how many distinct values of
* the given variable satisfy all current constraints.
*
* \param var The variable to count satisfying values for.
* \param max_count Maximum number of solutions to enumerate.
* \param min_consecutive Minimum consecutive count requirement.
* \return The number of distinct values that satisfy the constraints, or a negative error code.
*/
TVM_DLL int64_t CountSatisfyingValues(const Var& var, int64_t max_count = 2048,
int64_t min_consecutive = 1);
private:
friend class AnalyzerObj;
friend class Analyzer;
explicit Z3Prover(AnalyzerObj* parent);
TVM_DLL ~Z3Prover();
void CopyFrom(const Z3Prover& other);
class Impl;
Impl* impl_;
};
/*!
* \brief Analyzer that contains bunch of sub-analyzers.
*
* Each sub-analyzer can make use of another sub-analyzer
* by weak reference of this.
*
* NOTE for sub-analyzer developers:
* If the analyzer uses memoization, we need to clear the internal
* cache when information about a Var has been overridden.
*/
class TVM_DLL AnalyzerObj : public ffi::Object {
public:
/*! \brief sub-analyzer: const integer bound */
ConstIntBoundAnalyzer const_int_bound;
/*! \brief sub-analyzer: modular set */
ModularSetAnalyzer modular_set;
/*! \brief sub-analyzer rewrite simplify */
RewriteSimplifier rewrite_simplify;
/*! \brief sub-analyzer canonical simplify */
CanonicalSimplifier canonical_simplify;
/*! \brief sub-analyzer: int set */
IntSetAnalyzer int_set;
/*! \brief sub-analyzer transitive comparisons */
TransitiveComparisonAnalyzer transitive_comparisons;
/*! \brief sub-analyzer using Z3 */
Z3Prover z3_prover;
/*! \brief constructor */
AnalyzerObj();
/*!
* \brief Mark the value as non-negative value globally in analyzer.
*
* Only call this function if the non-neg condition is global and
* not context-dependent.
*
* This function does best-effort propagations to the sub-analyzers
*
* A canonical use of MarkGlobalNonNegValue is to record a non-negativity
* fact at a Var's definition site. Because each Var identity is defined
* exactly once in canonical IR, the fact is globally valid for that identity.
*
* \note We expose this function because non-negative global values,
* such as symbolic buffer shapes in function arguments are really
* important to ensure the best simplification, and usually they
* can be handled in a simpler way than the generic constraints.
*
* This function may call into the Update function of the sub-analyzers.
*/
void MarkGlobalNonNegValue(const PrimExpr& value);
/*!
* \brief Notify all the sub-analyzers that var
* is created and binded to expr.
*
* Each var can only be bound once.
*
* \param var The variable.
* \param expr The expression we bind to.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*!
* \brief Notify all the sub-analyzers that var
* is created and bound to a range.
*
* Each var can only be bound once.
*
* \param var The variable.
* \param range The range we bind to.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Bind all the vars in the Map
*
* \param variables The {variable -> range} map.
* \param allow_override Whether we allow overriding an existing var's
* expression. This option should not be used if there is any dependency
* between variables.
*/
void Bind(const ffi::Map<Var, Range>& variables, bool allow_override = false);
/*!
* \brief Whether can we prove expr >= val.
* Non-negative proof is very useful in integer analysis
* to lower divisions and mods given difference in trunc and ceil mode.
*
* \param expr The expression.
* \param lower_bound The lower bound.
* \return Whether we can prove it.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveGreaterEqual(const PrimExpr& expr, int64_t lower_bound);
/*!
* \brief Whether can we prove expr < val.
* Non-negative proof is very useful in integer analysis
* to lower divisions and mods given difference in trunc and ceil mode.
*
* \param expr The expression.
* \param upper_bound The upper bound.
* \return Whether we can prove it.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveLess(const PrimExpr& expr, int64_t upper_bound);
/*!
* \brief Whether can we prove lhs == rhs.
*
* \param lhs The input lhs.
* \param rhs The input rhs.
* \return Whether we can prove lhs == rhs.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
bool CanProveEqual(const PrimExpr& lhs, const PrimExpr& rhs);
/*!
* \brief Whether we can prove lhs is smaller than possibly symbolic shape.
*
* By calling this function, the caller gives an extra hint that shape > 0,
* because it appeared in buffer shape.
*
* This is useful to prove condition such as 32 <= 32 * n where the 32 * n
* is known to be a shape. Use this routine to reduce the symbolic comparisons
* in buffer compaction.
*
* The underlying analyzer will use the kSymbolicBound proof.
*
* \param lhs The input lhs.
* \param shape The symbolic shape.
* \return Whether we can prove lhs <= shape.
*/
bool CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs, const PrimExpr& shape);
/*!
* \brief Whether can we prove condition.
*
* \param cond The expression to be proved.
* \param strength the strength of the prove.
*
* \return The result.
*
* \note Analyzer will call into sub-analyzers to get the result.
* Do not use strength beyond default in sub-analyzers and
* only use it in top-level predicate analysis.
*/
bool CanProve(const PrimExpr& cond, ProofStrength strength = ProofStrength::kDefault);
/*!
* \brief Simplify expr.
*
* \param expr The expression to be simplified.
* \param steps The simplification runs in the order of
* rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
* rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
* param steps controls how many steps to run.
* Default is 2, i.e., rewrite_simplify + canonical_simplify.
* \return The result.
*
* \note Analyzer will call into sub-analyzers to get the result.
*/
PrimExpr Simplify(const PrimExpr& expr, int steps = 2);
/*!
* \brief Deep-copy this analyzer into a new, independent Analyzer.
*
* The returned analyzer carries the same accumulated facts (variable
* bounds, modular sets, rewrite/canonical bindings, integer-set domains,
* literal constraints and transitive comparisons) as this one, but owns
* its own state: binding or simplifying on either analyzer afterwards does
* not affect the other. This is the deep copy that handle-copying an
* Analyzer does not provide.
*
* \note Do not call this while a `With<ConstraintContext>` scope is active
* on this analyzer. The clone would inherit the scoped constraints
* but not the recovery functions that pop them on scope exit, so the
* constraints would leak as if they were global facts. Clone at a
* point where no constraint scope is in effect.
*
* \return A new Analyzer holding an independent copy of the facts.
*/
Analyzer Clone() const;
/*!
* \brief Analyzer methods update facts, constraints, caches, and stats.
*
* Marking the object mutable makes the `Analyzer` ObjectRef expose a
* non-const `operator->`, so APIs can take `const Analyzer&` while still
* allowing calls such as `analyzer->Bind(...)`.
* `const Analyzer&` keeps the handle itself from being rebound; it does
* not make the underlying AnalyzerObj immutable.
*/
static constexpr bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.Analyzer", AnalyzerObj, ffi::Object);
};
/*!
* \brief Managed reference to AnalyzerObj.
*
* Analyzer is a lightweight, reference-counted handle around a heap-allocated
* AnalyzerObj. Because it is now a first-class FFI object, an Analyzer can be
* passed across the tvm-ffi boundary (e.g. handed from Python into a C++ pass)
* and shared, so that accumulated bindings/constraints persist across calls.
* Copying an Analyzer copies the handle, and both handles share the same
* mutable AnalyzerObj state.
* This is not a deep copy of analyzer facts or caches.
*
* \sa AnalyzerObj
*/
class Analyzer : public ffi::ObjectRef {
public:
/*! \brief Default-construct a fresh analyzer (allocates an AnalyzerObj). */
Analyzer() : Analyzer(ffi::make_object<AnalyzerObj>()) {}
explicit Analyzer(ffi::ObjectPtr<AnalyzerObj> n) : ffi::ObjectRef(std::move(n)) {
TVM_FFI_ICHECK(this->get() != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Analyzer, ffi::ObjectRef, AnalyzerObj);
};
/*!
* \brief Constraint context.
*
* \code
*
* Var x("x");
* arith::Analyzer analyzer;
* {
* With<arith::ConstraintContext> scope(analyzer, tvm::floormod(x, 3) == 0);
* TVM_FFI_ICHECK_EQ(analyzer->modular_set(x)->coeff, 3);
* }
* // constraint no longer in effect.
* TVM_FFI_ICHECK_NE(analyzer->modular_set(x)->coeff, 3);
*
* \endcode
*/
class ConstraintContext {
private:
// declare friend to enable with.
friend class With<ConstraintContext>;
/*!
* \brief Construct a constraint context.
* \param analyzer The analyzer whose context is updated. The context
* keeps a reference to the analyzer while the scope is active.
* \param constraint The constraint to be applied.
*/
ConstraintContext(const Analyzer& analyzer, PrimExpr constraint)
: ConstraintContext(analyzer, std::move(constraint), false) {}
/*!
* \brief Construct a constraint context.
* \param analyzer The analyzer whose context is updated. The context
* keeps a reference to the analyzer while the scope is active.
* \param constraint The constraint to be applied.
* \param is_assume Whether the constraint comes from an assumption.
*/
ConstraintContext(const Analyzer& analyzer, PrimExpr constraint, bool is_assume)
: analyzer_(analyzer), constraint_(std::move(constraint)), is_assume_(is_assume) {}
/*!
* \brief Construct a constraint context from a borrowed analyzer object.
* \param analyzer The borrowed analyzer object.
* \param constraint The constraint to be applied.
*
* This overload is for internal callers that already operate on AnalyzerObj*.
*/
ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint)
: ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), false) {}
/*!
* \brief Construct a constraint context from a borrowed analyzer object.
* \param analyzer The borrowed analyzer object.
* \param constraint The constraint to be applied.
* \param is_assume Whether the constraint comes from an assumption.
*/
ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint, bool is_assume)
: ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), is_assume) {}
// enter the scope.
void EnterWithScope();
// exit the scope.
void ExitWithScope();
/*! \brief Analyzer kept alive while the context is active. */
Analyzer analyzer_;
/*! \brief The constraint */
PrimExpr constraint_;
/*! \brief functions to be called in recovery */
std::vector<std::function<void()>> recovery_functions_;
/*! \brief Whether the constraint comes from an assumption. */
bool is_assume_;
};
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_ANALYZER_H_
include/tvm/arith/bound.h
+85
View File
@@ -0,0 +1,85 @@
/*
* 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/bound.h
* \brief Bound deducers.
*/
#ifndef TVM_ARITH_BOUND_H_
#define TVM_ARITH_BOUND_H_
#include <tvm/arith/int_set.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt.h>
#include <unordered_map>
namespace tvm {
namespace arith {
using tirx::Region;
using tirx::Stmt;
using tirx::Var;
using tirx::VarNode;
/*!
* \brief Deduce the bound of the target variable in a expression,
* give the domain of each variables. Return undefined IntSet to
* represent failure.
*
* \note The returned set may be smaller than set that
* contains all possible values of v that satisfies the bound.
*
* \param v The target variable to be deduced.
* \param cond The conditional expression.
* \param hint_map The domain of variable, used to help deduce.
* \param relax_map The domain of each variable, used to relax the domain,
* The deduce bound must implies e for all value in relax_map
* \return An integer set that always satisfies the condition.
*/
IntSet DeduceBound(PrimExpr v, PrimExpr cond, const ffi::Map<Var, IntSet>& hint_map,
const ffi::Map<Var, IntSet>& relax_map);
/*!
* \brief Same as DeduceBound with unordered_map signature.
*
* \param v The target variable to be deduced.
* \param cond The conditional expression.
* \param hint_map The domain of variable, used to help deduce.
* \param relax_map The domain of each variable, used to relax the domain,
* The deduce bound mush implies e for all value in relax_map
* \return An integer set that always satisfies the condition.
*/
IntSet DeduceBound(PrimExpr v, PrimExpr cond,
const std::unordered_map<const VarNode*, IntSet>& hint_map,
const std::unordered_map<const VarNode*, IntSet>& relax_map);
/*!
* \brief Infer a regular domain that covers all the calls or provides within the given statement.
* \param body The given statement.
* \param buffer The buffer to check the access info.
* \param consider_loads If loads are considered.
* \param consider_stores If stores are considered.
* \return The domain that covers all the calls or provides within the given statement.
*/
Region DomainTouched(const Stmt& body, const tirx::Buffer& buffer, bool consider_loads,
bool consider_stores);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_BOUND_H_
include/tvm/arith/int_set.h
+339
View File
@@ -0,0 +1,339 @@
/*
* 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_set.h
* \brief Integer set
*/
#ifndef TVM_ARITH_INT_SET_H_
#define TVM_ARITH_INT_SET_H_
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
#include <unordered_map>
namespace tvm {
namespace arith {
using tirx::IterVar;
using tirx::Var;
using tirx::VarNode;
class AnalyzerObj;
class Analyzer;
//-----------------------------------------------
// Integer set data structure.
//
// This is a API build on top of the base
// integer analysis API to provide set analysis.
//------------------------------------------------
/*!
* \brief Sign type of an integer expression.
*/
enum SignType { kPositive, kNegative, kZero, kUnknown };
/*!
* \brief Base class of all Integer set containers.
* represent a set of integers in one dimension.
* \sa IntSet
*/
class IntSetNode : public ffi::Object {
public:
TVM_FFI_DECLARE_OBJECT_INFO("ir.IntSet", IntSetNode, ffi::Object);
};
/*!
* \brief Managed reference to IntSetNode.
* \sa IntSetNode
*/
class IntSet : public ffi::ObjectRef {
public:
/*!
* \brief Find a range that covers the region.
* \param max_range The range to be covered.
* \return The covering range.
*/
Range CoverRange(Range max_range) const;
/*! \return Lower bound of the set */
PrimExpr min() const;
/*! \return upper bound of the set */
PrimExpr max() const;
/*! \return The sign of the elements in the integer set */
SignType GetSignType() const;
/*! \return Whether the set represent nothing */
bool IsNothing() const;
/*! \return Whether the set represent everything */
bool IsEverything() const;
/*! \return Whether the set is a single point */
bool IsSinglePoint() const;
/*!
* \brief Check if we can prove it is a single point.
*
* Unlike IsSinglePoint, which only checks ptr equality
* this function will invoke analyzer to do stonger proofs
* but also takes longer time.
*
* Use this function in some of the primitives but do not
* use it in the inner loop of simplification.
*
* \param ana Analyzer used in the proof.
* \return Whether we can prove it is a single point
*/
bool CanProveSinglePoint(const Analyzer& ana) const;
// TODO(tvm-team): update all CanProve to explicitly take
// analyzer to encourage more analyzer reuse
/*! \return Whether the set is proved to be bigger than 0 */
bool CanProvePositive() const;
/*! \return Whether the set is proved to be smaller than 0 */
bool CanProveNegative() const;
/*! \return Whether the set is proved to be smaller than or equal to 0 */
bool CanProveNonPositive() const;
/*! \return Whether the set is proved to be larger than or equal to 0 */
bool CanProveNonNegative() const;
/*! \return Whether the set has upper bound. */
bool HasUpperBound() const;
/*! \return Whether the set has lower bound. */
bool HasLowerBound() const;
/*!
* \brief The single point value, call only if IsSinglePoint is true
* \return The point value.
*/
PrimExpr PointValue() const;
/*!
* \brief Try to match IntSet with range r.
*
* \note It is guanrateed that IntSet::FromRange(r).MatchRange(r) == true
* \return true if we can prove they are the same.
*/
bool MatchRange(const tvm::Range& r) const;
/*! \return The set contains nothing */
static IntSet Nothing();
/*! \return The set contains everything */
static IntSet Everything();
/*!
* \brief construct a point set.
* \param point The point in the set.
* \return construct a single point set
*/
static IntSet SinglePoint(PrimExpr point);
/*!
* \brief construct a integer set from vector expression.
* \param vec The vector expression, can also be single point.
* \return The result set containing the indices in the vector.
*/
static IntSet Vector(PrimExpr vec);
/*!
* \brief Construct a set representing a range [min, min + extent).
* \param min The minimum of the range range
* \param extent The extent of the range.
* \return The constructed set.
*/
static IntSet FromMinExtent(PrimExpr min, PrimExpr extent);
/*!
* \brief Construct a set representing a range.
* \param r The range
* \return The constructed set.
*/
static IntSet FromRange(tvm::Range r);
/*!
* \brief Construct a set representing a interval.
* \param min The minimum value of the interval.
* \param max The maximum value of the interval.
* \return The constructed set.
*/
static IntSet Interval(PrimExpr min, PrimExpr max);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntSet, ffi::ObjectRef, IntSetNode);
};
//-----------------------------------------------
// Integer set legacy API.
//------------------------------------------------
/*!
* \brief Convert std::unordered_map<const VarNode*, IntSet> to ffi::Map<Var, IntSet>
*
* \param dom_map The domain map to convert.
* \return The converted map.
*/
ffi::Map<Var, IntSet> ConvertDomMap(const std::unordered_map<const VarNode*, IntSet>& dom_map);
/*!
* \brief Find an symbolic integer set that contains all possible values of
* e given the domain of each iteration variables.
*
* \param e The expression to be evaluated.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values of e.
*/
IntSet EvalSet(PrimExpr e, const ffi::Map<IterVar, IntSet>& dom_map);
/*!
* \brief Find an symbolic integer set that contains all possible values of
* e given the domain of each variables.
*
* \param e The expression to be evaluated.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values of e.
*/
IntSet EvalSet(PrimExpr e, const ffi::Map<Var, IntSet>& dom_map);
/*!
* \brief Same as EvalSet, but takes unordered_map
*
* \param e The expression to be evaluated.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values of e.
*/
IntSet EvalSet(PrimExpr e, const std::unordered_map<const tirx::VarNode*, IntSet>& dom_map);
/*!
* \brief Find an symbolic integer set that contains is union over
* all the possible conditional values in dom_map.
*
* \param r The initial range.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values.
*/
IntSet EvalSet(Range r, const ffi::Map<IterVar, IntSet>& dom_map);
/*!
* \brief Find an symbolic integer set that contains is union over
* all the possible conditional values in dom_map.
*
* \param s The initial set.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values.
*/
IntSet EvalSet(IntSet s, const std::unordered_map<const VarNode*, IntSet>& dom_map);
/*!
* \brief Same as EvalSet, but takes unordered_map
*
* \param r The range to be evaluated.
* \param dom_map The domain of each variable.
* \return An integer set that can cover all the possible values of e.
*/
IntSet EvalSet(Range r, const std::unordered_map<const VarNode*, IntSet>& dom_map);
/*!
* \brief Same as EvalSet, but takes ffi::Array<Range>
*
* \param region The range to be evaluated.
* \param dom_map The domain of each variable.
* \return An array of integer sets that can cover all the possible values.
*/
ffi::Array<IntSet> EvalSet(const ffi::Array<Range>& region, const ffi::Map<Var, IntSet>& dom_map);
/*! \brief Map from Expr to IntSet */
using ExprIntSetMap = std::unordered_map<PrimExpr, IntSet, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
/*!
* \brief Find the integer set of every sub-expression, given the
* domain of each iteration variables.
*
* \param e The expression to be evaluated.
* \param dom_map The domain of each variable.
* \return the map from the expression to its possible value.
*/
ExprIntSetMap EvalSetForEachSubExpr(PrimExpr e,
const std::unordered_map<const VarNode*, IntSet>& dom_map);
/*!
* \brief Create a union set of all sets, possibly relaxed
* \param sets The sets to be combined
* \return the set after union
*/
IntSet Union(const ffi::Array<IntSet>& sets);
/*!
* \brief The union of N-dimensional integer sets
* \param nd_int_sets A list of N-dimensional integer sets
* \return An N-dimensional integer set as the result of union
*/
ffi::Array<IntSet> UnionRegion(const ffi::Array<ffi::Array<IntSet>>& nd_int_sets);
/*!
* \brief Create a lower-bound of union set, where some of the segments may be dropped
* \param sets The sets to be combined
* \return the set after union
*/
IntSet UnionLowerBound(const ffi::Array<IntSet>& sets);
/*!
* \brief The union of N-dimensional integer sets
* \param nd_int_sets A list of N-dimensional integer sets
* \return An N-dimensional integer set as the result of union
*/
ffi::Array<IntSet> UnionRegionLowerBound(const ffi::Array<ffi::Array<IntSet>>& nd_int_sets);
/*!
* \brief Create an intersected set of all sets
* \param sets The sets to be intersected
* \return the set after intersected
*/
IntSet Intersect(const ffi::Array<IntSet>& sets);
/*!
* \brief Converts the Ranges to IntSets
* \param var_dom The ranges of variables
* \return The integer sets of the variables
*/
ffi::Map<Var, arith::IntSet> AsIntSet(const ffi::Map<Var, Range>& var_dom);
/*!
* \brief Analyze the region with affine map, given the domain of variables and their predicate.
* The result should be strict, i.e. no region is discarded or relaxed.
* \param region The region to be analyzed
* \param var_dom The ranges of the variables
* \param predicate The predicate for the affine map
* \param analyzer The analyzer used
* \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of
* analysis
*/
TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
const arith::Analyzer& analyzer);
/*!
* \brief Analyze the region with affine map, given the domain of variables and their predicate.
* Some subregion may be discarded during the lower-bound analysis.
* \param region The region to be analyzed
* \param var_dom The ranges of the variables
* \param predicate The predicate for the affine map
* \param analyzer The analyzer used
* \return std::nullopt if the detection fails, or an array of arith::IntSet as the result of
* analysis
*/
TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionLowerBound(
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
const arith::Analyzer& analyzer);
/*!
* \brief Analyze the region with affine map, given the domain of variables and their predicate
* Relaxation of the region may be used in upper-bound analysis, i.e. some extra region may be added
* to the result.
* \param region The region to be analyzed
* \param var_dom The ranges of the variables
* \param predicate The predicate for the affine map
* \param analyzer The analyzer used
* \return an array of arith::IntSet as the result of analysis
*/
TVM_DLL ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region,
const ffi::Map<Var, Range>& var_dom,
const PrimExpr& predicate,
const arith::Analyzer& analyzer);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_INT_SET_H_
include/tvm/arith/int_solver.h
+338
View File
@@ -0,0 +1,338 @@
/*
* 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/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <unordered_map>
#include <utility>
#include <vector>
#include "analyzer.h"
namespace tvm {
namespace arith {
using tirx::IterVar;
using tirx::Var;
using tirx::VarNode;
// 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 Perform substitution on all components of the struct.
*/
IntGroupBounds Substitute(const ffi::Map<Var, PrimExpr>& subst) const;
/*!
* \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<Var> 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<Var> 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<Var>& 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_
include/tvm/arith/iter_affine_map.h
+425
View File
@@ -0,0 +1,425 @@
/*
* 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/iter_affine_map.h
* \brief Iterator quasi-affine mapping patterns.
*
* This file defines a collection of mapping patterns
* maps a collection of independent iterators to another
* collection of independent iterators.
*
* There are two main kinds of mapping patterns:
*
* - Fuse: fuse a collection of iterators into a single one
*
* domain(x0) = [0, 4), domain(x1) = [0, 3), domain(x2) = [0, 2)
* fuse(x0, x1, x2): y = x2 * 12 + x1 * 4 + x0
* domain(y) = [0, 24)
*
* - Split: split an iterator into multiple ones
*
* domain(x) = [0, 24)
* split(x, 3, 12): [y0, y1, y2] = [x % 3, (x % 12) / 3, x / 12]
* domain(y0) = [0, 3), domain(y1) = [0, 4), domain(y2) = [0, 2)
*
* We use the name "(quasi)affine" to be consistent with
* the terminology used in the polyhedral compilation.
* Notably, fuse is an affine transformation,
* while split corresponds to additional floordiv/mod operations
* that can appear in quasi-affine transformations.
*/
#ifndef TVM_ARITH_ITER_AFFINE_MAP_H_
#define TVM_ARITH_ITER_AFFINE_MAP_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/var.h>
namespace tvm {
namespace arith {
/*!
* \brief Base class of all iter map expressions.
*
* An IterMapExpr is a special expression to store
* the result of IterMapDetection.
* It should not appear in a legal TIR PrimFunc.
*/
class IterMapExprNode : public ExprNode {
public:
static constexpr const uint32_t _type_child_slots = 2;
TVM_FFI_DECLARE_OBJECT_INFO("arith.IterMapExpr", IterMapExprNode, ExprNode);
};
/*!
* \brief Managed reference to IterMapExprNode.
* \sa IterMapExprNode
*/
class IterMapExpr : public PrimExpr {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMapExpr, PrimExpr, IterMapExprNode);
static constexpr bool _type_container_is_exact = true;
};
/*!
* \brief Mark the source as an iterator in [0, extent).
*
* IterMark is used to mark source expression as a valid
* iterator to make future analysis easy.
*/
class IterMarkNode : public ffi::Object {
public:
/*!
* \brief The source expression, can either be
* a IterSumExpr or a Var.
*/
PrimExpr source;
/*!
* \brief The extent of the iteration.
*/
PrimExpr extent;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterMarkNode>()
.def_ro("source", &IterMarkNode::source)
.def_ro("extent", &IterMarkNode::extent);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMark", IterMarkNode, ffi::Object);
};
/*!
* \brief Managed reference to IterMarkExprNode.
* \sa IterMarkExprNode
*/
class IterMark : public ffi::ObjectRef {
public:
/*!
* \brief constructor.
* \param source The source expression.
* \param extent The extent of the iterator.
*/
TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMark, ffi::ObjectRef, IterMarkNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode);
};
/*!
* \brief Split of an iterator.
*
* result = floormod(floordiv(source, lower_factor), extent) * scale
*/
class IterSplitExprNode : public IterMapExprNode {
public:
/*! \brief The source marked iterator. */
IterMark source;
/*! \brief The lower factor to split the source. */
PrimExpr lower_factor;
/*! \brief The extent of the split. */
PrimExpr extent;
/*! \brief Additional scale. */
PrimExpr scale;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterSplitExprNode>()
.def_ro("source", &IterSplitExprNode::source)
.def_ro("lower_factor", &IterSplitExprNode::lower_factor)
.def_ro("extent", &IterSplitExprNode::extent)
.def_ro("scale", &IterSplitExprNode::scale);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSplitExpr", IterSplitExprNode, IterMapExprNode);
};
/*!
* \brief Managed reference to IterSplitExprNode.
* \sa IterSplitExprNode
*/
class IterSplitExpr : public IterMapExpr {
public:
/*!
* \brief constructor from just source.
* \param source The source expression.
*/
TVM_DLL explicit IterSplitExpr(IterMark source);
/*!
* \brief constructor from just source.
* \param source The source expression.
* \param scale The additional scaling factor.
*/
TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
/*!
* \brief constructor
* \param source The source expression.
* \param lower_factor The lower factor to split the source.
* \param extent The extent of the split.
* \param scale The additional scaling factor.
*/
TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
PrimExpr scale);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSplitExpr, IterMapExpr, IterSplitExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode);
};
/*!
* \brief Fuse multiple iterators by summing them with scaling.
*
* result = sum(args) + base
*/
class IterSumExprNode : public IterMapExprNode {
public:
/*! \brief The args to the sum. */
ffi::Array<IterSplitExpr> args;
/*! \brief The base offset. */
PrimExpr base;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterSumExprNode>()
.def_ro("args", &IterSumExprNode::args)
.def_ro("base", &IterSumExprNode::base);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSumExpr", IterSumExprNode, IterMapExprNode);
};
/*!
* \brief Managed reference to IterSumExprNode.
* \sa IterSumExprNode
*/
class IterSumExpr : public IterMapExpr {
public:
/*!
* \brief constructor.
* \param args The args to the sum.
* \param base The base offset.
*/
TVM_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSumExpr, IterMapExpr, IterSumExprNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode);
};
} // namespace arith
namespace ffi {
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSplitExprNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSumExprNode> = true;
} // namespace ffi
namespace arith {
/*! \brief Mapping level for iterators. */
enum IterMapLevel {
// Require the mapping to be bijective.
Bijective = 0,
// Require the mapping to be surjective.
Surjective = 1,
// No mapping safety check.
NoCheck = 3
};
/*!
* \brief Result of DetectIterMap.
*/
class IterMapResultNode : public ffi::Object {
public:
// The detected pattern if a match exists.
ffi::Array<IterSumExpr> indices;
// Any errors that occurred while converting the input indices. If
// the array is empty, the conversion was successful.
ffi::Array<ffi::String> errors;
/*! \brief Boolean expression indicating if a specific value w
*
* `padding_predicate` evaluates to true for a set of indices that
* are outside the bounds of the provided index iterators, but
* inside the bounds of the returned index iterators. This
* expression is in terms of the variables provided in
* `input_iters`.
*/
PrimExpr padding_predicate;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterMapResultNode>()
.def_ro("indices", &IterMapResultNode::indices)
.def_ro("errors", &IterMapResultNode::errors)
.def_ro("padding_predicate", &IterMapResultNode::padding_predicate);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, ffi::Object);
};
/*!
* \brief Managed reference to IterMapResultNode.
* \sa IterMapResultNode
*/
class IterMapResult : public ffi::ObjectRef {
public:
// constructor
IterMapResult() { data_ = ffi::make_object<IterMapResultNode>(); }
/*! \return mutable pointers to the node. */
IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
};
/*!
* \brief Detect if indices can be written as
* [y_0 + c_0, y_1 + c_1, ..., y_n + c_n]
*
* Here y = some-quasi-affine-iter-map(input_iters)
* and c are symbolic constants.
*
* We also requires that y_i and y_j to be independent for i != j.
*
* For returned value rv, the following is always true:
* - rv[i]->args.size() <=1: only one iterator per element.
*
* \param indices The indices to detect pattern for.
* \param input_iters Map from variable to iterator's range.
* \param predicate The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with extent of
* 1 will be replaced with a constant value.
*
* \return The detected iteration result.
* The return object's .indices is empty on failure.
*/
IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
IterMapLevel check_level, const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Use IterVarMap detector to rewrite and simplify the indices
*
* \param indices The indices to detect pattern for.
* \param input_iters Map from variable to iterator's range.
* \param input_pred The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with unit extents are simplified
* \return The indices after rewrite
*/
ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
const ffi::Map<Var, Range>& input_iters,
const PrimExpr& input_pred, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Apply the inverse of the affine transformation to the outputs.
*
* Similar to the back-propagation, starting from the outputs, it visits the DAG of the expressions
* in reverse topology order and applies the inverse of the affine transformation until it reaches
* the input. The affine iter map is required to be bijective.
*
* For example, iter_map = [l0 // 16, l0 % 16], outputs = [output_0, output_1],
* the affine transformation specified by `iter_map` will be applied to `outputs` and the result
* will be {l0: ((output_0*16) + output_1)}.
*
* The range of `outputs` should be the same as the output range of the affine transmation.
*
* \sa DetectIterMap
*
* \param iter_map The bijective affine iter map.
* \param outputs The outputs of the affine transformation.
*
* \return The map from the input to the transformed result.
*/
ffi::Map<Var, PrimExpr> InverseAffineIterMap(const ffi::Array<IterSumExpr>& iter_map,
const ffi::Array<PrimExpr> outputs);
/*!
* \brief Detect if bindings can be written as
* [a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]
*
* where a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
* b = some-quasi-affine-iter-map(sub_iters)
* c is constant symbols
* e is the extent of b
*
* For example, z*12 + y*3 + x + c = (z*4+y)*3 + x, if sub_iters={x}
*
* \param bindings The input bindings
* \param input_iters Map from variable to iterator's range.
* \param sub_iters Iterators of subspace.
* \param predicate The predicate constraints on the input iterators
* \param check_level The iter mapping checking level.
* \param analyzer Analyzer used to get context information.
* \param simplify_trivial_iterators If true, iterators with extent of
* 1 will be replaced with a constant value.
*
* \return The result list has length len(bindings) + 1
[0, len(bindings)): The iter map matching result. The inner list is of length 2.
The first expr is the basis of the quotient space.
The second expr is the basis of the subspace.
len(bindings): the predicate of outer space and inner space
Empty array if no match can be found.
*/
ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
const ffi::Map<Var, Range>& input_iters,
const ffi::Array<Var>& sub_iters,
const PrimExpr& predicate, IterMapLevel check_level,
const arith::Analyzer& analyzer,
bool simplify_trivial_iterators = true);
/*!
* \brief Given an expression that may contain IterMapExpr, transform it to normal PrimExpr.
* \param expr The input expression, which may contain IterMapExpr.
* \return The corresponding normal PrimExpr.
*/
PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr);
/*!
* \brief Rewrite index as IterSumExpr
*
* ((i0 // b0) % a0) * s0 + ((i0 // b1) % a1) * s1 ... + base
*
* The iterators are ordered such that s0 > s1 ...
* if we can prove the relation.
*
* Note that base may contain expressions that cannot be detected
* as the right pattern.
*
* \param index The input index
* \param input_iters The input iterators.
* \param analyzer The input analyzer.
* \note This function is useful to detect iterator stride patterns.
*/
IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
const arith::Analyzer& analyzer);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_ITER_AFFINE_MAP_H_
include/tvm/arith/pattern.h
+54
View File
@@ -0,0 +1,54 @@
/*
* 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/pattern.h
* \brief Expression pattern detectors.
*/
#ifndef TVM_ARITH_PATTERN_H_
#define TVM_ARITH_PATTERN_H_
#include <tvm/ir/expr.h>
#include <tvm/tirx/expr.h>
namespace tvm {
namespace arith {
/*!
* \brief Detect if e can be rewritten as e = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]
* Where coeff[i] and base are invariant of var[j] for all i and j.
*
* \param e The expression to be detected.
* \param vars List of variables to be used in detection.
* \return [coeff[i]] if it is possible, empty array if it is not.
*/
ffi::Array<PrimExpr> DetectLinearEquation(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
/*!
* \brief Detect if expression corresponds to clip bound of the vars
*
* \param e The expression to be detected.
* \param vars List of variables to be used in detection.
* \return concat([min_value[i], max_value[i]]), None is returned if there is no min or max value
* return empty if the e does not match the pattern.
*/
ffi::Array<PrimExpr> DetectClipBound(const PrimExpr& e, const ffi::Array<tirx::Var>& vars);
} // namespace arith
} // namespace tvm
#endif // TVM_ARITH_PATTERN_H_