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/script/ir_builder/base.h
+309
View File
@@ -0,0 +1,309 @@
/*
* 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_SCRIPT_IR_BUILDER_BASE_H_
#define TVM_SCRIPT_IR_BUILDER_BASE_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/node_functor.h>
#include <vector>
namespace tvm {
namespace script {
namespace ir_builder {
////////////////////////////// IRBuilderFrame //////////////////////////////
/*!
* \brief A stack frame of the IRBuilder used to keep track of the current scope.
* Furthermore, the information stored in each stack frame can be useful for context-dependent
* IR construction.
*
* \example
*
* The `T::MatchBuffer` below adds an element in `PrimFuncNode::buffer_map`:
*
* \code {.cpp}
*
* using T = tvm::script::ir_builder::tirx;
* With <PrimFuncFrame> _(...);
* Buffer buffer = T::MatchBuffer(...);
*
* \endcode
*
* The `T::MatchBuffer` below instead generates `MatchBufferRegion` in a TIR block:
*
* \code {.cpp}
*
* using T = tvm::script::ir_builder::tirx;
* With <PrimFuncFrame> _(...);
* {
* With<SBlockFrame> _2(...);
* Buffer buffer = T::MatchBuffer(...);
* }
*
* \endcode
*/
class IRBuilderFrameNode : public ffi::Object {
public:
/*! \brief A list of callbacks used when exiting the frame. */
std::vector<ffi::TypedFunction<void()>> callbacks;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IRBuilderFrameNode>();
// `callbacks` is not registered as it's not visited.
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("script.ir_builder.IRBuilderFrame", IRBuilderFrameNode, ffi::Object);
public:
/*! \brief Default destructor. */
virtual ~IRBuilderFrameNode() = default;
/*!
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
virtual void EnterWithScope();
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
virtual void ExitWithScope();
/*!
* \brief Add a callback method invoked when exiting the RAII scope.
* \param callback The callback to be added.
*/
void AddCallback(ffi::TypedFunction<void()> callback);
};
/*!
* \brief Managed reference to an IRBuilderFrameNode.
* \sa IRBuilderFrameNode
*/
class IRBuilderFrame : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilderFrame, ffi::ObjectRef, IRBuilderFrameNode);
protected:
/*! \brief Disallow direct construction of this object. */
IRBuilderFrame() = default;
explicit IRBuilderFrame(ffi::ObjectPtr<IRBuilderFrameNode> data) : ffi::ObjectRef(data) {}
public:
/*!
* \brief Redirected to `IRBuilderFrameNode::EnterWithScope`.
* \sa IRBuilderFrameNode::EnterWithScope
*/
inline void EnterWithScope() {
TVM_FFI_ICHECK(data_ != nullptr);
static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
}
/*!
* \brief Redirected to `IRBuilderFrameNode::ExitWithScope`.
* \sa IRBuilderFrameNode::ExitWithScope
*/
inline void ExitWithScope() {
TVM_FFI_ICHECK(data_ != nullptr);
static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
data_.reset();
}
};
////////////////////////////// IRBuilder //////////////////////////////
/*!
* \brief A dialect-agnostic IRBuilder that constructs any IR of TVM.
* An idiomatic use of this class is to put this inside the RAII with-scope,
* call dialect-specific methods accordingly. Upon exiting the scope.
*
* \code
*
* PrimFunc ConstructPrimFunc() {
* using tvm::script::ir_builder::IRBuilder;
* using T = tvm::script::ir_builder::tirx;
* IRBuilder builder;
* // Step 1. Place IRBuilder inside the with-scope.
* {
* With<IRBuilder> _(builder);
* // Step 2. Call dialect-specific methods.
* With<T::PrimFuncFrame> _2(...);
* T::MatchBuffer(...);
* }
* // Step 3. Return the constructed PrimFunc.
* return builder->Get<PrimFunc>();
* }
*
* \endcode
*/
class IRBuilderNode : public ffi::Object {
public:
/*! \brief A stack of context frames in the IRBuilder */
ffi::Array<IRBuilderFrame> frames;
/*! \brief The outcome of IR construction */
ffi::Optional<ffi::ObjectRef> result;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IRBuilderNode>()
.def_ro("frames", &IRBuilderNode::frames)
.def_ro("result", &IRBuilderNode::result);
}
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRBuilder", IRBuilderNode, ffi::Object);
public:
/*!
* \brief Find a frame of the given type in the stack `this->frames` from top to bottom.
* \tparam T The type of the frame to find.
* \return The frame if found, otherwise std::nullopt.
*/
template <typename TFrame>
inline ffi::Optional<TFrame> FindFrame() const;
/*!
* \brief Get the frame on top of the stack `this->frames` if its type is `TFrame`.
* \tparam TFrame The assumed type of the last frame on stack.
* \return The frame if the stack is non-empty and the top of the stack is of type `TFrame`.
* Otherwise std::nullopt.
*/
template <typename TFrame>
inline ffi::Optional<TFrame> GetLastFrame() const;
/*!
* \brief Get the IR being constructed.
* \tparam TObjectRef The type of the IR being constructed.
* \return The resulting IR. Throw an exception if the IR is not constructed yet.
*/
template <typename TObjectRef>
inline TObjectRef Get() const;
};
/*!
* \brief Managed reference to an IRBuilderNode.
* \sa IRBuilderNode
*/
class IRBuilder : public ffi::ObjectRef {
public:
/*! \brief Creates an IRBuilder. */
IRBuilder();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRBuilder, ffi::ObjectRef, IRBuilderNode);
public:
/*!
* \brief Puts the current IRBuilder into a thread-local scope, which can be retrieved using
* `IRBuilder::Current()`.
*
* \code {.cpp}
* IRBuilder builder;
* {
* With<IRBuilder> _(builder);
* // IRBuilder::Current() == builder
* }
* // IRBuilder::Current() == nullptr
* \endcode
*
* \sa IRBuilder::Current
* \sa IRBuilder::ExitWithScope
* \sa tvm::support::With
*/
void EnterWithScope();
/*!
* \brief Exit the RAII scope.
* \sa IRBuilder::EnterWithScope
* \sa IRBuilder::Current
* \sa tvm::support::With
*/
void ExitWithScope();
/*!
* \brief Get the current IRBuilder in the current thread-local scope.
* \return The current IRBuilder.
* \sa IRBuilder::EnterWithScope
* \sa IRBuilder::ExitWithScope
* \sa tvm::support::With
*/
static IRBuilder Current();
/*! \brief See if the current thread-local scope has an IRBuilder. */
static bool IsInScope();
/*!
* \brief Give a string name to the `obj`
* \tparam TObjectRef The type of the object to name.
* \param name The name to give to the object.
* \param obj The object to name.
*/
template <class TObjectRef>
inline static TObjectRef Name(ffi::String name, TObjectRef obj);
};
////////////////////////////// Details //////////////////////////////
namespace details {
class Namer {
public:
using FType = NodeFunctor<void(const ffi::ObjectRef&, ffi::String)>;
static FType& vtable();
static void Name(ffi::ObjectRef node, ffi::String name);
};
} // namespace details
template <class TObjectRef>
inline TObjectRef IRBuilder::Name(ffi::String name, TObjectRef obj) {
details::Namer::Name(obj, name);
return obj.template as_or_throw<TObjectRef>();
}
template <typename TFrame>
inline ffi::Optional<TFrame> IRBuilderNode::FindFrame() const {
using TFrameNode = typename TFrame::ContainerType;
for (auto it = frames.rbegin(); it != frames.rend(); ++it) {
if (const TFrameNode* p = (*it).template as<TFrameNode>()) {
return ffi::GetRef<TFrame>(p);
}
}
return std::nullopt;
}
template <typename TFrame>
inline ffi::Optional<TFrame> IRBuilderNode::GetLastFrame() const {
using TFrameNode = typename TFrame::ContainerType;
if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
return frames.back().as_or_throw<TFrame>();
}
return std::nullopt;
}
template <typename TObjectRef>
inline TObjectRef IRBuilderNode::Get() const {
using TObject = typename TObjectRef::ContainerType;
TVM_FFI_CHECK(result.has_value(), IndexError) << "No result exists in IRBuilder yet";
const auto* n = result.as<TObject>();
TVM_FFI_CHECK(n != nullptr, TypeError)
<< "IRBuilder result is not of type: " << TObject::_type_key;
return ffi::GetRef<TObjectRef>(n);
}
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_SCRIPT_IR_BUILDER_BASE_H_
include/tvm/script/ir_builder/ir/frame.h
+87
View File
@@ -0,0 +1,87 @@
/*
* 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_SCRIPT_IR_BUILDER_IR_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/ir/module.h>
#include <tvm/script/ir_builder/base.h>
#include <vector>
namespace tvm {
namespace script {
namespace ir_builder {
namespace ir {
/*!
* \brief A frame that represents the IRModule frame with functions and global variables.
*
* \sa IRModuleFrame
*/
class IRModuleFrameNode : public IRBuilderFrameNode {
public:
/*! \brief A map from string names to global variables that ensures global uniqueness. */
ffi::Map<ffi::String, GlobalVar> global_var_map;
/*!
* \brief A map from GlobalVar to all global functions.
* \note Only defined functions are in the map, while declared functions are not included.
*/
ffi::Map<GlobalVar, BaseFunc> functions;
/*! \brief IRModule's attributes. */
ffi::Map<ffi::String, Any> attrs;
/*! \brief IRModule's global_infos */
ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IRModuleFrameNode>()
.def_ro("global_vars", &IRModuleFrameNode::global_var_map)
.def_ro("functions", &IRModuleFrameNode::functions)
.def_ro("attrs", &IRModuleFrameNode::attrs)
.def_ro("global_infos", &IRModuleFrameNode::global_infos);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.IRModuleFrame", IRModuleFrameNode,
IRBuilderFrameNode);
public:
void ExitWithScope() final;
};
/*!
* \brief Managed reference to IRModuleFrameNode.
*
* \sa IRModuleFrameNode
*/
class IRModuleFrame : public IRBuilderFrame {
public:
explicit IRModuleFrame(ffi::ObjectPtr<IRModuleFrameNode> data) : IRBuilderFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IRModuleFrame, IRBuilderFrame, IRModuleFrameNode);
};
} // namespace ir
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_SCRIPT_IR_BUILDER_IR_FRAME_H_
include/tvm/script/ir_builder/ir/ir.h
+61
View File
@@ -0,0 +1,61 @@
/*
* 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_SCRIPT_IR_BUILDER_IR_IR_H_
#define TVM_SCRIPT_IR_BUILDER_IR_IR_H_
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <vector>
namespace tvm {
namespace script {
namespace ir_builder {
namespace ir {
/*!
* \brief The IRModule declaration statement.
* \return The IRModuleFrame.
*/
TVM_DLL IRModuleFrame IRModule();
/*!
* \brief Declare a Function without given the specific function implementation.
* \note It is usually used in cross-function call. And we can specify the function by `DefFunction`
* \param func_name The function unique name.
* \param func_signature A Function w/o body, which used to specify the function signature
* (i.e. func params and func return type/shape).
* \return The corresponding GlobalVar.
*/
TVM_DLL GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature);
/*!
* \brief Define the function which is declared before.
* \param func_name The function unique name.
* \param func The given function implementation
*/
TVM_DLL void DefFunction(const ffi::String& func_name, const BaseFunc& func);
} // namespace ir
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_SCRIPT_IR_BUILDER_IR_IR_H_