diff --git a/.githash b/.githash index b8a30771..8f358435 100644 --- a/.githash +++ b/.githash @@ -1 +1 @@ -dd8623e18ce0b61be51c3a90c8e7927d642a7574 +1fd9472c5bdee6909bb67f928361efdcc94f468e diff --git a/lib/cunumeric_jl_wrapper/src/ndarray.cpp b/lib/cunumeric_jl_wrapper/src/ndarray.cpp index 43dbc69f..fcd35dd7 100644 --- a/lib/cunumeric_jl_wrapper/src/ndarray.cpp +++ b/lib/cunumeric_jl_wrapper/src/ndarray.cpp @@ -98,6 +98,39 @@ void nda_add(CN_NDArray* rhs1, CN_NDArray* rhs2, CN_NDArray* out) { cupynumeric::add(rhs1->obj, rhs2->obj, out->obj); } +// NEW + +CN_NDArray* nda_unique(CN_NDArray* arr) { + NDArray result = cupynumeric::unique(arr->obj); + return new CN_NDArray{NDArray(std::move(result))}; +} + +CN_NDArray* nda_ravel(CN_NDArray* arr) { + NDArray result = cupynumeric::ravel(arr->obj, "C"); + return new CN_NDArray{NDArray(std::move(result))}; +} + +CN_NDArray* nda_trace(CN_NDArray* arr, int32_t offset, int32_t a1, int32_t a2, + CN_Type type) { + NDArray result = cupynumeric::trace(arr->obj, offset, a1, a2, type.obj); + return new CN_NDArray{NDArray(std::move(result))}; +} + +CN_NDArray* nda_eye(int32_t rows, CN_Type type) { + NDArray result = cupynumeric::eye(rows, rows, 0, type.obj); + return new CN_NDArray{NDArray(std::move(result))}; +} + +CN_NDArray* nda_diag(CN_NDArray* arr, int32_t k) { + NDArray result = cupynumeric::diag(arr->obj, k); + return new CN_NDArray{NDArray(std::move(result))}; +} + +CN_NDArray* nda_transpose(CN_NDArray* arr) { + NDArray result = cupynumeric::transpose(arr->obj); + return new CN_NDArray{NDArray(std::move(result))}; +} + CN_NDArray* nda_multiply_scalar(CN_NDArray* rhs1, CN_Type type, const void* value) { Scalar s(type.obj, value, true); diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl index 2c0861c1..5ecc10ea 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -221,11 +221,25 @@ function nda_array_equal(rhs1::NDArray{T,N}, rhs2::NDArray{T,N}) where {T,N} return NDArray(ptr; T=Bool, n_dim=1) end -function nda_multiply(rhs1::NDArray, rhs2::NDArray, out::NDArray) - ccall((:nda_multiply, libnda), - Cvoid, (NDArray_t, NDArray_t, NDArray_t), - rhs1.ptr, rhs2.ptr, out.ptr) - return out +function nda_diag(arr::NDArray, k::Int32) + ptr = ccall((:nda_diag, libnda), + NDArray_t, (NDArray_t, Int32), + arr.ptr, k) + return NDArray(ptr) +end + +function nda_unique(arr::NDArray) + ptr = ccall((:nda_unique, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + return NDArray(ptr) +end + +function nda_ravel(arr::NDArray) + ptr = ccall((:nda_ravel, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + return NDArray(ptr) end function nda_add(rhs1::NDArray, rhs2::NDArray, out::NDArray) @@ -267,6 +281,32 @@ function nda_dot(rhs1::NDArray, rhs2::NDArray) return NDArray(ptr) end +function nda_eye(rows::Int32, ::Type{T}) where {T} + legate_type = Legate.to_legate_type(T) + ptr = ccall((:nda_eye, libnda), + NDArray_t, (Int32, Legate.LegateTypeAllocated), + rows, legate_type) + return NDArray(ptr; T=T, n_dim=2) +end + +function nda_trace( + arr::NDArray, offset::Int32, a1::Int32, a2::Int32, ::Type{T} +) where {T} + legate_type = Legate.to_legate_type(T) + ptr = ccall((:nda_trace, libnda), + NDArray_t, + (NDArray_t, Int32, Int32, Int32, Legate.LegateTypeAllocated), + arr.ptr, offset, a1, a2, legate_type) + return NDArray(ptr; T=T, n_dim=1) +end + +function nda_transpose(arr::NDArray) + ptr = ccall((:nda_transpose, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + return NDArray(ptr) +end + function nda_attach_external(arr::AbstractArray{T,N}) where {T,N} ptr = Base.unsafe_convert(Ptr{Cvoid}, arr) nbytes = sizeof(T) * length(arr) diff --git a/src/ndarray/ndarray.jl b/src/ndarray/ndarray.jl index c2c5e8b4..02ec68e5 100644 --- a/src/ndarray/ndarray.jl +++ b/src/ndarray/ndarray.jl @@ -19,6 +19,60 @@ export unwrap +@doc""" + cuNumeric.transpose(arr::NDArray) + +Return a new `NDArray` that is the transpose of the input `arr`. +""" +function transpose(arr::NDArray) + return nda_transpose(arr) +end + +@doc""" + cuNumeric.eye(rows::Int; T=Float32) + +Create a 2D identity `NDArray` of size `rows x rows` with element type `T`. +""" +function eye(rows::Int; T::Type{S}=Float64) where {S} + return nda_eye(Int32(rows), S) +end + +@doc""" + cuNumeric.trace(arr::NDArray; offset=0, a1=0, a2=1, T=Float32) + +Compute the trace of the `NDArray` along the specified axes. +""" +function trace(arr::NDArray; offset::Int=0, a1::Int=0, a2::Int=1, T::Type{S}=Float32) where {S} + return nda_trace(arr, Int32(offset), Int32(a1), Int32(a2), S) +end + +@doc""" + cuNumeric.diag(arr::NDArray; k=0) + +Extract the k-th diagonal from a 2D `NDArray`. +""" +function diag(arr::NDArray; k::Int=0) + return nda_diag(arr, Int32(k)) +end + +@doc""" + cuNumeric.ravel(arr::NDArray) + +Return a flattened 1D view of the input `NDArray`. +""" +function ravel(arr::NDArray) + return nda_ravel(arr) +end + +@doc""" + cuNumeric.unique(arr::NDArray) + +Return a new `NDArray` containing the unique elements of the input `arr`. +""" +function unique(arr::NDArray) + return nda_unique(arr) +end + @doc""" Base.copy(arr::NDArray) diff --git a/test/runtests.jl b/test/runtests.jl index 2f7e3d1e..629ae5ab 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,7 +23,7 @@ using Random import Random: rand const VERBOSE = get(ENV, "VERBOSE", "1") != "0" -const run_gpu_tests = get(ENV, "GPUTESTS", "1") != "0" +const run_gpu_tests = (get(ENV, "GPUTESTS", "1") != "0") && (get(ENV, "NO_CUDA", "OFF") != "ON") @info "Run GPU Tests: $(run_gpu_tests)" if run_gpu_tests @@ -70,6 +70,10 @@ end @testset elementwise() end +@testset verbose = true "Linear Algebra Tests" begin + include("tests/linalg.jl") +end + @testset verbose = true "GEMM" begin N = 50 M = 25 diff --git a/test/tests/linalg.jl b/test/tests/linalg.jl new file mode 100644 index 00000000..7e3b2097 --- /dev/null +++ b/test/tests/linalg.jl @@ -0,0 +1,104 @@ +#= Copyright 2025 Northwestern University, + * Carnegie Mellon University University + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * Author(s): David Krasowska + * Ethan Meitz + * Nader Rahal +=# + +@testset "transpose" begin + A = rand(Float64, 4, 3) + nda = cuNumeric.NDArray(A) + + ref = transpose(A) + out = cuNumeric.transpose(nda) + + allowscalar() do + @test cuNumeric.compare(ref, out, atol(Float64), rtol(Float64)) + end +end + +@testset "eye" begin + for T in (Float32, Float64, Int32) + n = 5 + ref = Matrix{T}(I, n, n) + out = cuNumeric.eye(n; T=T) + allowscalar() do + @test cuNumeric.compare(ref, out, atol(T), rtol(T)) + end + end +end + +@testset "trace" begin + A = rand(Float64, 6, 6) + nda = cuNumeric.NDArray(A) + + ref = tr(A) + out = cuNumeric.trace(nda) + + allowscalar() do + @test ref ≈ out[1] atol=atol(Float32) rtol=rtol(Float32) + end +end + +@testset "trace with offset" begin + A = rand(Float32, 5, 5) + nda = cuNumeric.NDArray(A) + + for k in (-2, -1, 0, 1, 2) + ref = sum(diag(A, k)) + out = cuNumeric.trace(nda; offset=k) + + allowscalar() do + @test ref ≈ out[1] atol=atol(Float32) rtol=rtol(Float32) + end + end +end + +@testset "diag" begin + A = rand(Int, 6, 6) + nda = cuNumeric.NDArray(A) + + for k in (-2, 0, 3) + ref = diag(A, k) + out = cuNumeric.diag(nda; k=k) + + allowscalar() do + @test cuNumeric.compare(ref, out, atol(Int32), rtol(Int32)) + end + end +end + +# @testset "ravel" begin +# A = reshape(collect(1:12), 3, 4) +# nda = cuNumeric.NDArray(A) + +# ref = vec(A) +# out = cuNumeric.ravel(nda) + +# allowscalar() do +# @test cuNumeric.compare(ref, out, atol(Int32), rtol(Int32)) +# end +# end + +@testset "unique" begin + A = [1, 2, 2, 3, 4, 4, 4, 5] + nda = cuNumeric.NDArray(A) + + ref = unique(A) + out = cuNumeric.unique(nda) + + @test sort(Array(out)) == sort(ref) +end