Optimizes fill_halo_regions_open.jl

[jagoosw]

I was doing some profiling on a model with no open boundaries and discovered that this function was causing a big slow down. I guess this is because the compiler isn't managing to work out its just a load of nothing operations but this change appears to make it completely go away.

[simone-silvestri]

what was the slowdown?

[glwagner]

Might be good to profile with fill_open_boundary_regions!(fields::NTuple, boundary_conditions, indices, loc, grid, args...; kwargs...) = nothing also to check this

[glwagner]

Suggested change

	fill_open_boundary_regions!(fields::NTuple, boundary_conditions, indices, loc, grid, args...; kwargs...) =
	ntuple(n->fill_open_boundary_regions!(fields[n], boundary_conditions[n], indices, loc[n], grid, args...; kwargs...), Val(length(fields)))
	function fill_open_boundary_regions!(fields::NTuple, bcs, indices, loc, args...; kw...)
	N = length(fields)
	return ntuple(n -> fill_open_boundary_regions!(fields[n], bcs[n], indices, loc[n], args...; kw...), Val(N))
	end

[glwagner]

If we can't read the code, we've little to show for all our hard work

[glwagner]

The profiles should be documented here. I'll tentatively approve but the documentation is important

[jagoosw]

Ah so I've realised this isn't the fix we needed, and I was just hiding it from myself in the profile because I replaced the function by writing it in the REPL. I made an MWE:

using Oceananigans

grid = RectilinearGrid(GPU(), topology = (Flat, Flat, Bounded), size = (100, ), extent = (400, ))

model = HydrostaticFreeSurfaceModel(; grid, velocities = PrescribedVelocityFields(), momentum_advection=nothing, buoyancy=nothing, tracers = ntuple(n->Symbol(:T, n), Val(30)))

You can see from this profile that `fill_open_boundary_regions!` takes a lot longer than `fill_halo_event!`, even though there are no velocity open boundaries. This is because it is launching a load of zero size kernels where as `fill_halo_event!` just returns nothing instead.

I've fixed this now and get this from the profile instead:

In numbers, the original version benchmarks time_step! at around 4.074 ms ± 581.472 μs and the new version 2.438 ms ± 501.642 μs

[glwagner]

This is a little convoluted. Why do you need a function to determine the size? Why can't we just use loc?

[jagoosw]

This comes from normal halo filling and I think its main purpose is for windowed fields:

Oceananigans.jl/src/BoundaryConditions/fill_halo_regions.jl

Lines 349 to 433 in 4dbec26

	@inline fill_halo_size(::Tuple, ::WEB, args...) = :yz
	@inline fill_halo_size(::Tuple, ::SNB, args...) = :xz
	@inline fill_halo_size(::Tuple, ::TBB, args...) = :xy
	# If indices are colon, and locations are _not_ Nothing, fill the whole boundary plane!
	# If locations are _Nothing_, then the kwarg `reduced_dimensions` will allow the size `:xz`
	# to be correctly interpreted inside `launch!`.
	@inline fill_halo_size(::OffsetArray, ::WEB, ::Tuple{<:Any, <:Colon, <:Colon}, args...) = :yz
	@inline fill_halo_size(::OffsetArray, ::SNB, ::Tuple{<:Colon, <:Any, <:Colon}, args...) = :xz
	@inline fill_halo_size(::OffsetArray, ::TBB, ::Tuple{<:Colon, <:Colon, <:Any}, args...) = :xy
	# If the index is a Colon and the location is _NOT_ a `Nothing` (i.e. not a `ReducedField`),
	# then fill the whole boundary, otherwise fill the size of the corresponding array
	@inline whole_halo(idx, loc) = false
	@inline whole_halo(idx, ::Nothing) = false
	@inline whole_halo(::Colon, ::Nothing) = false
	@inline whole_halo(::Colon, loc) = true
	# Calculate kernel size for windowed fields. This code is only called when
	# one or more of the elements of `idx` is not Colon in the two direction perpendicular
	# to the halo region and `bc` is not `PeriodicBoundaryCondition`.
	@inline function fill_halo_size(c::OffsetArray, ::WEB, idx, bc, loc, grid)
	@inbounds begin
	whole_y_halo = whole_halo(idx[2], loc[2])
	whole_z_halo = whole_halo(idx[3], loc[3])
	end
	_, Ny, Nz = size(grid)
	_, Cy, Cz = size(c)
	Sy = ifelse(whole_y_halo, Ny, Cy)
	Sz = ifelse(whole_z_halo, Nz, Cz)
	return (Sy, Sz)
	end
	@inline function fill_halo_size(c::OffsetArray, ::SNB, idx, bc, loc, grid)
	@inbounds begin
	whole_x_halo = whole_halo(idx[1], loc[1])
	whole_z_halo = whole_halo(idx[3], loc[3])
	end
	Nx, _, Nz = size(grid)
	Cx, _, Cz = size(c)
	Sx = ifelse(whole_x_halo, Nx, Cx)
	Sz = ifelse(whole_z_halo, Nz, Cz)
	return (Sx, Sz)
	end
	@inline function fill_halo_size(c::OffsetArray, ::TBB, idx, bc, loc, grid)
	@inbounds begin
	whole_x_halo = whole_halo(idx[1], loc[1])
	whole_y_halo = whole_halo(idx[2], loc[2])
	end
	Nx, Ny, _ = size(grid)
	Cx, Cy, _ = size(c)
	Sx = ifelse(whole_x_halo, Nx, Cx)
	Sy = ifelse(whole_y_halo, Ny, Cy)
	return (Sx, Sy)
	end
	# Remember that Periodic BCs also fill halo points!
	@inline fill_halo_size(c::OffsetArray, ::WEB, idx, ::PBC, args...) = tuple(size(c, 2), size(c, 3))
	@inline fill_halo_size(c::OffsetArray, ::SNB, idx, ::PBC, args...) = tuple(size(c, 1), size(c, 3))
	@inline fill_halo_size(c::OffsetArray, ::TBB, idx, ::PBC, args...) = tuple(size(c, 1), size(c, 2))
	@inline function fill_halo_size(c::OffsetArray, ::WEB, ::Tuple{<:Any, <:Colon, <:Colon}, ::PBC, args...)
	_, Cy, Cz = size(c)
	return (Cy, Cz)
	end
	@inline function fill_halo_size(c::OffsetArray, ::SNB, ::Tuple{<:Colon, <:Any, <:Colon}, ::PBC, args...)
	Cx, _, Cz = size(c)
	return (Cx, Cz)
	end
	@inline function fill_halo_size(c::OffsetArray, ::TBB, ::Tuple{<:Colon, <:Colon, <:Any}, ::PBC, args...)
	Cx, Cy, _ = size(c)
	return (Cx, Cy)
	end

[glwagner]

Suggested change

	fill_open_halo_event!(open_fill, field, left_bc, right_bc, fill_size, loc, arch, grid, args)
	fill_open_halo!(filling_function, field, left_bc, right_bc, fill_size, loc, arch, grid, args)

[glwagner]

I think it's good to go.

The only thing that doesn't quite make sense to me is why

fill_size = fill_halo_size(field, regular_fill_function, indices, boundary_conditions, loc, grid)

depends on regular_fill_function, since

fill_function, regular_fill_function = get_open_halo_filling_functions(loc) 

and loc is an argument to both functions. It doesn't seem that from a purely logical point of view we need regular_fill_function at all here.

[jagoosw]

I'll stream line those two line before I merge

[simone-silvestri]

Is this PR close to merging?

[tomchor]

cc @ali-ramadhan

[simone-silvestri]

I think this has been superseded by #3834. @jagoosw let us know if it is the case or not.

[jagoosw]

I think this has been superseded by #3834. @jagoosw let us know if it is the case or not.

Yeah definitely superseded!