Metal Performance Shaders Graph

MPSGraphs

Operations on MtlArrays implemented with Apple's Metal Performance Shaders Graph framework, such as MPSGraphs.reduce! and MPSGraphs.matmul!, and low-level wrappers of the framework itself, which follow Apple's API closely.

MPSGraphs.supports(f, args...; kwargs...)::Bool

Return whether the MPSGraph operation f supports being called as f(args...; kwargs...): whether the arguments' types, shapes, element types, memory layout and options are ones that f handles on the current device. Calling f with arguments for which this returns false throws an ArgumentError or DimensionMismatch.

The check only inspects the arguments; it does not build or run any graph. It can be used to decide whether to use an MPSGraph operation or fall back to another implementation, without duplicating the operation's requirements.

if MPSGraphs.supports(MPSGraphs.reduce!, +, R, A)    MPSGraphs.reduce!(+, R, A)else    sum!(R, A)end

MPSGraphs.reduce!(op, R, A; init=nothing) -> R

Reduce A with op along every dimension where R has extent 1 and A does not, and store the result in R, overwriting its contents.

op is one of +, *, max and min (or Base.add_sum and Base.mul_prod). R and A are MtlArrays of the same rank and the same real element type T: Float32, Float16, BFloat16, or a signed or unsigned integer of 8 to 64 bits. Each dimension of R has the extent of A in that dimension, or 1. R must not overlap A.

If init is given, it must be a T, and every element of R is op(init, x), where x is the reduction of the corresponding elements of A; init is applied exactly once. If A is empty along a reduced dimension, the result is init, or without init zero for + and one for * (max and min then require init).

The results can differ from Base.reduce:

• floating-point sums and products are associated differently, and may round differently;
• integer sums and products wrap around on overflow in T; they are not widened;
• max and min propagate NaN like Base, but the sign of a zero result is unspecified (max(0.0, -0.0) may return -0.0), and a sum of negative zeros may return 0.0.

The reduction executes asynchronously, ordered with the other work submitted from the current task (kernels, copies, other MPSGraph operations).

Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.accumulate!(op, B, A; dims=nothing, init=nothing, inclusive=true, rev=false) -> B

Compute the cumulative op of A along dimension dims, and store it in B. dims can only be omitted for vectors.

op is + or * (or Base.add_sum and Base.mul_prod), and for integers also max or min. B and A are MtlArrays of the same size and the same real element type T: Float32, Float16, BFloat16, or a signed or unsigned integer of 8 to 64 bits. B may be A, but must not otherwise overlap it.

With inclusive=true, element i along dims combines elements 1 through i of A; with inclusive=false it combines elements 1 through i-1, so the first is the identity of op. rev=true accumulates from the end instead. If init is given, it must be a T, and every result x becomes op(init, x); for an exclusive scan, the first element is init.

As with MPSGraphs.reduce!, floating-point results can round differently from Base, integer results wrap around on overflow, and the operation executes asynchronously, ordered with the other work submitted from the current task.

Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.sort!(B, A; dims=nothing, rev=false) -> B

Sort A along dimension dims in ascending order (descending with rev=true), and store the result in B. dims can only be omitted for vectors.

B and A are MtlArrays of the same size and the same real element type. B must not overlap A. Elements are ordered with isless, so NaNs sort last.

The sort executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.sortperm!(I, A; dims=nothing, rev=false) -> I

Store in I the linear indices that sort A along dimension dims, in ascending order (descending with rev=true), like sortperm(A; dims). dims can only be omitted for vectors.

A is an MtlArray with a real element type, and I an MtlArray{Int32} or MtlArray{Int64} of the same size. Elements are ordered with isless, so NaNs sort last. The permutation is not guaranteed to be stable: equal elements can appear in any order.

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.matmul!(C, A, B, α=true, β=false; tA='N', tB='N') -> C

Compute C = α * opA(A) * opB(B) + β * C, where opA and opB are the identity ('N'), transpose ('T') or adjoint ('C'), as selected by tA and tB.

A, B and C are MtlArray matrices, or 3D arrays holding equally many matrices along the last dimension, which are multiplied pairwise. A and B have the same element type, which together with that of C is one of the combinations in MPSGraphs.MPSGRAPH_VALID_MATMUL_TYPES; integer inputs are multiplied in the output type. α and β are converted to the element types of A and C. As with BLAS, C is not read when β is zero. C must not overlap A or B.

The multiplication executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.matvecmul!(y, A, x, α=true, β=false; tA='N') -> y

Compute y = α * opA(A) * x + β * y, where opA is the identity ('N'), transpose ('T') or adjoint ('C'), as selected by tA.

A is an MtlArray matrix and x and y are MtlArray vectors. A and x have the same element type, which together with that of y is one of the combinations in MPSGraphs.MPSGRAPH_VALID_MATVECMUL_TYPES. Otherwise, this behaves like MPSGraphs.matmul!.

MPSGraphs.MPSGRAPH_VALID_MATMUL_TYPES

The combinations of input element type (of A and B) and output element type (of C), as (input, output) tuples, that MPSGraphs.matmul! supports.

MPSGraphs.MPSGRAPH_VALID_MATVECMUL_TYPES

The combinations of input element type (of A and x) and output element type (of y), as (input, output) tuples, that MPSGraphs.matvecmul! supports.

MPSGraphs.softmax!(y, x; dims=1) -> y

Compute the softmax exp.(x) ./ sum(exp.(x); dims) along dimension dims of x, and store it in y.

x and y are MtlArrays of the same size with element type Float16 or Float32. The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.logsoftmax!(y, x; dims=1) -> y

Compute the log-softmax x .- log.(sum(exp.(x); dims)) along dimension dims of x, and store it in y. The requirements of MPSGraphs.softmax! apply.

MPSGraphs.softmax_grad!(dx, dy, y; dims=1) -> dx

Compute the gradient dx of MPSGraphs.softmax! with respect to its input, given its output y and the gradient dy with respect to that output. All arrays have the same size, and the requirements of softmax! apply.

MPSGraphs.logsoftmax_grad!(dx, dy, y; dims=1) -> dx

Compute the gradient dx of MPSGraphs.logsoftmax! with respect to its input, given its output y and the gradient dy with respect to that output. All arrays have the same size, and the requirements of softmax! apply.

MPSGraphs.conv!(y, x, w; stride=(1, 1), padding=(0, 0, 0, 0), dilation=(1, 1),
                groups=1, flipkernel=false) -> y

Compute the 2D convolution of x with the kernel w, and store it in y.

The arrays use NNlib's layout: x and y are (width, height, channels, batch), and the kernel w is (width, height, input channels ÷ groups, output channels). padding is (left, right, top, bottom). By default the kernel is flipped, i.e. this computes a convolution rather than a cross-correlation; flipkernel=true computes the latter. All arrays are 4D MtlArrays with element type Float16 or Float32.

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.conv_data_grad!(dx, dy, w; stride=(1, 1), padding=(0, 0, 0, 0),
                          dilation=(1, 1), groups=1, flipkernel=false) -> dx

Compute the gradient dx of MPSGraphs.conv! with respect to its input x, given the gradient dy with respect to its output and the kernel w.

The arrays use NNlib's layout: x and y are (width, height, channels, batch), and the kernel w is (width, height, input channels ÷ groups, output channels). padding is (left, right, top, bottom). By default the kernel is flipped, i.e. this computes a convolution rather than a cross-correlation; flipkernel=true computes the latter. All arrays are 4D MtlArrays with element type Float16 or Float32.

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.conv_filter_grad!(dw, x, dy; stride=(1, 1), padding=(0, 0, 0, 0),
                            dilation=(1, 1), groups=1, flipkernel=false) -> dw

Compute the gradient dw of MPSGraphs.conv! with respect to its kernel, given its input x and the gradient dy with respect to its output.

The arrays use NNlib's layout: x and y are (width, height, channels, batch), and the kernel w is (width, height, input channels ÷ groups, output channels). padding is (left, right, top, bottom). By default the kernel is flipped, i.e. this computes a convolution rather than a cross-correlation; flipkernel=true computes the latter. All arrays are 4D MtlArrays with element type Float16 or Float32.

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.maxpool!(y, x; kernel, stride=kernel, padding=(0, 0, 0, 0),
                   dilation=(1, 1)) -> y

Compute the maximum of x over each 2D window of size kernel, and store it in y.

The arrays use NNlib's layout: x and y are 4D MtlArrays of (width, height, channels, batch) with element type Float16 or Float32. kernel, stride and dilation are (width, height), and padding is (left, right, top, bottom).

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.maxpool_grad!(dx, dy, x; kernel, stride=kernel, padding=(0, 0, 0, 0),
                        dilation=(1, 1)) -> dx

Compute the gradient dx of MPSGraphs.maxpool! with respect to its input x, given the gradient dy with respect to its output.

The arrays use NNlib's layout: x and y are 4D MtlArrays of (width, height, channels, batch) with element type Float16 or Float32. kernel, stride and dilation are (width, height), and padding is (left, right, top, bottom).

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.meanpool!(y, x; kernel, stride=kernel, padding=(0, 0, 0, 0),
                    dilation=(1, 1), count_include_pad=true) -> y

Compute the mean of x over each 2D window of size kernel, and store it in y. With count_include_pad, padding counts towards the number of averaged elements.

The arrays use NNlib's layout: x and y are 4D MtlArrays of (width, height, channels, batch) with element type Float16 or Float32. kernel, stride and dilation are (width, height), and padding is (left, right, top, bottom).

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.

MPSGraphs.meanpool_grad!(dx, dy, x; kernel, stride=kernel, padding=(0, 0, 0, 0),
                         dilation=(1, 1), count_include_pad=true) -> dx

Compute the gradient dx of MPSGraphs.meanpool! with respect to its input x, given the gradient dy with respect to its output.

The arrays use NNlib's layout: x and y are 4D MtlArrays of (width, height, channels, batch) with element type Float16 or Float32. kernel, stride and dilation are (width, height), and padding is (left, right, top, bottom).

The operation executes asynchronously, ordered with the other work submitted from the current task. Use MPSGraphs.supports to check whether a call is supported without throwing.