Merge pull request #319 from christiangnrd/minor-cleanup

Minor cleanup

Author: maleadt

.gitignore

@@ -4,3 +4,4 @@ Manifest.toml
 *.trace
 wip.*
 dev
+.vscode

Project.toml

@@ -27,6 +27,12 @@ SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
+[weakdeps]
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+
+[extensions]
+SpecialFunctionsExt = "SpecialFunctions"
+
 [compat]
 Adapt = "4"
 Artifacts = "1"
@@ -47,11 +53,5 @@ SHA = "0.7"
 StaticArrays = "1"
 julia = "1.8"
 
-[extensions]
-SpecialFunctionsExt = "SpecialFunctions"
-
 [extras]
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-
-[weakdeps]
-SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

lib/mtl/buffer.jl

@@ -33,13 +33,13 @@ function MTLBuffer(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer;
     return MTLBuffer(ptr)
 end
 
-function MTLBuffer(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer, ptr::Ptr;
+function MTLBuffer(dev::MTLDevice, bytesize::Integer, ptr::Ptr;
                    storage=Managed, hazard_tracking=DefaultTracking,
                    cache_mode=DefaultCPUCache)
     storage == Private && error("Can't create a Private copy-allocated buffer.")
     opts =  convert(MTLResourceOptions, storage) | hazard_tracking | cache_mode
 
-    @assert 0 < bytesize <= dev.maxBufferLength # XXX: not supported by MTLHeap
+    @assert 0 < bytesize <= dev.maxBufferLength
     ptr = alloc_buffer(dev, bytesize, opts, ptr)
 
     return MTLBuffer(ptr)
@@ -58,10 +58,6 @@ alloc_buffer(dev::MTLDevice, bytesize, opts, ptr::Ptr) =
 alloc_buffer(dev::MTLHeap, bytesize, opts) =
     @objc [dev::id{MTLHeap} newBufferWithLength:bytesize::NSUInteger
                             options:opts::MTLResourceOptions]::id{MTLBuffer}
-alloc_buffer(dev::MTLHeap, bytesize, opts, ptr::Ptr) =
-    @objc [dev::id{MTLHeap} newBufferWithBytes:ptr::Ptr{Cvoid}
-                            length:bytesize::NSUInteger
-                            options:opts::MTLResourceOptions]::id{MTLBuffer}
 
 """
     DidModifyRange!(buf::MTLBuffer, range::UnitRange)

src/array.jl

@@ -169,7 +169,7 @@ Base.sizeof(x::MtlArray) = Base.elsize(x) * length(x)
 
 @inline function Base.pointer(x::MtlArray{T}, i::Integer=1; storage=Private) where {T}
   PT = if storage == Private
-    MtlPointer{T}
+    MtlPtr{T}
   elseif storage == Shared || storage == Managed
     Ptr{T}
   else
@@ -179,9 +179,9 @@ Base.sizeof(x::MtlArray) = Base.elsize(x) * length(x)
 end
 
 
-function Base.unsafe_convert(::Type{MtlPointer{T}}, x::MtlArray) where {T}
+function Base.unsafe_convert(::Type{MtlPtr{T}}, x::MtlArray) where {T}
    buf = x.data[]
-   MtlPointer{T}(buf, x.offset*Base.elsize(x))
+   MtlPtr{T}(buf, x.offset*Base.elsize(x))
  end
 
 function Base.unsafe_convert(::Type{Ptr{S}}, x::MtlArray{T}) where {S, T}
@@ -479,15 +479,15 @@ Base.unsafe_convert(::Type{MTL.MTLBuffer}, A::PermutedDimsArray) =
 ## unsafe_wrap
 
 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)
+  return unsafe_wrap(Array{T,N}, arr.data[], dims; own)
 end
 
 function Base.unsafe_wrap(t::Type{<:Array{T}}, buf::MTLBuffer, dims; own=false) where T
     ptr = convert(Ptr{T}, buf)
     return unsafe_wrap(t, ptr, dims; own)
 end
 
-function Base.unsafe_wrap(t::Type{<:Array{T}}, ptr::MtlPointer{T}, dims; own=false) where T
+function Base.unsafe_wrap(t::Type{<:Array{T}}, ptr::MtlPtr{T}, dims; own=false) where T
     return unsafe_wrap(t, convert(Ptr{T}, ptr), dims; own)
 end
 
@@ -513,7 +513,7 @@ function Base.resize!(A::MtlVector{T}, n::Integer) where T
   # replace the data with a new one. this 'unshares' the array.
   # as a result, we can safely support resizing unowned buffers.
   buf = alloc(device(A), bufsize; storage=storagemode(A))
-  ptr = MtlPointer{T}(buf)
+  ptr = MtlPtr{T}(buf)
   m = min(length(A), n)
   if m > 0
     unsafe_copyto!(device(A), ptr, pointer(A), m)

src/compiler/execution.jl

@@ -112,7 +112,7 @@ function Adapt.adapt_storage(to::Adaptor, buf::MTLBuffer)
     end
     reinterpret(Core.LLVMPtr{Nothing,AS.Device}, buf.gpuAddress)
 end
-function Adapt.adapt_storage(to::Adaptor, ptr::MtlPointer{T}) where {T}
+function Adapt.adapt_storage(to::Adaptor, ptr::MtlPtr{T}) where {T}
     reinterpret(Core.LLVMPtr{T,AS.Device}, adapt(to, ptr.buffer)) + ptr.offset
 end
 
@@ -209,7 +209,7 @@ const _kernel_instances = Dict{UInt, Any}()
     end
 
     # the arguments passed into this function have not been `mtlconvert`ed, because we need
-    # to retain the top-level MTLBuffer and MtlPointer objects. eager conversion of nested
+    # to retain the top-level MTLBuffer and MtlPtr objects. eager conversion of nested
     # such objects to LLVMPtr seems fine, somehow.
     # TODO: can we just convert everything eagerly and support top-level LLVMPtrs?
 
@@ -219,7 +219,7 @@ const _kernel_instances = Dict{UInt, Any}()
         if argtyp <: MTLBuffer
             # top-level buffers are passed as a pointer-valued argument
             push!(ex.args, :(set_buffer!(cce, $argex, 0, $idx)))
-        elseif argtyp <: MtlPointer
+        elseif argtyp <: MtlPtr
             # the same as a buffer, but with an offset
             push!(ex.args, :(set_buffer!(cce, $argex.buffer, $argex.offset, $idx)))
         elseif isghosttype(argtyp) || Core.Compiler.isconstType(argtyp)
@@ -253,7 +253,7 @@ end
     end
 
     # pass by reference, in an argument buffer
-    argument_buffer = alloc(kernel.pipeline.device, sizeof(argtyp), storage=Shared)
+    argument_buffer = alloc(kernel.pipeline.device, sizeof(argtyp); storage=Shared)
     argument_buffer.label = "MTLBuffer for kernel argument"
     unsafe_store!(convert(Ptr{argtyp}, argument_buffer), arg)
     return argument_buffer

src/memory.jl

@@ -7,39 +7,39 @@
 # we cannot take a MTLBuffer's handle and work with that as it were a pointer to memory.
 # instead, the Metal APIs always take the original handle and an offset parameter.
 
-struct MtlPointer{T}
+struct MtlPtr{T}
     buffer::MTLBuffer
     offset::UInt    # in bytes
 
-    function MtlPointer{T}(buffer::MTLBuffer, offset=0) where {T}
+    function MtlPtr{T}(buffer::MTLBuffer, offset=0) where {T}
         new(buffer, offset)
     end
 end
 
-Base.eltype(::Type{<:MtlPointer{T}}) where {T} = T
+Base.eltype(::Type{<:MtlPtr{T}}) where {T} = T
 
 # limited arithmetic
-Base.:(+)(x::MtlPointer{T}, y::Integer) where {T} = MtlPointer{T}(x.buffer, x.offset+y)
-Base.:(-)(x::MtlPointer{T}, y::Integer) where {T} = MtlPointer{T}(x.buffer, x.offset-y)
-Base.:(+)(x::Integer, y::MtlPointer{T}) where {T} = MtlPointer{T}(x.buffer, y+x.offset)
+Base.:(+)(x::MtlPtr{T}, y::Integer) where {T} = MtlPtr{T}(x.buffer, x.offset+y)
+Base.:(-)(x::MtlPtr{T}, y::Integer) where {T} = MtlPtr{T}(x.buffer, x.offset-y)
+Base.:(+)(x::Integer, y::MtlPtr{T}) where {T} = MtlPtr{T}(x.buffer, y+x.offset)
 
-Base.convert(::Type{Ptr{T}}, ptr::MtlPointer) where {T} =
+Base.convert(::Type{Ptr{T}}, ptr::MtlPtr) where {T} =
     convert(Ptr{T}, ptr.buffer) + ptr.offset
 
 
 ## operations
 
 # CPU -> GPU
-function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T}, src::Ptr{T}, N::Integer;
+function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPtr{T}, src::Ptr{T}, N::Integer;
                              queue::MTLCommandQueue=global_queue(dev), async::Bool=false) where T
     storage_type = dst.buffer.storageMode
     if storage_type == MTL.MTLStorageModePrivate
         # stage through a shared buffer
-        # shared = alloc(dev, N*sizeof(T), src, storage=Shared)
+        # shared = alloc(dev, N*sizeof(T), src; storage=Shared)
         # unsafe_copyto!(dev, dst, pointer(shared), N; queue, async=false)
         # free(shared)
-        tmp_buf = alloc(dev, N*sizeof(T), src, storage=Shared) #CPU -> GPU (Shared)
-        unsafe_copyto!(dev, MtlPointer{T}(dst.buffer, dst.offset), MtlPointer{T}(tmp_buf, 0), N; queue, async=false) # GPU (Shared) -> GPU (Private)
+        tmp_buf = alloc(dev, N*sizeof(T), src; storage=Shared) #CPU -> GPU (Shared)
+        unsafe_copyto!(dev, MtlPtr{T}(dst.buffer, dst.offset), MtlPtr{T}(tmp_buf, 0), N; queue, async=false) # GPU (Shared) -> GPU (Private)
         free(tmp_buf)
     elseif storage_type == MTL.MTLStorageModeShared
         unsafe_copyto!(convert(Ptr{T}, dst), src, N)
@@ -51,13 +51,13 @@ function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T}, src::Ptr{T}, N:
 end
 
 # GPU -> CPU
-function Base.unsafe_copyto!(dev::MTLDevice, dst::Ptr{T}, src::MtlPointer{T}, N::Integer;
+function Base.unsafe_copyto!(dev::MTLDevice, dst::Ptr{T}, src::MtlPtr{T}, N::Integer;
                              queue::MTLCommandQueue=global_queue(dev), async::Bool=false) where T
     storage_type = src.buffer.storageMode
     if storage_type ==  MTL.MTLStorageModePrivate
         # stage through a shared buffer
-        shared = alloc(dev, N*sizeof(T), storage=Shared)
-        unsafe_copyto!(dev, MtlPointer{T}(shared, 0), MtlPointer{T}(src.buffer, src.offset), N; queue, async=false)
+        shared = alloc(dev, N*sizeof(T); storage=Shared)
+        unsafe_copyto!(dev, MtlPtr{T}(shared, 0), MtlPtr{T}(src.buffer, src.offset), N; queue, async=false)
         unsafe_copyto!(dst, convert(Ptr{T}, shared), N)
         free(shared)
     elseif storage_type ==  MTL.MTLStorageModeShared
@@ -69,8 +69,8 @@ function Base.unsafe_copyto!(dev::MTLDevice, dst::Ptr{T}, src::MtlPointer{T}, N:
 end
 
 # GPU -> GPU
-@autoreleasepool function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T},
-                                              src::MtlPointer{T}, N::Integer;
+@autoreleasepool function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPtr{T},
+                                              src::MtlPtr{T}, N::Integer;
                                               queue::MTLCommandQueue=global_queue(dev),
                                               async::Bool=false) where T
     cmdbuf = MTLCommandBuffer(queue)
@@ -81,7 +81,7 @@ end
     async || wait_completed(cmdbuf)
 end
 
-@autoreleasepool function unsafe_fill!(dev::MTLDevice, ptr::MtlPointer{T},
+@autoreleasepool function unsafe_fill!(dev::MTLDevice, ptr::MtlPtr{T},
                                        value::Union{UInt8,Int8}, N::Integer) where T
     cmdbuf = MTLCommandBuffer(global_queue(dev))
     MTLBlitCommandEncoder(cmdbuf) do enc

src/pool.jl

@@ -51,11 +51,11 @@ The storage kwarg controls where the buffer is stored. Possible values are:
 Note that `Private` buffers can't be directly accessed from the CPU, therefore you cannot
 use this option if you pass a ptr to initialize the memory.
 """
-function alloc(dev::Union{MTLDevice,MTLHeap}, sz::Integer, args...; storage, kwargs...)
+function alloc(dev::Union{MTLDevice,MTLHeap}, sz::Integer, args...; kwargs...)
     @signpost_event log=log_array() "Allocate" "Size=$(Base.format_bytes(sz))"
 
     time = Base.@elapsed begin
-        buf = @autoreleasepool MTLBuffer(dev, sz, args...; storage, kwargs...)
+        buf = @autoreleasepool MTLBuffer(dev, sz, args...; kwargs...)
     end
 
     Base.@atomic alloc_stats.alloc_count + 1

test/array.jl

@@ -72,7 +72,7 @@ check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
     N = length(dim)
 
     # mtl
-    let arr = mtl(rand(2,2), storage= SM)
+    let arr = mtl(rand(2,2); storage= SM)
         @test check_storagemode(arr,  SM)
     end
 
@@ -134,7 +134,7 @@ check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
         end
 
         let b = rand(Float32, 10)
-            arr_mtl = mtl(b, storage=Private)
+            arr_mtl = mtl(b; storage=Private)
             @test_throws ErrorException arr_mtl[1]
             @test Metal.@allowscalar arr_mtl[1] == b[1]
         end
@@ -145,7 +145,7 @@ check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
         end
 
         let b = rand(Float32, 10)
-            arr_mtl = mtl(b, storage=Shared)
+            arr_mtl = mtl(b; storage=Shared)
             @test arr_mtl[1] == b[1]
         end
     elseif SM == Metal.Managed
@@ -155,7 +155,7 @@ check_storagemode(arr, smode) = Metal.storagemode(arr) == smode
         end
 
         let b = rand(Float32, 10)
-            arr_mtl = mtl(b, storage=Managed)
+            arr_mtl = mtl(b; storage=Managed)
             @test arr_mtl[1] == b[1]
         end
     end

test/metal.jl

@@ -454,7 +454,7 @@ end
         arr = Metal.zeros(T, 4)
 
         buf = Base.unsafe_convert(MTL.MTLBuffer, arr)
-        Metal.unsafe_fill!(current_device(), Metal.MtlPointer{T}(buf, 0), T(val), 4)
+        Metal.unsafe_fill!(current_device(), Metal.MtlPtr{T}(buf, 0), T(val), 4)
 
         @test all(Array(arr) .== val)
     end