Merge #829

829: Improve performance of column operators r=charleskawczynski a=charleskawczynski This PR improves the performance of column operations by: - Adding ``@inline`` and ``@inbounds`` annotations - Fixing some type inference issues that prevented inlining (in `idxin`) Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>
CliMA · Jul 26, 2022 · 50322bd · 50322bd
2 parents 000daee + 2f28f6a
commit 50322bd
Show file tree

Hide file tree

Showing 19 changed files with 565 additions and 393 deletions.
diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl
@@ -137,7 +137,7 @@ function Base.copyto!(dest::D, src::D) where {D <: AbstractData}
     return dest
 end
 
-function ncomponents(data::AbstractData{S}) where {S}
+@inline function ncomponents(data::AbstractData{S}) where {S}
     typesize(eltype(parent(data)), S)
 end
 
@@ -607,7 +607,7 @@ end
     @boundscheck (1 <= i <= Nij && 1 <= j <= Nij) ||
                  throw(BoundsError(data, (i, j)))
     dataview = @inbounds view(parent(data), i, j, :)
-    get_struct(dataview, S)
+    @inbounds get_struct(dataview, S)
 end
 
 @inline function Base.setindex!(
@@ -711,7 +711,7 @@ end
 @inline function Base.getindex(data::IF{S, Ni}, i::Integer) where {S, Ni}
     @boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,)))
     dataview = @inbounds view(parent(data), i, :)
-    get_struct(dataview, S)
+    @inbounds get_struct(dataview, S)
 end
 
 @inline function Base.setindex!(data::IF{S, Ni}, val, i::Integer) where {S, Ni}
@@ -950,15 +950,15 @@ end
 end
 
 @propagate_inbounds function Base.getindex(data::VIJFH, I::CartesianIndex{5})
-    data[I[1], I[2], I[4]]
+    @inbounds data[I[1], I[2], I[4]]
 end
 
 @propagate_inbounds function Base.setindex!(
     data::VIJFH,
     val,
     I::CartesianIndex{5},
 )
-    data[I[1], I[2], I[4]] = val
+    @inbounds data[I[1], I[2], I[4]] = val
 end
 
 function gather(

diff --git a/src/DataLayouts/broadcast.jl b/src/DataLayouts/broadcast.jl
@@ -183,7 +183,7 @@ end
     bc::Base.Broadcast.Broadcasted{DS},
     inds...,
 ) where {N, DS <: Union{Data1DXStyle{N}, Data2DXStyle{N}}}
-    _args = column_args(bc.args, inds...)
+    @inbounds _args = column_args(bc.args, inds...)
     _axes = nothing
     Base.Broadcast.Broadcasted{DataColumnStyle(DS)}(bc.f, _args, _axes)
 end
@@ -200,15 +200,15 @@ end
     i,
     h,
 )
-    slab(bc, h)[i]
+    @inbounds slab(bc, h)[i]
 end
 @propagate_inbounds function column(
     bc::Union{Data1D, Base.Broadcast.Broadcasted{<:Data1D}},
     i,
     j,
     h,
 )
-    slab(bc, h)[i]
+    @inbounds slab(bc, h)[i]
 end
 
 @propagate_inbounds function column(
@@ -217,7 +217,7 @@ end
     j,
     h,
 )
-    slab(bc, h)[i, j]
+    @inbounds slab(bc, h)[i, j]
 end
 
 function Base.similar(

diff --git a/src/DataLayouts/struct.jl b/src/DataLayouts/struct.jl
@@ -119,10 +119,13 @@ promote_parent_array_type(
 Create an instance of type `T` from a tuple of field values `args`, bypassing
 possible internal constructors. `T` should be a concrete type.
 """
-@generated function bypass_constructor(::Type{T}, args) where {T}
+Base.@propagate_inbounds @generated function bypass_constructor(
+    ::Type{T},
+    args,
+) where {T}
     vars = ntuple(_ -> gensym(), fieldcount(T))
     assign = [
-        :($var::$(fieldtype(T, i)) = getfield(args, $i)) for
+        :(@inbounds $var::$(fieldtype(T, i)) = getfield(args, $i)) for
         (i, var) in enumerate(vars)
     ]
     construct = Expr(:new, :T, vars...)
@@ -166,85 +169,76 @@ end
 
 Construct an object of type `S` from the values of `array`, optionally offset by `offset` from the start of the array.
 """
-function get_struct(array::AbstractArray{T}, ::Type{S}, offset) where {T, S}
-    if @generated
-        tup = :(())
-        for i in 1:fieldcount(S)
-            push!(
-                tup.args,
-                :(get_struct(
-                    array,
-                    fieldtype(S, $i),
-                    offset + fieldtypeoffset(T, S, $i),
-                )),
-            )
-        end
-        return quote
-            Base.@_propagate_inbounds_meta
-            bypass_constructor(S, $tup)
-        end
-    else
+Base.@propagate_inbounds @generated function get_struct(
+    array::AbstractArray{T},
+    ::Type{S},
+    offset,
+) where {T, S}
+    tup = :(())
+    for i in 1:fieldcount(S)
+        push!(
+            tup.args,
+            :(get_struct(
+                array,
+                fieldtype(S, $i),
+                offset + fieldtypeoffset(T, S, $i),
+            )),
+        )
+    end
+    return quote
         Base.@_propagate_inbounds_meta
-        args = ntuple(fieldcount(S)) do i
-            get_struct(array, fieldtype(S, i), offset + fieldtypeoffset(T, S, i))
-        end
-        return bypass_constructor(S, args)
+        @inbounds bypass_constructor(S, $tup)
     end
 end
 
 # recursion base case: hit array type is the same as the struct leaf type
-@propagate_inbounds function get_struct(
+Base.@propagate_inbounds function get_struct(
     array::AbstractArray{S},
     ::Type{S},
     offset,
 ) where {S}
-    return array[offset + 1]
+    return @inbounds array[offset + 1]
 end
 
-@inline function get_struct(array::AbstractArray{T}, ::Type{S}) where {T, S}
+Base.@propagate_inbounds function get_struct(
+    array::AbstractArray{T},
+    ::Type{S},
+) where {T, S}
     @inbounds get_struct(array, S, 0)
 end
 
-function set_struct!(array::AbstractArray{T}, val::S, offset) where {T, S}
-    if @generated
-        ex = quote
-            Base.@_propagate_inbounds_meta
-        end
-        for i in 1:fieldcount(S)
-            push!(
-                ex.args,
-                :(set_struct!(
-                    array,
-                    getfield(val, $i),
-                    offset + fieldtypeoffset(T, S, $i),
-                )),
-            )
-        end
-        push!(ex.args, :(return val))
-        return ex
-    else
+Base.@propagate_inbounds @generated function set_struct!(
+    array::AbstractArray{T},
+    val::S,
+    offset,
+) where {T, S}
+    ex = quote
         Base.@_propagate_inbounds_meta
-        for i in 1:fieldcount(S)
-            set_struct!(
+    end
+    for i in 1:fieldcount(S)
+        push!(
+            ex.args,
+            :(set_struct!(
                 array,
-                getfield(val, i),
-                offset + fieldtypeoffset(T, S, i),
-            )
-        end
-        return val
+                getfield(val, $i),
+                offset + fieldtypeoffset(T, S, $i),
+            )),
+        )
     end
+    push!(ex.args, :(return val))
+    return ex
 end
 
-@propagate_inbounds function set_struct!(
+Base.@propagate_inbounds function set_struct!(
     array::AbstractArray{S},
     val::S,
     offset,
 ) where {S}
-    array[offset + 1] = val
+    @inbounds array[offset + 1] = val
     val
 end
 
-@inline function set_struct!(array, val)
+Base.@propagate_inbounds function set_struct!(array, val)
     @inbounds set_struct!(array, val, 0)
 end
 

diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -129,9 +129,12 @@ const ColumnField{V, S} =
 slab(field::Field, inds...) =
     Field(slab(field_values(field), inds...), slab(axes(field), inds...))
 
-column(field::Field, inds...) =
-    Field(column(field_values(field), inds...), column(axes(field), inds...))
-column(field::FiniteDifferenceField, inds...) = field
+@inline function column(field::Field, inds...)
+    @inbounds vals = column(field_values(field), inds...)
+    @inbounds spaces = column(axes(field), inds...)
+    @inbounds Field(vals, spaces)
+end
+@inline column(field::FiniteDifferenceField, inds...) = field
 
 
 

diff --git a/src/Fields/broadcast.jl b/src/Fields/broadcast.jl
@@ -92,14 +92,14 @@ function slab(
     Base.Broadcast.Broadcasted{Style}(bc.f, _args, _axes)
 end
 
-function column(
+@inline function column(
     bc::Base.Broadcast.Broadcasted{Style},
     i,
     j,
     h,
 ) where {Style <: AbstractFieldStyle}
-    _args = column_args(bc.args, i, j, h)
-    _axes = column(axes(bc), i, j, h)
+    @inbounds _args = column_args(bc.args, i, j, h)
+    @inbounds _axes = column(axes(bc), i, j, h)
     Base.Broadcast.Broadcasted{Style}(bc.f, _args, _axes)
 end
 

diff --git a/src/Fields/indices.jl b/src/Fields/indices.jl
@@ -17,16 +17,22 @@ end
 
 Base.getindex(field::Field, colidx::ColumnIndex) = column(field, colidx)
 
-function column(field::SpectralElementField1D, colidx::ColumnIndex{1})
+@inline function column(field::SpectralElementField1D, colidx::ColumnIndex{1})
     column(field, colidx.ij[1], colidx.h)
 end
-function column(field::ExtrudedFiniteDifferenceField, colidx::ColumnIndex{1})
+@inline function column(
+    field::ExtrudedFiniteDifferenceField,
+    colidx::ColumnIndex{1},
+)
     column(field, colidx.ij[1], colidx.h)
 end
-function column(field::SpectralElementField2D, colidx::ColumnIndex{2})
+@inline function column(field::SpectralElementField2D, colidx::ColumnIndex{2})
     column(field, colidx.ij[1], colidx.ij[2], colidx.h)
 end
-function column(field::ExtrudedFiniteDifferenceField, colidx::ColumnIndex{2})
+@inline function column(
+    field::ExtrudedFiniteDifferenceField,
+    colidx::ColumnIndex{2},
+)
     column(field, colidx.ij[1], colidx.ij[2], colidx.h)
 end