add support for Metal.jl (#34)

Author: lkdvos

Project.toml

@@ -1,7 +1,7 @@
 name = "StridedViews"
 uuid = "4db3bf67-4bd7-4b4e-b153-31dc3fb37143"
-authors = ["Lukas Devos <lukas.devos@ugent.be>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 version = "0.4.6"
+authors = ["Lukas Devos <ldevos98@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -10,8 +10,9 @@ PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-PtrArrays = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+PtrArrays = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
 
 [extensions]
 StridedViewsAMDGPUExt = "AMDGPU"
@@ -26,6 +27,7 @@ CUDA = "4,5"
 JET = "0.9, 0.10, 0.11"
 JLArrays = "0.3.1"
 LinearAlgebra = "1.6"
+Metal = "1.9.3"
 PackageExtensionCompat = "1"
 PtrArrays = "1.2.0"
 Random = "1.6"
@@ -38,9 +40,10 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 PtrArrays = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Random", "Aqua", "JET", "PtrArrays", "CUDA", "AMDGPU", "JLArrays"]
+test = ["Test", "Random", "Aqua", "JET", "PtrArrays", "CUDA", "AMDGPU", "JLArrays", "Metal"]

ext/StridedViewsMetalExt.jl

@@ -0,0 +1,27 @@
+module StridedViewsMetalExt
+
+using StridedViews
+using Metal
+using Metal: Adapt, MtlPtr
+
+const MtlStridedView{T, N, A <: MtlArray{T}} = StridedView{T, N, A}
+
+function Adapt.adapt_structure(to, A::MtlStridedView)
+    return StridedView(
+        Adapt.adapt_structure(to, parent(A)),
+        A.size, A.strides, A.offset, A.op
+    )
+end
+
+function Base.pointer(x::MtlStridedView{T}) where {T}
+    return Base.unsafe_convert(MtlPtr{T}, pointer(x.parent, x.offset + 1))
+end
+function Base.unsafe_convert(::Type{MtlPtr{T}}, a::MtlStridedView{T}) where {T}
+    return convert(MtlPtr{T}, pointer(a))
+end
+
+function Base.print_array(io::IO, X::MtlStridedView)
+    return Base.print_array(io, Adapt.adapt_structure(Array, X))
+end
+
+end # module

test/runtests.jl

@@ -315,7 +315,7 @@ if !is_buildkite
     end
 end
 
-using CUDA, AMDGPU
+using CUDA, AMDGPU, Metal
 
 if CUDA.functional()
     @testset "CuArrays with StridedView" begin
@@ -344,3 +344,17 @@ if AMDGPU.functional()
         end
     end
 end
+
+if Metal.functional()
+    @testset "MtlArrays with StridedView" begin
+        @testset for T in (Float32, ComplexF32)
+            A = MtlArray(randn(T, 10, 10, 10, 10))
+            @test isstrided(A)
+            B = StridedView(A)
+            @test B isa StridedView
+            Metal.@allowscalar begin
+                @test B == A
+            end
+        end
+    end
+end