[IR] Move "lower_access" to after "offload" (#1194)

taichi/transforms/compile_to_offloads.cpp

@@ -85,34 +85,32 @@ void compile_to_offloads(IRNode *ir,
   print("Access flagged I");
   irpass::analysis::verify(ir);
 
-  if (lower_global_access) {
-    irpass::lower_access(ir, true);
-    print("Access lowered");
-    irpass::analysis::verify(ir);
-
-    irpass::die(ir);
-    print("DIE");
-    irpass::analysis::verify(ir);
-  }
-
   irpass::full_simplify(ir);
   print("Simplified II");
   irpass::analysis::verify(ir);
 
-  irpass::flag_access(ir);
-  print("Access flagged II");
-  irpass::analysis::verify(ir);
-
   irpass::constant_fold(ir);
   print("Constant folded");
 
   irpass::offload(ir);
   print("Offloaded");
   irpass::analysis::verify(ir);
 
-  if (!lower_global_access) {
+  irpass::flag_access(ir);
+  print("Access flagged II");
+  irpass::analysis::verify(ir);
+
+  if (lower_global_access) {
+    irpass::lower_access(ir, true);
+    print("Access lowered");
+    irpass::analysis::verify(ir);
+
+    irpass::die(ir);
+    print("DIE");
+    irpass::analysis::verify(ir);
+
     irpass::flag_access(ir);
-    print("Access flagged after offloading");
+    print("Access flagged III");
     irpass::analysis::verify(ir);
   }
 
@@ -123,6 +121,12 @@ void compile_to_offloads(IRNode *ir,
   print("Atomics demoted");
   irpass::analysis::verify(ir);
 
+  irpass::full_simplify(ir);
+  print("Simplified III");
+
+  irpass::extract_constant(ir);
+  print("Constant extracted III");
+
   irpass::variable_optimization(ir, true);
   print("Store forwarded II");
 
@@ -131,7 +135,7 @@ void compile_to_offloads(IRNode *ir,
   irpass::analysis::verify(ir);
 
   irpass::full_simplify(ir);
-  print("Simplified III");
+  print("Simplified IV");
 
   // Final field registration correctness & type checking
   irpass::typecheck(ir);