[Bug] [ir] Fix the IdentifyValuesUsedInOtherOffloads pass (#3597)

taichi/transforms/offload.cpp

@@ -345,13 +345,6 @@ class IdentifyValuesUsedInOtherOffloads : public BasicStmtVisitor {
     TI_ASSERT(current_offloaded);
   }
 
-  void visit(RangeForStmt *stmt) override {
-    test_and_allocate(stmt->begin);
-    test_and_allocate(stmt->end);
-    if (stmt->body)
-      stmt->body->accept(this);
-  }
-
   void test_and_allocate(Stmt *stmt) {
     if (stmt == nullptr)
       return;
@@ -371,14 +364,22 @@ class IdentifyValuesUsedInOtherOffloads : public BasicStmtVisitor {
     }
   }
 
-  void visit(Stmt *stmt) override {
+  void generic_visit(Stmt *stmt) {
     int n_op = stmt->num_operands();
     for (int i = 0; i < n_op; i++) {
       auto op = stmt->operand(i);
       test_and_allocate(op);
     }
   }
 
+  void preprocess_container_stmt(Stmt *stmt) override {
+    generic_visit(stmt);
+  }
+
+  void visit(Stmt *stmt) override {
+    generic_visit(stmt);
+  }
+
   static StmtToOffsetMap run(
       IRNode *root,
       const CompileConfig &config,

tests/python/test_offload_cross.py

@@ -120,3 +120,15 @@ def run(a: ti.i32):
                 print('OK')
 
     run(2)
+
+
+@ti.test()
+def test_offload_with_cross_if_inside_for():
+    @ti.kernel
+    def run(a: ti.i32):
+        b = a > 2
+        for x in range(1):
+            if b:
+                print('OK')
+
+    run(2)