More types for MPSGraph reduction/scan/sort support

lib/mpsgraphs/reductions.jl

@@ -1,4 +1,5 @@
-const MPSGRAPH_VALID_REDUCTION_TYPES = Union{Float16, Float32}
+
+const MPSGRAPH_VALID_REDUCTION_TYPES = filter(T -> T <: Real && T != Bool, (MPS.jl_mps_to_typ |> values |> collect))
 
 function reductionSumWithTensor(graph::MPSGraph, tensor::MPSGraphTensor,
                                 axes::NSArray, name = "reduction_sum")
@@ -100,8 +101,8 @@ const reduction_graph_cache = Dict{ReductionGraphKey, CachedReductionGraph}()
 const reduction_graph_cache_lock = ReentrantLock()
 
 function check_reduction_args(out::MtlArray{T}, input::MtlArray{T}) where {T}
-    T <: MPSGRAPH_VALID_REDUCTION_TYPES ||
-        throw(ArgumentError("MPSGraph reduction supports Float16 and Float32"))
+    T <: Union{MPSGRAPH_VALID_REDUCTION_TYPES...} ||
+        throw(ArgumentError("MPSGraph reduction supports $(join(MPSGraphs.MPSGRAPH_VALID_REDUCTION_TYPES,", ", " and "))"))
     dims = reduction_axes(size(out), size(input))
     check_mpsgraph_offsets(out, input)
     return dims

lib/mpsgraphs/scan.jl

@@ -1,4 +1,4 @@
-const MPSGRAPH_VALID_SCAN_TYPES = Union{Float16, Float32}
+const MPSGRAPH_VALID_SCAN_TYPES = filter(T -> T <: Real && T != Bool, (MPS.jl_mps_to_typ |> values |> collect))
 
 function cumulativeSumWithTensor(graph::MPSGraph, tensor::MPSGraphTensor,
                                  axis::Integer, exclusive::Bool,
@@ -40,21 +40,27 @@ function cumulativeMinimumWithTensor(graph::MPSGraph, tensor::MPSGraphTensor,
                                                       name:name::id{NSString}]::MPSGraphTensor
 end
 
-scan_operation(::typeof(+)) = :sum
-scan_operation(::typeof(Base.add_sum)) = :sum
-scan_operation(::typeof(*)) = :product
-scan_operation(::typeof(Base.mul_prod)) = :product
-scan_operation(op) =
-    throw(ArgumentError("MPSGraph scan supports + and *"))
+scan_operation(::DataType, ::typeof(+)) = :sum
+scan_operation(::DataType, ::typeof(Base.add_sum)) = :sum
+scan_operation(::DataType, ::typeof(*)) = :product
+scan_operation(::DataType, ::typeof(Base.mul_prod)) = :product
+scan_operation(::Type{<:Integer}, ::typeof(min)) = :minimum
+scan_operation(::Type{<:Integer}, ::typeof(max)) = :maximum
+scan_operation(_T, _op) =
+    throw(ArgumentError("MPSGraph scan supports + and * on all real MPS types, and min/max with integer types only"))
 
 function scanWithTensor(graph::MPSGraph, op::Symbol, tensor::MPSGraphTensor,
                         axis::Integer, name = "scan")
     if op === :sum
         cumulativeSumWithTensor(graph, tensor, axis, false, false, name)
     elseif op === :product
         cumulativeProductWithTensor(graph, tensor, axis, false, false, name)
+    elseif op === :minimum
+        cumulativeMinimumWithTensor(graph, tensor, axis, false, false, name)
+    elseif op === :maximum
+        cumulativeMaximumWithTensor(graph, tensor, axis, false, false, name)
     else
-        throw(ArgumentError("MPSGraph scan supports + and *"))
+        throw(ArgumentError("MPSGraph scan supports + and * on all real MPS types, and min/max with integer types only"))
     end
 end
 
@@ -82,8 +88,8 @@ const scan_graph_cache = Dict{ScanGraphKey, CachedScanGraph}()
 const scan_graph_cache_lock = ReentrantLock()
 
 function check_scan_args(out::MtlArray{T}, input::MtlArray{T}, dim::Integer) where {T}
-    T <: MPSGRAPH_VALID_SCAN_TYPES ||
-        throw(ArgumentError("MPSGraph scan supports Float16 and Float32"))
+    T <: Union{MPSGRAPH_VALID_SCAN_TYPES...} ||
+        throw(ArgumentError("MPSGraph scan supports $(join(MPSGraphs.MPSGRAPH_VALID_SCAN_TYPES,", ", " and "))"))
     size(out) == size(input) ||
         throw(DimensionMismatch("output has dimensions $(size(out)), input has dimensions $(size(input))"))
     1 <= dim <= ndims(input) || throw(ArgumentError("dimension out of range"))
@@ -95,7 +101,7 @@ end
                                       dim::Integer = 1) where {T}
     dim = check_scan_args(out, input, dim)
     isempty(input) && return copyto!(out, input)
-    key = ScanGraphKey{T}(size(input), dim, scan_operation(op))
+    key = ScanGraphKey{T}(size(input), dim, scan_operation(T, op))
     cached = @lock scan_graph_cache_lock get!(scan_graph_cache, key) do
         CachedScanGraph(key)
     end

lib/mpsgraphs/sort.jl

@@ -1,4 +1,4 @@
-const MPSGRAPH_VALID_SORT_TYPES = Union{Float16, Float32, Int32, Int64}
+const MPSGRAPH_VALID_SORT_TYPES = filter(T -> T <: Real, (MPS.jl_mps_to_typ |> values |> collect))
 const MPSGRAPH_SORTPERM_INDEX_TYPES = Union{Int32, Int64}
 
 function sortWithTensor(graph::MPSGraph, tensor::MPSGraphTensor, axis::Integer,
@@ -102,7 +102,7 @@ const sortperm_graph_cache = Dict{SortPermGraphKey, CachedSortPermGraph}()
 const sortperm_graph_cache_lock = ReentrantLock()
 
 function check_sort_args(out::MtlArray{T}, input::MtlArray{T}, dim::Integer) where {T}
-    T <: MPSGRAPH_VALID_SORT_TYPES || throw(ArgumentError("MPSGraph sort supports Float16, Float32, Int32, and Int64"))
+    T <: Union{MPSGRAPH_VALID_SORT_TYPES...} || throw(ArgumentError("MPSGraph sort supports $(join(MPSGraphs.MPSGRAPH_VALID_SORT_TYPES,", ", " and "))"))
     size(out) == size(input) ||
         throw(DimensionMismatch("output has dimensions $(size(out)), input has dimensions $(size(input))"))
     1 <= dim <= ndims(input) || throw(ArgumentError("dimension out of range"))
@@ -112,7 +112,7 @@ end
 
 function check_sortperm_args(index::MtlArray{Ti}, input::MtlArray{T},
                              dim::Integer) where {Ti, T}
-    T <: MPSGRAPH_VALID_SORT_TYPES || throw(ArgumentError("MPSGraph sortperm supports Float16, Float32, Int32, and Int64 inputs"))
+    T <: Union{MPSGRAPH_VALID_SORT_TYPES...} || throw(ArgumentError("MPSGraph sortperm supports $(join(MPSGraphs.MPSGRAPH_VALID_SORT_TYPES,", ", " and ")) inputs"))
     Ti <: MPSGRAPH_SORTPERM_INDEX_TYPES || throw(ArgumentError("MPSGraph sortperm supports Int32 and Int64 indices"))
     size(index) == size(input) ||
         throw(DimensionMismatch("index output has dimensions $(size(index)), input has dimensions $(size(input))"))

perf/array.jl

@@ -148,11 +148,11 @@ let group = addgroup!(group, "random")
     end
 end
 
-# let group = addgroup!(group, "sorting")
-#     group["1d"] = @benchmarkable Metal.@sync sort($gpu_vec)
-#     group["2d"] = @benchmarkable Metal.@sync sort($gpu_mat; dims=1)
-#     group["by"] = @benchmarkable Metal.@sync sort($gpu_vec; by=sin)
-# end
+let group = addgroup!(group, "sorting")
+    group["1d"] = @benchmarkable Metal.@sync sort($gpu_vec)
+    group["2d"] = @benchmarkable Metal.@sync sort($gpu_mat; dims=1)
+    # group["by"] = @benchmarkable Metal.@sync sort($gpu_vec; by=sin)
+end
 
 let group = addgroup!(group, "permutedims")
     group["2d"] = @benchmarkable Metal.@sync permutedims($gpu_mat, (2,1))

src/accumulate.jl

@@ -172,18 +172,29 @@ end
 const scan_alg = ScopedValue(:auto)
 const mpsgraph_scan_threshold = 64 * 1024
 
+# MPSGraph has no efficient 64-bit-integer cumulative kernel: on the
+# `accumulate(+, rand(Int64, 3, 10^6); dims=1)` shape it is ~2.6× slower than the
+# native scan (vs ~0.85× for ≤32-bit ints and floats, which MPSGraph handles
+# well), and the slowdown grows as the scanned dimension shrinks. So keep 64-bit
+# integers on the native scan in `:auto`; they remain correct (just slow) and
+# available under an explicit `:MPSGraph` request.
+mpsgraph_scan_worthwhile(::Type{T}) where {T} = !(T === Int64 || T === UInt64)
+
 # MPSGraph cumulative max/min ignore NaNs while Base accumulate(max/min)
-# propagates them, so don't use MPSGraph scan for these operations.
-mpsgraph_scan_operation(::typeof(+)) = :sum
-mpsgraph_scan_operation(::typeof(Base.add_sum)) = :sum
-mpsgraph_scan_operation(::typeof(*)) = :product
-mpsgraph_scan_operation(::typeof(Base.mul_prod)) = :product
-mpsgraph_scan_operation(op) = nothing
+# propagates them, so don't use MPSGraph scan for these operations on
+# Float inputs
+mpsgraph_scan_operation(::DataType, ::typeof(+)) = :sum
+mpsgraph_scan_operation(::DataType, ::typeof(Base.add_sum)) = :sum
+mpsgraph_scan_operation(::DataType, ::typeof(*)) = :product
+mpsgraph_scan_operation(::DataType, ::typeof(Base.mul_prod)) = :product
+mpsgraph_scan_operation(::Type{<:Integer}, ::typeof(min)) = :minimum
+mpsgraph_scan_operation(::Type{<:Integer}, ::typeof(max)) = :maximum
+mpsgraph_scan_operation(_T, _op) = nothing
 
 function mpsgraph_scan_supported(op, output::MtlArray{T}, input::MtlArray{T},
                             dims::Integer, init::Nothing) where {T}
-    mpsgraph_scan_operation(op) === nothing && return false
-    T <: MPSGraphs.MPSGRAPH_VALID_SCAN_TYPES || return false
+    mpsgraph_scan_operation(T, op) === nothing && return false
+    T <: Union{MPSGraphs.MPSGRAPH_VALID_SCAN_TYPES...} || return false
     axes(output) == axes(input) || return false
     1 <= dims <= ndims(input) || return false
     return output.offset == 0 && input.offset == 0
@@ -208,7 +219,8 @@ function scan_with_algorithm!(op, output::WrappedMtlArray, input::WrappedMtlArra
     supported = mpsgraph_scan_supported(op, output, input, dims, init)
     if alg === :MPSGraph
         return mpsgraph_scan!(op, output, input; dims, init)
-    elseif alg === :auto && supported && length(input) >= mpsgraph_scan_threshold
+    elseif alg === :auto && supported && mpsgraph_scan_worthwhile(eltype(input)) &&
+           length(input) >= mpsgraph_scan_threshold
         return MPSGraphs.graph_scan!(op, output, input; dim=dims)
     elseif alg === :auto || alg === :native
         return scan!(op, output, input; dims, init)

src/mapreduce.jl

@@ -210,7 +210,7 @@ function mpsgraph_reduce_dimensions(f, op, R::MtlArray{T}, A::MtlArray{T},
                                init) where {T}
     f === identity || return nothing
     mpsgraph_reduction_operation(op) === nothing && return nothing
-    T <: MPSGraphs.MPSGRAPH_VALID_REDUCTION_TYPES || return nothing
+    T <: Union{MPSGraphs.MPSGRAPH_VALID_REDUCTION_TYPES...} || return nothing
     mpsgraph_reduction_init_supported(op, T, init) || return nothing
     R.offset == 0 && A.offset == 0 || return nothing
     return reduced_dimensions(R, A)

src/sorting.jl

@@ -13,15 +13,15 @@ function sort_descending(lt, by, rev::Union{Bool,Nothing}, order::Base.Order.Ord
     return rev === true ? !descending : descending
 end
 
-function mps_sort_descending(::Type{T}, lt, by, rev::Union{Bool,Nothing},
+function mpsgraph_sort_descending(::Type{T}, lt, by, rev::Union{Bool,Nothing},
                              order::Base.Order.Ordering) where {T}
-    T <: MPSGraphs.MPSGRAPH_VALID_SORT_TYPES || return nothing
+    T <: Union{MPSGraphs.MPSGRAPH_VALID_SORT_TYPES...} || return nothing
     return sort_descending(lt, by, rev, order)
 end
 
-mps_sort_descending(A::MtlArray{T}, lt, by, rev::Union{Bool,Nothing},
+mpsgraph_sort_descending(A::MtlArray{T}, lt, by, rev::Union{Bool,Nothing},
                     order::Base.Order.Ordering) where {T} =
-    A.offset == 0 ? mps_sort_descending(T, lt, by, rev, order) : nothing
+    A.offset == 0 ? mpsgraph_sort_descending(T, lt, by, rev, order) : nothing
 
 function invoke_base_sort!(v::AbstractVector{T}; alg, lt, by, rev, order, scratch) where {T}
     return invoke(Base.sort!, Tuple{AbstractVector{T}}, v; alg, lt, by, rev, order,
@@ -72,7 +72,7 @@ function check_sortperm_dim(A::MtlArray, dims)
     end
 end
 
-function mps_sort!(A::MtlArray; dim::Integer, rev::Bool)
+function mpsgraph_sort!(A::MtlArray; dim::Integer, rev::Bool)
     tmp = similar(A)
     MPSGraphs.graph_sort!(tmp, A; dim, rev)
     copyto!(A, tmp)
@@ -86,10 +86,10 @@ function Base.sort!(v::MtlVector{T};
                     rev::Union{Bool,Nothing}=nothing,
                     order::Base.Order.Ordering=Base.Order.Forward,
                     scratch::Union{Vector{T}, Nothing}=nothing) where {T}
-    descending = mps_sort_descending(v, lt, by, rev, order)
+    descending = mpsgraph_sort_descending(v, lt, by, rev, order)
     descending === nothing &&
         return invoke_base_sort!(v; alg, lt, by, rev, order, scratch)
-    return mps_sort!(v; dim=1, rev=descending)
+    return mpsgraph_sort!(v; dim=1, rev=descending)
 end
 
 function Base.sort!(A::MtlArray{T};
@@ -101,10 +101,10 @@ function Base.sort!(A::MtlArray{T};
                     order::Base.Order.Ordering=Base.Order.Forward,
                     scratch::Union{Vector{T}, Nothing}=nothing) where {T}
     dim = check_sort_dim(A, dims)
-    descending = mps_sort_descending(A, lt, by, rev, order)
+    descending = mpsgraph_sort_descending(A, lt, by, rev, order)
     descending === nothing &&
         return invoke_base_sort!(A; dims=dim, alg, lt, by, rev, order, scratch)
-    return mps_sort!(A; dim, rev=descending)
+    return mpsgraph_sort!(A; dim, rev=descending)
 end
 
 Base.sort(v::MtlVector; kws...) = sort!(copy(v); kws...)
@@ -129,7 +129,7 @@ function Base.sortperm(A::MtlArray;
                        order::Base.Order.Ordering=Base.Order.Forward,
                        scratch::Union{Vector{<:Integer}, Nothing}=nothing,
                        dims=nothing)
-    descending = mps_sort_descending(A, lt, by, rev, order)
+    descending = mpsgraph_sort_descending(A, lt, by, rev, order)
     descending === nothing &&
         return invoke_base_sortperm(A; alg, lt, by, rev, order, scratch, dims)
     index = similar(A, Int)
@@ -149,7 +149,7 @@ function Base.sortperm!(index::MtlArray{Ti}, A::MtlArray;
     axes(index) == axes(A) ||
         throw(ArgumentError("index array must have the same axes as the source array"))
     dim = check_sortperm_dim(A, dims)
-    descending = index.offset == 0 ? mps_sort_descending(A, lt, by, rev, order) : nothing
+    descending = index.offset == 0 ? mpsgraph_sort_descending(A, lt, by, rev, order) : nothing
     if descending === nothing || !(Ti <: MPSGraphs.MPSGRAPH_SORTPERM_INDEX_TYPES)
         basedims = dims === nothing ? nothing : dim
         return invoke_base_sortperm!(index, A; alg, lt, by, rev, order,

test/array.jl

@@ -655,14 +655,22 @@ end
     end
 
     @with (Metal.scan_alg => :MPSGraph) begin
-        int_input = Int32[1, 2, 3, 4]
-        @test_throws ArgumentError accumulate(+, MtlArray(int_input))
-        @test_throws ArgumentError accumulate(+, MtlArray(scan_input); dims=1,
-                                              init=1.0f0)
-
+        # min, max using MPSGraph supported for Integers
+        int_input = Int32.(scan_input)
+        @test Array(accumulate(max, MtlArray(int_input); dims=1)) ≈
+            accumulate(max, int_input; dims=1)
+        @test Array(accumulate(min, MtlArray(int_input); dims=2)) ≈
+            accumulate(min, int_input; dims=2)
+
+        # but not for floating point
         nan_input = Float32[1, NaN, 0.5, 2]
         @test_throws ArgumentError accumulate(max, MtlArray(nan_input))
         @test_throws ArgumentError accumulate(min, MtlArray(nan_input))
+
+        # invalid eltype and kwargs
+        @test_throws ArgumentError accumulate(+, MtlArray(Complex{Int32}.(int_input)))
+        @test_throws ArgumentError accumulate(+, MtlArray(scan_input); dims=1,
+                                              init=1.0f0)
     end
 
     large_nan_input = ones(Float32, Metal.mpsgraph_scan_threshold + 1)

test/mpsgraphs/reductions.jl

@@ -1,5 +1,9 @@
-@testset "reductions ($T)" for T in (Float16, Float32)
-    A = reshape(T.(1:24) ./ T(10), 3, 4, 2)
+@testset "reductions ($T)" for T in MPSGraphs.MPSGRAPH_VALID_REDUCTION_TYPES
+    A = if T <: Integer
+        rand(T.(1:3), 3, 4, 2)
+    else
+        reshape(T.(1:24) ./ T(10), 3, 4, 2)
+    end
 
     for dim in 1:3
         out_size = Base.setindex(size(A), 1, dim)
@@ -23,10 +27,12 @@
 end
 
 @testset "reduction unsupported input" begin
-    A = MtlArray(reshape(Int32.(1:6), 3, 2))
+    # unsupported input types
+    A = MtlArray(reshape(Complex{Int32}.(1:6), 3, 2))
     out = similar(A, (1, 2))
     @test_throws ArgumentError MPSGraphs.graph_mapreducedim!(+, out, A)
 
+    # offset input
     parent = MtlArray(Float32[1, 2, 3])
     offset_input = unsafe_wrap(MtlArray, pointer(parent, 2), 2)
     offset_out = similar(offset_input, (1,))

test/mpsgraphs/scan.jl

@@ -1,21 +1,23 @@
-@testset "scan ($T)" for T in (Float16, Float32)
-    A = reshape(T[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37], 3, 4)
+@testset "scan ($T)" for T in MPSGraphs.MPSGRAPH_VALID_SCAN_TYPES
+    A, ops = if T <: Integer
+        rand(T.(1:3), 3, 4), (+, *, min, max)
+    else
+        reshape(T[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37], 3, 4), (+, *)
+    end
 
-    for op in (+, *), dim in 1:2
+    for op in ops, dim in 1:2
         out = similar(MtlArray(A))
         MPSGraphs.graph_scan!(op, out, MtlArray(A); dim)
         @test Array(out) ≈ accumulate(op, A; dims=dim)
     end
 end
 
 @testset "scan unsupported input" begin
-    A = MtlArray(Int32[1, 2])
+    # unsupported input
+    A = MtlArray(Complex{Int32}[1, 2])
     @test_throws ArgumentError MPSGraphs.graph_scan!(+, similar(A), A)
 
-    B = MtlArray(Float32[1, 2])
-    @test_throws ArgumentError MPSGraphs.graph_scan!(max, similar(B), B)
-    @test_throws ArgumentError MPSGraphs.graph_scan!(min, similar(B), B)
-
+    # offset input
     parent = MtlArray(Float32[1, 2, 3])
     offset_input = unsafe_wrap(MtlArray, pointer(parent, 2), 2)
     @test_throws ArgumentError MPSGraphs.graph_scan!(+, similar(offset_input),

test/mpsgraphs/sort.jl

@@ -1,5 +1,5 @@
-@testset "sort ($T)" for T in (Float16, Float32, Int32)
-    A = reshape(T[7, 2, 5, 4, 9, 1, 6, 3, 8, 0, 10, 11], 3, 4)
+@testset "sort ($T)" for T in MPSGraphs.MPSGRAPH_VALID_SORT_TYPES
+    A = rand(T, 3, 4)
 
     for dim in 1:2
         out = similar(MtlArray(A))
@@ -11,7 +11,7 @@
     end
 end
 
-@testset "sort NaN ordering ($T)" for T in (Float16, Float32)
+@testset "sort NaN ordering ($T)" for T in filter(T -> T <: AbstractFloat, MPSGraphs.MPSGRAPH_VALID_SORT_TYPES)
     A = T[1 NaN 2; -1 0 NaN]
 
     for dim in 1:2
@@ -27,17 +27,19 @@ end
 end
 
 @testset "sort unsupported input" begin
-    A = MtlArray(Int16[2, 1])
+    # unsupported input types
+    A = MtlArray(Complex{Int16}[2, 1])
     out = similar(A)
     @test_throws ArgumentError MPSGraphs.graph_sort!(out, A)
 
+    # offset input
     parent = MtlArray(Float32[3, 2, 1])
     offset_input = unsafe_wrap(MtlArray, pointer(parent, 2), 2)
     @test_throws ArgumentError MPSGraphs.graph_sort!(similar(offset_input), offset_input)
 end
 
-@testset "sortperm ($T)" for T in (Float16, Float32, Int32)
-    A = reshape(T[7, 2, 5, 4, 9, 1, 6, 3, 8, 0, 10, 11], 3, 4)
+@testset "sortperm ($T)" for T in MPSGraphs.MPSGRAPH_VALID_SORT_TYPES
+    A = rand(T, 3, 4)
 
     for dim in 1:2
         index = similar(MtlArray(A), Int)