Arrays and memory

MtlArray{T,N,S} <: AbstractGPUArray{T,N}

An N-dimensional array of elements of type T in GPU memory with storage mode S, either Metal.SharedStorage or Metal.PrivateStorage. The storage mode defaults to shared, which can be changed with the default_storage preference.

MtlArray{T}(undef, dims...)MtlArray{T,N,S}(undef, dims...)MtlArray(A::AbstractArray)MtlArray{T}(A::AbstractArray)

Allocate an uninitialized array, or copy A to a new array, converting its elements to T when given. Unlike mtl, MtlArray(A) keeps the element type of A. Element types have to be stored inline, and cannot contain Float64, Int128 or UInt128 values, which Metal doesn't support.

julia> a = MtlArray([1f0, 2f0, 3f0])3-element MtlVector{Float32, Metal.SharedStorage}: 1.0 2.0 3.0julia> Array(a .* 2)3-element Vector{Float32}: 2.0 4.0 6.0

MtlVector{T,S}

One-dimensional MtlArray with elements of type T and storage mode S; an alias for MtlArray{T,1,S}.

MtlMatrix{T,S}

Two-dimensional MtlArray with elements of type T and storage mode S; an alias for MtlArray{T,2,S}.

MtlVecOrMat{T,S}

The union of MtlVector{T,S} and MtlMatrix{T,S}, like VecOrMat for Arrays.

mtl(A; storage=Metal.SharedStorage)

Copy A to the GPU, with storage mode storage (Metal.SharedStorage or Metal.PrivateStorage).

This is an opinionated GPU array adaptor, which may alter the element type T of arrays:

• For T<:AbstractFloat, it makes a MtlArray{Float32} for performance and compatibility reasons (except for Float16).
• For T<:Complex{<:AbstractFloat} it makes a MtlArray{ComplexF32}.
• For composite element types (e.g. SVector{2,Float64}), Adapt.jl recursively converts Float64 scalars to Float32 via adapt_structure/adapt_storage dispatch.
• For other isbitstype(T), it makes a MtlArray{T}.

By contrast, MtlArray(A) never changes the element type.

Uses Adapt.jl to act inside some wrapper structs.

Examples

julia> mtl(ones(3)')1×3 adjoint(::MtlVector{Float32, Metal.SharedStorage}) with eltype Float32: 1.0  1.0  1.0julia> mtl(zeros(1,3); storage=Metal.SharedStorage)1×3 MtlMatrix{Float32, Metal.SharedStorage}: 0.0  0.0  0.0julia> mtl(1:3)1:3julia> MtlArray(1:3)3-element MtlVector{Int64, Metal.SharedStorage}: 1 2 3

Metal.zeros(T=Float32, dims...; storage)

Create an MtlArray with element type T and storage mode storage (the default storage mode, unless given), filled with zeros.

Metal.ones(T=Float32, dims...; storage)

Create an MtlArray with element type T and storage mode storage (the default storage mode, unless given), filled with ones.

Metal.fill(x, dims...; storage)

Create an MtlArray with storage mode storage (the default storage mode, unless given), filled with x.

Metal.allowscalar(allow::Bool)

Allow or disallow scalar indexing of private-storage MtlArrays, like GPUArrays.allowscalar. Arrays with shared storage, the default, can always be indexed from the CPU, so disallowing scalar indexing only affects private arrays.

Metal.rand([T=Float32], dims...; storage)

Create an MtlArray of element type T and storage mode storage (the default storage mode, unless given), filled with uniformly distributed random numbers from Metal.default_rng. Without dims, return a single number.

Metal.randn([T=Float32], dims...; storage)

Like Metal.rand, but with normally distributed random numbers.

Metal.seed!([seed])

Seed the random number generators of the current task, Metal.default_rng and Metal.mps_rng, with seed, or with a random seed.

Metal.default_rng()

Return the random number generator that Metal.rand, Metal.randn and rand!/randn! on MtlArrays use: a GPUArrays.jl Philox RNG that generates numbers on the GPU. Every task has its own, per device.

See also Metal.mps_rng.

Metal.mps_rng()

Return the task's MPS.RNG, which generates random numbers with Metal Performance Shaders. It is seeded together with Metal.default_rng by Metal.seed!.

Metal.matmul_alg

A ScopedValue that selects the algorithm LinearAlgebra.mul! (and *) uses for MtlArray matrix-matrix and matrix-vector products:

• :auto (default): MPSGraph where it supports the operands, otherwise :native.
• :MPSGraph: MPSGraphs.matmul! and MPSGraphs.matvecmul!.
• :MPS: MPS.matmul! and MPS.matvecmul!.
• :native: the best of Metal.jl's own kernels for the operands, of :tensor, :simd and :scalar.
• :tensor: a Metal 4 tensor-operations kernel, for plain C = A*B products of Float16, Float32 or BFloat16 matrices; only matrix-matrix products.
• :simd: a SIMD-group matrix kernel.
• :scalar: a kernel without SIMD-group matrices, for any element type.
• :GPUArrays: the generic GPUArrays.jl kernel.

An algorithm that doesn't support the operands throws an error. Set it for a block of code with with, from Base.ScopedValues (on Julia 1.10, from the ScopedValues.jl package):

using LinearAlgebra, Base.ScopedValuesA = Metal.rand(Float32, 64, 64); B = Metal.rand(Float32, 64, 64)C = with(Metal.matmul_alg => :MPS) do    A * Bend

Metal.SharedStorage

The shared storage mode (MTLStorageModeShared), the default for MtlArrays: the memory is accessible from both the CPU and the GPU. Indexing such an array from the CPU, or wrapping it as an Array, accesses the memory directly, waiting for GPU work that uses the array to finish first.

See Choosing a resource storage mode for Apple GPUs, and Metal.PrivateStorage.

Metal.PrivateStorage

The private storage mode (MTLStorageModePrivate): the memory is only accessible from the GPU, and indexing such an array from the CPU makes a slow copy, subject to allowscalar. Use it for data that only the GPU accesses, or to make sure that operations on an array execute on the GPU. Create private arrays with, e.g., MtlArray{Float32,1,Metal.PrivateStorage}(undef, n), or mtl(A; storage=Metal.PrivateStorage).

See also Metal.SharedStorage.

Metal.ManagedStorage

The managed storage mode (MTLStorageModeManaged), which keeps separate CPU and GPU copies of the memory. Apple silicon GPUs don't need it, and MtlArrays don't support it: use Metal.SharedStorage instead.

is_shared(A::MtlArray)::Bool

Returns true if A has storage mode Metal.SharedStorage.

See also is_private and is_managed.

is_private(A::MtlArray)::Bool

Returns true if A has storage mode Metal.PrivateStorage.

See also is_shared.

is_managed(A::MtlArray)::Bool

Returns true if A has storage mode Metal.ManagedStorage.

See also is_shared and is_private.

unsafe_wrap(MtlArray, a::Array, [dims]; [dev], [hazard_tracking], [cache_mode])
unsafe_wrap(MtlArray{T,N}, a::Array, [dims]; ...)
unsafe_wrap(MtlArray, ptr::Ptr{T}, dims; ...)

Wrap an MtlArray around host memory, without copying, so that it can be used on the GPU, e.g., in kernels or broadcasts. Apple silicon has unified memory, so the GPU operates directly on the original memory, and changes are visible from both sides. The array has shared storage.

The keyword arguments select the device (dev, defaulting to device()) and are passed to the Metal buffer: hazard_tracking and cache_mode, which take MTLResourceOptions values from Metal.MTL.

When wrapping an Array, the returned MtlArray keeps it alive, until the GPU is done using it. When wrapping a pointer, the caller has to make sure the memory stays valid for as long as the MtlArray is used. In both cases, the memory must not be freed or reallocated while it is wrapped (e.g., by calling resize! on the original array), and resizing the wrapper detaches it from the original memory.

GPU operations execute asynchronously, so synchronize (e.g., using Metal.synchronize()) before accessing the original memory on the host.

a = rand(Float32, 1024)b = unsafe_wrap(MtlArray, a)b .= sin.(b)        # executes on the GPU, updating `a`Metal.synchronize()

unsafe_wrap(Array, arr::MtlArray, [dims])

Wrap a Julia Array around the memory that backs arr, without copying. This is only possible for arrays with SharedStorage, including host memory that was itself wrapped using unsafe_wrap(MtlArray, ...).

Wrapping waits for pending GPU operations on arr. GPU operations execute asynchronously, so synchronize again (e.g., using Metal.synchronize()) before accessing the returned array after using arr on the GPU.

resize!(a::MtlVector, n::Integer)

Resize a to contain n elements. If n is smaller than the current collection length, the first n elements will be retained. If n is larger, the new elements are not guaranteed to be initialized.

total_memory(dev=device())

Return the amount of memory, in bytes, that the GPU dev can use without affecting its performance: Metal's recommendedMaxWorkingSetSize. Apple GPUs share system memory with the CPU, so this is a part of it, not dedicated memory.

See also free_memory.

free_memory(dev=device())

Return the amount of memory, in bytes, that remains of total_memory after the allocations of this process on dev.

Metal.@allocated ex

Evaluate ex, discarding its value, and return the number of bytes of GPU memory that Metal.jl allocated during the evaluation.

Metal.@time ex

Evaluate ex, and print the time it took and the CPU and GPU memory it allocated, like Base.@time. The GPU is synchronized right before and after evaluating ex, so the time includes the GPU work of ex and excludes earlier work.

alloc(device, bytesize, [ptr]; storage=PrivateStorage, hazard_tracking, cache_mode)

Allocate an MTLBuffer of bytesize bytes on device, initialized with the memory at the CPU pointer ptr if given, or with zeros otherwise. Free it with free. storage is Metal.PrivateStorage or Metal.SharedStorage; a buffer initialized from ptr cannot be private. hazard_tracking and cache_mode take MTLResourceOptions values from Metal.MTL.

Without a ptr, the buffer is padded to whole 4-byte words (and at least one): Metal does not support empty buffers, and GPUCompiler implements 8- and 16-bit atomics on the containing 32-bit word. Buffers initialized from, or wrapping, host memory must be non-empty.

free(buffer::MTLBuffer)

Free a buffer allocated with alloc. Freeing a buffer twice is undefined behavior.