Fix type instability in KA.allocate (#665)

* Fix type instability

* Better fix

Author: christiangnrd

src/MetalKernels.jl

@@ -22,7 +22,8 @@ The `KernelAbstractions` backend for running on Metal GPUs.
 struct MetalBackend <: KA.GPU
 end
 
-KA.allocate(::MetalBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T = MtlArray{T, length(dims), unified ? SharedStorage : DefaultStorageMode}(undef, dims)
+# Ensure type stability. See JuliaGPU/KernelAbstractions#634
+@inline KA.allocate(::MetalBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T = MtlArray{T, length(dims), unified ? SharedStorage : DefaultStorageMode}(undef, dims)
 KA.zeros(::MetalBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T = Metal.zeros(T, dims; storage=unified ? SharedStorage : DefaultStorageMode)
 KA.ones(::MetalBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T = Metal.ones(T, dims; storage=unified ? SharedStorage : DefaultStorageMode)