wave equation: add partial gpu support, warm up jit cache

Author: tkarna

Diff for: examples/wave_equation.py

@@ -126,8 +126,10 @@ def ind_arr(shape, columns=False):
     T_shape = (nx, ny)
     U_shape = (nx + 1, ny)
     V_shape = (nx, ny + 1)
+    sync()
     x_t_2d = xmin + ind_arr(T_shape, True) * dx + dx / 2
     y_t_2d = ymin + ind_arr(T_shape) * dy + dy / 2
+    sync()
 
     dofs_T = int(numpy.prod(numpy.asarray(T_shape)))
     dofs_U = int(numpy.prod(numpy.asarray(U_shape)))
@@ -151,6 +153,8 @@ def ind_arr(shape, columns=False):
     u2 = create_full(U_shape, 0.0, dtype)
     v2 = create_full(V_shape, 0.0, dtype)
 
+    sync()
+
     def exact_elev(t, x_t_2d, y_t_2d, lx, ly):
         """
         Exact solution for elevation field.
@@ -224,7 +228,7 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
     sync()
 
     # initial solution
-    e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly)
+    e[:, :] = exact_elev(0.0, x_t_2d, y_t_2d, lx, ly).to_device(device)
     u[:, :] = create_full(U_shape, 0.0, dtype)
     v[:, :] = create_full(V_shape, 0.0, dtype)
     sync()
@@ -240,9 +244,22 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
         t = i * dt
 
         if t >= next_t_export - 1e-8:
-            _elev_max = np.max(e, all_axes)
-            _u_max = np.max(u, all_axes)
-            _total_v = np.sum(e + h, all_axes)
+            if device:
+                # FIXME gpu.memcpy to host requires identity layout
+                # FIXME reduction on gpu
+                # e_host = e.to_device()
+                # u_host = u.to_device()
+                # h_host = h.to_device()
+                # _elev_max = np.max(e_host, all_axes)
+                # _u_max = np.max(u_host, all_axes)
+                # _total_v = np.sum(e_host + h, all_axes)
+                _elev_max = 0
+                _u_max = 0
+                _total_v = 0
+            else:
+                _elev_max = np.max(e, all_axes)
+                _u_max = np.max(u, all_axes)
+                _total_v = np.sum(e + h, all_axes)
 
             elev_max = float(_elev_max)
             u_max = float(_u_max)
@@ -277,12 +294,19 @@ def step(u, v, e, u1, v1, e1, u2, v2, e2):
     duration = time_mod.perf_counter() - tic
     info(f"Duration: {duration:.2f} s")
 
-    e_exact = exact_elev(t, x_t_2d, y_t_2d, lx, ly)
-    err2 = (e_exact - e) * (e_exact - e) * dx * dy / lx / ly
-    err_L2 = math.sqrt(float(np.sum(err2, all_axes)))
+    if device:
+        # FIXME gpu.memcpy to host requires identity layout
+        # FIXME reduction on gpu
+        # err2_host = err2.to_device()
+        # err_L2 = math.sqrt(float(np.sum(err2_host, all_axes)))
+        err_L2 = 0
+    else:
+        e_exact = exact_elev(t, x_t_2d, y_t_2d, lx, ly)
+        err2 = (e_exact - e) * (e_exact - e) * dx * dy / lx / ly
+        err_L2 = math.sqrt(float(np.sum(err2, all_axes)))
     info(f"L2 error: {err_L2:7.5e}")
 
-    if nx == 128 and ny == 128 and not benchmark_mode:
+    if nx == 128 and ny == 128 and not benchmark_mode and not device:
         if datatype == "f32":
             assert numpy.allclose(err_L2, 7.2235471e-03, rtol=1e-4)
         else: