Fix MPS.synchronize_state

JuliaGPU · Oct 1, 2024 · 69c8d2a · 69c8d2a
1 parent f605bcb
commit 69c8d2a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/lib/mps/matrixrandom.jl b/lib/mps/matrixrandom.jl
@@ -107,7 +107,7 @@ for R in [:MPSMatrixRandomMTGP32, :MPSMatrixRandomPhilox]
 end
 
 synchronize_state(kern::MPSMatrixRandomMTGP32, cmdbuf::MTLCommandBuffer) =
-    @objc [obj::id{MPSMatrixRandomMTGP32} synchronizeStateOnCommandBuffer:cmdbuf::id{MTLCommandBuffer}]::Nothing
+    @objc [kern::id{MPSMatrixRandomMTGP32} synchronizeStateOnCommandBuffer:cmdbuf::id{MTLCommandBuffer}]::Nothing
 
 
 @inline function _mpsmat_rand!(randkern::MPSMatrixRandom, dest::MtlArray{T}, ::Type{T2};

diff --git a/test/mps/matrix.jl b/test/mps/matrix.jl
@@ -170,3 +170,10 @@ using .MPS: MPSMatrixFindTopK
     @test topk.sourceColumns == cols
     @test topk.sourceRows == rows
 end
+
+# Ensure that the function does not error
+@testset "MPSMatrixRandom sync state" begin
+    cmdbuf = MTL.MTLCommandBuffer(global_queue(device()))
+    rng = MPS.MPSMatrixRandomMTGP32(device())
+    @test isnothing(MPS.synchronize_state(rng, cmdbuf))
+end