[refactor][ir] Simplify ir.h to reduce compile-time (#934)

* try remove statements.h in ir.h

* remove visitor

re-add

* fix

* passes & analysis

* I believe that expr is frontend

* balance ir.cpp and expr.cpp

* fix import error

* move RangeForStmt and StructForStmt definition to ir.cpp

* [skip ci] add comment

* [skip ci] enforce code format

* [skip ci] passes.h -> transoforms.h

* namespace irpass::analysis

fix

Co-authored-by: Taichi Gardener <[email protected]>

taichi/analysis/check_fields_registered.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/count_statements.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/detect_fors_with_break.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 #include "taichi/ir/frontend_ir.h"
 
 #include <unordered_set>

taichi/analysis/gather_deactivations.cpp

@@ -1,5 +1,6 @@
 #include "taichi/ir/ir.h"
-#include <set>
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/gather_statements.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/gather_used_atomics.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 #include <unordered_set>
 
 TLANG_NAMESPACE_BEGIN

taichi/analysis/has_store_or_atomic.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/last_store_or_atomic.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/same_statements.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 #include <unordered_map>
 #include <unordered_set>
 

taichi/analysis/value_diff.cpp

@@ -1,6 +1,8 @@
 // This pass analyzes compile-time known offsets for two values.
 
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/analysis/verify.cpp

@@ -1,4 +1,6 @@
 #include "taichi/ir/ir.h"
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/visitors.h"
 #include <vector>
 #include <unordered_set>
 

taichi/backends/metal/codegen_metal.cpp

@@ -5,6 +5,7 @@
 
 #include "taichi/backends/metal/constants.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/transforms.h"
 #include "taichi/util/line_appender.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/backends/opengl/codegen_opengl.cpp

@@ -7,6 +7,7 @@
 #include "taichi/backends/opengl/opengl_data_types.h"
 #include "taichi/backends/opengl/opengl_kernel_util.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/transforms.h"
 #include "taichi/util/line_appender.h"
 #include "taichi/util/macros.h"
 

taichi/codegen/codegen.cpp

@@ -8,7 +8,7 @@
 #include "taichi/backends/cuda/codegen_cuda.h"
 #endif
 #include "taichi/system/timer.h"
-#include "taichi/system/timer.h"
+#include "taichi/ir/analysis.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/ir/analysis.h

@@ -0,0 +1,70 @@
+#pragma once
+
+#include "taichi/ir/ir.h"
+#include <atomic>
+#include <unordered_set>
+#include <unordered_map>
+
+TLANG_NAMESPACE_BEGIN
+
+class DiffRange {
+ private:
+  bool related;
+
+ public:
+  int coeff;
+  int low, high;
+
+  DiffRange() : DiffRange(false, 0) {
+  }
+
+  DiffRange(bool related, int coeff) : DiffRange(related, 0, 0) {
+    TI_ASSERT(related == false);
+  }
+
+  DiffRange(bool related, int coeff, int low)
+      : DiffRange(related, coeff, low, low + 1) {
+  }
+
+  DiffRange(bool related, int coeff, int low, int high)
+      : related(related), coeff(coeff), low(low), high(high) {
+    if (!related) {
+      this->low = this->high = 0;
+    }
+  }
+
+  bool related_() const {
+    return related;
+  }
+
+  bool linear_related() const {
+    return related && coeff == 1;
+  }
+
+  bool certain() {
+    TI_ASSERT(related);
+    return high == low + 1;
+  }
+};
+
+// IR Analysis
+namespace irpass::analysis {
+
+void check_fields_registered(IRNode *root);
+int count_statements(IRNode *root);
+std::unordered_set<Stmt *> detect_fors_with_break(IRNode *root);
+std::unordered_set<Stmt *> detect_loops_with_continue(IRNode *root);
+std::unordered_set<SNode *> gather_deactivations(IRNode *root);
+std::vector<Stmt *> gather_statements(IRNode *root,
+                                      const std::function<bool(Stmt *)> &test);
+std::unique_ptr<std::unordered_set<AtomicOpStmt *>> gather_used_atomics(
+    IRNode *root);
+bool has_store_or_atomic(IRNode *root, const std::vector<Stmt *> &vars);
+std::pair<bool, Stmt *> last_store_or_atomic(IRNode *root, Stmt *var);
+bool same_statements(IRNode *root1, IRNode *root2);
+DiffRange value_diff(Stmt *stmt, int lane, Stmt *alloca);
+void verify(IRNode *root);
+
+}  // namespace irpass::analysis
+
+TLANG_NAMESPACE_END

taichi/ir/expr.cpp

@@ -156,4 +156,51 @@ void Expr::operator/=(const Expr &o) {
   (*this) = (*this) / load_if_ptr(o);
 }
 
+void Cache(int v, const Expr &var) {
+  dec.scratch_opt.push_back(std::make_pair(v, var.snode()));
+}
+
+void CacheL1(const Expr &var) {
+  dec.scratch_opt.push_back(std::make_pair(1, var.snode()));
+}
+
+Expr load_if_ptr(const Expr &ptr) {
+  if (ptr.is<GlobalPtrExpression>()) {
+    return load(ptr);
+  } else if (ptr.is<GlobalVariableExpression>()) {
+    TI_ASSERT(ptr.cast<GlobalVariableExpression>()->snode->num_active_indices ==
+              0);
+    return load(ptr[ExprGroup()]);
+  } else
+    return ptr;
+}
+
+Expr load(const Expr &ptr) {
+  TI_ASSERT(ptr.is<GlobalPtrExpression>());
+  return Expr::make<GlobalLoadExpression>(ptr);
+}
+
+Expr ptr_if_global(const Expr &var) {
+  if (var.is<GlobalVariableExpression>()) {
+    // singleton global variable
+    TI_ASSERT(var.snode()->num_active_indices == 0);
+    return var[ExprGroup()];
+  } else {
+    // may be any local or global expr
+    return var;
+  }
+}
+
+Expr Var(const Expr &x) {
+  auto var = Expr(std::make_shared<IdExpression>());
+  current_ast_builder().insert(std::make_unique<FrontendAllocaStmt>(
+      std::static_pointer_cast<IdExpression>(var.expr)->id, DataType::unknown));
+  var = x;
+  return var;
+}
+
+void Print_(const Expr &a, const std::string &str) {
+  current_ast_builder().insert(std::make_unique<FrontendPrintStmt>(a, str));
+}
+
 TLANG_NAMESPACE_END

taichi/ir/expr.h

@@ -123,4 +123,19 @@ Expr bit_cast(const Expr &input) {
   return taichi::lang::bit_cast(input, get_data_type<T>());
 }
 
+Expr load_if_ptr(const Expr &ptr);
+Expr load(const Expr &ptr);
+Expr ptr_if_global(const Expr &var);
+
+inline Expr smart_load(const Expr &var) {
+  return load_if_ptr(ptr_if_global(var));
+}
+
+// Begin: legacy frontend functions
+void Print_(const Expr &a, const std::string &str);
+void Cache(int v, const Expr &var);
+void CacheL1(const Expr &var);
+Expr Var(const Expr &x);
+// End: legacy frontend functions
+
 TLANG_NAMESPACE_END