[AutoDiff] Automatically determine AdStack's size

taichi/backends/cc/codegen_cc.cpp

@@ -533,6 +533,9 @@ class CCTransformer : public IRVisitor {
 
   void visit(AdStackAllocaStmt *stmt) override {
     TI_ASSERT(stmt->width() == 1);
+    TI_ASSERT_INFO(
+        stmt->max_size > 0,
+        "Adaptive autodiff stack's size should have been determined.");
 
     const auto &var_name = stmt->raw_name();
     emit("Ti_u8 {}[{}];", var_name, stmt->size_in_bytes() + sizeof(uint32_t));

taichi/backends/metal/codegen_metal.cpp

@@ -109,10 +109,10 @@ class KernelCodegenImpl : public IRVisitor {
         kernel_(kernel),
         compiled_structs_(compiled_structs),
         needs_root_buffer_(compiled_structs_->root_size > 0),
-        ctx_attribs_(*kernel_),
         print_strtab_(print_strtab),
         cgen_config_(config),
-        offloaded_(offloaded) {
+        offloaded_(offloaded),
+        ctx_attribs_(*kernel_) {
     ti_kernel_attribs_.name = taichi_kernel_name;
     ti_kernel_attribs_.is_jit_evaluator = kernel->is_evaluator;
     // allow_undefined_visitor = true;
@@ -737,6 +737,9 @@ class KernelCodegenImpl : public IRVisitor {
 
   void visit(AdStackAllocaStmt *stmt) override {
     TI_ASSERT(stmt->width() == 1);
+    TI_ASSERT_INFO(
+        stmt->max_size > 0,
+        "Adaptive autodiff stack's size should have been determined.");
 
     const auto &var_name = stmt->raw_name();
     emit("byte {}[{}];", var_name, stmt->size_in_bytes());
@@ -1456,7 +1459,7 @@ class KernelCodegenImpl : public IRVisitor {
   }
 
   template <typename... Args>
-  void emit(std::string f, Args &&... args) {
+  void emit(std::string f, Args &&...args) {
     current_appender().append(std::move(f), std::forward<Args>(args)...);
   }
 

taichi/codegen/codegen_llvm.cpp

@@ -1864,6 +1864,8 @@ void CodeGenLLVM::visit(InternalFuncStmt *stmt) {
 
 void CodeGenLLVM::visit(AdStackAllocaStmt *stmt) {
   TI_ASSERT(stmt->width() == 1);
+  TI_ASSERT_INFO(stmt->max_size > 0,
+                 "Adaptive autodiff stack's size should have been determined.");
   auto type = llvm::ArrayType::get(llvm::Type::getInt8Ty(*llvm_context),
                                    stmt->size_in_bytes());
   auto alloca = create_entry_block_alloca(type, sizeof(int64));

taichi/ir/control_flow_graph.cpp

@@ -1,6 +1,7 @@
 #include "taichi/ir/control_flow_graph.h"
 
 #include <queue>
+#include <unordered_set>
 
 #include "taichi/ir/analysis.h"
 #include "taichi/ir/statements.h"
@@ -863,4 +864,103 @@ std::unordered_set<SNode *> ControlFlowGraph::gather_loaded_snodes() {
   return snodes;
 }
 
+void ControlFlowGraph::determine_ad_stack_size(int max_ad_stack_size) {
+  const int num_nodes = size();
+  std::unordered_map<AdStackAllocaStmt *, int> max_size;
+  std::vector<std::unordered_map<AdStackAllocaStmt *, int>>
+      max_size_at_node_begin(num_nodes);
+  std::vector<std::unordered_map<AdStackAllocaStmt *, int>> max_increased_size(
+      num_nodes);
+  std::vector<std::unordered_map<AdStackAllocaStmt *, int>> increased_size(
+      num_nodes);
+  std::queue<int> to_visit;
+  std::vector<bool> in_queue(num_nodes);
+  std::unordered_map<CFGNode *, int> node_ids;
+  std::unordered_set<AdStackAllocaStmt *> oversized_stacks;
+
+  for (int i = 0; i < num_nodes; i++)
+    node_ids[nodes[i].get()] = i;
+
+  for (int i = 0; i < num_nodes; i++) {
+    for (int j = nodes[i]->begin_location; j < nodes[i]->end_location; j++) {
+      Stmt *stmt = nodes[i]->block->statements[j].get();
+      if (auto *stack_push = stmt->cast<AdStackPushStmt>()) {
+        auto *stack = stack_push->stack->as<AdStackAllocaStmt>();
+        if (stack->max_size == 0 /*adaptive*/) {
+          increased_size[i][stack]++;
+          if (increased_size[i][stack] > max_increased_size[i][stack]) {
+            max_increased_size[i][stack] = increased_size[i][stack];
+          }
+        }
+      } else if (auto *stack_pop = stmt->cast<AdStackPopStmt>()) {
+        auto *stack = stack_pop->stack->as<AdStackAllocaStmt>();
+        if (stack->max_size == 0 /*adaptive*/) {
+          increased_size[i][stack]--;
+        }
+      }
+    }
+    to_visit.push(i);
+    in_queue[i] = true;
+  }
+
+  while (!to_visit.empty()) {
+    int node_id = to_visit.front();
+    to_visit.pop();
+    in_queue[node_id] = false;
+    CFGNode *now = nodes[node_id].get();
+
+    for (auto &it : max_increased_size[node_id]) {
+      auto *stack = it.first;
+      // Inside this CFGNode
+      auto current_max_size =
+          max_size_at_node_begin[node_id][stack] + it.second;
+      if (current_max_size > max_ad_stack_size) {
+        current_max_size = max_ad_stack_size;
+        oversized_stacks.insert(stack);
+      }
+      if (current_max_size > max_size[stack]) {
+        max_size[stack] = current_max_size;
+      }
+    }
+    for (auto &it : increased_size[node_id]) {
+      auto *stack = it.first;
+      // At the end of this CFGNode
+      auto current_size = max_size_at_node_begin[node_id][stack] + it.second;
+      if (current_size > max_ad_stack_size) {
+        current_size = max_ad_stack_size;  // avoid infinite loop
+      }
+      for (auto *next_node : now->next) {
+        int next_node_id = node_ids[next_node];
+        if (current_size > max_size_at_node_begin[next_node_id][stack]) {
+          max_size_at_node_begin[next_node_id][stack] = current_size;
+          if (!in_queue[next_node_id]) {
+            to_visit.push(next_node_id);
+            in_queue[next_node_id] = true;
+          }
+        }
+      }
+    }
+  }
+
+  if (!oversized_stacks.empty()) {
+    std::vector<std::string> oversized_stacks_name;
+    oversized_stacks_name.reserve(oversized_stacks.size());
+    for (auto &stack : oversized_stacks) {
+      oversized_stacks_name.push_back(stack->name());
+    }
+    TI_WARN(
+        "Unable to determine capacity for autodiff stacks: {}. "
+        "Use default capacity {} instead.",
+        fmt::join(oversized_stacks_name, ", "), max_ad_stack_size);
+  }
+
+  for (auto &it : max_size) {
+    auto *stack = it.first;
+    TI_WARN_IF(it.second == 0,
+               "Unused autodiff stack {} should have been eliminated.",
+               stack->name());
+    stack->max_size = 16;
+  }
+}
+
 TLANG_NAMESPACE_END

taichi/ir/control_flow_graph.h

@@ -121,6 +121,14 @@ class ControlFlowGraph {
 
   // Gather the SNodes this offload reads.
   std::unordered_set<SNode *> gather_loaded_snodes();
+
+  /**
+   * Determine all adaptive AD-stacks' capacity with the worklist algorithm.
+   * @param max_ad_stack_size
+   * The maximum allowed AD stack size. This parameter is set to prevent
+   * infinite loops of the algorithm.
+   */
+  void determine_ad_stack_size(int max_ad_stack_size);
 };
 
 TLANG_NAMESPACE_END

taichi/ir/statements.h

@@ -1222,7 +1222,7 @@ class InternalFuncStmt : public Stmt {
 class AdStackAllocaStmt : public Stmt {
  public:
   DataType dt;
-  std::size_t max_size;  // TODO: 0 = adaptive
+  std::size_t max_size;  // 0 = adaptive
 
   AdStackAllocaStmt(const DataType &dt, std::size_t max_size)
       : dt(dt), max_size(max_size) {

taichi/ir/transforms.h

@@ -71,6 +71,7 @@ bool lower_access(IRNode *root,
 void auto_diff(IRNode *root,
                const CompileConfig &config,
                bool use_stack = false);
+bool determine_ad_stack_size(IRNode *root, const CompileConfig &config);
 bool constant_fold(IRNode *root,
                    const CompileConfig &config,
                    const ConstantFoldPass::Args &args);
@@ -124,6 +125,7 @@ void offload_to_executable(IRNode *ir,
                            const CompileConfig &config,
                            Kernel *kernel,
                            bool verbose,
+                           bool determine_ad_stack_size,
                            bool lower_global_access,
                            bool make_thread_local,
                            bool make_block_local);

taichi/program/async_engine.cpp

@@ -154,6 +154,7 @@ void ExecutionQueue::enqueue(const TaskLaunchRecord &ker) {
           auto ir = stmt;
           offload_to_executable(
               ir, config, kernel, /*verbose=*/false,
+              /*determine_ad_stack_size=*/true,
               /*lower_global_access=*/true,
               /*make_thread_local=*/true,
               /*make_block_local=*/

taichi/program/compile_config.cpp

@@ -45,7 +45,7 @@ CompileConfig::CompileConfig() {
   cpu_max_num_threads = std::thread::hardware_concurrency();
   random_seed = 0;
 
-  ad_stack_size = 16;
+  ad_stack_size = 0;  // 0 = adaptive
 
   // LLVM backend options:
   print_struct_llvm_ir = false;

taichi/transforms/compile_to_offloads.cpp

@@ -145,6 +145,7 @@ void offload_to_executable(IRNode *ir,
                            const CompileConfig &config,
                            Kernel *kernel,
                            bool verbose,
+                           bool determine_ad_stack_size,
                            bool lower_global_access,
                            bool make_thread_local,
                            bool make_block_local) {
@@ -224,6 +225,11 @@ void offload_to_executable(IRNode *ir,
   irpass::full_simplify(ir, config, {lower_global_access, kernel->program});
   print("Simplified IV");
 
+  if (determine_ad_stack_size) {
+    irpass::determine_ad_stack_size(ir, config);
+    print("Autodiff stack size determined");
+  }
+
   if (is_extension_supported(config.arch, Extension::quant)) {
     irpass::optimize_bit_struct_stores(ir, config, amgr.get());
     print("Bit struct stores optimized");
@@ -250,8 +256,10 @@ void compile_to_executable(IRNode *ir,
   compile_to_offloads(ir, config, kernel, verbose, vectorize, grad,
                       ad_use_stack, start_from_ast);
 
-  offload_to_executable(ir, config, kernel, verbose, lower_global_access,
-                        make_thread_local, make_block_local);
+  offload_to_executable(ir, config, kernel, verbose,
+                        /*determine_ad_stack_size=*/grad && ad_use_stack,
+                        lower_global_access, make_thread_local,
+                        make_block_local);
 }
 
 void compile_inline_function(IRNode *ir,

taichi/transforms/determine_ad_stack_size.cpp

@@ -0,0 +1,34 @@
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/control_flow_graph.h"
+#include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
+#include "taichi/ir/transforms.h"
+
+#include <queue>
+#include <unordered_map>
+
+namespace taichi {
+namespace lang {
+
+namespace irpass {
+
+bool determine_ad_stack_size(IRNode *root, const CompileConfig &config) {
+  if (irpass::analysis::gather_statements(root, [&](Stmt *s) {
+        if (auto ad_stack = s->cast<AdStackAllocaStmt>()) {
+          return ad_stack->max_size == 0;  // adaptive
+        }
+        return false;
+      }).empty()) {
+    return false;  // no AD-stacks with adaptive size
+  }
+  auto cfg = analysis::build_cfg(root);
+  cfg->simplify_graph();
+  constexpr int kMaxAdStackSize = 32;
+  cfg->determine_ad_stack_size(kMaxAdStackSize);
+  return true;
+}
+
+}  // namespace irpass
+
+}  // namespace lang
+}  // namespace taichi