[IR] fix Yield (#214)

bytedance · Mar 29, 2023 · 6e0dcf7 · 6e0dcf7
1 parent 1cc6296
commit 6e0dcf7
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 2 deletions.
diff --git a/include/matxscript/ir/printer/doc.h b/include/matxscript/ir/printer/doc.h
@@ -945,6 +945,38 @@ class SliceDoc : public Doc {
   MATXSCRIPT_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SliceDoc, Doc, SliceDocNode);
 };
 
+/*!
+ * \brief Doc that represents yield statement.
+ *
+ * \sa YieldDoc
+ */
+class YieldDocNode : public ExprDocNode {
+ public:
+  Optional<ExprDoc> value;
+
+  void VisitAttrs(AttrVisitor* v) {
+    ExprDocNode::VisitAttrs(v);
+    v->Visit("value", &value);
+  }
+
+  static constexpr const char* _type_key = "ir.printer.YieldDoc";
+  MATXSCRIPT_DECLARE_FINAL_OBJECT_INFO(YieldDocNode, ExprDocNode);
+};
+
+/*!
+ * \brief Reference type of YieldDocNode.
+ *
+ * \sa YieldDocNode
+ */
+class YieldDoc : public ExprDoc {
+ public:
+  /*!
+   * \brief Constructor of YieldDoc.
+   */
+  explicit YieldDoc(Optional<ExprDoc> value);
+  MATXSCRIPT_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(YieldDoc, ExprDoc, YieldDocNode);
+};
+
 /*!
  * \brief Doc that represents assign statement.
  *

diff --git a/src/ir/printer/doc.cc b/src/ir/printer/doc.cc
@@ -178,6 +178,12 @@ SliceDoc::SliceDoc(Optional<ExprDoc> start, Optional<ExprDoc> stop, Optional<Exp
   this->data_ = std::move(n);
 }
 
+YieldDoc::YieldDoc(Optional<ExprDoc> value) {
+  ObjectPtr<YieldDocNode> n = make_object<YieldDocNode>();
+  n->value = std::move(value);
+  this->data_ = std::move(n);
+}
+
 AssignDoc::AssignDoc(ExprDoc lhs, Optional<ExprDoc> rhs, Optional<ExprDoc> annotation) {
   MXCHECK(rhs.defined() || annotation.defined())
       << "ValueError: At least one of rhs and annotation needs to be non-null for AssignDoc.";
@@ -463,6 +469,11 @@ MATXSCRIPT_REGISTER_GLOBAL("ir.printer.SliceDoc")
       return SliceDoc(std::move(start), std::move(stop), std::move(step));
     });
 
+MATXSCRIPT_REGISTER_NODE_TYPE(YieldDocNode);
+MATXSCRIPT_REGISTER_GLOBAL("ir.printer.YieldDoc").set_body_typed([](Optional<ExprDoc> value) {
+  return YieldDoc(std::move(value));
+});
+
 MATXSCRIPT_REGISTER_NODE_TYPE(AssignDocNode);
 MATXSCRIPT_REGISTER_GLOBAL("ir.printer.AssignDoc")
     .set_body_typed([](ExprDoc lhs, Optional<ExprDoc> rhs, Optional<ExprDoc> annotation) {

diff --git a/src/ir/printer/doc_printer/base_doc_printer.cc b/src/ir/printer/doc_printer/base_doc_printer.cc
@@ -357,6 +357,8 @@ void DocPrinter::PrintDoc(const Doc& doc) {
     PrintTypedDoc(GetRef<DictCompDoc>(doc_node));
   } else if (const auto* doc_node = doc.as<SliceDocNode>()) {
     PrintTypedDoc(GetRef<SliceDoc>(doc_node));
+  } else if (const auto* doc_node = doc.as<YieldDocNode>()) {
+    PrintTypedDoc(GetRef<YieldDoc>(doc_node));
   } else if (const auto* doc_node = doc.as<StmtBlockDocNode>()) {
     PrintTypedDoc(GetRef<StmtBlockDoc>(doc_node));
   } else if (const auto* doc_node = doc.as<AssignDocNode>()) {

diff --git a/src/ir/printer/doc_printer/base_doc_printer.h b/src/ir/printer/doc_printer/base_doc_printer.h
@@ -177,6 +177,11 @@ class DocPrinter {
    */
   virtual void PrintTypedDoc(const SliceDoc& doc) = 0;
 
+  /*!
+   * \brief Virtual method to print a YieldDoc
+   */
+  virtual void PrintTypedDoc(const YieldDoc& doc) = 0;
+
   /*!
    * \brief Virtual method to print a StmtBlockDoc
    */

diff --git a/src/ir/printer/doc_printer/python_doc_printer.cc b/src/ir/printer/doc_printer/python_doc_printer.cc
@@ -171,6 +171,7 @@ class PythonDocPrinter : public DocPrinter {
   void PrintTypedDoc(const SetCompDoc& doc) final;
   void PrintTypedDoc(const DictCompDoc& doc) final;
   void PrintTypedDoc(const SliceDoc& doc) final;
+  void PrintTypedDoc(const YieldDoc& doc) final;
   void PrintTypedDoc(const StmtBlockDoc& doc) final;
   void PrintTypedDoc(const AssignDoc& doc) final;
   void PrintTypedDoc(const IfDoc& doc) final;
@@ -598,6 +599,14 @@ void PythonDocPrinter::PrintTypedDoc(const SliceDoc& doc) {
   }
 }
 
+void PythonDocPrinter::PrintTypedDoc(const YieldDoc& doc) {
+  output_ << "yield";
+  if (doc->value.defined()) {
+    output_ << " ";
+    PrintDoc(doc->value.value());
+  }
+}
+
 void PythonDocPrinter::PrintTypedDoc(const StmtBlockDoc& doc) {
   for (const StmtDoc& stmt : doc->stmts) {
     PrintDoc(stmt);

diff --git a/src/ir/stmt.cc b/src/ir/stmt.cc
@@ -900,8 +900,8 @@ MATXSCRIPT_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
     .set_dispatch<ir::HLOYield>(  //
         "",
         [](ir::HLOYield stmt, ObjectPath p, IRDocsifier d) -> Doc {
-          // TODO: support HLOYieldDoc;
-          return CommentDoc("TODO");
+          ExprDoc value = d->AsDoc<ExprDoc>(stmt->symbol, p->Attr("symbol"));
+          return ExprStmtDoc(YieldDoc(value));
         });
 
 }  // namespace ir