[tvm] branch main updated: tir printing (#10805)

[masahi]

This is an automated email from the ASF dual-hosted git repository.

masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 9b5bbf8  tir printing (#10805)
9b5bbf8 is described below

commit 9b5bbf87bf108bdcb8624a64a51e66d4f431ad14
Author: AndrewZhaoLuo <andrew.zhao....@gmail.com>
AuthorDate: Mon Mar 28 19:51:13 2022 -0700

    tir printing (#10805)
---
 src/printer/tir_text_printer.cc | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/printer/tir_text_printer.cc b/src/printer/tir_text_printer.cc
index 16d4772..1ef62c2 100644
--- a/src/printer/tir_text_printer.cc
+++ b/src/printer/tir_text_printer.cc
@@ -91,9 +91,19 @@ Doc TIRTextPrinter::PrintPrimFunc(const PrimFunc& prim_func) 
{
   // collect buffers in buffer_map
   memo_var_.clear();
   memo_buf_.clear();
-  for (const auto& it : op->buffer_map) {
-    memo_buf_[it.second] = AllocBuf(it.second);
+
+  // ordered vars associated with buffers, for consistent printing
+  std::vector<Var> buffer_vars_ordered;
+
+  for (Var v : op->params) {
+    auto buffer_map_find = op->buffer_map.find(v);
+    if (buffer_map_find != op->buffer_map.end()) {
+      auto map_data = *buffer_map_find;
+      buffer_vars_ordered.push_back(map_data.first);
+      memo_buf_[map_data.second] = AllocBuf(map_data.second);
+    }
   }
+
   // print PrimFunc
   Doc doc;
   doc << "primfn"
@@ -122,8 +132,8 @@ Doc TIRTextPrinter::PrintPrimFunc(const PrimFunc& 
prim_func) {
   if (memo_buf_.size() != 0) {
     Doc buffer_doc;
     std::vector<Doc> buffer_docs;
-    for (const auto& it : memo_buf_) {
-      const auto& buf = it.first;
+    for (const Var& v : buffer_vars_ordered) {
+      const Buffer buf = op->buffer_map[v];
       buffer_docs.push_back(BufferNode2Doc(buf.get(), Print(buf)));
     }
     buffer_doc << Doc::NewLine() << "buffers = {";
@@ -134,8 +144,9 @@ Doc TIRTextPrinter::PrintPrimFunc(const PrimFunc& 
prim_func) {
   if (op->buffer_map.size() != 0) {
     // print buffer_map
     std::vector<Doc> buffer_map_doc;
-    for (const auto& it : op->buffer_map) {
-      buffer_map_doc.push_back(Print(it.first) << ": " << Print(it.second));
+    for (const Var& v : buffer_vars_ordered) {
+      const Buffer buf = op->buffer_map[v];
+      buffer_map_doc.push_back(Print(v) << ": " << Print(buf));
     }
     doc << Doc::Indent(
         2, Doc::NewLine() << "buffer_map = {" << PrintSep(buffer_map_doc, 
Doc::Text(", ")) << "}");