Replace `copy_stack_field_down!` with broadcasts of reshaped MPIStateArrays #1301

blallen · 2020-06-29T16:09:31Z

Description

After #1071 is merged, we no longer need a dedicated kernel to copy the vertical velocity at the top of the ocean (wz0) down the entire stack for use in the source term for η. We can also use this PR to discuss how we want the basic_grid_info function to work.

I have

Written and run all necessary tests with CLIMA by including tests/runtests.jl
Followed all necessary style guidelines and run julia .dev/climaformat.jl .
Updated the documentation to reflect changes from this PR.

For review by CLIMA developers

There are no open pull requests for this already
CLIMA developers with relevant expertise have been assigned to review this submission
The code conforms to the style guidelines and has consistent naming conventions. julia .dev/format.jl has been run in a separate commit.
This code does what it is technically intended to do (all numerics make sense physically and/or computationally)

blallen · 2020-06-29T16:11:30Z

@jkozdon if this works on the GPU, should I just remove copy_stack_field_down! entirely?

blallen · 2020-06-29T16:11:35Z

bors try

bors · 2020-06-29T16:54:16Z

try

Build failed:

ci/slurmci

blallen · 2020-06-29T17:11:49Z

@leios can you check why the ocean tests failed on the GPU on the slurm CI? this is a practical usecase and it seems to not work

jkozdon

No real opinion.

jkozdon · 2020-06-29T18:55:59Z

@leios can you check why the ocean tests failed on the GPU on the slurm CI? this is a practical usecase and it seems to not work

If you cannot get it working with reshapes of the underlying mpistate array, you can also just do it with the data directly.

leios · 2020-06-29T19:02:22Z

Yeah, so looking at this further, the types for the broadcast that are failing are:

1. Base.ReshapedArray{Float64,5,SubArray{Float64,3,CUDA.CuArray{Float64,3,CUDA.CuArray{Float64,3,Nothing}},Tuple{Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}
2.
Base.ReshapedArray{Float64,5,SubArray{Float64,3,Base.ReshapedArray{Float64,5,SubArray{Float64,3,CUDA.CuArray{Float64,3,CUDA.CuArray{Float64,3,Nothing}},Tuple{Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}},Tuple{Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}}},false},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}

PR #1071 allows for broadcasting of reshaped MPIStateArrays, but this might be more closely related to an issue with adapt, itself

jkozdon · 2020-06-29T19:05:49Z

Yeah, so looking at this further, the types for the broadcast that are failing are:

Right. Just saw that. This has nothing to with with #1071, since it's broadcast the data based on the vars views of the data.

but this might be more closely related to an issue with adapt, itself

This seems likely.

blallen · 2020-06-29T19:13:58Z

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

+    boxy_w = reshape(A.w, Nq^2, Nqk, 1, nelemv, nelemh)
+    flat_w = @view boxy_w[:, end, :, end, :]
+    flat_wz0 = reshape(flat_w, Nq^2, 1, 1, 1, nelemh)
+    boxy_wz0 = reshape(A.wz0, Nq^2, Nqk, 1, nelemv, nelemh)
+    boxy_wz0 .= flat_wz0


Would adding .data after A.w and A.wz0 solve the issue? The split-explicit code we developed with @jm-c does this quite a lot.

It's worth a shot. If that doesn't work, we can export the MPIStateArray functions we created to find .data of adapted MPIStateArrays.

jkozdon · 2020-06-29T19:16:27Z

but this might be more closely related to an issue with adapt, itself

Nothing to do with CLIMA or MPIStateArrays.

using CUDA;
CUDA.allowscalar(false);
Nq = 5;
Nqk = 5;
nelemv = 10;
nelemh = 100;
nelemh_ghost = 10;
nstate = 5;
Q = CuArray{Float64, 3}(undef, nstate, Nq * Nq * Nqk, nelemv * nelemh);
A = CuArray{Float64, 3}(undef, nstate, Nq * Nq * Nqk, nelemv * nelemh);
w = view(Q, 3, :, :);
wz0 = view(A, 1, :, :);
boxy_w = reshape(w, Nq^2, Nqk, 1, nelemv, nelemh);
flat_w = @view boxy_w[:, end, :, end, :];
flat_wz0 = reshape(flat_w, Nq^2, 1, 1, 1, nelemh);
boxy_wz0 = reshape(wz0, Nq^2, Nqk, 1, nelemv, nelemh);
boxy_wz0 .= flat_wz0;

will produce:

julia> using CUDA;

julia> CUDA.allowscalar(false);

julia> Nq = 5;

julia> Nqk = 5;

julia> nelemv = 10;

julia> nelemh = 100;

julia> nelemh_ghost = 10;

julia> nstate = 5;

julia> Q = CuArray{Float64, 3}(undef, nstate, Nq * Nq * Nqk, nelemv * nelemh);

julia> A = CuArray{Float64, 3}(undef, nstate, Nq * Nq * Nqk, nelemv * nelemh);

julia> w = view(Q, 3, :, :);

julia> wz0 = view(A, 1, :, :);

julia> boxy_w = reshape(w, Nq^2, Nqk, 1, nelemv, nelemh);

julia> flat_w = @view boxy_w[:, end, :, end, :];

julia> flat_wz0 = reshape(flat_w, Nq^2, 1, 1, 1, nelemh);

julia> boxy_wz0 = reshape(wz0, Nq^2, Nqk, 1, nelemv, nelemh);

julia> boxy_wz0 .= flat_wz0;
ERROR: scalar getindex is disallowed
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] assertscalar(::String) at /home/jekozdon/.julia/packages/GPUArrays/X4SqE/src/host/indexing.jl:41
 [3] getindex at /home/jekozdon/.julia/packages/GPUArrays/X4SqE/src/host/indexing.jl:96 [inlined]
 [4] getindex at ./subarray.jl:245 [inlined]
 [5] _unsafe_getindex_rs at ./reshapedarray.jl:249 [inlined]
 [6] _unsafe_getindex at ./reshapedarray.jl:246 [inlined]
 [7] getindex at ./reshapedarray.jl:234 [inlined]
 [8] getindex at ./subarray.jl:236 [inlined]
 [9] _unsafe_getindex_rs(::SubArray{Float64,3,Base.ReshapedArray{Float64,5,SubArray{Float64,2,CuArray{Float64,3,Nothing},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true},Tuple{}},Tuple{Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}}},false}, ::Tuple{Int64,Int64,Int64}) at ./reshapedarray.jl:249
 [10] _unsafe_getindex(::Base.ReshapedArray{Float64,5,SubArray{Float64,3,Base.ReshapedArray{Float64,5,SubArray{Float64,2,CuArray{Float64,3,Nothing},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true},Tuple{}},Tuple{Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}}},false},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}, ::Int64, ::Int64, ::Int64, ::Int64, ::Int64) at ./reshapedarray.jl:246
 [11] getindex at ./reshapedarray.jl:234 [inlined]
 [12] _getindex at ./abstractarray.jl:1020 [inlined]
 [13] getindex at ./abstractarray.jl:980 [inlined]
 [14] _broadcast_getindex at ./broadcast.jl:597 [inlined]
 [15] _getindex at ./broadcast.jl:628 [inlined]
 [16] _broadcast_getindex at ./broadcast.jl:603 [inlined]
 [17] getindex at ./broadcast.jl:564 [inlined]
 [18] macro expansion at ./broadcast.jl:910 [inlined]
 [19] macro expansion at ./simdloop.jl:77 [inlined]
 [20] copyto! at ./broadcast.jl:909 [inlined]
 [21] copyto! at ./broadcast.jl:864 [inlined]
 [22] materialize!(::Base.ReshapedArray{Float64,5,SubArray{Float64,2,CuArray{Float64,3,Nothing},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true},Tuple{}}, ::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{5},Nothing,typeof(identity),Tuple{Base.ReshapedArray{Float64,5,SubArray{Float64,3,Base.ReshapedArray{Float64,5,SubArray{Float64,2,CuArray{Float64,3,Nothing},Tuple{Int64,Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}}},true},Tuple{}},Tuple{Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}},Int64,Base.Slice{Base.OneTo{Int64}}},false},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64},Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}}}) at ./broadcast.jl:823
 [23] top-level scope at REPL[69]:1

leios · 2020-06-29T19:19:40Z

yeah, then we have to mess around with Adapt to get this to work... Unless there is a way to simplify the types (specifically flat_w)? There was a similar situation I ran into where we couldn't broadcast from a view to a non-view array.

jkozdon · 2020-06-29T19:26:04Z

If you reshape before the views this would work:

using CUDA;
CUDA.allowscalar(false);
Nq = 5;
Nqk = 5;
nelemv = 10;
nelemh = 100;
nelemh_ghost = 10;
nstate = 5;
A = CuArray{Float64, 3}(undef, nstate, Nq * Nq * Nqk, nelemv * nelemh);
B = reshape(A, nstate, Nq^2, Nqk, nelemv, nelemh);
w = view(B, 3, :, :, :, :);
wz0 = view(B, 1, :, 1:1, 1:1, :);
w .= wz0;

To do this you need to access the data then reshape and pull out the indices you want.

blallen · 2020-06-29T20:45:07Z

@jkozdon we ended up using varsindex and .data to get things to work, see d8e0d92. It's not elegant, but it's functional.

jkozdon · 2020-06-29T21:00:43Z

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

+    # project w(z=0) down the stack
+    boxy_w = reshape(A.data[:,index_w,:], Nq^2, Nqk, 1, nelemv, nelemh)
+    flat_w = @view boxy_w[:, end, :, end, :]
+    flat_wz0 = reshape(flat_w, Nq^2, 1, 1, 1, nelemh)
+    boxy_wz0 = reshape(A.data[:,index_wz0,:], Nq^2, Nqk, 1, nelemv, nelemh)


I don't think this will work, you are creating an array slice not a view, so you are not modifying the original array but a newly created array.

Also, you haven't accounted for the ghost elements, so I don't think it will work in parallel.

Possibly it would work in parrallel once you fix the views with:

Nq, Nqk, _, _, nelemv, nelemh, _, _ = basic_grid_info(dg)

above, though you will be copying down the ghost elements too.

ah, I think you are right. This worked on my GPU, but we need to think of a more general solution.

I think you need to reshape, then do the views (like I did here: #1301 (comment))

The reshape returns a new CuArray (pointing to the same data) which is why you miss the double indirection problem.

Looks good to me lurking - is there an extra ,1, after num_aux_state or is that needed for some more magic?

Ah it is an extra ,1,! that's what my parenthetical was about

o.o I think @jkozdon is editting it, the indices just jumped

This (original code) doesn't work, not even on a CPU. so we need to make some kind of further change for sure

Couple typos for sure. Updated my suggestion.

@jkozdon can you refresh me on the MPIStateArray.data layout? I thought it was ["GL points in element", "states", "elements"] but that doesn't match up with your views (or even the first reshape, shouldn't the 1 be a :?).

data ordering is [dof, state, elements]

just pushed a commit based on that fix, worked on the CPU, about to try it on the GPU

hmm didn't succeed on GPU due to the reshape dimensions not matching the array.data: https://gist.github.com/climabot/4a0358efe092e8b8015870c489a0e2df#file-out_10192381

any ideas @jkozdon? it works fine on CPU

src/Numerics/DGMethods/DGModel.jl

blallen · 2020-06-29T22:45:32Z

bors try

bors · 2020-06-29T23:42:57Z

try

Build failed:

ci/slurmci

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

blallen · 2020-06-30T15:31:31Z

bors try

bors · 2020-06-30T17:22:38Z

try

Build failed:

docs-build

blallen · 2020-06-30T17:25:20Z

try

Build failed:

docs-build

Does this mean slurm/CI succeeded or that it wasn't tried at all?

leios · 2020-06-30T17:26:04Z

I think it is still trying those

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl

blallen · 2020-06-30T18:37:05Z

@jkozdon if this works on the GPU, should I just remove copy_stack_field_down! entirely?

@jkozdon thoughts? since this seems to be working finally

blallen · 2020-06-30T18:37:15Z

bors try

jkozdon · 2020-06-30T19:28:11Z

@jkozdon thoughts? since this seems to be working finally

Fine by me to remove it. I believe that you are the only ones using this functionality anyway.

modifying HBM to no longer use vars beautification make basic_grid_info return NamedTuple use views of `.data` instead of reshapes of slices account for ghost elements clean up and doc fix Co-authored-by: leios <jrs.schloss@gmail.com> Co-authored-by: Jeremy E Kozdon <jekozdon@nps.edu>

bors · 2020-06-30T20:23:12Z

try

Build succeeded:

blallen · 2020-06-30T21:14:59Z

@jkozdon thoughts? since this seems to be working finally

Fine by me to remove it. I believe that you are the only ones using this functionality anyway.

okay, I went ahead and removed them in a separate commit.

blallen · 2020-06-30T21:19:32Z

bors r+

jkozdon · 2020-06-30T21:29:01Z

okay, I went ahead and removed them in a separate commit.

Two commits. Love it!

blallen · 2020-06-30T21:45:17Z

okay, I went ahead and removed them in a separate commit.

Two commits. Love it!

haha I know that both commits will pass CI :) But we had a conversation at MIT earlier today about this and sometimes multiple commits are justified for clarity. Also, I think we said we could have the auto-bug checker thing run only on bors commits, right?

1301: Replace `copy_stack_field_down!` with broadcasts of reshaped MPIStateArrays r=blallen a=blallen # Description After #1071 is merged, we no longer need a dedicated kernel to copy the vertical velocity at the top of the ocean (`wz0`) down the entire stack for use in the source term for `η`. We can also use this PR to discuss how we want the `basic_grid_info` function to work. Co-authored-by: Brandon Allen <ballen@mit.edu>

jkozdon · 2020-06-30T21:52:52Z

haha I know that both commits will pass CI :) But we had a conversation at MIT earlier today about this and sometimes multiple commits are justified for clarity. Also, I think we said we could have the auto-bug checker thing run only on bors commits, right? Yeah there is a way to get git-bisect to run on just bors commits.

bors · 2020-06-30T22:23:10Z

Build failed:

docs-build

charleskawczynski · 2020-06-30T23:23:53Z

bors r+

bors · 2020-07-01T01:26:31Z

Build succeeded:

blallen added Defensive programming Makes code safer Ocean 🌊 "How inappropriate to call this planet Earth when it is quite clearly Ocean" labels Jun 29, 2020

blallen requested review from jkozdon and thomasgibson June 29, 2020 16:09

blallen self-assigned this Jun 29, 2020

blallen mentioned this pull request Jun 29, 2020

Couple the HBModel and SWModel for use in the Split-Explicit Stepper #1268

Merged

7 tasks

bors bot added a commit that referenced this pull request Jun 29, 2020

Try #1301:

1ecb94b

vchuravy changed the base branch from master to jrs/MPI_wrapper June 29, 2020 17:04

vchuravy changed the base branch from jrs/MPI_wrapper to master June 29, 2020 17:04

jkozdon reviewed Jun 29, 2020

View reviewed changes

leios closed this Jun 29, 2020

leios reopened this Jun 29, 2020

blallen commented Jun 29, 2020

View reviewed changes

jkozdon reviewed Jun 29, 2020

View reviewed changes

src/Numerics/DGMethods/DGModel.jl Outdated Show resolved Hide resolved

bors bot added a commit that referenced this pull request Jun 29, 2020

Try #1301:

bac79e6

jkozdon reviewed Jun 30, 2020

View reviewed changes

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl Outdated Show resolved Hide resolved

blallen force-pushed the bla/wz0_broadcast branch from 1fbfd6e to ed1c682 Compare June 30, 2020 15:30

bors bot added a commit that referenced this pull request Jun 30, 2020

Try #1301:

e3cdea7

blallen force-pushed the bla/wz0_broadcast branch from b6468b6 to bc15c1c Compare June 30, 2020 18:34

blallen commented Jun 30, 2020

View reviewed changes

src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl Show resolved Hide resolved

bors bot added a commit that referenced this pull request Jun 30, 2020

Try #1301:

3bcd81d

blallen and others added 2 commits June 30, 2020 16:00

delete copy_stack_field_down! code

59f4828

blallen force-pushed the bla/wz0_broadcast branch from bc15c1c to 59f4828 Compare June 30, 2020 20:33

bors bot merged commit 78fd04b into CliMA:master Jul 1, 2020

blallen deleted the bla/wz0_broadcast branch July 23, 2020 19:21

Replace copy_stack_field_down! with broadcasts of reshaped MPIStateArrays #1301

Replace copy_stack_field_down! with broadcasts of reshaped MPIStateArrays #1301

Conversation

blallen commented Jun 29, 2020 • edited Loading

Description

For review by CLIMA developers

blallen commented Jun 29, 2020

blallen commented Jun 29, 2020

bors bot commented Jun 29, 2020

try

blallen commented Jun 29, 2020

jkozdon left a comment

Choose a reason for hiding this comment

jkozdon commented Jun 29, 2020

leios commented Jun 29, 2020

jkozdon commented Jun 29, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jkozdon commented Jun 29, 2020

leios commented Jun 29, 2020

jkozdon commented Jun 29, 2020

blallen commented Jun 29, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

blallen commented Jun 29, 2020

bors bot commented Jun 29, 2020

try

blallen commented Jun 30, 2020

bors bot commented Jun 30, 2020

try

blallen commented Jun 30, 2020

try

leios commented Jun 30, 2020

blallen commented Jun 30, 2020

blallen commented Jun 30, 2020

jkozdon commented Jun 30, 2020

bors bot commented Jun 30, 2020

try

blallen commented Jun 30, 2020

blallen commented Jun 30, 2020

jkozdon commented Jun 30, 2020

blallen commented Jun 30, 2020

jkozdon commented Jun 30, 2020 via email

bors bot commented Jun 30, 2020

charleskawczynski commented Jun 30, 2020

bors bot commented Jul 1, 2020

Replace `copy_stack_field_down!` with broadcasts of reshaped MPIStateArrays #1301

Replace `copy_stack_field_down!` with broadcasts of reshaped MPIStateArrays #1301

blallen commented Jun 29, 2020 •

edited

Loading