Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds support for using Base.mapreducedim! with AbstractField and AbstractOperations #1611

Merged
merged 67 commits into from
Apr 27, 2021

Conversation

glwagner
Copy link
Member

@glwagner glwagner commented Apr 24, 2021

This PR adds support for in-place reduction operations where the target is an AbstractReducedField and the source is either an AbstractDataField, AbstractOperation, or array of some kind. These work on the GPU. In the end the solution was simple since we now subtype AbstractArray; we only need to pass a view into the interior indices of the target to Base.mapreducedim! (which on the GPU ends up at GPUArrays.mapreducedim!).

The result is that AveragedField(op::AbstractOperation) no longer needs to allocate memory for the three-dimensional result of computing op. Instead, op is reduced in a kernel. This is both faster (much much faster, I think --- though a benchmark is a good idea) and more memory efficient. It also greatly simplifies compute!(field::AveragedField):

function compute!(avg::AveragedField, time=nothing)
    compute_at!(avg.operand, time)
    mean!(avg, operand)
    return nothing
end

So, resolves #1422.

I also took the liberty of resolving #1610 and nuking interiorparent...

@glwagner glwagner changed the title Adds support for using Base.mapreducedim! with AbstractField and AbstractOperations Adds support for using Base.mapreducedim! with AbstractField and AbstractOperations Apr 24, 2021
Copy link
Collaborator

@tomchor tomchor left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like a very positive change

using Oceananigans.Fields: architecture, tracernames
using Oceananigans.Architectures: device
using Oceananigans.Utils: work_layout
using Oceananigans.Utils: work_layout, new_data
using Oceananigans.Grids: new_data
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't know if it's relevant, just caught the fact that new_data is being imported from two different modules. Is that correct?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not correct -- nice catch!

mean!(avg.data.parent, operand_sans_halos)

avg .= 0 # should we have to do this?
mean!(avg, avg.operand)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Much simpler!

src/Forcings/continuous_forcing.jl Outdated Show resolved Hide resolved
@glwagner glwagner merged commit 6e6a1b8 into master Apr 27, 2021
@glwagner glwagner deleted the glw/mapreduce-abstract-ops branch April 27, 2021 16:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Passing abstract operations (and ComputedField?) to mapreduce
2 participants