Fixes + more dont limit

CliMA · Feb 17, 2025 · eea802c · eea802c
1 parent 5038012
commit eea802c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 11 deletions.
diff --git a/ext/cuda/data_layouts_threadblock.jl b/ext/cuda/data_layouts_threadblock.jl
@@ -309,22 +309,12 @@ end
 @inline function fd_stencil_universal_index(space::Spaces.AbstractSpace, us)
     (tv,) = CUDA.threadIdx()
     (h, bv, ij) = CUDA.blockIdx()
-    vid = tv + (bv - 1) * CUDA.blockDim().x
+    v = tv + (bv - 1) * CUDA.blockDim().x
     (Nq, _, _, _, _) = DataLayouts.universal_size(us)
     if Nq * Nq < ij
         return CartesianIndex((-1, -1, 1, -1, -1))
     end
     @inbounds (i, j) = CartesianIndices((Nq, Nq))[ij].I
-    v =
-        if space isa Spaces.FaceExtrudedFiniteDifferenceSpace ||
-           space isa Spaces.FaceFiniteDifferenceSpace
-            v = vid - half
-        elseif space isa Spaces.CenterExtrudedFiniteDifferenceSpace ||
-               space isa Spaces.CenterFiniteDifferenceSpace
-            v = vid
-        else
-            error("Invalid space")
-        end
     return CartesianIndex((i, j, 1, v, h))
 end
 @inline fd_stencil_is_valid_index(I::CI5, us::UniversalSize) =

diff --git a/ext/cuda/operators_fd_shmem_common.jl b/ext/cuda/operators_fd_shmem_common.jl
@@ -286,4 +286,10 @@ if hasfield(Method, :recursion_relation)
     for m in methods(_fd_resolve_shmem!)
         m.recursion_relation = dont_limit
     end
+    for m in methods(_fd_allocate_shmem)
+        m.recursion_relation = dont_limit
+    end
+    for m in methods(fd_allocate_shmem)
+        m.recursion_relation = dont_limit
+    end
 end