make bycolumn a no-op on CUDA

CliMA · Jun 21, 2023 · b1bdf1e · b1bdf1e
1 parent eb6f155
commit b1bdf1e
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 6 deletions.
diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -426,6 +426,8 @@ Base.@propagate_inbounds function level(
     Field(data, hspace)
 end
 
+Base.getindex(field::Field, ::Colon) = field
+
 Base.@propagate_inbounds Base.getindex(field::PointField) =
     getindex(field_values(field))
 Base.@propagate_inbounds Base.setindex!(field::PointField, val) =

diff --git a/src/Fields/indices.jl b/src/Fields/indices.jl
@@ -46,8 +46,14 @@ end
 """
     Fields.bycolumn(fn, space)
 
-Call `fn(colidx)` to every [`ColumnIndex`](@ref) `colidx` of `space`. This can be used to apply
-multiple column-wise operations in a single pass, making use of multiple threads.
+Call `fn(colidx)` to every [`ColumnIndex`](@ref) `colidx` of `space`. This can
+be used to apply multiple column-wise operations in a single pass, making use of
+multiple threads.
+
+!!! note
+
+    On GPUs this will simply evaluate `f` once with `colidx=:` (i.e. it doesn't
+    perform evaluation by columns). This may change in future.
 
 # Example
 
@@ -61,7 +67,15 @@ bycolumn(axes(f)) do colidx
 end
 ```
 """
-function bycolumn(fn, space::Spaces.SpectralElementSpace1D)
+function bycolumn(fn, space::Spaces.AbstractSpace)
+    bycolumn(fn, space, ClimaComms.device(space))
+end
+
+function bycolumn(
+    fn,
+    space::Spaces.SpectralElementSpace1D,
+    ::ClimaComms.CPUDevice,
+)
     Nh = Topologies.nlocalelems(space)
     Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space))
     @inbounds begin
@@ -81,7 +95,11 @@ function bycolumn(fn, space::Spaces.SpectralElementSpace1D)
     end
     return nothing
 end
-function bycolumn(fn, space::Spaces.SpectralElementSpace2D)
+function bycolumn(
+    fn,
+    space::Spaces.SpectralElementSpace2D,
+    ::ClimaComms.CPUDevice,
+)
     Nh = Topologies.nlocalelems(space)
     Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space))
     @inbounds begin
@@ -101,8 +119,18 @@ function bycolumn(fn, space::Spaces.SpectralElementSpace2D)
     end
     return nothing
 end
-bycolumn(fn, space::Spaces.ExtrudedFiniteDifferenceSpace) =
-    bycolumn(fn, space.horizontal_space)
+bycolumn(
+    fn,
+    space::Spaces.ExtrudedFiniteDifferenceSpace,
+    device::ClimaComms.CPUDevice,
+) = bycolumn(fn, space.horizontal_space, device)
+
+
+function bycolumn(fn, space::AbstractSpace, ::ClimaComms.CUDADevice)
+    fn(:)
+end
+
+
 
 """
     ncolumns(::Field)