Skip to content

Commit

Permalink
Try #1441:
Browse files Browse the repository at this point in the history
  • Loading branch information
bors[bot] authored Sep 18, 2023
2 parents 72c1063 + 13ffca3 commit 2938e93
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 1 deletion.
8 changes: 7 additions & 1 deletion src/Fields/indices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,13 @@ end

Base.@propagate_inbounds Base.getindex(field::Field, colidx::ColumnIndex) =
column(field, colidx)

Base.@propagate_inbounds function Base.getindex(
fv::FieldVector{T},
colidx::ColumnIndex,
) where {T}
values = map(x -> x[colidx], _values(fv))
return FieldVector{T, typeof(values)}(values)
end
Base.@propagate_inbounds function column(
field::SpectralElementField1D,
colidx::ColumnIndex{1},
Expand Down
37 changes: 37 additions & 0 deletions test/Fields/field.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
using Test
using JET

using ClimaComms
using OrderedCollections
using StaticArrays, IntervalSets
Expand All @@ -12,6 +14,7 @@ using LinearAlgebra: norm
using Statistics: mean
using ForwardDiff
using CUDA
using CUDA: @allowscalar

include(
joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"),
Expand Down Expand Up @@ -267,6 +270,40 @@ end
@test Y.k.z === 3.0
end

function call_getcolumn(fv, colidx)
@allowscalar fvcol = fv[colidx]
nothing
end
function call_getproperty(fv)
fva = fv.c.a
nothing
end
@testset "FieldVector getindex" begin
cspace = TU.CenterExtrudedFiniteDifferenceSpace(Float32)
fspace = Spaces.ExtrudedFiniteDifferenceSpace{Spaces.CellFace}(cspace)
c = fill((a = Float32(1), b = Float32(2)), cspace)
f = fill((x = Float32(1), y = Float32(2)), fspace)
fv = Fields.FieldVector(; c, f)
colidx = Fields.ColumnIndex((1, 1), 1) # arbitrary index

@allowscalar @test all(parent(fv.c.a[colidx]) .== Float32(1))
@allowscalar @test all(parent(fv.f.y[colidx]) .== Float32(2))
@allowscalar @test propertynames(fv) == propertynames(fv[colidx])

# JET tests
# prerequisite
call_getproperty(fv) # compile first
@test_opt call_getproperty(fv)

call_getcolumn(fv, colidx) # compile first
@test_opt call_getcolumn(fv, colidx)
p = @allocated call_getcolumn(fv, colidx)
device = ClimaComms.device()
if ClimaComms.SingletonCommsContext(device) isa ClimaComms.AbstractCPUDevice
@test p 32
end
end

@testset "FieldVector array_type" begin
device = ClimaComms.device()
context = ClimaComms.SingletonCommsContext(device)
Expand Down

0 comments on commit 2938e93

Please sign in to comment.