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/arith/analyzer.h>
#include <tvm/ir/diagnostic.h>
#include <tvm/ir/module.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/struct_info.h>
#include <tvm/tir/function.h>
#include <tvm/tir/index_map.h>
 
#include <functional>
#include <utility>
 
namespace tvm {
namespace relax {
//-----------------------------------
// Shape expression analysis
//----------------------------------
TVM_DLL bool CanProveShapeEqual(const ffi::Array<PrimExpr>& lhs, const ffi::Array<PrimExpr>& rhs,
                                arith::Analyzer* ana);
 
TVM_DLL bool CanProveShapeEqual(const Expr& lhs, const Expr& rhs, arith::Analyzer* ana);
 
//-----------------------------------
// Foundational StructInfo analysis
//-----------------------------------
TVM_DLL Type GetStaticType(const StructInfo& info);
 
TVM_DLL StructInfo StructInfoFromType(const Type& type);
 
TVM_DLL StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Call& call,
                                           const BlockBuilder& ctx, arith::Analyzer* ana = nullptr);
 
TVM_DLL StructInfo EraseToWellDefined(
    const StructInfo& info,
    std::function<ffi::Optional<PrimExpr>(const tir::Var& var)> f_shape_var_map = nullptr,
    std::function<ffi::Optional<Expr>(const Var& var)> f_var_map = nullptr,
    arith::Analyzer* ana = nullptr);
 
TVM_DLL StructInfo EraseToWellDefined(const StructInfo& info,
                                      ffi::Map<tir::Var, PrimExpr> shape_var_map,
                                      ffi::Map<Var, Expr> var_map, arith::Analyzer* ana = nullptr);
 
enum class BaseCheckResult {
  kFailL0 = 0,
  kFailL1 = 1,
  kFailL2 = 2,
  kPass = 3
};
 
TVM_DLL BaseCheckResult StructInfoBaseCheck(const StructInfo& base, const StructInfo& derived,
                                            arith::Analyzer* ana = nullptr);
 
TVM_DLL bool IsBaseOf(const StructInfo& base, const StructInfo& derived,
                      arith::Analyzer* ana = nullptr);
 
TVM_DLL PrimExpr StructInfoBaseCheckPrecondition(const StructInfo& base, const StructInfo& derived);
 
TVM_DLL StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs,
                                 arith::Analyzer* ana = nullptr);
 
TVM_DLL ffi::Array<tir::Var> TIRVarsInStructInfo(const StructInfo& sinfo);
 
TVM_DLL ffi::Array<tir::Var> DefinableTIRVarsInStructInfo(const StructInfo& sinfo);
 
TVM_DLL ffi::Array<PrimExpr> CollectNonNegativeExpressions(const StructInfo& sinfo);
 
TVM_DLL ffi::Array<tir::Var> DefinedSymbolicVars(const Expr& expr);
 
TVM_DLL ffi::Array<tir::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 Expr RemoveAllUnused(Expr expr);
 
TVM_DLL OpPatternKind AnalyzeOpPatternKind(const tir::PrimFunc& func);
 
TVM_DLL bool HasReshapePattern(const tir::PrimFunc& 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 bool WellFormed(ffi::Variant<IRModule, Function> obj, bool check_struct_info = true);
 
TVM_DLL ffi::Map<tir::Block, ffi::Map<ObjectRef, tir::IndexMap>> SuggestLayoutTransforms(
    const Function& fn, ffi::Array<tir::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_