Massive clean up for conditional operations

[glwagner]

This PR greatly simplifies the logic behind conditional_operation and how its extended for immersd boundary grid. The end result besides source code clean up seems to be an improvement in type inference.

Resolves #3750 I think. But @ali-ramadhan please test.

EDIT: I decided not to work on #3791 (here) because this requires fixing a bunch of tests and is a bigger effort.

[glwagner]

Here's an example:

using Oceananigans

Nx, Ny, Nz = 100, 100, 100
latitude = longitude = z = (0, 1)
underlying_grid = LatitudeLongitudeGrid(size=(Nx, Ny, Nz); latitude, longitude, z)
grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom((λ, φ) -> 0.5))

ci = CenterField(grid)
ciw = view(ci, 1:Nx, 1:Ny, 1:Nz)

cu = CenterField(underlying_grid)
cuw = view(cu, 1:Nx, 1:Ny, 1:Nz)

for n = 1:10
    @time minimum(ci)
    @time minimum(ciw)
    @time minimum(cu)
    @time minimum(cuw)
end

Previously reducing ciw was slow, basically because getindex for ConditionalOperand was not type inferred. This PR fixes that and things speed up. This PR could also improve reduction on the GPU.

The basic lesson is to use explicit dispatch with type-based information, rather than using functions like truefunc and identity.

[glwagner]

AUUUGHHH that was hard but finally found the source of the type instability:

Oceananigans.jl/src/AbstractOperations/AbstractOperations.jl

Lines 37 to 44 in 4dbec26

	function Base.axes(f::AbstractOperation)
	idx = indices(f)
	if idx === (:, : ,:)
	return Base.OneTo.(size(f))
	else
	return Tuple(idx[i] isa Colon ? Base.OneTo(size(f, i)) : idx[i] for i = 1:3)
	end
	end

(axes is called with mapreducedim! at some point)

fixing that finally made the difference. Just required generalizing the one we already have for field (supposed to fix the same problem, except only for Field and not abstract operations):

Oceananigans.jl/src/Fields/field.jl

Lines 401 to 413 in 4dbec26

	@inline axis(::Colon, N) = Base.OneTo(N)
	@inline axis(index::UnitRange, N) = index
	@inline function Base.axes(f::Field)
	Nx, Ny, Nz = size(f)
	ix, iy, iz = f.indices
	ax = axis(ix, Nx)
	ay = axis(iy, Ny)
	az = axis(iz, Nz)
	return (ax, ay, az)
	end

[navidcy]

AUUUGHHH that was hard but finally found the source of the type instability:

which type instability; is this related to discussion in #3750?

[glwagner]

AUUUGHHH that was hard but finally found the source of the type instability:

which type instability; is this related to discussion in #3750?

Yes. That issue documents slow reductions for windowed fields on immersed boundary grids. I hypothesized that it was due to a failure of type inference. Looking into it further I see that axes(op::AbstractOperation) cannot be type inferred when indices is not (:, :, :) because of the tuple generator. We found the same problem with axes for Field and fixed it but didn't fix it for AbstractOperations. Reducing windowed fields on immersed boundary grids requires this because they are wrapped in ConditionalOperation in order to mask the immersed regions during the reduction. This PR extends the fix we implemented for Field to also encompass AbstractOperation. It also cleans up conditional operations quite a bit.

[glwagner]

It's not really a bug btw

[glwagner]

But definitely clean up and performance. Just to clarify fixing type inference doesn't change the result of the reduction, it just makes it go much faster

[ali-ramadhan]

Was OOO last week but can review this later today!

This PR does resolve the slowness in #3750!

julia> @time minimum(u2)
  0.059865 seconds (665 allocations: 54.930 KiB)
0.0

julia> @time minimum(u2.data)
  0.011223 seconds (3 allocations: 1.562 KiB)
0.0

Some overhead which is expected I think, but way better than 2000x slower lol.

[glwagner]

@ali-ramadhan let's focus first on #3801 and then we can revisit this PR. There are some additional challenges to resolving #3791 unfortunately that will also require some test refactoring.