Fix strided matmul & minor cleanup (#677)

Author: pxl-th

src/AMDGPU.jl

@@ -73,7 +73,9 @@ include("utils.jl")
 
 include(joinpath("hsa", "HSA.jl"))
 include(joinpath("hip", "HIP.jl"))
-import .HIP: HIPContext, HIPDevice, HIPStream
+
+using .HIP
+using .HIP: HIPContext, HIPDevice, HIPStream
 export HIPContext, HIPDevice, HIPStream
 
 include("cache.jl")

src/array.jl

@@ -5,21 +5,19 @@ mutable struct ROCArray{T, N, B} <: AbstractGPUArray{T, N}
 
     function ROCArray{T, N, B}(
         ::UndefInitializer, dims::Dims{N},
-    ) where {T, N, B <: Union{Mem.HIPBuffer, Mem.HostBuffer}}
+    ) where {T, N, B <: Mem.AbstractAMDBuffer}
         @assert isbitstype(T) "ROCArray only supports bits types"
         data = DataRef(pool_free, pool_alloc(B, prod(dims) * sizeof(T)))
         xs = new{T, N, B}(data, dims, 0)
-        finalizer(unsafe_free!, xs)
-        return xs
+        return finalizer(unsafe_free!, xs)
     end
 
     function ROCArray{T, N}(
         buf::DataRef{Managed{B}}, dims::Dims{N}; offset::Integer = 0,
-    ) where {T, N, B <: Union{Mem.HIPBuffer, Mem.HostBuffer}}
+    ) where {T, N, B <: Mem.AbstractAMDBuffer}
         @assert isbitstype(T) "ROCArray only supports bits types"
         xs = new{T, N, B}(buf, dims, offset)
-        finalizer(unsafe_free!, xs)
-        return xs
+        return finalizer(unsafe_free!, xs)
     end
 end
 
@@ -117,11 +115,8 @@ ROCArray(A::AbstractArray{T,N}) where {T,N} = ROCArray{T,N}(A)
 ROCArray{T}(xs::AbstractArray{S,N}) where {T,N,S} = ROCArray{T,N}(xs)
 (::Type{ROCArray{T,N} where T})(x::AbstractArray{S,N}) where {S,N} = ROCArray{S,N}(x)
 
-# idempotency
 ROCArray{T,N}(xs::ROCArray{T,N}) where {T,N} = copy(xs)
 
-## conversions
-
 Base.convert(::Type{T}, x::T) where T <: ROCArray = x
 
 ## memory operations
@@ -176,7 +171,6 @@ function Base.copyto!(
     dest
 end
 
-# TODO: Workaround for hanging copy() broadcast kernel
 function Base.copy(X::ROCArray{T}) where T
     Xnew = ROCArray{T}(undef, size(X))
     copyto!(Xnew, 1, X, 1, length(X))
@@ -220,10 +214,8 @@ struct Float32Adaptor end
 
 Adapt.adapt_storage(::Float32Adaptor, xs::AbstractArray) =
     isbits(xs) ? xs : convert(ROCArray, xs)
-
 Adapt.adapt_storage(::Float32Adaptor, xs::AbstractArray{<:AbstractFloat}) =
     isbits(xs) ? xs : convert(ROCArray{Float32}, xs)
-
 Adapt.adapt_storage(::Float32Adaptor, xs::AbstractArray{<:Complex{<:AbstractFloat}}) =
     isbits(xs) ? xs : convert(ROCArray{ComplexF32}, xs)
 

src/blas/wrappers.jl

@@ -536,17 +536,18 @@ end
 
 # Level 3
 ## (GE) general matrix-matrix multiplication
-for (fname, elty) in
-        ((:rocblas_dgemm,:Float64),
-         (:rocblas_sgemm,:Float32),
-         (:rocblas_hgemm,:Float16),
-         (:rocblas_zgemm,:ComplexF64),
-         (:rocblas_cgemm,:ComplexF32))
+for (fname, elty) in (
+    (:rocblas_dgemm,:Float64),
+    (:rocblas_sgemm,:Float32),
+    (:rocblas_hgemm,:Float16),
+    (:rocblas_zgemm,:ComplexF64),
+    (:rocblas_cgemm,:ComplexF32),
+)
     @eval begin
         function gemm!(
             transA::Char, transB::Char, alpha::($elty),
-            A::ROCVecOrMat{$elty}, B::ROCVecOrMat{$elty}, beta::($elty),
-            C::ROCVecOrMat{$elty},
+            A::StridedROCVecOrMat{$elty}, B::StridedROCVecOrMat{$elty}, beta::($elty),
+            C::StridedROCVecOrMat{$elty},
         )
             m = size(A, transA == 'N' ? 1 : 2)
             k = size(A, transA == 'N' ? 2 : 1)

src/dnn/batchnorm.jl

@@ -104,10 +104,9 @@ function derive_beta_gamma_descriptors(
     dtype, dims, stride = unpack(handle, ndims(handle))
 
     bndesc = TensorDescriptor(handle, dtype)
-    finalizer(bndesc) do d_
-        miopenDestroyTensorDescriptor(d_.handle)
+    return finalizer(bndesc) do d
+        miopenDestroyTensorDescriptor(d.handle)
     end
-    bndesc
 end
 
 # Unsqueeze dimensions at the beginning:

src/dnn/descriptors.jl

@@ -35,11 +35,9 @@ function TensorDescriptor(
     handle = handle_ref[]
     miopenSetTensorDescriptor(handle, dtype, dims, sizes, strides)
     d = TensorDescriptor(handle, dtype)
-
-    finalizer(d) do d_
-        miopenDestroyTensorDescriptor(d_.handle)
+    return finalizer(d) do d
+        miopenDestroyTensorDescriptor(d.handle)
     end
-    d
 end
 
 function TensorDescriptor(x::ROCArray{T}) where T
@@ -96,11 +94,9 @@ function ConvolutionDescriptor(
         handle, n_dims, padding, stride, dilation, miopenConvolution)
     miopenSetConvolutionGroupCount(handle, groups)
     d = ConvolutionDescriptor(handle)
-
-    finalizer(d) do d_
-        miopenDestroyConvolutionDescriptor(d_.handle)
+    return finalizer(d) do d
+        miopenDestroyConvolutionDescriptor(d.handle)
     end
-    d
 end
 
 """
@@ -150,11 +146,9 @@ function PoolingDescriptor(
         handle, mode, n_dims, dims, padding, stride)
     miopenSetPoolingIndexType(handle, miopenIndexUint32)
     d = PoolingDescriptor(handle)
-
-    finalizer(d) do d_
-        miopenDestroyPoolingDescriptor(d_.handle)
+    return finalizer(d) do d
+        miopenDestroyPoolingDescriptor(d.handle)
     end
-    d
 end
 
 """
@@ -192,10 +186,9 @@ function ActivationDescriptor()
     miopenCreateActivationDescriptor(handle_ref)
     handle = handle_ref[]
     d = ActivationDescriptor(handle)
-    finalizer(d) do d_
-        miopenDestroyActivationDescriptor(d_.handle)
+    return finalizer(d) do d
+        miopenDestroyActivationDescriptor(d.handle)
     end
-    d
 end
 
 function set!(

src/fft/fft.jl

@@ -51,8 +51,7 @@ mutable struct cROCFFTPlan{T, K, inplace, N} <: ROCFFTPlan{T, K, inplace}
             rocfft_execution_info_set_work_buffer(info, workarea, length(workarea))
         end
         p = new(handle, stream, workarea, info, size(X), sizey, xtype, region)
-        finalizer(AMDGPU.unsafe_free!, p)
-        p
+        return finalizer(AMDGPU.unsafe_free!, p)
     end
 end
 
@@ -81,8 +80,7 @@ mutable struct rROCFFTPlan{T,K,inplace,N} <: ROCFFTPlan{T,K,inplace}
             rocfft_execution_info_set_work_buffer(info, workarea, length(workarea))
         end
         p = new(handle, stream, workarea, info, size(X), sizey, xtype, region)
-        finalizer(unsafe_free!, p)
-        p
+        return finalizer(unsafe_free!, p)
     end
 end
 

src/highlevel.jl

@@ -1,10 +1,3 @@
-"""
-    devices()
-
-Get list of all devices.
-"""
-devices() = HIP.devices()
-
 """
     device_id() -> Int
     device_id(device::HIPDevice) -> Int
@@ -22,18 +15,12 @@ Sets the current device to `AMDGPU.devices()[idx]`. See
 """
 device_id!(idx::Integer) = device!(devices()[idx])
 
-# Contexts
-
 function device(context::HIPContext)
     return HIP.context!(context) do
         HIP.device()
     end
 end
 
-# Streams.
-
-default_stream() = HIP.default_stream()
-
 device(stream::HIPStream) = stream.device
 device(idx::Integer) = devices()[idx]
 
@@ -91,12 +78,6 @@ macro sync(ex)
     end
 end
 
-"""
-Blocks until all kernels on all streams have completed.
-Uses currently active device.
-"""
-device_synchronize() = HIP.device_synchronize()
-
 """
     rocconvert(x)
 
@@ -184,11 +165,8 @@ macro roc(ex...)
     end)
 end
 
-function launch_configuration(
-    kern::Runtime.HIPKernel; shmem::Integer = 0, max_block_size::Integer = 0,
-)
+launch_configuration(kern::Runtime.HIPKernel; shmem::Integer = 0, max_block_size::Integer = 0) =
     HIP.launch_configuration(kern.fun; shmem, max_block_size)
-end
 
 """
     @elapsed ex

src/hip/HIP.jl

@@ -1,5 +1,6 @@
 module HIP
 export HIPError
+export devices, device_synchronize, default_stream
 
 using CEnum
 
@@ -74,6 +75,10 @@ include("event.jl")
 include("pool.jl")
 include("module.jl")
 
+"""
+Blocks until all kernels on all streams have completed.
+Uses currently active device.
+"""
 function device_synchronize()
     AMDGPU.maybe_collect(; blocking=true)
     hipDeviceSynchronize()

src/hip/device.jl

@@ -104,6 +104,11 @@ function ndevices()
     count_ref[]
 end
 
+"""
+    devices()
+
+Get list of all devices.
+"""
 function devices()
     isempty(ALL_DEVICES) || return copy(ALL_DEVICES)
 

src/hip/event.jl

@@ -44,25 +44,20 @@ end
 wait(event::HIPEvent) = hipEventSynchronize(event)
 
 function synchronize(event::HIPEvent)
-    if !non_blocking_synchronize(event)
-        AMDGPU.maybe_collect(; blocking=true)
-    end
+    non_blocking_synchronize(event) || AMDGPU.maybe_collect(; blocking=true)
     wait(event)
     return
 end
 
 function HIPEvent(stream::hipStream_t; do_record::Bool = true, timing=false)
     event_ref = Ref{hipEvent_t}()
-    if !timing
+    timing ?
+        hipEventCreate(event_ref) :
         hipEventCreateWithFlags(event_ref, hipEventDisableTiming)
-    else
-        hipEventCreate(event_ref)
-    end
     event = HIPEvent(event_ref[], stream)
     do_record && record(event)
 
-    finalizer(hipEventDestroy, event)
-    return event
+    return finalizer(hipEventDestroy, event)
 end
 HIPEvent(stream::HIPStream; kwargs...) = HIPEvent(stream.stream; kwargs...)
 
@@ -76,5 +71,5 @@ See also [`@elapsed`](@ref).
 function elapsed(start::HIPEvent, stop::HIPEvent)
     time_ref = Ref{Cfloat}()
     hipEventElapsedTime(time_ref, start, stop)
-    return time_ref[]/1000
+    return time_ref[] / 1000
 end

src/hip/module.jl

@@ -8,9 +8,7 @@ mutable struct HIPModule
         mod_ref = Ref{hipModule_t}()
         hipModuleLoadData(mod_ref, data)
         mod = new(mod_ref[])
-
-        finalizer(hipModuleUnload, mod)
-        mod
+        return finalizer(hipModuleUnload, mod)
     end
 end
 

src/hip/stream.jl

@@ -27,13 +27,12 @@ function HIPStream(priority::Symbol = :normal)
     hipStreamCreateWithPriority(stream_ref, 0, priority_int)
     d = device()
     stream = HIPStream(stream_ref[], priority, d, HIPContext(d), true)
-    finalizer(stream) do s
+    return finalizer(stream) do s
         AMDGPU.context!(s.ctx) do
             hipStreamDestroy(s.stream)
         end
         Base.@atomic s.valid = false
     end
-    return stream
 end
 
 isvalid(s::HIPStream) = s.valid

src/rand/random.jl

@@ -8,8 +8,7 @@ mutable struct RNG <: Random.AbstractRNG
         handle = Ref{rocrand_generator}()
         rocrand_create_generator(handle, typ)
         obj = new(handle[], typ)
-        finalizer(unsafe_destroy!, obj)
-        return obj
+        return finalizer(unsafe_destroy!, obj)
     end
 end
 

src/runtime/Runtime.jl

@@ -24,7 +24,7 @@ module Mem
     using Preferences
 
     import AMDGPU
-    import AMDGPU: HIP, HSA, Runtime
+    import AMDGPU: HIP, Runtime
     import .HIP: HIPDevice, HIPContext
     import .Runtime: ROCDim, ROCDim3
 

src/runtime/error.jl

@@ -35,4 +35,3 @@ function check(result::HSA.Status)
         throw(HSAError(result))
     end
 end
-