diff --git a/docs/src/lib/operations.md b/docs/src/lib/operations.md index 4c6bae92..88429dc5 100644 --- a/docs/src/lib/operations.md +++ b/docs/src/lib/operations.md @@ -51,7 +51,7 @@ between two tiles is matrix multiplication, not an element-wise product. | `ct.insert(tile, index, value)` | Replace a non-overlapping sub-tile (Tile IR v13.4+) | | `ct.cat((a, b), axis)` | Concatenate tiles | | `ct.broadcast_to(tile, shape)` | Broadcast to target shape | -| `dropdims(tile; dims)` | Remove singleton dimensions | +| `dropdims(tile; dims)` | Remove one or more singleton dimensions | ## Matrix multiplication @@ -83,20 +83,23 @@ Requires Blackwell. The supported operand/scale/accumulator dtypes and block siz | Operation | Description | |-----------|-------------| -| `sum(tile; dims)` | Sum along axis | -| `prod(tile; dims)` | Product along axis | -| `maximum(tile; dims, propagate_nan)` | Maximum along axis | -| `minimum(tile; dims, propagate_nan)` | Minimum along axis | -| `any(tile; dims)` | Logical OR along axis | -| `all(tile; dims)` | Logical AND along axis | -| `count(tile; dims)` | Count `true` elements along axis | +| `sum(tile; dims)` | Sum over dimensions | +| `prod(tile; dims)` | Product over dimensions | +| `maximum(tile; dims, propagate_nan)` | Maximum over dimensions | +| `minimum(tile; dims, propagate_nan)` | Minimum over dimensions | +| `any(tile; dims)` | Logical OR over dimensions | +| `all(tile; dims)` | Logical AND over dimensions | +| `count(tile; dims)` | Count `true` elements over dimensions | | `argmax(tile; dims, propagate_nan)` | 1-based index of maximum along axis | | `argmin(tile; dims, propagate_nan)` | 1-based index of minimum along axis | | `cumsum(tile; dims, rev)` | Cumulative sum | | `cumprod(tile; dims, rev)` | Cumulative product | -Following `Base` semantics, reductions keep the reduced dimension as size 1. Use `dropdims` -to remove it. +`dims` must be a compile-time constant. It may be an integer, an iterable of integers, or `:`; +order and repetitions do not matter. Following `Base` semantics, reductions retain reduced +dimensions with size 1. `dims=:` instead returns a scalar. Use `dropdims` with an integer or +tuple of integers to remove retained dimensions. `argmax`, `argmin`, and scans remain +single-axis operations. ## Higher-order functions @@ -105,8 +108,8 @@ to remove it. |-----------|-------------| | `map(f, tiles...)` | Apply function element-wise (same shape) | | `f.(tiles...)`, `broadcast(f, tiles...)` | Apply function with shape broadcasting | -| `reduce(f, tile; dims, init)` | Reduction with arbitrary function | -| `mapreduce(f, op, tile; dims, init)` | Map then reduce | +| `reduce(f, tile; dims, init)` | Reduction over dimensions with arbitrary function | +| `mapreduce(f, op, tile; dims, init)` | Map then reduce over dimensions | | `accumulate(f, tile; dims, init, rev)` | Scan/prefix-sum with arbitrary function | Any function that works on scalars "just works" when broadcast over tiles, so these cover diff --git a/docs/src/man/comparison.md b/docs/src/man/comparison.md index cdd44451..1a634917 100644 --- a/docs/src/man/comparison.md +++ b/docs/src/man/comparison.md @@ -188,14 +188,19 @@ Python reductions (`ct.sum`, `ct.max`, etc.) drop the reduced dimension by defau # Python result = ct.sum(tile, axis=1) # (M, N) → (M,) result = ct.sum(tile, axis=1, keepdims=True) # (M, N) → (M, 1) +result = ct.sum(tile, axis=(0, 1)) # (M, N) → () ``` ```julia # Julia result = sum(tile; dims=2) # (M, N) → (M, 1) result = dropdims(sum(tile; dims=2); dims=2) # (M, N) → (M,) +result = sum(tile; dims=(1, 2)) # (M, N) → (1, 1) +result = sum(tile; dims=:) # (M, N) → scalar ``` +Julia `dims` is a region: its order and repeated entries do not affect the result. + ## Scalar access and 0-D tiles diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl index 225f122d..2c0f0edf 100644 --- a/src/compiler/intrinsics/core.jl +++ b/src/compiler/intrinsics/core.jl @@ -795,8 +795,12 @@ function emit_reduce!(ctx::CGCtx, args) input_shape = tile_tvs[1].shape isempty(input_shape) && throw(IRError("Cannot reduce scalar tile")) - # Flip axis from Julia 0-indexed to Tile IR order + # Julia and Tile IR use opposite dimension orders. ndim = length(input_shape) + julia_axis >= 0 || + throw(IRError("region dimension(s) must be ≥ 1, got $(julia_axis + 1)")) + julia_axis < ndim || + throw(IRError("reduce() dimension must be in 1:$ndim; got $(julia_axis + 1)")) axis = ndim - 1 - julia_axis # ReduceOp removes the dimension; we'll reshape after to reintroduce it as size 1 diff --git a/src/compiler/utils.jl b/src/compiler/utils.jl index 7510d8ea..9272ecfb 100644 --- a/src/compiler/utils.jl +++ b/src/compiler/utils.jl @@ -587,6 +587,18 @@ function resolve_tuple(ctx::CGCtx, @nospecialize(arg), name::AbstractString) ] end + if tv.v isa Vector{Value} + T_jl = CC.widenconst(tv.jltype) + T_jl isa DataType && T_jl <: Tuple || throw(IRError("$name must be a tuple")) + return CGVal[ + let elem_T = T_jl.parameters[i] + CGVal(tv.v[i], tile_type_for_julia!(ctx, elem_T), elem_T, + RowMajorShape(extract_tile_shape(elem_T))) + end + for i in eachindex(tv.v) + ] + end + tv.tuple === nothing && throw(IRError("$name must be a tuple")) return CGVal[ let comp = emit_value!(ctx, ref) diff --git a/src/language/operations.jl b/src/language/operations.jl index b9982a2e..1ac5399a 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,13 +1397,54 @@ result = map(+, a, b) # Element-wise addition (same shape required) Intrinsics.from_scalar(f(Intrinsics.to_scalar(a), map(Intrinsics.to_scalar, rest)...), S) end +reduction_dims(dim::Integer, ::Tuple) = Int(dim) +function reduction_dims(dims, shape::Tuple) + isempty(dims) && return () + eltype(dims) <: Integer || throw(ArgumentError("reduced dimension(s) must be integers")) + invalid = min(Int(minimum(dims)), Int(maximum(dims))) + # Pass invalid dimensions through to Intrinsics.reduce for validation. + invalid < 1 && return invalid + reduction_dims(dims, shape, 1) +end +reduction_dims(::Colon, ::Tuple) = Colon() +reduction_dims(::Any, ::Tuple{}, ::Integer) = () +function reduction_dims(dims, shape::Tuple, dim::Integer) + rest = reduction_dims(dims, Base.tail(shape), dim + 1) + dim in dims ? (dim, rest...) : rest +end + +function reduce_tiles(tiles::Tuple, dim::Integer, f, init::Tuple) + dim > ndims(first(tiles)) && return tiles + Intrinsics.reduce(tiles, dim - 1, f, init) +end +reduce_tiles(tiles::Tuple, ::Tuple{}, ::Any, ::Tuple) = tiles +reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) = + reduce_tiles(Intrinsics.reduce(tiles, first(dims) - 1, f, init), + Base.tail(dims), f, init) +function reduce_tiles(tiles::Tuple, ::Colon, f, init::Tuple) + shape = size(first(tiles)) + reduced = reduce_tiles(tiles, reduction_dims(1:length(shape), shape), f, init) + map(tile -> Intrinsics.to_scalar(reshape(tile, ())), reduced) +end + +function reduce_tile(tile::Tile, dim::Integer, f, init) + dim > ndims(tile) && return tile + Intrinsics.reduce((tile,), dim - 1, f, (init,))[1] +end +reduce_tile(tile::Tile, dims, f, init) = reduce_tiles((tile,), dims, f, (init,))[1] + """ - mapreduce(identity, f, tile::Tile{T,S}; dims, init) -> Tile{T, reduced_shape} - mapreduce(f, op, tile::Tile{T,S}; dims, init) -> Tile{T, reduced_shape} - mapreduce(identity, f, tile1, tile2, ...; dims, init) -> Tuple{Tile...} + mapreduce(identity, f, tile::Tile{T,S}; dims, init) + mapreduce(f, op, tile::Tile{T,S}; dims, init) + mapreduce(identity, f, tile1, tile2, ...; dims, init) Reduce one or more tiles along `dims` using binary function `f`. +`dims` must be a compile-time constant and may be an integer, an iterable of +integers, or `:`. Dimensions form a region, so order and repetitions do not +matter and dimensions above `ndims(tile)` are no-ops. Reduced dimensions are +retained with size 1; `:` returns a scalar. + The single-tile form reduces a single tile along `dims` with identity element `init`. The multi-tile form reduces several same-shaped tiles simultaneously: `f` receives two tuples `(accumulators...)` and `(elements...)` and must return @@ -1418,6 +1459,7 @@ When a non-identity map function is provided, it is applied element-wise via sums = mapreduce(identity, +, tile; dims=2, init=zero(Float32)) sum_of_squares = mapreduce(x -> x * x, +, tile; dims=2, init=zero(Float32)) sum_of_abs = mapreduce(abs, +, tile; dims=2, init=zero(Float32)) +total = mapreduce(identity, +, tile; dims=(1, 2), init=zero(Float32)) # Simultaneous reduction of two tiles vals, idxs = mapreduce(identity, reducer, vals_tile, idx_tile; @@ -1425,42 +1467,42 @@ vals, idxs = mapreduce(identity, reducer, vals_tile, idx_tile; ``` """ @inline function Base.mapreduce(::typeof(identity), f, tile::Tile{T,S}; - dims::Integer, init) where {T<:Number, S} - Intrinsics.reduce((tile,), dims - 1, f, (T(init),))[1] + dims, init) where {T<:Number, S} + reduce_tile(tile, reduction_dims(dims, size(tile)), f, T(init)) end @inline function Base.mapreduce(f, op, tile::Tile{T,S}; - dims::Integer, init) where {T<:Number, S} + dims, init) where {T<:Number, S} reduce(op, map(f, tile); dims, init) end @inline function Base.mapreduce(::typeof(identity), f, tile1::Tile{<:Any,S}, tile2::Tile{<:Any,S}, tiles::Tile{<:Any,S}...; - dims::Integer, + dims, init::Tuple{Any, Any, Vararg{Any}}) where {S} all_tiles = (tile1, tile2, tiles...) function _combiner(args...) f(_deinterleave_accs(args...), _deinterleave_elems(args...)) end - Intrinsics.reduce(all_tiles, dims - 1, _combiner, init) + reduce_tiles(all_tiles, reduction_dims(dims, size(tile1)), _combiner, init) end """ - reduce(f, tile::Tile{T,S}; dims::Integer, init) -> Tile{T, reduced_shape} + reduce(f, tile::Tile{T,S}; dims, init) -Reduce a tile along the specified dimension using binary function `f` with -identity element `init`. The `dims` axis is 1-indexed. The reduced dimension -becomes size 1 (Julia semantics). +Reduce a tile along the specified dimensions using binary function `f` with +identity element `init`. See [`mapreduce`](@ref) for supported `dims` values. Supported functions: `+`, `*`, `max`, `min`. # Example ```julia sums = reduce(+, tile; dims=2, init=zero(Float32)) +total = reduce(+, tile; dims=(1, 2), init=zero(Float32)) ``` """ -@inline function Base.reduce(f, tile::Tile{T,S}; dims::Integer, init) where {T<:Number, S} +@inline function Base.reduce(f, tile::Tile{T,S}; dims, init) where {T<:Number, S} mapreduce(identity, f, tile; dims, init) end @@ -1487,9 +1529,9 @@ for (f, dim_helper, all_helper, ignore_op, propagate_op, integer_op, init_expr) (:minimum, :_minimum_dim, :_minimum_all, :_min_ignore_nan, :_min_propagate_nan, :min, :(typemax(T))), ] - @eval @inline $dim_helper(tile::Tile{T}, dims::Integer, ::Val{P}) where {T<:AbstractFloat, P} = + @eval @inline $dim_helper(tile::Tile{T}, dims, ::Val{P}) where {T<:AbstractFloat, P} = reduce(P ? $propagate_op : $ignore_op, tile; dims, init=$init_expr) - @eval @inline $dim_helper(tile::Tile{T}, dims::Integer, ::Val) where {T<:Integer} = + @eval @inline $dim_helper(tile::Tile{T}, dims, ::Val) where {T<:Integer} = reduce($integer_op, tile; dims, init=$init_expr) @eval @inline $all_helper(tile::Tile{T,Tuple{}}, ::Val) where {T<:Number} = @@ -1518,35 +1560,33 @@ for f in (:sum, :prod) end """ - any(tile::Tile{Bool,S}; dims) -> Tile{Bool, reduced_shape} + any(tile::Tile{Bool,S}; dims) -Logical OR reduction along the specified axis (1-indexed). -Reduced dimensions become size 1. +Logical OR reduction over `dims`. """ -@inline Base.any(tile::Tile{Bool,S}; dims::Integer) where {S} = +@inline Base.any(tile::Tile{Bool,S}; dims) where {S} = reduce(|, tile; dims, init=false) """ - all(tile::Tile{Bool,S}; dims) -> Tile{Bool, reduced_shape} + all(tile::Tile{Bool,S}; dims) -Logical AND reduction along the specified axis (1-indexed). -Reduced dimensions become size 1. +Logical AND reduction over `dims`. """ -@inline Base.all(tile::Tile{Bool,S}; dims::Integer) where {S} = +@inline Base.all(tile::Tile{Bool,S}; dims) where {S} = reduce(&, tile; dims, init=true) """ - count(tile::Tile{Bool,S}; dims) -> Tile{Int32, reduced_shape} + count(tile::Tile{Bool,S}; dims) -Count true elements along `dims` (1-indexed). Apply predicates via -broadcasting before calling count: +Count true elements along `dims`. Apply predicates via broadcasting before +calling `count`: # Example ```julia n_positive = count(tile .> 0.0f0; dims=1) ``` """ -@inline function Base.count(tile::Tile{Bool,S}; dims::Integer) where {S} +@inline function Base.count(tile::Tile{Bool,S}; dims) where {S} sum(convert(Tile{Int32}, tile); dims) end @@ -1643,25 +1683,39 @@ end end +check_drop_dims(tile::Tile, dims::Tuple) = check_drop_dims(tile, dims, ()) +check_drop_dims(::Tile, ::Tuple{}, ::Tuple) = nothing +function check_drop_dims(tile::Tile, dims::Tuple, seen::Tuple) + dim = first(dims) + 1 <= dim <= ndims(tile) || + throw(ArgumentError("dropped dims must be in range 1:ndims(A)")) + size(tile, dim) == 1 || throw(ArgumentError("dropped dims must all be size 1")) + dim in seen && throw(ArgumentError("dropped dims must be unique")) + check_drop_dims(tile, Base.tail(dims), (seen..., dim)) +end + +drop_shape(shape::Tuple, dims::Tuple) = drop_shape(shape, dims, 1) +drop_shape(::Tuple{}, ::Tuple, ::Integer) = () +function drop_shape(shape::Tuple, dims::Tuple, dim::Integer) + rest = drop_shape(Base.tail(shape), dims, dim + 1) + dim in dims ? rest : (first(shape), rest...) +end + """ dropdims(tile::Tile{T,S}; dims) -> Tile{T, squeezed_shape} -Remove singleton dimensions from a tile. The specified dimensions must have -size 1. This is the inverse of the dimension-preserving behavior of `sum`, -`prod`, `maximum`, and `minimum`. +Remove the singleton dimensions in `dims` from a tile. # Example ```julia -sums = sum(tile; dims=2) # (64, 1) -squeezed = dropdims(sums; dims=2) # (64,) +sums = sum(tile; dims=(1, 2)) +scalar_tile = dropdims(sums; dims=(1, 2)) ``` """ -@inline Base.dropdims(tile::Tile; dims::Integer) = - _dropdims(tile, Val(dims)) - -@inline function _dropdims(tile::Tile, ::Val{D}) where {D} - new_shape = ntuple(i -> i < D ? size(tile, i) : size(tile, i + 1), Val(ndims(tile) - 1)) - reshape(tile, new_shape) +function Base.dropdims(tile::Tile; dims::Union{Integer,Dims}) + region = dims isa Integer ? (Int(dims),) : dims + check_drop_dims(tile, region) + reshape(tile, drop_shape(size(tile), region)) end #============================================================================= diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 8084e8e2..6c61d4e6 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -71,6 +71,14 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray, reduce_dims:: dim_order = (filter(d -> d in reduce_dims, 1:N)..., filter(d -> !(d in reduce_dims), 1:N)...) ts = _compute_tile_sizes(size(A), dim_order) + # No reduced dimensions: map each tile independently. + if isempty(reduce_dims) + grid = ntuple(d -> cld(size(A, d), ts[d]), N) + reduce_stride = ntuple(d -> Int32(1), N) + return _launch_mapreduce!(grid, cuTileconvert(R), src_ta, f, op, ts, reduce_dims, + init, PaddingMode.Zero, reduce_stride, nothing) + end + pad_mode = _padding_for_neutral(init) if pad_mode === nothing has_oob = any(i -> size(A, i) % ts[i] != 0, 1:N) @@ -140,8 +148,13 @@ function _mapreduce(f, op, A::AbstractArray; dims, init) neutral = neutral_element(op, ET) N = ndims(A) - reduce_dims = dims === Colon() ? ntuple(identity, N) : - dims isa Integer ? (dims,) : Tuple(dims) + region = dims === Colon() ? ntuple(identity, N) : + dims isa Integer ? (dims,) : Tuple(dims) + for d in region + d isa Integer || throw(ArgumentError("reduced dimension(s) must be integers")) + Int(d) < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $d")) + end + reduce_dims = Tuple(d for d in 1:N if d in region) out_size = ntuple(d -> d in reduce_dims ? 1 : size(A, d), N) R = similar(A, ET, out_size) _mapreducedim!(f, op, R, A, reduce_dims; init=neutral) diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 3530dd6f..7f666efd 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -184,7 +184,6 @@ spec4d = ct.ArraySpec{4}(16, true) end @testset "dropdims" begin - # dropdims on dim 1: (1, 8) -> dropdims(; dims=2) -> (8,) @test @filecheck begin @check_label "entry" code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}, ct.TileArray{Float32,1,spec1d}}) do a, b @@ -197,7 +196,6 @@ spec4d = ct.ArraySpec{4}(16, true) end end - # dropdims on dim 2: (8, 1) -> dropdims(; dims=2) -> (8,) @test @filecheck begin @check_label "entry" code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}, ct.TileArray{Float32,1,spec1d}}) do a, b @@ -209,6 +207,23 @@ spec4d = ct.ArraySpec{4}(16, true) return end end + + @test @filecheck begin + @check_label "entry" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}) do a + tile = ct.load(a, (1, 1), (1, 1)) + @check "reshape" + Base.donotdelete(dropdims(tile; dims=(2, 1))) + return + end + end + + @test_throws "dropped dims must be unique" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + tile = ct.load(a, (1, 1), (1, 1)) + Base.donotdelete(dropdims(tile; dims=(1, 1))) + return + end end @testset "permutedims" begin @@ -1025,6 +1040,83 @@ spec4d = ct.ArraySpec{4}(16, true) end end + @test @filecheck begin + @check_label "entry" + @check "reduce {{.*}} dim=1" + @check "reduce {{.*}} dim=0" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(sum(tile; dims=(2, 1, 2))) + return + end + end + + @test @filecheck begin + @check_label "entry" + @check_count 4 "reduce" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(sum(tile; dims=1:2)) + Base.donotdelete(sum(tile; dims=:)) + return + end + end + + @test @filecheck begin + @check_label "entry" + code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + @check_not "reduce" + Base.donotdelete(sum(tile; dims=())) + Base.donotdelete(sum(tile; dims=3)) + Base.donotdelete(sum(tile; dims=(3, 4))) + return + end + end + + @test_throws "region dimension(s) must be ≥ 1, got 0" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(sum(tile; dims=0)) + return + end + + @test_throws "region dimension(s) must be ≥ 1, got 0" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(sum(tile; dims=(0,))) + return + end + + @test_throws "reduced dimension(s) must be integers" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(sum(tile; dims=(1.0,))) + return + end + + @test_throws "region dimension(s) must be ≥ 1, got 0" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(ct.Intrinsics.reduce((tile,), -1, +, (0.0f0,))) + return + end + + @test_throws "reduce() dimension must be in 1:2; got 3" code_tiled( + Tuple{ct.TileArray{Float32,2,spec2d}}) do a + pid = ct.bid(1) + tile = ct.load(a, pid, (4, 16)) + Base.donotdelete(ct.Intrinsics.reduce((tile,), 2, +, (0.0f0,))) + return + end + # any @test @filecheck begin @check_label "entry" @@ -1032,9 +1124,9 @@ spec4d = ct.ArraySpec{4}(16, true) pid = ct.bid(1) tile = ct.load(a, pid, (4, 16)) mask = tile .> 0.0f0 - @check "reduce" + @check_count 2 "reduce" @check "ori" - Base.donotdelete(any(mask; dims=2)) + Base.donotdelete(any(mask; dims=(1, 2))) return end end @@ -1046,9 +1138,9 @@ spec4d = ct.ArraySpec{4}(16, true) pid = ct.bid(1) tile = ct.load(a, pid, (4, 16)) mask = tile .> 0.0f0 - @check "reduce" + @check_count 2 "reduce" @check "andi" - Base.donotdelete(all(mask; dims=2)) + Base.donotdelete(all(mask; dims=(1, 2))) return end end @@ -1060,9 +1152,9 @@ spec4d = ct.ArraySpec{4}(16, true) pid = ct.bid(1) tile = ct.load(a, pid, (4, 16)) @check "exti" - @check "reduce" + @check_count 2 "reduce" @check "addi" - Base.donotdelete(count(tile .> 0.0f0; dims=2)) + Base.donotdelete(count(tile .> 0.0f0; dims=(1, 2))) return end end diff --git a/test/device/reductions.jl b/test/device/reductions.jl index 404003a2..c2177095 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -133,6 +133,142 @@ end end end +@testset "reduction dimensions" begin + function sum_dims_empty_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + tile = ct.load(a, (1, 1), (8, 16)) + ct.store(b, (1, 1), sum(tile; dims=())) + return + end + + function sum_dims_order_kernel(a::ct.TileArray{Float32,2}, + b12::ct.TileArray{Float32,2}, + b21::ct.TileArray{Float32,2}) + tile = ct.load(a, (1, 1), (2, 2)) + ct.store(b12, (1, 1), sum(tile; dims=(1, 2))) + ct.store(b21, (1, 1), sum(tile; dims=(2, 1))) + return + end + + function sum_dims_range_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + tile = ct.load(a, (1, 1), (8, 16)) + ct.store(b, (1, 1), sum(tile; dims=1:2)) + return + end + + function sum_dims_colon_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,1}) + tile = ct.load(a, (1, 1), (8, 16)) + b[1] = sum(tile; dims=:) + return + end + + function sum_dims23_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) + tile = ct.load(a, (1, 1), (8, 16)) + ct.store(b, (1, 1), sum(tile; dims=(2, 3))) + return + end + + m, n = 8, 16 + a = CUDA.rand(Float32, m, n) + a_cpu = Array(a) + + b = CUDA.zeros(Float32, m, n) + @cuda backend=cuTile sum_dims_empty_kernel(a, b) + @test Array(b) == a_cpu + + order_input = CuArray(Float32[-1.0f10 -1; 1.0f10 -1]) + b12 = CUDA.zeros(Float32, 1, 1) + b21 = similar(b12) + @cuda backend=cuTile sum_dims_order_kernel(order_input, b12, b21) + @test Array(b12) == Array(b21) == sum(Array(order_input); dims=(1, 2)) + + b = CUDA.zeros(Float32, 1, 1) + @cuda backend=cuTile sum_dims_range_kernel(a, b) + @test Array(b) ≈ sum(a_cpu; dims=1:2) rtol=1e-3 + + b_scalar = CUDA.zeros(Float32, 1) + @cuda backend=cuTile sum_dims_colon_kernel(a, b_scalar) + @test Array(b_scalar)[1] ≈ sum(a_cpu; dims=:) rtol=1e-3 + + b = CUDA.zeros(Float32, m, 1) + @cuda backend=cuTile sum_dims23_kernel(a, b) + @test Array(b) ≈ sum(a_cpu; dims=(2, 3)) rtol=1e-3 +end + +@testset "3D reduction dimensions" begin + function max_dims23_kernel(a::ct.TileArray{Float32,3}, b::ct.TileArray{Float32,3}) + tile = ct.load(a, (1, 1, 1), (4, 8, 16)) + ct.store(b, (1, 1, 1), maximum(tile; dims=(2, 3))) + return + end + + function mapreduce_dims13_kernel(a::ct.TileArray{Float32,3}, b::ct.TileArray{Float32,3}) + tile = ct.load(a, (1, 1, 1), (4, 8, 16)) + ct.store(b, (1, 1, 1), mapreduce(abs, +, tile; dims=(1, 3), init=0.0f0)) + return + end + + function sum_dims_colon_3d_kernel(a::ct.TileArray{Float32,3}, b::ct.TileArray{Float32,1}) + tile = ct.load(a, (1, 1, 1), (4, 8, 16)) + b[1] = sum(tile; dims=:) + return + end + + d1, d2, d3 = 4, 8, 16 + a = CUDA.rand(Float32, d1, d2, d3) .- 0.5f0 + a_cpu = Array(a) + + b = CUDA.zeros(Float32, d1, 1, 1) + @cuda backend=cuTile max_dims23_kernel(a, b) + @test Array(b) ≈ maximum(a_cpu; dims=(2, 3)) + + b = CUDA.zeros(Float32, 1, d2, 1) + @cuda backend=cuTile mapreduce_dims13_kernel(a, b) + @test Array(b) ≈ mapreduce(abs, +, a_cpu; dims=(1, 3)) rtol=1e-3 + + b = CUDA.zeros(Float32, 1) + @cuda backend=cuTile sum_dims_colon_3d_kernel(a, b) + @test Array(b)[1] ≈ sum(a_cpu) rtol=1e-3 +end + +@testset "count dimensions" begin + function count_dims12_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Int32,2}) + tile = ct.load(a, (1, 1), (8, 16)) + ct.store(b, (1, 1), count(tile .> 0.0f0; dims=(1, 2))) + return + end + + a = CUDA.rand(Float32, 8, 16) .- 0.5f0 + b = CUDA.zeros(Int32, 1, 1) + @cuda backend=cuTile count_dims12_kernel(a, b) + @test Array(b)[1, 1] == count(>(0.0f0), Array(a)) +end + +@testset "multi-tile mapreduce dimensions" begin + function pair_reduce_kernel(a::ct.TileArray{Float32,2}, + bv::ct.TileArray{Float32,2}, bt::ct.TileArray{Float32,2}) + tile = ct.load(a, (1, 1), (8, 16)) + tag = tile .* 2.0f0 + function reducer(accs, elems) + cond = accs[1] >= elems[1] + (ifelse(cond, accs[1], elems[1]), ifelse(cond, accs[2], elems[2])) + end + vals, tags = mapreduce(identity, reducer, tile, tag; + dims=(1, 2), init=(typemin(Float32), 0.0f0)) + ct.store(bv, (1, 1), vals) + ct.store(bt, (1, 1), tags) + return + end + + a = CUDA.rand(Float32, 8, 16) + bv = CUDA.zeros(Float32, 1, 1) + bt = CUDA.zeros(Float32, 1, 1) + @cuda backend=cuTile pair_reduce_kernel(a, bv, bt) + + a_cpu = Array(a) + @test Array(bv)[1, 1] == maximum(a_cpu) + @test Array(bt)[1, 1] == 2.0f0 * maximum(a_cpu) +end + @testset "map(abs, tile)" begin function map_abs_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) pid = ct.bid(1) diff --git a/test/host/mapreduce.jl b/test/host/mapreduce.jl index f1709510..43538156 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -27,6 +27,37 @@ using CUDA @test Array(R) ≈ Array(sum(A; dims=2)) end + @testset "sum with multiple dims" begin + A = CUDA.rand(Float32, 128, 256) + for dims in ((2, 1, 2), 1:2, [2, 1]) + @test Array(sum(ct.Tiled(A); dims)) ≈ Array(sum(A; dims)) + end + end + + @testset "non-aligned no-op dims" begin + A = CUDA.fill(Int32(2), 17) + @test Array(prod(ct.Tiled(A); dims=())) == Array(prod(A; dims=())) + @test Array(maximum(ct.Tiled(A); dims=())) == Array(maximum(A; dims=())) + @test Array(prod(ct.Tiled(A); dims=2)) == Array(prod(A; dims=2)) + end + + @testset "out-of-range dims" begin + A = CUDA.rand(Float32, 64, 32) + R = sum(ct.Tiled(A); dims=(2, 3)) + @test Array(R) ≈ Array(sum(A; dims=2)) + @test_throws InexactError sum(ct.Tiled(A); dims=big(typemax(Int)) + 1) + @test_throws ArgumentError sum(ct.Tiled(A); dims=0) + @test_throws ArgumentError("reduced dimension(s) must be integers") begin + sum(ct.Tiled(A); dims=(1.0,)) + end + end + + @testset "maximum 3D dims=(1,3)" begin + A = CUDA.rand(Float32, 32, 16, 8) + R = maximum(ct.Tiled(A); dims=(1, 3)) + @test Array(R) ≈ Array(maximum(A; dims=(1, 3))) + end + @testset "sum 2D full" begin A = CUDA.rand(Float32, 128, 256) @test sum(ct.Tiled(A)) ≈ sum(A)