[misc] Add "make_block_local" option to "CompileConfig"

taichi/program/async_engine.cpp

@@ -179,7 +179,8 @@ void ExecutionQueue::enqueue(KernelLaunchRecord &&ker) {
             /*lower_global_access=*/true,
             /*make_thread_local=*/true,
             /*make_block_local=*/
-            is_extension_supported(config.arch, Extension::bls));
+            is_extension_supported(config.arch, Extension::bls) &&
+                config.make_block_local);
       }
       auto codegen = KernelCodeGen::create(kernel->arch, kernel, stmt);
       auto func = codegen->codegen();

taichi/program/compile_config.cpp

@@ -32,6 +32,7 @@ CompileConfig::CompileConfig() {
   async_mode = false;
   flatten_if = false;
   make_thread_local = true;
+  make_block_local = true;
 
   saturating_grid_dim = 0;
   max_block_dim = 0;

taichi/program/compile_config.h

@@ -31,6 +31,7 @@ struct CompileConfig {
   bool async_mode;
   bool flatten_if;
   bool make_thread_local;
+  bool make_block_local;
   DataType default_fp;
   DataType default_ip;
   std::string extra_flags;

taichi/program/kernel.cpp

@@ -8,6 +8,7 @@
 #include "taichi/backends/cuda/cuda_driver.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/util/action_recorder.h"
+#include "taichi/program/extension.h"
 
 TLANG_NAMESPACE_BEGIN
 
@@ -81,7 +82,10 @@ void Kernel::lower(bool to_executable) {  // TODO: is a "Lowerer" class
       irpass::compile_to_executable(
           ir.get(), config, /*vectorize*/ arch_is_cpu(arch), grad,
           /*ad_use_stack*/ true, verbose, /*lower_global_access*/ to_executable,
-          /*make_thread_local*/ true, /*make_block_local*/ arch == Arch::cuda);
+          /*make_thread_local*/ config.make_thread_local,
+          /*make_block_local*/
+          is_extension_supported(config.arch, Extension::bls) &&
+              config.make_block_local);
     } else {
       irpass::compile_to_offloads(ir.get(), config, verbose,
                                   /*vectorize=*/arch_is_cpu(arch), grad,

taichi/python/export_lang.cpp

@@ -121,6 +121,7 @@ void export_lang(py::module &m) {
       .def_readwrite("async_mode", &CompileConfig::async_mode)
       .def_readwrite("flatten_if", &CompileConfig::flatten_if)
       .def_readwrite("make_thread_local", &CompileConfig::make_thread_local)
+      .def_readwrite("make_block_local", &CompileConfig::make_block_local)
       .def_readwrite("cc_compile_cmd", &CompileConfig::cc_compile_cmd)
       .def_readwrite("cc_link_cmd", &CompileConfig::cc_link_cmd);
 

taichi/transforms/constant_fold.cpp

@@ -1,12 +1,12 @@
-#include "taichi/ir/ir.h"
-#include "taichi/ir/transforms.h"
-#include "taichi/ir/visitors.h"
-#include "taichi/program/program.h"
 #include <deque>
 #include <set>
 #include <cmath>
 #include <thread>
 
+#include "taichi/ir/ir.h"
+#include "taichi/ir/transforms.h"
+#include "taichi/ir/visitors.h"
+#include "taichi/program/program.h"
 #include "taichi/ir/ir.h"
 #include "taichi/program/program.h"
 #include "taichi/ir/snode.h"

taichi/transforms/insert_scratch_pad.cpp

@@ -6,7 +6,7 @@
 
 TLANG_NAMESPACE_BEGIN
 
-// Figure out accessed snodes, and their ranges in this for stmt
+// Figure out accessed SNodes, and their ranges in this for stmt
 class AccessAnalysis : public BasicStmtVisitor {
   using BasicStmtVisitor::visit;
 
@@ -59,14 +59,6 @@ class AccessAnalysis : public BasicStmtVisitor {
     auto ptr = stmt->as<GlobalPtrStmt>();
     for (int l = 0; l < stmt->width(); l++) {
       auto snode = ptr->snodes[l];
-      // std::vector<SNode *> snodes;
-      /*
-      for (auto it: pads->pads) {
-        //snodes.push_back(it.first);
-        TI_P(it.first->node_type_name);
-      }
-      TI_P(snode->node_type_name);
-      */
       if (!pads->has(snode)) {
         continue;
       }
@@ -80,27 +72,11 @@ class AccessAnalysis : public BasicStmtVisitor {
         if (diff.linear_related()) {
           offsets[i].first = diff.low;
           offsets[i].second = diff.high;
-          /*
-          TI_P(ptr->name());
-          TI_P(diff.low);
-          TI_P(diff.high);
-          */
         } else {
-          /*
-          TI_P(i);
-          TI_P(for_stmt->loop_vars[i]->raw_name());
-          TI_P(ptr->indices[i]->raw_name());
-          */
           matching_indices = false;
         }
       }
       if (matching_indices) {
-        /*
-        TI_INFO("Detected regular access");
-        for (int i = 0; i < num_indices; i++) {
-          TI_P(offsets[i]);
-        }
-        */
         for (const auto &bind : block_indices) {
           std::function<void(std::vector<int>, int)> visit =
               [&](std::vector<int> ind, int depth) {