Try #829:

src/DataLayouts/struct.jl

@@ -193,12 +193,12 @@ function get_struct(array::AbstractArray{T}, ::Type{S}, offset) where {T, S}
 end
 
 # recursion base case: hit array type is the same as the struct leaf type
-@propagate_inbounds function get_struct(
+@inline 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}
@@ -235,12 +235,8 @@ function set_struct!(array::AbstractArray{T}, val::S, offset) where {T, S}
     end
 end
 
-@propagate_inbounds function set_struct!(
-    array::AbstractArray{S},
-    val::S,
-    offset,
-) where {S}
-    array[offset + 1] = val
+@inline function set_struct!(array::AbstractArray{S}, val::S, offset) where {S}
+    @inbounds array[offset + 1] = val
     val
 end
 

src/Fields/Fields.jl

@@ -129,9 +129,9 @@ const ColumnField{V, S} =
 slab(field::Field, inds...) =
     Field(slab(field_values(field), inds...), slab(axes(field), inds...))
 
-column(field::Field, inds...) =
+@inline column(field::Field, inds...) =
     Field(column(field_values(field), inds...), column(axes(field), inds...))
-column(field::FiniteDifferenceField, inds...) = field
+@inline column(field::FiniteDifferenceField, inds...) = field
 
 
 

src/Fields/broadcast.jl

@@ -92,7 +92,7 @@ function slab(
     Base.Broadcast.Broadcasted{Style}(bc.f, _args, _axes)
 end
 
-function column(
+@inline function column(
     bc::Base.Broadcast.Broadcasted{Style},
     i,
     j,

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
 

src/Geometry/axistensors.jl

@@ -27,19 +27,19 @@ ClimaCore.Geometry.CartesianAxis{(1, 2, 3)}()
 function dual end
 
 struct CovariantAxis{I} <: AbstractAxis{I} end
-symbols(::CovariantAxis) = (:u₁, :u₂, :u₃)
+@inline symbols(::CovariantAxis) = (:u₁, :u₂, :u₃)
 
 struct ContravariantAxis{I} <: AbstractAxis{I} end
-symbols(::ContravariantAxis) = (:u¹, :u², :u³)
+@inline symbols(::ContravariantAxis) = (:u¹, :u², :u³)
 dual(::CovariantAxis{I}) where {I} = ContravariantAxis{I}()
 dual(::ContravariantAxis{I}) where {I} = CovariantAxis{I}()
 
 struct LocalAxis{I} <: AbstractAxis{I} end
-symbols(::LocalAxis) = (:u, :v, :w)
+@inline symbols(::LocalAxis) = (:u, :v, :w)
 dual(::LocalAxis{I}) where {I} = LocalAxis{I}()
 
 struct CartesianAxis{I} <: AbstractAxis{I} end
-symbols(::CartesianAxis) = (:u1, :u2, :u3)
+@inline symbols(::CartesianAxis) = (:u1, :u2, :u3)
 dual(::CartesianAxis{I}) where {I} = CartesianAxis{I}()
 
 coordinate_axis(::Type{<:XPoint}) = (1,)
@@ -59,21 +59,46 @@ coordinate_axis(::Type{<:LatLongPoint}) = (1, 2)
 
 coordinate_axis(coord::AbstractPoint) = coordinate_axis(typeof(coord))
 
-@inline function idxin(I::Tuple{Vararg{Int}}, i::Int)
-    N = length(I)
-    for n in 1:N
-        if I[n] == i
-            return n
+@inline idxin(I::Tuple{Int}, i::Int) = 1
+
+@inline function idxin(I::Tuple{Int, Int}, i::Int)
+    @inbounds begin
+        if I[1] == i
+            return 1
+        else
+            return 2
+        end
+    end
+end
+
+@inline function idxin(I::Tuple{Int, Int, Int}, i::Int)
+    @inbounds begin
+        if I[1] == i
+            return 1
+        elseif I[2] == i
+            return 2
+        else
+            return 3
         end
     end
-    return nothing
 end
 
+#= For avoiding JET failures =#
+error_on_no_name_found() = true
+
 @inline function symidx(ax::AbstractAxis{I}, name::Symbol) where {I}
     S = symbols(ax)
-    name == S[1] ? idxin(I, 1) :
-    name == S[2] ? idxin(I, 2) :
-    name == S[3] ? idxin(I, 3) : error("$ax has no symbol $name")
+    if name == S[1]
+        return idxin(I, 1)
+    elseif name == S[2]
+        return idxin(I, 2)
+    elseif name == S[3]
+        return idxin(I, 3)
+    elseif error_on_no_name_found()
+        error("$ax has no symbol $name")
+    else
+        return -1 # for type stability
+    end
 end
 
 # most of these are required for printing
@@ -222,9 +247,13 @@ const CartesianVector{T, I, S} = AxisVector{T, CartesianAxis{I}, S}
 const LocalVector{T, I, S} = AxisVector{T, LocalAxis{I}, S}
 
 Base.propertynames(x::AxisVector) = symbols(axes(x, 1))
-function Base.getproperty(x::AxisVector, name::Symbol)
+@inline function Base.getproperty(x::AxisVector, name::Symbol)
     n = symidx(axes(x, 1), name)
-    isnothing(n) ? zero(eltype(x)) : components(x)[n]
+    if isnothing(n)
+        zero(eltype(x))
+    else
+        @inbounds components(x)[n]
+    end
 end
 
 
@@ -349,7 +378,7 @@ function Base.:(-)(A::Axis2Tensor, b::LinearAlgebra.UniformScaling)
     AxisTensor(axes(A), components(A) - b)
 end
 
-function _transform(
+@inline function _transform(
     ato::Ato,
     x::AxisVector{T, Afrom, SVector{N, T}},
 ) where {Ato <: AbstractAxis{I}, Afrom <: AbstractAxis{I}} where {I, T, N}
@@ -417,7 +446,7 @@ end
     return :(AxisVector(ato, SVector($(vals...))))
 end
 
-function _transform(
+@inline function _transform(
     ato::Ato,
     x::Axis2Tensor{T, Tuple{Afrom, A2}},
 ) where {
@@ -508,10 +537,11 @@ end
     ))
 end
 
-transform(ato::CovariantAxis, v::CovariantTensor) = _transform(ato, v)
-transform(ato::ContravariantAxis, v::ContravariantTensor) = _transform(ato, v)
-transform(ato::CartesianAxis, v::CartesianTensor) = _transform(ato, v)
-transform(ato::LocalAxis, v::LocalTensor) = _transform(ato, v)
+@inline transform(ato::CovariantAxis, v::CovariantTensor) = _transform(ato, v)
+@inline transform(ato::ContravariantAxis, v::ContravariantTensor) =
+    _transform(ato, v)
+@inline transform(ato::CartesianAxis, v::CartesianTensor) = _transform(ato, v)
+@inline transform(ato::LocalAxis, v::LocalTensor) = _transform(ato, v)
 
 project(ato::CovariantAxis, v::CovariantTensor) = _project(ato, v)
 project(ato::ContravariantAxis, v::ContravariantTensor) = _project(ato, v)

src/Geometry/conversions.jl

@@ -133,36 +133,41 @@ LocalVector(u::ContravariantVector{<:Any, (3,)}, ::LocalGeometry{(1, 2)}) =
     AxisVector(WAxis(), components(u))
 
 
-covariant1(u::AxisVector, local_geometry::LocalGeometry) =
+@inline covariant1(u::AxisVector, local_geometry::LocalGeometry) =
     CovariantVector(u, local_geometry).u₁
-covariant2(u::AxisVector, local_geometry::LocalGeometry) =
+@inline covariant2(u::AxisVector, local_geometry::LocalGeometry) =
     CovariantVector(u, local_geometry).u₂
-covariant3(u::AxisVector, local_geometry::LocalGeometry) =
+@inline covariant3(u::AxisVector, local_geometry::LocalGeometry) =
     CovariantVector(u, local_geometry).u₃
 
-contravariant1(u::AxisVector, local_geometry::LocalGeometry) =
+@inline contravariant1(u::AxisVector, local_geometry::LocalGeometry) =
     transform(Contravariant123Axis(), u, local_geometry).u¹
-contravariant2(u::AxisVector, local_geometry::LocalGeometry) =
+@inline contravariant2(u::AxisVector, local_geometry::LocalGeometry) =
     transform(Contravariant123Axis(), u, local_geometry).u²
-contravariant3(u::AxisVector, local_geometry::LocalGeometry) =
+@inline contravariant3(u::AxisVector, local_geometry::LocalGeometry) =
     transform(Contravariant123Axis(), u, local_geometry).u³
 
-contravariant1(u::Axis2Tensor, local_geometry::LocalGeometry) =
-    transform(Contravariant123Axis(), u, local_geometry)[1, :]
-contravariant2(u::Axis2Tensor, local_geometry::LocalGeometry) =
-    transform(Contravariant123Axis(), u, local_geometry)[2, :]
-contravariant3(u::Axis2Tensor, local_geometry::LocalGeometry) =
-    transform(Contravariant123Axis(), u, local_geometry)[3, :]
+@inline function contravariant1(u::Axis2Tensor, local_geometry::LocalGeometry)
+    @inbounds transform(Contravariant123Axis(), u, local_geometry)[1, :]
+end
+@inline function contravariant2(u::Axis2Tensor, local_geometry::LocalGeometry)
+    @inbounds transform(Contravariant123Axis(), u, local_geometry)[2, :]
+end
+@inline function contravariant3(u::Axis2Tensor, local_geometry::LocalGeometry)
+    @inbounds transform(Contravariant123Axis(), u, local_geometry)[3, :]
+end
 
-Jcontravariant3(u::AxisTensor, local_geometry::LocalGeometry) =
+@inline Jcontravariant3(u::AxisTensor, local_geometry::LocalGeometry) =
     local_geometry.J * contravariant3(u, local_geometry)
 
 # required for curl-curl
-covariant3(u::Contravariant3Vector, local_geometry::LocalGeometry{(1, 2)}) =
-    contravariant3(u, local_geometry)
+@inline covariant3(
+    u::Contravariant3Vector,
+    local_geometry::LocalGeometry{(1, 2)},
+) = contravariant3(u, local_geometry)
 
 # workarounds for using a Covariant12Vector/Covariant123Vector in a UW space:
-function LocalVector(
+@inline function LocalVector(
     vector::CovariantVector{<:Any, (1, 2, 3)},
     local_geometry::LocalGeometry{(1, 3)},
 )
@@ -171,23 +176,23 @@ function LocalVector(
     u, w = components(transform(LocalAxis{(1, 3)}(), vector2, local_geometry))
     return UVWVector(u, v, w)
 end
-function contravariant1(
+@inline function contravariant1(
     vector::CovariantVector{<:Any, (1, 2, 3)},
     local_geometry::LocalGeometry{(1, 3)},
 )
     u₁, _, u₃ = components(vector)
     vector2 = Covariant13Vector(u₁, u₃)
     return transform(Contravariant13Axis(), vector2, local_geometry).u¹
 end
-function contravariant3(
+@inline function contravariant3(
     vector::CovariantVector{<:Any, (1, 2)},
     local_geometry::LocalGeometry{(1, 3)},
 )
     u₁, _ = components(vector)
     vector2 = Covariant13Vector(u₁, zero(u₁))
     return transform(Contravariant13Axis(), vector2, local_geometry).u³
 end
-function ContravariantVector(
+@inline function ContravariantVector(
     vector::CovariantVector{<:Any, (1, 2)},
     local_geometry::LocalGeometry{(1, 3)},
 )

src/Geometry/localgeometry.jl

@@ -17,7 +17,7 @@ struct LocalGeometry{I, C <: AbstractPoint, FT, S}
     ∂ξ∂x::Axis2Tensor{FT, Tuple{ContravariantAxis{I}, LocalAxis{I}}, S}
 end
 
-LocalGeometry(coordinates, J, WJ, ∂x∂ξ) =
+@inline LocalGeometry(coordinates, J, WJ, ∂x∂ξ) =
     LocalGeometry(coordinates, J, WJ, ∂x∂ξ, inv(∂x∂ξ))
 
 """