Files
apache--tvm/include/tvm/relax/script/builder/ir.h
T
wehub-resource-sync 26446540fa
Lint / lint (push) Has been cancelled
CI / MacOS (push) Has been cancelled
CI / Windows (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

145 lines
4.4 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.
*/
#ifndef TVM_RELAX_SCRIPT_BUILDER_IR_H_
#define TVM_RELAX_SCRIPT_BUILDER_IR_H_
#include <tvm/relax/expr.h>
#include <tvm/relax/script/builder/frame.h>
#include <tvm/relax/type.h>
#include <tvm/script/ir_builder/base.h>
namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {
/////////////////////////////// Function ////////////////////////////////
/*!
* \brief Start a function frame.
* \param is_pure Whether the function is annotated as pure.
* \param is_private Whether the function is annotated as private.
* \return The created ir_builder Function frame.
*/
TVM_DLL FunctionFrame Function(bool is_pure, bool is_private);
/*!
* \brief Add a parameter to the last function frame.
* \param name The name of the parameter.
* \param ty The ty of the parameter.
* \return The created function parameter var.
*/
TVM_DLL tvm::relax::Var Arg(const ffi::String& name, const tvm::Type& ty);
/*!
* \brief Specify the name of the last function frame.
* \param name The function name.
*/
TVM_DLL void FuncName(const ffi::String& name);
/*!
* \brief Specify the attrs of the last function frame.
* \param attrs The function attrs.
*/
TVM_DLL void FuncAttrs(ffi::Map<ffi::String, Any> attrs);
/*!
* \brief Specify the return type of the last function frame.
* \param ret_ty The return type.
*/
TVM_DLL void FuncRetType(const tvm::Type& ret_ty);
/*!
* \brief Specify the return value of the last function frame.
* \param value The return value.
*/
TVM_DLL void FuncRetValue(const tvm::relax::Expr& value);
///////////////////////////// BindingBlock //////////////////////////////
/*!
* \brief Start a binding block frame.
* \return The created ir_builder Block frame.
*/
TVM_DLL BindingBlockFrame BindingBlock();
/*!
* \brief Start a dataflow binding block frame.
* \return The created ir_builder Block frame.
*/
TVM_DLL BindingBlockFrame Dataflow();
/*!
* \brief Expose the dataflow block output variables as global ones
* \param vars The output variables of a dataflow block
*/
TVM_DLL void DataflowBlockOutput(const ffi::Array<tvm::relax::Var>& vars);
////////////////////////////// Bindings ////////////////////////////////
/*!
* \brief Emit a binding to the last binding block frame.
* \param value The right side value of the bindings to be emitted.
* \param annotate_ty The optional type annotation for the emitted value.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var Emit(const tvm::relax::Expr& value,
const ffi::Optional<tvm::Type>& annotate_ty = std::nullopt);
/*!
* \brief Emit a match_cast binding to the last binding block frame.
* \param value The value of the MatchCast to be emitted.
* \param ty The type of the MatchCast to be emitted.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var EmitMatchCast(const tvm::relax::Expr& value, const tvm::Type& ty);
/*!
* \brief Emit a binding to the last binding block frame.
* \param binding The binding to be emitted.
* \return The left side var of the emitted binding.
*/
TVM_DLL tvm::relax::Var EmitVarBinding(const tvm::relax::VarBinding& binding);
///////////////////////////// If Then Else /////////////////////////////
/*!
* \brief Create an if statement.
* \param condition The condition of if statement.
* \return The result IfFrame.
*/
IfFrame If(tvm::relax::Expr condition);
/*!
* \brief Create a then.
* \return The result ThenFrame.
*/
ThenFrame Then();
/*!
* \brief Create an else.
* \return The result ElseFrame.
*/
ElseFrame Else();
} // namespace relax
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_RELAX_SCRIPT_BUILDER_IR_H_