analysis.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef TVM_RELAX_ANALYSIS_H_
#define TVM_RELAX_ANALYSIS_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/s_tir/stmt.h>
#include <tvm/sym/analyzer.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
//----------------------------------
TVM_DLL bool CanProveShapeEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs,
                                const sym::Analyzer& ana);
 
TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, const sym::Analyzer& ana);
 
//-----------------------------------
// Foundational Type analysis
//-----------------------------------
TVM_DLL Type GetStaticType(const Type& info);
 
TVM_DLL Type TypeFromStaticType(const Type& type);
 
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx);
TVM_DLL Type DeriveCallRetType(const FuncType& finfo, const Call& call, const BlockBuilder& ctx,
                               const sym::Analyzer& ana);
 
TVM_DLL Type EraseToWellDefined(
    const Type& info, std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr);
TVM_DLL Type EraseToWellDefined(const Type& info,
                                std::function<ffi::Optional<Expr>(const Var& var)> f_var_map,
                                const sym::Analyzer& ana);
 
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map);
TVM_DLL Type EraseToWellDefined(const Type& info, ffi::Map<Var, Expr> var_map,
                                const sym::Analyzer& ana);
 
enum class BaseCheckResult {
  kFailL0 = 0,
  kFailL1 = 1,
  kFailL2 = 2,
  kPass = 3
};
 
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived);
TVM_DLL BaseCheckResult TypeBaseCheck(const Type& base, const Type& derived,
                                      const sym::Analyzer& ana);
 
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived);
TVM_DLL bool IsBaseOf(const Type& base, const Type& derived, const sym::Analyzer& ana);
 
TVM_DLL PrimExpr TypeBaseCheckPrecondition(const Type& base, const Type& derived);
 
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs);
TVM_DLL Type TypeLCA(const Type& lhs, const Type& rhs, const sym::Analyzer& ana);
 
TVM_DLL ffi::Array<tvm::Var> TIRVarsInType(const Type& ty);
 
TVM_DLL ffi::Array<tvm::Var> DefinableTIRVarsInType(const Type& ty);
 
TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const Type& ty);
 
TVM_DLL ffi::Array<tvm::Var> DefinedSymbolicVars(const Expr& expr);
 
TVM_DLL ffi::Array<tvm::Var> FreeSymbolicVars(const Expr& expr);
//-----------------------------------
// General IR analysis
//-----------------------------------
TVM_DLL tvm::ffi::Array<Var> BoundVars(const Expr& expr);
 
TVM_DLL tvm::ffi::Array<Var> FreeVars(const Expr& expr);
 
TVM_DLL tvm::ffi::Array<Var> AllVars(const Expr& expr);
 
TVM_DLL tvm::ffi::Array<GlobalVar> AllGlobalVars(const Expr& expr);
 
TVM_DLL tvm::ffi::Array<tvm::ffi::Array<GlobalVar>> DetectRecursion(const IRModule& m);
 
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const IRModule& m);
 
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const Expr& expr);
 
TVM_DLL ffi::Map<Var, Expr> AnalyzeVar2Value(const DataflowBlock& dfb);
 
TVM_DLL ffi::Map<ffi::String, ffi::Array<Binding>> NameToBinding(const Function& fn);
 
TVM_DLL ffi::Map<Var, ffi::Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb);
 
std::pair<ffi::Map<Var, ffi::Array<Var>>, ffi::Array<Var>> FunctionUseDef(const Expr& expr);
 
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;
};
 
VarUsageInfo CollectVarUsage(const Expr& expr);
 
TVM_DLL std::set<const VarNode*> GetUsedVars(const Expr& expr);
 
TVM_DLL Expr RemoveAllUnused(Expr expr);
 
TVM_DLL OpPatternKind AnalyzeOpPatternKind(const tirx::Function& func);
 
TVM_DLL bool HasReshapePattern(const tirx::Function& func);
 
TVM_DLL ffi::Optional<Expr> FindImpureCall(
    const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
 
TVM_DLL bool ContainsImpureCall(
    const Expr& expr, const ffi::Optional<Expr>& own_name = ffi::Optional<Expr>(std::nullopt));
 
TVM_DLL void WellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
 
TVM_DLL bool CheckWellFormed(ffi::Variant<IRModule, Function> obj, bool check_ty = true);
 
TVM_DLL ffi::Map<s_tir::SBlock, ffi::Map<ffi::ObjectRef, tirx::IndexMap>> SuggestLayoutTransforms(
    const Function& fn, ffi::Array<tirx::IndexMap> write_tensor_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_