diff --git a/src/broadcast.jl b/src/broadcast.jl index 4894dcd1..61df207d 100644 --- a/src/broadcast.jl +++ b/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) diff --git a/src/language/operations.jl b/src/language/operations.jl index ff39cc96..a164d438 100644 --- a/src/language/operations.jl +++ b/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) diff --git a/src/language/types.jl b/src/language/types.jl index 884aaf3f..2e58e6c3 100644 --- a/src/language/types.jl +++ b/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(typeof(arr)) +Base.ndims(arr::AbstractTileArray) = ndims(typeof(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 diff --git a/src/mapreduce.jl b/src/mapreduce.jl index cd155ac9..8084e8e2 100644 --- a/src/mapreduce.jl +++ b/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