[autodiff] Support control flow for forward mode

ghstack-source-id: 06dfa909c4e0a4410a08693a54da15d35ecf84d8
Pull Request resolved: #5231

Author: erizmr

taichi/transforms/auto_diff.cpp

@@ -1137,12 +1137,39 @@ class MakeDual : public ADTransform {
       // d (x * y) = y * dx + x * dy
       accumulate(bin, mul(bin->lhs, dual(bin->rhs)));
       accumulate(bin, mul(bin->rhs, dual(bin->lhs)));
+    } else if (is_comparison(bin->op_type) || is_bit_op(bin->op_type)) {
+      // do nothing
     } else {
       TI_WARN("gradient of binary op {}", binary_op_type_name(bin->op_type));
       TI_NOT_IMPLEMENTED
     }
   }
 
+  void visit(IfStmt *if_stmt) override {
+    if (if_stmt->true_statements) {
+      std::vector<Stmt *> true_statements;
+      for (auto &stmt : if_stmt->true_statements->statements) {
+        true_statements.push_back(stmt.get());
+      }
+
+      for (auto stmt : true_statements) {
+        current_stmt = stmt;
+        stmt->accept(this);
+      }
+    }
+    if (if_stmt->false_statements) {
+      std::vector<Stmt *> false_statements;
+      for (auto &stmt : if_stmt->false_statements->statements) {
+        false_statements.push_back(stmt.get());
+      }
+
+      for (auto stmt : false_statements) {
+        current_stmt = stmt;
+        stmt->accept(this);
+      }
+    }
+  }
+
   void visit(RangeForStmt *for_stmt) override {
     std::vector<Stmt *> statements;
     // always make a copy since the list can be modified.

tests/python/test_ad_if_fwd.py

@@ -0,0 +1,82 @@
+from taichi.lang import impl
+from taichi.lang.misc import get_host_arch_list
+
+import taichi as ti
+from tests import test_utils
+
+
+@test_utils.test(arch=[ti.cpu, ti.gpu])
+def test_ad_if_simple_fwd():
+    x = ti.field(ti.f32, shape=())
+    y = ti.field(ti.f32, shape=())
+
+    @ti.kernel
+    def func():
+        if x[None] > 0.:
+            y[None] = x[None]
+
+    x[None] = 1
+    with ti.ad.FwdMode(loss=y, parameters=x, seed=[1.0]):
+        func()
+
+    assert y.grad[None] == 1
+
+
+@test_utils.test(arch=[ti.cpu, ti.gpu])
+def test_ad_if():
+    x = ti.field(ti.f32, shape=2)
+    y = ti.field(ti.f32, shape=2)
+
+    @ti.kernel
+    def func(i: ti.i32):
+        if x[i] > 0:
+            y[i] = x[i]
+        else:
+            y[i] = 2 * x[i]
+
+    x[0] = 0
+    x[1] = 1
+    with ti.ad.FwdMode(loss=y, parameters=x, seed=[1.0, 1.0]):
+        func(0)
+        func(1)
+    assert y.grad[0] == 2
+    assert y.grad[1] == 1
+
+
+@test_utils.test(arch=[ti.cpu, ti.gpu])
+def test_ad_if_nested():
+    n = 20
+    x = ti.field(ti.f32, shape=n)
+    y = ti.field(ti.f32, shape=n)
+    z = ti.field(ti.f32, shape=n)
+
+    ti.root.lazy_grad()
+
+    @ti.kernel
+    def func():
+        for i in x:
+            if x[i] < 2:
+                if x[i] == 0:
+                    y[i] = 0
+                else:
+                    y[i] = z[i] * 1
+            else:
+                if x[i] == 2:
+                    y[i] = z[i] * 2
+                else:
+                    y[i] = z[i] * 3
+
+    z.fill(1)
+
+    for i in range(n):
+        x[i] = i % 4
+
+    func()
+    for i in range(n):
+        assert y[i] == i % 4
+
+    with ti.ad.FwdMode(loss=y, parameters=z, seed=[1.0 for _ in range(n)]):
+        func()
+
+    for i in range(n):
+        assert y.grad[i] == i % 4