transform.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_RELAY_TRANSFORM_H_
#define TVM_RELAY_TRANSFORM_H_

#include <tvm/ir/transform.h>
#include <tvm/relay/attrs/transform.h>
#include <tvm/relay/expr.h>
#include <tvm/relay/function.h>
#include <tvm/relay/op.h>
#include <tvm/relay/op_attr_types.h>
#include <tvm/target/compilation_config.h>
#include <tvm/target/target.h>
#include <tvm/target/virtual_device.h>

#include <string>

namespace tvm {
namespace relay {
namespace transform {

using Pass = tvm::transform::Pass;
using PassNode = tvm::transform::PassNode;
using PassInfo = tvm::transform::PassInfo;
using PassInfoNode = tvm::transform::PassInfoNode;
using PassContext = tvm::transform::PassContext;
using PassContextNode = tvm::transform::PassContextNode;
using Sequential = tvm::transform::Sequential;

/*
 * \brief Create a function pass.
 *
 * \param pass_func The packed function that contains the optimization.
 * \param opt_level The optimization level of the function pass.
 * \param name The name of the function pass.
 * \param required The list of the passes that the function pass is dependent on.
 *
 * \return The created function pass.
 */
TVM_DLL Pass CreateFunctionPass(
    const runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)>& pass_func,
    int opt_level, String name, tvm::Array<String> required);

TVM_DLL Pass DeadCodeElimination(bool inline_once = false, bool ignore_purity = false);

TVM_DLL Pass LazyGradientInit();

TVM_DLL Pass FoldConstant(bool fold_qnn = false);

TVM_DLL Pass SplitArgs(int max_function_args);

TVM_DLL Pass FuseOps(int fuse_opt_level = -1);

TVM_DLL Pass DefuseOps();

TVM_DLL Pass RewriteAnnotatedOps(int fallback_device);

TVM_DLL Pass ToBasicBlockNormalForm();

TVM_DLL Pass ToANormalForm();

TVM_DLL Expr ToANormalForm(const Expr& expr);

TVM_DLL Pass ToCPS();

TVM_DLL Pass ToGraphNormalForm();

TVM_DLL Pass PartialEval();

TVM_DLL Pass SimplifyInference();

TVM_DLL Pass FastMath();

TVM_DLL Pass DynamicToStatic();

TVM_DLL Pass InferType();

TVM_DLL Type InferTypeLocal(const Expr& expr);

TVM_DLL Pass EliminateCommonSubexpr(runtime::PackedFunc fskip = nullptr);

TVM_DLL Pass CombineParallelConv2D(uint64_t min_num_branches = 3);

TVM_DLL Pass CombineParallelDense(uint64_t min_num_branches = 3, bool to_batch_matmul = true);

TVM_DLL Pass CombineParallelBatchMatmul(uint64_t min_num_branches = 3);

TVM_DLL Pass BackwardFoldScaleAxis();

TVM_DLL Pass ForwardFoldScaleAxis();

TVM_DLL Pass FoldScaleAxis();

TVM_DLL Pass CanonicalizeOps();

TVM_DLL Pass AlterOpLayout();

TVM_DLL Pass AutoSchedulerLayoutRewrite();

TVM_DLL Pass MetaScheduleLayoutRewrite();

TVM_DLL Pass ConvertLayout(const Map<String, Array<String>>& desired_layouts);

TVM_DLL Pass Legalize(const String& legalize_map_attr_name = "FTVMLegalize");

TVM_DLL Pass CanonicalizeCast();

TVM_DLL Pass EtaExpand(bool expand_constructor, bool expand_global_var);

TVM_DLL Pass PartitionGraph();

TVM_DLL Pass Inline();

TVM_DLL Pass RemoveUnusedFunctions(Array<runtime::String> entry_functions);

TVM_DLL Pass SimplifyExpr();

TVM_DLL Pass SimplifyExprPostAlterOp();

TVM_DLL Pass RelayToTIRTargetHook(CompilationConfig config);

TVM_DLL Pass ManifestAlloc(VirtualDevice cpu_virtual_device);

TVM_DLL Pass ManifestLifetimes();

TVM_DLL Pass PlanDevices(CompilationConfig config);

TVM_DLL Pass FlattenAtrousConv();

TVM_DLL Pass AnnotateUsedMemory();

TVM_DLL Pass CapturePostDfsIndexInSpans();

TVM_DLL Pass AnnotateMemoryScope();

TVM_DLL Pass RemoveStandaloneReshapes();

}  // namespace transform

TVM_DLL Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& binds);

TVM_DLL Function SubstituteBoundVars(const Function& func, const tvm::Map<Var, Expr>& binds);

TVM_DLL Expr ForwardRewrite(const Expr& expr, const String& rewrite_map_attr_name,
                            std::function<ObjectRef(const Call&)> fcontext = nullptr,
                            std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);

TVM_DLL Expr ForwardRewrite(const Expr& expr, const FForwardRewrite& rewrite_func,
                            std::function<ObjectRef(const Call&)> fcontext = nullptr,
                            std::function<Expr(const Expr&)> fmulti_ref_trigger = nullptr);

TVM_DLL Expr RewriteAnnotatedOps(const Expr& expr, int fallback_device);

TVM_DLL Function ToCPS(const Function& f, const IRModule& mod);

TVM_DLL Function UnCPS(const Function& f);

TVM_DLL Expr DeDup(const Expr& e);

namespace legalize {
TVM_DLL Expr Legalize(const Expr& expr, const std::string& legalize_map_attr_name);
}  // namespace legalize

}  // namespace relay
}  // namespace tvm

#endif  // TVM_RELAY_TRANSFORM_H_