use unified memory to check decomposition status

JuliaGPU · Mar 11, 2024 · 69fa85d · 69fa85d
1 parent b5c4140
commit 69fa85d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lib/mps/linalg.jl b/lib/mps/linalg.jl
@@ -174,7 +174,7 @@ function LinearAlgebra.lu(A::MtlMatrix{T}; check::Bool = true) where {T<:MtlFloa
     end
 
     P = MtlMatrix{UInt32}(undef, 1, min(N, M))
-    status = MtlArray{MPSMatrixDecompositionStatus}(undef)
+    status = MtlArray{MPSMatrixDecompositionStatus,0,Shared}(undef)
 
     cmdbuf_lu = MTLCommandBuffer(queue) do cmdbuf
         mps_p = MPSMatrix(P)
@@ -196,7 +196,7 @@ function LinearAlgebra.lu(A::MtlMatrix{T}; check::Bool = true) where {T<:MtlFloa
 
     wait_completed(cmdbuf_lu)
 
-    status = convert(LinearAlgebra.BlasInt, Metal.@allowscalar status[])
+    status = convert(LinearAlgebra.BlasInt, status[])
     check && checknonsingular(status)
 
     return LinearAlgebra.LU(B, p, status)
@@ -219,7 +219,7 @@ function LinearAlgebra.lu!(A::MtlMatrix{T}; check::Bool = true) where {T<:MtlFlo
     end
 
     P = MtlMatrix{UInt32}(undef, 1, min(N, M))
-    status = MtlArray{MPSMatrixDecompositionStatus}(undef)
+    status = MtlArray{MPSMatrixDecompositionStatus,0,Shared}(undef)
 
     cmdbuf_lu = MTLCommandBuffer(queue) do cmdbuf
         mps_p = MPSMatrix(P)
@@ -237,7 +237,7 @@ function LinearAlgebra.lu!(A::MtlMatrix{T}; check::Bool = true) where {T<:MtlFlo
 
     wait_completed(cmdbuf_lu)
 
-    status = convert(LinearAlgebra.BlasInt, Metal.@allowscalar status[])
+    status = convert(LinearAlgebra.BlasInt, status[])
     check && checknonsingular(status)
 
     return LinearAlgebra.LU(A, p, status)