Better API for Unified Memory (#305)

Author: tgymnich

docs/src/usage/array.md

@@ -28,16 +28,16 @@ can construct `MtlArray`s in the same way as regular `Array` objects:
 
 ```jldoctest
 julia> MtlArray{Int}(undef, 2)
-2-element MtlVector{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+2-element MtlVector{Int64, Private}:
  0
  0
 
 julia> MtlArray{Int}(undef, (1,2))
-1×2 MtlMatrix{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+1×2 MtlMatrix{Int64, Private}:
  0  0
 
 julia> similar(ans)
-1×2 MtlMatrix{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+1×2 MtlMatrix{Int64, Private}:
  0  0
 ```
 
@@ -46,7 +46,7 @@ Copying memory to or from the GPU can be expressed using constructors as well, o
 
 ```jldoctest
 julia> a = MtlArray([1,2])
-2-element MtlVector{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+2-element MtlVector{Int64, Private}:
  1
  2
 
@@ -73,11 +73,11 @@ perform simple element-wise operations you can use `map` or `broadcast`:
 julia> a = MtlArray{Float32}(undef, (1,2));
 
 julia> a .= 5
-1×2 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+1×2 MtlMatrix{Float32, Private}:
  5.0  5.0
 
 julia> map(sin, a)
-1×2 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+1×2 MtlMatrix{Float32, Private}:
  -0.958924  -0.958924
 ```
 
@@ -86,23 +86,23 @@ To reduce the dimensionality of arrays, Metal.jl implements the various flavours
 
 ```jldoctest
 julia> a = Metal.ones(2,3)
-2×3 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+2×3 MtlMatrix{Float32, Private}:
  1.0  1.0  1.0
  1.0  1.0  1.0
 
 julia> reduce(+, a)
 6.0f0
 
 julia> mapreduce(sin, *, a; dims=2)
-2×1 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+2×1 MtlMatrix{Float32, Private}:
  0.59582335
  0.59582335
 
 julia> b = Metal.zeros(1)
-1-element MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+1-element MtlVector{Float32, Private}:
  0.0
 
 julia> Base.mapreducedim!(identity, +, b, a)
-1×1 MtlMatrix{Float32, Metal.MTL.MTLResourceStorageModePrivate}:
+1×1 MtlMatrix{Float32, Private}:
  6.0
 ```

lib/mtl/buffer.jl

@@ -15,7 +15,7 @@ end
 Base.sizeof(buf::MTLBuffer) = Int(buf.length)
 
 function Base.convert(::Type{Ptr{T}}, buf::MTLBuffer) where {T}
-    buf.storageMode == Private && error("Cannot access the contents of a private buffer")
+    buf.storageMode == MTLStorageModePrivate && error("Cannot access the contents of a private buffer")
     convert(Ptr{T}, buf.contents)
 end
 
@@ -25,7 +25,7 @@ end
 function MTLBuffer(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer;
                    storage=Private, hazard_tracking=DefaultTracking,
                    cache_mode=DefaultCPUCache)
-    opts = storage | hazard_tracking | cache_mode
+    opts = convert(MTLResourceOptions, storage) | hazard_tracking | cache_mode
 
     @assert 0 < bytesize <= dev.maxBufferLength # XXX: not supported by MTLHeap
     ptr = alloc_buffer(dev, bytesize, opts)
@@ -37,7 +37,7 @@ function MTLBuffer(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer, ptr::Ptr;
                    storage=Managed, hazard_tracking=DefaultTracking,
                    cache_mode=DefaultCPUCache)
     storage == Private && error("Can't create a Private copy-allocated buffer.")
-    opts =  storage | hazard_tracking | cache_mode
+    opts =  convert(MTLResourceOptions, storage) | hazard_tracking | cache_mode
 
     @assert 0 < bytesize <= dev.maxBufferLength # XXX: not supported by MTLHeap
     ptr = alloc_buffer(dev, bytesize, opts, ptr)

lib/mtl/storage_type.jl

@@ -1,38 +1,30 @@
-abstract type StorageMode end
-export Shared, Managed, Private
+export Shared, Managed, Private, CPUStorage
 export ReadUsage, WriteUsage, ReadWriteUsage
 
 # Metal Has 4 storage types
 # Shared  -> Buffer in Host memory, accessed by the GPU. Requires no sync
 # Managed -> Mirrored memory buffers in host and GPU. Requires syncing
 # Private -> Memory in Device, not accessible by Host.
 # Memoryless -> iOS stuff. ignore it
-module AS
-import ..StorageMode
 
+abstract type StorageMode end
 struct Shared      <: StorageMode end
 struct Managed     <: StorageMode end
 struct Private     <: StorageMode end
 struct Memoryless  <: StorageMode end
-end
-
-const CPUStorage = Union{AS.Shared,AS.Managed}
-Base.convert(::Type{MTLStorageMode}, ::Type{AS.Shared})     = MTLStorageModeShared
-Base.convert(::Type{MTLStorageMode}, ::Type{AS.Managed})    = MTLStorageModeManaged
-Base.convert(::Type{MTLStorageMode}, ::Type{AS.Private})    = MTLStorageModePrivate
-Base.convert(::Type{MTLStorageMode}, ::Type{AS.Memoryless}) = MTLStorageModeMemoryless
 
-Base.convert(::Type{MTLResourceOptions}, ::Type{AS.Shared})     = MTLResourceStorageModeShared
-Base.convert(::Type{MTLResourceOptions}, ::Type{AS.Managed})    = MTLResourceStorageModeManaged
-Base.convert(::Type{MTLResourceOptions}, ::Type{AS.Private})    = MTLResourceStorageModePrivate
-Base.convert(::Type{MTLResourceOptions}, ::Type{AS.Memoryless}) = MTLResourceStorageModeMemoryless
+const CPUStorage = Union{Shared,Managed}
+Base.convert(::Type{MTLStorageMode}, ::Type{Shared})     = MTLStorageModeShared
+Base.convert(::Type{MTLStorageMode}, ::Type{Managed})    = MTLStorageModeManaged
+Base.convert(::Type{MTLStorageMode}, ::Type{Private})    = MTLStorageModePrivate
+Base.convert(::Type{MTLStorageMode}, ::Type{Memoryless}) = MTLStorageModeMemoryless
 
-Base.convert(::Type{MTLResourceOptions}, SM::MTL.MTLStorageMode)     = MTLResourceOptions(UInt(SM) << 4)
+Base.convert(::Type{MTLResourceOptions}, ::Type{Shared})     = MTLResourceStorageModeShared
+Base.convert(::Type{MTLResourceOptions}, ::Type{Managed})    = MTLResourceStorageModeManaged
+Base.convert(::Type{MTLResourceOptions}, ::Type{Private})    = MTLResourceStorageModePrivate
+Base.convert(::Type{MTLResourceOptions}, ::Type{Memoryless}) = MTLResourceStorageModeMemoryless
 
-const Shared                = MTLResourceStorageModeShared
-const Managed               = MTLResourceStorageModeManaged
-const Private               = MTLResourceStorageModePrivate
-const Memoryless            = MTLResourceStorageModeMemoryless
+Base.convert(::Type{MTLResourceOptions}, SM::MTLStorageMode)     = MTLResourceOptions(UInt(SM) << 4)
 
 const DefaultCPUCache       = MTLResourceCPUCacheModeDefaultCache
 const CombinedWriteCPUCache = MTLResourceCPUCacheModeWriteCombined

src/array.jl

@@ -1,6 +1,6 @@
 # host array
 
-export MtlArray, MtlVector, MtlMatrix, MtlVecOrMat, mtl
+export MtlArray, MtlVector, MtlMatrix, MtlVecOrMat, mtl, is_shared, is_managed, is_private
 
 function hasfieldcount(@nospecialize(dt))
     try
@@ -77,8 +77,16 @@ mutable struct MtlArray{T,N,S} <: AbstractGPUArray{T,N}
   function MtlArray{T,N}(data::DataRef{<:MTLBuffer}, dims::Dims{N};
                          maxsize::Int=prod(dims) * sizeof(T), offset::Int=0) where {T,N}
       check_eltype(T)
-      S = convert(MTL.MTLResourceOptions, data[].storageMode)
-      obj = new{T,N,S}(copy(data), maxsize, offset, dims)
+      storagemode = data[].storageMode
+      if storagemode == MTL.MTLStorageModeShared
+        obj = new{T,N,Shared}(copy(data), maxsize, offset, dims)
+      elseif storagemode == MTL.MTLStorageModeManaged
+        obj = new{T,N,Managed}(copy(data), maxsize, offset, dims)
+      elseif storagemode == MTL.MTLStorageModePrivate
+        obj = new{T,N,Private}(copy(data), maxsize, offset, dims)
+      elseif storagemode == MTL.MTLStorageModeMemoryless
+        obj = new{T,N,Memoryless}(copy(data), maxsize, offset, dims)
+      end
       finalizer(unsafe_free!, obj)
   end
 end
@@ -90,6 +98,10 @@ device(A::MtlArray) = A.data[].device
 storagemode(x::MtlArray) = storagemode(typeof(x))
 storagemode(::Type{<:MtlArray{<:Any,<:Any,S}}) where {S} = S
 
+is_shared(a::MtlArray) = storagemode(a) == Shared
+is_managed(a::MtlArray) = storagemode(a) == Managed
+is_private(a::MtlArray) = storagemode(a) == Private
+is_memoryless(a::MtlArray) = storagemode(a) == Memoryless
 
 ## convenience constructors
 
@@ -144,15 +156,42 @@ Base.elsize(::Type{<:MtlArray{T}}) where {T} = sizeof(T)
 Base.size(x::MtlArray) = x.dims
 Base.sizeof(x::MtlArray) = Base.elsize(x) * length(x)
 
-Base.pointer(x::MtlArray{T}) where {T} = Base.unsafe_convert(MtlPointer{T}, x)
-@inline function Base.pointer(x::MtlArray{T}, i::Integer) where T
-    Base.unsafe_convert(MtlPointer{T}, x) + Base._memory_offset(x, i)
+@inline function Base.pointer(x::MtlArray{T}, i::Integer=1; storage=Private) where {T}
+  PT = if storage == Private
+    MtlPointer{T}
+  elseif storage == Shared || storage == Managed
+    Ptr{T}
+  else
+    error("unknown memory type")
+  end
+  Base.unsafe_convert(PT, x) + Base._memory_offset(x, i)
 end
 
-Base.unsafe_convert(::Type{Ptr{S}}, x::MtlArray{T}) where {S, T} =
-  throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
-Base.unsafe_convert(::Type{MtlPointer{T}}, x::MtlArray) where {T} =
-  MtlPointer{T}(x.data[], x.offset*Base.elsize(x))
+
+function Base.unsafe_convert(::Type{MtlPointer{T}}, x::MtlArray) where {T}
+   buf = x.data[]
+   MtlPointer{T}(buf, x.offset*Base.elsize(x))
+ end
+
+function Base.unsafe_convert(::Type{Ptr{S}}, x::MtlArray{T}) where {S, T}
+  buf = x.data[]
+  if is_private(x)
+    throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
+  end
+  convert(Ptr{T}, buf) + x.offset*Base.elsize(x)
+end
+
+
+## indexing
+function Base.getindex(x::MtlArray{T,N,S}, I::Int) where {T,N,S<:Union{Shared,Managed}}
+  @boundscheck checkbounds(x, I)
+  unsafe_load(pointer(x, I; storage=S))
+end
+
+function Base.setindex!(x::MtlArray{T,N,S}, v, I::Int) where {T,N,S<:Union{Shared,Managed}}
+  @boundscheck checkbounds(x, I)
+  unsafe_store!(pointer(x, I; storage=S), v)
+end
 
 
 ## interop with other arrays
@@ -354,7 +393,7 @@ Uses Adapt.jl to act inside some wrapper structs.
 
 ```jldoctests
 julia> mtl(ones(3)')
-1×3 adjoint(::MtlVector{Float32, Metal.MTL.MTLResourceStorageModePrivate}) with eltype Float32:
+1×3 adjoint(::MtlVector{Float32, Private}) with eltype Float32:
  1.0  1.0  1.0
 
 julia> mtl(zeros(1,3); storage=Shared)
@@ -365,13 +404,13 @@ julia> mtl(1:3)
 1:3
 
 julia> MtlArray(1:3)
-3-element MtlVector{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+3-element MtlVector{Int64, Private}:
  1
  2
  3
 
 julia> mtl[1,2,3]
-3-element MtlVector{Int64, Metal.MTL.MTLResourceStorageModePrivate}:
+3-element MtlVector{Int64, Private}:
  1
  2
  3
@@ -433,8 +472,9 @@ Base.unsafe_convert(::Type{MTL.MTLBuffer}, A::PermutedDimsArray) =
 
 ## unsafe_wrap
 
-Base.unsafe_wrap(t::Type{<:Array}, arr::MtlArray, dims; own=false) =
-  unsafe_wrap(t, arr.data[], dims; own=own)
+function Base.unsafe_wrap(::Type{<:Array}, arr::MtlArray{T,N}, dims=size(arr); own=false) where {T,N}
+  return unsafe_wrap(Array{T,N}, arr.data[], dims; own=own)
+end
 
 function Base.unsafe_wrap(t::Type{<:Array{T}}, buf::MTLBuffer, dims; own=false) where T
     ptr = convert(Ptr{T}, buf)

src/memory.jl

@@ -45,7 +45,7 @@ function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T}, src::Ptr{T}, N:
         unsafe_copyto!(convert(Ptr{T}, dst), src, N)
     elseif storage_type == MTL.MTLStorageModeManaged
         unsafe_copyto!(convert(Ptr{T}, dst), src, N)
-        MTL.DidModifyRange!(dst, 1:N)
+        MTL.DidModifyRange!(dst.buffer, 1:N)
     end
     return dst
 end

src/random.jl

@@ -7,18 +7,18 @@ Random.rand!(A::MtlArray) = Random.rand!(gpuarrays_rng(), A)
 Random.randn!(A::MtlArray) = Random.randn!(gpuarrays_rng(), A)
 
 # GPUArrays out-of-place
-rand(T::Type, dims::Dims; storage::MTL.MTLResourceOptions=DefaultStorageMode) = Random.rand!(MtlArray{T,length(dims),storage}(undef, dims...))
-randn(T::Type, dims::Dims; storage::MTL.MTLResourceOptions=DefaultStorageMode, kwargs...) = Random.randn!(MtlArray{T,length(dims),storage}(undef, dims...); kwargs...)
+rand(T::Type, dims::Dims; storage=DefaultStorageMode) = Random.rand!(MtlArray{T,length(dims),storage}(undef, dims...))
+randn(T::Type, dims::Dims; storage=DefaultStorageMode, kwargs...) = Random.randn!(MtlArray{T,length(dims),storage}(undef, dims...); kwargs...)
 
 # support all dimension specifications
-rand(T::Type, dim1::Integer, dims::Integer...; storage::MTL.MTLResourceOptions=DefaultStorageMode) =
+rand(T::Type, dim1::Integer, dims::Integer...; storage=DefaultStorageMode) =
     Random.rand!(MtlArray{T,length(dims)+1,storage}(undef, dim1, dims...))
-randn(T::Type, dim1::Integer, dims::Integer...; storage::MTL.MTLResourceOptions=DefaultStorageMode, kwargs...) =
+randn(T::Type, dim1::Integer, dims::Integer...; storage=DefaultStorageMode, kwargs...) =
     Random.randn!(MtlArray{T,length(dims)+1,storage}(undef, dim1, dims...); kwargs...)
 
 # untyped out-of-place
-rand(dim1::Integer, dims::Integer...; storage::MTL.MTLResourceOptions=DefaultStorageMode) = Random.rand!(MtlArray{Float32,length(dims)+1,storage}(undef, dim1, dims...))
-randn(dim1::Integer, dims::Integer...; storage::MTL.MTLResourceOptions=DefaultStorageMode, kwargs...) = Random.randn!(MtlArray{Float32,length(dims)+1,storage}(undef, dim1, dims...); kwargs...)
+rand(dim1::Integer, dims::Integer...; storage=DefaultStorageMode) = Random.rand!(MtlArray{Float32,length(dims)+1,storage}(undef, dim1, dims...))
+randn(dim1::Integer, dims::Integer...; storage=DefaultStorageMode, kwargs...) = Random.randn!(MtlArray{Float32,length(dims)+1,storage}(undef, dim1, dims...); kwargs...)
 
 # seeding
 seed!(seed=Base.rand(UInt64)) = Random.seed!(gpuarrays_rng(), seed)

test/array.jl

@@ -129,8 +129,35 @@ check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
     # private storage errors.
     if SM == Metal.Private
         let arr_mtl = Metal.zeros(Float32, dim...; storage=Private)
+            @test is_private(arr_mtl) && !is_shared(arr_mtl) && !is_managed(arr_mtl)
             @test_throws "Cannot access the contents of a private buffer" arr_cpu = unsafe_wrap(Array{Float32}, arr_mtl, dim)
         end
+
+        let b = rand(Float32, 10)
+            arr_mtl = mtl(b, storage=Private)
+            @test_throws ErrorException arr_mtl[1]
+            @test Metal.@allowscalar arr_mtl[1] == b[1]
+        end
+    elseif SM == Metal.Shared
+        let arr_mtl = Metal.zeros(Float32, dim...; storage=Shared)
+            @test !is_private(arr_mtl) && is_shared(arr_mtl) && !is_managed(arr_mtl)
+            @test unsafe_wrap(Array{Float32}, arr_mtl) isa Array{Float32}
+        end
+
+        let b = rand(Float32, 10)
+            arr_mtl = mtl(b, storage=Shared)
+            @test arr_mtl[1] == b[1]
+        end
+    elseif SM == Metal.Managed
+        let arr_mtl = Metal.zeros(Float32, dim...; storage=Managed)
+            @test !is_private(arr_mtl) && !is_shared(arr_mtl) && is_managed(arr_mtl)
+            @test unsafe_wrap(Array{Float32}, arr_mtl) isa Array{Float32}
+        end
+
+        let b = rand(Float32, 10)
+            arr_mtl = mtl(b, storage=Managed)
+            @test arr_mtl[1] == b[1]
+        end
     end
 end