From 0a436472630ba6417a8c4ddf92236123a75849c1 Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 00:49:42 -0400 Subject: [PATCH 1/6] finish svd --- 0 | 0 7:26 | 0 Show | 0 ext/CUDAExt/cuda.jl | 5 -- src/cuNumeric.jl | 2 +- src/ndarray/detail/ndarray.jl | 6 ++ src/ndarray/linalg.jl | 75 ++++++++++++++++++++ test/tests/linalg.jl | 127 ++++++++++++++++++++++++++++++++++ using | 0 9 files changed, 209 insertions(+), 6 deletions(-) create mode 100644 0 create mode 100644 7:26 create mode 100644 Show create mode 100644 using diff --git a/0 b/0 new file mode 100644 index 00000000..e69de29b diff --git a/7:26 b/7:26 new file mode 100644 index 00000000..e69de29b diff --git a/Show b/Show new file mode 100644 index 00000000..e69de29b diff --git a/ext/CUDAExt/cuda.jl b/ext/CUDAExt/cuda.jl index 33421640..c8fd5c01 100644 --- a/ext/CUDAExt/cuda.jl +++ b/ext/CUDAExt/cuda.jl @@ -74,11 +74,6 @@ function check_sz(arr, maxshape) end end -function nda_to_logical_array(arr::NDArray{T,N}) where {T,N} - st_handle = cuNumeric.get_store(arr) - return Legate.LogicalArray{T,N}(st_handle, size(arr)) -end - function Launch(kernel::cuNumeric.CUDATask, inputs::Tuple{Vararg{NDArray}}, outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; blocks, threads) diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 0da59605..2449f2cd 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -64,7 +64,7 @@ const SUPPORTED_NUMERIC_TYPES = Union{ # solve has no integer backend kernel const SUPPORTED_SOLVE_TYPES = Union{SUPPORTED_FLOAT_TYPES,SUPPORTED_COMPLEX_TYPES} - +const SUPPORTED_SVD_TYPES = Union{SUPPORTED_FLOAT_TYPES,SUPPORTED_COMPLEX_TYPES} const SUPPORTED_ARRAY_TYPES = Union{Bool,SUPPORTED_NUMERIC_TYPES} const SUPPORTED_TYPES = Union{SUPPORTED_ARRAY_TYPES,String} diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl index 9d630fcc..4102327b 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -511,3 +511,9 @@ function nda_to_logical_store(arr::NDArray{T,N}) where {T,N} st_handle = Legate.data(Legate.LogicalArray{T,N}(la_handle, size(arr))) return Legate.LogicalStore{T,N}(st_handle, size(arr)) end + + +function nda_to_logical_array(arr::NDArray{T,N}) where {T,N} + st_handle = cuNumeric.get_store(arr) + return Legate.LogicalArray{T,N}(st_handle, size(arr)) +end diff --git a/src/ndarray/linalg.jl b/src/ndarray/linalg.jl index 1c01b44a..5d905df2 100644 --- a/src/ndarray/linalg.jl +++ b/src/ndarray/linalg.jl @@ -117,3 +117,78 @@ end function _solve(a::NDArray{T,N}, b::NDArray{S,M}) where {T,N,S,M} throw(ArgumentError("Batched matrices require signature (...,m,m),(...,m,n)->(...,m,n)")) end + +# ── svd ─────────────────────────────────────────────────────────────────────── + +function svd_single(a::NDArray{T,N}, u::NDArray, s::NDArray, vh::NDArray) where {T,N} + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + task = Legate.create_auto_task(rt, lib, cuNumeric.SVD) + + l_a = nda_to_logical_array(a) + l_u = nda_to_logical_array(u) + l_s = nda_to_logical_array(s) + l_vh = nda_to_logical_array(vh) + + Legate.add_input(task, l_a) + Legate.add_output(task, l_u) + Legate.add_output(task, l_s) + Legate.add_output(task, l_vh) + + Legate.add_broadcast(task, l_a) + Legate.add_broadcast(task, l_u) + Legate.add_broadcast(task, l_s) + Legate.add_broadcast(task, l_vh) + + Legate.submit_auto_task(rt, task) +end + +function _svd(a::NDArray{T,2}, full_matrices::Bool) where {T} + m, n = size(a) + k = min(m, n) + S = real(T) + # cuSolver requires full square buffers regardless of full_matrices + u_buf = zeros(T, m, m) + s = zeros(S, k) + vh_buf = zeros(T, n, n) + svd_single(a, u_buf, s, vh_buf) + # materialize transposed copies to fix cuSolver column-major layout + u_full = copy(cuNumeric.transpose(u_buf)) + vh_full = copy(cuNumeric.transpose(vh_buf)) + u = full_matrices ? u_full : u_full[:, 1:k] + vh = full_matrices ? vh_full : vh_full[1:k, :] + return u, s, vh +end + +# svd runs on float/complex only — no integer backend +const _SVD_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} +const _SVD_ACCEPTED = Union{SUPPORTED_SVD_TYPES,_SVD_PROMOTABLE} +_svd_eltype(::Type{T}) where {T<:_SVD_PROMOTABLE} = Float64 +_svd_eltype(::Type{T}) where {T<:SUPPORTED_SVD_TYPES} = T + +function svd(a::NDArray{<:_SVD_ACCEPTED}, full_matrices::Bool=true) + A = eltype(a) + O = _svd_eltype(A) + A <: _SVD_PROMOTABLE && assertpromotion(svd, A, O) + return _svd_check_dims(unchecked_promote_arr(a, O), full_matrices) +end + +function svd(a::NDArray, full_matrices::Bool=true) + throw(ArgumentError("array type $(eltype(a)) is unsupported in svd")) +end + +function _svd_check_dims(a::NDArray{<:Any,0}, full_matrices::Bool) + throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) +end + +function _svd_check_dims(a::NDArray{<:Any,1}, full_matrices::Bool) + throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) +end + +function _svd_check_dims(a::NDArray{<:Any,2}, full_matrices::Bool) + return _svd(a, full_matrices) +end + +function _svd_check_dims(a::NDArray, full_matrices::Bool) + throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) +end \ No newline at end of file diff --git a/test/tests/linalg.jl b/test/tests/linalg.jl index 8d5ededb..f8a643bb 100644 --- a/test/tests/linalg.jl +++ b/test/tests/linalg.jl @@ -185,3 +185,130 @@ end end end end + + +function check_svd_reconstruction(ref_A::AbstractMatrix, u, s, vh, tol_a, tol_r) + U = Array(u) + S = Array(s) + Vh = Array(vh) + A_rec = U * Diagonal(S) * Vh + return isapprox(ref_A, A_rec; atol=tol_a, rtol=tol_r) +end + +function check_svd_orthonormality(u, vh, tol_a, tol_r) + U = Array(u) + Vh = Array(vh) + ku = size(U, 2) + kv = size(Vh, 1) + ok_u = isapprox(U' * U, Matrix{eltype(U)}(I, ku, ku); atol=tol_a, rtol=tol_r) + ok_vh = isapprox(Vh * Vh', Matrix{eltype(Vh)}(I, kv, kv); atol=tol_a, rtol=tol_r) + return ok_u && ok_vh +end + +@testset "svd square matrix" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + A_ref = my_rand(T, 5, 5) + nda = cuNumeric.NDArray(A_ref) + u, s, vh = cuNumeric.svd(nda) + allowscalar() do + @test check_svd_reconstruction(A_ref, u, s, vh, atol(T), rtol(T)) + @test check_svd_orthonormality(u, vh, atol(T), rtol(T)) + end + end +end + +@testset "svd tall matrix (m > n)" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + A_ref = my_rand(T, 6, 4) + nda = cuNumeric.NDArray(A_ref) + u, s, vh = cuNumeric.svd(nda, false) # thin SVD for reconstruction test + allowscalar() do + @test check_svd_reconstruction(A_ref, u, s, vh, atol(T), rtol(T)) + @test check_svd_orthonormality(u, vh, atol(T), rtol(T)) + end + end +end + +@testset "svd thin output shapes (full_matrices=false)" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + m, n = 6, 4 + k = min(m, n) + A_ref = my_rand(T, m, n) + nda = cuNumeric.NDArray(A_ref) + u, s, vh = cuNumeric.svd(nda, false) + allowscalar() do + @test size(Array(u)) == (m, k) + @test size(Array(s)) == (k,) + @test size(Array(vh)) == (k, n) + @test check_svd_reconstruction(A_ref, u, s, vh, atol(T), rtol(T)) + end + end +end + +@testset "svd full output shapes (full_matrices=true)" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + m, n = 6, 4 + A_ref = my_rand(T, m, n) + nda = cuNumeric.NDArray(A_ref) + u, s, vh = cuNumeric.svd(nda, true) + allowscalar() do + @test size(Array(u)) == (m, m) + @test size(Array(s)) == (min(m, n),) + @test size(Array(vh)) == (n, n) + end + end +end + +@testset "svd singular values non-negative and sorted" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + A_ref = my_rand(T, 5, 5) + nda = cuNumeric.NDArray(A_ref) + _, s, _ = cuNumeric.svd(nda) + allowscalar() do + sv = Array(s) + @test all(sv .>= 0) + @test issorted(sv; rev=true) + end + end +end + +@testset "svd identity matrix" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + n = 4 + A_ref = Matrix{T}(I, n, n) + nda = cuNumeric.NDArray(A_ref) + _, s, _ = cuNumeric.svd(nda) + allowscalar() do + @test cuNumeric.compare(ones(T, n), s, atol(T), rtol(T)) + end + end +end + +@testset "svd rank-1 matrix" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_SVD_TYPES) + # outer product of two vectors: exactly one nonzero singular value + # 5x4 satisfies the M >= N constraint + v1 = T.(collect(1:5)) + v2 = T.(collect(1:4)) + A_ref = v1 * v2' + nda = cuNumeric.NDArray(A_ref) + _, s, _ = cuNumeric.svd(nda) + allowscalar() do + sv = Array(s) + @test sv[1] > atol(T) + @test all(sv[2:end] .< sqrt(atol(T)) * 100) + end + end +end + +@testset "svd promotion" begin + @testset verbose=true for T in (Int32, Int64, Bool) + vals = T == Bool ? T[1 0; 0 1] : reshape(T.(collect(1:4)), 2, 2) + A = cuNumeric.NDArray(vals) + @test_throws "Implicit promotion" cuNumeric.svd(A) + allowpromotion() do + u, s, vh = cuNumeric.svd(A) + @test eltype(Array(u)) == Float64 + end + end +end \ No newline at end of file diff --git a/using b/using new file mode 100644 index 00000000..e69de29b From 0c381c2909a4f5e564e555d5096317ba27bbff6b Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 13:40:04 -0400 Subject: [PATCH 2/6] forgot add the type lol --- lib/cunumeric_jl_wrapper/src/types.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/cunumeric_jl_wrapper/src/types.cpp b/lib/cunumeric_jl_wrapper/src/types.cpp index f181dc2e..998a6d75 100644 --- a/lib/cunumeric_jl_wrapper/src/types.cpp +++ b/lib/cunumeric_jl_wrapper/src/types.cpp @@ -166,4 +166,5 @@ void wrap_binary_ops(jlcxx::Module& mod) { void wrap_linalg_ops(jlcxx::Module& mod) { mod.set_const("SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); mod.set_const("MP_SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); + mod.set_const("SVD", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SVD}); } \ No newline at end of file From 4d7f9b860ce4e58da0eff5860dbf11381ccc224f Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 13:43:40 -0400 Subject: [PATCH 3/6] remove empty files --- 0 | 0 7:26 | 0 Show | 0 using | 0 4 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 0 delete mode 100644 7:26 delete mode 100644 Show delete mode 100644 using diff --git a/0 b/0 deleted file mode 100644 index e69de29b..00000000 diff --git a/7:26 b/7:26 deleted file mode 100644 index e69de29b..00000000 diff --git a/Show b/Show deleted file mode 100644 index e69de29b..00000000 diff --git a/using b/using deleted file mode 100644 index e69de29b..00000000 From 022b774fce76b13b6d0b1afb7134d090d6114068 Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 13:48:36 -0400 Subject: [PATCH 4/6] add types for other linalg ops --- lib/cunumeric_jl_wrapper/src/types.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/lib/cunumeric_jl_wrapper/src/types.cpp b/lib/cunumeric_jl_wrapper/src/types.cpp index 998a6d75..2d197d5f 100644 --- a/lib/cunumeric_jl_wrapper/src/types.cpp +++ b/lib/cunumeric_jl_wrapper/src/types.cpp @@ -167,4 +167,7 @@ void wrap_linalg_ops(jlcxx::Module& mod) { mod.set_const("SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); mod.set_const("MP_SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); mod.set_const("SVD", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SVD}); + mod.set_const("QR", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_QR}); + mod.set_const("SYEV", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SYEV}); + mod.set_const("GEEV", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_GEEV}); } \ No newline at end of file From 402a4fd906ac2035721d938f04c448bd22deaea5 Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 14:18:33 -0400 Subject: [PATCH 5/6] implement QR single gpu --- lib/cunumeric_jl_wrapper/src/types.cpp | 2 +- src/cuNumeric.jl | 1 + src/ndarray/linalg.jl | 69 +++++++++++++++++++++++++- test/tests/linalg.jl | 46 +++++++++++++++++ 4 files changed, 115 insertions(+), 3 deletions(-) diff --git a/lib/cunumeric_jl_wrapper/src/types.cpp b/lib/cunumeric_jl_wrapper/src/types.cpp index 2d197d5f..22a9b928 100644 --- a/lib/cunumeric_jl_wrapper/src/types.cpp +++ b/lib/cunumeric_jl_wrapper/src/types.cpp @@ -167,7 +167,7 @@ void wrap_linalg_ops(jlcxx::Module& mod) { mod.set_const("SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); mod.set_const("MP_SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); mod.set_const("SVD", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SVD}); - mod.set_const("QR", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_QR}); + mod.set_const("CQR", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_QR}); mod.set_const("SYEV", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SYEV}); mod.set_const("GEEV", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_GEEV}); } \ No newline at end of file diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 2449f2cd..43f07e5d 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -65,6 +65,7 @@ const SUPPORTED_NUMERIC_TYPES = Union{ # solve has no integer backend kernel const SUPPORTED_SOLVE_TYPES = Union{SUPPORTED_FLOAT_TYPES,SUPPORTED_COMPLEX_TYPES} const SUPPORTED_SVD_TYPES = Union{SUPPORTED_FLOAT_TYPES,SUPPORTED_COMPLEX_TYPES} +const SUPPORTED_QR_TYPES = Union{SUPPORTED_FLOAT_TYPES,SUPPORTED_COMPLEX_TYPES} const SUPPORTED_ARRAY_TYPES = Union{Bool,SUPPORTED_NUMERIC_TYPES} const SUPPORTED_TYPES = Union{SUPPORTED_ARRAY_TYPES,String} diff --git a/src/ndarray/linalg.jl b/src/ndarray/linalg.jl index 5d905df2..63e2896b 100644 --- a/src/ndarray/linalg.jl +++ b/src/ndarray/linalg.jl @@ -118,8 +118,6 @@ function _solve(a::NDArray{T,N}, b::NDArray{S,M}) where {T,N,S,M} throw(ArgumentError("Batched matrices require signature (...,m,m),(...,m,n)->(...,m,n)")) end -# ── svd ─────────────────────────────────────────────────────────────────────── - function svd_single(a::NDArray{T,N}, u::NDArray, s::NDArray, vh::NDArray) where {T,N} rt = Legate.get_runtime() lib = cuNumeric.get_lib() @@ -191,4 +189,71 @@ end function _svd_check_dims(a::NDArray, full_matrices::Bool) throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) +end + +# qr + +function qr_single(a::NDArray{T,N}, q::NDArray, r::NDArray) where {T,N} + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + task = Legate.create_auto_task(rt, lib, cuNumeric.CQR) + + l_a = nda_to_logical_array(a) + l_q = nda_to_logical_array(q) + l_r = nda_to_logical_array(r) + + Legate.add_input(task, l_a) + Legate.add_output(task, l_q) + Legate.add_output(task, l_r) + + Legate.add_broadcast(task, l_a) + Legate.add_broadcast(task, l_q) + Legate.add_broadcast(task, l_r) + + Legate.submit_auto_task(rt, task) +end + +function _qr(a::NDArray{T,2}) where {T} + m, n = size(a) + k = min(m, n) + # cuSolver requires full square buffers regardless of output shape + q_buf = zeros(T, m, m) + r_buf = zeros(T, n, n) + qr_single(a, q_buf, r_buf) + # materialize transposed copies to fix cuSolver column-major layout + q = copy(cuNumeric.transpose(q_buf))[:, 1:k] + r = copy(cuNumeric.transpose(r_buf))[1:k, :] + return q, r +end + +const _QR_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} +const _QR_ACCEPTED = Union{SUPPORTED_QR_TYPES,_QR_PROMOTABLE} +_qr_eltype(::Type{T}) where {T<:_QR_PROMOTABLE} = Float64 +_qr_eltype(::Type{T}) where {T<:SUPPORTED_QR_TYPES} = T + +function qr(a::NDArray{<:_QR_ACCEPTED}) + A = eltype(a) + O = _qr_eltype(A) + A <: _QR_PROMOTABLE && assertpromotion(qr, A, O) + return _qr_check_dims(unchecked_promote_arr(a, O)) +end + +function qr(a::NDArray) + throw(ArgumentError("array type $(eltype(a)) is unsupported in qr")) +end + +function _qr_check_dims(a::NDArray{<:Any,0}) + throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) +end + +function _qr_check_dims(a::NDArray{<:Any,1}) + throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) +end + +function _qr_check_dims(a::NDArray{<:Any,2}) + return _qr(a) +end + +function _qr_check_dims(a::NDArray) + throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) end \ No newline at end of file diff --git a/test/tests/linalg.jl b/test/tests/linalg.jl index f8a643bb..10703f7c 100644 --- a/test/tests/linalg.jl +++ b/test/tests/linalg.jl @@ -311,4 +311,50 @@ end @test eltype(Array(u)) == Float64 end end +end + +@testset "qr reconstruction" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_QR_TYPES) + A_ref = my_rand(T, 6, 4) + nda = cuNumeric.NDArray(A_ref) + q, r = cuNumeric.qr(nda) + allowscalar() do + Q = Array(q) + R = Array(r) + @test size(Q) == (6, 4) + @test size(R) == (4, 4) + @test isapprox(A_ref, Q * R; atol=atol(T), rtol=rtol(T)) + @test isapprox(Q' * Q, Matrix{eltype(Q)}(I, 4, 4); atol=atol(T), rtol=rtol(T)) + end + end +end + +@testset "qr square matrix" begin + @testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_QR_TYPES) + A_ref = my_rand(T, 5, 5) + nda = cuNumeric.NDArray(A_ref) + q, r = cuNumeric.qr(nda) + allowscalar() do + Q = Array(q) + R = Array(r) + @test size(Q) == (5, 5) + @test size(R) == (5, 5) + @test isapprox(A_ref, Q * R; atol=atol(T), rtol=rtol(T)) + end + end +end + +@testset "qr promotion" begin + @testset verbose=true for T in (Int32, Int64, Bool) + vals = T == Bool ? T[1 0; 0 1] : reshape(T.(collect(1:4)), 2, 2) + A = cuNumeric.NDArray(vals) + @test_throws "Implicit promotion" cuNumeric.qr(A) + allowpromotion() do + q, r = cuNumeric.qr(A) + allowscalar() do + @test eltype(Array(q)) == Float64 + @test isapprox(Float64.(vals), Array(q) * Array(r); atol=atol(Float64), rtol=rtol(Float64)) + end + end + end end \ No newline at end of file From 05827519cbc45e7949fadfcc14ec825f17c3cd43 Mon Sep 17 00:00:00 2001 From: Nader-Rahhal Date: Thu, 4 Jun 2026 14:34:27 -0400 Subject: [PATCH 6/6] reorganize --- src/cuNumeric.jl | 1 + src/ndarray/detail/linalg.jl | 222 ++++++++++++++++++++++++++++++++++ src/ndarray/linalg.jl | 223 ----------------------------------- 3 files changed, 223 insertions(+), 223 deletions(-) create mode 100644 src/ndarray/detail/linalg.jl diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 43f07e5d..239b97ce 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -142,6 +142,7 @@ include("warnings.jl") # NDArray internal include("ndarray/detail/ndarray.jl") +include("ndarray/detail/linalg.jl") # NDArray interface include("ndarray/promotion.jl") diff --git a/src/ndarray/detail/linalg.jl b/src/ndarray/detail/linalg.jl new file mode 100644 index 00000000..d9cff68e --- /dev/null +++ b/src/ndarray/detail/linalg.jl @@ -0,0 +1,222 @@ +function choose_nd_color_shape(shape::NTuple{N,Int}) where {N} + color_shape = Base.ones(Int, N) + if N > 2 + color_shape[1] = Legate.num_procs() + done = false + while !done && color_shape[1] % 2 == 0 + weight_per_dim = [shape[i] / color_shape[i] for i in 1:(N - 2)] + max_weight, idx = findmax(weight_per_dim) + if weight_per_dim[idx] > 2 * weight_per_dim[1] + color_shape[1] ÷= 2 + color_shape[idx] *= 2 + else + done = true + end + end + end + return Tuple(color_shape) +end + +function prepare_manual_task_for_batched_matrices(full_shape::NTuple{N,Int}) where {N} + initial_color_shape = choose_nd_color_shape(full_shape) + tilesize = Tuple( + (full_shape[i] + initial_color_shape[i] - 1) ÷ initial_color_shape[i] for i in 1:N + ) + color_shape = Tuple((full_shape[i] + tilesize[i] - 1) ÷ tilesize[i] for i in 1:N) + return tilesize, color_shape +end + +function solve_batched(a::NDArray{T,N}, b::NDArray, x::NDArray) where {T,N} + nrhs = size(b)[end] + full_shape = size(a) + tilesize_a, color_shape = prepare_manual_task_for_batched_matrices(full_shape) + tilesize_b = (tilesize_a[1:(end - 1)]..., nrhs) + + store_a = nda_to_logical_store(a) + store_b = nda_to_logical_store(b) + store_x = nda_to_logical_store(x) + + tiled_a = Legate.partition_by_tiling(store_a, collect(tilesize_a)) + tiled_b = Legate.partition_by_tiling(store_b, collect(tilesize_b)) + tiled_x = Legate.partition_by_tiling(store_x, collect(tilesize_b)) + + rt = Legate.get_runtime() + domain = Legate.domain_from_shape(Legate.Shape(Legate.to_cxx_vector(color_shape))) + lib = cuNumeric.get_lib() + task = Legate.create_manual_task(rt, lib, cuNumeric.SOLVE, domain) + + Legate.add_input(task, tiled_a) + Legate.add_input(task, tiled_b) + Legate.add_output(task, tiled_x) + + Legate.submit_manual_task(rt, task) +end + +# solve runs in floating point: +# int/bool inputs promote to Float64 (matching cupynumeric) +const _SOLVE_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} +const _SOLVE_ACCEPTED = Union{SUPPORTED_SOLVE_TYPES,_SOLVE_PROMOTABLE} +_solve_eltype(::Type{T}) where {T<:_SOLVE_PROMOTABLE} = Float64 +_solve_eltype(::Type{T}) where {T<:SUPPORTED_SOLVE_TYPES} = T + +# `a` must be at least 2D, `b` at least 1D. +function _solve_check_a_dims(a::NDArray{<:Any,0}, b::NDArray) + throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) +end +function _solve_check_a_dims(a::NDArray{<:Any,1}, b::NDArray) + throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) +end +_solve_check_a_dims(a::NDArray, b::NDArray) = _solve_check_b_dims(a, b) + +function _solve_check_b_dims(a::NDArray, b::NDArray{<:Any,0}) + throw(ArgumentError("0-dimensional array given. Array must be at least one-dimensional")) +end +_solve_check_b_dims(a::NDArray, b::NDArray) = _solve(a, b) + +# 2D case: (m,m),(m)->(m). +# Backend needs rhs "b" to be 2D. We reshape b from (n,) to (n,1) +function _solve(a::NDArray{T,2}, b::NDArray{S,1}) where {T,S} + m = size(b)[1] + return reshape(_solve(a, reshape(b, (m, 1))), (m,)) +end + +# 2D (m,m),(m,n)->(m,n) and batched (...,m,m),(...,m,n)->(...,m,n) +function _solve(a::NDArray{T,N}, b::NDArray{S,N}) where {T,S,N} + size(a)[end - 1] != size(a)[end] && + throw(ArgumentError("Last 2 dimensions of the array must be square")) + size(a)[end] != size(b)[end - 1] && + throw( + ArgumentError( + "Input operand 1 has a mismatch in its dimension " * + "$(N-2), with signature (...,m,m),(...,m,n)->(...,m,n)" * + " (size $(size(b)[end-1]) is different from $(size(a)[end]))", + ), + ) + prod(size(a)) == 0 || prod(size(b)) == 0 && return zeros(T, size(b)...) + x = zeros(T, size(b)...) + solve_batched(a, b, x) + return x +end + +# Mismatched batch dimensions +function _solve(a::NDArray{T,N}, b::NDArray{S,M}) where {T,N,S,M} + throw(ArgumentError("Batched matrices require signature (...,m,m),(...,m,n)->(...,m,n)")) +end + +function svd_single(a::NDArray{T,N}, u::NDArray, s::NDArray, vh::NDArray) where {T,N} + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + task = Legate.create_auto_task(rt, lib, cuNumeric.SVD) + + l_a = nda_to_logical_array(a) + l_u = nda_to_logical_array(u) + l_s = nda_to_logical_array(s) + l_vh = nda_to_logical_array(vh) + + Legate.add_input(task, l_a) + Legate.add_output(task, l_u) + Legate.add_output(task, l_s) + Legate.add_output(task, l_vh) + + Legate.add_broadcast(task, l_a) + Legate.add_broadcast(task, l_u) + Legate.add_broadcast(task, l_s) + Legate.add_broadcast(task, l_vh) + + Legate.submit_auto_task(rt, task) +end + +function _svd(a::NDArray{T,2}, full_matrices::Bool) where {T} + m, n = size(a) + k = min(m, n) + S = real(T) + # cuSolver requires full square buffers regardless of full_matrices + u_buf = zeros(T, m, m) + s = zeros(S, k) + vh_buf = zeros(T, n, n) + svd_single(a, u_buf, s, vh_buf) + # materialize transposed copies to fix cuSolver column-major layout + u_full = copy(cuNumeric.transpose(u_buf)) + vh_full = copy(cuNumeric.transpose(vh_buf)) + u = full_matrices ? u_full : u_full[:, 1:k] + vh = full_matrices ? vh_full : vh_full[1:k, :] + return u, s, vh +end + +# svd runs on float/complex only — no integer backend +const _SVD_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} +const _SVD_ACCEPTED = Union{SUPPORTED_SVD_TYPES,_SVD_PROMOTABLE} +_svd_eltype(::Type{T}) where {T<:_SVD_PROMOTABLE} = Float64 +_svd_eltype(::Type{T}) where {T<:SUPPORTED_SVD_TYPES} = T + +function _svd_check_dims(a::NDArray{<:Any,0}, full_matrices::Bool) + throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) +end + +function _svd_check_dims(a::NDArray{<:Any,1}, full_matrices::Bool) + throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) +end + +function _svd_check_dims(a::NDArray{<:Any,2}, full_matrices::Bool) + return _svd(a, full_matrices) +end + +function _svd_check_dims(a::NDArray, full_matrices::Bool) + throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) +end + +# qr + +function qr_single(a::NDArray{T,N}, q::NDArray, r::NDArray) where {T,N} + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + task = Legate.create_auto_task(rt, lib, cuNumeric.CQR) + + l_a = nda_to_logical_array(a) + l_q = nda_to_logical_array(q) + l_r = nda_to_logical_array(r) + + Legate.add_input(task, l_a) + Legate.add_output(task, l_q) + Legate.add_output(task, l_r) + + Legate.add_broadcast(task, l_a) + Legate.add_broadcast(task, l_q) + Legate.add_broadcast(task, l_r) + + Legate.submit_auto_task(rt, task) +end + +function _qr(a::NDArray{T,2}) where {T} + m, n = size(a) + k = min(m, n) + # cuSolver requires full square buffers regardless of output shape + q_buf = zeros(T, m, m) + r_buf = zeros(T, n, n) + qr_single(a, q_buf, r_buf) + # materialize transposed copies to fix cuSolver column-major layout + q = copy(cuNumeric.transpose(q_buf))[:, 1:k] + r = copy(cuNumeric.transpose(r_buf))[1:k, :] + return q, r +end + +const _QR_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} +const _QR_ACCEPTED = Union{SUPPORTED_QR_TYPES,_QR_PROMOTABLE} +_qr_eltype(::Type{T}) where {T<:_QR_PROMOTABLE} = Float64 +_qr_eltype(::Type{T}) where {T<:SUPPORTED_QR_TYPES} = T + +function _qr_check_dims(a::NDArray{<:Any,0}) + throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) +end + +function _qr_check_dims(a::NDArray{<:Any,1}) + throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) +end + +function _qr_check_dims(a::NDArray{<:Any,2}) + return _qr(a) +end + +function _qr_check_dims(a::NDArray) + throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) +end diff --git a/src/ndarray/linalg.jl b/src/ndarray/linalg.jl index 63e2896b..016466b8 100644 --- a/src/ndarray/linalg.jl +++ b/src/ndarray/linalg.jl @@ -1,64 +1,3 @@ -function choose_nd_color_shape(shape::NTuple{N,Int}) where {N} - color_shape = Base.ones(Int, N) - if N > 2 - color_shape[1] = Legate.num_procs() - done = false - while !done && color_shape[1] % 2 == 0 - weight_per_dim = [shape[i] / color_shape[i] for i in 1:(N - 2)] - max_weight, idx = findmax(weight_per_dim) - if weight_per_dim[idx] > 2 * weight_per_dim[1] - color_shape[1] ÷= 2 - color_shape[idx] *= 2 - else - done = true - end - end - end - return Tuple(color_shape) -end - -function prepare_manual_task_for_batched_matrices(full_shape::NTuple{N,Int}) where {N} - initial_color_shape = choose_nd_color_shape(full_shape) - tilesize = Tuple( - (full_shape[i] + initial_color_shape[i] - 1) ÷ initial_color_shape[i] for i in 1:N - ) - color_shape = Tuple((full_shape[i] + tilesize[i] - 1) ÷ tilesize[i] for i in 1:N) - return tilesize, color_shape -end - -function solve_batched(a::NDArray{T,N}, b::NDArray, x::NDArray) where {T,N} - nrhs = size(b)[end] - full_shape = size(a) - tilesize_a, color_shape = prepare_manual_task_for_batched_matrices(full_shape) - tilesize_b = (tilesize_a[1:(end - 1)]..., nrhs) - - store_a = nda_to_logical_store(a) - store_b = nda_to_logical_store(b) - store_x = nda_to_logical_store(x) - - tiled_a = Legate.partition_by_tiling(store_a, collect(tilesize_a)) - tiled_b = Legate.partition_by_tiling(store_b, collect(tilesize_b)) - tiled_x = Legate.partition_by_tiling(store_x, collect(tilesize_b)) - - rt = Legate.get_runtime() - domain = Legate.domain_from_shape(Legate.Shape(Legate.to_cxx_vector(color_shape))) - lib = cuNumeric.get_lib() - task = Legate.create_manual_task(rt, lib, cuNumeric.SOLVE, domain) - - Legate.add_input(task, tiled_a) - Legate.add_input(task, tiled_b) - Legate.add_output(task, tiled_x) - - Legate.submit_manual_task(rt, task) -end - -# solve runs in floating point: -# int/bool inputs promote to Float64 (matching cupynumeric) -const _SOLVE_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} -const _SOLVE_ACCEPTED = Union{SUPPORTED_SOLVE_TYPES,_SOLVE_PROMOTABLE} -_solve_eltype(::Type{T}) where {T<:_SOLVE_PROMOTABLE} = Float64 -_solve_eltype(::Type{T}) where {T<:SUPPORTED_SOLVE_TYPES} = T - # Type/dim guards dispatch on one argument at a time, then forward to `_solve`. function solve(a::NDArray{<:_SOLVE_ACCEPTED}, b::NDArray{<:_SOLVE_ACCEPTED}) A, B = eltype(a), eltype(b) @@ -74,96 +13,6 @@ function solve(a::NDArray, b::NDArray) throw(ArgumentError("array type $bad is unsupported in solve")) end -# `a` must be at least 2D, `b` at least 1D. -function _solve_check_a_dims(a::NDArray{<:Any,0}, b::NDArray) - throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) -end -function _solve_check_a_dims(a::NDArray{<:Any,1}, b::NDArray) - throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) -end -_solve_check_a_dims(a::NDArray, b::NDArray) = _solve_check_b_dims(a, b) - -function _solve_check_b_dims(a::NDArray, b::NDArray{<:Any,0}) - throw(ArgumentError("0-dimensional array given. Array must be at least one-dimensional")) -end -_solve_check_b_dims(a::NDArray, b::NDArray) = _solve(a, b) - -# 2D case: (m,m),(m)->(m). -# Backend needs rhs "b" to be 2D. We reshape b from (n,) to (n,1) -function _solve(a::NDArray{T,2}, b::NDArray{S,1}) where {T,S} - m = size(b)[1] - return reshape(_solve(a, reshape(b, (m, 1))), (m,)) -end - -# 2D (m,m),(m,n)->(m,n) and batched (...,m,m),(...,m,n)->(...,m,n) -function _solve(a::NDArray{T,N}, b::NDArray{S,N}) where {T,S,N} - size(a)[end - 1] != size(a)[end] && - throw(ArgumentError("Last 2 dimensions of the array must be square")) - size(a)[end] != size(b)[end - 1] && - throw( - ArgumentError( - "Input operand 1 has a mismatch in its dimension " * - "$(N-2), with signature (...,m,m),(...,m,n)->(...,m,n)" * - " (size $(size(b)[end-1]) is different from $(size(a)[end]))", - ), - ) - prod(size(a)) == 0 || prod(size(b)) == 0 && return zeros(T, size(b)...) - x = zeros(T, size(b)...) - solve_batched(a, b, x) - return x -end - -# Mismatched batch dimensions -function _solve(a::NDArray{T,N}, b::NDArray{S,M}) where {T,N,S,M} - throw(ArgumentError("Batched matrices require signature (...,m,m),(...,m,n)->(...,m,n)")) -end - -function svd_single(a::NDArray{T,N}, u::NDArray, s::NDArray, vh::NDArray) where {T,N} - rt = Legate.get_runtime() - lib = cuNumeric.get_lib() - task = Legate.create_auto_task(rt, lib, cuNumeric.SVD) - - l_a = nda_to_logical_array(a) - l_u = nda_to_logical_array(u) - l_s = nda_to_logical_array(s) - l_vh = nda_to_logical_array(vh) - - Legate.add_input(task, l_a) - Legate.add_output(task, l_u) - Legate.add_output(task, l_s) - Legate.add_output(task, l_vh) - - Legate.add_broadcast(task, l_a) - Legate.add_broadcast(task, l_u) - Legate.add_broadcast(task, l_s) - Legate.add_broadcast(task, l_vh) - - Legate.submit_auto_task(rt, task) -end - -function _svd(a::NDArray{T,2}, full_matrices::Bool) where {T} - m, n = size(a) - k = min(m, n) - S = real(T) - # cuSolver requires full square buffers regardless of full_matrices - u_buf = zeros(T, m, m) - s = zeros(S, k) - vh_buf = zeros(T, n, n) - svd_single(a, u_buf, s, vh_buf) - # materialize transposed copies to fix cuSolver column-major layout - u_full = copy(cuNumeric.transpose(u_buf)) - vh_full = copy(cuNumeric.transpose(vh_buf)) - u = full_matrices ? u_full : u_full[:, 1:k] - vh = full_matrices ? vh_full : vh_full[1:k, :] - return u, s, vh -end - -# svd runs on float/complex only — no integer backend -const _SVD_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} -const _SVD_ACCEPTED = Union{SUPPORTED_SVD_TYPES,_SVD_PROMOTABLE} -_svd_eltype(::Type{T}) where {T<:_SVD_PROMOTABLE} = Float64 -_svd_eltype(::Type{T}) where {T<:SUPPORTED_SVD_TYPES} = T - function svd(a::NDArray{<:_SVD_ACCEPTED}, full_matrices::Bool=true) A = eltype(a) O = _svd_eltype(A) @@ -175,62 +24,6 @@ function svd(a::NDArray, full_matrices::Bool=true) throw(ArgumentError("array type $(eltype(a)) is unsupported in svd")) end -function _svd_check_dims(a::NDArray{<:Any,0}, full_matrices::Bool) - throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) -end - -function _svd_check_dims(a::NDArray{<:Any,1}, full_matrices::Bool) - throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) -end - -function _svd_check_dims(a::NDArray{<:Any,2}, full_matrices::Bool) - return _svd(a, full_matrices) -end - -function _svd_check_dims(a::NDArray, full_matrices::Bool) - throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) -end - -# qr - -function qr_single(a::NDArray{T,N}, q::NDArray, r::NDArray) where {T,N} - rt = Legate.get_runtime() - lib = cuNumeric.get_lib() - task = Legate.create_auto_task(rt, lib, cuNumeric.CQR) - - l_a = nda_to_logical_array(a) - l_q = nda_to_logical_array(q) - l_r = nda_to_logical_array(r) - - Legate.add_input(task, l_a) - Legate.add_output(task, l_q) - Legate.add_output(task, l_r) - - Legate.add_broadcast(task, l_a) - Legate.add_broadcast(task, l_q) - Legate.add_broadcast(task, l_r) - - Legate.submit_auto_task(rt, task) -end - -function _qr(a::NDArray{T,2}) where {T} - m, n = size(a) - k = min(m, n) - # cuSolver requires full square buffers regardless of output shape - q_buf = zeros(T, m, m) - r_buf = zeros(T, n, n) - qr_single(a, q_buf, r_buf) - # materialize transposed copies to fix cuSolver column-major layout - q = copy(cuNumeric.transpose(q_buf))[:, 1:k] - r = copy(cuNumeric.transpose(r_buf))[1:k, :] - return q, r -end - -const _QR_PROMOTABLE = Union{SUPPORTED_INT_TYPES,Bool} -const _QR_ACCEPTED = Union{SUPPORTED_QR_TYPES,_QR_PROMOTABLE} -_qr_eltype(::Type{T}) where {T<:_QR_PROMOTABLE} = Float64 -_qr_eltype(::Type{T}) where {T<:SUPPORTED_QR_TYPES} = T - function qr(a::NDArray{<:_QR_ACCEPTED}) A = eltype(a) O = _qr_eltype(A) @@ -241,19 +34,3 @@ end function qr(a::NDArray) throw(ArgumentError("array type $(eltype(a)) is unsupported in qr")) end - -function _qr_check_dims(a::NDArray{<:Any,0}) - throw(ArgumentError("0-dimensional array given. Array must be at least two-dimensional")) -end - -function _qr_check_dims(a::NDArray{<:Any,1}) - throw(ArgumentError("1-dimensional array given. Array must be at least two-dimensional")) -end - -function _qr_check_dims(a::NDArray{<:Any,2}) - return _qr(a) -end - -function _qr_check_dims(a::NDArray) - throw(ArgumentError("cuNumeric does not yet support stacked 2d arrays")) -end \ No newline at end of file