Refactor2023: Remove dependencies on Program::this_thread_config() in…

… irpass::scalarize
taichi-dev · Dec 7, 2022 · 40d4b33 · 40d4b33
1 parent b8deaf0
commit 40d4b33
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/taichi/ir/transforms.h b/taichi/ir/transforms.h
@@ -29,7 +29,7 @@ namespace irpass {
 
 void re_id(IRNode *root);
 void flag_access(IRNode *root);
-void scalarize(IRNode *root);
+void scalarize(IRNode *root, bool dynamic_index);
 void lower_matrix_ptr(IRNode *root);
 bool die(IRNode *root);
 bool simplify(IRNode *root, const CompileConfig &config);

diff --git a/taichi/transforms/compile_to_offloads.cpp b/taichi/transforms/compile_to_offloads.cpp
@@ -61,7 +61,7 @@ void compile_to_offloads(IRNode *ir,
   }
 
   if (config.real_matrix && config.real_matrix_scalarize) {
-    irpass::scalarize(ir);
+    irpass::scalarize(ir, config.dynamic_index);
 
     // Remove redundant MatrixInitStmt inserted during scalarization
     irpass::die(ir);

diff --git a/taichi/transforms/scalarize.cpp b/taichi/transforms/scalarize.cpp
@@ -612,11 +612,10 @@ class ScalarizePointers : public BasicStmtVisitor {
 };
 
 namespace irpass {
-
-void scalarize(IRNode *root) {
+void scalarize(IRNode *root, bool dynamic_index) {
   TI_AUTO_PROF;
   Scalarize scalarize_pass(root);
-  if (!root->get_kernel()->program->this_thread_config().dynamic_index) {
+  if (!dynamic_index) {
     ScalarizePointers scalarize_pointers_pass(root);
   }
 }

diff --git a/tests/cpp/transforms/scalarize_test.cpp b/tests/cpp/transforms/scalarize_test.cpp
@@ -45,7 +45,8 @@ TEST(Scalarize, ScalarizeGlobalStore) {
 
   block->push_back<GlobalStoreStmt>(dest_stmt, matrix_init_stmt);
 
-  irpass::scalarize(block.get());
+  irpass::scalarize(block.get(),
+                    test_prog.prog()->this_thread_config().dynamic_index);
   irpass::lower_matrix_ptr(block.get());
   irpass::die(block.get());
 
@@ -102,7 +103,8 @@ TEST(Scalarize, ScalarizeGlobalLoad) {
   // Without this GlobalStoreStmt, nothing survives irpass::die()
   block->push_back<GlobalStoreStmt>(src_stmt, load_stmt);
 
-  irpass::scalarize(block.get());
+  irpass::scalarize(block.get(),
+                    test_prog.prog()->this_thread_config().dynamic_index);
   irpass::lower_matrix_ptr(block.get());
   irpass::die(block.get());
 
@@ -163,7 +165,8 @@ TEST(Scalarize, ScalarizeLocalStore) {
   // LocalStoreStmt survives irpass::die()
   block->push_back<LocalStoreStmt>(dest_stmt, matrix_init_stmt);
 
-  irpass::scalarize(block.get());
+  irpass::scalarize(block.get(),
+                    test_prog.prog()->this_thread_config().dynamic_index);
   irpass::die(block.get());
 
   EXPECT_EQ(block->size(), 2 /*const*/ + 4 /*alloca*/ + 4 /*store*/);
@@ -211,7 +214,8 @@ TEST(Scalarize, ScalarizeLocalLoad) {
   // Without this GlobalStoreStmt, nothing survives irpass::die()
   block->push_back<GlobalStoreStmt>(src_stmt, load_stmt);
 
-  irpass::scalarize(block.get());
+  irpass::scalarize(block.get(),
+                    test_prog.prog()->this_thread_config().dynamic_index);
   irpass::die(block.get());
 
   EXPECT_EQ(block->size(), 4 /*alloca*/ + 4 /*load*/ + 4 /*store*/);