wehub-resource-sync
321 lines
11 KiB
C++
321 lines
11 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/block_builder.h
|
|
* \brief The utility for constructing Relax binding blocks.
|
|
*/
|
|
#ifndef TVM_RELAX_BLOCK_BUILDER_H_
|
|
#define TVM_RELAX_BLOCK_BUILDER_H_
|
|
|
|
#include <tvm/arith/analyzer.h>
|
|
#include <tvm/ir/unique_name_supply.h>
|
|
#include <tvm/relax/expr.h>
|
|
#include <tvm/relax/utils.h>
|
|
#include <tvm/runtime/base.h>
|
|
|
|
namespace tvm {
|
|
namespace relax {
|
|
|
|
/*!
|
|
* \brief A builder to build Relax binding blocks.
|
|
*
|
|
* BlockBuilder provides the following three categories
|
|
* of main functionalities for IR building and transformations:
|
|
*
|
|
* - Global context management: manages the IRModule,
|
|
* allowing query, update the surrounding global context.
|
|
* Provide context tools for analysis.
|
|
* - Scope management:
|
|
* - Manages block scopes for bulding nested blocks.
|
|
* - Emit bindings to the current scope.
|
|
* - Construct blocks by calling EndScope.
|
|
* - Normalization: Take an Expr, normalize it
|
|
* to deduce shape/type, turn things into normal forms.
|
|
*
|
|
* Importantly, these three categories of features can be dependent
|
|
* on each other. For example, when we emit into scope we will call
|
|
* normalize to ensure the code is in normal form. Similarly, when we
|
|
* normalize we could choose to emit into the current context.
|
|
*
|
|
* We would encourage the developers to keep these three category
|
|
* in mind when using and developing BlockBuilder, we can group
|
|
* the code in a logically clean way.
|
|
*
|
|
* BlockBuilderNode is implemented as a virtual interface to
|
|
* allow logically grouped implementation and internal data
|
|
* structures that are hidden from the users.
|
|
*/
|
|
class BlockBuilderNode : public ffi::Object {
|
|
public:
|
|
//-------------------------------
|
|
// Global Context management
|
|
//-------------------------------
|
|
/*!
|
|
* \brief Get the unique name supply for generating unique names.
|
|
*
|
|
* \return The unique name supply.
|
|
*/
|
|
virtual UniqueNameSupply name_supply() = 0;
|
|
|
|
/*!
|
|
* \brief Get the context IRModule in this builder.
|
|
*
|
|
* \note The context
|
|
* \return The IRModule in this BlockBuilder.
|
|
*/
|
|
virtual IRModule GetContextIRModule() const = 0;
|
|
|
|
/*!
|
|
* \brief Finalize the building process and return the result IRModule. Possibly rename
|
|
* GlobalVars in the IRModule to ensure name uniqueness and the invariant:
|
|
* every public function has the same name as its "global_symbol" attribute.
|
|
*
|
|
* \note this method should be called only once at the end of the building process, since it may
|
|
* invalidate global vars previously returned by this builder. See also
|
|
* transform::NormalizeGlobalVar.
|
|
*
|
|
* \return The result IRModule.
|
|
*/
|
|
virtual IRModule Finalize() = 0;
|
|
|
|
/*!
|
|
* \brief Add a Relax function or a TIR PrimFunc to internal context module.
|
|
* \param func The function to be added.
|
|
* \param func_name_hint The name hint of the function to be added.
|
|
* \note If the function to be added already exists, return its
|
|
* GlobalVar directly.
|
|
* \return The global var bound to the added function.
|
|
*/
|
|
virtual GlobalVar AddFunction(const BaseFunc& func, ffi::String func_name_hint) = 0;
|
|
|
|
/*!
|
|
* \brief Update a Relax function or a TIR PrimFunc in the internal context module.
|
|
* \param gv The global var referring the function to be updated.
|
|
* \param function The updated function.
|
|
*/
|
|
virtual void UpdateFunction(const GlobalVar& gv, BaseFunc function) = 0;
|
|
|
|
//-------------------------------
|
|
// Scope management
|
|
//-------------------------------
|
|
/*!
|
|
* \brief Lookup the binding value that var binds to in the current emitted sequences.
|
|
* \param var The input var.
|
|
* \return The Expr bound to the input \p var.
|
|
* \note For function parameters, this function returns std::nullopt.
|
|
*/
|
|
virtual ffi::Optional<Expr> LookupBinding(const Var& var) = 0;
|
|
|
|
/*!
|
|
* \brief Begin a new scope, with optional parameters that
|
|
* are visible within the scope.
|
|
*
|
|
* Symbolic variables from the parent scope are not available.
|
|
*
|
|
* \param params Parameters that are visible within the scope.
|
|
*
|
|
* \note This function should be called when new scope is introduced
|
|
* (e.g. function bodies) to properly track the variable
|
|
* availability and help the best effort deduction.
|
|
*
|
|
* \sa EndScope
|
|
*/
|
|
virtual void BeginScope(ffi::Optional<ffi::Array<Var>> params) = 0;
|
|
|
|
/*!
|
|
* \brief Begin a new scope, which inherits visible parameters from
|
|
* its parent scope.
|
|
*
|
|
* Symbolic variables from the parent scope are available.
|
|
*
|
|
* \note This function should be called when an inner scope is
|
|
* introduced (e.g. conditional branches) to properly track
|
|
* the variable availability and help the best effort
|
|
* deduction.
|
|
*
|
|
* \sa EndScope
|
|
*/
|
|
virtual void BeginInnerScope() = 0;
|
|
|
|
/*!
|
|
* \brief Append a definition to the current scope.
|
|
*
|
|
* \param var A variable within the current scope.
|
|
*
|
|
* \note This function should be called when a new variable is
|
|
* defined that may impact struct inference (e.g. MatchCast)
|
|
* to properly track the variable availability and help the
|
|
* best effort deduction.
|
|
*
|
|
* \sa EndScope
|
|
*/
|
|
virtual void AddDefinitionToScope(Var var) = 0;
|
|
|
|
/*! \brief End the previously defined scope. */
|
|
virtual void EndScope() = 0;
|
|
|
|
/*! \brief Begin to build a DataflowBlock. */
|
|
virtual void BeginDataflowBlock() = 0;
|
|
|
|
/*! \brief Begin to build a BindingBlock. */
|
|
virtual void BeginBindingBlock() = 0;
|
|
/*!
|
|
* \brief End building a BindingBlock.
|
|
* \return The BindingBlock being built.
|
|
*/
|
|
virtual BindingBlock EndBlock() = 0;
|
|
|
|
/*!
|
|
* \brief Check if the block being built is DataflowBlock or not.
|
|
* \return A boolean that indicates if the block being built is DataflowBlock or not.
|
|
*/
|
|
virtual bool CurrentBlockIsDataFlow() = 0;
|
|
|
|
/*!
|
|
* \brief Emits an Expr, and returns the variable it is bound to.
|
|
* \param expr The Expr to be emitted.
|
|
* \param name_hint Name hint for the bound variable.
|
|
* \return The new variable that \p expr is bound to.
|
|
*
|
|
* \note This Emit function normalizes the \p expr, and
|
|
* performs shape and type deductions by calling Normalize.
|
|
*/
|
|
virtual Var Emit(Expr expr, ffi::String name_hint = "") = 0;
|
|
|
|
/*!
|
|
* \brief Emit a MatchCast.
|
|
* \param value The input value.
|
|
* \param ty The type to be matched.
|
|
* \param name_hint Name hint for the bound variable.
|
|
* \return The variable bound to the MatchCast.
|
|
*/
|
|
virtual Var EmitMatchCast(Expr value, Type ty, ffi::String name_hint = "") = 0;
|
|
|
|
/*!
|
|
* \brief Generate an output for the current dataflow block.
|
|
* \param output The output variable of the block.
|
|
* \param name_hint Name hint for the bound variable.
|
|
* \return The variable bound to \p output.
|
|
*/
|
|
virtual Var EmitOutput(Expr output, ffi::String name_hint = "") = 0;
|
|
|
|
/*!
|
|
* \brief Emit a binding that is already normalized.
|
|
*
|
|
* \param normalized_binding A binding whose value is already normalized.
|
|
*
|
|
* \note This function requires binding to be pre-normalized.
|
|
*/
|
|
virtual void EmitNormalized(Binding normalized_binding) = 0;
|
|
|
|
/*!
|
|
* \brief Convert an expression to normal form, and try to eagerly infer types and shapes.
|
|
* \param expr The input expression.
|
|
* \return The normalized expression.
|
|
*
|
|
* \note Invariant: If any of the sub expr have ty field.
|
|
* they must have already been normalized.
|
|
*/
|
|
virtual Expr Normalize(const Expr& expr) = 0;
|
|
|
|
/*!
|
|
* \brief Normalize argument to a call or another IRNode.
|
|
* \param expr The input expression.
|
|
* \return The normalized expression.
|
|
*
|
|
* \note This function will create a binding var for non-leaf expressions such as Call.
|
|
*/
|
|
virtual Expr NormalizeArgument(const Expr& expr) = 0;
|
|
|
|
/*!
|
|
* \brief Get the analyzer of the BlockBuilder.
|
|
* \return The BlockBuilder's arithmetic analyzer.
|
|
*/
|
|
virtual arith::Analyzer GetAnalyzer() = 0;
|
|
|
|
static constexpr const bool _type_mutable = true;
|
|
TVM_FFI_DECLARE_OBJECT_INFO("relax.BlockBuilder", BlockBuilderNode, ffi::Object);
|
|
};
|
|
|
|
class BlockBuilder : public ffi::ObjectRef {
|
|
public:
|
|
/*!
|
|
* \brief Create a BlockBuilder.
|
|
*
|
|
* \param ctx_mod Optional before-transformation context module for rewriting.
|
|
*
|
|
* \return The created BlockBuilder.
|
|
*
|
|
* \note When rewriting an existing IRModule, it is important to pass it in as
|
|
* ctx_mod so you can lookup the context functions for cross function
|
|
* call analysis.
|
|
*/
|
|
TVM_DLL static BlockBuilder Create(ffi::Optional<IRModule> ctx_mod);
|
|
|
|
/*! \brief A marker struct to disable FNormalize
|
|
*
|
|
* This struct is used as a marker to disable the use of FNormalize
|
|
* by this block builder. This should only be used for TVMScript
|
|
* parsing, which may require producing un-normalized Relax IR for
|
|
* testing purposes, and to ensure that round-trips are unchanged.
|
|
*
|
|
* The name is deliberately verbose to draw attention during a code
|
|
* review. The explicit default constructor prevents aggregate
|
|
* initialization, ensuring that the full name of the marker struct
|
|
* appears at the callsite.
|
|
*
|
|
* This constructor is marked as no-lint to allow a zero-parameter
|
|
* constructor to be marked as explicit. The constructor must be
|
|
* explicit in order to disable aggregate initialization in C++17.
|
|
* While C++20 disables aggregate initialization when a
|
|
* user-declared constructor is present, C++17 only disables
|
|
* aggregate initialization when a user-defined constructor is
|
|
* present. Therefore, we need to mark the zero-parameter
|
|
* constructor as explicit in order to prevent aggregate
|
|
* initialization, and to ensure that the name appears at all
|
|
* callsites.
|
|
*/
|
|
struct DisableOperatorSpecificNormalizationForTVMScript {
|
|
explicit DisableOperatorSpecificNormalizationForTVMScript() = default; // NOLINT(*)
|
|
};
|
|
/*!
|
|
* \brief Create a BlockBuilder.
|
|
*
|
|
* \param ctx_mod Optional before-transformation context module for rewriting.
|
|
*
|
|
* \param tag An instance of DisableOperatorSpecificNormalizationForTVMScript
|
|
*
|
|
* \return The created BlockBuilder.
|
|
*
|
|
* \note When rewriting an existing IRModule, it is important to pass it in as
|
|
* ctx_mod so you can lookup the context functions for cross function
|
|
* call analysis.
|
|
*/
|
|
TVM_DLL static BlockBuilder Create(ffi::Optional<IRModule> ctx_mod,
|
|
DisableOperatorSpecificNormalizationForTVMScript tag);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BlockBuilder, ffi::ObjectRef, BlockBuilderNode);
|
|
};
|
|
|
|
} // namespace relax
|
|
} // namespace tvm
|
|
|
|
#endif // TVM_RELAX_BLOCK_BUILDER_H_
|