Expose storage mode in allocator

jvkersch · Nov 16, 2023 · 08ea259 · 08ea259
1 parent e8f16bf
commit 08ea259
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/MetalKernels.jl b/src/MetalKernels.jl
@@ -15,11 +15,13 @@ import Adapt
 export MetalBackend
 
 struct MetalBackend <: KA.GPU
+    storage::MTL.MTLResourceOptions
 end
+MetalBackend(;storage=DefaultStorageMode) = MetalBackend(storage)
 
-KA.allocate(::MetalBackend, ::Type{T}, dims::Tuple) where T = MtlArray{T}(undef, dims)
-KA.zeros(::MetalBackend, ::Type{T}, dims::Tuple) where T = Metal.zeros(T, dims)
-KA.ones(::MetalBackend, ::Type{T}, dims::Tuple) where T = Metal.ones(T, dims)
+KA.allocate(backend::MetalBackend, ::Type{T}, dims::Dims{N}) where {T,N} = MtlArray{T,N,backend.storage}(undef, dims)
+KA.zeros(backend::MetalBackend, ::Type{T}, dims::Tuple) where T = Metal.zeros(T, dims; storage=backend.storage)
+KA.ones(backend::MetalBackend, ::Type{T}, dims::Tuple) where T = Metal.ones(T, dims; storage=backend.storage)
 
 KA.get_backend(::MtlArray) = MetalBackend()
 KA.synchronize(::MetalBackend) = synchronize()

diff --git a/src/array.jl b/src/array.jl
@@ -1,6 +1,6 @@
 # host array
 
-export MtlArray, MtlVector, MtlMatrix, MtlVecOrMat, mtl
+export MtlArray, MtlVector, MtlMatrix, MtlVecOrMat, mtl, DefaultStorageMode
 
 function hasfieldcount(@nospecialize(dt))
     try