transpose

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<uint64_t> 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<int64_t> 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);

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)

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
+
+
+
+