Apply the Arg::check_bounds to the other kernels.

lattice · Jan 3, 2025 · b0b4355 · b0b4355
1 parent df97e1a
commit b0b4355
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 30 deletions.
diff --git a/include/targets/cuda/kernel.h b/include/targets/cuda/kernel.h
@@ -25,12 +25,16 @@ namespace quda
 
     auto i = threadIdx.x + blockIdx.x * blockDim.x;
 
-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {
+      while (i < arg.threads.x) {
+        f(i);
+        if (grid_stride)
+          i += gridDim.x * blockDim.x;
+        else
+          break;
+      }
+    } else {
       f(i);
-      if (grid_stride)
-        i += gridDim.x * blockDim.x;
-      else
-        break;
     }
   }
 
@@ -161,15 +165,20 @@ namespace quda
     auto i = threadIdx.x + blockIdx.x * blockDim.x;
     auto j = threadIdx.y + blockIdx.y * blockDim.y;
     auto k = threadIdx.z + blockIdx.z * blockDim.z;
-    if (j >= arg.threads.y) return;
-    if (k >= arg.threads.z) return;
 
-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {
+      if (j >= arg.threads.y) return;
+      if (k >= arg.threads.z) return;
+
+      while (i < arg.threads.x) {
+        f(i, j, k);
+        if (grid_stride)
+          i += gridDim.x * blockDim.x;
+        else
+          break;
+      }
+    } else {
       f(i, j, k);
-      if (grid_stride)
-        i += gridDim.x * blockDim.x;
-      else
-        break;
     }
   }
 

diff --git a/include/targets/hip/kernel.h b/include/targets/hip/kernel.h
@@ -13,12 +13,16 @@ namespace quda
 
     auto i = threadIdx.x + blockIdx.x * blockDim.x;
 
-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {
+      while (i < arg.threads.x) {
+        f(i);
+        if (grid_stride)
+          i += gridDim.x * blockDim.x;
+        else
+          break;
+      }
+    } else {
       f(i);
-      if (grid_stride)
-        i += gridDim.x * blockDim.x;
-      else
-        break;
     }
   }
 
@@ -43,14 +47,19 @@ namespace quda
 
     auto i = threadIdx.x + blockIdx.x * blockDim.x;
     auto j = threadIdx.y + blockIdx.y * blockDim.y;
-    if (j >= arg.threads.y) return;
 
-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {
+      if (j >= arg.threads.y) return;
+
+      while (i < arg.threads.x) {
+        f(i, j);
+        if (grid_stride)
+          i += gridDim.x * blockDim.x;
+        else
+          break;
+      }
+    } else {
       f(i, j);
-      if (grid_stride)
-        i += gridDim.x * blockDim.x;
-      else
-        break;
     }
   }
 
@@ -76,15 +85,20 @@ namespace quda
     auto i = threadIdx.x + blockIdx.x * blockDim.x;
     auto j = threadIdx.y + blockIdx.y * blockDim.y;
     auto k = threadIdx.z + blockIdx.z * blockDim.z;
-    if (j >= arg.threads.y) return;
-    if (k >= arg.threads.z) return;
 
-    while (i < arg.threads.x) {
+    if constexpr (Arg::check_bounds) {
+      if (j >= arg.threads.y) return;
+      if (k >= arg.threads.z) return;
+
+      while (i < arg.threads.x) {
+        f(i, j, k);
+        if (grid_stride)
+          i += gridDim.x * blockDim.x;
+        else
+          break;
+      }
+    } else {
       f(i, j, k);
-      if (grid_stride)
-        i += gridDim.x * blockDim.x;
-      else
-        break;
     }
   }