wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tvm/relax/script/builder/frame.h
+332
View File
@@ -0,0 +1,332 @@
/*
* 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_FRAME_H_
#define TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/script/ir_builder/ir/ir.h>
#include <utility>
namespace tvm {
namespace script {
namespace ir_builder {
namespace relax {
/*! \brief The base ir_builder frame for the relax dialect. */
class RelaxFrameNode : public IRBuilderFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<RelaxFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.RelaxFrame", RelaxFrameNode,
IRBuilderFrameNode);
};
class RelaxFrame : public IRBuilderFrame {
public:
explicit RelaxFrame(ffi::ObjectPtr<RelaxFrameNode> data) : IRBuilderFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RelaxFrame, IRBuilderFrame, RelaxFrameNode);
protected:
RelaxFrame() = default;
};
/*! \brief The base ir_builder frame for frames with SeqExpr
i.e. Functions, If branches
*/
class SeqExprFrameNode : public RelaxFrameNode {
public:
/*! \brief The binding blocks inside the frame. */
ffi::Array<tvm::relax::BindingBlock> binding_blocks;
/*! \brief The frame output expr. `std::nullopt` when undefined. */
ffi::Optional<tvm::relax::Expr> output;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SeqExprFrameNode>()
.def_ro("binding_blocks", &SeqExprFrameNode::binding_blocks)
.def_ro("output", &SeqExprFrameNode::output);
}
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.relax.SeqExprFrame", SeqExprFrameNode,
RelaxFrameNode);
public:
void EnterWithScope() override;
void ExitWithScope() override;
};
class SeqExprFrame : public RelaxFrame {
public:
explicit SeqExprFrame(ffi::ObjectPtr<SeqExprFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SeqExprFrame, RelaxFrame, SeqExprFrameNode);
};
/*! \brief The ir_builder frame for the relax function. */
class FunctionFrameNode : public SeqExprFrameNode {
public:
/*!
* \brief The function name.
* \note The name will not be specified in constructor, so it is "Optional",
* However, we must specify the name by `R.func_name` before exit this frame.
*/
ffi::Optional<ffi::String> name;
/*! \brief The function params. */
ffi::Array<tvm::relax::Var> params;
/*!
* \brief The function return type.
* \note Usually the function return type can be deduced by the function body.
* But we can use this field to specify a more "accurate" return type.
* i.e. If the `ret_ty` is None, try to use the deduced type from body
* If the `ret_ty` is not None, we can still take body.ty
* if we ret_ty is base of body.ty. If not, we will
* take the specified `ret_ty`.
*/
ffi::Optional<tvm::Type> ret_ty;
/*! \brief Whether the function is annotated as pure */
ffi::Optional<bool> is_pure;
/*! \brief Whether the function is annotated as private */
ffi::Optional<bool> is_private;
/*! \brief The function attributes. */
ffi::Map<ffi::String, Any> attrs;
/*! \brief The block builder to create Relax function. */
tvm::relax::BlockBuilder block_builder;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<FunctionFrameNode>()
.def_ro("name", &FunctionFrameNode::name)
.def_ro("params", &FunctionFrameNode::params)
.def_ro("ret_ty", &FunctionFrameNode::ret_ty)
.def_ro("is_pure", &FunctionFrameNode::is_pure)
.def_ro("attrs", &FunctionFrameNode::attrs);
// `binding_blocks` and `output` are inherited from SeqExprFrameNode.
// `block_builder` is not registered as it's not visited.
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.FunctionFrame", FunctionFrameNode,
SeqExprFrameNode);
public:
void EnterWithScope() final;
void ExitWithScope() final;
};
class FunctionFrame : public SeqExprFrame {
public:
explicit FunctionFrame(ffi::ObjectPtr<FunctionFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionFrame, SeqExprFrame, FunctionFrameNode);
};
/*! \brief The ir_builder frame for relax binding blocks. */
class BindingBlockFrameNode : public RelaxFrameNode {
public:
/*! \brief The flag that indicates whether the block is a dataflow block. */
bool is_dataflow;
/*! \brief The variables emitted in this block. */
ffi::Array<tvm::relax::Var> emitted_vars;
/*!
* \brief A boolean indicating if the dataflow block is ended of construction.
* If it is true, any new binding trying to be emitted into this block will cause an error.
* \note Only used for a dataflow block.
*/
bool block_ended;
/*!
* \brief The output vars of the dataflow block.
* \note Only used for a dataflow block.
*/
ffi::Array<tvm::relax::Var> output_vars;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BindingBlockFrameNode>()
.def_ro("is_dataflow", &BindingBlockFrameNode::is_dataflow)
.def_ro("emitted_vars", &BindingBlockFrameNode::emitted_vars)
.def_ro("output_vars", &BindingBlockFrameNode::output_vars);
// `block_ended` is not registered as it's not visited.
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.BindingBlockFrame",
BindingBlockFrameNode, RelaxFrameNode);
public:
void EnterWithScope() final;
void ExitWithScope() final;
};
class BindingBlockFrame : public RelaxFrame {
public:
explicit BindingBlockFrame(ffi::ObjectPtr<BindingBlockFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BindingBlockFrame, RelaxFrame,
BindingBlockFrameNode);
};
/*!
* \brief A frame that represents if statement.
*
* \sa IfFrame
*/
class IfFrameNode : public RelaxFrameNode {
public:
/*! \brief The condition of the if statement. */
tvm::relax::Expr condition;
/*! \brief The Bindings in the true branch. */
ffi::Optional<tvm::relax::Expr> then_expr;
/*! \brief The Bindings in the false branch. */
ffi::Optional<tvm::relax::Expr> else_expr;
/*! \brief The Binding var. */
tvm::relax::Var var;
/*! \brief The binding var name. */
ffi::String var_name;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IfFrameNode>()
.def_ro("condition", &IfFrameNode::condition)
.def_ro("then_expr", &IfFrameNode::then_expr)
.def_ro("else_expr", &IfFrameNode::else_expr)
.def_ro("var", &IfFrameNode::var)
.def_ro("var_name", &IfFrameNode::var_name);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.IfFrame", IfFrameNode, RelaxFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to IfFrameNode.
*
* \sa IfFrameNode
*/
class IfFrame : public RelaxFrame {
public:
explicit IfFrame(ffi::ObjectPtr<IfFrameNode> data) : RelaxFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfFrame, RelaxFrame, IfFrameNode);
};
/*!
* \brief A frame that represents then.
*
* \sa ThenFrame
*/
class ThenFrameNode : public SeqExprFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ThenFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ThenFrame", ThenFrameNode,
SeqExprFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ThenFrameNode.
*
* \sa ThenFrameNode
*/
class ThenFrame : public SeqExprFrame {
public:
explicit ThenFrame(ffi::ObjectPtr<ThenFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThenFrame, SeqExprFrame, ThenFrameNode);
};
/*!
* \brief A frame that represents else.
*
* \sa ElseFrame
*/
class ElseFrameNode : public SeqExprFrameNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ElseFrameNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.relax.ElseFrame", ElseFrameNode,
SeqExprFrameNode);
public:
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ElseFrameNode.
*
* \sa ElseFrameNode
*/
class ElseFrame : public SeqExprFrame {
public:
explicit ElseFrame(ffi::ObjectPtr<ElseFrameNode> data) : SeqExprFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, SeqExprFrame, ElseFrameNode);
};
} // namespace relax
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_RELAX_SCRIPT_BUILDER_FRAME_H_
include/tvm/relax/script/builder/ir.h
+144
View File
@@ -0,0 +1,144 @@
/*
* 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_