Adapt ROCDeviceArray to DenseArray interface (#692)

* Adapt ROCDeviceArray to DenseArray interface

* Disable opaque pointers CI for now

Author: pxl-th

.buildkite/pipeline.yml

@@ -62,26 +62,26 @@ steps:
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"
       JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
-  - label: "Julia 1.11 opaque pointers"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "1.11"
-      - JuliaCI/julia-test#v1:
-      - JuliaCI/julia-coverage#v1:
-          codecov: true
-    agents:
-      queue: "juliagpu"
-      rocm: "*"
-      rocmgpu: "*"
-    if: build.message !~ /\[skip tests\]/
-    command: "julia --project -e 'using Pkg; Pkg.update()'"
-    timeout_in_minutes: 180
-    env:
-      JULIA_LLVM_ARGS: "-opaque-pointers"
-      JULIA_NUM_THREADS: 4
-      JULIA_AMDGPU_CORE_MUST_LOAD: "1"
-      JULIA_AMDGPU_HIP_MUST_LOAD: "1"
-      JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
+  # - label: "Julia 1.11 opaque pointers"
+  #   plugins:
+  #     - JuliaCI/julia#v1:
+  #         version: "1.11"
+  #     - JuliaCI/julia-test#v1:
+  #     - JuliaCI/julia-coverage#v1:
+  #         codecov: true
+  #   agents:
+  #     queue: "juliagpu"
+  #     rocm: "*"
+  #     rocmgpu: "*"
+  #   if: build.message !~ /\[skip tests\]/
+  #   command: "julia --project -e 'using Pkg; Pkg.update()'"
+  #   timeout_in_minutes: 180
+  #   env:
+  #     JULIA_LLVM_ARGS: "-opaque-pointers"
+  #     JULIA_NUM_THREADS: 4
+  #     JULIA_AMDGPU_CORE_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_HIP_MUST_LOAD: "1"
+  #     JULIA_AMDGPU_DISABLE_ARTIFACTS: "1"
 
   - label: "GPU-less environment"
     plugins:

Project.toml

@@ -35,7 +35,7 @@ UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 
 [compat]
 AbstractFFTs = "1.0"
-AcceleratedKernels = "0.1.0"
+AcceleratedKernels = "0.1, 0.2"
 Adapt = "4"
 Atomix = "0.1"
 CEnum = "0.4, 0.5"

src/ROCKernels.jl

@@ -17,8 +17,10 @@ struct ROCBackend <: KA.GPU end
 Adapt.adapt_storage(::ROCBackend, a::Array) = Adapt.adapt(AMDGPU.ROCArray, a)
 Adapt.adapt_storage(::ROCBackend, a::AMDGPU.ROCArray) = a
 Adapt.adapt_storage(::KA.CPU, a::AMDGPU.ROCArray) = convert(Array, a)
-Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) where T =
-    AMDGPU.ROCDeviceArray(a.shape, LLVM.Interop.addrspacecast(Core.LLVMPtr{T,AMDGPU.Device.AS.Constant}, a.ptr))
+function Adapt.adapt_storage(::KA.ConstAdaptor, a::AMDGPU.ROCDeviceArray{T}) where T
+    ptr = LLVM.Interop.addrspacecast(Core.LLVMPtr{T,AMDGPU.Device.AS.Constant}, a.ptr)
+    AMDGPU.ROCDeviceArray(a.dims, ptr)
+end
 
 KA.argconvert(::KA.Kernel{ROCBackend}, arg) = AMDGPU.rocconvert(arg)
 KA.get_backend(::AMDGPU.ROCArray) = ROCBackend()

src/device/gcn/array.jl

@@ -25,14 +25,15 @@ ROCDeviceArray
 # NOTE: we can't support the typical `tuple or series of integer` style
 # construction, because we're currently requiring a trailing pointer argument.
 
-struct ROCDeviceArray{T,N,A} <: AbstractArray{T,N}
-    shape::Dims{N}
+struct ROCDeviceArray{T,N,A} <: DenseArray{T,N}
+    # NOTE: `dims` is a mandatory name for things like `strides(xd)` to work.
+    dims::Dims{N}
     ptr::LLVMPtr{T,A}
     len::Int
 
     # inner constructors, fully parameterized, exact types (ie. Int not <:Integer)
-    function ROCDeviceArray{T,N,A}(shape::Dims{N}, ptr::LLVMPtr{T,A}) where {T,A,N}
-        new(shape, ptr, prod(shape))
+    function ROCDeviceArray{T,N,A}(dims::Dims{N}, ptr::LLVMPtr{T,A}) where {T,A,N}
+        new(dims, ptr, prod(dims))
     end
 end
 
@@ -65,7 +66,8 @@ Base.pointer(a::ROCDeviceArray, i::Integer) =
     pointer(a) + (i - 1) * Base.elsize(a) # TODO use _memory_offset(a, i)
 
 Base.elsize(::Type{<:ROCDeviceArray{T}}) where {T} = sizeof(T)
-Base.size(g::ROCDeviceArray) = g.shape
+Base.size(g::ROCDeviceArray) = g.dims
+Base.sizeof(x::ROCDeviceArray) = Base.elsize(x) * length(x)
 Base.length(g::ROCDeviceArray) = g.len
 
 # conversions
@@ -96,14 +98,14 @@ Base.IndexStyle(::Type{<:ROCDeviceArray}) = Base.IndexLinear()
 # comparisons
 
 Base.isequal(a1::R1, a2::R2) where {R1<:ROCDeviceArray,R2<:ROCDeviceArray} =
-    R1 == R2 && a1.shape == a2.shape && a1.ptr == a2.ptr
+    R1 == R2 && a1.dims == a2.dims && a1.ptr == a2.ptr
 
 # other
 
 Base.show(io::IO, a::ROCDeviceVector) =
     print(io, "$(length(a))-element device array at $(pointer(a))")
 Base.show(io::IO, a::ROCDeviceArray) =
-    print(io, "$(join(a.shape, '×')) device array at $(pointer(a))")
+    print(io, "$(join(a.dims, '×')) device array at $(pointer(a))")
 
 Base.show(io::IO, a::SubArray{T,D,P,I,F}) where {T,D,P<:ROCDeviceVector,I,F} =
     print(io, "$(length(a.indices[1]))-element device array view(::$P at $(pointer(parent(a))), $(a.indices[1])) with eltype $T")
@@ -113,12 +115,12 @@ Base.show(io::IO, a::SubArray{T,D,P,I,F}) where {T,D,P<:ROCDeviceArray,I,F} =
 Base.show(io::IO, a::S) where S<:AnyROCDeviceVector =
     print(io, "$(length(a))-element device array wrapper $S at $(pointer(parent(a)))")
 Base.show(io::IO, a::S) where S<:AnyROCDeviceArray =
-    print(io, "$(join(parent(a).shape, '×')) device array wrapper $S at $(pointer(parent(a)))")
+    print(io, "$(join(parent(a).dims, '×')) device array wrapper $S at $(pointer(parent(a)))")
 
 Base.show(io::IO, mime::MIME"text/plain", a::S) where S<:AnyROCDeviceArray = show(io, a)
 
 @inline function Base.unsafe_view(A::ROCDeviceVector{T}, I::Vararg{Base.ViewIndex,1}) where {T}
-    ptr = pointer(A) + (I[1].start-1)*sizeof(T)
+    ptr = pointer(A) + (I[1].start - 1) * sizeof(T)
     len = I[1].stop - I[1].start + 1
 
     return ROCDeviceArray(len, ptr)

test/device/array.jl

@@ -1,3 +1,19 @@
+@testset "ROCDeviceArray array interface" begin
+    x = ROCArray(zeros(Int, 1, 2, 3, 4))
+    xd = rocconvert(x)
+
+    @test typeof(xd) <: DenseArray
+    @test hasproperty(xd, :dims)
+
+    @test size(xd) == size(x)
+    @test length(xd) == length(x)
+    @test strides(xd) == strides(x)
+    @test Base.elsize(xd) == Base.elsize(x)
+    @test sizeof(xd) == sizeof(x)
+
+    @test Int(pointer(xd)) == Int(pointer(x))
+end
+
 @testset "ROCDeviceArray" begin
     RA = ROCArray(rand(4,4))
     RD = rocconvert(RA)