diff --git a/lib/cunumeric_jl_wrapper/src/ndarray.cpp b/lib/cunumeric_jl_wrapper/src/ndarray.cpp index 152055dd..0ebfee02 100644 --- a/lib/cunumeric_jl_wrapper/src/ndarray.cpp +++ b/lib/cunumeric_jl_wrapper/src/ndarray.cpp @@ -116,8 +116,7 @@ CN_NDArray* nda_zeros_array(int32_t dim, const uint64_t* shape, CN_Type type) { return new CN_NDArray{NDArray(std::move(result))}; } -CN_NDArray* nda_full_array(int32_t dim, const uint64_t* shape, CN_Type type, - const void* value) { +CN_NDArray* nda_full_array(int32_t dim, const uint64_t* shape, CN_Type type, const void* value) { std::vector shp(shape, shape + dim); Scalar s(type.obj, value, true); NDArray result = full(shp, s); @@ -132,8 +131,7 @@ CN_NDArray* nda_random_array(int32_t dim, const uint64_t* shape) { return new CN_NDArray{NDArray(std::move(result))}; } -CN_NDArray* nda_reshape_array(CN_NDArray* arr, int32_t dim, - const uint64_t* shape) { +CN_NDArray* nda_reshape_array(CN_NDArray* arr, int32_t dim, const uint64_t* shape) { std::vector shp(shape, shape + dim); NDArray result = cupynumeric::reshape(arr->obj, shp, "C"); return new CN_NDArray{NDArray(std::move(result))}; @@ -171,6 +169,39 @@ void nda_add(CN_NDArray* rhs1, CN_NDArray* rhs2, CN_NDArray* out) { cupynumeric::add(rhs1->obj, rhs2->obj, out->obj); } +// NEW + +CN_NDArray* unique(CN_NDArray* arr){ + NDArray result = cunumeric::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))}; +} + +void nda_trace(CN_NDArray* arr, int32_t offset, int32_t a1, int32_t a2, CN_Type type, CN_NDArray* out){ + cupynumeric::trace(arr->obj, offset, a1, a2, type, out); +} + +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* diag(CN_NDArray* arr, int32_t k){ + NDArray result = cunumeric::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..555944b2 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -228,6 +228,27 @@ function nda_multiply(rhs1::NDArray, rhs2::NDArray, out::NDArray) return out end +function nda_diag(arr::NDArray, k::Int32) + ptr = ccall((: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) ccall((:nda_add, libnda), Cvoid, (NDArray_t, NDArray_t, NDArray_t), @@ -267,6 +288,37 @@ 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((:eye, libnda), + NDArray_t, (Int32, Legate.LegateTypeAllocated), + rows, legate_type) + return NDArray(ptr; T=T, n_dim=2) +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 +end + +function nda_trace(arr::NDArray, offset::Int32, a1::Int32, a2::Int32, ::Type{T}, out::NDArray) where {T} + legate_type = Legate.to_legate_type(T) + ccall((:nda_trace, libnda), + Cvoid, + (NDArray_t, Int32, Int32, Int32, Legate.LegateTypeAllocated, NDArray_t), + arr.ptr, offset, a1, a2, legate_type, out.ptr) + return out +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..e9f819ce 100644 --- a/src/ndarray/ndarray.jl +++ b/src/ndarray/ndarray.jl @@ -19,6 +19,64 @@ 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(rows, cuNumeric.Type(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} + out = cuNumeric.zeros(S) + nda_trace(arr, offset, a1, a2, cuNumeric.Type(S), out) + return out +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, 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) @@ -708,3 +766,7 @@ end function Base.isapprox(arr::NDArray{T}, arr2::NDArray{T}; atol=0, rtol=0) where {T} return compare(arr, arr2, atol, rtol) end + + + +