Fixes

Author: charleskawczynski

ext/cuda/data_layouts.jl

@@ -87,7 +87,38 @@ function CUDA.CuArray(fa::DL.FieldArray{FD}) where {FD}
     return DL.FieldArray{FD}(arrays)
 end
 
-DL.field_array(
-    array::CUDA.CuArray,
-    as::ArraySize
-) = CUDA.CuArray(DL.field_array(Array(array), as))
+DL.field_array(array::CUDA.CuArray, as::ArraySize) =
+    CUDA.CuArray(DL.field_array(Array(array), as))
+
+
+# TODO: this could be improved, but it's not typically used at runtime
+function copyto_field_array_knl!(x::DL.FieldArray{FD}, y) where {FD}
+    gidx =
+        CUDA.threadIdx().x + (CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x
+    I = cart_ind(size(y), gidx)
+    x[I] = y[I]
+    return nothing
+end
+
+@inline function Base.copyto!(
+    x::DL.FieldArray{FD, NT},
+    y::CUDA.CuArray,
+) where {FD, NT <: NTuple}
+    if ndims(eltype(NT)) == ndims(y)
+        @inbounds for i in 1:DL.tuple_length(NT)
+            Base.copyto!(x.arrays[i], y)
+        end
+    elseif ndims(eltype(NT)) + 1 == ndims(y)
+        n = prod(size(y))
+        kernel =
+            CUDA.@cuda always_inline = true launch = false copyto_field_array_knl!(
+                x,
+                y,
+            )
+        config = CUDA.launch_configuration(kernel.fun)
+        threads = min(n, config.threads)
+        blocks = cld(n, threads)
+        kernel(x, y; threads, blocks)
+    end
+    x
+end

ext/cuda/data_layouts_copyto.jl

@@ -30,9 +30,13 @@ function knl_copyto_linear!(dest::AbstractData, bc, us)
     return nothing
 end
 
-function knl_copyto_linear!(dest::DataF, bc, us)
-    tidx = thread_index()
-    @inbounds dest[] = bc[tidx]
+function knl_copyto_linear!(dest::DataF{S},bc,us) where {S}
+    @inbounds begin
+        tidx = thread_index()
+        if tidx ≤ get_N(us)
+            dest[] = bc[tidx]
+        end
+    end
     return nothing
 end
 
@@ -48,13 +52,17 @@ function knl_copyto_flat!(dest::AbstractData, bc, us)
     return nothing
 end
 
-function knl_copyto_flat!(dest::DataF, bc, us)
+function knl_copyto_flat!(
+    dest::DataF{S},
+    bc::DataLayouts.BroadcastedUnionDataF{S},
+    us,
+) where {S}
     @inbounds begin
         tidx = thread_index()
         if tidx ≤ get_N(us)
             n = size(dest)
-            I = kernel_indexes(tidx, n)
-            dest[] = bc[I]
+            # I = kernel_indexes(tidx, n)
+            dest[] = bc[]
         end
     end
     return nothing
@@ -67,12 +75,7 @@ function cuda_copyto!(dest::AbstractData, bc)
     if Nv > 0 && Nh > 0
         if has_uniform_datalayouts(bc)
             bc′ = to_non_extruded_broadcasted(bc)
-            auto_launch!(
-                knl_copyto_linear!,
-                (dest, bc′, us),
-                n;
-                auto = true,
-            )
+            auto_launch!(knl_copyto_linear!, (dest, bc′, us), n; auto = true)
         else
             auto_launch!(knl_copyto_flat!, (dest, bc, us), n; auto = true)
         end

src/DataLayouts/DataLayouts.jl

@@ -564,6 +564,10 @@ end
     @inbounds col[]
 end
 
+@propagate_inbounds function Base.getindex(col::Data0D, I::Integer)
+    @inbounds col[]
+end
+
 Base.@propagate_inbounds function Base.setindex!(data::DataF{S}, val) where {S}
     @inbounds set_struct!(
         field_array(data),
@@ -581,6 +585,14 @@ end
     @inbounds col[] = val
 end
 
+@propagate_inbounds function Base.setindex!(
+    col::Data0D,
+    val,
+    I::Integer,
+)
+    @inbounds col[] = val
+end
+
 # ======================
 # DataSlab2D DataLayout
 # ======================
@@ -1306,7 +1318,7 @@ type parameters.
 #! format: on
 
 # Skip DataF here, since we want that to MethodError.
-for DL in (:IJKFVH, :IJFH, :IFH, :IJF, :IF, :VF, :VIJFH, :VIFH)
+for DL in (:IJKFVH, :IJFH, :IFH, :IJF, :IF, :DataF, :VF, :VIJFH, :VIFH)
     @eval @propagate_inbounds Base.getindex(data::$(DL), I::Integer) =
         linear_getindex(data, I)
     @eval @propagate_inbounds Base.setindex!(data::$(DL), val, I::Integer) =

src/DataLayouts/copyto.jl

@@ -40,8 +40,6 @@ end
 
 # broadcasting scalar assignment
 # Performance optimization for the common identity scalar case: dest .= val
-# And this is valid for the CPU or GPU, since the broadcasted object
-# is a scalar type.
 function Base.copyto!(
     dest::AbstractData,
     bc::Base.Broadcast.Broadcasted{Style},

src/DataLayouts/field_array.jl

@@ -395,6 +395,9 @@ function field_array(
     FieldArray{FD}(EmptyArray(scalar_array))
 end
 
+# Base.show(io::IO, fa::FieldArray{FD}) where {FD} = print(io, "$(arrays_type(typeof(fa)))(", Array(fa), ")")
+Base.show(io::IO, fa::FieldArray{FD}) where {FD} = print(io, "FieldArray{$FD}(", Array(fa), ")")
+
 # Warning: this method is type-unstable.
 function Base.view(fa::FieldArray{FD}, inds...) where {FD}
     AI = dropat(inds, Val(FD))

test/DataLayouts/data0d.jl

@@ -48,7 +48,7 @@ end
     array = zeros(Float64, 3)
     data = DataF{S}(array)
     @test data[][2] == zero(Float64)
-    @test_throws MethodError data[1]
+    # @test_throws MethodError data[1] # this no longer can throw an error
 end
 
 @testset "DataF type safety" begin