Enable MPS rand for more array sizes (#677)

Author: christiangnrd

lib/mps/matrixrandom.jl

@@ -93,12 +93,6 @@ synchronize_state(kern::MPSMatrixRandomMTGP32, cmdbuf::MTLCommandBuffer) =
     byteoffset = dest.offset * sizeof(T)
     bytesize = sizeof(dest)
 
-    # Even though `append_copy`` seems to work with any size or offset values, the documentation at
-    # https://developer.apple.com/documentation/metal/mtlblitcommandencoder/1400767-copyfrombuffer?language=objc
-    # mentions that both must be multiples of 4 bytes in MacOS so error when they are not
-    (bytesize % 4 == 0) || error(lazy"Destination buffer bytesize ($(bytesize)) must be a multiple of 4.")
-    (byteoffset % 4 == 0) || error(lazy"Destination buffer offset ($(byteoffset)) must be a multiple of 4.")
-
     cmdbuf = if bytesize % 16 == 0 && dest.offset == 0
         MTLCommandBuffer(queue) do cmdbuf
             vecDesc = MPSVectorDescriptor(bytesize ÷ sizeof(T2), T2)

lib/mps/random.jl

@@ -116,5 +116,5 @@ Random.randn(rng::RNG, dim1::Integer, dims::Integer...; storage=DefaultStorageMo
     Random.randn!(rng, MtlArray{Float32,length(dims) + 1,storage}(undef, dim1, dims...))
 
 # scalars
-Random.rand(rng::RNG, T::UniformType=Float32; storage=SharedStorage) = rand(rng, T, 4; storage)[1]
-Random.randn(rng::RNG, T::NormalType=Float32; storage=SharedStorage) = randn(rng, T, 4; storage)[1]
+Random.rand(rng::RNG, T::UniformType=Float32; storage=SharedStorage) = rand(rng, T, 1; storage)[1]
+Random.randn(rng::RNG, T::NormalType=Float32; storage=SharedStorage) = randn(rng, T, 1; storage)[1]

src/random.jl

@@ -9,28 +9,20 @@ mpsrand_rng() = MPS.default_rng()
 Random.rand!(A::MtlArray) = Random.rand!(gpuarrays_rng(), A)
 Random.randn!(A::MtlArray) = Random.randn!(gpuarrays_rng(), A)
 
-@inline function can_use_mpsrandom(A::MtlArray{T}) where {T}
-    return A.offset * sizeof(T) % 4 == 0 && sizeof(A) % 4 == 0
-end
-
 # Use MPS random functionality where possible
 function Random.rand!(A::MPS.UniformArray)
-    rng = can_use_mpsrandom(A) ? mpsrand_rng() : gpuarrays_rng()
-    return Random.rand!(rng, A)
+    return Random.rand!(mpsrand_rng(), A)
 end
 function Random.randn!(A::MPS.NormalArray)
-    rng = can_use_mpsrandom(A) ? mpsrand_rng() : gpuarrays_rng()
-    return Random.randn!(rng, A)
+    return Random.randn!(mpsrand_rng(), A)
 end
 
 # GPUArrays out-of-place
 function rand(T::MPS.UniformType, dims::Dims; storage=DefaultStorageMode)
-    rng =  prod(dims) * sizeof(T) % 4 == 0 ? mpsrand_rng() : gpuarrays_rng()
-    return Random.rand!(rng, MtlArray{T,length(dims),storage}(undef, dims...))
+    return Random.rand!(mpsrand_rng(), MtlArray{T,length(dims),storage}(undef, dims...))
 end
 function randn(T::MPS.NormalType, dims::Dims; storage=DefaultStorageMode)
-    rng =  prod(dims) * sizeof(T) % 4 == 0 ? mpsrand_rng() : gpuarrays_rng()
-    return Random.randn!(rng, MtlArray{T,length(dims),storage}(undef, dims...))
+    return Random.randn!(mpsrand_rng(), MtlArray{T,length(dims),storage}(undef, dims...))
 end
 rand(T::Type, dims::Dims; storage=DefaultStorageMode) =
     Random.rand!(gpuarrays_rng(), MtlArray{T,length(dims),storage}(undef, dims...))
@@ -39,12 +31,10 @@ randn(T::Type, dims::Dims; storage=DefaultStorageMode) =
 
 # support all dimension specifications
 function rand(T::MPS.UniformType, dim1::Integer, dims::Integer...; storage=DefaultStorageMode)
-    rng = (dim1 * prod(dims) * sizeof(T)) % 4 == 0 ? mpsrand_rng() : gpuarrays_rng()
-    return Random.rand!(rng, MtlArray{T,length(dims) + 1,storage}(undef, dim1, dims...))
+    return Random.rand!(mpsrand_rng(), MtlArray{T,length(dims) + 1,storage}(undef, dim1, dims...))
 end
 function randn(T::MPS.NormalType, dim1::Integer, dims::Integer...; storage=DefaultStorageMode)
-    rng = (dim1 * prod(dims) * sizeof(T)) % 4 == 0 ? mpsrand_rng() : gpuarrays_rng()
-    return Random.randn!(rng, MtlArray{T,length(dims) + 1,storage}(undef, dim1, dims...))
+    return Random.randn!(mpsrand_rng(), MtlArray{T,length(dims) + 1,storage}(undef, dim1, dims...))
 end
 
 rand(T::Type, dim1::Integer, dims::Integer...; storage=DefaultStorageMode) =
@@ -59,8 +49,8 @@ randn(dim1::Integer, dims::Integer...; storage=DefaultStorageMode) =
     Random.randn!(mpsrand_rng(), MtlArray{Float32,length(dims) + 1,storage}(undef, dim1, dims...))
 
 # scalars
-rand(T::Type=Float32; storage=SharedStorage) = rand(T, 4; storage)[1]
-randn(T::Type=Float32; storage=SharedStorage) = randn(T, 4; storage)[1]
+rand(T::Type=Float32; storage=SharedStorage) = rand(T, 1; storage)[1]
+randn(T::Type=Float32; storage=SharedStorage) = randn(T, 1; storage)[1]
 
 # seeding
 function seed!(seed=Base.rand(UInt64))

test/random.jl

@@ -24,12 +24,9 @@ const OOPLACE_TUPLES = [[(Metal.rand, rand, T) for T in RAND_TYPES];
                 # specified MPS rng
                 if T != Float16
                     fill!(A, T(0))
-                    if Metal.can_use_mpsrandom(A)
-                        f(rng, A)
-                        @test !iszero(collect(A))
-                    else
-                        @test_throws "Destination buffer" f(rng, A)
-                    end
+
+                    f(rng, A)
+                    @test !iszero(collect(A))
                 end
             end
 
@@ -44,12 +41,9 @@ const OOPLACE_TUPLES = [[(Metal.rand, rand, T) for T in RAND_TYPES];
                 # specified MPS rng
                 if T != Float16
                     fill!(A, T(0))
-                    if Metal.can_use_mpsrandom(A)
-                        f(rng, A)
-                        @test Array(A) == fill(1, 0)
-                    else
-                        @test_throws "Destination buffer" f(rng, A)
-                    end
+
+                    f(rng, A)
+                    @test Array(A) == fill(1, 0)
                 end
             end
         end
@@ -131,30 +125,22 @@ const OOPLACE_TUPLES = [[(Metal.rand, rand, T) for T in RAND_TYPES];
                 idx = 4:50
                 view_A = @view A[idx]
 
-                # Errors in Julia before crashing whole process
-                if Metal.can_use_mpsrandom(view_A)
-                    f(rng, view_A)
+                f(rng, view_A)
 
-                    cpuA = collect(A)
-                    @test !iszero(cpuA[idx])
-                    @test iszero(cpuA[1:100 .∉ Ref(idx)]) broken=(sizeof(view_A) % 4 != 0)
-                else
-                    @test_throws "Destination buffer" f(rng, view_A)
-                end
+                cpuA = collect(A)
+                @test !iszero(cpuA[idx])
+                @test iszero(cpuA[1:100 .∉ Ref(idx)])
 
                 ## Offset == 0
                 fill!(A, T(0))
                 idx = 1:51
                 view_A = @view A[idx]
-                if Metal.can_use_mpsrandom(view_A)
-                    f(rng, view_A)
-
-                    cpuA = collect(A)
-                    @test !iszero(cpuA[idx])
-                    @test iszero(cpuA[1:100 .∉ Ref(idx)])
-                else
-                    @test_throws "Destination buffer" f(rng, view_A)
-                end
+
+                f(rng, view_A)
+
+                cpuA = collect(A)
+                @test !iszero(cpuA[idx])
+                @test iszero(cpuA[1:100 .∉ Ref(idx)])
             end
         end
     end
@@ -200,12 +186,8 @@ const OOPLACE_TUPLES = [[(Metal.rand, rand, T) for T in RAND_TYPES];
 
                 # specified MPS rng
                 if T != Float16
-                    if length(zeros(args...)) * sizeof(T) % 4 == 0
-                        B = fr(rng, args...)
-                        @test eltype(B) == T
-                    else
-                        @test_throws "Destination buffer" fr(rng, args...)
-                    end
+                    B = fr(rng, args...)
+                    @test eltype(B) == T
                 end
             end
 
@@ -228,12 +210,9 @@ const OOPLACE_TUPLES = [[(Metal.rand, rand, T) for T in RAND_TYPES];
             # d == 2 and d == 3 are to hit the test cases where sizeof(A) <= 4
             @testset "$d" for d in (2, 3, (3, 3), (3, 3, 3), 16, (16, 16), (16, 16, 16), (1000,), (1000,1000), 16384, 16385)
                 A = zeros(T, d)
-                if (prod(d) * sizeof(T)) % 4 == 0
-                    f(rng, A)
-                    @test !iszero(A)
-                else
-                    @test_throws "Destination buffer" f(rng, A)
-                end
+
+                f(rng, A)
+                @test !iszero(A)
             end
         end
     end