Make lu results have same storage mode as input (#435)

JuliaGPU · Sep 30, 2024 · 3bbda32 · 3bbda32
1 parent 8652754
commit 3bbda32
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/lib/mps/linalg.jl b/lib/mps/linalg.jl
@@ -128,7 +128,7 @@ LinearAlgebra.ipiv2perm(v::MtlVector{T}, maxi::Integer) where T =
         encode!(cbuf, kernel, descriptor)
     end
 
-    P = MtlMatrix{UInt32}(undef, 1, min(N, M))
+    P = similar(A, UInt32, 1, min(N, M))
     status = MtlArray{MPSMatrixDecompositionStatus}(undef)
 
     commitAndContinue!(cmdbuf) do cbuf
@@ -137,7 +137,7 @@ LinearAlgebra.ipiv2perm(v::MtlVector{T}, maxi::Integer) where T =
         encode!(cbuf, kernel, mps_at, mps_at, mps_p, status)
     end
 
-    B = MtlMatrix{T}(undef, M, N)
+    B = similar(A, M, N)
 
     commit!(cmdbuf) do cbuf
         mps_b = MPSMatrix(B)
@@ -186,7 +186,7 @@ end
         encode!(cbuf, kernel, descriptor)
     end
 
-    P = MtlMatrix{UInt32}(undef, 1, min(N, M))
+    P = similar(A, UInt32, 1, min(N, M))
     status = MtlArray{MPSMatrixDecompositionStatus}(undef)
 
     commitAndContinue!(cmdbuf) do cbuf

diff --git a/test/mps/linalg.jl b/test/mps/linalg.jl
@@ -190,6 +190,7 @@ end
     end
 end
 
+using Metal: storagemode
 @testset "decompositions" begin
     A = MtlMatrix(rand(Float32, 1024, 1024))
     lua = lu(A)
@@ -211,6 +212,13 @@ end
 
     A = MtlMatrix{Float32}([1 2; 0 0])
     @test_throws SingularException lu(A)
+
+    altStorage = Metal.DefaultStorageMode != Metal.PrivateStorage ? Metal.PrivateStorage : Metal.SharedStorage
+    A = MtlMatrix{Float32,altStorage}(rand(Float32, 1024, 1024))
+    lua = lu(A)
+    @test storagemode(lua.factors) == storagemode(lua.ipiv) == storagemode(A)
+    lua = lu!(A)
+    @test storagemode(lua.factors) == storagemode(lua.ipiv) == storagemode(A)
 end
 
 using .MPS: MPSMatrixSoftMax, MPSMatrixLogSoftMax