[example] Add cuda options for sparse matrix examples (taichi-dev#6785)

Issue: taichi-dev#2906
quadpixels · May 13, 2023 · 38d86f9 · 38d86f9
1 parent e355cee
commit 38d86f9
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 12 deletions.
diff --git a/python/taichi/examples/simulation/implicit_mass_spring.py b/python/taichi/examples/simulation/implicit_mass_spring.py
@@ -237,20 +237,32 @@ def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)):
 
 
 def main():
-    ti.init(arch=ti.cpu)
-    h = 0.01
-    cloth = Cloth(N=5)
-
-    pause = False
     parser = argparse.ArgumentParser()
     parser.add_argument('-g',
                         '--use-ggui',
                         action='store_true',
                         help='Display with GGUI')
+    parser.add_argument('-a',
+                        '--arch',
+                        required=False,
+                        default="cpu",
+                        dest='arch',
+                        type=str,
+                        help='The arch (backend) to run this example on')
     args, unknowns = parser.parse_known_args()
-    use_ggui = False
-    use_ggui = args.use_ggui
+    arch = args.arch
+    if arch in ["x64", "cpu", "arm64"]:
+        ti.init(arch=ti.cpu)
+    elif arch in ["cuda", "gpu"]:
+        ti.init(arch=ti.cuda)
+    else:
+        raise ValueError('Only CPU and CUDA backends are supported for now.')
 
+    h = 0.01
+    pause = False
+    cloth = Cloth(N=5)
+
+    use_ggui = args.use_ggui
     if not use_ggui:
         gui = ti.GUI('Implicit Mass Spring System', res=(500, 500))
         while gui.running:

diff --git a/python/taichi/examples/simulation/stable_fluid.py b/python/taichi/examples/simulation/stable_fluid.py
@@ -18,9 +18,16 @@
                     '--use-sp-mat',
                     action='store_true',
                     help='Solve Poisson\'s equation by using a sparse matrix')
+parser.add_argument('-a',
+                    '--arch',
+                    required=False,
+                    default="cpu",
+                    dest='arch',
+                    type=str,
+                    help='The arch (backend) to run this example on')
 args, unknowns = parser.parse_known_args()
 
-res = 512
+res = 256
 dt = 0.03
 p_jacobi_iters = 500  # 40 for a quicker but less accurate result
 f_strength = 10000.0
@@ -32,12 +39,17 @@
 debug = False
 
 use_sparse_matrix = args.use_sp_mat
+arch = args.arch
+if arch in ["x64", "cpu", "arm64"]:
+    ti.init(arch=ti.cpu)
+elif arch in ["cuda", "gpu"]:
+    ti.init(arch=ti.cuda)
+else:
+    raise ValueError('Only CPU and CUDA backends are supported for now.')
 
 if use_sparse_matrix:
-    ti.init(arch=ti.x64)
     print('Using sparse matrix')
 else:
-    ti.init(arch=ti.gpu)
     print('Using jacobi iteration')
 
 _velocities = ti.Vector.field(2, float, shape=(res, res))
@@ -86,7 +98,7 @@ def fill_laplacian_matrix(A: ti.types.sparse_matrix_builder()):
 
     N = res * res
     K = ti.linalg.SparseMatrixBuilder(N, N, max_num_triplets=N * 6)
-    F_b = ti.field(ti.f32, shape=N)
+    F_b = ti.ndarray(ti.f32, shape=N)
 
     fill_laplacian_matrix(K)
     L = K.build()
@@ -236,7 +248,7 @@ def enhance_vorticity(vf: ti.template(), cf: ti.template()):
 
 
 @ti.kernel
-def copy_divergence(div_in: ti.template(), div_out: ti.template()):
+def copy_divergence(div_in: ti.template(), div_out: ti.types.ndarray()):
     for I in ti.grouped(div_in):
         div_out[I[0] * res + I[1]] = -div_in[I]