Add AbstractTileArray

Author: AntonOresten

src/broadcast.jl

@@ -17,8 +17,9 @@ BroadcastStyle(::TiledStyle{N}, ::CuArrayStyle{M}) where {N,M} = TiledStyle{max(
 ## broadcast interface
 
 function Base.Broadcast.materialize!(dest::Tiled, bc::Broadcasted)
-    _tiled_broadcast!(parent(dest), bc)
-    return dest
+    arr = parent(dest)
+    _tiled_broadcast!(arr, bc)
+    return arr
 end
 
 function Base.copy(bc::Broadcasted{TiledStyle{N}}) where N
@@ -43,8 +44,8 @@ end
 
 ## kernel wrapper
 
-_to_tiled_bc(t::Tiled) = TileArray(parent(t))
-_to_tiled_bc(arr::AbstractArray) = TileArray(arr)
+_to_tiled_bc(arr::AbstractArray) = cuTileconvert(arr)
+_to_tiled_bc(t::Tiled) = _to_tiled_bc(parent(t))
 _to_tiled_bc(x::Number) = x
 _to_tiled_bc(x) = x  # fallback for other types
 function _to_tiled_bc(bc::Broadcasted)
@@ -53,7 +54,7 @@ function _to_tiled_bc(bc::Broadcasted)
 end
 
 function _tiled_broadcast!(dest::AbstractArray{T,N}, bc::Broadcasted) where {T, N}
-    dest_ta = TileArray(dest)
+    dest_ta = _to_tiled_bc(dest)
     tiled_bc = _to_tiled_bc(bc)
 
     ts = _compute_tile_sizes(size(dest))
@@ -67,21 +68,24 @@ end
 
 ## kernel
 
-@inline _eval_bc(arr::TileArray, bid, tile_size) = cuTile.load(arr, bid, tile_size)
+@inline _eval_bc(arr::AbstractTileArray, bid, tile_size) = cuTile.load(arr, bid, tile_size)
 @inline _eval_bc(x::Number, bid, tile_size) = x
 
+@inline _bc_func(f) = f
+@inline _bc_func(::Constant{Type{T}, T}) where {T} = T
+
 @inline function _eval_bc(bc::Broadcasted, bid, tile_size)
     args = _eval_bc_args(bc.args, bid, tile_size)
     # Use broadcast to get element-wise semantics (not direct call, which
     # would dispatch to e.g. matmul for * on tiles)
-    broadcast(bc.f, args...)
+    broadcast(_bc_func(bc.f), args...)
 end
 
 @inline _eval_bc_args(::Tuple{}, bid, tile_size) = ()
 @inline _eval_bc_args(args::Tuple, bid, tile_size) =
     (_eval_bc(args[1], bid, tile_size), _eval_bc_args(Base.tail(args), bid, tile_size)...)
 
-@generated function broadcast_kernel(dest::TileArray{T, N}, bc, tile_size, overflow_grids) where {T, N}
+@generated function broadcast_kernel(dest::AbstractTileArray{T, N}, bc, tile_size, overflow_grids) where {T, N}
     quote
         bids = _unflatten_bids(Val{$N}(), overflow_grids)
         result = _eval_bc(bc, bids, tile_size)

src/language/operations.jl

@@ -313,6 +313,10 @@ Axis is 1-indexed. Equivalent to cld(size(arr, axis), shape[axis]).
     Intrinsics.get_index_space_shape(pv, axis - One())  # convert to 0-indexed
 end
 
+@inline function num_tiles(arr::AbstractTileArray, axis::Integer, shape::NTuple{<:Any, Int})
+    Int32(cld(size(arr, axis), shape[axis]))
+end
+
 # Match a shape tuple to a target rank N by padding trailing 1s or squeezing trailing singletons.
 @generated function _match_shape(::Val{Shape}, ::Val{N}) where {Shape, N}
     M = length(Shape)

src/language/types.jl

@@ -1,7 +1,20 @@
-public TileArray, Tile, Constant, TFloat32, similar_type,
+public AbstractTileArray, TileArray, Tile, Constant, TFloat32, similar_type,
        ScalarInt, ScalarFloat, IntTile, FloatTile, TileOrInt, TileOrFloat,
        TileOrScalar
 
+"""
+    AbstractTileArray{T, N}
+
+Supertype for N-dimensional kernel array arguments with element type `T`.
+"""
+abstract type AbstractTileArray{T, N} end
+
+Base.eltype(::Type{<:AbstractTileArray{T}}) where T = T
+Base.ndims(::Type{<:AbstractTileArray{<:Any,N}}) where N = N
+Base.eltype(arr::AbstractTileArray) = eltype(arr)
+Base.ndims(arr::AbstractTileArray) = ndims(arr)
+
+
 """
     ArraySpec{N}
 
@@ -167,17 +180,11 @@ specializations (e.g., aligned vs unaligned) produce different cubins.
 - `sizes::NTuple{N, Int32}`: Size in each dimension
 - `strides::NTuple{N, Int32}`: Stride in each dimension (in elements)
 """
-struct TileArray{T, N, S}
+struct TileArray{T, N, S} <: AbstractTileArray{T, N}
     ptr::Ptr{T}
     sizes::NTuple{N, Int32}
     strides::NTuple{N, Int32}
 end
-
-# Type accessors for TileArray
-Base.eltype(::Type{<:TileArray{T}}) where {T} = T
-Base.eltype(::TileArray{T}) where {T} = T
-Base.ndims(::Type{<:TileArray{<:Any, N}}) where {N} = N
-Base.ndims(::TileArray{<:Any, N}) where {N} = N
 Base.size(arr::TileArray) = arr.sizes
 function Base.size(arr::TileArray{<:Any, N}, d::Integer) where N
     d < 1 && error("arraysize: dimension out of range") # from Array method

src/mapreduce.jl

@@ -11,7 +11,7 @@ _atomic_op(::typeof(&), ::Type{<:Integer}) = atomic_and
 _atomic_op(_, ::Type) = nothing
 
 @generated function mapreduce_kernel(
-    dest::TileArray{TD, N}, src::TileArray{TS, N},
+    dest::AbstractTileArray{TD, N}, src::AbstractTileArray{TS, N},
     f, op, tile_size, reduce_dims, overflow_grids, init_val, pad_mode,
     reduce_stride, atomic_op
 ) where {TD, TS, N}
@@ -65,7 +65,7 @@ end
 
 function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray, reduce_dims::Tuple; init)
     N = ndims(A)
-    src_ta = TileArray(A)
+    src_ta = cuTileconvert(A)
 
     # Reduced dims first (larger tiles => better hardware reduction)
     dim_order = (filter(d -> d in reduce_dims, 1:N)..., filter(d -> !(d in reduce_dims), 1:N)...)
@@ -103,19 +103,19 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray, reduce_dims::
 
         if atomic_op !== nothing
             fill!(R, init)
-            _launch_mapreduce!(grid, TileArray(R), src_ta, f, op, ts, reduce_dims,
+            _launch_mapreduce!(grid, cuTileconvert(R), src_ta, f, op, ts, reduce_dims,
                                init, pad_mode, reduce_stride, atomic_op)
         else
             # Two-pass: parallelize along par_dim, then reduce partials
             tmp = similar(A, eltype(R), ntuple(_dim_size, N))
-            _launch_mapreduce!(grid, TileArray(tmp), src_ta, f, op, ts, reduce_dims,
+            _launch_mapreduce!(grid, cuTileconvert(tmp), src_ta, f, op, ts, reduce_dims,
                                init, pad_mode, reduce_stride, nothing)
             _mapreducedim!(identity, op, R, tmp, (par_dim,); init)
         end
     else
         grid = ntuple(d -> d in reduce_dims ? 1 : cld(size(A, d), ts[d]), N)
         reduce_stride = ntuple(d -> Int32(1), N)
-        _launch_mapreduce!(grid, TileArray(R), src_ta, f, op, ts, reduce_dims,
+        _launch_mapreduce!(grid, cuTileconvert(R), src_ta, f, op, ts, reduce_dims,
                            init, pad_mode, reduce_stride, nothing)
     end
 end