From 6779e0dc312f3d955da7254e1eddc384d8ec4130 Mon Sep 17 00:00:00 2001
From: Yuanming Hu <yuanming-hu@users.noreply.github.com>
Date: Sat, 23 Jan 2021 22:43:14 -0500
Subject: [PATCH] [async] Add option async_opt_fusion_max_iter (#2170)

---
 taichi/program/async_engine.cpp | 11 ++++++++---
 taichi/program/compile_config.h |  1 +
 taichi/python/export_lang.cpp   |  2 ++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/taichi/program/async_engine.cpp b/taichi/program/async_engine.cpp
index 931a7bf53c007..8a168aebdb1ff 100644
--- a/taichi/program/async_engine.cpp
+++ b/taichi/program/async_engine.cpp
@@ -279,9 +279,14 @@ void AsyncEngine::flush() {
     }
     sfg->verify();
     if (program->config.async_opt_fusion) {
-      while (sfg->fuse()) {
-        debug_sfg("fuse");
-        modified = true;
+      auto max_iter = program->config.async_opt_fusion_max_iter;
+      for (int iter = 0; max_iter == 0 || iter < max_iter; iter++) {
+        if (sfg->fuse()) {
+          debug_sfg("fuse");
+          modified = true;
+        } else {
+          break;
+        }
       }
     }
     sfg->verify();
diff --git a/taichi/program/compile_config.h b/taichi/program/compile_config.h
index d978de3a9f677..be42f80218362 100644
--- a/taichi/program/compile_config.h
+++ b/taichi/program/compile_config.h
@@ -65,6 +65,7 @@ struct CompileConfig {
   // Async options
   int async_opt_passes{3};
   bool async_opt_fusion{true};
+  int async_opt_fusion_max_iter{0};  // 0 means unlimited
   bool async_opt_listgen{true};
   bool async_opt_activation_demotion{true};
   bool async_opt_dse{true};
diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
index 27c7e17d54339..9ea4312ef0ee8 100644
--- a/taichi/python/export_lang.cpp
+++ b/taichi/python/export_lang.cpp
@@ -168,6 +168,8 @@ void export_lang(py::module &m) {
       .def_readwrite("cc_link_cmd", &CompileConfig::cc_link_cmd)
       .def_readwrite("async_opt_passes", &CompileConfig::async_opt_passes)
       .def_readwrite("async_opt_fusion", &CompileConfig::async_opt_fusion)
+      .def_readwrite("async_opt_fusion_max_iter",
+                     &CompileConfig::async_opt_fusion_max_iter)
       .def_readwrite("async_opt_listgen", &CompileConfig::async_opt_listgen)
       .def_readwrite("async_opt_activation_demotion",
                      &CompileConfig::async_opt_activation_demotion)