Add support for nonuniform data structures [perf]

Author: dennisYatunin

docs/src/APIs/datalayouts_api.md

@@ -18,4 +18,14 @@ DataLayouts.IHF
 DataLayouts.IJHF
 DataLayouts.VIHF
 DataLayouts.VIJHF
+DataLayouts.parent_array_type
+DataLayouts.promote_parent_array_type
+DataLayouts.default_basetype
+DataLayouts.check_basetype
+DataLayouts.checked_valid_basetype
+DataLayouts.storage_length
+DataLayouts.struct_field_view
+DataLayouts.set_struct!
+DataLayouts.get_struct
+DataLayouts.bitcast_struct
 ```

docs/src/APIs/utilities_api.md

@@ -7,6 +7,7 @@ CurrentModule = ClimaCore
 ```@docs
 Utilities.PlusHalf
 Utilities.half
+Utilities.replace_type_parameter
 ```
 
 ## Utilities.Cache

ext/cuda/data_layouts.jl

@@ -12,51 +12,27 @@ import ClimaCore.DataLayouts: fused_copyto!
 import Adapt
 import CUDA
 
-parent_array_type(::Type{<:CUDA.CuArray{T, N, B} where {N}}) where {T, B} =
-    CUDA.CuArray{T, N, B} where {N}
-
-# allow on-device use of lazy broadcast objects
+# Ensure that all CuArrays have the same memory buffer type.
 parent_array_type(
-    ::Type{<:CUDA.CuDeviceArray{T, N, A} where {N}},
-) where {T, A} = CUDA.CuDeviceArray{T, N, A} where {N}
-
-# Ensure that both parent array types have the same memory buffer type.
-promote_parent_array_type(
-    ::Type{CUDA.CuArray{T1, N, B} where {N}},
-    ::Type{CUDA.CuArray{T2, N, B} where {N}},
-) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), N, B} where {N}
-
-# allow on-device use of lazy broadcast objects
+    ::Type{<:CUDA.CuArray{<:Any, <:Any, B}},
+    ::Type{T},
+) where {T, B} = CUDA.CuArray{T, <:Any, B}
 promote_parent_array_type(
-    ::Type{CUDA.CuDeviceArray{T1, N, B} where {N}},
-    ::Type{CUDA.CuDeviceArray{T2, N, B} where {N}},
-) where {T1, T2, B} = CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N}
+    ::Type{CUDA.CuArray{T1, <:Any, B}},
+    ::Type{CUDA.CuArray{T2, <:Any, B}},
+) where {T1, T2, B} = CUDA.CuArray{promote_type(T1, T2), <:Any, B}
 
-# allow on-device use of lazy broadcast objects with different type params
-promote_parent_array_type(
-    ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}},
-    ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}},
-) where {T1, T2, B1, B2} =
-    CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}
-
-# allow on-device use of lazy broadcast objects with different type params
+# Allow on-device use of lazy broadcast objects.
+parent_array_type(::Type{<:CUDA.CuDeviceArray}, ::Type{T}) where {T} =
+    CUDA.CuDeviceArray{T}
 promote_parent_array_type(
     ::Type{CUDA.CuDeviceArray{T1}},
-    ::Type{CUDA.CuDeviceArray{T2, N, B2} where {N}},
-) where {T1, T2, B2} =
-    CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}
-
-promote_parent_array_type(
-    ::Type{CUDA.CuDeviceArray{T1, N, B1} where {N}},
-    ::Type{CUDA.CuDeviceArray{T2} where {N}},
-) where {T1, T2, B1} =
-    CUDA.CuDeviceArray{promote_type(T1, T2), N, B} where {N, B}
+    ::Type{CUDA.CuDeviceArray{T2}},
+) where {T1, T2} = CUDA.CuDeviceArray{promote_type(T1, T2)}
 
 # Make `similar` accept our special `UnionAll` parent array type for CuArray.
-Base.similar(
-    ::Type{CUDA.CuArray{T, N′, B} where {N′}},
-    dims::Dims{N},
-) where {T, N, B} = similar(CUDA.CuArray{T, N, B}, dims)
+Base.similar(::Type{CUDA.CuArray{T, <:Any, B}}, dims::Dims{N}) where {T, N, B} =
+    similar(CUDA.CuArray{T, N, B}, dims)
 
 unval(::Val{CI}) where {CI} = CI
 unval(CI) = CI

ext/cuda/operators_sem_shmem.jl

@@ -92,7 +92,7 @@ Base.@propagate_inbounds function operator_shmem(
     Nq = Quadratures.degrees_of_freedom(QS)
     # allocate temp output
     RT = operator_return_eltype(op, eltype(arg))
-    Nf = DataLayouts.typesize(FT, RT)
+    Nf = DataLayouts.storage_length(FT, RT)
     WJv¹ = CUDA.CuStaticSharedArray(RT, (Nq, Nvt))
     return (WJv¹,)
 end
@@ -107,7 +107,7 @@ Base.@propagate_inbounds function operator_shmem(
     Nq = Quadratures.degrees_of_freedom(QS)
     # allocate temp output
     RT = operator_return_eltype(op, eltype(arg))
-    Nf = DataLayouts.typesize(FT, RT)
+    Nf = DataLayouts.storage_length(FT, RT)
     WJv¹ = CUDA.CuStaticSharedArray(RT, (Nq, Nq, Nvt))
     WJv² = CUDA.CuStaticSharedArray(RT, (Nq, Nq, Nvt))
     return (WJv¹, WJv²)