Test all storage modes (#314)

Includes workarounds for known bugs that still need to be investigated.

Author: christiangnrd

.buildkite/pipeline.yml

@@ -39,6 +39,31 @@ steps:
                 julia: "1.11"
               soft_fail: true
 
+  # Test Storage modes Shared and Managed
+  - group: ":floppy_disk: Storage mode"
+    key: "storage"
+    steps:
+      - label: "MtlArray with {{matrix.storage}} storage"
+        plugins:
+          - JuliaCI/julia#v1:
+              version: "1.10"
+          - JuliaCI/julia-test#v1:
+              test_args: "--quickfail"
+        agents:
+          queue: "juliaecosystem"
+          os: "macos"
+          arch: "aarch64"
+        if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip storage\]/ && !build.pull_request.draft
+        timeout_in_minutes: 60
+        matrix:
+          setup:
+            storage:
+              - "Shared"
+              - "Managed"
+        commands: |
+          echo -e "[Metal]\ndefault_storage = \"{{matrix.storage}}\"" >LocalPreferences.toml
+
+
   # special tests
   - group: ":eyes: Special"
     depends_on: "julia"

.gitignore

@@ -1,5 +1,4 @@
 Manifest.toml
-LocalPreferences.toml
 *.DS_Store
 *.gputrace
 *.trace

LocalPreferences.toml

@@ -0,0 +1,4 @@
+[Metal]
+# which storage mode unspecified allocations should default to.
+# possible values: "Private", "Shared", "Managed"
+#default_storage = "Private"

Project.toml

@@ -17,6 +17,7 @@ LLVMDowngrader_jll = "f52de702-fb25-5922-94ba-81dd59b07444"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ObjectFile = "d8793406-e978-5875-9003-1fc021f44a92"
 ObjectiveC = "e86c9b32-1129-44ac-8ea0-90d5bb39ded9"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Python_jll = "93d3a430-8e7c-50da-8e8d-3dfcfb3baf05"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -26,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"
@@ -45,11 +52,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"

src/Metal.jl

@@ -7,6 +7,7 @@ using GPUCompiler
 using LLVM
 using LLVM.Interop
 import LLVMDowngrader_jll
+using Preferences: @load_preference, load_preference
 using Python_jll
 using ObjectFile
 using ExprTools: splitdef, combinedef

src/array.jl

@@ -110,7 +110,18 @@ const MtlMatrix{T,S} = MtlArray{T,2,S}
 const MtlVecOrMat{T,S} = Union{MtlVector{T,S},MtlMatrix{T,S}}
 
 # default to private memory
-const DefaultStorageMode = Private
+const DefaultStorageMode = let str = @load_preference("default_storage", "Private")
+  if str == "Private"
+    Private
+  elseif str == "Shared"
+    Shared
+  elseif str == "Managed"
+    Managed
+  else
+    error("unknown default storage mode: $default_storage")
+  end
+end
+
 MtlArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N} =
   MtlArray{T,N,DefaultStorageMode}(undef, dims)
 
@@ -170,14 +181,16 @@ end
 
 function Base.unsafe_convert(::Type{MtlPointer{T}}, x::MtlArray) where {T}
    buf = x.data[]
+   synchronize()
    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
+  synchronize()
+  buf = x.data[]
   convert(Ptr{T}, buf) + x.offset*Base.elsize(x)
 end
 
@@ -237,7 +250,7 @@ Base.convert(::Type{T}, x::T) where T <: MtlArray = x
 Base.unsafe_convert(::Type{<:Ptr}, x::MtlArray) =
   throw(ArgumentError("cannot take the host address of a $(typeof(x))"))
 
-Base.unsafe_convert(t::Type{MTL.MTLBuffer}, x::MtlArray) = x.data[]
+Base.unsafe_convert(::Type{MTL.MTLBuffer}, x::MtlArray) = x.data[]
 
 
 ## interop with ObjC libraries

src/utilities.jl

@@ -48,6 +48,19 @@ function versioninfo(io::IO=stdout)
         println(io)
     end
 
+    prefs = [
+        "default_storage" => load_preference(Metal, "default_storage"),
+    ]
+    if any(x->!isnothing(x[2]), prefs)
+        println(io, "Preferences:")
+        for (key, val) in prefs
+            if !isnothing(val)
+                println(io, "- $key: $val")
+            end
+        end
+        println(io)
+    end
+
     devs = devices()
     if isempty(devs)
         println(io, "No Metal devices.")

test/array.jl

@@ -1,4 +1,4 @@
-STORAGEMODES = [Private, Shared]#, Managed]
+STORAGEMODES = [Private, Shared, Managed]
 
 @testset "array" begin
 
@@ -43,9 +43,9 @@ end
     end
 
     # test the regular adaptor
-    @test Adapt.adapt(MtlArray, [1 2;3 4]) isa MtlArray{Int, 2, Private}
-    @test Adapt.adapt(MtlArray{Float32}, [1 2;3 4]) isa MtlArray{Float32, 2, Private}
-    @test Adapt.adapt(MtlArray{Float32, 2}, [1 2;3 4]) isa MtlArray{Float32, 2, Private}
+    @test Adapt.adapt(MtlArray, [1 2;3 4]) isa MtlArray{Int, 2, Metal.DefaultStorageMode}
+    @test Adapt.adapt(MtlArray{Float32}, [1 2;3 4]) isa MtlArray{Float32, 2, Metal.DefaultStorageMode}
+    @test Adapt.adapt(MtlArray{Float32, 2}, [1 2;3 4]) isa MtlArray{Float32, 2, Metal.DefaultStorageMode}
     @test Adapt.adapt(MtlArray{Float32, 2, Shared}, [1 2;3 4]) isa MtlArray{Float32, 2, Shared}
     @test Adapt.adapt(MtlMatrix{ComplexF32, Shared}, [1 2;3 4]) isa MtlArray{ComplexF32, 2, Shared}
     @test Adapt.adapt(MtlArray{Float16}, Float64[1]) isa MtlArray{Float16}

test/runtests.jl

@@ -75,6 +75,10 @@ for (rootpath, dirs, files) in walkdir(@__DIR__)
 end
 ## GPUArrays testsuite
 for name in keys(TestSuite.tests)
+    if Metal.DefaultStorageMode != Private && name == "indexing scalar"
+        # GPUArrays' scalar indexing tests assume that indexing is not supported
+        continue
+    end
     push!(tests, "gpuarrays$(Base.Filesystem.path_separator)$name")
     test_runners["gpuarrays$(Base.Filesystem.path_separator)$name"] = ()->TestSuite.tests[name](MtlArray)
 end