transpose

.githash

@@ -1 +1 @@
-dd8623e18ce0b61be51c3a90c8e7927d642a7574
+1fd9472c5bdee6909bb67f928361efdcc94f468e

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

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)

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)
 

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

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 <[email protected]>
+ *            Ethan Meitz <[email protected]>
+ *            Nader Rahal <[email protected]>
+=#
+
+@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