CUDA.@allowscalar is not just for tests

Author: navidcy

src/utils.jl

@@ -248,16 +248,16 @@ end
 Returns an array, of the ArrayType of the device `grid` lives on, that contains the values of
 function `func` evaluated on the `grid`.
 """
-on_grid(func, grid::OneDGrid) = @. func(grid.x)
+on_grid(func, grid::OneDGrid) = CUDA.@allowscalar @. func(grid.x)
 
 function on_grid(func, grid::TwoDGrid)
   x, y = gridpoints(grid)
-  return func.(x, y)
+  return CUDA.@allowscalar func.(x, y)
 end
 
 function on_grid(func, grid::ThreeDGrid)
   x, y, z = gridpoints(grid)
-  return func.(x, y, z)
+  return CUDA.@allowscalar func.(x, y, z)
 end
 
 """

test/test_utils.jl

@@ -157,5 +157,5 @@ function test_ongrid(dev::Device)
   X₃, Y₃, Z₃ = gridpoints(g₃)
   f₃(x, y, z) = x^2 - y^3 + sin(z)
 
-  return CUDA.@allowscalar (FourierFlows.on_grid(f₁, g₁) == f₁.(X₁) && FourierFlows.on_grid(f₂, g₂) == f₂.(X₂, Y₂) && FourierFlows.on_grid(f₃, g₃) == f₃.(X₃, Y₃, Z₃))
+  return (FourierFlows.on_grid(f₁, g₁) == f₁.(X₁) && FourierFlows.on_grid(f₂, g₂) == f₂.(X₂, Y₂) && FourierFlows.on_grid(f₃, g₃) == f₃.(X₃, Y₃, Z₃))
 end