tvm
/workspace/include/tvm/script/ir_builder/base.h

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.The T::MatchBuffer below adds an element in PrimFuncNode::buffer_map:

With <PrimFuncFrame> _(...);
Buffer buffer = T::MatchBuffer(...);
Definition: frame.h:29
Buffer MatchBuffer(ObjectRef param, Array< PrimExpr > shape, DataType dtype=DataType::Float(32), Optional< Var > data=NullOpt, Array< PrimExpr > strides={}, PrimExpr elem_offset=PrimExpr(), String storage_scope="global", int align=-1, int offset_factor=0, String buffer_type="default", Array< IntImm > axis_separators={})
The buffer match statement.

The T::MatchBuffer below instead generates MatchBufferRegion in a TIR block:

With <PrimFuncFrame> _(...);
{
With<BlockFrame> _2(...);
Buffer buffer = T::MatchBuffer(...);
}
/*
* 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/ir/expr.h>
#include <tvm/node/node.h>
#include <vector>
namespace tvm {
namespace script {
namespace ir_builder {
class IRBuilderFrameNode : public runtime::Object {
public:
std::vector<runtime::TypedPackedFunc<void()>> callbacks;
// `callbacks` is not visited.
}
static constexpr const char* _type_key = "script.ir_builder.IRBuilderFrame";
TVM_DECLARE_BASE_OBJECT_INFO(IRBuilderFrameNode, runtime::Object);
public:
virtual ~IRBuilderFrameNode() = default;
virtual void EnterWithScope();
virtual void ExitWithScope();
void AddCallback(runtime::TypedPackedFunc<void()> callback);
};
class IRBuilderFrame : public runtime::ObjectRef {
public:
virtual ~IRBuilderFrame() = default;
protected:
IRBuilderFrame() = default;
public:
inline void EnterWithScope() {
ICHECK(data_ != nullptr);
static_cast<IRBuilderFrameNode*>(data_.get())->EnterWithScope();
}
inline void ExitWithScope() {
ICHECK(data_ != nullptr);
static_cast<IRBuilderFrameNode*>(data_.get())->ExitWithScope();
data_.reset();
}
};
class IRBuilderNode : public runtime::Object {
public:
runtime::Array<IRBuilderFrame> frames;
Optional<ObjectRef> result;
v->Visit("frames", &frames);
v->Visit("result", &result);
}
static constexpr const char* _type_key = "script.ir_builder.IRBuilder";
TVM_DECLARE_FINAL_OBJECT_INFO(IRBuilderNode, runtime::Object);
public:
template <typename TFrame>
inline Optional<TFrame> FindFrame() const;
template <typename TFrame>
inline Optional<TFrame> GetLastFrame() const;
template <typename TObjectRef>
inline TObjectRef Get() const;
};
class IRBuilder : public runtime::ObjectRef {
public:
public:
void ExitWithScope();
static IRBuilder Current();
static bool IsInScope();
template <class TObjectRef>
inline static TObjectRef Name(String name, TObjectRef obj);
};
namespace details {
class Namer {
public:
using FType = NodeFunctor<void(const ObjectRef&, String)>;
static FType& vtable();
static void Name(ObjectRef node, String name);
};
} // namespace details
template <class TObjectRef>
inline TObjectRef IRBuilder::Name(String name, TObjectRef obj) {
return Downcast<TObjectRef>(obj);
}
template <typename TFrame>
inline 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 GetRef<TFrame>(p);
}
}
return NullOpt;
}
template <typename TFrame>
inline Optional<TFrame> IRBuilderNode::GetLastFrame() const {
using TFrameNode = typename TFrame::ContainerType;
if (!frames.empty() && frames.back()->IsInstance<TFrameNode>()) {
return Downcast<TFrame>(frames.back());
}
return NullOpt;
}
template <typename TObjectRef>
inline TObjectRef IRBuilderNode::Get() const {
using TObject = typename TObjectRef::ContainerType;
CHECK(result.defined()) << "IndexError: No result exists in IRBuilder yet";
const auto* n = result.as<TObject>();
CHECK(n != nullptr) << "TypeError: IRBuilder result is not of type: " << TObject::_type_key;
return GetRef<TObjectRef>(n);
}
} // namespace ir_builder
} // namespace script
} // namespace tvm
#endif // TVM_SCRIPT_IR_BUILDER_BASE_H_
Visitor class to get the attributes of an AST/IR node. The content is going to be called for each fie...
Definition: reflection.h:52
ObjectRef()=default
default constructor
ObjectPtr< Object > data_
Internal pointer that backs the reference.
Definition: object.h:605
void AddCallback(runtime::TypedPackedFunc< void()> callback)
Add a callback method invoked when exiting the RAII scope.
void VisitAttrs(tvm::AttrVisitor *v)
Definition: base.h:69
TVM_DECLARE_BASE_OBJECT_INFO(IRBuilderFrameNode, runtime::Object)
static constexpr const char * _type_key
Definition: base.h:73
virtual ~IRBuilderFrameNode()=default
Default destructor.
virtual void ExitWithScope()
The method called when exiting RAII scope.
std::vector< runtime::TypedPackedFunc< void()> > callbacks
A list of callbacks used when exiting the frame.
Definition: base.h:67
virtual void EnterWithScope()
The method called when entering RAII scope.
void EnterWithScope()
Redirected to IRBuilderFrameNode::EnterWithScope.
Definition: base.h:115
IRBuilderFrame()=default
Disallow direct construction of this object.
void ExitWithScope()
Redirected to IRBuilderFrameNode::ExitWithScope.
Definition: base.h:123
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilderFrame, ObjectRef, IRBuilderFrameNode)
virtual ~IRBuilderFrame()=default
Default destructor.
void VisitAttrs(tvm::AttrVisitor *v)
Definition: base.h:163
TObjectRef Get() const
Get the IR being constructed.
Definition: base.h:292
Optional< TFrame > GetLastFrame() const
Get the frame on top of the stack this->frames if its type is TFrame.
Definition: base.h:283
static constexpr const char * _type_key
Definition: base.h:168
TVM_DECLARE_FINAL_OBJECT_INFO(IRBuilderNode, runtime::Object)
Optional< TFrame > FindFrame() const
Find a frame of the given type in the stack this->frames from top to bottom.
Definition: base.h:272
Optional< ObjectRef > result
The outcome of IR construction.
Definition: base.h:161
runtime::Array< IRBuilderFrame > frames
A stack of context frames in the IRBuilder.
Definition: base.h:159
void ExitWithScope()
Exit the RAII scope.
static IRBuilder Current()
Get the current IRBuilder in the current thread-local scope.
void EnterWithScope()
Puts the current IRBuilder into a thread-local scope, which can be retrieved using IRBuilder::Current...
static bool IsInScope()
See if the current thread-local scope has an IRBuilder.
static TObjectRef Name(String name, TObjectRef obj)
Give a string name to the obj
Definition: base.h:266
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRBuilder, ObjectRef, IRBuilderNode)
IRBuilder()
Creates an IRBuilder.
static void Name(ObjectRef node, String name)
NodeFunctor< void(const ObjectRef &, String)> FType
Definition: base.h:258
Base expr nodes in TVM.
Function nodes.
runtime implementation for LibTorch/TorchScript.
Definition: analyzer.h:36
constexpr runtime::NullOptType NullOpt
Definition: optional.h:169
Definitions and helper macros for IR/AST nodes.