Merge #917

917: Do inlining and inbounds more carefully r=charleskawczynski a=charleskawczynski This PR is an attempt to partially revert (and complete) #829, by: - Exchanging the use of ``@inline`` and ``@inbounds`` - Removing ``@inline`` for "expensive" functions that do not have or uses indices which require boundscheck elision - Applying inlining and ``@inbounds`` to slab functions The only thing that I'm not 100% sure about is the ``@inbounds`` around the threading loops, does what I've done look okay? cc `@simonbyrne` Supersedes #891. Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>
CliMA · Aug 26, 2022 · 3d15515 · 3d15515
2 parents 208d8d8 + 21d9bfc
commit 3d15515
Show file tree

Hide file tree

Showing 21 changed files with 587 additions and 420 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -108,6 +108,7 @@ steps:
       - "julia -O0 --color=yes --check-bounds=yes --project=test test/Operators/finitedifference/implicit_stencils.jl --float_type Float32"
     agents:
       config: cpu
+      slurm_mem: 30GB
       queue: central
       slurm_ntasks: 1
 
@@ -117,6 +118,7 @@ steps:
       - "julia -O0 --color=yes --check-bounds=yes --project=test test/Operators/finitedifference/implicit_stencils.jl --float_type Float64"
     agents:
       config: cpu
+      slurm_mem: 30GB
       queue: central
       slurm_ntasks: 1
 
@@ -126,6 +128,7 @@ steps:
       - "julia --color=yes --check-bounds=yes --project=test test/Operators/finitedifference/opt_implicit_stencils.jl"
     agents:
       config: cpu
+      slurm_mem: 30GB
       queue: central
       slurm_ntasks: 1
 

diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl
@@ -270,7 +270,7 @@ end
 
 function Base.fill!(data::IJFH, val)
     (_, _, _, _, Nh) = size(data)
-    for h in 1:Nh
+    @inbounds for h in 1:Nh
         fill!(slab(data, h), val)
     end
     return data
@@ -408,7 +408,7 @@ end
 
 function Base.fill!(data::IFH, val)
     (_, _, _, _, Nh) = size(data)
-    for h in 1:Nh
+    @inbounds for h in 1:Nh
         fill!(slab(data, h), val)
     end
     return data
@@ -419,15 +419,16 @@ end
     dataview = @inbounds view(parent(data), :, :, h)
     IF{S, Ni}(dataview)
 end
-slab(data::IFH, v::Integer, h::Integer) = slab(data, h)
+Base.@propagate_inbounds slab(data::IFH, v::Integer, h::Integer) = slab(data, h)
 
 @inline function column(data::IFH{S, Ni}, i, h) where {S, Ni}
     @boundscheck (1 <= h <= length(data) && 1 <= i <= Ni) ||
                  throw(BoundsError(data, (i, h)))
     dataview = @inbounds view(parent(data), i, :, h)
     DataF{S}(dataview)
 end
-@inline column(data::IFH{S, Ni}, i, j, h) where {S, Ni} = column(data, i, h)
+Base.@propagate_inbounds column(data::IFH{S, Ni}, i, j, h) where {S, Ni} =
+    column(data, i, h)
 
 @generated function _property_view(
     data::IFH{S, Ni, A},
@@ -480,7 +481,7 @@ function DataF{S}(ArrayType) where {S}
 end
 
 function Base.fill!(data::DataF, val)
-    data[] = val
+    @inbounds data[] = val
     return data
 end
 @inline function Base.getproperty(data::DataF{S}, i::Integer) where {S}
@@ -505,15 +506,15 @@ end
     return :(DataF{$SS}(@inbounds view(parent(data), $field_byterange)))
 end
 
-@inline function Base.getindex(data::DataF{S}) where {S}
+Base.@propagate_inbounds function Base.getindex(data::DataF{S}) where {S}
     @inbounds get_struct(parent(data), S)
 end
 
 @propagate_inbounds function Base.getindex(col::Data0D, I::CartesianIndex{5})
-    col[]
+    @inbounds col[]
 end
 
-@inline function Base.setindex!(data::DataF{S}, val) where {S}
+Base.@propagate_inbounds function Base.setindex!(data::DataF{S}, val) where {S}
     @inbounds set_struct!(parent(data), convert(S, val))
 end
 
@@ -522,7 +523,7 @@ end
     val,
     I::CartesianIndex{5},
 )
-    col[] = val
+    @inbounds col[] = val
 end
 
 Base.copy(data::DataF{S}) where {S} = DataF{S}(copy(parent(data)))
@@ -535,15 +536,15 @@ Base.copy(data::DataF{S}) where {S} = DataF{S}(copy(parent(data)))
     slab::DataSlab2D{S},
     I::CartesianIndex,
 ) where {S}
-    slab[I[1], I[2]]
+    @inbounds slab[I[1], I[2]]
 end
 
 @propagate_inbounds function Base.setindex!(
     slab::DataSlab2D{S},
     val,
     I::CartesianIndex,
 ) where {S}
-    slab[I[1], I[2]] = val
+    @inbounds slab[I[1], I[2]] = val
 end
 
 Base.size(::DataSlab2D{S, Nij}) where {S, Nij} = (Nij, Nij, 1, 1, 1)
@@ -589,8 +590,8 @@ function Base.size(data::IJF{S, Nij}) where {S, Nij}
     return (Nij, Nij, 1, 1, 1)
 end
 function Base.fill!(data::IJF{S, Nij}, val) where {S, Nij}
-    for j in 1:Nij, i in 1:Nij
-        @inbounds data[i, j] = val
+    @inbounds for j in 1:Nij, i in 1:Nij
+        data[i, j] = val
     end
     return data
 end
@@ -654,15 +655,15 @@ end
 # ======================
 
 @propagate_inbounds function Base.getindex(slab::DataSlab1D, I::CartesianIndex)
-    slab[I[1]]
+    @inbounds slab[I[1]]
 end
 
 @propagate_inbounds function Base.setindex!(
     slab::DataSlab1D,
     val,
     I::CartesianIndex,
 )
-    slab[I[1]] = val
+    @inbounds slab[I[1]] = val
 end
 
 function Base.size(::DataSlab1D{<:Any, Ni}) where {Ni}
@@ -707,8 +708,8 @@ function replace_basetype(data::IF{S, Ni}, ::Type{T}) where {S, Ni, T}
 end
 
 function Base.fill!(data::IF{S, Ni}, val) where {S, Ni}
-    for i in 1:Ni
-        @inbounds data[i] = val
+    @inbounds for i in 1:Ni
+        data[i] = val
     end
     return data
 end
@@ -803,8 +804,8 @@ function Base.size(data::VF)
 end
 function Base.fill!(data::VF, val)
     Nv = size(parent(data), 1)
-    for v in 1:Nv
-        @inbounds data[v] = val
+    @inbounds for v in 1:Nv
+        data[v] = val
     end
     return data
 end
@@ -839,15 +840,15 @@ end
     col::DataColumn,
     I::CartesianIndex{5},
 )
-    col[I[4]]
+    @inbounds col[I[4]]
 end
 
 @propagate_inbounds function Base.setindex!(
     col::DataColumn,
     val,
     I::CartesianIndex{5},
 )
-    col[I[4]] = val
+    @inbounds col[I[4]] = val
 end
 
 @inline function Base.setindex!(data::VF{S}, val, v::Integer) where {S}
@@ -919,7 +920,7 @@ end
 
 function Base.fill!(data::VIJFH, val)
     (Ni, Nj, _, Nv, Nh) = size(data)
-    for h in 1:Nh, v in 1:Nv
+    @inbounds for h in 1:Nh, v in 1:Nv
         fill!(slab(data, v, h), val)
     end
     return data
@@ -1060,7 +1061,7 @@ function Base.length(data::VIFH)
 end
 function Base.fill!(data::VIFH, val)
     (Ni, _, _, Nv, Nh) = size(data)
-    for h in 1:Nh, v in 1:Nv
+    @inbounds for h in 1:Nh, v in 1:Nv
         fill!(slab(data, v, h), val)
     end
     return data

diff --git a/src/DataLayouts/broadcast.jl b/src/DataLayouts/broadcast.jl
@@ -161,7 +161,7 @@ Base.Broadcast.BroadcastStyle(
 
 Base.Broadcast.broadcastable(data::AbstractData) = data
 
-@inline function slab(
+Base.@propagate_inbounds function slab(
     bc::Base.Broadcast.Broadcasted{DS},
     inds...,
 ) where {Ni, DS <: Union{Data1DStyle{Ni}, Data1DXStyle{Ni}}}
@@ -170,7 +170,7 @@ Base.Broadcast.broadcastable(data::AbstractData) = data
     Base.Broadcast.Broadcasted{DataSlab1DStyle(DS)}(bc.f, _args, _axes)
 end
 
-@inline function slab(
+Base.@propagate_inbounds function slab(
     bc::Base.Broadcast.Broadcasted{DS},
     inds...,
 ) where {Nij, DS <: Union{Data2DStyle{Nij}, Data2DXStyle{Nij}}}
@@ -179,11 +179,11 @@ end
     Base.Broadcast.Broadcasted{DataSlab2DStyle(DS)}(bc.f, _args, _axes)
 end
 
-@inline function column(
+Base.@propagate_inbounds function column(
     bc::Base.Broadcast.Broadcasted{DS},
     inds...,
 ) where {N, DS <: Union{Data1DXStyle{N}, Data2DXStyle{N}}}
-    @inbounds _args = column_args(bc.args, inds...)
+    _args = column_args(bc.args, inds...)
     _axes = nothing
     Base.Broadcast.Broadcasted{DataColumnStyle(DS)}(bc.f, _args, _axes)
 end
@@ -195,29 +195,29 @@ end
     bc
 end
 
-@propagate_inbounds function column(
+Base.@propagate_inbounds function column(
     bc::Union{Data1D, Base.Broadcast.Broadcasted{<:Data1D}},
     i,
     h,
 )
-    @inbounds slab(bc, h)[i]
+    slab(bc, h)[i]
 end
-@propagate_inbounds function column(
+Base.@propagate_inbounds function column(
     bc::Union{Data1D, Base.Broadcast.Broadcasted{<:Data1D}},
     i,
     j,
     h,
 )
-    @inbounds slab(bc, h)[i]
+    slab(bc, h)[i]
 end
 
-@propagate_inbounds function column(
+Base.@propagate_inbounds function column(
     bc::Union{Data2D, Base.Broadcast.Broadcasted{<:Data2D}},
     i,
     j,
     h,
 )
-    @inbounds slab(bc, h)[i, j]
+    slab(bc, h)[i, j]
 end
 
 function Base.similar(
@@ -414,14 +414,15 @@ end
     bc = Base.Broadcast.instantiate(
         Base.Broadcast.Broadcasted{Style}(bc.f, bc.args, ()),
     )
-    fill!(dest, bc[])
+    @inbounds bc0 = bc[]
+    fill!(dest, bc0)
 end
 
 @inline function Base.copyto!(
     dest::DataF{S},
     bc::Union{DataF{S, A}, Base.Broadcast.Broadcasted{DataFStyle{A}}},
 ) where {S, A}
-    dest[] = convert(S, bc[])
+    @inbounds dest[] = convert(S, bc[])
     return dest
 end
 
@@ -508,12 +509,14 @@ function _threaded_copyto!(
 ) where {S, Ni, A}
     _, _, _, _, Nh = size(dest)
     # parallelize over elements
-    Threads.@threads for h in 1:Nh
-        # copy contiguous columns
-        @inbounds for i in 1:Ni
-            col_dest = column(dest, i, h)
-            col_bc = column(bc, i, h)
-            copyto!(col_dest, col_bc)
+    @inbounds begin
+        Threads.@threads for h in 1:Nh
+            # copy contiguous columns
+            for i in 1:Ni
+                col_dest = column(dest, i, h)
+                col_bc = column(bc, i, h)
+                copyto!(col_dest, col_bc)
+            end
         end
     end
     return dest
@@ -556,12 +559,14 @@ function _threaded_copyto!(
 ) where {S, Nij, A}
     _, _, _, _, Nh = size(dest)
     # parallelize over elements
-    Threads.@threads for h in 1:Nh
-        # copy contiguous columns
-        @inbounds for j in 1:Nij, i in 1:Nij
-            col_dest = column(dest, i, j, h)
-            col_bc = column(bc, i, j, h)
-            copyto!(col_dest, col_bc)
+    @inbounds begin
+        Threads.@threads for h in 1:Nh
+            # copy contiguous columns
+            for j in 1:Nij, i in 1:Nij
+                col_dest = column(dest, i, j, h)
+                col_bc = column(bc, i, j, h)
+                copyto!(col_dest, col_bc)
+            end
         end
     end
     return dest

diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -126,13 +126,11 @@ const SlabField2D{V, S} = Field{
 const ColumnField{V, S} =
     Field{V, S} where {V <: DataLayouts.DataColumn, S <: Spaces.AbstractSpace}
 
-slab(field::Field, inds...) =
+Base.@propagate_inbounds slab(field::Field, inds...) =
     Field(slab(field_values(field), inds...), slab(axes(field), inds...))
 
-@inline function column(field::Field, inds...)
-    @inbounds vals = column(field_values(field), inds...)
-    @inbounds spaces = column(axes(field), inds...)
-    @inbounds Field(vals, spaces)
+Base.@propagate_inbounds function column(field::Field, inds...)
+    Field(column(field_values(field), inds...), column(axes(field), inds...))
 end
 @inline column(field::FiniteDifferenceField, inds...) = field
 
@@ -356,7 +354,7 @@ function Spaces.create_ghost_buffer(field::Field)
 end
 
 
-function level(
+Base.@propagate_inbounds function level(
     field::Union{
         CenterFiniteDifferenceField,
         CenterExtrudedFiniteDifferenceField,
@@ -367,17 +365,19 @@ function level(
     data = level(field_values(field), v)
     Field(data, hspace)
 end
-function level(
+Base.@propagate_inbounds function level(
     field::Union{FaceFiniteDifferenceField, FaceExtrudedFiniteDifferenceField},
     v::PlusHalf,
 )
     hspace = level(axes(field), v)
-    data = level(field_values(field), v.i + 1)
+    @inbounds data = level(field_values(field), v.i + 1)
     Field(data, hspace)
 end
 
-Base.getindex(field::PointField) = getindex(field_values(field))
-Base.setindex!(field::PointField, val) = setindex!(field_values(field), val)
+Base.@propagate_inbounds Base.getindex(field::PointField) =
+    getindex(field_values(field))
+Base.@propagate_inbounds Base.setindex!(field::PointField, val) =
+    setindex!(field_values(field), val)
 
 """
     set!(f::Function, field::Field, args = ())