Files
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
CI / MacOS (push) Waiting to run
CI / Windows (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

686 lines
26 KiB
C++

/*
* 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/relax/analysis.h
* \brief The set of Relax specific analysis on IR.
*/
#ifndef TVM_RELAX_ANALYSIS_H_
#define TVM_RELAX_ANALYSIS_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/type.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/index_map.h>
#include <functional>
#include <set>
#include <utility>
namespace tvm {
namespace relax {
//-----------------------------------
// Shape expression analysis
//----------------------------------
/*!
* \brief Can prove the two symbolic shape arrays equals to each other.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana The analyzer used for integer analysis.
* \return The prove result.
*
* \note This function does best effort prove, which means
* if result is false, there is still possibility that
* two shapes equals to each other during runtime.
*/
TVM_DLL bool CanProveShapeEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs,
const arith::Analyzer& ana);
/*!
* \brief Can prove the two symbolic shape expressions equals to each other.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana The analyzer used for integer analysis.
*
* \note This function does best effort prove, which means
* if result is false, there is still possibility that
* two shapes equals to each other during runtime.
*/
TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const arith::Analyzer& ana);
//-----------------------------------
// Foundational Type analysis
//-----------------------------------
/*!
* \brief Get the corresponding static type from a given type.
* \param info The type.
* \return the corresponding static type.
*/
TVM_DLL Type GetStaticType(const Type& info);
/*!
* \brief Get the corresponding type from static type.
* \param type The input type
* \return the corresponding type.
*/
TVM_DLL Type TypeFromStaticType(const Type& type);
/*!
* \return Derive the call's ret value type from inputs.
* \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \return The derived type of the call.
* \note call->op field is ignored during derivation and we only rely on information
* presented by func_ty.
*/
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx);
/*!
* \brief Derive the call's ret value type using a caller-provided analyzer.
* \param finfo The function type.
* \param call The call expression to be derived.
* \param ctx The builder context.
* \param ana Context analyzer to prove symbolic expression equality.
* \return The derived type of the call.
*/
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx,
const arith::Analyzer& ana);
/*!
* \brief Erase the info to a corresponding more coarse grained
* type that is still well-defined(with all the vars in scope).
*
* When we are returning a Type to another scope,
* it is important to remember that Type may carry
* dependencies on var that is not defined the other scope.
*
* In such cases, it is important to call EraseToWellDefined to get
* another Type that **only** contains the vars that are defined
* in the target scope.
*
* For example, consider the following function
*
* \code
*
* @R.function
* def f(x: R.Tensor[(n, m)]):
* k = tirx.Var("k", "int64")
* v0 = opaque_fn(x)
* v1 = match_cast(v0, R.Tensor[(n, k)])
* v2 : R.Tensor[(n + 1, k + 2)] = pad(v1)
* return v2
*
* \endcode
*
* In the above code, the return value y have shape `(n + 1, k + 2)`,
* However, at the level of function signature, only n, m are defined,
* k is undefined here.
*
* When we call EraseToWellDefined(R.Tensor[(n + 1, k + 2)], fshape_var_map={n: n, m: m}),
* we will obtain R.Tensor(ndim=2), which is an erased info that does not depend
* on k(which is undefined from parameter signature).
*
* However, if we call EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: n, m: m}),
* Then the return value will be R.Tensor[(n + 1, m)], because both n and m are defined.
*
* We can also make these var map to return a different expression.
* For example, EraseToWellDefined(R.Tensor[(n + 1, m)], fshape_var_map={n: 2, m: m})
* will give us R.Tensor[(3, m)], where n get replaced by 2.
*
* Use this function in the following scenarios:
* - Decide the ty of expr with sub-scopes, such as If, SeqExpr
* - Decide the deduced return ty of a function that can be fully decided by params.
*
* \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param f_var_map callback function to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(
const Type& info,
std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map = nullptr,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr);
/*!
* \brief EraseToWellDefined overload using a caller-provided analyzer.
* \param info The type.
* \param f_shape_var_map callback function to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param f_var_map callback function to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(
const Type& info, std::function<ffi::Optional<PrimExpr>(const tirx::Var& var)> f_shape_var_map,
std::function<ffi::Optional<Expr>(const Var& var)> f_var_map, const arith::Analyzer& ana);
/*!
* \brief EraseToWellDefined variant with map.
* \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param var_map map to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
*
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
ffi::Map<Var, Expr> var_map);
/*!
* \brief EraseToWellDefined map overload using a caller-provided analyzer.
* \param info The type.
* \param shape_var_map map to specify
* whether a symbolic shape var is defined and the value it maps to,
* return nullopt if var is undefined.
* \param var_map map to specify
* whether a var is defined in the target scope and the value it maps to,
* return nullopt if var is undefined.
* \param ana Context analyzer to prove symbolic expression equality.
* \return the corresponding erased type.
*/
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<tirx::Var, PrimExpr> shape_var_map,
ffi::Map<Var, Expr> var_map, const arith::Analyzer& ana);
/*!
* \brief Fine grained result of base check.
*
* This analysis comes with different levels of checking failures
* that can help to customize the compilation decisions.
*
* For a given pair of lhs_ty, rhs_ty. We adopt
* the following terminology:
* - LSet = {value | value matches lhs_ty}
* - RSet = {value | value matches rhs_ty}
*
* See the definition of each level below.
*/
enum class BaseCheckResult {
/*!
* \brief The two value sets have no intersection at all: Interset(LSet, RSet) = empty
*/
kFailL0 = 0,
/*!
* \brief LSet is not superset of RSet by only looking at static information.
*
* \note This level will trigger static type checking error when lhs is param and rhs is arg.
*/
kFailL1 = 1,
/*!
* \brief WLSet is not superset of RSet because of mismatch in value information.
*
* L1-level mismatches in params of FuncType is categorized as
* If lhs is FuncType, then L1-level mismatch in its params
* is categorized as L2-level mismatch for lhs.
*
* Design considerations for functions:
* - (a) We want to be able to erase type/value in function signature
* when we unify function type and preserve simpler representations.
* - (b) We automatically insert match_cast at function boundary, so
* we can erase (int)->int argument as (object)->int.
* The input shape/type mismatch will be detected by runtime checks at function boundary.
* This behavior is also consistent with the ffi::Function behavior.
*
* \note This level means there is no problem about static known information.
* It is OK for the checker to do best effort and return this value.
*/
kFailL2 = 2,
/*! \brief LSet is superset of RSet. */
kPass = 3
};
/*!
* \brief Run a base check to see if base subsumes derived.
*
* This function returns fine-grained base-check result on reasons of failure.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived);
/*!
* \brief Run a base check using a caller-provided analyzer.
* \param base The base type.
* \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*
* \sa BaseCheckResult
*/
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived,
const arith::Analyzer& ana);
/*!
* \brief Check the relation of two type to see if one subsumes another one.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether the relation holds.
*/
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived);
/*!
* \brief Check whether one type subsumes another using a caller-provided analyzer.
* \param base The base type.
* \param derived The derived type.
* \param ana Context analyzer to prove symbolic expression equality.
* \return Whether the relation holds.
*/
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const arith::Analyzer& ana);
/*!
* \brief Return the condition for which base is a superset of derived
*
* This function returns finer-grained conditions for kFailL2 than TypeBaseCheck
*
* If the returned expression is true, or simplifies to true, then
* base is a superset of derived. If the returned expression is
* false, or simplifies to false, then base is not a superset of
* derived.
*
* If the returned expression is neither true nor false, it is an
* expression in terms of the symbolic variables available in `base`
* and `derived`.
*
* \param base The base type.
* \param derived The derived type.
* \return Whether base is a base of derived.
*
* \sa BaseCheckResult
*/
TVM_DLL PrimExpr TypeBaseCheckPrecondition(const Type& base, const Type& derived);
/*!
* \brief Unify the two type to their least common ancestor.
*
* \param lhs The left operand.
* \param rhs The right operand.
* \return The unified information.
*/
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs);
/*!
* \brief Unify two types using a caller-provided analyzer.
* \param lhs The left operand.
* \param rhs The right operand.
* \param ana Context analyzer to prove symbolic expression equality.
* \return The unified information.
*/
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const arith::Analyzer& ana);
/*!
* \brief Get the TIR variables that appear in the input type.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param ty The type object to be analyzed.
* \return The list of TIR variables that appear in the input type.
*/
TVM_DLL ffi::Array<tirx::Var> TIRVarsInType(const Type& ty);
/*!
* \brief Get the TIR variables that appear in the input type.
*
* Returns all symbolic variables that are definable based on, and
* used within, the Type.
*
* \param ty The type object to be analyzed.
*
* \return A tuple of (definable,used) TIR variables. Both lists are
* deduplicated, each TIR variable will appear at most once, and in
* order of occurrence.
*/
TVM_DLL ffi::Array<tirx::Var> DefinableTIRVarsInType(const Type& ty);
/*! \brief Collect expressions whose usage requires them to be non-negative
*
* Any PrimExpr that is used as a tensor shape, or as an element in a
* ShapeExpr, may not be negative. This utility function can be used
* to generate assertions prior to calling a kernel, or to provide
* assumptions within a kernel that may be useful for simplification.
*
* \param ty The type to be analyzed
*
* \return A list of non-negative expressions.
*/
TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const Type& ty);
/*!
* \brief Get the TIR variables that defined in the input function.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param expr The relax expression (e.g. a Function) to be analyzed.
* \return The list of TIR variables that are defined in the input function.
*/
TVM_DLL ffi::Array<tirx::Var> DefinedSymbolicVars(const Expr& expr);
/*!
* \brief Get the TIR variables that are used but not defined in the input function.
* The returned list is deduplicated - each TIR variable will appear at most once.
* \param expr The relax expression (e.g. a Function) to be analyzed.
* \return The list of TIR variables that are used but not defined in the input function.
*/
TVM_DLL ffi::Array<tirx::Var> FreeSymbolicVars(const Expr& expr);
//-----------------------------------
// General IR analysis
//-----------------------------------
/*!
* \brief Get all bound variables from expression expr.
*
* Bound variables are all variables that are declared in the expr.
* They only have meaning inside that expr, and can only be used in it.
*
* \param expr the expression.
*
* \return List of bound vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> BoundVars(const Expr& expr);
/*!
* \brief Get free type parameters from expression expr.
*
* Free variables are variables that are not bound by a
* varbinding or a function parameter in the context.
*
* \param expr the expression.
*
* \return List of free vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> FreeVars(const Expr& expr);
/*!
* \brief Get all variables from expression expr.
*
* \param expr the expression.
*
* \return List of all vars, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<Var> AllVars(const Expr& expr);
/*!
* \brief Get all global variables from expression expr.
*
* AllVars is a superset of BoundVars and FreeVars.
* The union of BoundVars and FreeVars is Allvars.
*
* \param expr the expression.
*
* \return List of all global variables, in the PostDFS order in the expression.
*/
TVM_DLL tvm::ffi::Array<GlobalVar> AllGlobalVars(const Expr& expr);
/*!
* \brief Find all sets of recursive or mutually recursive functions in the module.
*
* Two or more functions are mutually recursive if there is some cycle of references
* among them. For example, if there are two functions A and B, they are
* mutually recursive if A calls B and B calls A. Another case would be with
* three functions A, B, and C, where A calls B, B calls C, and C calls A.
*
* (Note that functions do not have to call each other to reference each other.
* For example, if a function returns another function, that is still a reference
* that could potentially be recursive, even without a call.)
*
* If a function is simply recursive and not mutually recursive with any other,
* it will be reported as a group by itself.
*
* \param m The module
*
* \return List of all groups of mutually recursive functions.
* Each member of the result is a list of functions in the module
* that are all mutually recursive.
* If a function is simply recursive and not mutually recursive with any other,
* then it will be listed as a group by itself.
*/
TVM_DLL tvm::ffi::Array<tvm::ffi::Array<GlobalVar>> DetectRecursion(const IRModule& m);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param m The IRModule to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const IRModule& m);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param expr The expression to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const Expr& expr);
/*!
* \brief Analyze var -> value mapping from VarBindings.
*
* \param dfb The dataflow block to check.
* \return Var -> Value (Expr)
*/
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const DataflowBlock& dfb);
/*!
* \brief Return a mapping from variable name to its Bindings.
*
* \param fn The function to be analyzed.
* \return A mapping from variable name to its Bindings.
*/
TVM_DLL ffi::Map<ffi::String, ffi::Array<Binding>> NameToBinding(const Function& fn);
/*!
* \brief Get the use-def chain of variables inside a dataflow block.
*
* \param dfb The dataflow block to be analyzed.
* \return A map mapping variable definitions to a set of uses.
*/
TVM_DLL ffi::Map<Var, ffi::Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb);
/*!
* \brief Get the use-def chain of variables inside a function.
*
* \param expr The expression to be analyzed.
*
* \return A tuple of variable usage and variable outputs. The first
* element is a map from variable definitions to the set of downstream
* users of that definition. The second element is a list of
* variables whose usage occurs outside of any variable binding,
* typically the output body of a relax::Function or a relax::SeqExpr.
*/
std::pair<ffi::Map<Var, ffi::Array<Var>>, ffi::Array<Var>> FunctionUseDef(const Expr& expr);
/*! \brief A utility struct returned by CollectVarUsage
*/
struct VarUsageInfo {
/* \brief A map from variables to the bound expression.
*
* This is equivalent to the output of AnalyzeVar2Value
*/
ffi::Map<Var, Expr> bound_values;
/* \brief The map from variables to downstream usages of the variable
*
* This is equivalent to the first output of FunctionUseDef.
*/
ffi::Map<Var, ffi::Array<Var>> downstream_usage;
/* \brief A list of variables produced as output
*
* This is equivalent to the second output of FunctionUseDef
*/
ffi::Array<Var> outputs;
};
/*! \brief Collect variable bindings and usage
*
* This function is equivalent to calling both FunctionUseDef and
* AnalyzeVar2Value, but requires only a single traversal of the
* expression.
*
* \param expr The expression to analyze
*
* \return The collected information
*/
VarUsageInfo CollectVarUsage(const Expr& expr);
/*!
* \brief Get the used variables in an expression.
*
* This function collects all variables that are referenced within the given expression.
*
* \param expr The expression to analyze
*
* \return A set of variable nodes that are used in the expression
*/
TVM_DLL std::set<const VarNode*> GetUsedVars(const Expr& expr);
/*!
* \brief Remove unused statements inside DataflowBlocks.
*
* \param expr The expression (typically a relax::Function) from which
* to remove unused statements.
*
* \return The updated function with no unused statements in DataflowBlock.
*/
TVM_DLL Expr RemoveAllUnused(Expr expr);
/*!
* \brief Annotate Op Pattern Kind for PrimFunc, which is used in relax FuseOps.
*
* \param func The PrimFunc to be analyzed.
* \return The Op Pattern Kind.
*
* \note This analysis applies on TIR function but is primarily used by relax passes.
* As a result we place it under the relax namespace.
*/
TVM_DLL OpPatternKind AnalyzeOpPatternKind(const tirx::PrimFunc& func);
/*!
* \brief Check if the given PrimFunc is essentially doing a reshape operation.
* The reshape operation also includes expand_dims, squeeze, flatten, etc.
* \details Here the allowed reshape pattern is: for example, assume the operation is
* `B[l_0, l_1, ..., l_b] = A[r_0, r_1, ..., r_a]`, we check if we can prove that the flattened
* index of l_0, ..., l_b under buffer B equals to the flattened index of r_0, ..., r_a under
* buffer A.
* \param func The function to be examined.
* \return A boolean indicating if the given PrimFunc is doing a reshape.
* \note According to the description above, the returned result can only be false-negative and
* cannot be false-positive, since whenever we cannot prove the equality, we return false. This
* property guarantees the safety of this function.
*/
TVM_DLL bool HasReshapePattern(const tirx::PrimFunc& func);
/*!
* \brief Check if the given expression (likely a function body) contains any impure calls.
* \param expr The expression to be examined. If expr is a function, we check the body.
* \param own_name (Optional.) If we are checking a recursive function body,
* the caller can pass the function's name so recursive calls
* can be ignored in the check (must be a Var or GlobalVar).
* \return The impure expression, if one exists within the given
* expression. Otherwise, std::nullopt.
* \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
TVM_DLL ffi::Optional<Expr> FindImpureCall(
const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
/*!
* \brief Check if the given expression (likely a function body) contains any impure calls.
* \param expr The expression to be examined. If expr is a function, we check the body.
* \param own_name (Optional.) If we are checking a recursive function body,
* the caller can pass the function's name so recursive calls
* can be ignored in the check (must be a Var or GlobalVar).
* \return A boolean indicating if the expression contains any impure calls.
* \note Relies on Type annotations, so ensure that the module has been normalized first.
* Also, an impure call in a *nested* function does *not* mean that the outer expression contains
* an impure call--it only does if the nested function is *later called*.
*/
TVM_DLL bool ContainsImpureCall(
const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
/*!
* \brief Check if an IRModule or Function is well-formed.
*
* Throws an ffi::Error on the first well-formedness violation. The error is
* seeded with the offending node so a pass runner can resolve a precise access
* path. Use \ref CheckWellFormed for a boolean answer.
*
* \param obj The IRModule or relax::Function to check.
* \param check_ty If true, verify that every Expr has ty populated.
* \note By default the type information is always checked. It is only in test cases
* where `check_ty` might be false, so that other well-formed requirements
* will be well tested and will not be blocked by not having type information.
*/
TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Return whether an IRModule or Function is well-formed.
*
* Wraps \ref WellFormed, returning false instead of throwing on the first
* violation.
*
* \param obj The IRModule or relax::Function to check.
* \param check_ty If true, verify that every Expr has ty populated.
* \return true if the object is well-formed, false otherwise.
*/
TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
/*!
* \brief Using the layout transforms on the outputs, suggest layout transformation on the blocks
* and buffers for the PrimFunc.
*
* \param fn The PrimFunc to be analyzed.
* \param write_buffer_transformations Array of IndexMap transformations on PrimFunc outputs.
* \return Suggested transforms per block in `fn`. For each block the returned value is a map
* from the object (block or buffer) to it's index map transformation.
*/
TVM_DLL ffi::Map<tirx::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
const Function& fn, ffi::Array<tirx::IndexMap> write_buffer_transformations);
/* \brief Collect variables whose value can be computed at compile-time
*
* If a function has the `kNumInput` attribute, then the first
* `kNumInput` parameters are provided at run-time, while all
* remaining parameters may be known at compile-time. This utility
* collects all variable bindings that only depend, directly or
* indirectly, on the parameters known at compile-time.
*
* \param func The relax::Function to analyze
*
* \return The set of variables that can be computed at compile-time,
* in order of their occurrence within the function.
*/
TVM_DLL ffi::Array<Var> ComputableAtCompileTime(const Function& func);
} // namespace relax
} // namespace tvm
#endif // TVM_RELAX_ANALYSIS_H_