From f49530e0a4208c923c528600478cdd6be0053f57 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 19 May 2020 11:36:30 -0700 Subject: [PATCH 1/6] Make let processing iterative --- python/tvm/relay/transform/memory_plan.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/python/tvm/relay/transform/memory_plan.py b/python/tvm/relay/transform/memory_plan.py index e359a9e3def0..d24c9530e3a0 100644 --- a/python/tvm/relay/transform/memory_plan.py +++ b/python/tvm/relay/transform/memory_plan.py @@ -330,6 +330,18 @@ def visit_function(self, fn): fn.type_params, fn.attrs) + def visit_let(self, let): + def _each_binding(var, val): + new_var = self.visit(let.var) + new_val = self.visit(let.value) + return new_var, new_val + + def _after(bindings, body): + new_body = self.visit(let.body) + return mk_let(bindings, new_body) + + return iterative_let(let, _each_binding, _after) + @function_pass(opt_level=0) class MemoryPlan: """An explicit pass wrapper around StorageCoalesce.""" From 05cdd954cce1b112a894ed61e7b21f75ad17297f Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 19 May 2020 11:57:52 -0700 Subject: [PATCH 2/6] Try again --- python/tvm/relay/transform/memory_plan.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/python/tvm/relay/transform/memory_plan.py b/python/tvm/relay/transform/memory_plan.py index d24c9530e3a0..853dd29d0aa1 100644 --- a/python/tvm/relay/transform/memory_plan.py +++ b/python/tvm/relay/transform/memory_plan.py @@ -331,16 +331,18 @@ def visit_function(self, fn): fn.attrs) def visit_let(self, let): - def _each_binding(var, val): + bindings = [] + i = 0 + while isinstance(let, expr.Let): + print(i) new_var = self.visit(let.var) new_val = self.visit(let.value) - return new_var, new_val + bindings.append((new_var, new_val)) + let = let.body + i += 1 - def _after(bindings, body): - new_body = self.visit(let.body) - return mk_let(bindings, new_body) - - return iterative_let(let, _each_binding, _after) + new_body = self.visit(let) + return mk_let(bindings, new_body) @function_pass(opt_level=0) class MemoryPlan: From 03cd8e8aa9ba097fc4afd10ad37ae493043e8527 Mon Sep 17 00:00:00 2001 From: Jared Roesch Date: Tue, 19 May 2020 22:50:12 -0700 Subject: [PATCH 3/6] Fix pretty printer overflow --- src/printer/relay_text_printer.cc | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc index 981d0c357e24..ee95c210ecf8 100644 --- a/src/printer/relay_text_printer.cc +++ b/src/printer/relay_text_printer.cc @@ -365,11 +365,19 @@ Doc RelayTextPrinter::VisitExpr_(const IfNode* op) { Doc RelayTextPrinter::VisitExpr_(const LetNode* op) { Doc doc; - doc << "let " << AllocVar(op->var) << " = " << Print(op->value, false, true) << ";" - << Doc::NewLine(); + Expr let = GetRef(op); + while (auto let_node = let.as()) { + doc << "let " + << AllocVar(let_node->var) + << " = " + << Print(let_node->value, false, true) + << ";" + << Doc::NewLine(); + let = let_node->body; + } // we use a scope here so GNF hoisting doesn't escape too far // and nested, unique lets are not hoisted - doc << PrintScope(op->body); + doc << PrintScope(let); return doc; } From 504092e51a912aa8af9fbd88aa11cc138e811063 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Tue, 9 Jun 2020 19:11:16 +0000 Subject: [PATCH 4/6] cleanup --- python/tvm/relay/transform/memory_plan.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/python/tvm/relay/transform/memory_plan.py b/python/tvm/relay/transform/memory_plan.py index 853dd29d0aa1..8f21af9292a9 100644 --- a/python/tvm/relay/transform/memory_plan.py +++ b/python/tvm/relay/transform/memory_plan.py @@ -332,14 +332,11 @@ def visit_function(self, fn): def visit_let(self, let): bindings = [] - i = 0 while isinstance(let, expr.Let): - print(i) new_var = self.visit(let.var) new_val = self.visit(let.value) bindings.append((new_var, new_val)) let = let.body - i += 1 new_body = self.visit(let) return mk_let(bindings, new_body) From 06ae06d7e2e0167f67f1f7c3708974646a2f0b76 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Wed, 10 Jun 2020 21:41:01 +0000 Subject: [PATCH 5/6] fix lint --- src/printer/relay_text_printer.cc | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc index ee95c210ecf8..c771626f4e40 100644 --- a/src/printer/relay_text_printer.cc +++ b/src/printer/relay_text_printer.cc @@ -367,11 +367,7 @@ Doc RelayTextPrinter::VisitExpr_(const LetNode* op) { Doc doc; Expr let = GetRef(op); while (auto let_node = let.as()) { - doc << "let " - << AllocVar(let_node->var) - << " = " - << Print(let_node->value, false, true) - << ";" + doc << "let " << AllocVar(let_node->var) << " = " << Print(let_node->value, false, true) << ";" << Doc::NewLine(); let = let_node->body; } From 2411e2b6860b7dc53b8e27a33166de649adcccd6 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Thu, 11 Jun 2020 01:13:47 +0000 Subject: [PATCH 6/6] Fix text printer --- src/printer/relay_text_printer.cc | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc index c771626f4e40..a09e24b12429 100644 --- a/src/printer/relay_text_printer.cc +++ b/src/printer/relay_text_printer.cc @@ -364,16 +364,21 @@ Doc RelayTextPrinter::VisitExpr_(const IfNode* op) { } Doc RelayTextPrinter::VisitExpr_(const LetNode* op) { - Doc doc; + int n = 0; Expr let = GetRef(op); while (auto let_node = let.as()) { + Doc doc; doc << "let " << AllocVar(let_node->var) << " = " << Print(let_node->value, false, true) << ";" << Doc::NewLine(); + doc_stack_.push_back(doc); let = let_node->body; + ++n; + } + Doc doc = PrintScope(let); + for (int i = 0; i < n; ++i) { + doc = doc_stack_.back() << doc; + doc_stack_.pop_back(); } - // we use a scope here so GNF hoisting doesn't escape too far - // and nested, unique lets are not hoisted - doc << PrintScope(let); return doc; }