Skip to content

Commit

Permalink
Streamline IR-building Process (apache#35)
Browse files Browse the repository at this point in the history
* upd

* upd

* Add more intrinsics

* add POC
  • Loading branch information
junrushao committed Jul 4, 2022
1 parent 8b9d720 commit 7ceb2c8
Show file tree
Hide file tree
Showing 14 changed files with 406 additions and 96 deletions.
33 changes: 30 additions & 3 deletions src/script/builder/builder.cc
Original file line number Diff line number Diff line change
Expand Up @@ -22,30 +22,57 @@ namespace tvm {
namespace script {
namespace builder {

Builder::Builder() {
ObjectPtr<BuilderNode> n = make_object<BuilderNode>();
n->frames.clear();
n->result = NullOpt;
data_ = n;
}

std::vector<Builder>* ThreadLocalBuilderStack() {
thread_local std::vector<Builder> stack;
return &stack;
}

void Builder::EnterWithScope() {
BuilderNode* n = this->get();
CHECK(n->frames.empty()) << "ValueError: There are frame(s) left in the builder: "
<< n->frames.size()
<< ". Please use a fresh new builder every time building IRs";
n->frames.push_back(IRModuleFrame());
std::vector<Builder>* stack = ThreadLocalBuilderStack();
stack->push_back(*this);
}

void Builder::ExitWithScope() {
BuilderNode* n = this->get();
ICHECK_EQ(n->frames.size(), 1);
IRModuleFrame frame = Downcast<IRModuleFrame>(n->frames.back());
n->frames.pop_back();
std::vector<Builder>* stack = ThreadLocalBuilderStack();
CHECK(!stack->empty());
ICHECK(!stack->empty());
stack->pop_back();
if (!frame->stmts.empty()) {
ICHECK(frame->global_vars.empty());
ICHECK(frame->functions.empty());
n->result = frame->stmts;
} else {
Map<GlobalVar, BaseFunc> func_map;
ICHECK_EQ(frame->functions.size(), frame->global_vars.size());
int m = frame->functions.size();
for (int i = 0; i < m; ++i) {
func_map.Set(frame->global_vars[i], frame->functions[i]);
}
}
}

Builder Builder::Current() {
std::vector<Builder>* stack = ThreadLocalBuilderStack();
CHECK(!stack->empty());
CHECK(!stack->empty()) << "ValueError: No builder in current scope";
return stack->back();
}

TVM_REGISTER_NODE_TYPE(BuilderNode);
TVM_REGISTER_NODE_TYPE(FrameNode);

} // namespace builder
} // namespace script
Expand Down
88 changes: 31 additions & 57 deletions src/script/builder/builder.h
Original file line number Diff line number Diff line change
Expand Up @@ -21,90 +21,64 @@

#include <tvm/node/node.h>

#include "./frame.h"

namespace tvm {
namespace script {
namespace builder {

class FrameNode : public runtime::Object {
public:
std::vector<runtime::TypedPackedFunc<void()>> callbacks;

void VisitAttrs(tvm::AttrVisitor* v) {
// `callbacks` is not visited.
}

void AddCallback(runtime::TypedPackedFunc<void()> callback) { callbacks.push_back(callback); }

static constexpr const char* _type_key = "script.builder.Frame";
TVM_DECLARE_BASE_OBJECT_INFO(FrameNode, runtime::Object);

public:
virtual void EnterWithScope() {}

virtual void ExitWithScope() {}

virtual ~FrameNode() {
for (auto it = callbacks.rbegin(); it != callbacks.rend(); ++it) {
(*it)();
}
}
};

class Frame : public runtime::ObjectRef {
public:
void EnterWithScope() {
ICHECK(data_ != nullptr);
static_cast<FrameNode*>(data_.get())->EnterWithScope();
}

void ExitWithScope() {
ICHECK(data_ != nullptr);
static_cast<FrameNode*>(data_.get())->ExitWithScope();
data_.reset();
}

TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Frame, ObjectRef, FrameNode);

protected:
Frame() = default;
};

class BuilderNode : public runtime::Object {
public:
runtime::Array<Frame> frames;
Optional<ObjectRef> result;

void VisitAttrs(tvm::AttrVisitor* v) {
v->Visit("frames", &frames); //
v->Visit("frames", &frames);
v->Visit("result", &result);
}

static constexpr const char* _type_key = "script.builder.Builder";
TVM_DECLARE_FINAL_OBJECT_INFO(BuilderNode, runtime::Object);

public:
template <typename TFrame>
Optional<TFrame> 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;
}
inline Optional<TFrame> FindFrame() const;

template <typename TObjectRef>
inline TObjectRef Get() const;
};

class Builder : public runtime::ObjectRef {
public:
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Builder, ObjectRef, BuilderNode);
Builder();
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Builder, ObjectRef, BuilderNode);

public:
void EnterWithScope();

void ExitWithScope();

static Builder Current();
};

template <typename TFrame>
inline Optional<TFrame> BuilderNode::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 TObjectRef>
inline TObjectRef BuilderNode::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) << "IndexError: IRBuilder result is not of type: " << TObject::_type_key;
return GetRef<TObjectRef>(n);
}

} // namespace builder
} // namespace script
} // namespace tvm
Expand Down
50 changes: 50 additions & 0 deletions src/script/builder/frame.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
/*
* 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.
*/
#include "./builder.h"

namespace tvm {
namespace script {
namespace builder {

void FrameNode::EnterWithScope() {
// Push to the current builder
Builder::Current()->frames.push_back(GetRef<Frame>(this));
}

void FrameNode::ExitWithScope() {
for (auto it = callbacks.rbegin(); it != callbacks.rend(); ++it) {
(*it)();
}
this->callbacks.clear();
Builder::Current()->frames.pop_back();
}

IRModuleFrame::IRModuleFrame() {
ObjectPtr<IRModuleFrameNode> n = make_object<IRModuleFrameNode>();
n->global_vars.clear();
n->functions.clear();
n->stmts.clear();
data_ = std::move(n);
}

TVM_REGISTER_NODE_TYPE(FrameNode);

} // namespace builder
} // namespace script
} // namespace tvm
97 changes: 97 additions & 0 deletions src/script/builder/frame.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
/*
* 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_BUILDER_FRAME_H_
#define TVM_SCRIPT_BUILDER_FRAME_H_

#include <tvm/ir/module.h>
#include <tvm/node/node.h>

namespace tvm {
namespace script {
namespace builder {

class FrameNode : public runtime::Object {
public:
std::vector<runtime::TypedPackedFunc<void()>> callbacks;

void VisitAttrs(tvm::AttrVisitor* v) {
// `callbacks` is not visited.
}

static constexpr const char* _type_key = "script.builder.Frame";
TVM_DECLARE_BASE_OBJECT_INFO(FrameNode, runtime::Object);

public:
virtual ~FrameNode() = default;
virtual void EnterWithScope();
virtual void ExitWithScope();
};

class Frame : public runtime::ObjectRef {
public:
virtual ~Frame() = default;
TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Frame, ObjectRef, FrameNode);

protected:
Frame() = default;

public:
inline void EnterWithScope();
inline void ExitWithScope();
};

class IRModuleFrameNode : public FrameNode {
public:
Array<GlobalVar> global_vars;
Array<BaseFunc> functions;
Array<ObjectRef> stmts;

void VisitAttrs(tvm::AttrVisitor* v) {
FrameNode::VisitAttrs(v);
v->Visit("global_vars", &global_vars);
v->Visit("functions", &functions);
v->Visit("stmts", &stmts);
}

static constexpr const char* _type_key = "script.builder.IRModuleFrame";
TVM_DECLARE_FINAL_OBJECT_INFO(IRModuleFrameNode, FrameNode);
};

class IRModuleFrame : public Frame {
public:
IRModuleFrame();
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(IRModuleFrame, Frame, IRModuleFrameNode);
};

inline void Frame::EnterWithScope() {
ICHECK(data_ != nullptr);
static_cast<FrameNode*>(data_.get())->EnterWithScope();
}

inline void Frame::ExitWithScope() {
ICHECK(data_ != nullptr);
static_cast<FrameNode*>(data_.get())->ExitWithScope();
data_.reset();
}

} // namespace builder
} // namespace script
} // namespace tvm

#endif // TVM_SCRIPT_BUILDER_FRAME_H_
62 changes: 62 additions & 0 deletions src/script/builder/tir/base.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* 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.
*/
#include "./base.h"

#include <tvm/support/with.h>

#include "./block_frame.h"
#include "./for_frame.h"
#include "./prim_func_frame.h"
#include "./var.h"

namespace tvm {
namespace script {
namespace builder {
namespace tir {

TVM_REGISTER_NODE_TYPE(TIRFrameNode);

void TestPOC() {
namespace T = tvm::script::builder::tir;
using namespace ::tvm::tir;

With<Builder> builder;
{
With<PrimFuncFrame> _{T::PrimFunc_("main")};
Buffer A = T::Buffer_({128, 128, 128}, DataType::Float(32));
Buffer B = T::Buffer_({128, 128, 128}, DataType::Float(32));
{
With<ForFrame> _{T::Grid({128, 128, 128})};
Var i = _()->vars[0];
Var j = _()->vars[1];
Var k = _()->vars[2];
{
With<BlockFrame> _{T::Block_("block")};
IterVar vi = T::axis::Spatial(Range(0, 128), i);
IterVar vj = T::axis::Spatial(Range(0, 128), j);
IterVar vk = T::axis::Spatial(Range(0, 128), k);
}
}
}
}

} // namespace tir
} // namespace builder
} // namespace script
} // namespace tvm
Loading

0 comments on commit 7ceb2c8

Please sign in to comment.