Support reductions without dims arg. (#118)

Author: maleadt

src/language/operations.jl

@@ -200,6 +200,16 @@ end
     store(arr, (index,), tile; kwargs...)
 end
 
+# Scalar value → wrap in 1-element tile and store
+@inline function store(arr::TileArray{T}, index, val::T; kwargs...) where {T}
+    shape = ntuple(_ -> 1, Val(ndims(arr)))
+    tile = reshape(Intrinsics.from_scalar(val, Tuple{}), shape)
+    store(arr, index, tile; kwargs...)
+end
+@inline function store(arr::TileArray{T}, index::Integer, val::T; kwargs...) where {T}
+    store(arr, (index,), val; kwargs...)
+end
+
 @inline function _store_reshaped(arr::TileArray{T}, tile::Tile{T},
                                  order, latency, allow_tma, indices::NTuple{<:Any, <:Integer}) where {T}
     tv = Intrinsics.make_tensor_view(arr)
@@ -221,7 +231,7 @@ end
 # would DCE the entire call as a pure function with unused result.
 Base.Experimental.@consistent_overlay cuTileMethodTable function Base.setindex!(arr::TileArray{T, N}, val::T, indices::Vararg{Integer, N}) where {T, N}
     shape = ntuple(_ -> 1, Val(N))
-    tile = reshape(Intrinsics.from_scalar(val, Val(Tuple{})), shape)
+    tile = reshape(Intrinsics.from_scalar(val, Tuple{}), shape)
     store(arr, indices, tile)
     return
 end
@@ -697,6 +707,16 @@ Reduced dimensions become size 1.
 @inline Base.minimum(tile::Tile{T,S}; dims) where {T<:Number, S} =
     reduce(min, tile; dims, init=typemax(T))
 
+# sum/prod/max/min without dims — reduce all dimensions, return scalar T
+# Recursive: reduce dim 1, dropdims, recurse. Base case: 0D tile → scalar.
+for f in (:sum, :prod, :maximum, :minimum)
+    @eval @inline Base.$f(tile::Tile{T,Tuple{}}) where {T<:Number} =
+        Intrinsics.to_scalar(tile)
+
+    @eval @inline Base.$f(tile::Tile{T,S}) where {T<:Number, S<:Tuple{Any,Vararg}} =
+        $f(dropdims($f(tile; dims=1); dims=1))
+end
+
 """
     any(tile::Tile{Bool,S}; dims) -> Tile{Bool, reduced_shape}
 
@@ -730,6 +750,22 @@ n_positive = count(tile .> 0.0f0; dims=1)
     sum(convert(Tile{Int32}, tile); dims)
 end
 
+# any/all without dims — return scalar Bool
+for f in (:any, :all)
+    @eval @inline Base.$f(tile::Tile{Bool,Tuple{}}) =
+        Intrinsics.to_scalar(tile)
+
+    @eval @inline Base.$f(tile::Tile{Bool,S}) where {S<:Tuple{Any,Vararg}} =
+        $f(dropdims($f(tile; dims=1); dims=1))
+end
+
+# count without dims — return scalar Int32
+# count reduces to Int32 (not Bool), so after first dim use sum for remaining.
+@inline Base.count(tile::Tile{Bool,Tuple{}}) =
+    Intrinsics.to_scalar(tile)
+@inline Base.count(tile::Tile{Bool,S}) where {S<:Tuple{Any,Vararg}} =
+    sum(dropdims(count(tile; dims=1); dims=1))
+
 """
     argmax(tile::Tile{T,S}; dims) -> Tile{Int32, reduced_shape}
 

test/execution/basic.jl

@@ -151,6 +151,22 @@ end
     end
 end
 
+@testset "scalar store" begin
+    function scalar_store_1d(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, tileSz::Int)
+        tile = ct.load(a, ct.bid(1), (tileSz,))
+        ct.store(b, ct.bid(1), sum(tile))
+        return nothing
+    end
+
+    sz = 32; N = 1024
+    a = CUDA.rand(Float32, N)
+    b = CUDA.zeros(Float32, cld(N, sz))
+    ct.launch(scalar_store_1d, cld(N, sz), a, b, ct.Constant(sz))
+
+    a_cpu = reshape(Array(a), sz, :)
+    @test Array(b) ≈ vec(sum(a_cpu; dims=1)) rtol=1e-3
+end
+
 @testset "transpose" begin
     function transpose_kernel(x::ct.TileArray{Float32,2}, y::ct.TileArray{Float32,2})
         bidx = ct.bid(1)

test/execution/reductions.jl

@@ -567,3 +567,131 @@ end
         @test Array(b_min_rand)[row, 1] == expected_min
     end
 end
+
+@testset "sum without dims (1D)" begin
+    function sum_no_dims_1d(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, tileSz::Int)
+        tile = ct.load(a, ct.bid(1), (tileSz,))
+        b[ct.bid(1)] = sum(tile)
+        return nothing
+    end
+
+    sz = 32; N = 1024
+    a = CUDA.rand(Float32, N)
+    b = CUDA.zeros(Float32, cld(N, sz))
+    ct.launch(sum_no_dims_1d, cld(N, sz), a, b, ct.Constant(sz))
+
+    a_cpu = reshape(Array(a), sz, :)
+    @test Array(b) ≈ vec(sum(a_cpu; dims=1)) rtol=1e-3
+end
+
+@testset "sum without dims (2D)" begin
+    function sum_no_dims_2d(a::ct.TileArray{Float32,2}, b::ct.TileArray{Float32,1})
+        pid = ct.bid(1)
+        tile = ct.load(a, (pid, 1), (1, 128))
+        b[pid] = sum(tile)
+        return nothing
+    end
+
+    m, n = 64, 128
+    a = CUDA.rand(Float32, m, n)
+    b = CUDA.zeros(Float32, m)
+    ct.launch(sum_no_dims_2d, m, a, b)
+
+    a_cpu = Array(a)
+    for i in 1:m
+        @test Array(b)[i] ≈ sum(a_cpu[i, :]) rtol=1e-3
+    end
+end
+
+@testset "any without dims (1D)" begin
+    function any_no_dims_1d(a::ct.TileArray{Float32,1}, b::ct.TileArray{Int32,1}, tileSz::Int)
+        tile = ct.load(a, ct.bid(1), (tileSz,))
+        # Store as Int32 since we can't store Bool scalars directly
+        b[ct.bid(1)] = Int32(any(tile .> 0.0f0))
+        return nothing
+    end
+
+    sz = 32; N = 1024
+    a = CUDA.rand(Float32, N) .- 0.5f0  # some positive, some negative
+    b = CUDA.zeros(Int32, cld(N, sz))
+    ct.launch(any_no_dims_1d, cld(N, sz), a, b, ct.Constant(sz))
+
+    a_cpu = reshape(Array(a), sz, :)
+    for i in 1:cld(N, sz)
+        @test Array(b)[i] == Int32(any(a_cpu[:, i] .> 0.0f0))
+    end
+end
+
+@testset "all without dims (1D)" begin
+    function all_no_dims_1d(a::ct.TileArray{Float32,1}, b::ct.TileArray{Int32,1}, tileSz::Int)
+        tile = ct.load(a, ct.bid(1), (tileSz,))
+        b[ct.bid(1)] = Int32(all(tile .> 0.0f0))
+        return nothing
+    end
+
+    sz = 32; N = 1024
+    a = CUDA.rand(Float32, N)  # all positive
+    b = CUDA.zeros(Int32, cld(N, sz))
+    ct.launch(all_no_dims_1d, cld(N, sz), a, b, ct.Constant(sz))
+
+    a_cpu = reshape(Array(a), sz, :)
+    for i in 1:cld(N, sz)
+        @test Array(b)[i] == Int32(all(a_cpu[:, i] .> 0.0f0))
+    end
+end
+
+@testset "sum without dims (3D)" begin
+    function sum_no_dims_3d(a::ct.TileArray{Float32,3}, b::ct.TileArray{Float32,1})
+        pid = ct.bid(1)
+        tile = ct.load(a, (pid, 1, 1), (1, 8, 16))
+        b[pid] = sum(tile)
+        return nothing
+    end
+
+    d1, d2, d3 = 4, 8, 16
+    a = CUDA.rand(Float32, d1, d2, d3)
+    b = CUDA.zeros(Float32, d1)
+    ct.launch(sum_no_dims_3d, d1, a, b)
+
+    a_cpu = Array(a)
+    for i in 1:d1
+        @test Array(b)[i] ≈ sum(a_cpu[i, :, :]) rtol=1e-3
+    end
+end
+
+@testset "maximum without dims (3D)" begin
+    function maximum_no_dims_3d(a::ct.TileArray{Float32,3}, b::ct.TileArray{Float32,1})
+        pid = ct.bid(1)
+        tile = ct.load(a, (pid, 1, 1), (1, 8, 16))
+        b[pid] = maximum(tile)
+        return nothing
+    end
+
+    d1, d2, d3 = 4, 8, 16
+    a = CUDA.rand(Float32, d1, d2, d3)
+    b = CUDA.zeros(Float32, d1)
+    ct.launch(maximum_no_dims_3d, d1, a, b)
+
+    a_cpu = Array(a)
+    for i in 1:d1
+        @test Array(b)[i] ≈ maximum(a_cpu[i, :, :]) rtol=1e-3
+    end
+end
+
+@testset "count without dims (1D)" begin
+    function count_no_dims_1d(a::ct.TileArray{Float32,1}, b::ct.TileArray{Int32,1}, tileSz::Int)
+        tile = ct.load(a, ct.bid(1), (tileSz,))
+        b[ct.bid(1)] = count(tile .> 0.0f0)
+        return nothing
+    end
+
+    sz = 32; N = 1024
+    a = CUDA.rand(Float32, N) .- 0.5f0
+    b = CUDA.zeros(Int32, cld(N, sz))
+    ct.launch(count_no_dims_1d, cld(N, sz), a, b, ct.Constant(sz))
+
+    a_cpu = reshape(Array(a), sz, :)
+    for i in 1:cld(N, sz)
+        @test Array(b)[i] == count(a_cpu[:, i] .> 0.0f0)
+    end
+end