From ced4e7d42617644ff30743bfa476e7018ef24f2f Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Wed, 5 Aug 2026 13:53:30 +0200 Subject: [PATCH 01/11] Add multi-dim reduction support --- src/compiler/intrinsics/core.jl | 2 + src/compiler/utils.jl | 15 +++++ src/language/operations.jl | 59 ++++++++++++----- src/mapreduce.jl | 12 ++++ test/codegen/operations.jl | 36 ++++++++++ test/device/reductions.jl | 114 ++++++++++++++++++++++++++++++++ test/host/mapreduce.jl | 39 +++++++++++ 7 files changed, 262 insertions(+), 15 deletions(-) diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl index 225f122d..350ee7eb 100644 --- a/src/compiler/intrinsics/core.jl +++ b/src/compiler/intrinsics/core.jl @@ -797,6 +797,8 @@ function emit_reduce!(ctx::CGCtx, args) # Flip axis from Julia 0-indexed to Tile IR order ndim = length(input_shape) + 0 <= 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..c3dd7fbe 100644 --- a/src/compiler/utils.jl +++ b/src/compiler/utils.jl @@ -587,6 +587,21 @@ function resolve_tuple(ctx::CGCtx, @nospecialize(arg), name::AbstractString) ] end + # Multi-value tuple result consumed whole (e.g. one intrinsic's tuple + # output fed into another, as in chained reduce): split into component + # CGVals the same way getfield extraction does. + 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_type = T_jl.parameters[i] + CGVal(tv.v[i], tile_type_for_julia!(ctx, elem_type), elem_type, + RowMajorShape(extract_tile_shape(elem_type))) + end + for i in 1:length(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..c91683d6 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,6 +1397,28 @@ 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 +# `dims` for tile reductions: a single 1-indexed axis or a tuple of axes. +# Tuples reduce each listed axis in turn (one `cuda_tile.reduce` per axis); +# the empty tuple reduces nothing, returning the input unchanged. Following +# Base, axes beyond `ndims` are implicit size-1 dimensions and reduce as a +# no-op; axes < 1 error during codegen. +const ReduceDims = Union{Integer, Tuple{Vararg{Integer}}} +@inline _reduce_dims_tuple(dims::Integer) = (dims,) +@inline _reduce_dims_tuple(dims::Tuple{Vararg{Integer}}) = dims + +Base.@constprop :aggressive @inline _drop_implicit_dims(::Tuple{}, n::Int) = () +Base.@constprop :aggressive @inline _drop_implicit_dims(dims::Tuple, n::Int) = + first(dims) <= n ? (first(dims), _drop_implicit_dims(Base.tail(dims), n)...) : + _drop_implicit_dims(Base.tail(dims), n) + +Base.@constprop :aggressive @inline _canon_reduce_dims(dims::ReduceDims, n::Int) = + _drop_implicit_dims(_reduce_dims_tuple(dims), n) + +Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, ::Tuple{}, f, init::Tuple) = tiles +Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) = + _reduce_tiles(Intrinsics.reduce(tiles, first(dims) - 1, f, init), + Base.tail(dims), f, init) + """ 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} @@ -1404,6 +1426,10 @@ end Reduce one or more tiles along `dims` using binary function `f`. +`dims` is a 1-indexed axis or a tuple of axes; each reduced axis becomes +size 1 (Julia semantics). `dims=()` reduces nothing and returns the +(mapped) input unchanged. + 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 +1444,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 +1452,44 @@ 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::ReduceDims, init) where {T<:Number, S} + _reduce_tiles((tile,), _canon_reduce_dims(dims, ndims(tile)), f, (T(init),))[1] end @inline function Base.mapreduce(f, op, tile::Tile{T,S}; - dims::Integer, init) where {T<:Number, S} + dims::ReduceDims, 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::ReduceDims, 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, _canon_reduce_dims(dims, ndims(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) -> Tile{T, reduced_shape} -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`. `dims` is a 1-indexed axis or a tuple of axes; each +reduced axis becomes size 1 (Julia semantics). `dims=()` reduces nothing and +returns the tile unchanged. 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::ReduceDims, init) where {T<:Number, S} mapreduce(identity, f, tile; dims, init) end @@ -1487,9 +1516,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::ReduceDims, ::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::ReduceDims, ::Val) where {T<:Integer} = reduce($integer_op, tile; dims, init=$init_expr) @eval @inline $all_helper(tile::Tile{T,Tuple{}}, ::Val) where {T<:Number} = @@ -1523,7 +1552,7 @@ end Logical OR reduction along the specified axis (1-indexed). Reduced dimensions become size 1. """ -@inline Base.any(tile::Tile{Bool,S}; dims::Integer) where {S} = +@inline Base.any(tile::Tile{Bool,S}; dims::ReduceDims) where {S} = reduce(|, tile; dims, init=false) """ @@ -1532,7 +1561,7 @@ Reduced dimensions become size 1. Logical AND reduction along the specified axis (1-indexed). Reduced dimensions become size 1. """ -@inline Base.all(tile::Tile{Bool,S}; dims::Integer) where {S} = +@inline Base.all(tile::Tile{Bool,S}; dims::ReduceDims) where {S} = reduce(&, tile; dims, init=true) """ @@ -1546,7 +1575,7 @@ broadcasting before calling count: 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::ReduceDims) where {S} sum(convert(Tile{Int32}, tile); dims) end diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 8084e8e2..7b301a6e 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -80,6 +80,14 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray, reduce_dims:: pad_mode = PaddingMode.Zero end + # dims=(): nothing to reduce; each block maps its tile through op(init, f(x)) + 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, pad_mode, reduce_stride, nothing) + end + # Pick the largest reduced dim for potential parallelization non_reduce_blocks = prod(cld(size(A, d), ts[d]) for d in 1:N if !(d in reduce_dims); init=1) par_dim = reduce_dims[argmax(map(d -> cld(size(A, d), ts[d]), reduce_dims))] @@ -142,6 +150,10 @@ function _mapreduce(f, op, A::AbstractArray; dims, init) reduce_dims = dims === Colon() ? ntuple(identity, N) : dims isa Integer ? (dims,) : Tuple(dims) + all(>=(1), reduce_dims) || + throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) + # Following Base, dims beyond ndims are implicit size-1 dimensions: no-op + reduce_dims = filter(<=(N), reduce_dims) 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..ddae4c45 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1025,6 +1025,42 @@ spec4d = ct.ArraySpec{4}(16, true) end end + # Tuple dims lower to one reduce per axis + @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 "reduce" + @check "reduce" + Base.donotdelete(sum(tile; dims=(1, 2))) + return + end + end + + # dims=() and axes beyond ndims reduce nothing (Base parity) + @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 + + # dims < 1 is a compile-time error + @test_throws "reduce() dimension must be in 1:2" 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 + # any @test @filecheck begin @check_label "entry" diff --git a/test/device/reductions.jl b/test/device/reductions.jl index 404003a2..eb19233c 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -133,6 +133,120 @@ end end end +@testset "reduce with tuple dims" begin + function sum_dims12_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_dims2_tuple_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,))) + return + end + + 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 + + m, n = 8, 16 + a = CUDA.rand(Float32, m, n) + a_cpu = Array(a) + + b = CUDA.zeros(Float32, 1, 1) + @cuda backend=cuTile sum_dims12_kernel(a, b) + @test Array(b)[1, 1] ≈ sum(a_cpu) rtol=1e-3 + + b = CUDA.zeros(Float32, m, 1) + @cuda backend=cuTile sum_dims2_tuple_kernel(a, b) + @test Array(b) ≈ sum(a_cpu; dims=2) rtol=1e-3 + + b = CUDA.zeros(Float32, m, n) + @cuda backend=cuTile sum_dims_empty_kernel(a, b) + @test Array(b) == a_cpu + + # Base parity: axes beyond ndims are size-1, reducing them is a no-op + 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 + + 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 "tuple dims reductions (3D)" 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 + + 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 +end + +@testset "count with tuple dims" 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 with tuple dims" begin + # Simultaneous (value, tag) reduction over both dims: the reducer keeps + # the pair whose first component is larger, so tags must follow values. + 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..505aee83 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -27,6 +27,45 @@ using CUDA @test Array(R) ≈ Array(sum(A; dims=2)) end + @testset "sum 2D dims=(1,2)" begin + A = CUDA.rand(Float32, 128, 256) + R = sum(ct.Tiled(A); dims=(1, 2)) + @test size(R) == (1, 1) + @test Array(R) ≈ Array(sum(A; dims=(1, 2))) + end + + @testset "sum 2D dims=(2,)" begin + A = CUDA.rand(Float32, 128, 256) + R = sum(ct.Tiled(A); dims=(2,)) + @test size(R) == (128, 1) + @test Array(R) ≈ Array(sum(A; dims=2)) + end + + @testset "sum 2D dims=()" begin + A = CUDA.rand(Float32, 128, 256) + R = sum(ct.Tiled(A); dims=()) + @test size(R) == (128, 256) + @test Array(R) ≈ Array(A) + end + + @testset "dims beyond ndims (Base parity)" begin + A = CUDA.rand(Float32, 64, 32) + R = sum(ct.Tiled(A); dims=3) + @test size(R) == (64, 32) + @test Array(R) ≈ Array(A) + R = sum(ct.Tiled(A); dims=(2, 3)) + @test size(R) == (64, 1) + @test Array(R) ≈ Array(sum(A; dims=2)) + @test_throws ArgumentError sum(ct.Tiled(A); dims=0) + end + + @testset "maximum 3D dims=(1,3)" begin + A = CUDA.rand(Float32, 32, 16, 8) + R = maximum(ct.Tiled(A); dims=(1, 3)) + @test size(R) == (1, 16, 1) + @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) From 914fd3081c41e3539ab58f6ab27052ca2ea07b0e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 18:53:30 +0200 Subject: [PATCH 02/11] Canonicalize reduction dimensions --- src/language/operations.jl | 18 +++++++++++++----- test/codegen/operations.jl | 18 ++++++++++++++---- test/device/reductions.jl | 15 +++++++++++++++ 3 files changed, 42 insertions(+), 9 deletions(-) diff --git a/src/language/operations.jl b/src/language/operations.jl index c91683d6..6ce3b1a8 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1406,13 +1406,21 @@ const ReduceDims = Union{Integer, Tuple{Vararg{Integer}}} @inline _reduce_dims_tuple(dims::Integer) = (dims,) @inline _reduce_dims_tuple(dims::Tuple{Vararg{Integer}}) = dims -Base.@constprop :aggressive @inline _drop_implicit_dims(::Tuple{}, n::Int) = () -Base.@constprop :aggressive @inline _drop_implicit_dims(dims::Tuple, n::Int) = - first(dims) <= n ? (first(dims), _drop_implicit_dims(Base.tail(dims), n)...) : - _drop_implicit_dims(Base.tail(dims), n) +Base.@constprop :aggressive @inline _insert_reduce_dim(d, ::Tuple{}) = (d,) +Base.@constprop :aggressive @inline function _insert_reduce_dim(d, dims::Tuple) + head = first(dims) + d == head && return dims + d < head ? (d, dims...) : (head, _insert_reduce_dim(d, Base.tail(dims))...) +end + +Base.@constprop :aggressive @inline _canonical_reduce_dims(::Tuple{}, n::Int) = () +Base.@constprop :aggressive @inline function _canonical_reduce_dims(dims::Tuple, n::Int) + tail = _canonical_reduce_dims(Base.tail(dims), n) + first(dims) <= n ? _insert_reduce_dim(first(dims), tail) : tail +end Base.@constprop :aggressive @inline _canon_reduce_dims(dims::ReduceDims, n::Int) = - _drop_implicit_dims(_reduce_dims_tuple(dims), n) + _canonical_reduce_dims(_reduce_dims_tuple(dims), n) Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, ::Tuple{}, f, init::Tuple) = tiles Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) = diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index ddae4c45..436ae9ff 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1025,15 +1025,25 @@ spec4d = ct.ArraySpec{4}(16, true) end end - # Tuple dims lower to one reduce per axis + # Tuple dims are sorted and deduplicated before lowering @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 "reduce" - @check "reduce" - Base.donotdelete(sum(tile; dims=(1, 2))) + @check "reduce {{.*}} dim=1" + @check "reduce {{.*}} dim=0" + Base.donotdelete(sum(tile; dims=(2, 1))) + return + end + end + @test @filecheck begin + @check_label "entry" + @check_count 1 "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, 1))) return end end diff --git a/test/device/reductions.jl b/test/device/reductions.jl index eb19233c..faca66f8 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -152,6 +152,15 @@ end 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 + m, n = 8, 16 a = CUDA.rand(Float32, m, n) a_cpu = Array(a) @@ -168,6 +177,12 @@ end @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)) + # Base parity: axes beyond ndims are size-1, reducing them is a no-op function sum_dims23_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) tile = ct.load(a, (1, 1), (8, 16)) From 041f91ef4b5cdf5e483d353a67e910f660f4e277 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 18:55:14 +0200 Subject: [PATCH 03/11] Handle no-op reductions before padding checks --- src/mapreduce.jl | 16 ++++++++-------- test/host/mapreduce.jl | 7 +++++++ 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 7b301a6e..469f6bf8 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) + # dims=(): nothing to reduce; each block maps its tile through op(init, f(x)) + 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) @@ -80,14 +88,6 @@ function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray, reduce_dims:: pad_mode = PaddingMode.Zero end - # dims=(): nothing to reduce; each block maps its tile through op(init, f(x)) - 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, pad_mode, reduce_stride, nothing) - end - # Pick the largest reduced dim for potential parallelization non_reduce_blocks = prod(cld(size(A, d), ts[d]) for d in 1:N if !(d in reduce_dims); init=1) par_dim = reduce_dims[argmax(map(d -> cld(size(A, d), ts[d]), reduce_dims))] diff --git a/test/host/mapreduce.jl b/test/host/mapreduce.jl index 505aee83..caab5adf 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -48,6 +48,13 @@ using CUDA @test Array(R) ≈ Array(A) 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 "dims beyond ndims (Base parity)" begin A = CUDA.rand(Float32, 64, 32) R = sum(ct.Tiled(A); dims=3) From 16dd7dd07c71da5b02c43d44c292f93a9331ec6c Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 19:15:33 +0200 Subject: [PATCH 04/11] Support general device reduction regions --- src/language/operations.jl | 104 ++++++++++++++++++++++--------------- test/codegen/operations.jl | 24 +++++++++ test/device/reductions.jl | 20 +++++++ 3 files changed, 105 insertions(+), 43 deletions(-) diff --git a/src/language/operations.jl b/src/language/operations.jl index 6ce3b1a8..cde4b7a4 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,46 +1397,60 @@ 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 -# `dims` for tile reductions: a single 1-indexed axis or a tuple of axes. -# Tuples reduce each listed axis in turn (one `cuda_tile.reduce` per axis); -# the empty tuple reduces nothing, returning the input unchanged. Following -# Base, axes beyond `ndims` are implicit size-1 dimensions and reduce as a -# no-op; axes < 1 error during codegen. -const ReduceDims = Union{Integer, Tuple{Vararg{Integer}}} -@inline _reduce_dims_tuple(dims::Integer) = (dims,) -@inline _reduce_dims_tuple(dims::Tuple{Vararg{Integer}}) = dims - -Base.@constprop :aggressive @inline _insert_reduce_dim(d, ::Tuple{}) = (d,) -Base.@constprop :aggressive @inline function _insert_reduce_dim(d, dims::Tuple) +# `dims` for tile reductions: a 1-indexed axis, a tuple or integer range of axes, +# or `:`. +# Each effective axis lowers to one `cuda_tile.reduce`; an empty region reduces +# nothing. Following Base, axes beyond `ndims` are implicit size-1 dimensions +# and reduce as a no-op; axes < 1 error during codegen. +const ReduceDims = Union{Integer, Tuple{Vararg{Integer}}, AbstractRange{<:Integer}} +const ReduceRegion = Union{Colon, ReduceDims} +_reduce_dims_tuple(dims::Integer) = (dims,) +_reduce_dims_tuple(dims::Tuple{Vararg{Integer}}) = dims + +_insert_reduce_dim(d, ::Tuple{}) = (d,) +function _insert_reduce_dim(d, dims::Tuple) head = first(dims) d == head && return dims d < head ? (d, dims...) : (head, _insert_reduce_dim(d, Base.tail(dims))...) end -Base.@constprop :aggressive @inline _canonical_reduce_dims(::Tuple{}, n::Int) = () -Base.@constprop :aggressive @inline function _canonical_reduce_dims(dims::Tuple, n::Int) +_canonical_reduce_dims(::Tuple{}, n::Int) = () +function _canonical_reduce_dims(dims::Tuple, n::Int) tail = _canonical_reduce_dims(Base.tail(dims), n) first(dims) <= n ? _insert_reduce_dim(first(dims), tail) : tail end -Base.@constprop :aggressive @inline _canon_reduce_dims(dims::ReduceDims, n::Int) = +_canon_reduce_dims(dims::ReduceDims, n::Int) = _canonical_reduce_dims(_reduce_dims_tuple(dims), n) +function _canon_reduce_dims(dims::AbstractRange{<:Integer}, n::Int) + invalid = !isempty(dims) && min(first(dims), last(dims)) < 1 + invalid && return (min(first(dims), last(dims)),) + candidates = ntuple(d -> d in dims ? d : n + 1, n) + _canonical_reduce_dims(candidates, n) +end +_canon_reduce_dims(::Colon, ::Int) = Colon() -Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, ::Tuple{}, f, init::Tuple) = tiles -Base.@constprop :aggressive @inline _reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) = +_reduce_tiles(tiles::Tuple, ::Tuple{}, f, init::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) +_reduction_scalar(tile::Tile) = Intrinsics.to_scalar(reshape(tile, ())) +function _reduce_tiles(tiles::Tuple, ::Colon, f, init::Tuple) + dims = ntuple(identity, ndims(first(tiles))) + map(_reduction_scalar, _reduce_tiles(tiles, dims, f, init)) +end """ - 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, tile::Tile{T,S}; dims, init) -> Tile or scalar + mapreduce(f, op, tile::Tile{T,S}; dims, init) -> Tile or scalar mapreduce(identity, f, tile1, tile2, ...; dims, init) -> Tuple{Tile...} Reduce one or more tiles along `dims` using binary function `f`. -`dims` is a 1-indexed axis or a tuple of axes; each reduced axis becomes -size 1 (Julia semantics). `dims=()` reduces nothing and returns the -(mapped) input unchanged. +`dims` is a 1-indexed axis, a tuple or integer range of axes, or `:`. Each +reduced axis becomes size 1 (Julia semantics). `dims=:` reduces all axes and +returns a scalar. `dims=()` reduces nothing and returns the (mapped) input +unchanged. The single-tile form reduces a single tile along `dims` with identity element `init`. The multi-tile form reduces several same-shaped tiles simultaneously: @@ -1460,19 +1474,19 @@ vals, idxs = mapreduce(identity, reducer, vals_tile, idx_tile; ``` """ @inline function Base.mapreduce(::typeof(identity), f, tile::Tile{T,S}; - dims::ReduceDims, init) where {T<:Number, S} + dims::ReduceRegion, init) where {T<:Number, S} _reduce_tiles((tile,), _canon_reduce_dims(dims, ndims(tile)), f, (T(init),))[1] end @inline function Base.mapreduce(f, op, tile::Tile{T,S}; - dims::ReduceDims, init) where {T<:Number, S} + dims::ReduceRegion, 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::ReduceDims, + dims::ReduceRegion, init::Tuple{Any, Any, Vararg{Any}}) where {S} all_tiles = (tile1, tile2, tiles...) function _combiner(args...) @@ -1482,12 +1496,13 @@ end end """ - reduce(f, tile::Tile{T,S}; dims, init) -> Tile{T, reduced_shape} + reduce(f, tile::Tile{T,S}; dims, init) -> Tile or scalar Reduce a tile along the specified dimensions using binary function `f` with -identity element `init`. `dims` is a 1-indexed axis or a tuple of axes; each -reduced axis becomes size 1 (Julia semantics). `dims=()` reduces nothing and -returns the tile unchanged. +identity element `init`. `dims` is a 1-indexed axis, a tuple or integer range +of axes, or `:`. Each reduced axis becomes size 1 (Julia semantics). `dims=:` +reduces all axes and returns a scalar. `dims=()` reduces nothing and returns +the tile unchanged. Supported functions: `+`, `*`, `max`, `min`. @@ -1497,7 +1512,7 @@ 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::ReduceDims, init) where {T<:Number, S} +@inline function Base.reduce(f, tile::Tile{T,S}; dims::ReduceRegion, init) where {T<:Number, S} mapreduce(identity, f, tile; dims, init) end @@ -1524,9 +1539,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::ReduceDims, ::Val{P}) where {T<:AbstractFloat, P} = + @eval @inline $dim_helper(tile::Tile{T}, dims::ReduceRegion, ::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::ReduceDims, ::Val) where {T<:Integer} = + @eval @inline $dim_helper(tile::Tile{T}, dims::ReduceRegion, ::Val) where {T<:Integer} = reduce($integer_op, tile; dims, init=$init_expr) @eval @inline $all_helper(tile::Tile{T,Tuple{}}, ::Val) where {T<:Number} = @@ -1555,35 +1570,38 @@ for f in (:sum, :prod) end """ - any(tile::Tile{Bool,S}; dims) -> Tile{Bool, reduced_shape} + any(tile::Tile{Bool,S}; dims) -> Tile or Bool -Logical OR reduction along the specified axis (1-indexed). -Reduced dimensions become size 1. +Logical OR reduction over `dims`, which accepts a 1-indexed axis, a tuple or +integer range of axes, or `:`. Reduced dimensions become size 1; `dims=:` +returns a scalar. """ -@inline Base.any(tile::Tile{Bool,S}; dims::ReduceDims) where {S} = +@inline Base.any(tile::Tile{Bool,S}; dims::ReduceRegion) where {S} = reduce(|, tile; dims, init=false) """ - all(tile::Tile{Bool,S}; dims) -> Tile{Bool, reduced_shape} + all(tile::Tile{Bool,S}; dims) -> Tile or Bool -Logical AND reduction along the specified axis (1-indexed). -Reduced dimensions become size 1. +Logical AND reduction over `dims`, which accepts a 1-indexed axis, a tuple or +integer range of axes, or `:`. Reduced dimensions become size 1; `dims=:` +returns a scalar. """ -@inline Base.all(tile::Tile{Bool,S}; dims::ReduceDims) where {S} = +@inline Base.all(tile::Tile{Bool,S}; dims::ReduceRegion) where {S} = reduce(&, tile; dims, init=true) """ - count(tile::Tile{Bool,S}; dims) -> Tile{Int32, reduced_shape} + count(tile::Tile{Bool,S}; dims) -> Tile or Int32 -Count true elements along `dims` (1-indexed). Apply predicates via -broadcasting before calling count: +Count true elements along `dims`, which accepts a 1-indexed axis, a tuple or +integer range of axes, or `:`. 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::ReduceDims) where {S} +@inline function Base.count(tile::Tile{Bool,S}; dims::ReduceRegion) where {S} sum(convert(Tile{Int32}, tile); dims) end diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 436ae9ff..bc670f5a 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1048,6 +1048,30 @@ spec4d = ct.ArraySpec{4}(16, true) end end + # Constant integer ranges lower like their equivalent tuples + @test @filecheck begin + @check_label "entry" + @check_count 2 "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)) + return + end + end + + # Colon reduces all dimensions to a scalar + @test @filecheck begin + @check_label "entry" + @check_count 2 "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=:)) + return + end + end + # dims=() and axes beyond ndims reduce nothing (Base parity) @test @filecheck begin @check_label "entry" diff --git a/test/device/reductions.jl b/test/device/reductions.jl index faca66f8..5c6aacd4 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -161,6 +161,18 @@ end 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 + m, n = 8, 16 a = CUDA.rand(Float32, m, n) a_cpu = Array(a) @@ -183,6 +195,14 @@ end @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 + # Base parity: axes beyond ndims are size-1, reducing them is a no-op function sum_dims23_kernel(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,2}) tile = ct.load(a, (1, 1), (8, 16)) From ac0c2a5a80f2ed682b0e903f841d8e008f9d5e02 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 19:18:05 +0200 Subject: [PATCH 05/11] Validate host reduction dimensions --- src/mapreduce.jl | 6 ++++-- test/host/mapreduce.jl | 3 +++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 469f6bf8..771d8537 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -150,8 +150,10 @@ function _mapreduce(f, op, A::AbstractArray; dims, init) reduce_dims = dims === Colon() ? ntuple(identity, N) : dims isa Integer ? (dims,) : Tuple(dims) - all(>=(1), reduce_dims) || - throw(ArgumentError("region dimension(s) must be ≥ 1, got $dims")) + for d in reduce_dims + 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 # Following Base, dims beyond ndims are implicit size-1 dimensions: no-op reduce_dims = filter(<=(N), reduce_dims) out_size = ntuple(d -> d in reduce_dims ? 1 : size(A, d), N) diff --git a/test/host/mapreduce.jl b/test/host/mapreduce.jl index caab5adf..bee4f053 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -64,6 +64,9 @@ using CUDA @test size(R) == (64, 1) @test Array(R) ≈ Array(sum(A; dims=2)) @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 From f21c88de1fd99fca662978e0730070aed8b436ae Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 19:53:36 +0200 Subject: [PATCH 06/11] Simplify reduction dimension handling --- src/compiler/intrinsics/core.jl | 2 +- src/compiler/utils.jl | 11 +-- src/language/operations.jl | 128 ++++++++++++-------------------- src/mapreduce.jl | 13 ++-- test/codegen/operations.jl | 63 +++++++--------- test/device/reductions.jl | 43 +++-------- test/host/mapreduce.jl | 30 ++------ 7 files changed, 103 insertions(+), 187 deletions(-) diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl index 350ee7eb..d4a8229b 100644 --- a/src/compiler/intrinsics/core.jl +++ b/src/compiler/intrinsics/core.jl @@ -795,7 +795,7 @@ 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) 0 <= julia_axis < ndim || throw(IRError("reduce() dimension must be in 1:$ndim; got $(julia_axis + 1)")) diff --git a/src/compiler/utils.jl b/src/compiler/utils.jl index c3dd7fbe..9272ecfb 100644 --- a/src/compiler/utils.jl +++ b/src/compiler/utils.jl @@ -587,18 +587,15 @@ function resolve_tuple(ctx::CGCtx, @nospecialize(arg), name::AbstractString) ] end - # Multi-value tuple result consumed whole (e.g. one intrinsic's tuple - # output fed into another, as in chained reduce): split into component - # CGVals the same way getfield extraction does. 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_type = T_jl.parameters[i] - CGVal(tv.v[i], tile_type_for_julia!(ctx, elem_type), elem_type, - RowMajorShape(extract_tile_shape(elem_type))) + 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 1:length(tv.v) + for i in eachindex(tv.v) ] end diff --git a/src/language/operations.jl b/src/language/operations.jl index cde4b7a4..b8cd8fe0 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,60 +1397,38 @@ 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 -# `dims` for tile reductions: a 1-indexed axis, a tuple or integer range of axes, -# or `:`. -# Each effective axis lowers to one `cuda_tile.reduce`; an empty region reduces -# nothing. Following Base, axes beyond `ndims` are implicit size-1 dimensions -# and reduce as a no-op; axes < 1 error during codegen. -const ReduceDims = Union{Integer, Tuple{Vararg{Integer}}, AbstractRange{<:Integer}} -const ReduceRegion = Union{Colon, ReduceDims} -_reduce_dims_tuple(dims::Integer) = (dims,) -_reduce_dims_tuple(dims::Tuple{Vararg{Integer}}) = dims - -_insert_reduce_dim(d, ::Tuple{}) = (d,) -function _insert_reduce_dim(d, dims::Tuple) - head = first(dims) - d == head && return dims - d < head ? (d, dims...) : (head, _insert_reduce_dim(d, Base.tail(dims))...) -end - -_canonical_reduce_dims(::Tuple{}, n::Int) = () -function _canonical_reduce_dims(dims::Tuple, n::Int) - tail = _canonical_reduce_dims(Base.tail(dims), n) - first(dims) <= n ? _insert_reduce_dim(first(dims), tail) : tail -end - -_canon_reduce_dims(dims::ReduceDims, n::Int) = - _canonical_reduce_dims(_reduce_dims_tuple(dims), n) -function _canon_reduce_dims(dims::AbstractRange{<:Integer}, n::Int) - invalid = !isempty(dims) && min(first(dims), last(dims)) < 1 - invalid && return (min(first(dims), last(dims)),) - candidates = ntuple(d -> d in dims ? d : n + 1, n) - _canonical_reduce_dims(candidates, n) -end -_canon_reduce_dims(::Colon, ::Int) = Colon() - -_reduce_tiles(tiles::Tuple, ::Tuple{}, f, init::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) -_reduction_scalar(tile::Tile) = Intrinsics.to_scalar(reshape(tile, ())) -function _reduce_tiles(tiles::Tuple, ::Colon, f, init::Tuple) - dims = ntuple(identity, ndims(first(tiles))) - map(_reduction_scalar, _reduce_tiles(tiles, dims, f, init)) -end - -""" - mapreduce(identity, f, tile::Tile{T,S}; dims, init) -> Tile or scalar - mapreduce(f, op, tile::Tile{T,S}; dims, init) -> Tile or scalar - mapreduce(identity, f, tile1, tile2, ...; dims, init) -> Tuple{Tile...} +function reduction_dims(dims, n) + isempty(dims) && return () + eltype(dims) <: Integer || throw(ArgumentError("reduced dimension(s) must be integers")) + invalid = minimum(dims; init=1) + invalid < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $invalid")) + foldl(1:n; init=()) do region, dim + dim in dims ? (region..., dim) : region + end +end +reduction_dims(::Colon, _) = Colon() + +function reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) + foldl(dims; init=tiles) do reduced, dim + Intrinsics.reduce(reduced, dim - 1, f, init) + end +end +function reduce_tiles(tiles::Tuple, ::Colon, f, init::Tuple) + reduced = reduce_tiles(tiles, ntuple(identity, ndims(first(tiles))), f, init) + map(tile -> Intrinsics.to_scalar(reshape(tile, ())), reduced) +end + +""" + 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` is a 1-indexed axis, a tuple or integer range of axes, or `:`. Each -reduced axis becomes size 1 (Julia semantics). `dims=:` reduces all axes and -returns a scalar. `dims=()` reduces nothing and returns the (mapped) input -unchanged. +`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: @@ -1474,35 +1452,32 @@ vals, idxs = mapreduce(identity, reducer, vals_tile, idx_tile; ``` """ @inline function Base.mapreduce(::typeof(identity), f, tile::Tile{T,S}; - dims::ReduceRegion, init) where {T<:Number, S} - _reduce_tiles((tile,), _canon_reduce_dims(dims, ndims(tile)), f, (T(init),))[1] + dims, init) where {T<:Number, S} + reduce_tiles((tile,), reduction_dims(dims, ndims(tile)), f, (T(init),))[1] end @inline function Base.mapreduce(f, op, tile::Tile{T,S}; - dims::ReduceRegion, 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::ReduceRegion, + 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 - _reduce_tiles(all_tiles, _canon_reduce_dims(dims, ndims(tile1)), _combiner, init) + reduce_tiles(all_tiles, reduction_dims(dims, ndims(tile1)), _combiner, init) end """ - reduce(f, tile::Tile{T,S}; dims, init) -> Tile or scalar + reduce(f, tile::Tile{T,S}; dims, init) Reduce a tile along the specified dimensions using binary function `f` with -identity element `init`. `dims` is a 1-indexed axis, a tuple or integer range -of axes, or `:`. Each reduced axis becomes size 1 (Julia semantics). `dims=:` -reduces all axes and returns a scalar. `dims=()` reduces nothing and returns -the tile unchanged. +identity element `init`. See [`mapreduce`](@ref) for supported `dims` values. Supported functions: `+`, `*`, `max`, `min`. @@ -1512,7 +1487,7 @@ 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::ReduceRegion, 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 @@ -1539,9 +1514,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::ReduceRegion, ::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::ReduceRegion, ::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} = @@ -1570,38 +1545,33 @@ for f in (:sum, :prod) end """ - any(tile::Tile{Bool,S}; dims) -> Tile or Bool + any(tile::Tile{Bool,S}; dims) -Logical OR reduction over `dims`, which accepts a 1-indexed axis, a tuple or -integer range of axes, or `:`. Reduced dimensions become size 1; `dims=:` -returns a scalar. +Logical OR reduction over `dims`. """ -@inline Base.any(tile::Tile{Bool,S}; dims::ReduceRegion) where {S} = +@inline Base.any(tile::Tile{Bool,S}; dims) where {S} = reduce(|, tile; dims, init=false) """ - all(tile::Tile{Bool,S}; dims) -> Tile or Bool + all(tile::Tile{Bool,S}; dims) -Logical AND reduction over `dims`, which accepts a 1-indexed axis, a tuple or -integer range of axes, or `:`. Reduced dimensions become size 1; `dims=:` -returns a scalar. +Logical AND reduction over `dims`. """ -@inline Base.all(tile::Tile{Bool,S}; dims::ReduceRegion) where {S} = +@inline Base.all(tile::Tile{Bool,S}; dims) where {S} = reduce(&, tile; dims, init=true) """ - count(tile::Tile{Bool,S}; dims) -> Tile or Int32 + count(tile::Tile{Bool,S}; dims) -Count true elements along `dims`, which accepts a 1-indexed axis, a tuple or -integer range of axes, or `:`. 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::ReduceRegion) where {S} +@inline function Base.count(tile::Tile{Bool,S}; dims) where {S} sum(convert(Tile{Int32}, tile); dims) end diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 771d8537..0d9f7b02 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -71,7 +71,7 @@ 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) - # dims=(): nothing to reduce; each block maps its tile through op(init, f(x)) + # 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) @@ -148,14 +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) - for d in reduce_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")) + d < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $d")) end - # Following Base, dims beyond ndims are implicit size-1 dimensions: no-op - reduce_dims = filter(<=(N), reduce_dims) + 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 bc670f5a..ae710d95 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1025,54 +1025,30 @@ spec4d = ct.ArraySpec{4}(16, true) end end - # Tuple dims are sorted and deduplicated before lowering @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)) - @check "reduce {{.*}} dim=1" - @check "reduce {{.*}} dim=0" - Base.donotdelete(sum(tile; dims=(2, 1))) - return - end - end - @test @filecheck begin - @check_label "entry" - @check_count 1 "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, 1))) + Base.donotdelete(sum(tile; dims=(2, 1, 2))) return end end - # Constant integer ranges lower like their equivalent tuples @test @filecheck begin @check_label "entry" - @check_count 2 "reduce" + @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)) - return - end - end - - # Colon reduces all dimensions to a scalar - @test @filecheck begin - @check_label "entry" - @check_count 2 "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=:)) return end end - # dims=() and axes beyond ndims reduce nothing (Base parity) @test @filecheck begin @check_label "entry" code_tiled(Tuple{ct.TileArray{Float32,2,spec2d}}) do a @@ -1086,8 +1062,7 @@ spec4d = ct.ArraySpec{4}(16, true) end end - # dims < 1 is a compile-time error - @test_throws "reduce() dimension must be in 1:2" code_tiled( + @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)) @@ -1095,6 +1070,22 @@ spec4d = ct.ArraySpec{4}(16, true) 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 "reduce() dimension must be in 1:2" 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 + # any @test @filecheck begin @check_label "entry" @@ -1102,9 +1093,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 @@ -1116,9 +1107,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 @@ -1130,9 +1121,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 5c6aacd4..662351a9 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -133,19 +133,7 @@ end end end -@testset "reduce with tuple dims" begin - function sum_dims12_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_dims2_tuple_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,))) - return - 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=())) @@ -173,18 +161,16 @@ end 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, 1, 1) - @cuda backend=cuTile sum_dims12_kernel(a, b) - @test Array(b)[1, 1] ≈ sum(a_cpu) rtol=1e-3 - - b = CUDA.zeros(Float32, m, 1) - @cuda backend=cuTile sum_dims2_tuple_kernel(a, b) - @test Array(b) ≈ sum(a_cpu; dims=2) rtol=1e-3 - b = CUDA.zeros(Float32, m, n) @cuda backend=cuTile sum_dims_empty_kernel(a, b) @test Array(b) == a_cpu @@ -203,19 +189,12 @@ end @cuda backend=cuTile sum_dims_colon_kernel(a, b_scalar) @test Array(b_scalar)[1] ≈ sum(a_cpu; dims=:) rtol=1e-3 - # Base parity: axes beyond ndims are size-1, reducing them is a no-op - 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 - 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 "tuple dims reductions (3D)" begin +@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))) @@ -241,7 +220,7 @@ end @test Array(b) ≈ mapreduce(abs, +, a_cpu; dims=(1, 3)) rtol=1e-3 end -@testset "count with tuple dims" begin +@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))) @@ -254,9 +233,7 @@ end @test Array(b)[1, 1] == count(>(0.0f0), Array(a)) end -@testset "multi-tile mapreduce with tuple dims" begin - # Simultaneous (value, tag) reduction over both dims: the reducer keeps - # the pair whose first component is larger, so tags must follow values. +@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)) diff --git a/test/host/mapreduce.jl b/test/host/mapreduce.jl index bee4f053..f3b7468e 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -27,25 +27,11 @@ using CUDA @test Array(R) ≈ Array(sum(A; dims=2)) end - @testset "sum 2D dims=(1,2)" begin + @testset "sum with multiple dims" begin A = CUDA.rand(Float32, 128, 256) - R = sum(ct.Tiled(A); dims=(1, 2)) - @test size(R) == (1, 1) - @test Array(R) ≈ Array(sum(A; dims=(1, 2))) - end - - @testset "sum 2D dims=(2,)" begin - A = CUDA.rand(Float32, 128, 256) - R = sum(ct.Tiled(A); dims=(2,)) - @test size(R) == (128, 1) - @test Array(R) ≈ Array(sum(A; dims=2)) - end - - @testset "sum 2D dims=()" begin - A = CUDA.rand(Float32, 128, 256) - R = sum(ct.Tiled(A); dims=()) - @test size(R) == (128, 256) - @test Array(R) ≈ Array(A) + 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 @@ -55,14 +41,11 @@ using CUDA @test Array(prod(ct.Tiled(A); dims=2)) == Array(prod(A; dims=2)) end - @testset "dims beyond ndims (Base parity)" begin + @testset "out-of-range dims" begin A = CUDA.rand(Float32, 64, 32) - R = sum(ct.Tiled(A); dims=3) - @test size(R) == (64, 32) - @test Array(R) ≈ Array(A) R = sum(ct.Tiled(A); dims=(2, 3)) - @test size(R) == (64, 1) @test Array(R) ≈ Array(sum(A; dims=2)) + @test Array(sum(ct.Tiled(A); dims=big(typemax(Int)) + 1)) ≈ Array(A) @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,)) @@ -72,7 +55,6 @@ using CUDA @testset "maximum 3D dims=(1,3)" begin A = CUDA.rand(Float32, 32, 16, 8) R = maximum(ct.Tiled(A); dims=(1, 3)) - @test size(R) == (1, 16, 1) @test Array(R) ≈ Array(maximum(A; dims=(1, 3))) end From 7c57cb1d17f630dd36733f3f081eaccca2947f85 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 19:58:29 +0200 Subject: [PATCH 07/11] Support tuple dimensions in dropdims --- src/language/operations.jl | 31 ++++++++++++++++++++----------- test/codegen/operations.jl | 19 +++++++++++++++++-- 2 files changed, 37 insertions(+), 13 deletions(-) diff --git a/src/language/operations.jl b/src/language/operations.jl index b8cd8fe0..c24bdeb6 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1668,25 +1668,34 @@ end end +check_drop_dims(::Tile, ::Tuple{}, ::Any) = nothing +function check_drop_dims(tile::Tile, dims::Tuple, seen=()) + 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 + """ 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) + shape = foldl(1:ndims(tile); init=()) do shape, dim + dim in region ? shape : (shape..., size(tile, dim)) + end + reshape(tile, shape) end #============================================================================= diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index ae710d95..9f2b9cf0 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 From a58f3df94ae0518c98bb1e8ed2c2ce47ad82a99c Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 19:59:04 +0200 Subject: [PATCH 08/11] Document reduction regions --- docs/src/lib/operations.md | 27 +++++++++++++++------------ docs/src/man/comparison.md | 5 +++++ 2 files changed, 20 insertions(+), 12 deletions(-) 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 From 1606c19a51b5e4facb6b49f1e2ccab35d0e263cb Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 20:27:48 +0200 Subject: [PATCH 09/11] Preserve static reduction tuple lengths --- src/compiler/intrinsics/core.jl | 4 ++- src/language/operations.jl | 55 ++++++++++++++++++++++----------- test/codegen/operations.jl | 10 +++++- test/device/reductions.jl | 10 ++++++ 4 files changed, 59 insertions(+), 20 deletions(-) diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl index d4a8229b..2c0f0edf 100644 --- a/src/compiler/intrinsics/core.jl +++ b/src/compiler/intrinsics/core.jl @@ -797,7 +797,9 @@ function emit_reduce!(ctx::CGCtx, args) # Julia and Tile IR use opposite dimension orders. ndim = length(input_shape) - 0 <= julia_axis < ndim || + 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 diff --git a/src/language/operations.jl b/src/language/operations.jl index c24bdeb6..44e262bf 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,27 +1397,41 @@ 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 -function reduction_dims(dims, n) +reduction_dims(dim::Integer, ::Tuple) = dim +function reduction_dims(dims, shape::Tuple) isempty(dims) && return () eltype(dims) <: Integer || throw(ArgumentError("reduced dimension(s) must be integers")) invalid = minimum(dims; init=1) invalid < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $invalid")) - foldl(1:n; init=()) do region, dim - dim in dims ? (region..., dim) : region - end + 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 -reduction_dims(::Colon, _) = Colon() -function reduce_tiles(tiles::Tuple, dims::Tuple, f, init::Tuple) - foldl(dims; init=tiles) do reduced, dim - Intrinsics.reduce(reduced, dim - 1, f, init) - 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) - reduced = reduce_tiles(tiles, ntuple(identity, ndims(first(tiles))), f, init) + 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) mapreduce(f, op, tile::Tile{T,S}; dims, init) @@ -1453,7 +1467,7 @@ vals, idxs = mapreduce(identity, reducer, vals_tile, idx_tile; """ @inline function Base.mapreduce(::typeof(identity), f, tile::Tile{T,S}; dims, init) where {T<:Number, S} - reduce_tiles((tile,), reduction_dims(dims, ndims(tile)), f, (T(init),))[1] + reduce_tile(tile, reduction_dims(dims, size(tile)), f, T(init)) end @inline function Base.mapreduce(f, op, tile::Tile{T,S}; @@ -1470,7 +1484,7 @@ end function _combiner(args...) f(_deinterleave_accs(args...), _deinterleave_elems(args...)) end - reduce_tiles(all_tiles, reduction_dims(dims, ndims(tile1)), _combiner, init) + reduce_tiles(all_tiles, reduction_dims(dims, size(tile1)), _combiner, init) end """ @@ -1668,8 +1682,9 @@ end end -check_drop_dims(::Tile, ::Tuple{}, ::Any) = nothing -function check_drop_dims(tile::Tile, dims::Tuple, seen=()) +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)")) @@ -1678,6 +1693,13 @@ function check_drop_dims(tile::Tile, dims::Tuple, seen=()) 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} @@ -1692,10 +1714,7 @@ scalar_tile = dropdims(sums; dims=(1, 2)) function Base.dropdims(tile::Tile; dims::Union{Integer,Dims}) region = dims isa Integer ? (Int(dims),) : dims check_drop_dims(tile, region) - shape = foldl(1:ndims(tile); init=()) do shape, dim - dim in region ? shape : (shape..., size(tile, dim)) - end - reshape(tile, shape) + reshape(tile, drop_shape(size(tile), region)) end #============================================================================= diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 9f2b9cf0..6d753c5d 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1093,7 +1093,7 @@ spec4d = ct.ArraySpec{4}(16, true) return end - @test_throws "reduce() dimension must be in 1:2" code_tiled( + @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)) @@ -1101,6 +1101,14 @@ spec4d = ct.ArraySpec{4}(16, true) 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" diff --git a/test/device/reductions.jl b/test/device/reductions.jl index 662351a9..c2177095 100644 --- a/test/device/reductions.jl +++ b/test/device/reductions.jl @@ -207,6 +207,12 @@ end 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) @@ -218,6 +224,10 @@ end 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 From 1d49f04998d3c473acce196cf73933ec5f3e1954 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 5 Aug 2026 22:46:12 +0200 Subject: [PATCH 10/11] Fix reduction validation on Julia 1.11 --- src/language/operations.jl | 5 +++-- test/codegen/operations.jl | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/language/operations.jl b/src/language/operations.jl index 44e262bf..f050b43c 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1401,8 +1401,9 @@ reduction_dims(dim::Integer, ::Tuple) = dim function reduction_dims(dims, shape::Tuple) isempty(dims) && return () eltype(dims) <: Integer || throw(ArgumentError("reduced dimension(s) must be integers")) - invalid = minimum(dims; init=1) - invalid < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $invalid")) + invalid = minimum(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() diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 6d753c5d..7f666efd 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -1085,6 +1085,14 @@ spec4d = ct.ArraySpec{4}(16, true) 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) From a43f244f49db0af488eadc67b74483824ef84440 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 6 Aug 2026 07:33:23 +0200 Subject: [PATCH 11/11] Match Base reduction dimension validation --- src/language/operations.jl | 4 ++-- src/mapreduce.jl | 2 +- test/host/mapreduce.jl | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/language/operations.jl b/src/language/operations.jl index f050b43c..1ac5399a 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1397,11 +1397,11 @@ 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) = dim +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 = minimum(dims) + 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) diff --git a/src/mapreduce.jl b/src/mapreduce.jl index 0d9f7b02..6c61d4e6 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -152,7 +152,7 @@ function _mapreduce(f, op, A::AbstractArray; dims, init) dims isa Integer ? (dims,) : Tuple(dims) for d in region d isa Integer || throw(ArgumentError("reduced dimension(s) must be integers")) - d < 1 && throw(ArgumentError("region dimension(s) must be ≥ 1, got $d")) + 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) diff --git a/test/host/mapreduce.jl b/test/host/mapreduce.jl index f3b7468e..43538156 100644 --- a/test/host/mapreduce.jl +++ b/test/host/mapreduce.jl @@ -45,7 +45,7 @@ using CUDA A = CUDA.rand(Float32, 64, 32) R = sum(ct.Tiled(A); dims=(2, 3)) @test Array(R) ≈ Array(sum(A; dims=2)) - @test Array(sum(ct.Tiled(A); dims=big(typemax(Int)) + 1)) ≈ Array(A) + @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,))