Refactor ASTTransformer.visit_For and fix a bug on grouped ndrange lo…

…ops (#648)

* New extension 'adstack'

* [skip ci] enforce format

* Simplify init in mgpcg_advanced.py (#645)

* put static/range/struct for into separate functions

* make decorators look clearer

* [skip ci] raise syntax error for more cases

* [skip ci] dynamic grouped ndrange?

* [skip ci] debug

* fix

* add more tests

* [skip ci] enforce code format

* further simplify mgpcg_advanced.py and make it able to run with self.dim=2

* [skip ci] enforce code format

* [skip ci] rename ijk to indices

* move components of visit_For outside

Co-authored-by: Yuanming Hu <yuanmhu@gmail.com>
Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>
Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

CMakeLists.txt

@@ -11,12 +11,12 @@ SET(TI_VERSION_MINOR 5)
 SET(TI_VERSION_PATCH 8)
 
 execute_process(
-  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} 
+  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
   COMMAND git rev-parse --short HEAD
   RESULT_VARIABLE SHORT_HASH_RESULT
   OUTPUT_VARIABLE TI_COMMIT_SHORT_HASH)
 execute_process(
-  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} 
+  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
   COMMAND git rev-parse HEAD
   RESULT_VARIABLE SHORT_HASH_RESULT
   OUTPUT_VARIABLE TI_COMMIT_HASH)
@@ -71,4 +71,3 @@ FILE(WRITE ${CMAKE_CURRENT_LIST_DIR}/taichi/common/version.h
         "#define TI_CUDAVERSION \"${CUDA_VERSION}\"\n"
         "#define TI_CUDAROOT_DIR \"${CUDA_TOOLKIT_ROOT_DIR}\"\n"
         )
-

examples/mgpcg_advanced.py

@@ -35,26 +35,28 @@ def __init__(self):
         self.pixels = ti.var(dt=real,
                              shape=(self.N_gui, self.N_gui))  # image buffer
 
-        self.grid = ti.root.pointer(ti.ijk, [self.N_tot // 4]).dense(
-            ti.ijk, 4).place(self.x, self.p, self.Ap)
+        indices = ti.ijk if self.dim == 3 else ti.ij
+        self.grid = ti.root.pointer(indices,
+                                    [self.N_tot // 4]).dense(indices, 4).place(
+                                        self.x, self.p, self.Ap)
 
         for l in range(self.n_mg_levels):
-            self.grid = ti.root.pointer(ti.ijk,
+            self.grid = ti.root.pointer(indices,
                                         [self.N_tot // (4 * 2**l)]).dense(
-                                            ti.ijk,
-                                            4).place(self.r[l], self.z[l])
+                                            indices, 4).place(
+                                                self.r[l], self.z[l])
 
         ti.root.place(self.alpha, self.beta, self.sum)
 
     @ti.kernel
     def init(self):
-        for I in ti.grouped(ti.ndrange((self.N_ext, self.N_tot - self.N_ext),
-                                       (self.N_ext, self.N_tot - self.N_ext),
-                                       (self.N_ext, self.N_tot - self.N_ext))):
+        for I in ti.grouped(
+                ti.ndrange(*(
+                    (self.N_ext, self.N_tot - self.N_ext), ) * self.dim)):
             self.r[0][I] = 1.0
             for i in ti.static(range(self.dim)):
-                self.r[0][I] *= ti.sin(2.0 * np.pi *
-                                       (i - self.N_ext) * 2.0 / self.N_tot)
+                self.r[0][I] *= ti.sin(2.0 * np.pi * (i - self.N_ext) * 2.0 /
+                                       self.N_tot)
             self.z[0][I] = 0.0
             self.Ap[I] = 0.0
             self.p[I] = 0.0
@@ -98,7 +100,7 @@ def update_p(self):
     @ti.kernel
     def restrict(self, l: ti.template()):
         for I in ti.grouped(self.r[l]):
-            res = self.r[l][I] - (6.0 * self.z[l][I] -
+            res = self.r[l][I] - (2.0 * self.dim * self.z[l][I] -
                                   self.neighbor_sum(self.z[l], I))
             self.r[l + 1][I // 2] += res * 0.5
 
@@ -112,8 +114,8 @@ def smooth(self, l: ti.template(), phase: ti.template()):
         # phase = red/black Gauss-Seidel phase
         for I in ti.grouped(self.r[l]):
             if (I.sum()) & 1 == phase:
-                self.z[l][I] = (self.r[l][I] +
-                                self.neighbor_sum(self.z[l], I)) / 6.0
+                self.z[l][I] = (self.r[l][I] + self.neighbor_sum(
+                    self.z[l], I)) / (2.0 * self.dim)
 
     def apply_preconditioner(self):
         self.z[0].fill(0)
@@ -137,11 +139,12 @@ def apply_preconditioner(self):
 
     @ti.kernel
     def paint(self):
-        kk = self.N_tot * 3 // 8
-        for i, j in self.pixels:
-            ii = int(i * self.N / self.N_gui) + self.N_ext
-            jj = int(j * self.N / self.N_gui) + self.N_ext
-            self.pixels[i, j] = self.x[ii, jj, kk] / self.N_tot
+        if ti.static(self.dim == 3):
+            kk = self.N_tot * 3 // 8
+            for i, j in self.pixels:
+                ii = int(i * self.N / self.N_gui) + self.N_ext
+                jj = int(j * self.N / self.N_gui) + self.N_ext
+                self.pixels[i, j] = self.x[ii, jj, kk] / self.N_tot
 
     def run(self):
         gui = ti.GUI("Multigrid Preconditioned Conjugate Gradients",