wehub-resource-sync
310 lines
9.2 KiB
C++
310 lines
9.2 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_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_
|