Streamline IR-building Process (apache#35)

* upd

* upd

* Add more intrinsics

* add POC

src/script/builder/builder.cc

@@ -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

src/script/builder/builder.h

@@ -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

src/script/builder/frame.cc

@@ -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

src/script/builder/frame.h

@@ -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_

src/script/builder/tir/base.cc

@@ -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