Repository: taichi-dev/taichi (language: python)

Issue:
[RFC] Replace JIT evaluator (on device) in constant folding pass with a minimal CHI IR interpreter on host
## Background:

Currently when we do constant folding we have[ this JIT evaluator](https://github.com/taichi-dev/taichi/blob/master/taichi/transforms/constant_fold.cpp#L25-L72) which creates another taichi kernel and then compile and run on the fly and then inject the evaluated result back to the original compilation process. And this JIT evaluator kernel runs on the exact same device as initiated taichi arch. 

## Goal:

We want to replace this JIT evaluator on device with a minimal CHI IR interpreter running on host. The main motivation is: 
- Currently JIT evaluator only handles unary cast op and binary ops, these ops are pretty lightweight, creating a taichi kernel and evaluating them on the device often creates overhead. 
- Well the main motivation is actually simplify the codebase, esp. compiling JIT evaluator kernels requires a different compileconfig than its parent compilation and it has been a headache to maintain this. Replacing it with a small interpreter on host can greatly simplify the process. And we have this [arithmetic interpreter prototype](https://github.com/taichi-dev/taichi/blob/master/taichi/analysis/arithmetic_interpretor.cpp) from @k-ye. 

## Planned steps

- [ ] Add tests for arithmetic interpreter to make sure it's still behaving correctly
- [ ] Expand binary op coverage in arithmetic interpreter so that we can replace the binary jit evaluator kernels
- [ ] Support unary cast in arithmetic interpreter so that we can replace the unary jit evaluator kernels
- [ ] There might be some precision issues since we now evaluate on host instead of taichi device, we think the precision issue is tolerable and we'll just "fix" the tests failures as they show up. 


Discussion / hints:
Do we even need an interpreter to do this? IMO we can just do something similar to `alg_simp` and do some pattern matching. We can reference NIR's simplify pass design
> Do we even need an interpreter to do this? IMO we can just do something similar to alg_simp and do some pattern matching. We can reference NIR's simplify pass design

I agree that we don't need an interpreter here. It seems that the arithmetic interpreter prototype mentioned above has two main parts:
1. Skeleton code of an interpreter at inner-block level;
2. Pattern matching rules for arithmetic evaluation on constants.

IMO part 2 can be useful; part 1 is unnecessary. So my suggestion here is to complete part 2, extract it out and use it in the constant folding pass.
Yup that sounds good to me! Code isn't complicated so we just port them over. 

Fix patch:
diff --git a/taichi/ir/transforms.h b/taichi/ir/transforms.h
index 2cfe4068d1226..7346ff05464a7 100644
--- a/taichi/ir/transforms.h
+++ b/taichi/ir/transforms.h
@@ -106,9 +106,7 @@ void differentiation_validation_check(IRNode *root,
  * AD-stacks before this pass.
  */
 bool determine_ad_stack_size(IRNode *root, const CompileConfig &config);
-bool constant_fold(IRNode *root,
-                   const CompileConfig &config,
-                   const ConstantFoldPass::Args &args);
+bool constant_fold(IRNode *root);
 void offload(IRNode *root, const CompileConfig &config);
 bool transform_statements(
     IRNode *root,
diff --git a/taichi/transforms/constant_fold.cpp b/taichi/transforms/constant_fold.cpp
index a7ebe55a88f7f..051e3afb30ff7 100644
--- a/taichi/transforms/constant_fold.cpp
+++ b/taichi/transforms/constant_fold.cpp
@@ -29,15 +29,6 @@ class ConstantFold : public BasicStmtVisitor {
  public:
   using BasicStmtVisitor::visit;
   DelayedIRModifier modifier;
-  Program *program;
-  CompileConfig compile_config;
-
-  explicit ConstantFold(Program *program, const CompileConfig &compile_config)
-      : program(program), compile_config(compile_config) {
-    this->compile_config.advanced_optimization = false;
-    this->compile_config.constant_folding = false;
-    this->compile_config.external_optimization_level = 0;
-  }
 
   static bool is_good_type(DataType dt) {
     // ConstStmt of `bad` types like `i8` is not supported by LLVM.
@@ -237,10 +228,8 @@ class ConstantFold : public BasicStmtVisitor {
     return;
   }
 
-  static bool run(IRNode *node,
-                  Program *program,
-                  const CompileConfig &compile_config) {
-    ConstantFold folder(program, compile_config);
+  static bool run(IRNode *node) {
+    ConstantFold folder;
     bool modified = false;
 
     while (true) {
@@ -268,13 +257,9 @@ const PassID ConstantFoldPass::id = "ConstantFoldPass";
 
 namespace irpass {
 
-bool constant_fold(IRNode *root,
-                   const CompileConfig &compile_config,
-                   const ConstantFoldPass::Args &args) {
+bool constant_fold(IRNode *root) {
   TI_AUTO_PROF;
-  if (!compile_config.advanced_optimization)
-    return false;
-  return ConstantFold::run(root, args.program, compile_config);
+  return ConstantFold::run(root);
 }
 
 }  // namespace irpass
diff --git a/taichi/transforms/simplify.cpp b/taichi/transforms/simplify.cpp
index 3838cbc35424a..0bf75e1866c55 100644
--- a/taichi/transforms/simplify.cpp
+++ b/taichi/transforms/simplify.cpp
@@ -523,8 +523,7 @@ void full_simplify(IRNode *root,
         modified = true;
       if (binary_op_simplify(root, config))
         modified = true;
-      if (config.constant_folding &&
-          constant_fold(root, config, {args.program}))
+      if (config.constant_folding && constant_fold(root))
         modified = true;
       if (die(root))
         modified = true;
@@ -553,7 +552,7 @@ void full_simplify(IRNode *root,
     return;
   }
   if (config.constant_folding) {
-    constant_fold(root, config, {args.program});
+    constant_fold(root);
     die(root);
   }
   simplify(root, config);


Test patch:
diff --git a/tests/cpp/transforms/alg_simp_test.cpp b/tests/cpp/transforms/alg_simp_test.cpp
index 34f7fef38b66e..9472ca158d1a0 100644
--- a/tests/cpp/transforms/alg_simp_test.cpp
+++ b/tests/cpp/transforms/alg_simp_test.cpp
@@ -123,8 +123,7 @@ TEST_F(AlgebraicSimplicationTest, SimplifyMultiplyZeroFastMath) {
   irpass::type_check(block.get(), config_without_fast_math);  // insert 2 casts
   EXPECT_EQ(block->size(), 10);
 
-  irpass::constant_fold(block.get(), config_without_fast_math,
-                        {&prog()});  // should change 2 casts into const
+  irpass::constant_fold(block.get());  // should change 2 casts into const
   irpass::alg_simp(block.get(),
                    config_without_fast_math);  // should not eliminate
   irpass::die(block.get());                    // should eliminate 2 const
diff --git a/tests/cpp/transforms/binary_op_simplify_test.cpp b/tests/cpp/transforms/binary_op_simplify_test.cpp
index d7e1463524ed1..4d771fba9410d 100644
--- a/tests/cpp/transforms/binary_op_simplify_test.cpp
+++ b/tests/cpp/transforms/binary_op_simplify_test.cpp
@@ -33,7 +33,7 @@ TEST_F(BinaryOpSimplifyTest, MultiplyPOT) {
   // -> (x << 5) << 3
   irpass::binary_op_simplify(ir_block, CompileConfig());
   // -> x << (5 + 3)
-  irpass::constant_fold(ir_block, CompileConfig(), {tp_.prog()});
+  irpass::constant_fold(ir_block);
   // -> x << 8
   irpass::die(ir_block);
 
@@ -83,7 +83,7 @@ TEST_F(BinaryOpSimplifyTest, ModPOT) {
   EXPECT_EQ(ir_block->size(), 5);
 
   // -> x & 7
-  irpass::constant_fold(ir_block, CompileConfig(), {tp_.prog()});
+  irpass::constant_fold(ir_block);
   irpass::die(ir_block);
   EXPECT_EQ(ir_block->size(), 4);
   EXPECT_EQ(ir_block->statements[0].get(), x);
diff --git a/tests/cpp/transforms/constant_fold_test.cpp b/tests/cpp/transforms/constant_fold_test.cpp
index 9b0a9454c43c0..070a69bc2bffd 100644
--- a/tests/cpp/transforms/constant_fold_test.cpp
+++ b/tests/cpp/transforms/constant_fold_test.cpp
@@ -9,24 +9,19 @@ namespace taichi::lang {
 
 class ConstantFoldTest : public ::testing::Test {
  protected:
-  void SetUp() override {
-    tp_.setup();
-  }
-
   void run_constant_fold() {
     ir = builder.extract_ir();
     ASSERT_TRUE(ir->is<Block>());
     auto *ir_block = ir->as<Block>();
     irpass::type_check(ir_block, CompileConfig());
 
-    irpass::constant_fold(ir_block, CompileConfig(), {tp_.prog()});
+    irpass::constant_fold(ir_block);
     irpass::die(ir_block);
 
     EXPECT_EQ(ir_block->size(), 4);
     EXPECT_TRUE(ir_block->statements[1]->is<ConstStmt>());
   }
 
-  TestProgram tp_;
   IRBuilder builder;
   std::unique_ptr<Block> ir;
 };
diff --git a/tests/cpp/transforms/simplify_test.cpp b/tests/cpp/transforms/simplify_test.cpp
index be8c7f954888b..71adb1c87c1cb 100644
--- a/tests/cpp/transforms/simplify_test.cpp
+++ b/tests/cpp/transforms/simplify_test.cpp
@@ -40,8 +40,7 @@ TEST(Simplify, SimplifyLinearizedWithTrivialInputs) {
       test_prog.prog()->compile_config());  // should lower linearized
   // EXPECT_EQ(block->size(), 11);  // not required to check size here
 
-  irpass::constant_fold(block.get(), test_prog.prog()->compile_config(),
-                        {test_prog.prog()});
+  irpass::constant_fold(block.get());
   irpass::alg_simp(block.get(), test_prog.prog()->compile_config());
   irpass::die(block.get());  // should eliminate consts
   irpass::simplify(block.get(), test_prog.prog()->compile_config());