[opt] [async] Improve full_simplify and optimize_dead_store (#2160)

taichi/program/ir_bank.cpp

@@ -198,15 +198,21 @@ std::pair<IRHandle, bool> IRBank::optimize_dse(
 
   if (verbose) {
     TI_INFO("  DSE: before CFG");
+    std::cout << std::flush;
+    for (auto &s : snodes)
+      std::cout << s->get_node_type_name_hinted() << std::endl;
     irpass::print(new_ir.get());
+    std::cout << std::flush;
   }
   ControlFlowGraph::LiveVarAnalysisConfig lva_config;
   lva_config.eliminable_snodes = {snodes.begin(), snodes.end()};
   const bool modified = irpass::cfg_optimization(
       new_ir.get(), /*after_lower_access=*/false, lva_config);
   if (verbose) {
     TI_INFO("  DSE: after CFG, modified={}", modified);
+    std::cout << std::flush;
     irpass::print(new_ir.get());
+    std::cout << std::flush;
   }
 
   if (!modified) {
@@ -215,6 +221,10 @@ std::pair<IRHandle, bool> IRBank::optimize_dse(
     return std::make_pair(ret_handle, false);
   }
 
+  // Remove unused global pointers (possibly with activate == true).
+  irpass::flag_access(new_ir.get());
+  irpass::die(new_ir.get());
+
   ret_handle = IRHandle(new_ir.get(), get_hash(new_ir.get()));
   insert(std::move(new_ir), ret_handle.hash());
   return std::make_pair(ret_handle, false);

taichi/program/state_flow_graph.cpp

@@ -1257,18 +1257,20 @@ bool StateFlowGraph::optimize_dead_store() {
         continue;
       }
       auto *snode = s.snode();
-      if (!snode->is_scalar()) {
-        // TODO: handle non-scalar SNodes, i.e. num_active_indices > 0.
-        continue;
-      }
       bool used = false;
       for (auto &other : task->output_edges[s]) {
         if (task->has_state_flow(s, other.second)) {
-          // Check if this is a RAW dependency. For scalar SNodes, a WAW flow
-          // edge decades to a dependency edge.
+          // Check if this is a RAW dependency. For scalar SNodes, a
+          // A WAW flow edge decades to a dependency edge.
           used = true;
         } else {
           // Note that a dependency edge does not count as an data usage
+          if (!snode->is_scalar() &&
+              other.second->meta->element_wise.count(snode) == 0) {
+            // If the SNode is not element-wise written, the value stored
+            // in this task can be used later.
+            used = true;
+          }
         }
       }
       // This state is used by some other node, so it cannot be erased
@@ -1325,8 +1327,8 @@ bool StateFlowGraph::optimize_dead_store() {
     }
   }
 
-  if (!to_delete.empty()) {
-    modified = true;
+  // Task metas can change even if the number of tasks doesn't change.
+  if (modified) {
     delete_nodes(to_delete);
     rebuild_graph(/*sort=*/false);
   }

taichi/transforms/simplify.cpp

@@ -609,10 +609,10 @@ void full_simplify(IRNode *root, bool after_lower_access, Kernel *kernel) {
         modified = true;
       if (die(root))
         modified = true;
-      // Don't do these time-consuming optimization passes again if the IR is
-      // not modified.
-      if ((first_iteration || modified) && whole_kernel_cse(root))
+      if (whole_kernel_cse(root))
         modified = true;
+      // Don't do this time-consuming optimization pass again if the IR is
+      // not modified.
       if ((first_iteration || modified) &&
           root->get_config().cfg_optimization &&
           cfg_optimization(root, after_lower_access))