From bd196a9f99ed4b9e8dc603df459887b2a24e9b7f Mon Sep 17 00:00:00 2001 From: Matt Gallagher Date: Fri, 8 May 2026 05:16:12 +1000 Subject: [PATCH 01/39] Add AdamW bias correction parameter. (#400) --- Source/MLXOptimizers/Optimizers.swift | 80 ++++++++++++++++++----- Tests/MLXTests/OptimizerTests.swift | 38 +++++++++++ skills/mlx-swift/references/optimizers.md | 6 +- 3 files changed, 106 insertions(+), 18 deletions(-) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index fb9c5d963..5d4a08fb1 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -107,6 +107,32 @@ public struct TupleState: Updatable { } } +/// State container for Adam-style optimizers that need first and second moments +/// plus an update step for optional bias correction. +public struct AdamState: Updatable { + let values: (MLXArray, MLXArray) + var step: MLXArray + + init(_ values: (MLXArray, MLXArray), step: MLXArray) { + self.values = values + self.step = step + } + + init(_ a: MLXArray, _ b: MLXArray, step: MLXArray) { + self.values = (a, b) + self.step = step + } + + init(zeros array: MLXArray) { + self.values = (MLXArray.zeros(like: array), MLXArray.zeros(like: array)) + self.step = MLXArray(0) + } + + public func innerState() -> [MLXArray] { + [values.0, values.1, step] + } +} + /// Stochastic gradient descent optimizer. /// /// ### See Also @@ -302,14 +328,14 @@ open class AdaDelta: OptimizerBase { /// The Adam optimizer [1]. /// -/// Our Adam implementation follows the original paper and omits the bias -/// correction in the first and second moment estimates. In detail, +/// By default our Adam implementation omits bias correction in the first and +/// second moment estimates. Set `biasCorrection` to `true` to apply it. In detail, /// /// [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic optimization. ICLR 2015. /// /// ### See Also /// - -open class Adam: OptimizerBase { +open class Adam: OptimizerBase { /// The learning rate public var learningRate: Float @@ -317,41 +343,60 @@ open class Adam: OptimizerBase { public var betas: (Float, Float) = (0.9, 0.999) /// The epsilon added to the denominator to improve numerical stability public var eps: Float = 1e-8 + /// If `true`, apply bias correction to the first and second moments + public var biasCorrection = false /// Initialize the optimizer. /// - Parameters: /// - learningRate: the learning rate /// - betas: coefficients used for computing running averages of the gradient and its square /// - eps: the epsilon added to the denominator to improve numerical stability - public init(learningRate: Float, betas: (Float, Float) = (0.9, 0.999), eps: Float = 1e-8) { + /// - biasCorrection: if `true`, apply bias correction to the first and second moments + public init( + learningRate: Float, betas: (Float, Float) = (0.9, 0.999), eps: Float = 1e-8, + biasCorrection: Bool = false + ) { self.learningRate = learningRate self.betas = betas self.eps = eps + self.biasCorrection = biasCorrection } - override open func newState(parameter: MLXArray) -> TupleState { - TupleState(zeros: parameter) + override open func newState(parameter: MLXArray) -> AdamState { + AdamState(zeros: parameter) } - override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( - MLXArray, TupleState + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: AdamState) -> ( + MLXArray, AdamState ) { let (b1, b2) = betas + var state = state var (m, v) = state.values + state.step = state.step + 1 m = b1 * m + (1 - b1) * gradient v = b2 * v + (1 - b2) * square(gradient) - return (parameter - learningRate * m / (sqrt(v) + eps), TupleState(m, v)) + let update: MLXArray + if biasCorrection { + let step = state.step + let c1 = learningRate / (1 - pow(b1, step)) + let c2 = rsqrt(1 - pow(b2, step)) + update = (c1 * m) / (sqrt(v) * c2 + eps) + } else { + update = learningRate * m / (sqrt(v) + eps) + } + + return (parameter - update, AdamState(m, v, step: state.step)) } } /// The AdamW optimizer [1]. /// -/// Following the above convention, in contrast with [1], we do not use bias -/// correction in the first and second moments for AdamW. We update the weights -/// with a `weightDecay` lambda value: +/// By default AdamW omits bias correction in the first and second moments, to +/// match the longstanding MLX Swift behavior. Set `biasCorrection` to `true` +/// to apply it. We update the weights with a `weightDecay` lambda value: /// /// [1]: Loshchilov, I. and Hutter, F., 2019. Decoupled weight decay regularization. ICLR 2019. /// @@ -368,16 +413,19 @@ open class AdamW: Adam { /// - betas: coefficients used for computing running averages of the gradient and its square /// - eps: the epsilon added to the denominator to improve numerical stability /// - weightDecay:the weight decay + /// - biasCorrection: if `true`, apply bias correction to the first and second moments public init( learningRate: Float, betas: (Float, Float) = (0.9, 0.999), eps: Float = 1e-8, - weightDecay: Float = 0.01 + weightDecay: Float = 0.01, biasCorrection: Bool = false ) { self.weightDecay = weightDecay - super.init(learningRate: learningRate, betas: betas, eps: eps) + super.init( + learningRate: learningRate, betas: betas, eps: eps, + biasCorrection: biasCorrection) } - override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: TupleState) -> ( - MLXArray, TupleState + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: AdamState) -> ( + MLXArray, AdamState ) { return super.applySingle( gradient: gradient, parameter: parameter * (1 - learningRate * weightDecay), diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index 74eda3b00..d340d46fd 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -151,12 +151,50 @@ class OptimizerTests: XCTestCase { checkTrain(optimizer: Adam(learningRate: 0.1), compile: true) } + func testAdamBiasCorrection() { + let parameter = MLXArray([1.0 as Float, -2.0, 3.0]) + let gradient = MLXArray([0.5 as Float, -0.25, 2.0]) + let optimizer = Adam(learningRate: 0.1 as Float, biasCorrection: true) + + let result = optimizer.applySingle( + gradient: gradient, parameter: parameter, + state: optimizer.newState(parameter: parameter)) + + let step = MLXArray(1.0 as Float) + let c1 = Float(0.1) / (1 - pow(Float(0.9), step)) + let c2 = rsqrt(1 - pow(Float(0.999), step)) + let m = (1 - Float(0.9)) * gradient + let v = (1 - Float(0.999)) * square(gradient) + let expected = parameter - (c1 * m) / (sqrt(v) * c2 + Float(1e-8)) + assertEqual(result.0, expected, atol: 1e-6) + } + func testAdamW() { checkShape(optimizer: AdamW(learningRate: 0.1)) checkTrain(optimizer: AdamW(learningRate: 0.1)) checkTrain(optimizer: AdamW(learningRate: 0.1), compile: true) } + func testAdamWBiasCorrection() { + let parameter = MLXArray([1.0 as Float, -2.0, 3.0]) + let gradient = MLXArray([0.5 as Float, -0.25, 2.0]) + let optimizer = AdamW( + learningRate: 0.1 as Float, weightDecay: 0.01 as Float, biasCorrection: true) + + let result = optimizer.applySingle( + gradient: gradient, parameter: parameter, + state: optimizer.newState(parameter: parameter)) + + let decayed = parameter * (1 - Float(0.1) * Float(0.01)) + let step = MLXArray(1.0 as Float) + let c1 = Float(0.1) / (1 - pow(Float(0.9), step)) + let c2 = rsqrt(1 - pow(Float(0.999), step)) + let m = (1 - Float(0.9)) * gradient + let v = (1 - Float(0.999)) * square(gradient) + let expected = decayed - (c1 * m) / (sqrt(v) * c2 + Float(1e-8)) + assertEqual(result.0, expected, atol: 1e-6) + } + func testAdamax() { checkShape(optimizer: Adamax(learningRate: 0.1)) checkTrain(optimizer: Adamax(learningRate: 0.1)) diff --git a/skills/mlx-swift/references/optimizers.md b/skills/mlx-swift/references/optimizers.md index 28de697c1..73858f457 100644 --- a/skills/mlx-swift/references/optimizers.md +++ b/skills/mlx-swift/references/optimizers.md @@ -80,7 +80,8 @@ Adaptive Moment Estimation: let optimizer = Adam( learningRate: 0.001, betas: (0.9, 0.999), // (β1, β2) for moment estimates - eps: 1e-8 // Numerical stability + eps: 1e-8, // Numerical stability + biasCorrection: false // Optional bias correction ) ``` @@ -101,7 +102,8 @@ let optimizer = AdamW( learningRate: 0.001, betas: (0.9, 0.999), eps: 1e-8, - weightDecay: 0.01 // Weight decay factor + weightDecay: 0.01, // Weight decay factor + biasCorrection: false // Optional bias correction ) ``` From 89cece7dcb513a4f3a20a83ece3767ef23665887 Mon Sep 17 00:00:00 2001 From: Vernon Stinebaker Date: Fri, 8 May 2026 07:09:10 +0800 Subject: [PATCH 02/39] fix: check mlx_detail_compile and mlx_closure_apply return values (#398) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit mlx_detail_compile and mlx_closure_apply both return int (0=success, 1=failure) but their return values were silently ignored. When an error fires inside a withError scope the MLX error handler stores the error in an ErrorBox instead of calling fatalError; execution then continues past the failed call, innerCall returns an empty result vector, and the single/two/three-array compile overloads crash with a Swift 'Index out of range' trap — bypassing withError entirely. Fix: capture both return values and early-return [] from innerCall on failure. The placeholder return from the compile overloads is never observed by the caller because withError throws before the value is used. Adds three regression tests covering the single-array, two-array, and [MLXArray]->[MLXArray] compile overloads. --- Source/MLX/Transforms+Compile.swift | 28 ++++++++++--- Tests/MLXTests/TransformTests.swift | 61 +++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 5 deletions(-) diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index 1423f848f..e52835380 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -88,12 +88,20 @@ final class CompiledFunction: @unchecked (Sendable) { // but will be able to re-evaluate with fresh state if needed evalLock.lock() var compiled = mlx_closure_new() - mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0) + let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0) defer { mlx_closure_free(compiled) evalLock.unlock() } + // mlx_error was already dispatched on failure: + // • outside withError — fatalError was called; we won't reach here + // • inside withError — error is stored in the ErrorBox; return [] so + // withError can throw instead of crashing downstream + guard compileStatus == 0 else { + return [] + } + let innerInputs = arguments + stateInputs let innerInputsVector = new_mlx_vector_array(innerInputs) defer { mlx_vector_array_free(innerInputsVector) } @@ -101,9 +109,13 @@ final class CompiledFunction: @unchecked (Sendable) { // will compile the function (if needed) and evaluate the // compiled graph var resultVector = mlx_vector_array_new() - mlx_closure_apply(&resultVector, compiled, innerInputsVector) + let applyStatus = mlx_closure_apply(&resultVector, compiled, innerInputsVector) defer { mlx_vector_array_free(resultVector) } + guard applyStatus == 0 else { + return [] + } + let resultsPlusStateOutput = mlx_vector_array_values(resultVector) // push the stateOutput into the state @@ -164,7 +176,11 @@ public func compile( } return { a in - compileState.call([a])[0] + let r = compileState.call([a]) + // r is empty only when an MLX error fired inside a withError scope — the + // error is already stored in the ErrorBox. Return a placeholder so that + // withError can throw instead of crashing with "Index out of range". + return r.isEmpty ? MLXArray(0) : r[0] } } @@ -185,7 +201,8 @@ public func compile( } return { a, b in - compileState.call([a, b])[0] + let r = compileState.call([a, b]) + return r.isEmpty ? MLXArray(0) : r[0] } } @@ -206,7 +223,8 @@ public func compile( } return { a, b, c in - compileState.call([a, b, c])[0] + let r = compileState.call([a, b, c]) + return r.isEmpty ? MLXArray(0) : r[0] } } diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift index fbfbaac5e..d06999c7d 100644 --- a/Tests/MLXTests/TransformTests.swift +++ b/Tests/MLXTests/TransformTests.swift @@ -412,4 +412,65 @@ class TransformTests: XCTestCase { // Note: OptimizerTests contains additional integration tests of compile() + func testCompileSingleArrayErrorPropagatesViaWithError() throws { + // A compiled function that produces a broadcast shape mismatch at evaluation time. + // The constant has shape [3]; passing an input with shape [2] triggers an MLX error + // inside mlx_closure_apply. + // + // Before fix: mlx_closure_apply's non-zero return value was silently ignored, + // causing innerCall to return [], and the single-array overload to crash with + // "Fatal error: Index out of range" — a Swift trap that bypasses withError. + // + // After fix: innerCall returns [] early on error, the overload returns a placeholder, + // and withError properly surfaces the MLXError. + let compiled = compile { (x: MLXArray) -> MLXArray in + let constant = MLXArray([Float](repeating: 1, count: 3)) + return x + constant + } + let x = MLXArray([Float](repeating: 0, count: 2)) + + XCTAssertThrowsError( + try withError { + eval(compiled(x)) + } + ) { error in + XCTAssertTrue(error is MLXError, "expected MLXError, got \(error)") + } + } + + func testCompileMultiArrayErrorPropagatesViaWithError() throws { + // Same scenario with the [MLXArray] -> [MLXArray] overload. + let compiled = compile { (inputs: [MLXArray]) -> [MLXArray] in + let constant = MLXArray([Float](repeating: 1, count: 3)) + return [inputs[0] + constant] + } + let x = MLXArray([Float](repeating: 0, count: 2)) + + XCTAssertThrowsError( + try withError { + eval(compiled([x])) + } + ) { error in + XCTAssertTrue(error is MLXError, "expected MLXError, got \(error)") + } + } + + func testCompileTwoArrayErrorPropagatesViaWithError() throws { + // Same scenario with the (MLXArray, MLXArray) -> MLXArray overload. + let compiled = compile { (a: MLXArray, _: MLXArray) -> MLXArray in + let constant = MLXArray([Float](repeating: 1, count: 3)) + return a + constant + } + let x = MLXArray([Float](repeating: 0, count: 2)) + let y = MLXArray([Float](repeating: 0, count: 2)) + + XCTAssertThrowsError( + try withError { + eval(compiled(x, y)) + } + ) { error in + XCTAssertTrue(error is MLXError, "expected MLXError, got \(error)") + } + } + } From 058eda67e964b3c00960bd4f57a567b78181f1a6 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Thu, 21 May 2026 13:13:30 -0700 Subject: [PATCH 03/39] hold evalLock while computing tostring (#410) * hold evalLock while computing tostring - tostring internally calls eval, which is not thread safe * remove show-sdk-version -- ci breakage --- .github/workflows/pull_request.yml | 1 - Source/MLX/Device.swift | 4 +++- Source/MLX/MLXArray.swift | 4 +++- Source/MLX/Stream.swift | 4 +++- 4 files changed, 9 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index c00c2c19c..de86d599e 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -103,7 +103,6 @@ jobs: shell: sh run: | xcodebuild -version - xcrun --show-sdk-build-version swift --version rm -rf ~/Library/Developer/Xcode/DerivedData/* xcodebuild build-for-testing -scheme mlx-swift-Package -destination 'platform=macOS' diff --git a/Source/MLX/Device.swift b/Source/MLX/Device.swift index 3a262c553..203c1e3f8 100644 --- a/Source/MLX/Device.swift +++ b/Source/MLX/Device.swift @@ -173,8 +173,10 @@ public final class Device: @unchecked Sendable, Equatable { extension Device: CustomStringConvertible { public var description: String { var s = mlx_string_new() - mlx_device_tostring(&s, ctx) defer { mlx_string_free(s) } + _ = evalLock.withLock { + mlx_device_tostring(&s, ctx) + } return String(cString: mlx_string_data(s), encoding: .utf8)! } } diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift index 44014aad5..9e6148e00 100644 --- a/Source/MLX/MLXArray.swift +++ b/Source/MLX/MLXArray.swift @@ -616,8 +616,10 @@ extension MLXArray: Updatable, Evaluatable { extension MLXArray: CustomStringConvertible { public var description: String { var s = mlx_string_new() - mlx_array_tostring(&s, ctx) defer { mlx_string_free(s) } + _ = evalLock.withLock { + mlx_array_tostring(&s, ctx) + } return String(cString: mlx_string_data(s), encoding: .utf8)! } } diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index c276d6dec..6f2d2aa40 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -161,8 +161,10 @@ public final class Stream: @unchecked Sendable, Equatable { extension Stream: CustomStringConvertible { public var description: String { var s = mlx_string_new() - mlx_stream_tostring(&s, ctx) defer { mlx_string_free(s) } + _ = evalLock.withLock { + mlx_stream_tostring(&s, ctx) + } return String(cString: mlx_string_data(s), encoding: .utf8)! } } From dc43e62d7055353c7f99fa071a4e71d29dfddc44 Mon Sep 17 00:00:00 2001 From: RNT56 <108626871+RNT56@users.noreply.github.com> Date: Tue, 26 May 2026 17:48:59 +0200 Subject: [PATCH 04/39] Expose nuclear norm order in linalg API (#411) --- Source/MLX/Linalg.swift | 11 +++++++---- Source/MLXLinalg/Linalg.swift | 4 ++-- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/Source/MLX/Linalg.swift b/Source/MLX/Linalg.swift index 1d48f7c1d..da889492e 100644 --- a/Source/MLX/Linalg.swift +++ b/Source/MLX/Linalg.swift @@ -13,6 +13,9 @@ public enum MLXLinalg { public enum NormKind: String, Sendable { /// Frobenius norm case fro + + /// Nuclear norm, the sum of singular values + case nuc } /// Matrix or vector norm. @@ -36,7 +39,7 @@ public enum MLXLinalg { /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// - /// > Nuclear norm and norms based on singular values are not yet implemented. + /// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 @@ -91,7 +94,7 @@ public enum MLXLinalg { /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// - /// > Nuclear norm and norms based on singular values are not yet implemented. + /// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 @@ -548,7 +551,7 @@ public enum MLXLinalg { /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// -/// > Nuclear norm and norms based on singular values are not yet implemented. +/// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 @@ -597,7 +600,7 @@ public func norm( /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// -/// > Nuclear norm and norms based on singular values are not yet implemented. +/// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 diff --git a/Source/MLXLinalg/Linalg.swift b/Source/MLXLinalg/Linalg.swift index 7a287bf65..7d6950fc2 100644 --- a/Source/MLXLinalg/Linalg.swift +++ b/Source/MLXLinalg/Linalg.swift @@ -31,7 +31,7 @@ public let deprecationWarning: Void = () /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// -/// > Nuclear norm and norms based on singular values are not yet implemented. +/// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 @@ -81,7 +81,7 @@ public func norm( /// -2 | smallest singular value | as below /// other | -- | sum(abs(x)**ord)**(1./ord) /// -/// > Nuclear norm and norms based on singular values are not yet implemented. +/// Nuclear norm and norms based on singular values are implemented by the linalg backend. /// /// The Frobenius norm is given by G. H. Golub and C. F. Van Loan, *Matrix Computations*, /// Baltimore, MD, Johns Hopkins University Press, 1985, pg. 15 From 1cd3ed5efb547a9f4952afea9ce25b7527f166a7 Mon Sep 17 00:00:00 2001 From: tsushanth <78000697+tsushanth@users.noreply.github.com> Date: Thu, 11 Jun 2026 13:25:07 -0700 Subject: [PATCH 05/39] fix(Device): resolve initial default from the MLX C++ core, not hard-coded GPU (#422) Closes #395. `_resolveGlobalDefaultDevice()` returned `.gpu` whenever no explicit `setDefault(...)` had been called, so on a CPU-only host (Linux with no Metal / no CUDA) `Device.defaultDevice()` and `StreamOrDevice.default` both returned a device that simply doesn't exist, causing every subsequent op routed to the default stream to fail. The deprecated `Device.init()` already shows the right pattern: call `mlx_get_default_device()` and let the C++ core pick the device its build supports. Use that path instead of the hard-coded `.gpu`. --- Source/MLX/Device.swift | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/Source/MLX/Device.swift b/Source/MLX/Device.swift index 203c1e3f8..f092ec158 100644 --- a/Source/MLX/Device.swift +++ b/Source/MLX/Device.swift @@ -102,7 +102,18 @@ public final class Device: @unchecked Sendable, Equatable { private static func _resolveGlobalDefaultDevice() -> Device { _lock.withLock { - _defaultDevice ?? .gpu + if let device = _defaultDevice { + return device + } + // Ask the underlying MLX C++ core for its default device rather + // than hard-coding `.gpu`. On Apple platforms with Metal this + // still resolves to GPU; on a CPU-only host (Linux without + // CUDA / no Metal) it correctly resolves to CPU. Hard-coding GPU + // here meant `defaultDevice()` / `StreamOrDevice.default` + // returned an unavailable device on those hosts. + var ctx = mlx_device_new() + mlx_get_default_device(&ctx) + return Device(ctx) } } From e23ae6b2cf96737bbd31a34ed304123162f8c409 Mon Sep 17 00:00:00 2001 From: Gabriele Mondada Date: Thu, 18 Jun 2026 01:54:14 +0200 Subject: [PATCH 06/39] Allowing SPM to compile on Linux with CUDA (#413) * Support CUDA compilation with SPM * Enable CUDA by default on Linux * Add code generation capabilities in CudaBuild plugin * Use `tools/update-mlx.sh` to generate `cuda_jit_sources.h` --- CMakeLists.txt | 4 + Package.swift | 323 +- Plugins/CudaBuild/plugin.swift | 197 + Source/Cmlx/CudaBuild.json | 31 + .../mlx-generated/cuda/cuda_jit_sources.h | 4295 +++++++++++++++++ Source/Encuda/encuda-compile.swift | 96 + Source/Encuda/encuda-error.swift | 13 + Source/Encuda/encuda-link.swift | 124 + Source/Encuda/encuda-utils.swift | 25 + Source/Encuda/encuda.swift | 8 + Source/MLX/GPU+CUDA.swift | 23 + Source/MLX/MLXFastKernel.swift | 347 +- tools/update-mlx.sh | 11 +- xcode/MLX.xcodeproj/project.pbxproj | 1 + 14 files changed, 5234 insertions(+), 264 deletions(-) create mode 100644 Plugins/CudaBuild/plugin.swift create mode 100644 Source/Cmlx/CudaBuild.json create mode 100644 Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h create mode 100644 Source/Encuda/encuda-compile.swift create mode 100644 Source/Encuda/encuda-error.swift create mode 100644 Source/Encuda/encuda-link.swift create mode 100644 Source/Encuda/encuda-utils.swift create mode 100644 Source/Encuda/encuda.swift create mode 100644 Source/MLX/GPU+CUDA.swift diff --git a/CMakeLists.txt b/CMakeLists.txt index 18f350bf5..ce75855f6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -64,6 +64,10 @@ if(NOT MLX_BUILD_METAL) ${CMAKE_CURRENT_LIST_DIR}/Source/MLX/MLXArray+Metal.swift) endif() +if(NOT MLX_BUILD_CUDA) + list(REMOVE_ITEM MLX-src ${CMAKE_CURRENT_LIST_DIR}/Source/MLX/GPU+CUDA.swift) +endif() + add_library(MLX STATIC ${MLX-src}) target_include_directories(MLX PUBLIC ${CMAKE_CURRENT_LIST_DIR}/Source/Cmlx/include) diff --git a/Package.swift b/Package.swift index 17a4178f7..e36e0d862 100644 --- a/Package.swift +++ b/Package.swift @@ -1,89 +1,197 @@ -// swift-tools-version: 5.12 +// swift-tools-version: 6.3;(experimentalCGen) // The swift-tools-version declares the minimum version of Swift required to build this package. // Copyright © 2024 Apple Inc. import PackageDescription -#if os(Linux) - let platformExcludes: [String] = [ - // Linux specific excludes - "framework", - "include-framework", - "metal-cpp", - // Exclude Metal backend files on Linux, but keep no_metal.cpp for stubs - "mlx/mlx/backend/metal/allocator.cpp", - "mlx/mlx/backend/metal/binary.cpp", - "mlx/mlx/backend/metal/compiled.cpp", - "mlx/mlx/backend/metal/conv.cpp", - "mlx/mlx/backend/metal/copy.cpp", - "mlx/mlx/backend/metal/custom_kernel.cpp", - "mlx/mlx/backend/metal/device.cpp", - "mlx/mlx/backend/metal/device_info.cpp", - "mlx/mlx/backend/metal/distributed.cpp", - "mlx/mlx/backend/metal/eval.cpp", - "mlx/mlx/backend/metal/event.cpp", - "mlx/mlx/backend/metal/fence.cpp", - "mlx/mlx/backend/metal/fft.cpp", - "mlx/mlx/backend/metal/hadamard.cpp", - "mlx/mlx/backend/metal/indexing.cpp", - "mlx/mlx/backend/metal/jit_kernels.cpp", - "mlx/mlx/backend/metal/logsumexp.cpp", - "mlx/mlx/backend/metal/matmul.cpp", - "mlx/mlx/backend/metal/metal.cpp", - "mlx/mlx/backend/metal/normalization.cpp", - "mlx/mlx/backend/metal/primitives.cpp", - "mlx/mlx/backend/metal/quantized.cpp", - "mlx/mlx/backend/metal/reduce.cpp", - "mlx/mlx/backend/metal/resident.cpp", - "mlx/mlx/backend/metal/rope.cpp", - "mlx/mlx/backend/metal/scaled_dot_product_attention.cpp", - "mlx/mlx/backend/metal/scan.cpp", - "mlx/mlx/backend/metal/slicing.cpp", - "mlx/mlx/backend/metal/softmax.cpp", - "mlx/mlx/backend/metal/sort.cpp", - "mlx/mlx/backend/metal/ternary.cpp", - "mlx/mlx/backend/metal/unary.cpp", - "mlx/mlx/backend/metal/utils.cpp", - "mlx/mlx/backend/metal/kernels", // Exclude kernels directory - "mlx/mlx/backend/metal/jit", // Exclude jit directory - - "mlx/mlx/backend/gpu", // Exclude GPU backend on Linux, use no_gpu instead - "mlx/mlx/backend/no_cpu", // Exclude no_cpu backend on Linux, use cpu instead - "mlx/mlx/backend/cpu/gemms/bnns.cpp", // macOS Accelerate version - "mlx-conditional", - "mlx-c/mlx/c/metal.cpp", - - "mlx-c/mlx/c/fast.cpp", // Exclude on Linux - calls metal_kernel unconditionally - ] - - let cxxSettings: [CXXSetting] = [] - - let linkerSettings: [LinkerSetting] = [ - .linkedLibrary("gfortran", .when(platforms: [.linux])), - .linkedLibrary("blas", .when(platforms: [.linux])), - .linkedLibrary("lapack", .when(platforms: [.linux])), - .linkedLibrary("openblas", .when(platforms: [.linux])), - ] +let noMetalCmlxExcludes = [ + // Exclude Metal backend files, but keep no_metal.cpp for stubs + // "mlx/mlx/backend/metal/no_metal.cpp", + "mlx/mlx/backend/metal/allocator.cpp", + "mlx/mlx/backend/metal/binary.cpp", + "mlx/mlx/backend/metal/compiled.cpp", + "mlx/mlx/backend/metal/conv.cpp", + "mlx/mlx/backend/metal/copy.cpp", + "mlx/mlx/backend/metal/custom_kernel.cpp", + "mlx/mlx/backend/metal/device.cpp", + "mlx/mlx/backend/metal/device_info.cpp", + "mlx/mlx/backend/metal/distributed.cpp", + "mlx/mlx/backend/metal/eval.cpp", + "mlx/mlx/backend/metal/event.cpp", + "mlx/mlx/backend/metal/fence.cpp", + "mlx/mlx/backend/metal/fft.cpp", + "mlx/mlx/backend/metal/hadamard.cpp", + "mlx/mlx/backend/metal/indexing.cpp", + "mlx/mlx/backend/metal/jit_kernels.cpp", + "mlx/mlx/backend/metal/logsumexp.cpp", + "mlx/mlx/backend/metal/matmul.cpp", + "mlx/mlx/backend/metal/metal.cpp", + "mlx/mlx/backend/metal/normalization.cpp", + "mlx/mlx/backend/metal/primitives.cpp", + "mlx/mlx/backend/metal/quantized.cpp", + "mlx/mlx/backend/metal/reduce.cpp", + "mlx/mlx/backend/metal/resident.cpp", + "mlx/mlx/backend/metal/rope.cpp", + "mlx/mlx/backend/metal/scaled_dot_product_attention.cpp", + "mlx/mlx/backend/metal/scan.cpp", + "mlx/mlx/backend/metal/slicing.cpp", + "mlx/mlx/backend/metal/softmax.cpp", + "mlx/mlx/backend/metal/sort.cpp", + "mlx/mlx/backend/metal/ternary.cpp", + "mlx/mlx/backend/metal/unary.cpp", + "mlx/mlx/backend/metal/utils.cpp", + "mlx/mlx/backend/metal/kernels", // Exclude kernels directory + "mlx/mlx/backend/metal/jit", // Exclude jit directory +] + +let noCudaCmlxExcludes = [ + // Exclude CUDA backend files, but keep no_cuda.cpp for stubs + // mlx/mlx/backend/cuda/no_cuda.cpp + "mlx/mlx/backend/cuda/allocator.cpp", + "mlx/mlx/backend/cuda/compiled.cpp", + "mlx/mlx/backend/cuda/conv.cpp", + "mlx/mlx/backend/cuda/cublas_utils.cpp", + "mlx/mlx/backend/cuda/cudnn_utils.cpp", + "mlx/mlx/backend/cuda/custom_kernel.cpp", + "mlx/mlx/backend/cuda/delayload.cpp", + "mlx/mlx/backend/cuda/device_info.cpp", + "mlx/mlx/backend/cuda/device.cpp", + "mlx/mlx/backend/cuda/eval.cpp", + "mlx/mlx/backend/cuda/fence.cpp", + "mlx/mlx/backend/cuda/indexing.cpp", + "mlx/mlx/backend/cuda/jit_module.cpp", + "mlx/mlx/backend/cuda/load.cpp", + "mlx/mlx/backend/cuda/matmul.cpp", + "mlx/mlx/backend/cuda/primitives.cpp", + "mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp", + "mlx/mlx/backend/cuda/slicing.cpp", + "mlx/mlx/backend/cuda/utils.cpp", + "mlx/mlx/backend/cuda/worker.cpp", + "mlx/mlx/backend/cuda/binary", + "mlx/mlx/backend/cuda/conv", + "mlx/mlx/backend/cuda/copy", + "mlx/mlx/backend/cuda/device", + "mlx/mlx/backend/cuda/gemms", + "mlx/mlx/backend/cuda/quantized", + "mlx/mlx/backend/cuda/reduce", + "mlx/mlx/backend/cuda/steel", + "mlx/mlx/backend/cuda/unary", +] - let mlxSwiftExcludes: [String] = [ - "GPU+Metal.swift", - "MLXArray+Metal.swift", - "MLXFast.swift", - "MLXFastKernel.swift", - ] +#if os(Linux) + let platformExcludes: [String] + let cxxSettings: [CXXSetting] + let linkerSettings: [LinkerSetting] + let mlxSwiftExcludes: [String] + + if Context.environment["SPM_CUDA"] != "0" { + // Linux with CUDA + + platformExcludes = + [ + "framework", + "include-framework", + "metal-cpp", + + "mlx/mlx/backend/no_gpu", + "mlx/mlx/backend/cuda/no_cuda.cpp", + "mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp", + "mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp", + "mlx/mlx/backend/no_cpu", // Exclude no_cpu backend on Linux, use cpu instead + "mlx/mlx/backend/cpu/gemms/bnns.cpp", // macOS Accelerate version + "mlx-conditional", + "mlx-c/mlx/c/metal.cpp", + + "mlx/mlx/backend/cuda/delayload.cpp", // For Windows + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmv.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh", + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmm.h", + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmm.cu", + "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu", + "mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu", + ] + noMetalCmlxExcludes + + cxxSettings = [ + .unsafeFlags(["-I/usr/local/cuda/include"]), + .unsafeFlags(["-I/usr/local/cuda/include/cccl"]), + .define("MLX_CCCL_DIR", to: "\"/usr/local/cuda/include/cccl\""), + ] + + linkerSettings = [ + .linkedLibrary("gfortran", .when(platforms: [.linux])), + .linkedLibrary("blas", .when(platforms: [.linux])), + .linkedLibrary("lapack", .when(platforms: [.linux])), + .linkedLibrary("openblas", .when(platforms: [.linux])), + .unsafeFlags(["-L/usr/local/cuda/lib64"]), + .unsafeFlags(["-L/usr/local/cuda/lib64/stubs"]), + .linkedLibrary("cudnn"), + .linkedLibrary("cublas"), + .linkedLibrary("cublasLt"), + .linkedLibrary("nvrtc"), + .linkedLibrary("cudart"), + .linkedLibrary("cuda"), + ] + + mlxSwiftExcludes = [ + "GPU+Metal.swift", + "MLXArray+Metal.swift", + ] + } else { + // Linux without CUDA (CPU only) + + platformExcludes = + [ + "framework", + "include-framework", + "metal-cpp", + + "mlx/mlx/backend/gpu", // Exclude GPU backend on Linux, use no_gpu instead + "mlx/mlx/backend/no_cpu", // Exclude no_cpu backend on Linux, use cpu instead + "mlx/mlx/backend/cpu/gemms/bnns.cpp", // macOS Accelerate version + "mlx-conditional", + "mlx-c/mlx/c/metal.cpp", + + "mlx-c/mlx/c/fast.cpp", // Exclude on Linux - calls metal_kernel unconditionally + + ] + noMetalCmlxExcludes + noCudaCmlxExcludes + + cxxSettings = [] + + linkerSettings = [ + .linkedLibrary("gfortran", .when(platforms: [.linux])), + .linkedLibrary("blas", .when(platforms: [.linux])), + .linkedLibrary("lapack", .when(platforms: [.linux])), + .linkedLibrary("openblas", .when(platforms: [.linux])), + ] + + mlxSwiftExcludes = [ + "GPU+Metal.swift", + "GPU+CUDA.swift", + "MLXArray+Metal.swift", + "MLXFast.swift", + "MLXFastKernel.swift", + ] + } #else - let platformExcludes: [String] = [ - "mlx/mlx/backend/cpu/compiled.cpp", + // Apple's platforms with Metal - // opt-out of these backends (using metal) - "mlx/mlx/backend/no_gpu", - "mlx/mlx/backend/no_cpu", - "mlx/mlx/backend/metal/no_metal.cpp", + let platformExcludes: [String] = + [ + "mlx/mlx/backend/cpu/compiled.cpp", - // bnns instead of simd (accelerate) - "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp", - "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp", - ] + // opt-out of these backends (using metal) + "mlx/mlx/backend/no_gpu", + "mlx/mlx/backend/no_cpu", + "mlx/mlx/backend/metal/no_metal.cpp", + + // bnns instead of simd (accelerate) + "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp", + "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp", + ] + noCudaCmlxExcludes let cxxSettings: [CXXSetting] = [ .headerSearchPath("metal-cpp"), @@ -101,7 +209,9 @@ import PackageDescription .linkedFramework("Accelerate"), ] - let mlxSwiftExcludes: [String] = [] + let mlxSwiftExcludes: [String] = [ + "GPU+CUDA.swift" + ] #endif let cmlx = Target.target( @@ -147,40 +257,6 @@ let cmlx = Target.target( "mlx/setup.py", "mlx/tests", - // special handling for cuda -- we need to keep one file: - // mlx/mlx/backend/cuda/no_cuda.cpp - - "mlx/mlx/backend/cuda/allocator.cpp", - "mlx/mlx/backend/cuda/compiled.cpp", - "mlx/mlx/backend/cuda/conv.cpp", - "mlx/mlx/backend/cuda/cublas_utils.cpp", - "mlx/mlx/backend/cuda/cudnn_utils.cpp", - "mlx/mlx/backend/cuda/custom_kernel.cpp", - "mlx/mlx/backend/cuda/delayload.cpp", - "mlx/mlx/backend/cuda/device_info.cpp", - "mlx/mlx/backend/cuda/device.cpp", - "mlx/mlx/backend/cuda/eval.cpp", - "mlx/mlx/backend/cuda/fence.cpp", - "mlx/mlx/backend/cuda/indexing.cpp", - "mlx/mlx/backend/cuda/jit_module.cpp", - "mlx/mlx/backend/cuda/load.cpp", - "mlx/mlx/backend/cuda/matmul.cpp", - "mlx/mlx/backend/cuda/primitives.cpp", - "mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp", - "mlx/mlx/backend/cuda/slicing.cpp", - "mlx/mlx/backend/cuda/utils.cpp", - "mlx/mlx/backend/cuda/worker.cpp", - - "mlx/mlx/backend/cuda/binary", - "mlx/mlx/backend/cuda/conv", - "mlx/mlx/backend/cuda/copy", - "mlx/mlx/backend/cuda/device", - "mlx/mlx/backend/cuda/gemms", - "mlx/mlx/backend/cuda/quantized", - "mlx/mlx/backend/cuda/reduce", - "mlx/mlx/backend/cuda/steel", - "mlx/mlx/backend/cuda/unary", - // build variants (we are opting _out_ of these) "mlx/mlx/io/no_safetensors.cpp", "mlx/mlx/io/gguf.cpp", @@ -203,6 +279,7 @@ let cmlx = Target.target( cSettings: [ .headerSearchPath("mlx"), .headerSearchPath("mlx-c"), + .headerSearchPath("mlx-generated/cuda"), ], cxxSettings: cxxSettings + [ .headerSearchPath("mlx"), @@ -211,7 +288,10 @@ let cmlx = Target.target( .headerSearchPath("fmt/include"), .define("MLX_VERSION", to: "\"0.31.1\""), ], - linkerSettings: linkerSettings + linkerSettings: linkerSettings, + plugins: [ + .plugin(name: "CudaBuild") + ], ) let package = Package( @@ -236,7 +316,8 @@ let package = Package( ], dependencies: [ // for Complex type - .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0") + .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0"), + .package(url: "https://github.com/apple/swift-argument-parser", from: "1.0.0"), ], targets: [ cmlx, @@ -333,6 +414,20 @@ let package = Package( path: "Source/Examples", sources: ["CustomFunctionExampleSimple.swift"] ), + .executableTarget( + name: "encuda", + dependencies: [ + .product(name: "ArgumentParser", package: "swift-argument-parser") + ], + path: "Source/Encuda", + ), + .plugin( + name: "CudaBuild", + capability: .buildTool(), + dependencies: [ + .target(name: "encuda") + ], + ), ], cxxLanguageStandard: .gnucxx20 ) diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift new file mode 100644 index 000000000..f7d12dd4c --- /dev/null +++ b/Plugins/CudaBuild/plugin.swift @@ -0,0 +1,197 @@ +import Foundation +import PackagePlugin + +@main +struct CudaBuild: BuildToolPlugin { + + struct CodeGeneration: Codable { + var tool: String + var inputs: [String] = [] + var outputs: [String] = [] + } + + struct Settings: Codable { + var headerSearchPaths: [String] = [] + var exclude: [String] = [] + var cppLanguageStandard: String? = nil + var verbose: Bool = false + var codeGeneration: [CodeGeneration] = [] + + init() {} + + init(from decoder: Decoder) throws { + let c = try decoder.container(keyedBy: CodingKeys.self) + headerSearchPaths = + try c.decodeIfPresent([String].self, forKey: .headerSearchPaths) ?? [] + exclude = try c.decodeIfPresent([String].self, forKey: .exclude) ?? [] + cppLanguageStandard = try c.decodeIfPresent(String.self, forKey: .cppLanguageStandard) + verbose = try c.decodeIfPresent(Bool.self, forKey: .verbose) ?? false + codeGeneration = + try c.decodeIfPresent([CodeGeneration].self, forKey: .codeGeneration) ?? [] + } + } + + func createBuildCommands(context: PluginContext, target: Target) async throws -> [Command] { + + print("CUDA Build Plugin") + + guard isCudaEnabled() else { + print("CUDA is disabled") + return [] + } + + guard let clangUrl = try? context.tool(named: "clang++") else { + fatalError("clang++ not found") + } + + print("Use clang++ at: \(clangUrl.url.path)") + + let sourceDir = target.directoryURL + + print("Source directory: \(sourceDir.path)") + + // Read settings from CudaBuild.json if it exists + + let settingsFile = sourceDir.appendingPathComponent("CudaBuild.json") + let settings: Settings + if let data = try? Data(contentsOf: settingsFile) { + print("Found settings file: \(settingsFile.path)") + settings = try JSONDecoder().decode(Settings.self, from: data) + } else { + settings = Settings() + } + + // Scan source directory for .cu files + + let sourceDirPath = sourceDir.path.hasSuffix("/") ? sourceDir.path : sourceDir.path + "/" + var sourceCuFiles: [URL] = [] + if let enumerator = FileManager.default.enumerator( + at: sourceDir, includingPropertiesForKeys: nil) + { + while let inputUrl = enumerator.nextObject() as? URL { + if inputUrl.pathExtension == "cu" { + guard inputUrl.path.hasPrefix(sourceDirPath) else { + fatalError( + "Input file \(inputUrl.path) is not under source directory \(sourceDirPath)" + ) + } + let relativePath = String(inputUrl.path.dropFirst(sourceDirPath.count)) + if isExcluded(settings: settings, relativePath: relativePath) { + print("Excluding \(relativePath)") + continue + } + sourceCuFiles.append(URL(string: relativePath, relativeTo: sourceDir)!) + } + } + } + + sourceCuFiles.sort { $0.path < $1.path } // make it deterministic + + print("Source files: \(sourceCuFiles.map { $0.relativePath })") + + let outputDir = context.pluginWorkDirectoryURL + + var commands: [Command] = [] + + // Invoke code generation tools + + var generatedCuFiles: [URL] = [] + + for (index, codeGen) in settings.codeGeneration.enumerated() { + let tool = try context.tool(named: codeGen.tool) + let codeGenOutputDir = outputDir.appendingPathComponent("gen\(index)-\(codeGen.tool)") + try FileManager.default.createDirectory( + at: codeGenOutputDir, withIntermediateDirectories: true) + let inputUrls = codeGen.inputs.map { sourceDir.appendingPathComponent($0) } + let outputUrls = codeGen.outputs.map { codeGenOutputDir.appendingPathComponent($0) } + commands.append( + .buildCommand( + displayName: "Running \(codeGen.tool)", + executable: tool.url, + arguments: inputUrls.map { $0.path } + outputUrls.flatMap { ["-o", $0.path] }, + inputFiles: inputUrls, + outputFiles: outputUrls + ) + ) + generatedCuFiles += outputUrls.filter { $0.pathExtension == "cu" } + } + + print("Generated source files: \(generatedCuFiles.map { $0.relativePath })") + + // Ask encuda to skip work if the output is already up to date. + // This should normally be handled by the build system, but apparently it is not. + + let incrementalFlag = ["--incremental"] + + // Invoke `encuda compile` to compile each .cu file to .cpp + + let encuda = try context.tool(named: "encuda") + let verboseFlag = settings.verbose ? ["-v"] : [] + let headerSearchPathArgs = settings.headerSearchPaths.flatMap { + ["-I", sourceDirPath + $0] + } + let stdArgs = settings.cppLanguageStandard.map { ["--std", $0] } ?? [] + + for inputFile in sourceCuFiles + generatedCuFiles { + let outputCpp = URL(string: inputFile.relativePath, relativeTo: outputDir)! + .deletingPathExtension().appendingPathExtension("cpp") + commands.append( + .buildCommand( + displayName: + "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)", + executable: encuda.url, + arguments: ["compile"] + verboseFlag + stdArgs + incrementalFlag + [ + "--clangpp", clangUrl.url.path, + "-I", sourceDir.path, + ] + headerSearchPathArgs + [ + inputFile.path, + "-o", outputCpp.path, + ], + inputFiles: [inputFile], + outputFiles: [outputCpp] + ) + ) + } + + // Invoke `encuda link` with all .cpp files + + let outputCpps = (sourceCuFiles + generatedCuFiles).map { inputFile in + URL(string: inputFile.relativePath, relativeTo: outputDir)! + .deletingPathExtension() + .appendingPathExtension("cpp") + } + + let linkOutput = outputDir.appendingPathComponent("__cuda_link.cpp") + + commands.append( + .buildCommand( + displayName: "Linking CUDA objects", + executable: encuda.url, + arguments: ["link"] + verboseFlag + stdArgs + incrementalFlag + [ + "--clangpp", clangUrl.url.path, + ] + outputCpps.map { $0.path } + ["-o", linkOutput.path], + inputFiles: outputCpps, + outputFiles: [linkOutput] + ) + ) + + return commands + } + + func isExcluded(settings: Settings, relativePath: String) -> Bool { + func isInFolder(path: String, folderPath: String) -> Bool { + let prefix = folderPath.hasSuffix("/") ? folderPath : (folderPath + "/") + return path.hasPrefix(prefix) + } + return settings.exclude.contains(relativePath) + || settings.exclude.contains { isInFolder(path: relativePath, folderPath: $0) } + } + + func isCudaEnabled() -> Bool { + #if os(Linux) + return ProcessInfo.processInfo.environment["SPM_CUDA"] != "0" + #else + return false + #endif + } +} diff --git a/Source/Cmlx/CudaBuild.json b/Source/Cmlx/CudaBuild.json new file mode 100644 index 000000000..3ba88fc9d --- /dev/null +++ b/Source/Cmlx/CudaBuild.json @@ -0,0 +1,31 @@ +{ + "headerSearchPaths": [ + "mlx", + "mlx-c", + "json/single_include/nlohmann", + "fmt/include" + ], + "exclude": [ + "framework", + "include-framework", + "metal-cpp", + "mlx/mlx/backend/metal/kernels", + "mlx/mlx/backend/metal/jit", + "mlx/mlx/backend/no_cpu", + "mlx-conditional", + "mlx-c/examples", + "json", + "fmt/test", + "fmt/doc", + "fmt/support", + "mlx/benchmarks", + "mlx/cmake", + "mlx/docs", + "mlx/examples", + "mlx/python", + "mlx/tests", + "mlx/mlx/backend/metal/kernels" + ], + "cppLanguageStandard": "c++20", + "verbose": false +} diff --git a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h new file mode 100644 index 000000000..7d14d04bc --- /dev/null +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -0,0 +1,4295 @@ +#pragma once + +namespace mlx::core { + +constexpr char jit_source_atomic_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, + 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, + 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x66, 0x70, 0x31, 0x36, + 0x5f, 0x6d, 0x61, 0x74, 0x68, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x3e, 0x0a, 0x0a, + 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, + 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, + 0x7b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x54, + 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x72, 0x65, 0x66, 0x3c, 0x54, + 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x3e, 0x20, 0x72, 0x65, 0x66, 0x28, 0x2a, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x66, 0x20, 0x2b, 0x3d, + 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, + 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, + 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, + 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x61, 0x74, 0x6f, 0x6d, 0x69, + 0x63, 0x5f, 0x72, 0x65, 0x66, 0x3c, 0x54, 0x2c, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x73, 0x63, + 0x6f, 0x70, 0x65, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x3e, 0x20, + 0x72, 0x65, 0x66, 0x28, 0x2a, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x54, 0x20, 0x6f, 0x6c, 0x64, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x66, + 0x2e, 0x6c, 0x6f, 0x61, 0x64, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x77, + 0x68, 0x69, 0x6c, 0x65, 0x20, 0x28, 0x21, 0x72, 0x65, 0x66, 0x2e, 0x63, + 0x6f, 0x6d, 0x70, 0x61, 0x72, 0x65, 0x5f, 0x65, 0x78, 0x63, 0x68, 0x61, + 0x6e, 0x67, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x6f, 0x6e, 0x67, 0x28, 0x6f, + 0x6c, 0x64, 0x2c, 0x20, 0x6f, 0x6c, 0x64, 0x20, 0x2a, 0x20, 0x76, 0x61, + 0x6c, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x72, 0x65, 0x66, 0x3c, 0x54, 0x2c, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x5f, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x3e, 0x20, 0x72, 0x65, 0x66, 0x28, 0x2a, 0x6f, 0x75, 0x74, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x66, 0x2e, 0x66, 0x65, 0x74, 0x63, + 0x68, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x54, + 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x72, 0x65, 0x66, 0x3c, 0x54, + 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x3e, 0x20, 0x72, 0x65, 0x66, 0x28, 0x2a, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x66, 0x2e, 0x66, 0x65, + 0x74, 0x63, 0x68, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x76, 0x61, 0x6c, 0x29, + 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x53, 0x6f, 0x6d, 0x65, + 0x68, 0x6f, 0x77, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x72, 0x65, 0x66, 0x20, 0x64, 0x6f, 0x65, + 0x73, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, + 0x65, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x20, 0x61, 0x64, 0x64, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x66, 0x6f, 0x6c, 0x6c, 0x6f, 0x77, 0x69, + 0x6e, 0x67, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, + 0x63, 0x5f, 0x61, 0x64, 0x64, 0x5f, 0x67, 0x65, 0x6e, 0x65, 0x72, 0x61, + 0x6c, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, + 0x76, 0x61, 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x72, 0x65, + 0x66, 0x3c, 0x54, 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x73, 0x63, 0x6f, 0x70, 0x65, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x3e, 0x20, 0x72, 0x65, 0x66, 0x28, + 0x2a, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x6f, + 0x6c, 0x64, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x66, 0x2e, 0x6c, 0x6f, 0x61, + 0x64, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x77, 0x68, 0x69, 0x6c, 0x65, + 0x20, 0x28, 0x21, 0x72, 0x65, 0x66, 0x2e, 0x63, 0x6f, 0x6d, 0x70, 0x61, + 0x72, 0x65, 0x5f, 0x65, 0x78, 0x63, 0x68, 0x61, 0x6e, 0x67, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x6f, 0x6e, 0x67, 0x28, 0x6f, 0x6c, 0x64, 0x2c, 0x20, + 0x6f, 0x6c, 0x64, 0x20, 0x2b, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x69, 0x6e, 0x6c, + 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, 0x74, 0x6f, 0x6d, + 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, + 0x66, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, + 0x6c, 0x66, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x41, 0x64, 0x64, 0x28, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, + 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x36, + 0x34, 0x5f, 0x74, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x5f, + 0x67, 0x65, 0x6e, 0x65, 0x72, 0x61, 0x6c, 0x28, 0x6f, 0x75, 0x74, 0x2c, + 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x69, 0x6e, + 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x61, 0x74, 0x6f, + 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x5f, 0x5f, 0x6e, 0x76, + 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x31, 0x36, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x7b, + 0x0a, 0x23, 0x69, 0x66, 0x20, 0x5f, 0x5f, 0x43, 0x55, 0x44, 0x41, 0x5f, + 0x41, 0x52, 0x43, 0x48, 0x5f, 0x5f, 0x20, 0x3c, 0x20, 0x38, 0x30, 0x30, + 0x0a, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, + 0x64, 0x5f, 0x67, 0x65, 0x6e, 0x65, 0x72, 0x61, 0x6c, 0x28, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6c, + 0x73, 0x65, 0x0a, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x41, + 0x64, 0x64, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, + 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, 0x66, 0x0a, 0x7d, 0x0a, 0x0a, + 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, + 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_binary_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x6e, 0x61, 0x72, 0x79, + 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3e, + 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, + 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, + 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x41, 0x64, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2b, 0x20, 0x79, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x46, 0x6c, 0x6f, 0x6f, 0x72, 0x44, 0x69, + 0x76, 0x69, 0x64, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, + 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, + 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x72, 0x75, 0x6e, 0x63, + 0x28, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x44, 0x69, 0x76, 0x69, 0x64, 0x65, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, + 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, 0x6e, 0x64, 0x65, 0x72, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, + 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x5f, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, + 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x78, + 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, 0x20, 0x21, 0x3d, 0x20, 0x30, + 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, 0x3c, 0x20, 0x30, 0x20, 0x21, + 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x20, + 0x2b, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x72, 0x20, 0x3d, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, + 0x6d, 0x6f, 0x64, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, 0x20, 0x21, + 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, 0x3c, 0x20, + 0x30, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x20, 0x3d, 0x20, 0x72, 0x20, 0x2b, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x71, 0x75, 0x61, 0x6c, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, 0x3d, 0x20, + 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4e, 0x61, 0x4e, 0x45, 0x71, 0x75, + 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x20, 0x7c, + 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x28, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, + 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x69, + 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x73, 0x6e, + 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, + 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, + 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, + 0x29, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, + 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, + 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, + 0x3d, 0x20, 0x79, 0x20, 0x7c, 0x7c, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x78, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x79, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, 0x72, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, + 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, 0x72, 0x45, + 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, + 0x20, 0x3e, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, + 0x73, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, + 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x73, 0x73, 0x45, + 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, + 0x20, 0x3c, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, + 0x67, 0x41, 0x64, 0x64, 0x45, 0x78, 0x70, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, + 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, + 0x71, 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, + 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, 0x71, + 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x78, 0x20, + 0x3d, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3e, + 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3f, 0x20, + 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x20, 0x3d, 0x20, + 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3c, 0x20, 0x79, + 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3f, 0x20, 0x78, 0x20, + 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, + 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x21, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x28, 0x6d, 0x69, 0x6e, 0x5f, + 0x72, 0x65, 0x61, 0x6c, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x6d, 0x69, + 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3d, 0x3d, 0x20, 0x6d, 0x61, + 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6d, + 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3c, 0x20, 0x30, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, 0x6e, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x4c, + 0x6f, 0x67, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, 0x7b, 0x7d, 0x28, 0x6d, + 0x69, 0x6e, 0x29, 0x20, 0x2b, 0x20, 0x45, 0x78, 0x70, 0x7b, 0x7d, 0x28, + 0x6d, 0x61, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x4c, + 0x6f, 0x67, 0x31, 0x70, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, 0x7b, 0x7d, + 0x28, 0x6d, 0x69, 0x6e, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, 0x29, 0x29, + 0x20, 0x2b, 0x20, 0x6d, 0x61, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x20, 0x7c, + 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x71, 0x75, + 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x54, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x76, 0x61, + 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x78, 0x2c, 0x20, 0x79, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x28, 0x6d, 0x69, 0x6e, 0x76, 0x61, 0x6c, 0x20, + 0x3d, 0x3d, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, + 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x20, 0x7c, 0x7c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x3d, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, + 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, + 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, + 0x74, 0x79, 0x28, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x3f, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x3a, + 0x20, 0x54, 0x28, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x2b, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6c, + 0x6f, 0x67, 0x31, 0x70, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x78, 0x70, 0x28, 0x6d, 0x69, 0x6e, 0x76, + 0x61, 0x6c, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x29, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x4d, 0x61, 0x78, 0x69, 0x6d, 0x75, 0x6d, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, + 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, + 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, + 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, + 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4d, 0x69, 0x6e, 0x69, 0x6d, 0x75, 0x6d, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, + 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x78, 0x2c, + 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, + 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, 0x79, 0x20, + 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4d, 0x75, 0x6c, 0x74, 0x69, 0x70, + 0x6c, 0x79, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, + 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2a, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4e, 0x6f, 0x74, 0x45, 0x71, 0x75, 0x61, 0x6c, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, + 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x2e, + 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x78, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, + 0x21, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x50, 0x6f, 0x77, 0x65, 0x72, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, + 0x28, 0x54, 0x20, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x54, 0x20, 0x65, + 0x78, 0x70, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, + 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x54, 0x20, 0x72, 0x65, 0x73, 0x20, 0x3d, 0x20, 0x31, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x61, 0x69, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x61, 0x6e, 0x20, 0x69, 0x6e, 0x74, 0x65, + 0x67, 0x65, 0x72, 0x20, 0x74, 0x6f, 0x20, 0x61, 0x20, 0x6e, 0x65, 0x67, + 0x61, 0x74, 0x69, 0x76, 0x65, 0x20, 0x70, 0x6f, 0x77, 0x65, 0x72, 0x20, + 0x69, 0x73, 0x20, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x64, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x28, 0x65, 0x78, 0x70, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x77, 0x68, 0x69, + 0x6c, 0x65, 0x20, 0x28, 0x65, 0x78, 0x70, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x65, + 0x78, 0x70, 0x20, 0x26, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x20, + 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x65, 0x78, 0x70, 0x20, 0x3e, 0x3e, 0x3d, 0x20, 0x31, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x61, + 0x73, 0x65, 0x20, 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x65, 0x73, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, + 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x70, 0x6f, + 0x77, 0x28, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x65, 0x78, 0x70, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x70, 0x6f, 0x77, 0x28, 0x62, 0x61, 0x73, 0x65, 0x2c, + 0x20, 0x65, 0x78, 0x70, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x53, 0x75, 0x62, 0x74, 0x72, 0x61, 0x63, 0x74, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, + 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2d, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, 0x6c, 0x41, 0x6e, 0x64, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, + 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x26, 0x20, 0x79, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, 0x6c, 0x4f, + 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, + 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, 0x7c, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, 0x65, + 0x41, 0x6e, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x20, 0x79, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, + 0x65, 0x4f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, 0x20, 0x79, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, + 0x65, 0x58, 0x6f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x5e, 0x20, + 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x66, 0x74, 0x53, + 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x3c, + 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x52, 0x69, 0x67, 0x68, + 0x74, 0x53, 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, + 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, + 0x3e, 0x3e, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x41, 0x72, + 0x63, 0x54, 0x61, 0x6e, 0x32, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x79, 0x2c, 0x20, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x74, 0x61, 0x6e, + 0x32, 0x28, 0x79, 0x2c, 0x20, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x44, 0x69, 0x76, 0x4d, 0x6f, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, + 0x61, 0x79, 0x3c, 0x54, 0x2c, 0x20, 0x32, 0x3e, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, + 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x46, 0x6c, 0x6f, 0x6f, + 0x72, 0x44, 0x69, 0x76, 0x69, 0x64, 0x65, 0x7b, 0x7d, 0x28, 0x78, 0x2c, + 0x20, 0x79, 0x29, 0x2c, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, 0x6e, 0x64, + 0x65, 0x72, 0x7b, 0x7d, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x7d, 0x20, + 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, + 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_cast_op[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x5f, 0x62, 0x66, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, + 0x5f, 0x66, 0x70, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, + 0x0a, 0x2f, 0x2f, 0x20, 0x41, 0x6e, 0x20, 0x6f, 0x70, 0x20, 0x74, 0x68, + 0x61, 0x74, 0x20, 0x64, 0x6f, 0x65, 0x73, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x2c, 0x20, 0x77, 0x69, 0x74, + 0x68, 0x20, 0x63, 0x75, 0x73, 0x74, 0x6f, 0x6d, 0x20, 0x63, 0x6f, 0x6e, + 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x73, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x73, 0x6f, 0x6d, 0x65, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x72, 0x63, 0x54, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x44, + 0x73, 0x74, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x3d, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x3e, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x61, 0x62, 0x6c, + 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, + 0x74, 0x69, 0x62, 0x6c, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, + 0x2c, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x44, + 0x73, 0x74, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x53, 0x72, 0x63, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x44, 0x73, 0x74, 0x54, 0x3e, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x43, 0x61, 0x73, + 0x74, 0x69, 0x6e, 0x67, 0x73, 0x20, 0x62, 0x65, 0x74, 0x77, 0x65, 0x65, + 0x6e, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x20, 0x61, 0x6e, + 0x64, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x65, 0x61, 0x6e, 0x2e, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x2c, + 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, + 0x63, 0x61, 0x73, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, + 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, + 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, 0x69, + 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, + 0x2c, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, + 0x54, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, 0x65, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, + 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x78, 0x20, 0x3f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x31, 0x2c, 0x20, 0x31, 0x7d, 0x20, + 0x3a, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, + 0x54, 0x3e, 0x7b, 0x30, 0x2c, 0x20, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6e, + 0x76, 0x65, 0x72, 0x74, 0x69, 0x6e, 0x67, 0x20, 0x61, 0x20, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x20, 0x6e, 0x75, 0x6d, 0x62, 0x65, 0x72, + 0x20, 0x74, 0x6f, 0x20, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x6e, 0x75, 0x6d, + 0x62, 0x65, 0x72, 0x20, 0x64, 0x69, 0x73, 0x63, 0x61, 0x72, 0x64, 0x73, + 0x20, 0x74, 0x68, 0x65, 0x20, 0x69, 0x6d, 0x61, 0x67, 0x69, 0x6e, 0x61, + 0x72, 0x79, 0x20, 0x70, 0x61, 0x72, 0x74, 0x2e, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, + 0x2c, 0x20, 0x44, 0x73, 0x74, 0x54, 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, + 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x21, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x44, 0x73, 0x74, + 0x54, 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x69, 0x62, 0x6c, 0x65, 0x5f, 0x76, + 0x3c, 0x54, 0x2c, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x44, 0x73, 0x74, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x61, 0x73, + 0x73, 0x65, 0x72, 0x74, 0x28, 0x21, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x44, 0x73, 0x74, 0x54, 0x3e, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x3c, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x28, 0x78, 0x2e, 0x72, 0x65, + 0x61, 0x6c, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x41, 0x6c, 0x6c, 0x6f, 0x77, 0x20, + 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x69, 0x6e, 0x67, 0x20, 0x61, + 0x20, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x6e, 0x75, 0x6d, 0x62, 0x65, 0x72, + 0x20, 0x74, 0x6f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x20, + 0x6e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x53, + 0x72, 0x63, 0x54, 0x2c, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, + 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x21, 0x69, 0x73, 0x5f, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, + 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x63, 0x6f, + 0x6e, 0x76, 0x65, 0x72, 0x74, 0x69, 0x62, 0x6c, 0x65, 0x5f, 0x76, 0x3c, + 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x54, 0x3e, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x53, 0x72, 0x63, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x61, 0x73, 0x73, + 0x65, 0x72, 0x74, 0x28, 0x21, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x3e, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, + 0x3e, 0x7b, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x3c, 0x54, 0x3e, 0x28, 0x78, 0x29, 0x2c, 0x20, 0x30, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x44, 0x6f, 0x20, 0x6e, 0x6f, 0x74, 0x68, 0x69, 0x6e, 0x67, 0x20, 0x77, + 0x68, 0x65, 0x6e, 0x20, 0x6e, 0x6f, 0x20, 0x63, 0x61, 0x73, 0x74, 0x69, + 0x6e, 0x67, 0x20, 0x69, 0x73, 0x20, 0x6e, 0x65, 0x65, 0x64, 0x65, 0x64, + 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x72, 0x63, + 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x44, 0x73, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x44, + 0x73, 0x74, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, + 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, + 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x44, + 0x73, 0x74, 0x54, 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, + 0x63, 0x61, 0x73, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, + 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x53, 0x72, 0x63, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x53, + 0x72, 0x63, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x49, 0x6e, + 0x20, 0x43, 0x55, 0x44, 0x41, 0x20, 0x31, 0x31, 0x20, 0x74, 0x68, 0x65, + 0x20, 0x68, 0x61, 0x6c, 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x20, + 0x64, 0x6f, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x64, 0x65, 0x66, 0x69, 0x6e, + 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, + 0x73, 0x20, 0x62, 0x65, 0x74, 0x77, 0x65, 0x65, 0x6e, 0x20, 0x73, 0x6f, + 0x6d, 0x65, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2c, 0x0a, 0x2f, 0x2f, + 0x20, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x20, 0x66, 0x61, 0x6c, + 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x73, 0x20, 0x68, 0x65, 0x72, 0x65, 0x2e, + 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, + 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, 0x20, 0x3c, 0x20, 0x31, 0x32, + 0x30, 0x30, 0x30, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, + 0x72, 0x63, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x44, 0x73, 0x74, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, + 0x61, 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x21, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x63, 0x6f, + 0x6e, 0x76, 0x65, 0x72, 0x74, 0x69, 0x62, 0x6c, 0x65, 0x5f, 0x76, 0x3c, + 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x20, + 0x26, 0x26, 0x20, 0x21, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x3e, 0x20, 0x26, + 0x26, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x44, 0x73, 0x74, 0x54, + 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, 0x20, 0x7c, 0x7c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, + 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x44, 0x73, 0x74, 0x54, 0x2c, + 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x3e, 0x29, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, + 0x63, 0x61, 0x73, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, + 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x44, 0x73, 0x74, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x53, + 0x72, 0x63, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x44, 0x73, 0x74, 0x54, + 0x28, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, + 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x78, 0x29, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x43, 0x61, 0x73, 0x74, + 0x4f, 0x70, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x53, 0x72, 0x63, 0x54, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x44, 0x73, 0x74, 0x54, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, + 0x5f, 0x74, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x21, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x69, 0x62, + 0x6c, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x44, + 0x73, 0x74, 0x54, 0x3e, 0x20, 0x26, 0x26, 0x20, 0x21, 0x69, 0x73, 0x5f, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x53, 0x72, + 0x63, 0x54, 0x3e, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x21, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, + 0x3c, 0x44, 0x73, 0x74, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, + 0x66, 0x3e, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x21, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, + 0x44, 0x73, 0x74, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x3e, 0x20, 0x26, 0x26, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, + 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, + 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, 0x20, 0x7c, 0x7c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, + 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x72, 0x63, 0x54, 0x2c, 0x20, 0x5f, + 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, + 0x3e, 0x29, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, + 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x44, 0x73, 0x74, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x53, 0x72, 0x63, + 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x44, 0x73, 0x74, 0x54, 0x28, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, 0x66, + 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, 0x56, + 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, 0x20, 0x3c, 0x20, 0x31, 0x32, 0x30, + 0x30, 0x30, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x48, 0x65, 0x6c, 0x70, 0x65, + 0x72, 0x20, 0x74, 0x6f, 0x20, 0x64, 0x65, 0x64, 0x75, 0x63, 0x65, 0x20, + 0x74, 0x68, 0x65, 0x20, 0x53, 0x72, 0x63, 0x54, 0x2e, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x44, 0x73, 0x74, 0x54, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x72, 0x63, 0x54, + 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x61, + 0x73, 0x74, 0x5f, 0x74, 0x6f, 0x28, 0x53, 0x72, 0x63, 0x54, 0x20, 0x78, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x43, 0x61, 0x73, 0x74, 0x4f, 0x70, 0x3c, 0x53, 0x72, 0x63, 0x54, + 0x2c, 0x20, 0x44, 0x73, 0x74, 0x54, 0x3e, 0x7b, 0x7d, 0x28, 0x78, 0x29, + 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_complex[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x2f, + 0x2f, 0x20, 0x4d, 0x61, 0x6b, 0x65, 0x20, 0x6d, 0x75, 0x6c, 0x74, 0x69, + 0x70, 0x6c, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x61, 0x6e, + 0x64, 0x20, 0x64, 0x69, 0x76, 0x69, 0x73, 0x69, 0x6f, 0x6e, 0x20, 0x66, + 0x61, 0x73, 0x74, 0x65, 0x72, 0x2e, 0x0a, 0x23, 0x64, 0x65, 0x66, 0x69, + 0x6e, 0x65, 0x20, 0x4c, 0x49, 0x42, 0x43, 0x55, 0x44, 0x41, 0x43, 0x58, + 0x58, 0x5f, 0x45, 0x4e, 0x41, 0x42, 0x4c, 0x45, 0x5f, 0x53, 0x49, 0x4d, + 0x50, 0x4c, 0x49, 0x46, 0x49, 0x45, 0x44, 0x5f, 0x43, 0x4f, 0x4d, 0x50, + 0x4c, 0x45, 0x58, 0x5f, 0x4f, 0x50, 0x45, 0x52, 0x41, 0x54, 0x49, 0x4f, + 0x4e, 0x53, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, + 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, + 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x2f, 0x73, + 0x74, 0x64, 0x2f, 0x74, 0x79, 0x70, 0x65, 0x5f, 0x74, 0x72, 0x61, 0x69, + 0x74, 0x73, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, + 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x54, + 0x4f, 0x44, 0x4f, 0x3a, 0x20, 0x43, 0x6f, 0x6e, 0x73, 0x69, 0x64, 0x65, + 0x72, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x61, 0x20, 0x66, 0x61, + 0x73, 0x74, 0x65, 0x72, 0x20, 0x69, 0x6d, 0x70, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x61, 0x73, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x20, 0x68, 0x61, 0x73, 0x20, 0x74, 0x6f, + 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x6f, 0x6e, 0x66, 0x6f, 0x72, 0x6d, 0x20, + 0x74, 0x6f, 0x20, 0x43, 0x2b, 0x2b, 0x20, 0x73, 0x74, 0x61, 0x6e, 0x64, + 0x61, 0x72, 0x64, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x3c, 0x54, 0x3e, 0x3b, 0x0a, 0x0a, 0x75, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x36, 0x34, + 0x5f, 0x74, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x74, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3b, 0x0a, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x31, 0x32, 0x38, 0x5f, 0x74, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x64, 0x6f, 0x75, 0x62, 0x6c, 0x65, + 0x3e, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x69, 0x73, 0x5f, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x20, 0x3a, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, 0x61, 0x6c, + 0x73, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x20, 0x7b, 0x7d, 0x3b, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x3c, + 0x54, 0x3e, 0x3e, 0x20, 0x3a, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x72, 0x75, 0x65, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x20, 0x7b, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x76, 0x20, 0x3d, 0x20, 0x69, 0x73, 0x5f, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x20, 0x69, 0x73, 0x20, 0x6d, 0x69, 0x73, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x73, 0x6f, 0x6d, 0x65, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x25, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, + 0x20, 0x61, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x62, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x61, 0x2e, + 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x2d, 0x20, 0x66, 0x6c, 0x6f, + 0x6f, 0x72, 0x28, 0x61, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, + 0x2f, 0x20, 0x62, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, + 0x2a, 0x20, 0x62, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x54, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x61, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x20, 0x2d, 0x20, 0x66, 0x6c, 0x6f, 0x6f, 0x72, + 0x28, 0x61, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x2f, 0x20, + 0x62, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x2a, 0x20, + 0x62, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x7b, 0x72, 0x2c, 0x20, 0x69, + 0x7d, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, + 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x3e, 0x28, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x61, + 0x2c, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, + 0x54, 0x3e, 0x20, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x28, 0x61, 0x2e, 0x72, 0x65, 0x61, 0x6c, + 0x28, 0x29, 0x20, 0x3e, 0x20, 0x62, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, + 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x28, 0x61, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x62, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x20, 0x26, 0x26, 0x20, 0x61, 0x2e, 0x69, 0x6d, 0x61, + 0x67, 0x28, 0x29, 0x20, 0x3e, 0x20, 0x62, 0x2e, 0x69, 0x6d, 0x61, 0x67, + 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, + 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x3c, 0x28, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, + 0x20, 0x61, 0x2c, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, + 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x3e, 0x28, 0x62, 0x2c, 0x20, 0x61, 0x29, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, + 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x3c, 0x3d, 0x28, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x61, 0x2c, 0x20, + 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, + 0x20, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x21, 0x28, 0x61, 0x20, 0x3e, 0x20, 0x62, 0x29, 0x3b, + 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x3e, 0x3d, 0x28, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x61, 0x2c, + 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, + 0x3e, 0x20, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x21, 0x28, 0x61, 0x20, 0x3c, 0x20, 0x62, 0x29, + 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_config[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x54, 0x68, 0x69, 0x73, 0x20, 0x66, 0x69, 0x6c, 0x65, 0x20, 0x69, 0x73, + 0x20, 0x75, 0x73, 0x65, 0x64, 0x20, 0x62, 0x79, 0x20, 0x62, 0x6f, 0x74, + 0x68, 0x20, 0x43, 0x55, 0x44, 0x41, 0x20, 0x6b, 0x65, 0x72, 0x6e, 0x65, + 0x6c, 0x20, 0x63, 0x6f, 0x64, 0x65, 0x20, 0x61, 0x6e, 0x64, 0x20, 0x68, + 0x6f, 0x73, 0x74, 0x2d, 0x6f, 0x6e, 0x6c, 0x79, 0x20, 0x43, 0x2b, 0x2b, + 0x20, 0x63, 0x6f, 0x64, 0x65, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, + 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x2f, 0x2f, + 0x20, 0x54, 0x68, 0x65, 0x20, 0x6d, 0x61, 0x78, 0x69, 0x6d, 0x75, 0x6d, + 0x20, 0x64, 0x69, 0x6d, 0x65, 0x6e, 0x73, 0x69, 0x6f, 0x6e, 0x73, 0x20, + 0x6f, 0x66, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x20, 0x70, 0x61, 0x73, 0x73, 0x65, 0x64, 0x20, + 0x61, 0x73, 0x20, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, 0x70, 0x61, + 0x72, 0x61, 0x6d, 0x65, 0x74, 0x65, 0x72, 0x73, 0x2e, 0x0a, 0x23, 0x64, + 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x41, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x20, 0x31, 0x30, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x41, 0x6c, + 0x6c, 0x20, 0x65, 0x78, 0x69, 0x73, 0x74, 0x69, 0x6e, 0x67, 0x20, 0x4e, + 0x56, 0x49, 0x44, 0x49, 0x41, 0x20, 0x68, 0x61, 0x72, 0x64, 0x77, 0x61, + 0x72, 0x65, 0x20, 0x68, 0x61, 0x73, 0x20, 0x61, 0x20, 0x66, 0x69, 0x78, + 0x65, 0x64, 0x20, 0x33, 0x32, 0x20, 0x77, 0x61, 0x72, 0x70, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x2e, 0x20, 0x54, 0x68, 0x6f, 0x75, 0x67, 0x68, 0x20, + 0x61, 0x20, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x2d, 0x69, 0x6e, 0x0a, 0x2f, + 0x2f, 0x20, 0x77, 0x61, 0x72, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x20, 0x76, + 0x61, 0x72, 0x69, 0x61, 0x62, 0x6c, 0x65, 0x20, 0x65, 0x78, 0x69, 0x73, + 0x74, 0x73, 0x2c, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x69, 0x74, + 0x20, 0x77, 0x6f, 0x75, 0x6c, 0x64, 0x20, 0x70, 0x72, 0x65, 0x76, 0x65, + 0x6e, 0x74, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, 0x65, 0x2d, 0x74, + 0x69, 0x6d, 0x65, 0x20, 0x6f, 0x70, 0x74, 0x69, 0x6d, 0x69, 0x7a, 0x61, + 0x74, 0x69, 0x6f, 0x6e, 0x73, 0x2e, 0x0a, 0x23, 0x64, 0x65, 0x66, 0x69, + 0x6e, 0x65, 0x20, 0x57, 0x41, 0x52, 0x50, 0x5f, 0x53, 0x49, 0x5a, 0x45, + 0x20, 0x33, 0x32, 0x0a, 0x00, +}; + +constexpr char jit_source_fp16_math[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x5f, 0x62, 0x66, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, + 0x5f, 0x66, 0x70, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x73, 0x74, 0x64, 0x2f, 0x74, 0x79, 0x70, 0x65, 0x5f, 0x74, 0x72, 0x61, + 0x69, 0x74, 0x73, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x42, 0x69, 0x6e, 0x61, + 0x72, 0x79, 0x20, 0x6f, 0x70, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x68, + 0x61, 0x6c, 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x23, 0x64, 0x65, 0x66, + 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, + 0x4e, 0x45, 0x5f, 0x42, 0x49, 0x4e, 0x41, 0x52, 0x59, 0x5f, 0x4f, 0x50, + 0x28, 0x4e, 0x41, 0x4d, 0x45, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x5f, + 0x4f, 0x50, 0x29, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x4e, 0x41, 0x4d, 0x45, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, + 0x79, 0x29, 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, + 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, + 0x29, 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x48, 0x41, 0x4c, + 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, + 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x3e, 0x29, 0x20, 0x7b, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x78, + 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x3a, 0x3a, 0x4e, + 0x41, 0x4d, 0x45, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, + 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x42, 0x49, 0x4e, 0x41, 0x52, 0x59, + 0x5f, 0x4f, 0x50, 0x28, 0x6d, 0x61, 0x78, 0x2c, 0x20, 0x5f, 0x5f, 0x68, + 0x6d, 0x61, 0x78, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, + 0x49, 0x4e, 0x45, 0x5f, 0x42, 0x49, 0x4e, 0x41, 0x52, 0x59, 0x5f, 0x4f, + 0x50, 0x28, 0x6d, 0x69, 0x6e, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x6d, 0x69, + 0x6e, 0x29, 0x0a, 0x0a, 0x23, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x20, 0x4d, + 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x42, 0x49, + 0x4e, 0x41, 0x52, 0x59, 0x5f, 0x4f, 0x50, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x41, 0x64, 0x64, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x20, 0x43, 0x2b, 0x2b, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x20, 0x6f, 0x76, 0x65, 0x72, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x62, 0x65, 0x74, 0x77, 0x65, 0x65, + 0x6e, 0x20, 0x68, 0x61, 0x6c, 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, + 0x20, 0x61, 0x6e, 0x64, 0x20, 0x6e, 0x61, 0x74, 0x69, 0x76, 0x65, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x55, + 0x3e, 0x0a, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, + 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x65, 0x78, 0x63, 0x65, 0x70, 0x74, 0x20, + 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, + 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x20, 0x26, 0x26, + 0x20, 0x21, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, + 0x2c, 0x20, 0x55, 0x3e, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x55, 0x3e, 0x0a, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x61, + 0x72, 0x69, 0x74, 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x5f, 0x65, 0x78, + 0x63, 0x65, 0x70, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x61, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x5f, + 0x76, 0x3c, 0x54, 0x3e, 0x20, 0x26, 0x26, 0x20, 0x21, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, + 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x55, 0x3e, 0x3b, + 0x0a, 0x0a, 0x23, 0x64, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x4c, + 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, + 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x48, 0x41, 0x4c, 0x46, 0x2c, 0x20, 0x48, + 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x2c, 0x20, 0x46, + 0x4c, 0x4f, 0x41, 0x54, 0x32, 0x48, 0x41, 0x4c, 0x46, 0x2c, 0x20, 0x4f, + 0x50, 0x29, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, + 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, + 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x65, 0x78, 0x63, 0x65, 0x70, 0x74, + 0x3c, 0x54, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x3e, 0x3e, 0x3e, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x20, 0x4f, 0x50, 0x28, + 0x48, 0x41, 0x4c, 0x46, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, + 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x32, 0x48, 0x41, 0x4c, 0x46, 0x28, + 0x48, 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x28, 0x78, + 0x29, 0x20, 0x4f, 0x50, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, + 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, + 0x79, 0x29, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, + 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, + 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x65, 0x78, 0x63, 0x65, 0x70, 0x74, + 0x3c, 0x54, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x3e, 0x3e, 0x3e, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x20, 0x4f, 0x50, 0x28, + 0x54, 0x20, 0x78, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x20, 0x79, 0x29, + 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x32, 0x48, 0x41, 0x4c, 0x46, 0x28, + 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x78, 0x29, 0x20, 0x4f, 0x50, + 0x20, 0x48, 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, 0x4f, 0x41, 0x54, 0x28, + 0x79, 0x29, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x5c, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x23, 0x64, 0x65, 0x66, 0x69, + 0x6e, 0x65, 0x20, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, + 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x48, + 0x41, 0x4c, 0x46, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, + 0x4f, 0x41, 0x54, 0x2c, 0x20, 0x4f, 0x50, 0x29, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x5f, + 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x69, 0x73, 0x5f, 0x61, 0x72, 0x69, 0x74, + 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x5f, 0x65, 0x78, 0x63, 0x65, 0x70, + 0x74, 0x3c, 0x54, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x3e, 0x3e, 0x3e, + 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x20, 0x4f, 0x50, + 0x28, 0x48, 0x41, 0x4c, 0x46, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, + 0x29, 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, 0x4f, + 0x41, 0x54, 0x28, 0x78, 0x29, 0x20, 0x4f, 0x50, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x3e, 0x28, 0x79, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x5f, + 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x69, 0x73, 0x5f, 0x61, 0x72, 0x69, 0x74, + 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x5f, 0x65, 0x78, 0x63, 0x65, 0x70, + 0x74, 0x3c, 0x54, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x3e, 0x3e, 0x3e, + 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x20, 0x4f, 0x50, + 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x20, 0x79, + 0x29, 0x20, 0x7b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, + 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x79, + 0x29, 0x20, 0x4f, 0x50, 0x20, 0x48, 0x41, 0x4c, 0x46, 0x32, 0x46, 0x4c, + 0x4f, 0x41, 0x54, 0x28, 0x78, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x5c, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, + 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, + 0x66, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x32, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x2b, 0x29, 0x0a, 0x4d, 0x4c, + 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, + 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, + 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x68, + 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x2d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, + 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, + 0x4f, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x5f, + 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, + 0x20, 0x5f, 0x5f, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x68, 0x61, 0x6c, + 0x66, 0x2c, 0x20, 0x2a, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, + 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, + 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x5f, 0x5f, 0x68, + 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, + 0x5f, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x68, 0x61, 0x6c, 0x66, 0x2c, + 0x20, 0x2f, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, + 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, 0x28, 0x5f, + 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, + 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, + 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x32, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, + 0x2c, 0x20, 0x2b, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, + 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, 0x28, + 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, 0x5f, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x2c, 0x20, 0x2d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, + 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, + 0x28, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, 0x5f, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x2c, 0x20, 0x2a, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, + 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, + 0x50, 0x28, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x31, 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x5f, + 0x5f, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x32, 0x62, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x31, 0x36, 0x2c, 0x20, 0x2f, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, + 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, + 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, + 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x2c, 0x20, 0x3c, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, + 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, + 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x5f, 0x5f, 0x68, + 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3e, + 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, + 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, + 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, + 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3c, 0x3d, 0x29, 0x0a, + 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, + 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, + 0x6c, 0x66, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3e, 0x3d, 0x29, 0x0a, 0x4d, 0x4c, + 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, + 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, + 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x2c, 0x20, 0x3d, 0x3d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, + 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, + 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x2c, 0x20, + 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x2c, 0x20, 0x21, 0x3d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, + 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, + 0x50, 0x28, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x31, 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3c, + 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, + 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, + 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x2c, + 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x32, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3e, 0x29, 0x0a, 0x4d, 0x4c, + 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, + 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x2c, 0x20, 0x3c, 0x3d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, + 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, + 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x31, 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, + 0x3e, 0x3d, 0x29, 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, + 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, + 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x2c, 0x20, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x3d, 0x3d, 0x29, + 0x0a, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, + 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x43, 0x4d, 0x50, 0x28, 0x5f, 0x5f, 0x6e, + 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x2c, 0x20, + 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x32, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x21, 0x3d, 0x29, 0x0a, 0x0a, 0x23, + 0x75, 0x6e, 0x64, 0x65, 0x66, 0x20, 0x4d, 0x4c, 0x58, 0x5f, 0x44, 0x45, + 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, 0x4f, 0x50, + 0x0a, 0x23, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x20, 0x4d, 0x4c, 0x58, 0x5f, + 0x44, 0x45, 0x46, 0x49, 0x4e, 0x45, 0x5f, 0x48, 0x41, 0x4c, 0x46, 0x5f, + 0x43, 0x4d, 0x50, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_gather[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, + 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x4e, 0x49, 0x44, 0x58, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, 0x44, + 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, + 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, + 0x69, 0x64, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, + 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, + 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, + 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, + 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, + 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, + 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, + 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, + 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, + 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, + 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, + 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, + 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, + 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x25, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x65, + 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, + 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, + 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, + 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, + 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, + 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, + 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, + 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x73, 0x72, + 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, + 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, + 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_gather_axis[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, + 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x72, 0x63, 0x43, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x49, 0x64, 0x78, 0x43, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, + 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, + 0x2a, 0x20, 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x20, 0x69, + 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x70, 0x72, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, + 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, + 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, + 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, + 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, + 0x3e, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, + 0x4d, 0x3e, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, + 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x5f, 0x61, 0x78, 0x69, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, + 0x6f, 0x63, 0x54, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, + 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, + 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, + 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x20, 0x2a, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, + 0x69, 0x73, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x78, + 0x2c, 0x20, 0x79, 0x2c, 0x20, 0x7a, 0x5d, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x5f, 0x74, 0x6f, 0x5f, 0x64, 0x69, 0x6d, 0x73, 0x28, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x29, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x7a, 0x20, 0x2a, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, + 0x74, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, + 0x61, 0x78, 0x69, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x49, 0x64, + 0x78, 0x43, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, + 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, + 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, + 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, + 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, + 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x53, 0x72, 0x63, 0x43, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, + 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, + 0x4d, 0x3e, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x2b, 0x20, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, + 0x6f, 0x63, 0x54, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x3d, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, + 0x20, 0x78, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, + 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_hadamard[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, + 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, + 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, + 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x68, 0x61, 0x64, 0x61, + 0x6d, 0x61, 0x72, 0x64, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x5f, 0x6d, + 0x28, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2a, 0x20, 0x78, 0x29, 0x3b, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x4e, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x50, 0x6f, 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x61, 0x73, 0x73, + 0x65, 0x72, 0x74, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, + 0x4e, 0x20, 0x3e, 0x20, 0x30, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x28, + 0x4e, 0x20, 0x26, 0x20, 0x28, 0x4e, 0x20, 0x2d, 0x20, 0x31, 0x29, 0x29, + 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x22, 0x4e, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x62, 0x65, + 0x20, 0x61, 0x20, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x76, 0x65, 0x20, + 0x70, 0x6f, 0x77, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x20, 0x74, 0x77, 0x6f, + 0x2e, 0x22, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x20, 0x3d, 0x20, + 0x31, 0x20, 0x2b, 0x20, 0x50, 0x6f, 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, + 0x3c, 0x4e, 0x20, 0x2f, 0x20, 0x32, 0x3e, 0x3a, 0x3a, 0x76, 0x61, 0x6c, + 0x75, 0x65, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x50, 0x6f, 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, 0x3c, + 0x31, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x52, 0x3e, 0x0a, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x66, 0x6f, 0x72, 0x63, 0x65, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x68, 0x61, 0x64, 0x61, + 0x6d, 0x61, 0x72, 0x64, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x5f, 0x70, + 0x6f, 0x77, 0x32, 0x28, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2a, 0x20, 0x78, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x4c, 0x6f, 0x67, + 0x52, 0x20, 0x3d, 0x20, 0x50, 0x6f, 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, + 0x3c, 0x52, 0x3e, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x68, 0x20, 0x3d, 0x20, 0x31, 0x3b, + 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, + 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x73, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x73, 0x20, + 0x3c, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x52, 0x3b, 0x20, 0x2b, 0x2b, 0x73, + 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, + 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x52, 0x20, 0x2f, 0x20, 0x32, + 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x69, + 0x20, 0x26, 0x20, 0x28, 0x68, 0x20, 0x2d, 0x20, 0x31, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6a, 0x20, + 0x3d, 0x20, 0x28, 0x28, 0x69, 0x20, 0x2d, 0x20, 0x6b, 0x29, 0x20, 0x3c, + 0x3c, 0x20, 0x31, 0x29, 0x20, 0x2b, 0x20, 0x6b, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x61, 0x20, + 0x3d, 0x20, 0x78, 0x5b, 0x6a, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x62, 0x20, 0x3d, 0x20, + 0x78, 0x5b, 0x6a, 0x20, 0x2b, 0x20, 0x68, 0x5d, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x78, 0x5b, 0x6a, 0x5d, 0x20, 0x3d, 0x20, 0x61, + 0x20, 0x2b, 0x20, 0x62, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x78, 0x5b, 0x6a, 0x20, 0x2b, 0x20, 0x68, 0x5d, 0x20, 0x3d, 0x20, 0x61, + 0x20, 0x2d, 0x20, 0x62, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x68, 0x20, 0x3c, 0x3c, 0x3d, 0x20, 0x31, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, + 0x69, 0x78, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x65, 0x61, 0x64, + 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x3d, 0x20, 0x31, 0x3e, 0x0a, + 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, + 0x6f, 0x69, 0x64, 0x0a, 0x68, 0x61, 0x64, 0x61, 0x6d, 0x61, 0x72, 0x64, + 0x5f, 0x6e, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, + 0x69, 0x6e, 0x2c, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, + 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6e, + 0x75, 0x6d, 0x5f, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x4e, 0x75, + 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x3d, 0x20, 0x4e, + 0x20, 0x2f, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, + 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x4e, 0x20, + 0x3d, 0x20, 0x50, 0x6f, 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, 0x3c, 0x4e, + 0x3e, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x52, 0x20, 0x3d, 0x20, 0x50, 0x6f, + 0x77, 0x32, 0x4c, 0x6f, 0x67, 0x32, 0x3c, 0x6d, 0x61, 0x78, 0x5f, 0x72, + 0x61, 0x64, 0x69, 0x78, 0x3e, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, + 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x4e, 0x75, 0x6d, 0x53, 0x74, + 0x65, 0x70, 0x73, 0x20, 0x3d, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x4e, 0x20, + 0x2f, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x52, 0x3b, 0x0a, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x20, 0x3d, + 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x4e, 0x20, 0x25, 0x20, 0x6b, 0x4c, 0x6f, + 0x67, 0x52, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x46, 0x69, 0x6e, + 0x61, 0x6c, 0x52, 0x61, 0x64, 0x69, 0x78, 0x20, 0x3d, 0x20, 0x31, 0x20, + 0x3c, 0x3c, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x46, 0x69, 0x6e, 0x61, 0x6c, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x20, 0x3e, 0x3d, 0x20, + 0x6b, 0x4e, 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x73, 0x68, 0x61, 0x72, 0x65, 0x64, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x62, + 0x75, 0x66, 0x5b, 0x4e, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x49, + 0x64, 0x78, 0x2e, 0x78, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, + 0x74, 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x3d, 0x20, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x3b, 0x20, + 0x74, 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x3c, 0x20, + 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, + 0x6d, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, + 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x2b, 0x3d, 0x20, + 0x67, 0x72, 0x69, 0x64, 0x44, 0x69, 0x6d, 0x2e, 0x78, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, + 0x6e, 0x67, 0x20, 0x62, 0x61, 0x73, 0x65, 0x20, 0x3d, 0x20, 0x28, 0x74, + 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x2f, 0x20, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x29, 0x20, 0x2a, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x6c, 0x6f, 0x6e, + 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x3e, 0x28, 0x4e, 0x29, 0x20, 0x2a, + 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x2b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x74, 0x72, 0x61, 0x6e, 0x73, + 0x66, 0x6f, 0x72, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x3d, 0x3d, 0x20, 0x31, 0x29, 0x20, + 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, + 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6a, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x6a, 0x20, 0x3c, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, + 0x61, 0x64, 0x69, 0x78, 0x20, 0x2f, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x6a, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x6a, 0x20, + 0x2a, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, + 0x20, 0x2a, 0x20, 0x6b, 0x4e, 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x72, 0x65, 0x61, + 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x72, 0x20, 0x3c, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, + 0x74, 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x75, 0x66, + 0x5b, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x5b, 0x62, 0x61, 0x73, 0x65, 0x20, 0x2b, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, + 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x6a, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6a, 0x20, 0x3c, 0x20, + 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x3b, 0x20, 0x2b, + 0x2b, 0x6a, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x62, 0x75, 0x66, 0x5b, 0x6a, 0x20, 0x2a, 0x20, 0x6b, 0x4e, + 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2b, 0x20, + 0x69, 0x5d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x5b, 0x62, 0x61, 0x73, 0x65, + 0x20, 0x2b, 0x20, 0x28, 0x6a, 0x20, 0x2a, 0x20, 0x6b, 0x4e, 0x75, 0x6d, + 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2b, 0x20, 0x69, 0x29, + 0x20, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x78, 0x5b, + 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x68, 0x20, 0x3d, 0x20, + 0x31, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, + 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x73, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x73, 0x20, 0x3c, 0x20, 0x6b, 0x4e, 0x75, 0x6d, 0x53, + 0x74, 0x65, 0x70, 0x73, 0x3b, 0x20, 0x2b, 0x2b, 0x73, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, + 0x20, 0x3d, 0x20, 0x69, 0x20, 0x26, 0x20, 0x28, 0x68, 0x20, 0x2d, 0x20, + 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6a, 0x20, 0x3d, 0x20, 0x28, 0x28, 0x69, 0x20, 0x2d, 0x20, + 0x6b, 0x29, 0x20, 0x3c, 0x3c, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x52, 0x29, + 0x20, 0x2b, 0x20, 0x6b, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, + 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x72, 0x20, 0x3c, 0x20, + 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x3b, 0x20, 0x2b, + 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x78, 0x5b, 0x72, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x3e, 0x28, 0x62, 0x75, 0x66, 0x5b, 0x6a, 0x20, 0x2b, 0x20, + 0x68, 0x20, 0x2a, 0x20, 0x72, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x68, 0x61, 0x64, 0x61, 0x6d, 0x61, 0x72, 0x64, 0x5f, 0x72, 0x61, 0x64, + 0x69, 0x78, 0x5f, 0x70, 0x6f, 0x77, 0x32, 0x3c, 0x6d, 0x61, 0x78, 0x5f, + 0x72, 0x61, 0x64, 0x69, 0x78, 0x3e, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x0a, + 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, + 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x72, 0x20, 0x3c, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, + 0x69, 0x78, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x75, 0x66, 0x5b, 0x6a, + 0x20, 0x2b, 0x20, 0x68, 0x20, 0x2a, 0x20, 0x72, 0x5d, 0x20, 0x3d, 0x20, + 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x54, 0x3e, 0x28, 0x78, 0x5b, 0x72, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x68, 0x20, 0x3c, 0x3c, 0x3d, 0x20, 0x6b, 0x4c, 0x6f, 0x67, 0x52, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, + 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x6b, 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x52, 0x61, 0x64, 0x69, + 0x78, 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x74, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x74, 0x20, + 0x3c, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x20, + 0x2f, 0x20, 0x6b, 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x52, 0x61, 0x64, 0x69, + 0x78, 0x3b, 0x20, 0x2b, 0x2b, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x20, 0x2b, 0x20, 0x74, 0x20, + 0x2a, 0x20, 0x6b, 0x4e, 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x26, 0x20, 0x28, 0x68, 0x20, 0x2d, 0x20, 0x31, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6a, 0x20, 0x3d, 0x20, 0x28, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, + 0x2d, 0x20, 0x6b, 0x29, 0x20, 0x3c, 0x3c, 0x20, 0x6b, 0x4c, 0x6f, 0x67, + 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x29, 0x20, 0x2b, 0x20, 0x6b, 0x3b, 0x0a, + 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, + 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x72, 0x20, 0x3c, 0x20, 0x6b, 0x46, 0x69, 0x6e, 0x61, + 0x6c, 0x52, 0x61, 0x64, 0x69, 0x78, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x78, 0x5b, 0x72, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x3e, 0x28, 0x62, 0x75, 0x66, 0x5b, 0x6a, 0x20, 0x2b, 0x20, 0x68, + 0x20, 0x2a, 0x20, 0x72, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x68, 0x61, 0x64, 0x61, 0x6d, 0x61, 0x72, 0x64, 0x5f, + 0x72, 0x61, 0x64, 0x69, 0x78, 0x5f, 0x70, 0x6f, 0x77, 0x32, 0x3c, 0x6b, + 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x52, 0x61, 0x64, 0x69, 0x78, 0x3e, 0x28, + 0x78, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, + 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x72, 0x20, 0x3c, 0x20, + 0x6b, 0x46, 0x69, 0x6e, 0x61, 0x6c, 0x52, 0x61, 0x64, 0x69, 0x78, 0x3b, + 0x20, 0x2b, 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x75, 0x66, 0x5b, 0x6a, 0x20, + 0x2b, 0x20, 0x68, 0x20, 0x2a, 0x20, 0x72, 0x5d, 0x20, 0x3d, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x54, + 0x3e, 0x28, 0x78, 0x5b, 0x72, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, + 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x3d, 0x3d, + 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, + 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x6a, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6a, 0x20, 0x3c, 0x20, 0x6d, + 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x20, 0x2f, 0x20, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, 0x20, 0x2b, + 0x2b, 0x6a, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, + 0x3d, 0x20, 0x6a, 0x20, 0x2a, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, + 0x69, 0x64, 0x74, 0x68, 0x20, 0x2a, 0x20, 0x6b, 0x4e, 0x75, 0x6d, 0x54, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, + 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, + 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, + 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x20, 0x72, 0x20, 0x3c, 0x20, 0x72, 0x65, 0x61, 0x64, + 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x76, 0x61, 0x6c, 0x20, 0x3d, + 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, + 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x62, 0x75, 0x66, 0x5b, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, + 0x75, 0x74, 0x5b, 0x62, 0x61, 0x73, 0x65, 0x20, 0x2b, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x20, 0x3d, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x54, + 0x3e, 0x28, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x63, 0x61, 0x6c, + 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x23, + 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, + 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6a, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x6a, 0x20, 0x3c, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x61, 0x64, 0x69, + 0x78, 0x3b, 0x20, 0x2b, 0x2b, 0x6a, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x62, 0x61, + 0x73, 0x65, 0x20, 0x2b, 0x20, 0x28, 0x6a, 0x20, 0x2a, 0x20, 0x6b, 0x4e, + 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2b, 0x20, + 0x69, 0x29, 0x20, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5d, + 0x20, 0x3d, 0x20, 0x62, 0x75, 0x66, 0x5b, 0x6a, 0x20, 0x2a, 0x20, 0x6b, + 0x4e, 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2b, + 0x20, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4d, 0x2c, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, + 0x74, 0x68, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x0a, 0x68, 0x61, 0x64, 0x61, + 0x6d, 0x61, 0x72, 0x64, 0x5f, 0x6d, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x54, 0x2a, 0x20, 0x69, 0x6e, 0x2c, 0x20, 0x54, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x73, 0x63, + 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, + 0x6e, 0x67, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x61, 0x73, 0x6b, 0x73, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x54, 0x61, 0x73, + 0x6b, 0x73, 0x50, 0x65, 0x72, 0x42, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, + 0x20, 0x4e, 0x20, 0x2f, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, + 0x64, 0x74, 0x68, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x74, + 0x61, 0x73, 0x6b, 0x20, 0x3d, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, + 0x64, 0x78, 0x2e, 0x78, 0x20, 0x2a, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, + 0x44, 0x69, 0x6d, 0x2e, 0x78, 0x20, 0x2b, 0x20, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x3b, 0x20, 0x74, 0x61, 0x73, + 0x6b, 0x20, 0x3c, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x61, 0x73, 0x6b, + 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x61, + 0x73, 0x6b, 0x20, 0x2b, 0x3d, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x44, + 0x69, 0x6d, 0x2e, 0x78, 0x20, 0x2a, 0x20, 0x67, 0x72, 0x69, 0x64, 0x44, + 0x69, 0x6d, 0x2e, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x69, 0x20, + 0x3d, 0x20, 0x74, 0x61, 0x73, 0x6b, 0x20, 0x25, 0x20, 0x6b, 0x54, 0x61, + 0x73, 0x6b, 0x73, 0x50, 0x65, 0x72, 0x42, 0x61, 0x74, 0x63, 0x68, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, + 0x6e, 0x67, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x74, + 0x61, 0x73, 0x6b, 0x20, 0x2f, 0x20, 0x6b, 0x54, 0x61, 0x73, 0x6b, 0x73, + 0x50, 0x65, 0x72, 0x42, 0x61, 0x74, 0x63, 0x68, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, 0x20, + 0x62, 0x61, 0x73, 0x65, 0x20, 0x3d, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x20, 0x2a, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x3c, 0x6c, 0x6f, 0x6e, 0x67, 0x20, 0x6c, 0x6f, 0x6e, 0x67, + 0x3e, 0x28, 0x4d, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x78, 0x5b, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x5d, 0x5b, 0x4d, + 0x5d, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, + 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x63, 0x20, 0x3d, 0x20, 0x30, + 0x3b, 0x20, 0x63, 0x20, 0x3c, 0x20, 0x4d, 0x3b, 0x20, 0x2b, 0x2b, 0x63, + 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, + 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, + 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x72, 0x20, 0x3c, 0x20, 0x72, 0x65, 0x61, + 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x72, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x78, 0x5b, 0x72, 0x5d, 0x5b, 0x63, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x74, + 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x69, 0x6e, 0x5b, 0x62, 0x61, 0x73, 0x65, + 0x20, 0x2b, 0x20, 0x63, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, + 0x20, 0x2a, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, + 0x68, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, + 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x72, + 0x20, 0x3c, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, 0x64, 0x74, + 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x68, 0x61, 0x64, 0x61, 0x6d, 0x61, 0x72, 0x64, + 0x5f, 0x72, 0x61, 0x64, 0x69, 0x78, 0x5f, 0x6d, 0x28, 0x78, 0x5b, 0x72, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x23, + 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, + 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x63, 0x20, + 0x3c, 0x20, 0x4d, 0x3b, 0x20, 0x2b, 0x2b, 0x63, 0x29, 0x20, 0x7b, 0x0a, + 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, + 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x72, 0x20, 0x3c, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x77, 0x69, + 0x64, 0x74, 0x68, 0x3b, 0x20, 0x2b, 0x2b, 0x72, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, + 0x62, 0x61, 0x73, 0x65, 0x20, 0x2b, 0x20, 0x63, 0x20, 0x2a, 0x20, 0x4e, + 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x77, 0x69, 0x64, 0x74, 0x68, 0x20, 0x2b, 0x20, 0x72, 0x5d, 0x20, 0x3d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, + 0x3c, 0x54, 0x3e, 0x28, 0x78, 0x5b, 0x72, 0x5d, 0x5b, 0x63, 0x5d, 0x20, + 0x2a, 0x20, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, + 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, + 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, + 0x00, +}; + +constexpr char jit_source_indexing[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x73, 0x74, 0x64, 0x2f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3e, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x74, 0x79, 0x70, 0x65, 0x5f, 0x74, + 0x72, 0x61, 0x69, 0x74, 0x73, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x2f, + 0x2f, 0x20, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x20, 0x61, 0x6e, + 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x74, 0x6f, 0x20, 0x70, 0x6f, 0x73, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x73, 0x20, 0x69, 0x6e, 0x20, 0x61, 0x20, 0x33, 0x64, + 0x20, 0x67, 0x72, 0x69, 0x64, 0x2c, 0x20, 0x61, 0x73, 0x73, 0x75, 0x6d, + 0x69, 0x6e, 0x67, 0x20, 0x74, 0x68, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x20, 0x69, 0x73, 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x61, 0x6c, 0x63, + 0x75, 0x6c, 0x61, 0x74, 0x65, 0x64, 0x20, 0x77, 0x69, 0x74, 0x68, 0x3a, + 0x0a, 0x2f, 0x2f, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, + 0x78, 0x20, 0x2a, 0x20, 0x64, 0x69, 0x6d, 0x31, 0x20, 0x2a, 0x20, 0x64, + 0x69, 0x6d, 0x32, 0x20, 0x2b, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x64, 0x69, + 0x6d, 0x32, 0x20, 0x2b, 0x20, 0x7a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, + 0x3c, 0x54, 0x2c, 0x20, 0x54, 0x2c, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x5f, 0x74, 0x6f, 0x5f, 0x64, 0x69, 0x6d, 0x73, 0x28, + 0x54, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x64, + 0x69, 0x6d, 0x31, 0x2c, 0x20, 0x54, 0x20, 0x64, 0x69, 0x6d, 0x32, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x78, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2f, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x31, + 0x20, 0x2a, 0x20, 0x64, 0x69, 0x6d, 0x32, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x54, 0x20, 0x79, 0x20, 0x3d, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x25, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x31, 0x20, 0x2a, 0x20, 0x64, + 0x69, 0x6d, 0x32, 0x29, 0x29, 0x20, 0x2f, 0x20, 0x64, 0x69, 0x6d, 0x32, + 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x7a, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x25, 0x20, 0x64, 0x69, 0x6d, 0x32, 0x3b, 0x0a, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x2c, + 0x20, 0x7a, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x47, + 0x65, 0x74, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x20, 0x70, + 0x6f, 0x73, 0x73, 0x69, 0x62, 0x6c, 0x65, 0x20, 0x6e, 0x65, 0x67, 0x61, + 0x74, 0x69, 0x76, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2e, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, + 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x61, 0x62, 0x73, + 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x69, 0x6e, + 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x75, 0x6e, 0x73, 0x69, + 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x3e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, + 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, + 0x74, 0x3e, 0x28, 0x69, 0x64, 0x78, 0x20, 0x3c, 0x20, 0x30, 0x20, 0x3f, + 0x20, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, + 0x3a, 0x20, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, + 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, + 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, + 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, + 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, + 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x49, 0x44, + 0x58, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, + 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, + 0x74, 0x65, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, 0x70, 0x64, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x70, + 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, + 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, + 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, + 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, + 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, + 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, + 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x25, 0x20, 0x75, 0x70, 0x64, + 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x75, 0x70, 0x64, 0x5f, + 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x49, 0x44, 0x58, + 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, + 0x29, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, + 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, + 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, + 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, + 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, + 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, + 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, + 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, + 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, + 0x63, 0x5d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x3d, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, + 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, + 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2b, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2a, 0x20, 0x75, 0x70, 0x64, + 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, + 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, + 0x74, 0x61, 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, 0x20, 0x2b, + 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x75, 0x70, + 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x29, 0x3b, + 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, + 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, + 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter_axis[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, + 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, + 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, + 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, + 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, + 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, + 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x55, 0x70, 0x64, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x49, 0x64, 0x78, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, + 0x74, 0x65, 0x72, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, + 0x70, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x20, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, + 0x72, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, + 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, + 0x73, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, + 0x49, 0x4d, 0x3e, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, + 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, + 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, + 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, + 0x20, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, + 0x69, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, + 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, + 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, + 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x20, 0x2a, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, + 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, + 0x6f, 0x73, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x78, 0x2c, 0x20, 0x79, + 0x2c, 0x20, 0x7a, 0x5d, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x5f, 0x74, 0x6f, 0x5f, 0x64, 0x69, 0x6d, 0x73, 0x28, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x5f, 0x61, 0x78, 0x69, 0x73, 0x2c, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x72, 0x65, 0x29, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x3d, 0x20, 0x7a, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, 0x61, 0x78, 0x69, + 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x49, 0x64, 0x78, 0x43, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, 0x78, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, + 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, + 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, + 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, + 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, + 0x5d, 0x2c, 0x20, 0x61, 0x78, 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x79, 0x20, 0x2a, + 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x5f, + 0x61, 0x78, 0x69, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x55, 0x70, + 0x64, 0x43, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x61, 0x78, 0x69, 0x73, 0x20, 0x2b, 0x20, + 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, + 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, + 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, + 0x20, 0x2b, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, + 0x75, 0x70, 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, + 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, + 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x61, 0x74, 0x6f, + 0x6d, 0x69, 0x63, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, + 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, + 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x41, 0x73, 0x73, 0x69, 0x67, + 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x2a, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x53, 0x75, + 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x64, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x61, 0x78, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x6f, 0x75, 0x74, + 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, + 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x69, 0x6e, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, + 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, + 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, + 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_ternary_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x23, 0x70, 0x72, 0x61, + 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, 0x65, 0x6c, 0x65, + 0x63, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x63, 0x6f, 0x6e, 0x64, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x2c, 0x20, 0x54, + 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x6f, + 0x6e, 0x64, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x3f, 0x20, 0x78, 0x20, + 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_unary_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x66, 0x70, 0x31, + 0x36, 0x5f, 0x6d, 0x61, 0x74, 0x68, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, + 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, + 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, + 0x5f, 0x66, 0x70, 0x38, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, + 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x6d, 0x61, 0x74, 0x68, 0x5f, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x73, 0x2e, 0x68, 0x3e, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x63, 0x6d, 0x61, 0x74, 0x68, + 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, + 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x41, 0x62, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, + 0x75, 0x6e, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x54, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x62, + 0x73, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x41, 0x72, 0x63, 0x43, 0x6f, 0x73, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x63, 0x6f, 0x73, 0x28, + 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x41, 0x72, 0x63, 0x43, 0x6f, + 0x73, 0x68, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x61, 0x63, 0x6f, 0x73, 0x68, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x41, 0x72, 0x63, 0x53, 0x69, 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x73, 0x69, 0x6e, 0x28, 0x78, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x41, 0x72, 0x63, 0x53, 0x69, 0x6e, 0x68, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x73, + 0x69, 0x6e, 0x68, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x41, + 0x72, 0x63, 0x54, 0x61, 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x61, 0x74, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x41, 0x72, 0x63, 0x54, 0x61, 0x6e, 0x68, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x74, 0x61, 0x6e, + 0x68, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, + 0x77, 0x69, 0x73, 0x65, 0x49, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7e, 0x78, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x43, 0x65, 0x69, 0x6c, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, + 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x54, 0x7b, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, + 0x65, 0x69, 0x6c, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x29, 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x63, 0x65, 0x69, 0x6c, 0x28, 0x78, 0x2e, 0x69, 0x6d, 0x61, + 0x67, 0x28, 0x29, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x65, 0x69, 0x6c, + 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x43, 0x6f, 0x6e, 0x6a, 0x75, 0x67, 0x61, 0x74, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, + 0x3c, 0x54, 0x3e, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, + 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x6e, 0x6a, 0x28, + 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x43, 0x6f, 0x73, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x73, 0x28, + 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x43, 0x6f, 0x73, 0x68, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x73, + 0x68, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x72, 0x66, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, + 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x65, 0x72, 0x66, 0x28, 0x5f, 0x5f, + 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, + 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x65, 0x72, 0x66, 0x28, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x32, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x65, 0x72, 0x66, 0x28, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x72, 0x66, + 0x49, 0x6e, 0x76, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, + 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, + 0x61, 0x6c, 0x66, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x65, 0x72, 0x66, + 0x69, 0x6e, 0x76, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, + 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x65, 0x72, 0x66, 0x69, 0x6e, 0x76, + 0x28, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x32, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x65, 0x72, 0x66, 0x69, 0x6e, 0x76, 0x28, 0x78, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x78, 0x70, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x78, + 0x70, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x78, 0x70, + 0x6d, 0x31, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x65, 0x78, 0x70, 0x6d, 0x31, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x46, 0x6c, 0x6f, 0x6f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, + 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, + 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, + 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x54, 0x7b, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, 0x6c, 0x6f, + 0x6f, 0x72, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, + 0x2c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x66, 0x6c, 0x6f, 0x6f, 0x72, 0x28, 0x78, 0x2e, 0x69, 0x6d, 0x61, + 0x67, 0x28, 0x29, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, 0x6c, 0x6f, 0x6f, + 0x72, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x49, 0x6d, 0x61, 0x67, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, + 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, 0x69, 0x6d, + 0x61, 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6c, 0x6f, + 0x67, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, + 0x32, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x79, 0x20, 0x3d, 0x20, 0x4c, 0x6f, 0x67, 0x7b, 0x7d, 0x28, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, + 0x29, 0x20, 0x2f, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, 0x4c, + 0x4e, 0x32, 0x5f, 0x46, 0x2c, 0x20, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, + 0x28, 0x29, 0x20, 0x2f, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, + 0x4c, 0x4e, 0x32, 0x5f, 0x46, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6c, 0x6f, 0x67, + 0x32, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x31, 0x30, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, + 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6c, 0x6f, 0x67, 0x31, 0x30, 0x28, + 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x31, 0x70, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x7a, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, + 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x20, 0x78, 0x20, 0x3d, 0x20, 0x7a, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x20, 0x79, 0x20, 0x3d, 0x20, 0x7a, 0x2e, 0x69, 0x6d, 0x61, + 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x7a, 0x61, 0x62, 0x73, 0x20, 0x3d, 0x20, + 0x41, 0x62, 0x73, 0x7b, 0x7d, 0x28, 0x7a, 0x29, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x74, 0x68, 0x65, 0x74, 0x61, 0x20, 0x3d, + 0x20, 0x61, 0x74, 0x61, 0x6e, 0x32, 0x66, 0x28, 0x79, 0x2c, 0x20, 0x78, + 0x20, 0x2b, 0x20, 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x7a, 0x61, 0x62, 0x73, 0x20, 0x3c, 0x20, + 0x30, 0x2e, 0x35, 0x66, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x72, 0x20, + 0x3d, 0x20, 0x78, 0x20, 0x2a, 0x20, 0x28, 0x32, 0x20, 0x2b, 0x20, 0x78, + 0x29, 0x20, 0x2b, 0x20, 0x79, 0x20, 0x2a, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, + 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x20, 0x2f, 0x2f, 0x20, + 0x68, 0x61, 0x6e, 0x64, 0x6c, 0x65, 0x20, 0x75, 0x6e, 0x64, 0x65, 0x72, + 0x66, 0x6c, 0x6f, 0x77, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x78, + 0x2c, 0x20, 0x74, 0x68, 0x65, 0x74, 0x61, 0x7d, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, + 0x30, 0x2e, 0x35, 0x66, 0x20, 0x2a, 0x20, 0x6c, 0x6f, 0x67, 0x31, 0x70, + 0x66, 0x28, 0x72, 0x29, 0x2c, 0x20, 0x74, 0x68, 0x65, 0x74, 0x61, 0x7d, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x7a, 0x30, 0x20, 0x3d, 0x20, + 0x68, 0x79, 0x70, 0x6f, 0x74, 0x66, 0x28, 0x78, 0x20, 0x2b, 0x20, 0x31, + 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x6c, 0x6f, + 0x67, 0x66, 0x28, 0x7a, 0x30, 0x29, 0x2c, 0x20, 0x74, 0x68, 0x65, 0x74, + 0x61, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6c, 0x6f, 0x67, 0x31, 0x70, 0x28, 0x7a, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, + 0x63, 0x61, 0x6c, 0x4e, 0x6f, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x21, + 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4e, 0x65, 0x67, 0x61, 0x74, 0x69, + 0x76, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x54, 0x7b, 0x30, 0x2c, 0x20, 0x30, 0x7d, 0x20, + 0x2d, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, 0x78, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x52, 0x65, 0x61, 0x6c, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, + 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x72, 0x69, 0x6e, + 0x74, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x2c, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x72, 0x69, 0x6e, 0x74, 0x28, 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, + 0x29, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x72, 0x69, 0x6e, 0x74, 0x28, 0x78, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x53, 0x69, 0x67, 0x6d, 0x6f, 0x69, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, + 0x20, 0x79, 0x20, 0x3d, 0x20, 0x31, 0x20, 0x2f, 0x20, 0x28, 0x31, 0x20, + 0x2b, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x65, 0x78, 0x70, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x62, 0x73, 0x28, 0x78, 0x29, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x28, 0x78, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x20, 0x3f, 0x20, 0x79, + 0x20, 0x3a, 0x20, 0x31, 0x20, 0x2d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x69, 0x67, 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x75, 0x6e, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, + 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x21, 0x3d, 0x20, + 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x78, 0x2e, 0x72, + 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x20, 0x26, + 0x26, 0x20, 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x3d, + 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2f, 0x20, 0x41, + 0x62, 0x73, 0x28, 0x29, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, + 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x3e, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, + 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x28, 0x78, 0x20, 0x3e, + 0x20, 0x54, 0x28, 0x30, 0x2e, 0x66, 0x29, 0x29, 0x20, 0x2d, 0x20, 0x28, + 0x78, 0x20, 0x3c, 0x20, 0x54, 0x28, 0x30, 0x2e, 0x66, 0x29, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x28, 0x78, 0x20, 0x3e, 0x20, 0x54, 0x28, 0x30, + 0x29, 0x29, 0x20, 0x2d, 0x20, 0x28, 0x78, 0x20, 0x3c, 0x20, 0x54, 0x28, + 0x30, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x53, 0x69, 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x73, 0x69, 0x6e, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x53, 0x69, 0x6e, 0x68, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, + 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x73, 0x69, 0x6e, 0x68, 0x28, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x53, 0x71, 0x75, 0x61, 0x72, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2a, + 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, 0x71, 0x72, 0x74, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x73, 0x71, 0x72, + 0x74, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x52, 0x73, 0x71, + 0x72, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x31, 0x2e, 0x30, 0x66, 0x20, 0x2f, 0x20, 0x53, + 0x71, 0x72, 0x74, 0x7b, 0x7d, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, + 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, + 0x73, 0x71, 0x72, 0x74, 0x28, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x32, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, + 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x73, 0x71, 0x72, 0x74, + 0x28, 0x5f, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x32, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x72, 0x73, 0x71, 0x72, 0x74, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x54, 0x61, 0x6e, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, + 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x61, 0x6e, + 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x54, 0x61, 0x6e, 0x68, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x61, + 0x6e, 0x68, 0x28, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x54, 0x6f, + 0x46, 0x50, 0x38, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x38, + 0x5f, 0x74, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x5f, 0x5f, 0x6e, 0x76, + 0x5f, 0x66, 0x70, 0x38, 0x5f, 0x65, 0x34, 0x6d, 0x33, 0x28, 0x78, 0x29, + 0x2e, 0x5f, 0x5f, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x46, 0x72, 0x6f, + 0x6d, 0x46, 0x50, 0x38, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, + 0x28, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, 0x20, 0x78, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x28, 0x2a, 0x28, 0x5f, 0x5f, 0x6e, + 0x76, 0x5f, 0x66, 0x70, 0x38, 0x5f, 0x65, 0x34, 0x6d, 0x33, 0x2a, 0x29, + 0x28, 0x26, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_utils[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x54, 0x68, 0x69, 0x73, 0x20, 0x66, 0x69, 0x6c, 0x65, 0x20, 0x6d, 0x75, + 0x73, 0x74, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x69, 0x6e, 0x63, 0x6c, 0x75, + 0x64, 0x65, 0x20, 0x61, 0x6e, 0x79, 0x20, 0x68, 0x6f, 0x73, 0x74, 0x2d, + 0x6f, 0x6e, 0x6c, 0x79, 0x20, 0x63, 0x6f, 0x64, 0x65, 0x2c, 0x20, 0x75, + 0x74, 0x69, 0x6c, 0x69, 0x74, 0x69, 0x65, 0x73, 0x20, 0x74, 0x68, 0x61, + 0x74, 0x20, 0x77, 0x6f, 0x72, 0x6b, 0x20, 0x75, 0x6e, 0x64, 0x65, 0x72, + 0x20, 0x62, 0x6f, 0x74, 0x68, 0x0a, 0x2f, 0x2f, 0x20, 0x68, 0x6f, 0x73, + 0x74, 0x20, 0x61, 0x6e, 0x64, 0x20, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x20, 0x63, 0x61, 0x6e, 0x20, 0x62, 0x65, 0x20, 0x70, 0x75, 0x74, 0x20, + 0x68, 0x65, 0x72, 0x65, 0x2e, 0x0a, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, + 0x53, 0x65, 0x65, 0x20, 0x6d, 0x6f, 0x72, 0x65, 0x20, 0x61, 0x62, 0x6f, + 0x75, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x72, 0x65, 0x71, 0x75, 0x69, + 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x20, 0x61, 0x74, 0x3a, 0x0a, + 0x2f, 0x2f, 0x20, 0x68, 0x74, 0x74, 0x70, 0x73, 0x3a, 0x2f, 0x2f, 0x64, + 0x6f, 0x63, 0x73, 0x2e, 0x6e, 0x76, 0x69, 0x64, 0x69, 0x61, 0x2e, 0x63, + 0x6f, 0x6d, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, 0x6e, 0x76, 0x72, 0x74, + 0x63, 0x2f, 0x23, 0x6c, 0x61, 0x6e, 0x67, 0x75, 0x61, 0x67, 0x65, 0x0a, + 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, + 0x65, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, + 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, + 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x2f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x2e, 0x63, 0x75, 0x68, + 0x22, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, + 0x6d, 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, + 0x63, 0x75, 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, + 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x2e, 0x68, 0x22, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x5f, 0x62, 0x66, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, + 0x5f, 0x66, 0x70, 0x31, 0x36, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x73, 0x74, 0x64, 0x2f, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3e, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, + 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, + 0x63, 0x75, 0x64, 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x74, 0x75, 0x70, + 0x6c, 0x65, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, + 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x43, 0x55, 0x44, 0x41, 0x20, + 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x2f, 0x2f, + 0x20, 0x54, 0x6f, 0x20, 0x70, 0x61, 0x73, 0x73, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x74, + 0x6f, 0x20, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x73, 0x20, 0x76, 0x69, + 0x61, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x20, 0x6d, + 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x2c, 0x20, 0x74, 0x68, 0x65, 0x69, 0x72, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x62, + 0x65, 0x0a, 0x2f, 0x2f, 0x20, 0x6b, 0x6e, 0x6f, 0x77, 0x6e, 0x20, 0x61, + 0x74, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, 0x65, 0x20, 0x74, 0x69, + 0x6d, 0x65, 0x2e, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x68, + 0x61, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, + 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4d, 0x41, 0x58, 0x5f, + 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x3b, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x3d, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, + 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, + 0x20, 0x4d, 0x41, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x3b, 0x0a, + 0x0a, 0x2f, 0x2f, 0x20, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x69, 0x7a, + 0x65, 0x64, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x2f, 0x73, 0x74, 0x6f, 0x72, + 0x65, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x61, 0x6c, 0x69, 0x67, 0x6e, 0x61, 0x73, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x28, 0x54, 0x29, 0x20, 0x2a, 0x20, 0x4e, + 0x29, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x76, 0x61, + 0x6c, 0x5b, 0x4e, 0x5d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x26, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x5b, 0x5d, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x76, 0x61, 0x6c, 0x5b, 0x69, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x5b, 0x5d, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x76, + 0x61, 0x6c, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x69, 0x73, 0x5f, 0x61, 0x6c, 0x69, 0x67, 0x6e, + 0x65, 0x64, 0x28, 0x54, 0x2a, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x28, 0x72, 0x65, 0x69, + 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x70, 0x74, 0x72, 0x5f, 0x74, 0x3e, + 0x28, 0x78, 0x29, 0x20, 0x25, 0x20, 0x28, 0x4e, 0x20, 0x2a, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x28, 0x54, 0x29, 0x29, 0x29, 0x20, 0x3d, + 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, + 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, + 0x4e, 0x3e, 0x20, 0x75, 0x6e, 0x73, 0x61, 0x66, 0x65, 0x5f, 0x6c, 0x6f, + 0x61, 0x64, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, + 0x70, 0x74, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, + 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, + 0x66, 0x72, 0x6f, 0x6d, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, + 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, + 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, + 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x5b, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5d, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, + 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x41, + 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, + 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x5f, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x70, 0x74, 0x72, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x61, 0x6c, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x3c, 0x4e, 0x3e, 0x28, 0x70, 0x74, 0x72, + 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x2a, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x20, 0x3d, 0x20, 0x72, 0x65, + 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x3c, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x41, 0x6c, 0x69, + 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, + 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x66, 0x72, 0x6f, 0x6d, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, + 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, + 0x20, 0x76, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, + 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x3b, 0x20, 0x2b, 0x2b, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x76, + 0x5b, 0x69, 0x5d, 0x20, 0x3d, 0x20, 0x70, 0x74, 0x72, 0x5b, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x76, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x69, 0x7a, 0x65, + 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x41, 0x6c, 0x69, + 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, + 0x2c, 0x20, 0x4e, 0x3e, 0x0a, 0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x76, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, + 0x2a, 0x20, 0x70, 0x74, 0x72, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, + 0x53, 0x69, 0x7a, 0x65, 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, + 0x54, 0x20, 0x66, 0x61, 0x6c, 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x61, + 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x3c, 0x4e, 0x3e, 0x28, 0x70, 0x74, + 0x72, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x20, 0x2b, 0x20, 0x31, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x3c, + 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x66, 0x72, 0x6f, 0x6d, + 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, + 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, + 0x70, 0x74, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x5b, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, + 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x20, 0x76, 0x3b, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, + 0x4e, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x76, 0x5b, 0x69, 0x5d, 0x20, 0x3d, 0x20, 0x28, + 0x4e, 0x20, 0x2a, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, + 0x20, 0x69, 0x29, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3f, + 0x20, 0x70, 0x74, 0x72, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, + 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x5d, 0x20, 0x3a, 0x20, 0x66, + 0x61, 0x6c, 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x76, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, + 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x53, 0x69, 0x7a, 0x65, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, + 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x20, 0x6c, + 0x6f, 0x61, 0x64, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, + 0x20, 0x70, 0x74, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, + 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x53, 0x69, 0x7a, 0x65, 0x54, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x66, 0x61, 0x6c, + 0x6c, 0x62, 0x61, 0x63, 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x61, 0x6c, 0x69, 0x67, 0x6e, 0x65, + 0x64, 0x3c, 0x4e, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x20, 0x26, 0x26, + 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x3d, 0x3d, 0x20, 0x31, + 0x20, 0x26, 0x26, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, + 0x2b, 0x20, 0x31, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x3c, 0x3d, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x20, 0x3d, + 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, + 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, + 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, + 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x66, 0x72, 0x6f, 0x6d, 0x5b, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, 0x69, 0x67, + 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, + 0x20, 0x4e, 0x3e, 0x20, 0x76, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, + 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x3b, + 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x76, 0x5b, 0x69, 0x5d, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x4e, 0x20, 0x2a, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x20, + 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3f, 0x20, 0x70, 0x74, 0x72, + 0x5b, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x2a, 0x20, 0x28, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, + 0x69, 0x29, 0x5d, 0x20, 0x3a, 0x20, 0x66, 0x61, 0x6c, 0x6c, 0x62, 0x61, + 0x63, 0x6b, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x76, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x0a, 0x75, + 0x6e, 0x73, 0x61, 0x66, 0x65, 0x5f, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x54, 0x2a, 0x20, 0x70, 0x74, + 0x72, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x26, 0x20, 0x76, + 0x65, 0x63, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x2a, 0x20, 0x74, 0x6f, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, + 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, + 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, + 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x74, 0x6f, 0x5b, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x5d, 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x0a, 0x73, 0x74, 0x6f, 0x72, 0x65, + 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x54, 0x2a, 0x20, 0x70, + 0x74, 0x72, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x26, 0x20, + 0x76, 0x65, 0x63, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x69, 0x73, 0x5f, 0x61, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x3c, + 0x4e, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x74, 0x6f, 0x20, + 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, + 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x41, 0x6c, 0x69, 0x67, 0x6e, + 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, + 0x4e, 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x74, 0x6f, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x5d, 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x23, 0x70, 0x72, 0x61, + 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, + 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x70, 0x74, 0x72, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x5d, 0x20, 0x3d, + 0x20, 0x76, 0x65, 0x63, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, + 0x4e, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x53, 0x69, 0x7a, 0x65, 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, + 0x65, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x2a, 0x20, 0x70, 0x74, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x26, 0x20, + 0x76, 0x65, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x53, 0x69, 0x7a, + 0x65, 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x61, 0x6c, 0x69, 0x67, + 0x6e, 0x65, 0x64, 0x3c, 0x4e, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, 0x20, + 0x26, 0x26, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, + 0x20, 0x31, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x3c, 0x3d, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x74, 0x6f, 0x20, 0x3d, 0x20, 0x72, 0x65, + 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x3c, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, + 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, + 0x28, 0x70, 0x74, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, + 0x6f, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5d, 0x20, 0x3d, 0x20, + 0x76, 0x65, 0x63, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x28, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, + 0x2b, 0x20, 0x69, 0x29, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, + 0x26, 0x26, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x3b, 0x20, 0x2b, 0x2b, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x70, + 0x74, 0x72, 0x5b, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2a, 0x20, + 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x5d, 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, + 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x69, 0x7a, 0x65, + 0x54, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, + 0x64, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x76, 0x65, 0x63, 0x74, + 0x6f, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x70, + 0x74, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x41, + 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, + 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x26, 0x20, 0x76, 0x65, 0x63, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x53, 0x69, 0x7a, 0x65, 0x54, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x73, 0x5f, + 0x61, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x3c, 0x4e, 0x3e, 0x28, 0x70, + 0x74, 0x72, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x20, 0x2b, 0x20, 0x31, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x20, + 0x3c, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x26, 0x26, 0x20, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x20, 0x3d, 0x3d, 0x20, 0x31, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, + 0x74, 0x6f, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, + 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x56, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, + 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, 0x70, 0x74, 0x72, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x6f, 0x5b, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x5d, 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x20, + 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x26, 0x26, 0x20, 0x69, 0x20, + 0x3c, 0x20, 0x4e, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x70, 0x74, 0x72, 0x5b, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x20, 0x2a, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x5d, + 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x54, + 0x79, 0x70, 0x65, 0x20, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x75, + 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x76, + 0x6f, 0x69, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x54, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, + 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, + 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, + 0x61, 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, + 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x3e, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, + 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x64, 0x6f, 0x75, 0x62, 0x6c, 0x65, 0x3e, + 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x61, + 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, + 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, + 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, + 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x28, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, + 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, + 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, + 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, + 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, + 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, + 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, + 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x43, 0x55, 0x44, 0x41, 0x20, 0x31, 0x31, 0x20, 0x64, 0x6f, 0x65, 0x73, + 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x68, 0x61, 0x76, 0x65, 0x20, 0x68, 0x6f, + 0x73, 0x74, 0x20, 0x73, 0x69, 0x64, 0x65, 0x20, 0x61, 0x72, 0x69, 0x74, + 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x68, 0x61, 0x6c, + 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, + 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, + 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, + 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, + 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, + 0x36, 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, + 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, + 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, + 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, + 0x44, 0x41, 0x52, 0x54, 0x5f, 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, + 0x20, 0x3c, 0x20, 0x31, 0x32, 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, 0x20, + 0x5f, 0x5f, 0x43, 0x55, 0x44, 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, 0x5f, + 0x5f, 0x20, 0x3c, 0x20, 0x38, 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, + 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, + 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6c, 0x73, 0x65, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, + 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, + 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, + 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, 0x66, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, + 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, + 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, + 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, + 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, + 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, + 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, 0x20, 0x3c, 0x20, 0x31, 0x32, + 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, 0x20, 0x5f, 0x5f, 0x43, 0x55, 0x44, + 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, 0x5f, 0x5f, 0x20, 0x3c, 0x20, 0x38, + 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, + 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, + 0x6c, 0x73, 0x65, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, 0x77, + 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, + 0x66, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, + 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, + 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x74, + 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, + 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x66, 0x61, 0x6c, 0x73, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, + 0x74, 0x73, 0x3c, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, + 0x3c, 0x54, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, + 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x6f, + 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x6d, + 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x4c, 0x69, 0x6d, 0x69, 0x74, + 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x2c, + 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, + 0x6d, 0x61, 0x78, 0x28, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, + 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, + 0x3c, 0x54, 0x3e, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, + 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, + 0x69, 0x6e, 0x28, 0x29, 0x2c, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, + 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x49, 0x6e, 0x64, 0x65, + 0x78, 0x69, 0x6e, 0x67, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, + 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x0a, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, + 0x6f, 0x63, 0x28, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, + 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x20, 0x3e, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, 0x63, 0x20, + 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x2a, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, + 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x2f, 0x2f, 0x20, 0x4f, 0x70, 0x74, 0x69, 0x6d, 0x69, 0x7a, 0x65, 0x20, + 0x77, 0x68, 0x65, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x6e, 0x64, 0x69, + 0x6d, 0x20, 0x69, 0x73, 0x20, 0x6b, 0x6e, 0x6f, 0x77, 0x6e, 0x20, 0x61, + 0x74, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, 0x65, 0x20, 0x74, 0x69, + 0x6d, 0x65, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, + 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, + 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x49, 0x64, 0x78, 0x54, 0x0a, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x28, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x23, + 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, + 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, + 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, + 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x2a, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, 0x7d, + 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, + 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, + 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, + 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, + 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, + 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, + 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, + 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, + 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, 0x63, + 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, + 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, + 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, + 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, + 0x6e, 0x64, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, + 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, + 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x63, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, + 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, + 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, + 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, + 0x20, 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, + 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, + 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, 0x6c, 0x6f, 0x63, + 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, 0x63, + 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x63, 0x5f, 0x6c, + 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, + 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, + 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, + 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, + 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, + 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, + 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, + 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, + 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, + 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, + 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, + 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, + 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x29, + 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, + 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, + 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, + 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x2a, 0x20, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x64, + 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, + 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, + 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, + 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, + 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, + 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x63, + 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x45, 0x6c, 0x65, + 0x6d, 0x20, 0x74, 0x6f, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x69, 0x6e, 0x20, + 0x61, 0x20, 0x6c, 0x6f, 0x6f, 0x70, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, + 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x47, 0x65, + 0x6e, 0x65, 0x72, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, 0x65, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x74, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, + 0x4c, 0x6f, 0x63, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, + 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, 0x44, + 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x2c, 0x20, 0x47, 0x65, 0x6e, 0x65, + 0x72, 0x61, 0x6c, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, + 0x3e, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x3b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, 0x3b, + 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, + 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x3a, 0x20, 0x64, 0x69, + 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, 0x2c, 0x20, 0x69, 0x6e, 0x6e, 0x65, + 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x28, 0x64, 0x69, 0x6d, + 0x20, 0x2d, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, + 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, + 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, + 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, + 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, + 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, + 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, + 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6e, + 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, + 0x20, 0x31, 0x5d, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, + 0x31, 0x5d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3d, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2d, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x65, 0x78, + 0x74, 0x72, 0x61, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x6e, + 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6e, 0x65, + 0x78, 0x74, 0x28, 0x31, 0x20, 0x2b, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, + 0x20, 0x2f, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, + 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x78, 0x74, 0x72, + 0x61, 0x20, 0x3d, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x25, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, + 0x31, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, + 0x70, 0x65, 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, + 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3e, 0x20, 0x30, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6e, 0x65, + 0x78, 0x74, 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x2c, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x74, + 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, + 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, + 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, + 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x29, + 0x20, 0x3a, 0x20, 0x64, 0x69, 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, 0x20, + 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, + 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, + 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, + 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, 0x5d, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, + 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, + 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, + 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, + 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, + 0x6c, 0x6f, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, + 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x66, 0x61, + 0x6c, 0x73, 0x65, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, + 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, + 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x29, 0x20, + 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, + 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, + 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x5b, 0x30, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, 0x6f, + 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +} // namespace mlx::core diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift new file mode 100644 index 000000000..42bd75313 --- /dev/null +++ b/Source/Encuda/encuda-compile.swift @@ -0,0 +1,96 @@ +import ArgumentParser +import Foundation + +extension Encuda { + struct Compile: ParsableCommand { + static let configuration = CommandConfiguration( + commandName: "compile" + ) + + @Option(name: .customLong("nvcc"), help: "Path to nvcc") + var nvccPath: String? = nil + + @Option(name: .customLong("clangpp"), help: "Path to clang++") + var clangppPath: String? = nil + + @Option( + name: .customShort("I"), parsing: .unconditionalSingleValue, help: "Include directories" + ) + var includeDirs: [String] = [] + + @Option(name: .customShort("o"), help: "Output file path") + var output: String + + @Argument(help: "Input .cu files to compile") + var inputFiles: [String] + + @Option(name: .customLong("std"), help: "C++ standard to use (e.g., c++17, c++20)") + var std: String? = nil + + @Flag(name: .customShort("v"), help: "Enable verbose output") + var verbose: Bool = false + + @Flag(name: .customLong("incremental"), help: "Skip compilation if output is up to date") + var incremental: Bool = false + + mutating func run() throws { + + if verbose { + print("Running encuda compile") + } + + if incremental && isUpToDate() { + if verbose { + print("Output is up to date, skipping compilation") + } + return + } + + let resolvedNvcc: String + if let path = nvccPath { + resolvedNvcc = path + } else if let found = searchForCommand("nvcc") { + resolvedNvcc = found.path + } else { + fatalError("nvcc not found") + } + + let includeArgs = includeDirs.flatMap { ["-I", $0] } + let ccbinArgs = clangppPath.map { ["-ccbin=\($0)"] } ?? [] + let stdArgs = std.map { ["-std=\($0)"] } ?? [] + let archArgs = + ProcessInfo.processInfo.environment["CUDA_ARCH"].map { ["-arch", $0] } ?? [] + + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedNvcc) + process.arguments = + ["-cuda", "-rdc=true", "--expt-relaxed-constexpr"] + stdArgs + ccbinArgs + archArgs + + (verbose ? ["-v"] : []) + includeArgs + inputFiles + ["-o", output] + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.nvccFailed(process.terminationStatus) + } + } + + private func isUpToDate() -> Bool { + let fm = FileManager.default + let outputURL = URL(fileURLWithPath: output) + guard fm.fileExists(atPath: output), + let outputMod = + (try? outputURL.resourceValues(forKeys: [.contentModificationDateKey]))? + .contentModificationDate + else { return false } + for input in inputFiles { + guard fm.fileExists(atPath: input), + let inputMod = + (try? URL(fileURLWithPath: input).resourceValues(forKeys: [ + .contentModificationDateKey + ]))?.contentModificationDate + else { return false } + if inputMod >= outputMod { return false } + } + return true + } + } +} diff --git a/Source/Encuda/encuda-error.swift b/Source/Encuda/encuda-error.swift new file mode 100644 index 000000000..e3dd42c9f --- /dev/null +++ b/Source/Encuda/encuda-error.swift @@ -0,0 +1,13 @@ +import Foundation + +enum EncudaError: Error, CustomStringConvertible { + case nvccFailed(Int32) + case clangFailed(Int32) + + var description: String { + switch self { + case .nvccFailed(let code): return "nvcc failed with exit code \(code)" + case .clangFailed(let code): return "clang++ failed with exit code \(code)" + } + } +} diff --git a/Source/Encuda/encuda-link.swift b/Source/Encuda/encuda-link.swift new file mode 100644 index 000000000..92710e320 --- /dev/null +++ b/Source/Encuda/encuda-link.swift @@ -0,0 +1,124 @@ +import ArgumentParser +import Foundation + +extension Encuda { + struct Link: ParsableCommand { + static let configuration = CommandConfiguration( + commandName: "link" + ) + + @Option(name: .customLong("nvcc"), help: "Path to nvcc") + var nvccPath: String? = nil + + @Option(name: .customLong("clangpp"), help: "Path to clang++") + var clangppPath: String? = nil + + @Argument(help: "Input .cpp files to link") + var inputFiles: [String] + + @Option(name: .customShort("o"), help: "Output file path") + var output: String + + @Option(name: .customLong("std"), help: "C++ standard to use (e.g., c++17, c++20)") + var std: String? = nil + + @Flag(name: .customShort("v"), help: "Enable verbose output") + var verbose: Bool = false + + @Flag(name: .customLong("incremental"), help: "Skip link if output is up to date") + var incremental: Bool = false + + mutating func run() throws { + if verbose { + print("CUDA Link") + print("Input files: \(inputFiles)") + print("Output file: \(output)") + } + + if incremental && isUpToDate() { + if verbose { + print("Output is up to date, skipping link") + } + return + } + + let stdArgs = std.map { ["-std=\($0)"] } ?? [] + let archArgs = + ProcessInfo.processInfo.environment["CUDA_ARCH"].map { ["-arch", $0] } ?? [] + + for input in inputFiles { + try clang(args: stdArgs + ["-c", input, "-o", input + ".cudalink.o"]) + } + + try nvcc( + args: ["--device-link"] + stdArgs + archArgs + inputFiles.map { $0 + ".cudalink.o" } + + ["-o", output, "-Xcompiler", "-E"]) + + for input in inputFiles { + try? FileManager.default.removeItem(atPath: input + ".cudalink.o") + } + } + + var resolvedNvcc: String { + if let path = nvccPath { + return path + } else if let found = searchForCommand("nvcc") { + return found.path + } else { + fatalError("nvcc not found") + } + } + + var resolvedClangpp: String { + if let path = clangppPath { + return path + } else if let found = searchForCommand("clang++") { + return found.path + } else { + fatalError("clang++ not found") + } + } + + func nvcc(args: [String]) throws { + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedNvcc) + process.arguments = ["-ccbin=\(resolvedClangpp)"] + (verbose ? ["-v"] : []) + args + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.nvccFailed(process.terminationStatus) + } + } + + func clang(args: [String]) throws { + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedClangpp) + process.arguments = (verbose ? ["-v"] : []) + args + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.clangFailed(process.terminationStatus) + } + } + + private func isUpToDate() -> Bool { + let fm = FileManager.default + let outputURL = URL(fileURLWithPath: output) + guard fm.fileExists(atPath: output), + let outputMod = + (try? outputURL.resourceValues(forKeys: [.contentModificationDateKey]))? + .contentModificationDate + else { return false } + for input in inputFiles { + guard fm.fileExists(atPath: input), + let inputMod = + (try? URL(fileURLWithPath: input).resourceValues(forKeys: [ + .contentModificationDateKey + ]))?.contentModificationDate + else { return false } + if inputMod >= outputMod { return false } + } + return true + } + } +} diff --git a/Source/Encuda/encuda-utils.swift b/Source/Encuda/encuda-utils.swift new file mode 100644 index 000000000..4f08a2033 --- /dev/null +++ b/Source/Encuda/encuda-utils.swift @@ -0,0 +1,25 @@ +import Foundation + +func searchForCommand(_ name: String) -> URL? { + let path = ProcessInfo.processInfo.environment["PATH"] ?? "" + for folder in path.split(separator: ":") { + let url = URL(fileURLWithPath: String(folder)).appendingPathComponent(name) + if FileManager.default.isExecutableFile(atPath: url.path) { + return url + } + } + return nil +} + +extension Process { + /** + * When running on aarch64 Ubuntu 24 docker on Mac, it happens that + * the `waitUntilExit()` method never returns. + * We adopted this workaround. + */ + func waitUntilExitWorkaround() { + while isRunning { + Thread.sleep(forTimeInterval: 0.05) + } + } +} diff --git a/Source/Encuda/encuda.swift b/Source/Encuda/encuda.swift new file mode 100644 index 000000000..c6497d37f --- /dev/null +++ b/Source/Encuda/encuda.swift @@ -0,0 +1,8 @@ +import ArgumentParser + +@main +struct Encuda: ParsableCommand { + static let configuration = CommandConfiguration( + subcommands: [Compile.self, Link.self] + ) +} diff --git a/Source/MLX/GPU+CUDA.swift b/Source/MLX/GPU+CUDA.swift new file mode 100644 index 000000000..325006316 --- /dev/null +++ b/Source/MLX/GPU+CUDA.swift @@ -0,0 +1,23 @@ +// Copyright © 2024 Apple Inc. + +import Cmlx +import Foundation + +/// API for controlling GPU related features. +/// +/// Note: previously this also had properties that managed the buffer use -- those are now +/// found in ``Memory`` but remain here as deprecated. +/// +/// ### See Also +/// - +/// - ``Memory`` +public enum GPU { + + /// Returns GPU's recommended working set size in bytes as an `Int`. + /// + /// This value is derived from ``DeviceInfo/maxRecommendedWorkingSetSize`` and + /// is clamped to `Int.max` when necessary. Returns `nil` when unavailable. + public static func maxRecommendedWorkingSetBytes() -> Int? { + return nil + } +} diff --git a/Source/MLX/MLXFastKernel.swift b/Source/MLX/MLXFastKernel.swift index 037149130..47dc66698 100644 --- a/Source/MLX/MLXFastKernel.swift +++ b/Source/MLX/MLXFastKernel.swift @@ -16,170 +16,219 @@ extension Bool: KernelTemplateArg {} extension Int: KernelTemplateArg {} extension DType: KernelTemplateArg {} -extension MLXFast { - - /// Container for a kernel created by - /// ``callAsFunction(_:template:grid:threadGroup:outputShapes:outputDTypes:initValue:verbose:stream:)`` - /// - /// The ``callAsFunction(_:template:grid:threadGroup:outputShapes:outputDTypes:initValue:verbose:stream:)`` can be used to evaluate the kernel with inputs: - /// - /// ```swift - /// let a = normal([2, 2]) - /// let kernel = MLXFast.metalKernel( - /// name: "basic", - /// inputNames: ["a"], - /// outputNames: ["out1"], - /// source: """ - /// uint elem = thread_position_in_grid.x; - /// out1[elem] = a[elem]; - /// """, - /// grid: (4, 1, 1), - /// threadGroup: (2, 1, 1), - /// outputShapes: [[2, 2]], - /// outputDTypes: [.float32]) - /// - /// let out = kernel([a]) - /// ``` - final public class MLXFastKernel: @unchecked Sendable { - let kernel: mlx_fast_metal_kernel - public let outputNames: [String] - - init( - name: String, inputNames: some Sequence, outputNames: some Sequence, - source: String, header: String = "", - ensureRowContiguous: Bool = true, - atomicOutputs: Bool = false - ) { - self.outputNames = Array(outputNames) +#if os(macOS) || os(iOS) || os(tvOS) || os(visionOS) + + extension MLXFast { + + /// Container for a kernel created by + /// ``callAsFunction(_:template:grid:threadGroup:outputShapes:outputDTypes:initValue:verbose:stream:)`` + /// + /// The ``callAsFunction(_:template:grid:threadGroup:outputShapes:outputDTypes:initValue:verbose:stream:)`` can be used to evaluate the kernel with inputs: + /// + /// ```swift + /// let a = normal([2, 2]) + /// let kernel = MLXFast.metalKernel( + /// name: "basic", + /// inputNames: ["a"], + /// outputNames: ["out1"], + /// source: """ + /// uint elem = thread_position_in_grid.x; + /// out1[elem] = a[elem]; + /// """, + /// grid: (4, 1, 1), + /// threadGroup: (2, 1, 1), + /// outputShapes: [[2, 2]], + /// outputDTypes: [.float32]) + /// + /// let out = kernel([a]) + /// ``` + final public class MLXFastKernel: @unchecked Sendable { + let kernel: mlx_fast_metal_kernel + public let outputNames: [String] + + init( + name: String, inputNames: some Sequence, outputNames: some Sequence, + source: String, header: String = "", + ensureRowContiguous: Bool = true, + atomicOutputs: Bool = false + ) { + self.outputNames = Array(outputNames) + + let input_names = mlx_vector_string_new() + defer { mlx_vector_string_free(input_names) } + for name in inputNames { + mlx_vector_string_append_value(input_names, name) + } + + let output_names = mlx_vector_string_new() + defer { mlx_vector_string_free(output_names) } + for name in self.outputNames { + mlx_vector_string_append_value(output_names, name) + } - let input_names = mlx_vector_string_new() - defer { mlx_vector_string_free(input_names) } - for name in inputNames { - mlx_vector_string_append_value(input_names, name) + self.kernel = mlx_fast_metal_kernel_new( + name.cString(using: .utf8), + input_names, output_names, + source.cString(using: .utf8), + header.cString(using: .utf8), + ensureRowContiguous, atomicOutputs) } - let output_names = mlx_vector_string_new() - defer { mlx_vector_string_free(output_names) } - for name in self.outputNames { - mlx_vector_string_append_value(output_names, name) + deinit { + mlx_fast_metal_kernel_free(kernel) } - self.kernel = mlx_fast_metal_kernel_new( - name.cString(using: .utf8), - input_names, output_names, - source.cString(using: .utf8), - header.cString(using: .utf8), - ensureRowContiguous, atomicOutputs) - } + /// Call the prepared metal kernel. + /// + /// See ``MLXFastKernel`` for example. Use + /// ``metalKernel(name:inputNames:outputNames:source:header:ensureRowContiguous:atomicOutputs:)`` + /// to create an instance. + /// + /// - Parameters: + /// - inputs: inputs passed to the metal kernel + /// - template: template arguments + /// - grid: 3-tuple specifying the grid to launch the kernel with + /// - threadGroup: 3-tuple specifying the threadgroup size to use + /// - outputShapes: list of shapes for each output in ``outputNames`` + /// - outputDTypes: list of data types for each output in ``outputNames`` + /// - initValue: optional value to use to initialize all of the output arrays + /// - verbose: if true will print the full generated source code of the kernel when run + /// - stream: stream to run on + /// - Returns: array of `MLXArray` + public func callAsFunction( + _ inputs: [any ScalarOrArray], + template: [(String, any KernelTemplateArg)]? = nil, + grid: (Int, Int, Int), + threadGroup: (Int, Int, Int), + outputShapes: some Sequence<[Int]>, + outputDTypes: some Sequence, + initValue: Float? = nil, + verbose: Bool = false, + stream: StreamOrDevice = .default + ) -> [MLXArray] { + let config = mlx_fast_metal_kernel_config_new() + defer { mlx_fast_metal_kernel_config_free(config) } + + if let template { + for (name, arg) in template { + switch arg { + case let value as Bool: + mlx_fast_metal_kernel_config_add_template_arg_bool(config, name, value) + + case let value as Int: + mlx_fast_metal_kernel_config_add_template_arg_int( + config, name, Int32(value)) + + case let value as DType: + mlx_fast_metal_kernel_config_add_template_arg_dtype( + config, name, value.cmlxDtype) + + default: + fatalError( + "Unable to handle KernelTemplateArg \(name) with type: \(type(of: arg))." + ) + } + } + } + + mlx_fast_metal_kernel_config_set_grid( + config, grid.0.int32, grid.1.int32, grid.2.int32) + mlx_fast_metal_kernel_config_set_thread_group( + config, threadGroup.0.int32, threadGroup.1.int32, threadGroup.2.int32) + + for (shape, dtype) in zip(outputShapes, outputDTypes) { + mlx_fast_metal_kernel_config_add_output_arg( + config, shape.map { Int32($0) }, shape.count, dtype.cmlxDtype) + } + + if let initValue { + mlx_fast_metal_kernel_config_set_init_value(config, initValue) + } + + mlx_fast_metal_kernel_config_set_verbose(config, verbose) + + let inputs = new_mlx_vector_array(inputs.map { $0.asMLXArray(dtype: nil) }) + defer { mlx_vector_array_free(inputs) } - deinit { - mlx_fast_metal_kernel_free(kernel) + var result = mlx_vector_array_new() + mlx_fast_metal_kernel_apply(&result, kernel, inputs, config, stream.ctx) + defer { mlx_vector_array_free(result) } + + return mlx_vector_array_values(result) + } } - /// Call the prepared metal kernel. - /// - /// See ``MLXFastKernel`` for example. Use - /// ``metalKernel(name:inputNames:outputNames:source:header:ensureRowContiguous:atomicOutputs:)`` - /// to create an instance. + /// A jit-compiled custom Metal kernel defined from a source string. /// /// - Parameters: - /// - inputs: inputs passed to the metal kernel - /// - template: template arguments - /// - grid: 3-tuple specifying the grid to launch the kernel with - /// - threadGroup: 3-tuple specifying the threadgroup size to use - /// - outputShapes: list of shapes for each output in ``outputNames`` - /// - outputDTypes: list of data types for each output in ``outputNames`` - /// - initValue: optional value to use to initialize all of the output arrays - /// - verbose: if true will print the full generated source code of the kernel when run - /// - stream: stream to run on - /// - Returns: array of `MLXArray` - public func callAsFunction( - _ inputs: [any ScalarOrArray], - template: [(String, any KernelTemplateArg)]? = nil, - grid: (Int, Int, Int), - threadGroup: (Int, Int, Int), - outputShapes: some Sequence<[Int]>, - outputDTypes: some Sequence, - initValue: Float? = nil, - verbose: Bool = false, - stream: StreamOrDevice = .default - ) -> [MLXArray] { - let config = mlx_fast_metal_kernel_config_new() - defer { mlx_fast_metal_kernel_config_free(config) } - - if let template { - for (name, arg) in template { - switch arg { - case let value as Bool: - mlx_fast_metal_kernel_config_add_template_arg_bool(config, name, value) - - case let value as Int: - mlx_fast_metal_kernel_config_add_template_arg_int( - config, name, Int32(value)) - - case let value as DType: - mlx_fast_metal_kernel_config_add_template_arg_dtype( - config, name, value.cmlxDtype) - - default: - fatalError( - "Unable to handle KernelTemplateArg \(name) with type: \(type(of: arg))." - ) - } - } - } + /// - name: name for the kernel + /// - inputNames: parameter names of the inputs in the function signature + /// - outputNames: parameter names of the outputs in the function signature + /// - source: source code -- this is the body of a function in Metal, + /// the function signature will be automatically generated. + /// - header: header source code to include before the main function. Useful + /// for helper functions or includes that should live outside of the main function body. + /// - ensureRowContiguous: whether to ensure the inputs are row contiguous + /// before the kernel runs (at a performance cost) + /// - atomicOutputs: whether to use atomic outputs in the function signature, + /// e.g. `device atomic` + /// - Returns: an ``MLXFastKernel`` -- see that for information on how to call it + public static func metalKernel( + name: String, inputNames: some Sequence, outputNames: some Sequence, + source: String, header: String = "", ensureRowContiguous: Bool = true, + atomicOutputs: Bool = false + ) -> MLXFastKernel { + MLXFastKernel( + name: name, inputNames: inputNames, outputNames: outputNames, + source: source, header: header, + ensureRowContiguous: ensureRowContiguous, atomicOutputs: atomicOutputs + ) + } - mlx_fast_metal_kernel_config_set_grid(config, grid.0.int32, grid.1.int32, grid.2.int32) - mlx_fast_metal_kernel_config_set_thread_group( - config, threadGroup.0.int32, threadGroup.1.int32, threadGroup.2.int32) + } // MLXFast - for (shape, dtype) in zip(outputShapes, outputDTypes) { - mlx_fast_metal_kernel_config_add_output_arg( - config, shape.map { Int32($0) }, shape.count, dtype.cmlxDtype) - } +#else - if let initValue { - mlx_fast_metal_kernel_config_set_init_value(config, initValue) - } + // Stub allowing compilation for platforms without Metal. - mlx_fast_metal_kernel_config_set_verbose(config, verbose) + extension MLXFast { - let inputs = new_mlx_vector_array(inputs.map { $0.asMLXArray(dtype: nil) }) - defer { mlx_vector_array_free(inputs) } + final public class MLXFastKernel: @unchecked Sendable { + public let outputNames: [String] - var result = mlx_vector_array_new() - mlx_fast_metal_kernel_apply(&result, kernel, inputs, config, stream.ctx) - defer { mlx_vector_array_free(result) } + init( + name: String, inputNames: some Sequence, outputNames: some Sequence, + source: String, header: String = "", + ensureRowContiguous: Bool = true, + atomicOutputs: Bool = false + ) { + self.outputNames = [] + fatalError("MLXFastKernel is not available without Metal") + } - return mlx_vector_array_values(result) + public func callAsFunction( + _ inputs: [any ScalarOrArray], + template: [(String, any KernelTemplateArg)]? = nil, + grid: (Int, Int, Int), + threadGroup: (Int, Int, Int), + outputShapes: some Sequence<[Int]>, + outputDTypes: some Sequence, + initValue: Float? = nil, + verbose: Bool = false, + stream: StreamOrDevice = .default + ) -> [MLXArray] { + fatalError("MLXFastKernel is not available without Metal") + } } - } - - /// A jit-compiled custom Metal kernel defined from a source string. - /// - /// - Parameters: - /// - name: name for the kernel - /// - inputNames: parameter names of the inputs in the function signature - /// - outputNames: parameter names of the outputs in the function signature - /// - source: source code -- this is the body of a function in Metal, - /// the function signature will be automatically generated. - /// - header: header source code to include before the main function. Useful - /// for helper functions or includes that should live outside of the main function body. - /// - ensureRowContiguous: whether to ensure the inputs are row contiguous - /// before the kernel runs (at a performance cost) - /// - atomicOutputs: whether to use atomic outputs in the function signature, - /// e.g. `device atomic` - /// - Returns: an ``MLXFastKernel`` -- see that for information on how to call it - public static func metalKernel( - name: String, inputNames: some Sequence, outputNames: some Sequence, - source: String, header: String = "", ensureRowContiguous: Bool = true, - atomicOutputs: Bool = false - ) -> MLXFastKernel { - MLXFastKernel( - name: name, inputNames: inputNames, outputNames: outputNames, - source: source, header: header, - ensureRowContiguous: ensureRowContiguous, atomicOutputs: atomicOutputs - ) - } - -} // MLXFast + + public static func metalKernel( + name: String, inputNames: some Sequence, outputNames: some Sequence, + source: String, header: String = "", ensureRowContiguous: Bool = true, + atomicOutputs: Bool = false + ) -> MLXFastKernel { + fatalError("MLXFastKernel is not available without Metal") + } + + } // MLXFast + +#endif diff --git a/tools/update-mlx.sh b/tools/update-mlx.sh index 940a3f333..fc1ab9615 100755 --- a/tools/update-mlx.sh +++ b/tools/update-mlx.sh @@ -14,7 +14,7 @@ fi rm -f Source/Cmlx/include/mlx/c/* cp Source/Cmlx/mlx-c/mlx/c/*.h Source/Cmlx/include/mlx/c -# run the command to do the build-time code generation +# run the command to do the build-time code generation for Metal mkdir build cd build @@ -73,12 +73,21 @@ make \ cd ../../.. make cpu_compiled_preamble +# run the command to do the build-time code generation for CUDA +cmake \ + -DMLX_SOURCE_ROOT="../Source/Cmlx/mlx/mlx/backend/cuda" \ + -DMLX_JIT_SOURCES="device/atomic_ops.cuh:device/binary_ops.cuh:device/cast_op.cuh:device/complex.cuh:device/config.h:device/fp16_math.cuh:device/gather.cuh:device/gather_axis.cuh:device/hadamard.cuh:device/indexing.cuh:device/scatter.cuh:device/scatter_axis.cuh:device/scatter_ops.cuh:device/ternary_ops.cuh:device/unary_ops.cuh:device/utils.cuh" \ + -P "../Source/Cmlx/mlx/mlx/backend/cuda/bin2h.cmake" + cd .. rm -rf Source/Cmlx/mlx-generated/metal +rm -rf Source/Cmlx/mlx-generated/cuda rm -f Source/Cmlx/mlx-generated/* +mkdir -p Source/Cmlx/mlx-generated/cuda cp build/mlx/backend/metal/jit/* Source/Cmlx/mlx-generated cp build/mlx/backend/cpu/compiled_preamble.cpp Source/Cmlx/mlx-generated +cp build/gen/cuda_jit_sources.h Source/Cmlx/mlx-generated/cuda # we don't need the cmake build directory any more rm -rf build diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 3bfc24928..87b93ba14 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -102,6 +102,7 @@ isa = PBXFileSystemSynchronizedBuildFileExceptionSet; membershipExceptions = ( Documentation.docc, + "GPU+CUDA.swift", ); target = C3AE8EC32EAAA15F000BD280 /* MLX */; }; From 683df31c5b81f13c95bec2b7c8f36d7f869934c3 Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Tue, 30 Jun 2026 11:46:18 -0500 Subject: [PATCH 07/39] Add MultiOptimizer (port from Python mlx.optimizers) (#431) Port `mlx.optimizers.MultiOptimizer` to Swift: a meta-optimizer that routes each model parameter to the first sub-optimizer whose filter predicate matches its flattened key path, with the last optimizer acting as the unconditioned fallback. This allows using different optimizers or hyperparameters for different subsets of a model's weights. `update(model:gradients:)` partitions the flattened gradient tree by first-matching filter and delegates each subset to the corresponding child optimizer; `innerState()` surfaces the children's state. It conforms to the `Optimizer` protocol directly since it holds no per-parameter state of its own. Closes #228 Co-authored-by: Claude Opus 4.8 (1M context) --- Source/MLXOptimizers/Optimizers.swift | 64 +++++++++++++++++++++++++++ Tests/MLXTests/OptimizerTests.swift | 26 +++++++++++ 2 files changed, 90 insertions(+) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index 5d4a08fb1..5b3ce3968 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -15,6 +15,70 @@ public protocol Optimizer: Updatable, Evaluatable { func update(model: Module, gradients: ModuleParameters) } +/// An optimizer that delegates to several sub-optimizers, routing each parameter to the first +/// optimizer whose filter matches. +/// +/// This makes it easy to use different optimizers (or different hyperparameters) for different +/// weights. The `filters` are predicates over a parameter's flattened key path and its gradient +/// and must be one fewer than `optimizers` -- the last optimizer is the fallback and receives +/// every parameter not matched by an earlier filter. +/// +/// ```swift +/// // Adam for everything except biases, which use SGD. +/// let optimizer = MultiOptimizer( +/// optimizers: [SGD(learningRate: 0.1), Adam(learningRate: 1e-3)], +/// filters: [{ key, _ in key.hasSuffix(".bias") }]) +/// ``` +/// +/// This is a port of Python `mlx.optimizers.MultiOptimizer`. +/// +/// ### See Also +/// - +open class MultiOptimizer: Optimizer { + + /// The sub-optimizers; the last one is the fallback for unmatched parameters. + public let optimizers: [Optimizer] + + /// Predicates over `(key, gradient)`, one per optimizer. The provided filters select the first + /// `optimizers.count - 1`; the last is an implicit always-true fallback. + private let filters: [(String, MLXArray) -> Bool] + + /// - Parameters: + /// - optimizers: the optimizers to delegate to (at least one) + /// - filters: predicates selecting which parameters route to each optimizer; there must be + /// `optimizers.count - 1` of them (the last optimizer is the unconditioned fallback) + public init(optimizers: [Optimizer], filters: [(String, MLXArray) -> Bool] = []) { + precondition(!optimizers.isEmpty, "MultiOptimizer requires at least one optimizer") + precondition( + filters.count == optimizers.count - 1, + "MultiOptimizer given \(filters.count) filters but \(optimizers.count - 1) needed") + self.optimizers = optimizers + self.filters = filters + [{ _, _ in true }] + } + + /// Partition the flattened gradients into one group per optimizer, by first matching filter. + private func split(_ gradients: ModuleParameters) -> [[(String, MLXArray)]] { + var parts = Array(repeating: [(String, MLXArray)](), count: optimizers.count) + for (key, gradient) in gradients.flattened() { + for (i, filter) in filters.enumerated() where filter(key, gradient) { + parts[i].append((key, gradient)) + break + } + } + return parts + } + + public func update(model: Module, gradients: ModuleParameters) { + for (optimizer, part) in zip(optimizers, split(gradients)) where !part.isEmpty { + optimizer.update(model: model, gradients: ModuleParameters.unflattened(part)) + } + } + + public func innerState() -> [MLXArray] { + optimizers.flatMap { $0.innerState() } + } +} + /// The base class for all optimizers. It allows us to implement an optimizer on a per-parameter basis /// and apply it to a parameter tree. /// diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index d340d46fd..bcee1b322 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -213,4 +213,30 @@ class OptimizerTests: XCTestCase { checkTrain(optimizer: Adafactor(learningRate: 0.1), compile: true) } + class TwoParameterModel: Module { + let weight = MLXArray.zeros([3]) + let bias = MLXArray.zeros([3]) + } + + func testMultiOptimizer() { + let model = TwoParameterModel() + let grads = model.parameters().mapValues { MLXArray.ones(like: $0) } + + // Route `bias` to a high learning-rate SGD; every other parameter falls back to the + // low learning-rate SGD. Distinct results prove each parameter was routed correctly. + let optimizer = MultiOptimizer( + optimizers: [SGD(learningRate: 1.0), SGD(learningRate: 0.1)], + filters: [{ key, _ in key == "bias" }]) + + optimizer.update(model: model, gradients: grads) + eval(model) + + // plain SGD from zeros with unit gradients: new = -learningRate + assertEqual(model.bias, MLXArray([Float(-1.0), -1.0, -1.0]), atol: 1e-6) + assertEqual(model.weight, MLXArray([Float(-0.1), -0.1, -0.1]), atol: 1e-6) + + // innerState surfaces the sub-optimizers' state. + XCTAssertEqual(optimizer.innerState().count, 2) + } + } From bb0399d30ae1aef9e0ca1d40959673e7d4f06c3f Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Tue, 30 Jun 2026 09:50:54 -0700 Subject: [PATCH 08/39] Add MLXArray and DType extensions (#429) * Add MLXArray and DType extensions * Simplify magic numbers * Return MLXArray instead of Double --- Source/MLX/DType.swift | 71 ++++++++++++++---------------- Source/MLX/MLXArray+maskFill.swift | 14 ++++++ skills/mlx-swift/SKILL.md | 34 ++++++++++++++ 3 files changed, 80 insertions(+), 39 deletions(-) create mode 100644 Source/MLX/MLXArray+maskFill.swift diff --git a/Source/MLX/DType.swift b/Source/MLX/DType.swift index 772e619e4..8912d3ca4 100644 --- a/Source/MLX/DType.swift +++ b/Source/MLX/DType.swift @@ -113,6 +113,21 @@ public enum DType: Hashable, Sendable, CaseIterable { mlx_dtype_size(cmlxDtype) } + /// A scalar `MLXArray` containing the largest finite value for this dtype's + /// real floating-point component. + /// + /// Use this when the dtype is statically known to be floating point, e.g. + /// `DType.float16.greatestFiniteMagnitudeArray`. + /// + /// - Precondition: `self` is a floating-point dtype. Calling this on a + /// non-floating-point dtype is a programming error and traps. + public var greatestFiniteMagnitudeArray: MLXArray { + guard let info = finfo else { + fatalError("greatestFiniteMagnitudeArray requires a floating-point dtype, got \(self)") + } + return MLXArray(info.max, dtype: self) + } + /// For floating point values return the floating point info, similar to `numpy.finfo`. public var finfo: FInfo? { isFloatingPoint ? FInfo(dtype: self) : nil @@ -130,33 +145,17 @@ public enum DType: Hashable, Sendable, CaseIterable { #if !arch(x86_64) case .float16: Double(Float16.ulpOfOne) #else - case .float16: 0.000977 + case .float16: 0x1p-10 // 2^-10 #endif case .float32: Double(Float.ulpOfOne) - case .bfloat16: 0.0078125 - case .complex64: Double.ulpOfOne - case .float64: Double.ulpOfOne - default: - fatalError("\(dtype) is not a floating point type") + case .bfloat16: 0x1p-7 // 2^-7 (7 mantissa bits) + case .complex64, .float64: Double.ulpOfOne + default: fatalError("\(dtype) is not a floating point type") } } /// The smallest representable number - public var min: Double { - switch dtype { - #if !arch(x86_64) - case .float16: -Double(Float16.greatestFiniteMagnitude) - #else - case .float16: -65500.0 - #endif - case .float32: -Double(Float.greatestFiniteMagnitude) - case .bfloat16: -3.3895313892515355e+38 - case .complex64: -Double.greatestFiniteMagnitude - case .float64: -Double.greatestFiniteMagnitude - default: - fatalError("\(dtype) is not a floating point type") - } - } + public var min: Double { -max } /// The largest representable number public var max: Double { @@ -164,14 +163,12 @@ public enum DType: Hashable, Sendable, CaseIterable { #if !arch(x86_64) case .float16: Double(Float16.greatestFiniteMagnitude) #else - case .float16: 65500.0 + case .float16: 0x1.FFCp15 // 65504 #endif + case .bfloat16: Double(Float(bitPattern: 0x7F7F_0000)) // bf16 = high 16 bits of f32 case .float32: Double(Float.greatestFiniteMagnitude) - case .bfloat16: 3.3895313892515355e+38 - case .complex64: Double.greatestFiniteMagnitude - case .float64: Double.greatestFiniteMagnitude - default: - fatalError("\(dtype) is not a floating point type") + case .complex64, .float64: Double.greatestFiniteMagnitude + default: fatalError("\(dtype) is not a floating point type") } } @@ -181,14 +178,12 @@ public enum DType: Hashable, Sendable, CaseIterable { #if !arch(x86_64) case .float16: Double(Float16.leastNormalMagnitude) #else - case .float16: 6.104e-05 + case .float16: 0x1p-14 // 2^-14 #endif case .float32: Double(Float.leastNormalMagnitude) - case .bfloat16: 1.1754943508222875e-38 - case .complex64: Double.leastNormalMagnitude - case .float64: Double.leastNormalMagnitude - default: - fatalError("\(dtype) is not a floating point type") + case .bfloat16: 0x1p-126 // 2^-126 (same exponent range as f32) + case .complex64, .float64: Double.leastNormalMagnitude + default: fatalError("\(dtype) is not a floating point type") } } @@ -198,14 +193,12 @@ public enum DType: Hashable, Sendable, CaseIterable { #if !arch(x86_64) case .float16: Double(Float16.leastNonzeroMagnitude) #else - case .float16: 6e-08 + case .float16: 0x1p-24 // 2^-24 #endif case .float32: Double(Float.leastNonzeroMagnitude) - case .bfloat16: 1.1754943508222875e-38 - case .complex64: Double.leastNonzeroMagnitude - case .float64: Double.leastNonzeroMagnitude - default: - fatalError("\(dtype) is not a floating point type") + case .bfloat16: 0x1p-133 // 2^-133 = 2^-126 · 2^-7 + case .complex64, .float64: Double.leastNonzeroMagnitude + default: fatalError("\(dtype) is not a floating point type") } } } diff --git a/Source/MLX/MLXArray+maskFill.swift b/Source/MLX/MLXArray+maskFill.swift new file mode 100644 index 000000000..c05d88903 --- /dev/null +++ b/Source/MLX/MLXArray+maskFill.swift @@ -0,0 +1,14 @@ +// Copyright © 2026 Apple Inc. + +import Foundation + +extension MLXArray { + /// Scalar masked-score fill (`-finfo(dtype).max`) constructed directly in + /// `dtype`, so masked positions vanish under softmax — no `asType` to forget. + public static func maskFill(for dtype: DType) -> MLXArray { + guard let info = dtype.finfo else { + fatalError("maskFill requires a floating-point dtype, got \(dtype)") + } + return -dtype.greatestFiniteMagnitudeArray + } +} diff --git a/skills/mlx-swift/SKILL.md b/skills/mlx-swift/SKILL.md index 3606a85b6..ba6fb71c4 100644 --- a/skills/mlx-swift/SKILL.md +++ b/skills/mlx-swift/SKILL.md @@ -49,6 +49,8 @@ Cmlx (C/C++ bindings, Metal GPU) | Custom kernels | Source/MLX/MLXFastKernel.swift | | Wired memory coordinator | Source/MLX/WiredMemory.swift | | GPU working-set helper | Source/MLX/GPU+Metal.swift | +| DType & numeric limits (finfo) | Source/MLX/DType.swift | +| Attention mask fill | Source/MLX/MLXArray+maskFill.swift | ## Quick Start @@ -79,6 +81,38 @@ array.dtype // .int64 array.count // 3 (first dimension) ``` +### DType Numeric Limits (`finfo`) + +Floating-point limits for a `DType`, analogous to `numpy.finfo`. `finfo` is **optional** — `nil` for non-floating-point dtypes, so integer dtypes are rejected safely instead of returning a bogus value: + +```swift +if let info = DType.float16.finfo { // FInfo? — nil for non-float dtypes + info.max // 65504.0 (largest finite value) + info.min // -65504.0 (== -max) + info.eps // ulp of 1.0 (float16: 2^-10) + info.smallestNormal // (float16: 2^-14) + info.smallestSubnormal // (float16: 2^-24) +} + +DType.int32.finfo // nil +``` + +When the dtype is statically known to be floating point, use the non-optional `greatestFiniteMagnitude` (mirrors `Float.greatestFiniteMagnitude`; **traps** on a non-float dtype): + +```swift +let bound = DType.float16.greatestFiniteMagnitude // 65504.0 (Double) +``` + +### Attention Mask Fill (`maskFill`) + +`MLXArray.maskFill(for:)` builds the value for masked-out positions before a softmax — `-finfo(dtype).max`, constructed directly in `dtype` so masked scores vanish under softmax with no `asType` to forget: + +```swift +scores = MLX.where(causalMask, scores, MLXArray.maskFill(for: scores.dtype)) +``` + +Requires a floating-point dtype (traps otherwise) — attention scores are float even when the KV cache is quantized. + ### Basic Operations ```swift From 1a864ce7579d53f99cdba33ea003ba8fe60ac636 Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Tue, 30 Jun 2026 12:18:12 -0500 Subject: [PATCH 09/39] Add Muon optimizer (port from Python mlx.optimizers) (#432) Port `mlx.optimizers.Muon` (MomentUm Orthogonalized by Newton-Schulz): a momentum optimizer that orthogonalizes the update for >=2D parameters via a 5-step quintic Newton-Schulz iteration, scaling the learning rate by sqrt(max(1, rows/cols)). 0D/1D parameters fall back to a plain momentum/Nesterov update. Mirrors the reference (weight decay, Nesterov default on, ns_steps=5). Tests pin the 1D momentum/Nesterov math exactly and exercise the 2D orthogonalization path (shape preserved, finite, parameter moved). Co-authored-by: Claude Opus 4.8 (1M context) --- Source/MLXOptimizers/Optimizers.swift | 95 +++++++++++++++++++++++++++ Tests/MLXTests/OptimizerTests.swift | 29 ++++++++ 2 files changed, 124 insertions(+) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index 5b3ce3968..a01083b85 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -767,6 +767,101 @@ open class Adafactor: OptimizerBase { } } +/// The Muon (MomentUm Orthogonalized by Newton-schulz) optimizer. +/// +/// Follows the original implementation: [Muon: An optimizer for hidden layers in +/// neural networks](https://kellerjordan.github.io/posts/muon/). +/// +/// - Note: Muon may be sub-optimal for the embedding layer, the final fully +/// connected layer, or any 0D/1D parameters; optimize those with a different +/// method (e.g. ``AdamW``). For parameters with more than 2 dimensions (e.g. 4D +/// convolution filters) the trailing dimensions are flattened. +/// +/// ### See Also +/// - +open class Muon: OptimizerBaseArrayState { + + /// The learning rate + public var learningRate: Float + /// The momentum strength + public var momentum: Float = 0.95 + /// The weight decay (L2 penalty) + public var weightDecay: Float = 0.01 + /// Enables Nesterov momentum (recommended) + public var nesterov = true + /// Number of Newton-Schulz iteration steps for orthogonalization + public var nsSteps: Int = 5 + + /// Initialize the optimizer. + /// - Parameters: + /// - learningRate: the learning rate + /// - momentum: the momentum strength + /// - weightDecay: the weight decay (L2 penalty) + /// - nesterov: enables Nesterov momentum + /// - nsSteps: number of Newton-Schulz iteration steps + public init( + learningRate: Float, momentum: Float = 0.95, weightDecay: Float = 0.01, + nesterov: Bool = true, nsSteps: Int = 5 + ) { + self.learningRate = learningRate + self.momentum = momentum + self.weightDecay = weightDecay + self.nesterov = nesterov + self.nsSteps = nsSteps + } + + /// Orthogonalize a 2D matrix via a quintic Newton-Schulz iteration. + private func zeropowerViaNewtonSchulz5(_ input: MLXArray, steps: Int) -> MLXArray { + precondition(input.ndim == 2, "Newton-Schulz iteration expects a 2D array") + let (a, b, c): (Float, Float, Float) = (3.4445, -4.7750, 2.0315) + let transposeNeeded = input.dim(-2) > input.dim(-1) + + var X = transposeNeeded ? input.transposed(1, 0) : input + // Frobenius-normalize so the iteration converges. + X = X / (sqrt((X * X).sum(keepDims: true)) + 1e-7) + + for _ in 0 ..< steps { + let A = matmul(X, X.transposed(1, 0)) + let B = b * A + c * matmul(A, A) + X = a * X + matmul(B, X) + } + + return transposeNeeded ? X.transposed(1, 0) : X + } + + override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: MLXArray) -> ( + MLXArray, MLXArray + ) { + var gradient = gradient + if weightDecay != 0 { + gradient = gradient + weightDecay * parameter + } + + let v = momentum * state + (1 - momentum) * gradient + + var update = nesterov ? (gradient * (1 - momentum) + v * momentum) : v + var lr = learningRate + + if update.ndim >= 2 { + let originalShape = update.shape + let reshapeNeeded = update.ndim > 2 + if reshapeNeeded { + update = update.reshaped([update.dim(0), -1]) + } + update = zeropowerViaNewtonSchulz5(update, steps: nsSteps) + if reshapeNeeded { + update = update.reshaped(originalShape) + } + // Scale the learning rate by sqrt(max(1, rows / cols)). + let rows = Float(update.dim(-2)) + let cols = Float(update.dim(-1)) + lr *= pow(max(1, rows / cols), 0.5) + } + + return (parameter - lr * update, v) + } +} + /// Clips the global norm of the gradients. /// /// This function ensures that the global norm of the gradients does not exceed diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index bcee1b322..d773a5563 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -239,4 +239,33 @@ class OptimizerTests: XCTestCase { XCTAssertEqual(optimizer.innerState().count, 2) } + func testMuon() { + checkShape(optimizer: Muon(learningRate: 0.1)) + + // 1D (and 0D) parameters skip the Newton-Schulz orthogonalization and + // use a plain momentum/Nesterov update — pin that math exactly. + // v = 0.95*0 + 0.05*g = 0.05*g; nesterov update = g*(1-m) + v*m; + // param' = param - lr*update. + let opt = Muon(learningRate: 0.1, momentum: 0.95, weightDecay: 0) + let (p, v) = opt.applySingle( + gradient: MLXArray([1.0, 1.0] as [Float]), + parameter: MLXArray([1.0, 2.0] as [Float]), + state: MLXArray([0.0, 0.0] as [Float])) + eval(p, v) + // update = 1*0.05 + 0.05*0.95 = 0.0975 ; param0' = 1 - 0.1*0.0975 + XCTAssertEqual(p[0].item(Float.self), 0.99025, accuracy: 1e-5) + XCTAssertEqual(p[1].item(Float.self), 1.99025, accuracy: 1e-5) + XCTAssertEqual(v[0].item(Float.self), 0.05, accuracy: 1e-6) + + // 2D parameters take the orthogonalization path: it must run, preserve + // shape, and produce finite values that move the parameter. + let g2 = MLXArray(converting: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]) + let (p2, _) = opt.applySingle( + gradient: g2, parameter: MLXArray.zeros([2, 3]), state: MLXArray.zeros([2, 3])) + eval(p2) + XCTAssertEqual(p2.shape, [2, 3]) + XCTAssertTrue(p2.sum().item(Float.self).isFinite) + XCTAssertGreaterThan(abs(p2).sum().item(Float.self), 0) + } + } From eb889e01b7267bc6d9c01fde91bc517b09506c2d Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Tue, 30 Jun 2026 12:31:45 -0500 Subject: [PATCH 10/39] Add learning-rate schedules (port from Python mlx.optimizers.schedulers) (#433) Port the mlx.optimizers schedule functions to Swift: exponentialDecay, stepDecay, cosineDecay, linearSchedule, and joinSchedules. Each returns a (Int) -> Float mapping a step to a learning rate, composing with the mutable `learningRate` exposed by the optimizers (assign optimizer.learningRate = schedule(step) each step). mlx-swift previously had no LR schedules. Tests pin exact values against the reference docstring examples. Co-authored-by: Claude Opus 4.8 (1M context) --- Source/MLXOptimizers/Schedulers.swift | 86 +++++++++++++++++++++++++++ Tests/MLXTests/SchedulerTests.swift | 59 ++++++++++++++++++ 2 files changed, 145 insertions(+) create mode 100644 Source/MLXOptimizers/Schedulers.swift create mode 100644 Tests/MLXTests/SchedulerTests.swift diff --git a/Source/MLXOptimizers/Schedulers.swift b/Source/MLXOptimizers/Schedulers.swift new file mode 100644 index 000000000..3cee0cd6b --- /dev/null +++ b/Source/MLXOptimizers/Schedulers.swift @@ -0,0 +1,86 @@ +// Copyright © 2024 Apple Inc. + +import Foundation + +// Learning-rate schedules. Ports of `mlx.optimizers.schedulers`. +// +// Each function returns a closure mapping a step count to a learning rate. +// Use it by assigning the optimizer's `learningRate` each step, e.g. +// +// ```swift +// let schedule = cosineDecay(1e-1, decaySteps: 1000) +// let optimizer = SGD(learningRate: schedule(0)) +// for step in 0 ..< 1000 { +// optimizer.learningRate = schedule(step) +// optimizer.update(model: model, gradients: grads) +// } +// ``` + +/// Make an exponential-decay schedule: `initial * decayRate ** step`. +/// +/// - Parameters: +/// - initial: initial value +/// - decayRate: multiplicative factor to decay by each step +public func exponentialDecay(_ initial: Float, decayRate: Float) -> (Int) -> Float { + { step in initial * pow(decayRate, Float(step)) } +} + +/// Make a step-decay schedule: `initial * decayRate ** (step / stepSize)`. +/// +/// - Parameters: +/// - initial: initial value +/// - decayRate: multiplicative factor to decay by +/// - stepSize: decay every `stepSize` steps +public func stepDecay(_ initial: Float, decayRate: Float, stepSize: Int) -> (Int) -> Float { + { step in initial * pow(decayRate, Float(step / stepSize)) } +} + +/// Make a cosine-decay schedule. The value is constant (`end`) for steps beyond +/// `decaySteps`. +/// +/// - Parameters: +/// - initial: initial value +/// - decaySteps: number of steps to decay over +/// - end: final value to decay to (default 0) +public func cosineDecay(_ initial: Float, decaySteps: Int, end: Float = 0.0) -> (Int) -> Float { + { step in + let s = Float(min(step, decaySteps)) + let decay = 0.5 * (1.0 + cos((Float.pi / Float(decaySteps)) * s)) + return end + decay * (initial - end) + } +} + +/// Make a linear schedule from `initial` to `end` over `steps`. The value is +/// `end` for any steps beyond `steps`. +/// +/// - Parameters: +/// - initial: initial value +/// - end: final value +/// - steps: number of steps to apply the schedule over (must be >= 1) +public func linearSchedule(_ initial: Float, end: Float, steps: Int) -> (Int) -> Float { + precondition(steps >= 1, "steps must be greater than 0, but got \(steps).") + return { step in + let s = Float(min(step, steps)) + return s * ((end - initial) / Float(steps)) + initial + } +} + +/// Join multiple schedules into one. Schedule `i + 1` receives a step count +/// indicating the number of steps since the `i`-th boundary. +/// +/// - Parameters: +/// - schedules: the schedules to join +/// - boundaries: `schedules.count - 1` step counts marking transitions +public func joinSchedules(_ schedules: [(Int) -> Float], boundaries: [Int]) -> (Int) -> Float { + precondition(!schedules.isEmpty, "Must provide at least 1 schedule to join.") + precondition( + schedules.count == boundaries.count + 1, + "Received \(boundaries.count) boundaries but expected \(schedules.count - 1).") + return { step in + var output = schedules[0](step) + for (boundary, schedule) in zip(boundaries, schedules.dropFirst()) where step >= boundary { + output = schedule(step - boundary) + } + return output + } +} diff --git a/Tests/MLXTests/SchedulerTests.swift b/Tests/MLXTests/SchedulerTests.swift new file mode 100644 index 000000000..6666efb68 --- /dev/null +++ b/Tests/MLXTests/SchedulerTests.swift @@ -0,0 +1,59 @@ +// Copyright © 2024 Apple Inc. + +import Foundation +import XCTest + +@testable import MLXOptimizers + +class SchedulerTests: XCTestCase { + + func testExponentialDecay() { + let s = exponentialDecay(0.1, decayRate: 0.9) + XCTAssertEqual(s(0), 0.1, accuracy: 1e-6) + XCTAssertEqual(s(1), 0.09, accuracy: 1e-6) + XCTAssertEqual(s(2), 0.081, accuracy: 1e-6) + XCTAssertEqual(s(5), 0.1 * 0.9 * 0.9 * 0.9 * 0.9 * 0.9, accuracy: 1e-6) + } + + func testStepDecay() { + let s = stepDecay(0.1, decayRate: 0.9, stepSize: 10) + XCTAssertEqual(s(0), 0.1, accuracy: 1e-6) + XCTAssertEqual(s(9), 0.1, accuracy: 1e-6) // constant within a step window + XCTAssertEqual(s(10), 0.09, accuracy: 1e-6) + XCTAssertEqual(s(21), 0.081, accuracy: 1e-6) // 0.1 * 0.9^2 + } + + func testCosineDecay() { + let s = cosineDecay(0.1, decaySteps: 1000) + XCTAssertEqual(s(0), 0.1, accuracy: 1e-6) + XCTAssertEqual(s(500), 0.05, accuracy: 1e-6) // half-way: 0.5*(1+cos(pi/2))*0.1 + XCTAssertEqual(s(1000), 0.0, accuracy: 1e-6) + XCTAssertEqual(s(5000), 0.0, accuracy: 1e-6) // clamped beyond decaySteps + + // non-zero end value + let s2 = cosineDecay(0.1, decaySteps: 100, end: 0.01) + XCTAssertEqual(s2(0), 0.1, accuracy: 1e-6) + XCTAssertEqual(s2(100), 0.01, accuracy: 1e-6) + } + + func testLinearSchedule() { + let s = linearSchedule(0.0, end: 0.1, steps: 100) + XCTAssertEqual(s(0), 0.0, accuracy: 1e-6) + XCTAssertEqual(s(50), 0.05, accuracy: 1e-6) + XCTAssertEqual(s(100), 0.1, accuracy: 1e-6) + XCTAssertEqual(s(200), 0.1, accuracy: 1e-6) // clamped beyond steps + } + + func testJoinSchedules() { + // warm-up linearly for 10 steps, then cosine-decay. + let linear = linearSchedule(0.0, end: 0.1, steps: 10) + let cosine = cosineDecay(0.1, decaySteps: 200) + let s = joinSchedules([linear, cosine], boundaries: [10]) + + XCTAssertEqual(s(0), 0.0, accuracy: 1e-6) // linear leg + XCTAssertEqual(s(5), 0.05, accuracy: 1e-6) // linear leg + XCTAssertEqual(s(10), 0.1, accuracy: 1e-6) // cosine leg at its step 0 + // at step 11 we are 1 step into the cosine schedule + XCTAssertEqual(s(11), cosine(1), accuracy: 1e-6) + } +} From 041cc7c2662b9564c0666a168867e6dc04b9f509 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Tue, 30 Jun 2026 11:15:44 -0700 Subject: [PATCH 11/39] Move .complex64 to .float32 (#434) --- Source/MLX/DType.swift | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/Source/MLX/DType.swift b/Source/MLX/DType.swift index 8912d3ca4..3aa0534c7 100644 --- a/Source/MLX/DType.swift +++ b/Source/MLX/DType.swift @@ -147,9 +147,9 @@ public enum DType: Hashable, Sendable, CaseIterable { #else case .float16: 0x1p-10 // 2^-10 #endif - case .float32: Double(Float.ulpOfOne) + case .float32, .complex64: Double(Float.ulpOfOne) case .bfloat16: 0x1p-7 // 2^-7 (7 mantissa bits) - case .complex64, .float64: Double.ulpOfOne + case .float64: Double.ulpOfOne default: fatalError("\(dtype) is not a floating point type") } } @@ -166,8 +166,8 @@ public enum DType: Hashable, Sendable, CaseIterable { case .float16: 0x1.FFCp15 // 65504 #endif case .bfloat16: Double(Float(bitPattern: 0x7F7F_0000)) // bf16 = high 16 bits of f32 - case .float32: Double(Float.greatestFiniteMagnitude) - case .complex64, .float64: Double.greatestFiniteMagnitude + case .float32, .complex64: Double(Float.greatestFiniteMagnitude) + case .float64: Double.greatestFiniteMagnitude default: fatalError("\(dtype) is not a floating point type") } } @@ -180,9 +180,9 @@ public enum DType: Hashable, Sendable, CaseIterable { #else case .float16: 0x1p-14 // 2^-14 #endif - case .float32: Double(Float.leastNormalMagnitude) + case .float32, .complex64: Double(Float.leastNormalMagnitude) case .bfloat16: 0x1p-126 // 2^-126 (same exponent range as f32) - case .complex64, .float64: Double.leastNormalMagnitude + case .float64: Double.leastNormalMagnitude default: fatalError("\(dtype) is not a floating point type") } } @@ -195,9 +195,9 @@ public enum DType: Hashable, Sendable, CaseIterable { #else case .float16: 0x1p-24 // 2^-24 #endif - case .float32: Double(Float.leastNonzeroMagnitude) + case .float32, .complex64: Double(Float.leastNonzeroMagnitude) case .bfloat16: 0x1p-133 // 2^-133 = 2^-126 · 2^-7 - case .complex64, .float64: Double.leastNonzeroMagnitude + case .float64: Double.leastNonzeroMagnitude default: fatalError("\(dtype) is not a floating point type") } } From 0bb916c67f4b9e5c682cbe02a42c701c93ab5021 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:29:35 -0700 Subject: [PATCH 12/39] guard use of Process (#437) - does not build on iOS, etc. - fix #436 --- Source/Encuda/encuda-compile.swift | 66 +++++++++++++++--------------- Source/Encuda/encuda-link.swift | 36 ++++++++-------- Source/Encuda/encuda-utils.swift | 22 +++++----- 3 files changed, 66 insertions(+), 58 deletions(-) diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift index 42bd75313..d5ab1f5da 100644 --- a/Source/Encuda/encuda-compile.swift +++ b/Source/Encuda/encuda-compile.swift @@ -34,43 +34,45 @@ extension Encuda { var incremental: Bool = false mutating func run() throws { - - if verbose { - print("Running encuda compile") - } - - if incremental && isUpToDate() { + #if os(macOS) || os(Linux) if verbose { - print("Output is up to date, skipping compilation") + print("Running encuda compile") } - return - } - let resolvedNvcc: String - if let path = nvccPath { - resolvedNvcc = path - } else if let found = searchForCommand("nvcc") { - resolvedNvcc = found.path - } else { - fatalError("nvcc not found") - } + if incremental && isUpToDate() { + if verbose { + print("Output is up to date, skipping compilation") + } + return + } - let includeArgs = includeDirs.flatMap { ["-I", $0] } - let ccbinArgs = clangppPath.map { ["-ccbin=\($0)"] } ?? [] - let stdArgs = std.map { ["-std=\($0)"] } ?? [] - let archArgs = - ProcessInfo.processInfo.environment["CUDA_ARCH"].map { ["-arch", $0] } ?? [] + let resolvedNvcc: String + if let path = nvccPath { + resolvedNvcc = path + } else if let found = searchForCommand("nvcc") { + resolvedNvcc = found.path + } else { + fatalError("nvcc not found") + } - let process = Process() - process.executableURL = URL(fileURLWithPath: resolvedNvcc) - process.arguments = - ["-cuda", "-rdc=true", "--expt-relaxed-constexpr"] + stdArgs + ccbinArgs + archArgs - + (verbose ? ["-v"] : []) + includeArgs + inputFiles + ["-o", output] - try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.nvccFailed(process.terminationStatus) - } + let includeArgs = includeDirs.flatMap { ["-I", $0] } + let ccbinArgs = clangppPath.map { ["-ccbin=\($0)"] } ?? [] + let stdArgs = std.map { ["-std=\($0)"] } ?? [] + let archArgs = + ProcessInfo.processInfo.environment["CUDA_ARCH"].map { ["-arch", $0] } ?? [] + + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedNvcc) + process.arguments = + ["-cuda", "-rdc=true", "--expt-relaxed-constexpr"] + stdArgs + ccbinArgs + + archArgs + + (verbose ? ["-v"] : []) + includeArgs + inputFiles + ["-o", output] + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.nvccFailed(process.terminationStatus) + } + #endif } private func isUpToDate() -> Bool { diff --git a/Source/Encuda/encuda-link.swift b/Source/Encuda/encuda-link.swift index 92710e320..2220d1a12 100644 --- a/Source/Encuda/encuda-link.swift +++ b/Source/Encuda/encuda-link.swift @@ -80,25 +80,29 @@ extension Encuda { } func nvcc(args: [String]) throws { - let process = Process() - process.executableURL = URL(fileURLWithPath: resolvedNvcc) - process.arguments = ["-ccbin=\(resolvedClangpp)"] + (verbose ? ["-v"] : []) + args - try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.nvccFailed(process.terminationStatus) - } + #if os(macOS) || os(Linux) + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedNvcc) + process.arguments = ["-ccbin=\(resolvedClangpp)"] + (verbose ? ["-v"] : []) + args + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.nvccFailed(process.terminationStatus) + } + #endif } func clang(args: [String]) throws { - let process = Process() - process.executableURL = URL(fileURLWithPath: resolvedClangpp) - process.arguments = (verbose ? ["-v"] : []) + args - try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.clangFailed(process.terminationStatus) - } + #if os(macOS) || os(Linux) + let process = Process() + process.executableURL = URL(fileURLWithPath: resolvedClangpp) + process.arguments = (verbose ? ["-v"] : []) + args + try process.run() + process.waitUntilExitWorkaround() + guard process.terminationStatus == 0 else { + throw EncudaError.clangFailed(process.terminationStatus) + } + #endif } private func isUpToDate() -> Bool { diff --git a/Source/Encuda/encuda-utils.swift b/Source/Encuda/encuda-utils.swift index 4f08a2033..ff20be2b4 100644 --- a/Source/Encuda/encuda-utils.swift +++ b/Source/Encuda/encuda-utils.swift @@ -11,15 +11,17 @@ func searchForCommand(_ name: String) -> URL? { return nil } -extension Process { - /** - * When running on aarch64 Ubuntu 24 docker on Mac, it happens that - * the `waitUntilExit()` method never returns. - * We adopted this workaround. - */ - func waitUntilExitWorkaround() { - while isRunning { - Thread.sleep(forTimeInterval: 0.05) +#if os(macOS) || os(Linux) + extension Process { + /** + * When running on aarch64 Ubuntu 24 docker on Mac, it happens that + * the `waitUntilExit()` method never returns. + * We adopted this workaround. + */ + func waitUntilExitWorkaround() { + while isRunning { + Thread.sleep(forTimeInterval: 0.05) + } } } -} +#endif From 7a8d05a2d70d8b02298a8248a00943afc4b5391b Mon Sep 17 00:00:00 2001 From: Lucas Newman Date: Tue, 7 Jul 2026 17:03:30 -0400 Subject: [PATCH 13/39] Add fp8 conversion utilities (#435) * Add FP8 conversion utilities. --- Source/MLX/Ops.swift | 37 ++++++++++++++++++++++++ Tests/MLXTests/FP8Tests.swift | 54 +++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+) create mode 100644 Tests/MLXTests/FP8Tests.swift diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index ec83c0440..b7e95a330 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -1374,6 +1374,28 @@ public func expm1(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXAr return MLXArray(result) } +/// Convert an E4M3 FP8 array to the given floating-point ``DType``. +/// +/// The input array is interpreted as raw FP8 bytes. +/// +/// - Parameters: +/// - array: input array containing E4M3 FP8 values +/// - dtype: data type of the output array +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``toFP8(_:stream:)`` +/// - +public func fromFP8( + _ array: MLXArray, + dtype: DType = .bfloat16, + stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_from_fp8(&result, array.ctx, dtype.cmlxDtype, stream.ctx) + return MLXArray(result) +} + @available(*, deprecated, renamed: "gatherMM(_:_:lhsIndices:rhsIndices:sortedIndices:stream:)") public func gatherMatmul( _ a: MLXArray, _ b: MLXArray, lhsIndices: MLXArray? = nil, rhsIndices: MLXArray? = nil, @@ -3076,6 +3098,21 @@ public func tiled(_ array: MLXArray, repetitions: Int, stream: StreamOrDevice = return MLXArray(result) } +/// Convert a floating-point array to E4M3 FP8 bytes. +/// +/// - Parameters: +/// - array: input array +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - ``fromFP8(_:dtype:stream:)`` +/// - +public func toFP8(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { + var result = mlx_array_new() + mlx_to_fp8(&result, array.ctx, stream.ctx) + return MLXArray(result) +} + /// Returns the `k` largest elements from the input along a given axis. /// /// The elements will not necessarily be in sorted order. diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift new file mode 100644 index 000000000..9ec00c401 --- /dev/null +++ b/Tests/MLXTests/FP8Tests.swift @@ -0,0 +1,54 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +class FP8Tests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + func testFromFP8DecodesE4M3Bytes() { + let allBytes = (0 ... 255).map { UInt8($0) } + let bytes = MLXArray(allBytes) + let decoded = fromFP8(bytes, dtype: .float32) + + XCTAssertEqual(decoded.dtype, .float32) + XCTAssertEqual(decoded.shape, [256]) + + let expected = MLXArray(allBytes.map { e4m3Value($0) }) + assertEqual(decoded, expected) + } + + func testToFP8EncodesE4M3Bytes() { + let values = MLXArray( + [ + 0.0, -0.0, 1.0, -2.0, 0.125, 0.015625, 448.0, -448.0, + ] as [Float]) + + let encoded = toFP8(values) + + XCTAssertEqual(encoded.dtype, .uint8) + XCTAssertEqual(encoded.shape, values.shape) + XCTAssertEqual( + encoded.asArray(UInt8.self), + [0x00, 0x80, 0x38, 0xC0, 0x20, 0x08, 0x7E, 0xFE]) + } +} + +private func e4m3Value(_ value: UInt8) -> Float { + let sign: Float = (value & 0x80) == 0 ? 1 : -1 + let exponent = Int((value >> 3) & 0x0F) + let mantissa = Int(value & 0x07) + + if exponent == 0 { + if mantissa == 0 { + return sign < 0 ? -0.0 : 0.0 + } + return sign * Float(mantissa) / 512 + } + + return sign * Float(8 + mantissa) * powf(2, Float(exponent - 10)) +} From de3342cefe687116afbdd4a422d5bc8a19d21506 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:19:15 -0700 Subject: [PATCH 14/39] build iOS in CI (#440) - add an iOS build in CI to catch iOS-only build breaks - move the docs build into the same macOS runner --- .github/workflows/pull_request.yml | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index de86d599e..fad6ac5f9 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -69,17 +69,6 @@ jobs: run: | pre-commit run --all || (echo "Style checks failed, please install pre-commit and run pre-commit run --all and push the change"; echo ""; git --no-pager diff; exit 1) - docs: - needs: lint - if: github.repository == 'ml-explore/mlx-swift' - runs-on: [self-hosted, macos] - steps: - - uses: actions/checkout@v6 - with: - submodules: recursive - - - name: Verify documentation - run: scripts/verify-docs.sh mac_build_and_test: needs: lint @@ -99,6 +88,17 @@ jobs: run: | brew install cmake ninja + - name: Verify documentation + run: scripts/verify-docs.sh + + - name: Build SwiftPM (Xcode, iOS) + shell: sh + run: | + xcodebuild -version + swift --version + rm -rf ~/Library/Developer/Xcode/DerivedData/* + xcodebuild build -scheme mlx-swift-Package -destination 'generic/platform=iOS' + - name: Build SwiftPM (Xcode, macOS) shell: sh run: | From 09051ed1986b6f2c69b6fda4e32ca42dd7dcdec3 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Tue, 14 Jul 2026 00:20:12 +0400 Subject: [PATCH 15/39] Add global scale support to quantized layers (#426) * Add global scale support to quantized layers Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations. Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there. Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface. * Keep nvfp4 global scale on the quantized path for Metal Since MLX encodes nvfp4 group scales pre-multiplied by (448 * 6) / globalScale, the global scale factors out of dequantization linearly. Exploit this by always running quantizedMM (supported on Metal without globalScale) and applying globalScale / (448 * 6) to the output, instead of dequantizing the full weight and falling back to a plain matmul. This keeps weights quantized end-to-end on all backends, removes the dequantizedWeight fallback from QuantizedLinear and QuantizedEmbedding, and applies the scale before the bias add so the bias is unaffected. Add Metal tests that validate the GPU quantized path against CPU dequantized(globalScale:) references for QuantizedLinear (matmul + bias) and QuantizedEmbedding (asLinear and index lookup). --- Source/MLXNN/Quantized.swift | 52 ++++++++-- Tests/MLXTests/QuantizationTests.swift | 126 +++++++++++++++++++++++++ 2 files changed, 168 insertions(+), 10 deletions(-) diff --git a/Source/MLXNN/Quantized.swift b/Source/MLXNN/Quantized.swift index 076e91ca0..ca5913ee7 100644 --- a/Source/MLXNN/Quantized.swift +++ b/Source/MLXNN/Quantized.swift @@ -27,6 +27,12 @@ public protocol Quantized: Module { var mode: QuantizationMode { get } } +private func applyNVFP4GlobalScale(_ value: MLXArray, globalScale: MLXArray?) -> MLXArray { + guard let globalScale else { return value } + // NVFP4 encodes its E4M3 group scales with (448 * 6) / globalScale. + return (value * (globalScale / (448 * 6))).asType(value.dtype) +} + /// Quantize any ``Quantizable`` layer that is not already quantized. public func quantizeSingle( layer: Module, groupSize: Int = 64, bits: Int = 4, mode: QuantizationMode = .affine @@ -159,6 +165,7 @@ open class QuantizedEmbedding: Embedding, Quantized { public let mode: QuantizationMode public let scales: MLXArray public let biases: MLXArray? + @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray? open override var shape: (Int, Int) { let (embeddingCount, dimensions) = super.shape @@ -184,14 +191,16 @@ open class QuantizedEmbedding: Embedding, Quantized { public init( weight: MLXArray, groupSize: Int = 64, bits: Int = 4, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode + self._globalScale.wrappedValue = globalScale let (quantizedWeight, scales, biases) = MLX.quantized( - weight, groupSize: groupSize, bits: bits, mode: mode) + weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale) self.scales = scales self.biases = biases @@ -201,19 +210,36 @@ open class QuantizedEmbedding: Embedding, Quantized { self.freeze() } + /// Initializer meant for subclasses to provide arrays directly. + public init( + weight: MLXArray, scales: MLXArray, biases: MLXArray?, + groupSize: Int, bits: Int, + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil + ) { + self.groupSize = groupSize + self.bits = bits + self.mode = mode + self.scales = scales + self.biases = biases + self._globalScale.wrappedValue = globalScale + super.init(weight: weight) + } + open override func callAsFunction(_ x: MLXArray) -> MLXArray { let s = x.shape let x = x.flattened() let out = dequantized( weight[x], scales: scales[x], biases: biases == nil ? nil : biases![x], groupSize: groupSize, bits: bits, mode: mode) - return out.reshaped(s + [-1]) + return applyNVFP4GlobalScale(out, globalScale: globalScale).reshaped(s + [-1]) } open override func asLinear(_ x: MLXArray) -> MLXArray { - quantizedMM( + let result = quantizedMM( x, weight, scales: scales, biases: biases, transpose: true, groupSize: groupSize, bits: bits, mode: mode) + return applyNVFP4GlobalScale(result, globalScale: globalScale) } } @@ -243,6 +269,7 @@ open class QuantizedLinear: Linear, Quantized { public let mode: QuantizationMode public let scales: MLXArray public let biases: MLXArray? + @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray? open override var shape: (Int, Int) { let shape = weight.shape2 @@ -292,14 +319,16 @@ open class QuantizedLinear: Linear, Quantized { /// Initialize a ``QuantizedLinear`` with non-quantized weights and bias. public init( weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode + self._globalScale.wrappedValue = globalScale let (quantizedWeight, scales, biases) = MLX.quantized( - weight, groupSize: groupSize, bits: bits, mode: mode) + weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale) self.scales = scales self.biases = biases @@ -316,13 +345,15 @@ open class QuantizedLinear: Linear, Quantized { public init( weight: MLXArray, bias: MLXArray? = nil, scales: MLXArray, biases: MLXArray?, groupSize: Int, bits: Int, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode self.scales = scales self.biases = biases + self._globalScale.wrappedValue = globalScale super.init(weight: weight, bias: bias) } @@ -334,7 +365,7 @@ open class QuantizedLinear: Linear, Quantized { } open override func callAsFunction(_ x: MLXArray) -> MLXArray { - var x = quantizedMM( + var result = quantizedMM( x, weight, scales: scales, @@ -344,10 +375,11 @@ open class QuantizedLinear: Linear, Quantized { bits: bits, mode: mode ) + result = applyNVFP4GlobalScale(result, globalScale: globalScale) if let bias { - x = x + bias + result = result + bias } - return x + return result } /// Returns a QuantizedLinear layer that applies the same linear transformation up to the quantization error. diff --git a/Tests/MLXTests/QuantizationTests.swift b/Tests/MLXTests/QuantizationTests.swift index 0edbd5452..357ac81e3 100644 --- a/Tests/MLXTests/QuantizationTests.swift +++ b/Tests/MLXTests/QuantizationTests.swift @@ -6,6 +6,74 @@ import MLXNN import XCTest class QuantizationTests: XCTestCase { + func testQuantizedLinearAppliesNVFP4GlobalScaleOnMetal() { + let (input, quantizedWeight, scales, biases, globalScale, bias, expected) = + Device.withDefaultDevice(.cpu) { + MLXRandom.seed(7) + let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32]) + let input = MLXRandom.uniform(low: -1, high: 1, [4, 32]) + let globalScale = MLX.max(abs(weight)).asType(.float32) + let bias = MLXRandom.uniform(low: -1, high: 1, [32]) + let (quantizedWeight, scales, biases) = MLX.quantized( + weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + let expected = + matmul( + input, + dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale, dtype: .float32 + ).T) + + bias + eval(quantizedWeight, scales, expected) + return (input, quantizedWeight, scales, biases, globalScale, bias, expected) + } + + let actual = Device.withDefaultDevice(.gpu) { + QuantizedLinear( + weight: quantizedWeight, bias: bias, scales: scales, biases: biases, + groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)(input) + } + + assertEqual(actual, expected, rtol: 1e-5, atol: 1e-5) + } + + func testQuantizedEmbeddingAppliesNVFP4GlobalScaleOnMetal() { + let ( + input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear, + expectedLookup + ) = Device.withDefaultDevice(.cpu) { + MLXRandom.seed(11) + let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32]) + let input = MLXRandom.uniform(low: -1, high: 1, [4, 32]) + let indices = MLXArray([0, 3, 7, 31]) + let globalScale = MLX.max(abs(weight)).asType(.float32) + let (quantizedWeight, scales, biases) = MLX.quantized( + weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + let dequantizedWeight = dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale, dtype: .float32) + let expectedLinear = matmul(input, dequantizedWeight.T) + let expectedLookup = dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale)[indices].asType(.bfloat16) + eval(quantizedWeight, scales, expectedLinear, expectedLookup) + return ( + input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear, + expectedLookup + ) + } + + let (actualLinear, actualLookup) = Device.withDefaultDevice(.gpu) { + let embedding = QuantizedEmbedding( + weight: quantizedWeight, scales: scales, biases: biases, + groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + return (embedding.asLinear(input), embedding(indices)) + } + + assertEqual(actualLinear, expectedLinear, rtol: 1e-5, atol: 1e-5) + assertEqual(actualLookup, expectedLookup, rtol: 1e-5, atol: 1e-5) + } + func testQuantizedLinearShapeDesc() { let linear1 = Linear(512, 1024) let quantized1 = linear1.toQuantized(groupSize: 64, bits: 4) @@ -39,4 +107,62 @@ class QuantizationTests: XCTestCase { let quantized = QuantizedLinear(64, 64, groupSize: 32, bits: 4, mode: .mxfp4) XCTAssertNil(quantized.biases) } + + func testQuantizedLinearStoresGlobalScale() { + let globalScale = MLXArray(1.0, dtype: .float32) + let quantized = QuantizedLinear( + weight: MLXArray.zeros([8, 4], dtype: .uint32), + bias: nil, + scales: MLXArray.ones([8, 4], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4, + globalScale: globalScale) + + XCTAssertNotNil(quantized.globalScale) + XCTAssertEqual(quantized.globalScale?.dtype, .float32) + XCTAssertNotNil(quantized.parameters()["global_scale"]) + XCTAssertNil(quantized.parameters()["globalScale"]) + } + + func testQuantizedEmbeddingStoresGlobalScale() { + let globalScale = MLXArray(1.0, dtype: .float32) + let quantized = QuantizedEmbedding( + weight: MLXArray.zeros([8, 2], dtype: .uint32), + scales: MLXArray.ones([8, 2], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4, + globalScale: globalScale) + + XCTAssertNotNil(quantized.globalScale) + XCTAssertEqual(quantized.globalScale?.dtype, .float32) + XCTAssertNotNil(quantized.parameters()["global_scale"]) + XCTAssertNil(quantized.parameters()["globalScale"]) + } + + func testQuantizedGlobalScaleIsOptionalParameter() { + let linear = QuantizedLinear( + weight: MLXArray.zeros([8, 4], dtype: .uint32), + bias: nil, + scales: MLXArray.ones([8, 4], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4) + let embedding = QuantizedEmbedding( + weight: MLXArray.zeros([8, 2], dtype: .uint32), + scales: MLXArray.ones([8, 2], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4) + + XCTAssertNil(linear.globalScale) + XCTAssertNil(linear.parameters()["global_scale"]) + XCTAssertNil(embedding.globalScale) + XCTAssertNil(embedding.parameters()["global_scale"]) + } } From 9a09a0388098ac2bc545932720cebbeee274f0f8 Mon Sep 17 00:00:00 2001 From: agerjura <136923563+agerjura@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:41:33 +0300 Subject: [PATCH 16/39] Fix FinalizerCaptureState leak in MLXArray(rawPointer:_:dtype:finalizer:) (#448) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The initializer retains a box holding the finalizer closure and hands it to mlx as an opaque payload, but `finalizerTrampoline` read it back with `takeUnretainedValue`, which never consumes the `passRetained` +1. The local `free(ptr:)` that would have balanced it was never passed anywhere — only the trampoline is — so it was dead code and the box leaked. The box owns the closure, so the closure's captures leak with it. The initializer's own documented example captures an `IOSurface` and comments that it "implicitly releases it when it returns"; it did not, and the surface was pinned for the lifetime of the process. `testIOSurface` covers the path but only prints, so nothing caught it. `takeRetainedValue` consumes the +1 exactly once. mlx-c documents the dtor as "Callback for when the buffer is no longer needed" and calls it once per buffer, and the trampoline has exactly one caller, so a single consume is correct. Removes the dead `free(ptr:)`. Adds a differential regression test: with the bug the finalizer runs but its captures are never released (witness releases 0); with the fix they are (1). The test drains mlx's allocator cache first, since a freed buffer's dtor is not guaranteed to have run at scope exit, and uses lock-guarded counters because mlx may invoke the dtor from its own thread. Full suite: 534 tests, 0 failures. Co-authored-by: Fred Jura --- Source/MLX/MLXArray+Init.swift | 17 +++++-- Tests/MLXTests/MLXArray+InitTests.swift | 67 +++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/Source/MLX/MLXArray+Init.swift b/Source/MLX/MLXArray+Init.swift index 01a8daa21..9ca79d333 100644 --- a/Source/MLX/MLXArray+Init.swift +++ b/Source/MLX/MLXArray+Init.swift @@ -31,7 +31,16 @@ private class FinalizerCaptureState { func finalizerTrampoline( payload: UnsafeMutableRawPointer? ) { - let state = Unmanaged.fromOpaque(payload!).takeUnretainedValue() + // `takeRetainedValue` consumes the +1 from the `passRetained` in + // `init(rawPointer:_:dtype:finalizer:)`. mlx-c documents this dtor as + // "Callback for when the buffer is no longer needed" and invokes it exactly + // once, so consuming the reference here balances the retain exactly. + // + // Using `takeUnretainedValue` would read the box without consuming it, + // leaking the `FinalizerCaptureState` and permanently pinning everything the + // finalizer closure captured — the `IOSurface` in the initializer's own + // documented example, for instance. + let state = Unmanaged.fromOpaque(payload!).takeRetainedValue() state.f() } @@ -80,10 +89,8 @@ extension MLXArray { _ shape: (some Collection)? = [Int]?.none, dtype: DType, finalizer: @escaping () -> Void ) { - func free(ptr: UnsafeMutableRawPointer?) { - Unmanaged.fromOpaque(ptr!).release() - } - + // Balanced by the `takeRetainedValue` in `finalizerTrampoline`, which + // mlx-c invokes once when the buffer is no longer needed. let payload = Unmanaged.passRetained(FinalizerCaptureState(finalizer)).toOpaque() self.init( diff --git a/Tests/MLXTests/MLXArray+InitTests.swift b/Tests/MLXTests/MLXArray+InitTests.swift index 19574da08..acd487306 100644 --- a/Tests/MLXTests/MLXArray+InitTests.swift +++ b/Tests/MLXTests/MLXArray+InitTests.swift @@ -267,6 +267,73 @@ class MLXArrayInitTests: XCTestCase { XCTAssertEqual(e.dtype, .float64) } + /// The finalizer's captures must be released once mlx has called it. + /// + /// `init(rawPointer:_:dtype:finalizer:)` retains a box holding the closure and + /// hands it to mlx as an opaque payload; `finalizerTrampoline` is the only + /// thing that can release it. If the trampoline reads the box without + /// consuming the reference, the closure — and everything it captured — is + /// pinned for the lifetime of the process, once per adopted buffer. + /// + /// `testIOSurface` below documents the intended behaviour in a comment + /// ("implicitly releases it when it returns") but never asserts it, which is + /// what let this go unnoticed. + func testFinalizerCapturesAreReleasedAfterFinalizerRuns() { + // mlx may invoke the dtor from its own thread, so the counters cannot be + // plain locals read across that boundary. + final class Counter: @unchecked Sendable { + private let lock = NSLock() + private var value = 0 + func increment() { + lock.lock() + value += 1 + lock.unlock() + } + var current: Int { + lock.lock() + defer { lock.unlock() } + return value + } + } + final class Witness { + let onDeinit: () -> Void + init(_ onDeinit: @escaping () -> Void) { self.onDeinit = onDeinit } + deinit { onDeinit() } + } + + let finalizerRuns = Counter() + let witnessReleases = Counter() + + do { + let buffer = UnsafeMutableRawPointer.allocate( + byteCount: 4 * MemoryLayout.stride, alignment: 16) + buffer.initializeMemory(as: Float.self, repeating: 1, count: 4) + + let witness = Witness { witnessReleases.increment() } + let array = MLXArray(rawPointer: buffer, [4], dtype: .float32) { + [witness] in + // Holds the witness exactly the way the IOSurface example holds + // its surface. + _ = witness + finalizerRuns.increment() + buffer.deallocate() + } + XCTAssertEqual(array.sum().item(Float.self), 4) + } + + // mlx keeps freed buffers in its allocator cache, so the dtor has not + // necessarily run at scope exit. + Memory.clearCache() + let deadline = Date().addingTimeInterval(5) + while finalizerRuns.current == 0, Date() < deadline { usleep(1_000) } + + XCTAssertEqual(finalizerRuns.current, 1, "mlx must call the finalizer exactly once") + XCTAssertEqual( + witnessReleases.current, 1, + "the finalizer's captures must be released along with it — otherwise every adopted " + + "buffer permanently pins whatever the closure held") + } + #if canImport(IOSurface) func testIOSurface() { let height = 100 From da318704cc0e972b61dcca43c62cd15e545362ae Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Mon, 10 Aug 2026 11:42:56 -0500 Subject: [PATCH 17/39] Derive the jit-source lists in update-mlx.sh instead of hand-listing them (#445) `tools/update-mlx.sh` hand-tracked two lists that mlx's own CMake derives, so each one silently went stale whenever the vendored mlx submodule gained a jit source. The CUDA `MLX_JIT_SOURCES` argument named 16 `device/` headers explicitly, where `mlx/backend/cuda/CMakeLists.txt` uses `file(GLOB)` over `device/*.h` and `device/*.cuh`. A bump that adds one regenerated a `cuda_jit_sources.h` missing its array, failing later as `use of undeclared identifier 'jit_source_*'`: mlx v0.31.2 needs 17 and v0.32.0 needs 22, against the 16 the frozen list produced. The Metal step named 46 `make` targets explicitly, where `mlx/backend/metal/CMakeLists.txt` declares them via `make_jit_source()`. A bump that adds one left `jit/includes.h` declaring a `mlx::core::metal::()` that no generated source defined, i.e. an undefined symbol at link time on macOS and iOS: v0.31.2 declares 47 and v0.32.0 declares 48, against the same 46. Both lists are now derived from the vendored tree at run time. Against the currently pinned mlx (v0.31.1) the derived values are identical to the removed literals, and the regenerated `cuda_jit_sources.h` is byte-for-byte identical to the committed one, so this is a no-op at the current pin. Each derivation is guarded: with `set -e` an empty result would not fail on its own, and bare `make` with no targets would quietly build `all` instead. --- tools/update-mlx.sh | 84 ++++++++++++++++++--------------------------- 1 file changed, 34 insertions(+), 50 deletions(-) diff --git a/tools/update-mlx.sh b/tools/update-mlx.sh index fc1ab9615..d69d5e381 100755 --- a/tools/update-mlx.sh +++ b/tools/update-mlx.sh @@ -10,6 +10,8 @@ then exit 1 fi +mlx_source="$PWD/Source/Cmlx/mlx" + # copy mlx-c headers to build area rm -f Source/Cmlx/include/mlx/c/* cp Source/Cmlx/mlx-c/mlx/c/*.h Source/Cmlx/include/mlx/c @@ -22,62 +24,44 @@ cmake ../Source/Cmlx/mlx -DMLX_METAL_JIT=ON -DMACOS_VERSION=14.0 # run the cmake build to generate the source files cd mlx/backend/metal -make \ - arange \ - binary \ - binary_ops \ - binary_two \ - conv \ - copy \ - fft \ - fp_quantized \ - fp_quantized_nax \ - gather \ - gather_axis \ - gather_front \ - gemm \ - gemm_nax \ - gemv_masked \ - hadamard \ - logsumexp \ - masked_scatter \ - quantized \ - quantized_nax \ - quantized_utils \ - reduce \ - reduce_utils \ - scan \ - scatter \ - scatter_axis \ - softmax \ - sort \ - steel_attention \ - steel_attention_nax \ - steel_conv \ - steel_conv_3d \ - steel_conv_general \ - steel_gemm_fused \ - steel_gemm_fused_nax \ - steel_gemm_gather \ - steel_gemm_gather_nax \ - steel_gemm_masked \ - steel_gemm_segmented \ - steel_gemm_splitk \ - steel_gemm_splitk_nax \ - ternary \ - ternary_ops \ - unary \ - unary_ops \ - utils + +# one target per make_jit_source() in mlx/backend/metal/CMakeLists.txt, which is +# where the set of Metal jit sources is declared +metal_jit_targets=(${(f)"$(tr '\n' ' ' < "$mlx_source/mlx/backend/metal/CMakeLists.txt" \ + | grep -oE 'make_jit_source\( *[a-zA-Z_0-9/]+' \ + | sed -E 's/make_jit_source\( *//' \ + | sed 's|.*/||' \ + | sort -u)"}) + +if (( ! $#metal_jit_targets )) +then + echo "Found no make_jit_source() targets in mlx/backend/metal/CMakeLists.txt" + exit 1 +fi + +make $metal_jit_targets cd ../../.. make cpu_compiled_preamble # run the command to do the build-time code generation for CUDA + +cuda_source="$mlx_source/mlx/backend/cuda" + +# the same set that the file(GLOB) in mlx/backend/cuda/CMakeLists.txt embeds +cuda_jit_sources=(${cuda_source}/device/*.(h|cuh)(N)) +cuda_jit_sources=(${cuda_jit_sources#$cuda_source/}) + +if (( ! $#cuda_jit_sources )) +then + echo "Found no jit sources in mlx/backend/cuda/device" + exit 1 +fi + cmake \ - -DMLX_SOURCE_ROOT="../Source/Cmlx/mlx/mlx/backend/cuda" \ - -DMLX_JIT_SOURCES="device/atomic_ops.cuh:device/binary_ops.cuh:device/cast_op.cuh:device/complex.cuh:device/config.h:device/fp16_math.cuh:device/gather.cuh:device/gather_axis.cuh:device/hadamard.cuh:device/indexing.cuh:device/scatter.cuh:device/scatter_axis.cuh:device/scatter_ops.cuh:device/ternary_ops.cuh:device/unary_ops.cuh:device/utils.cuh" \ - -P "../Source/Cmlx/mlx/mlx/backend/cuda/bin2h.cmake" + -DMLX_SOURCE_ROOT="$cuda_source" \ + -DMLX_JIT_SOURCES="${(j.:.)cuda_jit_sources}" \ + -P "$cuda_source/bin2h.cmake" cd .. From 215785259a2fc19a4e4efb67fa858c02ab4193be Mon Sep 17 00:00:00 2001 From: Joannis Orlandos Date: Fri, 14 Aug 2026 10:10:31 -0600 Subject: [PATCH 18/39] Improve CUDA NVCC spawn performance (on Linux) (#451) * Reap encuda's child processes with waitpid instead of polling isRunning `waitUntilExitWorkaround()` polled `Process.isRunning` in a 50ms sleep loop, added because `waitUntilExit()` "never returns" on aarch64 Linux. It did not fix that hang -- `isRunning` reads the same Foundation bookkeeping that stalls, so the workaround only changed a blocked call into an infinite sleep loop. Observed during a parallel SwiftPM build of the CUDA backend for linux/arm64 (aarch64-unknown-linux-gnu, CUDA 13, sm_121): `nvcc` finishes and writes its output in full, then becomes a zombie that is never reaped, while `encuda-tool` spins in `hrtimer_nanosleep` holding a completed 15MB transpiled .cpp on disk. Invoking encuda directly on the same file with identical arguments completes in 2-3 seconds, so the work itself is fine -- the process simply never notices the child finished. It is load-dependent: `--jobs 2` transpiled all 96 .cu files without incident, while an unconstrained build (SwiftPM defaults to every core) deadlocked. That is consistent with pipe file descriptors leaking between concurrently spawned siblings -- the hung `encuda-tool` held both the read and the write end of the same pipes, so no reader could ever see EOF. `waitpid` answers the only question this code has: did the child exit, and with what status. It does not depend on Foundation's reaper, and it collects the zombie. On ECHILD -- Foundation's own reaper winning the race -- it falls back to `terminationStatus`, which is authoritative in exactly that case. Exit status is decoded without the C macros, which Swift does not import: the low 7 bits carry the terminating signal and the next 8 the exit code, an encoding shared by Linux and Darwin. Signals are reported as `128 + signal`, following the shell convention, so a killed compiler is still a non-zero failure rather than a silent success. Verified on macOS: exit codes 0, 7 and 42 round-trip exactly; SIGKILL and SIGTERM report 137 and 143; no reapable child remains afterwards. The aarch64 Linux deadlock itself has not been reproduced locally -- it needs a CUDA toolchain and a parallel build to trigger -- so this is a fix for a diagnosed root cause rather than one observed to clear the symptom. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01K1rzxACJqmxVqgRsa4xj3R * Don't read terminationStatus before Foundation has published it The waitpid change cleared the deadlock it was written for -- an unconstrained SwiftPM build of mlx's CUDA backend for aarch64 Linux (CUDA 13.2, sm_121, 16 jobs) now runs through the transpile step at full speed instead of stalling forever on a zombie nvcc. But it replaced the hang with a crash. `waitpid` and Foundation's own reaper thread race for every child. When Foundation wins, our waitpid returns ECHILD and the fallback read `terminationStatus` immediately. That is not safe: Foundation publishes the result in two steps, and reading the status inside the window before `isRunning` goes false traps the process. *** Program crashed: System trap at 0x0000ffff90b6fb50 *** Thread 0 "encuda-tool" crashed: 0 Process.terminationStatus.getter + 80 in libFoundation.so 1 Process.waitForExitStatus() + 107 in encuda-tool at Source/Encuda/encuda-utils.swift:73:24 2 Encuda.Compile.run() + 1883 in encuda-tool Roughly one transpile in 250 died this way, at a different .cu file each run. The build failure it produces is unusually opaque: the crash report goes to the transpiler's own output, so `swift build` exits 1 with no diagnostic at all, and the last thing in the log is an ordinary "Compiling scan.cu to scan.cpp" progress line. So wait for `isRunning` to clear before reading the status -- but on a deadline, since an unbounded wait is the hang this function exists to avoid. On expiry it reports 254 rather than trapping or spinning: the caller turns any non-zero status into a build error, which is diagnosable, and neither a trap nor a hang is. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01RdegWrNm5PKBqeXdkgmn9Q --------- Co-authored-by: Claude Opus 5 (1M context) --- Source/Encuda/encuda-compile.swift | 6 +- Source/Encuda/encuda-link.swift | 12 +-- Source/Encuda/encuda-utils.swift | 113 +++++++++++++++++++++++++++-- 3 files changed, 114 insertions(+), 17 deletions(-) diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift index d5ab1f5da..65ea7b0f6 100644 --- a/Source/Encuda/encuda-compile.swift +++ b/Source/Encuda/encuda-compile.swift @@ -68,9 +68,9 @@ extension Encuda { + archArgs + (verbose ? ["-v"] : []) + includeArgs + inputFiles + ["-o", output] try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.nvccFailed(process.terminationStatus) + let status = process.waitForExitStatus() + guard status == 0 else { + throw EncudaError.nvccFailed(status) } #endif } diff --git a/Source/Encuda/encuda-link.swift b/Source/Encuda/encuda-link.swift index 2220d1a12..3879e4ee7 100644 --- a/Source/Encuda/encuda-link.swift +++ b/Source/Encuda/encuda-link.swift @@ -85,9 +85,9 @@ extension Encuda { process.executableURL = URL(fileURLWithPath: resolvedNvcc) process.arguments = ["-ccbin=\(resolvedClangpp)"] + (verbose ? ["-v"] : []) + args try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.nvccFailed(process.terminationStatus) + let status = process.waitForExitStatus() + guard status == 0 else { + throw EncudaError.nvccFailed(status) } #endif } @@ -98,9 +98,9 @@ extension Encuda { process.executableURL = URL(fileURLWithPath: resolvedClangpp) process.arguments = (verbose ? ["-v"] : []) + args try process.run() - process.waitUntilExitWorkaround() - guard process.terminationStatus == 0 else { - throw EncudaError.clangFailed(process.terminationStatus) + let status = process.waitForExitStatus() + guard status == 0 else { + throw EncudaError.clangFailed(status) } #endif } diff --git a/Source/Encuda/encuda-utils.swift b/Source/Encuda/encuda-utils.swift index ff20be2b4..eae4bf00a 100644 --- a/Source/Encuda/encuda-utils.swift +++ b/Source/Encuda/encuda-utils.swift @@ -1,5 +1,13 @@ import Foundation +#if canImport(Glibc) + import Glibc +#elseif canImport(Musl) + import Musl +#elseif canImport(Darwin) + import Darwin +#endif + func searchForCommand(_ name: String) -> URL? { let path = ProcessInfo.processInfo.environment["PATH"] ?? "" for folder in path.split(separator: ":") { @@ -13,15 +21,104 @@ func searchForCommand(_ name: String) -> URL? { #if os(macOS) || os(Linux) extension Process { - /** - * When running on aarch64 Ubuntu 24 docker on Mac, it happens that - * the `waitUntilExit()` method never returns. - * We adopted this workaround. - */ - func waitUntilExitWorkaround() { - while isRunning { - Thread.sleep(forTimeInterval: 0.05) + /// Waits for the spawned tool to finish and returns its exit status, + /// reaping the child directly rather than going through Foundation. + /// + /// **Why this does not use `waitUntilExit()`.** On aarch64 Linux that + /// call was observed never returning, and the previous workaround here + /// — polling `isRunning` in a 50ms sleep loop — did not fix it. It + /// only changed how the hang presented, because `isRunning` reads the + /// same Foundation bookkeeping that stalls. + /// + /// Observed during a parallel SwiftPM build of the CUDA backend for + /// linux/arm64: `nvcc` finishes and writes its output in full, then + /// becomes a zombie that is never reaped, and `encuda-tool` spins in + /// `hrtimer_nanosleep` indefinitely holding a completed 15MB result on + /// disk. It is load-dependent — `--jobs 2` transpiled all 96 `.cu` + /// files without incident, while an unconstrained build (SwiftPM + /// defaults to every core) deadlocked — which is consistent with pipe + /// file descriptors leaking between concurrently spawned siblings: + /// `encuda-tool` was holding both the read and the write end of the + /// same pipes, so no reader could ever observe EOF. + /// + /// `waitpid` answers the only question this code actually has — did + /// the child exit, and with what status — without depending on any of + /// that. + /// + /// - Returns: the child's exit code, or `128 + signal` if it was + /// terminated by a signal, matching the shell's convention. + func waitForExitStatus() -> Int32 { + var status: Int32 = 0 + while true { + let reaped = waitpid(processIdentifier, &status, 0) + if reaped == processIdentifier { + // Decode the wait status without the C macros, which are + // not imported into Swift. This encoding is shared by + // Linux and Darwin: the low 7 bits hold the terminating + // signal (0 when the process exited normally), and the + // next 8 bits hold the exit code. + let terminatingSignal = status & 0x7F + if terminatingSignal == 0 { + return (status >> 8) & 0xFF + } + return 128 + terminatingSignal + } + if reaped == -1 && errno == EINTR { + continue + } + // `waitpid` has nothing to reap. Overwhelmingly this is ECHILD: + // Foundation runs its own reaper thread over every `Process` it + // spawns, and on a parallel build it sometimes gets there + // first. Whatever the reason, the child is gone and Foundation's + // bookkeeping is now the only available source of truth. + return foundationTerminationStatus() + } + } + + /// Foundation's view of the exit status, waited for rather than read + /// straight away. + /// + /// `terminationStatus` is not safe to read the instant `waitpid` + /// returns ECHILD. Foundation publishes the result in two steps, and + /// reading the status inside the window before `isRunning` goes false + /// **traps the process**: + /// + /// ``` + /// *** Program crashed: System trap at 0x0000ffff90b6fb50 *** + /// Thread 0 "encuda-tool" crashed: + /// 0 Process.terminationStatus.getter + 80 in libFoundation.so + /// 1 Process.waitForExitStatus() + 107 in encuda-tool + /// ``` + /// + /// Observed on aarch64 Linux with Swift 6.3.1 during an unconstrained + /// SwiftPM build of mlx's CUDA backend: one transpile in ~250 dies this + /// way, which kills `swift build` with exit code 1 and *no diagnostic*, + /// because the crash report goes to the transpiler's own output rather + /// than surfacing as a build error. + /// + /// So wait for the flag — but on a deadline. An unbounded wait here + /// would reintroduce exactly the hang this whole function exists to + /// avoid. If the deadline expires, report a non-zero status instead: + /// the caller turns that into a build failure, which is diagnosable, + /// whereas a hang and a trap are not. + private func foundationTerminationStatus() -> Int32 { + let deadline = Date().addingTimeInterval(reaperSettleTimeout) + while isRunning && Date() < deadline { + usleep(1000) + } + guard !isRunning else { + return reaperTimedOutStatus } + return terminationStatus } + + /// How long to let Foundation finish publishing a child's exit status + /// before giving up on it. Generous, because exceeding it means a build + /// failure: the window being waited on is normally sub-millisecond. + private var reaperSettleTimeout: TimeInterval { 30 } + + /// Reported when Foundation never publishes a status. Chosen from the + /// range no compiler returns, so it is recognisable in a build log. + private var reaperTimedOutStatus: Int32 { 254 } } #endif From 3a24fa39e059ea5de65cdefc09d3d351b1053f52 Mon Sep 17 00:00:00 2001 From: "Dr. Thorsten Fuchs" Date: Fri, 14 Aug 2026 18:11:53 +0200 Subject: [PATCH 19/39] Condition CUDA build plugin on Linux host (#447) --- .github/workflows/pull_request.yml | 3 ++ Package.swift | 54 +++++++++++++++++----------- scripts/verify-spm-manifest.py | 58 ++++++++++++++++++++++++++++++ 3 files changed, 94 insertions(+), 21 deletions(-) create mode 100755 scripts/verify-spm-manifest.py diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index fad6ac5f9..56452955c 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -83,6 +83,9 @@ jobs: shell: bash run: xcodebuild -showComponent MetalToolchain + - name: Verify SwiftPM manifest CUDA containment + run: scripts/verify-spm-manifest.py + - name: Setup cmake shell: sh run: | diff --git a/Package.swift b/Package.swift index e36e0d862..941a108a5 100644 --- a/Package.swift +++ b/Package.swift @@ -4,6 +4,35 @@ import PackageDescription +#if os(Linux) + let cudaBuildPlugins: [Target.PluginUsage] = [ + .plugin(name: "CudaBuild") + ] + let cudaPackageDependencies: [Package.Dependency] = [ + .package(url: "https://github.com/apple/swift-argument-parser", from: "1.0.0") + ] + let cudaTargets: [Target] = [ + .executableTarget( + name: "encuda", + dependencies: [ + .product(name: "ArgumentParser", package: "swift-argument-parser") + ], + path: "Source/Encuda", + ), + .plugin( + name: "CudaBuild", + capability: .buildTool(), + dependencies: [ + .target(name: "encuda") + ], + ), + ] +#else + let cudaBuildPlugins: [Target.PluginUsage] = [] + let cudaPackageDependencies: [Package.Dependency] = [] + let cudaTargets: [Target] = [] +#endif + let noMetalCmlxExcludes = [ // Exclude Metal backend files, but keep no_metal.cpp for stubs // "mlx/mlx/backend/metal/no_metal.cpp", @@ -289,9 +318,7 @@ let cmlx = Target.target( .define("MLX_VERSION", to: "\"0.31.1\""), ], linkerSettings: linkerSettings, - plugins: [ - .plugin(name: "CudaBuild") - ], + plugins: cudaBuildPlugins, ) let package = Package( @@ -316,9 +343,8 @@ let package = Package( ], dependencies: [ // for Complex type - .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0"), - .package(url: "https://github.com/apple/swift-argument-parser", from: "1.0.0"), - ], + .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0") + ] + cudaPackageDependencies, targets: [ cmlx, .testTarget( @@ -414,21 +440,7 @@ let package = Package( path: "Source/Examples", sources: ["CustomFunctionExampleSimple.swift"] ), - .executableTarget( - name: "encuda", - dependencies: [ - .product(name: "ArgumentParser", package: "swift-argument-parser") - ], - path: "Source/Encuda", - ), - .plugin( - name: "CudaBuild", - capability: .buildTool(), - dependencies: [ - .target(name: "encuda") - ], - ), - ], + ] + cudaTargets, cxxLanguageStandard: .gnucxx20 ) diff --git a/scripts/verify-spm-manifest.py b/scripts/verify-spm-manifest.py new file mode 100755 index 000000000..d3d3d0f62 --- /dev/null +++ b/scripts/verify-spm-manifest.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 + +import json +import platform +import subprocess + + +def require(condition: bool, message: str) -> None: + if not condition: + raise SystemExit(message) + + +manifest = json.loads( + subprocess.run( + ["swift", "package", "dump-package"], + check=True, + capture_output=True, + text=True, + ).stdout +) + +targets = {target["name"]: target for target in manifest["targets"]} +dependencies = { + dependency["sourceControl"][0]["identity"] + for dependency in manifest["dependencies"] + if "sourceControl" in dependency +} + +cuda_targets = {"CudaBuild", "encuda"} +cuda_plugin_usages = targets["Cmlx"].get("pluginUsages") or [] + +if platform.system() == "Linux": + require( + cuda_targets <= targets.keys(), "Linux manifest is missing CUDA build targets" + ) + require( + "swift-argument-parser" in dependencies, + "Linux manifest is missing the CUDA helper dependency", + ) + require( + cuda_plugin_usages, + "Linux Cmlx target is missing the CUDA build plugin", + ) +else: + require( + cuda_targets.isdisjoint(targets), + "non-Linux manifest unexpectedly contains CUDA build targets", + ) + require( + "swift-argument-parser" not in dependencies, + "non-Linux manifest unexpectedly contains the CUDA helper dependency", + ) + require( + not cuda_plugin_usages, + "non-Linux Cmlx target unexpectedly uses the CUDA build plugin", + ) + +print(f"SwiftPM manifest CUDA containment passed for {platform.system()}.") From 7b745d2b736e117d3aaa7effe4136b2c004fe6f3 Mon Sep 17 00:00:00 2001 From: Valentin Roussellet Date: Mon, 17 Aug 2026 10:24:02 -0700 Subject: [PATCH 20/39] Nested initializers (#456) * nested initializers --- Source/MLX/Documentation.docc/MLX.md | 1 + .../Organization/initialization.md | 28 ++++- Source/MLX/MLXArray+Init.swift | 87 ++++++++++++++- Source/MLX/NestedArrayElement.swift | 95 ++++++++++++++++ Tests/MLXTests/MLXArray+InitTests.swift | 105 ++++++++++++++++++ skills/mlx-swift/SKILL.md | 8 +- skills/mlx-swift/references/arrays.md | 4 + 7 files changed, 321 insertions(+), 7 deletions(-) create mode 100644 Source/MLX/NestedArrayElement.swift diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md index 26bfe4577..5e2f98c5c 100644 --- a/Source/MLX/Documentation.docc/MLX.md +++ b/Source/MLX/Documentation.docc/MLX.md @@ -109,6 +109,7 @@ See for configuration, best practices, and policy guidance. - ``DType`` - ``HasDType`` - ``ScalarOrArray`` +- ``NestedArrayElement`` ### Parameter Types diff --git a/Source/MLX/Documentation.docc/Organization/initialization.md b/Source/MLX/Documentation.docc/Organization/initialization.md index 5e9e82e2a..4b6fb0ba7 100644 --- a/Source/MLX/Documentation.docc/Organization/initialization.md +++ b/Source/MLX/Documentation.docc/Organization/initialization.md @@ -87,7 +87,7 @@ let a2 = MLXArray(int64: 500) XCTAssertEqual(a2.dtype, .int64) ``` -All of the `Int` initializers (e.g. `[Int]` and `Sequence`) work +All of the `Int` initializers (e.g. `[Int]`, `Sequence` and nested arrays) work the same way and all have the `int64:` variant. ### Double @@ -120,6 +120,20 @@ let v1 = MLXArray(0 ..< 12) let v2 = MLXArray(stride(from: Float(0.5), to: Float(1.5), by: Float(0.1))) ``` +Multiple dimensions can be given directly as a nested array, in which case, the +shape is inferred from the nesting. + +```swift +// create an array of Int32 with shape [2, 3] +let v1 = MLXArray([[1, 2, 3], + [4, 5, 6]]) + +// any depth of nesting works -- shape [2, 2, 2] +let v2 = MLXArray([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) +``` + +Nested arrays must be tensors: every element at a given level must have the same shape. + If you have `Data` or a `UnsafePointer` (of various kinds) you can also create an `MLXArray` from that: @@ -159,7 +173,7 @@ let identity = MLXArray.identity(5) ### Complex Values -``MLXArray`` supports complex numbers, specifically a real and imaginary `Float32` +``MLXArray`` supports complex numbers, specifically a real and imaginary `Float32` as ``DType/complex64``. MLX uses [swift-numerics](https://github.com/apple/swift-numerics/tree/main) to represent the `Complex` type, though there are a few functions for manipulating @@ -212,9 +226,11 @@ there are specific initializers to request it: - ``MLXArray/init(_:)-(Int)`` - ``MLXArray/init(_:_:)-([Int],_)`` -- ``MLXArray/init(int64:)`` +- ``MLXArray/init(_:)-([[N]])`` +- ``MLXArray/init(int64:)-(Int)`` - ``MLXArray/init(int64:_:)-([Int],_)`` - ``MLXArray/init(int64:_:)-(Sequence,_)`` +- ``MLXArray/init(int64:)-([[N]])`` ### MLXArray Array Initializers @@ -225,6 +241,12 @@ there are specific initializers to request it: - ``MLXArray/init(_:_:type:)-(UnsafeRawBufferPointer,_,_)`` - ``MLXArray/init(_:_:type:)-(Data,_,_)`` +### MLXArray Nested Array Initializers + +- ``MLXArray/init(_:)-([[N]])`` +- ``MLXArray/init(int64:)-([[N]])`` +- ``MLXArray/init(converting:)`` + ### MLXArray Complex Initializers - ``MLXArray/init(real:imaginary:)`` diff --git a/Source/MLX/MLXArray+Init.swift b/Source/MLX/MLXArray+Init.swift index 9ca79d333..9aa04d109 100644 --- a/Source/MLX/MLXArray+Init.swift +++ b/Source/MLX/MLXArray+Init.swift @@ -18,6 +18,17 @@ private func shapePrecondition(shape: (some Collection)?, byteCount: Int, t } } +// flatten a nested array into its shape and its scalars in row major order +private func flattenNested(_ value: [[N]]) -> ( + shape: [Int], values: [N.Scalar] +) { + let shape = value.nestedShape + var values = [N.Scalar]() + values.reserveCapacity(shape.reduce(1, *)) + value.appendNestedScalars(to: &values) + return (shape, values) +} + // holds reference to `finalizer` as capture state private class FinalizerCaptureState { let f: () -> Void @@ -121,11 +132,11 @@ extension MLXArray { /// ``` /// /// Note: if the value is out of bounds for an `Int32` the precondition will fail. If you - /// need an `Int` (`Int64`) scalar, please use ``init(int64:)``. + /// need an `Int` (`Int64`) scalar, please use ``init(int64:)-(Int)``. /// /// ### See Also /// - - /// - ``init(int64:)`` + /// - ``init(int64:)-(Int)`` public convenience init(_ value: Int) { precondition( (Int(Int32.min) ... Int(Int32.max)).contains(value), @@ -451,6 +462,78 @@ extension MLXArray { fatalError("unavailable") } + /// Initializer allowing creation of a multi dimensional `MLXArray` from a nested array of + /// `HasDType` values. + /// + /// The shape is inferred directly from the provided array. Note that the nesting must + /// be rectangular (all arrays of a given level must have the same count). + /// Failing to meet this constraint will hit a precondition failure. + /// + /// ```swift + /// // shape [2, 3] + /// let twoByThree = MLXArray([[1, 2, 3], + /// [4, 5, 6]]) + /// + /// // shape [2, 2, 2] -- any depth of nesting works + /// let cube = MLXArray([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + /// ``` + /// + /// Note: as with the one dimensional initializers, `Int` values produce a ``DType/int32`` + /// result and the precondition will fail if a value is out of bounds. See + /// ``init(int64:)-([[N]])`` if `.int64` is required. + /// + /// ### See Also + /// - + /// - ``init(int64:)-([[N]])`` + public convenience init(_ value: [[N]]) { + let (shape, values) = flattenNested(value) + if N.Scalar.self == Int.self { + // Similarly to the Sequence initializer below, + // having an override for Int is ambiguous so we + // do a runtime check and force it to the [Int] variant + self.init(values as! [Int], shape) + } else { + self.init(values, shape) + } + } + + /// Initializer allowing creation of a multi dimensional `MLXArray` from a nested array of + /// `Int` as a `DType.int64`. + /// + /// ```swift + /// let a = MLXArray(int64: [[1, 2, 3], [4, 5, 6]]) + /// ``` + /// + /// Note ``init(_:)-([[N]])`` (producing an `int32` result) is preferred. + /// + /// ### See Also + /// - + /// - ``init(_:)-([[N]])`` + public convenience init(int64 value: [[N]]) where N.Scalar == Int { + let (shape, values) = flattenNested(value) + self.init(int64: values, shape) + } + + /// Initializer allowing creation of a multi dimensional `MLXArray` from a nested array of + /// `Double` values. + /// + /// Note: this converts the values to `Float`, matching ``init(converting:_:)``. Unlike the + /// one dimensional case, ``init(_:)-([[N]])`` also accepts `[[Double]]` and produces a + /// ``DType/float64`` result. + /// + /// ```swift + /// let a = MLXArray(converting: [[0.5, 0.9], [1.5, 1.9]]) + /// ``` + /// + /// ### See Also + /// - + /// - ``init(_:)-([[N]])`` + public convenience init(converting value: [[N]]) + where N.Scalar == Double { + let (shape, values) = flattenNested(value) + self.init(converting: values, shape) + } + /// Initializer allowing creation of `MLXArray` from a sequence of `HasDType` values with /// an optional shape. /// diff --git a/Source/MLX/NestedArrayElement.swift b/Source/MLX/NestedArrayElement.swift new file mode 100644 index 000000000..8ec81d759 --- /dev/null +++ b/Source/MLX/NestedArrayElement.swift @@ -0,0 +1,95 @@ +// Copyright © 2026 Apple Inc. + +import Numerics + +/// A value that can appear in a nested array used to create an ``MLXArray``. +/// +/// Conformance is recursive: the built-in ``HasDType`` scalars conform and so does an `Array` of +/// conforming values, adding one dimension. A nested array therefore describes both the shape +/// and the contents of an ``MLXArray``: +/// +/// ```swift +/// // shape [2, 3] +/// let a = MLXArray([[1, 2, 3], [4, 5, 6]]) +/// ``` +/// +/// This exists to express the nested array initializers, e.g. ``MLXArray/init(_:)-([[N]])`` -- +/// there should be no need to conform your own types to it. +/// +/// ### See Also +/// - +public protocol NestedArrayElement { + + /// The scalar type at the leaves of the nesting, e.g. `Int` for `[[Int]]`. + associatedtype Scalar: HasDType + + /// The number of dimensions this type contributes -- `0` for a scalar, `1` for `[Scalar]`. + static var nestedRank: Int { get } + + /// The shape of the receiver. + /// + /// Note: the nesting must be rectangular -- every element at a given level must have the + /// same shape -- or the precondition will fail. + /// + /// Note: an empty level says nothing about the levels below it, so those are reported as + /// `0` -- `[[Int]]()` has shape `[0, 0]`. + var nestedShape: [Int] { get } + + /// Append the receiver's scalars to `values` in row major order. + func appendNestedScalars(to values: inout [Scalar]) +} + +extension NestedArrayElement where Self: HasDType { + public static var nestedRank: Int { 0 } + + public var nestedShape: [Int] { [] } + + public func appendNestedScalars(to values: inout [Self]) { + values.append(self) + } +} + +// the scalars, providing the base case of the recursion. HasDType deliberately does not refine +// this protocol: conditional conformances to HasDType (Complex is one) would not imply it, +// so that would break any such conformance outside this package. Keep this list in sync with +// the HasDType conformances in DType.swift. +extension Bool: NestedArrayElement {} +extension Int: NestedArrayElement {} +extension Int8: NestedArrayElement {} +extension Int16: NestedArrayElement {} +extension Int32: NestedArrayElement {} +extension Int64: NestedArrayElement {} +extension UInt8: NestedArrayElement {} +extension UInt16: NestedArrayElement {} +extension UInt32: NestedArrayElement {} +extension UInt64: NestedArrayElement {} +extension UInt: NestedArrayElement {} +#if !arch(x86_64) + extension Float16: NestedArrayElement {} +#endif +extension Float32: NestedArrayElement {} +extension Float64: NestedArrayElement {} +extension Complex: NestedArrayElement {} + +extension Array: NestedArrayElement where Element: NestedArrayElement { + public typealias Scalar = Element.Scalar + + public static var nestedRank: Int { 1 + Element.nestedRank } + + public var nestedShape: [Int] { + // an empty level tells us nothing about the levels below it, but the static rank does + let inner = first?.nestedShape ?? [Int](repeating: 0, count: Element.nestedRank) + for element in dropFirst() { + precondition( + element.nestedShape == inner, + "nested array is not rectangular: \(element.nestedShape) != \(inner)") + } + return [count] + inner + } + + public func appendNestedScalars(to values: inout [Element.Scalar]) { + for element in self { + element.appendNestedScalars(to: &values) + } + } +} diff --git a/Tests/MLXTests/MLXArray+InitTests.swift b/Tests/MLXTests/MLXArray+InitTests.swift index acd487306..8a6cde37e 100644 --- a/Tests/MLXTests/MLXArray+InitTests.swift +++ b/Tests/MLXTests/MLXArray+InitTests.swift @@ -89,6 +89,111 @@ class MLXArrayInitTests: XCTestCase { XCTAssertEqual(a.ndim, 2) } + // MARK: - Nested (multi dimensional) arrays + + func testArrayCreationNested2D() { + let a = MLXArray([ + [1, 2, 3], + [4, 5, 6], + ]) + XCTAssertEqual(a.dtype, .int32) + XCTAssertEqual(a.shape, [2, 3]) + XCTAssertEqual(a.ndim, 2) + XCTAssertEqual(a.size, 6) + + // count is the size of dimension 0 + XCTAssertEqual(a.count, 2) + + assertEqual(a, MLXArray(1 ... 6, [2, 3])) + } + + func testArrayCreationNested3D() { + // literal nesting and the flat + shape form must agree + let a = MLXArray([ + [[0, 1], [2, 3], [4, 5]], + [[6, 7], [8, 9], [10, 11]], + ]) + XCTAssertEqual(a.dtype, .int32) + XCTAssertEqual(a.shape, [2, 3, 2]) + assertEqual(a, MLXArray(0 ..< 12, [2, 3, 2])) + + // also works with a computed (non literal) nested array + let nested = (0 ..< 2).map { i in + (0 ..< 3).map { j in + (0 ..< 4).map { k in i * 12 + j * 4 + k } + } + } + let b = MLXArray(nested) + XCTAssertEqual(b.shape, [2, 3, 4]) + assertEqual(b, MLXArray(0 ..< 24, [2, 3, 4])) + } + + func testArrayCreationNested4D() { + let a = MLXArray([ + [[[0, 1], [2, 3]], [[4, 5], [6, 7]]], + [[[8, 9], [10, 11]], [[12, 13], [14, 15]]], + ]) + XCTAssertEqual(a.shape, [2, 2, 2, 2]) + XCTAssertEqual(a.ndim, 4) + assertEqual(a, MLXArray(0 ..< 16, [2, 2, 2, 2])) + } + + func testArrayCreationNestedDTypes() { + // the dtype of a nested array matches the 1d initializers for the same leaf type + XCTAssertEqual(MLXArray([[1, 2], [3, 4]]).dtype, .int32) + XCTAssertEqual(MLXArray([[Int32(1), 2], [3, 4]]).dtype, .int32) + XCTAssertEqual(MLXArray([[Int16(1), 2], [3, 4]]).dtype, .int16) + XCTAssertEqual(MLXArray([[UInt8(1), 2], [3, 4]]).dtype, .uint8) + XCTAssertEqual(MLXArray([[Float(1), 2], [3, 4]]).dtype, .float32) + XCTAssertEqual(MLXArray([[true, false], [false, true]]).dtype, .bool) + + // Int64 is a distinct type from Int and must not take the Int -> .int32 path + XCTAssertEqual(MLXArray([[Int64(1), 2], [3, 4]]).dtype, .int64) + + // casting the literal as a whole is a different inference path than casting the + // elements and must stay unambiguous + XCTAssertEqual( + MLXArray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] as [[Float32]]).dtype, .float32) + XCTAssertEqual(MLXArray([[[1.0, 2.0]]] as [[[Float32]]]).dtype, .float32) + + // as with `MLXArray([Double])` this produces .float64 rather than promoting + let doubles: [[Double]] = [[1, 2], [3, 4]] + XCTAssertEqual(MLXArray(doubles).dtype, .float64) + + let complex: [[Complex]] = [[Complex(2, 7)], [Complex(3, 8)]] + let c = MLXArray(complex) + XCTAssertEqual(c.dtype, .complex64) + XCTAssertEqual(c.shape, [2, 1]) + assertEqual(c, MLXArray([Complex(2, 7), Complex(3, 8)], [2, 1])) + } + + func testArrayCreationNestedInt64() { + // a value that does not fit in an Int32 -- the whole point of the int64: variant + let big = Int(Int32.max) + 10 + let a = MLXArray(int64: [[1, 2, 3], [4, 5, big]]) + XCTAssertEqual(a.dtype, .int64) + XCTAssertEqual(a.shape, [2, 3]) + assertEqual(a, MLXArray(int64: [1, 2, 3, 4, 5, big], [2, 3])) + } + + func testArrayCreationNestedConverting() { + // non square so a transposed flattening could not pass + let a = MLXArray(converting: [[0.1, 0.5, 0.9], [1.3, 1.7, 2.1]]) + XCTAssertEqual(a.dtype, .float32) + XCTAssertEqual(a.shape, [2, 3]) + assertEqual(a, MLXArray(converting: [0.1, 0.5, 0.9, 1.3, 1.7, 2.1], [2, 3])) + } + + func testArrayCreationNestedEmpty() { + let a = MLXArray([[Int](), [Int]()]) + XCTAssertEqual(a.shape, [2, 0]) + XCTAssertEqual(a.size, 0) + + let b = MLXArray([[Int]]()) + XCTAssertEqual(b.shape, [0, 0]) + XCTAssertEqual(b.size, 0) + } + func testArrayCreationClosedRange() { let a = MLXArray(Int16(3) ... Int16(6)) XCTAssertEqual(a.dtype, .int16) diff --git a/skills/mlx-swift/SKILL.md b/skills/mlx-swift/SKILL.md index ba6fb71c4..70681b257 100644 --- a/skills/mlx-swift/SKILL.md +++ b/skills/mlx-swift/SKILL.md @@ -62,8 +62,12 @@ import MLX // Create arrays let a = MLXArray([1, 2, 3, 4]) let b = MLXArray(0 ..< 12, [3, 4]) // Shape [3, 4] -let c = MLXArray.zeros([2, 3]) -let d = MLXArray.ones([4, 4], dtype: .float32) + +// Nested arrays give the shape directly (any depth) +let c = MLXArray([[1, 2, 3], [4, 5, 6]]) // Shape [2, 3] + +let d = MLXArray.zeros([2, 3]) +let e = MLXArray.ones([4, 4], dtype: .float32) // Random arrays (use MLXRandom namespace or free functions) let uniform = MLXRandom.uniform(0.0 ..< 1.0, [3, 3]) diff --git a/skills/mlx-swift/references/arrays.md b/skills/mlx-swift/references/arrays.md index 34ef2c497..806f4ace4 100644 --- a/skills/mlx-swift/references/arrays.md +++ b/skills/mlx-swift/references/arrays.md @@ -22,6 +22,10 @@ let c = MLXArray(0 ..< 10) let d = MLXArray(0 ..< 12, [3, 4]) // 3 rows, 4 columns let e = MLXArray([1, 2, 3, 4, 5, 6], [2, 3]) +// Multi-dimensional from a nested array -- the nesting gives the shape +let f = MLXArray([[1, 2, 3], [4, 5, 6]]) // shape [2, 3] +let g = MLXArray([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) // shape [2, 2, 2] + // Complex numbers let complex = MLXArray(real: 1.0, imaginary: 2.0) ``` From 97cf19efeaa4e929415e75982e999adb34f62c0d Mon Sep 17 00:00:00 2001 From: GoodOlClint Date: Mon, 17 Aug 2026 16:36:51 -0500 Subject: [PATCH 21/39] [CUDA] Build the CUDA argument lists in steps (#442) Three argument lists in the CUDA build are assembled as long `+` chains that the type checker struggles with. On a Swift 6.2 toolchain they exceed its budget outright: encuda-compile.swift:66:17: error: the compiler is unable to type-check this expression in reasonable time; try breaking up the expression into distinct sub-expressions plugin.swift:172:13: error: the compiler is unable to type-check this expression in reasonable time; ... Appends into local arrays instead, which is behaviourally identical and is what the diagnostic suggests. The plugin's two chains move into `compileArguments` and `linkArguments` helpers to keep `createBuildCommands` readable, and adding a further argument to either list no longer risks tipping the type checker over. Note this does *not* make the package buildable on Swift 6.2: SwiftPM 6.2 ignores non-Swift files produced by a build-tool plugin ("Only Swift is supported for generated plugin source files at this time"), so every generated .cpp is dropped and the CUDA symbols go undefined at link time. The 6.3 floor is load-bearing. This is purely a type-check-cost and readability change. Co-authored-by: Claude Opus 5 --- Plugins/CudaBuild/plugin.swift | 71 +++++++++++++++++++++++++----- Source/Encuda/encuda-compile.swift | 18 ++++++-- 2 files changed, 75 insertions(+), 14 deletions(-) diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift index f7d12dd4c..15ecd2e88 100644 --- a/Plugins/CudaBuild/plugin.swift +++ b/Plugins/CudaBuild/plugin.swift @@ -140,13 +140,16 @@ struct CudaBuild: BuildToolPlugin { displayName: "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)", executable: encuda.url, - arguments: ["compile"] + verboseFlag + stdArgs + incrementalFlag + [ - "--clangpp", clangUrl.url.path, - "-I", sourceDir.path, - ] + headerSearchPathArgs + [ - inputFile.path, - "-o", outputCpp.path, - ], + arguments: compileArguments( + clangpp: clangUrl.url.path, + sourceDir: sourceDir.path, + verboseFlag: verboseFlag, + stdArgs: stdArgs, + incrementalFlag: incrementalFlag, + headerSearchPathArgs: headerSearchPathArgs, + input: inputFile.path, + output: outputCpp.path + ), inputFiles: [inputFile], outputFiles: [outputCpp] ) @@ -167,9 +170,14 @@ struct CudaBuild: BuildToolPlugin { .buildCommand( displayName: "Linking CUDA objects", executable: encuda.url, - arguments: ["link"] + verboseFlag + stdArgs + incrementalFlag + [ - "--clangpp", clangUrl.url.path, - ] + outputCpps.map { $0.path } + ["-o", linkOutput.path], + arguments: linkArguments( + clangpp: clangUrl.url.path, + verboseFlag: verboseFlag, + stdArgs: stdArgs, + incrementalFlag: incrementalFlag, + inputs: outputCpps.map { $0.path }, + output: linkOutput.path + ), inputFiles: outputCpps, outputFiles: [linkOutput] ) @@ -178,6 +186,49 @@ struct CudaBuild: BuildToolPlugin { return commands } + /// Assembled in steps rather than as one `+` chain, which is expensive to + /// type-check and gets more so with each argument added. + func compileArguments( + clangpp: String, + sourceDir: String, + verboseFlag: [String], + stdArgs: [String], + incrementalFlag: [String], + headerSearchPathArgs: [String], + input: String, + output: String + ) -> [String] { + var arguments: [String] = ["compile"] + arguments += verboseFlag + arguments += stdArgs + arguments += incrementalFlag + arguments += ["--clangpp", clangpp] + arguments += ["-I", sourceDir] + arguments += headerSearchPathArgs + arguments += [input] + arguments += ["-o", output] + return arguments + } + + /// See `compileArguments` for why this is not a single expression. + func linkArguments( + clangpp: String, + verboseFlag: [String], + stdArgs: [String], + incrementalFlag: [String], + inputs: [String], + output: String + ) -> [String] { + var arguments: [String] = ["link"] + arguments += verboseFlag + arguments += stdArgs + arguments += incrementalFlag + arguments += ["--clangpp", clangpp] + arguments += inputs + arguments += ["-o", output] + return arguments + } + func isExcluded(settings: Settings, relativePath: String) -> Bool { func isInFolder(path: String, folderPath: String) -> Bool { let prefix = folderPath.hasSuffix("/") ? folderPath : (folderPath + "/") diff --git a/Source/Encuda/encuda-compile.swift b/Source/Encuda/encuda-compile.swift index 65ea7b0f6..34ded3f97 100644 --- a/Source/Encuda/encuda-compile.swift +++ b/Source/Encuda/encuda-compile.swift @@ -61,12 +61,22 @@ extension Encuda { let archArgs = ProcessInfo.processInfo.environment["CUDA_ARCH"].map { ["-arch", $0] } ?? [] + let verboseArgs: [String] = verbose ? ["-v"] : [] + + // Built up in steps rather than as one `+` chain, which is expensive + // to type-check and gets more so with each argument added. + var arguments: [String] = ["-cuda", "-rdc=true", "--expt-relaxed-constexpr"] + arguments += stdArgs + arguments += ccbinArgs + arguments += archArgs + arguments += verboseArgs + arguments += includeArgs + arguments += inputFiles + arguments += ["-o", output] + let process = Process() process.executableURL = URL(fileURLWithPath: resolvedNvcc) - process.arguments = - ["-cuda", "-rdc=true", "--expt-relaxed-constexpr"] + stdArgs + ccbinArgs - + archArgs - + (verbose ? ["-v"] : []) + includeArgs + inputFiles + ["-o", output] + process.arguments = arguments try process.run() let status = process.waitForExitStatus() guard status == 0 else { From df9ae2623680ee644a9a532247bd60f5b0ccd5e8 Mon Sep 17 00:00:00 2001 From: DanielMercado60660 <153678010+DanielMercado60660@users.noreply.github.com> Date: Mon, 24 Aug 2026 13:55:31 -0400 Subject: [PATCH 22/39] Fix a deadlock between `CompiledFunction.lock` and `evalLock`, and a data race on the compiler cache (#461) * add a deterministic repro for the compile/eval lock-order inversion CompiledFunction.call takes the per-instance lock (#226) and then, in innerCall, the process-global evalLock (#339), holding it across mlx_closure_apply. On a cold cache entry that runs compile_trace, which calls back into the Swift body being compiled; if that body calls another compiled function -- every MLXNN activation does, e.g. silu calls the file-scope compiledSilu -- the nested call takes a second per-instance lock while evalLock is held. thread H: L_cf(outer) -> evalLock -> wants L_cf(shared) thread W: L_cf(shared) -> wants evalLock The orders are inverted and the two threads deadlock permanently. CompileLockRepro owns its threads, uses fresh compile() results per round so the compiler cache is cold, signals from inside the traced body so W is ordered after H enters the trace, and holds the window open with a usleep in that body. A per-round deadline turns the hang into a hard failure. It is its own executable because a wedged round cannot be recovered and would take a shared test process down with it. Co-Authored-By: Claude Fable 5 * take evalLock outermost in CompiledFunction.call Establishes one global lock order, evalLock -> CompiledFunction.lock, so the inversion described in the previous commit cannot form. The reverse order is not avoidable on its own: the trace has to run under evalLock (the compiler cache is an unguarded global and the tracing flag is a plain static, not thread_local), the trace calls back into Swift, and the traced body calls other compiled functions -- every MLXNN activation is a file-scope CompiledFunction. So evalLock -> lock happens regardless. The only fix available is to make it the *only* order, which is what taking evalLock at the top of call(_:) does. evalLock is recursive, so the nested acquisition from inside a trace is free. innerCall already took evalLock before doing any real work and held it through mlx_closure_apply, so this widens the held region only by the state gathering and closure creation at the top of innerCall. lock becomes an NSRecursiveLock so that a compiled function whose body re-enters that same compiled function on one thread no longer self- deadlocks on a non-recursive NSLock. CompileLockRepro: 20/20 rounds pass; without this it wedges on round 1. Co-Authored-By: Claude Fable 5 * skip the per-instance lock when a compiled function has no state #226 added CompiledFunction.lock to guard the observed state. The _updateInternal swaps in innerCall are the only mutation of anything reachable from self -- every other stored property is let, and id is written once at init -- so when inputs and outputs are both empty the lock has nothing to guard. With the previous commit evalLock is held for the whole call, which already serializes the compiler cache and every other caller, so dropping the acquisition changes nothing observable. Optional on top of the previous commit; the deadlock is fixed without it. It removes a mutex round-trip from the hot path of the common case: every compiled function in MLXNN is stateless. Co-Authored-By: Claude Fable 5 * check the evalLock -> CompiledFunction.lock ordering in debug builds A passing CompileLockRepro only shows that the race did not fire, so assert the ordering itself. evalLock is recursive, and a successful try() on a recursive lock cannot distinguish 'this thread owns it' from 'nobody owns it', so ownership has to be recorded explicitly. withEvalLock is now the single choke point for evalLock; in debug builds it keeps a per-thread depth (pthread-specific, so the bookkeeping itself needs no lock -- it runs while holding evalLock). CompiledFunction acquires its per-instance lock through withInstanceLock, which requires that depth to be non-zero. Moving the per-instance lock back outside evalLock -- the shape that deadlocks -- trips it on the first compiled call, with no race to lose. All of it compiles out of release builds. CompileLockOrderingTests covers: that the bookkeeping can report 'not held' and is per thread (otherwise a violation count of zero would prove nothing), that a stateful compiled function really reaches the checked acquisition, and the nested case where one compiled function is called from inside another's trace. Verified by inverting the order locally: the assertion fires immediately with 'CompiledFunction.lock was acquired without holding evalLock'. Co-Authored-By: Claude Fable 5 * erase the compiler cache under evalLock in CompiledFunction.deinit `mlx_detail_compile_erase` is `compiler_cache().erase(fun_id)` on the process-global `mlx::core::detail::CompilerCache`, which has no synchronization of its own. `CompiledFunction.deinit` called it with no lock held, so it could run concurrently with another thread's `CompilerCache::find` -- `cache_[fun_id]`, which can itself insert and rehash, followed by `entries.push_back` -- and with the whole trace that fills in the entry `find` returned a reference to. Unlike the lock-order inversion fixed in the preceding commits, this is a data race rather than a deadlock, and it predates them. Every other use of that cache already happens under `evalLock`, so the erase takes it too. `evalLock` is recursive, so a deinit that fires on a thread already inside a compiled call or an `eval` -- a traced body dropping the last reference to some other compiled function, for instance -- re-enters for free, and a thread holding a `CompiledFunction.lock` holds `evalLock` already, so no new edge is added to the lock graph. A deinit on a thread with no MLX lock simply waits for the call in flight. Nothing in the library hands a `CompiledFunction` reference to a back-end worker thread, so this erase can never be what an `evalLock` holder is waiting for. The erase stays synchronous rather than deferred: `id` is the object's own address, so a later `CompiledFunction` can be allocated at the same address and receive the same id, and a deferred erase could delete that new function's entries. Two tests. The first is the interleaving itself: one thread is parked inside `compile_trace` holding `evalLock` (it signals a semaphore from inside the traced body), the other drops the last reference to a warmed compiled function at that instant, and the erase is required to complete only after the first thread has left the traced body -- which is what serialization against `find` means. The second runs a stream of erases against a stream of cold compiles, covering the insert-and-rehash window too, and checks the results are still correct. Both assert the debug counters so a pass cannot be vacuous. --- Package.swift | 12 + .../CompileLockRepro/CompileLockRepro.swift | 216 +++++++++++ Source/MLX/Device.swift | 2 +- Source/MLX/IO.swift | 6 +- Source/MLX/MLXArray.swift | 4 +- Source/MLX/Memory.swift | 2 +- Source/MLX/Stream.swift | 10 +- Source/MLX/Transforms+Compile.swift | 144 ++++++- Source/MLX/Transforms+Eval.swift | 113 +++++- Source/MLX/Transforms+Internal.swift | 2 +- Source/MLX/Transforms+Vmap.swift | 2 +- Source/MLX/Transforms.swift | 4 +- Source/MLX/WiredMemory.swift | 2 +- Tests/MLXTests/CompileLockOrderingTests.swift | 352 ++++++++++++++++++ 14 files changed, 844 insertions(+), 27 deletions(-) create mode 100644 Source/CompileLockRepro/CompileLockRepro.swift create mode 100644 Tests/MLXTests/CompileLockOrderingTests.swift diff --git a/Package.swift b/Package.swift index 941a108a5..04180bf6a 100644 --- a/Package.swift +++ b/Package.swift @@ -440,6 +440,18 @@ let package = Package( path: "Source/Examples", sources: ["CustomFunctionExampleSimple.swift"] ), + // Deterministic reproduction of the CompiledFunction/evalLock + // lock-order inversion. A wedged run deadlocks its threads + // permanently, so it is deliberately kept out of the test bundle and + // built as its own executable. + .executableTarget( + name: "CompileLockRepro", + dependencies: ["MLX"], + path: "Source/CompileLockRepro", + swiftSettings: [ + .enableExperimentalFeature("StrictConcurrency") + ] + ), ] + cudaTargets, cxxLanguageStandard: .gnucxx20 ) diff --git a/Source/CompileLockRepro/CompileLockRepro.swift b/Source/CompileLockRepro/CompileLockRepro.swift new file mode 100644 index 000000000..2f4577b44 --- /dev/null +++ b/Source/CompileLockRepro/CompileLockRepro.swift @@ -0,0 +1,216 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX + +/// Deterministic reproduction of a cross-thread lock-order inversion between +/// the per-instance `CompiledFunction.lock` and the process-global `evalLock`. +/// +/// ## The cycle +/// +/// `CompiledFunction.call` takes the per-instance `lock` (added by #226 to guard +/// compiled-function state) and then, inside `innerCall`, takes the global +/// `evalLock` (added by #339 to guard tracing) and holds it across +/// `mlx_closure_apply`. On a *cold* cache entry `mlx_closure_apply` runs +/// `mlx::core::detail::compile_trace`, which calls back into the Swift body of +/// the function being compiled. If that body calls a *different* compiled +/// function — which every `MLXNN` activation does, e.g. `silu` uses the +/// file-scope `compiledSilu` — the nested `CompiledFunction.call` takes a second +/// per-instance lock *while already holding `evalLock`*. +/// +/// Thread H: L_cf(outer) -> evalLock -> wants L_cf(shared) +/// Thread W: L_cf(shared) -> wants evalLock +/// +/// The two acquisition orders are inverted, so H and W deadlock permanently. +/// `evalLock` being an `NSRecursiveLock` does not help: this is not same-thread +/// re-entry. +/// +/// ## Why this is deterministic and not a race we hope to lose +/// +/// * Fresh `compile()` results per round guarantee a **cold** compiler cache. +/// The re-entry into Swift only happens while tracing, i.e. on first use, so a +/// warm instance would simply never open the window. +/// * A semaphore signalled **from inside the traced Swift body** releases W only +/// once H is provably inside `compile_trace` holding `evalLock`. +/// * A `usleep` in that same traced body holds the window open for 200 ms — +/// roughly a million times longer than the few instructions W needs — so W +/// always reaches `L_cf(shared)` first. +/// * Threads are owned here rather than borrowed from a test runner, so nothing +/// depends on runner parallelism. +/// +/// ## Why it is its own executable +/// +/// A wedged round deadlocks the threads permanently and they cannot be +/// unwedged, so this cannot live in a shared test process without taking the +/// whole lane down. The deadline turns the hang into a hard, fast failure. +/// +/// ## Running it +/// +/// swift build --product CompileLockRepro +/// .build/debug/CompileLockRepro # --device cpu|gpu +/// +/// SwiftPM (command line) cannot build the Metal shaders, and the MLX scheduler +/// opens a GPU stream when it is constructed, so a `swift build` product needs a +/// `mlx.metallib` colocated with the binary even for `--device cpu`. Copy one +/// built by `xcodebuild -project xcode/MLX.xcodeproj -scheme MLX` (it lands at +/// `Build/Products/Debug/Cmlx.framework/Versions/A/Resources/default.metallib`) +/// into `$(swift build --show-bin-path)/mlx.metallib`. +/// +/// Exit codes: `0` no deadlock, `1` deadlock (the bug is present), `2` +/// inconclusive (compilation was disabled, so no trace ever ran). +@main +struct CompileLockRepro { + + /// Per-round completion deadline. A round that has not finished by now is + /// wedged: every thread is parked in `__psynch_mutexwait` and no amount of + /// further waiting will help. + static let deadline: TimeInterval = 10 + + /// How long the traced body holds the race window open. + static let widenMicroseconds: UInt32 = 200_000 + + static let rounds = 20 + + /// Minimal thread-safe boolean; the repro must not itself depend on any + /// MLX locking. + final class Flag: @unchecked Sendable { + private let lock = NSLock() + private var value = false + func set() { + lock.lock() + value = true + lock.unlock() + } + var isSet: Bool { + lock.lock() + defer { lock.unlock() } + return value + } + } + + static func die(_ code: Int32, _ message: String) -> Never { + print(message) + fflush(stdout) + // The wedged threads hold `evalLock` and a `CompiledFunction.lock`. + // A normal `exit()` would run atexit handlers that can block on exactly + // those locks, turning a reported failure back into a hang. + _exit(code) + } + + static func main() { + // The deadlock is entirely in the Swift locking + the backend-independent + // `mlx::core::detail::compile` machinery, so the repro runs on either + // device. It defaults to `.cpu` because SwiftPM (command line) cannot + // build the Metal shaders, so a plain `swift run` has no metallib. + var device = Device.cpu + if let index = CommandLine.arguments.firstIndex(of: "--device"), + index + 1 < CommandLine.arguments.count + { + switch CommandLine.arguments[index + 1].lowercased() { + case "cpu": device = .cpu + case "gpu": device = .gpu + default: die(2, "unknown --device \(CommandLine.arguments[index + 1])") + } + } + Device.setDefault(device: device) + + print( + "compile-lock repro: device=\(device), \(rounds) rounds, \(Int(deadline))s deadline per round" + ) + + // Keep every compiled function alive for the whole run: `id` is the + // object address and `deinit` erases the cache entry, so recycling an + // address could alias two rounds. + var keepAlive: [Any] = [] + var tracedAtLeastOnce = false + + for round in 1 ... rounds { + tracedAtLeastOnce = runRound(round, keepAlive: &keepAlive) || tracedAtLeastOnce + } + + if !tracedAtLeastOnce { + die( + 2, + """ + INCONCLUSIVE: the traced body never ran, so no compile trace happened \ + and the deadlock window was never opened. Is MLX_DISABLE_COMPILE set? + """) + } + + print("PASS: \(rounds) rounds completed, no deadlock") + } + + /// Runs one round. Returns whether the traced body actually executed. + static func runRound(_ round: Int, keepAlive: inout [Any]) -> Bool { + let gate = DispatchSemaphore(value: 0) // H (inside the trace) -> W + let hDone = DispatchSemaphore(value: 0) + let wDone = DispatchSemaphore(value: 0) + let traced = Flag() // the traced Swift body really ran + let stop = Flag() // main -> W: H is finished, you may stop + + // Fresh instances: a cold compiler cache is what makes `compile_trace` + // — and therefore the re-entry into Swift under `evalLock` — happen. + let shared: @Sendable (MLXArray) -> MLXArray = compile(shapeless: true) { x in + x * 2 + 1 + } + let outer: @Sendable (MLXArray) -> MLXArray = compile(shapeless: true) { x in + // This body runs on thread H, on the C++ side of + // `mlx_closure_apply`, inside `compile_trace`, with `evalLock` held. + traced.set() + gate.signal() // W may now take L_cf(shared) + usleep(widenMicroseconds) // hold the window open + return shared(x) + 3 // H blocks here on L_cf(shared) + } + keepAlive.append(shared) + keepAlive.append(outer) + + let h = Thread { + let a = MLXArray([1.0, 2.0, 3.0] as [Float]) + let r = outer(a) + eval(r) + hDone.signal() + } + h.name = "H-nested-compiled-call" + + let w = Thread { + gate.wait() + let b = MLXArray([4.0, 5.0, 6.0] as [Float]) + var iterations = 0 + repeat { + let r = shared(b) // W takes L_cf(shared), then wants evalLock + eval(r) + iterations += 1 + } while !stop.isSet && iterations < 1000 + wDone.signal() + } + w.name = "W-plain-compiled-call" + + h.start() + w.start() + + if hDone.wait(timeout: .now() + deadline) == .timedOut { + die( + 1, + """ + FAIL: round \(round): thread H did not complete within \(Int(deadline))s. + H holds evalLock (taken in CompiledFunction.innerCall) and is blocked + acquiring the per-instance lock of the nested compiled function; + W holds that per-instance lock and is blocked acquiring evalLock. + enteredTrace=\(traced.isSet) + This is the lock-order inversion; the process is deadlocked. + """) + } + stop.set() + if wDone.wait(timeout: .now() + deadline) == .timedOut { + die( + 1, + """ + FAIL: round \(round): thread W did not complete within \(Int(deadline))s \ + after H finished. Blocked acquiring evalLock. + """) + } + + print(" round \(round): ok (enteredTrace=\(traced.isSet))") + return traced.isSet + } +} diff --git a/Source/MLX/Device.swift b/Source/MLX/Device.swift index f092ec158..da5fc697e 100644 --- a/Source/MLX/Device.swift +++ b/Source/MLX/Device.swift @@ -185,7 +185,7 @@ extension Device: CustomStringConvertible { public var description: String { var s = mlx_string_new() defer { mlx_string_free(s) } - _ = evalLock.withLock { + _ = withEvalLock { mlx_device_tostring(&s, ctx) } return String(cString: mlx_string_data(s), encoding: .utf8)! diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index aee73a568..fa84d2848 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -36,7 +36,7 @@ public func save(array: MLXArray, url: URL, stream: StreamOrDevice = .default) t switch url.pathExtension { case "npy": _ = try withError { - _ = evalLock.withLock { + _ = withEvalLock { mlx_save(path.cString(using: .utf8), array.ctx) } } @@ -74,7 +74,7 @@ public func save( switch url.pathExtension { case "safetensors": _ = try withError { - _ = evalLock.withLock { + _ = withEvalLock { mlx_save_safetensors(path.cString(using: .utf8), mlx_arrays, mlx_metadata) } } @@ -288,7 +288,7 @@ public func saveToData( defer { mlx_io_writer_free(writer) } _ = try withError { - _ = evalLock.withLock { + _ = withEvalLock { mlx_save_safetensors_writer(writer, mlx_arrays, mlx_metadata) } } diff --git a/Source/MLX/MLXArray.swift b/Source/MLX/MLXArray.swift index 9e6148e00..46b260e35 100644 --- a/Source/MLX/MLXArray.swift +++ b/Source/MLX/MLXArray.swift @@ -554,7 +554,7 @@ public final class MLXArray { /// MLX is lazy and arrays are not fully realized until they are evaluated. This method is typically /// not needed as all reads ensure the contents are evaluated. public func eval() { - _ = evalLock.withLock { + _ = withEvalLock { mlx_array_eval(ctx) } } @@ -617,7 +617,7 @@ extension MLXArray: CustomStringConvertible { public var description: String { var s = mlx_string_new() defer { mlx_string_free(s) } - _ = evalLock.withLock { + _ = withEvalLock { mlx_array_tostring(&s, ctx) } return String(cString: mlx_string_data(s), encoding: .utf8)! diff --git a/Source/MLX/Memory.swift b/Source/MLX/Memory.swift index 89baf9cc6..7036633c0 100644 --- a/Source/MLX/Memory.swift +++ b/Source/MLX/Memory.swift @@ -354,7 +354,7 @@ public enum Memory { /// Cause all cached buffers to be deallocated. public static func clearCache() { - _ = evalLock.withLock { + _ = withEvalLock { mlx_clear_cache() } } diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index 6f2d2aa40..c266d50c7 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -118,7 +118,7 @@ public final class Stream: @unchecked Sendable, Equatable { @available(*, deprecated, message: "use init(Device) -- index not supported") public init(index: Int32, _ device: Device) { - self.ctx = evalLock.withLock { + self.ctx = withEvalLock { mlx_stream_new_device(device.ctx) } } @@ -127,20 +127,20 @@ public final class Stream: @unchecked Sendable, Equatable { /// /// See also ``withNewDefaultStream(device:_:)-5bwc3`` public init(_ device: Device) { - self.ctx = evalLock.withLock { + self.ctx = withEvalLock { mlx_stream_new_device(device.ctx) } } deinit { - _ = evalLock.withLock { + _ = withEvalLock { mlx_stream_free(ctx) } } /// Synchronize with the given stream public func synchronize() { - _ = evalLock.withLock { + _ = withEvalLock { mlx_synchronize(ctx) } } @@ -162,7 +162,7 @@ extension Stream: CustomStringConvertible { public var description: String { var s = mlx_string_new() defer { mlx_string_free(s) } - _ = evalLock.withLock { + _ = withEvalLock { mlx_stream_tostring(&s, ctx) } return String(cString: mlx_string_data(s), encoding: .utf8)! diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index e52835380..0bae50dd3 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -9,7 +9,14 @@ final class CompiledFunction: @unchecked (Sendable) { /// unique (for the lifetime of the object) identifier for the compiled function private var id: UInt! - let lock = NSLock() + /// guards the observed state (``inputs`` / ``outputs``) -- see #226 + /// + /// This is always acquired *inside* `evalLock`; see ``call(_:)`` for the + /// ordering rule. It is recursive because a compiled function's body can + /// legitimately re-enter that same compiled function on the same thread + /// while it is being traced (a nested call with tracer inputs bypasses the + /// cache and simply runs the body again). + let lock = NSRecursiveLock() /// the function to compile let f: ([MLXArray]) -> [MLXArray] @@ -31,17 +38,139 @@ final class CompiledFunction: @unchecked (Sendable) { self.id = UInt(bitPattern: Unmanaged.passUnretained(self).toOpaque()) } + /// Remove the compiled structure from the back end. + /// + /// ### Why this takes `evalLock` + /// + /// `mlx_detail_compile_erase` is `compiler_cache().erase(fun_id)`, i.e. + /// `std::unordered_map::erase` on the *same* process-global + /// `mlx::core::detail::CompilerCache` that ``innerCall(_:)`` reads and + /// inserts into via `CompilerCache::find` (`cache_[fun_id]`, which itself + /// can insert and rehash, followed by `entries.push_back`). That cache has + /// no synchronization of its own, and `find` hands back a *reference* to an + /// entry that the caller then fills in and reads from across the whole + /// trace. An unlocked erase can therefore rehash the map out from under a + /// concurrent `find`, or destroy the very entry another thread is compiling + /// into. + /// + /// Every other use of that cache happens under `evalLock`, so the erase has + /// to as well; otherwise the lock does not actually make the cache + /// single-threaded, it only makes most of it single-threaded. + /// + /// ### Why blocking here is safe + /// + /// `deinit` runs on whichever thread drops the last reference: + /// + /// - A thread that holds no MLX lock simply waits for the current eval or + /// compiled call to finish. Bounded, since `evalLock` is only ever held + /// for the duration of one call. + /// - A thread already inside a compiled call or an `eval` -- for instance a + /// traced body that drops the last reference to some other compiled + /// function -- already holds `evalLock`, and it is an `NSRecursiveLock`, + /// so re-entry is free. + /// - A thread holding some other ``CompiledFunction``'s ``lock`` also holds + /// `evalLock` already (see ``call(_:)``), so that is the same case; no new + /// edge is added to the lock graph. + /// + /// Within this library the only releases of a ``CompiledFunction`` are the + /// caller's own (the closure returned by ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` + /// holds the sole strong reference) and the `mlx_closure_free` of the + /// closure built in ``innerCall(_:)``, which runs on the calling thread with + /// `evalLock` already held. Nothing hands a reference to a back-end worker + /// thread, so the erase can never be the thing an `evalLock` holder is + /// waiting for. + /// + /// The erase stays synchronous rather than being deferred to a queue on + /// purpose: ``id`` is the object's own address, so once this object is + /// freed a new ``CompiledFunction`` can be allocated at the same address and + /// receive the same id. A deferred erase could delete that new function's + /// cache entries. deinit { - // remove the compiled structure from the back end - mlx_detail_compile_erase(id) + let functionID = id! + withEvalLock { + #if DEBUG + // Trivially true two lines below the acquisition -- it is here to + // catch a future change that moves the erase back out from under + // the lock, the same way `withInstanceLock` guards `call(_:)`. + EvalLockOwnership.requireHeldByCurrentThread( + "the compiler cache was erased without holding evalLock; " + + "see CompiledFunction.deinit") + EvalLockOwnership.recordCompileErase() + #endif + _ = mlx_detail_compile_erase(functionID) + } } + /// Evaluate the compiled function. + /// + /// ### Lock ordering + /// + /// `evalLock` is acquired **outermost**, before the per-instance ``lock``. + /// That order is not a preference, it is forced: + /// + /// - Tracing mutates process-global state that has no synchronization of its + /// own -- `mlx::core::detail::CompilerCache` is an unguarded singleton and + /// `detail::InTracing::trace_stack()` is a plain function-local `static`, + /// not `thread_local` -- so `evalLock` has to span the whole trace (#339). + /// - The trace calls back into Swift, and the traced body routinely calls + /// *other* compiled functions: every `MLXNN` activation is a file-scope + /// `CompiledFunction`, so e.g. `silu` inside a traced body enters + /// `compiledSilu.call` and takes *its* per-instance lock. + /// + /// So `evalLock` → ``lock`` happens whether or not we ask for it. Taking + /// ``lock`` first here would add the reverse order ``lock`` → `evalLock` for + /// any thread entering from the top, and two threads -- one nested inside a + /// trace, one not -- would deadlock on a cold compiler cache: + /// + /// thread H: L_cf(outer) -> evalLock -> wants L_cf(shared) + /// thread W: L_cf(shared) -> wants evalLock + /// + /// Acquiring `evalLock` here gives the process a single global order, which + /// makes the inversion impossible. `Source/CompileLockRepro` reproduces the + /// deadlock deterministically without this. + /// + /// The cost is small: ``innerCall(_:)`` used to take `evalLock` itself and + /// hold it through `mlx_closure_apply` -- i.e. across the actual evaluation + /// -- so this only moves the acquisition earlier by the state gathering and + /// closure creation at the top of ``innerCall(_:)``. func call(_ arguments: [MLXArray]) -> [MLXArray] { - lock.withLock { - innerCall(arguments) + withEvalLock { + // ``lock`` exists only to guard the observed state (#226): the + // `_updateInternal` swaps in `innerCall` are the sole mutation of + // anything reachable from `self` -- every other stored property is + // `let`, and `id` is written once at init. With no state there is + // nothing to guard, and `evalLock` above already serializes the + // compiler cache and every other caller. + // + // This is the common case: every compiled function in `MLXNN` is + // stateless. + if inputs.isEmpty && outputs.isEmpty { + return innerCall(arguments) + } + return withInstanceLock { + innerCall(arguments) + } } } + /// Acquire ``lock``, checking the ordering invariant first. + /// + /// The check is where the invariant actually lives, not in ``call(_:)``: a + /// change that moved the per-instance lock back outside `evalLock` -- the + /// shape that produced the deadlock -- trips this on the first compiled + /// call, deterministically and without needing to lose a race. + private func withInstanceLock(_ body: () throws -> R) rethrows -> R { + #if DEBUG + EvalLockOwnership.requireHeldByCurrentThread( + "CompiledFunction.lock was acquired without holding evalLock; " + + "that is the lock-order inversion, see CompiledFunction.call") + #endif + return try lock.withLock(body) + } + + /// - Precondition: `evalLock` is held by the caller, and ``lock`` is held + /// too unless this function has no observed state; see ``call(_:)``, + /// which is the only caller. func innerCall(_ arguments: [MLXArray]) -> [MLXArray] { let stateInputs = inputs.flatMap { $0.innerState() } let argumentsCount = arguments.count @@ -86,12 +215,13 @@ final class CompiledFunction: @unchecked (Sendable) { // note: this will use the cached compile (via the id) // but will be able to re-evaluate with fresh state if needed - evalLock.lock() + // + // evalLock is already held by call(_:) and covers everything through + // mlx_closure_apply below, which is where the trace runs. var compiled = mlx_closure_new() let compileStatus = mlx_detail_compile(&compiled, innerClosure, id, shapeless, [], 0) defer { mlx_closure_free(compiled) - evalLock.unlock() } // mlx_error was already dispatched on failure: diff --git a/Source/MLX/Transforms+Eval.swift b/Source/MLX/Transforms+Eval.swift index 46fe9c593..64e76098a 100644 --- a/Source/MLX/Transforms+Eval.swift +++ b/Source/MLX/Transforms+Eval.swift @@ -6,15 +6,122 @@ import Foundation /// lock to be held while doing any eval or asyncEval. This is /// a recursive lock to handle any cases where a closure might /// call back into eval. +/// +/// Acquire it through ``withEvalLock(_:)`` rather than directly: debug builds +/// record ownership there so that the `evalLock` → `CompiledFunction.lock` +/// ordering can be checked. +/// +/// `evalLock` is the outermost lock in the library. Nothing may be acquired +/// before it and then be waited on from inside it; see `CompiledFunction.call`. let evalLock = NSRecursiveLock() +#if DEBUG + + /// Per-thread `evalLock` ownership, recorded so the ordering invariant in + /// `CompiledFunction.call` can be checked in debug builds. + /// + /// Explicit bookkeeping is required: `evalLock` is recursive, so a + /// successful `try()` cannot distinguish "this thread owns it" from "nobody + /// owns it". The depth is thread-local, so it needs no lock of its own -- + /// which matters, because this runs while holding `evalLock`. + enum EvalLockOwnership { + + private static let depthKey: pthread_key_t = { + var key = pthread_key_t() + pthread_key_create(&key, nil) + return key + }() + + static var isHeldByCurrentThread: Bool { + depth > 0 + } + + static var depth: Int { + Int(bitPattern: pthread_getspecific(depthKey)) + } + + static func entered() { + pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth + 1)) + } + + static func exited() { + pthread_setspecific(depthKey, UnsafeRawPointer(bitPattern: depth - 1)) + } + + // Counters, so a test can assert both that the invariant held and that + // it was actually exercised -- a check count of zero would mean the test + // proved nothing. + private static let counterLock = NSLock() + nonisolated(unsafe) private static var checkCount = 0 + nonisolated(unsafe) private static var violationCount = 0 + nonisolated(unsafe) private static var compileEraseCount = 0 + + /// Records, and in debug builds asserts, that `evalLock` is held by the + /// current thread at a point where an inner lock is about to be taken. + static func requireHeldByCurrentThread(_ message: @autoclosure () -> String) { + let held = isHeldByCurrentThread + counterLock.withLock { + checkCount += 1 + if !held { violationCount += 1 } + } + assert(held, message()) + } + + /// Records a `CompiledFunction.deinit` erase of the compiler cache, so a + /// test can tell "the erase ran under the lock" from "the deinit never + /// fired". + static func recordCompileErase() { + counterLock.withLock { compileEraseCount += 1 } + } + + /// (checks, violations) since ``resetCounters()``. + static var counters: (checks: Int, violations: Int) { + counterLock.withLock { (checkCount, violationCount) } + } + + /// Compiler-cache erases since ``resetCounters()``. + static var compileErases: Int { + counterLock.withLock { compileEraseCount } + } + + static func resetCounters() { + counterLock.withLock { + checkCount = 0 + violationCount = 0 + compileEraseCount = 0 + } + } + } + +#endif + +/// Acquire ``evalLock`` for the duration of `body`. +/// +/// This is the single choke point for ``evalLock``; debug builds record +/// ownership here so that `CompiledFunction` can check that its per-instance +/// lock is never taken outside it. +@inline(__always) +func withEvalLock(_ body: () throws -> R) rethrows -> R { + evalLock.lock() + #if DEBUG + EvalLockOwnership.entered() + #endif + defer { + #if DEBUG + EvalLockOwnership.exited() + #endif + evalLock.unlock() + } + return try body() +} + /// Evaluate one or more `MLXArray` /// /// ### See Also /// - public func eval(_ arrays: MLXArray...) { let vector_array = new_mlx_vector_array(arrays) - _ = evalLock.withLock { + _ = withEvalLock { mlx_eval(vector_array) } mlx_vector_array_free(vector_array) @@ -26,7 +133,7 @@ public func eval(_ arrays: MLXArray...) { /// - public func eval(_ arrays: some Collection) { let vector_array = new_mlx_vector_array(arrays) - _ = evalLock.withLock { + _ = withEvalLock { mlx_eval(vector_array) } mlx_vector_array_free(vector_array) @@ -39,7 +146,7 @@ public func eval(_ arrays: some Collection) { /// - ``asyncEval(_:)-(Collection)`` public func asyncEval(_ arrays: some Collection) { let vector_array = new_mlx_vector_array(arrays) - _ = evalLock.withLock { + _ = withEvalLock { mlx_async_eval(vector_array) } mlx_vector_array_free(vector_array) diff --git a/Source/MLX/Transforms+Internal.swift b/Source/MLX/Transforms+Internal.swift index 394aa4e98..c73ecf189 100644 --- a/Source/MLX/Transforms+Internal.swift +++ b/Source/MLX/Transforms+Internal.swift @@ -16,7 +16,7 @@ private func valueAndGradient( var r0 = mlx_vector_array_new() var r1 = mlx_vector_array_new() - _ = evalLock.withLock { + _ = withEvalLock { mlx_closure_value_and_grad_apply(&r0, &r1, valueAndGrad, input_vector) } diff --git a/Source/MLX/Transforms+Vmap.swift b/Source/MLX/Transforms+Vmap.swift index bd885c148..d1cf92099 100644 --- a/Source/MLX/Transforms+Vmap.swift +++ b/Source/MLX/Transforms+Vmap.swift @@ -29,7 +29,7 @@ public func vmap( var traceInputs = mlx_vector_array_new() var traceOutputs = mlx_vector_array_new() - evalLock.withLock { + withEvalLock { let closure = new_mlx_closure(f) _ = inAxes32.withUnsafeBufferPointer { inAxesBuf in mlx_detail_vmap_trace( diff --git a/Source/MLX/Transforms.swift b/Source/MLX/Transforms.swift index 3e73a8f8b..e9e8c7b6c 100644 --- a/Source/MLX/Transforms.swift +++ b/Source/MLX/Transforms.swift @@ -28,7 +28,7 @@ public func jvp( let closure = new_mlx_closure(f) - _ = evalLock.withLock { + _ = withEvalLock { mlx_jvp(&r0, &r1, closure, primals_mlx, tangents_mlx) } @@ -65,7 +65,7 @@ public func vjp( let closure = new_mlx_closure(f) - _ = evalLock.withLock { + _ = withEvalLock { mlx_vjp(&r0, &r1, closure, primals_mlx, cotangents_mlx) } diff --git a/Source/MLX/WiredMemory.swift b/Source/MLX/WiredMemory.swift index fff688e35..455ab0a11 100644 --- a/Source/MLX/WiredMemory.swift +++ b/Source/MLX/WiredMemory.swift @@ -24,7 +24,7 @@ private enum WiredMemoryBackend { guard isSupported else { return false } guard limit >= 0 else { return false } var previous: size_t = 0 - let result = evalLock.withLock { + let result = withEvalLock { mlx_set_wired_limit(&previous, size_t(limit)) } return result == 0 diff --git a/Tests/MLXTests/CompileLockOrderingTests.swift b/Tests/MLXTests/CompileLockOrderingTests.swift new file mode 100644 index 000000000..745e14b2d --- /dev/null +++ b/Tests/MLXTests/CompileLockOrderingTests.swift @@ -0,0 +1,352 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import XCTest + +@testable import MLX + +// The ownership bookkeeping these tests read is debug-only. +#if DEBUG + + /// The `evalLock` → `CompiledFunction.lock` ordering invariant. + /// + /// `Source/CompileLockRepro` shows that the reverse order deadlocks. It + /// lives in its own executable because a failure wedges the process, and a + /// passing run of it only shows that the race did not fire -- not that the + /// ordering is right. These tests assert the ordering itself, cannot hang, + /// and are safe to run in the shared test process. + /// + /// A violation trips an `assert` at the point of acquisition, so a + /// regression traps there rather than reaching the counter assertions + /// below. The counters are what make a *pass* meaningful: they show the + /// check was reached at all, rather than the code path never running. + class CompileLockOrderingTests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + /// The instrumentation must be able to report "not held", otherwise a + /// violation count of zero would mean nothing. + func testOwnershipTrackingIsNotVacuous() { + XCTAssertFalse(EvalLockOwnership.isHeldByCurrentThread) + XCTAssertEqual(EvalLockOwnership.depth, 0) + + withEvalLock { + XCTAssertTrue(EvalLockOwnership.isHeldByCurrentThread) + XCTAssertEqual(EvalLockOwnership.depth, 1) + + // evalLock is recursive; nesting is expected and must be counted + withEvalLock { + XCTAssertEqual(EvalLockOwnership.depth, 2) + } + XCTAssertEqual(EvalLockOwnership.depth, 1) + } + + XCTAssertFalse(EvalLockOwnership.isHeldByCurrentThread) + XCTAssertEqual(EvalLockOwnership.depth, 0) + } + + /// Ownership is per thread: another thread must not believe it holds the + /// lock this one is holding. (The other thread only reads its own + /// bookkeeping -- it never waits on `evalLock` -- so this cannot hang.) + func testOwnershipIsPerThread() { + final class Result: @unchecked Sendable { + var heldOnOtherThread = true + } + let result = Result() + let done = DispatchSemaphore(value: 0) + + withEvalLock { + let thread = Thread { + result.heldOnOtherThread = EvalLockOwnership.isHeldByCurrentThread + done.signal() + } + thread.start() + XCTAssertEqual(done.wait(timeout: .now() + 10), .success) + } + + XCTAssertFalse(result.heldOnOtherThread) + } + + /// A compiled function with observed state really takes the per-instance + /// lock -- a stateless one has nothing to guard and skips it -- so this + /// exercises the checked acquisition. + func testEvalLockIsHeldWhenTheInstanceLockIsTaken() { + EvalLockOwnership.resetCounters() + + let state = MLXArray(2) + func body(_ x: [MLXArray]) -> [MLXArray] { + [x[0] + state] + } + let compiled = compile(inputs: [state], outputs: [state], body(_:)) + + let r = compiled([MLXArray(5)]) + eval(r) + XCTAssertEqual(r[0].item(Float.self), 7) + + let counters = EvalLockOwnership.counters + XCTAssertGreaterThan( + counters.checks, 0, + "the ordering invariant was never checked, so this test proves nothing") + XCTAssertEqual( + counters.violations, 0, + "CompiledFunction.lock was taken without evalLock held") + } + + /// The deadlock shape, single threaded: one compiled function called + /// from inside another's trace. The nested acquisition happens while an + /// outer `evalLock` is already held, which is exactly the ordering the + /// invariant describes. + func testNestedCompiledCallKeepsTheOrdering() { + EvalLockOwnership.resetCounters() + + let innerState = MLXArray(1) + func innerBody(_ x: [MLXArray]) -> [MLXArray] { + [x[0] * 2 + innerState] + } + let inner = compile(inputs: [innerState], outputs: [innerState], innerBody(_:)) + + let outerState = MLXArray(3) + nonisolated(unsafe) var tracedNested = false + func outerBody(_ x: [MLXArray]) -> [MLXArray] { + tracedNested = true + return [inner([x[0]])[0] + outerState] + } + let outer = compile(inputs: [outerState], outputs: [outerState], outerBody(_:)) + + let r = outer([MLXArray(5)]) + eval(r) + XCTAssertEqual(r[0].item(Float.self), 14) + + XCTAssertTrue( + tracedNested, + "the traced body never ran, so no nested compiled call was exercised") + + let counters = EvalLockOwnership.counters + XCTAssertGreaterThan(counters.checks, 0) + XCTAssertEqual( + counters.violations, 0, + "a nested CompiledFunction.lock was taken without evalLock held") + } + + // MARK: - CompiledFunction.deinit vs. a concurrent compiled call + + /// A `CompiledFunction` whose only strong reference is `release()`-able + /// on demand, so a test can make its `deinit` -- and therefore the + /// compiler-cache erase -- happen at a chosen instant on a chosen + /// thread. + private final class Victim: @unchecked Sendable { + private var compiled: (@Sendable (MLXArray) -> MLXArray)? + + init() { + // A fresh compile() per instance, so the erase has a populated + // cache entry to remove. The unique constant keeps this from + // sharing an entry with anything else. + let bias = MLXArray(Float.random(in: 1 ... 2)) + compiled = MLX.compile { (x: MLXArray) in x * 3 + bias } + } + + /// Populate the cache entry that `deinit` will erase. + func warm() { + let r = compiled!(MLXArray(1)) + eval(r) + } + + /// Drop the last reference. Returns once `deinit` has completed. + func release() { + compiled = nil + } + } + + /// The interleaving this is about: thread R's `deinit` erase lands while + /// thread C is provably inside `compile_trace`, holding `evalLock` and + /// holding a live reference into the cache entry it is filling in. + /// + /// Unlike `Source/CompileLockRepro`, this is safe in the shared test + /// process: with the fix there is a single lock order, `evalLock` is + /// only ever held for the duration of one call, and R's `deinit` cannot + /// be what C is waiting for -- so R blocks for the length of the traced + /// body and no longer. The deadline turns a regression into a failure + /// rather than a hung lane. + func testDeinitEraseSerializesAgainstATraceInFlight() { + EvalLockOwnership.resetCounters() + + final class Shared: @unchecked Sendable { + let lock = NSLock() + private var _insideTrace = false + var insideTrace: Bool { + get { lock.withLock { _insideTrace } } + set { lock.withLock { _insideTrace = newValue } } + } + var tracedRan = false + var insideTraceAfterErase = true + var erasesAroundRelease = 0 + var result: Float = .nan + } + let shared = Shared() + + // Signalled from *inside* the traced body, so R only starts its + // release once C is known to be inside the trace holding evalLock. + let cIsInsideTrace = DispatchSemaphore(value: 0) + let cDone = DispatchSemaphore(value: 0) + let rDone = DispatchSemaphore(value: 0) + + // Built on this thread; released on R. + let victim = Victim() + victim.warm() + + // C's own compiled function is created and retained *here*, not on + // C, so that C never destroys a `CompiledFunction` of its own -- + // the erase counted around `release()` below can then only be the + // victim's. Compiling does not trace; the trace happens on the + // first call, which is on C. + let state = MLXArray(Float(4)) + func body(_ x: [MLXArray]) -> [MLXArray] { + shared.tracedRan = true + shared.insideTrace = true + cIsInsideTrace.signal() + // Hold the window open long enough that R is certainly parked + // in its erase rather than merely scheduled. + usleep(50_000) + shared.insideTrace = false + return [x[0] + state] + } + let cCompiled = compile(inputs: [state], outputs: [state], body(_:)) + + let c = Thread { + let r = cCompiled([MLXArray(Float(6))]) + eval(r) + shared.result = r[0].item(Float.self) + cDone.signal() + } + + let r = Thread { + // Only proceed once C is inside the trace. + guard cIsInsideTrace.wait(timeout: .now() + 30) == .success else { + rDone.signal() + return + } + let before = EvalLockOwnership.compileErases + victim.release() + let after = EvalLockOwnership.compileErases + // Read the flag the instant the erase returns. With the erase + // under `evalLock` this can only be observed false: R cannot + // have acquired the lock until C left the traced body. + shared.insideTraceAfterErase = shared.insideTrace + shared.erasesAroundRelease = after - before + rDone.signal() + } + + c.start() + r.start() + + XCTAssertEqual(cDone.wait(timeout: .now() + 60), .success, "the compiled call wedged") + XCTAssertEqual(rDone.wait(timeout: .now() + 60), .success, "the deinit erase wedged") + + XCTAssertTrue(shared.tracedRan, "the traced body never ran, so nothing was interleaved") + XCTAssertEqual( + shared.erasesAroundRelease, 1, + "dropping the last reference did not run CompiledFunction.deinit, " + + "so no erase was interleaved with the trace") + XCTAssertFalse( + shared.insideTraceAfterErase, + "the compiler-cache erase completed while another thread was still inside " + + "compile_trace -- the erase is not serialized against find/insert") + + XCTAssertEqual(shared.result, 10) + withExtendedLifetime(cCompiled) {} + + let counters = EvalLockOwnership.counters + XCTAssertGreaterThan(counters.checks, 0) + XCTAssertEqual( + counters.violations, 0, + "the compiler cache was mutated without evalLock held") + } + + /// The same hazard without choreography: a stream of `deinit` erases + /// against a stream of cold compiles, so erase runs against `find`'s + /// `cache_[fun_id]` insert-and-rehash as well as against tracing. + /// + /// This one is a stress loop, so a pass is not by itself proof. What it + /// adds is that it cannot pass *vacuously*: the counters assert the + /// erases happened and that every checked acquisition held `evalLock`. + func testConcurrentDeinitErasesAndCompiledCalls() { + EvalLockOwnership.resetCounters() + + let rounds = 200 + let start = DispatchSemaphore(value: 0) + let ready = DispatchSemaphore(value: 0) + let done = DispatchSemaphore(value: 0) + + final class Outcome: @unchecked Sendable { + var callerValues = [Float]() + var releases = 0 + } + let outcome = Outcome() + + let caller = Thread { + ready.signal() + start.wait() + var values = [Float]() + for i in 0 ..< rounds { + // A fresh compile each round keeps the cache cold, so every + // round really performs an insert. + let compiled = MLX.compile { (x: MLXArray) in x * 2 + MLXArray(Float(i)) } + let r = compiled(MLXArray(Float(1))) + eval(r) + values.append(r.item(Float.self)) + } + outcome.callerValues = values + done.signal() + } + + let releaser = Thread { + ready.signal() + start.wait() + var released = 0 + for _ in 0 ..< rounds { + let victim = Victim() + victim.warm() + victim.release() + released += 1 + } + outcome.releases = released + done.signal() + } + + caller.start() + releaser.start() + ready.wait() + ready.wait() + start.signal() + start.signal() + + XCTAssertEqual(done.wait(timeout: .now() + 300), .success, "a thread wedged") + XCTAssertEqual(done.wait(timeout: .now() + 300), .success, "a thread wedged") + + XCTAssertEqual(outcome.releases, rounds) + XCTAssertEqual(outcome.callerValues.count, rounds) + // Results must survive the erase traffic: a clobbered cache entry + // would show up as a wrong value rather than only as a crash. + for (i, value) in outcome.callerValues.enumerated() { + XCTAssertEqual(value, Float(2 + i), "round \(i) produced \(value)") + } + + XCTAssertGreaterThanOrEqual( + EvalLockOwnership.compileErases, rounds, + "the deinit erase path was never exercised") + + // Everything compiled here is stateless, so it skips the + // per-instance lock: these checks all come from the erase path. + let counters = EvalLockOwnership.counters + XCTAssertGreaterThan( + counters.checks, 0, + "no erase was checked for evalLock ownership, so this proves nothing") + XCTAssertEqual( + counters.violations, 0, + "the compiler cache was mutated without evalLock held") + } + } + +#endif From d737cd0780dfae5120f6f7ba524a5abc43e4ab64 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Thu, 27 Aug 2026 18:36:15 +0400 Subject: [PATCH 23/39] Fix StreamOrDevice.stream(_:) ignoring its argument (#463) StreamOrDevice.stream(_:) discarded the stream it was given and returned the default stream of the default device instead, so operations that passed an explicit stream silently ran somewhere else. Return the given stream and add a regression test. --- Source/MLX/Stream.swift | 2 +- Tests/MLXTests/StreamTests.swift | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index c266d50c7..8c0e929d2 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -52,7 +52,7 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable { public static let gpu = device(.gpu) public static func stream(_ stream: Stream) -> StreamOrDevice { - StreamOrDevice(Device.defaultStream()) + StreamOrDevice(stream) } /// Internal context -- used with Cmlx calls. diff --git a/Tests/MLXTests/StreamTests.swift b/Tests/MLXTests/StreamTests.swift index ec3a35de7..32040e9df 100644 --- a/Tests/MLXTests/StreamTests.swift +++ b/Tests/MLXTests/StreamTests.swift @@ -74,6 +74,16 @@ class StreamTests: XCTestCase { print("here") } + func testStreamOrDeviceStreamUsesGivenStream() { + // StreamOrDevice.stream(_:) used to ignore its argument and return + // the default stream instead + let stream = Stream(.cpu) + XCTAssertEqual(StreamOrDevice.stream(stream).stream, stream) + + XCTAssertEqual(StreamOrDevice.stream(.cpu).stream, Stream.cpu) + XCTAssertEqual(StreamOrDevice.stream(.gpu).stream, Stream.gpu) + } + func disabledTestCreateStream() { // see https://github.com/ml-explore/mlx/issues/2118 for _ in 1 ..< 10000 { From ac4ae9a39af573d4135109d1640817147c79a5e7 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Mon, 31 Aug 2026 14:43:29 -0700 Subject: [PATCH 24/39] test Module parameters and compile (#465) * test Module parameters and compile --- .../Articles/compilation.md | 62 +++++++++++++++++++ Source/MLXNN/Module.swift | 7 +++ Tests/MLXTests/TransformTests.swift | 59 ++++++++++++++++++ skills/mlx-swift/SKILL.md | 8 ++- skills/mlx-swift/references/optimizers.md | 3 +- skills/mlx-swift/references/transforms.md | 38 +++++++++--- 6 files changed, 167 insertions(+), 10 deletions(-) diff --git a/Source/MLX/Documentation.docc/Articles/compilation.md b/Source/MLX/Documentation.docc/Articles/compilation.md index 87ae57faf..736851f29 100644 --- a/Source/MLX/Documentation.docc/Articles/compilation.md +++ b/Source/MLX/Documentation.docc/Articles/compilation.md @@ -240,6 +240,68 @@ let c2b = c2(bias) XCTAssertFalse(allClose(c2a, c2b).item()) ``` +The same rule applies to a `Module` that is captured (closed over) rather than +passed as an explicit function argument. `Module` conforms to ``Updatable``, so +its parameters count as implicit state, just like `MLXRandom.globalState` above. +If you compile a function that reads a module's parameters without listing that +module in `inputs:`/`outputs:`, the values are baked into the compiled graph the +first time it traces and later updates to the module (for example swapping in +new LoRA weights with `Module.update(parameters:)`) will be silently ignored: + +```swift +class ScaleModule: Module, UnaryLayer { + var scale: MLXArray + + init(_ value: Float) { + self.scale = MLXArray(value) + } + + func callAsFunction(_ x: MLXArray) -> MLXArray { + x * scale + } +} + +let model = ScaleModule(2.0) + +// no inputs/outputs -- the module's parameters are captured as constants +let compiled = compile { (x: MLXArray) -> MLXArray in + model(x) +} + +print(compiled(MLXArray(Float(3)))) // 6 + +// simulate a LoRA-style in place update of the module's weights +model.update(parameters: .unflattened([("scale", MLXArray(Float(10)))])) +eval(model) + +// still 6 -- the compiled function never sees the update! +print(compiled(MLXArray(Float(3)))) +``` + +To make the compiled function observe the update, pass the module via `inputs:`: + +```swift +let compiled = compile(inputs: [model]) { (x: MLXArray) -> MLXArray in + model(x) +} + +print(compiled(MLXArray(Float(3)))) // 6 + +model.update(parameters: .unflattened([("scale", MLXArray(Float(10)))])) +eval(model) + +print(compiled(MLXArray(Float(3)))) // 30 -- now it sees the update +``` + +`inputs:` is what makes the compiled function observe external updates: on every +call the *current* values of the listed state are fed back into the compiled +graph. `outputs:` is the reverse direction -- it writes values back out of the +graph into the real arrays, and is only needed when the compiled closure itself +*mutates* the captured state (for example a training step that calls +`optimizer.update(model:gradients:)` inside the compiled function, as shown +below). Since `model(x)` above only reads `scale`, `inputs: [model]` alone is +enough; adding `outputs: [model]` would be harmless but unnecessary. + ## Compiling Training Graphs This section will step through how to use ``compile(inputs:outputs:shapeless:_:)-([Updatable],[Updatable],Bool,([MLXArray])->[MLXArray])`` diff --git a/Source/MLXNN/Module.swift b/Source/MLXNN/Module.swift index 7cba9b12b..39407c263 100644 --- a/Source/MLXNN/Module.swift +++ b/Source/MLXNN/Module.swift @@ -426,6 +426,13 @@ open class Module { /// The ``apply(filter:map:)`` can be used for similar purposes to apply changes /// in-place. /// + /// > If this module is used inside a compiled closure (`compile()`) that does not + /// list the module in `inputs:`, calling `update(parameters:)` after the first compiled call + /// will not be observed by the compiled function -- `inputs:` is what makes `compile` re-read + /// current parameter values on every call. `outputs:` is only needed in addition when the + /// compiled closure itself mutates this module's parameters and the mutation should be + /// written back. + /// /// If a parameter is missing from the update and validation indicates `.allModelKeysSet` this /// will call ``updateMissing(parameter:verify:path:modulePath:)`` which will /// throw an error. Subclasses can override this if needed. diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift index d06999c7d..1ae105c22 100644 --- a/Tests/MLXTests/TransformTests.swift +++ b/Tests/MLXTests/TransformTests.swift @@ -217,6 +217,65 @@ class TransformTests: XCTestCase { XCTAssertEqual(state.o!.item(Float.self), -8) } + // https://github.com/ml-explore/mlx-swift-lm/issues/586 -- + // compile can capture Module parameters + class ScaleModule: Module, UnaryLayer { + var scale: MLXArray + + init(_ value: Float) { + self.scale = MLXArray(value) + } + + func callAsFunction(_ x: MLXArray) -> MLXArray { + x * scale + } + } + + func testCompileCapturesModuleParameterWithoutState() { + // verify the capture path happens + + let model = ScaleModule(2.0) + + let compiled = compile { (x: MLXArray) -> MLXArray in + model(x) + } + + let r1 = compiled(MLXArray(Float(3))) + XCTAssertEqual(r1.item(Float.self), 6) + + // simulate a LoRA-style in place update of the module's weights + model.update(parameters: .unflattened([("scale", MLXArray(Float(10)))])) + eval(model) + + // Verify capture of Module parameters -- it doesn't see the update + let r2 = compiled(MLXArray(Float(3))) + XCTAssertEqual( + r2.item(Float.self), 6, + "compiled function incorrectly ignored the updated module parameter") + } + + func testCompileWithStateObservesModuleParameterUpdates() { + // verify inputs correctly handles updated + // Module parameters + + let model = ScaleModule(2.0) + + let compiled = compile(inputs: [model]) { (x: MLXArray) -> MLXArray in + model(x) + } + + let r1 = compiled(MLXArray(Float(3))) + XCTAssertEqual(r1.item(Float.self), 6) + + // simulate a LoRA-style in place update of the module's weights + model.update(parameters: .unflattened([("scale", MLXArray(Float(10)))])) + eval(model) + + // with state:, the compiled function observes the updated parameter + let r2 = compiled(MLXArray(Float(3))) + XCTAssertEqual(r2.item(Float.self), 30) + } + func testCompiledRandom() { func f(_ bias: MLXArray) -> MLXArray { MLXRandom.uniform(0 ..< 1, [4]) + bias diff --git a/skills/mlx-swift/SKILL.md b/skills/mlx-swift/SKILL.md index 70681b257..f3bc120f5 100644 --- a/skills/mlx-swift/SKILL.md +++ b/skills/mlx-swift/SKILL.md @@ -344,8 +344,12 @@ let compiledOp = compile { (a: MLXArray, b: MLXArray) -> MLXArray in // Use compiled version let output = compiledOp(arrayA, arrayB) -// Note: compile() works best with pure MLXArray functions. -// For models, call model methods directly (they can use internal compilation). +// compile() also supports compiling a Module/Optimizer training step: pass +// them via inputs:/outputs: (they conform to Updatable) so compile observes +// in-place updates (e.g. optimizer.update(model:gradients:), LoRA weight swaps). +// If a captured Module/Optimizer is left out of inputs:/outputs:, compile() +// silently freezes its parameter values at trace time -- see +// references/transforms.md for the full pitfall and example. ``` ## Quaternary Workflow: Wired Memory Coordination diff --git a/skills/mlx-swift/references/optimizers.md b/skills/mlx-swift/references/optimizers.md index 73858f457..6958f6eca 100644 --- a/skills/mlx-swift/references/optimizers.md +++ b/skills/mlx-swift/references/optimizers.md @@ -304,7 +304,8 @@ import MLXOptimizers let model = MLP(inputDim: 784, hiddenDim: 256, outputDim: 10) let optimizer = AdamW(learningRate: 0.001, weightDecay: 0.01) -// Training loop (compile only works with pure MLXArray functions, not models) +// Training loop. compile() can compile the full step, but only if the +// model/optimizer are declared via inputs:/outputs: -- see references/transforms.md. func loss(model: MLP, x: MLXArray, y: MLXArray) -> MLXArray { let logits = model(x) return crossEntropy(logits: logits, targets: y, reduction: .mean) diff --git a/skills/mlx-swift/references/transforms.md b/skills/mlx-swift/references/transforms.md index 67e8c2dbe..25696fa84 100644 --- a/skills/mlx-swift/references/transforms.md +++ b/skills/mlx-swift/references/transforms.md @@ -203,12 +203,13 @@ let compiledF = compile(f) let result = compiledF(a, b) ``` -### compile Limitations +### compile Limitations and Module/Optimizer State -**Important:** `compile()` only works with pure `MLXArray` functions. It does NOT support: -- Model objects as parameters -- Optimizer objects as parameters -- Any non-MLXArray inputs +`compile()` traces a function that takes/returns `MLXArray` values (or tuples/arrays +of them) -- it can't take a `Module` or `Optimizer` as a parameter directly. But it +CAN compile closures that capture and mutate a `Module`/`Optimizer`, as long as any +mutable state they touch is declared via `inputs:`/`outputs:`. Both `Module` and +`Optimizer` conform to `Updatable`, so they can be passed there directly: ```swift // CORRECT: compile with pure MLXArray functions @@ -217,12 +218,35 @@ let compiledOp = compile { (a: MLXArray, b: MLXArray) -> MLXArray in return sum(x * x) } -// For model training, call model methods directly without compile: +// ALSO CORRECT: compile a full training step by capturing model/optimizer +// and declaring them as inputs/outputs so compile observes their updates func loss(model: MyModel, x: MLXArray, y: MLXArray) -> MLXArray { mseLoss(predictions: model(x), targets: y, reduction: .mean) } let lossAndGradFn = valueAndGrad(model: model, loss) -let (lossValue, grads) = lossAndGradFn(model, x, y) + +let step = compile(inputs: [model, optimizer], outputs: [model, optimizer]) { x, y in + let (loss, grads) = lossAndGradFn(model, x, y) + optimizer.update(model: model, gradients: grads) + return loss +} + +// PITFALL: if you close over a Module but omit it from inputs/outputs, +// compile() bakes in its parameter values at trace time. Later in-place +// updates to the module (e.g. swapping LoRA weights) will be silently +// ignored by the compiled function until it is recompiled for another +// reason (e.g. a shape change). +let compiled = compile { (x: MLXArray) -> MLXArray in model(x) } +_ = compiled(x) +model.update(parameters: newParameters) +eval(model) +_ = compiled(x) // still uses the OLD parameters + +// FIX: `inputs:` is what makes compile() re-read current state on every call; +// `outputs:` is only needed when the compiled closure itself mutates the +// state (as `optimizer.update(...)` does in the `step` example above) and +// you need that mutation written back to the real arrays. +let fixed = compile(inputs: [model]) { (x: MLXArray) -> MLXArray in model(x) } optimizer.update(model: model, gradients: grads) eval(model, optimizer) ``` From ab924c82ead3b970caaa1c0ac11171de23f0305a Mon Sep 17 00:00:00 2001 From: Charlie Le Date: Tue, 1 Sep 2026 08:00:52 -0700 Subject: [PATCH 25/39] update for mlx v0.32.2 (#450) * update for mlx v0.32.0 mlx v0.32.2 - carry `Source/Cmlx/mlx.patch` for mlx's `Device::operator<` / `Stream::operator<`, which are not strict weak orderings and now run on every `eval()` via `std::set`. libc++ aborts on them under the Xcode 27 SDK. Still broken upstream in v0.32.0 and on main, no issue filed. Applied by `tools/patch-mlx.sh` from `update-mlx.sh` and CI, since `git submodule update` restores a pristine checkout. - streams became thread affine in v0.31.2, which broke the multi-threaded tests. Adopt mlx's `new_thread_unsafe_stream` (added upstream for language bindings that lock themselves -- mlx-swift serializes on `evalLock`) via a small C shim, as mlx-c does not bind it. - FFT ops gained a `norm` parameter: add `FFTNorm` and thread it through, plus `fftfreq`/`rfftfreq`. - wrap the new reducing slice-updates in `ArrayAtIndices` and add `exportToDot`/`graphDescription`/`NodeNamer`. - fix `SEEK_END` in the in-memory IO stream -- it computed from the read position rather than the end of the buffer, which v0.32.0's `load_safetensors` size check exposed. - route `cpu/jit_compiler.cpp` through `mlx-conditional`: it shells out via `std::system()`, unavailable on iOS. - `gemv` is JIT compiled as of v0.32.0, so drop `gemv.metal` from the metallib in `fix-metal-includes.sh`. The jit source lists themselves no longer need touching -- #445 derives them from the vendored tree, and the derived sets match what a hand-bump would have produced here (48 Metal targets, 22 CUDA sources). * disable cuda runs for now -- need clear_streams exposed in mlx-c --------- Co-authored-by: David Koski --- .github/scripts/build-linux-cuda-cmake.sh | 5 +- CMakeLists.txt | 2 +- Package.swift | 25 +- Source/Cmlx/include-framework/Cmlx.h | 4 + Source/Cmlx/include-framework/Metal.hpp | 2 +- Source/Cmlx/include-framework/mlx-allocator.h | 2 + Source/Cmlx/include-framework/mlx-array.h | 11 + .../mlx-backend-common-metal_kernel.h | 29 + .../mlx-backend-common-utils.h | 33 + .../mlx-backend-cpu-encoder.h | 14 +- .../include-framework/mlx-backend-gpu-eval.h | 5 +- .../mlx-backend-metal-device.h | 225 +- .../mlx-backend-metal-metal.h | 5 + .../mlx-backend-metal-resident.h | 125 + Source/Cmlx/include-framework/mlx-c-compile.h | 12 +- .../mlx-c-distributed_group.h | 32 +- Source/Cmlx/include-framework/mlx-c-fast.h | 7 + Source/Cmlx/include-framework/mlx-c-fft.h | 20 + .../include-framework/mlx-c-graph_utils.h | 61 + Source/Cmlx/include-framework/mlx-c-io.h | 5 + .../Cmlx/include-framework/mlx-c-io_types.h | 46 + Source/Cmlx/include-framework/mlx-c-linalg.h | 6 + Source/Cmlx/include-framework/mlx-c-metal.h | 2 + Source/Cmlx/include-framework/mlx-c-mlx.h | 1 + Source/Cmlx/include-framework/mlx-c-ops.h | 191 +- Source/Cmlx/include-framework/mlx-c-stream.h | 12 + Source/Cmlx/include-framework/mlx-device.h | 14 +- .../mlx-distributed-distributed.h | 34 +- Source/Cmlx/include-framework/mlx-error.h | 51 + Source/Cmlx/include-framework/mlx-event.h | 32 + Source/Cmlx/include-framework/mlx-export.h | 18 +- .../Cmlx/include-framework/mlx-export_impl.h | 19 +- Source/Cmlx/include-framework/mlx-fast.h | 9 +- Source/Cmlx/include-framework/mlx-fft.h | 142 +- .../Cmlx/include-framework/mlx-graph_utils.h | 69 + Source/Cmlx/include-framework/mlx-io-load.h | 8 +- Source/Cmlx/include-framework/mlx-linalg.h | 4 + Source/Cmlx/include-framework/mlx-ops.h | 228 +- .../Cmlx/include-framework/mlx-primitives.h | 90 +- Source/Cmlx/include-framework/mlx-random.h | 16 +- Source/Cmlx/include-framework/mlx-scheduler.h | 141 +- Source/Cmlx/include-framework/mlx-stream.h | 39 +- .../include-framework/mlx-types-half_types.h | 25 +- Source/Cmlx/include-framework/mlx-utils.h | 57 +- Source/Cmlx/include-framework/mlx-version.h | 4 +- Source/Cmlx/include/mlx/c/compile.h | 12 +- Source/Cmlx/include/mlx/c/distributed_group.h | 32 +- Source/Cmlx/include/mlx/c/fast.h | 7 + Source/Cmlx/include/mlx/c/fft.h | 20 + Source/Cmlx/include/mlx/c/graph_utils.h | 61 + Source/Cmlx/include/mlx/c/io.h | 5 + Source/Cmlx/include/mlx/c/io_types.h | 46 + Source/Cmlx/include/mlx/c/linalg.h | 6 + Source/Cmlx/include/mlx/c/metal.h | 2 + Source/Cmlx/include/mlx/c/mlx.h | 1 + Source/Cmlx/include/mlx/c/ops.h | 191 +- Source/Cmlx/include/mlx/c/stream.h | 12 + Source/Cmlx/mlx | 2 +- Source/Cmlx/mlx-c | 2 +- .../jit_compiler_conditional.cpp | 12 + Source/Cmlx/mlx-generated/binary_ops.cpp | 101 +- .../Cmlx/mlx-generated/compiled_preamble.cpp | 55 +- Source/Cmlx/mlx-generated/conv.cpp | 170 +- Source/Cmlx/mlx-generated/copy.cpp | 46 +- .../mlx-generated/cuda/cuda_jit_sources.h | 6925 +++++++++++++---- Source/Cmlx/mlx-generated/fft.cpp | 1190 ++- Source/Cmlx/mlx-generated/fp_quantized.cpp | 603 +- .../Cmlx/mlx-generated/fp_quantized_nax.cpp | 236 +- Source/Cmlx/mlx-generated/gemm.cpp | 120 +- Source/Cmlx/mlx-generated/gemm_nax.cpp | 1406 ++-- Source/Cmlx/mlx-generated/gemv.cpp | 2352 ++++++ Source/Cmlx/mlx-generated/gemv_masked.cpp | 4 +- .../Cmlx/mlx-generated/metal/arg_reduce.metal | 8 +- Source/Cmlx/mlx-generated/metal/atomic.h | 25 +- Source/Cmlx/mlx-generated/metal/binary_ops.h | 101 +- Source/Cmlx/mlx-generated/metal/complex.h | 268 +- Source/Cmlx/mlx-generated/metal/conv.metal | 20 +- Source/Cmlx/mlx-generated/metal/copy.h | 46 +- Source/Cmlx/mlx-generated/metal/dot.h | 68 + Source/Cmlx/mlx-generated/metal/dot.metal | 18 + Source/Cmlx/mlx-generated/metal/fft.h | 225 +- Source/Cmlx/mlx-generated/metal/fft/radix.h | 307 +- .../Cmlx/mlx-generated/metal/fft/readwrite.h | 333 +- Source/Cmlx/mlx-generated/metal/fp4.h | 11 +- Source/Cmlx/mlx-generated/metal/fp8.h | 25 +- .../Cmlx/mlx-generated/metal/fp_quantized.h | 567 +- .../mlx-generated/metal/fp_quantized_nax.h | 200 +- .../metal/{gemv.metal => gemv.h} | 392 +- Source/Cmlx/mlx-generated/metal/gemv_masked.h | 4 +- .../mlx-generated/metal/indexing/scatter.h | 80 + Source/Cmlx/mlx-generated/metal/logging.h | 2 +- Source/Cmlx/mlx-generated/metal/quantized.h | 406 +- .../Cmlx/mlx-generated/metal/quantized_nax.h | 289 +- Source/Cmlx/mlx-generated/metal/random.metal | 4 +- .../Cmlx/mlx-generated/metal/reduction/ops.h | 84 +- .../metal/reduction/reduce_all.h | 4 +- .../metal/reduction/reduce_col.h | 14 +- .../metal/reduction/reduce_row.h | 9 +- .../Cmlx/mlx-generated/metal/rms_norm.metal | 176 +- .../metal/scaled_dot_product_attention.metal | 15 + Source/Cmlx/mlx-generated/metal/scan.h | 70 +- Source/Cmlx/mlx-generated/metal/sdpa_vector.h | 143 + .../Cmlx/mlx-generated/metal/searchsorted.h | 34 + .../mlx-generated/metal/searchsorted.metal | 28 + Source/Cmlx/mlx-generated/metal/sort.h | 32 +- .../steel/attn/kernels/steel_attention.h | 11 +- .../steel/attn/kernels/steel_attention.metal | 4 + .../steel/attn/kernels/steel_attention_nax.h | 684 +- .../mlx-generated/metal/steel/attn/loader.h | 28 +- .../Cmlx/mlx-generated/metal/steel/attn/mma.h | 58 +- .../Cmlx/mlx-generated/metal/steel/attn/nax.h | 960 +-- .../metal/steel/attn/transforms.h | 8 +- .../metal/steel/conv/kernels/steel_conv.h | 2 +- .../metal/steel/conv/kernels/steel_conv_3d.h | 2 +- .../steel/conv/kernels/steel_conv_general.h | 4 +- .../steel/conv/loaders/loader_channel_l.h | 48 +- .../steel/conv/loaders/loader_channel_n.h | 16 +- .../metal/steel/conv/loaders/loader_general.h | 20 +- .../mlx-generated/metal/steel/gemm/gemm_nax.h | 94 +- .../steel/gemm/kernels/steel_gemm_fused_nax.h | 77 +- .../gemm/kernels/steel_gemm_gather_nax.h | 15 +- .../steel/gemm/kernels/steel_gemm_masked.h | 4 +- .../gemm/kernels/steel_gemm_segmented_nax.h | 125 + .../gemm/kernels/steel_gemm_splitk_nax.h | 13 +- .../mlx-generated/metal/steel/gemm/loader.h | 14 +- .../Cmlx/mlx-generated/metal/steel/gemm/mma.h | 89 +- .../Cmlx/mlx-generated/metal/steel/gemm/nax.h | 944 +-- .../metal/steel/gemm/transforms.h | 8 +- .../metal/steel/utils/integral_constant.h | 11 +- Source/Cmlx/mlx-generated/metal/ternary_ops.h | 2 +- Source/Cmlx/mlx-generated/metal/unary_ops.h | 164 +- Source/Cmlx/mlx-generated/metal/utils.h | 67 +- Source/Cmlx/mlx-generated/quantized.cpp | 406 +- Source/Cmlx/mlx-generated/quantized_nax.cpp | 287 +- Source/Cmlx/mlx-generated/reduce.cpp | 27 +- Source/Cmlx/mlx-generated/reduce_utils.cpp | 109 +- Source/Cmlx/mlx-generated/scan.cpp | 171 +- Source/Cmlx/mlx-generated/scatter.cpp | 80 + Source/Cmlx/mlx-generated/searchsorted.cpp | 53 + Source/Cmlx/mlx-generated/sort.cpp | 32 +- Source/Cmlx/mlx-generated/steel_attention.cpp | 114 +- .../mlx-generated/steel_attention_nax.cpp | 1793 +++-- Source/Cmlx/mlx-generated/steel_conv.cpp | 2 +- Source/Cmlx/mlx-generated/steel_conv_3d.cpp | 2 +- .../Cmlx/mlx-generated/steel_conv_general.cpp | 24 +- .../mlx-generated/steel_gemm_fused_nax.cpp | 77 +- .../mlx-generated/steel_gemm_gather_nax.cpp | 15 +- .../Cmlx/mlx-generated/steel_gemm_masked.cpp | 4 +- .../steel_gemm_segmented_nax.cpp | 144 + .../mlx-generated/steel_gemm_splitk_nax.cpp | 13 +- Source/Cmlx/mlx-generated/ternary_ops.cpp | 2 +- Source/Cmlx/mlx-generated/unary_ops.cpp | 189 +- Source/Cmlx/mlx-generated/utils.cpp | 336 +- Source/MLX/ArrayAt.swift | 136 +- Source/MLX/Cmlx+Util.swift | 4 + .../Articles/converting-python.md | 4 +- .../Organization/arithmetic.md | 1 + .../Organization/cumulative.md | 8 +- .../Organization/initialization.md | 8 +- .../Organization/logical.md | 1 + .../MLX/Documentation.docc/free-functions.md | 15 +- Source/MLX/FFT.swift | 306 +- Source/MLX/Factory.swift | 80 +- Source/MLX/GPU+Metal.swift | 21 + Source/MLX/GraphUtils.swift | 122 + Source/MLX/IO.swift | 4 +- Source/MLX/Linalg.swift | 75 + Source/MLX/MLXArray+Indexing.swift | 71 +- Source/MLX/MLXArray+Ops.swift | 65 +- Source/MLX/MLXFast.swift | 29 +- Source/MLX/Ops+Array.swift | 41 +- Source/MLX/Ops.swift | 162 +- Source/MLX/Stream.swift | 33 +- Source/MLX/Transforms+Compile.swift | 6 +- Tests/MLXTests/ArrayAtTests.swift | 61 + Tests/MLXTests/FFTTests.swift | 73 + Tests/MLXTests/GraphUtilsTests.swift | 69 + Tests/MLXTests/SaveTests.swift | 16 + tools/fix-metal-includes.sh | 3 +- tools/update-mlx-xcodeproj.sh | 4 + xcode/MLX.xcodeproj/project.pbxproj | 153 +- xcode/xcconfig/Cmlx.xcconfig | 2 +- 182 files changed, 20616 insertions(+), 7818 deletions(-) create mode 100644 Source/Cmlx/include-framework/mlx-backend-common-metal_kernel.h create mode 100644 Source/Cmlx/include-framework/mlx-backend-metal-resident.h create mode 100644 Source/Cmlx/include-framework/mlx-c-graph_utils.h create mode 100644 Source/Cmlx/include-framework/mlx-error.h create mode 100644 Source/Cmlx/include-framework/mlx-graph_utils.h create mode 100644 Source/Cmlx/include/mlx/c/graph_utils.h create mode 100644 Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp create mode 100644 Source/Cmlx/mlx-generated/gemv.cpp create mode 100644 Source/Cmlx/mlx-generated/metal/dot.h create mode 100644 Source/Cmlx/mlx-generated/metal/dot.metal rename Source/Cmlx/mlx-generated/metal/{gemv.metal => gemv.h} (71%) create mode 100644 Source/Cmlx/mlx-generated/metal/searchsorted.h create mode 100644 Source/Cmlx/mlx-generated/metal/searchsorted.metal create mode 100644 Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h create mode 100644 Source/Cmlx/mlx-generated/searchsorted.cpp create mode 100644 Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp create mode 100644 Source/MLX/GraphUtils.swift create mode 100644 Tests/MLXTests/FFTTests.swift create mode 100644 Tests/MLXTests/GraphUtilsTests.swift diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index c84e02f27..2a571ea6d 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.github/scripts/build-linux-cuda-cmake.sh @@ -17,6 +17,7 @@ mkdir -p build pushd build cmake -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=ON -DMLX_C_BUILD_EXAMPLES=OFF .. -G Ninja ninja -./example1 --device gpu -./tutorial --device gpu +# TODO dkoski -- disabled for now until clear_streams is available +# ./example1 --device gpu +# ./tutorial --device gpu popd diff --git a/CMakeLists.txt b/CMakeLists.txt index ce75855f6..41404a01d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -30,7 +30,7 @@ endif() FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" - GIT_TAG "v0.6.0") + GIT_TAG "c74db5307cc8ce122f48d97ef951b30578674e7f") FetchContent_MakeAvailable(mlx-c) # swift-numerics diff --git a/Package.swift b/Package.swift index 04180bf6a..4a7ee0076 100644 --- a/Package.swift +++ b/Package.swift @@ -85,8 +85,10 @@ let noCudaCmlxExcludes = [ "mlx/mlx/backend/cuda/delayload.cpp", "mlx/mlx/backend/cuda/device_info.cpp", "mlx/mlx/backend/cuda/device.cpp", + "mlx/mlx/backend/cuda/dirs.cpp", "mlx/mlx/backend/cuda/eval.cpp", "mlx/mlx/backend/cuda/fence.cpp", + "mlx/mlx/backend/cuda/fft.cu", "mlx/mlx/backend/cuda/indexing.cpp", "mlx/mlx/backend/cuda/jit_module.cpp", "mlx/mlx/backend/cuda/load.cpp", @@ -132,16 +134,9 @@ let noCudaCmlxExcludes = [ "mlx-c/mlx/c/metal.cpp", "mlx/mlx/backend/cuda/delayload.cpp", // For Windows - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmv.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh", - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmm.h", - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmm.cu", - "mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu", - "mlx/mlx/backend/cuda/quantized/qmm/fp_qmv.cu", + + // built by the CudaBuild plugin (nvcc), not by SwiftPM + "mlx/mlx/backend/cuda/quantized/qmm", ] + noMetalCmlxExcludes cxxSettings = [ @@ -210,7 +205,10 @@ let noCudaCmlxExcludes = [ let platformExcludes: [String] = [ + // selected conditionally by mlx-conditional -- these are #included by + // the wrappers there so they must not also be built directly "mlx/mlx/backend/cpu/compiled.cpp", + "mlx/mlx/backend/cpu/jit_compiler.cpp", // opt-out of these backends (using metal) "mlx/mlx/backend/no_gpu", @@ -299,11 +297,8 @@ let cmlx = Target.target( "mlx/mlx/distributed/mpi/mpi.cpp", "mlx/mlx/distributed/ring/ring.cpp", "mlx/mlx/distributed/nccl/nccl.cpp", - "mlx/mlx/distributed/nccl/nccl_stub", "mlx/mlx/distributed/jaccl/jaccl.cpp", - "mlx/mlx/distributed/jaccl/mesh.cpp", - "mlx/mlx/distributed/jaccl/ring.cpp", - "mlx/mlx/distributed/jaccl/utils.cpp", + "mlx/mlx/distributed/jaccl/lib", ], cSettings: [ .headerSearchPath("mlx"), @@ -315,7 +310,7 @@ let cmlx = Target.target( .headerSearchPath("mlx-c"), .headerSearchPath("json/single_include/nlohmann"), .headerSearchPath("fmt/include"), - .define("MLX_VERSION", to: "\"0.31.1\""), + .define("MLX_VERSION", to: "\"0.32.0\""), ], linkerSettings: linkerSettings, plugins: cudaBuildPlugins, diff --git a/Source/Cmlx/include-framework/Cmlx.h b/Source/Cmlx/include-framework/Cmlx.h index 8c0077864..47eba4c25 100644 --- a/Source/Cmlx/include-framework/Cmlx.h +++ b/Source/Cmlx/include-framework/Cmlx.h @@ -12,9 +12,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -38,6 +40,8 @@ #include #include #include +#include +#include #include #include #include diff --git a/Source/Cmlx/include-framework/Metal.hpp b/Source/Cmlx/include-framework/Metal.hpp index e2bd5fe90..717aadc89 100644 --- a/Source/Cmlx/include-framework/Metal.hpp +++ b/Source/Cmlx/include-framework/Metal.hpp @@ -2,7 +2,7 @@ // // Metal.hpp.in // -// Autogenerated from commit 09bcf14a0c4d70b39c72e2d0b7d500c1e31efe4c. +// Autogenerated from commit 68ad1c14e2bae5b8cca5aefe3a5eb91650c0147a. // // Copyright 2020-2024 Apple Inc. // diff --git a/Source/Cmlx/include-framework/mlx-allocator.h b/Source/Cmlx/include-framework/mlx-allocator.h index 57069a4bd..b014aee0d 100644 --- a/Source/Cmlx/include-framework/mlx-allocator.h +++ b/Source/Cmlx/include-framework/mlx-allocator.h @@ -73,5 +73,7 @@ inline void release(Buffer buffer) { allocator().release(buffer); } +MLX_API bool can_reuse_alien_buffer(void* ptr); + } // namespace mlx::core::allocator #endif diff --git a/Source/Cmlx/include-framework/mlx-array.h b/Source/Cmlx/include-framework/mlx-array.h index 042692bc0..9c6819c2e 100644 --- a/Source/Cmlx/include-framework/mlx-array.h +++ b/Source/Cmlx/include-framework/mlx-array.h @@ -427,6 +427,7 @@ class MLX_API array { } void detach_event() const { + array_desc_->event.check_error(); array_desc_->event = Event{}; } @@ -444,6 +445,16 @@ class MLX_API array { size_t data_size, Strides strides, Flags flags, + Deleter d) { + set_data(buffer, data_size, std::move(strides), flags, 0, std::move(d)); + } + + void set_data( + allocator::Buffer buffer, + size_t data_size, + Strides strides, + Flags flags, + int64_t offset = 0, Deleter d = allocator::free); void copy_shared_buffer( diff --git a/Source/Cmlx/include-framework/mlx-backend-common-metal_kernel.h b/Source/Cmlx/include-framework/mlx-backend-common-metal_kernel.h new file mode 100644 index 000000000..658f0360e --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-backend-common-metal_kernel.h @@ -0,0 +1,29 @@ +#ifdef __cplusplus +// Copyright © 2026 Apple Inc. + +#pragma once + +namespace mlx::core { + +enum class MathMode { + Safe = 0, + Relaxed = 1, + Fast = 2, +}; + +struct CompileOptions { + MathMode math_mode = MathMode::Safe; + + CompileOptions() = default; + bool operator==(const CompileOptions&) const = default; + + // A simple way to make export work, needs more work when adding new options. + using Data = int; + CompileOptions(Data data) : math_mode(static_cast(data)) {} + Data serialize() const { + return static_cast(math_mode); + } +}; + +} // namespace mlx::core +#endif diff --git a/Source/Cmlx/include-framework/mlx-backend-common-utils.h b/Source/Cmlx/include-framework/mlx-backend-common-utils.h index 660ced1a9..0b28db3b3 100644 --- a/Source/Cmlx/include-framework/mlx-backend-common-utils.h +++ b/Source/Cmlx/include-framework/mlx-backend-common-utils.h @@ -117,6 +117,39 @@ struct ContiguousIterator { loc += strides_[i]; } + void step(int64_t s) { + int dims = shape_.size(); + if (dims == 0) { + return; + } + int i = dims - 1; + while (s > 0) { + if (shape_[i] - pos_[i] > 1) { + int steps = static_cast( + std::min(static_cast(shape_[i] - pos_[i] - 1), s)); + pos_[i] += steps; + loc += strides_[i] * steps; + s -= steps; + } else { + while (pos_[i] == (shape_[i] - 1) && i > 0) { + pos_[i] = 0; + loc -= (shape_[i] - 1) * strides_[i]; + i--; + } + pos_[i]++; + loc += strides_[i]; + s--; + } + } + } + + int64_t contiguous_suffix() { + if (shape_.size() == 0) { + return 0; + } + return (strides_.back() == 1) ? shape_.back() : 0; + } + void seek(int64_t n) { loc = 0; for (int i = shape_.size() - 1; i >= 0; --i) { diff --git a/Source/Cmlx/include-framework/mlx-backend-cpu-encoder.h b/Source/Cmlx/include-framework/mlx-backend-cpu-encoder.h index cc0b0eabe..6e480b795 100644 --- a/Source/Cmlx/include-framework/mlx-backend-cpu-encoder.h +++ b/Source/Cmlx/include-framework/mlx-backend-cpu-encoder.h @@ -47,11 +47,10 @@ struct MLX_API CommandEncoder { auto task = std::bind(std::forward(f), std::forward(args)...); if (num_ops_ == 0) { scheduler::notify_new_task(stream_); - auto task_wrap = [s = stream_, task = std::move(task)]() mutable { - task(); - scheduler::notify_task_completion(s); - }; - scheduler::enqueue(stream_, std::move(task_wrap)); + scheduler::enqueue(stream_, std::move(task)); + // Notify completion separately as |task| may throw exception. + scheduler::enqueue( + stream_, [s = stream_] { scheduler::notify_task_completion(s); }); } else { scheduler::enqueue(stream_, std::move(task)); } @@ -63,7 +62,10 @@ struct MLX_API CommandEncoder { int num_ops_{0}; }; -MLX_API CommandEncoder& get_command_encoder(Stream stream); +MLX_API CommandEncoder& get_command_encoder(Stream s); + +std::unordered_map& get_command_encoders(); +std::unordered_map& get_global_command_encoders(); } // namespace mlx::core::cpu #endif diff --git a/Source/Cmlx/include-framework/mlx-backend-gpu-eval.h b/Source/Cmlx/include-framework/mlx-backend-gpu-eval.h index bedef920e..2e2872aa5 100644 --- a/Source/Cmlx/include-framework/mlx-backend-gpu-eval.h +++ b/Source/Cmlx/include-framework/mlx-backend-gpu-eval.h @@ -11,10 +11,13 @@ namespace mlx::core::gpu { -void new_stream(Stream stream); +void init(); +void new_stream(Stream s); +void new_thread_unsafe_stream(Stream s); void eval(array& arr); void finalize(Stream s); void synchronize(Stream s); +void clear_streams(); } // namespace mlx::core::gpu #endif diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-device.h b/Source/Cmlx/include-framework/mlx-backend-metal-device.h index bfb07deab..b7c06c1a4 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-device.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-device.h @@ -7,11 +7,12 @@ #include #include #include -#include #include #include #include +#include +#include #include namespace mlx::core::metal { @@ -19,10 +20,13 @@ namespace mlx::core::metal { using MTLFCList = std::vector>; -struct DeviceStream; +class Device; + +class MLX_API CommandEncoder { + public: + CommandEncoder(Device& d, int index, ResidencySets& residency_sets); + ~CommandEncoder(); -struct MLX_API CommandEncoder { - explicit CommandEncoder(DeviceStream& stream); CommandEncoder(const CommandEncoder&) = delete; CommandEncoder& operator=(const CommandEncoder&) = delete; @@ -41,29 +45,26 @@ struct MLX_API CommandEncoder { CommandEncoder& enc; }; + void set_buffer(const MTL::Buffer* buf, int idx, int64_t offset = 0); void set_input_array(const array& a, int idx, int64_t offset = 0); void set_output_array(array& a, int idx, int64_t offset = 0); void register_output_array(const array& a); + + void add_temporary(array arr); + void add_temporaries(std::vector arrays); + void dispatch_threadgroups(MTL::Size grid_dims, MTL::Size group_dims); void dispatch_threads(MTL::Size grid_dims, MTL::Size group_dims); void maybeInsertBarrier(); - void set_buffer(const MTL::Buffer* buf, int idx, int64_t offset = 0); void set_compute_pipeline_state(MTL::ComputePipelineState* kernel) { - enc_->setComputePipelineState(kernel); - } - - void wait_for_fence(MTL::Fence* fence) { - enc_->waitForFence(fence); - } - - void update_fence(MTL::Fence* fence) { - enc_->updateFence(fence); + get_command_encoder()->setComputePipelineState(kernel); } template >> void set_vector_bytes(const Vec& vec, size_t nelems, int idx) { - enc_->setBytes(vec.data(), nelems * sizeof(typename Vec::value_type), idx); + get_command_encoder()->setBytes( + vec.data(), nelems * sizeof(typename Vec::value_type), idx); } template >> void set_vector_bytes(const Vec& vec, int idx) { @@ -72,79 +73,76 @@ struct MLX_API CommandEncoder { template void set_bytes(const T* v, int n, int idx) { - return enc_->setBytes(v, n * sizeof(T), idx); + return get_command_encoder()->setBytes(v, n * sizeof(T), idx); } template void set_bytes(const T& v, int idx) { - return enc_->setBytes(&v, sizeof(T), idx); + return get_command_encoder()->setBytes(&v, sizeof(T), idx); } void set_threadgroup_memory_length(size_t length, int idx) { - enc_->setThreadgroupMemoryLength(length, idx); + get_command_encoder()->setThreadgroupMemoryLength(length, idx); } ConcurrentContext start_concurrent() { return ConcurrentContext(*this); } - ~CommandEncoder(); - - // Inputs to all kernels in the encoder including temporaries - std::unordered_set& inputs() { - return all_inputs_; - }; - - // Outputs of all kernels in the encoder including temporaries - std::unordered_set& outputs() { - return all_outputs_; - }; void barrier(); + void end_encoding(); + void wait_event(Event event, uint64_t value); + void signal_event(Event event, uint64_t value); + bool needs_commit() const; + void commit(std::function completion = nullptr); + void synchronize(); + + MTL::CommandBuffer* get_command_buffer() const { + return buffer_.get(); + } private: - DeviceStream& stream_; - MTL::ComputeCommandEncoder* enc_; + MTL::ComputeCommandEncoder* get_command_encoder(); + + Device& device_; + bool exiting_{false}; + + // Buffer that stores encoded commands. + NS::SharedPtr queue_; + NS::SharedPtr buffer_; + int buffer_ops_{0}; + size_t buffer_sizes_{0}; + + // The residency set and how many of its sets this queue has attached. + ResidencySets& residency_sets_; + uint64_t sets_attached_{0}; + + // The events hooked to current command buffer. + std::vector wait_events_; + std::vector> signal_events_; + + // Error from previous commited command buffer. + Error error_; + + // Encoder for issuing GPU commands. + // The members are used within a single ComputeCommandEncoder and will be + // reset after calling end_encoding(). + NS::SharedPtr encoder_; + NS::SharedPtr fence_; bool needs_barrier_{false}; bool concurrent_{false}; + std::vector temporaries_; + std::unordered_set prev_inputs_; std::unordered_set prev_outputs_; + std::unordered_set next_inputs_; std::unordered_set next_outputs_; std::unordered_set concurrent_outputs_; std::unordered_set all_inputs_; std::unordered_set all_outputs_; -}; -struct Fence { - Fence(MTL::Fence* fence) : fence(fence) {} - ~Fence() { - fence->release(); - } - MTL::Fence* fence; -}; - -struct DeviceStream { - DeviceStream(MTL::CommandQueue* queue) : queue(queue) {}; - ~DeviceStream() { - queue->release(); - if (buffer != nullptr) { - buffer->release(); - } - }; - MTL::CommandQueue* queue; - // A map of prior command encoder outputs to their corresponding fence - std::unordered_map> outputs; - // Used to allow thread-safe access to the outputs map - std::mutex fence_mtx; - - // Data updated between command buffers - MTL::CommandBuffer* buffer{nullptr}; - int buffer_ops{0}; - size_t buffer_sizes{0}; - - // The command encoder, fence, and temporaries are updated between command - // encoders - std::unique_ptr encoder{nullptr}; - std::shared_ptr fence; - std::vector temporaries; + // A map of prior command encoder outputs to their corresponding fence. + std::unordered_map> prev_ce_outputs_; + std::mutex outputs_mtx_; }; class MLX_API Device { @@ -155,26 +153,18 @@ class MLX_API Device { ~Device(); MTL::Device* mtl_device() { - return device_; + return device_.get(); }; - const std::string& get_architecture() { + const std::string& get_architecture() const { return arch_; } - int get_architecture_gen() const { return arch_gen_; } - - void new_queue(int index); - - MTL::CommandQueue* get_queue(Stream stream); - - MTL::CommandBuffer* get_command_buffer(int index); - bool command_buffer_needs_commit(int index); - void commit_command_buffer(int index); - CommandEncoder& get_command_encoder(int index); - void end_encoding(int index); + std::tuple get_max_ops_mb_per_buffer() const { + return std::make_tuple(max_ops_per_buffer_, max_mb_per_buffer_); + } MTL::Library* get_library( const std::string& name, @@ -182,8 +172,15 @@ class MLX_API Device { MTL::Library* get_library( const std::string& name, + const CompileOptions& compile_options, const std::function& builder); + MTL::Library* get_library( + const std::string& name, + const std::function& builder) { + return get_library(name, {}, builder); + } + void clear_library(const std::string& name); MTL::ComputePipelineState* get_kernel( @@ -199,40 +196,31 @@ class MLX_API Device { const MTLFCList& func_consts = {}, const std::vector& linked_functions = {}); - MTL::ArgumentEncoder* argument_encoder( - const std::vector& arg_descs) const; - - // Record temporary arrays for the given stream index - void add_temporary(array arr, int index); - void add_temporaries(std::vector arrays, int index); - - void set_residency_set(const MTL::ResidencySet* residency_set); - - private: - DeviceStream& get_stream_(int index) { - return stream_map_.find(index)->second; + ResidencySets& residency_sets() { + return residency_sets_; } - MTL::Library* get_library_cache_(const std::string& name); - MTL::Library* get_library_(const std::string& name); - MTL::Library* build_library_(const std::string& source_string); - - MTL::Function* get_function_(const std::string& name, MTL::Library* mtl_lib); + private: + NS::SharedPtr build_library_( + const std::string& source_string, + const CompileOptions& compile_options = {}); - MTL::Function* get_function_( + NS::SharedPtr get_function_( + const std::string& name, + MTL::Library* mtl_lib); + NS::SharedPtr get_function_( const std::string& name, const std::string& specialized_name, const MTLFCList& func_consts, MTL::Library* mtl_lib); - MTL::LinkedFunctions* get_linked_functions_( + NS::SharedPtr get_linked_functions_( const std::vector& funcs); - MTL::ComputePipelineState* get_kernel_( + NS::SharedPtr get_kernel_( const std::string& name, const MTL::Function* mtl_function); - - MTL::ComputePipelineState* get_kernel_( + NS::SharedPtr get_kernel_( const std::string& name, const MTL::Function* mtl_function, const MTL::LinkedFunctions* linked_functions); @@ -244,18 +232,17 @@ class MLX_API Device { const MTLFCList& func_consts = {}, const std::vector& linked_functions = {}); - MTL::Device* device_; - std::unordered_map stream_map_; + NS::SharedPtr device_; + ResidencySets residency_sets_; std::shared_mutex kernel_mtx_; std::shared_mutex library_mtx_; - std::unordered_map library_map_; - MTL::Library* default_library_; + std::unordered_map> library_map_; + NS::SharedPtr default_library_; std::unordered_map< MTL::Library*, - std::unordered_map> + std::unordered_map>> library_kernels_; - const MTL::ResidencySet* residency_set_{nullptr}; std::string arch_; int arch_gen_; int max_ops_per_buffer_; @@ -263,29 +250,13 @@ class MLX_API Device { }; MLX_API Device& device(mlx::core::Device); +MLX_API CommandEncoder& get_command_encoder(Stream s); -std::unique_ptr> new_scoped_memory_pool(); - -inline bool is_nax_available() { -#ifdef MLX_METAL_NO_NAX - return false; -#else - auto _check_nax = []() { - bool can_use_nax = false; - if (__builtin_available( - macOS 26.2, iOS 26.2, tvOS 26.2, visionOS 26.2, *)) { - can_use_nax = true; - } - auto& d = metal::device(mlx::core::Device::gpu); - auto arch = d.get_architecture().back(); - auto gen = d.get_architecture_gen(); - can_use_nax &= gen >= (arch == 'p' ? 18 : 17); - return can_use_nax; - }; - static bool is_nax_available_ = _check_nax(); - return is_nax_available_; -#endif -} +std::unordered_map& get_command_encoders(); +std::unordered_map& get_global_command_encoders(); +NS::SharedPtr new_scoped_memory_pool(); + +bool is_nax_available(); } // namespace mlx::core::metal #endif diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-metal.h b/Source/Cmlx/include-framework/mlx-backend-metal-metal.h index 5e1db21f0..06511de13 100644 --- a/Source/Cmlx/include-framework/mlx-backend-metal-metal.h +++ b/Source/Cmlx/include-framework/mlx-backend-metal-metal.h @@ -23,5 +23,10 @@ MLX_API const std::unordered_map>& device_info(); +/* Set a custom path to mlx.metallib. Must be called before any MLX operation. + */ +MLX_API void set_metallib_path(const std::string& path); +MLX_API const std::string& get_metallib_path(); + } // namespace mlx::core::metal #endif diff --git a/Source/Cmlx/include-framework/mlx-backend-metal-resident.h b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h new file mode 100644 index 000000000..a5c6b58f4 --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-backend-metal-resident.h @@ -0,0 +1,125 @@ +#ifdef __cplusplus +// Copyright © 2024 Apple Inc. + +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace mlx::core::metal { + +// Keeps allocations GPU-resident, up to the wired limit set by +// `set_wired_limit` (0 by default, i.e. nothing is wired unless asked for). +// +// Within that budget the allocations are distributed over several size-capped +// MTL::ResidencySets rather than one large one. macOS makes residency decisions +// per residency set, so when a set loses residency under GPU memory pressure +// only the allocations in that set have to be made resident again. Capping the +// size of each set bounds the cost of one such event. Every set holds a +// standing requestResidency() and is attached to every command queue. +// +// MLX_RESIDENCY_SET_MAX_PCT (env::residency_set_max_pct) sets the per-set cap. +// The total wired budget is unaffected by it: that is still `set_wired_limit`. +class ResidencySets { + public: + ResidencySets(MTL::Device* d); + ~ResidencySets(); + + ResidencySets(const ResidencySets&) = delete; + ResidencySets& operator=(const ResidencySets&) = delete; + + // Called with the allocator's mutex held. + void insert(MTL::Allocation* buf); + void erase(MTL::Allocation* buf); + + void resize(size_t size); + + bool enabled() const { + return enabled_; + } + + // Attaches the sets this queue has not seen yet. Called from the encoder + // thread before every command-buffer commit, because Metal locks a command + // buffer's residency at commit time. + void attach_new_sets(MTL::CommandQueue* q, uint64_t& attached); + + // Total bytes currently wired across all sets. + size_t wired_size() const { + std::lock_guard lk(mtx_); + return total_wired_; + } + size_t num_sets() const { + return num_sets_.load(std::memory_order_acquire); + } + + // Testing only: sets the per-set cap for subsequent inserts, bypassing the + // size floor so tests can reach the multi-set paths cheaply. 0 selects the + // single-set layout. + void set_max_bytes_per_set(size_t bytes) { + std::lock_guard lk(mtx_); + max_bytes_per_set_ = bytes; + } + size_t max_bytes_per_set() const { + std::lock_guard lk(mtx_); + return max_bytes_per_set_; + } + + private: + // A set id of kNoSet means the allocation is tracked but is not in any + // set, because it did not fit in the wired limit. + static constexpr uint32_t kNoSet = UINT32_MAX; + // A command queue accepts a limited number of residency sets and every set + // is attached to every queue, so the number of sets is capped. Once the cap + // is reached allocations go to the emptiest set, which then grows past + // max_bytes_per_set_. The limit is from the Metal feature set tables: + // https://developer.apple.com/metal-Metal-Feature-Set-Tables.pdf + static constexpr uint32_t kMaxSets = 32; + + struct Set { + NS::SharedPtr set; + size_t size{0}; + }; + + // The set an allocation lives in and the size it was inserted with. Sizes + // are recorded rather than read back from the allocation at erase, so the + // running totals cannot drift. + struct Placement { + uint32_t set_id; + size_t bytes; + }; + + // The following all require mtx_. add_set_locked reports whether the driver + // gave us a set; the new set is the last one. add_to_set_locked returns the + // set it used. Neither it nor remove_from_set_locked commits, so a bulk + // resize costs one commit per set instead of one per allocation. + bool add_set_locked(NS::Error** error = nullptr); + uint32_t choose_set_locked(size_t bytes); + uint32_t add_to_set_locked(const MTL::Allocation* buf, Placement& at); + void remove_from_set_locked(const MTL::Allocation* buf, Placement& at); + + MTL::Device* device_{nullptr}; + bool enabled_{false}; + + // Per-set cap in bytes (0 for a single set) and the total wired budget. + size_t max_bytes_per_set_{0}; + size_t capacity_{0}; + size_t total_wired_{0}; + + // Sets indexed by id. Append-only, so ids stay valid and dense. + std::vector sets_; + // Every tracked allocation, wired or not. + std::unordered_map buf_to_set_; + + // sets_.size(), published so a queue can check for new sets without + // taking the lock. + std::atomic num_sets_{0}; + mutable std::mutex mtx_; +}; + +} // namespace mlx::core::metal +#endif diff --git a/Source/Cmlx/include-framework/mlx-c-compile.h b/Source/Cmlx/include-framework/mlx-c-compile.h index bbc4c3e10..6ceb63d2e 100644 --- a/Source/Cmlx/include-framework/mlx-c-compile.h +++ b/Source/Cmlx/include-framework/mlx-c-compile.h @@ -34,6 +34,13 @@ typedef enum mlx_compile_mode_ { MLX_COMPILE_MODE_NO_FUSE, MLX_COMPILE_MODE_ENABLED } mlx_compile_mode; + +typedef struct mlx_compile_cache_ { + void* ctx; +} mlx_compile_cache; +mlx_compile_cache mlx_compile_cache_new(); +int mlx_compile_cache_free(mlx_compile_cache cache); + int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless); int mlx_detail_compile( mlx_closure* res, @@ -42,8 +49,9 @@ int mlx_detail_compile( bool shapeless, const uint64_t* constants, size_t constants_num); -int mlx_detail_compile_clear_cache(void); -int mlx_detail_compile_erase(uintptr_t fun_id); +int mlx_detail_compile_cache(mlx_compile_cache* res); +int mlx_detail_compile_clear_cache(const mlx_compile_cache cache); +int mlx_detail_compile_erase(const mlx_compile_cache cache, uintptr_t fun_id); int mlx_disable_compile(void); int mlx_enable_compile(void); int mlx_set_compile_mode(mlx_compile_mode mode); diff --git a/Source/Cmlx/include-framework/mlx-c-distributed_group.h b/Source/Cmlx/include-framework/mlx-c-distributed_group.h index dea429b35..c55ee98c0 100644 --- a/Source/Cmlx/include-framework/mlx-c-distributed_group.h +++ b/Source/Cmlx/include-framework/mlx-c-distributed_group.h @@ -23,6 +23,24 @@ typedef struct mlx_distributed_group_ { void* ctx; } mlx_distributed_group; +/** + * Create an empty group. + */ +mlx_distributed_group mlx_distributed_group_new(void); + +/** + * Free the group. + */ +int mlx_distributed_group_free(mlx_distributed_group group); + +/** + * Initialize distributed. + */ +int mlx_distributed_init( + mlx_distributed_group* res, + bool strict, + const char* bk /* may be null */); + /** * Get the rank. */ @@ -36,21 +54,17 @@ int mlx_distributed_group_size(mlx_distributed_group group); /** * Split the group. */ -mlx_distributed_group -mlx_distributed_group_split(mlx_distributed_group group, int color, int key); +int mlx_distributed_group_split( + mlx_distributed_group* res, + mlx_distributed_group group, + int color, + int key); /** * Check if distributed is available. */ bool mlx_distributed_is_available(const char* bk /* may be null */); -/** - * Initialize distributed. - */ -mlx_distributed_group mlx_distributed_init( - bool strict, - const char* bk /* may be null */); - /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/include-framework/mlx-c-fast.h b/Source/Cmlx/include-framework/mlx-c-fast.h index adfdcd0e1..eac10b396 100644 --- a/Source/Cmlx/include-framework/mlx-c-fast.h +++ b/Source/Cmlx/include-framework/mlx-c-fast.h @@ -28,6 +28,12 @@ extern "C" { */ /**@{*/ +int mlx_fast_cross_entropy( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s); + typedef struct mlx_fast_cuda_kernel_config_ { void* ctx; } mlx_fast_cuda_kernel_config; @@ -195,6 +201,7 @@ int mlx_fast_scaled_dot_product_attention( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include-framework/mlx-c-fft.h b/Source/Cmlx/include-framework/mlx-c-fft.h index a9dd4eb7c..c432b180e 100644 --- a/Source/Cmlx/include-framework/mlx-c-fft.h +++ b/Source/Cmlx/include-framework/mlx-c-fft.h @@ -28,11 +28,18 @@ extern "C" { */ /**@{*/ +typedef enum mlx_fft_norm_ { + MLX_FFT_NORM_BACKWARD, + MLX_FFT_NORM_ORTHO, + MLX_FFT_NORM_FORWARD +} mlx_fft_norm; + int mlx_fft_fft( mlx_array* res, const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_fft2( mlx_array* res, @@ -41,7 +48,9 @@ int mlx_fft_fft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); +int mlx_fft_fftfreq(mlx_array* res, int n, double d, const mlx_stream s); int mlx_fft_fftn( mlx_array* res, const mlx_array a, @@ -49,6 +58,7 @@ int mlx_fft_fftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_fftshift( mlx_array* res, @@ -61,6 +71,7 @@ int mlx_fft_ifft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifft2( mlx_array* res, @@ -69,6 +80,7 @@ int mlx_fft_ifft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifftn( mlx_array* res, @@ -77,6 +89,7 @@ int mlx_fft_ifftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifftshift( mlx_array* res, @@ -89,6 +102,7 @@ int mlx_fft_irfft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_irfft2( mlx_array* res, @@ -97,6 +111,7 @@ int mlx_fft_irfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_irfftn( mlx_array* res, @@ -105,12 +120,14 @@ int mlx_fft_irfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_rfft( mlx_array* res, const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_rfft2( mlx_array* res, @@ -119,7 +136,9 @@ int mlx_fft_rfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); +int mlx_fft_rfftfreq(mlx_array* res, int n, double d, const mlx_stream s); int mlx_fft_rfftn( mlx_array* res, const mlx_array a, @@ -127,6 +146,7 @@ int mlx_fft_rfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include-framework/mlx-c-graph_utils.h b/Source/Cmlx/include-framework/mlx-c-graph_utils.h new file mode 100644 index 000000000..f2269ffa9 --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-c-graph_utils.h @@ -0,0 +1,61 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#ifndef MLX_GRAPH_UTILS_H +#define MLX_GRAPH_UTILS_H + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * \defgroup graph_utils Graph Utils + */ +/**@{*/ + +typedef struct mlx_node_namer_ { + void* ctx; +} mlx_node_namer; + +mlx_node_namer mlx_node_namer_new(); +int mlx_node_namer_free(mlx_node_namer namer); +int mlx_node_namer_set_name( + mlx_node_namer namer, + const mlx_array arr, + const char* name); +int mlx_node_namer_get_name( + const char** name, + mlx_node_namer namer, + const mlx_array arr); + +int mlx_export_to_dot( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs); +int mlx_print_graph( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs); + +/**@}*/ + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/Source/Cmlx/include-framework/mlx-c-io.h b/Source/Cmlx/include-framework/mlx-c-io.h index 1c21a88e4..4cc653f19 100644 --- a/Source/Cmlx/include-framework/mlx-c-io.h +++ b/Source/Cmlx/include-framework/mlx-c-io.h @@ -33,6 +33,9 @@ int mlx_load_reader( mlx_io_reader in_stream, const mlx_stream s); int mlx_load(mlx_array* res, const char* file, const mlx_stream s); + +int mlx_load_gguf(mlx_io_gguf* gguf, const char* file, const mlx_stream s); + int mlx_load_safetensors_reader( mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, @@ -45,6 +48,8 @@ int mlx_load_safetensors( const mlx_stream s); int mlx_save_writer(mlx_io_writer out_stream, const mlx_array a); int mlx_save(const char* file, const mlx_array a); +int mlx_save_gguf(const char* file, mlx_io_gguf gguf); + int mlx_save_safetensors_writer( mlx_io_writer in_stream, const mlx_map_string_to_array param, diff --git a/Source/Cmlx/include-framework/mlx-c-io_types.h b/Source/Cmlx/include-framework/mlx-c-io_types.h index 52fc7fbc5..4eccc95a9 100644 --- a/Source/Cmlx/include-framework/mlx-c-io_types.h +++ b/Source/Cmlx/include-framework/mlx-c-io_types.h @@ -95,6 +95,52 @@ int mlx_io_writer_tostring(mlx_string* str_, mlx_io_writer io); */ int mlx_io_writer_free(mlx_io_writer io); +/** + * A MLX GGUF object. + */ +typedef struct mlx_io_gguf_ { + void* ctx; +} mlx_io_gguf; + +mlx_io_gguf mlx_io_gguf_new(void); +int mlx_io_gguf_free(mlx_io_gguf io); +int mlx_io_gguf_get_keys(mlx_vector_string* keys, mlx_io_gguf io); +int mlx_io_gguf_get_array(mlx_array* arr, mlx_io_gguf io, const char* key); +int mlx_io_gguf_get_metadata_array( + mlx_array* arr, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_get_metadata_string( + mlx_string* str, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_get_metadata_vector_string( + mlx_vector_string* vstr, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_has_metadata_array(bool* flag, mlx_io_gguf io, const char* key); +int mlx_io_gguf_has_metadata_string( + bool* flag, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_has_metadata_vector_string( + bool* flag, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_set_array(mlx_io_gguf io, const char* key, const mlx_array arr); +int mlx_io_gguf_set_metadata_array( + mlx_io_gguf io, + const char* key, + const mlx_array marr); +int mlx_io_gguf_set_metadata_string( + mlx_io_gguf io, + const char* key, + const char* mstr); +int mlx_io_gguf_set_metadata_vector_string( + mlx_io_gguf io, + const char* key, + const mlx_vector_string mvstr); + /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/include-framework/mlx-c-linalg.h b/Source/Cmlx/include-framework/mlx-c-linalg.h index 7d2941e6e..75acd1254 100644 --- a/Source/Cmlx/include-framework/mlx-c-linalg.h +++ b/Source/Cmlx/include-framework/mlx-c-linalg.h @@ -44,6 +44,7 @@ int mlx_linalg_cross( const mlx_array b, int axis, const mlx_stream s); +int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -97,6 +98,11 @@ int mlx_linalg_qr( mlx_array* res_1, const mlx_array a, const mlx_stream s); +int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s); int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/Source/Cmlx/include-framework/mlx-c-metal.h b/Source/Cmlx/include-framework/mlx-c-metal.h index 52fad31ee..7e52f30bf 100644 --- a/Source/Cmlx/include-framework/mlx-c-metal.h +++ b/Source/Cmlx/include-framework/mlx-c-metal.h @@ -28,7 +28,9 @@ extern "C" { */ /**@{*/ +int mlx_metal_get_metallib_path(mlx_string* res); int mlx_metal_is_available(bool* res); +int mlx_metal_set_metallib_path(const char* path); int mlx_metal_start_capture(const char* path); int mlx_metal_stop_capture(void); diff --git a/Source/Cmlx/include-framework/mlx-c-mlx.h b/Source/Cmlx/include-framework/mlx-c-mlx.h index d063b6f86..fb25573ef 100644 --- a/Source/Cmlx/include-framework/mlx-c-mlx.h +++ b/Source/Cmlx/include-framework/mlx-c-mlx.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include diff --git a/Source/Cmlx/include-framework/mlx-c-ops.h b/Source/Cmlx/include-framework/mlx-c-ops.h index a38f2f53b..7531fe08f 100644 --- a/Source/Cmlx/include-framework/mlx-c-ops.h +++ b/Source/Cmlx/include-framework/mlx-c-ops.h @@ -323,7 +323,38 @@ int mlx_conv_transpose3d( int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s); +int mlx_count_nonzero_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); +int mlx_cummax_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_cummax( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cummin_axis( mlx_array* res, const mlx_array a, int axis, @@ -331,25 +362,40 @@ int mlx_cummax( bool inclusive, const mlx_stream s); int mlx_cummin( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cumprod_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumprod( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s); +int mlx_cumsum_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumsum( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_depends( @@ -375,6 +421,12 @@ int mlx_diagonal( int axis1, int axis2, const mlx_stream s); +int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s); int mlx_divide( mlx_array* res, const mlx_array a, @@ -423,6 +475,18 @@ int mlx_flatten( int start_axis, int end_axis, const mlx_stream s); +int mlx_flip_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s); +int mlx_flip_axis( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_flip(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor_divide( mlx_array* res, @@ -486,6 +550,20 @@ int mlx_gather_qmm( const char* mode, bool sorted_indices, const mlx_stream s); +int mlx_gather_qqmm( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s); int mlx_greater( mlx_array* res, const mlx_array a, @@ -543,6 +621,14 @@ int mlx_less_equal( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_linspace_endpoint( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s); int mlx_linspace( mlx_array* res, double start, @@ -559,13 +645,19 @@ int mlx_logaddexp( const mlx_array a, const mlx_array b, const mlx_stream s); -int mlx_logcumsumexp( +int mlx_logcumsumexp_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s); +int mlx_logcumsumexp( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_logical_and( mlx_array* res, const mlx_array a, @@ -577,6 +669,11 @@ int mlx_logical_or( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s); int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -647,13 +744,24 @@ int mlx_mean( const mlx_array a, bool keepdims, const mlx_stream s); -int mlx_median( +int mlx_median_axes( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s); +int mlx_median_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_median( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); int mlx_meshgrid( mlx_vector_array* res, const mlx_vector_array arrays, @@ -721,6 +829,11 @@ int mlx_ones( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_ones_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_outer( mlx_array* res, @@ -757,6 +870,7 @@ int mlx_partition( const mlx_array a, int kth, const mlx_stream s); +int mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_power( mlx_array* res, const mlx_array a, @@ -956,6 +1070,12 @@ int mlx_scatter_prod_single( const mlx_array updates, int axis, const mlx_stream s); +int mlx_searchsorted( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s); int mlx_segmented_mm( mlx_array* res, const mlx_array a, @@ -1004,6 +1124,50 @@ int mlx_slice_update_dynamic( const int* axes, size_t axes_num, const mlx_stream s); +int mlx_slice_update_add( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_max( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_min( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_prod( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); int mlx_softmax_axes( mlx_array* res, const mlx_array a, @@ -1163,7 +1327,7 @@ int mlx_topk_axis( int axis, const mlx_stream s); int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s); -int mlx_trace( +int mlx_trace_axes( mlx_array* res, const mlx_array a, int offset, @@ -1171,6 +1335,7 @@ int mlx_trace( int axis2, mlx_dtype dtype, const mlx_stream s); +int mlx_trace(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_transpose_axes( mlx_array* res, const mlx_array a, @@ -1187,6 +1352,7 @@ int mlx_tri( const mlx_stream s); int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s); int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s); +int mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -1194,6 +1360,12 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unstack_axis( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_unstack(mlx_vector_array* res, const mlx_array a, const mlx_stream s); int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -1215,6 +1387,12 @@ int mlx_var( bool keepdims, int ddof, const mlx_stream s); +int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s); int mlx_view( mlx_array* res, const mlx_array a, @@ -1232,6 +1410,11 @@ int mlx_zeros( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_zeros_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include-framework/mlx-c-stream.h b/Source/Cmlx/include-framework/mlx-c-stream.h index a4aa39abc..039d4b60a 100644 --- a/Source/Cmlx/include-framework/mlx-c-stream.h +++ b/Source/Cmlx/include-framework/mlx-c-stream.h @@ -33,6 +33,18 @@ mlx_stream mlx_stream_new(void); * Returns a new stream on a device. */ mlx_stream mlx_stream_new_device(mlx_device dev); +/** + * Returns a new stream on a device that can be used from any thread. + * + * Streams are otherwise thread affine: a stream's GPU command encoder is + * registered per thread, so evaluating on a stream from a thread other than + * the one that created it fails. Streams returned here are registered + * globally instead. + * + * MLX applies no synchronization to these streams -- it is the caller's + * responsibility to ensure there are no data races on them. + */ +mlx_stream mlx_stream_new_thread_unsafe(mlx_device dev); /** * Set stream to provided src stream. */ diff --git a/Source/Cmlx/include-framework/mlx-device.h b/Source/Cmlx/include-framework/mlx-device.h index e7f9f2d19..aeea7b326 100644 --- a/Source/Cmlx/include-framework/mlx-device.h +++ b/Source/Cmlx/include-framework/mlx-device.h @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -24,15 +25,22 @@ struct MLX_API Device { DeviceType type; int index; + + // TODO: Use default three-way comparison when it gets supported in XCode. + bool operator==(const Device&) const = default; + bool operator<(const Device& rhs) const { + return std::tie(type, index) < std::tie(rhs.type, rhs.index); + } }; +inline bool operator==(const Device& device, Device::DeviceType type) { + return device.type == type; +} + MLX_API const Device& default_device(); MLX_API void set_default_device(const Device& d); -MLX_API bool operator==(const Device& lhs, const Device& rhs); -MLX_API bool operator!=(const Device& lhs, const Device& rhs); - MLX_API bool is_available(const Device& d); /** Get the number of available devices for the given device type. */ diff --git a/Source/Cmlx/include-framework/mlx-distributed-distributed.h b/Source/Cmlx/include-framework/mlx-distributed-distributed.h index 7fdb08fa7..0c3fa034f 100644 --- a/Source/Cmlx/include-framework/mlx-distributed-distributed.h +++ b/Source/Cmlx/include-framework/mlx-distributed-distributed.h @@ -3,6 +3,8 @@ #pragma once +#include +#include #include #include @@ -14,7 +16,26 @@ namespace mlx::core::distributed { // Forward declaration of the base group implementation. namespace detail { class GroupImpl; -}; +} + +/** + * Byte-level all-gather function used by the JACCL side channel. + * + * Args: + * src: Pointer to this rank's data of size n_bytes. + * dst: Pointer to an output buffer of size size_ * n_bytes. After the call, + * dst[rank * n_bytes, (rank+1) * n_bytes] contains the data from rank. + * n_bytes: The number of bytes contributed by each rank. + */ +using AllGatherFn = std::function; + +/** + * Factory that produces a per-rank side-channel all-gather function. + * + * The factory receives the rank and size of the group and returns the + * AllGatherFn that will be used for the side channel. + */ +using AllGatherFactory = std::function; /* Check if a communication backend is available */ MLX_API bool is_available(); @@ -59,5 +80,16 @@ struct MLX_API Group { */ MLX_API Group init(bool strict = false, const std::string& bk = "any"); +/** + * Initialize the distributed backend using a custom JACCL side channel. + * + * This behaves like init(strict, bk) but requires bk to be "jaccl". + */ +MLX_API Group +init(bool strict, const std::string& bk, AllGatherFactory factory); + +/** Clear the distributed backend cache. */ +MLX_API void clear_backends(); + } // namespace mlx::core::distributed #endif diff --git a/Source/Cmlx/include-framework/mlx-error.h b/Source/Cmlx/include-framework/mlx-error.h new file mode 100644 index 000000000..128ae3027 --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-error.h @@ -0,0 +1,51 @@ +#ifdef __cplusplus +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace mlx::core { + +class Error { + public: + // TODO: Use std::atomic when it gets supported in Xcode. + using Message = std::shared_ptr; + + void set_message(Message msg) { + std::atomic_store(&message_, std::move(msg)); + } + + bool valid() const { + auto msg = std::atomic_load(&message_); + return msg.get(); + } + + // If |ptr| is a valid event, copy and return true. + bool store_if_valid(const Error* ptr) { + if (ptr && this != ptr) { + Message msg = std::atomic_load(&ptr->message_); + if (msg) { + set_message(std::move(msg)); + return true; + } + } + return false; + } + + // If current error is valid, throw and clear. + void check() { + auto msg = std::atomic_exchange(&message_, {}); + if (msg) { + throw std::runtime_error(*msg); + } + } + + private: + Message message_; +}; + +} // namespace mlx::core +#endif diff --git a/Source/Cmlx/include-framework/mlx-event.h b/Source/Cmlx/include-framework/mlx-event.h index 2b20196b1..e13ef6f50 100644 --- a/Source/Cmlx/include-framework/mlx-event.h +++ b/Source/Cmlx/include-framework/mlx-event.h @@ -6,6 +6,7 @@ #include #include +#include #include namespace mlx::core { @@ -27,6 +28,26 @@ class Event { // Check if the event has been signaled at its current value bool is_signaled() const; + // Associate an error to the event + void set_error(Error& err) { + error().store(&err); + } + + // Get the error associated with the event + Error* load_error() const { + if (!valid()) { + return nullptr; + } + return error().load(); + } + + // Throw and clear the associated error + void check_error() { + if (auto* p = load_error(); p) { + p->check(); + } + } + // Check if the event is valid bool valid() const { return event_ != nullptr; @@ -48,7 +69,18 @@ class Event { return stream_; } + template + auto& cast() const { + return *static_cast(event_.get()); + } + private: + std::atomic& error(); + + const std::atomic& error() const { + return const_cast(this)->error(); + } + // Default constructed stream should never be used // since the event is not yet valid Stream stream_{0, Device::cpu}; diff --git a/Source/Cmlx/include-framework/mlx-export.h b/Source/Cmlx/include-framework/mlx-export.h index 99d2ea279..07bf03911 100644 --- a/Source/Cmlx/include-framework/mlx-export.h +++ b/Source/Cmlx/include-framework/mlx-export.h @@ -51,17 +51,20 @@ struct FunctionExporter; MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API FunctionExporter exporter( const std::string& path, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); /** * Export a function to a file. @@ -70,20 +73,23 @@ MLX_API void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); MLX_API void export_function( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless = false); + bool shapeless = false, + const std::string& metadata = ""); struct ImportedFunction; diff --git a/Source/Cmlx/include-framework/mlx-export_impl.h b/Source/Cmlx/include-framework/mlx-export_impl.h index f5de6df3b..be82f7e7e 100644 --- a/Source/Cmlx/include-framework/mlx-export_impl.h +++ b/Source/Cmlx/include-framework/mlx-export_impl.h @@ -28,17 +28,20 @@ struct MLX_API FunctionExporter { friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&, const Kwargs&)>&, - bool shapeless); + bool shapeless, + const std::string& metadata); friend MLX_API FunctionExporter exporter( const ExportCallback&, @@ -58,7 +61,8 @@ struct MLX_API FunctionExporter { FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless); + bool shapeless, + std::string metadata); FunctionExporter( const ExportCallback& callback, @@ -78,6 +82,7 @@ struct MLX_API FunctionExporter { int count{0}; bool closed{false}; std::shared_ptr ftable; + std::string metadata_; }; struct MLX_API ImportedFunction { @@ -89,12 +94,18 @@ struct MLX_API ImportedFunction { std::vector operator()(const Kwargs& kwargs) const; std::vector operator()(const Args& args, const Kwargs& kwargs) const; + // The metadata saved with the function when it was exported. + const std::string& metadata() const { + return metadata_; + } + private: ImportedFunction(const std::string& file); friend MLX_API ImportedFunction import_function(const std::string&); ImportedFunction(); std::shared_ptr ftable; + std::string metadata_; }; } // namespace mlx::core diff --git a/Source/Cmlx/include-framework/mlx-fast.h b/Source/Cmlx/include-framework/mlx-fast.h index dde88191c..345f9fd6d 100644 --- a/Source/Cmlx/include-framework/mlx-fast.h +++ b/Source/Cmlx/include-framework/mlx-fast.h @@ -7,6 +7,7 @@ #include #include +#include #include namespace mlx::core::fast { @@ -24,6 +25,10 @@ MLX_API array layer_norm( float eps, StreamOrDevice s = {}); +/** Fused cross entropy with class indices as targets. */ +MLX_API array +cross_entropy(const array& logits, const array& targets, StreamOrDevice s = {}); + MLX_API array rope( const array& x, int dims, @@ -53,6 +58,7 @@ MLX_API array scaled_dot_product_attention( const std::string& mask_mode = "", std::optional mask_arr = {}, const std::optional& sinks = {}, + bool force_fused = false, StreamOrDevice s = {}); using TemplateArg = std::variant; @@ -76,7 +82,8 @@ MLX_API CustomKernelFunction metal_kernel( const std::string& source, const std::string& header = "", bool ensure_row_contiguous = true, - bool atomic_outputs = false); + bool atomic_outputs = false, + const CompileOptions& compile_options = {}); MLX_API CustomKernelFunction cuda_kernel( const std::string& name, diff --git a/Source/Cmlx/include-framework/mlx-fft.h b/Source/Cmlx/include-framework/mlx-fft.h index d726997ec..628113909 100644 --- a/Source/Cmlx/include-framework/mlx-fft.h +++ b/Source/Cmlx/include-framework/mlx-fft.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -12,40 +13,74 @@ namespace mlx::core::fft { +enum class FFTNorm { + Backward, + Ortho, + Forward, +}; + /** Compute the n-dimensional Fourier Transform. */ MLX_API array fftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}); +MLX_API array fftn( + const array& a, + const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); MLX_API array -fftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); -MLX_API array fftn(const array& a, StreamOrDevice s = {}); +fftn(const array& a, FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); /** Compute the n-dimensional inverse Fourier Transform. */ MLX_API array ifftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}); +MLX_API array ifftn( + const array& a, + const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); MLX_API array -ifftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); -MLX_API array ifftn(const array& a, StreamOrDevice s = {}); +ifftn(const array& a, FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); /** Compute the one-dimensional Fourier Transform. */ -inline array fft(const array& a, int n, int axis, StreamOrDevice s = {}) { - return fftn(a, {n}, {axis}, s); +inline array fft( + const array& a, + int n, + int axis, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return fftn(a, {n}, {axis}, norm, s); } -inline array fft(const array& a, int axis = -1, StreamOrDevice s = {}) { - return fftn(a, {axis}, s); +inline array fft( + const array& a, + int axis = -1, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return fftn(a, {axis}, norm, s); } /** Compute the one-dimensional inverse Fourier Transform. */ -inline array ifft(const array& a, int n, int axis, StreamOrDevice s = {}) { - return ifftn(a, {n}, {axis}, s); +inline array ifft( + const array& a, + int n, + int axis, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return ifftn(a, {n}, {axis}, norm, s); } -inline array ifft(const array& a, int axis = -1, StreamOrDevice s = {}) { - return ifftn(a, {axis}, s); +inline array ifft( + const array& a, + int axis = -1, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return ifftn(a, {axis}, norm, s); } /** Compute the two-dimensional Fourier Transform. */ @@ -53,14 +88,16 @@ inline array fft2( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return fftn(a, n, axes, s); + return fftn(a, n, axes, norm, s); } inline array fft2( const array& a, const std::vector& axes = {-2, -1}, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return fftn(a, axes, s); + return fftn(a, axes, norm, s); } /** Compute the two-dimensional inverse Fourier Transform. */ @@ -68,14 +105,16 @@ inline array ifft2( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return ifftn(a, n, axes, s); + return ifftn(a, n, axes, norm, s); } inline array ifft2( const array& a, const std::vector& axes = {-2, -1}, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return ifftn(a, axes, s); + return ifftn(a, axes, norm, s); } /** Compute the n-dimensional Fourier Transform on a real input. */ @@ -83,34 +122,62 @@ MLX_API array rfftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}); +MLX_API array rfftn( + const array& a, + const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); MLX_API array -rfftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); -MLX_API array rfftn(const array& a, StreamOrDevice s = {}); +rfftn(const array& a, FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); /** Compute the n-dimensional inverse of `rfftn`. */ MLX_API array irfftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}); +MLX_API array irfftn( + const array& a, + const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); MLX_API array -irfftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); -MLX_API array irfftn(const array& a, StreamOrDevice s = {}); +irfftn(const array& a, FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}); /** Compute the one-dimensional Fourier Transform on a real input. */ -inline array rfft(const array& a, int n, int axis, StreamOrDevice s = {}) { - return rfftn(a, {n}, {axis}, s); +inline array rfft( + const array& a, + int n, + int axis, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return rfftn(a, {n}, {axis}, norm, s); } -inline array rfft(const array& a, int axis = -1, StreamOrDevice s = {}) { - return rfftn(a, {axis}, s); +inline array rfft( + const array& a, + int axis = -1, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return rfftn(a, {axis}, norm, s); } /** Compute the one-dimensional inverse of `rfft`. */ -inline array irfft(const array& a, int n, int axis, StreamOrDevice s = {}) { - return irfftn(a, {n}, {axis}, s); +inline array irfft( + const array& a, + int n, + int axis, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return irfftn(a, {n}, {axis}, norm, s); } -inline array irfft(const array& a, int axis = -1, StreamOrDevice s = {}) { - return irfftn(a, {axis}, s); +inline array irfft( + const array& a, + int axis = -1, + FFTNorm norm = FFTNorm::Backward, + StreamOrDevice s = {}) { + return irfftn(a, {axis}, norm, s); } /** Compute the two-dimensional Fourier Transform on a real input. */ @@ -118,14 +185,16 @@ inline array rfft2( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return rfftn(a, n, axes, s); + return rfftn(a, n, axes, norm, s); } inline array rfft2( const array& a, const std::vector& axes = {-2, -1}, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return rfftn(a, axes, s); + return rfftn(a, axes, norm, s); } /** Compute the two-dimensional inverse of `rfft2`. */ @@ -133,15 +202,24 @@ inline array irfft2( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return irfftn(a, n, axes, s); + return irfftn(a, n, axes, norm, s); } inline array irfft2( const array& a, const std::vector& axes = {-2, -1}, + FFTNorm norm = FFTNorm::Backward, StreamOrDevice s = {}) { - return irfftn(a, axes, s); + return irfftn(a, axes, norm, s); } +/** Compute the discrete Fourier Transform sample frequencies. */ +MLX_API array fftfreq(int n, double d = 1.0, StreamOrDevice s = {}); + +/** Compute the discrete Fourier Transform sample frequencies for + * `rfft`/`irfft`. */ +MLX_API array rfftfreq(int n, double d = 1.0, StreamOrDevice s = {}); + /** Shift the zero-frequency component to the center of the spectrum. */ MLX_API array fftshift(const array& a, StreamOrDevice s = {}); diff --git a/Source/Cmlx/include-framework/mlx-graph_utils.h b/Source/Cmlx/include-framework/mlx-graph_utils.h new file mode 100644 index 000000000..31cbf92c1 --- /dev/null +++ b/Source/Cmlx/include-framework/mlx-graph_utils.h @@ -0,0 +1,69 @@ +#ifdef __cplusplus +// Copyright © 2023 Apple Inc. + +#pragma once + +#include + +#include +#include + +namespace mlx::core { + +struct MLX_API NodeNamer { + std::unordered_map names; + + const std::string& get_name(const array& x); + void set_name(const array& x, std::string n); +}; + +MLX_API void print_graph( + std::ostream& os, + NodeNamer namer, + const std::vector& outputs); + +inline void print_graph(std::ostream& os, const std::vector& outputs) { + print_graph(os, NodeNamer{}, outputs); +} + +template > +inline void print_graph(std::ostream& os, Arrays&&... outputs) { + print_graph( + os, NodeNamer{}, std::vector{std::forward(outputs)...}); +} + +template > +inline void +print_graph(std::ostream& os, NodeNamer namer, Arrays&&... outputs) { + print_graph( + os, + std::move(namer), + std::vector{std::forward(outputs)...}); +} + +MLX_API void export_to_dot( + std::ostream& os, + NodeNamer namer, + const std::vector& outputs); + +inline void export_to_dot(std::ostream& os, const std::vector& outputs) { + export_to_dot(os, NodeNamer{}, outputs); +} + +template > +inline void export_to_dot(std::ostream& os, Arrays&&... outputs) { + export_to_dot( + os, NodeNamer{}, std::vector{std::forward(outputs)...}); +} + +template > +inline void +export_to_dot(std::ostream& os, NodeNamer namer, Arrays&&... outputs) { + export_to_dot( + os, + std::move(namer), + std::vector{std::forward(outputs)...}); +} + +} // namespace mlx::core +#endif diff --git a/Source/Cmlx/include-framework/mlx-io-load.h b/Source/Cmlx/include-framework/mlx-io-load.h index 823f7f42a..fc9a942ad 100644 --- a/Source/Cmlx/include-framework/mlx-io-load.h +++ b/Source/Cmlx/include-framework/mlx-io-load.h @@ -84,7 +84,9 @@ class ParallelFileReader : public Reader { void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { if (way == std::ios_base::beg) { - lseek(fd_, off, 0); + lseek(fd_, off, SEEK_SET); + } else if (way == std::ios_base::end) { + lseek(fd_, off, SEEK_END); } else { lseek(fd_, off, SEEK_CUR); } @@ -144,7 +146,9 @@ class FileWriter : public Writer { void seek(int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { if (way == std::ios_base::beg) { - lseek(fd_, off, 0); + lseek(fd_, off, SEEK_SET); + } else if (way == std::ios_base::end) { + lseek(fd_, off, SEEK_END); } else { lseek(fd_, off, SEEK_CUR); } diff --git a/Source/Cmlx/include-framework/mlx-linalg.h b/Source/Cmlx/include-framework/mlx-linalg.h index 38879f39c..43db9c794 100644 --- a/Source/Cmlx/include-framework/mlx-linalg.h +++ b/Source/Cmlx/include-framework/mlx-linalg.h @@ -113,5 +113,9 @@ eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {}); MLX_API std::pair eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {}); +MLX_API array det(const array& a, StreamOrDevice s = {}); + +MLX_API std::pair slogdet(const array& a, StreamOrDevice s = {}); + } // namespace mlx::core::linalg #endif diff --git a/Source/Cmlx/include-framework/mlx-ops.h b/Source/Cmlx/include-framework/mlx-ops.h index 37500fd5c..96e3ecfe2 100644 --- a/Source/Cmlx/include-framework/mlx-ops.h +++ b/Source/Cmlx/include-framework/mlx-ops.h @@ -39,13 +39,25 @@ MLX_API array arange(int start, int stop, int step, StreamOrDevice s = {}); MLX_API array arange(int start, int stop, StreamOrDevice s = {}); MLX_API array arange(int stop, StreamOrDevice s = {}); -/** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */ +/** + * A 1D array of `num` evenly spaced numbers in the range `[start, stop]`, or + * in the half-open range `[start, stop)` when `endpoint` is false. + */ MLX_API array linspace( double start, double stop, - int num = 50, + int num, + bool endpoint, Dtype dtype = float32, StreamOrDevice s = {}); +inline array linspace( + double start, + double stop, + int num = 50, + Dtype dtype = float32, + StreamOrDevice s = {}) { + return linspace(start, stop, num, true, dtype, s); +} /** Convert an array to the given data type. */ MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {}); @@ -90,14 +102,18 @@ MLX_API array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array zeros(const Shape& shape, StreamOrDevice s = {}) { return zeros(shape, float32, s); } +/** Create an array of zeros with the shape of `a`. */ MLX_API array zeros_like(const array& a, StreamOrDevice s = {}); +MLX_API array zeros_like(const array& a, Dtype dtype, StreamOrDevice s = {}); /** Fill an array of the given shape with ones. */ MLX_API array ones(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array ones(const Shape& shape, StreamOrDevice s = {}) { return ones(shape, float32, s); } +/** Create an array of ones with the shape of `a`. */ MLX_API array ones_like(const array& a, StreamOrDevice s = {}); +MLX_API array ones_like(const array& a, Dtype dtype, StreamOrDevice s = {}); /** Fill an array of the given shape (n,m) with ones in the specified diagonal * k, and zeros everywhere else. */ @@ -172,6 +188,16 @@ MLX_API array expand_dims( /** Add a singleton dimension at the given axis. */ MLX_API array expand_dims(const array& a, int axis, StreamOrDevice s = {}); +/** Reverse the order of the elements along the given axes. */ +MLX_API array +flip(const array& a, const std::vector& axes, StreamOrDevice s = {}); + +/** Reverse the order of the elements along the given axis. */ +MLX_API array flip(const array& a, int axis, StreamOrDevice s = {}); + +/** Reverse the order of the elements along all axes. */ +MLX_API array flip(const array& a, StreamOrDevice s = {}); + /** Slice an array. */ MLX_API array slice( const array& a, @@ -225,6 +251,78 @@ MLX_API array slice_update( std::vector axes, StreamOrDevice s = {}); +/** Slice update and add updates to given slice. */ +MLX_API array slice_update_add( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s = {}); + +/** Slice update and add updates to given slice with stride 1 in each dimension. + */ +MLX_API array slice_update_add( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s = {}); + +/** Slice update and prod updates to given slice. */ +MLX_API array slice_update_prod( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s = {}); + +/** Slice update and prod updates to given slice with stride 1 in each + * dimension. */ +MLX_API array slice_update_prod( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s = {}); + +/** Slice update and max updates to given slice. */ +MLX_API array slice_update_max( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s = {}); + +/** Slice update and max updates to given slice with stride 1 in each dimension. + */ +MLX_API array slice_update_max( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s = {}); + +/** Slice update and min updates to given slice. */ +MLX_API array slice_update_min( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s = {}); + +/** Slice update and min updates to given slice with stride 1 in each dimension. + */ +MLX_API array slice_update_min( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s = {}); + /** Split an array into sub-arrays along a given axis. */ MLX_API std::vector split(const array& a, int num_splits, int axis, StreamOrDevice s = {}); @@ -235,6 +333,11 @@ split(const array& a, const Shape& indices, int axis, StreamOrDevice s = {}); MLX_API std::vector split(const array& a, const Shape& indices, StreamOrDevice s = {}); +/** Split an array into a sequence of arrays along an axis, removing it. */ +MLX_API std::vector +unstack(const array& a, int axis, StreamOrDevice s = {}); +MLX_API std::vector unstack(const array& a, StreamOrDevice s = {}); + /** A vector of coordinate arrays from coordinate vectors. */ MLX_API std::vector meshgrid( const std::vector& arrays, @@ -532,6 +635,20 @@ sum(const array& a, MLX_API array sum(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); +/** Count the number of non-zero elements in an array. */ +MLX_API array +count_nonzero(const array& a, bool keepdims = false, StreamOrDevice s = {}); +MLX_API array count_nonzero( + const array& a, + int axis, + bool keepdims = false, + StreamOrDevice s = {}); +MLX_API array count_nonzero( + const array& a, + const std::vector& axes, + bool keepdims = false, + StreamOrDevice s = {}); + /** Computes the mean of the elements of an array. */ MLX_API array mean(const array& a, bool keepdims, StreamOrDevice s = {}); inline array mean(const array& a, StreamOrDevice s = {}) { @@ -699,16 +816,28 @@ inline array argmax(const array& a, StreamOrDevice s = {}) { MLX_API array argmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); -/** Returns a sorted copy of the flattened array. */ +/** + * Returns a sorted copy of the flattened array. + * The sort is stable and NaN values are placed at the end. + */ MLX_API array sort(const array& a, StreamOrDevice s = {}); -/** Returns a sorted copy of the array along a given axis. */ +/** + * Returns a sorted copy of the array along a given axis. + * The sort is stable and NaN values are placed at the end. + */ MLX_API array sort(const array& a, int axis, StreamOrDevice s = {}); -/** Returns indices that sort the flattened array. */ +/** + * Returns indices that sort the flattened array. + * The sort is stable and NaN values are placed at the end. + */ MLX_API array argsort(const array& a, StreamOrDevice s = {}); -/** Returns indices that sort the array along a given axis. */ +/** + * Returns indices that sort the array along a given axis. + * The sort is stable and NaN values are placed at the end. + */ MLX_API array argsort(const array& a, int axis, StreamOrDevice s = {}); /** @@ -737,6 +866,13 @@ MLX_API array argpartition(const array& a, int kth, StreamOrDevice s = {}); MLX_API array argpartition(const array& a, int kth, int axis, StreamOrDevice s = {}); +/** Find the indices of `values` in `sorted_sequence`. */ +MLX_API array searchsorted( + const array& sorted_sequence, + const array& values, + const std::string& side = "left", + StreamOrDevice s = {}); + /** Returns topk elements of the flattened array. */ MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); @@ -781,6 +917,9 @@ MLX_API array logsumexp( /** Absolute value of elements in an array. */ MLX_API array abs(const array& a, StreamOrDevice s = {}); +/** Unary plus — return a copy of the array unchanged. */ +MLX_API array positive(const array& a, StreamOrDevice s = {}); + /** Negate an array. */ MLX_API array negative(const array& a, StreamOrDevice s = {}); MLX_API array operator-(const array& a); @@ -800,6 +939,10 @@ MLX_API array operator&&(const array& a, const array& b); MLX_API array logical_or(const array& a, const array& b, StreamOrDevice s = {}); MLX_API array operator||(const array& a, const array& b); +/** Logical exclusive or of two arrays */ +MLX_API array +logical_xor(const array& a, const array& b, StreamOrDevice s = {}); + /** The reciprocal (1/x) of the elements in an array. */ MLX_API array reciprocal(const array& a, StreamOrDevice s = {}); @@ -877,6 +1020,9 @@ MLX_API array floor(const array& a, StreamOrDevice s = {}); /** Ceil the element of an array. **/ MLX_API array ceil(const array& a, StreamOrDevice s = {}); +/** Truncate the elements of an array towards zero. **/ +MLX_API array trunc(const array& a, StreamOrDevice s = {}); + /** Square the elements of an array. */ MLX_API array square(const array& a, StreamOrDevice s = {}); @@ -1231,7 +1377,12 @@ MLX_API array cumsum( const array& a, bool reverse = false, bool inclusive = true, + std::optional dtype = std::nullopt, StreamOrDevice s = {}); +inline array +cumsum(const array& a, bool reverse, bool inclusive, StreamOrDevice s) { + return cumsum(a, reverse, inclusive, std::nullopt, s); +} /** Cumulative sum of an array along the given axis. */ MLX_API array cumsum( @@ -1239,14 +1390,28 @@ MLX_API array cumsum( int axis, bool reverse = false, bool inclusive = true, + std::optional dtype = std::nullopt, StreamOrDevice s = {}); +inline array cumsum( + const array& a, + int axis, + bool reverse, + bool inclusive, + StreamOrDevice s) { + return cumsum(a, axis, reverse, inclusive, std::nullopt, s); +} /** Cumulative product of an array. */ MLX_API array cumprod( const array& a, bool reverse = false, bool inclusive = true, + std::optional dtype = std::nullopt, StreamOrDevice s = {}); +inline array +cumprod(const array& a, bool reverse, bool inclusive, StreamOrDevice s) { + return cumprod(a, reverse, inclusive, std::nullopt, s); +} /** Cumulative product of an array along the given axis. */ MLX_API array cumprod( @@ -1254,7 +1419,16 @@ MLX_API array cumprod( int axis, bool reverse = false, bool inclusive = true, + std::optional dtype = std::nullopt, StreamOrDevice s = {}); +inline array cumprod( + const array& a, + int axis, + bool reverse, + bool inclusive, + StreamOrDevice s) { + return cumprod(a, axis, reverse, inclusive, std::nullopt, s); +} /** Cumulative max of an array. */ MLX_API array cummax( @@ -1286,6 +1460,10 @@ MLX_API array cummin( bool inclusive = true, StreamOrDevice s = {}); +/** The n-th discrete difference along the given axis. */ +MLX_API array +diff(const array& a, int n = 1, int axis = -1, StreamOrDevice s = {}); + /** General convolution with a filter */ MLX_API array conv_general( array input, @@ -1388,10 +1566,10 @@ MLX_API array conv_transpose3d( /** Quantized matmul multiplies x with a quantized matrix w*/ MLX_API array quantized_matmul( - array x, - array w, - array scales, - std::optional biases = std::nullopt, + const array& x, + const array& w, + const array& scales, + const std::optional& biases = std::nullopt, bool transpose = true, std::optional group_size = std::nullopt, std::optional bits = std::nullopt, @@ -1420,15 +1598,15 @@ MLX_API array dequantize( StreamOrDevice s = {}); MLX_API array qqmm( - array x, // input activations - array w, // maybe quantized weights - const std::optional w_scales = std::nullopt, // optional scales if w - // is quantized + const array& x, // input activations + const array& w, // maybe quantized weights + const std::optional& w_scales = std::nullopt, // optional scales if w + // is quantized std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "nvfp4", - const std::optional global_scale_x = std::nullopt, - const std::optional global_scale_w = std::nullopt, + const std::optional& global_scale_x = std::nullopt, + const std::optional& global_scale_w = std::nullopt, StreamOrDevice s = {}); /** Convert an E4M3 float8 to the given floating point dtype. */ @@ -1452,6 +1630,20 @@ MLX_API array gather_qmm( bool sorted_indices = false, StreamOrDevice s = {}); +MLX_API array gather_qqmm( + const array& x, + const array& w, + const std::optional& scales_w = std::nullopt, + const std::optional& lhs_indices = std::nullopt, + const std::optional& rhs_indices = std::nullopt, + std::optional group_size = std::nullopt, + std::optional bits = std::nullopt, + const std::string& mode = "nvfp4", + const std::optional& global_scale_x = std::nullopt, + const std::optional& global_scale_w = std::nullopt, + bool sorted_indices = false, + StreamOrDevice s = {}); + /** Returns a contraction of a and b over multiple dimensions. */ MLX_API array tensordot( const array& a, @@ -1472,6 +1664,10 @@ MLX_API array outer(const array& a, const array& b, StreamOrDevice s = {}); /** Compute the inner product of two vectors. */ MLX_API array inner(const array& a, const array& b, StreamOrDevice s = {}); +/** Compute a vector dot product along an axis. */ +MLX_API array +vecdot(const array& a, const array& b, int axis = -1, StreamOrDevice s = {}); + /** Compute D = beta * C + alpha * (A @ B) */ MLX_API array addmm( array c, diff --git a/Source/Cmlx/include-framework/mlx-primitives.h b/Source/Cmlx/include-framework/mlx-primitives.h index dcb719e05..33d2bb094 100644 --- a/Source/Cmlx/include-framework/mlx-primitives.h +++ b/Source/Cmlx/include-framework/mlx-primitives.h @@ -201,6 +201,7 @@ class AddMM : public UnaryPrimitive { DEFINE_NAME(AddMM) bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; std::pair state() const { return {alpha_, beta_}; }; @@ -543,6 +544,7 @@ class GatherMM : public UnaryPrimitive { DEFINE_NAME(GatherMM) bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; auto state() const { return std::make_pair(left_sorted_, right_sorted_); } @@ -698,6 +700,7 @@ class Conjugate : public UnaryPrimitive { void eval_gpu(const std::vector& inputs, array& out) override; DEFINE_VMAP() + DEFINE_GRADS() DEFINE_NAME(Conjugate) DEFINE_DEFAULT_IS_EQUIVALENT() DEFINE_INPUT_OUTPUT_SHAPE() @@ -973,9 +976,9 @@ class Equal : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS() - DEFINE_DEFAULT_IS_EQUIVALENT() DEFINE_INPUT_OUTPUT_SHAPE() + bool is_equivalent(const Primitive& other) const override; const char* name() const override { if (equal_nan_) { return "NaNEqual"; @@ -1323,9 +1326,9 @@ class Log : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS() - DEFINE_DEFAULT_IS_EQUIVALENT() DEFINE_INPUT_OUTPUT_SHAPE() + bool is_equivalent(const Primitive& other) const override; Base state() const { return base_; }; @@ -1699,6 +1702,7 @@ class GatherQMM : public UnaryPrimitive { DEFINE_GRADS() DEFINE_NAME(GatherQMM) bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; auto state() const { return std::make_tuple( group_size_, bits_, mode_, transpose_, left_sorted_, right_sorted_); @@ -1713,6 +1717,41 @@ class GatherQMM : public UnaryPrimitive { bool right_sorted_; }; +class GatherQQMM : public UnaryPrimitive { + public: + explicit GatherQQMM( + Stream stream, + int group_size, + int bits, + QuantizationMode mode, + bool left_sorted = false, + bool right_sorted = false) + : UnaryPrimitive(stream), + group_size_(group_size), + bits_(bits), + mode_(mode), + left_sorted_(left_sorted), + right_sorted_(right_sorted) {} + + void eval_cpu(const std::vector& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_NAME(GatherQQMM) + bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; + auto state() const { + return std::make_tuple( + group_size_, bits_, mode_, left_sorted_, right_sorted_); + } + + private: + int group_size_; + int bits_; + QuantizationMode mode_; + bool left_sorted_; + bool right_sorted_; +}; + class RandomBits : public UnaryPrimitive { public: explicit RandomBits(Stream stream, const Shape& shape, int width) @@ -2056,12 +2095,16 @@ class Slice : public UnaryPrimitive { class SliceUpdate : public UnaryPrimitive { public: + enum ReduceType { Max, Min, Sum, Prod, None }; + explicit SliceUpdate( Stream stream, + ReduceType reduce_type, const Shape& start_indices, const Shape& end_indices, const Shape& strides) : UnaryPrimitive(stream), + reduce_type_(reduce_type), start_indices_(start_indices), end_indices_(end_indices), strides_(strides) {} @@ -2071,14 +2114,32 @@ class SliceUpdate : public UnaryPrimitive { DEFINE_VMAP() DEFINE_GRADS() - DEFINE_NAME(SliceUpdate) + + const char* name() const override { + switch (reduce_type_) { + case Sum: + return "SliceUpdate Sum"; + case Prod: + return "SliceUpdate Prod"; + case Min: + return "SliceUpdate Min"; + case Max: + return "SliceUpdate Max"; + case None: + return "SliceUpdate"; + } + return ""; + } + bool is_equivalent(const Primitive& other) const override; DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { - return std::make_tuple(start_indices_, end_indices_, strides_); + return std::make_tuple( + reduce_type_, start_indices_, end_indices_, strides_); } private: + ReduceType reduce_type_; Shape start_indices_; Shape end_indices_; Shape strides_; @@ -2151,6 +2212,27 @@ class Softmax : public UnaryPrimitive { bool precise_; }; +class SearchSorted : public UnaryPrimitive { + public: + explicit SearchSorted(Stream stream, bool right) + : UnaryPrimitive(stream), right_(right) {} + + void eval_cpu(const std::vector& inputs, array& out) override; + void eval_gpu(const std::vector& inputs, array& out) override; + + DEFINE_VMAP() + DEFINE_GRADS() + DEFINE_NAME(SearchSorted) + bool is_equivalent(const Primitive& other) const override; + std::vector output_shapes(const std::vector& inputs) override; + auto state() const { + return right_; + } + + private: + bool right_; +}; + class Sort : public UnaryPrimitive { public: explicit Sort(Stream stream, int axis) diff --git a/Source/Cmlx/include-framework/mlx-random.h b/Source/Cmlx/include-framework/mlx-random.h index 36db67574..a1ef40911 100644 --- a/Source/Cmlx/include-framework/mlx-random.h +++ b/Source/Cmlx/include-framework/mlx-random.h @@ -20,20 +20,20 @@ class MLX_API KeySequence { void seed(uint64_t seed); array next(); - // static default + // Each thread has its own random key to avoid race condition. static KeySequence& default_() { - static KeySequence ks(get_current_time_seed()); + static auto time_seed = []() { + auto now = std::chrono::system_clock::now(); + return std::chrono::duration_cast( + now.time_since_epoch()) + .count(); + }(); + static thread_local KeySequence ks(time_seed); return ks; } private: array key_; - static uint64_t get_current_time_seed() { - auto now = std::chrono::system_clock::now(); - return std::chrono::duration_cast( - now.time_since_epoch()) - .count(); - } }; /** Get a PRNG key from a seed. */ diff --git a/Source/Cmlx/include-framework/mlx-scheduler.h b/Source/Cmlx/include-framework/mlx-scheduler.h index 3c2896d2b..a4953f847 100644 --- a/Source/Cmlx/include-framework/mlx-scheduler.h +++ b/Source/Cmlx/include-framework/mlx-scheduler.h @@ -4,75 +4,24 @@ #pragma once #include -#include #include -#include +#include #include #include #include #include #include +#include namespace mlx::core::scheduler { -struct StreamThread { - std::mutex mtx; - std::queue> q; - std::condition_variable cond; - bool stop; - std::thread thread; - - StreamThread() : stop(false), thread(&StreamThread::thread_fn, this) {} - - ~StreamThread() { - { - std::lock_guard lk(mtx); - stop = true; - } - cond.notify_one(); - thread.join(); - } - - void thread_fn() { - while (true) { - std::function task; - { - std::unique_lock lk(mtx); - cond.wait(lk, [this] { return !this->q.empty() || this->stop; }); - if (q.empty() && stop) { - return; - } - task = std::move(q.front()); - q.pop(); - } - - task(); - } - } +class StreamThread; - template - void enqueue(F&& f) { - { - std::lock_guard lk(mtx); - if (stop) { - throw std::runtime_error( - "Cannot enqueue work after stream is stopped."); - } - q.emplace(std::forward(f)); - } - cond.notify_one(); - } -}; - -class Scheduler { +class MLX_API Scheduler { public: - Scheduler() : n_active_tasks_(0) { - if (is_available(Device::gpu)) { - default_streams_.insert({Device::gpu, new_stream(Device::gpu)}); - } - default_streams_.insert({Device::cpu, new_stream(Device::cpu)}); - } + Scheduler(); + ~Scheduler(); // Not copyable or moveable Scheduler(const Scheduler&) = delete; @@ -80,33 +29,10 @@ class Scheduler { Scheduler& operator=(const Scheduler&) = delete; Scheduler& operator=(Scheduler&&) = delete; - Stream new_stream(const Device& d) { - streams_.emplace_back(streams_.size(), d); - if (d == Device::gpu) { - threads_.push_back(nullptr); - gpu::new_stream(streams_.back()); - } else { - threads_.push_back(new StreamThread{}); - } - return streams_.back(); - } - - template - void enqueue(const Stream& stream, F&& f); - - Stream get_default_stream(const Device& d) const { - return default_streams_.at(d.type); - } - Stream get_stream(int index) const { - return streams_.at(index); - } - std::vector get_streams() const { - return streams_; - } - - void set_default_stream(const Stream& s) { - default_streams_.at(s.device.type) = s; - } + void enqueue(Stream s, std::function task); + void wait_event(Stream s, Event event, std::function task); + void signal_event(Stream s, Event event, std::function task); + void check_error(Stream s); void notify_new_task(const Stream& stream) { { @@ -138,40 +64,39 @@ class Scheduler { } } - ~Scheduler() { - for (auto s : streams_) { - try { - synchronize(s); - } catch (const std::runtime_error&) { - // ignore errors if synch fails - } - } - for (auto t : threads_) { - if (t != nullptr) { - delete t; - } - } - } - private: - int n_active_tasks_; - std::vector threads_; - std::vector streams_; - std::unordered_map default_streams_; + friend Stream mlx::core::new_stream(Device d); + + StreamThread& get_thread(Stream s); + + int n_active_tasks_{0}; + std::unordered_map> threads_; + std::shared_mutex threads_mtx_; std::condition_variable completion_cv; std::mutex mtx; }; +MLX_API Scheduler& scheduler(); + template -void Scheduler::enqueue(const Stream& stream, F&& f) { - threads_[stream.index]->enqueue(std::forward(f)); +inline void enqueue(Stream s, F&& f) { + scheduler().enqueue(s, std::forward(f)); } -MLX_API Scheduler& scheduler(); +// Like enqueue but the task is used for processing the passed event. +template +inline void wait_event(Stream s, Event event, F&& f) { + scheduler().wait_event(s, std::move(event), std::forward(f)); +} template -void enqueue(const Stream& stream, F&& f) { - scheduler().enqueue(stream, std::forward(f)); +inline void signal_event(Stream s, Event event, F&& f) { + scheduler().signal_event(s, std::move(event), std::forward(f)); +} + +// Throw and clear the error stored in the stream, if any. +inline void check_error(Stream s) { + scheduler().check_error(s); } inline int n_active_tasks() { diff --git a/Source/Cmlx/include-framework/mlx-stream.h b/Source/Cmlx/include-framework/mlx-stream.h index 3ff0b398f..9b9bae946 100644 --- a/Source/Cmlx/include-framework/mlx-stream.h +++ b/Source/Cmlx/include-framework/mlx-stream.h @@ -3,6 +3,7 @@ #pragma once +#include #include #include @@ -14,30 +15,38 @@ struct MLX_API Stream { int index; Device device; explicit Stream(int index, Device device) : index(index), device(device) {} + + // TODO: Use default three-way comparison when it gets supported in XCode. + bool operator==(const Stream&) const = default; + bool operator<(const Stream& rhs) const { + return std::tie(device, index) < std::tie(rhs.device, rhs.index); + } +}; + +struct MLX_API ThreadLocalStream : public Stream { + using Stream::Stream; }; -/** Get the default stream for the given device. */ +/** Get the default stream of current thread for the given device. */ MLX_API Stream default_stream(Device d); -/** Make the stream the default for its device. */ +/** Make the stream the default for its device on current thread. */ MLX_API void set_default_stream(Stream s); /** Make a new stream on the given device. */ MLX_API Stream new_stream(Device d); -/** Get the stream with the given index. */ -MLX_API Stream get_stream(int index); +/** Make a new stream that can be used in any thread. */ +MLX_API Stream new_thread_unsafe_stream(Device d); -/** Get all available streams. */ -MLX_API std::vector get_streams(); +/** Make a new stream that will be unique per thread. */ +MLX_API ThreadLocalStream new_thread_local_stream(Device d); -inline bool operator==(const Stream& lhs, const Stream& rhs) { - return lhs.index == rhs.index; -} +/** Get the stream for current thread from ThreadLocalStream. */ +MLX_API Stream stream_from_thread_local_stream(ThreadLocalStream tls); -inline bool operator!=(const Stream& lhs, const Stream& rhs) { - return !(lhs == rhs); -} +/** Get all available streams. */ +MLX_API std::vector get_streams(); /* Synchronize with the default stream. */ MLX_API void synchronize(); @@ -45,5 +54,11 @@ MLX_API void synchronize(); /* Synchronize with the provided stream. */ MLX_API void synchronize(Stream); +/* Synchronize with the stream corresponding to the current thread. */ +MLX_API void synchronize(ThreadLocalStream); + +/* Destroy all streams created in current thread. */ +MLX_API void clear_streams(); + } // namespace mlx::core #endif diff --git a/Source/Cmlx/include-framework/mlx-types-half_types.h b/Source/Cmlx/include-framework/mlx-types-half_types.h index f9a8622cb..b6869f528 100644 --- a/Source/Cmlx/include-framework/mlx-types-half_types.h +++ b/Source/Cmlx/include-framework/mlx-types-half_types.h @@ -37,10 +37,11 @@ typedef struct _MLX_BFloat16 bfloat16_t; #endif // __ARM_FEATURE_BF16 -#ifdef ADD_HALF_BINOPS namespace mlx::core { - // clang-format off + +#ifdef ADD_HALF_BINOPS + #define fp16_bf16_binop_helper(__op__, __operator__) \ inline float __operator__(float16_t lhs, bfloat16_t rhs) { \ return static_cast(lhs) __op__ static_cast(rhs); \ @@ -53,8 +54,24 @@ fp16_bf16_binop_helper(+, operator+) fp16_bf16_binop_helper(-, operator-) fp16_bf16_binop_helper(*, operator*) fp16_bf16_binop_helper(/, operator/) -// clang-format on -} // namespace mlx::core #endif + +template +constexpr bool is_complex = false; + +template +constexpr bool is_complex().real())>> = + true; + +template +constexpr bool is_signed_v = std::is_signed_v || + std::is_same_v || std::is_same_v; + +template +constexpr bool is_floating_point_v = std::is_floating_point_v || + std::is_same_v || std::is_same_v; + +// clang-format on +} // namespace mlx::core #endif diff --git a/Source/Cmlx/include-framework/mlx-utils.h b/Source/Cmlx/include-framework/mlx-utils.h index a22cd558c..a0d8478a4 100644 --- a/Source/Cmlx/include-framework/mlx-utils.h +++ b/Source/Cmlx/include-framework/mlx-utils.h @@ -4,6 +4,10 @@ #pragma once #include +#include +#include +#include +#include #include #include @@ -14,7 +18,12 @@ namespace mlx::core { -using StreamOrDevice = std::variant; +using StreamOrDevice = std::variant< + std::monostate, + Stream, + ThreadLocalStream, + Device, + Device::DeviceType>; MLX_API Stream to_stream(StreamOrDevice s); MLX_API Stream to_stream(StreamOrDevice s, Device default_); @@ -39,6 +48,10 @@ struct StreamContext { Stream _stream; }; +struct MLX_API PrintOptions { + int precision{-1}; +}; + struct PrintFormatter { inline void print(std::ostream& os, bool val); inline void print(std::ostream& os, int16_t val); @@ -54,20 +67,25 @@ struct PrintFormatter { inline void print(std::ostream& os, complex64_t val); bool capitalize_bool{false}; + PrintOptions format_options; }; +MLX_API void set_printoptions(PrintOptions options); + MLX_API PrintFormatter& get_global_formatter(); -/** Print the exception and then abort. */ -MLX_API void abort_with_exception(const std::exception& error); +/** Return whether current thread is the first one that called this function. */ +bool is_main_thread(); /** Holds information about floating-point types. */ struct MLX_API finfo { explicit finfo(Dtype dtype); Dtype dtype; + int bits; double min; double max; double eps; + double smallest_normal; }; /** Holds information about integral types. */ @@ -89,6 +107,24 @@ MLX_API Dtype result_type(const std::vector& arrays); MLX_API Shape broadcast_shapes(const Shape& s1, const Shape& s2); +template +inline ShapeElem safe_cast(T dim, std::string_view op = "") { + constexpr int64_t lo = std::numeric_limits::min(); + constexpr int64_t hi = std::numeric_limits::max(); + auto v = static_cast(dim); + if (v < lo || v > hi) { + std::ostringstream msg; + if (!op.empty()) { + msg << "[" << op << "] "; + } + msg << "Shape dimension " << v << " is outside the supported range [" << lo + << ", " << hi + << "]. MLX currently uses 32-bit integers for shape dimensions."; + throw std::overflow_error(msg.str()); + } + return static_cast(v); +} + /** * Returns the axis normalized to be in the range [0, ndim). */ @@ -156,6 +192,21 @@ inline int max_mb_per_buffer(int default_value) { return max_mb_per_buffer_; } +// Per-set residency-set size, as a percentage of the device's recommended +// max working-set size. Controls only how wired memory is distributed across +// residency sets, never how much is wired; see metal::ResidencySets. A value +// <= 0 or >= 100 puts everything in a single set. +inline int residency_set_max_pct() { + static int residency_set_max_pct_ = get_var("MLX_RESIDENCY_SET_MAX_PCT", 5); + return residency_set_max_pct_; +} + +// Log each residency set as it is created. +inline bool residency_debug() { + static bool residency_debug_ = get_var("MLX_RESIDENCY_DEBUG", 0); + return residency_debug_; +} + inline bool metal_fast_synch() { static bool metal_fast_synch = get_var("MLX_METAL_FAST_SYNCH", 0); return metal_fast_synch; diff --git a/Source/Cmlx/include-framework/mlx-version.h b/Source/Cmlx/include-framework/mlx-version.h index 86d15b658..3b3eb3f2a 100644 --- a/Source/Cmlx/include-framework/mlx-version.h +++ b/Source/Cmlx/include-framework/mlx-version.h @@ -6,8 +6,8 @@ #include #define MLX_VERSION_MAJOR 0 -#define MLX_VERSION_MINOR 31 -#define MLX_VERSION_PATCH 1 +#define MLX_VERSION_MINOR 32 +#define MLX_VERSION_PATCH 2 #define MLX_VERSION_NUMERIC \ (100000 * MLX_VERSION_MAJOR + 1000 * MLX_VERSION_MINOR + MLX_VERSION_PATCH) diff --git a/Source/Cmlx/include/mlx/c/compile.h b/Source/Cmlx/include/mlx/c/compile.h index 04567fb3a..68605ce5f 100644 --- a/Source/Cmlx/include/mlx/c/compile.h +++ b/Source/Cmlx/include/mlx/c/compile.h @@ -34,6 +34,13 @@ typedef enum mlx_compile_mode_ { MLX_COMPILE_MODE_NO_FUSE, MLX_COMPILE_MODE_ENABLED } mlx_compile_mode; + +typedef struct mlx_compile_cache_ { + void* ctx; +} mlx_compile_cache; +mlx_compile_cache mlx_compile_cache_new(); +int mlx_compile_cache_free(mlx_compile_cache cache); + int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless); int mlx_detail_compile( mlx_closure* res, @@ -42,8 +49,9 @@ int mlx_detail_compile( bool shapeless, const uint64_t* constants, size_t constants_num); -int mlx_detail_compile_clear_cache(void); -int mlx_detail_compile_erase(uintptr_t fun_id); +int mlx_detail_compile_cache(mlx_compile_cache* res); +int mlx_detail_compile_clear_cache(const mlx_compile_cache cache); +int mlx_detail_compile_erase(const mlx_compile_cache cache, uintptr_t fun_id); int mlx_disable_compile(void); int mlx_enable_compile(void); int mlx_set_compile_mode(mlx_compile_mode mode); diff --git a/Source/Cmlx/include/mlx/c/distributed_group.h b/Source/Cmlx/include/mlx/c/distributed_group.h index 43aa2ae56..bfbaa80eb 100644 --- a/Source/Cmlx/include/mlx/c/distributed_group.h +++ b/Source/Cmlx/include/mlx/c/distributed_group.h @@ -23,6 +23,24 @@ typedef struct mlx_distributed_group_ { void* ctx; } mlx_distributed_group; +/** + * Create an empty group. + */ +mlx_distributed_group mlx_distributed_group_new(void); + +/** + * Free the group. + */ +int mlx_distributed_group_free(mlx_distributed_group group); + +/** + * Initialize distributed. + */ +int mlx_distributed_init( + mlx_distributed_group* res, + bool strict, + const char* bk /* may be null */); + /** * Get the rank. */ @@ -36,21 +54,17 @@ int mlx_distributed_group_size(mlx_distributed_group group); /** * Split the group. */ -mlx_distributed_group -mlx_distributed_group_split(mlx_distributed_group group, int color, int key); +int mlx_distributed_group_split( + mlx_distributed_group* res, + mlx_distributed_group group, + int color, + int key); /** * Check if distributed is available. */ bool mlx_distributed_is_available(const char* bk /* may be null */); -/** - * Initialize distributed. - */ -mlx_distributed_group mlx_distributed_init( - bool strict, - const char* bk /* may be null */); - /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/include/mlx/c/fast.h b/Source/Cmlx/include/mlx/c/fast.h index c825d00e5..914fd2a28 100644 --- a/Source/Cmlx/include/mlx/c/fast.h +++ b/Source/Cmlx/include/mlx/c/fast.h @@ -28,6 +28,12 @@ extern "C" { */ /**@{*/ +int mlx_fast_cross_entropy( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s); + typedef struct mlx_fast_cuda_kernel_config_ { void* ctx; } mlx_fast_cuda_kernel_config; @@ -195,6 +201,7 @@ int mlx_fast_scaled_dot_product_attention( const char* mask_mode, const mlx_array mask_arr /* may be null */, const mlx_array sinks /* may be null */, + bool force_fused, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include/mlx/c/fft.h b/Source/Cmlx/include/mlx/c/fft.h index 779803e9b..7140b601f 100644 --- a/Source/Cmlx/include/mlx/c/fft.h +++ b/Source/Cmlx/include/mlx/c/fft.h @@ -28,11 +28,18 @@ extern "C" { */ /**@{*/ +typedef enum mlx_fft_norm_ { + MLX_FFT_NORM_BACKWARD, + MLX_FFT_NORM_ORTHO, + MLX_FFT_NORM_FORWARD +} mlx_fft_norm; + int mlx_fft_fft( mlx_array* res, const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_fft2( mlx_array* res, @@ -41,7 +48,9 @@ int mlx_fft_fft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); +int mlx_fft_fftfreq(mlx_array* res, int n, double d, const mlx_stream s); int mlx_fft_fftn( mlx_array* res, const mlx_array a, @@ -49,6 +58,7 @@ int mlx_fft_fftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_fftshift( mlx_array* res, @@ -61,6 +71,7 @@ int mlx_fft_ifft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifft2( mlx_array* res, @@ -69,6 +80,7 @@ int mlx_fft_ifft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifftn( mlx_array* res, @@ -77,6 +89,7 @@ int mlx_fft_ifftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_ifftshift( mlx_array* res, @@ -89,6 +102,7 @@ int mlx_fft_irfft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_irfft2( mlx_array* res, @@ -97,6 +111,7 @@ int mlx_fft_irfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_irfftn( mlx_array* res, @@ -105,12 +120,14 @@ int mlx_fft_irfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_rfft( mlx_array* res, const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s); int mlx_fft_rfft2( mlx_array* res, @@ -119,7 +136,9 @@ int mlx_fft_rfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); +int mlx_fft_rfftfreq(mlx_array* res, int n, double d, const mlx_stream s); int mlx_fft_rfftn( mlx_array* res, const mlx_array a, @@ -127,6 +146,7 @@ int mlx_fft_rfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include/mlx/c/graph_utils.h b/Source/Cmlx/include/mlx/c/graph_utils.h new file mode 100644 index 000000000..81eec7a87 --- /dev/null +++ b/Source/Cmlx/include/mlx/c/graph_utils.h @@ -0,0 +1,61 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#ifndef MLX_GRAPH_UTILS_H +#define MLX_GRAPH_UTILS_H + +#include +#include +#include + +#include "mlx/c/array.h" +#include "mlx/c/closure.h" +#include "mlx/c/distributed_group.h" +#include "mlx/c/io_types.h" +#include "mlx/c/map.h" +#include "mlx/c/stream.h" +#include "mlx/c/string.h" +#include "mlx/c/vector.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * \defgroup graph_utils Graph Utils + */ +/**@{*/ + +typedef struct mlx_node_namer_ { + void* ctx; +} mlx_node_namer; + +mlx_node_namer mlx_node_namer_new(); +int mlx_node_namer_free(mlx_node_namer namer); +int mlx_node_namer_set_name( + mlx_node_namer namer, + const mlx_array arr, + const char* name); +int mlx_node_namer_get_name( + const char** name, + mlx_node_namer namer, + const mlx_array arr); + +int mlx_export_to_dot( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs); +int mlx_print_graph( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs); + +/**@}*/ + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/Source/Cmlx/include/mlx/c/io.h b/Source/Cmlx/include/mlx/c/io.h index 6eb205c9a..e1a0c0a7d 100644 --- a/Source/Cmlx/include/mlx/c/io.h +++ b/Source/Cmlx/include/mlx/c/io.h @@ -33,6 +33,9 @@ int mlx_load_reader( mlx_io_reader in_stream, const mlx_stream s); int mlx_load(mlx_array* res, const char* file, const mlx_stream s); + +int mlx_load_gguf(mlx_io_gguf* gguf, const char* file, const mlx_stream s); + int mlx_load_safetensors_reader( mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, @@ -45,6 +48,8 @@ int mlx_load_safetensors( const mlx_stream s); int mlx_save_writer(mlx_io_writer out_stream, const mlx_array a); int mlx_save(const char* file, const mlx_array a); +int mlx_save_gguf(const char* file, mlx_io_gguf gguf); + int mlx_save_safetensors_writer( mlx_io_writer in_stream, const mlx_map_string_to_array param, diff --git a/Source/Cmlx/include/mlx/c/io_types.h b/Source/Cmlx/include/mlx/c/io_types.h index 88349b57c..5382e7342 100644 --- a/Source/Cmlx/include/mlx/c/io_types.h +++ b/Source/Cmlx/include/mlx/c/io_types.h @@ -95,6 +95,52 @@ int mlx_io_writer_tostring(mlx_string* str_, mlx_io_writer io); */ int mlx_io_writer_free(mlx_io_writer io); +/** + * A MLX GGUF object. + */ +typedef struct mlx_io_gguf_ { + void* ctx; +} mlx_io_gguf; + +mlx_io_gguf mlx_io_gguf_new(void); +int mlx_io_gguf_free(mlx_io_gguf io); +int mlx_io_gguf_get_keys(mlx_vector_string* keys, mlx_io_gguf io); +int mlx_io_gguf_get_array(mlx_array* arr, mlx_io_gguf io, const char* key); +int mlx_io_gguf_get_metadata_array( + mlx_array* arr, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_get_metadata_string( + mlx_string* str, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_get_metadata_vector_string( + mlx_vector_string* vstr, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_has_metadata_array(bool* flag, mlx_io_gguf io, const char* key); +int mlx_io_gguf_has_metadata_string( + bool* flag, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_has_metadata_vector_string( + bool* flag, + mlx_io_gguf io, + const char* key); +int mlx_io_gguf_set_array(mlx_io_gguf io, const char* key, const mlx_array arr); +int mlx_io_gguf_set_metadata_array( + mlx_io_gguf io, + const char* key, + const mlx_array marr); +int mlx_io_gguf_set_metadata_string( + mlx_io_gguf io, + const char* key, + const char* mstr); +int mlx_io_gguf_set_metadata_vector_string( + mlx_io_gguf io, + const char* key, + const mlx_vector_string mvstr); + /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/include/mlx/c/linalg.h b/Source/Cmlx/include/mlx/c/linalg.h index 91d5d661e..6d5b5ae15 100644 --- a/Source/Cmlx/include/mlx/c/linalg.h +++ b/Source/Cmlx/include/mlx/c/linalg.h @@ -44,6 +44,7 @@ int mlx_linalg_cross( const mlx_array b, int axis, const mlx_stream s); +int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -97,6 +98,11 @@ int mlx_linalg_qr( mlx_array* res_1, const mlx_array a, const mlx_stream s); +int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s); int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/Source/Cmlx/include/mlx/c/metal.h b/Source/Cmlx/include/mlx/c/metal.h index 5877b224b..f65d15b62 100644 --- a/Source/Cmlx/include/mlx/c/metal.h +++ b/Source/Cmlx/include/mlx/c/metal.h @@ -28,7 +28,9 @@ extern "C" { */ /**@{*/ +int mlx_metal_get_metallib_path(mlx_string* res); int mlx_metal_is_available(bool* res); +int mlx_metal_set_metallib_path(const char* path); int mlx_metal_start_capture(const char* path); int mlx_metal_stop_capture(void); diff --git a/Source/Cmlx/include/mlx/c/mlx.h b/Source/Cmlx/include/mlx/c/mlx.h index ffadac89a..2aa9077c6 100644 --- a/Source/Cmlx/include/mlx/c/mlx.h +++ b/Source/Cmlx/include/mlx/c/mlx.h @@ -14,6 +14,7 @@ #include "mlx/c/export.h" #include "mlx/c/fast.h" #include "mlx/c/fft.h" +#include "mlx/c/graph_utils.h" #include "mlx/c/half.h" #include "mlx/c/io.h" #include "mlx/c/io_types.h" diff --git a/Source/Cmlx/include/mlx/c/ops.h b/Source/Cmlx/include/mlx/c/ops.h index 64d70e2f4..741778570 100644 --- a/Source/Cmlx/include/mlx/c/ops.h +++ b/Source/Cmlx/include/mlx/c/ops.h @@ -323,7 +323,38 @@ int mlx_conv_transpose3d( int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s); +int mlx_count_nonzero_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); +int mlx_cummax_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_cummax( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cummin_axis( mlx_array* res, const mlx_array a, int axis, @@ -331,25 +362,40 @@ int mlx_cummax( bool inclusive, const mlx_stream s); int mlx_cummin( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cumprod_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumprod( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s); +int mlx_cumsum_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumsum( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_depends( @@ -375,6 +421,12 @@ int mlx_diagonal( int axis1, int axis2, const mlx_stream s); +int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s); int mlx_divide( mlx_array* res, const mlx_array a, @@ -423,6 +475,18 @@ int mlx_flatten( int start_axis, int end_axis, const mlx_stream s); +int mlx_flip_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s); +int mlx_flip_axis( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_flip(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor_divide( mlx_array* res, @@ -486,6 +550,20 @@ int mlx_gather_qmm( const char* mode, bool sorted_indices, const mlx_stream s); +int mlx_gather_qqmm( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s); int mlx_greater( mlx_array* res, const mlx_array a, @@ -543,6 +621,14 @@ int mlx_less_equal( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_linspace_endpoint( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s); int mlx_linspace( mlx_array* res, double start, @@ -559,13 +645,19 @@ int mlx_logaddexp( const mlx_array a, const mlx_array b, const mlx_stream s); -int mlx_logcumsumexp( +int mlx_logcumsumexp_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s); +int mlx_logcumsumexp( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_logical_and( mlx_array* res, const mlx_array a, @@ -577,6 +669,11 @@ int mlx_logical_or( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s); int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -647,13 +744,24 @@ int mlx_mean( const mlx_array a, bool keepdims, const mlx_stream s); -int mlx_median( +int mlx_median_axes( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s); +int mlx_median_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_median( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); int mlx_meshgrid( mlx_vector_array* res, const mlx_vector_array arrays, @@ -721,6 +829,11 @@ int mlx_ones( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_ones_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_outer( mlx_array* res, @@ -757,6 +870,7 @@ int mlx_partition( const mlx_array a, int kth, const mlx_stream s); +int mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_power( mlx_array* res, const mlx_array a, @@ -956,6 +1070,12 @@ int mlx_scatter_prod_single( const mlx_array updates, int axis, const mlx_stream s); +int mlx_searchsorted( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s); int mlx_segmented_mm( mlx_array* res, const mlx_array a, @@ -1004,6 +1124,50 @@ int mlx_slice_update_dynamic( const int* axes, size_t axes_num, const mlx_stream s); +int mlx_slice_update_add( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_max( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_min( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); +int mlx_slice_update_prod( + mlx_array* res, + const mlx_array src, + const mlx_array update, + const int* start, + size_t start_num, + const int* stop, + size_t stop_num, + const int* strides, + size_t strides_num, + const mlx_stream s); int mlx_softmax_axes( mlx_array* res, const mlx_array a, @@ -1163,7 +1327,7 @@ int mlx_topk_axis( int axis, const mlx_stream s); int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s); -int mlx_trace( +int mlx_trace_axes( mlx_array* res, const mlx_array a, int offset, @@ -1171,6 +1335,7 @@ int mlx_trace( int axis2, mlx_dtype dtype, const mlx_stream s); +int mlx_trace(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_transpose_axes( mlx_array* res, const mlx_array a, @@ -1187,6 +1352,7 @@ int mlx_tri( const mlx_stream s); int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s); int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s); +int mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -1194,6 +1360,12 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unstack_axis( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_unstack(mlx_vector_array* res, const mlx_array a, const mlx_stream s); int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -1215,6 +1387,12 @@ int mlx_var( bool keepdims, int ddof, const mlx_stream s); +int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s); int mlx_view( mlx_array* res, const mlx_array a, @@ -1232,6 +1410,11 @@ int mlx_zeros( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_zeros_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s); /**@}*/ diff --git a/Source/Cmlx/include/mlx/c/stream.h b/Source/Cmlx/include/mlx/c/stream.h index d5865b806..abd26701b 100644 --- a/Source/Cmlx/include/mlx/c/stream.h +++ b/Source/Cmlx/include/mlx/c/stream.h @@ -33,6 +33,18 @@ mlx_stream mlx_stream_new(void); * Returns a new stream on a device. */ mlx_stream mlx_stream_new_device(mlx_device dev); +/** + * Returns a new stream on a device that can be used from any thread. + * + * Streams are otherwise thread affine: a stream's GPU command encoder is + * registered per thread, so evaluating on a stream from a thread other than + * the one that created it fails. Streams returned here are registered + * globally instead. + * + * MLX applies no synchronization to these streams -- it is the caller's + * responsibility to ensure there are no data races on them. + */ +mlx_stream mlx_stream_new_thread_unsafe(mlx_device dev); /** * Set stream to provided src stream. */ diff --git a/Source/Cmlx/mlx b/Source/Cmlx/mlx index ce45c5250..1f8e74e3f 160000 --- a/Source/Cmlx/mlx +++ b/Source/Cmlx/mlx @@ -1 +1 @@ -Subproject commit ce45c52505c8158ea48d2a54e8caae05efd86bfe +Subproject commit 1f8e74e3f12f31365464a6867c6579f0e9b29d85 diff --git a/Source/Cmlx/mlx-c b/Source/Cmlx/mlx-c index 0726ca922..c74db5307 160000 --- a/Source/Cmlx/mlx-c +++ b/Source/Cmlx/mlx-c @@ -1 +1 @@ -Subproject commit 0726ca922fc902c4c61ef9c27d94132be418e945 +Subproject commit c74db5307cc8ce122f48d97ef951b30578674e7f diff --git a/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp b/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp new file mode 100644 index 000000000..91e81980f --- /dev/null +++ b/Source/Cmlx/mlx-conditional/jit_compiler_conditional.cpp @@ -0,0 +1,12 @@ +// Copyright © 2025 Apple Inc. + +#include + +// `JitCompiler` shells out via `std::system()`, which is unavailable on iOS and +// visionOS. It is only referenced from `backend/cpu/compiled.cpp`, and +// `compiled_conditional.cpp` selects `no_cpu/compiled.cpp` on those platforms, +// so it isn't needed there. Keep this condition in sync with +// `compiled_conditional.cpp`. +#if !(TARGET_OS_IOS || TARGET_OS_VISION) +#include "../mlx/mlx/backend/cpu/jit_compiler.cpp" +#endif diff --git a/Source/Cmlx/mlx-generated/binary_ops.cpp b/Source/Cmlx/mlx-generated/binary_ops.cpp index 651cb006d..2c8157725 100644 --- a/Source/Cmlx/mlx-generated/binary_ops.cpp +++ b/Source/Cmlx/mlx-generated/binary_ops.cpp @@ -21,33 +21,40 @@ constant mlx::os_log logger("mlx", "binary_ops"); struct Add { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x + y; } }; struct FloorDivide { template - T operator()(T x, T y) { + metal::enable_if_t & !metal::is_signed_v, T> + operator()(T x, T y) thread { return x / y; } - template <> - float operator()(float x, float y) { - return trunc(x / y); + template + metal::enable_if_t & metal::is_signed_v, T> + operator()(T x, T y) thread { + auto q = x / y; + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + return q; } - template <> - half operator()(half x, half y) { - return trunc(x / y); + template + metal::enable_if_t, T> operator()(T x, T y) thread { + return floor(x / y); } template <> - bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { - return trunc(x / y); + complex64_t operator()(complex64_t x, complex64_t y) thread { + // Complex is not supported, simply make compiler happy. + return x / y; } }; struct Divide { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x / y; } }; @@ -55,12 +62,12 @@ struct Divide { struct Remainder { template metal::enable_if_t & !metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { return x % y; } template metal::enable_if_t & metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { auto r = x % y; if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -68,7 +75,7 @@ struct Remainder { return r; } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { T r = fmod(x, y); if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -76,25 +83,25 @@ struct Remainder { return r; } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { return x % y; } }; struct Equal { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y; } }; struct NaNEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y || (metal::isnan(x) && metal::isnan(y)); } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x == y || (metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) && metal::isnan(y.imag)) || @@ -105,35 +112,35 @@ struct NaNEqual { struct Greater { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x > y; } }; struct GreaterEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x >= y; } }; struct Less { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x < y; } }; struct LessEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x <= y; } }; struct LogAddExp { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { if (metal::isnan(x) || metal::isnan(y)) { return metal::numeric_limits::quiet_NaN(); } @@ -145,7 +152,7 @@ struct LogAddExp { : (maxval + log1p(metal::exp(minval - maxval))); }; - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) || metal::isnan(y.imag)) { return metal::numeric_limits::quiet_NaN(); @@ -166,12 +173,12 @@ struct LogAddExp { struct Maximum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::max(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -179,7 +186,7 @@ struct Maximum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -189,12 +196,12 @@ struct Maximum { struct Minimum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::min(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -202,7 +209,7 @@ struct Minimum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -212,30 +219,32 @@ struct Minimum { struct Multiply { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x * y; } }; struct NotEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x != y; } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x.real != y.real || x.imag != y.imag; } }; struct Power { template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { return metal::pow(base, exp); } template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { T res = 1; // Undefined to raise integer to negative power if (exp < 0) { @@ -255,7 +264,7 @@ struct Power { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (x.real == 0 && x.imag == 0) { if (metal::isnan(y.real) || metal::isnan(y.imag)) { auto nan = metal::numeric_limits::quiet_NaN(); @@ -273,70 +282,70 @@ struct Power { struct Subtract { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x - y; } }; struct LogicalAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x && y; }; }; struct LogicalOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x || y; }; }; struct BitwiseAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x & y; }; }; struct BitwiseOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x | y; }; }; struct BitwiseXor { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x ^ y; }; }; struct LeftShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x << y; }; }; struct RightShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x >> y; }; }; struct ArcTan2 { template - T operator()(T y, T x) { + T operator()(T y, T x) thread { return metal::precise::atan2(y, x); } }; struct DivMod { template - metal::array operator()(T x, T y) { + metal::array operator()(T x, T y) thread { return {FloorDivide{}(x, y), Remainder{}(x, y)}; }; }; diff --git a/Source/Cmlx/mlx-generated/compiled_preamble.cpp b/Source/Cmlx/mlx-generated/compiled_preamble.cpp index b0804ed70..8b89f66a8 100644 --- a/Source/Cmlx/mlx-generated/compiled_preamble.cpp +++ b/Source/Cmlx/mlx-generated/compiled_preamble.cpp @@ -1,4 +1,4 @@ -const char* get_kernel_preamble() { +const char* get_prebuilt_preamble() { return R"preamble( #include #include @@ -76,6 +76,17 @@ inline float operator+(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator-(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) - static_cast(rhs); } inline float operator*(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator*(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) * static_cast(rhs); } inline float operator/(float16_t lhs, bfloat16_t rhs) { return static_cast(lhs) / static_cast(rhs); } inline float operator/(bfloat16_t lhs, float16_t rhs) { return static_cast(lhs) / static_cast(rhs); } +template +constexpr bool is_complex = false; +template +constexpr bool is_complex().real())>> = + true; +template +constexpr bool is_signed_v = std::is_signed_v || + std::is_same_v || std::is_same_v; +template +constexpr bool is_floating_point_v = std::is_floating_point_v || + std::is_same_v || std::is_same_v; } namespace mlx::core { struct complex64_t; @@ -169,11 +180,6 @@ void store(T* dst, Simd x) { } *(Simd*)dst = x; } -template -constexpr bool is_complex = false; -template -constexpr bool is_complex().real())>> = - true; template Simd rint(Simd in) { if constexpr (is_complex) { @@ -193,7 +199,6 @@ Simd recip(Simd in) { } template Simd operator-(Simd in) { return std::negate{}(in.value); } template Simd operator!(Simd in) { return std::logical_not{}(in.value); } -template Simd abs(Simd in) { return std::abs(in.value); } template Simd acos(Simd in) { return std::acos(in.value); } template Simd acosh(Simd in) { return std::acosh(in.value); } template Simd asin(Simd in) { return std::asin(in.value); } @@ -212,6 +217,14 @@ template Simd sqrt(Simd in) { return std::sqrt(in.value template Simd tan(Simd in) { return std::tan(in.value); } template Simd tanh(Simd in) { return std::tanh(in.value); } template +Simd abs(Simd in) { + if constexpr (std::is_unsigned_v) { + return in; + } else { + return std::abs(in.value); + } +} +template Simd log1p(Simd in) { if constexpr (is_complex) { auto x = in.value.real(); @@ -283,7 +296,7 @@ Simd remainder(Simd a_, Simd b_) { } else { r = std::remainder(a, b); } - if constexpr (std::is_signed_v) { + if constexpr (is_signed_v) { if (r != 0 && (r < 0 != b < 0)) { r += b; } @@ -320,6 +333,11 @@ Simd pow(Simd a, Simd b) { return std::pow(base, exp); } else { T res = 1; + if constexpr (std::is_signed_v) { + if (exp < 0) { + return 0; + } + } while (exp) { if (exp & 1) { res *= base; @@ -422,12 +440,25 @@ Simd sincos(Simd in) { return select(sign_mask_cos, yc, -yc); } } +template +Simd sincos_checked(Simd x) { + Simd xf = x; + if (any(abs(xf) > Simd(8388608.0f))) { + Simd out; + for (int i = 0; i < N; ++i) { + float v = xf[i]; + out[i] = static_cast(Sine ? std::sin(v) : std::cos(v)); + } + return out; + } + return sincos(x); +} template Simd sin(Simd x) { if constexpr (is_complex) { return std::sin(x.value); } else { - return sincos(x); + return sincos_checked(x); } } template @@ -435,7 +466,7 @@ Simd cos(Simd x) { if constexpr (is_complex) { return std::cos(x.value); } else { - return sincos(x); + return sincos_checked(x); } } template @@ -685,8 +716,6 @@ struct Select { } }; } -const char* get_kernel_preamble(); -using namespace mlx::core; -using namespace mlx::core::detail; +const char* get_prebuilt_preamble(); )preamble"; } diff --git a/Source/Cmlx/mlx-generated/conv.cpp b/Source/Cmlx/mlx-generated/conv.cpp index 9c3e07417..25c6d5433 100644 --- a/Source/Cmlx/mlx-generated/conv.cpp +++ b/Source/Cmlx/mlx-generated/conv.cpp @@ -241,10 +241,10 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -280,7 +280,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -308,7 +308,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -390,10 +390,10 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -464,7 +464,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -489,7 +489,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -571,11 +571,11 @@ struct Conv2DWeightBlockLoader { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -585,7 +585,7 @@ struct Conv2DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -612,7 +612,7 @@ struct Conv2DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_hw < (params->wS[1] * params->wS[0])) { src += weight_step; return; @@ -681,10 +681,10 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -736,7 +736,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -765,7 +765,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -860,10 +860,10 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -954,7 +954,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t d_mask = mask_t(1) << weight_d; mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -981,7 +981,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -1074,11 +1074,11 @@ struct Conv3DWeightBlockLoader { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -1088,7 +1088,7 @@ struct Conv3DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -1115,7 +1115,7 @@ struct Conv3DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_dhw < (params->wS[0] * params->wS[1] * params->wS[2])) { src += weight_step; return; @@ -1234,10 +1234,10 @@ struct Conv2DInputBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -1266,7 +1266,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { @@ -1322,7 +1322,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; @@ -1378,11 +1378,11 @@ struct Conv2DWeightBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld), params(params_), @@ -1391,7 +1391,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (bi >= BROWS || bj >= BCOLS) return; @@ -1445,7 +1445,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; @@ -1491,7 +1491,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -1507,14 +1507,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast( x * static_cast(alpha) + (static_cast(beta) * c)); } @@ -1623,13 +1623,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -1659,7 +1655,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -1828,7 +1825,7 @@ struct BaseMMAFrag { for (short j = 0; j < kElemCols; j++) { if ((off_x + i) < lim_x && (off_y + j) < lim_y) { dst[i * kElemCols + j] = - static_cast(src[(off_x + i) * str_x + (off_x + j) * str_y]); + static_cast(src[(off_x + i) * str_x + (off_y + j) * str_y]); } else { dst[i * kElemCols + j] = T(0); } @@ -1973,24 +1970,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -1999,16 +1997,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2025,7 +2023,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2042,7 +2040,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2057,7 +2055,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2072,8 +2070,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2092,8 +2092,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2116,7 +2118,7 @@ struct MMATile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -2226,7 +2228,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -2244,7 +2246,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -2271,7 +2273,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -2284,8 +2286,11 @@ struct BlockMMA { Ctile.template store(D, ldd); } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -2305,7 +2310,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -2324,7 +2329,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -2338,7 +2344,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -2367,7 +2373,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -2412,7 +2418,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -2445,7 +2451,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -2558,7 +2564,7 @@ struct BlockMMA< /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -2578,7 +2584,7 @@ struct BlockMMA< /* Karatsuba MMA: 3 real MMAs per K-chunk */ METAL_FUNC void mma( const threadgroup complex64_t* As, - const threadgroup complex64_t* Bs) { + const threadgroup complex64_t* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -2636,7 +2642,7 @@ struct BlockMMA< } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Adjust for simdgroup and thread location D += sm * ldd + sn; @@ -2655,8 +2661,11 @@ struct BlockMMA< } } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { D += sm * ldd + sn; start -= short2(sn, sm); stop -= short2(sn, sm); @@ -2687,7 +2696,7 @@ struct BlockMMA< } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { D += sm * ldd + sn; dst_tile_dims -= short2(sn, sm); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) @@ -2713,7 +2722,8 @@ struct BlockMMA< /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) { complex64_t out = epilogue_op.apply( @@ -2729,7 +2739,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -2761,7 +2771,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -2810,7 +2820,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -2845,7 +2855,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/Source/Cmlx/mlx-generated/copy.cpp b/Source/Cmlx/mlx-generated/copy.cpp index e8d5f39c0..21090d887 100644 --- a/Source/Cmlx/mlx-generated/copy.cpp +++ b/Source/Cmlx/mlx-generated/copy.cpp @@ -22,11 +22,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } } @@ -40,11 +40,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } } @@ -59,11 +59,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } } @@ -78,11 +78,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } } @@ -94,7 +94,7 @@ template constant const int64_t& src_stride [[buffer(3)]], uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); - dst[index] = static_cast(src[src_idx]); + dst[index] = cast_to(src[src_idx]); } template @@ -106,7 +106,7 @@ template uint2 grid_dim [[threads_per_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * index.y; - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -119,7 +119,7 @@ template auto src_idx = elem_to_loc_3(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -136,14 +136,14 @@ template if (N == 1) { IdxT dst_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); return; } auto xshape = src_shape[ndim - 1]; IdxT dst_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z); auto src_xstride = src_strides[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[dst_idx + i] = static_cast(src[src_idx]); + dst[dst_idx + i] = cast_to(src[src_idx]); src_idx += src_xstride; } } @@ -157,7 +157,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -169,7 +169,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -181,7 +181,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -200,14 +200,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } @@ -224,7 +224,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -238,7 +238,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -252,7 +252,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -275,14 +275,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } diff --git a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h index 7d14d04bc..fe469a89c 100644 --- a/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -169,296 +169,283 @@ constexpr char jit_source_binary_ops[] = { 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, - 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, - 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x72, 0x75, 0x6e, 0x63, - 0x28, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, - 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x44, 0x69, 0x76, 0x69, 0x64, 0x65, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, - 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, - 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, - 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, - 0x74, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, 0x6e, 0x64, 0x65, 0x72, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, - 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, - 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, - 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, - 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, - 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, - 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, - 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, - 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, - 0x69, 0x73, 0x5f, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, - 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x20, 0x3d, 0x20, 0x78, - 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, 0x20, 0x21, 0x3d, 0x20, 0x30, - 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, 0x3c, 0x20, 0x30, 0x20, 0x21, - 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x20, - 0x2b, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, - 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x72, 0x20, 0x3d, 0x20, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, - 0x6d, 0x6f, 0x64, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, 0x20, 0x21, - 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, 0x3c, 0x20, - 0x30, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x20, 0x3d, 0x20, 0x72, 0x20, 0x2b, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x71, 0x20, 0x3d, + 0x20, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x78, 0x20, 0x25, + 0x20, 0x79, 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x28, + 0x78, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x28, 0x79, + 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x71, 0x20, 0x2d, 0x3d, 0x20, + 0x31, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x71, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6d, 0x70, 0x6c, 0x65, + 0x78, 0x20, 0x69, 0x73, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x73, 0x75, 0x70, + 0x70, 0x6f, 0x72, 0x74, 0x65, 0x64, 0x2c, 0x20, 0x73, 0x69, 0x6d, 0x70, + 0x6c, 0x79, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x20, 0x63, 0x6f, 0x6d, 0x70, + 0x69, 0x6c, 0x65, 0x72, 0x20, 0x68, 0x61, 0x70, 0x70, 0x79, 0x2e, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x66, 0x6c, 0x6f, + 0x6f, 0x72, 0x28, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x44, 0x69, 0x76, 0x69, + 0x64, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, + 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2f, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, 0x6e, 0x64, 0x65, + 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, + 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, + 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, + 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, + 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x72, 0x20, 0x3d, + 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, 0x20, 0x21, 0x3d, + 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, 0x3c, 0x20, 0x30, + 0x20, 0x21, 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x20, 0x2b, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, - 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x71, 0x75, 0x61, 0x6c, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, - 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, 0x3d, 0x20, - 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, - 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4e, 0x61, 0x4e, 0x45, 0x71, 0x75, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, + 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x25, 0x20, 0x79, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x72, 0x20, + 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x66, 0x6d, 0x6f, 0x64, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x72, + 0x20, 0x21, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x28, 0x72, 0x20, + 0x3c, 0x20, 0x30, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x20, 0x3c, 0x20, 0x30, + 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x20, 0x3d, 0x20, 0x72, 0x20, 0x2b, 0x20, 0x79, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x45, 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, - 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x20, 0x7c, - 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x28, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, - 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, - 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, - 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x69, - 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x73, 0x6e, - 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, - 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, - 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, - 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, - 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, - 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, - 0x29, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, - 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, - 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, - 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, 0x69, 0x6d, - 0x61, 0x67, 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, 0x69, 0x6d, - 0x61, 0x67, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, - 0x3d, 0x20, 0x79, 0x20, 0x7c, 0x7c, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, - 0x6e, 0x28, 0x78, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, - 0x6e, 0x28, 0x79, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, 0x72, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, - 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, - 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, - 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, 0x72, 0x45, + 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4e, 0x61, 0x4e, 0x45, 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, + 0x61, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, + 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, + 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3d, 0x3d, 0x20, 0x79, + 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, + 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, + 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, + 0x29, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, + 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, + 0x28, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, + 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, + 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, 0x61, + 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x28, 0x69, 0x73, 0x6e, 0x61, + 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, + 0x26, 0x26, 0x20, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x72, + 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, 0x26, 0x26, 0x20, 0x78, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, - 0x20, 0x3e, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, - 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, - 0x73, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, - 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, - 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, - 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, - 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, - 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, - 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x73, 0x73, 0x45, - 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x20, 0x3d, 0x3d, 0x20, 0x79, 0x20, 0x7c, 0x7c, 0x20, 0x28, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x20, 0x26, 0x26, 0x20, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, + 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, + 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x47, 0x72, 0x65, 0x61, 0x74, 0x65, + 0x72, 0x45, 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x4c, 0x65, 0x73, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, - 0x20, 0x3c, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, - 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, - 0x67, 0x41, 0x64, 0x64, 0x45, 0x78, 0x70, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, - 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, - 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, - 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, - 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, - 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, - 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, - 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, - 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, - 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, - 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, - 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, - 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, - 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x69, 0x6d, - 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, - 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, - 0x71, 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, - 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, - 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, 0x71, - 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x7d, 0x3b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x78, 0x20, - 0x3d, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3e, - 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3f, 0x20, - 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x20, 0x3d, 0x20, - 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3c, 0x20, 0x79, - 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3f, 0x20, 0x78, 0x20, - 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, - 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, - 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, - 0x6f, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3d, - 0x20, 0x6d, 0x61, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x21, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, - 0x73, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x28, 0x6d, 0x69, 0x6e, 0x5f, - 0x72, 0x65, 0x61, 0x6c, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x6d, 0x69, - 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3d, 0x3d, 0x20, 0x6d, 0x61, - 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6d, - 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3c, 0x20, 0x30, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, 0x6e, 0x3b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, - 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x4c, - 0x6f, 0x67, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, 0x7b, 0x7d, 0x28, 0x6d, - 0x69, 0x6e, 0x29, 0x20, 0x2b, 0x20, 0x45, 0x78, 0x70, 0x7b, 0x7d, 0x28, - 0x6d, 0x61, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x4c, - 0x6f, 0x67, 0x31, 0x70, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, 0x7b, 0x7d, - 0x28, 0x6d, 0x69, 0x6e, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, 0x29, 0x29, - 0x20, 0x2b, 0x20, 0x6d, 0x61, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, - 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, - 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, - 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x20, 0x7c, - 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x29, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, - 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, - 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x71, 0x75, - 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x54, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, - 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x76, 0x61, - 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x78, 0x2c, 0x20, 0x79, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x28, 0x6d, 0x69, 0x6e, 0x76, 0x61, 0x6c, 0x20, - 0x3d, 0x3d, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, - 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, - 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, - 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x20, 0x7c, 0x7c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x3d, 0x20, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, - 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, - 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, - 0x74, 0x79, 0x28, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x3f, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x3a, - 0x20, 0x54, 0x28, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x2b, 0x20, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6c, - 0x6f, 0x67, 0x31, 0x70, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, - 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x78, 0x70, 0x28, 0x6d, 0x69, 0x6e, 0x76, - 0x61, 0x6c, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x29, - 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, - 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, - 0x74, 0x20, 0x4d, 0x61, 0x78, 0x69, 0x6d, 0x75, 0x6d, 0x20, 0x7b, 0x0a, + 0x20, 0x3c, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x73, + 0x73, 0x45, 0x71, 0x75, 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x78, 0x20, 0x3c, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x4c, 0x6f, 0x67, 0x41, 0x64, 0x64, 0x45, 0x78, 0x70, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, - 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, - 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, - 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, - 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, - 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x29, 0x20, - 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, - 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x69, 0x6d, - 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, - 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, - 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, - 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, - 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x4d, 0x69, 0x6e, 0x69, 0x6d, 0x75, 0x6d, 0x20, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, + 0x28, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, + 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x20, 0x7c, 0x7c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x2e, + 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, + 0x3a, 0x3a, 0x71, 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, + 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, + 0x3a, 0x71, 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, + 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, + 0x78, 0x20, 0x3d, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, + 0x20, 0x3e, 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, + 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x20, + 0x3d, 0x20, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3c, + 0x20, 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x3f, 0x20, + 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, + 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x2e, 0x72, 0x65, 0x61, + 0x6c, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x21, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x28, 0x6d, 0x69, + 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, + 0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3d, 0x3d, 0x20, + 0x6d, 0x61, 0x78, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x6d, 0x69, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x6c, 0x20, 0x3c, 0x20, + 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, + 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x4c, 0x6f, 0x67, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, 0x7b, 0x7d, + 0x28, 0x6d, 0x69, 0x6e, 0x29, 0x20, 0x2b, 0x20, 0x45, 0x78, 0x70, 0x7b, + 0x7d, 0x28, 0x6d, 0x61, 0x78, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x4c, 0x6f, 0x67, 0x31, 0x70, 0x7b, 0x7d, 0x28, 0x45, 0x78, 0x70, + 0x7b, 0x7d, 0x28, 0x6d, 0x69, 0x6e, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, + 0x29, 0x29, 0x20, 0x2b, 0x20, 0x6d, 0x61, 0x78, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x29, + 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x79, 0x29, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, + 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, + 0x71, 0x75, 0x69, 0x65, 0x74, 0x5f, 0x4e, 0x61, 0x4e, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, + 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x78, 0x2c, + 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x28, 0x6d, 0x69, 0x6e, 0x76, 0x61, + 0x6c, 0x20, 0x3d, 0x3d, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, + 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, + 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x20, 0x7c, + 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, 0x3d, + 0x3d, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, + 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x3f, 0x20, 0x6d, 0x61, 0x78, 0x76, + 0x61, 0x6c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x3a, 0x20, 0x54, 0x28, 0x6d, 0x61, 0x78, 0x76, 0x61, 0x6c, 0x20, + 0x2b, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6c, 0x6f, 0x67, 0x31, 0x70, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x78, 0x70, 0x28, 0x6d, 0x69, + 0x6e, 0x76, 0x61, 0x6c, 0x20, 0x2d, 0x20, 0x6d, 0x61, 0x78, 0x76, 0x61, + 0x6c, 0x29, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4d, 0x61, 0x78, 0x69, 0x6d, 0x75, 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, @@ -469,7 +456,7 @@ constexpr char jit_source_binary_ops[] = { 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x78, 0x2c, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, @@ -483,7 +470,7 @@ constexpr char jit_source_binary_ops[] = { 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, + 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, @@ -491,187 +478,223 @@ constexpr char jit_source_binary_ops[] = { 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, 0x79, 0x20, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3e, 0x20, 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, - 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4d, 0x75, 0x6c, 0x74, 0x69, 0x70, - 0x6c, 0x79, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, - 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, - 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, - 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, - 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2a, 0x20, 0x79, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x4e, 0x6f, 0x74, 0x45, 0x71, 0x75, 0x61, 0x6c, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, - 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, - 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, - 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x2e, - 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x2e, - 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x78, 0x2e, - 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x2e, - 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, - 0x21, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, - 0x63, 0x74, 0x20, 0x50, 0x6f, 0x77, 0x65, 0x72, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, - 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, - 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, - 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, - 0x28, 0x54, 0x20, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x54, 0x20, 0x65, - 0x78, 0x70, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4d, 0x69, 0x6e, 0x69, 0x6d, 0x75, + 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, + 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x54, 0x20, 0x72, 0x65, 0x73, 0x20, 0x3d, 0x20, 0x31, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x61, 0x69, - 0x73, 0x69, 0x6e, 0x67, 0x20, 0x61, 0x6e, 0x20, 0x69, 0x6e, 0x74, 0x65, - 0x67, 0x65, 0x72, 0x20, 0x74, 0x6f, 0x20, 0x61, 0x20, 0x6e, 0x65, 0x67, - 0x61, 0x74, 0x69, 0x76, 0x65, 0x20, 0x70, 0x6f, 0x77, 0x65, 0x72, 0x20, - 0x69, 0x73, 0x20, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x64, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, - 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, - 0x20, 0x28, 0x65, 0x78, 0x70, 0x20, 0x3c, 0x20, 0x30, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x77, 0x68, 0x69, - 0x6c, 0x65, 0x20, 0x28, 0x65, 0x78, 0x70, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x65, - 0x78, 0x70, 0x20, 0x26, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x20, - 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x69, 0x6e, 0x28, + 0x78, 0x2c, 0x20, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, + 0x28, 0x29, 0x29, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, + 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, 0x79, 0x20, 0x3f, 0x20, + 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x6e, 0x61, 0x6e, 0x28, 0x78, + 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x65, 0x78, 0x70, 0x20, 0x3e, 0x3e, 0x3d, 0x20, 0x31, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x61, - 0x73, 0x65, 0x20, 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x65, 0x73, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, - 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, - 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, - 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x70, 0x6f, - 0x77, 0x28, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x65, 0x78, 0x70, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, - 0x64, 0x3a, 0x3a, 0x70, 0x6f, 0x77, 0x28, 0x62, 0x61, 0x73, 0x65, 0x2c, - 0x20, 0x65, 0x78, 0x70, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x20, + 0x79, 0x20, 0x3f, 0x20, 0x78, 0x20, 0x3a, 0x20, 0x79, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4d, 0x75, 0x6c, 0x74, + 0x69, 0x70, 0x6c, 0x79, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, + 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2a, 0x20, + 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4e, 0x6f, 0x74, 0x45, 0x71, 0x75, + 0x61, 0x6c, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, 0x6c, + 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x78, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, + 0x79, 0x2e, 0x72, 0x65, 0x61, 0x6c, 0x28, 0x29, 0x20, 0x7c, 0x7c, 0x20, + 0x78, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x20, 0x21, 0x3d, 0x20, + 0x79, 0x2e, 0x69, 0x6d, 0x61, 0x67, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x78, 0x20, 0x21, 0x3d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, + 0x72, 0x75, 0x63, 0x74, 0x20, 0x50, 0x6f, 0x77, 0x65, 0x72, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, + 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, + 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x54, 0x20, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x54, + 0x20, 0x65, 0x78, 0x70, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x5f, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x72, 0x61, 0x6c, + 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x20, 0x72, 0x65, 0x73, 0x20, 0x3d, 0x20, 0x31, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, + 0x61, 0x69, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x61, 0x6e, 0x20, 0x69, 0x6e, + 0x74, 0x65, 0x67, 0x65, 0x72, 0x20, 0x74, 0x6f, 0x20, 0x61, 0x20, 0x6e, + 0x65, 0x67, 0x61, 0x74, 0x69, 0x76, 0x65, 0x20, 0x70, 0x6f, 0x77, 0x65, + 0x72, 0x20, 0x69, 0x73, 0x20, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x69, 0x6e, + 0x65, 0x64, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, + 0x5f, 0x73, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x54, 0x3e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x65, 0x78, 0x70, 0x20, 0x3c, 0x20, 0x30, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x30, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x77, + 0x68, 0x69, 0x6c, 0x65, 0x20, 0x28, 0x65, 0x78, 0x70, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x65, 0x78, 0x70, 0x20, 0x26, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x73, 0x20, 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x78, 0x70, 0x20, 0x3e, 0x3e, 0x3d, + 0x20, 0x31, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x62, 0x61, 0x73, 0x65, 0x20, 0x2a, 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, + 0x65, 0x73, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x28, 0x69, 0x73, 0x5f, 0x63, 0x6f, 0x6d, 0x70, + 0x6c, 0x65, 0x78, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x70, 0x6f, 0x77, 0x28, 0x62, 0x61, 0x73, 0x65, 0x2c, 0x20, 0x65, 0x78, + 0x70, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x70, 0x6f, 0x77, 0x28, 0x62, 0x61, 0x73, + 0x65, 0x2c, 0x20, 0x65, 0x78, 0x70, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, 0x75, 0x62, 0x74, 0x72, 0x61, + 0x63, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, + 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2d, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x53, 0x75, 0x62, 0x74, 0x72, 0x61, 0x63, 0x74, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, - 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, - 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x2d, 0x20, 0x79, 0x3b, 0x0a, 0x20, - 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, - 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, 0x6c, 0x41, 0x6e, 0x64, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, - 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, - 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x26, 0x20, 0x79, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, 0x6c, 0x4f, - 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, - 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, - 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, - 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, - 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, - 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, 0x7c, 0x20, 0x79, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, - 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, 0x65, - 0x41, 0x6e, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, - 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, - 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, - 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, - 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x20, 0x79, + 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, 0x6c, 0x41, + 0x6e, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, + 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x26, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, - 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, - 0x65, 0x4f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x67, 0x69, 0x63, 0x61, + 0x6c, 0x4f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, 0x20, 0x79, - 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, - 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, 0x73, - 0x65, 0x58, 0x6f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, - 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, - 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, - 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x5e, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, 0x7c, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, - 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x66, 0x74, 0x53, - 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, - 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, - 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, - 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x3c, 0x3c, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, 0x69, + 0x73, 0x65, 0x41, 0x6e, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x26, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, - 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x52, 0x69, 0x67, 0x68, - 0x74, 0x53, 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, + 0x69, 0x73, 0x65, 0x4f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, 0x7c, + 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, + 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x42, 0x69, 0x74, 0x77, + 0x69, 0x73, 0x65, 0x58, 0x6f, 0x72, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, - 0x3e, 0x3e, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, - 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x41, 0x72, - 0x63, 0x54, 0x61, 0x6e, 0x32, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, - 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6f, - 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, - 0x79, 0x2c, 0x20, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x74, 0x61, 0x6e, - 0x32, 0x28, 0x79, 0x2c, 0x20, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, - 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, - 0x44, 0x69, 0x76, 0x4d, 0x6f, 0x64, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, + 0x5e, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x65, 0x66, + 0x74, 0x53, 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, - 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, - 0x61, 0x79, 0x3c, 0x54, 0x2c, 0x20, 0x32, 0x3e, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, 0x78, 0x2c, - 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x46, 0x6c, 0x6f, 0x6f, - 0x72, 0x44, 0x69, 0x76, 0x69, 0x64, 0x65, 0x7b, 0x7d, 0x28, 0x78, 0x2c, - 0x20, 0x79, 0x29, 0x2c, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, 0x6e, 0x64, - 0x65, 0x72, 0x7b, 0x7d, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, 0x7d, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x7d, 0x20, - 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, - 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, - 0x63, 0x75, 0x0a, 0x00, + 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, + 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x78, 0x20, + 0x3c, 0x3c, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x52, 0x69, + 0x67, 0x68, 0x74, 0x53, 0x68, 0x69, 0x66, 0x74, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x54, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, + 0x28, 0x54, 0x20, 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x78, 0x20, 0x3e, 0x3e, 0x20, 0x79, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x41, 0x72, 0x63, 0x54, 0x61, 0x6e, 0x32, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, + 0x54, 0x20, 0x79, 0x2c, 0x20, 0x54, 0x20, 0x78, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x74, + 0x61, 0x6e, 0x32, 0x28, 0x79, 0x2c, 0x20, 0x78, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x44, 0x69, 0x76, 0x4d, 0x6f, 0x64, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, + 0x72, 0x72, 0x61, 0x79, 0x3c, 0x54, 0x2c, 0x20, 0x32, 0x3e, 0x20, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x20, + 0x78, 0x2c, 0x20, 0x54, 0x20, 0x79, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x46, 0x6c, + 0x6f, 0x6f, 0x72, 0x44, 0x69, 0x76, 0x69, 0x64, 0x65, 0x7b, 0x7d, 0x28, + 0x78, 0x2c, 0x20, 0x79, 0x29, 0x2c, 0x20, 0x52, 0x65, 0x6d, 0x61, 0x69, + 0x6e, 0x64, 0x65, 0x72, 0x7b, 0x7d, 0x28, 0x78, 0x2c, 0x20, 0x79, 0x29, + 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, + 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_cast_op[] = { @@ -1125,6 +1148,512 @@ constexpr char jit_source_config[] = { 0x20, 0x33, 0x32, 0x0a, 0x00, }; +constexpr char jit_source_cute_dequant[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, + 0x65, 0x2f, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x2f, 0x6e, 0x75, + 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, + 0x68, 0x70, 0x70, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, + 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x65, 0x2f, 0x74, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x2e, 0x68, 0x70, 0x70, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, + 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, + 0x73, 0x2f, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x63, 0x6f, + 0x6e, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x2e, 0x68, 0x3e, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x20, 0x7b, 0x0a, 0x0a, + 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x62, + 0x5f, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x65, 0x72, + 0x5f, 0x73, 0x75, 0x62, 0x62, 0x79, 0x74, 0x65, 0x3c, 0x33, 0x2c, 0x20, + 0x66, 0x61, 0x6c, 0x73, 0x65, 0x3e, 0x3b, 0x0a, 0x75, 0x73, 0x69, 0x6e, + 0x67, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x35, 0x62, 0x5f, 0x74, 0x20, 0x3d, + 0x20, 0x69, 0x6e, 0x74, 0x65, 0x67, 0x65, 0x72, 0x5f, 0x73, 0x75, 0x62, + 0x62, 0x79, 0x74, 0x65, 0x3c, 0x35, 0x2c, 0x20, 0x66, 0x61, 0x6c, 0x73, + 0x65, 0x3e, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x46, 0x6c, + 0x6f, 0x61, 0x74, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x53, 0x74, 0x79, 0x6c, + 0x65, 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x41, + 0x72, 0x72, 0x61, 0x79, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x65, + 0x72, 0x3c, 0x54, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x62, 0x5f, + 0x74, 0x2c, 0x20, 0x4e, 0x2c, 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, + 0x61, 0x73, 0x73, 0x65, 0x72, 0x74, 0x28, 0x4e, 0x20, 0x25, 0x20, 0x38, + 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, + 0x74, 0x79, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, + 0x69, 0x6e, 0x67, 0x20, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, + 0x79, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x75, 0x69, 0x6e, 0x74, 0x33, 0x62, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x3e, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, + 0x5f, 0x48, 0x4f, 0x53, 0x54, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, + 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x72, 0x65, + 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x20, 0x63, 0x6f, + 0x6e, 0x76, 0x65, 0x72, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x26, + 0x20, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x73, 0x5f, 0x62, + 0x61, 0x73, 0x65, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, + 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, + 0x2a, 0x3e, 0x28, 0x26, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, + 0x5f, 0x50, 0x52, 0x41, 0x47, 0x4d, 0x41, 0x5f, 0x55, 0x4e, 0x52, 0x4f, + 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, + 0x20, 0x3c, 0x20, 0x4e, 0x20, 0x2f, 0x20, 0x38, 0x3b, 0x20, 0x2b, 0x2b, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x5f, 0x62, + 0x61, 0x73, 0x65, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x33, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, + 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x5d, 0x20, 0x3d, 0x20, 0x54, + 0x28, 0x73, 0x5b, 0x30, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x37, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, + 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, 0x2b, 0x20, + 0x31, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x30, 0x5d, + 0x20, 0x26, 0x20, 0x30, 0x78, 0x33, 0x38, 0x29, 0x20, 0x3e, 0x3e, 0x20, + 0x33, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, 0x2b, + 0x20, 0x32, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x30, + 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x63, 0x30, 0x29, 0x20, 0x3e, 0x3e, + 0x20, 0x36, 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, + 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x31, 0x29, 0x20, 0x3c, 0x3c, + 0x20, 0x32, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, + 0x2b, 0x20, 0x33, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, + 0x31, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x65, 0x29, 0x20, 0x3e, + 0x3e, 0x20, 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, + 0x20, 0x2b, 0x20, 0x34, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, + 0x5b, 0x31, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x37, 0x30, 0x29, 0x20, + 0x3e, 0x3e, 0x20, 0x34, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, + 0x38, 0x20, 0x2b, 0x20, 0x35, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, + 0x73, 0x5b, 0x31, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x38, 0x30, 0x29, + 0x20, 0x3e, 0x3e, 0x20, 0x37, 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, + 0x73, 0x5b, 0x32, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x33, 0x29, + 0x20, 0x3c, 0x3c, 0x20, 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, + 0x20, 0x38, 0x20, 0x2b, 0x20, 0x36, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, + 0x28, 0x73, 0x5b, 0x32, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x31, 0x63, + 0x29, 0x20, 0x3e, 0x3e, 0x20, 0x32, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, + 0x2a, 0x20, 0x38, 0x20, 0x2b, 0x20, 0x37, 0x5d, 0x20, 0x3d, 0x20, 0x54, + 0x28, 0x28, 0x73, 0x5b, 0x32, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x65, + 0x30, 0x29, 0x20, 0x3e, 0x3e, 0x20, 0x35, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, + 0x53, 0x5f, 0x48, 0x4f, 0x53, 0x54, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, + 0x45, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, + 0x79, 0x70, 0x65, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, + 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x73, 0x6f, 0x75, + 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x26, 0x20, 0x73, 0x29, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x6f, 0x6e, 0x76, + 0x65, 0x72, 0x74, 0x28, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x46, 0x6c, 0x6f, + 0x61, 0x74, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x53, 0x74, 0x79, 0x6c, 0x65, + 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x4e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x41, 0x72, + 0x72, 0x61, 0x79, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x65, 0x72, + 0x3c, 0x54, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x35, 0x62, 0x5f, 0x74, + 0x2c, 0x20, 0x4e, 0x2c, 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x61, + 0x73, 0x73, 0x65, 0x72, 0x74, 0x28, 0x4e, 0x20, 0x25, 0x20, 0x38, 0x20, + 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x75, 0x73, + 0x69, 0x6e, 0x67, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, + 0x79, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x54, 0x2c, 0x20, 0x4e, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x75, + 0x69, 0x6e, 0x74, 0x35, 0x62, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x3e, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, + 0x48, 0x4f, 0x53, 0x54, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, 0x0a, + 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x72, 0x65, 0x73, + 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x20, 0x63, 0x6f, 0x6e, + 0x76, 0x65, 0x72, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x73, + 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x26, 0x20, + 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, + 0x65, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x73, 0x5f, 0x62, 0x61, + 0x73, 0x65, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, + 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, 0x2a, + 0x3e, 0x28, 0x26, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, + 0x50, 0x52, 0x41, 0x47, 0x4d, 0x41, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, + 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, + 0x3c, 0x20, 0x4e, 0x20, 0x2f, 0x20, 0x38, 0x3b, 0x20, 0x2b, 0x2b, 0x69, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x2a, 0x20, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x5f, 0x62, 0x61, + 0x73, 0x65, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x35, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, + 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, + 0x73, 0x5b, 0x30, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x31, 0x66, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, + 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, 0x2b, 0x20, 0x31, + 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x30, 0x5d, 0x20, + 0x26, 0x20, 0x30, 0x78, 0x65, 0x30, 0x29, 0x20, 0x3e, 0x3e, 0x20, 0x35, + 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, 0x5d, 0x20, + 0x26, 0x20, 0x30, 0x78, 0x30, 0x33, 0x29, 0x20, 0x3c, 0x3c, 0x20, 0x33, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, + 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, 0x2b, 0x20, + 0x32, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, 0x5d, + 0x20, 0x26, 0x20, 0x30, 0x78, 0x37, 0x63, 0x29, 0x20, 0x3e, 0x3e, 0x20, + 0x32, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, 0x2b, + 0x20, 0x33, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, + 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x38, 0x30, 0x29, 0x20, 0x3e, 0x3e, + 0x20, 0x37, 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x32, + 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x66, 0x29, 0x20, 0x3c, 0x3c, + 0x20, 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, 0x20, + 0x2b, 0x20, 0x34, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, + 0x32, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x66, 0x30, 0x29, 0x20, 0x3e, + 0x3e, 0x20, 0x34, 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, + 0x33, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x31, 0x29, 0x20, 0x3c, + 0x3c, 0x20, 0x34, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x38, + 0x20, 0x2b, 0x20, 0x35, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, + 0x5b, 0x33, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x33, 0x65, 0x29, 0x20, + 0x3e, 0x3e, 0x20, 0x31, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, + 0x38, 0x20, 0x2b, 0x20, 0x36, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, + 0x73, 0x5b, 0x33, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x63, 0x30, 0x29, + 0x20, 0x3e, 0x3e, 0x20, 0x36, 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, + 0x73, 0x5b, 0x34, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x37, 0x29, + 0x20, 0x3c, 0x3c, 0x20, 0x32, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, + 0x20, 0x38, 0x20, 0x2b, 0x20, 0x37, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, + 0x28, 0x73, 0x5b, 0x34, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x66, 0x38, + 0x29, 0x20, 0x3e, 0x3e, 0x20, 0x33, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, + 0x5f, 0x48, 0x4f, 0x53, 0x54, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, + 0x0a, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, + 0x29, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x73, 0x6f, 0x75, 0x72, + 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x26, 0x20, 0x73, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, + 0x72, 0x74, 0x28, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x2c, 0x20, 0x46, 0x6c, 0x6f, 0x61, + 0x74, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x53, 0x74, 0x79, 0x6c, 0x65, 0x20, + 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x4e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x41, 0x72, 0x72, + 0x61, 0x79, 0x43, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x65, 0x72, 0x3c, + 0x54, 0x2c, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x36, 0x62, 0x5f, 0x74, 0x2c, + 0x20, 0x4e, 0x2c, 0x20, 0x52, 0x6f, 0x75, 0x6e, 0x64, 0x3e, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x61, 0x73, + 0x73, 0x65, 0x72, 0x74, 0x28, 0x4e, 0x20, 0x25, 0x20, 0x34, 0x20, 0x3d, + 0x3d, 0x20, 0x30, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x54, + 0x2c, 0x20, 0x4e, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, + 0x67, 0x20, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, + 0x65, 0x20, 0x3d, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x75, 0x69, + 0x6e, 0x74, 0x36, 0x62, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x3e, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, 0x48, + 0x4f, 0x53, 0x54, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, 0x0a, 0x20, + 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x72, 0x65, 0x73, 0x75, + 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x76, + 0x65, 0x72, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x73, 0x6f, + 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x26, 0x20, 0x73, + 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, 0x65, + 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x2a, 0x20, 0x73, 0x5f, 0x62, 0x61, 0x73, + 0x65, 0x20, 0x3d, 0x20, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, + 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, 0x2a, 0x3e, + 0x28, 0x26, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, 0x50, + 0x52, 0x41, 0x47, 0x4d, 0x41, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, + 0x20, 0x4e, 0x20, 0x2f, 0x20, 0x34, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x2a, 0x20, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x5f, 0x62, 0x61, 0x73, + 0x65, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x33, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5b, + 0x69, 0x20, 0x2a, 0x20, 0x34, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x73, + 0x5b, 0x30, 0x5d, 0x20, 0x26, 0x20, 0x30, 0x78, 0x33, 0x66, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, 0x6c, + 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x34, 0x20, 0x2b, 0x20, 0x31, 0x5d, + 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x30, 0x5d, 0x20, 0x3e, + 0x3e, 0x20, 0x36, 0x29, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x33, 0x29, + 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, 0x5d, 0x20, 0x26, + 0x20, 0x30, 0x78, 0x30, 0x66, 0x29, 0x20, 0x3c, 0x3c, 0x20, 0x32, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, 0x75, + 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x34, 0x20, 0x2b, 0x20, 0x32, + 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x31, 0x5d, 0x20, + 0x3e, 0x3e, 0x20, 0x34, 0x29, 0x20, 0x26, 0x20, 0x30, 0x78, 0x30, 0x66, + 0x29, 0x20, 0x2b, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x32, 0x5d, 0x20, + 0x26, 0x20, 0x30, 0x78, 0x30, 0x33, 0x29, 0x20, 0x3c, 0x3c, 0x20, 0x34, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x73, + 0x75, 0x6c, 0x74, 0x5b, 0x69, 0x20, 0x2a, 0x20, 0x34, 0x20, 0x2b, 0x20, + 0x33, 0x5d, 0x20, 0x3d, 0x20, 0x54, 0x28, 0x28, 0x73, 0x5b, 0x32, 0x5d, + 0x20, 0x3e, 0x3e, 0x20, 0x32, 0x29, 0x20, 0x26, 0x20, 0x30, 0x78, 0x33, + 0x66, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x72, 0x65, 0x73, + 0x75, 0x6c, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, + 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, 0x48, 0x4f, 0x53, 0x54, + 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, 0x0a, 0x20, 0x20, 0x72, 0x65, + 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x73, 0x6f, 0x75, 0x72, 0x63, 0x65, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x26, 0x20, 0x73, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x28, 0x73, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x7d, 0x20, + 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, + 0x20, 0x7b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x71, 0x75, 0x69, + 0x72, 0x65, 0x64, 0x20, 0x62, 0x79, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x64, + 0x20, 0x63, 0x6f, 0x70, 0x79, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x33, 0x2f, + 0x35, 0x2f, 0x36, 0x2d, 0x62, 0x69, 0x74, 0x20, 0x77, 0x65, 0x69, 0x67, + 0x68, 0x74, 0x73, 0x2e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x75, 0x69, 0x6e, 0x74, 0x32, 0x34, 0x5f, 0x74, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, + 0x74, 0x2c, 0x20, 0x33, 0x3e, 0x20, 0x62, 0x79, 0x74, 0x65, 0x73, 0x3b, + 0x0a, 0x7d, 0x3b, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x34, 0x30, 0x5f, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, + 0x72, 0x72, 0x61, 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, + 0x2c, 0x20, 0x35, 0x3e, 0x20, 0x62, 0x79, 0x74, 0x65, 0x73, 0x3b, 0x0a, + 0x7d, 0x3b, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x75, 0x69, + 0x6e, 0x74, 0x34, 0x38, 0x5f, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, + 0x72, 0x61, 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x38, 0x5f, 0x74, 0x2c, + 0x20, 0x36, 0x3e, 0x20, 0x62, 0x79, 0x74, 0x65, 0x73, 0x3b, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x75, 0x69, + 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x3c, 0x32, 0x34, 0x3e, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x20, 0x3d, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x32, 0x34, 0x5f, 0x74, + 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x75, 0x69, 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x3c, 0x34, 0x30, 0x3e, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x34, 0x30, + 0x5f, 0x74, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, + 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x3c, 0x34, + 0x38, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x34, 0x38, 0x5f, 0x74, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x7d, 0x20, + 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x63, 0x75, 0x74, 0x65, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x75, 0x73, + 0x69, 0x6e, 0x67, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x57, 0x68, 0x65, 0x74, 0x68, 0x65, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20, + 0x71, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x74, 0x79, 0x70, 0x65, 0x20, 0x69, + 0x73, 0x20, 0x61, 0x66, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x71, 0x75, 0x61, + 0x6e, 0x74, 0x69, 0x7a, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, + 0x0a, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, + 0x73, 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, 0x76, 0x20, 0x3d, 0x20, 0x21, + 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x68, 0x61, 0x73, + 0x5f, 0x6e, 0x65, 0x67, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x7a, 0x65, + 0x72, 0x6f, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x3b, + 0x0a, 0x0a, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, + 0x69, 0x7a, 0x65, 0x20, 0x43, 0x75, 0x54, 0x65, 0x20, 0x74, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x73, 0x20, 0x77, 0x69, 0x74, 0x68, 0x20, 0x6f, 0x75, + 0x74, 0x20, 0x3d, 0x20, 0x77, 0x20, 0x2a, 0x20, 0x73, 0x20, 0x2b, 0x20, + 0x7a, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x57, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x4f, 0x3e, 0x0a, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x44, + 0x45, 0x56, 0x49, 0x43, 0x45, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x0a, 0x63, + 0x75, 0x74, 0x65, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x69, 0x7a, + 0x65, 0x64, 0x5f, 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x28, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x57, 0x20, 0x77, 0x2c, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x20, 0x73, 0x2c, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x5a, 0x20, 0x7a, 0x2c, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x4f, 0x20, 0x6f, 0x75, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x20, + 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x6f, + 0x75, 0x74, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, 0x74, + 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, + 0x70, 0x65, 0x28, 0x77, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, + 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x53, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x62, + 0x65, 0x20, 0x6f, 0x6e, 0x65, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2e, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x53, 0x54, + 0x41, 0x54, 0x49, 0x43, 0x5f, 0x41, 0x53, 0x53, 0x45, 0x52, 0x54, 0x5f, + 0x56, 0x28, 0x63, 0x6f, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x73, 0x2e, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x29, 0x29, 0x20, 0x3d, 0x3d, 0x20, + 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x53, 0x54, 0x41, 0x54, 0x49, 0x43, + 0x5f, 0x41, 0x53, 0x53, 0x45, 0x52, 0x54, 0x5f, 0x56, 0x28, 0x63, 0x6f, + 0x73, 0x69, 0x7a, 0x65, 0x28, 0x7a, 0x2e, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x28, 0x29, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x49, 0x6e, 0x74, 0x3c, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x6d, 0x75, 0x73, 0x74, 0x20, 0x62, + 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x74, 0x69, 0x67, 0x75, 0x6f, 0x75, 0x73, + 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x61, 0x6c, 0x65, 0x73, + 0x63, 0x65, 0x28, 0x77, 0x2e, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x53, + 0x54, 0x41, 0x54, 0x49, 0x43, 0x5f, 0x41, 0x53, 0x53, 0x45, 0x52, 0x54, + 0x5f, 0x56, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x49, 0x6e, 0x74, + 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x55, 0x73, 0x65, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x73, + 0x69, 0x6f, 0x6e, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x20, + 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x26, 0x20, + 0x77, 0x5f, 0x76, 0x65, 0x63, 0x20, 0x3d, 0x20, 0x2a, 0x28, 0x72, 0x65, + 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, + 0x73, 0x74, 0x3c, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x63, 0x75, 0x74, + 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x41, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x4e, 0x3e, 0x2a, 0x3e, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x61, 0x77, 0x5f, 0x70, + 0x6f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x28, + 0x77, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x29, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x73, + 0x63, 0x61, 0x6c, 0x65, 0x7b, 0x73, 0x5b, 0x30, 0x5d, 0x7d, 0x3b, 0x0a, + 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x4e, + 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x41, 0x72, 0x72, 0x61, 0x79, 0x43, + 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x65, 0x72, 0x3c, 0x45, 0x6c, 0x65, + 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, + 0x20, 0x4e, 0x3e, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, 0x65, + 0x72, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x77, 0x5f, + 0x64, 0x71, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x6e, 0x76, 0x65, 0x72, 0x74, + 0x65, 0x72, 0x28, 0x77, 0x5f, 0x76, 0x65, 0x63, 0x29, 0x20, 0x2a, 0x20, + 0x73, 0x63, 0x61, 0x6c, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x71, + 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, 0x61, + 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x20, 0x7a, 0x65, 0x72, 0x6f, 0x5f, 0x70, 0x6f, 0x69, 0x6e, 0x74, + 0x7b, 0x7a, 0x5b, 0x30, 0x5d, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x77, 0x5f, 0x64, 0x71, 0x20, 0x3d, 0x20, 0x77, 0x5f, 0x64, 0x71, 0x20, + 0x2b, 0x20, 0x7a, 0x65, 0x72, 0x6f, 0x5f, 0x70, 0x6f, 0x69, 0x6e, 0x74, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, + 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x72, 0x6d, 0x65, 0x6d, 0x5f, 0x70, + 0x74, 0x72, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x28, + 0x26, 0x77, 0x5f, 0x64, 0x71, 0x29, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x2e, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x29, 0x29, 0x2c, 0x20, 0x6f, + 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x57, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x5a, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x4f, 0x3e, 0x0a, 0x43, + 0x55, 0x54, 0x45, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, 0x20, 0x76, + 0x6f, 0x69, 0x64, 0x0a, 0x63, 0x75, 0x74, 0x65, 0x5f, 0x6e, 0x61, 0x69, + 0x76, 0x65, 0x5f, 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x28, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x57, 0x20, 0x77, 0x2c, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x20, 0x73, 0x2c, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x5a, 0x20, 0x7a, 0x2c, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x4f, 0x20, 0x6f, 0x75, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x6f, + 0x75, 0x74, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, 0x74, + 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, + 0x70, 0x65, 0x28, 0x77, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, + 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, + 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, + 0x75, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x74, + 0x72, 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x28, 0x77, 0x2c, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x5b, 0x5d, 0x28, 0x51, 0x75, 0x61, 0x6e, + 0x74, 0x20, 0x71, 0x29, 0x20, 0x7b, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x28, 0x71, 0x29, + 0x3b, 0x20, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x74, 0x72, 0x61, 0x6e, + 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, + 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x5b, 0x5d, 0x28, 0x45, 0x6c, + 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x65, 0x2c, 0x20, 0x53, 0x63, 0x61, + 0x6c, 0x65, 0x20, 0x73, 0x29, 0x20, 0x7b, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x65, 0x20, 0x2a, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x28, 0x73, 0x29, 0x3b, 0x20, 0x7d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, + 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, + 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x72, + 0x61, 0x6e, 0x73, 0x66, 0x6f, 0x72, 0x6d, 0x28, 0x6f, 0x75, 0x74, 0x2c, + 0x20, 0x7a, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x70, 0x6c, 0x75, + 0x73, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x57, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x4f, 0x3e, 0x0a, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x44, 0x45, 0x56, + 0x49, 0x43, 0x45, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x63, 0x75, 0x74, + 0x65, 0x5f, 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x28, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x57, 0x20, 0x77, 0x2c, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x53, 0x20, 0x73, 0x2c, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x5a, 0x20, 0x7a, 0x2c, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x4f, 0x20, 0x6f, 0x75, 0x74, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x28, 0x63, 0x6f, 0x61, 0x6c, 0x65, 0x73, 0x63, 0x65, + 0x28, 0x77, 0x2e, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x29, 0x29, + 0x29, 0x20, 0x3d, 0x3d, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, + 0x7d, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x73, 0x5f, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x76, 0x3c, 0x64, + 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x2e, 0x6c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x28, 0x29, 0x29, 0x3e, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x65, 0x5f, 0x76, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x69, 0x7a, 0x65, 0x64, 0x5f, 0x64, 0x65, 0x71, 0x75, + 0x61, 0x6e, 0x74, 0x28, 0x77, 0x2c, 0x20, 0x73, 0x2c, 0x20, 0x7a, 0x2c, + 0x20, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, + 0x74, 0x65, 0x5f, 0x6e, 0x61, 0x69, 0x76, 0x65, 0x5f, 0x64, 0x65, 0x71, + 0x75, 0x61, 0x6e, 0x74, 0x28, 0x77, 0x2c, 0x20, 0x73, 0x2c, 0x20, 0x7a, + 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + constexpr char jit_source_fp16_math[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, @@ -1514,140 +2043,6 @@ constexpr char jit_source_fp16_math[] = { 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; -constexpr char jit_source_gather[] = { - 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, - 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, - 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, - 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, - 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, - 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, - 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, - 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, - 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, - 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, - 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, - 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x4e, 0x49, 0x44, 0x58, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, 0x44, - 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, - 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, - 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, - 0x69, 0x64, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x28, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, - 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, - 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, - 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, - 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, - 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, - 0x65, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, - 0x32, 0x5f, 0x74, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, - 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, - 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, - 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, - 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, - 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, - 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, - 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, - 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, - 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, - 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, - 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, - 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, - 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, - 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, - 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, - 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, - 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, - 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, - 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, - 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, - 0x20, 0x25, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, - 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, - 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, - 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, - 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x65, - 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, - 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, - 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, - 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, - 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, - 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, - 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, - 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, - 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, - 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, - 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, - 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, - 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, - 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, - 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, - 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, - 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, - 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, - 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, - 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, - 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x73, 0x72, - 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, - 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, - 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, - 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, - 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, - 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, - 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, -}; - constexpr char jit_source_gather_axis[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, @@ -1790,6 +2185,1010 @@ constexpr char jit_source_gather_axis[] = { 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; +constexpr char jit_source_gather[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, 0x69, 0x6c, 0x73, + 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, + 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x2e, + 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, 0x20, 0x63, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x4e, 0x49, 0x44, 0x58, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, 0x44, + 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x3e, 0x0a, 0x5f, + 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, + 0x69, 0x64, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, + 0x73, 0x72, 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, + 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, + 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, + 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, + 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, + 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, + 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, + 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, + 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, + 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, + 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, + 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, + 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x25, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x6f, 0x75, 0x74, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x65, + 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, + 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, + 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, + 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, + 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, + 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, + 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, + 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, + 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x2c, 0x20, 0x73, 0x72, + 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, + 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, + 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_gemm_sm70[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x65, 0x2f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x2e, 0x68, 0x70, 0x70, 0x3e, 0x0a, + 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x6c, 0x61, 0x6e, 0x67, 0x2d, 0x66, 0x6f, + 0x72, 0x6d, 0x61, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, + 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x53, 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x41, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, + 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x42, 0x3e, 0x0a, + 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, 0x68, 0x61, 0x72, 0x65, + 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x45, 0x6e, 0x67, 0x69, 0x6e, 0x65, + 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x63, 0x6f, + 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x6d, 0x65, 0x6d, 0x4c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x41, 0x3e, 0x3e, 0x20, 0x41, 0x3b, 0x0a, + 0x20, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x45, 0x6e, 0x67, 0x69, 0x6e, + 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x63, + 0x6f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x6d, 0x65, 0x6d, + 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x42, 0x3e, 0x3e, 0x20, 0x42, 0x3b, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, + 0x6f, 0x72, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4b, 0x3e, 0x0a, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x28, 0x54, 0x69, 0x6c, 0x65, 0x4d, 0x20, 0x62, 0x4d, 0x2c, + 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4b, 0x20, 0x62, 0x4b, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x4f, 0x44, 0x4f, 0x3a, 0x20, + 0x43, 0x61, 0x6c, 0x63, 0x75, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x73, 0x77, + 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x20, 0x62, 0x61, 0x73, 0x65, 0x64, 0x20, + 0x6f, 0x6e, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, + 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x20, 0x3d, 0x20, + 0x63, 0x6f, 0x6d, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x28, + 0x53, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x3c, 0x33, 0x2c, 0x33, 0x2c, + 0x33, 0x3e, 0x7b, 0x7d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x20, 0x3c, 0x5f, 0x38, 0x2c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, + 0x3c, 0x5f, 0x38, 0x2c, 0x20, 0x5f, 0x38, 0x3e, 0x3e, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x38, 0x2c, 0x53, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x36, 0x34, + 0x3e, 0x3e, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x74, 0x6f, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x77, 0x69, + 0x7a, 0x7a, 0x6c, 0x65, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x28, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4b, 0x29, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x73, + 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x6d, + 0x70, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x53, 0x77, 0x69, + 0x7a, 0x7a, 0x6c, 0x65, 0x3c, 0x33, 0x2c, 0x33, 0x2c, 0x33, 0x3e, 0x7b, + 0x7d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, + 0x31, 0x36, 0x2c, 0x5f, 0x38, 0x3e, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x3c, 0x5f, 0x38, 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x74, 0x6f, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x28, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x2c, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, + 0x6f, 0x72, 0x41, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, + 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, + 0x65, 0x72, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x28, 0x43, 0x74, 0x61, 0x54, + 0x69, 0x6c, 0x65, 0x72, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4e, + 0x6f, 0x74, 0x65, 0x3a, 0x20, 0x4b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, + 0x6c, 0x61, 0x75, 0x6e, 0x63, 0x68, 0x65, 0x72, 0x20, 0x61, 0x73, 0x73, + 0x75, 0x6d, 0x65, 0x73, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x20, 0x62, 0x65, 0x69, 0x6e, 0x67, 0x20, 0x73, + 0x61, 0x6d, 0x65, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x61, 0x6c, 0x6c, 0x20, + 0x70, 0x61, 0x72, 0x61, 0x6d, 0x65, 0x74, 0x65, 0x72, 0x73, 0x2e, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x62, 0x4d, 0x2c, 0x20, + 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x5d, 0x20, 0x3d, 0x20, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x3b, 0x0a, 0x20, 0x20, 0x61, + 0x75, 0x74, 0x6f, 0x20, 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, + 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x4b, 0x4d, 0x61, + 0x6a, 0x6f, 0x72, 0x41, 0x3e, 0x28, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4b, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x73, 0x42, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x3e, 0x28, + 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3a, 0x3a, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x73, + 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, 0x42, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x20, 0x3d, 0x20, 0x66, + 0x61, 0x6c, 0x73, 0x65, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x3d, + 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, + 0x72, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, + 0x6d, 0x6d, 0x61, 0x28, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x6f, 0x74, 0x65, 0x3a, + 0x20, 0x4b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, 0x6c, 0x61, 0x75, 0x6e, + 0x63, 0x68, 0x65, 0x72, 0x20, 0x61, 0x73, 0x73, 0x75, 0x6d, 0x65, 0x73, + 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, + 0x20, 0x62, 0x65, 0x69, 0x6e, 0x67, 0x20, 0x73, 0x61, 0x6d, 0x65, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x61, 0x6c, 0x6c, 0x20, 0x70, 0x61, 0x72, 0x61, + 0x6d, 0x65, 0x74, 0x65, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x73, + 0x69, 0x6e, 0x67, 0x20, 0x41, 0x74, 0x6f, 0x6d, 0x20, 0x3d, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, + 0x6e, 0x64, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x5f, 0x74, 0x3c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, 0x6f, 0x6e, 0x64, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x5f, 0x74, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, + 0x65, 0x5f, 0x76, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, + 0x20, 0x68, 0x61, 0x6c, 0x66, 0x5f, 0x74, 0x3e, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x53, 0x4d, 0x38, 0x30, + 0x5f, 0x31, 0x36, 0x78, 0x38, 0x78, 0x31, 0x36, 0x5f, 0x46, 0x33, 0x32, + 0x46, 0x31, 0x36, 0x46, 0x31, 0x36, 0x46, 0x33, 0x32, 0x5f, 0x54, 0x4e, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x63, + 0x6f, 0x6e, 0x64, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x5f, 0x74, + 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, + 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x62, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x31, 0x36, 0x5f, 0x74, 0x3e, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x53, 0x4d, 0x38, 0x30, 0x5f, 0x31, 0x36, 0x78, 0x38, 0x78, 0x31, + 0x36, 0x5f, 0x46, 0x33, 0x32, 0x42, 0x46, 0x31, 0x36, 0x42, 0x46, 0x31, + 0x36, 0x46, 0x33, 0x32, 0x5f, 0x54, 0x4e, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x55, + 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x46, 0x4d, 0x41, 0x3c, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x3e, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x55, 0x6e, + 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x46, 0x4d, 0x41, 0x3c, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x3e, + 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x21, 0x53, 0x4d, 0x38, 0x30, 0x20, + 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, + 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x41, 0x74, 0x6f, + 0x6d, 0x7b, 0x7d, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, + 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x31, 0x36, 0x2c, 0x5f, 0x38, + 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, + 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, 0x3e, 0x3d, + 0x20, 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x41, + 0x74, 0x6f, 0x6d, 0x7b, 0x7d, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x3c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x32, 0x2c, 0x5f, + 0x32, 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x54, 0x69, + 0x6c, 0x65, 0x3c, 0x5f, 0x33, 0x32, 0x2c, 0x5f, 0x33, 0x32, 0x2c, 0x5f, + 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, + 0x28, 0x41, 0x74, 0x6f, 0x6d, 0x7b, 0x7d, 0x2c, 0x20, 0x4c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x3c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x31, + 0x2c, 0x5f, 0x34, 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x54, 0x69, 0x6c, 0x65, 0x3c, 0x5f, 0x31, 0x36, 0x2c, 0x5f, 0x33, 0x32, + 0x2c, 0x5f, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4e, 0x75, 0x6d, 0x54, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x73, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4d, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x69, 0x6c, 0x65, + 0x4b, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x28, 0x4e, 0x75, 0x6d, 0x54, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x73, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x73, 0x2c, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4d, 0x20, 0x62, + 0x4d, 0x2c, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x4b, 0x20, 0x62, 0x4b, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x54, 0x4f, 0x44, 0x4f, + 0x3a, 0x20, 0x4f, 0x6e, 0x6c, 0x79, 0x20, 0x64, 0x6f, 0x20, 0x31, 0x2d, + 0x65, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x72, 0x65, 0x61, 0x64, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x20, 0x6f, 0x66, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, + 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6e, 0x5f, 0x72, + 0x65, 0x61, 0x64, 0x20, 0x3d, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x20, 0x3f, 0x20, 0x38, 0x20, 0x3a, 0x20, + 0x31, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x20, 0x3d, 0x20, 0x43, 0x6f, 0x70, 0x79, + 0x5f, 0x41, 0x74, 0x6f, 0x6d, 0x3c, 0x55, 0x6e, 0x69, 0x76, 0x65, 0x72, + 0x73, 0x61, 0x6c, 0x43, 0x6f, 0x70, 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, + 0x5f, 0x62, 0x69, 0x74, 0x5f, 0x74, 0x3c, 0x6e, 0x5f, 0x72, 0x65, 0x61, + 0x64, 0x20, 0x2a, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, + 0x69, 0x74, 0x73, 0x5f, 0x76, 0x3c, 0x54, 0x3e, 0x3e, 0x3e, 0x2c, 0x20, + 0x54, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x4b, 0x4d, + 0x61, 0x6a, 0x6f, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6b, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x20, 0x3d, 0x20, 0x62, 0x4b, 0x20, 0x2f, 0x20, 0x6e, 0x5f, + 0x72, 0x65, 0x61, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x49, 0x6e, 0x74, 0x3c, 0x6e, + 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2f, + 0x20, 0x6b, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x3e, 0x7b, + 0x7d, 0x2c, 0x20, 0x6b, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, + 0x29, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x52, 0x69, 0x67, + 0x68, 0x74, 0x7b, 0x7d, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, + 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x64, 0x29, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x5f, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x20, 0x3d, 0x20, 0x62, 0x4d, 0x20, 0x2f, 0x20, + 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x5f, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, + 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, + 0x2f, 0x20, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x3e, + 0x7b, 0x7d, 0x29, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x4c, + 0x65, 0x66, 0x74, 0x7b, 0x7d, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x28, 0x6e, 0x5f, 0x72, 0x65, 0x61, 0x64, 0x2c, 0x20, 0x49, + 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x29, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, + 0x61, 0x6a, 0x6f, 0x72, 0x41, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x2c, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x2c, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, + 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x42, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x43, 0x3e, 0x0a, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x44, 0x45, 0x56, 0x49, + 0x43, 0x45, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x67, 0x65, 0x6d, 0x6d, + 0x5f, 0x73, 0x6d, 0x37, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, + 0x6f, 0x70, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x74, 0x61, 0x54, + 0x69, 0x6c, 0x65, 0x72, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x41, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x42, 0x20, 0x67, 0x42, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x43, 0x20, + 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x6d, + 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, + 0x64, 0x75, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, + 0x74, 0x68, 0x65, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x20, 0x6f, 0x66, + 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x6e, 0x64, 0x73, 0x2e, 0x0a, 0x20, + 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x67, + 0x41, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, 0x74, 0x79, + 0x70, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x53, 0x68, + 0x69, 0x66, 0x74, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, + 0x6f, 0x20, 0x77, 0x65, 0x20, 0x68, 0x61, 0x6e, 0x64, 0x6c, 0x65, 0x20, + 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x20, 0x6f, 0x66, 0x20, 0x4b, + 0x20, 0x69, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x30, 0x74, 0x68, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x67, 0x41, 0x20, 0x3d, + 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, + 0x75, 0x65, 0x2c, 0x20, 0x30, 0x29, 0x2c, 0x20, 0x67, 0x41, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x64, 0x6f, 0x6d, 0x61, + 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, + 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x30, + 0x29, 0x2c, 0x20, 0x67, 0x42, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x73, 0x6d, 0x65, + 0x6d, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x2e, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x73, 0x41, 0x5f, 0x6c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x5d, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, + 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x41, 0x2c, 0x20, 0x4b, 0x4d, + 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x20, 0x6d, 0x65, 0x6d, 0x6f, + 0x72, 0x79, 0x20, 0x62, 0x75, 0x66, 0x66, 0x65, 0x72, 0x2e, 0x0a, 0x20, + 0x20, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x20, 0x5f, 0x5f, 0x73, 0x68, + 0x61, 0x72, 0x65, 0x64, 0x5f, 0x5f, 0x20, 0x63, 0x68, 0x61, 0x72, 0x20, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x62, 0x75, 0x66, 0x5b, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x68, 0x61, 0x72, + 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x20, 0x3d, 0x20, + 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, + 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x64, + 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x41, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x2c, 0x20, 0x64, 0x65, 0x63, 0x6c, + 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x29, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x53, 0x68, 0x61, 0x72, + 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x26, 0x20, 0x73, + 0x6d, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x2a, 0x72, 0x65, 0x69, 0x6e, 0x74, + 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, + 0x65, 0x2a, 0x3e, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x62, 0x75, 0x66, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x73, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x2e, + 0x41, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, 0x29, 0x2c, 0x20, + 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x73, 0x6d, 0x65, + 0x6d, 0x2e, 0x42, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, 0x29, + 0x2c, 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x4d, 0x41, 0x2e, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x6d, 0x61, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, + 0x5f, 0x6d, 0x6d, 0x61, 0x3c, 0x53, 0x4d, 0x38, 0x30, 0x2c, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x6d, + 0x61, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x5b, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x5d, + 0x20, 0x3d, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, + 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, 0x6f, 0x70, + 0x79, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, + 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x41, 0x2c, 0x20, 0x41, 0x6c, 0x69, + 0x67, 0x6e, 0x65, 0x64, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x73, 0x2c, 0x20, 0x62, 0x4d, 0x2c, 0x20, 0x62, + 0x4b, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, + 0x6f, 0x70, 0x79, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x2c, 0x20, 0x41, + 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x2c, 0x20, 0x62, 0x4e, 0x2c, + 0x20, 0x62, 0x4b, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x50, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x74, 0x68, + 0x65, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x69, 0x6e, 0x67, 0x20, 0x6f, 0x66, + 0x20, 0x41, 0x2f, 0x42, 0x2f, 0x43, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x73, + 0x20, 0x61, 0x63, 0x72, 0x6f, 0x73, 0x73, 0x20, 0x74, 0x68, 0x65, 0x20, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, + 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x61, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, + 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x41, 0x43, 0x50, 0x59, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, + 0x4d, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x73, + 0x41, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, + 0x5f, 0x44, 0x28, 0x73, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x41, 0x43, 0x50, 0x59, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, + 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x72, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, + 0x74, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, 0x74, 0x41, 0x73, 0x41, 0x29, + 0x3b, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, 0x59, + 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x41, 0x43, 0x50, 0x59, + 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, + 0x70, 0x79, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x62, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x67, + 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x42, 0x20, + 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, + 0x28, 0x67, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, + 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, + 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x73, 0x42, 0x20, 0x3d, 0x20, + 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, + 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x44, 0x28, 0x73, + 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, + 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, + 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x74, 0x42, 0x72, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x69, + 0x6b, 0x65, 0x28, 0x74, 0x42, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, + 0x59, 0x5f, 0x4d, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x4d, 0x41, 0x2e, 0x0a, 0x20, + 0x20, 0x54, 0x68, 0x72, 0x4d, 0x4d, 0x41, 0x20, 0x74, 0x68, 0x72, 0x5f, + 0x6d, 0x6d, 0x61, 0x20, 0x3d, 0x20, 0x6d, 0x6d, 0x61, 0x2e, 0x67, 0x65, + 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x73, 0x41, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, + 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x41, 0x28, 0x73, 0x41, 0x29, + 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x43, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, + 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x42, 0x28, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, + 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, + 0x67, 0x43, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x43, + 0x28, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x43, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x43, + 0x28, 0x74, 0x43, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, + 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x65, 0x73, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x6d, 0x2f, 0x6e, 0x20, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, + 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x70, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x41, 0x73, 0x41, 0x29, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x41, + 0x73, 0x41, 0x29, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x30, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, + 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x42, 0x70, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, + 0x6c, 0x3e, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x42, + 0x73, 0x42, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x74, 0x42, 0x73, 0x42, 0x29, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x30, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x5f, 0x4e, + 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x63, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, + 0x3e, 0x28, 0x73, 0x41, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x31, 0x3e, 0x28, 0x73, 0x41, 0x29, 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, + 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x63, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x69, 0x64, + 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x73, 0x42, 0x29, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x73, 0x42, + 0x29, 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x63, 0x43, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, + 0x79, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x30, 0x3e, 0x28, 0x67, 0x43, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x43, 0x29, 0x29, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x63, 0x41, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x63, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, + 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x42, 0x63, 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x63, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x43, 0x50, 0x59, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, + 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x63, 0x43, 0x20, 0x3d, 0x20, 0x74, + 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x43, 0x28, 0x63, 0x43, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, + 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, + 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x6d, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6d, 0x20, 0x3c, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x70, 0x41, + 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x41, 0x70, 0x41, 0x28, 0x6d, 0x2c, 0x30, 0x29, 0x20, + 0x3d, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x63, + 0x41, 0x28, 0x30, 0x2c, 0x6d, 0x2c, 0x30, 0x29, 0x29, 0x20, 0x3c, 0x20, + 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, + 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, + 0x28, 0x74, 0x42, 0x70, 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x42, 0x70, 0x42, 0x28, + 0x6e, 0x2c, 0x30, 0x29, 0x20, 0x3d, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x30, + 0x3e, 0x28, 0x74, 0x42, 0x63, 0x42, 0x28, 0x30, 0x2c, 0x6e, 0x2c, 0x30, + 0x29, 0x29, 0x20, 0x3c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x4d, 0x45, 0x4d, 0x20, 0x3d, 0x3e, 0x20, + 0x52, 0x4d, 0x45, 0x4d, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x20, + 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, + 0x2c, 0x20, 0x74, 0x41, 0x70, 0x41, 0x2c, 0x20, 0x74, 0x41, 0x67, 0x41, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, + 0x2c, 0x20, 0x74, 0x41, 0x72, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, + 0x79, 0x5f, 0x62, 0x2c, 0x20, 0x74, 0x42, 0x70, 0x42, 0x2c, 0x20, 0x74, + 0x42, 0x67, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x4d, + 0x45, 0x4d, 0x20, 0x3d, 0x3e, 0x20, 0x53, 0x4d, 0x45, 0x4d, 0x2e, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, + 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, + 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x74, 0x41, + 0x72, 0x41, 0x2c, 0x20, 0x74, 0x41, 0x73, 0x41, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x74, 0x42, 0x72, 0x42, + 0x2c, 0x20, 0x74, 0x42, 0x73, 0x42, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6c, 0x65, 0x61, 0x72, 0x20, 0x74, + 0x68, 0x65, 0x20, 0x72, 0x6d, 0x65, 0x6d, 0x20, 0x74, 0x69, 0x6c, 0x65, + 0x73, 0x20, 0x74, 0x6f, 0x20, 0x61, 0x63, 0x63, 0x6f, 0x75, 0x6e, 0x74, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, + 0x74, 0x65, 0x64, 0x20, 0x6f, 0x66, 0x66, 0x20, 0x6c, 0x6f, 0x61, 0x64, + 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, + 0x41, 0x72, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, + 0x72, 0x28, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x50, 0x72, 0x65, 0x66, 0x65, 0x74, 0x63, 0x68, 0x20, + 0x66, 0x69, 0x72, 0x73, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x67, + 0x41, 0x5f, 0x6b, 0x20, 0x3d, 0x20, 0x74, 0x41, 0x67, 0x41, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x42, 0x5f, 0x6b, + 0x20, 0x3d, 0x20, 0x74, 0x42, 0x67, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, + 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, + 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x30, + 0x3b, 0x20, 0x6b, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, + 0x3e, 0x28, 0x74, 0x41, 0x72, 0x41, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6b, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, + 0x67, 0x65, 0x74, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x41, 0x63, 0x41, 0x28, + 0x30, 0x2c, 0x30, 0x2c, 0x6b, 0x29, 0x29, 0x20, 0x3e, 0x3d, 0x20, 0x2d, + 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2c, 0x20, 0x74, + 0x41, 0x70, 0x41, 0x28, 0x5f, 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x41, + 0x67, 0x41, 0x5f, 0x6b, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, 0x29, 0x2c, + 0x20, 0x74, 0x41, 0x72, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, 0x29, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, + 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6b, 0x20, 0x3c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x42, 0x72, + 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x67, 0x65, 0x74, 0x3c, 0x31, + 0x3e, 0x28, 0x74, 0x42, 0x63, 0x42, 0x28, 0x30, 0x2c, 0x30, 0x2c, 0x6b, + 0x29, 0x29, 0x20, 0x3e, 0x3d, 0x20, 0x2d, 0x6b, 0x5f, 0x72, 0x65, 0x73, + 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x62, 0x2c, 0x20, 0x74, 0x42, 0x70, 0x42, 0x28, 0x5f, + 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x67, 0x42, 0x5f, 0x6b, 0x28, + 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x42, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x43, 0x6c, 0x65, 0x61, 0x72, 0x20, 0x61, 0x63, 0x63, 0x75, + 0x6d, 0x75, 0x6c, 0x61, 0x74, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, + 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, 0x43, 0x72, 0x43, 0x29, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x20, + 0x6f, 0x76, 0x65, 0x72, 0x20, 0x43, 0x54, 0x41, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x4b, + 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3d, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x33, 0x3e, 0x28, 0x74, 0x41, 0x67, 0x41, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, + 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x74, 0x69, 0x6c, + 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x74, 0x6f, + 0x72, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x41, 0x76, 0x6f, 0x69, 0x64, 0x20, + 0x66, 0x65, 0x74, 0x63, 0x68, 0x69, 0x6e, 0x67, 0x20, 0x66, 0x75, 0x6c, + 0x6c, 0x20, 0x30, 0x74, 0x68, 0x2d, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x77, + 0x68, 0x65, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x72, 0x65, 0x20, 0x69, 0x73, + 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2e, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, + 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, + 0x67, 0x6d, 0x65, 0x6d, 0x28, 0x28, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x2b, + 0x20, 0x31, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, + 0x4d, 0x41, 0x58, 0x29, 0x20, 0x3f, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, + 0x2b, 0x20, 0x31, 0x20, 0x3a, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, + 0x65, 0x6d, 0x6d, 0x28, 0x6d, 0x6d, 0x61, 0x2c, 0x20, 0x74, 0x43, 0x73, + 0x41, 0x2c, 0x20, 0x74, 0x43, 0x73, 0x42, 0x2c, 0x20, 0x74, 0x43, 0x72, + 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x45, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, + 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x28, 0x74, 0x43, 0x72, 0x43, 0x29, 0x3b, 0x20, + 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x28, 0x67, 0x65, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x74, + 0x43, 0x63, 0x43, 0x28, 0x69, 0x29, 0x29, 0x20, 0x3c, 0x20, 0x6d, 0x5f, + 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x20, 0x26, + 0x26, 0x20, 0x28, 0x67, 0x65, 0x74, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x43, + 0x63, 0x43, 0x28, 0x69, 0x29, 0x29, 0x20, 0x3c, 0x20, 0x6e, 0x5f, 0x6d, + 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x67, 0x43, 0x28, + 0x69, 0x29, 0x20, 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x28, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x3e, 0x0a, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x2c, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x74, + 0x65, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3a, 0x3a, + 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, + 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3a, 0x3a, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x63, 0x75, + 0x74, 0x65, 0x3a, 0x3a, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, + 0x2c, 0x20, 0x6d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, + 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x41, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x2c, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x65, 0x64, 0x2c, 0x20, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, + 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, + 0x61, 0x75, 0x6e, 0x63, 0x68, 0x5f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, + 0x5f, 0x5f, 0x28, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, + 0x73, 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x3c, 0x53, 0x4d, 0x38, 0x30, + 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x28, 0x43, + 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, 0x29, 0x29, 0x29, + 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, + 0x64, 0x20, 0x67, 0x65, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x37, 0x30, 0x5f, + 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2a, 0x20, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, + 0x20, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, + 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, 0x73, + 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x43, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x2c, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, + 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, + 0x64, 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, + 0x7a, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, + 0x20, 0x28, 0x6d, 0x69, 0x78, 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x41, + 0x3e, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, + 0x42, 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, + 0x6e, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, + 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, 0x29, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, 0x65, + 0x73, 0x65, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, 0x6c, + 0x6c, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, + 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, + 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x41, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x5f, + 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x42, 0x29, 0x2c, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, + 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, + 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, + 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, + 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x46, 0x6f, 0x72, 0x20, + 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x2c, 0x20, 0x75, 0x73, 0x65, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x6c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x20, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x69, 0x6e, 0x67, + 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, + 0x20, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x6c, + 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, + 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, + 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, + 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, 0x68, 0x73, 0x5f, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, + 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, + 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x6d, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, + 0x5f, 0x2c, 0x5f, 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, + 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, + 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, + 0x6d, 0x6e, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, + 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, + 0x20, 0x74, 0x68, 0x65, 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, + 0x69, 0x61, 0x74, 0x65, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x20, 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, + 0x65, 0x72, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, + 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x6d, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, + 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, + 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, + 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, + 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, + 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, + 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, + 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, + 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x20, 0x3d, 0x20, 0x6d, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x30, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, + 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, 0x5f, + 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, + 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, + 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, + 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, + 0x41, 0x29, 0x20, 0x2a, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x67, 0x41, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x67, 0x65, 0x6d, + 0x6d, 0x5f, 0x73, 0x6d, 0x37, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, + 0x6f, 0x6f, 0x70, 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x41, 0x2c, + 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x42, 0x2c, 0x20, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x65, 0x64, 0x2c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x3e, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, + 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, + 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, + 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, + 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, + 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + constexpr char jit_source_hadamard[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, @@ -2249,181 +3648,2338 @@ constexpr char jit_source_indexing[] = { 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, - 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x75, 0x6e, 0x73, 0x69, - 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x3e, - 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, - 0x72, 0x6e, 0x20, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, - 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, - 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, - 0x74, 0x3e, 0x28, 0x69, 0x64, 0x78, 0x20, 0x3c, 0x20, 0x30, 0x20, 0x3f, - 0x20, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, - 0x3a, 0x20, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, - 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, - 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, -}; - -constexpr char jit_source_scatter[] = { - 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, - 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, - 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, - 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, - 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, - 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, - 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, - 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, - 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, - 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, - 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, - 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, - 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, - 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, - 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, - 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, - 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, - 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, - 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x49, 0x44, - 0x58, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, - 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, - 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, - 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, - 0x74, 0x65, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, 0x70, 0x64, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, - 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, - 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, - 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, - 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x75, - 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, - 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x70, - 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, - 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, - 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, - 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, - 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, - 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, - 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, - 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, - 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, - 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, - 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, - 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, - 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, - 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, - 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, - 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, - 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, - 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, - 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, - 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, - 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, - 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, - 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, - 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, - 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, - 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, - 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x69, 0x66, 0x20, 0x28, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, - 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, - 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, - 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x25, 0x20, 0x75, 0x70, 0x64, - 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, - 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, 0x70, - 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x75, 0x70, 0x64, 0x5f, - 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, - 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, - 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, - 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, - 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x49, 0x44, 0x58, - 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, - 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, - 0x29, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, - 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, - 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, - 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, - 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, - 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, - 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, - 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, - 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, - 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, - 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, - 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, - 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, - 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, - 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, - 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, - 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, - 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, - 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, - 0x63, 0x5d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x3d, - 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, - 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, - 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, - 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2b, 0x20, 0x69, 0x64, - 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2a, 0x20, 0x75, 0x70, 0x64, - 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, - 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, - 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, - 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, - 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, - 0x74, 0x61, 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, - 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, 0x20, 0x2b, - 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x75, 0x70, - 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x29, 0x3b, - 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, - 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, - 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x75, 0x6e, 0x73, 0x69, + 0x67, 0x6e, 0x65, 0x64, 0x5f, 0x76, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x3e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, + 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, + 0x74, 0x3e, 0x28, 0x69, 0x64, 0x78, 0x20, 0x3c, 0x20, 0x30, 0x20, 0x3f, + 0x20, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, + 0x3a, 0x20, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_qmm_naive[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x63, 0x75, 0x74, 0x65, 0x5f, + 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, + 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, + 0x75, 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x67, + 0x65, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x37, 0x30, 0x2e, 0x63, 0x75, 0x68, + 0x22, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, + 0x6d, 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, + 0x63, 0x75, 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, + 0x75, 0x74, 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x63, 0x6d, 0x61, 0x74, 0x68, 0x3e, 0x0a, 0x0a, 0x2f, + 0x2f, 0x20, 0x63, 0x6c, 0x61, 0x6e, 0x67, 0x2d, 0x66, 0x6f, 0x72, 0x6d, + 0x61, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x20, 0x48, 0x61, 0x73, 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, + 0x65, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x4d, 0x38, 0x30, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, + 0x54, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x41, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x5a, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x43, 0x3e, 0x0a, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x44, 0x45, 0x56, 0x49, + 0x43, 0x45, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, + 0x6e, 0x61, 0x69, 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, + 0x6f, 0x70, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x74, 0x61, 0x54, + 0x69, 0x6c, 0x65, 0x72, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x41, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x42, 0x20, 0x67, 0x42, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x20, + 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x5a, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x43, 0x20, 0x67, 0x43, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x66, 0x6c, + 0x6f, 0x61, 0x74, 0x2a, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, + 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, + 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, + 0x64, 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, + 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, + 0x20, 0x6f, 0x66, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x6e, 0x64, 0x73, + 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x6c, + 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, + 0x65, 0x28, 0x67, 0x41, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, + 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, + 0x6e, 0x67, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, + 0x74, 0x79, 0x70, 0x65, 0x28, 0x67, 0x42, 0x29, 0x3a, 0x3a, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, + 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x20, + 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, + 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x67, 0x53, 0x29, 0x3a, + 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x53, 0x68, 0x69, 0x66, 0x74, + 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x6f, 0x20, 0x77, + 0x65, 0x20, 0x68, 0x61, 0x6e, 0x64, 0x6c, 0x65, 0x20, 0x72, 0x65, 0x73, + 0x69, 0x64, 0x75, 0x65, 0x20, 0x6f, 0x66, 0x20, 0x4b, 0x20, 0x69, 0x6e, + 0x20, 0x74, 0x68, 0x65, 0x20, 0x30, 0x74, 0x68, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x64, 0x6f, + 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, + 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, + 0x20, 0x30, 0x29, 0x2c, 0x20, 0x67, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, 0x69, 0x74, + 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x20, 0x25, + 0x20, 0x38, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x64, 0x6f, 0x6d, 0x61, + 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, + 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x30, + 0x29, 0x2c, 0x20, 0x67, 0x42, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x54, 0x4f, 0x44, 0x4f, 0x3a, 0x20, 0x46, 0x69, 0x67, 0x75, + 0x72, 0x65, 0x20, 0x6f, 0x75, 0x74, 0x20, 0x77, 0x68, 0x79, 0x20, 0x64, + 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x20, 0x69, 0x73, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x69, 0x6e, 0x67, 0x20, 0x77, 0x72, 0x6f, 0x6e, 0x67, 0x20, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x42, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x3d, 0x20, + 0x72, 0x65, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x70, 0x74, 0x72, 0x3c, 0x51, + 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x72, 0x61, 0x77, 0x5f, 0x70, 0x6f, + 0x69, 0x6e, 0x74, 0x65, 0x72, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x28, 0x67, + 0x42, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x20, 0x2b, 0x20, + 0x67, 0x42, 0x2e, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x29, 0x28, + 0x30, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, + 0x2c, 0x20, 0x30, 0x29, 0x20, 0x2a, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, 0x38, 0x2c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, 0x69, 0x74, 0x73, + 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, 0x2f, + 0x20, 0x38, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x67, + 0x53, 0x20, 0x3d, 0x20, 0x64, 0x6f, 0x6d, 0x61, 0x69, 0x6e, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, + 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x30, 0x29, 0x2c, 0x20, 0x67, + 0x53, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x71, 0x75, 0x61, 0x6e, + 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, 0x76, + 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x64, 0x6f, 0x6d, 0x61, + 0x69, 0x6e, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x30, 0x2c, 0x20, + 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x30, + 0x29, 0x2c, 0x20, 0x67, 0x5a, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, + 0x20, 0x73, 0x6d, 0x65, 0x6d, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x73, + 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, 0x42, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x5d, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x73, 0x3c, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, 0x4b, + 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x20, 0x6d, 0x65, 0x6d, 0x6f, + 0x72, 0x79, 0x20, 0x62, 0x75, 0x66, 0x66, 0x65, 0x72, 0x2e, 0x0a, 0x20, + 0x20, 0x65, 0x78, 0x74, 0x65, 0x72, 0x6e, 0x20, 0x5f, 0x5f, 0x73, 0x68, + 0x61, 0x72, 0x65, 0x64, 0x5f, 0x5f, 0x20, 0x63, 0x68, 0x61, 0x72, 0x20, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x62, 0x75, 0x66, 0x5b, 0x5d, 0x3b, 0x0a, + 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x68, 0x61, 0x72, + 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x20, 0x3d, 0x20, + 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, + 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x64, + 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x41, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x2c, 0x20, 0x64, 0x65, 0x63, 0x6c, + 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x29, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x53, 0x68, 0x61, 0x72, + 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x26, 0x20, 0x73, + 0x6d, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x2a, 0x72, 0x65, 0x69, 0x6e, 0x74, + 0x65, 0x72, 0x70, 0x72, 0x65, 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, + 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, + 0x65, 0x2a, 0x3e, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x62, 0x75, 0x66, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x73, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x2e, + 0x41, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, 0x29, 0x2c, 0x20, + 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, + 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x73, 0x6d, 0x65, + 0x6d, 0x2e, 0x42, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, 0x29, + 0x2c, 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x4d, 0x41, 0x2e, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x6d, 0x61, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, + 0x5f, 0x6d, 0x6d, 0x61, 0x3c, 0x53, 0x4d, 0x38, 0x30, 0x2c, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x6d, + 0x61, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x5b, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x5d, + 0x20, 0x3d, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, + 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, 0x6f, 0x70, + 0x79, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, + 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, 0x21, 0x48, 0x61, 0x73, 0x4b, 0x52, + 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x2c, 0x20, 0x62, 0x4d, 0x2c, + 0x20, 0x62, 0x4b, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, 0x65, + 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, + 0x2c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x74, 0x72, + 0x75, 0x65, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x73, 0x2c, 0x20, 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x29, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, 0x61, 0x72, 0x74, + 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x69, 0x6e, 0x67, 0x20, 0x6f, 0x66, 0x20, 0x41, 0x2f, 0x42, + 0x2f, 0x43, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x73, 0x20, 0x61, 0x63, 0x72, + 0x6f, 0x73, 0x73, 0x20, 0x74, 0x68, 0x65, 0x20, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, + 0x70, 0x79, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x61, 0x20, 0x3d, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x67, + 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x67, 0x41, 0x20, + 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, + 0x28, 0x67, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, + 0x50, 0x59, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x41, 0x43, + 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x73, 0x41, 0x20, 0x3d, 0x20, + 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x70, + 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x44, 0x28, 0x73, + 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, 0x59, + 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x41, 0x43, 0x50, 0x59, + 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x74, 0x41, 0x72, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x69, + 0x6b, 0x65, 0x28, 0x74, 0x41, 0x73, 0x41, 0x29, 0x3b, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, 0x59, 0x2c, 0x41, 0x43, 0x50, + 0x59, 0x5f, 0x4d, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x74, + 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x20, 0x3d, 0x20, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, + 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, + 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, 0x61, 0x72, + 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, 0x42, 0x29, + 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, + 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x73, 0x42, + 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x62, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, + 0x44, 0x28, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, + 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x42, 0x72, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x66, + 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x69, 0x6b, 0x65, + 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x74, 0x42, 0x73, 0x42, + 0x29, 0x3b, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, + 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x42, 0x43, 0x50, + 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x42, 0x72, 0x42, 0x5f, 0x64, 0x71, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, + 0x74, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, 0x74, 0x42, 0x73, 0x42, 0x29, + 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, + 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x53, 0x20, 0x3d, 0x20, + 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, + 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, + 0x53, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, + 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, + 0x72, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x66, 0x72, + 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, + 0x74, 0x42, 0x67, 0x53, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, + 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, + 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, + 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x74, 0x42, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, 0x5a, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, + 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, + 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x72, 0x5a, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, + 0x6e, 0x74, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, 0x74, 0x42, 0x67, 0x5a, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, + 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x4d, 0x41, 0x2e, 0x0a, 0x20, + 0x20, 0x54, 0x68, 0x72, 0x4d, 0x4d, 0x41, 0x20, 0x74, 0x68, 0x72, 0x5f, + 0x6d, 0x6d, 0x61, 0x20, 0x3d, 0x20, 0x6d, 0x6d, 0x61, 0x2e, 0x67, 0x65, + 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x73, 0x41, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, + 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x41, 0x28, 0x73, 0x41, 0x29, + 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x43, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, + 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x42, 0x28, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, + 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, + 0x67, 0x43, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x43, + 0x28, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x43, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x43, + 0x28, 0x74, 0x43, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, + 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x65, 0x73, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x6d, 0x2f, 0x6e, 0x20, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, + 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x70, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x41, 0x73, 0x41, 0x29, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x41, + 0x73, 0x41, 0x29, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x30, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, + 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x42, 0x70, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, + 0x6c, 0x3e, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x42, + 0x73, 0x42, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x74, 0x42, 0x73, 0x42, 0x29, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x30, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x5f, 0x4e, + 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x63, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, + 0x3e, 0x28, 0x73, 0x41, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x31, 0x3e, 0x28, 0x73, 0x41, 0x29, 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, + 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x63, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x69, 0x64, + 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x73, 0x42, 0x29, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x73, 0x42, + 0x29, 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x63, 0x43, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, + 0x79, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x30, 0x3e, 0x28, 0x67, 0x43, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, + 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x43, 0x29, 0x29, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x63, 0x41, 0x20, 0x3d, + 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x63, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, + 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x42, 0x63, 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, + 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x63, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x43, 0x50, 0x59, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, + 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x63, 0x43, 0x20, 0x3d, 0x20, 0x74, + 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x43, 0x28, 0x63, 0x43, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, + 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, + 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x6d, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6d, 0x20, 0x3c, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x70, 0x41, + 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x41, 0x70, 0x41, 0x28, 0x6d, 0x2c, 0x30, 0x29, 0x20, + 0x3d, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x63, + 0x41, 0x28, 0x30, 0x2c, 0x6d, 0x2c, 0x30, 0x29, 0x29, 0x20, 0x3c, 0x20, + 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, + 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, + 0x28, 0x74, 0x42, 0x70, 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x42, 0x70, 0x42, 0x28, + 0x6e, 0x2c, 0x30, 0x29, 0x20, 0x3d, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x30, + 0x3e, 0x28, 0x74, 0x42, 0x63, 0x42, 0x28, 0x30, 0x2c, 0x6e, 0x2c, 0x30, + 0x29, 0x29, 0x20, 0x3c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x4d, 0x45, 0x4d, 0x20, 0x3d, 0x3e, 0x20, + 0x52, 0x4d, 0x45, 0x4d, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x20, + 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, + 0x2c, 0x20, 0x74, 0x41, 0x70, 0x41, 0x2c, 0x20, 0x74, 0x41, 0x67, 0x41, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, + 0x2c, 0x20, 0x74, 0x41, 0x72, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, + 0x79, 0x5f, 0x62, 0x2c, 0x20, 0x74, 0x42, 0x70, 0x42, 0x2c, 0x20, 0x74, + 0x42, 0x67, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, + 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, + 0x74, 0x42, 0x72, 0x53, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6e, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, + 0x74, 0x42, 0x70, 0x42, 0x28, 0x6e, 0x2c, 0x30, 0x29, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, + 0x79, 0x28, 0x74, 0x42, 0x67, 0x53, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x5f, + 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x53, + 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x5f, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x71, 0x75, 0x61, 0x6e, + 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, 0x76, + 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, + 0x79, 0x28, 0x74, 0x42, 0x67, 0x5a, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x5f, + 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x5a, + 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x5f, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x4d, 0x45, 0x4d, + 0x20, 0x3d, 0x3e, 0x20, 0x53, 0x4d, 0x45, 0x4d, 0x2e, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x73, + 0x6d, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x74, 0x41, 0x72, 0x41, + 0x2c, 0x20, 0x74, 0x41, 0x73, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6b, 0x20, 0x3c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x42, 0x72, + 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, + 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x31, 0x3e, 0x28, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, + 0x6e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x75, 0x74, 0x65, 0x5f, 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, + 0x74, 0x28, 0x74, 0x42, 0x72, 0x42, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x6b, + 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x53, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, + 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x5a, 0x28, 0x5f, 0x2c, 0x6e, + 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x42, 0x5f, 0x64, 0x71, + 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x74, 0x42, 0x72, + 0x42, 0x5f, 0x64, 0x71, 0x2c, 0x20, 0x74, 0x42, 0x73, 0x42, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6c, 0x65, + 0x61, 0x72, 0x20, 0x74, 0x68, 0x65, 0x20, 0x72, 0x6d, 0x65, 0x6d, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x73, 0x20, 0x74, 0x6f, 0x20, 0x61, 0x63, 0x63, + 0x6f, 0x75, 0x6e, 0x74, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x70, 0x72, 0x65, + 0x64, 0x69, 0x63, 0x61, 0x74, 0x65, 0x64, 0x20, 0x6f, 0x66, 0x66, 0x20, + 0x6c, 0x6f, 0x61, 0x64, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x48, + 0x61, 0x73, 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, + 0x74, 0x41, 0x72, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, 0x42, + 0x72, 0x53, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6c, 0x65, + 0x61, 0x72, 0x28, 0x74, 0x42, 0x72, 0x5a, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, 0x72, 0x65, 0x66, + 0x65, 0x74, 0x63, 0x68, 0x20, 0x66, 0x69, 0x72, 0x73, 0x74, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x48, 0x61, 0x73, + 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x67, 0x41, 0x5f, 0x6b, 0x20, 0x3d, 0x20, 0x74, 0x41, 0x67, 0x41, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, + 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6b, + 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, + 0x41, 0x72, 0x41, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6b, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x67, + 0x65, 0x74, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x41, 0x63, 0x41, 0x28, 0x30, + 0x2c, 0x30, 0x2c, 0x6b, 0x29, 0x29, 0x20, 0x3e, 0x3d, 0x20, 0x2d, 0x6b, + 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2c, 0x20, + 0x74, 0x41, 0x70, 0x41, 0x28, 0x5f, 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, + 0x41, 0x67, 0x41, 0x5f, 0x6b, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, 0x29, + 0x2c, 0x20, 0x74, 0x41, 0x72, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6b, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x42, 0x5f, 0x6b, 0x20, + 0x3d, 0x20, 0x74, 0x42, 0x67, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, + 0x2c, 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x53, 0x5f, 0x6b, 0x20, 0x3d, + 0x20, 0x74, 0x42, 0x67, 0x53, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, + 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x42, 0x67, 0x5a, 0x5f, 0x6b, 0x20, 0x3d, 0x20, + 0x74, 0x42, 0x67, 0x5a, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, + 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x6b, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, + 0x32, 0x3e, 0x28, 0x74, 0x42, 0x72, 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, + 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x28, 0x67, 0x65, 0x74, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x42, + 0x63, 0x42, 0x28, 0x30, 0x2c, 0x30, 0x2c, 0x6b, 0x29, 0x29, 0x20, 0x3e, + 0x3d, 0x20, 0x2d, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, 0x66, 0x28, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x62, 0x2c, 0x20, 0x74, 0x42, 0x70, 0x42, 0x28, 0x5f, 0x2c, 0x6b, + 0x29, 0x2c, 0x20, 0x74, 0x42, 0x67, 0x42, 0x5f, 0x6b, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x42, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, + 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x20, 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x42, 0x72, 0x53, 0x29, 0x3b, 0x20, 0x2b, + 0x2b, 0x6e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x74, 0x42, 0x70, 0x42, + 0x28, 0x6e, 0x2c, 0x6b, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, + 0x79, 0x28, 0x74, 0x42, 0x67, 0x53, 0x5f, 0x6b, 0x28, 0x5f, 0x2c, 0x6e, + 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, 0x72, 0x53, 0x28, 0x5f, 0x2c, + 0x6e, 0x2c, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x71, 0x75, 0x61, + 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, + 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x74, 0x42, 0x67, 0x5a, 0x5f, + 0x6b, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x42, + 0x72, 0x5a, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x6b, 0x29, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x28, 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6c, 0x65, 0x61, 0x72, 0x20, + 0x61, 0x63, 0x63, 0x75, 0x6d, 0x75, 0x6c, 0x61, 0x74, 0x6f, 0x72, 0x73, + 0x2e, 0x0a, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, 0x43, + 0x72, 0x43, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4c, + 0x6f, 0x6f, 0x70, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x43, 0x54, 0x41, + 0x20, 0x74, 0x69, 0x6c, 0x65, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, + 0x58, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x33, 0x3e, 0x28, + 0x74, 0x41, 0x67, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3c, 0x20, 0x4b, + 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x3b, 0x20, 0x2b, + 0x2b, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x28, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x48, 0x61, 0x73, + 0x4b, 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x41, 0x76, 0x6f, + 0x69, 0x64, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x69, 0x6e, 0x67, 0x20, + 0x66, 0x75, 0x6c, 0x6c, 0x20, 0x30, 0x74, 0x68, 0x2d, 0x74, 0x69, 0x6c, + 0x65, 0x20, 0x77, 0x68, 0x65, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x72, 0x65, + 0x20, 0x69, 0x73, 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2e, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x4b, + 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3e, 0x20, + 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x28, + 0x28, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3c, 0x20, + 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x29, 0x20, + 0x3f, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3a, + 0x20, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x28, 0x28, 0x74, + 0x69, 0x6c, 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3c, 0x20, 0x4b, 0x5f, + 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x29, 0x20, 0x3f, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3a, 0x20, 0x74, + 0x69, 0x6c, 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x6d, 0x6d, 0x28, 0x6d, 0x6d, 0x61, + 0x2c, 0x20, 0x74, 0x43, 0x73, 0x41, 0x2c, 0x20, 0x74, 0x43, 0x73, 0x42, + 0x2c, 0x20, 0x74, 0x43, 0x72, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x45, 0x70, 0x69, 0x6c, 0x6f, + 0x67, 0x75, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, + 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x69, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x74, 0x43, + 0x72, 0x43, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x28, 0x67, 0x65, 0x74, + 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x43, 0x63, 0x43, 0x28, 0x69, 0x29, 0x29, + 0x20, 0x3c, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x29, 0x20, 0x26, 0x26, 0x20, 0x28, 0x67, 0x65, 0x74, 0x3c, + 0x31, 0x3e, 0x28, 0x74, 0x43, 0x63, 0x43, 0x28, 0x69, 0x29, 0x29, 0x20, + 0x3c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x29, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x51, 0x75, + 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, + 0x3a, 0x3a, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x5f, 0x65, 0x32, 0x6d, 0x31, + 0x5f, 0x74, 0x3e, 0x20, 0x26, 0x26, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, + 0x76, 0x3c, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x63, 0x75, 0x74, + 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x5f, + 0x65, 0x34, 0x6d, 0x33, 0x5f, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4f, 0x6e, + 0x6c, 0x79, 0x20, 0x6e, 0x76, 0x66, 0x70, 0x34, 0x20, 0x73, 0x75, 0x70, + 0x70, 0x6f, 0x72, 0x74, 0x73, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x20, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x67, 0x6c, 0x6f, 0x62, + 0x61, 0x6c, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, + 0x67, 0x43, 0x28, 0x69, 0x29, 0x20, 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, + 0x65, 0x6e, 0x74, 0x28, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x20, + 0x2a, 0x20, 0x28, 0x2a, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x73, + 0x63, 0x61, 0x6c, 0x65, 0x20, 0x2f, 0x20, 0x28, 0x46, 0x38, 0x45, 0x34, + 0x4d, 0x33, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x2a, 0x20, 0x46, 0x34, 0x45, + 0x32, 0x4d, 0x31, 0x5f, 0x4d, 0x41, 0x58, 0x29, 0x29, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x74, 0x69, 0x6e, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x67, 0x43, + 0x28, 0x69, 0x29, 0x20, 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x28, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, + 0x72, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6c, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x20, 0x3d, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, + 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, + 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x28, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, + 0x5f, 0x30, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, + 0x2a, 0x20, 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, + 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x20, 0x6e, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x30, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, + 0x2c, 0x20, 0x6e, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x20, + 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, + 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, + 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x48, 0x61, 0x73, 0x4b, 0x52, 0x65, + 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x53, 0x4d, 0x38, 0x30, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, + 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, + 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x0a, + 0x5f, 0x5f, 0x6c, 0x61, 0x75, 0x6e, 0x63, 0x68, 0x5f, 0x62, 0x6f, 0x75, + 0x6e, 0x64, 0x73, 0x5f, 0x5f, 0x28, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, + 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x28, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x3c, 0x53, + 0x4d, 0x38, 0x30, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x3e, 0x28, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, + 0x29, 0x29, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x29, 0x0a, + 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, 0x69, + 0x76, 0x65, 0x5f, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, 0x6c, 0x65, + 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, + 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2a, 0x20, 0x53, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x5a, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x66, 0x6c, 0x6f, + 0x61, 0x74, 0x2a, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x73, + 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, + 0x2a, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, + 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, + 0x20, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x6d, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, + 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, + 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, + 0x64, 0x78, 0x2e, 0x79, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, + 0x7a, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, + 0x20, 0x28, 0x6d, 0x69, 0x78, 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, + 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x3c, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x3e, + 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, + 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x64, 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x63, 0x61, + 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x47, + 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x4b, 0x4d, + 0x61, 0x6a, 0x6f, 0x72, 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, + 0x6e, 0x64, 0x6c, 0x65, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, + 0x73, 0x74, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x28, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, + 0x32, 0x3e, 0x28, 0x64, 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x29, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, + 0x65, 0x73, 0x65, 0x6e, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, + 0x6c, 0x6c, 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, + 0x6d, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, + 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, + 0x70, 0x74, 0x72, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, + 0x29, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, + 0x64, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, + 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, + 0x43, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, + 0x20, 0x6e, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, + 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, + 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x46, 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, + 0x68, 0x65, 0x72, 0x2c, 0x20, 0x75, 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x20, 0x6c, 0x6f, 0x6f, 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, + 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, + 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, + 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, + 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, + 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, + 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, + 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, + 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, + 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, + 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, + 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, + 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x61, + 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, 0x20, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x74, 0x68, + 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, 0x62, 0x6c, + 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, 0x3d, 0x20, + 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, 0x3b, 0x0a, + 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, + 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, 0x6e, 0x2c, + 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, + 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, + 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x6c, 0x6f, + 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x42, 0x2c, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, + 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, + 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, + 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x43, + 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, + 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, + 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, 0x31, 0x2c, + 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, + 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, + 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, + 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, + 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x5b, 0x26, + 0x5d, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, + 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, + 0x61, 0x73, 0x5f, 0x76, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, + 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, + 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, + 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, + 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x75, 0x6d, 0x6d, 0x79, + 0x20, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3b, 0x20, 0x6e, 0x6f, 0x2d, + 0x62, 0x69, 0x61, 0x73, 0x20, 0x70, 0x61, 0x74, 0x68, 0x73, 0x20, 0x6e, + 0x65, 0x76, 0x65, 0x72, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, + 0x6f, 0x72, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x67, 0x5a, 0x2e, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x67, 0x53, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x7d, 0x28, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, 0x65, + 0x20, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, + 0x6e, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x6d, + 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, + 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, + 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, 0x2a, 0x20, + 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x20, 0x2a, 0x20, + 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x0a, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x20, 0x3d, 0x20, 0x6e, 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x31, 0x3e, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x29, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, + 0x5f, 0x4e, 0x20, 0x2a, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, + 0x69, 0x64, 0x75, 0x65, 0x20, 0x3d, 0x20, 0x6b, 0x20, 0x2d, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x67, 0x41, 0x29, 0x20, 0x2a, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x67, 0x41, 0x29, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x6e, 0x61, 0x69, + 0x76, 0x65, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x3c, + 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x20, 0x48, 0x61, 0x73, 0x4b, + 0x52, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x2c, 0x20, 0x53, 0x4d, 0x38, + 0x30, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, + 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, + 0x73, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x2c, 0x20, 0x6e, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x2c, 0x20, 0x6b, 0x5f, 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, + 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, + 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, + 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_qmm_sm80[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x63, 0x75, 0x74, 0x65, 0x5f, + 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, + 0x63, 0x75, 0x64, 0x61, 0x2f, 0x63, 0x6d, 0x61, 0x74, 0x68, 0x3e, 0x0a, + 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x6c, 0x61, 0x6e, 0x67, 0x2d, 0x66, 0x6f, + 0x72, 0x6d, 0x61, 0x74, 0x20, 0x6f, 0x66, 0x66, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, + 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, + 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, + 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x53, 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x41, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x6d, 0x65, + 0x6d, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x42, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x53, 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x43, 0x3e, 0x0a, 0x75, 0x6e, 0x69, 0x6f, 0x6e, 0x20, + 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, + 0x45, 0x6e, 0x67, 0x69, 0x6e, 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x2c, 0x20, 0x63, 0x6f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x76, + 0x3c, 0x53, 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x41, + 0x3e, 0x3e, 0x20, 0x41, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x41, 0x72, + 0x72, 0x61, 0x79, 0x45, 0x6e, 0x67, 0x69, 0x6e, 0x65, 0x3c, 0x51, 0x75, + 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x73, 0x69, 0x7a, + 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x6d, 0x65, 0x6d, 0x4c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x42, 0x3e, 0x3e, 0x20, 0x42, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x20, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x3b, 0x0a, 0x20, + 0x20, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x41, 0x72, 0x72, 0x61, 0x79, 0x45, 0x6e, 0x67, 0x69, 0x6e, + 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x63, + 0x6f, 0x73, 0x69, 0x7a, 0x65, 0x5f, 0x76, 0x3c, 0x53, 0x6d, 0x65, 0x6d, + 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x43, 0x3e, 0x3e, 0x20, 0x43, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, + 0x65, 0x3b, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, + 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x73, 0x28, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, + 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x6f, 0x74, 0x65, 0x3a, + 0x20, 0x4b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, 0x6c, 0x61, 0x75, 0x6e, + 0x63, 0x68, 0x65, 0x72, 0x20, 0x61, 0x73, 0x73, 0x75, 0x6d, 0x65, 0x73, + 0x20, 0x63, 0x6f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x62, 0x65, 0x69, 0x6e, + 0x67, 0x20, 0x73, 0x61, 0x6d, 0x65, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x61, + 0x6c, 0x6c, 0x20, 0x4b, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2e, 0x0a, 0x20, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x62, 0x4d, 0x2c, 0x20, 0x62, + 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x5d, 0x20, 0x3d, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x74, 0x68, 0x65, + 0x20, 0x41, 0x2f, 0x42, 0x20, 0x73, 0x6d, 0x65, 0x6d, 0x20, 0x6c, 0x61, + 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x73, 0x74, 0x61, 0x74, 0x69, + 0x63, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x73, + 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x5f, 0x61, 0x62, 0x20, 0x3d, 0x20, + 0x63, 0x6f, 0x6d, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x28, + 0x53, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x3c, 0x33, 0x2c, 0x33, 0x2c, + 0x33, 0x3e, 0x7b, 0x7d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3c, 0x53, 0x68, 0x61, + 0x70, 0x65, 0x20, 0x3c, 0x5f, 0x38, 0x2c, 0x53, 0x68, 0x61, 0x70, 0x65, + 0x20, 0x3c, 0x5f, 0x38, 0x2c, 0x20, 0x5f, 0x38, 0x3e, 0x3e, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x38, + 0x2c, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, + 0x36, 0x34, 0x3e, 0x3e, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x62, 0x50, 0x20, 0x3d, 0x20, 0x49, 0x6e, + 0x74, 0x3c, 0x33, 0x3e, 0x7b, 0x7d, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x70, + 0x69, 0x70, 0x65, 0x6c, 0x69, 0x6e, 0x65, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x20, 0x3d, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x74, 0x6f, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, + 0x5f, 0x61, 0x62, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x28, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4b, 0x2c, 0x20, + 0x62, 0x50, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, + 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x74, 0x6f, 0x5f, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x28, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x5f, 0x61, + 0x62, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x62, 0x4e, 0x2c, 0x20, 0x62, 0x4b, 0x2c, 0x20, 0x62, 0x50, + 0x29, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x20, 0x74, 0x68, 0x65, 0x20, 0x43, 0x20, 0x73, + 0x6d, 0x65, 0x6d, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, + 0x28, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x29, 0x2e, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x54, 0x4f, 0x44, 0x4f, 0x3a, 0x20, 0x46, 0x69, 0x6e, + 0x64, 0x20, 0x61, 0x20, 0x62, 0x65, 0x74, 0x74, 0x65, 0x72, 0x20, 0x73, + 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x73, 0x43, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x20, 0x3d, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x74, 0x6f, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x77, 0x69, 0x7a, 0x7a, 0x6c, 0x65, + 0x5f, 0x61, 0x62, 0x2c, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x28, 0x62, 0x4d, 0x2c, 0x20, 0x62, 0x4e, 0x29, 0x29, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x63, 0x75, 0x74, 0x65, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x75, 0x70, 0x6c, 0x65, 0x28, 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x73, 0x43, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x68, + 0x61, 0x6c, 0x66, 0x5f, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x54, + 0x69, 0x6c, 0x65, 0x4d, 0x20, 0x3d, 0x20, 0x33, 0x32, 0x3e, 0x0a, 0x69, + 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4e, 0x6f, 0x74, + 0x65, 0x3a, 0x20, 0x4b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x20, 0x6c, 0x61, + 0x75, 0x6e, 0x63, 0x68, 0x65, 0x72, 0x20, 0x61, 0x73, 0x73, 0x75, 0x6d, + 0x65, 0x73, 0x20, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x20, 0x62, 0x65, 0x69, 0x6e, 0x67, 0x20, 0x73, 0x61, 0x6d, + 0x65, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x61, 0x6c, 0x6c, 0x20, 0x70, 0x61, + 0x72, 0x61, 0x6d, 0x65, 0x74, 0x65, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, + 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x41, 0x74, 0x6f, 0x6d, 0x20, 0x3d, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x63, 0x6f, 0x6e, 0x64, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x61, 0x6c, 0x5f, + 0x74, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, + 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2c, 0x20, 0x68, 0x61, 0x6c, 0x66, 0x5f, 0x74, 0x3e, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x5f, 0x31, + 0x36, 0x78, 0x38, 0x78, 0x31, 0x36, 0x5f, 0x46, 0x33, 0x32, 0x46, 0x31, + 0x36, 0x46, 0x31, 0x36, 0x46, 0x33, 0x32, 0x5f, 0x54, 0x4e, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x5f, 0x31, + 0x36, 0x78, 0x38, 0x78, 0x31, 0x36, 0x5f, 0x46, 0x33, 0x32, 0x42, 0x46, + 0x31, 0x36, 0x42, 0x46, 0x31, 0x36, 0x46, 0x33, 0x32, 0x5f, 0x54, 0x4e, + 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x54, 0x69, 0x6c, 0x65, 0x4d, + 0x20, 0x3e, 0x3d, 0x20, 0x33, 0x32, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, + 0x41, 0x74, 0x6f, 0x6d, 0x7b, 0x7d, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x3c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x32, 0x2c, + 0x5f, 0x32, 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x54, + 0x69, 0x6c, 0x65, 0x3c, 0x5f, 0x33, 0x32, 0x2c, 0x5f, 0x33, 0x32, 0x2c, + 0x5f, 0x31, 0x36, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, 0x6d, 0x61, 0x28, 0x41, 0x74, + 0x6f, 0x6d, 0x7b, 0x7d, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x3c, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x34, + 0x2c, 0x5f, 0x31, 0x3e, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x54, 0x69, 0x6c, + 0x65, 0x3c, 0x5f, 0x31, 0x36, 0x2c, 0x5f, 0x33, 0x32, 0x2c, 0x5f, 0x31, + 0x36, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x62, 0x69, 0x74, 0x73, 0x2c, 0x20, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x55, 0x3e, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x41, 0x74, 0x6f, 0x6d, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4e, 0x75, 0x6d, 0x54, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x4e, 0x75, + 0x6d, 0x54, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x6e, 0x75, 0x6d, + 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x43, 0x6f, 0x70, 0x79, + 0x5f, 0x41, 0x74, 0x6f, 0x6d, 0x3c, 0x41, 0x74, 0x6f, 0x6d, 0x3c, 0x75, + 0x69, 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x5f, 0x74, 0x3c, 0x62, 0x69, + 0x74, 0x73, 0x3e, 0x3e, 0x2c, 0x20, 0x54, 0x3e, 0x7b, 0x7d, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x28, 0x49, 0x6e, 0x74, 0x3c, 0x6e, 0x75, 0x6d, + 0x5f, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x2f, 0x20, 0x38, + 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x38, 0x3e, 0x7b, + 0x7d, 0x29, 0x2c, 0x20, 0x4c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x52, 0x69, + 0x67, 0x68, 0x74, 0x7b, 0x7d, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x28, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x49, + 0x6e, 0x74, 0x3c, 0x62, 0x69, 0x74, 0x73, 0x20, 0x2f, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, 0x69, 0x74, 0x73, 0x5f, 0x76, 0x3c, + 0x54, 0x3e, 0x3e, 0x7b, 0x7d, 0x29, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, + 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x41, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x42, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x5a, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x43, + 0x3e, 0x0a, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, + 0x45, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, + 0x6d, 0x38, 0x30, 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, + 0x65, 0x72, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x41, 0x20, 0x67, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x42, 0x20, 0x67, 0x42, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x53, 0x20, 0x67, 0x53, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x5a, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x43, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, + 0x78, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, + 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x20, + 0x6f, 0x66, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x6e, 0x64, 0x73, 0x2e, + 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x6c, 0x65, + 0x6d, 0x65, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, + 0x28, 0x67, 0x41, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, + 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, + 0x67, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x20, 0x3d, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, + 0x79, 0x70, 0x65, 0x28, 0x67, 0x42, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, + 0x75, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x20, 0x3d, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x64, 0x65, + 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x67, 0x53, 0x29, 0x3a, 0x3a, + 0x76, 0x61, 0x6c, 0x75, 0x65, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, + 0x20, 0x73, 0x6d, 0x65, 0x6d, 0x20, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x5b, 0x73, + 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, 0x42, + 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x73, 0x43, 0x5f, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x5d, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x73, 0x28, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x72, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x53, 0x68, + 0x61, 0x72, 0x65, 0x64, 0x20, 0x6d, 0x65, 0x6d, 0x6f, 0x72, 0x79, 0x20, + 0x62, 0x75, 0x66, 0x66, 0x65, 0x72, 0x2e, 0x0a, 0x20, 0x20, 0x65, 0x78, + 0x74, 0x65, 0x72, 0x6e, 0x20, 0x5f, 0x5f, 0x73, 0x68, 0x61, 0x72, 0x65, + 0x64, 0x5f, 0x5f, 0x20, 0x63, 0x68, 0x61, 0x72, 0x20, 0x73, 0x6d, 0x65, + 0x6d, 0x5f, 0x62, 0x75, 0x66, 0x5b, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, + 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x20, 0x3d, 0x20, 0x53, 0x68, 0x61, + 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x3c, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, + 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, + 0x65, 0x28, 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, + 0x28, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, + 0x73, 0x43, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3e, 0x3b, + 0x0a, 0x20, 0x20, 0x53, 0x68, 0x61, 0x72, 0x65, 0x64, 0x53, 0x74, 0x6f, + 0x72, 0x61, 0x67, 0x65, 0x26, 0x20, 0x73, 0x6d, 0x65, 0x6d, 0x20, 0x3d, + 0x20, 0x2a, 0x72, 0x65, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x70, 0x72, 0x65, + 0x74, 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x53, 0x68, 0x61, 0x72, 0x65, + 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x2a, 0x3e, 0x28, 0x73, + 0x6d, 0x65, 0x6d, 0x5f, 0x62, 0x75, 0x66, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x41, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, + 0x72, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x2e, 0x6d, 0x61, 0x69, 0x6e, 0x6c, + 0x6f, 0x6f, 0x70, 0x2e, 0x41, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, + 0x29, 0x29, 0x2c, 0x20, 0x73, 0x41, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, + 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, + 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x2e, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, + 0x6f, 0x70, 0x2e, 0x42, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, + 0x29, 0x2c, 0x20, 0x73, 0x42, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, + 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, + 0x73, 0x6d, 0x65, 0x6d, 0x2e, 0x65, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, + 0x65, 0x2e, 0x43, 0x2e, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x28, 0x29, 0x29, + 0x2c, 0x20, 0x73, 0x43, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x4d, 0x4d, 0x41, 0x2e, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x6d, 0x61, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, + 0x5f, 0x6d, 0x6d, 0x61, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x29, 0x3e, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6e, 0x75, 0x6d, 0x5f, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x28, 0x6d, 0x6d, 0x61, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, + 0x70, 0x79, 0x20, 0x61, 0x74, 0x6f, 0x6d, 0x73, 0x2e, 0x0a, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x62, 0x69, + 0x74, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, + 0x62, 0x69, 0x74, 0x73, 0x5f, 0x76, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x71, 0x75, 0x61, + 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x73, 0x20, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, 0x69, 0x74, 0x73, 0x5f, 0x76, 0x3c, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x71, 0x6c, 0x6f, 0x61, 0x64, 0x20, 0x3d, 0x20, 0x31, 0x32, 0x38, 0x20, + 0x2f, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x62, + 0x69, 0x74, 0x73, 0x20, 0x2f, 0x20, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, + 0x62, 0x69, 0x74, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, + 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x3c, + 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x31, 0x32, 0x38, + 0x2c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x5f, 0x43, 0x50, 0x5f, 0x41, 0x53, + 0x59, 0x4e, 0x43, 0x5f, 0x43, 0x41, 0x43, 0x48, 0x45, 0x41, 0x4c, 0x57, + 0x41, 0x59, 0x53, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, + 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x3c, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x71, 0x6c, 0x6f, 0x61, 0x64, + 0x2c, 0x20, 0x53, 0x4d, 0x38, 0x30, 0x5f, 0x43, 0x50, 0x5f, 0x41, 0x53, + 0x59, 0x4e, 0x43, 0x5f, 0x43, 0x41, 0x43, 0x48, 0x45, 0x41, 0x4c, 0x57, + 0x41, 0x59, 0x53, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, + 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x3c, + 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x31, 0x32, 0x38, + 0x2c, 0x20, 0x55, 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x43, + 0x6f, 0x70, 0x79, 0x3e, 0x28, 0x6e, 0x75, 0x6d, 0x5f, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x73, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x43, 0x6f, + 0x70, 0x79, 0x5f, 0x41, 0x74, 0x6f, 0x6d, 0x3c, 0x53, 0x4d, 0x37, 0x35, + 0x5f, 0x55, 0x33, 0x32, 0x78, 0x34, 0x5f, 0x4c, 0x44, 0x53, 0x4d, 0x5f, + 0x4e, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x20, + 0x73, 0x32, 0x72, 0x5f, 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x61, 0x3b, 0x0a, + 0x20, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x5f, 0x41, 0x74, 0x6f, 0x6d, 0x3c, + 0x55, 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x43, 0x6f, 0x70, + 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x5f, 0x74, + 0x3c, 0x32, 0x20, 0x2a, 0x20, 0x71, 0x75, 0x61, 0x6e, 0x74, 0x5f, 0x62, + 0x69, 0x74, 0x73, 0x3e, 0x3e, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, + 0x3e, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x62, + 0x3b, 0x0a, 0x20, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x5f, 0x41, 0x74, 0x6f, + 0x6d, 0x3c, 0x55, 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x43, + 0x6f, 0x70, 0x79, 0x3c, 0x75, 0x69, 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, + 0x5f, 0x74, 0x3c, 0x32, 0x20, 0x2a, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x5f, 0x62, 0x69, 0x74, 0x73, 0x3e, 0x3e, 0x2c, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x20, 0x72, 0x32, 0x73, 0x5f, + 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x63, 0x3b, 0x0a, 0x20, 0x20, 0x43, 0x6f, + 0x70, 0x79, 0x5f, 0x41, 0x74, 0x6f, 0x6d, 0x3c, 0x55, 0x6e, 0x69, 0x76, + 0x65, 0x72, 0x73, 0x61, 0x6c, 0x43, 0x6f, 0x70, 0x79, 0x3c, 0x53, 0x63, + 0x61, 0x6c, 0x65, 0x3e, 0x2c, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x3e, + 0x20, 0x67, 0x32, 0x72, 0x5f, 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x73, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, 0x61, 0x72, 0x74, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x63, 0x6f, 0x70, + 0x79, 0x69, 0x6e, 0x67, 0x20, 0x6f, 0x66, 0x20, 0x41, 0x2f, 0x42, 0x2f, + 0x43, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x73, 0x20, 0x61, 0x63, 0x72, 0x6f, + 0x73, 0x73, 0x20, 0x74, 0x68, 0x65, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, + 0x79, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, + 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, + 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x41, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x67, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, + 0x59, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x41, 0x43, 0x50, + 0x59, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x74, 0x41, 0x73, 0x41, 0x20, 0x3d, 0x20, 0x67, + 0x32, 0x73, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, + 0x44, 0x28, 0x73, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, + 0x43, 0x50, 0x59, 0x2c, 0x41, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x41, + 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x50, 0x49, 0x50, 0x45, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x67, + 0x32, 0x73, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x62, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, + 0x5f, 0x62, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, + 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x42, 0x67, 0x42, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x74, 0x68, + 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x70, 0x61, 0x72, + 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, 0x42, 0x29, + 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, + 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, + 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x74, 0x42, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x44, 0x28, + 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, + 0x50, 0x59, 0x2c, 0x42, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x2c, 0x42, 0x43, + 0x50, 0x59, 0x5f, 0x4b, 0x2c, 0x50, 0x49, 0x50, 0x45, 0x29, 0x0a, 0x0a, + 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, 0x32, + 0x67, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, + 0x20, 0x3d, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x63, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x32, + 0x67, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x20, 0x3d, 0x20, 0x73, 0x32, 0x67, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x73, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x43, 0x50, + 0x59, 0x2c, 0x43, 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x43, 0x50, + 0x59, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, 0x67, 0x43, 0x20, 0x3d, + 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, + 0x79, 0x5f, 0x63, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, + 0x6e, 0x5f, 0x44, 0x28, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x43, 0x43, 0x50, 0x59, 0x2c, 0x43, 0x43, 0x50, 0x59, 0x5f, 0x4d, + 0x2c, 0x43, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x4d, 0x4d, 0x41, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x68, + 0x72, 0x4d, 0x4d, 0x41, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, + 0x20, 0x3d, 0x20, 0x6d, 0x6d, 0x61, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, + 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x41, 0x20, 0x3d, 0x20, 0x74, 0x68, + 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, + 0x5f, 0x41, 0x28, 0x73, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x73, + 0x42, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x42, 0x28, + 0x73, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, + 0x6c, 0x69, 0x6b, 0x65, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, + 0x74, 0x43, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, + 0x42, 0x5f, 0x64, 0x71, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x69, 0x6b, + 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x28, 0x74, + 0x43, 0x73, 0x42, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x43, 0x67, 0x43, 0x20, 0x3d, 0x20, 0x74, 0x68, + 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x5f, 0x43, 0x28, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, + 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, 0x6c, + 0x69, 0x6b, 0x65, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x28, 0x74, + 0x43, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x66, 0x72, 0x61, 0x67, 0x6d, 0x65, 0x6e, 0x74, 0x5f, + 0x6c, 0x69, 0x6b, 0x65, 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x3e, 0x28, 0x74, 0x43, 0x67, 0x43, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, + 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x29, 0x0a, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, + 0x67, 0x53, 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x42, + 0x28, 0x67, 0x53, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, + 0x72, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, 0x74, 0x43, + 0x67, 0x53, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x67, 0x5a, + 0x20, 0x3d, 0x20, 0x74, 0x68, 0x72, 0x5f, 0x6d, 0x6d, 0x61, 0x2e, 0x70, + 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x42, 0x28, 0x67, + 0x5a, 0x29, 0x3b, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x72, 0x5a, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x5f, 0x6c, 0x69, 0x6b, 0x65, 0x28, 0x74, 0x43, 0x67, 0x5a, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x30, 0x29, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x4d, 0x41, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x41, 0x74, 0x6f, 0x6d, + 0x20, 0x72, 0x65, 0x74, 0x69, 0x6c, 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, + 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, + 0x32, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x41, 0x28, 0x73, 0x32, 0x72, 0x5f, 0x61, 0x74, + 0x6f, 0x6d, 0x5f, 0x61, 0x2c, 0x20, 0x6d, 0x6d, 0x61, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, 0x32, + 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, + 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x61, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x32, + 0x72, 0x5f, 0x74, 0x43, 0x73, 0x41, 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x73, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, + 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4b, 0x2c, 0x50, 0x49, 0x50, 0x45, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x72, + 0x41, 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x72, 0x65, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x44, 0x28, 0x74, 0x43, 0x72, 0x41, 0x29, 0x3b, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x41, 0x43, 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, + 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, + 0x32, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x42, 0x28, 0x73, 0x32, 0x72, 0x5f, 0x61, 0x74, + 0x6f, 0x6d, 0x5f, 0x62, 0x2c, 0x20, 0x6d, 0x6d, 0x61, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x73, 0x32, + 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, + 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x62, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x32, + 0x72, 0x5f, 0x74, 0x43, 0x73, 0x42, 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, + 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, + 0x73, 0x42, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, + 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4b, 0x2c, 0x50, 0x49, 0x50, 0x45, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x72, + 0x42, 0x20, 0x3d, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2e, 0x72, 0x65, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x44, 0x28, 0x74, 0x43, 0x72, 0x42, 0x29, 0x3b, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, + 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, + 0x20, 0x54, 0x69, 0x6c, 0x65, 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x72, + 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x20, 0x3d, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x43, 0x28, 0x72, 0x32, 0x73, 0x5f, 0x61, 0x74, + 0x6f, 0x6d, 0x5f, 0x63, 0x2c, 0x20, 0x6d, 0x6d, 0x61, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x54, 0x68, 0x72, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x72, 0x32, + 0x73, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, + 0x20, 0x3d, 0x20, 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x63, 0x2e, 0x67, 0x65, 0x74, 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, + 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x72, 0x32, + 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, 0x20, 0x3d, 0x20, 0x72, 0x32, 0x73, + 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2e, + 0x72, 0x65, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x53, 0x28, 0x74, 0x43, 0x72, + 0x43, 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x43, 0x50, + 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, + 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x20, 0x3d, 0x20, 0x72, + 0x32, 0x73, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, + 0x63, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, + 0x44, 0x28, 0x73, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, + 0x43, 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4d, 0x2c, 0x4d, 0x4d, + 0x41, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x69, 0x6c, 0x65, + 0x64, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x73, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x42, + 0x28, 0x67, 0x32, 0x72, 0x5f, 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x73, 0x2c, + 0x20, 0x6d, 0x6d, 0x61, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x68, 0x72, + 0x43, 0x6f, 0x70, 0x79, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x68, 0x72, + 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x20, 0x3d, 0x20, 0x67, 0x32, + 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x2e, 0x67, 0x65, 0x74, + 0x5f, 0x73, 0x6c, 0x69, 0x63, 0x65, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x67, + 0x53, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, + 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, 0x28, 0x67, 0x53, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, + 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x32, 0x72, + 0x5f, 0x74, 0x43, 0x72, 0x53, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x72, 0x5f, + 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x2e, 0x72, + 0x65, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x44, 0x28, 0x74, 0x43, 0x72, 0x53, + 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, + 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, + 0x32, 0x72, 0x5f, 0x74, 0x43, 0x67, 0x5a, 0x20, 0x3d, 0x20, 0x67, 0x32, + 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, + 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x53, + 0x28, 0x67, 0x5a, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x43, + 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, 0x2c, 0x4d, 0x4d, 0x41, + 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x72, 0x5a, 0x20, + 0x3d, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, 0x6f, + 0x70, 0x79, 0x5f, 0x73, 0x2e, 0x72, 0x65, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x44, 0x28, 0x74, 0x43, 0x72, 0x5a, 0x29, 0x3b, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x42, 0x43, 0x50, 0x59, 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4e, + 0x2c, 0x4d, 0x4d, 0x41, 0x5f, 0x4b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x50, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x65, 0x73, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x6d, 0x20, 0x62, 0x6f, 0x75, 0x6e, 0x64, + 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x70, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, + 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x74, 0x41, 0x73, 0x41, 0x29, + 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x41, + 0x73, 0x41, 0x29, 0x29, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x3c, 0x5f, 0x31, 0x2c, 0x5f, 0x30, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x43, 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, + 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, + 0x70, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x28, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, + 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, + 0x73, 0x43, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, + 0x28, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x29, 0x2c, + 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x5f, + 0x30, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, + 0x50, 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x63, 0x41, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, + 0x69, 0x74, 0x79, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, + 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x73, 0x41, 0x29, 0x2c, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x3c, 0x31, 0x3e, 0x28, 0x73, 0x41, 0x29, 0x29, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, + 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x63, 0x43, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x69, 0x64, 0x65, 0x6e, 0x74, 0x69, 0x74, 0x79, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, + 0x28, 0x73, 0x43, 0x29, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, + 0x3e, 0x28, 0x73, 0x43, 0x29, 0x29, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, + 0x41, 0x63, 0x41, 0x20, 0x3d, 0x20, 0x67, 0x32, 0x73, 0x5f, 0x74, 0x68, + 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x61, 0x2e, 0x70, 0x61, 0x72, + 0x74, 0x69, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x44, 0x28, 0x63, 0x41, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x2c, 0x43, 0x50, + 0x59, 0x5f, 0x4d, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4b, 0x29, 0x0a, 0x20, + 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x74, 0x43, 0x63, 0x43, + 0x20, 0x3d, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x68, 0x72, 0x5f, 0x63, + 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2e, 0x70, 0x61, 0x72, 0x74, 0x69, 0x74, + 0x69, 0x6f, 0x6e, 0x5f, 0x44, 0x28, 0x63, 0x43, 0x29, 0x3b, 0x20, 0x2f, + 0x2f, 0x20, 0x28, 0x43, 0x50, 0x59, 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4d, + 0x2c, 0x43, 0x50, 0x59, 0x5f, 0x4e, 0x29, 0x0a, 0x20, 0x20, 0x43, 0x55, + 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, + 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x20, 0x6d, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x70, 0x41, 0x29, 0x3b, 0x20, 0x2b, + 0x2b, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x41, + 0x70, 0x41, 0x28, 0x6d, 0x2c, 0x30, 0x29, 0x20, 0x3d, 0x20, 0x67, 0x65, + 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x41, 0x63, 0x41, 0x28, 0x30, 0x2c, + 0x6d, 0x2c, 0x30, 0x29, 0x29, 0x20, 0x3c, 0x20, 0x6d, 0x5f, 0x6d, 0x61, + 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, + 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x6d, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6d, 0x20, 0x3c, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x43, 0x70, + 0x43, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x74, 0x43, 0x70, 0x43, 0x28, 0x6d, 0x2c, 0x30, 0x29, + 0x20, 0x3d, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x30, 0x3e, 0x28, 0x74, 0x43, + 0x63, 0x43, 0x28, 0x30, 0x2c, 0x6d, 0x2c, 0x30, 0x29, 0x29, 0x20, 0x3c, + 0x20, 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x4b, 0x5f, 0x50, 0x49, 0x50, 0x45, 0x5f, 0x4d, 0x41, 0x58, + 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x33, 0x3e, 0x28, 0x74, + 0x41, 0x73, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, + 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, 0x72, 0x65, + 0x61, 0x64, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, + 0x77, 0x72, 0x69, 0x74, 0x65, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x41, 0x2f, + 0x42, 0x3a, 0x20, 0x47, 0x4d, 0x45, 0x4d, 0x20, 0x3d, 0x3e, 0x20, 0x53, + 0x4d, 0x45, 0x4d, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x20, 0x3d, + 0x20, 0x5b, 0x26, 0x5d, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, + 0x79, 0x5f, 0x69, 0x66, 0x28, 0x67, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, + 0x79, 0x5f, 0x61, 0x2c, 0x20, 0x74, 0x41, 0x70, 0x41, 0x2c, 0x20, 0x74, + 0x41, 0x67, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x74, 0x41, 0x73, 0x41, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x5f, 0x2c, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, + 0x65, 0x5f, 0x77, 0x72, 0x69, 0x74, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x67, 0x32, 0x73, 0x5f, + 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x62, 0x2c, 0x20, 0x74, 0x42, 0x67, 0x42, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, + 0x2c, 0x20, 0x74, 0x42, 0x73, 0x42, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, + 0x2c, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, 0x77, + 0x72, 0x69, 0x74, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x70, 0x5f, 0x61, 0x73, 0x79, 0x6e, 0x63, 0x5f, 0x66, 0x65, 0x6e, + 0x63, 0x65, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x73, 0x6d, + 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, 0x77, 0x72, 0x69, 0x74, + 0x65, 0x20, 0x3d, 0x20, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, + 0x70, 0x65, 0x5f, 0x77, 0x72, 0x69, 0x74, 0x65, 0x20, 0x2b, 0x20, 0x31, + 0x29, 0x20, 0x25, 0x20, 0x4b, 0x5f, 0x50, 0x49, 0x50, 0x45, 0x5f, 0x4d, + 0x41, 0x58, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x20, 0x53, 0x2f, 0x5a, 0x3a, 0x20, + 0x47, 0x4d, 0x45, 0x4d, 0x20, 0x3d, 0x3e, 0x20, 0x52, 0x4d, 0x45, 0x4d, + 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x66, 0x65, 0x74, + 0x63, 0x68, 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x20, 0x3d, 0x20, + 0x5b, 0x26, 0x5d, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, + 0x28, 0x67, 0x32, 0x72, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x2c, + 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x67, 0x53, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x5f, 0x2c, 0x74, 0x69, 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x67, + 0x32, 0x72, 0x5f, 0x74, 0x43, 0x72, 0x53, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x3a, 0x3a, 0x71, 0x75, 0x61, 0x6e, 0x74, + 0x5f, 0x68, 0x61, 0x73, 0x5f, 0x62, 0x69, 0x61, 0x73, 0x5f, 0x76, 0x3c, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x67, 0x32, 0x72, + 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x73, 0x2c, 0x20, 0x67, 0x32, 0x72, + 0x5f, 0x74, 0x43, 0x67, 0x5a, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x5f, 0x2c, + 0x74, 0x69, 0x6c, 0x65, 0x29, 0x2c, 0x20, 0x67, 0x32, 0x72, 0x5f, 0x74, + 0x43, 0x72, 0x5a, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6f, + 0x70, 0x79, 0x20, 0x41, 0x2f, 0x42, 0x3a, 0x20, 0x53, 0x4d, 0x45, 0x4d, + 0x20, 0x3d, 0x3e, 0x20, 0x52, 0x4d, 0x45, 0x4d, 0x2e, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x73, + 0x6d, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x5b, 0x26, 0x5d, 0x28, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, 0x73, 0x32, 0x72, + 0x5f, 0x61, 0x74, 0x6f, 0x6d, 0x5f, 0x61, 0x2c, 0x20, 0x73, 0x32, 0x72, + 0x5f, 0x74, 0x43, 0x73, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x6c, + 0x6f, 0x63, 0x6b, 0x2c, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, + 0x65, 0x5f, 0x72, 0x65, 0x61, 0x64, 0x29, 0x2c, 0x20, 0x73, 0x32, 0x72, + 0x5f, 0x74, 0x43, 0x72, 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x6c, + 0x6f, 0x63, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x70, 0x79, 0x28, 0x73, 0x32, 0x72, 0x5f, 0x61, 0x74, 0x6f, 0x6d, + 0x5f, 0x62, 0x2c, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x73, 0x42, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2c, 0x73, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, 0x72, 0x65, 0x61, + 0x64, 0x29, 0x2c, 0x20, 0x73, 0x32, 0x72, 0x5f, 0x74, 0x43, 0x72, 0x42, + 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, + 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x20, 0x3d, 0x20, 0x30, + 0x3b, 0x20, 0x6e, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x31, + 0x3e, 0x28, 0x74, 0x43, 0x72, 0x42, 0x29, 0x3b, 0x20, 0x2b, 0x2b, 0x6e, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x6c, + 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x3a, + 0x3a, 0x63, 0x75, 0x74, 0x65, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, + 0x69, 0x7a, 0x65, 0x64, 0x5f, 0x64, 0x65, 0x71, 0x75, 0x61, 0x6e, 0x74, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x43, 0x72, 0x42, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x62, 0x6c, 0x6f, + 0x63, 0x6b, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, 0x53, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, 0x5a, 0x28, 0x5f, + 0x2c, 0x6e, 0x2c, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, + 0x42, 0x5f, 0x64, 0x71, 0x28, 0x5f, 0x2c, 0x6e, 0x2c, 0x62, 0x6c, 0x6f, + 0x63, 0x6b, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x20, + 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x33, 0x3e, 0x28, 0x74, 0x41, + 0x67, 0x41, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x4b, 0x5f, 0x42, 0x4c, 0x4f, 0x43, 0x4b, 0x5f, 0x4d, 0x41, 0x58, 0x20, + 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x32, 0x3e, 0x28, 0x74, 0x43, + 0x72, 0x41, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, + 0x72, 0x65, 0x66, 0x65, 0x74, 0x63, 0x68, 0x20, 0x62, 0x65, 0x67, 0x69, + 0x6e, 0x6e, 0x69, 0x6e, 0x67, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x73, 0x2e, + 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x70, 0x69, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, + 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, + 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x3b, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x50, + 0x49, 0x50, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x2d, 0x20, 0x31, 0x3b, + 0x20, 0x2b, 0x2b, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x70, 0x69, 0x70, 0x65, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, + 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x28, 0x74, 0x69, 0x6c, 0x65, 0x5f, + 0x70, 0x69, 0x70, 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, 0x6c, 0x65, 0x61, 0x72, 0x20, 0x61, + 0x63, 0x63, 0x75, 0x6d, 0x75, 0x6c, 0x61, 0x74, 0x6f, 0x72, 0x73, 0x2e, + 0x0a, 0x20, 0x20, 0x63, 0x6c, 0x65, 0x61, 0x72, 0x28, 0x74, 0x43, 0x72, + 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x50, 0x72, 0x65, 0x66, 0x65, 0x74, 0x63, 0x68, 0x20, + 0x66, 0x69, 0x72, 0x73, 0x74, 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, + 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, + 0x78, 0x70, 0x72, 0x20, 0x28, 0x4b, 0x5f, 0x42, 0x4c, 0x4f, 0x43, 0x4b, + 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x70, 0x5f, 0x61, 0x73, 0x79, 0x6e, 0x63, + 0x5f, 0x77, 0x61, 0x69, 0x74, 0x3c, 0x4b, 0x5f, 0x50, 0x49, 0x50, 0x45, + 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x2d, 0x20, 0x32, 0x3e, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, + 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x63, 0x61, 0x6c, + 0x65, 0x73, 0x28, 0x30, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x66, + 0x65, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x28, 0x49, 0x6e, + 0x74, 0x3c, 0x30, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x20, + 0x6f, 0x76, 0x65, 0x72, 0x20, 0x43, 0x54, 0x41, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x54, 0x49, + 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x74, 0x69, + 0x6c, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x55, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x20, 0x4d, 0x4d, 0x41, 0x20, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, 0x4c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x42, 0x4c, + 0x4f, 0x43, 0x4b, 0x5f, 0x4d, 0x41, 0x58, 0x3b, 0x20, 0x2b, 0x2b, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x57, 0x61, 0x69, 0x74, 0x20, 0x66, 0x6f, + 0x72, 0x20, 0x6c, 0x61, 0x73, 0x74, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x2e, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x62, + 0x6c, 0x6f, 0x63, 0x6b, 0x20, 0x3d, 0x3d, 0x20, 0x4b, 0x5f, 0x42, 0x4c, + 0x4f, 0x43, 0x4b, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x2d, 0x20, 0x31, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x5f, 0x72, 0x65, 0x61, + 0x64, 0x20, 0x3d, 0x20, 0x28, 0x73, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x69, + 0x70, 0x65, 0x5f, 0x72, 0x65, 0x61, 0x64, 0x20, 0x2b, 0x20, 0x31, 0x29, + 0x20, 0x25, 0x20, 0x4b, 0x5f, 0x50, 0x49, 0x50, 0x45, 0x5f, 0x4d, 0x41, + 0x58, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x70, 0x5f, 0x61, 0x73, 0x79, 0x6e, 0x63, 0x5f, 0x77, 0x61, 0x69, 0x74, + 0x3c, 0x4b, 0x5f, 0x50, 0x49, 0x50, 0x45, 0x5f, 0x4d, 0x41, 0x58, 0x20, + 0x2d, 0x20, 0x32, 0x3e, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x5f, 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x73, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x73, + 0x63, 0x61, 0x6c, 0x65, 0x73, 0x28, 0x28, 0x74, 0x69, 0x6c, 0x65, 0x20, + 0x2b, 0x20, 0x31, 0x20, 0x3c, 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, + 0x5f, 0x4d, 0x41, 0x58, 0x29, 0x20, 0x3f, 0x20, 0x74, 0x69, 0x6c, 0x65, + 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3a, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, 0x72, 0x65, 0x66, 0x65, + 0x74, 0x63, 0x68, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x20, 0x62, 0x6c, 0x6f, + 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x66, 0x65, + 0x74, 0x63, 0x68, 0x5f, 0x73, 0x6d, 0x65, 0x6d, 0x28, 0x28, 0x62, 0x6c, + 0x6f, 0x63, 0x6b, 0x20, 0x2b, 0x20, 0x31, 0x29, 0x20, 0x25, 0x20, 0x4b, + 0x5f, 0x42, 0x4c, 0x4f, 0x43, 0x4b, 0x5f, 0x4d, 0x41, 0x58, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x50, 0x72, + 0x65, 0x66, 0x65, 0x74, 0x63, 0x68, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x20, + 0x74, 0x69, 0x6c, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x20, 0x3d, 0x3d, + 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x66, 0x65, 0x74, 0x63, 0x68, 0x5f, 0x67, 0x6d, 0x65, 0x6d, + 0x28, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x69, 0x6c, + 0x65, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x20, 0x3d, 0x20, 0x28, 0x74, 0x69, + 0x6c, 0x65, 0x5f, 0x70, 0x69, 0x70, 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, + 0x3c, 0x20, 0x4b, 0x5f, 0x54, 0x49, 0x4c, 0x45, 0x5f, 0x4d, 0x41, 0x58, + 0x29, 0x20, 0x3f, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x5f, 0x70, 0x69, 0x70, + 0x65, 0x20, 0x2b, 0x20, 0x31, 0x20, 0x3a, 0x20, 0x74, 0x69, 0x6c, 0x65, + 0x5f, 0x70, 0x69, 0x70, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x4d, 0x4d, 0x41, 0x2e, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, + 0x65, 0x6d, 0x6d, 0x28, 0x6d, 0x6d, 0x61, 0x2c, 0x20, 0x74, 0x43, 0x72, + 0x41, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, + 0x2c, 0x20, 0x74, 0x43, 0x72, 0x42, 0x5f, 0x64, 0x71, 0x28, 0x5f, 0x2c, + 0x5f, 0x2c, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x29, 0x2c, 0x20, 0x74, 0x43, + 0x72, 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x2f, + 0x2f, 0x20, 0x45, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x2e, 0x0a, + 0x20, 0x20, 0x43, 0x55, 0x54, 0x45, 0x5f, 0x55, 0x4e, 0x52, 0x4f, 0x4c, + 0x4c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x69, 0x20, 0x3c, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x28, 0x74, 0x43, 0x72, 0x43, 0x5f, 0x61, 0x63, + 0x63, 0x75, 0x29, 0x3b, 0x20, 0x69, 0x2b, 0x2b, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x74, 0x43, 0x72, 0x43, 0x28, 0x69, 0x29, 0x20, + 0x3d, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x28, 0x74, 0x43, + 0x72, 0x43, 0x5f, 0x61, 0x63, 0x63, 0x75, 0x28, 0x69, 0x29, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x28, + 0x72, 0x32, 0x73, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, 0x2c, 0x20, + 0x72, 0x32, 0x73, 0x5f, 0x74, 0x43, 0x72, 0x43, 0x2c, 0x20, 0x72, 0x32, + 0x73, 0x5f, 0x74, 0x43, 0x73, 0x43, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x5f, + 0x5f, 0x73, 0x79, 0x6e, 0x63, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x73, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x69, + 0x66, 0x28, 0x73, 0x32, 0x67, 0x5f, 0x63, 0x6f, 0x70, 0x79, 0x5f, 0x63, + 0x2c, 0x20, 0x74, 0x43, 0x70, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, + 0x74, 0x43, 0x73, 0x43, 0x2c, 0x20, 0x73, 0x32, 0x67, 0x5f, 0x74, 0x43, + 0x67, 0x43, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, + 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, + 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, + 0x69, 0x6e, 0x65, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6b, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x3d, 0x20, 0x49, 0x6e, 0x74, + 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x3e, 0x7b, + 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x6b, 0x20, 0x2f, 0x20, + 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x2c, + 0x20, 0x6c, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, + 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x2c, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x3c, 0x5f, 0x30, + 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, + 0x6b, 0x20, 0x2f, 0x20, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, + 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, + 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2c, 0x20, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, + 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, + 0x6c, 0x5f, 0x5f, 0x0a, 0x5f, 0x5f, 0x6c, 0x61, 0x75, 0x6e, 0x63, 0x68, + 0x5f, 0x62, 0x6f, 0x75, 0x6e, 0x64, 0x73, 0x5f, 0x5f, 0x28, 0x64, 0x65, + 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x73, 0x69, 0x7a, 0x65, 0x28, + 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x64, 0x5f, 0x6d, + 0x6d, 0x61, 0x28, 0x29, 0x29, 0x29, 0x3a, 0x3a, 0x76, 0x61, 0x6c, 0x75, + 0x65, 0x29, 0x0a, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x71, 0x6d, 0x6d, 0x5f, + 0x73, 0x6d, 0x38, 0x30, 0x5f, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x41, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x51, 0x75, 0x61, + 0x6e, 0x74, 0x2a, 0x20, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x53, 0x63, 0x61, 0x6c, 0x65, 0x2a, 0x20, + 0x53, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x5a, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, + 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x6c, 0x68, 0x73, + 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x75, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x2a, 0x20, 0x72, 0x68, 0x73, 0x5f, 0x69, 0x6e, + 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2a, 0x20, 0x43, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6d, 0x2c, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6b, 0x2c, 0x20, + 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, + 0x6f, 0x6f, 0x6c, 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, + 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, + 0x49, 0x64, 0x78, 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, + 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, + 0x2e, 0x78, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, + 0x28, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x79, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6c, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x28, 0x62, 0x6c, + 0x6f, 0x63, 0x6b, 0x49, 0x64, 0x78, 0x2e, 0x7a, 0x29, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x44, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x20, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x73, 0x20, 0x28, 0x6d, 0x69, 0x78, + 0x65, 0x64, 0x29, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, + 0x64, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, + 0x31, 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6b, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4b, 0x2c, + 0x64, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, + 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x28, 0x6b, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, + 0x3e, 0x7b, 0x7d, 0x2c, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x6b, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4e, 0x2c, 0x64, 0x4b, 0x2c, 0x64, + 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x64, 0x43, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x28, 0x6e, 0x2c, 0x20, 0x49, 0x6e, 0x74, 0x3c, 0x31, 0x3e, + 0x7b, 0x7d, 0x2c, 0x20, 0x6d, 0x20, 0x2a, 0x20, 0x6e, 0x29, 0x3b, 0x20, + 0x2f, 0x2f, 0x20, 0x28, 0x64, 0x4d, 0x2c, 0x64, 0x4e, 0x2c, 0x64, 0x4c, + 0x29, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x53, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, + 0x5f, 0x73, 0x63, 0x61, 0x6c, 0x65, 0x73, 0x5f, 0x6c, 0x61, 0x79, 0x6f, + 0x75, 0x74, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x3e, 0x28, 0x6e, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x3b, 0x0a, + 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x48, 0x61, 0x6e, 0x64, 0x6c, 0x65, + 0x20, 0x62, 0x72, 0x6f, 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x69, 0x6e, + 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x62, 0x72, 0x6f, + 0x61, 0x64, 0x63, 0x61, 0x73, 0x74, 0x5f, 0x62, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x64, + 0x42, 0x29, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x67, 0x65, 0x74, 0x3c, 0x32, 0x3e, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x28, 0x53, 0x5f, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x29, + 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x2f, 0x2f, 0x20, 0x52, 0x65, 0x70, 0x72, 0x65, 0x73, 0x65, 0x6e, + 0x74, 0x20, 0x74, 0x68, 0x65, 0x20, 0x66, 0x75, 0x6c, 0x6c, 0x20, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x73, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, + 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x5f, 0x6d, 0x6b, 0x6c, 0x20, + 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, + 0x70, 0x74, 0x72, 0x28, 0x41, 0x29, 0x2c, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, + 0x64, 0x41, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, + 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x3c, + 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x28, 0x42, 0x29, 0x2c, 0x20, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6e, 0x2c, + 0x20, 0x6b, 0x2c, 0x20, 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x42, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x2c, 0x4c, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x5f, + 0x6d, 0x6e, 0x6c, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, + 0x6d, 0x65, 0x6d, 0x5f, 0x70, 0x74, 0x72, 0x28, 0x43, 0x29, 0x2c, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x28, 0x6d, 0x2c, 0x20, 0x6e, 0x2c, 0x20, + 0x6c, 0x29, 0x2c, 0x20, 0x64, 0x43, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4e, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, + 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x5f, 0x6e, 0x6b, 0x6c, + 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, + 0x5f, 0x70, 0x74, 0x72, 0x28, 0x53, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, + 0x61, 0x79, 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x20, 0x3d, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, + 0x28, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x67, 0x6d, 0x65, 0x6d, 0x5f, 0x70, + 0x74, 0x72, 0x28, 0x5a, 0x29, 0x2c, 0x20, 0x53, 0x5f, 0x6c, 0x61, 0x79, + 0x6f, 0x75, 0x74, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x2c, 0x4c, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x46, + 0x6f, 0x72, 0x20, 0x67, 0x61, 0x74, 0x68, 0x65, 0x72, 0x2c, 0x20, 0x75, + 0x73, 0x65, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x6c, 0x6f, 0x6f, + 0x6b, 0x75, 0x70, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x69, 0x6e, 0x70, 0x75, + 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, 0x6c, 0x69, 0x63, + 0x69, 0x6e, 0x67, 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x61, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, + 0x3d, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x20, 0x3f, 0x20, 0x6c, 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, + 0x63, 0x65, 0x73, 0x5b, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, + 0x20, 0x3a, 0x20, 0x6c, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, + 0x20, 0x20, 0x75, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x62, + 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x3d, 0x20, 0x72, 0x68, 0x73, + 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x20, 0x3f, 0x20, 0x72, + 0x68, 0x73, 0x5f, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5b, 0x6c, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x5d, 0x20, 0x3a, 0x20, 0x6c, 0x5f, + 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, + 0x20, 0x47, 0x65, 0x74, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x20, 0x73, + 0x6c, 0x69, 0x63, 0x65, 0x2e, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x20, 0x6d, 0x41, 0x20, 0x3d, 0x20, 0x6d, 0x41, 0x5f, 0x6d, + 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x61, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4d, 0x2c, 0x4b, + 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, + 0x42, 0x20, 0x3d, 0x20, 0x6d, 0x42, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, 0x5f, + 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x4b, 0x29, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x43, 0x20, 0x3d, 0x20, + 0x6d, 0x43, 0x5f, 0x6d, 0x6e, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x6c, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, + 0x28, 0x4d, 0x2c, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x6d, 0x53, 0x20, 0x3d, 0x20, 0x6d, 0x53, 0x5f, + 0x6e, 0x6b, 0x6c, 0x28, 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, + 0x28, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, + 0x4b, 0x2f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x29, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x6d, 0x5a, 0x20, 0x3d, 0x20, 0x6d, 0x5a, 0x5f, 0x6e, 0x6b, 0x6c, 0x28, + 0x5f, 0x2c, 0x5f, 0x2c, 0x62, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x29, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x4e, 0x2c, 0x28, 0x67, 0x72, 0x6f, + 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x4b, 0x2f, 0x67, 0x72, + 0x6f, 0x75, 0x70, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x29, 0x0a, 0x0a, + 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x47, 0x65, 0x74, 0x20, 0x74, 0x68, 0x65, + 0x20, 0x61, 0x70, 0x70, 0x72, 0x6f, 0x70, 0x72, 0x69, 0x61, 0x74, 0x65, + 0x20, 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x74, 0x68, 0x69, 0x73, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x20, + 0x62, 0x6c, 0x6f, 0x63, 0x6b, 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, + 0x6f, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x20, + 0x3d, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x7b, 0x7d, + 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x28, 0x6d, 0x5f, 0x63, 0x6f, + 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x6e, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x2c, 0x20, 0x5f, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x6d, 0x2c, + 0x6e, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x20, 0x67, 0x41, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, + 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x41, 0x2c, 0x20, 0x63, 0x74, + 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, + 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, + 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, + 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, + 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x20, 0x20, + 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x42, 0x20, 0x3d, 0x20, + 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, + 0x42, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, + 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, + 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, + 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, + 0x6b, 0x29, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, + 0x67, 0x43, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, + 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x43, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, + 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, + 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x5f, + 0x31, 0x2c, 0x5f, 0x31, 0x2c, 0x20, 0x58, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, + 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4d, 0x2c, 0x42, + 0x4c, 0x4b, 0x5f, 0x4e, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x65, 0x6e, + 0x73, 0x6f, 0x72, 0x20, 0x67, 0x53, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x63, + 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x28, 0x6d, 0x53, 0x2c, 0x20, + 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x20, 0x63, + 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, 0x20, 0x53, 0x74, + 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, 0x31, 0x2c, 0x5f, 0x31, 0x3e, + 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x28, 0x42, 0x4c, 0x4b, + 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, 0x4b, 0x2c, 0x6b, 0x29, 0x0a, + 0x20, 0x20, 0x54, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x20, 0x67, 0x5a, 0x20, + 0x3d, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x6c, 0x5f, 0x74, 0x69, 0x6c, 0x65, + 0x28, 0x6d, 0x5a, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x2c, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x63, 0x6f, 0x6f, 0x72, + 0x64, 0x2c, 0x20, 0x53, 0x74, 0x65, 0x70, 0x3c, 0x20, 0x58, 0x2c, 0x5f, + 0x31, 0x2c, 0x5f, 0x31, 0x3e, 0x7b, 0x7d, 0x29, 0x3b, 0x20, 0x2f, 0x2f, + 0x20, 0x28, 0x42, 0x4c, 0x4b, 0x5f, 0x4e, 0x2c, 0x42, 0x4c, 0x4b, 0x5f, + 0x4b, 0x2c, 0x6b, 0x29, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x43, + 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x65, 0x20, 0x74, 0x69, 0x6c, 0x65, 0x20, + 0x72, 0x65, 0x73, 0x69, 0x64, 0x75, 0x65, 0x73, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x70, 0x72, 0x65, 0x64, 0x69, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, + 0x2e, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x5f, 0x6d, + 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x20, 0x3d, 0x20, 0x6d, + 0x20, 0x2d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x3c, 0x30, 0x3e, 0x28, 0x67, + 0x41, 0x29, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x3b, 0x20, 0x2f, 0x2f, 0x20, 0x4d, 0x20, 0x2d, 0x20, 0x42, 0x4c, 0x4b, + 0x5f, 0x4d, 0x20, 0x2a, 0x20, 0x6d, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, + 0x0a, 0x0a, 0x20, 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x38, 0x30, + 0x5f, 0x6d, 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x28, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x74, 0x61, 0x5f, 0x74, 0x69, 0x6c, + 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x41, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x42, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x67, 0x53, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x67, 0x5a, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x67, 0x43, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x6d, 0x5f, 0x6d, 0x61, 0x78, 0x5f, 0x63, 0x6f, 0x6f, 0x72, 0x64, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x68, 0x72, 0x65, 0x61, + 0x64, 0x5f, 0x69, 0x64, 0x78, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, + 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_qmm_sm90[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x36, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x66, + 0x20, 0x64, 0x65, 0x66, 0x69, 0x6e, 0x65, 0x64, 0x28, 0x5f, 0x5f, 0x43, + 0x55, 0x44, 0x41, 0x43, 0x43, 0x5f, 0x52, 0x54, 0x43, 0x5f, 0x5f, 0x29, + 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, + 0x63, 0x75, 0x64, 0x61, 0x2f, 0x63, 0x6d, 0x61, 0x74, 0x68, 0x3e, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x73, 0x74, 0x64, 0x2f, 0x74, 0x79, 0x70, 0x65, 0x5f, + 0x74, 0x72, 0x61, 0x69, 0x74, 0x73, 0x3e, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, + 0x53, 0x6f, 0x6d, 0x65, 0x20, 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, + 0x20, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x73, 0x20, 0x75, 0x73, 0x65, + 0x20, 0x66, 0x6f, 0x6c, 0x6c, 0x6f, 0x77, 0x69, 0x6e, 0x67, 0x20, 0x73, + 0x74, 0x64, 0x20, 0x66, 0x75, 0x6e, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x73, + 0x20, 0x62, 0x75, 0x74, 0x20, 0x77, 0x65, 0x20, 0x63, 0x61, 0x6e, 0x27, + 0x74, 0x20, 0x75, 0x73, 0x65, 0x20, 0x53, 0x54, 0x4c, 0x20, 0x77, 0x68, + 0x65, 0x6e, 0x0a, 0x2f, 0x2f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, + 0x69, 0x6e, 0x67, 0x20, 0x77, 0x69, 0x74, 0x68, 0x20, 0x4e, 0x56, 0x52, + 0x54, 0x43, 0x2e, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x73, 0x74, 0x64, 0x20, 0x7b, 0x0a, 0x75, 0x73, 0x69, 0x6e, + 0x67, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x69, 0x73, 0x5f, 0x70, 0x6f, 0x69, 0x6e, 0x74, 0x65, 0x72, 0x5f, + 0x76, 0x3b, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x3b, + 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, + 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x76, 0x6f, 0x69, 0x64, 0x5f, 0x74, + 0x3b, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x73, 0x74, 0x64, 0x0a, 0x0a, 0x2f, 0x2f, + 0x20, 0x54, 0x68, 0x65, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, + 0x2f, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x69, 0x6e, 0x67, 0x5f, 0x70, 0x6f, + 0x69, 0x6e, 0x74, 0x5f, 0x6e, 0x76, 0x72, 0x74, 0x63, 0x2e, 0x68, 0x20, + 0x66, 0x69, 0x6c, 0x65, 0x20, 0x61, 0x73, 0x73, 0x75, 0x6d, 0x65, 0x73, + 0x20, 0x66, 0x6f, 0x6c, 0x6c, 0x6f, 0x77, 0x69, 0x6e, 0x67, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x73, 0x20, 0x6e, 0x6f, 0x74, + 0x0a, 0x2f, 0x2f, 0x20, 0x62, 0x65, 0x69, 0x6e, 0x67, 0x20, 0x64, 0x65, + 0x66, 0x69, 0x6e, 0x65, 0x64, 0x20, 0x62, 0x75, 0x74, 0x20, 0x74, 0x68, + 0x65, 0x79, 0x20, 0x61, 0x72, 0x65, 0x20, 0x69, 0x6e, 0x20, 0x74, 0x68, + 0x65, 0x20, 0x43, 0x55, 0x44, 0x41, 0x20, 0x73, 0x74, 0x64, 0x20, 0x68, + 0x65, 0x61, 0x64, 0x65, 0x72, 0x73, 0x2e, 0x0a, 0x23, 0x75, 0x6e, 0x64, + 0x65, 0x66, 0x20, 0x46, 0x50, 0x5f, 0x4e, 0x41, 0x4e, 0x0a, 0x23, 0x75, + 0x6e, 0x64, 0x65, 0x66, 0x20, 0x46, 0x50, 0x5f, 0x49, 0x4e, 0x46, 0x49, + 0x4e, 0x49, 0x54, 0x45, 0x0a, 0x23, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x20, + 0x46, 0x50, 0x5f, 0x5a, 0x45, 0x52, 0x4f, 0x0a, 0x23, 0x75, 0x6e, 0x64, + 0x65, 0x66, 0x20, 0x46, 0x50, 0x5f, 0x53, 0x55, 0x42, 0x4e, 0x4f, 0x52, + 0x4d, 0x41, 0x4c, 0x0a, 0x23, 0x75, 0x6e, 0x64, 0x65, 0x66, 0x20, 0x46, + 0x50, 0x5f, 0x4e, 0x4f, 0x52, 0x4d, 0x41, 0x4c, 0x0a, 0x0a, 0x23, 0x65, + 0x6e, 0x64, 0x69, 0x66, 0x20, 0x2f, 0x2f, 0x20, 0x64, 0x65, 0x66, 0x69, + 0x6e, 0x65, 0x64, 0x28, 0x5f, 0x5f, 0x43, 0x55, 0x44, 0x41, 0x43, 0x43, + 0x5f, 0x52, 0x54, 0x43, 0x5f, 0x5f, 0x29, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x65, 0x2f, + 0x74, 0x65, 0x6e, 0x73, 0x6f, 0x72, 0x2e, 0x68, 0x70, 0x70, 0x3e, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, + 0x74, 0x6c, 0x61, 0x73, 0x73, 0x2f, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, + 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, + 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x2f, 0x65, + 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x2f, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x2f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, + 0x63, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x62, 0x75, 0x69, 0x6c, 0x64, 0x65, + 0x72, 0x2e, 0x68, 0x70, 0x70, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, + 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, + 0x2f, 0x67, 0x65, 0x6d, 0x6d, 0x2f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x76, 0x65, 0x2f, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, + 0x69, 0x76, 0x65, 0x5f, 0x62, 0x75, 0x69, 0x6c, 0x64, 0x65, 0x72, 0x2e, + 0x68, 0x70, 0x70, 0x3e, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, + 0x65, 0x20, 0x3c, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x2f, 0x67, + 0x65, 0x6d, 0x6d, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x67, + 0x65, 0x6d, 0x6d, 0x5f, 0x75, 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, + 0x6c, 0x5f, 0x61, 0x64, 0x61, 0x70, 0x74, 0x65, 0x72, 0x2e, 0x68, 0x3e, + 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, + 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x2f, 0x67, 0x65, 0x6d, 0x6d, 0x2f, + 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x2f, 0x67, 0x65, 0x6d, 0x6d, 0x5f, + 0x75, 0x6e, 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x2e, 0x68, 0x70, + 0x70, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, + 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, + 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, + 0x75, 0x74, 0x65, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, + 0x70, 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, + 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x43, 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, + 0x43, 0x55, 0x54, 0x4c, 0x41, 0x53, 0x53, 0x5f, 0x48, 0x4f, 0x53, 0x54, + 0x5f, 0x44, 0x45, 0x56, 0x49, 0x43, 0x45, 0x20, 0x61, 0x75, 0x74, 0x6f, + 0x20, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, + 0x39, 0x30, 0x5f, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, + 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x6d, + 0x65, 0x6e, 0x74, 0x41, 0x20, 0x3d, 0x20, 0x31, 0x32, 0x38, 0x20, 0x2f, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, 0x62, 0x69, 0x74, 0x73, + 0x3c, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x3e, 0x3a, 0x3a, 0x76, + 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x6d, 0x65, 0x6e, 0x74, 0x42, 0x20, 0x3d, 0x20, 0x31, + 0x32, 0x38, 0x20, 0x2f, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x5f, + 0x62, 0x69, 0x74, 0x73, 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x3e, 0x3a, + 0x3a, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x41, 0x72, 0x63, 0x68, 0x20, 0x3d, 0x20, + 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x61, 0x72, 0x63, + 0x68, 0x3a, 0x3a, 0x53, 0x6d, 0x39, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x41, 0x63, 0x63, 0x75, 0x6d, 0x75, 0x6c, + 0x61, 0x74, 0x6f, 0x72, 0x20, 0x3d, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x3b, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x43, 0x6c, + 0x75, 0x73, 0x74, 0x65, 0x72, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x3d, + 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x3c, 0x5f, 0x31, 0x2c, 0x20, 0x5f, + 0x31, 0x2c, 0x20, 0x5f, 0x31, 0x3e, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x75, + 0x73, 0x69, 0x6e, 0x67, 0x20, 0x45, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, + 0x65, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x65, 0x70, + 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x63, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x3a, 0x3a, 0x43, 0x6f, 0x6c, 0x6c, + 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x42, 0x75, 0x69, 0x6c, 0x64, 0x65, + 0x72, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x72, 0x63, + 0x68, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, + 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x61, 0x72, 0x63, 0x68, 0x3a, 0x3a, + 0x4f, 0x70, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x54, 0x65, 0x6e, 0x73, 0x6f, + 0x72, 0x4f, 0x70, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x43, + 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x53, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x65, 0x70, 0x69, 0x6c, + 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, + 0x74, 0x69, 0x76, 0x65, 0x3a, 0x3a, 0x45, 0x70, 0x69, 0x6c, 0x6f, 0x67, + 0x75, 0x65, 0x54, 0x69, 0x6c, 0x65, 0x41, 0x75, 0x74, 0x6f, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x63, 0x63, 0x75, 0x6d, 0x75, + 0x6c, 0x61, 0x74, 0x6f, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x41, 0x63, 0x63, 0x75, 0x6d, 0x75, 0x6c, 0x61, 0x74, 0x6f, 0x72, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x43, 0x3a, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, + 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3a, 0x3a, 0x43, 0x6f, 0x6c, 0x75, + 0x6d, 0x6e, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x6d, 0x65, 0x6e, 0x74, + 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x44, 0x3a, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, + 0x73, 0x73, 0x3a, 0x3a, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, 0x3a, 0x3a, + 0x43, 0x6f, 0x6c, 0x75, 0x6d, 0x6e, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, + 0x6d, 0x65, 0x6e, 0x74, 0x41, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x65, 0x70, + 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x54, 0x6d, 0x61, 0x57, + 0x61, 0x72, 0x70, 0x53, 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x69, 0x7a, + 0x65, 0x64, 0x43, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, + 0x65, 0x3e, 0x3a, 0x3a, 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, + 0x76, 0x65, 0x4f, 0x70, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x2f, 0x2f, 0x20, + 0x4e, 0x6f, 0x74, 0x65, 0x20, 0x74, 0x68, 0x61, 0x74, 0x20, 0x41, 0x2f, + 0x42, 0x20, 0x61, 0x72, 0x65, 0x20, 0x73, 0x77, 0x61, 0x70, 0x70, 0x65, + 0x64, 0x20, 0x61, 0x6e, 0x64, 0x20, 0x74, 0x72, 0x61, 0x6e, 0x73, 0x70, + 0x6f, 0x73, 0x65, 0x64, 0x20, 0x74, 0x6f, 0x20, 0x75, 0x73, 0x65, 0x20, + 0x54, 0x4d, 0x41, 0x20, 0x65, 0x70, 0x69, 0x6c, 0x6f, 0x67, 0x75, 0x65, + 0x2e, 0x0a, 0x20, 0x20, 0x75, 0x73, 0x69, 0x6e, 0x67, 0x20, 0x4d, 0x61, + 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x20, 0x3d, 0x20, 0x74, 0x79, 0x70, + 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, + 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, 0x6d, 0x3a, 0x3a, 0x63, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x3a, 0x3a, 0x43, 0x6f, 0x6c, + 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x42, 0x75, 0x69, 0x6c, 0x64, + 0x65, 0x72, 0x3c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x72, + 0x63, 0x68, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, + 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x61, 0x72, 0x63, 0x68, 0x3a, + 0x3a, 0x4f, 0x70, 0x43, 0x6c, 0x61, 0x73, 0x73, 0x54, 0x65, 0x6e, 0x73, + 0x6f, 0x72, 0x4f, 0x70, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x41, 0x3a, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x74, 0x75, 0x70, 0x6c, 0x65, + 0x3c, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, + 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, + 0x3e, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, + 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x6c, 0x61, 0x79, 0x6f, 0x75, 0x74, + 0x3a, 0x3a, 0x52, 0x6f, 0x77, 0x4d, 0x61, 0x6a, 0x6f, 0x72, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, 0x69, 0x67, 0x6e, 0x6d, + 0x65, 0x6e, 0x74, 0x42, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x2f, 0x2f, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x42, 0x3a, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, + 0x6e, 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, + 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x6c, 0x61, 0x79, 0x6f, 0x75, + 0x74, 0x3a, 0x3a, 0x43, 0x6f, 0x6c, 0x75, 0x6d, 0x6e, 0x4d, 0x61, 0x6a, + 0x6f, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x41, 0x6c, + 0x69, 0x67, 0x6e, 0x6d, 0x65, 0x6e, 0x74, 0x41, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x41, 0x63, 0x63, 0x75, 0x6d, 0x75, 0x6c, 0x61, + 0x74, 0x6f, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x43, + 0x74, 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x43, 0x6c, 0x75, 0x73, 0x74, 0x65, 0x72, 0x53, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, 0x6d, + 0x3a, 0x3a, 0x63, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, + 0x3a, 0x3a, 0x53, 0x74, 0x61, 0x67, 0x65, 0x43, 0x6f, 0x75, 0x6e, 0x74, + 0x41, 0x75, 0x74, 0x6f, 0x43, 0x61, 0x72, 0x76, 0x65, 0x6f, 0x75, 0x74, + 0x3c, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, 0x74, + 0x3c, 0x69, 0x6e, 0x74, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x6f, 0x66, 0x28, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, 0x70, 0x69, + 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3a, 0x3a, 0x53, 0x68, 0x61, 0x72, 0x65, + 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x29, 0x29, 0x3e, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x74, 0x6c, 0x61, + 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, 0x6d, 0x3a, 0x3a, 0x4b, 0x65, + 0x72, 0x6e, 0x65, 0x6c, 0x54, 0x6d, 0x61, 0x57, 0x61, 0x72, 0x70, 0x53, + 0x70, 0x65, 0x63, 0x69, 0x61, 0x6c, 0x69, 0x7a, 0x65, 0x64, 0x43, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x3e, 0x3a, 0x3a, + 0x43, 0x6f, 0x6c, 0x6c, 0x65, 0x63, 0x74, 0x69, 0x76, 0x65, 0x4f, 0x70, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x63, 0x75, 0x74, 0x6c, 0x61, 0x73, 0x73, 0x3a, 0x3a, 0x67, 0x65, 0x6d, + 0x6d, 0x3a, 0x3a, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3a, 0x3a, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x47, 0x65, 0x6d, 0x6d, 0x55, 0x6e, + 0x69, 0x76, 0x65, 0x72, 0x73, 0x61, 0x6c, 0x3c, 0x53, 0x68, 0x61, 0x70, + 0x65, 0x3c, 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x2c, 0x20, + 0x69, 0x6e, 0x74, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x3e, 0x2c, 0x20, 0x4d, + 0x61, 0x69, 0x6e, 0x6c, 0x6f, 0x6f, 0x70, 0x2c, 0x20, 0x45, 0x70, 0x69, + 0x6c, 0x6f, 0x67, 0x75, 0x65, 0x3e, 0x7b, 0x7d, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, + 0x6e, 0x74, 0x20, 0x47, 0x72, 0x6f, 0x75, 0x70, 0x53, 0x69, 0x7a, 0x65, + 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x45, + 0x6c, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x43, 0x74, 0x61, + 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x0a, 0x75, 0x73, 0x69, 0x6e, 0x67, + 0x20, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, 0x6b, 0x65, + 0x72, 0x6e, 0x65, 0x6c, 0x5f, 0x74, 0x20, 0x3d, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x64, 0x65, 0x63, 0x6c, 0x74, 0x79, 0x70, 0x65, 0x28, 0x6d, 0x61, + 0x6b, 0x65, 0x5f, 0x71, 0x6d, 0x6d, 0x5f, 0x73, 0x6d, 0x39, 0x30, 0x5f, + 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x3c, 0x47, 0x72, 0x6f, 0x75, 0x70, + 0x53, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x45, 0x6c, 0x65, 0x6d, 0x65, 0x6e, + 0x74, 0x2c, 0x20, 0x51, 0x75, 0x61, 0x6e, 0x74, 0x2c, 0x20, 0x43, 0x74, + 0x61, 0x54, 0x69, 0x6c, 0x65, 0x72, 0x3e, 0x28, 0x29, 0x29, 0x3b, 0x0a, + 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, + 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, + 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_scatter_axis[] = { @@ -2558,90 +6114,599 @@ constexpr char jit_source_scatter_axis[] = { 0x5f, 0x6e, 0x64, 0x3c, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x20, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, - 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, - 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, - 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, - 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, - 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, - 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, - 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, - 0x0a, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, - 0x20, 0x2b, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, - 0x75, 0x70, 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, - 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, - 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, - 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, + 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x20, 0x2b, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x61, 0x78, + 0x69, 0x73, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x78, 0x3b, + 0x0a, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, + 0x20, 0x2b, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, + 0x75, 0x70, 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, + 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, + 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter_ops[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, + 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, + 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x61, 0x74, 0x6f, + 0x6d, 0x69, 0x63, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, + 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, + 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, + 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x41, 0x73, 0x73, 0x69, 0x67, + 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x2a, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x53, 0x75, + 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, + 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x64, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, + 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, + 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x61, 0x78, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, + 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, + 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x6f, 0x75, 0x74, + 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, + 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x69, 0x6e, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, + 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, + 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, + 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, + 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, + 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, +}; + +constexpr char jit_source_scatter[] = { + 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, + 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, + 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x69, 0x6e, 0x67, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x23, 0x69, 0x6e, + 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, + 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, 0x61, 0x2f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x63, 0x61, 0x74, 0x74, + 0x65, 0x72, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, + 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, + 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, + 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, 0x74, + 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, 0x69, + 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, + 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, + 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, + 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, 0x61, + 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, 0x3d, + 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, + 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, + 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x49, 0x44, + 0x58, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x49, + 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4c, 0x6f, + 0x63, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x63, 0x61, 0x74, + 0x74, 0x65, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x75, 0x70, 0x64, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x70, + 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, + 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3e, 0x20, 0x61, 0x78, + 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, + 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2a, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, + 0x3e, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, + 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, + 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, + 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, + 0x73, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, + 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x61, 0x72, 0x72, 0x61, 0x79, 0x3c, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x2c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x20, 0x2a, 0x20, 0x49, 0x44, + 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, + 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, + 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x75, 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, + 0x70, 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x25, 0x20, 0x75, 0x70, 0x64, + 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, + 0x64, 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2f, 0x20, 0x75, 0x70, 0x64, 0x5f, + 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, 0x2b, 0x20, 0x49, 0x44, 0x58, + 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, + 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, + 0x29, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, + 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, + 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x49, 0x44, 0x58, 0x3b, 0x20, 0x2b, 0x2b, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x4c, 0x6f, 0x63, + 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, + 0x6e, 0x64, 0x3c, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, 0x49, 0x4d, 0x3e, + 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x20, 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, + 0x44, 0x49, 0x4d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x64, 0x69, 0x63, 0x65, 0x73, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x20, + 0x2b, 0x20, 0x69, 0x20, 0x2a, 0x20, 0x49, 0x44, 0x58, 0x5f, 0x4e, 0x44, + 0x49, 0x4d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x33, 0x32, 0x5f, 0x74, 0x20, 0x61, 0x78, 0x69, 0x73, 0x20, 0x3d, 0x20, + 0x61, 0x78, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, + 0x6c, 0x20, 0x3d, 0x20, 0x61, 0x62, 0x73, 0x6f, 0x6c, 0x75, 0x74, 0x65, + 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x28, 0x69, 0x6e, 0x64, 0x69, 0x63, + 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x5b, 0x69, 0x64, 0x78, 0x5f, 0x6c, 0x6f, + 0x63, 0x5d, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x5b, 0x61, 0x78, 0x69, 0x73, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x3d, + 0x20, 0x69, 0x64, 0x78, 0x5f, 0x76, 0x61, 0x6c, 0x20, 0x2a, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x61, + 0x78, 0x69, 0x73, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x4c, 0x6f, 0x63, 0x54, 0x20, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, + 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2b, 0x20, 0x69, 0x64, + 0x78, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2a, 0x20, 0x75, 0x70, 0x64, + 0x5f, 0x70, 0x6f, 0x73, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5f, 0x73, 0x69, + 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, + 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, + 0x74, 0x61, 0x28, 0x29, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x75, 0x70, 0x64, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x4f, 0x70, 0x7b, 0x7d, 0x28, 0x6f, 0x75, 0x74, 0x20, 0x2b, + 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x75, 0x70, + 0x64, 0x5b, 0x75, 0x70, 0x64, 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x29, 0x3b, + 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x53, 0x72, 0x63, 0x43, 0x6f, + 0x6e, 0x74, 0x69, 0x67, 0x75, 0x6f, 0x75, 0x73, 0x2c, 0x20, 0x62, 0x6f, + 0x6f, 0x6c, 0x20, 0x44, 0x73, 0x74, 0x43, 0x6f, 0x6e, 0x74, 0x69, 0x67, + 0x75, 0x6f, 0x75, 0x73, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, + 0x6c, 0x6f, 0x62, 0x61, 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, + 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x65, 0x64, 0x5f, 0x73, 0x63, 0x61, 0x74, + 0x74, 0x65, 0x72, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x64, 0x73, 0x74, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x62, 0x6f, 0x6f, + 0x6c, 0x2a, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x5f, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x73, 0x72, + 0x63, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, + 0x74, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, + 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, + 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, + 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, + 0x6e, 0x74, 0x33, 0x32, 0x5f, 0x74, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x6e, + 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, + 0x70, 0x65, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, + 0x5f, 0x74, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, + 0x69, 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, + 0x72, 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x54, 0x20, 0x64, 0x73, 0x74, + 0x5f, 0x76, 0x61, 0x6c, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x44, 0x73, + 0x74, 0x43, 0x6f, 0x6e, 0x74, 0x69, 0x67, 0x75, 0x6f, 0x75, 0x73, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x76, + 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x64, 0x73, 0x74, 0x5b, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, + 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x64, 0x73, 0x74, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, + 0x2c, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, + 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x64, 0x73, 0x74, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, + 0x28, 0x29, 0x2c, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x6e, 0x64, 0x69, 0x6d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x76, + 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x64, 0x73, 0x74, 0x5b, 0x64, 0x73, 0x74, + 0x5f, 0x6c, 0x6f, 0x63, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x6d, 0x61, 0x73, 0x6b, 0x5b, 0x69, + 0x6e, 0x64, 0x65, 0x78, 0x5d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, + 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x28, + 0x73, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x5f, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x73, 0x5b, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5d, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x73, 0x72, 0x63, + 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3c, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x3d, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2f, 0x20, 0x6d, 0x61, + 0x73, 0x6b, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x53, + 0x72, 0x63, 0x43, 0x6f, 0x6e, 0x74, 0x69, 0x67, 0x75, 0x6f, 0x75, 0x73, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x6f, 0x75, 0x74, 0x5b, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5d, 0x20, 0x3d, + 0x20, 0x73, 0x72, 0x63, 0x5b, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x2a, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x73, 0x72, + 0x63, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5d, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3d, + 0x20, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, + 0x20, 0x73, 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x20, 0x2b, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, + 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x65, 0x6c, 0x65, + 0x6d, 0x2c, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, + 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, + 0x61, 0x28, 0x29, 0x2c, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x6e, 0x64, 0x69, + 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x6f, 0x75, 0x74, 0x5b, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5d, 0x20, 0x3d, + 0x20, 0x73, 0x72, 0x63, 0x5b, 0x73, 0x72, 0x63, 0x5f, 0x6c, 0x6f, 0x63, + 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x5d, + 0x20, 0x3d, 0x20, 0x64, 0x73, 0x74, 0x5f, 0x76, 0x61, 0x6c, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x5f, 0x52, 0x45, + 0x41, 0x44, 0x53, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, + 0x6c, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6d, 0x61, 0x73, + 0x6b, 0x65, 0x64, 0x5f, 0x73, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x5f, + 0x76, 0x65, 0x63, 0x5f, 0x63, 0x6f, 0x6e, 0x74, 0x69, 0x67, 0x75, 0x6f, + 0x75, 0x73, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x54, 0x2a, 0x20, 0x64, 0x73, 0x74, 0x2c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x2a, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x33, 0x32, 0x5f, + 0x74, 0x2a, 0x20, 0x73, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x5f, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, 0x2a, 0x20, 0x73, 0x72, 0x63, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x73, + 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x5f, 0x62, 0x61, 0x74, 0x63, + 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x76, 0x65, 0x63, 0x5f, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x20, 0x3d, 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, + 0x73, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, + 0x65, 0x61, 0x64, 0x5f, 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x61, 0x73, 0x65, 0x20, + 0x3d, 0x20, 0x76, 0x65, 0x63, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, + 0x2a, 0x20, 0x4e, 0x5f, 0x52, 0x45, 0x41, 0x44, 0x53, 0x3b, 0x0a, 0x20, + 0x20, 0x69, 0x66, 0x20, 0x28, 0x62, 0x61, 0x73, 0x65, 0x20, 0x3e, 0x3d, + 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x0a, 0x20, 0x20, 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6f, 0x75, 0x74, + 0x5f, 0x76, 0x65, 0x63, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x5f, + 0x76, 0x65, 0x63, 0x74, 0x6f, 0x72, 0x3c, 0x4e, 0x5f, 0x52, 0x45, 0x41, + 0x44, 0x53, 0x3e, 0x28, 0x64, 0x73, 0x74, 0x2c, 0x20, 0x76, 0x65, 0x63, + 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, + 0x2c, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x5f, 0x63, 0x61, 0x73, + 0x74, 0x3c, 0x54, 0x3e, 0x28, 0x30, 0x29, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x61, 0x75, 0x74, 0x6f, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x5f, 0x76, 0x65, + 0x63, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x76, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x3c, 0x4e, 0x5f, 0x52, 0x45, 0x41, 0x44, 0x53, 0x3e, + 0x28, 0x6d, 0x61, 0x73, 0x6b, 0x2c, 0x20, 0x76, 0x65, 0x63, 0x5f, 0x69, + 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, + 0x66, 0x61, 0x6c, 0x73, 0x65, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x61, 0x75, + 0x74, 0x6f, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x76, 0x65, + 0x63, 0x20, 0x3d, 0x20, 0x6c, 0x6f, 0x61, 0x64, 0x5f, 0x76, 0x65, 0x63, + 0x74, 0x6f, 0x72, 0x3c, 0x4e, 0x5f, 0x52, 0x45, 0x41, 0x44, 0x53, 0x3e, + 0x28, 0x73, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x5f, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x73, 0x2c, 0x20, 0x76, 0x65, 0x63, 0x5f, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x20, 0x30, + 0x29, 0x3b, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, + 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, + 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x20, 0x69, 0x20, 0x3c, 0x20, 0x4e, 0x5f, 0x52, 0x45, 0x41, 0x44, 0x53, + 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, + 0x3d, 0x20, 0x62, 0x61, 0x73, 0x65, 0x20, 0x2b, 0x20, 0x69, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x62, 0x72, 0x65, 0x61, 0x6b, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x6d, 0x61, 0x73, 0x6b, 0x5f, 0x76, 0x65, 0x63, + 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, + 0x5f, 0x63, 0x61, 0x73, 0x74, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x28, + 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x5f, 0x76, 0x65, 0x63, 0x5b, 0x69, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x28, 0x73, 0x72, 0x63, 0x5f, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, + 0x3c, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x61, 0x74, + 0x63, 0x68, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x20, 0x2f, 0x20, 0x6d, 0x61, 0x73, 0x6b, 0x5f, 0x62, 0x61, + 0x74, 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x76, 0x65, + 0x63, 0x5b, 0x69, 0x5d, 0x20, 0x3d, 0x20, 0x73, 0x72, 0x63, 0x5b, 0x62, + 0x61, 0x74, 0x63, 0x68, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x73, + 0x72, 0x63, 0x5f, 0x62, 0x61, 0x74, 0x63, 0x68, 0x5f, 0x73, 0x69, 0x7a, + 0x65, 0x20, 0x2b, 0x20, 0x73, 0x72, 0x63, 0x5f, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x73, 0x74, 0x6f, 0x72, 0x65, 0x5f, 0x76, 0x65, 0x63, 0x74, 0x6f, + 0x72, 0x3c, 0x4e, 0x5f, 0x52, 0x45, 0x41, 0x44, 0x53, 0x3e, 0x28, 0x6f, + 0x75, 0x74, 0x2c, 0x20, 0x76, 0x65, 0x63, 0x5f, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x76, 0x65, 0x63, 0x2c, 0x20, + 0x73, 0x69, 0x7a, 0x65, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x7d, 0x20, + 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, + 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, + 0x63, 0x75, 0x0a, 0x00, }; -constexpr char jit_source_scatter_ops[] = { +constexpr char jit_source_slice_update[] = { 0x2f, 0x2f, 0x20, 0x43, 0x6f, 0x70, 0x79, 0x72, 0x69, 0x67, 0x68, 0x74, 0x20, 0x20, 0x20, 0x20, 0x32, 0x30, 0x32, 0x35, 0x20, 0x41, 0x70, 0x70, 0x6c, 0x65, 0x20, 0x49, 0x6e, 0x63, 0x2e, 0x0a, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x6f, 0x6e, 0x63, 0x65, 0x0a, 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, 0x75, 0x64, - 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x61, 0x74, 0x6f, - 0x6d, 0x69, 0x63, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, - 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, - 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, - 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, - 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x41, 0x73, 0x73, 0x69, 0x67, - 0x6e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, - 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, - 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x2a, 0x6f, 0x75, 0x74, 0x20, 0x3d, 0x20, 0x76, 0x61, 0x6c, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, - 0x63, 0x74, 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x53, 0x75, - 0x6d, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, - 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, - 0x20, 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, - 0x65, 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, - 0x6f, 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x61, 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x61, 0x64, 0x64, 0x28, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, - 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x50, 0x72, 0x6f, 0x64, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x54, 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, - 0x72, 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, - 0x74, 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x70, 0x72, 0x6f, 0x64, 0x28, 0x6f, - 0x75, 0x74, 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, - 0x20, 0x53, 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x61, 0x78, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, - 0x3e, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, - 0x61, 0x74, 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, - 0x74, 0x2c, 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, - 0x6e, 0x73, 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, - 0x6f, 0x6d, 0x69, 0x63, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x6f, 0x75, 0x74, - 0x2c, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x7d, 0x3b, 0x0a, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x53, - 0x63, 0x61, 0x74, 0x74, 0x65, 0x72, 0x4d, 0x69, 0x6e, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, - 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, - 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, - 0x6f, 0x72, 0x28, 0x29, 0x28, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, - 0x20, 0x54, 0x20, 0x76, 0x61, 0x6c, 0x29, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x74, 0x6f, 0x6d, - 0x69, 0x63, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x6f, 0x75, 0x74, 0x2c, 0x20, - 0x76, 0x61, 0x6c, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, - 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, - 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x62, 0x69, 0x6e, + 0x61, 0x72, 0x79, 0x5f, 0x6f, 0x70, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, + 0x0a, 0x23, 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x22, 0x6d, + 0x6c, 0x78, 0x2f, 0x62, 0x61, 0x63, 0x6b, 0x65, 0x6e, 0x64, 0x2f, 0x63, + 0x75, 0x64, 0x61, 0x2f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x75, + 0x74, 0x69, 0x6c, 0x73, 0x2e, 0x63, 0x75, 0x68, 0x22, 0x0a, 0x0a, 0x23, + 0x69, 0x6e, 0x63, 0x6c, 0x75, 0x64, 0x65, 0x20, 0x3c, 0x63, 0x6f, 0x6f, + 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, 0x65, 0x5f, 0x67, 0x72, 0x6f, + 0x75, 0x70, 0x73, 0x2e, 0x68, 0x3e, 0x0a, 0x0a, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x20, 0x7b, 0x0a, 0x0a, 0x6e, + 0x61, 0x6d, 0x65, 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x63, 0x67, 0x20, + 0x3d, 0x20, 0x63, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x61, 0x74, 0x69, 0x76, + 0x65, 0x5f, 0x67, 0x72, 0x6f, 0x75, 0x70, 0x73, 0x3b, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, + 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x70, + 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x4f, + 0x55, 0x54, 0x5f, 0x52, 0x4f, 0x57, 0x5f, 0x43, 0x4f, 0x4e, 0x54, 0x49, + 0x47, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, + 0x55, 0x50, 0x44, 0x5f, 0x52, 0x4f, 0x57, 0x5f, 0x43, 0x4f, 0x4e, 0x54, + 0x49, 0x47, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x6f, 0x6f, 0x6c, + 0x20, 0x55, 0x50, 0x44, 0x5f, 0x53, 0x43, 0x41, 0x4c, 0x41, 0x52, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x57, 0x4f, + 0x52, 0x4b, 0x3e, 0x0a, 0x5f, 0x5f, 0x67, 0x6c, 0x6f, 0x62, 0x61, 0x6c, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x73, 0x6c, 0x69, 0x63, + 0x65, 0x5f, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x6f, 0x70, 0x28, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x54, + 0x2a, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x2a, 0x20, 0x6f, 0x75, 0x74, 0x2c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x75, + 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x73, 0x69, 0x7a, 0x65, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, + 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, + 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x68, 0x61, 0x70, 0x65, 0x20, 0x75, 0x70, + 0x64, 0x61, 0x74, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x5f, 0x5f, + 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x61, 0x6e, + 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x20, + 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x33, + 0x32, 0x5f, 0x74, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x6e, + 0x64, 0x69, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x5f, 0x5f, 0x67, 0x72, 0x69, 0x64, 0x5f, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x74, 0x5f, 0x5f, 0x20, 0x53, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x20, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x20, 0x6f, 0x75, 0x74, + 0x70, 0x75, 0x74, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x4f, 0x70, 0x20, 0x6f, 0x70, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x69, 0x64, 0x78, 0x20, 0x3d, + 0x20, 0x63, 0x67, 0x3a, 0x3a, 0x74, 0x68, 0x69, 0x73, 0x5f, 0x67, 0x72, + 0x69, 0x64, 0x28, 0x29, 0x2e, 0x74, 0x68, 0x72, 0x65, 0x61, 0x64, 0x5f, + 0x72, 0x61, 0x6e, 0x6b, 0x28, 0x29, 0x20, 0x2a, 0x20, 0x4e, 0x57, 0x4f, + 0x52, 0x4b, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x6f, + 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x69, 0x64, + 0x78, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x4f, 0x55, 0x54, 0x5f, + 0x52, 0x4f, 0x57, 0x5f, 0x43, 0x4f, 0x4e, 0x54, 0x49, 0x47, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x3d, 0x20, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x49, 0x64, + 0x78, 0x54, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, + 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, + 0x29, 0x2c, 0x20, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, 0x74, 0x61, 0x28, 0x29, + 0x2c, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x6e, 0x64, 0x69, + 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x69, + 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x28, 0x21, 0x55, 0x50, 0x44, 0x5f, 0x53, 0x43, 0x41, 0x4c, 0x41, 0x52, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x55, 0x50, + 0x44, 0x5f, 0x52, 0x4f, 0x57, 0x5f, 0x43, 0x4f, 0x4e, 0x54, 0x49, 0x47, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, + 0x64, 0x61, 0x74, 0x65, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x69, + 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, + 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, + 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, + 0x49, 0x64, 0x78, 0x54, 0x3e, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x64, 0x78, 0x2c, 0x20, 0x75, 0x70, + 0x64, 0x61, 0x74, 0x65, 0x5f, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2e, 0x64, + 0x61, 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, + 0x65, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2e, 0x64, 0x61, + 0x74, 0x61, 0x28, 0x29, 0x2c, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, + 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, + 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x0a, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6f, + 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x6a, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x6a, 0x20, 0x3c, + 0x20, 0x4e, 0x57, 0x4f, 0x52, 0x4b, 0x20, 0x26, 0x26, 0x20, 0x69, 0x64, + 0x78, 0x20, 0x3c, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x73, + 0x69, 0x7a, 0x65, 0x3b, 0x20, 0x6a, 0x2b, 0x2b, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5b, 0x6f, 0x75, 0x74, 0x5f, + 0x69, 0x64, 0x78, 0x5d, 0x20, 0x3d, 0x20, 0x6f, 0x70, 0x28, 0x6f, 0x75, + 0x74, 0x5b, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x2c, 0x20, + 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x73, 0x5b, 0x75, 0x70, 0x64, 0x61, + 0x74, 0x65, 0x5f, 0x69, 0x64, 0x78, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x64, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x28, 0x4f, 0x55, 0x54, 0x5f, 0x52, 0x4f, 0x57, 0x5f, + 0x43, 0x4f, 0x4e, 0x54, 0x49, 0x47, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, + 0x3d, 0x20, 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x75, 0x74, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2b, 0x3d, + 0x20, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x5b, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x6e, + 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x3b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x55, + 0x50, 0x44, 0x5f, 0x52, 0x4f, 0x57, 0x5f, 0x43, 0x4f, 0x4e, 0x54, 0x49, + 0x47, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x75, + 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, + 0x69, 0x64, 0x78, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, + 0x6c, 0x73, 0x65, 0x20, 0x69, 0x66, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x28, 0x21, 0x55, 0x50, 0x44, 0x5f, 0x53, + 0x43, 0x41, 0x4c, 0x41, 0x52, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x2b, 0x3d, 0x20, 0x75, 0x70, 0x64, 0x61, 0x74, 0x65, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x75, 0x70, 0x64, 0x61, + 0x74, 0x65, 0x5f, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; constexpr char jit_source_ternary_ops[] = { @@ -3598,301 +7663,349 @@ constexpr char jit_source_utils[] = { 0x69, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x20, 0x3c, 0x20, 0x73, 0x69, 0x7a, 0x65, 0x20, 0x26, 0x26, 0x20, 0x69, 0x20, - 0x3c, 0x20, 0x4e, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x70, 0x74, 0x72, 0x5b, 0x73, 0x74, - 0x72, 0x69, 0x64, 0x65, 0x20, 0x2a, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, - 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x5d, - 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x54, - 0x79, 0x70, 0x65, 0x20, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x75, - 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, 0x20, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, 0x76, - 0x6f, 0x69, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, - 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, - 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, - 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, - 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, - 0x54, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, - 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, - 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, - 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, - 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, - 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, - 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, - 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, - 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, - 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, - 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, - 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, - 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, - 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, - 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, - 0x61, 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, - 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, - 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, - 0x3e, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, - 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, - 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x64, 0x6f, 0x75, 0x62, 0x6c, 0x65, 0x3e, - 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, - 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, - 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x61, - 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, - 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, - 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, - 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, - 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x3c, 0x20, 0x4e, 0x3b, 0x20, 0x2b, 0x2b, 0x69, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x70, 0x74, 0x72, 0x5b, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x20, 0x2a, 0x20, 0x28, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x20, 0x2a, 0x20, 0x4e, 0x20, 0x2b, 0x20, 0x69, 0x29, 0x5d, + 0x20, 0x3d, 0x20, 0x76, 0x65, 0x63, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x0a, 0x0a, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x54, + 0x79, 0x70, 0x65, 0x20, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x75, + 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x0a, 0x0a, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, + 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x46, 0x38, 0x45, 0x34, 0x4d, 0x33, + 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3d, 0x20, 0x34, 0x34, 0x38, 0x2e, 0x30, + 0x66, 0x3b, 0x0a, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x20, 0x46, 0x34, 0x45, 0x32, 0x4d, + 0x31, 0x5f, 0x4d, 0x41, 0x58, 0x20, 0x3d, 0x20, 0x36, 0x2e, 0x30, 0x66, + 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x2c, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x3d, 0x20, + 0x76, 0x6f, 0x69, 0x64, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, + 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x54, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, + 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, + 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, - 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, - 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, - 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, - 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, - 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, - 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, - 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, - 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, - 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, - 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, - 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, - 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x20, - 0x43, 0x55, 0x44, 0x41, 0x20, 0x31, 0x31, 0x20, 0x64, 0x6f, 0x65, 0x73, - 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x68, 0x61, 0x76, 0x65, 0x20, 0x68, 0x6f, - 0x73, 0x74, 0x20, 0x73, 0x69, 0x64, 0x65, 0x20, 0x61, 0x72, 0x69, 0x74, - 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x20, 0x6f, 0x70, 0x65, 0x72, 0x61, - 0x74, 0x6f, 0x72, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x68, 0x61, 0x6c, - 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x74, 0x65, 0x6d, - 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, - 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, - 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, 0x62, - 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, - 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, - 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, 0x3e, - 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, - 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, - 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x31, - 0x36, 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, - 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, - 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, - 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, - 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, - 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, - 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, - 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, - 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, - 0x44, 0x41, 0x52, 0x54, 0x5f, 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, - 0x20, 0x3c, 0x20, 0x31, 0x32, 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, 0x20, - 0x5f, 0x5f, 0x43, 0x55, 0x44, 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, 0x5f, - 0x5f, 0x20, 0x3c, 0x20, 0x38, 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, - 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, - 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, - 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, - 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6c, 0x73, 0x65, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, - 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, - 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, - 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, 0x66, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, - 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, - 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, - 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, - 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, - 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, - 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, - 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, - 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, - 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, - 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, - 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, 0x20, 0x3c, 0x20, 0x31, 0x32, - 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, 0x20, 0x5f, 0x5f, 0x43, 0x55, 0x44, - 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, 0x5f, 0x5f, 0x20, 0x3c, 0x20, 0x38, - 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, - 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, - 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, - 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, - 0x6c, 0x73, 0x65, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, - 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, 0x77, - 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, - 0x66, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, - 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, 0x74, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, + 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, + 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, + 0x28, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, - 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, + 0x6e, 0x61, 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, + 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x66, 0x6c, 0x6f, 0x61, + 0x74, 0x3e, 0x20, 0x7c, 0x7c, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, + 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x64, 0x6f, 0x75, 0x62, 0x6c, 0x65, + 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, + 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, + 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, + 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, + 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, + 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, + 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, + 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, 0x69, 0x6e, + 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, + 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, + 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, + 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, + 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, + 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, + 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, + 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, + 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, 0x69, 0x6e, 0x69, + 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, + 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, + 0x54, 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, + 0x20, 0x43, 0x55, 0x44, 0x41, 0x20, 0x31, 0x31, 0x20, 0x64, 0x6f, 0x65, + 0x73, 0x20, 0x6e, 0x6f, 0x74, 0x20, 0x68, 0x61, 0x76, 0x65, 0x20, 0x68, + 0x6f, 0x73, 0x74, 0x20, 0x73, 0x69, 0x64, 0x65, 0x20, 0x61, 0x72, 0x69, + 0x74, 0x68, 0x6d, 0x65, 0x74, 0x69, 0x63, 0x20, 0x6f, 0x70, 0x65, 0x72, + 0x61, 0x74, 0x6f, 0x72, 0x73, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x68, 0x61, + 0x6c, 0x66, 0x20, 0x74, 0x79, 0x70, 0x65, 0x73, 0x2e, 0x0a, 0x74, 0x65, + 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, + 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, + 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x54, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x65, 0x6e, 0x61, + 0x62, 0x6c, 0x65, 0x5f, 0x69, 0x66, 0x5f, 0x74, 0x3c, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, + 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, + 0x5f, 0x76, 0x3c, 0x54, 0x2c, 0x20, 0x5f, 0x5f, 0x68, 0x61, 0x6c, 0x66, + 0x3e, 0x20, 0x7c, 0x7c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x69, 0x73, 0x5f, 0x73, 0x61, 0x6d, 0x65, 0x5f, 0x76, 0x3c, 0x54, 0x2c, + 0x20, 0x5f, 0x5f, 0x6e, 0x76, 0x5f, 0x62, 0x66, 0x6c, 0x6f, 0x61, 0x74, + 0x31, 0x36, 0x3e, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, - 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x62, - 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x74, - 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, + 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, + 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, + 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, 0x69, 0x74, 0x79, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x6d, + 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, + 0x55, 0x44, 0x41, 0x52, 0x54, 0x5f, 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, + 0x4e, 0x20, 0x3c, 0x20, 0x31, 0x32, 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, + 0x20, 0x5f, 0x5f, 0x43, 0x55, 0x44, 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, + 0x5f, 0x5f, 0x20, 0x3c, 0x20, 0x38, 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x2d, 0x63, 0x75, 0x64, + 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, + 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, + 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, + 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6c, 0x73, 0x65, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x2d, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, + 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x69, 0x6e, 0x66, 0x69, 0x6e, + 0x69, 0x74, 0x79, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, 0x69, + 0x66, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, + 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, + 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, 0x6d, 0x69, + 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, + 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, + 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, + 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x54, 0x20, 0x66, + 0x69, 0x6e, 0x69, 0x74, 0x65, 0x5f, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x23, 0x69, 0x66, 0x20, 0x43, 0x55, 0x44, 0x41, 0x52, 0x54, + 0x5f, 0x56, 0x45, 0x52, 0x53, 0x49, 0x4f, 0x4e, 0x20, 0x3c, 0x20, 0x31, + 0x32, 0x30, 0x30, 0x30, 0x20, 0x26, 0x26, 0x20, 0x5f, 0x5f, 0x43, 0x55, + 0x44, 0x41, 0x5f, 0x41, 0x52, 0x43, 0x48, 0x5f, 0x5f, 0x20, 0x3c, 0x20, + 0x38, 0x30, 0x30, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, 0x69, + 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x66, 0x6c, 0x6f, 0x61, 0x74, 0x3e, 0x3a, + 0x3a, 0x6c, 0x6f, 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, + 0x65, 0x6c, 0x73, 0x65, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, + 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, + 0x64, 0x3a, 0x3a, 0x6e, 0x75, 0x6d, 0x65, 0x72, 0x69, 0x63, 0x5f, 0x6c, + 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6c, 0x6f, + 0x77, 0x65, 0x73, 0x74, 0x28, 0x29, 0x3b, 0x0a, 0x23, 0x65, 0x6e, 0x64, + 0x69, 0x66, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x3e, 0x0a, 0x73, + 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, + 0x3c, 0x62, 0x6f, 0x6f, 0x6c, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, - 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, + 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x66, 0x61, 0x6c, 0x73, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, - 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, - 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, 0x3e, - 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, 0x69, - 0x74, 0x73, 0x3c, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, - 0x3c, 0x54, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, - 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, - 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x6f, - 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x6d, - 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x4c, 0x69, 0x6d, 0x69, 0x74, - 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x2c, - 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, - 0x6d, 0x61, 0x78, 0x28, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, - 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, - 0x3c, 0x54, 0x3e, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, - 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, - 0x69, 0x6e, 0x28, 0x29, 0x2c, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, - 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x7d, 0x3b, - 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, + 0x74, 0x72, 0x75, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, + 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x66, 0x61, 0x6c, 0x73, 0x65, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, + 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, + 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x54, + 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x69, 0x6d, + 0x69, 0x74, 0x73, 0x3c, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, + 0x74, 0x3c, 0x54, 0x3e, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x73, 0x74, + 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x65, 0x78, + 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, + 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, 0x74, 0x3c, 0x54, 0x3e, 0x20, + 0x6d, 0x61, 0x78, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x7b, 0x4c, 0x69, 0x6d, 0x69, + 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, + 0x2c, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, + 0x3a, 0x6d, 0x61, 0x78, 0x28, 0x29, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x73, 0x74, 0x61, 0x74, 0x69, 0x63, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x65, 0x78, 0x70, 0x72, 0x20, 0x5f, 0x5f, 0x68, 0x6f, + 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x5f, + 0x74, 0x3c, 0x54, 0x3e, 0x20, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x20, 0x7b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x7b, 0x4c, 0x69, 0x6d, 0x69, 0x74, 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, + 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x2c, 0x20, 0x4c, 0x69, 0x6d, 0x69, 0x74, + 0x73, 0x3c, 0x54, 0x3e, 0x3a, 0x3a, 0x6d, 0x69, 0x6e, 0x28, 0x29, 0x7d, + 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x49, 0x6e, 0x64, 0x65, - 0x78, 0x69, 0x6e, 0x67, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x49, 0x6e, 0x64, + 0x65, 0x78, 0x69, 0x6e, 0x67, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, - 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, - 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, - 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, - 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x49, 0x64, - 0x78, 0x54, 0x0a, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, - 0x6f, 0x63, 0x28, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, - 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, - 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, - 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, - 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x0a, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, + 0x6c, 0x6f, 0x63, 0x28, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, + 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, + 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x69, 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, + 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x3e, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, 0x63, + 0x20, 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, + 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x2a, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, + 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, + 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, + 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, 0x7d, 0x0a, + 0x0a, 0x2f, 0x2f, 0x20, 0x4f, 0x70, 0x74, 0x69, 0x6d, 0x69, 0x7a, 0x65, + 0x20, 0x77, 0x68, 0x65, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x6e, 0x64, + 0x69, 0x6d, 0x20, 0x69, 0x73, 0x20, 0x6b, 0x6e, 0x6f, 0x77, 0x6e, 0x20, + 0x61, 0x74, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, 0x65, 0x20, 0x74, + 0x69, 0x6d, 0x65, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, + 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, + 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, + 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x49, 0x64, 0x78, 0x54, 0x0a, 0x65, 0x6c, + 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, + 0x28, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, - 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, - 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, - 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x20, 0x26, 0x26, 0x20, 0x65, - 0x6c, 0x65, 0x6d, 0x20, 0x3e, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, - 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, 0x63, 0x20, - 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, - 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x2a, 0x20, 0x49, - 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, - 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, - 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, - 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, - 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, - 0x2f, 0x2f, 0x20, 0x4f, 0x70, 0x74, 0x69, 0x6d, 0x69, 0x7a, 0x65, 0x20, - 0x77, 0x68, 0x65, 0x6e, 0x20, 0x74, 0x68, 0x65, 0x20, 0x6e, 0x64, 0x69, - 0x6d, 0x20, 0x69, 0x73, 0x20, 0x6b, 0x6e, 0x6f, 0x77, 0x6e, 0x20, 0x61, - 0x74, 0x20, 0x63, 0x6f, 0x6d, 0x70, 0x69, 0x6c, 0x65, 0x20, 0x74, 0x69, - 0x6d, 0x65, 0x2e, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, - 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, - 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, - 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, - 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x49, 0x64, 0x78, 0x54, 0x0a, 0x65, 0x6c, 0x65, - 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x28, - 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, - 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, - 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x23, - 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, - 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, - 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, - 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, - 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, 0x6f, - 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, - 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, 0x2a, - 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, - 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, - 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, - 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, 0x7d, + 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, + 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, + 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, + 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, + 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6c, + 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x28, 0x65, 0x6c, 0x65, 0x6d, 0x20, + 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x29, 0x20, + 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6c, 0x6f, 0x63, 0x3b, 0x0a, + 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, + 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, + 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, + 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, + 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, + 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, + 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, + 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x28, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, + 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, + 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, + 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, + 0x30, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, + 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, + 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, + 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, + 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, + 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, @@ -3901,69 +8014,72 @@ constexpr char jit_source_utils[] = { 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, - 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, - 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, 0x6e, 0x64, 0x28, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, - 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, + 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, + 0x5f, 0x6e, 0x64, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, + 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, - 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, - 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, + 0x20, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, - 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, - 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, - 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, - 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, - 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, - 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, - 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, - 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, - 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, - 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, - 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, 0x63, - 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, - 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, - 0x6e, 0x74, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x74, 0x79, 0x70, - 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, - 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, - 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, - 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, - 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, - 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, - 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, - 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x5f, - 0x6e, 0x64, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, - 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, - 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, - 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, - 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, - 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, - 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x63, 0x5f, 0x6c, 0x6f, - 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, 0x67, - 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, 0x20, + 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x63, 0x5f, 0x6c, + 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x23, 0x70, 0x72, 0x61, + 0x67, 0x6d, 0x61, 0x20, 0x75, 0x6e, 0x72, 0x6f, 0x6c, 0x6c, 0x0a, 0x20, + 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, + 0x3d, 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, + 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, + 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, + 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, + 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, + 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, + 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, + 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, + 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, + 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, + 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x63, 0x5f, + 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, + 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, + 0x61, 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, + 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, + 0x6e, 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, + 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, + 0x70, 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, + 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, + 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x74, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, + 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, - 0x20, 0x4e, 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, + 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, @@ -3975,321 +8091,276 @@ constexpr char jit_source_utils[] = { 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, 0x6c, 0x6f, 0x63, - 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, - 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, 0x75, 0x64, - 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, 0x6b, 0x65, - 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, 0x6f, 0x63, - 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x63, 0x5f, 0x6c, - 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, - 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, - 0x6d, 0x65, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, - 0x74, 0x36, 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, - 0x65, 0x20, 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, - 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, - 0x6c, 0x65, 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, - 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, - 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, - 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, - 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, - 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, + 0x2f, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, + 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, + 0x6e, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, + 0x3a, 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, + 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, + 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, + 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, + 0x34, 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, + 0x5f, 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, + 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, + 0x3c, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, + 0x20, 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, + 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x49, 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, - 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, - 0x64, 0x78, 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, - 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, - 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x66, - 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x20, 0x3d, 0x20, - 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, 0x20, 0x69, 0x20, - 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, - 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, - 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, - 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, - 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x62, - 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, - 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x62, - 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, - 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, - 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, - 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, - 0x6d, 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, - 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x29, - 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, - 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, - 0x49, 0x64, 0x78, 0x54, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, - 0x5f, 0x74, 0x3e, 0x0a, 0x69, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x20, 0x5f, - 0x5f, 0x68, 0x6f, 0x73, 0x74, 0x5f, 0x5f, 0x20, 0x5f, 0x5f, 0x64, 0x65, - 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x63, 0x75, 0x64, 0x61, 0x3a, - 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x3c, - 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, 0x49, 0x64, 0x78, 0x54, 0x2c, 0x20, - 0x49, 0x64, 0x78, 0x54, 0x3e, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, - 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x28, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x49, - 0x64, 0x78, 0x54, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x2c, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, - 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, - 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, - 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, 0x73, - 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, - 0x74, 0x2a, 0x20, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, - 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x64, - 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, - 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, - 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, - 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, - 0x20, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, - 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x69, - 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x3b, - 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, 0x69, - 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, 0x6c, - 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, - 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, - 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, - 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, - 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, - 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, - 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, 0x6c, - 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, - 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, 0x73, - 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, 0x20, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x63, - 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, 0x61, - 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, 0x6c, - 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x63, - 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x2f, 0x2f, + 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, + 0x5f, 0x74, 0x2a, 0x20, 0x61, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, + 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x62, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, + 0x5f, 0x74, 0x2a, 0x20, 0x63, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x2c, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x6e, + 0x64, 0x69, 0x6d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x61, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, 0x54, 0x20, 0x62, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x49, 0x64, 0x78, + 0x54, 0x20, 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x3d, 0x20, 0x30, 0x3b, + 0x0a, 0x20, 0x20, 0x66, 0x6f, 0x72, 0x20, 0x28, 0x69, 0x6e, 0x74, 0x20, + 0x69, 0x20, 0x3d, 0x20, 0x6e, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, + 0x3b, 0x20, 0x69, 0x20, 0x3e, 0x3d, 0x20, 0x30, 0x3b, 0x20, 0x2d, 0x2d, + 0x69, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x3d, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x20, 0x25, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, + 0x69, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x61, 0x5f, 0x6c, 0x6f, + 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, + 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x61, 0x5f, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, + 0x64, 0x69, 0x6d, 0x5f, 0x69, 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, + 0x78, 0x54, 0x28, 0x62, 0x5f, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x5b, 0x69, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x63, 0x5f, + 0x6c, 0x6f, 0x63, 0x20, 0x2b, 0x3d, 0x20, 0x64, 0x69, 0x6d, 0x5f, 0x69, + 0x64, 0x78, 0x20, 0x2a, 0x20, 0x49, 0x64, 0x78, 0x54, 0x28, 0x63, 0x5f, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x69, 0x5d, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x20, 0x2f, 0x3d, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x69, 0x5d, 0x3b, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, + 0x63, 0x75, 0x64, 0x61, 0x3a, 0x3a, 0x73, 0x74, 0x64, 0x3a, 0x3a, 0x6d, + 0x61, 0x6b, 0x65, 0x5f, 0x74, 0x75, 0x70, 0x6c, 0x65, 0x28, 0x61, 0x5f, + 0x6c, 0x6f, 0x63, 0x2c, 0x20, 0x62, 0x5f, 0x6c, 0x6f, 0x63, 0x2c, 0x20, + 0x63, 0x5f, 0x6c, 0x6f, 0x63, 0x29, 0x3b, 0x0a, 0x7d, 0x0a, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x45, 0x6c, 0x65, - 0x6d, 0x20, 0x74, 0x6f, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x69, 0x6e, 0x20, - 0x61, 0x20, 0x6c, 0x6f, 0x6f, 0x70, 0x20, 0x75, 0x74, 0x69, 0x6c, 0x73, - 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x2f, 0x2f, 0x20, 0x45, 0x6c, + 0x65, 0x6d, 0x20, 0x74, 0x6f, 0x20, 0x6c, 0x6f, 0x63, 0x20, 0x69, 0x6e, + 0x20, 0x61, 0x20, 0x6c, 0x6f, 0x6f, 0x70, 0x20, 0x75, 0x74, 0x69, 0x6c, + 0x73, 0x0a, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, - 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, 0x65, - 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, 0x20, - 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x47, 0x65, - 0x6e, 0x65, 0x72, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, 0x65, - 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x3d, 0x20, 0x73, 0x69, 0x7a, - 0x65, 0x5f, 0x74, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, - 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, - 0x4c, 0x6f, 0x63, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, - 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, 0x44, - 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x2c, 0x20, 0x47, 0x65, 0x6e, 0x65, - 0x72, 0x61, 0x6c, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, - 0x3e, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, - 0x65, 0x72, 0x3b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, - 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, 0x3b, - 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, - 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, - 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, - 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x3a, 0x20, 0x64, 0x69, - 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, 0x2c, 0x20, 0x69, 0x6e, 0x6e, 0x65, - 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x28, 0x64, 0x69, 0x6d, - 0x20, 0x2d, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, - 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, - 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, - 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, - 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, 0x20, 0x20, + 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x2f, 0x0a, 0x0a, 0x74, + 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x69, 0x6e, 0x74, + 0x20, 0x44, 0x49, 0x4d, 0x2c, 0x20, 0x62, 0x6f, 0x6f, 0x6c, 0x20, 0x47, + 0x65, 0x6e, 0x65, 0x72, 0x61, 0x6c, 0x20, 0x3d, 0x20, 0x74, 0x72, 0x75, + 0x65, 0x2c, 0x20, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x3d, 0x20, 0x73, 0x69, + 0x7a, 0x65, 0x5f, 0x74, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, + 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, + 0x6f, 0x4c, 0x6f, 0x63, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x4c, 0x6f, 0x6f, 0x70, + 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, + 0x44, 0x49, 0x4d, 0x20, 0x2d, 0x20, 0x31, 0x2c, 0x20, 0x47, 0x65, 0x6e, + 0x65, 0x72, 0x61, 0x6c, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x3e, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, + 0x70, 0x65, 0x72, 0x3b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, + 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, + 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, + 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, + 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x20, 0x3a, 0x20, 0x64, + 0x69, 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, 0x2c, 0x20, 0x69, 0x6e, 0x6e, + 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x28, 0x64, 0x69, + 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x7d, 0x0a, 0x0a, 0x20, + 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, + 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, + 0x72, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, + 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, + 0x69, 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, + 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, + 0x65, 0x78, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, + 0x65, 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, + 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, + 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, + 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, + 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, + 0x6d, 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x64, 0x65, 0x78, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, - 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, - 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, - 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x69, - 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x6e, 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, + 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, + 0x2d, 0x20, 0x31, 0x5d, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, + 0x20, 0x31, 0x5d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3d, + 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2d, 0x20, 0x73, 0x68, 0x61, + 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x65, + 0x78, 0x74, 0x72, 0x61, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x29, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, - 0x78, 0x20, 0x3d, 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, - 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, 0x70, 0x65, - 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, - 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, - 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, - 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, - 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, 0x6e, - 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, - 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, - 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, - 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, - 0x20, 0x3d, 0x3d, 0x20, 0x30, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x3b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, - 0x65, 0x78, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, - 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6e, - 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, - 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, - 0x20, 0x31, 0x5d, 0x29, 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, - 0x66, 0x20, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3e, 0x3d, 0x20, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, - 0x31, 0x5d, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x69, 0x6e, 0x74, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3d, 0x20, - 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2d, 0x20, 0x73, 0x68, 0x61, 0x70, - 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x65, 0x78, - 0x74, 0x72, 0x61, 0x20, 0x3e, 0x3d, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, - 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x6e, - 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6e, 0x65, - 0x78, 0x74, 0x28, 0x31, 0x20, 0x2b, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, - 0x20, 0x2f, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, - 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, - 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x78, 0x74, 0x72, - 0x61, 0x20, 0x3d, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x25, 0x20, - 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, 0x20, - 0x31, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, - 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, - 0x70, 0x65, 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, - 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, 0x20, - 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, - 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, - 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, - 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3e, 0x20, 0x30, 0x29, 0x20, - 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6e, 0x65, - 0x78, 0x74, 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x2c, 0x20, 0x73, 0x68, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, + 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6e, + 0x65, 0x78, 0x74, 0x28, 0x31, 0x20, 0x2b, 0x20, 0x65, 0x78, 0x74, 0x72, + 0x61, 0x20, 0x2f, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, + 0x6d, 0x20, 0x2d, 0x20, 0x31, 0x5d, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, + 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x65, 0x78, 0x74, + 0x72, 0x61, 0x20, 0x3d, 0x20, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x25, + 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x5b, 0x64, 0x69, 0x6d, 0x20, 0x2d, + 0x20, 0x31, 0x5d, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, + 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, 0x5f, 0x6c, 0x6f, + 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, - 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x74, - 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, - 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, - 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, - 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, - 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, - 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, - 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, 0x20, - 0x69, 0x6e, 0x64, 0x65, 0x78, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, 0x20, + 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x3d, + 0x20, 0x30, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x69, 0x6e, 0x6e, 0x65, 0x72, + 0x5f, 0x6c, 0x6f, 0x6f, 0x70, 0x65, 0x72, 0x2e, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, + 0x20, 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x20, 0x3e, 0x20, 0x30, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6e, + 0x65, 0x78, 0x74, 0x28, 0x65, 0x78, 0x74, 0x72, 0x61, 0x2c, 0x20, 0x73, + 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, + 0x73, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, - 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, - 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x29, - 0x20, 0x3a, 0x20, 0x64, 0x69, 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, 0x20, - 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, - 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, 0x6e, - 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, - 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, - 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, 0x6c, - 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, 0x66, - 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, - 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, - 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, - 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, 0x5d, 0x29, - 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, - 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, - 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, - 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, - 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, - 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, - 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, - 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, - 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, 0x6f, 0x63, 0x61, + 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, + 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, 0x20, 0x3c, 0x74, 0x79, + 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, 0x63, 0x74, 0x20, 0x4c, + 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, 0x6f, 0x4c, + 0x6f, 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x74, 0x72, 0x75, 0x65, 0x2c, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x20, 0x7b, 0x0a, 0x20, + 0x20, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, 0x3b, 0x0a, 0x20, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x20, 0x20, 0x69, 0x6e, 0x74, + 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x7b, 0x30, 0x7d, 0x3b, 0x0a, 0x0a, + 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, + 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, 0x6d, 0x54, + 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x20, 0x64, 0x69, 0x6d, + 0x29, 0x20, 0x3a, 0x20, 0x64, 0x69, 0x6d, 0x28, 0x64, 0x69, 0x6d, 0x29, + 0x20, 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, + 0x65, 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x63, 0x6f, + 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, + 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2b, 0x2b, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, + 0x20, 0x3e, 0x20, 0x31, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, + 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, + 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, + 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, - 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, 0x6f, - 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, - 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x68, 0x61, - 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x2c, - 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, - 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, - 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x69, - 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, - 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, - 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, - 0x65, 0x5f, 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, - 0x6c, 0x6f, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, - 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, - 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, 0x65, - 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, 0x75, - 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, - 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x66, 0x61, - 0x6c, 0x73, 0x65, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, - 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, - 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, 0x3b, + 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, 0x5d, + 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, - 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, 0x65, - 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x29, 0x20, - 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, - 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, - 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, - 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, - 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, 0x30, - 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, - 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, - 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, 0x6e, 0x74, 0x20, - 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, - 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, - 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, - 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x20, 0x2a, 0x20, 0x4f, - 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, - 0x65, 0x73, 0x5b, 0x30, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, - 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, - 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, 0x6f, - 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, 0x20, - 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, 0x66, - 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, 0x3b, - 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, 0x73, - 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, 0x6f, - 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, + 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, + 0x28, 0x69, 0x6e, 0x74, 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, + 0x74, 0x20, 0x69, 0x6e, 0x74, 0x2a, 0x20, 0x73, 0x68, 0x61, 0x70, 0x65, + 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, 0x74, 0x36, + 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x69, 0x6e, 0x64, 0x65, + 0x78, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x69, 0x66, 0x20, 0x28, 0x64, 0x69, 0x6d, 0x20, 0x3e, 0x20, 0x31, 0x29, + 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, + 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, 0x65, 0x6c, 0x65, 0x6d, 0x5f, 0x74, + 0x6f, 0x5f, 0x6c, 0x6f, 0x63, 0x3c, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x3e, 0x28, 0x69, 0x6e, 0x64, 0x65, 0x78, 0x2c, 0x20, 0x73, 0x68, + 0x61, 0x70, 0x65, 0x2c, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, + 0x2c, 0x20, 0x64, 0x69, 0x6d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, + 0x7d, 0x20, 0x65, 0x6c, 0x73, 0x65, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, + 0x20, 0x20, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x20, 0x3d, 0x20, + 0x69, 0x6e, 0x64, 0x65, 0x78, 0x20, 0x2a, 0x20, 0x4f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, + 0x30, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x7d, 0x0a, 0x20, + 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, + 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, + 0x20, 0x6c, 0x6f, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, + 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, + 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x7d, 0x3b, 0x0a, 0x0a, 0x74, 0x65, 0x6d, 0x70, 0x6c, 0x61, 0x74, + 0x65, 0x20, 0x3c, 0x74, 0x79, 0x70, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x3e, 0x0a, 0x73, 0x74, 0x72, + 0x75, 0x63, 0x74, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, + 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x3c, 0x31, 0x2c, 0x20, 0x66, + 0x61, 0x6c, 0x73, 0x65, 0x2c, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, + 0x54, 0x3e, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, + 0x74, 0x54, 0x20, 0x6f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x7b, 0x30, 0x7d, + 0x3b, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, + 0x65, 0x5f, 0x5f, 0x20, 0x4c, 0x6f, 0x6f, 0x70, 0x65, 0x64, 0x45, 0x6c, + 0x65, 0x6d, 0x54, 0x6f, 0x4c, 0x6f, 0x63, 0x28, 0x69, 0x6e, 0x74, 0x29, + 0x20, 0x7b, 0x7d, 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, + 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, 0x6f, 0x69, 0x64, 0x20, 0x6e, + 0x65, 0x78, 0x74, 0x28, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x4f, 0x66, 0x66, 0x73, + 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, 0x64, 0x65, 0x73, 0x5b, + 0x30, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x0a, 0x20, 0x20, + 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x5f, 0x5f, 0x20, 0x76, + 0x6f, 0x69, 0x64, 0x20, 0x6e, 0x65, 0x78, 0x74, 0x28, 0x69, 0x6e, 0x74, + 0x20, 0x6e, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x2a, 0x2c, 0x20, 0x63, 0x6f, 0x6e, 0x73, 0x74, 0x20, 0x69, 0x6e, + 0x74, 0x36, 0x34, 0x5f, 0x74, 0x2a, 0x20, 0x73, 0x74, 0x72, 0x69, 0x64, + 0x65, 0x73, 0x29, 0x20, 0x7b, 0x0a, 0x20, 0x20, 0x20, 0x20, 0x6f, 0x66, + 0x66, 0x73, 0x65, 0x74, 0x20, 0x2b, 0x3d, 0x20, 0x6e, 0x20, 0x2a, 0x20, + 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x28, 0x73, 0x74, 0x72, 0x69, + 0x64, 0x65, 0x73, 0x5b, 0x30, 0x5d, 0x29, 0x3b, 0x0a, 0x20, 0x20, 0x7d, + 0x0a, 0x0a, 0x20, 0x20, 0x5f, 0x5f, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, + 0x5f, 0x5f, 0x20, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x54, 0x20, 0x6c, + 0x6f, 0x63, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x28, 0x29, 0x20, 0x7b, 0x0a, + 0x20, 0x20, 0x20, 0x20, 0x72, 0x65, 0x74, 0x75, 0x72, 0x6e, 0x20, 0x6f, + 0x66, 0x66, 0x73, 0x65, 0x74, 0x3b, 0x0a, 0x20, 0x20, 0x7d, 0x0a, 0x7d, + 0x3b, 0x0a, 0x0a, 0x7d, 0x20, 0x2f, 0x2f, 0x20, 0x6e, 0x61, 0x6d, 0x65, + 0x73, 0x70, 0x61, 0x63, 0x65, 0x20, 0x6d, 0x6c, 0x78, 0x3a, 0x3a, 0x63, + 0x6f, 0x72, 0x65, 0x3a, 0x3a, 0x63, 0x75, 0x0a, 0x00, }; } // namespace mlx::core diff --git a/Source/Cmlx/mlx-generated/fft.cpp b/Source/Cmlx/mlx-generated/fft.cpp index 561121e0c..d60d565b8 100644 --- a/Source/Cmlx/mlx-generated/fft.cpp +++ b/Source/Cmlx/mlx-generated/fft.cpp @@ -28,49 +28,64 @@ them into (n-1)=6,10,12 codelets. */ #include #include -METAL_FUNC float2 complex_mul(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); +// The codelets are templated over the scalar lane type so reduced-precision +// complex paths can reuse them; complex values are plain two-lane vectors. + +template +METAL_FUNC metal::vec complex_mul( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x}; } // Complex mul followed by conjugate -METAL_FUNC float2 complex_mul_conj(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x); +template +METAL_FUNC metal::vec complex_mul_conj( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x}; } // Compute an FFT twiddle factor -METAL_FUNC float2 get_twiddle(int k, int p) { - float theta = -2.0f * k * M_PI_F / p; - - float2 twiddle = {metal::fast::cos(theta), metal::fast::sin(theta)}; - return twiddle; +template +METAL_FUNC metal::vec get_twiddle(int k, int p) { + // Derive phase evaluation precision from the scalar lane and Metal's pi + // constant. Reduced lanes currently promote to float for fast trig, then + // narrow the result to the lane type explicitly. + using phase_T = decltype(M_PI_F * T(0)); + phase_T theta = -phase_T(2) * phase_T(k) * phase_T(M_PI_F) / phase_T(p); + return {T(metal::fast::cos(theta)), T(metal::fast::sin(theta))}; } -METAL_FUNC void radix2(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix2(thread metal::vec* x, thread metal::vec* y) { y[0] = x[0] + x[1]; y[1] = x[0] - x[1]; } -METAL_FUNC void radix3(thread float2* x, thread float2* y) { - float pi_2_3 = -0.8660254037844387; +template +METAL_FUNC void radix3(thread metal::vec* x, thread metal::vec* y) { + T pi_2_3 = -0.8660254037844387; - float2 a_1 = x[1] + x[2]; - float2 a_2 = x[1] - x[2]; + metal::vec a_1 = x[1] + x[2]; + metal::vec a_2 = x[1] - x[2]; y[0] = x[0] + a_1; - float2 b_1 = x[0] - 0.5 * a_1; - float2 b_2 = pi_2_3 * a_2; + metal::vec b_1 = x[0] - 0.5 * a_1; + metal::vec b_2 = pi_2_3 * a_2; - float2 b_2_j = {-b_2.y, b_2.x}; + metal::vec b_2_j = {-b_2.y, b_2.x}; y[1] = b_1 + b_2_j; y[2] = b_1 - b_2_j; } -METAL_FUNC void radix4(thread float2* x, thread float2* y) { - float2 z_0 = x[0] + x[2]; - float2 z_1 = x[0] - x[2]; - float2 z_2 = x[1] + x[3]; - float2 z_3 = x[1] - x[3]; - float2 z_3_i = {z_3.y, -z_3.x}; +template +METAL_FUNC void radix4(thread metal::vec* x, thread metal::vec* y) { + metal::vec z_0 = x[0] + x[2]; + metal::vec z_1 = x[0] - x[2]; + metal::vec z_2 = x[1] + x[3]; + metal::vec z_3 = x[1] - x[3]; + metal::vec z_3_i = {z_3.y, -z_3.x}; y[0] = z_0 + z_2; y[1] = z_1 + z_3_i; @@ -78,25 +93,26 @@ METAL_FUNC void radix4(thread float2* x, thread float2* y) { y[3] = z_1 - z_3_i; } -METAL_FUNC void radix5(thread float2* x, thread float2* y) { - float2 root_5_4 = 0.5590169943749475; - float2 sin_2pi_5 = 0.9510565162951535; - float2 sin_1pi_5 = 0.5877852522924731; - - float2 a_1 = x[1] + x[4]; - float2 a_2 = x[2] + x[3]; - float2 a_3 = x[1] - x[4]; - float2 a_4 = x[2] - x[3]; - - float2 a_5 = a_1 + a_2; - float2 a_6 = root_5_4 * (a_1 - a_2); - float2 a_7 = x[0] - a_5 / 4; - float2 a_8 = a_7 + a_6; - float2 a_9 = a_7 - a_6; - float2 a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; - float2 a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; - float2 a_10_j = {a_10.y, -a_10.x}; - float2 a_11_j = {a_11.y, -a_11.x}; +template +METAL_FUNC void radix5(thread metal::vec* x, thread metal::vec* y) { + T root_5_4 = 0.5590169943749475; + T sin_2pi_5 = 0.9510565162951535; + T sin_1pi_5 = 0.5877852522924731; + + metal::vec a_1 = x[1] + x[4]; + metal::vec a_2 = x[2] + x[3]; + metal::vec a_3 = x[1] - x[4]; + metal::vec a_4 = x[2] - x[3]; + + metal::vec a_5 = a_1 + a_2; + metal::vec a_6 = root_5_4 * (a_1 - a_2); + metal::vec a_7 = x[0] - a_5 / 4; + metal::vec a_8 = a_7 + a_6; + metal::vec a_9 = a_7 - a_6; + metal::vec a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; + metal::vec a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; + metal::vec a_10_j = {a_10.y, -a_10.x}; + metal::vec a_11_j = {a_11.y, -a_11.x}; y[0] = x[0] + a_5; y[1] = a_8 + a_10_j; @@ -105,23 +121,24 @@ METAL_FUNC void radix5(thread float2* x, thread float2* y) { y[4] = a_8 - a_10_j; } -METAL_FUNC void radix6(thread float2* x, thread float2* y) { - float sin_pi_3 = 0.8660254037844387; - float2 a_1 = x[2] + x[4]; - float2 a_2 = x[0] - a_1 / 2; - float2 a_3 = sin_pi_3 * (x[2] - x[4]); - float2 a_4 = x[5] + x[1]; - float2 a_5 = x[3] - a_4 / 2; - float2 a_6 = sin_pi_3 * (x[5] - x[1]); - float2 a_7 = x[0] + a_1; - - float2 a_3_i = {a_3.y, -a_3.x}; - float2 a_6_i = {a_6.y, -a_6.x}; - float2 a_8 = a_2 + a_3_i; - float2 a_9 = a_2 - a_3_i; - float2 a_10 = x[3] + a_4; - float2 a_11 = a_5 + a_6_i; - float2 a_12 = a_5 - a_6_i; +template +METAL_FUNC void radix6(thread metal::vec* x, thread metal::vec* y) { + T sin_pi_3 = 0.8660254037844387; + metal::vec a_1 = x[2] + x[4]; + metal::vec a_2 = x[0] - a_1 / 2; + metal::vec a_3 = sin_pi_3 * (x[2] - x[4]); + metal::vec a_4 = x[5] + x[1]; + metal::vec a_5 = x[3] - a_4 / 2; + metal::vec a_6 = sin_pi_3 * (x[5] - x[1]); + metal::vec a_7 = x[0] + a_1; + + metal::vec a_3_i = {a_3.y, -a_3.x}; + metal::vec a_6_i = {a_6.y, -a_6.x}; + metal::vec a_8 = a_2 + a_3_i; + metal::vec a_9 = a_2 - a_3_i; + metal::vec a_10 = x[3] + a_4; + metal::vec a_11 = a_5 + a_6_i; + metal::vec a_12 = a_5 - a_6_i; y[0] = a_7 + a_10; y[1] = a_8 - a_11; @@ -131,26 +148,28 @@ METAL_FUNC void radix6(thread float2* x, thread float2* y) { y[5] = a_9 - a_12; } -METAL_FUNC void radix7(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix7(thread metal::vec* x, thread metal::vec* y) { // Rader's algorithm - float2 inv = {1 / 6.0, -1 / 6.0}; + T inv_r = static_cast(1 / 6.0); + metal::vec inv = {inv_r, -inv_r}; // fft - float2 in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; - radix6(in1, y + 1); + metal::vec in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; + radix6(in1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(2.44013336, -1.02261879)); - y[3] = complex_mul_conj(y[3], float2(2.37046941, -1.17510629)); - y[4] = complex_mul_conj(y[4], float2(0, -2.64575131)); - y[5] = complex_mul_conj(y[5], float2(2.37046941, 1.17510629)); - y[6] = complex_mul_conj(y[6], float2(-2.44013336, -1.02261879)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(2.44013336, -1.02261879)); + y[3] = complex_mul_conj(y[3], metal::vec(2.37046941, -1.17510629)); + y[4] = complex_mul_conj(y[4], metal::vec(0, -2.64575131)); + y[5] = complex_mul_conj(y[5], metal::vec(2.37046941, 1.17510629)); + y[6] = complex_mul_conj(y[6], metal::vec(-2.44013336, -1.02261879)); // ifft - radix6(y + 1, x + 1); + radix6(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[5] = x[2] * inv + x[0]; @@ -160,79 +179,87 @@ METAL_FUNC void radix7(thread float2* x, thread float2* y) { y[3] = x[6] * inv + x[0]; } -METAL_FUNC void radix8(thread float2* x, thread float2* y) { - float cos_pi_4 = 0.7071067811865476; - float2 w_0 = {cos_pi_4, -cos_pi_4}; - float2 w_1 = {-cos_pi_4, -cos_pi_4}; - float2 temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; - radix4(temp, x); - radix4(temp + 4, x + 4); +template +METAL_FUNC void radix8(thread metal::vec* x, thread metal::vec* y) { + T cos_pi_4 = 0.7071067811865476; + metal::vec w_0 = {cos_pi_4, -cos_pi_4}; + metal::vec w_1 = {-cos_pi_4, -cos_pi_4}; + metal::vec temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; + radix4(temp, x); + radix4(temp + 4, x + 4); y[0] = x[0] + x[4]; y[4] = x[0] - x[4]; - float2 x_5 = complex_mul(x[5], w_0); + metal::vec x_5 = complex_mul(x[5], w_0); y[1] = x[1] + x_5; y[5] = x[1] - x_5; - float2 x_6 = {x[6].y, -x[6].x}; + metal::vec x_6 = {x[6].y, -x[6].x}; y[2] = x[2] + x_6; y[6] = x[2] - x_6; - float2 x_7 = complex_mul(x[7], w_1); + metal::vec x_7 = complex_mul(x[7], w_1); y[3] = x[3] + x_7; y[7] = x[3] - x_7; } -template -METAL_FUNC void radix10(thread float2* x, thread float2* y) { - float2 w[4]; - w[0] = {0.8090169943749475, -0.5877852522924731}; - w[1] = {0.30901699437494745, -0.9510565162951535}; +template +METAL_FUNC void radix10( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[4]; + w[0] = {T(0.8090169943749475), T(-0.5877852522924731)}; + w[1] = {T(0.30901699437494745), T(-0.9510565162951535)}; w[2] = {-w[1].x, w[1].y}; w[3] = {-w[0].x, w[0].y}; if (raders_perm) { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[3], x[4], x[8], x[2], x[1], x[7], x[9], x[6], x[5]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } else { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[2], x[4], x[6], x[8], x[1], x[3], x[5], x[7], x[9]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } y[0] = x[0] + x[5]; y[5] = x[0] - x[5]; for (int t = 1; t < 5; t++) { - float2 a = complex_mul(x[t + 5], w[t - 1]); + metal::vec a = complex_mul(x[t + 5], w[t - 1]); y[t] = x[t] + a; y[t + 5] = x[t] - a; } } -METAL_FUNC void radix11(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 10.0, -1 / 10.0}; +template +METAL_FUNC void radix11( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 10.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix10(x + 1, y + 1); + radix10(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(0.955301878, -3.17606649)); - y[3] = complex_mul_conj(y[3], float2(2.63610556, 2.01269656)); - y[4] = complex_mul_conj(y[4], float2(2.54127802, 2.13117479)); - y[5] = complex_mul_conj(y[5], float2(2.07016210, 2.59122150)); - y[6] = complex_mul_conj(y[6], float2(0, -3.31662479)); - y[7] = complex_mul_conj(y[7], float2(2.07016210, -2.59122150)); - y[8] = complex_mul_conj(y[8], float2(-2.54127802, 2.13117479)); - y[9] = complex_mul_conj(y[9], float2(2.63610556, -2.01269656)); - y[10] = complex_mul_conj(y[10], float2(-0.955301878, -3.17606649)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(0.955301878, -3.17606649)); + y[3] = complex_mul_conj(y[3], metal::vec(2.63610556, 2.01269656)); + y[4] = complex_mul_conj(y[4], metal::vec(2.54127802, 2.13117479)); + y[5] = complex_mul_conj(y[5], metal::vec(2.07016210, 2.59122150)); + y[6] = complex_mul_conj(y[6], metal::vec(0, -3.31662479)); + y[7] = complex_mul_conj(y[7], metal::vec(2.07016210, -2.59122150)); + y[8] = complex_mul_conj(y[8], metal::vec(-2.54127802, 2.13117479)); + y[9] = complex_mul_conj(y[9], metal::vec(2.63610556, -2.01269656)); + y[10] = + complex_mul_conj(y[10], metal::vec(-0.955301878, -3.17606649)); // ifft - radix10(y + 1, x + 1); + radix10(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[6] = x[2] * inv + x[0]; @@ -246,10 +273,12 @@ METAL_FUNC void radix11(thread float2* x, thread float2* y) { y[2] = x[10] * inv + x[0]; } -template -METAL_FUNC void radix12(thread float2* x, thread float2* y) { - float2 w[6]; - float sin_pi_3 = 0.8660254037844387; +template +METAL_FUNC void radix12( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[6]; + T sin_pi_3 = 0.8660254037844387; w[0] = {sin_pi_3, -0.5}; w[1] = {0.5, -sin_pi_3}; w[2] = {0, -1}; @@ -257,7 +286,7 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { w[4] = {-sin_pi_3, -0.5}; if (raders_perm) { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[3], x[2], @@ -270,10 +299,10 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[10], x[4], x[6]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } else { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[2], x[4], @@ -286,44 +315,50 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[7], x[9], x[11]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } y[0] = x[0] + x[6]; y[6] = x[0] - x[6]; for (int t = 1; t < 6; t++) { - float2 a = complex_mul(x[t + 6], w[t - 1]); + metal::vec a = complex_mul(x[t + 6], w[t - 1]); y[t] = x[t] + a; y[t + 6] = x[t] - a; } } -METAL_FUNC void radix13(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 12.0, -1 / 12.0}; +template +METAL_FUNC void radix13( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 12.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix12(x + 1, y + 1); + radix12(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(3.07497206, -1.88269669)); - y[3] = complex_mul_conj(y[3], float2(3.09912468, 1.84266823)); - y[4] = complex_mul_conj(y[4], float2(3.45084438, -1.04483161)); - y[5] = complex_mul_conj(y[5], float2(0.91083583, 3.48860690)); - y[6] = complex_mul_conj(y[6], float2(-3.60286363, 0.139189267)); - y[7] = complex_mul_conj(y[7], float2(3.60555128, 0)); - y[8] = complex_mul_conj(y[8], float2(3.60286363, 0.139189267)); - y[9] = complex_mul_conj(y[9], float2(0.91083583, -3.48860690)); - y[10] = complex_mul_conj(y[10], float2(-3.45084438, -1.04483161)); - y[11] = complex_mul_conj(y[11], float2(3.09912468, -1.84266823)); - y[12] = complex_mul_conj(y[12], float2(-3.07497206, -1.88269669)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(3.07497206, -1.88269669)); + y[3] = complex_mul_conj(y[3], metal::vec(3.09912468, 1.84266823)); + y[4] = complex_mul_conj(y[4], metal::vec(3.45084438, -1.04483161)); + y[5] = complex_mul_conj(y[5], metal::vec(0.91083583, 3.48860690)); + y[6] = complex_mul_conj(y[6], metal::vec(-3.60286363, 0.139189267)); + y[7] = complex_mul_conj(y[7], metal::vec(3.60555128, 0)); + y[8] = complex_mul_conj(y[8], metal::vec(3.60286363, 0.139189267)); + y[9] = complex_mul_conj(y[9], metal::vec(0.91083583, -3.48860690)); + y[10] = + complex_mul_conj(y[10], metal::vec(-3.45084438, -1.04483161)); + y[11] = complex_mul_conj(y[11], metal::vec(3.09912468, -1.84266823)); + y[12] = + complex_mul_conj(y[12], metal::vec(-3.07497206, -1.88269669)); // ifft - radix12(y + 1, x + 1); + radix12(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[7] = x[2] * inv + x[0]; @@ -339,6 +374,334 @@ METAL_FUNC void radix13(thread float2* x, thread float2* y) { y[2] = x[12] * inv + x[0]; } +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/bf16.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/bf16.h" +// Copyright © 2023 Apple Inc. + + +#include + +using namespace metal; + +typedef bfloat bfloat16_t; +inline uint16_t bfloat16_to_uint16(const bfloat16_t x) { + return as_type(x); +} + +inline bfloat16_t uint16_to_bfloat16(const uint16_t x) { + return as_type(x); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/complex.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/complex.h" +// Copyright © 2023 Apple Inc. + + +#include + + +using namespace metal; + +template +struct complex_t; + +template +static constexpr constant bool is_complex_v = false; + +template +static constexpr constant bool is_complex_v> = true; + +// Metal accepts explicit bfloat casts that is_convertible_v reports as false. +template +static constexpr constant bool is_lane_convertible_v = + is_convertible_v || + (is_same_v && is_convertible_v) || + (is_same_v && is_convertible_v); + +template +struct complex_t { + using value_type = T; + + T real; + T imag; + + // Constructors + constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {}; + constexpr complex_t() thread : real(0), imag(0) {}; + constexpr complex_t() threadgroup : real(0), imag(0) {}; + + // Conversions from scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) thread : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) threadgroup : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) device : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) constant : real(static_cast(x)), + imag(static_cast(0)) {} + + // Conversions between complex types + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) thread : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) threadgroup + : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) device : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + // Conversions to and from two-lane vectors (the FFT lane representation) + constexpr complex_t(vec v) thread : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) threadgroup : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) device : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) constant : real(v.x), imag(v.y) {}; + + constexpr operator vec() const thread { + return vec(real, imag); + } + + constexpr operator vec() const threadgroup { + return vec(real, imag); + } + + constexpr operator vec() const device { + return vec(real, imag); + } + + constexpr operator vec() const constant { + return vec(real, imag); + } + + // Conversions to scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const thread { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const threadgroup { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const device { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const constant { + return static_cast(real); + } +}; + +using complex32_t = complex_t; +using complex64_t = complex_t; + +static_assert(sizeof(complex32_t) == 2 * sizeof(half)); +static_assert(sizeof(complex64_t) == 2 * sizeof(float)); +static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); + +template +constexpr complex_t operator-(complex_t x) { + return {-x.real, -x.imag}; +} + +template +constexpr bool operator>=(complex_t a, complex_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); +} + +template +constexpr bool operator>(complex_t a, complex_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); +} + +template +constexpr bool operator<=(complex_t a, complex_t b) { + return operator>=(b, a); +} + +template +constexpr bool operator<(complex_t a, complex_t b) { + return operator>(b, a); +} + +template +constexpr bool operator==(complex_t a, complex_t b) { + return a.real == b.real && a.imag == b.imag; +} + +template +constexpr complex_t operator+(complex_t a, complex_t b) { + return {a.real + b.real, a.imag + b.imag}; +} + +template +constexpr thread complex_t& operator+=( + thread complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template +constexpr threadgroup complex_t& operator+=( + threadgroup complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template +constexpr device complex_t& operator+=( + device complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(U a, complex_t b) { + return {static_cast(a) + b.real, b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(complex_t a, U b) { + return {a.real + static_cast(b), a.imag}; +} + +template +constexpr complex_t operator-(complex_t a, complex_t b) { + return {a.real - b.real, a.imag - b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(U a, complex_t b) { + return {static_cast(a) - b.real, -b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(complex_t a, U b) { + return {a.real - static_cast(b), a.imag}; +} + +template +constexpr complex_t operator*(complex_t a, complex_t b) { + return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; +} + +template +constexpr complex_t operator/(complex_t a, complex_t b) { + auto denom = b.real * b.real + b.imag * b.imag; + auto x = a.real * b.real + a.imag * b.imag; + auto y = a.imag * b.real - a.real * b.imag; + return {x / denom, y / denom}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator/(U a, complex_t b) { + auto scalar = static_cast(a); + auto denom = b.real * b.real + b.imag * b.imag; + auto x = scalar * b.real; + auto y = -scalar * b.imag; + return {x / denom, y / denom}; +} + +template +constexpr complex_t operator%(complex_t a, complex_t b) { + auto real = a.real - (b.real * static_cast(a.real / b.real)); + auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); + if (real != 0 && (real < 0 != b.real < 0)) { + real += b.real; + } + if (imag != 0 && (imag < 0 != b.imag < 0)) { + imag += b.imag; + } + return {real, imag}; +} + +static_assert( + (complex_t{1.0h, 2.0h} * complex_t{3.0h, 4.0h}).real == -5.0h); +static_assert( + (complex_t{bfloat16_t(1.0f), bfloat16_t(2.0f)} * + complex_t{bfloat16_t(3.0f), bfloat16_t(4.0f)}) + .real == bfloat16_t(-5.0f)); + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/fft/readwrite.h" /////////////////////////////////////////////////////////////////////////////// @@ -372,14 +735,48 @@ Each with support for: using namespace metal; +// Derives the FFT scalar arithmetic type from a storage type. +template +struct FFTStorageTraits { + using scalar_T = storage_T; +}; + +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + +template +struct FFTIOTypeTraits { + using scalar_T = typename FFTStorageTraits::scalar_T; + + static_assert( + metal::is_same_v::scalar_T>, + "FFT input and output storage must share a scalar arithmetic type"); + static_assert( + sizeof(in_T) <= 16 && 16 % sizeof(in_T) == 0, + "FFT input storage must divide the 128-bit sequential access width"); + static_assert( + sizeof(out_T) <= 16 && 16 % sizeof(out_T) == 0, + "FFT output storage must divide the 128-bit sequential access width"); +}; + template < typename in_T, typename out_T, int step = 0, bool four_step_real = false> struct ReadWriter { + using scalar_T = typename FFTIOTypeTraits::scalar_T; + using complex_T = vec; + const device in_T* in; - threadgroup float2* buf; + threadgroup complex_T* buf; device out_T* out; int n; int batch_size; @@ -395,23 +792,22 @@ struct ReadWriter { METAL_FUNC ReadWriter( const device in_T* in_, - threadgroup float2* buf_, + threadgroup complex_T* buf_, device out_T* out_, const short n_, const int batch_size_, const short elems_per_thread_, const uint3 elem_, const uint3 grid_, - const bool inv_) - : in(in_), - buf(buf_), - out(out_), - n(n_), - batch_size(batch_size_), - elems_per_thread(elems_per_thread_), - elem(elem_), - grid(grid_), - inv(inv_) { + const bool inv_) thread : in(in_), + buf(buf_), + out(out_), + n(n_), + batch_size(batch_size_), + elems_per_thread(elems_per_thread_), + elem(elem_), + grid(grid_), + inv(inv_) { // Account for padding on last threadgroup threads_per_tg = elem.x == grid.x - 1 ? (batch_size - (grid.x - 1) * grid.y) * grid.z @@ -419,110 +815,139 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const { - return inv ? float2(elem.x, -elem.y) : elem; + METAL_FUNC complex_T post_in(complex_T elem) const thread { + return inv ? complex_T(elem.x, -elem.y) : elem; + } + + // Handle packed complex_t storage + METAL_FUNC complex_T post_in(complex_t elem) const thread { + return post_in(complex_T(elem.real, elem.imag)); } // Handle float case for generic RFFT alg - METAL_FUNC float2 post_in(float elem) const { - return float2(elem, 0); + METAL_FUNC complex_T post_in(scalar_T elem) const thread { + return complex_T(elem, 0); } - METAL_FUNC float2 pre_out(float2 elem) const { - return inv ? float2(elem.x / n, -elem.y / n) : elem; + METAL_FUNC complex_T pre_out(complex_T elem) const thread { + return pre_out(elem, n); } - METAL_FUNC float2 pre_out(float2 elem, int length) const { - return inv ? float2(elem.x / length, -elem.y / length) : elem; + METAL_FUNC complex_T pre_out(complex_T elem, int length) const thread { + if (!inv) { + return elem; + } + if constexpr (metal::is_same_v) { + return complex_T(elem.x / length, -elem.y / length); + } else { + // Compute the reciprocal in float before narrowing to the lane. + scalar_T inv_length = static_cast(1.0f / length); + return complex_T(elem.x * inv_length, -elem.y * inv_length); + } } - METAL_FUNC bool out_of_bounds() const { + METAL_FUNC bool out_of_bounds() const thread { // Account for possible extra threadgroups int grid_index = elem.x * grid.y + elem.y; return grid_index >= batch_size; } - METAL_FUNC void load() const { + METAL_FUNC void load(scalar_T scale) const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - // 2 complex64s = 128 bits - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + // Keep each thread's sequential access at 128 bits where possible. + constexpr int read_width = 16 / sizeof(in_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - buf[index] = post_in(in[batch_idx + index]); - buf[index + 1] = post_in(in[batch_idx + index + 1]); + for (short r = 0; r < read_width; r++) { + buf[index + r] = post_in(in[batch_idx + index + r]) * scale; + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); - buf[index] = post_in(in[batch_idx + index]); + buf[index] = post_in(in[batch_idx + index]) * scale; } } - METAL_FUNC void write() const { + METAL_FUNC void load() const thread { + load(static_cast(1)); + } + + METAL_FUNC void write() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + constexpr int read_width = 16 / sizeof(out_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - out[batch_idx + index] = pre_out(buf[index]); - out[batch_idx + index + 1] = pre_out(buf[index + 1]); + for (short r = 0; r < read_width; r++) { + out[batch_idx + index + r] = pre_out(buf[index + r]); + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); out[batch_idx + index] = pre_out(buf[index]); } } // Padded IO for Bluestein's algorithm - METAL_FUNC void load_padded(int length, const device float2* w_k) const { + METAL_FUNC void load_padded(int length, const device complex_T* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = post_in(in[batch_idx + index]); - seq_buf[index] = complex_mul(elem, w_k[index]); + complex_T elem = post_in(in[batch_idx + index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0.0; } } } - METAL_FUNC void write_padded(int length, const device float2* w_k) const { + METAL_FUNC void write_padded(int length, const device complex_T* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - float2 inv_factor = {1.0f / n, -1.0f / n}; + scalar_T inv_n = static_cast(1.0f / n); + complex_T inv_factor = {inv_n, -inv_n}; + if constexpr (!metal::is_same_v) { + // Reduced lanes applied the 1/n scale with the convolution, so only + // the conjugation remains here. + inv_factor = {scalar_T(1), -scalar_T(1)}; + } - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = seq_buf[index + length - 1] * inv_factor; - out[batch_idx + index] = pre_out(complex_mul(elem, w_k[index]), length); + complex_T elem = seq_buf[index + length - 1] * inv_factor; + out[batch_idx + index] = + pre_out(complex_mul(elem, w_k[index]), length); } } } // Strided IO for four step FFT - METAL_FUNC void compute_strided_indices(int stride, int overall_n) { + METAL_FUNC void compute_strided_indices(int stride, int overall_n) thread { // Use the batch threadgroup dimension to coalesce memory accesses: // e.g. stride = 12 // device | shared mem @@ -543,53 +968,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step - METAL_FUNC void load_strided(int stride, int overall_n) { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - buf[strided_shared_idx + e] = - post_in(in[strided_device_idx + e * stride]); + // Four-step FFT I/O + METAL_FUNC void load_strided(int stride, int overall_n) thread { + if constexpr (step == 1 && !four_step_real) { + // Do not invert between C2C four-step passes. + (void)stride; + (void)overall_n; + bool default_inv = inv; + inv = false; + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize inverse transforms on input so + // intermediates stay in range; see write_strided. + if (default_inv) { + load(static_cast(1.0f / n)); + } else { + load(); + } + } else { + load(); + } + inv = default_inv; + } else { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + complex_T input = post_in(in[strided_device_idx + e * stride]); + if constexpr (!metal::is_same_v) { + if (step == 0 && !four_step_real && inv) { + input *= static_cast(1.0f / n); + } + } + buf[strided_shared_idx + e] = input; + } } } - METAL_FUNC void write_strided(int stride, int overall_n) { - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - int combined_idx = (strided_device_idx + e * stride) % overall_n; - int ij = (combined_idx / stride) * (combined_idx % stride); - // Apply four step twiddles at end of first step - float2 twiddle = get_twiddle(ij, overall_n); - out[strided_device_idx + e * stride] = complex_mul(output, twiddle); + METAL_FUNC void write_strided(int stride, int overall_n) thread { + if constexpr (step == 1 && !four_step_real) { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + if constexpr (!metal::is_same_v) { + // Inverse C2C passes were normalized on input, so only the + // conjugation of the inverse transform remains here. + complex_T value = buf[strided_shared_idx + e]; + out[strided_device_idx + e * stride] = + inv ? complex_T(value.x, -value.y) : value; + } else { + out[strided_device_idx + e * stride] = + pre_out(buf[strided_shared_idx + e], overall_n); + } + } + } else { + for (int e = 0; e < elems_per_thread; e++) { + complex_T output = buf[strided_shared_idx + e]; + int combined_idx = (strided_device_idx + e * stride) % overall_n; + int ij = (combined_idx / stride) * (combined_idx % stride); + // Apply four step twiddles at end of first step + complex_T twiddle = get_twiddle(ij, overall_n); + out[strided_device_idx + e * stride] = + complex_mul(output, twiddle); + } } } }; -// Four Step FFT Second Step -template <> -METAL_FUNC void ReadWriter::load_strided( - int stride, - int overall_n) { - // Silence compiler warnings - (void)stride; - (void)overall_n; - // Don't invert between steps - bool default_inv = inv; - inv = false; - load(); - inv = default_inv; -} - -template <> -METAL_FUNC void ReadWriter::write_strided( - int stride, - int overall_n) { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - out[strided_device_idx + e * stride] = pre_out(output, overall_n); - } -} - +// Packed RFFT/IRFFT remains float-specific. This generic foundation covers +// C2C storage; reduced-precision real transforms need separate host dispatch. +// // For RFFT, we interleave batches of two real sequences into one complex one: // // z_k = x_k + j.y_k @@ -598,14 +1042,14 @@ METAL_FUNC void ReadWriter::write_strided( // // This roughly doubles the throughput over the regular FFT. template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for RFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -625,7 +1069,7 @@ METAL_FUNC void ReadWriter::load() const { } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = @@ -654,7 +1098,7 @@ METAL_FUNC void ReadWriter::write() const { float2 x_n_minus_k = seq_buf[n - index] * conj; out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -662,7 +1106,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -679,7 +1123,7 @@ METAL_FUNC void ReadWriter::load_padded( if (index < length) { float2 elem = float2(in[batch_idx + index], in[batch_idx + index + next_in]); - seq_buf[index] = complex_mul(elem, w_k[index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0; } @@ -689,7 +1133,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int length_over_2 = (length / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * length_over_2 * 2 + elem.y * length_over_2 * 2; @@ -712,18 +1156,18 @@ METAL_FUNC void ReadWriter::write_padded( // x_0 = z_0.real // y_0 = z_0.imag if (index == 0) { - float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); out[batch_idx + index] = float2(elem.x, 0); out[batch_idx + index + next_out] = float2(elem.y, 0); } else { - float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); - float2 x_n_minus_k = complex_mul( + float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 x_n_minus_k = complex_mul( w_k[length - index], seq_buf[length - index] * inv_factor); x_n_minus_k *= conj; // w_k should happen before this extraction out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -734,14 +1178,14 @@ METAL_FUNC void ReadWriter::write_padded( // x_k = Re(Z_k) // Y_k = Imag(Z_k) template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for IRFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * n_over_2 * 2 + elem.y * n_over_2 * 2; @@ -770,17 +1214,17 @@ METAL_FUNC void ReadWriter::load() const { x = float2(x.x, 0); y = float2(y.x, 0); } - seq_buf[index] = x + complex_mul(y, plus_j); + seq_buf[index] = x + complex_mul(y, plus_j); seq_buf[index].y = -seq_buf[index].y; if (index > 0 && !last_val) { - seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); + seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[n - index].y = -seq_buf[n - index].y; } } } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { int batch_idx = elem.x * grid.y * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -801,7 +1245,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int n_over_2 = (n / 2) + 1; int length_over_2 = (length / 2) + 1; @@ -831,12 +1275,12 @@ METAL_FUNC void ReadWriter::load_padded( x = float2(x.x, 0); y = float2(y.x, 0); } - float2 elem1 = x + complex_mul(y, plus_j); - seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); + float2 elem1 = x + complex_mul(y, plus_j); + seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); if (index > 0 && !last_val) { - float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); + float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[length - index] = - complex_mul(elem2 * conj, w_k[length - index]); + complex_mul(elem2 * conj, w_k[length - index]); } } else { short pad_index = metal::min(length + (index - length_over_2) * 2, n - 2); @@ -849,7 +1293,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n + length - 1; @@ -864,7 +1308,8 @@ METAL_FUNC void ReadWriter::write_padded( for (int e = 0; e < elems_per_thread; e++) { int index = fft_idx + e * m; if (index < length) { - float2 output = complex_mul(seq_buf[index] * inv_factor, w_k[index]); + float2 output = + complex_mul(seq_buf[index] * inv_factor, w_k[index]); out[batch_idx + index] = output.x / length; out[batch_idx + index + next_out] = output.y / -length; } @@ -876,7 +1321,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -891,7 +1336,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; int coalesce_width = grid.y; int tg_idx = elem.y * grid.z + elem.z; @@ -920,7 +1365,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; auto conj = float2(1, -1); @@ -945,7 +1390,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -959,7 +1404,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { @@ -1000,7 +1445,7 @@ using namespace metal; #define MAX_RADIX 13 // Reached when elems_per_thread_ = 6, max_radix = 13 -// and some threads have to do 3 radix 6s requiring 18 float2s. +// and some threads have to do 3 radix 6s requiring 18 complex values. #define MAX_OUTPUT_SIZE 18 // Specialize for a particular value of N at runtime @@ -1029,18 +1474,47 @@ STEEL_CONST int rader_5_steps_ [[function_constant(18)]]; STEEL_CONST int rader_4_steps_ [[function_constant(19)]]; STEEL_CONST int rader_3_steps_ [[function_constant(20)]]; STEEL_CONST int rader_2_steps_ [[function_constant(21)]]; +STEEL_CONST bool use_bluestein_twiddle_table_ [[function_constant(22)]]; + +// Generate the compact radix-8 Stockham twiddle table used by the fused +// 4096-point Bluestein convolution. The first p=1 stage needs no twiddles; +// the remaining stages contribute 8, 64, and 512 entries. +[[kernel]] void generate_bluestein_twiddles( + device float2* twiddles [[buffer(0)]], + uint index [[thread_position_in_grid]]) { + uint p; + uint k; + if (index < 8) { + p = 8; + k = index; + } else if (index < 72) { + p = 64; + k = index - 8; + } else { + p = 512; + k = index - 72; + } + twiddles[index] = get_twiddle(k, 8 * p); +} -// See "radix.h" for radix codelets -typedef void (*RadixFunc)(thread float2*, thread float2*); +// See "radix.h" for radix codelets. +template +using RadixFunc = void (*)(thread vec*, thread vec*); // Perform a single radix n butterfly with appropriate twiddles -template +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_butterfly( int i, int p, - thread float2* x, + int twiddle_offset, + thread vec* x, thread short* indices, - thread float2* y) { + thread vec* y, + const device vec* twiddles) { // i: the index in the overall DFT that we're processing. // p: the size of the DFTs we're merging at this step. // m: how many threads are working on this DFT. @@ -1059,14 +1533,19 @@ METAL_FUNC void radix_butterfly( // Apply twiddles if (p > 1) { - float2 twiddle_1 = get_twiddle(k, radix * p); - float2 twiddle = twiddle_1; - x[1] = complex_mul(x[1], twiddle); + vec twiddle_1; + if constexpr (use_twiddle_table) { + twiddle_1 = twiddles[twiddle_offset + k]; + } else { + twiddle_1 = get_twiddle(k, radix * p); + } + vec twiddle = twiddle_1; + x[1] = complex_mul(x[1], twiddle); STEEL_PRAGMA_UNROLL for (int t = 2; t < radix; t++) { - twiddle = complex_mul(twiddle, twiddle_1); - x[t] = complex_mul(x[t], twiddle); + twiddle = complex_mul(twiddle, twiddle_1); + x[t] = complex_mul(x[t], twiddle); } } @@ -1080,17 +1559,23 @@ METAL_FUNC void radix_butterfly( // Perform all the radix steps required for a // particular radix size n. -template +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_n_steps( int i, thread int* p, + thread int* twiddle_offset, int m, int n, int num_steps, - thread float2* inputs, + thread vec* inputs, thread short* indices, - thread float2* values, - threadgroup float2* buf) { + thread vec* values, + threadgroup vec* buf, + const device vec* twiddles) { int m_r = n / radix; // When combining different sized radices, we have to do // multiple butterflies in a single thread. @@ -1110,8 +1595,14 @@ METAL_FUNC void radix_n_steps( for (int r = 0; r < radix; r++) { inputs[r] = buf[index + r * m_r]; } - radix_butterfly( - index, *p, inputs, indices + t * radix, values + t * radix); + radix_butterfly( + index, + *p, + *twiddle_offset, + inputs, + indices + t * radix, + values + t * radix, + twiddles); } } @@ -1130,20 +1621,41 @@ METAL_FUNC void radix_n_steps( // Wait until all threads have written back to threadgroup mem threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (use_twiddle_table) { + if (*p > 1) { + *twiddle_offset += *p; + } + } *p *= radix; } } -#define RADIX_STEP(radix, radix_func, num_steps) \ - radix_n_steps( \ - fft_idx, p, m, n, num_steps, inputs, indices, values, buf); - -template -METAL_FUNC void -perform_fft(int fft_idx, thread int* p, int m, int n, threadgroup float2* buf) { - float2 inputs[MAX_RADIX]; +#define RADIX_STEP(radix, radix_func, num_steps) \ + radix_n_steps>( \ + fft_idx, \ + p, \ + &twiddle_offset, \ + m, \ + n, \ + num_steps, \ + inputs, \ + indices, \ + values, \ + buf, \ + twiddles); + +template +METAL_FUNC void perform_fft( + int fft_idx, + thread int* p, + int m, + int n, + threadgroup vec* buf, + const device vec* twiddles = nullptr) { + vec inputs[MAX_RADIX]; short indices[MAX_OUTPUT_SIZE]; - float2 values[MAX_OUTPUT_SIZE]; + vec values[MAX_OUTPUT_SIZE]; + int twiddle_offset = 0; RADIX_STEP(2, radix2, rader ? rader_2_steps_ : radix_2_steps_); RADIX_STEP(3, radix3, rader ? rader_3_steps_ : radix_3_steps_); @@ -1168,7 +1680,8 @@ template constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { - threadgroup float2 shared_in[tg_mem_size]; + using scalar_T = typename FFTIOTypeTraits::scalar_T; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -1192,9 +1705,9 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -1203,7 +1716,8 @@ template [[kernel]] void rader_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* raders_b_q [[buffer(2)]], + const device vec::scalar_T, 2>* + raders_b_q [[buffer(2)]], const device short* raders_g_q [[buffer(3)]], const device short* raders_g_minus_q [[buffer(4)]], constant const int& n, @@ -1211,9 +1725,10 @@ template constant const int& rader_n, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Use Rader's algorithm to compute fast FFTs // when a prime factor `p` of `n` is greater than 13 but - // has `p - 1` Stockham decomposable into to prime factors <= 13. + // has `p - 1` Stockham decomposable into prime factors <= 13. // // E.g. n = 102 // = 2 * 3 * 17 @@ -1234,7 +1749,7 @@ template // // Rader's uses fewer operations than Bluestein's and so // is more accurate. It's also faster in most cases. - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -1259,7 +1774,7 @@ template int fft_idx = elem.z; int tg_idx = elem.y * n; - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; // rader_m = n / rader_n; int rader_m = rader_m_; @@ -1271,10 +1786,10 @@ template // 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 short x_0_index = metal::min(fft_idx * elems_per_thread_ / (rader_n - 1), rader_m - 1); - float2 x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; + vec x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; // Do the Rader permutation in shared memory - float2 temp[MAX_RADIX]; + vec temp[MAX_RADIX]; int max_index = n - rader_m - 1; for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); @@ -1293,19 +1808,31 @@ template // Rader FFT on x[rader_m:] int p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); // x_1 + ... + x_n is computed for us in the first FFT step so // we save it in the first rader_m indices of the array for later. int x_sum_index = metal::min(fft_idx, rader_m - 1); buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)]; - float2 inv = {1.0f, -1.0f}; + // The convolution is conjugated and scaled by 1/(rader_n - 1) exactly once + // each. Float lanes apply the scale after the inverse FFT, preserving the + // existing float code generation; reduced lanes apply it before, so the + // intermediates cannot overflow the narrow lane range. + scalar_T rader_inv_r = static_cast(1.0f / (rader_n - 1)); + vec convolution_factor = {1.0f, -1.0f}; + vec convolution_inv_factor = {rader_inv_r, -rader_inv_r}; + if constexpr (!metal::is_same_v) { + convolution_factor = convolution_inv_factor; + convolution_inv_factor = {1.0f, -1.0f}; + } + for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); short interleaved_index = index / rader_m + (index % rader_m) * (rader_n - 1); - temp[e] = complex_mul( + temp[e] = complex_mul( buf[rader_m + interleaved_index], raders_b_q[interleaved_index % (rader_n - 1)]); } @@ -1314,25 +1841,24 @@ template for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); - buf[rader_m + index] = temp[e] * inv; + buf[rader_m + index] = temp[e] * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); // Rader IFFT on x[rader_m:] p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); - - float2 rader_inv_factor = {1.0f / (rader_n - 1), -1.0f / (rader_n - 1)}; + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, n - rader_m - 1); short diff_index = index / (rader_n - 1) - x_0_index; - temp[e] = buf[rader_m + index] * rader_inv_factor + x_0[diff_index]; + temp[e] = buf[rader_m + index] * convolution_inv_factor + x_0[diff_index]; } // Use the sum of elements that was computed in the first FFT - float2 x_sum = buf[x_0_index] + x_0[0]; + vec x_sum = buf[x_0_index] + x_0[0]; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -1349,7 +1875,7 @@ template threadgroup_barrier(mem_flags::mem_threadgroup); p = rader_n; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -1358,13 +1884,16 @@ template [[kernel]] void bluestein_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* w_q [[buffer(2)]], - const device float2* w_k [[buffer(3)]], + const device vec::scalar_T, 2>* w_q + [[buffer(2)]], + const device vec::scalar_T, 2>* w_k + [[buffer(3)]], constant const int& length, constant const int& n, constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Computes arbitrary length FFTs with Bluestein's algorithm // // In numpy: @@ -1374,7 +1903,7 @@ template // Where w_k and w_q are precomputed on CPU in high precision as: // w_k = np.exp(-1j * np.pi / n * (np.arange(-n + 1, n) ** 2)) // w_q = np.fft.fft(1/w_k[-n:]) - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -1398,22 +1927,50 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; + const device vec* twiddles = w_q + n; // fft - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } - float2 inv = float2(1.0f, -1.0f); + vec convolution_factor = {1.0f, -1.0f}; + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize the convolution before the inverse FFT so the + // intermediates cannot overflow; the write path then only conjugates. + scalar_T inv_n = static_cast(1.0f / n); + convolution_factor = {inv_n, -inv_n}; + } for (int t = 0; t < elems_per_thread_; t++) { int index = fft_idx + t * m; - buf[index] = complex_mul(buf[index], w_q[index]) * inv; + buf[index] = + complex_mul(buf[index], w_q[index]) * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); // ifft p = 1; - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } read_writer.write_padded(length, w_k); } @@ -1432,6 +1989,7 @@ template < constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Fast four step FFT implementation for powers of 2. int overall_n = n1 * n2; int n = step == 0 ? n1 : n2; @@ -1441,8 +1999,8 @@ template < int m = grid.z; int fft_idx = elem.z; - threadgroup float2 shared_in[tg_mem_size]; - threadgroup float2* buf = &shared_in[elem.y * n]; + threadgroup vec shared_in[tg_mem_size]; + threadgroup vec* buf = &shared_in[elem.y * n]; using read_writer_t = ReadWriter; read_writer_t read_writer = read_writer_t( @@ -1464,7 +2022,7 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); int p = 1; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write_strided(stride, overall_n); } diff --git a/Source/Cmlx/mlx-generated/fp_quantized.cpp b/Source/Cmlx/mlx-generated/fp_quantized.cpp index eb7389bbd..1c56cdb56 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized.cpp @@ -12,8 +12,11 @@ const char* fp_quantized() { #line 1 "mlx/backend/metal/kernels/fp4.h" +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { - fp4_e2m1(float x) { + fp4_e2m1(float x) thread { if (metal::isnan(x)) { bits = 0x7; return; @@ -42,17 +45,17 @@ struct fp4_e2m1 { bits |= sign_bit; } - operator float16_t() { + operator float16_t() thread { half converted = as_type(ushort((bits & 7) << 9)); converted *= 16384.0; return bits & 8 ? -converted : converted; } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } @@ -67,7 +70,7 @@ struct fp4_e2m1 { struct fp8_e4m3 { template - fp8_e4m3(T f) { + fp8_e4m3(T f) thread { // From PyTorch // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148 uint32_t fp8_max = 543 << 21; @@ -94,7 +97,7 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } - operator float16_t() { + operator float16_t() thread { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; @@ -102,11 +105,11 @@ struct fp8_e4m3 { return (sign ? -converted : converted); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } @@ -114,7 +117,7 @@ struct fp8_e4m3 { }; struct fp8_e8m0 { - fp8_e8m0(float x) { + fp8_e8m0(float x) thread { if (!metal::isfinite(x)) { bits = 0xFF; return; @@ -131,12 +134,12 @@ struct fp8_e8m0 { bits = static_cast(n + 127); } - operator bfloat16_t() { + operator bfloat16_t() thread { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } - operator float() { + operator float() thread { uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); return as_type(out); } @@ -144,12 +147,23 @@ struct fp8_e8m0 { uint8_t bits; }; +// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down +// leaves the block's largest elements outside the element range, where they +// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0. +inline float mx_scale_round_up(float x) { + fp8_e8m0 s(x); + if (s.bits < 0xFE && float(s) < x) { + s.bits += 1; + } + return float(s); +} + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/fp_quantized.h" /////////////////////////////////////////////////////////////////////////////// #line 1 "mlx/backend/metal/kernels/fp_quantized.h" -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -188,7 +202,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -199,7 +213,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -286,11 +300,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } @@ -335,15 +350,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -352,7 +367,7 @@ struct QuantizedBlockLoader { scales_ + bi * src_ld / group_size + (bj * pack_factor) / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -364,7 +379,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -390,7 +405,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -470,10 +485,16 @@ METAL_FUNC void fp_qmv_quad_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -483,7 +504,6 @@ METAL_FUNC void fp_qmv_fast_impl( uint simd_lid [[thread_index_in_simdgroup]]) { constexpr int packs_per_thread = 2; constexpr int num_simdgroups = 2; - constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; @@ -523,18 +543,28 @@ METAL_FUNC void fp_qmv_fast_impl( x += block_size; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } -template +template METAL_FUNC void fp_qmv_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -570,6 +600,11 @@ METAL_FUNC void fp_qmv_impl( return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // In this case we need to properly guard all our reads because there isn't // even 1 tile in the matrix if (out_vec_size < (num_simdgroups * results_per_simdgroup)) { @@ -589,7 +624,7 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - uint8_t s = sl[0]; + U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); } @@ -620,7 +655,11 @@ METAL_FUNC void fp_qmv_impl( row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } @@ -668,20 +707,145 @@ METAL_FUNC void fp_qmv_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } } -template +// Quantized matrix-vector for a small batch of input vectors, the M in +// [2, vector_limit) band between qmv (M==1) and qmm. Each thread owns one +// output row that k_lanes lanes reduce over K; the vecs_per_tg vectors are +// streamed so each weight group is dequantized once and reused across them. +template +METAL_FUNC void fp_qmv_wide_impl( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int pack_factor = get_pack_factor<32, bits>(); + constexpr int bytes_per_pack = get_bytes_per_pack<32>(); + constexpr int nf4 = group_size / 4; // float4 lanes per quant group + + typedef float U; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; // first input vector handled here + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor; + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device uint8_t* srow = scales + row * in_vec_size_g; + + // One device pointer per streamed vector; the clamp keeps an out-of-range + // tail slot reading a valid row (it is never written below). + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's quant groups: one group is + // dequantized to float4 lanes and dot()'d against each streamed vector. One + // group per iteration keeps weight-register pressure low for occupancy. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + const int k0 = g * group_size; + U s = dequantize_scale(srow[g]); + const device uint8_t* wg = wrow + k0 * bytes_per_pack / pack_factor; + + float4 w4[nf4]; + if constexpr (bits == 4) { + const device uint16_t* wq = (const device uint16_t*)wg; +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<4>{}(wq[i]), + Dequantize<4>{}(wq[i] >> 4), + Dequantize<4>{}(wq[i] >> 8), + Dequantize<4>{}(wq[i] >> 12)); + } + } else { +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<8>{}(wg[4 * i]), + Dequantize<8>{}(wg[4 * i + 1]), + Dequantize<8>{}(wg[4 * i + 2]), + Dequantize<8>{}(wg[4 * i + 3])); + } + } + +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device vec* xv4 = (const device vec*)(xv[v] + k0); + float acc = 0; +#pragma unroll + for (int j = 0; j < nf4; j++) { + acc += dot(w4[j], float4(xv4[j])); + } + result[v] += s * acc; + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); + } + } + } +} + +template METAL_FUNC void fp_qvm_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const int in_vec_size, const int out_vec_size, + const int in_vec_stride, uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { @@ -712,13 +876,18 @@ METAL_FUNC void fp_qvm_impl( int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid); ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w; scales += out_col / group_size + simd_lid * out_vec_size_g; - x += tid.x * in_vec_size + simd_lid; + x += tid.x * in_vec_stride + simd_lid; y += tid.x * out_vec_size + out_col; if (out_col >= out_vec_size) { return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // Loop over in_vec in blocks of block_size int remaining = in_vec_size % block_size; if (remaining == 0) { @@ -768,7 +937,11 @@ METAL_FUNC void fp_qvm_impl( if (simd_lid == 0) { #pragma clang loop unroll(full) for (int k = 0; k < tn * pack_factor; k++) { - y[k] = static_cast(result[k]); + if constexpr (has_global_scale) { + y[k] = static_cast(result[k] * inv_scale_enc); + } else { + y[k] = static_cast(result[k]); + } } } } @@ -791,6 +964,7 @@ METAL_FUNC void fp_qmm_t_impl( const constant int& K, const constant int& N, const constant int& M, + const constant int& K_eff, uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], @@ -844,7 +1018,7 @@ METAL_FUNC void fp_qmm_t_impl( if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -854,7 +1028,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -866,7 +1040,7 @@ METAL_FUNC void fp_qmm_t_impl( } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -876,7 +1050,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -900,11 +1074,11 @@ METAL_FUNC void fp_qmm_t_impl( template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + int BM = 32, + int BK = 32, + int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, @@ -1156,10 +1330,17 @@ template w, scales, x, y, in_vec_size, out_vec_size, tid, quad_gid, quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1191,14 +1372,35 @@ template s_strides, tid); } - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl< + T, + group_size, + bits, + has_global_scale, + results_per_simdgroup>( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1230,14 +1432,74 @@ template s_strides, tid); } - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void fp_qmv_wide( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + tid); + } + fp_qmv_wide_impl( + w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); +} + +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void fp_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1269,8 +1531,18 @@ template s_strides, tid); } - fp_qvm_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qvm_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + in_vec_size, + tid, + simd_gid, + simd_lid); } template @@ -1312,8 +1584,21 @@ template int in_vec_size_adj = tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size; - fp_qvm_impl( - w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid); + // The in_vec_stride is the full K dimension, not the partition size + int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; + + fp_qvm_impl( + w, + scales, + nullptr, + x, + y, + in_vec_size_adj, + out_vec_size, + in_vec_stride, + tid, + simd_gid, + simd_lid); } template < @@ -1368,17 +1653,18 @@ template < tid); } fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } template < typename T, - const int group_size, - const int bits, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1427,10 +1713,11 @@ template < w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void fp_gather_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1472,14 +1759,24 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1521,14 +1818,24 @@ template w_strides, s_strides, tid); - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1570,8 +1877,18 @@ template w_strides, s_strides, tid); - fp_qvm_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qvm_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + in_vec_size, + tid, + simd_gid, + simd_lid); } template < @@ -1635,16 +1952,71 @@ template < s_strides, tid); fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } template < typename T, const int group_size, const int bits, + const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32> +[[kernel]] void fp_qmm_t_splitk( + const device uint32_t* w [[buffer(0)]], + const device uint8_t* scales [[buffer(1)]], + const device T* x [[buffer(2)]], + device T* y [[buffer(3)]], + const constant int& K [[buffer(4)]], + const constant int& N [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& k_partition_size [[buffer(7)]], + const constant int& split_k_partition_stride [[buffer(8)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + (void)lid; + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + constexpr int pack_factor = get_pack_factor<8, bits>(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + const int k_start = tid.z * k_partition_size; + x += k_start; + + auto wl = (const device uint8_t*)w; + wl += k_start * bytes_per_pack / pack_factor; + scales += k_start / group_size; + y += tid.z * static_cast(split_k_partition_stride); + + fp_qmm_t_impl( + (const device uint32_t*)wl, + scales, + x, + y, + Xs, + Ws, + K, + N, + M, + k_partition_size, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1895,38 +2267,50 @@ template < } } -template +template [[kernel]] void fp_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], device uint8_t* scales [[buffer(2)]], + const device float* global_scale [[buffer(3)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + if constexpr (use_mx_scale) { + scale_dec_b = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); + auto s = ScaleType(scale_dec_b); uint8_t q_scale = s.bits; - scale = float(s); - size_t gindex = index / group_size; if (index % group_size == 0) { scales[gindex] = q_scale; } - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = float(s); + scale_enc_b = (scale_enc_b == 0) ? 0.0f : (scale_enc / scale_enc_b); + uint8_t output = Quantize{}(w_thread * scale_enc_b); if (bits == 4) { uint8_t sval = simd_shuffle_down(output, 1); output |= sval << bits; @@ -1937,10 +2321,11 @@ template } } -template +template [[kernel]] void fp_dequantize( const device uint8_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], + const device float* global_scale [[buffer(2)]], device T* out [[buffer(3)]], uint2 index [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]]) { @@ -1952,9 +2337,17 @@ template out += oindex; + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + using ScaleType = metal::conditional_t; auto q_scale = ((device ScaleType*)(scales))[gindex]; auto scale = float(q_scale); + if constexpr (has_global_scale) { + scale *= inv_scale_enc; + } uint val = w[offset]; #pragma clang loop unroll(full) @@ -1969,33 +2362,45 @@ template } } -template +template [[kernel]] void fp_quantize_dequantize( const device T* w [[buffer(0)]], - device T* out [[buffer(1)]], + const device float* global_scale [[buffer(1)]], + device T* out [[buffer(2)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + if constexpr (use_mx_scale) { + scale_dec_b = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); - scale = float(s); + auto scale = float(ScaleType(scale_dec_b)); - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = (scale == 0) ? 0.0f : (scale_enc / scale); + uint8_t output = Quantize{}(w_thread * scale_enc_b); - out[index] = static_cast(scale * Dequantize{}(output)); + out[index] = static_cast((scale / scale_enc) * Dequantize{}(output)); } /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp index 1e1cc5d7e..6b8473d85 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp @@ -12,8 +12,11 @@ const char* fp_quantized_nax() { #line 1 "mlx/backend/metal/kernels/fp4.h" +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { - fp4_e2m1(float x) { + fp4_e2m1(float x) thread { if (metal::isnan(x)) { bits = 0x7; return; @@ -42,17 +45,17 @@ struct fp4_e2m1 { bits |= sign_bit; } - operator float16_t() { + operator float16_t() thread { half converted = as_type(ushort((bits & 7) << 9)); converted *= 16384.0; return bits & 8 ? -converted : converted; } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } @@ -67,7 +70,7 @@ struct fp4_e2m1 { struct fp8_e4m3 { template - fp8_e4m3(T f) { + fp8_e4m3(T f) thread { // From PyTorch // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148 uint32_t fp8_max = 543 << 21; @@ -94,7 +97,7 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } - operator float16_t() { + operator float16_t() thread { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; @@ -102,11 +105,11 @@ struct fp8_e4m3 { return (sign ? -converted : converted); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } @@ -114,7 +117,7 @@ struct fp8_e4m3 { }; struct fp8_e8m0 { - fp8_e8m0(float x) { + fp8_e8m0(float x) thread { if (!metal::isfinite(x)) { bits = 0xFF; return; @@ -131,12 +134,12 @@ struct fp8_e8m0 { bits = static_cast(n + 127); } - operator bfloat16_t() { + operator bfloat16_t() thread { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } - operator float() { + operator float() thread { uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); return as_type(out); } @@ -144,6 +147,17 @@ struct fp8_e8m0 { uint8_t bits; }; +// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down +// leaves the block's largest elements outside the element range, where they +// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0. +inline float mx_scale_round_up(float x) { + fp8_e8m0 s(x); + if (s.bits < 0xFE && float(s) < x) { + s.bits += 1; + } + return float(s); +} + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/fp_quantized_nax.h" /////////////////////////////////////////////////////////////////////////////// @@ -188,7 +202,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -199,7 +213,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -210,11 +224,12 @@ struct Dequantize { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } @@ -261,14 +276,14 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -276,7 +291,7 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -292,7 +307,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -322,12 +337,12 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { scales += n_groups; } else { - scales += n_groups * group_stride; + scales += group_stride; } } }; @@ -395,16 +410,13 @@ METAL_FUNC void fp_qmm_t_impl( // Make the weight loader loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -412,22 +424,17 @@ METAL_FUNC void fp_qmm_t_impl( constexpr bool transpose_a = false; constexpr bool transpose_b = true; - const short sgp_sm = min(SM, short(M - (y_row + tm))); + const short sgp_sm = min(int(SM), M - (y_row + tm)); const bool is_unaligned_sm = (sgp_sm != SM); - const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn))); + const short sgp_sn = aligned_N ? SN : min(int(SN), N - (y_col + tn)); const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -446,8 +453,8 @@ METAL_FUNC void fp_qmm_t_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -549,16 +556,13 @@ METAL_FUNC void fp_qmm_n_impl( // const short num_outs = min(BN, N - y_col); loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -570,12 +574,7 @@ METAL_FUNC void fp_qmm_n_impl( using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -587,8 +586,8 @@ METAL_FUNC void fp_qmm_n_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -1018,24 +1017,19 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); - const short sgp_sn = - align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); + const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); @@ -1045,10 +1039,6 @@ template < using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - // Do as many matmuls as necessary uint32_t index; short offset; @@ -1070,8 +1060,11 @@ template < threadgroup_barrier(mem_flags::mem_none); // Prepare threadgroup mma operation - NAXTile Dtile; + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); const device T* xn = x + tm * K; @@ -1100,33 +1093,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -1140,38 +1135,37 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { diff --git a/Source/Cmlx/mlx-generated/gemm.cpp b/Source/Cmlx/mlx-generated/gemm.cpp index fa8491cd6..ee4e24cc4 100644 --- a/Source/Cmlx/mlx-generated/gemm.cpp +++ b/Source/Cmlx/mlx-generated/gemm.cpp @@ -72,18 +72,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -94,7 +94,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -103,7 +103,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -151,7 +151,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; @@ -235,7 +235,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -251,14 +251,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast( x * static_cast(alpha) + (static_cast(beta) * c)); } @@ -367,13 +367,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -403,7 +399,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -572,7 +569,7 @@ struct BaseMMAFrag { for (short j = 0; j < kElemCols; j++) { if ((off_x + i) < lim_x && (off_y + j) < lim_y) { dst[i * kElemCols + j] = - static_cast(src[(off_x + i) * str_x + (off_x + j) * str_y]); + static_cast(src[(off_x + i) * str_x + (off_y + j) * str_y]); } else { dst[i * kElemCols + j] = T(0); } @@ -717,24 +714,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -743,16 +741,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -769,7 +767,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -786,7 +784,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -801,7 +799,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -816,8 +814,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -836,8 +836,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -860,7 +862,7 @@ struct MMATile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -970,7 +972,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -988,7 +990,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -1015,7 +1017,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1028,8 +1030,11 @@ struct BlockMMA { Ctile.template store(D, ldd); } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1049,7 +1054,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1068,7 +1073,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1082,7 +1088,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -1111,7 +1117,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -1156,7 +1162,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1189,7 +1195,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1302,7 +1308,7 @@ struct BlockMMA< /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -1322,7 +1328,7 @@ struct BlockMMA< /* Karatsuba MMA: 3 real MMAs per K-chunk */ METAL_FUNC void mma( const threadgroup complex64_t* As, - const threadgroup complex64_t* Bs) { + const threadgroup complex64_t* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -1380,7 +1386,7 @@ struct BlockMMA< } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Adjust for simdgroup and thread location D += sm * ldd + sn; @@ -1399,8 +1405,11 @@ struct BlockMMA< } } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { D += sm * ldd + sn; start -= short2(sn, sm); stop -= short2(sn, sm); @@ -1431,7 +1440,7 @@ struct BlockMMA< } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { D += sm * ldd + sn; dst_tile_dims -= short2(sn, sm); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) @@ -1457,7 +1466,8 @@ struct BlockMMA< /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) { complex64_t out = epilogue_op.apply( @@ -1473,7 +1483,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -1505,7 +1515,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -1554,7 +1564,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1589,7 +1599,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/Source/Cmlx/mlx-generated/gemm_nax.cpp b/Source/Cmlx/mlx-generated/gemm_nax.cpp index a7797c7b5..d3eac6d9c 100644 --- a/Source/Cmlx/mlx-generated/gemm_nax.cpp +++ b/Source/Cmlx/mlx-generated/gemm_nax.cpp @@ -18,135 +18,6 @@ const char* gemm_nax() { #define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") #define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") -/////////////////////////////////////////////////////////////////////////////// -// Contents from "mlx/backend/metal/kernels/steel/utils.h" -/////////////////////////////////////////////////////////////////////////////// - -#line 1 "mlx/backend/metal/kernels/steel/utils.h" -// Copyright © 2024 Apple Inc. - - -#include - -METAL_FUNC ulong2 elem_to_loc_broadcast( - uint elem, - constant const int* shape, - constant const int64_t* a_strides, - constant const int64_t* b_strides, - int ndim) { - ulong loc_a{0}; - ulong loc_b{0}; - for (int i = ndim - 1; i >= 0 && elem > 0; --i) { - int pos_in_dim = (elem % shape[i]); - elem /= shape[i]; - loc_a += pos_in_dim * a_strides[i]; - loc_b += pos_in_dim * b_strides[i]; - } - return ulong2(loc_a, loc_b); -} - -METAL_FUNC ulong3 elem_to_loc_broadcast( - uint elem, - constant const int* shape, - constant const int64_t* a_strides, - constant const int64_t* b_strides, - constant const int64_t* c_strides, - int ndim) { - ulong loc_a{0}; - ulong loc_b{0}; - ulong loc_c{0}; - for (int i = ndim - 1; i >= 0 && elem > 0; --i) { - int pos_in_dim = (elem % shape[i]); - elem /= shape[i]; - loc_a += pos_in_dim * a_strides[i]; - loc_b += pos_in_dim * b_strides[i]; - loc_c += pos_in_dim * c_strides[i]; - } - return ulong3(loc_a, loc_b, loc_c); -} - -/////////////////////////////////////////////////////////////////////////////// -// Contents from "mlx/backend/metal/kernels/steel/gemm/transforms.h" -/////////////////////////////////////////////////////////////////////////////// - -#line 1 "mlx/backend/metal/kernels/steel/gemm/transforms.h" -// Copyright © 2024 Apple Inc. - - - -/////////////////////////////////////////////////////////////////////////////// -// Transforms and Epilogues -/////////////////////////////////////////////////////////////////////////////// - -namespace mlx { -namespace steel { - -template -struct TransformNone { - static METAL_FUNC OutT apply(InT x) { - return static_cast(x); - } - - static METAL_FUNC OutT apply(InT x, OutT) { - return static_cast(x); - } -}; - -template -struct TransformAdd { - TransformAdd(const float, const float) {} - - static METAL_FUNC OutT apply(InT x) { - return static_cast(x); - } - - static METAL_FUNC OutT apply(InT x, OutT c) { - return static_cast(x) + c; - } -}; - -template -struct TransformAxpby { - const float alpha; - const float beta; - - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} - - static METAL_FUNC OutT apply(InT x) { - return static_cast(x); - } - - METAL_FUNC OutT apply(InT x, OutT c) const { - return static_cast( - x * static_cast(alpha) + (static_cast(beta) * c)); - } -}; - -template -struct AccumHelper { - typedef float accum_type; -}; - -struct BlockSwizzle { - static METAL_FUNC int2 - swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) { - const int tid_x = (tid.x) >> swizzle_log; - const int tid_y = - ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1)); - return int2(tid_x, tid_y); - } -}; - -} // namespace steel -} // namespace mlx - -/////////////////////////////////////////////////////////////////////////////// -// Contents from "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.3.525.0.RrIjD2/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.850/include/metal/__exec/units.h" -/////////////////////////////////////////////////////////////////////////////// - -#line 1 "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.3.525.0.RrIjD2/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.850/include/metal/__exec/units.h" - /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/utils/type_traits.h" /////////////////////////////////////////////////////////////////////////////// @@ -232,13 +103,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -268,7 +135,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -345,6 +213,72 @@ METAL_FUNC constexpr auto sum(T x, Us... us) { #pragma METAL internals : disable +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsBase.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsBase.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.6.42.0.n257DK/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.883/include/metal/__exec/units.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.6.42.0.n257DK/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.883/include/metal/__exec/units.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTraits.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTraits.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTypes.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTypes.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsUtility.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsUtility.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsConvolution2dImpl.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsConvolution2dImpl.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsConvolution2d.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsConvolution2d.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsAvailability.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsAvailability.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsMatMul2dImpl.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsMatMul2dImpl.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsMatMul2d.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsMatMul2d.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MetalPerformancePrimitives.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MetalPerformancePrimitives.h" + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/gemm/nax.h" /////////////////////////////////////////////////////////////////////////////// @@ -422,10 +356,12 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -459,12 +395,15 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -479,7 +418,10 @@ struct BaseNAXFrag { } } else { - dst = dtype_frag_t(0); + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } } } } @@ -503,13 +445,17 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if ((r < lx) && ((c + j) < ly)) { dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j) * str_y]); } else { @@ -536,10 +482,12 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -575,12 +523,15 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -618,14 +569,18 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if (r < lx && (c + j) < ly) { dst[r * str_x + (c + j) * str_y] = static_cast(src[i * kElemCols + j]); } @@ -711,634 +666,424 @@ struct BaseNAXFrag { } } } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; + } + } }; template < typename T, - short kRows_, - short kCols_, - typename NAXFrag_t = BaseNAXFrag> -struct NAXSubTile { - STEEL_CONST short kRows = kRows_; - STEEL_CONST short kCols = kCols_; + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> +struct NAXTile { + using NAXFrag_t = NAXFrag_; + using elem_type = T; STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; - STEEL_CONST short kSubTileRows = kRows / kFragRows; - STEEL_CONST short kSubTileCols = kCols / kFragCols; + STEEL_CONST short kTileRows = kTileRows_; + STEEL_CONST short kTileCols = kTileCols_; - STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols; - STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag; + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; - STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows; - STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols; + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - using frag_type = typename NAXFrag_t::template dtype_frag_t; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - frag_type val_frags[kNumFrags]; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - METAL_FUNC constexpr void clear() { + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; + + METAL_FUNC NAXTile() thread {} + + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { - return val_frags[i * kSubTileCols + j]; + const short j) const thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& frag_at() { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } } - METAL_FUNC thread T* elems() { - return reinterpret_cast(val_frags); + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } } - METAL_FUNC const thread T* elems() const { - return reinterpret_cast(val_frags); + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { + for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + for (short j = 0; j < kTileCols; ++j) { NAXFrag_t::template row_reduce( - frag_at(i, j), &vals[i * kFragThrRows]); + frag_at(i, j), &vptr[i * kFragThrRows]); } } } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { + for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + for (short j = 0; j < kTileCols; ++j) { NAXFrag_t::template row_bin_op( - frag_at(i, j), &vals[i * kFragThrRows]); + frag_at(i, j), &vptr[i * kFragThrRows]); } } } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load( - SrcPtrType src, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + template + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( - frag_at(i, j), - src, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store( - DstPtrType dst, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store( - frag_at(i, j), - dst, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } - } - - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_rows( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_rows( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_safe( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_safe( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_safe( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_rows( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename StartX, - typename StopX, - typename StartY, - typename StopY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_slice( - DstPtrType dst, - StrX str_x, - StrY str_y, - StartX start_x, - StopX stop_x, - StartY start_y, - StopY stop_y, - OffX off_x = Int<0>{}, - OffY off_y = Int<0>{}) const { - const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) { - const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) { - NAXFrag_t::store_slice( - frag_at(), - dst, - str_x, - str_y, - start_x, - stop_x, - start_y, - stop_y, - off_x + idx_row * Int{}, - off_y + idx_col * Int{}); - }); - }); - } -}; - -template < - short RC, - short CC, - short RA, - short CA, - short RB, - short CB, - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - typename NAXFrag_t = BaseNAXFrag> -METAL_FUNC void subtile_matmad_nax( - thread NAXSubTile& C, - thread NAXSubTile& A, - metal::bool_constant, - thread NAXSubTile& B, - metal::bool_constant) { - // Static checks - constexpr short FMa = transpose_a ? CA : RA; - constexpr short FMc = RC; - static_assert(FMa == FMc, "NAX matmul: M dimensions do not match"); - - constexpr short FNb = transpose_b ? RB : CB; - constexpr short FNc = CC; - static_assert(FNb == FNc, "NAX matmul: N dimensions do not match"); - - constexpr short FKa = transpose_a ? RA : CA; - constexpr short FKb = transpose_b ? CB : RB; - static_assert(FKa == FKb, "NAX matmul: N dimensions do not match"); - - constexpr short FM = FMc; - constexpr short FN = FNc; - constexpr short FK = FKa; - - constexpr int TM = FM / 16; - constexpr int TN = FN / 16; - constexpr int TK = FK / 16; - - // Create Matmul descriptor - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - FM, - FN, - FK, - transpose_a, - transpose_b, - true, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); - - // Create matmul op - mpp::tensor_ops::matmul2d gemm_op; - - // Create matmul operands in registers - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - - // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - // Load A in to left operand registers - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_a ? kk : mm; - const short fj = transpose_a ? mm : kk; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i]; - } - } - } - - // Load B into right operand registers - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_b ? nn : kk; - const short fj = transpose_b ? kk : nn; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i]; - } - } - } - - // Load C into output registers (op handles accumulation) - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - ct_c[i] = C.elems()[i]; - } - - // Do matmul - gemm_op.run(ct_a, ct_b, ct_c); - - // Copy out results - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - C.elems()[i] = ct_c[i]; - } -} - -template -struct NAXTile { - using NAXSubTile_t = NAXSubTile_; - using elem_type = T; - STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows; - STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols; - STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile; - - STEEL_CONST short kTileRows = kTileRows_; - STEEL_CONST short kTileCols = kTileCols_; - - STEEL_CONST short kRows = kTileRows * kSubTileRows; - STEEL_CONST short kCols = kTileCols * kSubTileCols; - - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; - - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; - - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; - - NAXSubTile_t val_subtiles[kSubTiles]; - - METAL_FUNC NAXTile() thread {} - - METAL_FUNC constexpr void clear() { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTiles; ++i) { - val_subtiles[i].clear(); - } - } - - METAL_FUNC constexpr thread NAXSubTile_t& subtile_at( - const short i, - const short j) { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( - const short i, - const short j) const { - return val_subtiles[i * kTileCols + j]; - } - - template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC thread elem_type* elems() { - return reinterpret_cast(val_subtiles[0].elems()); - } - - METAL_FUNC const thread elem_type* elems() const { - return reinterpret_cast(val_subtiles[0].elems()); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - auto sub_rows = (thread metal::vec*)(&vals); - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_reduce(sub_rows[i]); - } - } - } - - template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - auto sub_rows = (thread metal::vec*)(&vals); - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_bin_op(sub_rows[i]); - } - } - } - - template - METAL_FUNC void load(const threadgroup U* src) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( + frag_at(), src, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(threadgroup U* dst) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), dst, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void load(const device U* src, const int ld) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( - &src[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(device U* dst, const int ld) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( - &dst[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_rows( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_rows( + frag_at(), + src, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_safe( + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_safe( + frag_at(), src, ld, Int<1>{}, src_tile_dims.y, src_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_rows( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), + dst, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_safe( + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), dst, ld, Int<1>{}, dst_tile_dims.y, dst_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template @@ -1346,10 +1091,11 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { - subtile_at().store_slice( + NAXFrag_t::store_slice( + frag_at(), dst, ld, Int<1>{}, @@ -1357,8 +1103,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1378,51 +1124,54 @@ METAL_FUNC void tile_matmad_nax( metal::bool_constant) { // Static checks constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; - constexpr short TMc = CTile::kTileRows; - static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match"); - - constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows; - constexpr short FMc = CTile::kSubTileRows; - static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match"); + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; - constexpr short TNc = CTile::kTileCols; - static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match"); - - constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols; - constexpr short FNc = CTile::kSubTileCols; - static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match"); + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; - constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows; - static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match"); - - constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols; - constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows; - static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match"); + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); - constexpr short TM = TMc; - constexpr short TN = TNc; - constexpr short TK = TKa; + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; - // Do matmul here - STEEL_PRAGMA_UNROLL - for (short i = 0; i < TM; ++i) { + if constexpr (TN == 1 && TM % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; mm += 2) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < TN; ++j) { + for (short mm = 0; mm < TM; ++mm) { STEEL_PRAGMA_UNROLL - for (short k = 0; k < TK; ++k) { - const short ra = transpose_a ? k : i; - const short ca = transpose_a ? i : k; - const short rb = transpose_b ? j : k; - const short cb = transpose_b ? k : j; - - subtile_matmad_nax( - C.subtile_at(i, j), - A.subtile_at(ra, ca), - metal::bool_constant{}, - B.subtile_at(rb, cb), - metal::bool_constant{}); + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } } } } @@ -1501,6 +1250,129 @@ struct GEMMAddMMParams { } // namespace steel } // namespace mlx +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils.h" +// Copyright © 2024 Apple Inc. + + +#include + +METAL_FUNC ulong2 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + } + return ulong2(loc_a, loc_b); +} + +METAL_FUNC ulong3 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + constant const int64_t* c_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + ulong loc_c{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + loc_c += pos_in_dim * c_strides[i]; + } + return ulong3(loc_a, loc_b, loc_c); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/transforms.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/transforms.h" +// Copyright © 2024 Apple Inc. + + + +/////////////////////////////////////////////////////////////////////////////// +// Transforms and Epilogues +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template +struct TransformNone { + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + static METAL_FUNC OutT apply(InT x, OutT) { + return static_cast(x); + } +}; + +template +struct TransformAdd { + TransformAdd(const float, const float) thread {} + + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + static METAL_FUNC OutT apply(InT x, OutT c) { + return static_cast(x) + c; + } +}; + +template +struct TransformAxpby { + const float alpha; + const float beta; + + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} + + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + METAL_FUNC OutT apply(InT x, OutT c) const thread { + return static_cast( + x * static_cast(alpha) + (static_cast(beta) * c)); + } +}; + +template +struct AccumHelper { + typedef float accum_type; +}; + +struct BlockSwizzle { + static METAL_FUNC int2 + swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) { + const int tid_x = (tid.x) >> swizzle_log; + const int tid_y = + ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1)); + return int2(tid_x, tid_y); + } +}; + +} // namespace steel +} // namespace mlx + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h" /////////////////////////////////////////////////////////////////////////////// @@ -1525,9 +1397,6 @@ template < bool kAlignedM, bool kAlignedN, bool kAlignedK, - short UM, - short UN, - short UK, typename AccumType = float> auto gemm_loop( const device T* A, @@ -1538,9 +1407,9 @@ auto gemm_loop( int gemm_k_iterations_aligned, const short sgp_sm, const short sgp_sn) { - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; constexpr int RA = transpose_a ? TK : TM; constexpr int CA = transpose_a ? TM : TK; @@ -1548,25 +1417,25 @@ auto gemm_loop( constexpr int RB = transpose_b ? TN : TK; constexpr int CB = transpose_b ? TK : TN; - using DSubTile = NAXSubTile; - using ASubTile = - NAXSubTile; - using BSubTile = - NAXSubTile; - - NAXTile Dtile; + NAXTile Dtile; Dtile.clear(); + const bool has_output = sgp_sm > 0 && sgp_sn > 0; + int gemm_k_iterations_ = gemm_k_iterations_aligned; STEEL_PRAGMA_NO_UNROLL for (int kk0 = 0; kk0 < gemm_k_iterations_; kk0++) { threadgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + continue; + } STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; const int k = kk1; volatile int compiler_barrier; @@ -1606,51 +1475,38 @@ auto gemm_loop( if constexpr (!kAlignedK) { simdgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + return Dtile; + } const short rem_bk = K - gemm_k_iterations_ * BK; STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < rem_bk; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; - STEEL_PRAGMA_UNROLL - for (int mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (int nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (int kk = 0; kk < TK; kk++) { - const int m = mm * UM; - const int n = nn * UN; - const int k = kk1 + kk * UK; - const short psk = max(0, rem_bk - k); - - const int A_offset = transpose_a ? (m + k * lda) : (m * lda + k); - const int B_offset = transpose_b ? (k + n * ldb) : (k * ldb + n); - - { - const short psm = kAlignedM ? SM : max(0, sgp_sm - m); - const short rmax = transpose_a ? psk : psm; - const short cmax = transpose_a ? psm : psk; - Atile.load_safe(A + A_offset, lda, short2(cmax, rmax)); - } - - { - const short psn = kAlignedN ? SN : max(0, sgp_sn - n); - const short rmax = transpose_b ? psn : psk; - const short cmax = transpose_b ? psk : psn; - Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax)); - } - - subtile_matmad_nax( - Dtile.subtile_at(mm, nn), - Atile.subtile_at(0, 0), - metal::bool_constant{}, - Btile.subtile_at(0, 0), - metal::bool_constant{}); - } - } - } + const int k = kk1; + const short psk = max(0, rem_bk - k); + + const short2 Aklims = + transpose_a ? short2(sgp_sm, psk) : short2(psk, sgp_sm); + const short2 Bklims = + transpose_b ? short2(psk, sgp_sn) : short2(sgp_sn, psk); + + const int A_offset = transpose_a ? k * lda : k; + const int B_offset = transpose_b ? k : k * ldb; + + Atile.load_safe(A + A_offset, lda, Aklims); + Btile.load_safe(B + B_offset, ldb, Bklims); + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); } } diff --git a/Source/Cmlx/mlx-generated/gemv.cpp b/Source/Cmlx/mlx-generated/gemv.cpp new file mode 100644 index 000000000..36ed559bb --- /dev/null +++ b/Source/Cmlx/mlx-generated/gemv.cpp @@ -0,0 +1,2352 @@ +namespace mlx::core::metal { + +const char* gemv() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/gemv.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/bf16.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/bf16.h" +// Copyright © 2023 Apple Inc. + + +#include + +using namespace metal; + +typedef bfloat bfloat16_t; +inline uint16_t bfloat16_to_uint16(const bfloat16_t x) { + return as_type(x); +} + +inline bfloat16_t uint16_to_bfloat16(const uint16_t x) { + return as_type(x); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/bf16_math.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/bf16_math.h" +// Copyright © 2023 Apple Inc. + + +/////////////////////////////////////////////////////////////////////////////// +// Metal math for bfloat16 +/////////////////////////////////////////////////////////////////////////////// + +/* + +Following the Metal Shading Language Specification (Metal 3.1) + +"bfloat is an extended itypeing point type that only allows implicit conversion + to a type of greater itypeing point rank. While bfloat can be implicitly + converted to itype, it cannot be implicitly converted to half, and neither + itype nor half can be implicitly converted to bfloat." + +Further, as far as I can tell, the stdlib math/simd functions are not defined +for bfloat and calling with an argument of type bfloat will result in that +argument getting implicitly converted to itype which then returns an output +that is (likely) a itype which cannot be implicitly converted into a bfloat + +This leads to situations where +bfloat a = 5.0bf; +bfloat b = metal::abs(a); // this will throw an error since abs return itype +bfloat c = static_cast(metal::abs(a)); // this is fine + +For the moment, I will be adding overloaded instantiations of the math +functions to accordingly automatically handle the casting + +*/ + +#define instantiate_metal_math_funcs(itype, otype, ctype, mfast) \ + \ + METAL_FUNC otype abs(itype x) { \ + return static_cast(__metal_fabs(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype acos(itype x) { \ + return static_cast(__metal_acos(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype acosh(itype x) { \ + return static_cast(__metal_acosh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype asin(itype x) { \ + return static_cast(__metal_asin(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype asinh(itype x) { \ + return static_cast(__metal_asinh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype atan(itype y_over_x) { \ + return static_cast( \ + __metal_atan(static_cast(y_over_x), mfast)); \ + } \ + METAL_FUNC otype atan2(itype y, itype x) { \ + return static_cast( \ + __metal_atan2(static_cast(y), static_cast(x), mfast)); \ + } \ + METAL_FUNC otype atanh(itype x) { \ + return static_cast(__metal_atanh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype ceil(itype x) { \ + return static_cast(__metal_ceil(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cos(itype x) { \ + return static_cast(__metal_cos(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cosh(itype x) { \ + return static_cast(__metal_cosh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cospi(itype x) { \ + return static_cast(__metal_cospi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype divide(itype x, itype y) { \ + return static_cast( \ + __metal_divide(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype exp(itype x) { \ + return static_cast(__metal_exp(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype exp10(itype x) { \ + return static_cast(__metal_exp10(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype exp2(itype x) { \ + return static_cast(__metal_exp2(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fabs(itype x) { \ + return static_cast(__metal_fabs(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fdim(itype x, itype y) { \ + ctype t = static_cast(x - y); \ + return static_cast(select(t, ctype(0), t < ctype(0) || x == y)); \ + } \ + METAL_FUNC otype floor(itype x) { \ + return static_cast(__metal_floor(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fma(itype x, itype y, itype z) { \ + return static_cast(__metal_fma( \ + static_cast(x), static_cast(y), static_cast(z))); \ + } \ + METAL_FUNC otype fmax(itype x, itype y) { \ + return static_cast( \ + __metal_fmax(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fmax3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmax3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmedian3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmedian3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmin(itype x, itype y) { \ + return static_cast( \ + __metal_fmin(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fmin3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmin3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmod(itype x, itype y) { \ + return static_cast( \ + __metal_fmod(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fract(itype x) { \ + return static_cast(__metal_fract(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype frexp(itype x, thread int& exp) { \ + return static_cast(__metal_frexp(static_cast(x), &exp)); \ + } \ + METAL_FUNC otype ldexp(itype x, int k) { \ + return static_cast(__metal_ldexp(static_cast(x), k, mfast)); \ + } \ + METAL_FUNC otype log(itype x) { \ + return static_cast(__metal_log(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype log10(itype x) { \ + return static_cast(__metal_log10(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype log2(itype x) { \ + return static_cast(__metal_log2(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype max(itype x, itype y) { \ + return static_cast( \ + __metal_fmax(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype max3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmax3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype median3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmedian3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype min(itype x, itype y) { \ + return static_cast( \ + __metal_fmin(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype min3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmin3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype nextafter(itype x, itype y) { \ + return static_cast( \ + __metal_nextafter(static_cast(x), static_cast(y))); \ + } \ + METAL_FUNC otype pow(itype x, itype y) { \ + return static_cast( \ + __metal_pow(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype powr(itype x, itype y) { \ + return static_cast( \ + __metal_powr(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype rint(itype x) { \ + return static_cast(__metal_rint(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype round(itype x) { \ + return static_cast(__metal_round(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype rsqrt(itype x) { \ + return static_cast(__metal_rsqrt(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sin(itype x) { \ + return static_cast(__metal_sin(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sinh(itype x) { \ + return static_cast(__metal_sinh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sinpi(itype x) { \ + return static_cast(__metal_sinpi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sqrt(itype x) { \ + return static_cast(__metal_sqrt(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tan(itype x) { \ + return static_cast(__metal_tan(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tanh(itype x) { \ + return static_cast(__metal_tanh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tanpi(itype x) { \ + return static_cast(__metal_tanpi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype trunc(itype x) { \ + return static_cast(__metal_trunc(static_cast(x), mfast)); \ + } + +namespace metal { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_MAYBE_FAST_MATH__); + +namespace fast { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_FAST_MATH__); + +} // namespace fast + +namespace precise { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_PRECISE_MATH__); + +} // namespace precise + +} // namespace metal + +/////////////////////////////////////////////////////////////////////////////// +// Metal simd for bfloat16 +/////////////////////////////////////////////////////////////////////////////// + +#define instantiate_metal_simd_comm_funcs( \ + itype, otype, ctype, itype_to_ctype, ctype_to_otype) \ + \ + METAL_FUNC otype simd_broadcast(itype data, ushort broadcast_lane_id) { \ + return ctype_to_otype( \ + __metal_simd_broadcast(itype_to_ctype(data), broadcast_lane_id)); \ + } \ + \ + METAL_FUNC otype simd_shuffle(itype data, ushort simd_lane_id) { \ + return ctype_to_otype( \ + __metal_simd_shuffle(itype_to_ctype(data), simd_lane_id)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_down( \ + itype data, itype filling_data, ushort delta, ushort modulo) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_down( \ + itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_down( \ + itype data, itype filling_data, ushort delta) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_down( \ + itype_to_ctype(data), \ + itype_to_ctype(filling_data), \ + delta, \ + __metal_get_simdgroup_size(ushort()))); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_up( \ + itype data, itype filling_data, ushort delta, ushort modulo) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_up( \ + itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_up( \ + itype data, itype filling_data, ushort delta) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_up( \ + itype_to_ctype(data), \ + itype_to_ctype(filling_data), \ + delta, \ + __metal_get_simdgroup_size(ushort()))); \ + } \ + \ + METAL_FUNC otype simd_shuffle_down(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_down(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_rotate_down(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_rotate_down(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_rotate_up(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_rotate_up(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_up(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_up(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_xor(itype data, ushort mask) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_xor(itype_to_ctype(data), mask)); \ + } + +#define instantiate_metal_simd_reduction_funcs(itype, otype, ctype) \ + \ + METAL_FUNC otype simd_max(itype data) { \ + return static_cast(__metal_simd_max(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_min(itype data) { \ + return static_cast(__metal_simd_min(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_exclusive_product(itype data) { \ + return static_cast( \ + __metal_simd_prefix_exclusive_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_exclusive_sum(itype data) { \ + return static_cast( \ + __metal_simd_prefix_exclusive_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_inclusive_product(itype data) { \ + return static_cast( \ + __metal_simd_prefix_inclusive_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_inclusive_sum(itype data) { \ + return static_cast( \ + __metal_simd_prefix_inclusive_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_product(itype data) { \ + return static_cast(__metal_simd_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_sum(itype data) { \ + return static_cast(__metal_simd_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_xor(itype data) { \ + return static_cast(__metal_simd_xor(static_cast(data))); \ + } + +namespace metal { + +instantiate_metal_simd_comm_funcs( + bfloat16_t, + bfloat16_t, + uint16_t, + bfloat16_to_uint16, + uint16_to_bfloat16); +instantiate_metal_simd_reduction_funcs(bfloat16_t, bfloat16_t, float); + +} // namespace metal + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/complex.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/complex.h" +// Copyright © 2023 Apple Inc. + + +#include + + +using namespace metal; + +template +struct complex_t; + +template +static constexpr constant bool is_complex_v = false; + +template +static constexpr constant bool is_complex_v> = true; + +// Metal accepts explicit bfloat casts that is_convertible_v reports as false. +template +static constexpr constant bool is_lane_convertible_v = + is_convertible_v || + (is_same_v && is_convertible_v) || + (is_same_v && is_convertible_v); + +template +struct complex_t { + using value_type = T; + + T real; + T imag; + + // Constructors + constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {}; + constexpr complex_t() thread : real(0), imag(0) {}; + constexpr complex_t() threadgroup : real(0), imag(0) {}; + + // Conversions from scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) thread : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) threadgroup : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) device : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) constant : real(static_cast(x)), + imag(static_cast(0)) {} + + // Conversions between complex types + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) thread : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) threadgroup + : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) device : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + // Conversions to and from two-lane vectors (the FFT lane representation) + constexpr complex_t(vec v) thread : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) threadgroup : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) device : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) constant : real(v.x), imag(v.y) {}; + + constexpr operator vec() const thread { + return vec(real, imag); + } + + constexpr operator vec() const threadgroup { + return vec(real, imag); + } + + constexpr operator vec() const device { + return vec(real, imag); + } + + constexpr operator vec() const constant { + return vec(real, imag); + } + + // Conversions to scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const thread { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const threadgroup { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const device { + return static_cast(real); + } + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const constant { + return static_cast(real); + } +}; + +using complex32_t = complex_t; +using complex64_t = complex_t; + +static_assert(sizeof(complex32_t) == 2 * sizeof(half)); +static_assert(sizeof(complex64_t) == 2 * sizeof(float)); +static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); + +template +constexpr complex_t operator-(complex_t x) { + return {-x.real, -x.imag}; +} + +template +constexpr bool operator>=(complex_t a, complex_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); +} + +template +constexpr bool operator>(complex_t a, complex_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); +} + +template +constexpr bool operator<=(complex_t a, complex_t b) { + return operator>=(b, a); +} + +template +constexpr bool operator<(complex_t a, complex_t b) { + return operator>(b, a); +} + +template +constexpr bool operator==(complex_t a, complex_t b) { + return a.real == b.real && a.imag == b.imag; +} + +template +constexpr complex_t operator+(complex_t a, complex_t b) { + return {a.real + b.real, a.imag + b.imag}; +} + +template +constexpr thread complex_t& operator+=( + thread complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template +constexpr threadgroup complex_t& operator+=( + threadgroup complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template +constexpr device complex_t& operator+=( + device complex_t& a, + complex_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(U a, complex_t b) { + return {static_cast(a) + b.real, b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(complex_t a, U b) { + return {a.real + static_cast(b), a.imag}; +} + +template +constexpr complex_t operator-(complex_t a, complex_t b) { + return {a.real - b.real, a.imag - b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(U a, complex_t b) { + return {static_cast(a) - b.real, -b.imag}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(complex_t a, U b) { + return {a.real - static_cast(b), a.imag}; +} + +template +constexpr complex_t operator*(complex_t a, complex_t b) { + return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; +} + +template +constexpr complex_t operator/(complex_t a, complex_t b) { + auto denom = b.real * b.real + b.imag * b.imag; + auto x = a.real * b.real + a.imag * b.imag; + auto y = a.imag * b.real - a.real * b.imag; + return {x / denom, y / denom}; +} + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator/(U a, complex_t b) { + auto scalar = static_cast(a); + auto denom = b.real * b.real + b.imag * b.imag; + auto x = scalar * b.real; + auto y = -scalar * b.imag; + return {x / denom, y / denom}; +} + +template +constexpr complex_t operator%(complex_t a, complex_t b) { + auto real = a.real - (b.real * static_cast(a.real / b.real)); + auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); + if (real != 0 && (real < 0 != b.real < 0)) { + real += b.real; + } + if (imag != 0 && (imag < 0 != b.imag < 0)) { + imag += b.imag; + } + return {real, imag}; +} + +static_assert( + (complex_t{1.0h, 2.0h} * complex_t{3.0h, 4.0h}).real == -5.0h); +static_assert( + (complex_t{bfloat16_t(1.0f), bfloat16_t(2.0f)} * + complex_t{bfloat16_t(3.0f), bfloat16_t(4.0f)}) + .real == bfloat16_t(-5.0f)); + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/defines.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/defines.h" +// Copyright © 2023 Apple Inc. + + +#if defined __METAL__ || defined MLX_METAL_JIT +#define MTL_CONST constant +#else +#define MTL_CONST +#endif + +static MTL_CONST constexpr int MAX_REDUCE_SPECIALIZED_DIMS = 4; +static MTL_CONST constexpr int REDUCE_N_READS = 4; +static MTL_CONST constexpr int REDUCE_N_WRITES = 4; +static MTL_CONST constexpr int SOFTMAX_N_READS = 4; +static MTL_CONST constexpr int RMS_N_READS = 4; +static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; + +// Instantiate a templated kernel. +// Extra args are used as template parameters: +// e.g. instantiate_kernel(binary_int, binary, a, b) -> +// [[host_name(binary_int)]] [kernel] binary +#define instantiate_kernel(name, func, ...) \ + template [[host_name( \ + name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>; + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/logging.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/logging.h" +// Copyright © 2025 Apple Inc. + + +#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320) +#include + +namespace mlx { +using os_log = metal::os_log; +} // namespace mlx + +#else + +namespace mlx { +struct os_log { + constexpr os_log(constant char*, constant char*) constant {} + + template + void log_debug(constant char*, Args...) const thread {} + + template + void log_debug(constant char*, Args...) const constant {} +}; +} // namespace mlx + +#endif + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/utils.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/utils.h" +// Copyright © 2023-2024 Apple Inc. + + +#include + + +typedef half float16_t; + +// Work per thread values for different types. The values here are expected to +// match get_work_per_thread in mlx/backend/metal/utils.h +template +struct WorkPerThread { + static_assert(sizeof(U) <= 8, "Type too large"); + static constexpr int constant n = 8 / sizeof(U); +}; + +/////////////////////////////////////////////////////////////////////////////// +// Type limits utils +/////////////////////////////////////////////////////////////////////////////// + +template +struct Limits { + static const constant U max = metal::numeric_limits::max(); + static const constant U min = metal::numeric_limits::min(); + static const constant U finite_max = metal::numeric_limits::max(); + static const constant U finite_min = metal::numeric_limits::min(); +}; + +#define instantiate_default_limit(type) \ + template <> \ + struct Limits { \ + static constexpr constant type max = metal::numeric_limits::max(); \ + static constexpr constant type min = metal::numeric_limits::min(); \ + static constexpr constant type finite_max = \ + metal::numeric_limits::max(); \ + static constexpr constant type finite_min = \ + metal::numeric_limits::min(); \ + }; + +instantiate_default_limit(uint8_t); +instantiate_default_limit(uint16_t); +instantiate_default_limit(uint32_t); +instantiate_default_limit(uint64_t); +instantiate_default_limit(int8_t); +instantiate_default_limit(int16_t); +instantiate_default_limit(int32_t); +instantiate_default_limit(int64_t); + +#define instantiate_float_limit(type) \ + template <> \ + struct Limits { \ + static constexpr constant type max = \ + metal::numeric_limits::infinity(); \ + static constexpr constant type min = \ + -metal::numeric_limits::infinity(); \ + static constexpr constant type finite_max = \ + metal::numeric_limits::max(); \ + static constexpr constant type finite_min = \ + -metal::numeric_limits::max(); \ + }; + +instantiate_float_limit(half); +instantiate_float_limit(float); +instantiate_float_limit(bfloat16_t); + +template <> +struct Limits { + static constexpr constant bool max = true; + static constexpr constant bool min = false; +}; + +template +struct Limits> { + inline static constexpr constant complex_t max = complex_t( + metal::numeric_limits::infinity(), + metal::numeric_limits::infinity()); + inline static constexpr constant complex_t min = complex_t( + -metal::numeric_limits::infinity(), + -metal::numeric_limits::infinity()); +}; + +/////////////////////////////////////////////////////////////////////////////// +// Indexing utils +/////////////////////////////////////////////////////////////////////////////// + +#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") + +/////////////////////////////////////////////////////////////////////////////// +// Single Array with generic dims + +template +METAL_FUNC IdxT elem_to_loc( + IdxT elem, + constant const int* shape, + constant const int64_t* strides, + int ndim) { + IdxT loc = 0; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + loc += (elem % shape[i]) * IdxT(strides[i]); + elem /= shape[i]; + } + return loc; +} + +// Non templated version to handle arbitrary dims +template +METAL_FUNC IdxT elem_to_loc( + uint3 elem, + constant const int* shape, + constant const int64_t* strides, + int ndim) { + IdxT loc = + elem.x * IdxT(strides[ndim - 1]) + elem.y * IdxT(strides[ndim - 2]); + for (int d = ndim - 3; d >= 0; --d) { + loc += (elem.z % shape[d]) * IdxT(strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Single Array with fixed N dims + +template +METAL_FUNC IdxT elem_to_loc_1(uint elem, constant const int64_t& stride) { + return elem * IdxT(stride); +} + +template +METAL_FUNC IdxT elem_to_loc_2(uint2 elem, constant const int64_t strides[2]) { + return elem.x * IdxT(strides[1]) + elem.y * IdxT(strides[0]); +} + +template +METAL_FUNC IdxT elem_to_loc_3(uint3 elem, constant const int64_t strides[3]) { + return elem.x * IdxT(strides[2]) + elem.y * IdxT(strides[1]) + + elem.z * IdxT(strides[0]); +} + +/////////////////////////////////////////////////////////////////////////////// +// Multiple Arrays with generic dims + +template +METAL_FUNC vec elem_to_loc_2_nd( + uint3 elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + int ndim) { + vec loc = { + IdxT( + elem.x * IdxT(a_strides[ndim - 1]) + + IdxT(elem.y) * IdxT(a_strides[ndim - 2])), + IdxT( + elem.x * IdxT(b_strides[ndim - 1]) + + elem.y * IdxT(b_strides[ndim - 2]))}; + for (int d = ndim - 3; d >= 0; --d) { + uint l = elem.z % shape[d]; + loc.x += l * IdxT(a_strides[d]); + loc.y += l * IdxT(b_strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +template +METAL_FUNC vec elem_to_loc_3_nd( + uint3 elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + constant const int64_t* c_strides, + int ndim) { + vec loc = { + IdxT(elem.x * IdxT(a_strides[ndim - 1])) + + IdxT(elem.y * IdxT(a_strides[ndim - 2])), + IdxT(elem.x * IdxT(b_strides[ndim - 1])) + + IdxT(elem.y * IdxT(b_strides[ndim - 2])), + IdxT(elem.x * IdxT(c_strides[ndim - 1])) + + IdxT(elem.y * IdxT(c_strides[ndim - 2]))}; + for (int d = ndim - 3; d >= 0; --d) { + uint l = elem.z % shape[d]; + loc.x += l * IdxT(a_strides[d]); + loc.y += l * IdxT(b_strides[d]); + loc.z += l * IdxT(c_strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Elem to loc in a loop utils +/////////////////////////////////////////////////////////////////////////////// + +template +struct LoopedElemToLoc { + int dim; + LoopedElemToLoc inner_looper; + OffsetT offset{0}; + int index{0}; + + LoopedElemToLoc(int dim) thread : dim(dim), inner_looper(dim - 1) {} + + void next(const constant int* shape, const constant int64_t* strides) thread { + if (dim == 0) { + return; + } + index++; + offset += OffsetT(strides[dim - 1]); + if (index >= shape[dim - 1]) { + index = 0; + inner_looper.next(shape, strides); + offset = inner_looper.offset; + } + } + + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { + if (dim == 0) { + return; + } + index += n; + offset += n * OffsetT(strides[dim - 1]); + + if (index >= shape[dim - 1]) { + int extra = index - shape[dim - 1]; + if (extra >= shape[dim - 1]) { + inner_looper.next(1 + extra / shape[dim - 1], shape, strides); + extra = extra % shape[dim - 1]; + } else { + inner_looper.next(shape, strides); + } + index = 0; + offset = inner_looper.offset; + if (extra > 0) { + next(extra, shape, strides); + } + } + } + + OffsetT location() thread { + return offset; + } +}; + +template +struct LoopedElemToLoc<1, OffsetT, true> { + int dim; + OffsetT offset{0}; + uint index{0}; + + LoopedElemToLoc(int dim) thread : dim(dim) {} + + void next(const constant int* shape, const constant int64_t* strides) thread { + index++; + if (dim > 1) { + offset = elem_to_loc(index, shape, strides, dim); + } else { + offset += OffsetT(strides[0]); + } + } + + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { + index += n; + if (dim > 1) { + offset = elem_to_loc(index, shape, strides, dim); + } else { + offset = index * OffsetT(strides[0]); + } + } + + OffsetT location() thread { + return offset; + } +}; + +template +struct LoopedElemToLoc<1, OffsetT, false> { + OffsetT offset{0}; + + LoopedElemToLoc(int) thread {} + + void next(const constant int*, const constant int64_t* strides) thread { + offset += OffsetT(strides[0]); + } + + void next(int n, const constant int*, const constant int64_t* strides) + thread { + offset += n * OffsetT(strides[0]); + } + + OffsetT location() thread { + return offset; + } +}; + +/////////////////////////////////////////////////////////////////////////////// +// Calculation utils +/////////////////////////////////////////////////////////////////////////////// + +/** Compute ceil((float)N/(float)M) */ +template +inline T ceildiv(T N, U M) { + return (N + M - 1) / M; +} + +// https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1202 +inline float log1p(float x) { + float xp1 = 1.0f + x; + if (xp1 == Limits::max) { + return Limits::max; + } + if (xp1 == 1.0f) { + return x; + } + + return x * (metal::log(xp1) / (xp1 - 1.0f)); +} + +inline bfloat16_t log1p(bfloat16_t x) { + float xp1 = 1.0f + static_cast(x); + if (xp1 == Limits::max) { + return Limits::max; + } + if (xp1 == 1.0f) { + return x; + } + + return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f))); +} + +inline complex64_t log1p(complex64_t in) { + float x = in.real; + float y = in.imag; + float zabs = metal::precise::sqrt(x * x + y * y); + float theta = metal::atan2(y, x + 1); + if (zabs < 0.5f) { + float r = x * (2 + x) + y * y; + if (r == 0) { // handle underflow + return {x, theta}; + } + return {0.5f * log1p(r), theta}; + } else { + auto z0 = metal::sqrt((x + 1) * (x + 1) + y * y); + return {metal::log(z0), theta}; + } +} + +/////////////////////////////////////////////////////////////////////////////// +// SIMD shuffle ops +/////////////////////////////////////////////////////////////////////////////// + +inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) { + return as_type( + metal::simd_shuffle_down(as_type(data), delta)); +} + +inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) { + return as_type( + metal::simd_shuffle_down(as_type(data), delta)); +} + +inline bool simd_shuffle_down(bool data, uint16_t delta) { + return simd_shuffle_down(static_cast(data), delta); +} + +inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) { + return complex64_t( + simd_shuffle_down(data.real, delta), simd_shuffle_down(data.imag, delta)); +} + +inline uint64_t simd_shuffle_up(uint64_t data, uint16_t delta) { + return as_type(metal::simd_shuffle_up(as_type(data), delta)); +} + +inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) { + return as_type(metal::simd_shuffle_up(as_type(data), delta)); +} + +inline bool simd_shuffle_up(bool data, uint16_t delta) { + return simd_shuffle_up(static_cast(data), delta); +} + +inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) { + return complex64_t( + simd_shuffle_up(data.real, delta), simd_shuffle_up(data.imag, delta)); +} + +inline uint64_t +simd_shuffle_and_fill_up(uint64_t data, uint64_t filling, uint16_t delta) { + return as_type(metal::simd_shuffle_and_fill_up( + as_type(data), as_type(filling), delta)); +} + +inline int64_t +simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) { + return as_type(metal::simd_shuffle_and_fill_up( + as_type(data), as_type(filling), delta)); +} + +inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) { + return simd_shuffle_and_fill_up( + static_cast(data), static_cast(filling), delta); +} + +inline complex64_t simd_shuffle_and_fill_up( + complex64_t data, + complex64_t filling, + uint16_t delta) { + return complex64_t( + simd_shuffle_and_fill_up(data.real, filling.real, delta), + simd_shuffle_and_fill_up(data.imag, filling.imag, delta)); +} + +inline uint64_t simd_shuffle(uint64_t data, uint16_t lane) { + return as_type(metal::simd_shuffle(as_type(data), lane)); +} + +inline int64_t simd_shuffle(int64_t data, uint16_t lane) { + return as_type(metal::simd_shuffle(as_type(data), lane)); +} + +inline bool simd_shuffle(bool data, uint16_t lane) { + return simd_shuffle(static_cast(data), lane); +} + +inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) { + return complex64_t( + simd_shuffle(data.real, lane), simd_shuffle(data.imag, lane)); +} + +// std::conditional is not included with Metal +template +struct ConditionalType { + using type = U; +}; + +template +struct ConditionalType { + using type = T; +}; + +/////////////////////////////////////////////////////////////////////////////// +// Type casting utils +/////////////////////////////////////////////////////////////////////////////// + +template +inline U cast_to(T val) { + return static_cast(val); +} + +template <> +inline bool cast_to(float val) { + return (as_type(val) & 0x7FFFFFFF) != 0; +} + +template <> +inline bool cast_to(bfloat16_t val) { + return (as_type(val) & 0x7FFF) != 0; +} + +template <> +inline bool cast_to(complex64_t val) { + return cast_to(val.real) || cast_to(val.imag); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/utils.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/utils.h" +// Copyright © 2024 Apple Inc. + + +#include + +METAL_FUNC ulong2 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + } + return ulong2(loc_a, loc_b); +} + +METAL_FUNC ulong3 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + constant const int64_t* c_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + ulong loc_c{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + loc_c += pos_in_dim * c_strides[i]; + } + return ulong3(loc_a, loc_b, loc_c); +} + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/gemv.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/gemv.h" +// Copyright © 2023-2024 Apple Inc. + + +#include +#include + + + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +/// Matrix vector multiplication +/////////////////////////////////////////////////////////////////////////////// + +#define MLX_MTL_CONST static constant constexpr const + +template +struct DefaultAccT { + using type = float; +}; +template <> +struct DefaultAccT { + using type = complex64_t; +}; + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN, /* Thread cols (in elements) */ + const bool kDoAxpby, /* Do out = alpha * out + beta * bias */ + typename AccT = typename DefaultAccT::type> +struct GEMVKernel { + using acc_type = AccT; + + MLX_MTL_CONST int threadsM = BM * SM; + MLX_MTL_CONST int threadsN = BN * SN; + + MLX_MTL_CONST int blockM = threadsM * TM; + MLX_MTL_CONST int blockN = threadsN * TN; + + static_assert(SM * SN == 32, "simdgroup can only have 32 threads"); + + static_assert( + SN == 4 || SN == 8 || SN == 16 || SN == 32, + "gemv block must have a width of 4, 8, 16, or 32"); + + // - The matrix of size (M = out_vec_size, K = in_vec_size) is divided up + // into blocks of (blockM, blockN) divided among threadgroups + // - Every thread works on a block of (TM, TN) + // - We assume each threadgroup has (threadsN, threadsM, 1) threads + // + // 1. A thread loads TN elements each from mat along TM rows + // and the corresponding scalar from the vector + // 2. The thread then multiplies and adds to accumulate its local result for + // the block + // 3. At the end, each thread has accumulated results over all blocks across + // the rows. These are then summed up across the threadgroup + // 4. Each threadgroup writes its accumulated blockM outputs + // + // Edge case handling: + // - The threadgroup with the largest tid has blocks that exceed the matrix + // * The blocks that start outside the matrix are never read (thread results + // remain zero) + // * The last thread that partially overlaps with the matrix is shifted + // inwards such that the thread block fits exactly in the matrix + + MLX_MTL_CONST short tgp_mem_size = BN > 1 ? BN*(blockM + TM) : 0; + MLX_MTL_CONST bool needs_tgp_reduction = BN > 1; + + template + static METAL_FUNC void + load_unsafe(const device T* src, thread U dst[TN], const int src_offset = 0) { + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + dst[tn] = static_cast(src[src_offset + tn]); + } + } + + template + static METAL_FUNC void load_safe( + const device T* src, + thread U dst[TN], + const int src_offset = 0, + const int src_size = TN) { + if (src_offset + TN <= src_size) { + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + dst[tn] = static_cast(src[src_offset + tn]); + } + } else { // Edgecase + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + dst[tn] = src_offset + tn < src_size + ? static_cast(src[src_offset + tn]) + : U(0); + } + } + } + + static METAL_FUNC void run( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& matrix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& bias_stride [[buffer(14)]], + threadgroup AccT* tgp_memory [[threadgroup(0)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + // Appease compiler + (void)lid; + + // Thread local accumulation results + thread AccT result[TM] = {0}; + thread T inter[TN]; + thread AccT v_coeff[TN]; + + const int thrM = SN != 32 ? simd_lid / SN : 0; + const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid); + + const int sgN = BN != 1 ? (simd_gid % BN) : 0; + + const int simdM = BN != 1 ? SM * (simd_gid / BN) : int(SM * simd_gid); + const int simdN = BN != 1 ? SN * (simd_gid % BN) : 0; + + int bm = (simdM + thrM) * TM; + int bn = (simdN + thrN) * TN; + + // Block position + int out_row = tid.x * blockM + bm; + + // Exit simdgroup if rows out of bound + if (out_row >= out_vec_size) + return; + + // Adjust tail simdgroup to ensure in bound reads + out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM; + + // Advance matrix + mat += size_t(out_row) * matrix_ld; + + constexpr const uniform loop_stride = make_uniform(blockN); + const uniform in_size = make_uniform(in_vec_size); + const uniform n_iter = in_size / loop_stride; + const uniform last_iter = loop_stride * n_iter; + const uniform leftover = in_size - last_iter; + + // Loop over in_vec in blocks of blockN + for (int i = 0; i < n_iter; ++i) { + load_unsafe(in_vec, v_coeff, bn); + + // Per thread work loop + int mat_offset = 0; + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + // Load for the row + load_unsafe(mat, inter, mat_offset + bn); + + // Accumulate results + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + result[tm] += inter[tn] * v_coeff[tn]; + } + + mat_offset += matrix_ld; + } + + bn += blockN; + } + + if (leftover > 0) { + load_safe(in_vec, v_coeff, bn, in_size); + + // Per thread work loop + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + // Load for the row + load_safe(&mat[tm * matrix_ld], inter, bn, in_size); + + // Accumulate results + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + result[tm] += inter[tn] * v_coeff[tn]; + } + } + } + + // Simdgroup accumulations + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + MLX_MTL_PRAGMA_UNROLL + for (ushort sn = (SN / 2); sn >= 1; sn >>= 1) { + result[tm] += simd_shuffle_down(result[tm], sn); + } + } + + // Threadgroup accumulation results + if (needs_tgp_reduction) { + threadgroup AccT* tgp_results = tgp_memory + sgN * (blockM + TM) + bm; + if (thrN == 0) { + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + tgp_results[tm] = result[tm]; + } + + threadgroup_barrier(mem_flags::mem_none); + + if (sgN == 0) { + MLX_MTL_PRAGMA_UNROLL + for (int sgn = 1; sgn < BN; sgn++) { + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + result[tm] += tgp_results[sgn * (blockM + TM) + tm]; + } + } + } + } + } + + // Write outputs + if (simdN == 0 && thrN == 0) { + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + if (kDoAxpby) { + out_vec[out_row + tm] = + static_cast(alpha) * static_cast(result[tm]) + + static_cast(beta) * bias[(out_row + tm) * bias_stride]; + } else { + out_vec[out_row + tm] = static_cast(result[tm]); + } + } + } + } +}; + +/////////////////////////////////////////////////////////////////////////////// +/// Vector matrix multiplication +/////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN, /* Thread cols (in elements) */ + const bool kDoAxpby, /* Do out = alpha * out + beta * bias */ + typename AccT = typename DefaultAccT::type> +struct GEMVTKernel { + using acc_type = AccT; + + MLX_MTL_CONST int threadsM = BM * SM; + MLX_MTL_CONST int threadsN = BN * SN; + + MLX_MTL_CONST int blockM = threadsM * TM; + MLX_MTL_CONST int blockN = threadsN * TN; + + static_assert(SM * SN == 32, "simdgroup can only have 32 threads"); + + // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up + // into blocks of (blockM, blockN) divided among threadgroups + // - Every thread works on a block of (TM, TN) + // - We assume each threadgroup has (threadsN, threadsM, 1) threads + // + // 1. A thread loads TN elements each from mat along TM contiguous rows + // and the corresponding scalar from the vector + // 2. The thread then accumulates its local result for the block + // 3. At the end, each thread has accumulated results over all blocks across + // the rows. These are then summed up across the threadgroup + // 4. Each threadgroup writes its accumulated BN * TN outputs + // + // Edge case handling: + // - The threadgroup with the largest tid has blocks that exceed the matrix + // * The blocks that start outside the matrix are never read (thread results + // remain zero) + // * The last thread that partially overlaps with the matrix is shifted + // inwards such that the thread block fits exactly in the matrix + + MLX_MTL_CONST short tgp_mem_size = BM > 1 ? BM*(blockN + TN) : 0; + MLX_MTL_CONST bool needs_tgp_reduction = BM > 1; + + static METAL_FUNC void run( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& marix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& bias_stride [[buffer(14)]], + threadgroup AccT* tgp_memory [[threadgroup(0)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + // Appease compiler + (void)lid; + + // Thread local accumulation results + AccT result[TN] = {0}; + T inter[TN]; + AccT v_coeff[TM]; + const int thrM = SN != 32 ? simd_lid / SN : 0; + const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid); + + const int sgM = BN != 1 ? (simd_gid / BN) : int(simd_gid); + const int sgN = BN != 1 ? (simd_gid % BN) : 0; + + const int simdM = SM * sgM; + const int simdN = SN * sgN; + + int cm = (simdM + thrM); + int cn = (simdN + thrN); + + int bm = cm * TM; + int bn = cn * TN; + + int out_col = tid.x * blockN + bn; + + constexpr const uniform loop_stride = make_uniform(blockM); + const uniform in_size = make_uniform(in_vec_size); + const uniform n_iter = in_size / loop_stride; + const uniform last_iter = loop_stride * n_iter; + const uniform leftover = in_size - last_iter; + + // Edgecase handling + if (out_col < out_vec_size) { + out_col = out_col + TN < out_vec_size ? out_col : out_vec_size - TN; + + // Per thread accumulation main loop + for (int i = 0; i < n_iter; ++i) { + // Adding a threadgroup_barrier improves performance slightly + // This is possibly it may help exploit cache better + threadgroup_barrier(mem_flags::mem_none); + + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + v_coeff[tm] = static_cast(in_vec[bm + tm]); + } + + MLX_MTL_PRAGMA_UNROLL + for (int tm = 0; tm < TM; tm++) { + auto vc = static_cast(v_coeff[tm]); + for (int tn = 0; tn < TN; tn++) { + inter[tn] = mat[size_t(bm + tm) * marix_ld + out_col + tn]; + } + for (int tn = 0; tn < TN; tn++) { + result[tn] += vc * inter[tn]; + } + } + + bm += blockM; + } + + if (leftover > 0) { + for (int tm = 0; tm < TM && bm + tm < in_vec_size; tm++) { + v_coeff[tm] = static_cast(in_vec[bm + tm]); + + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + inter[tn] = mat[size_t(bm + tm) * marix_ld + out_col + tn]; + } + + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + result[tn] += v_coeff[tm] * inter[tn]; + } + } + } + } + + // Simdgroup accumulations + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + MLX_MTL_PRAGMA_UNROLL + for (ushort sm = (SM / 2); sm >= 1; sm >>= 1) { + result[tn] += simd_shuffle_down(result[tn], SN * sm); + } + } + + // Threadgroup accumulation results + if (needs_tgp_reduction) { + threadgroup AccT* tgp_results = tgp_memory + sgM * (blockN + TN) + bn; + if (thrM == 0) { + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + tgp_results[tn] = result[tn]; + } + + threadgroup_barrier(mem_flags::mem_none); + + if (sgM == 0) { + MLX_MTL_PRAGMA_UNROLL + for (int sgm = 1; sgm < BM; sgm++) { + MLX_MTL_PRAGMA_UNROLL + for (int tn = 0; tn < TN; tn++) { + result[tn] += tgp_results[sgm * (blockN + TN) + tn]; + } + } + } + } + } + + // Threadgroup accumulation and writing out results + if (cm == 0 && out_col < out_vec_size) { + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < TN; j++) { + if (kDoAxpby) { + out_vec[out_col + j] = + static_cast(alpha) * static_cast(result[j]) + + static_cast(beta) * bias[(out_col + j) * bias_stride]; + } else { + out_vec[out_col + j] = static_cast(result[j]); + } + } + } + } +}; + +/////////////////////////////////////////////////////////////////////////////// +/// Matrix vector multiplication kernel +/////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN, /* Thread cols (in elements) */ + const bool kDoNCBatch, /* Batch ndim > 1 */ + const bool kDoAxpby> /* Do out = alpha * out + beta * bias */ +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& marix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& batch_ndim [[buffer(9)]], + const constant int* batch_shape [[buffer(10)]], + const constant int64_t* vector_batch_stride [[buffer(11)]], + const constant int64_t* matrix_batch_stride [[buffer(12)]], + const constant int64_t* bias_batch_stride [[buffer(13)]], + const constant int& bias_stride [[buffer(14)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + using gemv_kernel = GEMVKernel; + threadgroup typename gemv_kernel::acc_type tgp_memory + [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size]; + + // Update batch offsets + if (kDoNCBatch) { + in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim); + mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim); + + if (kDoAxpby) { + bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim); + } + + } else { + in_vec += tid.z * vector_batch_stride[0]; + mat += tid.z * matrix_batch_stride[0]; + + if (kDoAxpby) { + bias += tid.z * bias_batch_stride[0]; + } + } + + out_vec += tid.z * out_vec_size; + + gemv_kernel::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + marix_ld, + alpha, + beta, + bias_stride, + gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN> /* Thread cols (in elements) */ +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_gather( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& marix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& batch_ndim [[buffer(9)]], + const constant int* batch_shape [[buffer(10)]], + const constant int64_t* index_batch_strides [[buffer(11)]], + const constant int& vector_batch_ndim [[buffer(12)]], + const constant int* vector_batch_shape [[buffer(13)]], + const constant int64_t* vector_batch_stride [[buffer(14)]], + const constant int& matrix_batch_ndim [[buffer(15)]], + const constant int* matrix_batch_shape [[buffer(16)]], + const constant int64_t* matrix_batch_stride [[buffer(17)]], + const constant uint32_t* vec_indices [[buffer(18)]], + const constant uint32_t* mat_indices [[buffer(19)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + using gemv_kernel = GEMVKernel; + threadgroup typename gemv_kernel::acc_type tgp_memory + [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size]; + + uint32_t indx_vec; + uint32_t indx_mat; + + // Update batch offsets + if (batch_ndim > 1) { + const constant auto* veci_bstrides = index_batch_strides; + const constant auto* mati_bstrides = index_batch_strides + batch_ndim; + + ulong2 batch_offsets = elem_to_loc_broadcast( + tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim); + + indx_vec = vec_indices[batch_offsets.x]; + indx_mat = mat_indices[batch_offsets.y]; + + } else { + indx_vec = vec_indices[index_batch_strides[0] * tid.z]; + indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z]; + } + + if (vector_batch_ndim > 1) { + in_vec += elem_to_loc( + indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim); + } else { + in_vec += indx_vec * vector_batch_stride[0]; + } + + if (matrix_batch_ndim > 1) { + mat += elem_to_loc( + indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim); + } else { + mat += indx_mat * matrix_batch_stride[0]; + } + + out_vec += tid.z * out_vec_size; + + gemv_kernel::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + marix_ld, + alpha, + beta, + batch_ndim, // Not used + gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory, + tid, + lid, + simd_gid, + simd_lid); +} + +/////////////////////////////////////////////////////////////////////////////// +/// Vector matrix multiplication kernel +/////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN, /* Thread cols (in elements) */ + const bool kDoNCBatch, /* Batch ndim > 1 */ + const bool kDoAxpby> /* Do out = alpha * out + beta * bias */ +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& marix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& batch_ndim [[buffer(9)]], + const constant int* batch_shape [[buffer(10)]], + const constant int64_t* vector_batch_stride [[buffer(11)]], + const constant int64_t* matrix_batch_stride [[buffer(12)]], + const constant int64_t* bias_batch_stride [[buffer(13)]], + const constant int& bias_stride [[buffer(14)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + using gemv_kernel = GEMVTKernel; + threadgroup typename gemv_kernel::acc_type tgp_memory + [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size]; + + // Update batch offsets + if (kDoNCBatch) { + in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim); + mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim); + + if (kDoAxpby) { + bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim); + } + + } else { + in_vec += tid.z * vector_batch_stride[0]; + mat += tid.z * matrix_batch_stride[0]; + + if (kDoAxpby) { + bias += tid.z * bias_batch_stride[0]; + } + } + + out_vec += tid.z * out_vec_size; + + gemv_kernel::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + marix_ld, + alpha, + beta, + bias_stride, + gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + const int BM, /* Threadgroup rows (in simdgroups) */ + const int BN, /* Threadgroup cols (in simdgroups) */ + const int SM, /* Simdgroup rows (in threads) */ + const int SN, /* Simdgroup cols (in threads) */ + const int TM, /* Thread rows (in elements) */ + const int TN> /* Thread cols (in elements) */ +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_gather( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& marix_ld [[buffer(6)]], + const constant float& alpha [[buffer(7)]], + const constant float& beta [[buffer(8)]], + const constant int& batch_ndim [[buffer(9)]], + const constant int* batch_shape [[buffer(10)]], + const constant int64_t* index_batch_strides [[buffer(11)]], + const constant int& vector_batch_ndim [[buffer(12)]], + const constant int* vector_batch_shape [[buffer(13)]], + const constant int64_t* vector_batch_stride [[buffer(14)]], + const constant int& matrix_batch_ndim [[buffer(15)]], + const constant int* matrix_batch_shape [[buffer(16)]], + const constant int64_t* matrix_batch_stride [[buffer(17)]], + const constant uint32_t* vec_indices [[buffer(18)]], + const constant uint32_t* mat_indices [[buffer(19)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + using gemv_kernel = GEMVTKernel; + threadgroup typename gemv_kernel::acc_type tgp_memory + [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size]; + + uint32_t indx_vec; + uint32_t indx_mat; + + // Update batch offsets + if (batch_ndim > 1) { + const constant auto* veci_bstrides = index_batch_strides; + const constant auto* mati_bstrides = index_batch_strides + batch_ndim; + + ulong2 batch_offsets = elem_to_loc_broadcast( + tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim); + + indx_vec = vec_indices[batch_offsets.x]; + indx_mat = mat_indices[batch_offsets.y]; + + } else { + indx_vec = vec_indices[index_batch_strides[0] * tid.z]; + indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z]; + } + + if (vector_batch_ndim > 1) { + in_vec += elem_to_loc( + indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim); + } else { + in_vec += indx_vec * vector_batch_stride[0]; + } + + if (matrix_batch_ndim > 1) { + mat += elem_to_loc( + indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim); + } else { + mat += indx_mat * matrix_batch_stride[0]; + } + + out_vec += tid.z * out_vec_size; + + gemv_kernel::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + marix_ld, + alpha, + beta, + batch_ndim, // Not used, + gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory, + tid, + lid, + simd_gid, + simd_lid); +} + +/////////////////////////////////////////////////////////////////////////////// +/// Multi-vector matrix-vector multiplication (wide gemv) +/////////////////////////////////////////////////////////////////////////////// + +constant bool gemv_wide_has_batch [[function_constant(0)]]; +constant bool gemv_wide_do_axpby [[function_constant(1)]]; + +// out[M, N] = x[M, K] @ mat[N, K]^T for small M: each threadgroup streams a +// block of matrix rows once and applies it to vecs_per_tg input vectors, so +// the matrix is read ceil(M / vecs_per_tg) times instead of once per padded +// GEMM tile. k_lanes lanes reduce K for one row, 32 / k_lanes rows per +// simdgroup, k_lanes / 8 simdgroups per threadgroup. +template < + typename T, + const int vecs_per_tg, + const int k_lanes, + typename AccT = typename DefaultAccT::type> +struct GemvWide { + static constexpr constant int unroll = 8; + static constexpr constant int num_simdgroups = k_lanes / 8; + + static METAL_FUNC void run( + const device T* mat, + const device T* in_vec, + const device T* bias, + device T* out_vec, + int in_vec_size, + int out_vec_size, + int M, + int matrix_ld, + int vector_ld, + float alpha, + float beta, + int bias_ld, + int bias_fd, + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int rows_per_simdgroup = 32 / k_lanes; + constexpr int rows_per_tg = rows_per_simdgroup * num_simdgroups; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = + simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = + tid.y * rows_per_tg + rows_per_simdgroup * simd_gid + sg_row; + + // Clamped tail rows/vectors read valid memory; the guarded writes + // below never store them. + const int row = min(out_row, out_vec_size - 1); + const device T* wrow = mat + int64_t(row) * matrix_ld; + const device vec* w4 = (const device vec*)wrow; + const int n_v4 = in_vec_size / 4; + const int n_main = n_v4 - n_v4 % (k_lanes * unroll); + + // Vector chunks beyond grid.x round-robin onto the same threadgroups: + // re-walking the row block hits cache where an extra grid column would + // re-stream it from DRAM. + const device vec* x4 = (const device vec*)in_vec; + const int x_ld4 = vector_ld / 4; + const int n_chunks = (M + vecs_per_tg - 1) / vecs_per_tg; + for (int chunk = tid.x; chunk < n_chunks; chunk += tgpg.x) { + const int vec0 = chunk * vecs_per_tg; + + int x_off4[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + x_off4[v] = min(vec0 + v, M - 1) * x_ld4; + } + + AccT result[vecs_per_tg] = {0}; + + // Adjacent lanes read adjacent blocks so every transaction lands on + // whole cachelines; the unroll keeps loads in flight on short rows. + for (int base = 0; base < n_main; base += k_lanes * unroll) { + vec wf[unroll]; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + wf[i] = vec(w4[base + i * k_lanes + k_lane]); + } + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + AccT acc = 0; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + acc += + dot(wf[i], + vec(x4[x_off4[v] + base + i * k_lanes + k_lane])); + } + result[v] += acc; + } + } + for (int idx = n_main + k_lane; idx < n_v4; idx += k_lanes) { + const vec wf = vec(w4[idx]); + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + result[v] += dot(wf, vec(x4[x_off4[v] + idx])); + } + } + + // The halving shuffles reduce each row's k_lanes while rows sharing + // the simdgroup stay separate. + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + MLX_MTL_PRAGMA_UNROLL + for (ushort off = k_lanes / 2; off >= 1; off >>= 1) { + result[v] += simd_shuffle_down(result[v], off); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + int64_t out_idx = int64_t(vec0 + v) * out_vec_size + out_row; + if (gemv_wide_do_axpby) { + AccT bias_val = static_cast( + bias[int64_t(vec0 + v) * bias_ld + out_row * bias_fd]); + out_vec[out_idx] = + static_cast(alpha * result[v] + beta * bias_val); + } else { + out_vec[out_idx] = static_cast(result[v]); + } + } + } + } + } + } +}; + +template +[[kernel]] void gemv_wide( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2), function_constant(gemv_wide_do_axpby)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant float& alpha + [[buffer(9), function_constant(gemv_wide_do_axpby)]], + const constant float& beta + [[buffer(10), function_constant(gemv_wide_do_axpby)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* vector_batch_stride [[buffer(13)]], + const constant int64_t* matrix_batch_stride [[buffer(14)]], + const constant int64_t* bias_batch_stride + [[buffer(15), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_ld + [[buffer(16), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_fd + [[buffer(17), function_constant(gemv_wide_do_axpby)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (gemv_wide_has_batch) { + in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim); + mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim); + + if (gemv_wide_do_axpby) { + bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim); + } + } else { + in_vec += tid.z * vector_batch_stride[0]; + mat += tid.z * matrix_batch_stride[0]; + + if (gemv_wide_do_axpby) { + bias += tid.z * bias_batch_stride[0]; + } + } + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + alpha, + beta, + bias_ld, + bias_fd, + tid, + tgpg, + simd_gid, + simd_lid); +} + +template +[[kernel]] void gemv_wide_gather( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* index_batch_strides [[buffer(13)]], + const constant int& vector_batch_ndim [[buffer(14)]], + const constant int* vector_batch_shape [[buffer(15)]], + const constant int64_t* vector_batch_stride [[buffer(16)]], + const constant int& matrix_batch_ndim [[buffer(17)]], + const constant int* matrix_batch_shape [[buffer(18)]], + const constant int64_t* matrix_batch_stride [[buffer(19)]], + const constant uint32_t* vec_indices [[buffer(20)]], + const constant uint32_t* mat_indices [[buffer(21)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + uint32_t indx_vec; + uint32_t indx_mat; + + // Update batch offsets + if (batch_ndim > 1) { + const constant auto* veci_bstrides = index_batch_strides; + const constant auto* mati_bstrides = index_batch_strides + batch_ndim; + + ulong2 batch_offsets = elem_to_loc_broadcast( + tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim); + + indx_vec = vec_indices[batch_offsets.x]; + indx_mat = mat_indices[batch_offsets.y]; + + } else { + indx_vec = vec_indices[index_batch_strides[0] * tid.z]; + indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z]; + } + + if (vector_batch_ndim > 1) { + in_vec += elem_to_loc( + indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim); + } else { + in_vec += indx_vec * vector_batch_stride[0]; + } + + if (matrix_batch_ndim > 1) { + mat += elem_to_loc( + indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim); + } else { + mat += indx_mat * matrix_batch_stride[0]; + } + + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + 1.0f, + 0.0f, + 0, + 0, + tid, + tgpg, + simd_gid, + simd_lid); +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/gemv_masked.cpp b/Source/Cmlx/mlx-generated/gemv_masked.cpp index 36d22a1a5..8e2faebe0 100644 --- a/Source/Cmlx/mlx-generated/gemv_masked.cpp +++ b/Source/Cmlx/mlx-generated/gemv_masked.cpp @@ -69,7 +69,7 @@ using namespace metal; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -89,7 +89,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal index 3cd95c52f..ba0268dc9 100644 --- a/Source/Cmlx/mlx-generated/metal/arg_reduce.metal +++ b/Source/Cmlx/mlx-generated/metal/arg_reduce.metal @@ -16,7 +16,7 @@ template struct ArgMin { static constexpr constant U init = Limits::max; - IndexValPair reduce(IndexValPair best, IndexValPair current) { + IndexValPair reduce(IndexValPair best, IndexValPair current) thread { if (best.val > current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -27,7 +27,7 @@ struct ArgMin { template IndexValPair - reduce_many(IndexValPair best, thread U* vals, uint32_t offset) { + reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { if (vals[i] < best.val) { best.val = vals[i]; @@ -42,7 +42,7 @@ template struct ArgMax { static constexpr constant U init = Limits::min; - IndexValPair reduce(IndexValPair best, IndexValPair current) { + IndexValPair reduce(IndexValPair best, IndexValPair current) thread { if (best.val < current.val || (best.val == current.val && best.index > current.index)) { return current; @@ -53,7 +53,7 @@ struct ArgMax { template IndexValPair - reduce_many(IndexValPair best, thread U* vals, uint32_t offset) { + reduce_many(IndexValPair best, thread U* vals, uint32_t offset) thread { for (int i = 0; i < N; i++) { if (vals[i] > best.val) { best.val = vals[i]; diff --git a/Source/Cmlx/mlx-generated/metal/atomic.h b/Source/Cmlx/mlx-generated/metal/atomic.h index 93952c2cf..3835fff02 100644 --- a/Source/Cmlx/mlx-generated/metal/atomic.h +++ b/Source/Cmlx/mlx-generated/metal/atomic.h @@ -95,6 +95,19 @@ METAL_FUNC void mlx_atomic_fetch_mul_explicit( T expected = mlx_atomic_load_explicit(object, offset); while (!mlx_atomic_compare_exchange_weak_explicit( object, &expected, val * expected, offset)) { + // Workaround: Metal's atomic_compare_exchange_weak_explicit does + // not perform bitwise comparison as required by the C++ atomics spec. + // The compiler lowers the success check to `fcmp fast ueq` under + // no-nans-fp-math, which evaluates to false when either operand is NaN - + // even when the bit patterns are identical. With NaN in memory the CAS + // can never succeed, so the loop spins. Bail out instead: memory is + // already NaN and that is the correct reduction result regardless of + // this thread's update. + if constexpr (metal::is_floating_point_v) { + if (isnan(expected)) { + break; + } + } } } @@ -159,7 +172,7 @@ union uint_or_packed { template struct mlx_atomic_update_helper { - uint operator()(uint_or_packed init, T update, size_t elem_offset) { + uint operator()(uint_or_packed init, T update, size_t elem_offset) thread { Op op; init.val[elem_offset] = op(update, init.val[elem_offset]); return init.bits; @@ -196,7 +209,7 @@ struct __None { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { #pragma unused(b) return a; } @@ -210,7 +223,7 @@ struct __Add { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a + b; } }; @@ -222,7 +235,7 @@ struct __Mul { return b != 0; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a * b; } }; @@ -233,7 +246,7 @@ struct __Max { return a > b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return max(a, b); } }; @@ -244,7 +257,7 @@ struct __Min { return a < b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return min(a, b); } }; diff --git a/Source/Cmlx/mlx-generated/metal/binary_ops.h b/Source/Cmlx/mlx-generated/metal/binary_ops.h index 4e3d881fe..37650c7d9 100644 --- a/Source/Cmlx/mlx-generated/metal/binary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/binary_ops.h @@ -9,33 +9,40 @@ constant mlx::os_log logger("mlx", "binary_ops"); struct Add { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x + y; } }; struct FloorDivide { template - T operator()(T x, T y) { + metal::enable_if_t & !metal::is_signed_v, T> + operator()(T x, T y) thread { return x / y; } - template <> - float operator()(float x, float y) { - return trunc(x / y); + template + metal::enable_if_t & metal::is_signed_v, T> + operator()(T x, T y) thread { + auto q = x / y; + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + return q; } - template <> - half operator()(half x, half y) { - return trunc(x / y); + template + metal::enable_if_t, T> operator()(T x, T y) thread { + return floor(x / y); } template <> - bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { - return trunc(x / y); + complex64_t operator()(complex64_t x, complex64_t y) thread { + // Complex is not supported, simply make compiler happy. + return x / y; } }; struct Divide { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x / y; } }; @@ -43,12 +50,12 @@ struct Divide { struct Remainder { template metal::enable_if_t & !metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { return x % y; } template metal::enable_if_t & metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { auto r = x % y; if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -56,7 +63,7 @@ struct Remainder { return r; } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { T r = fmod(x, y); if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -64,25 +71,25 @@ struct Remainder { return r; } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { return x % y; } }; struct Equal { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y; } }; struct NaNEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y || (metal::isnan(x) && metal::isnan(y)); } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x == y || (metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) && metal::isnan(y.imag)) || @@ -93,35 +100,35 @@ struct NaNEqual { struct Greater { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x > y; } }; struct GreaterEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x >= y; } }; struct Less { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x < y; } }; struct LessEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x <= y; } }; struct LogAddExp { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { if (metal::isnan(x) || metal::isnan(y)) { return metal::numeric_limits::quiet_NaN(); } @@ -133,7 +140,7 @@ struct LogAddExp { : (maxval + log1p(metal::exp(minval - maxval))); }; - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) || metal::isnan(y.imag)) { return metal::numeric_limits::quiet_NaN(); @@ -154,12 +161,12 @@ struct LogAddExp { struct Maximum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::max(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -167,7 +174,7 @@ struct Maximum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -177,12 +184,12 @@ struct Maximum { struct Minimum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::min(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -190,7 +197,7 @@ struct Minimum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -200,30 +207,32 @@ struct Minimum { struct Multiply { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x * y; } }; struct NotEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x != y; } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x.real != y.real || x.imag != y.imag; } }; struct Power { template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { return metal::pow(base, exp); } template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { T res = 1; // Undefined to raise integer to negative power if (exp < 0) { @@ -243,7 +252,7 @@ struct Power { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (x.real == 0 && x.imag == 0) { if (metal::isnan(y.real) || metal::isnan(y.imag)) { auto nan = metal::numeric_limits::quiet_NaN(); @@ -261,70 +270,70 @@ struct Power { struct Subtract { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x - y; } }; struct LogicalAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x && y; }; }; struct LogicalOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x || y; }; }; struct BitwiseAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x & y; }; }; struct BitwiseOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x | y; }; }; struct BitwiseXor { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x ^ y; }; }; struct LeftShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x << y; }; }; struct RightShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x >> y; }; }; struct ArcTan2 { template - T operator()(T y, T x) { + T operator()(T y, T x) thread { return metal::precise::atan2(y, x); } }; struct DivMod { template - metal::array operator()(T x, T y) { + metal::array operator()(T x, T y) thread { return {FloorDivide{}(x, y), Remainder{}(x, y)}; }; }; diff --git a/Source/Cmlx/mlx-generated/metal/complex.h b/Source/Cmlx/mlx-generated/metal/complex.h index 6e391483d..629b6289b 100644 --- a/Source/Cmlx/mlx-generated/metal/complex.h +++ b/Source/Cmlx/mlx-generated/metal/complex.h @@ -4,163 +4,286 @@ #include +#include "bf16.h" + using namespace metal; -struct complex64_t; +template +struct complex_t; template -static constexpr constant bool can_convert_to_complex64 = - !is_same_v && is_convertible_v; +static constexpr constant bool is_complex_v = false; + +template +static constexpr constant bool is_complex_v> = true; + +// Metal accepts explicit bfloat casts that is_convertible_v reports as false. +template +static constexpr constant bool is_lane_convertible_v = + is_convertible_v || + (is_same_v && is_convertible_v) || + (is_same_v && is_convertible_v); template -static constexpr constant bool can_convert_from_complex64 = - !is_same_v && - (is_convertible_v || is_convertible_v); +struct complex_t { + using value_type = T; -struct complex64_t { - float real; - float imag; + T real; + T imag; // Constructors - constexpr complex64_t(float real, float imag) : real(real), imag(imag) {}; - constexpr complex64_t() : real(0), imag(0) {}; - constexpr complex64_t() threadgroup : real(0), imag(0) {}; + constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {}; + constexpr complex_t() thread : real(0), imag(0) {}; + constexpr complex_t() threadgroup : real(0), imag(0) {}; + + // Conversions from scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) thread : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) threadgroup : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) device : real(static_cast(x)), + imag(static_cast(0)) {} + + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) constant : real(static_cast(x)), + imag(static_cast(0)) {} - // Conversions to complex64_t + // Conversions between complex types template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) thread : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) thread : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) threadgroup : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) threadgroup + : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) device : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) device : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) constant : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + // Conversions to and from two-lane vectors (the FFT lane representation) + constexpr complex_t(vec v) thread : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) threadgroup : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) device : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) constant : real(v.x), imag(v.y) {}; + + constexpr operator vec() const thread { + return vec(real, imag); + } + + constexpr operator vec() const threadgroup { + return vec(real, imag); + } + + constexpr operator vec() const device { + return vec(real, imag); + } + + constexpr operator vec() const constant { + return vec(real, imag); + } - // Conversions from complex64_t + // Conversions to scalar types template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const thread { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const thread { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const threadgroup { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const threadgroup { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const device { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const device { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const constant { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const constant { + return static_cast(real); } }; -constexpr complex64_t operator-(complex64_t x) { +using complex32_t = complex_t; +using complex64_t = complex_t; + +static_assert(sizeof(complex32_t) == 2 * sizeof(half)); +static_assert(sizeof(complex64_t) == 2 * sizeof(float)); +static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); + +template +constexpr complex_t operator-(complex_t x) { return {-x.real, -x.imag}; } -constexpr bool operator>=(complex64_t a, complex64_t b) { +template +constexpr bool operator>=(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); } -constexpr bool operator>(complex64_t a, complex64_t b) { +template +constexpr bool operator>(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); } -constexpr bool operator<=(complex64_t a, complex64_t b) { +template +constexpr bool operator<=(complex_t a, complex_t b) { return operator>=(b, a); } -constexpr bool operator<(complex64_t a, complex64_t b) { +template +constexpr bool operator<(complex_t a, complex_t b) { return operator>(b, a); } -constexpr bool operator==(complex64_t a, complex64_t b) { +template +constexpr bool operator==(complex_t a, complex_t b) { return a.real == b.real && a.imag == b.imag; } -constexpr complex64_t operator+(complex64_t a, complex64_t b) { +template +constexpr complex_t operator+(complex_t a, complex_t b) { return {a.real + b.real, a.imag + b.imag}; } -constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) { +template +constexpr thread complex_t& operator+=( + thread complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr threadgroup complex64_t& operator+=( - threadgroup complex64_t& a, - complex64_t b) { +template +constexpr threadgroup complex_t& operator+=( + threadgroup complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) { +template +constexpr device complex_t& operator+=( + device complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr complex64_t operator+(float a, complex64_t b) { - return {a + b.real, b.imag}; +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(U a, complex_t b) { + return {static_cast(a) + b.real, b.imag}; } -constexpr complex64_t operator+(complex64_t a, float b) { - return {a.real + b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(complex_t a, U b) { + return {a.real + static_cast(b), a.imag}; } -constexpr complex64_t operator-(complex64_t a, complex64_t b) { +template +constexpr complex_t operator-(complex_t a, complex_t b) { return {a.real - b.real, a.imag - b.imag}; } -constexpr complex64_t operator-(float a, complex64_t b) { - return {a - b.real, -b.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(U a, complex_t b) { + return {static_cast(a) - b.real, -b.imag}; } -constexpr complex64_t operator-(complex64_t a, float b) { - return {a.real - b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(complex_t a, U b) { + return {a.real - static_cast(b), a.imag}; } -constexpr complex64_t operator*(complex64_t a, complex64_t b) { +template +constexpr complex_t operator*(complex_t a, complex_t b) { return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; } -constexpr complex64_t operator/(complex64_t a, complex64_t b) { +template +constexpr complex_t operator/(complex_t a, complex_t b) { auto denom = b.real * b.real + b.imag * b.imag; auto x = a.real * b.real + a.imag * b.imag; auto y = a.imag * b.real - a.real * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator/(float a, complex64_t b) { +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator/(U a, complex_t b) { + auto scalar = static_cast(a); auto denom = b.real * b.real + b.imag * b.imag; - auto x = a * b.real; - auto y = -a * b.imag; + auto x = scalar * b.real; + auto y = -scalar * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator%(complex64_t a, complex64_t b) { +template +constexpr complex_t operator%(complex_t a, complex_t b) { auto real = a.real - (b.real * static_cast(a.real / b.real)); auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); if (real != 0 && (real < 0 != b.real < 0)) { @@ -171,3 +294,10 @@ constexpr complex64_t operator%(complex64_t a, complex64_t b) { } return {real, imag}; } + +static_assert( + (complex_t{1.0h, 2.0h} * complex_t{3.0h, 4.0h}).real == -5.0h); +static_assert( + (complex_t{bfloat16_t(1.0f), bfloat16_t(2.0f)} * + complex_t{bfloat16_t(3.0f), bfloat16_t(4.0f)}) + .real == bfloat16_t(-5.0f)); diff --git a/Source/Cmlx/mlx-generated/metal/conv.metal b/Source/Cmlx/mlx-generated/metal/conv.metal index e6cc127c0..74027df73 100644 --- a/Source/Cmlx/mlx-generated/metal/conv.metal +++ b/Source/Cmlx/mlx-generated/metal/conv.metal @@ -20,6 +20,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -39,8 +40,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -83,6 +85,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -103,8 +106,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -150,6 +154,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); \ template \ [[host_name("naive_unfold_transpose_nd_" #name "_" #n)]] [[kernel]] void \ @@ -157,6 +162,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); #define instantiate_naive_unfold_nd_dims(name, itype) \ @@ -206,7 +212,7 @@ template threadgroup T ins[TGH * TGW * TGC]; - const int n_tgblocks_h = params.oS[0] / th; + const int n_tgblocks_h = (params.oS[0] + th - 1) / th; const int n = tid.z / n_tgblocks_h; const int tghid = tid.z % n_tgblocks_h; const int oh = tghid * th + lid.z; @@ -277,6 +283,10 @@ template } threadgroup_barrier(mem_flags::mem_none); + if (oh >= params.oS[0] || ow >= params.oS[1]) { + return; + } + out += n * params.out_strides[0] + oh * params.out_strides[1] + ow * params.out_strides[2]; out[c] = static_cast(o); diff --git a/Source/Cmlx/mlx-generated/metal/copy.h b/Source/Cmlx/mlx-generated/metal/copy.h index cf22347ee..95ed69b76 100644 --- a/Source/Cmlx/mlx-generated/metal/copy.h +++ b/Source/Cmlx/mlx-generated/metal/copy.h @@ -9,11 +9,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[0]); + dst[index + i] = cast_to(src[0]); } } } @@ -27,11 +27,11 @@ template ::n> index *= N; if (N > 1 && index + N > size) { for (int i = 0; index + i < size; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } else { for (int i = 0; i < N; ++i) { - dst[index + i] = static_cast(src[index + i]); + dst[index + i] = cast_to(src[index + i]); } } } @@ -46,11 +46,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[0]); + dst[offset + i] = cast_to(src[0]); } } } @@ -65,11 +65,11 @@ template ::n> int64_t offset = N * (index.x + grid_dim.x * int64_t(index.y)); if (N > 1 && offset + N > size) { for (int i = 0; offset + i < size; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } else { for (int i = 0; i < N; ++i) { - dst[offset + i] = static_cast(src[offset + i]); + dst[offset + i] = cast_to(src[offset + i]); } } } @@ -81,7 +81,7 @@ template constant const int64_t& src_stride [[buffer(3)]], uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); - dst[index] = static_cast(src[src_idx]); + dst[index] = cast_to(src[src_idx]); } template @@ -93,7 +93,7 @@ template uint2 grid_dim [[threads_per_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * index.y; - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -106,7 +106,7 @@ template auto src_idx = elem_to_loc_3(index, src_strides); IdxT dst_idx = index.x + IdxT(grid_dim.x) * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -123,14 +123,14 @@ template if (N == 1) { IdxT dst_idx = index.x + grid_dim.x * (index.y + IdxT(grid_dim.y) * index.z); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); return; } auto xshape = src_shape[ndim - 1]; IdxT dst_idx = N * index.x + xshape * (index.y + IdxT(grid_dim.y) * index.z); auto src_xstride = src_strides[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[dst_idx + i] = static_cast(src[src_idx]); + dst[dst_idx + i] = cast_to(src[src_idx]); src_idx += src_xstride; } } @@ -144,7 +144,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -156,7 +156,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -168,7 +168,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx] = static_cast(src[src_idx]); + dst[dst_idx] = cast_to(src[src_idx]); } template @@ -187,14 +187,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = static_cast(src[idx.x]); + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } @@ -211,7 +211,7 @@ template uint index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_1(index, src_stride); auto dst_idx = elem_to_loc_1(index, dst_stride); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -225,7 +225,7 @@ template uint2 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_2(index, src_strides); auto dst_idx = elem_to_loc_2(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -239,7 +239,7 @@ template uint3 index [[thread_position_in_grid]]) { auto src_idx = elem_to_loc_3(index, src_strides); auto dst_idx = elem_to_loc_3(index, dst_strides); - dst[dst_idx + dst_offset] = src[src_idx + src_offset]; + dst[dst_idx + dst_offset] = cast_to(src[src_idx + src_offset]); } template @@ -262,14 +262,14 @@ template dst_strides, ndim); if (N == 1) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); return; } IdxT src_xstride = src_strides[ndim - 1]; IdxT dst_xstride = dst_strides[ndim - 1]; auto xshape = src_shape[ndim - 1]; for (int i = 0; i < N && (int(N * index.x) + i) < xshape; ++i) { - dst[idx.y] = src[idx.x]; + dst[idx.y] = cast_to(src[idx.x]); idx.x += src_xstride; idx.y += dst_xstride; } diff --git a/Source/Cmlx/mlx-generated/metal/dot.h b/Source/Cmlx/mlx-generated/metal/dot.h new file mode 100644 index 000000000..a48ddf1d0 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/dot.h @@ -0,0 +1,68 @@ +#pragma once + +#include + +#include "utils.h" + +using namespace metal; + +template < + typename T, + const int ITEMS_PER_THREAD, + const int TG_SIZE, + const uint SIMD_GROUPS> +[[kernel]] void dot_product( + const device T* a [[buffer(0)]], + const device T* b [[buffer(1)]], + device float* output [[buffer(2)]], + const constant int& n [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd_id [[simdgroup_index_in_threadgroup]], + uint tg_id [[threadgroup_position_in_grid]]) { + constexpr int VEC = 16 / sizeof(T); + int start = (tg_id * TG_SIZE + simd_id * 32) * ITEMS_PER_THREAD + lane * VEC; + + float4 c = 0.0f; + + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < ITEMS_PER_THREAD; i += VEC) { + int idx = start + i * ITEMS_PER_THREAD; + if (idx + VEC <= n) { + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < VEC; j += 4) { + c += float4(*reinterpret_cast*>( + a + idx + j)) * + float4(*reinterpret_cast*>( + b + idx + j)); + } + } else { + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < VEC; ++j) { + int nidx = idx + j; + if (nidx < n) { + c[j & 3] += float(a[nidx]) * float(b[nidx]); + } + } + } + } + + threadgroup float smem[SIMD_GROUPS]; + + float sum = c[0] + c[1] + c[2] + c[3]; + sum = simd_sum(sum); + + if (lane == 0) { + smem[simd_id] = sum; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (tid < SIMD_GROUPS) { + sum = smem[tid]; + sum = simd_sum(sum); + if (tid == 0) { + output[tg_id] = sum; + } + } +} diff --git a/Source/Cmlx/mlx-generated/metal/dot.metal b/Source/Cmlx/mlx-generated/metal/dot.metal new file mode 100644 index 000000000..4b4d84a96 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/dot.metal @@ -0,0 +1,18 @@ +#include + +#include "dot.h" + +#define instantiate_dot_product_kernel( \ + name, itype, items_per_thread, tg_size, simd_groups) \ + instantiate_kernel( \ + "dot_product_" #name "_it" #items_per_thread "_tg" #tg_size \ + "_sg" #simd_groups, \ + dot_product, \ + itype, \ + items_per_thread, \ + tg_size, \ + simd_groups) + +instantiate_dot_product_kernel(float32, float, 32, 512, 16); +instantiate_dot_product_kernel(float16, half, 32, 512, 16); +instantiate_dot_product_kernel(bfloat16, bfloat16_t, 32, 512, 16); diff --git a/Source/Cmlx/mlx-generated/metal/fft.h b/Source/Cmlx/mlx-generated/metal/fft.h index 4f18730b8..ddddcd465 100644 --- a/Source/Cmlx/mlx-generated/metal/fft.h +++ b/Source/Cmlx/mlx-generated/metal/fft.h @@ -16,7 +16,7 @@ using namespace metal; #define MAX_RADIX 13 // Reached when elems_per_thread_ = 6, max_radix = 13 -// and some threads have to do 3 radix 6s requiring 18 float2s. +// and some threads have to do 3 radix 6s requiring 18 complex values. #define MAX_OUTPUT_SIZE 18 // Specialize for a particular value of N at runtime @@ -45,18 +45,47 @@ STEEL_CONST int rader_5_steps_ [[function_constant(18)]]; STEEL_CONST int rader_4_steps_ [[function_constant(19)]]; STEEL_CONST int rader_3_steps_ [[function_constant(20)]]; STEEL_CONST int rader_2_steps_ [[function_constant(21)]]; +STEEL_CONST bool use_bluestein_twiddle_table_ [[function_constant(22)]]; + +// Generate the compact radix-8 Stockham twiddle table used by the fused +// 4096-point Bluestein convolution. The first p=1 stage needs no twiddles; +// the remaining stages contribute 8, 64, and 512 entries. +[[kernel]] void generate_bluestein_twiddles( + device float2* twiddles [[buffer(0)]], + uint index [[thread_position_in_grid]]) { + uint p; + uint k; + if (index < 8) { + p = 8; + k = index; + } else if (index < 72) { + p = 64; + k = index - 8; + } else { + p = 512; + k = index - 72; + } + twiddles[index] = get_twiddle(k, 8 * p); +} -// See "radix.h" for radix codelets -typedef void (*RadixFunc)(thread float2*, thread float2*); +// See "radix.h" for radix codelets. +template +using RadixFunc = void (*)(thread vec*, thread vec*); // Perform a single radix n butterfly with appropriate twiddles -template +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_butterfly( int i, int p, - thread float2* x, + int twiddle_offset, + thread vec* x, thread short* indices, - thread float2* y) { + thread vec* y, + const device vec* twiddles) { // i: the index in the overall DFT that we're processing. // p: the size of the DFTs we're merging at this step. // m: how many threads are working on this DFT. @@ -75,14 +104,19 @@ METAL_FUNC void radix_butterfly( // Apply twiddles if (p > 1) { - float2 twiddle_1 = get_twiddle(k, radix * p); - float2 twiddle = twiddle_1; - x[1] = complex_mul(x[1], twiddle); + vec twiddle_1; + if constexpr (use_twiddle_table) { + twiddle_1 = twiddles[twiddle_offset + k]; + } else { + twiddle_1 = get_twiddle(k, radix * p); + } + vec twiddle = twiddle_1; + x[1] = complex_mul(x[1], twiddle); STEEL_PRAGMA_UNROLL for (int t = 2; t < radix; t++) { - twiddle = complex_mul(twiddle, twiddle_1); - x[t] = complex_mul(x[t], twiddle); + twiddle = complex_mul(twiddle, twiddle_1); + x[t] = complex_mul(x[t], twiddle); } } @@ -96,17 +130,23 @@ METAL_FUNC void radix_butterfly( // Perform all the radix steps required for a // particular radix size n. -template +template < + typename T, + bool use_twiddle_table, + int radix, + RadixFunc radix_func> METAL_FUNC void radix_n_steps( int i, thread int* p, + thread int* twiddle_offset, int m, int n, int num_steps, - thread float2* inputs, + thread vec* inputs, thread short* indices, - thread float2* values, - threadgroup float2* buf) { + thread vec* values, + threadgroup vec* buf, + const device vec* twiddles) { int m_r = n / radix; // When combining different sized radices, we have to do // multiple butterflies in a single thread. @@ -126,8 +166,14 @@ METAL_FUNC void radix_n_steps( for (int r = 0; r < radix; r++) { inputs[r] = buf[index + r * m_r]; } - radix_butterfly( - index, *p, inputs, indices + t * radix, values + t * radix); + radix_butterfly( + index, + *p, + *twiddle_offset, + inputs, + indices + t * radix, + values + t * radix, + twiddles); } } @@ -146,20 +192,41 @@ METAL_FUNC void radix_n_steps( // Wait until all threads have written back to threadgroup mem threadgroup_barrier(mem_flags::mem_threadgroup); + if constexpr (use_twiddle_table) { + if (*p > 1) { + *twiddle_offset += *p; + } + } *p *= radix; } } -#define RADIX_STEP(radix, radix_func, num_steps) \ - radix_n_steps( \ - fft_idx, p, m, n, num_steps, inputs, indices, values, buf); - -template -METAL_FUNC void -perform_fft(int fft_idx, thread int* p, int m, int n, threadgroup float2* buf) { - float2 inputs[MAX_RADIX]; +#define RADIX_STEP(radix, radix_func, num_steps) \ + radix_n_steps>( \ + fft_idx, \ + p, \ + &twiddle_offset, \ + m, \ + n, \ + num_steps, \ + inputs, \ + indices, \ + values, \ + buf, \ + twiddles); + +template +METAL_FUNC void perform_fft( + int fft_idx, + thread int* p, + int m, + int n, + threadgroup vec* buf, + const device vec* twiddles = nullptr) { + vec inputs[MAX_RADIX]; short indices[MAX_OUTPUT_SIZE]; - float2 values[MAX_OUTPUT_SIZE]; + vec values[MAX_OUTPUT_SIZE]; + int twiddle_offset = 0; RADIX_STEP(2, radix2, rader ? rader_2_steps_ : radix_2_steps_); RADIX_STEP(3, radix3, rader ? rader_3_steps_ : radix_3_steps_); @@ -184,7 +251,8 @@ template constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { - threadgroup float2 shared_in[tg_mem_size]; + using scalar_T = typename FFTIOTypeTraits::scalar_T; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -208,9 +276,9 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -219,7 +287,8 @@ template [[kernel]] void rader_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* raders_b_q [[buffer(2)]], + const device vec::scalar_T, 2>* + raders_b_q [[buffer(2)]], const device short* raders_g_q [[buffer(3)]], const device short* raders_g_minus_q [[buffer(4)]], constant const int& n, @@ -227,9 +296,10 @@ template constant const int& rader_n, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Use Rader's algorithm to compute fast FFTs // when a prime factor `p` of `n` is greater than 13 but - // has `p - 1` Stockham decomposable into to prime factors <= 13. + // has `p - 1` Stockham decomposable into prime factors <= 13. // // E.g. n = 102 // = 2 * 3 * 17 @@ -250,7 +320,7 @@ template // // Rader's uses fewer operations than Bluestein's and so // is more accurate. It's also faster in most cases. - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -275,7 +345,7 @@ template int fft_idx = elem.z; int tg_idx = elem.y * n; - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; // rader_m = n / rader_n; int rader_m = rader_m_; @@ -287,10 +357,10 @@ template // 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 short x_0_index = metal::min(fft_idx * elems_per_thread_ / (rader_n - 1), rader_m - 1); - float2 x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; + vec x_0[2] = {buf[x_0_index], buf[x_0_index + 1]}; // Do the Rader permutation in shared memory - float2 temp[MAX_RADIX]; + vec temp[MAX_RADIX]; int max_index = n - rader_m - 1; for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); @@ -309,19 +379,31 @@ template // Rader FFT on x[rader_m:] int p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); // x_1 + ... + x_n is computed for us in the first FFT step so // we save it in the first rader_m indices of the array for later. int x_sum_index = metal::min(fft_idx, rader_m - 1); buf[x_sum_index] = buf[rader_m + x_sum_index * (rader_n - 1)]; - float2 inv = {1.0f, -1.0f}; + // The convolution is conjugated and scaled by 1/(rader_n - 1) exactly once + // each. Float lanes apply the scale after the inverse FFT, preserving the + // existing float code generation; reduced lanes apply it before, so the + // intermediates cannot overflow the narrow lane range. + scalar_T rader_inv_r = static_cast(1.0f / (rader_n - 1)); + vec convolution_factor = {1.0f, -1.0f}; + vec convolution_inv_factor = {rader_inv_r, -rader_inv_r}; + if constexpr (!metal::is_same_v) { + convolution_factor = convolution_inv_factor; + convolution_inv_factor = {1.0f, -1.0f}; + } + for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); short interleaved_index = index / rader_m + (index % rader_m) * (rader_n - 1); - temp[e] = complex_mul( + temp[e] = complex_mul( buf[rader_m + interleaved_index], raders_b_q[interleaved_index % (rader_n - 1)]); } @@ -330,25 +412,24 @@ template for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, max_index); - buf[rader_m + index] = temp[e] * inv; + buf[rader_m + index] = temp[e] * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); // Rader IFFT on x[rader_m:] p = 1; - perform_fft(fft_idx, &p, m, n - rader_m, buf + rader_m); - - float2 rader_inv_factor = {1.0f / (rader_n - 1), -1.0f / (rader_n - 1)}; + perform_fft( + fft_idx, &p, m, n - rader_m, buf + rader_m); for (int e = 0; e < elems_per_thread_; e++) { short index = metal::min(fft_idx * elems_per_thread_ + e, n - rader_m - 1); short diff_index = index / (rader_n - 1) - x_0_index; - temp[e] = buf[rader_m + index] * rader_inv_factor + x_0[diff_index]; + temp[e] = buf[rader_m + index] * convolution_inv_factor + x_0[diff_index]; } // Use the sum of elements that was computed in the first FFT - float2 x_sum = buf[x_0_index] + x_0[0]; + vec x_sum = buf[x_0_index] + x_0[0]; threadgroup_barrier(mem_flags::mem_threadgroup); @@ -365,7 +446,7 @@ template threadgroup_barrier(mem_flags::mem_threadgroup); p = rader_n; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write(); } @@ -374,13 +455,16 @@ template [[kernel]] void bluestein_fft( const device in_T* in [[buffer(0)]], device out_T* out [[buffer(1)]], - const device float2* w_q [[buffer(2)]], - const device float2* w_k [[buffer(3)]], + const device vec::scalar_T, 2>* w_q + [[buffer(2)]], + const device vec::scalar_T, 2>* w_k + [[buffer(3)]], constant const int& length, constant const int& n, constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Computes arbitrary length FFTs with Bluestein's algorithm // // In numpy: @@ -390,7 +474,7 @@ template // Where w_k and w_q are precomputed on CPU in high precision as: // w_k = np.exp(-1j * np.pi / n * (np.arange(-n + 1, n) ** 2)) // w_q = np.fft.fft(1/w_k[-n:]) - threadgroup float2 shared_in[tg_mem_size]; + threadgroup vec shared_in[tg_mem_size]; thread ReadWriter read_writer = ReadWriter( in, @@ -414,22 +498,50 @@ template int fft_idx = elem.z; // Thread index in DFT int m = grid.z; // Threads per DFT int tg_idx = elem.y * n; // Index of this DFT in threadgroup - threadgroup float2* buf = &shared_in[tg_idx]; + threadgroup vec* buf = &shared_in[tg_idx]; + const device vec* twiddles = w_q + n; // fft - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } - float2 inv = float2(1.0f, -1.0f); + vec convolution_factor = {1.0f, -1.0f}; + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize the convolution before the inverse FFT so the + // intermediates cannot overflow; the write path then only conjugates. + scalar_T inv_n = static_cast(1.0f / n); + convolution_factor = {inv_n, -inv_n}; + } for (int t = 0; t < elems_per_thread_; t++) { int index = fft_idx + t * m; - buf[index] = complex_mul(buf[index], w_q[index]) * inv; + buf[index] = + complex_mul(buf[index], w_q[index]) * convolution_factor; } threadgroup_barrier(mem_flags::mem_threadgroup); // ifft p = 1; - perform_fft(fft_idx, &p, m, n, buf); + if constexpr ( + tg_mem_size == 4096 && metal::is_same_v && + metal::is_same_v) { + if (use_bluestein_twiddle_table_) { + perform_fft(fft_idx, &p, m, n, buf, twiddles); + } else { + perform_fft(fft_idx, &p, m, n, buf); + } + } else { + perform_fft(fft_idx, &p, m, n, buf); + } read_writer.write_padded(length, w_k); } @@ -448,6 +560,7 @@ template < constant const int& batch_size, uint3 elem [[thread_position_in_grid]], uint3 grid [[threads_per_grid]]) { + using scalar_T = typename FFTIOTypeTraits::scalar_T; // Fast four step FFT implementation for powers of 2. int overall_n = n1 * n2; int n = step == 0 ? n1 : n2; @@ -457,8 +570,8 @@ template < int m = grid.z; int fft_idx = elem.z; - threadgroup float2 shared_in[tg_mem_size]; - threadgroup float2* buf = &shared_in[elem.y * n]; + threadgroup vec shared_in[tg_mem_size]; + threadgroup vec* buf = &shared_in[elem.y * n]; using read_writer_t = ReadWriter; read_writer_t read_writer = read_writer_t( @@ -480,7 +593,7 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); int p = 1; - perform_fft(fft_idx, &p, m, n, buf); + perform_fft(fft_idx, &p, m, n, buf); read_writer.write_strided(stride, overall_n); } diff --git a/Source/Cmlx/mlx-generated/metal/fft/radix.h b/Source/Cmlx/mlx-generated/metal/fft/radix.h index bd61eef6d..9d742eb33 100644 --- a/Source/Cmlx/mlx-generated/metal/fft/radix.h +++ b/Source/Cmlx/mlx-generated/metal/fft/radix.h @@ -16,49 +16,64 @@ them into (n-1)=6,10,12 codelets. */ #include #include -METAL_FUNC float2 complex_mul(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x); +// The codelets are templated over the scalar lane type so reduced-precision +// complex paths can reuse them; complex values are plain two-lane vectors. + +template +METAL_FUNC metal::vec complex_mul( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x}; } // Complex mul followed by conjugate -METAL_FUNC float2 complex_mul_conj(float2 a, float2 b) { - return float2(a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x); +template +METAL_FUNC metal::vec complex_mul_conj( + metal::vec a, + metal::vec b) { + return {a.x * b.x - a.y * b.y, -a.x * b.y - a.y * b.x}; } // Compute an FFT twiddle factor -METAL_FUNC float2 get_twiddle(int k, int p) { - float theta = -2.0f * k * M_PI_F / p; - - float2 twiddle = {metal::fast::cos(theta), metal::fast::sin(theta)}; - return twiddle; +template +METAL_FUNC metal::vec get_twiddle(int k, int p) { + // Derive phase evaluation precision from the scalar lane and Metal's pi + // constant. Reduced lanes currently promote to float for fast trig, then + // narrow the result to the lane type explicitly. + using phase_T = decltype(M_PI_F * T(0)); + phase_T theta = -phase_T(2) * phase_T(k) * phase_T(M_PI_F) / phase_T(p); + return {T(metal::fast::cos(theta)), T(metal::fast::sin(theta))}; } -METAL_FUNC void radix2(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix2(thread metal::vec* x, thread metal::vec* y) { y[0] = x[0] + x[1]; y[1] = x[0] - x[1]; } -METAL_FUNC void radix3(thread float2* x, thread float2* y) { - float pi_2_3 = -0.8660254037844387; +template +METAL_FUNC void radix3(thread metal::vec* x, thread metal::vec* y) { + T pi_2_3 = -0.8660254037844387; - float2 a_1 = x[1] + x[2]; - float2 a_2 = x[1] - x[2]; + metal::vec a_1 = x[1] + x[2]; + metal::vec a_2 = x[1] - x[2]; y[0] = x[0] + a_1; - float2 b_1 = x[0] - 0.5 * a_1; - float2 b_2 = pi_2_3 * a_2; + metal::vec b_1 = x[0] - 0.5 * a_1; + metal::vec b_2 = pi_2_3 * a_2; - float2 b_2_j = {-b_2.y, b_2.x}; + metal::vec b_2_j = {-b_2.y, b_2.x}; y[1] = b_1 + b_2_j; y[2] = b_1 - b_2_j; } -METAL_FUNC void radix4(thread float2* x, thread float2* y) { - float2 z_0 = x[0] + x[2]; - float2 z_1 = x[0] - x[2]; - float2 z_2 = x[1] + x[3]; - float2 z_3 = x[1] - x[3]; - float2 z_3_i = {z_3.y, -z_3.x}; +template +METAL_FUNC void radix4(thread metal::vec* x, thread metal::vec* y) { + metal::vec z_0 = x[0] + x[2]; + metal::vec z_1 = x[0] - x[2]; + metal::vec z_2 = x[1] + x[3]; + metal::vec z_3 = x[1] - x[3]; + metal::vec z_3_i = {z_3.y, -z_3.x}; y[0] = z_0 + z_2; y[1] = z_1 + z_3_i; @@ -66,25 +81,26 @@ METAL_FUNC void radix4(thread float2* x, thread float2* y) { y[3] = z_1 - z_3_i; } -METAL_FUNC void radix5(thread float2* x, thread float2* y) { - float2 root_5_4 = 0.5590169943749475; - float2 sin_2pi_5 = 0.9510565162951535; - float2 sin_1pi_5 = 0.5877852522924731; - - float2 a_1 = x[1] + x[4]; - float2 a_2 = x[2] + x[3]; - float2 a_3 = x[1] - x[4]; - float2 a_4 = x[2] - x[3]; - - float2 a_5 = a_1 + a_2; - float2 a_6 = root_5_4 * (a_1 - a_2); - float2 a_7 = x[0] - a_5 / 4; - float2 a_8 = a_7 + a_6; - float2 a_9 = a_7 - a_6; - float2 a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; - float2 a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; - float2 a_10_j = {a_10.y, -a_10.x}; - float2 a_11_j = {a_11.y, -a_11.x}; +template +METAL_FUNC void radix5(thread metal::vec* x, thread metal::vec* y) { + T root_5_4 = 0.5590169943749475; + T sin_2pi_5 = 0.9510565162951535; + T sin_1pi_5 = 0.5877852522924731; + + metal::vec a_1 = x[1] + x[4]; + metal::vec a_2 = x[2] + x[3]; + metal::vec a_3 = x[1] - x[4]; + metal::vec a_4 = x[2] - x[3]; + + metal::vec a_5 = a_1 + a_2; + metal::vec a_6 = root_5_4 * (a_1 - a_2); + metal::vec a_7 = x[0] - a_5 / 4; + metal::vec a_8 = a_7 + a_6; + metal::vec a_9 = a_7 - a_6; + metal::vec a_10 = sin_2pi_5 * a_3 + sin_1pi_5 * a_4; + metal::vec a_11 = sin_1pi_5 * a_3 - sin_2pi_5 * a_4; + metal::vec a_10_j = {a_10.y, -a_10.x}; + metal::vec a_11_j = {a_11.y, -a_11.x}; y[0] = x[0] + a_5; y[1] = a_8 + a_10_j; @@ -93,23 +109,24 @@ METAL_FUNC void radix5(thread float2* x, thread float2* y) { y[4] = a_8 - a_10_j; } -METAL_FUNC void radix6(thread float2* x, thread float2* y) { - float sin_pi_3 = 0.8660254037844387; - float2 a_1 = x[2] + x[4]; - float2 a_2 = x[0] - a_1 / 2; - float2 a_3 = sin_pi_3 * (x[2] - x[4]); - float2 a_4 = x[5] + x[1]; - float2 a_5 = x[3] - a_4 / 2; - float2 a_6 = sin_pi_3 * (x[5] - x[1]); - float2 a_7 = x[0] + a_1; - - float2 a_3_i = {a_3.y, -a_3.x}; - float2 a_6_i = {a_6.y, -a_6.x}; - float2 a_8 = a_2 + a_3_i; - float2 a_9 = a_2 - a_3_i; - float2 a_10 = x[3] + a_4; - float2 a_11 = a_5 + a_6_i; - float2 a_12 = a_5 - a_6_i; +template +METAL_FUNC void radix6(thread metal::vec* x, thread metal::vec* y) { + T sin_pi_3 = 0.8660254037844387; + metal::vec a_1 = x[2] + x[4]; + metal::vec a_2 = x[0] - a_1 / 2; + metal::vec a_3 = sin_pi_3 * (x[2] - x[4]); + metal::vec a_4 = x[5] + x[1]; + metal::vec a_5 = x[3] - a_4 / 2; + metal::vec a_6 = sin_pi_3 * (x[5] - x[1]); + metal::vec a_7 = x[0] + a_1; + + metal::vec a_3_i = {a_3.y, -a_3.x}; + metal::vec a_6_i = {a_6.y, -a_6.x}; + metal::vec a_8 = a_2 + a_3_i; + metal::vec a_9 = a_2 - a_3_i; + metal::vec a_10 = x[3] + a_4; + metal::vec a_11 = a_5 + a_6_i; + metal::vec a_12 = a_5 - a_6_i; y[0] = a_7 + a_10; y[1] = a_8 - a_11; @@ -119,26 +136,28 @@ METAL_FUNC void radix6(thread float2* x, thread float2* y) { y[5] = a_9 - a_12; } -METAL_FUNC void radix7(thread float2* x, thread float2* y) { +template +METAL_FUNC void radix7(thread metal::vec* x, thread metal::vec* y) { // Rader's algorithm - float2 inv = {1 / 6.0, -1 / 6.0}; + T inv_r = static_cast(1 / 6.0); + metal::vec inv = {inv_r, -inv_r}; // fft - float2 in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; - radix6(in1, y + 1); + metal::vec in1[6] = {x[1], x[3], x[2], x[6], x[4], x[5]}; + radix6(in1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(2.44013336, -1.02261879)); - y[3] = complex_mul_conj(y[3], float2(2.37046941, -1.17510629)); - y[4] = complex_mul_conj(y[4], float2(0, -2.64575131)); - y[5] = complex_mul_conj(y[5], float2(2.37046941, 1.17510629)); - y[6] = complex_mul_conj(y[6], float2(-2.44013336, -1.02261879)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(2.44013336, -1.02261879)); + y[3] = complex_mul_conj(y[3], metal::vec(2.37046941, -1.17510629)); + y[4] = complex_mul_conj(y[4], metal::vec(0, -2.64575131)); + y[5] = complex_mul_conj(y[5], metal::vec(2.37046941, 1.17510629)); + y[6] = complex_mul_conj(y[6], metal::vec(-2.44013336, -1.02261879)); // ifft - radix6(y + 1, x + 1); + radix6(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[5] = x[2] * inv + x[0]; @@ -148,79 +167,87 @@ METAL_FUNC void radix7(thread float2* x, thread float2* y) { y[3] = x[6] * inv + x[0]; } -METAL_FUNC void radix8(thread float2* x, thread float2* y) { - float cos_pi_4 = 0.7071067811865476; - float2 w_0 = {cos_pi_4, -cos_pi_4}; - float2 w_1 = {-cos_pi_4, -cos_pi_4}; - float2 temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; - radix4(temp, x); - radix4(temp + 4, x + 4); +template +METAL_FUNC void radix8(thread metal::vec* x, thread metal::vec* y) { + T cos_pi_4 = 0.7071067811865476; + metal::vec w_0 = {cos_pi_4, -cos_pi_4}; + metal::vec w_1 = {-cos_pi_4, -cos_pi_4}; + metal::vec temp[8] = {x[0], x[2], x[4], x[6], x[1], x[3], x[5], x[7]}; + radix4(temp, x); + radix4(temp + 4, x + 4); y[0] = x[0] + x[4]; y[4] = x[0] - x[4]; - float2 x_5 = complex_mul(x[5], w_0); + metal::vec x_5 = complex_mul(x[5], w_0); y[1] = x[1] + x_5; y[5] = x[1] - x_5; - float2 x_6 = {x[6].y, -x[6].x}; + metal::vec x_6 = {x[6].y, -x[6].x}; y[2] = x[2] + x_6; y[6] = x[2] - x_6; - float2 x_7 = complex_mul(x[7], w_1); + metal::vec x_7 = complex_mul(x[7], w_1); y[3] = x[3] + x_7; y[7] = x[3] - x_7; } -template -METAL_FUNC void radix10(thread float2* x, thread float2* y) { - float2 w[4]; - w[0] = {0.8090169943749475, -0.5877852522924731}; - w[1] = {0.30901699437494745, -0.9510565162951535}; +template +METAL_FUNC void radix10( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[4]; + w[0] = {T(0.8090169943749475), T(-0.5877852522924731)}; + w[1] = {T(0.30901699437494745), T(-0.9510565162951535)}; w[2] = {-w[1].x, w[1].y}; w[3] = {-w[0].x, w[0].y}; if (raders_perm) { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[3], x[4], x[8], x[2], x[1], x[7], x[9], x[6], x[5]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } else { - float2 temp[10] = { + metal::vec temp[10] = { x[0], x[2], x[4], x[6], x[8], x[1], x[3], x[5], x[7], x[9]}; - radix5(temp, x); - radix5(temp + 5, x + 5); + radix5(temp, x); + radix5(temp + 5, x + 5); } y[0] = x[0] + x[5]; y[5] = x[0] - x[5]; for (int t = 1; t < 5; t++) { - float2 a = complex_mul(x[t + 5], w[t - 1]); + metal::vec a = complex_mul(x[t + 5], w[t - 1]); y[t] = x[t] + a; y[t + 5] = x[t] - a; } } -METAL_FUNC void radix11(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 10.0, -1 / 10.0}; +template +METAL_FUNC void radix11( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 10.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix10(x + 1, y + 1); + radix10(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(0.955301878, -3.17606649)); - y[3] = complex_mul_conj(y[3], float2(2.63610556, 2.01269656)); - y[4] = complex_mul_conj(y[4], float2(2.54127802, 2.13117479)); - y[5] = complex_mul_conj(y[5], float2(2.07016210, 2.59122150)); - y[6] = complex_mul_conj(y[6], float2(0, -3.31662479)); - y[7] = complex_mul_conj(y[7], float2(2.07016210, -2.59122150)); - y[8] = complex_mul_conj(y[8], float2(-2.54127802, 2.13117479)); - y[9] = complex_mul_conj(y[9], float2(2.63610556, -2.01269656)); - y[10] = complex_mul_conj(y[10], float2(-0.955301878, -3.17606649)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(0.955301878, -3.17606649)); + y[3] = complex_mul_conj(y[3], metal::vec(2.63610556, 2.01269656)); + y[4] = complex_mul_conj(y[4], metal::vec(2.54127802, 2.13117479)); + y[5] = complex_mul_conj(y[5], metal::vec(2.07016210, 2.59122150)); + y[6] = complex_mul_conj(y[6], metal::vec(0, -3.31662479)); + y[7] = complex_mul_conj(y[7], metal::vec(2.07016210, -2.59122150)); + y[8] = complex_mul_conj(y[8], metal::vec(-2.54127802, 2.13117479)); + y[9] = complex_mul_conj(y[9], metal::vec(2.63610556, -2.01269656)); + y[10] = + complex_mul_conj(y[10], metal::vec(-0.955301878, -3.17606649)); // ifft - radix10(y + 1, x + 1); + radix10(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[6] = x[2] * inv + x[0]; @@ -234,10 +261,12 @@ METAL_FUNC void radix11(thread float2* x, thread float2* y) { y[2] = x[10] * inv + x[0]; } -template -METAL_FUNC void radix12(thread float2* x, thread float2* y) { - float2 w[6]; - float sin_pi_3 = 0.8660254037844387; +template +METAL_FUNC void radix12( + thread metal::vec* x, + thread metal::vec* y) { + metal::vec w[6]; + T sin_pi_3 = 0.8660254037844387; w[0] = {sin_pi_3, -0.5}; w[1] = {0.5, -sin_pi_3}; w[2] = {0, -1}; @@ -245,7 +274,7 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { w[4] = {-sin_pi_3, -0.5}; if (raders_perm) { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[3], x[2], @@ -258,10 +287,10 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[10], x[4], x[6]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } else { - float2 temp[12] = { + metal::vec temp[12] = { x[0], x[2], x[4], @@ -274,44 +303,50 @@ METAL_FUNC void radix12(thread float2* x, thread float2* y) { x[7], x[9], x[11]}; - radix6(temp, x); - radix6(temp + 6, x + 6); + radix6(temp, x); + radix6(temp + 6, x + 6); } y[0] = x[0] + x[6]; y[6] = x[0] - x[6]; for (int t = 1; t < 6; t++) { - float2 a = complex_mul(x[t + 6], w[t - 1]); + metal::vec a = complex_mul(x[t + 6], w[t - 1]); y[t] = x[t] + a; y[t + 6] = x[t] - a; } } -METAL_FUNC void radix13(thread float2* x, thread float2* y) { - // Raders Algorithm - float2 inv = {1 / 12.0, -1 / 12.0}; +template +METAL_FUNC void radix13( + thread metal::vec* x, + thread metal::vec* y) { + // Rader's algorithm + T inv_r = static_cast(1 / 12.0); + metal::vec inv = {inv_r, -inv_r}; // fft - radix12(x + 1, y + 1); + radix12(x + 1, y + 1); y[0] = y[1] + x[0]; // b_q - y[1] = complex_mul_conj(y[1], float2(-1, 0)); - y[2] = complex_mul_conj(y[2], float2(3.07497206, -1.88269669)); - y[3] = complex_mul_conj(y[3], float2(3.09912468, 1.84266823)); - y[4] = complex_mul_conj(y[4], float2(3.45084438, -1.04483161)); - y[5] = complex_mul_conj(y[5], float2(0.91083583, 3.48860690)); - y[6] = complex_mul_conj(y[6], float2(-3.60286363, 0.139189267)); - y[7] = complex_mul_conj(y[7], float2(3.60555128, 0)); - y[8] = complex_mul_conj(y[8], float2(3.60286363, 0.139189267)); - y[9] = complex_mul_conj(y[9], float2(0.91083583, -3.48860690)); - y[10] = complex_mul_conj(y[10], float2(-3.45084438, -1.04483161)); - y[11] = complex_mul_conj(y[11], float2(3.09912468, -1.84266823)); - y[12] = complex_mul_conj(y[12], float2(-3.07497206, -1.88269669)); + y[1] = complex_mul_conj(y[1], metal::vec(-1, 0)); + y[2] = complex_mul_conj(y[2], metal::vec(3.07497206, -1.88269669)); + y[3] = complex_mul_conj(y[3], metal::vec(3.09912468, 1.84266823)); + y[4] = complex_mul_conj(y[4], metal::vec(3.45084438, -1.04483161)); + y[5] = complex_mul_conj(y[5], metal::vec(0.91083583, 3.48860690)); + y[6] = complex_mul_conj(y[6], metal::vec(-3.60286363, 0.139189267)); + y[7] = complex_mul_conj(y[7], metal::vec(3.60555128, 0)); + y[8] = complex_mul_conj(y[8], metal::vec(3.60286363, 0.139189267)); + y[9] = complex_mul_conj(y[9], metal::vec(0.91083583, -3.48860690)); + y[10] = + complex_mul_conj(y[10], metal::vec(-3.45084438, -1.04483161)); + y[11] = complex_mul_conj(y[11], metal::vec(3.09912468, -1.84266823)); + y[12] = + complex_mul_conj(y[12], metal::vec(-3.07497206, -1.88269669)); // ifft - radix12(y + 1, x + 1); + radix12(y + 1, x + 1); y[1] = x[1] * inv + x[0]; y[7] = x[2] * inv + x[0]; @@ -325,4 +360,4 @@ METAL_FUNC void radix13(thread float2* x, thread float2* y) { y[8] = x[10] * inv + x[0]; y[4] = x[11] * inv + x[0]; y[2] = x[12] * inv + x[0]; -} \ No newline at end of file +} diff --git a/Source/Cmlx/mlx-generated/metal/fft/readwrite.h b/Source/Cmlx/mlx-generated/metal/fft/readwrite.h index 4459d36f9..a063c896f 100644 --- a/Source/Cmlx/mlx-generated/metal/fft/readwrite.h +++ b/Source/Cmlx/mlx-generated/metal/fft/readwrite.h @@ -2,6 +2,7 @@ #include +#include "../complex.h" #include "../fft/radix.h" /* FFT helpers for reading and writing from/to device memory. @@ -27,14 +28,48 @@ Each with support for: using namespace metal; +// Derives the FFT scalar arithmetic type from a storage type. +template +struct FFTStorageTraits { + using scalar_T = storage_T; +}; + +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + +template +struct FFTStorageTraits> { + using scalar_T = T; +}; + +template +struct FFTIOTypeTraits { + using scalar_T = typename FFTStorageTraits::scalar_T; + + static_assert( + metal::is_same_v::scalar_T>, + "FFT input and output storage must share a scalar arithmetic type"); + static_assert( + sizeof(in_T) <= 16 && 16 % sizeof(in_T) == 0, + "FFT input storage must divide the 128-bit sequential access width"); + static_assert( + sizeof(out_T) <= 16 && 16 % sizeof(out_T) == 0, + "FFT output storage must divide the 128-bit sequential access width"); +}; + template < typename in_T, typename out_T, int step = 0, bool four_step_real = false> struct ReadWriter { + using scalar_T = typename FFTIOTypeTraits::scalar_T; + using complex_T = vec; + const device in_T* in; - threadgroup float2* buf; + threadgroup complex_T* buf; device out_T* out; int n; int batch_size; @@ -50,23 +85,22 @@ struct ReadWriter { METAL_FUNC ReadWriter( const device in_T* in_, - threadgroup float2* buf_, + threadgroup complex_T* buf_, device out_T* out_, const short n_, const int batch_size_, const short elems_per_thread_, const uint3 elem_, const uint3 grid_, - const bool inv_) - : in(in_), - buf(buf_), - out(out_), - n(n_), - batch_size(batch_size_), - elems_per_thread(elems_per_thread_), - elem(elem_), - grid(grid_), - inv(inv_) { + const bool inv_) thread : in(in_), + buf(buf_), + out(out_), + n(n_), + batch_size(batch_size_), + elems_per_thread(elems_per_thread_), + elem(elem_), + grid(grid_), + inv(inv_) { // Account for padding on last threadgroup threads_per_tg = elem.x == grid.x - 1 ? (batch_size - (grid.x - 1) * grid.y) * grid.z @@ -74,110 +108,139 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const { - return inv ? float2(elem.x, -elem.y) : elem; + METAL_FUNC complex_T post_in(complex_T elem) const thread { + return inv ? complex_T(elem.x, -elem.y) : elem; + } + + // Handle packed complex_t storage + METAL_FUNC complex_T post_in(complex_t elem) const thread { + return post_in(complex_T(elem.real, elem.imag)); } // Handle float case for generic RFFT alg - METAL_FUNC float2 post_in(float elem) const { - return float2(elem, 0); + METAL_FUNC complex_T post_in(scalar_T elem) const thread { + return complex_T(elem, 0); } - METAL_FUNC float2 pre_out(float2 elem) const { - return inv ? float2(elem.x / n, -elem.y / n) : elem; + METAL_FUNC complex_T pre_out(complex_T elem) const thread { + return pre_out(elem, n); } - METAL_FUNC float2 pre_out(float2 elem, int length) const { - return inv ? float2(elem.x / length, -elem.y / length) : elem; + METAL_FUNC complex_T pre_out(complex_T elem, int length) const thread { + if (!inv) { + return elem; + } + if constexpr (metal::is_same_v) { + return complex_T(elem.x / length, -elem.y / length); + } else { + // Compute the reciprocal in float before narrowing to the lane. + scalar_T inv_length = static_cast(1.0f / length); + return complex_T(elem.x * inv_length, -elem.y * inv_length); + } } - METAL_FUNC bool out_of_bounds() const { + METAL_FUNC bool out_of_bounds() const thread { // Account for possible extra threadgroups int grid_index = elem.x * grid.y + elem.y; return grid_index >= batch_size; } - METAL_FUNC void load() const { + METAL_FUNC void load(scalar_T scale) const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - // 2 complex64s = 128 bits - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + // Keep each thread's sequential access at 128 bits where possible. + constexpr int read_width = 16 / sizeof(in_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - buf[index] = post_in(in[batch_idx + index]); - buf[index + 1] = post_in(in[batch_idx + index + 1]); + for (short r = 0; r < read_width; r++) { + buf[index + r] = post_in(in[batch_idx + index + r]) * scale; + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); - buf[index] = post_in(in[batch_idx + index]); + buf[index] = post_in(in[batch_idx + index]) * scale; } } - METAL_FUNC void write() const { + METAL_FUNC void load() const thread { + load(static_cast(1)); + } + + METAL_FUNC void write() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n; short tg_idx = elem.y * grid.z + elem.z; - short max_index = grid.y * n - 2; - - constexpr int read_width = 2; - for (short e = 0; e < (elems_per_thread / read_width); e++) { + constexpr int read_width = 16 / sizeof(out_T); + short max_full_index = grid.y * n - read_width; + short full_width_reads = elems_per_thread / read_width; + for (short e = 0; e < full_width_reads; e++) { short index = read_width * tg_idx + read_width * threads_per_tg * e; - index = metal::min(index, max_index); + index = metal::min(index, max_full_index); // vectorized reads - out[batch_idx + index] = pre_out(buf[index]); - out[batch_idx + index + 1] = pre_out(buf[index + 1]); + for (short r = 0; r < read_width; r++) { + out[batch_idx + index + r] = pre_out(buf[index + r]); + } } - max_index += 1; - if (elems_per_thread % 2 != 0) { - short index = tg_idx + - read_width * threads_per_tg * (elems_per_thread / read_width); + short max_index = grid.y * n - 1; + for (short r = 0; r < elems_per_thread % read_width; r++) { + short index = tg_idx + r * threads_per_tg + + read_width * threads_per_tg * full_width_reads; index = metal::min(index, max_index); out[batch_idx + index] = pre_out(buf[index]); } } // Padded IO for Bluestein's algorithm - METAL_FUNC void load_padded(int length, const device float2* w_k) const { + METAL_FUNC void load_padded(int length, const device complex_T* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = post_in(in[batch_idx + index]); - seq_buf[index] = complex_mul(elem, w_k[index]); + complex_T elem = post_in(in[batch_idx + index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0.0; } } } - METAL_FUNC void write_padded(int length, const device float2* w_k) const { + METAL_FUNC void write_padded(int length, const device complex_T* w_k) + const thread { size_t batch_idx = size_t(elem.x * grid.y) * length + elem.y * length; int fft_idx = elem.z; int m = grid.z; - float2 inv_factor = {1.0f / n, -1.0f / n}; + scalar_T inv_n = static_cast(1.0f / n); + complex_T inv_factor = {inv_n, -inv_n}; + if constexpr (!metal::is_same_v) { + // Reduced lanes applied the 1/n scale with the convolution, so only + // the conjugation remains here. + inv_factor = {scalar_T(1), -scalar_T(1)}; + } - threadgroup float2* seq_buf = buf + elem.y * n; + threadgroup complex_T* seq_buf = buf + elem.y * n; for (int e = 0; e < elems_per_thread; e++) { int index = metal::min(fft_idx + e * m, n - 1); if (index < length) { - float2 elem = seq_buf[index + length - 1] * inv_factor; - out[batch_idx + index] = pre_out(complex_mul(elem, w_k[index]), length); + complex_T elem = seq_buf[index + length - 1] * inv_factor; + out[batch_idx + index] = + pre_out(complex_mul(elem, w_k[index]), length); } } } // Strided IO for four step FFT - METAL_FUNC void compute_strided_indices(int stride, int overall_n) { + METAL_FUNC void compute_strided_indices(int stride, int overall_n) thread { // Use the batch threadgroup dimension to coalesce memory accesses: // e.g. stride = 12 // device | shared mem @@ -198,53 +261,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step - METAL_FUNC void load_strided(int stride, int overall_n) { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - buf[strided_shared_idx + e] = - post_in(in[strided_device_idx + e * stride]); + // Four-step FFT I/O + METAL_FUNC void load_strided(int stride, int overall_n) thread { + if constexpr (step == 1 && !four_step_real) { + // Do not invert between C2C four-step passes. + (void)stride; + (void)overall_n; + bool default_inv = inv; + inv = false; + if constexpr (!metal::is_same_v) { + // Reduced lanes normalize inverse transforms on input so + // intermediates stay in range; see write_strided. + if (default_inv) { + load(static_cast(1.0f / n)); + } else { + load(); + } + } else { + load(); + } + inv = default_inv; + } else { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + complex_T input = post_in(in[strided_device_idx + e * stride]); + if constexpr (!metal::is_same_v) { + if (step == 0 && !four_step_real && inv) { + input *= static_cast(1.0f / n); + } + } + buf[strided_shared_idx + e] = input; + } } } - METAL_FUNC void write_strided(int stride, int overall_n) { - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - int combined_idx = (strided_device_idx + e * stride) % overall_n; - int ij = (combined_idx / stride) * (combined_idx % stride); - // Apply four step twiddles at end of first step - float2 twiddle = get_twiddle(ij, overall_n); - out[strided_device_idx + e * stride] = complex_mul(output, twiddle); + METAL_FUNC void write_strided(int stride, int overall_n) thread { + if constexpr (step == 1 && !four_step_real) { + compute_strided_indices(stride, overall_n); + for (int e = 0; e < elems_per_thread; e++) { + if constexpr (!metal::is_same_v) { + // Inverse C2C passes were normalized on input, so only the + // conjugation of the inverse transform remains here. + complex_T value = buf[strided_shared_idx + e]; + out[strided_device_idx + e * stride] = + inv ? complex_T(value.x, -value.y) : value; + } else { + out[strided_device_idx + e * stride] = + pre_out(buf[strided_shared_idx + e], overall_n); + } + } + } else { + for (int e = 0; e < elems_per_thread; e++) { + complex_T output = buf[strided_shared_idx + e]; + int combined_idx = (strided_device_idx + e * stride) % overall_n; + int ij = (combined_idx / stride) * (combined_idx % stride); + // Apply four step twiddles at end of first step + complex_T twiddle = get_twiddle(ij, overall_n); + out[strided_device_idx + e * stride] = + complex_mul(output, twiddle); + } } } }; -// Four Step FFT Second Step -template <> -METAL_FUNC void ReadWriter::load_strided( - int stride, - int overall_n) { - // Silence compiler warnings - (void)stride; - (void)overall_n; - // Don't invert between steps - bool default_inv = inv; - inv = false; - load(); - inv = default_inv; -} - -template <> -METAL_FUNC void ReadWriter::write_strided( - int stride, - int overall_n) { - compute_strided_indices(stride, overall_n); - for (int e = 0; e < elems_per_thread; e++) { - float2 output = buf[strided_shared_idx + e]; - out[strided_device_idx + e * stride] = pre_out(output, overall_n); - } -} - +// Packed RFFT/IRFFT remains float-specific. This generic foundation covers +// C2C storage; reduced-precision real transforms need separate host dispatch. +// // For RFFT, we interleave batches of two real sequences into one complex one: // // z_k = x_k + j.y_k @@ -253,14 +335,14 @@ METAL_FUNC void ReadWriter::write_strided( // // This roughly doubles the throughput over the regular FFT. template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for RFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { size_t batch_idx = size_t(elem.x * grid.y) * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -280,7 +362,7 @@ METAL_FUNC void ReadWriter::load() const { } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = @@ -309,7 +391,7 @@ METAL_FUNC void ReadWriter::write() const { float2 x_n_minus_k = seq_buf[n - index] * conj; out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -317,7 +399,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -334,7 +416,7 @@ METAL_FUNC void ReadWriter::load_padded( if (index < length) { float2 elem = float2(in[batch_idx + index], in[batch_idx + index + next_in]); - seq_buf[index] = complex_mul(elem, w_k[index]); + seq_buf[index] = complex_mul(elem, w_k[index]); } else { seq_buf[index] = 0; } @@ -344,7 +426,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int length_over_2 = (length / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * length_over_2 * 2 + elem.y * length_over_2 * 2; @@ -367,18 +449,18 @@ METAL_FUNC void ReadWriter::write_padded( // x_0 = z_0.real // y_0 = z_0.imag if (index == 0) { - float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 elem = complex_mul(w_k[index], seq_buf[index] * inv_factor); out[batch_idx + index] = float2(elem.x, 0); out[batch_idx + index + next_out] = float2(elem.y, 0); } else { - float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); - float2 x_n_minus_k = complex_mul( + float2 x_k = complex_mul(w_k[index], seq_buf[index] * inv_factor); + float2 x_n_minus_k = complex_mul( w_k[length - index], seq_buf[length - index] * inv_factor); x_n_minus_k *= conj; // w_k should happen before this extraction out[batch_idx + index] = (x_k + x_n_minus_k) / 2; out[batch_idx + index + next_out] = - complex_mul(((x_k - x_n_minus_k) / 2), minus_j); + complex_mul(((x_k - x_n_minus_k) / 2), minus_j); } } } @@ -389,14 +471,14 @@ METAL_FUNC void ReadWriter::write_padded( // x_k = Re(Z_k) // Y_k = Imag(Z_k) template <> -METAL_FUNC bool ReadWriter::out_of_bounds() const { +METAL_FUNC bool ReadWriter::out_of_bounds() const thread { int grid_index = elem.x * grid.y + elem.y; // We pack two sequences into one for IRFFTs return grid_index * 2 >= batch_size; } template <> -METAL_FUNC void ReadWriter::load() const { +METAL_FUNC void ReadWriter::load() const thread { short n_over_2 = (n / 2) + 1; size_t batch_idx = size_t(elem.x * grid.y) * n_over_2 * 2 + elem.y * n_over_2 * 2; @@ -425,17 +507,17 @@ METAL_FUNC void ReadWriter::load() const { x = float2(x.x, 0); y = float2(y.x, 0); } - seq_buf[index] = x + complex_mul(y, plus_j); + seq_buf[index] = x + complex_mul(y, plus_j); seq_buf[index].y = -seq_buf[index].y; if (index > 0 && !last_val) { - seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); + seq_buf[n - index] = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[n - index].y = -seq_buf[n - index].y; } } } template <> -METAL_FUNC void ReadWriter::write() const { +METAL_FUNC void ReadWriter::write() const thread { int batch_idx = elem.x * grid.y * n * 2 + elem.y * n * 2; threadgroup float2* seq_buf = buf + elem.y * n; @@ -456,7 +538,7 @@ METAL_FUNC void ReadWriter::write() const { template <> METAL_FUNC void ReadWriter::load_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { int n_over_2 = (n / 2) + 1; int length_over_2 = (length / 2) + 1; @@ -486,12 +568,12 @@ METAL_FUNC void ReadWriter::load_padded( x = float2(x.x, 0); y = float2(y.x, 0); } - float2 elem1 = x + complex_mul(y, plus_j); - seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); + float2 elem1 = x + complex_mul(y, plus_j); + seq_buf[index] = complex_mul(elem1 * conj, w_k[index]); if (index > 0 && !last_val) { - float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); + float2 elem2 = (x * conj) + complex_mul(y * conj, plus_j); seq_buf[length - index] = - complex_mul(elem2 * conj, w_k[length - index]); + complex_mul(elem2 * conj, w_k[length - index]); } } else { short pad_index = metal::min(length + (index - length_over_2) * 2, n - 2); @@ -504,7 +586,7 @@ METAL_FUNC void ReadWriter::load_padded( template <> METAL_FUNC void ReadWriter::write_padded( int length, - const device float2* w_k) const { + const device float2* w_k) const thread { size_t batch_idx = size_t(elem.x * grid.y) * length * 2 + elem.y * length * 2; threadgroup float2* seq_buf = buf + elem.y * n + length - 1; @@ -519,7 +601,8 @@ METAL_FUNC void ReadWriter::write_padded( for (int e = 0; e < elems_per_thread; e++) { int index = fft_idx + e * m; if (index < length) { - float2 output = complex_mul(seq_buf[index] * inv_factor, w_k[index]); + float2 output = + complex_mul(seq_buf[index] * inv_factor, w_k[index]); out[batch_idx + index] = output.x / length; out[batch_idx + index + next_out] = output.y / -length; } @@ -531,7 +614,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -546,7 +629,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; int coalesce_width = grid.y; int tg_idx = elem.y * grid.z + elem.z; @@ -575,7 +658,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { int overall_n_over_2 = overall_n / 2 + 1; auto conj = float2(1, -1); @@ -600,7 +683,7 @@ template <> METAL_FUNC void ReadWriter::load_strided( int stride, - int overall_n) { + int overall_n) thread { // Silence compiler warnings (void)stride; (void)overall_n; @@ -614,7 +697,7 @@ template <> METAL_FUNC void ReadWriter::write_strided( int stride, - int overall_n) { + int overall_n) thread { compute_strided_indices(stride, overall_n); for (int e = 0; e < elems_per_thread; e++) { diff --git a/Source/Cmlx/mlx-generated/metal/fp4.h b/Source/Cmlx/mlx-generated/metal/fp4.h index 25642f201..af4cbb02c 100644 --- a/Source/Cmlx/mlx-generated/metal/fp4.h +++ b/Source/Cmlx/mlx-generated/metal/fp4.h @@ -1,7 +1,10 @@ #pragma once +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { - fp4_e2m1(float x) { + fp4_e2m1(float x) thread { if (metal::isnan(x)) { bits = 0x7; return; @@ -30,17 +33,17 @@ struct fp4_e2m1 { bits |= sign_bit; } - operator float16_t() { + operator float16_t() thread { half converted = as_type(ushort((bits & 7) << 9)); converted *= 16384.0; return bits & 8 ? -converted : converted; } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } diff --git a/Source/Cmlx/mlx-generated/metal/fp8.h b/Source/Cmlx/mlx-generated/metal/fp8.h index 60d34be69..42c5ec128 100644 --- a/Source/Cmlx/mlx-generated/metal/fp8.h +++ b/Source/Cmlx/mlx-generated/metal/fp8.h @@ -2,7 +2,7 @@ struct fp8_e4m3 { template - fp8_e4m3(T f) { + fp8_e4m3(T f) thread { // From PyTorch // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148 uint32_t fp8_max = 543 << 21; @@ -29,7 +29,7 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } - operator float16_t() { + operator float16_t() thread { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; @@ -37,11 +37,11 @@ struct fp8_e4m3 { return (sign ? -converted : converted); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } @@ -49,7 +49,7 @@ struct fp8_e4m3 { }; struct fp8_e8m0 { - fp8_e8m0(float x) { + fp8_e8m0(float x) thread { if (!metal::isfinite(x)) { bits = 0xFF; return; @@ -66,15 +66,26 @@ struct fp8_e8m0 { bits = static_cast(n + 127); } - operator bfloat16_t() { + operator bfloat16_t() thread { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } - operator float() { + operator float() thread { uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); return as_type(out); } uint8_t bits; }; + +// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down +// leaves the block's largest elements outside the element range, where they +// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0. +inline float mx_scale_round_up(float x) { + fp8_e8m0 s(x); + if (s.bits < 0xFE && float(s) < x) { + s.bits += 1; + } + return float(s); +} diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized.h b/Source/Cmlx/mlx-generated/metal/fp_quantized.h index eef3f2cf1..733361958 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -39,7 +39,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -50,7 +50,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -137,11 +137,12 @@ inline void qouter(const thread uint8_t* w, U x, U scale, thread U* result) { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } @@ -186,15 +187,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -203,7 +204,7 @@ struct QuantizedBlockLoader { scales_ + bi * src_ld / group_size + (bj * pack_factor) / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -215,7 +216,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -241,7 +242,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -321,10 +322,16 @@ METAL_FUNC void fp_qmv_quad_impl( } } -template +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int results_per_simdgroup = 4> METAL_FUNC void fp_qmv_fast_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -334,7 +341,6 @@ METAL_FUNC void fp_qmv_fast_impl( uint simd_lid [[thread_index_in_simdgroup]]) { constexpr int packs_per_thread = 2; constexpr int num_simdgroups = 2; - constexpr int results_per_simdgroup = 4; constexpr int pack_factor = get_pack_factor<32, bits>(); constexpr int bytes_per_pack = get_bytes_per_pack<32>(); constexpr int values_per_thread = pack_factor * packs_per_thread; @@ -374,18 +380,28 @@ METAL_FUNC void fp_qmv_fast_impl( x += block_size; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } -template +template METAL_FUNC void fp_qmv_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -421,6 +437,11 @@ METAL_FUNC void fp_qmv_impl( return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // In this case we need to properly guard all our reads because there isn't // even 1 tile in the matrix if (out_vec_size < (num_simdgroups * results_per_simdgroup)) { @@ -440,7 +461,7 @@ METAL_FUNC void fp_qmv_impl( auto wl = (const device uint8_t*)(ws + row * in_vec_size_w); const device auto* sl = scales + row * in_vec_size_g; - uint8_t s = sl[0]; + U s = dequantize_scale(sl[0]); result[row] += qdot(wl, x_thread, s); } @@ -471,7 +492,11 @@ METAL_FUNC void fp_qmv_impl( row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } } } } @@ -519,20 +544,145 @@ METAL_FUNC void fp_qmv_impl( for (int row = 0; row < results_per_simdgroup; row++) { result[row] = simd_sum(result[row]); if (simd_lid == 0) { - y[row] = static_cast(result[row]); + if constexpr (has_global_scale) { + y[row] = static_cast(result[row] * inv_scale_enc); + } else { + y[row] = static_cast(result[row]); + } + } + } + } +} + +// Quantized matrix-vector for a small batch of input vectors, the M in +// [2, vector_limit) band between qmv (M==1) and qmm. Each thread owns one +// output row that k_lanes lanes reduce over K; the vecs_per_tg vectors are +// streamed so each weight group is dequantized once and reused across them. +template +METAL_FUNC void fp_qmv_wide_impl( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int pack_factor = get_pack_factor<32, bits>(); + constexpr int bytes_per_pack = get_bytes_per_pack<32>(); + constexpr int nf4 = group_size / 4; // float4 lanes per quant group + + typedef float U; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; // first input vector handled here + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bytes_per_pack / pack_factor; + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device uint8_t* srow = scales + row * in_vec_size_g; + + // One device pointer per streamed vector; the clamp keeps an out-of-range + // tail slot reading a valid row (it is never written below). + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's quant groups: one group is + // dequantized to float4 lanes and dot()'d against each streamed vector. One + // group per iteration keeps weight-register pressure low for occupancy. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + const int k0 = g * group_size; + U s = dequantize_scale(srow[g]); + const device uint8_t* wg = wrow + k0 * bytes_per_pack / pack_factor; + + float4 w4[nf4]; + if constexpr (bits == 4) { + const device uint16_t* wq = (const device uint16_t*)wg; +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<4>{}(wq[i]), + Dequantize<4>{}(wq[i] >> 4), + Dequantize<4>{}(wq[i] >> 8), + Dequantize<4>{}(wq[i] >> 12)); + } + } else { +#pragma unroll + for (int i = 0; i < nf4; i++) { + w4[i] = float4( + Dequantize<8>{}(wg[4 * i]), + Dequantize<8>{}(wg[4 * i + 1]), + Dequantize<8>{}(wg[4 * i + 2]), + Dequantize<8>{}(wg[4 * i + 3])); + } + } + +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device vec* xv4 = (const device vec*)(xv[v] + k0); + float acc = 0; +#pragma unroll + for (int j = 0; j < nf4; j++) { + acc += dot(w4[j], float4(xv4[j])); + } + result[v] += s * acc; + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); } } } } -template +template METAL_FUNC void fp_qvm_impl( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const int in_vec_size, const int out_vec_size, + const int in_vec_stride, uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { @@ -563,13 +713,18 @@ METAL_FUNC void fp_qvm_impl( int out_col = pack_factor * tn * (tid.y * num_simdgroups + simd_gid); ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w; scales += out_col / group_size + simd_lid * out_vec_size_g; - x += tid.x * in_vec_size + simd_lid; + x += tid.x * in_vec_stride + simd_lid; y += tid.x * out_vec_size + out_col; if (out_col >= out_vec_size) { return; } + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + // Loop over in_vec in blocks of block_size int remaining = in_vec_size % block_size; if (remaining == 0) { @@ -619,7 +774,11 @@ METAL_FUNC void fp_qvm_impl( if (simd_lid == 0) { #pragma clang loop unroll(full) for (int k = 0; k < tn * pack_factor; k++) { - y[k] = static_cast(result[k]); + if constexpr (has_global_scale) { + y[k] = static_cast(result[k] * inv_scale_enc); + } else { + y[k] = static_cast(result[k]); + } } } } @@ -642,6 +801,7 @@ METAL_FUNC void fp_qmm_t_impl( const constant int& K, const constant int& N, const constant int& M, + const constant int& K_eff, uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], @@ -695,7 +855,7 @@ METAL_FUNC void fp_qmm_t_impl( if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -705,7 +865,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -717,7 +877,7 @@ METAL_FUNC void fp_qmm_t_impl( } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -727,7 +887,7 @@ METAL_FUNC void fp_qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -751,11 +911,11 @@ METAL_FUNC void fp_qmm_t_impl( template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + int BM = 32, + int BK = 32, + int BN = 32> METAL_FUNC void fp_qmm_n_impl( const device uint32_t* w, const device uint8_t* scales, @@ -1007,10 +1167,17 @@ template w, scales, x, y, in_vec_size, out_vec_size, tid, quad_gid, quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1042,14 +1209,35 @@ template s_strides, tid); } - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl< + T, + group_size, + bits, + has_global_scale, + results_per_simdgroup>( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void fp_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1081,14 +1269,74 @@ template s_strides, tid); } - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); +} + +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void fp_qmv_wide( + const device uint32_t* w, + const device uint8_t* scales, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + tid); + } + fp_qmv_wide_impl( + w, scales, x, y, in_vec_size, out_vec_size, M, tid, simd_gid, simd_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void fp_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, device T* y, const constant int& in_vec_size, @@ -1120,8 +1368,18 @@ template s_strides, tid); } - fp_qvm_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qvm_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + in_vec_size, + tid, + simd_gid, + simd_lid); } template @@ -1163,8 +1421,21 @@ template int in_vec_size_adj = tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size; - fp_qvm_impl( - w, scales, x, y, in_vec_size_adj, out_vec_size, tid, simd_gid, simd_lid); + // The in_vec_stride is the full K dimension, not the partition size + int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; + + fp_qvm_impl( + w, + scales, + nullptr, + x, + y, + in_vec_size_adj, + out_vec_size, + in_vec_stride, + tid, + simd_gid, + simd_lid); } template < @@ -1219,17 +1490,18 @@ template < tid); } fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } template < typename T, - const int group_size, - const int bits, - const bool batched, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1278,10 +1550,11 @@ template < w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void fp_gather_qmv_fast( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1323,14 +1596,24 @@ template w_strides, s_strides, tid); - fp_qmv_fast_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_fast_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qmv( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1372,14 +1655,24 @@ template w_strides, s_strides, tid); - fp_qmv_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qmv_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + tid, + simd_gid, + simd_lid); } -template +template [[kernel]] void fp_gather_qvm( const device uint32_t* w, const device uint8_t* scales, + const device float* global_scale, const device T* x, const device uint32_t* lhs_indices, const device uint32_t* rhs_indices, @@ -1421,8 +1714,18 @@ template w_strides, s_strides, tid); - fp_qvm_impl( - w, scales, x, y, in_vec_size, out_vec_size, tid, simd_gid, simd_lid); + fp_qvm_impl( + w, + scales, + global_scale, + x, + y, + in_vec_size, + out_vec_size, + in_vec_size, + tid, + simd_gid, + simd_lid); } template < @@ -1486,16 +1789,71 @@ template < s_strides, tid); fp_qmm_t_impl( - w, scales, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, scales, x, y, Xs, Ws, K, N, M, K, tid, lid, simd_gid, simd_lid); } template < typename T, const int group_size, const int bits, + const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32> +[[kernel]] void fp_qmm_t_splitk( + const device uint32_t* w [[buffer(0)]], + const device uint8_t* scales [[buffer(1)]], + const device T* x [[buffer(2)]], + device T* y [[buffer(3)]], + const constant int& K [[buffer(4)]], + const constant int& N [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& k_partition_size [[buffer(7)]], + const constant int& split_k_partition_stride [[buffer(8)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + (void)lid; + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + constexpr int pack_factor = get_pack_factor<8, bits>(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + const int k_start = tid.z * k_partition_size; + x += k_start; + + auto wl = (const device uint8_t*)w; + wl += k_start * bytes_per_pack / pack_factor; + scales += k_start / group_size; + y += tid.z * static_cast(split_k_partition_stride); + + fp_qmm_t_impl( + (const device uint32_t*)wl, + scales, + x, + y, + Xs, + Ws, + K, + N, + M, + k_partition_size, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void fp_gather_qmm_n( const device uint32_t* w, const device uint8_t* scales, @@ -1746,38 +2104,50 @@ template < } } -template +template [[kernel]] void fp_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], device uint8_t* scales [[buffer(2)]], + const device float* global_scale [[buffer(3)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + if constexpr (use_mx_scale) { + scale_dec_b = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); + auto s = ScaleType(scale_dec_b); uint8_t q_scale = s.bits; - scale = float(s); - size_t gindex = index / group_size; if (index % group_size == 0) { scales[gindex] = q_scale; } - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = float(s); + scale_enc_b = (scale_enc_b == 0) ? 0.0f : (scale_enc / scale_enc_b); + uint8_t output = Quantize{}(w_thread * scale_enc_b); if (bits == 4) { uint8_t sval = simd_shuffle_down(output, 1); output |= sval << bits; @@ -1788,10 +2158,11 @@ template } } -template +template [[kernel]] void fp_dequantize( const device uint8_t* w [[buffer(0)]], const device uint8_t* scales [[buffer(1)]], + const device float* global_scale [[buffer(2)]], device T* out [[buffer(3)]], uint2 index [[thread_position_in_grid]], uint2 grid_dim [[threads_per_grid]]) { @@ -1803,9 +2174,17 @@ template out += oindex; + float inv_scale_enc = 1.0f; + if constexpr (has_global_scale) { + inv_scale_enc = *global_scale / (F8E4M3_MAX * F4E2M1_MAX); + } + using ScaleType = metal::conditional_t; auto q_scale = ((device ScaleType*)(scales))[gindex]; auto scale = float(q_scale); + if constexpr (has_global_scale) { + scale *= inv_scale_enc; + } uint val = w[offset]; #pragma clang loop unroll(full) @@ -1820,31 +2199,43 @@ template } } -template +template [[kernel]] void fp_quantize_dequantize( const device T* w [[buffer(0)]], - device T* out [[buffer(1)]], + const device float* global_scale [[buffer(1)]], + device T* out [[buffer(2)]], uint2 tidx [[thread_position_in_grid]], - uint2 grid_dim [[threads_per_grid]]) { + uint2 grid_dim [[threads_per_grid]], + uint simd_lid [[thread_index_in_simdgroup]]) { constexpr bool use_mx_scale = group_size == 32; size_t index = tidx.x + grid_dim.x * size_t(tidx.y); - float scale; + float scale_enc = 1.0f; + if constexpr (has_global_scale) { + scale_enc = (F8E4M3_MAX * F4E2M1_MAX) / *global_scale; + } + + float scale_dec_b; float w_thread = w[index]; - if (use_mx_scale) { - scale = simd_max(abs(w_thread)); + if constexpr (use_mx_scale) { + scale_dec_b = simd_max(abs(w_thread)); } else { - float w_max_l = simd_max(tidx.x < 16 ? abs(w_thread) : 0.0); - float w_max_r = simd_max(tidx.x >= 16 ? abs(w_thread) : 0.0); - scale = tidx.x < 16 ? w_max_l : w_max_r; + float w_max_l = simd_max(simd_lid < 16 ? abs(w_thread) : 0.0); + float w_max_r = simd_max(simd_lid >= 16 ? abs(w_thread) : 0.0); + scale_dec_b = simd_lid < 16 ? w_max_l : w_max_r; + } + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + if constexpr (has_global_scale) { + scale_dec_b *= scale_enc; + } else if constexpr (use_mx_scale) { + scale_dec_b = mx_scale_round_up(scale_dec_b); } - scale /= bits == 4 ? 6.0f : 448.0f; using ScaleType = metal::conditional_t; - auto s = ScaleType(scale); - scale = float(s); + auto scale = float(ScaleType(scale_dec_b)); - uint8_t output = Quantize{}(scale == 0 ? 0.0f : w_thread / scale); + float scale_enc_b = (scale == 0) ? 0.0f : (scale_enc / scale); + uint8_t output = Quantize{}(w_thread * scale_enc_b); - out[index] = static_cast(scale * Dequantize{}(output)); + out[index] = static_cast((scale / scale_enc) * Dequantize{}(output)); } diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h index 38d9fb654..35c61c2fa 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h @@ -39,7 +39,7 @@ static inline T dequantize_scale(uint8_t s) { template struct Quantize { - uint8_t operator()(float x) { + uint8_t operator()(float x) thread { if (bits == 8) { return fp8_e4m3(x).bits; } else { @@ -50,7 +50,7 @@ struct Quantize { template struct Dequantize { - U operator()(uint8_t x) { + U operator()(uint8_t x) thread { if constexpr (bits == 8) { return U(*(thread fp8_e4m3*)(&x)); } else { @@ -61,11 +61,12 @@ struct Dequantize { template inline void dequantize(uint8_t w, U scale, threadgroup U* w_local) { + const float s = float(scale); if constexpr (bits == 4) { - w_local[0] = scale * Dequantize<4, U>{}(w); - w_local[1] = scale * Dequantize<4, U>{}(w >> 4); + w_local[0] = static_cast(s * Dequantize<4, float>{}(w)); + w_local[1] = static_cast(s * Dequantize<4, float>{}(w >> 4)); } else { - w_local[0] = scale * Dequantize<8, U>{}(w); + w_local[0] = static_cast(s * Dequantize<8, float>{}(w)); } } @@ -112,14 +113,14 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -127,7 +128,7 @@ struct QuantizedBlockLoader { bj * bytes_per_pack), scales(scales_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -143,7 +144,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -173,12 +174,12 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { scales += n_groups; } else { - scales += n_groups * group_stride; + scales += group_stride; } } }; @@ -246,16 +247,13 @@ METAL_FUNC void fp_qmm_t_impl( // Make the weight loader loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -263,22 +261,17 @@ METAL_FUNC void fp_qmm_t_impl( constexpr bool transpose_a = false; constexpr bool transpose_b = true; - const short sgp_sm = min(SM, short(M - (y_row + tm))); + const short sgp_sm = min(int(SM), M - (y_row + tm)); const bool is_unaligned_sm = (sgp_sm != SM); - const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn))); + const short sgp_sn = aligned_N ? SN : min(int(SN), N - (y_col + tn)); const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -297,8 +290,8 @@ METAL_FUNC void fp_qmm_t_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -400,16 +393,13 @@ METAL_FUNC void fp_qmm_n_impl( // const short num_outs = min(BN, N - y_col); loader_w_t loader_w(wl, scales, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -421,12 +411,7 @@ METAL_FUNC void fp_qmm_n_impl( using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -438,8 +423,8 @@ METAL_FUNC void fp_qmm_n_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -869,24 +854,19 @@ template < wl += transpose ? y_col_long * K_w : y_col * bytes_per_pack / pack_factor; scales += transpose ? y_col_long * K_g : y_col / group_size; - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); - const short sgp_sm = - align_M ? SM : min(SM, short(max(0, (M - (y_row + tm))))); - const short sgp_sn = - align_N ? SN : min(SN, short(max(0, (N - (y_col + tn))))); + const short sgp_sm = align_M ? SM : min(int(SM), max(0, (M - (y_row + tm)))); + const short sgp_sn = align_N ? SN : min(int(SN), max(0, (N - (y_col + tn)))); const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); const bool is_unaligned_bn = align_N ? false : (tgp_bn != BN); @@ -896,10 +876,6 @@ template < using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - // Do as many matmuls as necessary uint32_t index; short offset; @@ -921,8 +897,11 @@ template < threadgroup_barrier(mem_flags::mem_none); // Prepare threadgroup mma operation - NAXTile Dtile; + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); const device T* xn = x + tm * K; @@ -951,33 +930,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); - } - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -991,38 +972,37 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - - if constexpr (transpose) { - Btile.template load( - Ws + tn * BK_padded + kk1); - } else { - Btile.template load( - Ws + tn + kk1 * BN_padded); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + + if constexpr (transpose) { + Btile.template load( + Ws + tn * BK_padded + kk1); + } else { + Btile.template load( + Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { diff --git a/Source/Cmlx/mlx-generated/metal/gemv.metal b/Source/Cmlx/mlx-generated/metal/gemv.h similarity index 71% rename from Source/Cmlx/mlx-generated/metal/gemv.metal rename to Source/Cmlx/mlx-generated/metal/gemv.h index 89403d3dc..7ac7388ab 100644 --- a/Source/Cmlx/mlx-generated/metal/gemv.metal +++ b/Source/Cmlx/mlx-generated/metal/gemv.h @@ -1,5 +1,7 @@ // Copyright © 2023-2024 Apple Inc. +#pragma once + #include #include @@ -148,7 +150,7 @@ struct GEMVKernel { out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM; // Advance matrix - mat += out_row * matrix_ld; + mat += size_t(out_row) * matrix_ld; constexpr const uniform loop_stride = make_uniform(blockN); const uniform in_size = make_uniform(in_vec_size); @@ -356,7 +358,7 @@ struct GEMVTKernel { for (int tm = 0; tm < TM; tm++) { auto vc = static_cast(v_coeff[tm]); for (int tn = 0; tn < TN; tn++) { - inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn]; + inter[tn] = mat[size_t(bm + tm) * marix_ld + out_col + tn]; } for (int tn = 0; tn < TN; tn++) { result[tn] += vc * inter[tn]; @@ -372,7 +374,7 @@ struct GEMVTKernel { MLX_MTL_PRAGMA_UNROLL for (int tn = 0; tn < TN; tn++) { - inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn]; + inter[tn] = mat[size_t(bm + tm) * marix_ld + out_col + tn]; } MLX_MTL_PRAGMA_UNROLL @@ -432,7 +434,7 @@ struct GEMVTKernel { }; /////////////////////////////////////////////////////////////////////////////// -/// Matrix vector multiplication +/// Matrix vector multiplication kernel /////////////////////////////////////////////////////////////////////////////// template < @@ -507,44 +509,6 @@ template < simd_lid); } -#define instantiate_gemv_helper( \ - name, itype, bm, bn, sm, sn, tm, tn, nc, axpby) \ - instantiate_kernel( \ - "gemv_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm \ - "_tn" #tn "_nc" #nc "_axpby" #axpby, \ - gemv, \ - itype, \ - bm, \ - bn, \ - sm, \ - sn, \ - tm, \ - tn, \ - nc, \ - axpby) - -// clang-format off -#define instantiate_gemv(name, itype, bm, bn, sm, sn, tm, tn) \ - instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 0) \ - instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 1) \ - instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 0) \ - instantiate_gemv_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 1) // clang-format on - -// clang-format off -#define instantiate_gemv_blocks(name, itype) \ - instantiate_gemv(name, itype, 1, 8, 1, 32, 4, 4) \ - instantiate_gemv(name, itype, 1, 8, 1, 32, 1, 4) \ - instantiate_gemv(name, itype, 1, 1, 8, 4, 4, 4) \ - instantiate_gemv(name, itype, 1, 1, 8, 4, 1, 4) \ - instantiate_gemv(name, itype, 4, 1, 1, 32, 1, 4) \ - instantiate_gemv(name, itype, 4, 1, 1, 32, 4, 4) \ - instantiate_gemv(name, itype, 8, 1, 1, 32, 4, 4) // clang-format on - -instantiate_gemv_blocks(float32, float); -instantiate_gemv_blocks(float16, half); -instantiate_gemv_blocks(bfloat16, bfloat16_t); -instantiate_gemv_blocks(complex64, complex64_t); - template < typename T, const int BM, /* Threadgroup rows (in simdgroups) */ @@ -635,25 +599,8 @@ template < simd_lid); } -// clang-format off -#define instantiate_gemv_bs_helper(nm, itype, bm, bn, sm, sn, tm, tn) \ - instantiate_kernel( \ - "gemv_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm \ - "_sn" #sn "_tm" #tm "_tn" #tn, \ - gemv_gather, itype, bm, bn, sm, sn, tm, tn) - -#define instantiate_gemv_bs_blocks(name, itype) \ - instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 1, 4) \ - instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 4, 4) \ - instantiate_gemv_bs_helper(name, itype, 8, 1, 1, 32, 4, 4) // clang-format on - -instantiate_gemv_bs_blocks(float32, float); -instantiate_gemv_bs_blocks(float16, half); -instantiate_gemv_bs_blocks(bfloat16, bfloat16_t); -instantiate_gemv_bs_blocks(complex64, complex64_t); - /////////////////////////////////////////////////////////////////////////////// -/// Vector matrix multiplication +/// Vector matrix multiplication kernel /////////////////////////////////////////////////////////////////////////////// template < @@ -728,34 +675,6 @@ template < simd_lid); } -// clang-format off -#define instantiate_gemv_t_helper( \ - name, itype, bm, bn, sm, sn, tm, tn, nc, axpby) \ - instantiate_kernel( \ - "gemv_t_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn \ - "_tm" #tm "_tn" #tn "_nc" #nc "_axpby" #axpby, \ - gemv_t, itype, bm, bn, sm, sn, tm, tn, nc, axpby) - -#define instantiate_gemv_t(name, itype, bm, bn, sm, sn, tm, tn) \ - instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 0) \ - instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 1) \ - instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 0) \ - instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 1) // clang-format on - -// clang-format off -#define instantiate_gemv_t_blocks(name, itype) \ - instantiate_gemv_t(name, itype, 1, 2, 8, 4, 4, 1) \ - instantiate_gemv_t(name, itype, 1, 2, 8, 4, 4, 4) \ - instantiate_gemv_t(name, itype, 1, 4, 8, 4, 4, 4) \ - instantiate_gemv_t(name, itype, 1, 16, 8, 4, 4, 4) \ - instantiate_gemv_t(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on - -// clang-format off -instantiate_gemv_t_blocks(float32, float); -instantiate_gemv_t_blocks(float16, half); -instantiate_gemv_t_blocks(bfloat16, bfloat16_t); -instantiate_gemv_t_blocks(complex64, complex64_t); // clang-format on - template < typename T, const int BM, /* Threadgroup rows (in simdgroups) */ @@ -846,23 +765,280 @@ template < simd_lid); } -// clang-format off -#define instantiate_gemv_t_bs_helper( \ - nm, itype, bm, bn, sm, sn, tm, tn) \ - instantiate_kernel( \ - "gemv_t_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm \ - "_sn" #sn "_tm" #tm "_tn" #tn, \ - gemv_t_gather, itype, bm, bn, sm, sn, tm, tn) - -#define instantiate_gemv_t_bs_blocks(name, itype) \ - instantiate_gemv_t_bs_helper(name, itype, 1, 2, 8, 4, 4, 1) \ - instantiate_gemv_t_bs_helper(name, itype, 1, 2, 8, 4, 4, 4) \ - instantiate_gemv_t_bs_helper(name, itype, 1, 4, 8, 4, 4, 4) \ - instantiate_gemv_t_bs_helper(name, itype, 1, 16, 8, 4, 4, 4) \ - instantiate_gemv_t_bs_helper(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on - -// clang-format off -instantiate_gemv_t_bs_blocks(float32, float); -instantiate_gemv_t_bs_blocks(float16, half); -instantiate_gemv_t_bs_blocks(bfloat16, bfloat16_t); -instantiate_gemv_t_bs_blocks(complex64, complex64_t); // clang-format on +/////////////////////////////////////////////////////////////////////////////// +/// Multi-vector matrix-vector multiplication (wide gemv) +/////////////////////////////////////////////////////////////////////////////// + +constant bool gemv_wide_has_batch [[function_constant(0)]]; +constant bool gemv_wide_do_axpby [[function_constant(1)]]; + +// out[M, N] = x[M, K] @ mat[N, K]^T for small M: each threadgroup streams a +// block of matrix rows once and applies it to vecs_per_tg input vectors, so +// the matrix is read ceil(M / vecs_per_tg) times instead of once per padded +// GEMM tile. k_lanes lanes reduce K for one row, 32 / k_lanes rows per +// simdgroup, k_lanes / 8 simdgroups per threadgroup. +template < + typename T, + const int vecs_per_tg, + const int k_lanes, + typename AccT = typename DefaultAccT::type> +struct GemvWide { + static constexpr constant int unroll = 8; + static constexpr constant int num_simdgroups = k_lanes / 8; + + static METAL_FUNC void run( + const device T* mat, + const device T* in_vec, + const device T* bias, + device T* out_vec, + int in_vec_size, + int out_vec_size, + int M, + int matrix_ld, + int vector_ld, + float alpha, + float beta, + int bias_ld, + int bias_fd, + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int rows_per_simdgroup = 32 / k_lanes; + constexpr int rows_per_tg = rows_per_simdgroup * num_simdgroups; + + const short k_lane = + simd_lid % k_lanes; // this lane's slot in the K reduction + const short sg_row = + simd_lid / k_lanes; // which output row of the simdgroup + + const int out_row = + tid.y * rows_per_tg + rows_per_simdgroup * simd_gid + sg_row; + + // Clamped tail rows/vectors read valid memory; the guarded writes + // below never store them. + const int row = min(out_row, out_vec_size - 1); + const device T* wrow = mat + int64_t(row) * matrix_ld; + const device vec* w4 = (const device vec*)wrow; + const int n_v4 = in_vec_size / 4; + const int n_main = n_v4 - n_v4 % (k_lanes * unroll); + + // Vector chunks beyond grid.x round-robin onto the same threadgroups: + // re-walking the row block hits cache where an extra grid column would + // re-stream it from DRAM. + const device vec* x4 = (const device vec*)in_vec; + const int x_ld4 = vector_ld / 4; + const int n_chunks = (M + vecs_per_tg - 1) / vecs_per_tg; + for (int chunk = tid.x; chunk < n_chunks; chunk += tgpg.x) { + const int vec0 = chunk * vecs_per_tg; + + int x_off4[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + x_off4[v] = min(vec0 + v, M - 1) * x_ld4; + } + + AccT result[vecs_per_tg] = {0}; + + // Adjacent lanes read adjacent blocks so every transaction lands on + // whole cachelines; the unroll keeps loads in flight on short rows. + for (int base = 0; base < n_main; base += k_lanes * unroll) { + vec wf[unroll]; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + wf[i] = vec(w4[base + i * k_lanes + k_lane]); + } + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + AccT acc = 0; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < unroll; i++) { + acc += + dot(wf[i], + vec(x4[x_off4[v] + base + i * k_lanes + k_lane])); + } + result[v] += acc; + } + } + for (int idx = n_main + k_lane; idx < n_v4; idx += k_lanes) { + const vec wf = vec(w4[idx]); + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + result[v] += dot(wf, vec(x4[x_off4[v] + idx])); + } + } + + // The halving shuffles reduce each row's k_lanes while rows sharing + // the simdgroup stay separate. + MLX_MTL_PRAGMA_UNROLL + for (int v = 0; v < vecs_per_tg; v++) { + MLX_MTL_PRAGMA_UNROLL + for (ushort off = k_lanes / 2; off >= 1; off >>= 1) { + result[v] += simd_shuffle_down(result[v], off); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + int64_t out_idx = int64_t(vec0 + v) * out_vec_size + out_row; + if (gemv_wide_do_axpby) { + AccT bias_val = static_cast( + bias[int64_t(vec0 + v) * bias_ld + out_row * bias_fd]); + out_vec[out_idx] = + static_cast(alpha * result[v] + beta * bias_val); + } else { + out_vec[out_idx] = static_cast(result[v]); + } + } + } + } + } + } +}; + +template +[[kernel]] void gemv_wide( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2), function_constant(gemv_wide_do_axpby)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant float& alpha + [[buffer(9), function_constant(gemv_wide_do_axpby)]], + const constant float& beta + [[buffer(10), function_constant(gemv_wide_do_axpby)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* vector_batch_stride [[buffer(13)]], + const constant int64_t* matrix_batch_stride [[buffer(14)]], + const constant int64_t* bias_batch_stride + [[buffer(15), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_ld + [[buffer(16), function_constant(gemv_wide_do_axpby)]], + const constant int& bias_fd + [[buffer(17), function_constant(gemv_wide_do_axpby)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (gemv_wide_has_batch) { + in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim); + mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim); + + if (gemv_wide_do_axpby) { + bias += elem_to_loc(tid.z, batch_shape, bias_batch_stride, batch_ndim); + } + } else { + in_vec += tid.z * vector_batch_stride[0]; + mat += tid.z * matrix_batch_stride[0]; + + if (gemv_wide_do_axpby) { + bias += tid.z * bias_batch_stride[0]; + } + } + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + alpha, + beta, + bias_ld, + bias_fd, + tid, + tgpg, + simd_gid, + simd_lid); +} + +template +[[kernel]] void gemv_wide_gather( + const device T* mat [[buffer(0)]], + const device T* in_vec [[buffer(1)]], + const device T* bias [[buffer(2)]], + device T* out_vec [[buffer(3)]], + const constant int& in_vec_size [[buffer(4)]], + const constant int& out_vec_size [[buffer(5)]], + const constant int& M [[buffer(6)]], + const constant int& matrix_ld [[buffer(7)]], + const constant int& vector_ld [[buffer(8)]], + const constant int& batch_ndim [[buffer(11)]], + const constant int* batch_shape [[buffer(12)]], + const constant int64_t* index_batch_strides [[buffer(13)]], + const constant int& vector_batch_ndim [[buffer(14)]], + const constant int* vector_batch_shape [[buffer(15)]], + const constant int64_t* vector_batch_stride [[buffer(16)]], + const constant int& matrix_batch_ndim [[buffer(17)]], + const constant int* matrix_batch_shape [[buffer(18)]], + const constant int64_t* matrix_batch_stride [[buffer(19)]], + const constant uint32_t* vec_indices [[buffer(20)]], + const constant uint32_t* mat_indices [[buffer(21)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tgpg [[threadgroups_per_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + uint32_t indx_vec; + uint32_t indx_mat; + + // Update batch offsets + if (batch_ndim > 1) { + const constant auto* veci_bstrides = index_batch_strides; + const constant auto* mati_bstrides = index_batch_strides + batch_ndim; + + ulong2 batch_offsets = elem_to_loc_broadcast( + tid.z, batch_shape, veci_bstrides, mati_bstrides, batch_ndim); + + indx_vec = vec_indices[batch_offsets.x]; + indx_mat = mat_indices[batch_offsets.y]; + + } else { + indx_vec = vec_indices[index_batch_strides[0] * tid.z]; + indx_mat = mat_indices[index_batch_strides[batch_ndim] * tid.z]; + } + + if (vector_batch_ndim > 1) { + in_vec += elem_to_loc( + indx_vec, vector_batch_shape, vector_batch_stride, vector_batch_ndim); + } else { + in_vec += indx_vec * vector_batch_stride[0]; + } + + if (matrix_batch_ndim > 1) { + mat += elem_to_loc( + indx_mat, matrix_batch_shape, matrix_batch_stride, matrix_batch_ndim); + } else { + mat += indx_mat * matrix_batch_stride[0]; + } + + out_vec += int64_t(tid.z) * M * out_vec_size; + + GemvWide::run( + mat, + in_vec, + bias, + out_vec, + in_vec_size, + out_vec_size, + M, + matrix_ld, + vector_ld, + 1.0f, + 0.0f, + 0, + 0, + tid, + tgpg, + simd_gid, + simd_lid); +} diff --git a/Source/Cmlx/mlx-generated/metal/gemv_masked.h b/Source/Cmlx/mlx-generated/metal/gemv_masked.h index 9d4fac234..30371eb06 100644 --- a/Source/Cmlx/mlx-generated/metal/gemv_masked.h +++ b/Source/Cmlx/mlx-generated/metal/gemv_masked.h @@ -10,7 +10,7 @@ using namespace metal; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -30,7 +30,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h index 99e65d205..3a97de9bb 100644 --- a/Source/Cmlx/mlx-generated/metal/indexing/scatter.h +++ b/Source/Cmlx/mlx-generated/metal/indexing/scatter.h @@ -57,3 +57,83 @@ METAL_FUNC void scatter_impl( op.atomic_update(out, updates[upd_idx], out_idx); } } + +template < + typename T, + typename IdxT, + typename Op, + bool OUT_ROW_CONTIG, + bool UPD_ROW_CONTIG, + bool UPD_SCALAR, + int NWORK, + int NDIM> +[[kernel]] void slice_update_op_impl( + const device T* updates [[buffer(0)]], + device T* out [[buffer(1)]], + const constant int* update_shape [[buffer(2)]], + const constant int64_t* update_strides [[buffer(3)]], + const constant int& update_ndim [[buffer(4)]], + const constant int64_t& update_size [[buffer(5)]], + const constant int64_t* output_strides [[buffer(6)]], + const constant int64_t& output_offset [[buffer(7)]], + uint3 gid [[thread_position_in_grid]], + uint3 gsize [[threads_per_grid]]) { + Op op; + + IdxT idx = + (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK + + IdxT(gid.x) * NWORK; + IdxT out_idx; + IdxT update_idx; + + if constexpr (OUT_ROW_CONTIG) { + out_idx = idx; + } else if constexpr (NDIM == 1) { + out_idx = NWORK * gid.x * output_strides[0]; + } else if constexpr (NDIM == 2) { + out_idx = gid.y * output_strides[0] + NWORK * gid.x * output_strides[1]; + } else if constexpr (NDIM == 3) { + out_idx = gid.z * output_strides[0] + gid.y * output_strides[1] + + NWORK * gid.x * output_strides[2]; + } else { + out_idx = elem_to_loc(idx, update_shape, output_strides, update_ndim); + } + + if constexpr (UPD_SCALAR) { + update_idx = 0; + } else if constexpr (UPD_ROW_CONTIG) { + update_idx = idx; + } else if constexpr (NDIM == 1) { + update_idx = NWORK * gid.x * update_strides[0]; + } else if constexpr (NDIM == 2) { + update_idx = gid.y * update_strides[0] + NWORK * gid.x * update_strides[1]; + } else if constexpr (NDIM == 3) { + update_idx = gid.z * update_strides[0] + gid.y * update_strides[1] + + NWORK * gid.x * update_strides[2]; + } else { + update_idx = + elem_to_loc(idx, update_shape, update_strides, update_ndim); + } + + out += output_offset; + + if constexpr (OUT_ROW_CONTIG && (UPD_ROW_CONTIG || UPD_SCALAR)) { + for (int j = 0; j < NWORK; j++) { + out[out_idx] = op(out[out_idx], updates[update_idx]); + out_idx++; + if constexpr (!UPD_SCALAR) { + update_idx++; + } + } + } else { + auto out_stride = output_strides[update_ndim - 1]; + auto update_stride = update_strides[update_ndim - 1]; + for (int j = 0; j < NWORK; j++) { + out[out_idx] = op(out[out_idx], updates[update_idx]); + out_idx += out_stride; + if constexpr (!UPD_SCALAR) { + update_idx += update_stride; + } + } + } +} diff --git a/Source/Cmlx/mlx-generated/metal/logging.h b/Source/Cmlx/mlx-generated/metal/logging.h index 7b3ee0467..d7446c722 100644 --- a/Source/Cmlx/mlx-generated/metal/logging.h +++ b/Source/Cmlx/mlx-generated/metal/logging.h @@ -16,7 +16,7 @@ struct os_log { constexpr os_log(constant char*, constant char*) constant {} template - void log_debug(constant char*, Args...) const {} + void log_debug(constant char*, Args...) const thread {} template void log_debug(constant char*, Args...) const constant {} diff --git a/Source/Cmlx/mlx-generated/metal/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index 5ac4c6e16..42855e14d 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized.h +++ b/Source/Cmlx/mlx-generated/metal/quantized.h @@ -480,25 +480,25 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { } } -template -inline void -dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { +// Decode one quantized block (scale * q + bias) into w_local. W (the output +// pointer type) serves the threadgroup block loader or a thread-local decode. +template +inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -507,22 +507,24 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -531,14 +533,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -546,16 +552,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } @@ -609,15 +617,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -625,7 +633,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -638,7 +646,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -668,7 +676,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -974,6 +982,103 @@ METAL_FUNC void qmv_impl( } } +// Affine analog of fp_qmv_wide. Weights carry a scale and bias per group, so +// each group is decoded in 8-value sub-chunks (scale * q + bias, registers +// bounded for any group_size) and reused across the vecs_per_tg vectors. +template +METAL_FUNC void qmv_wide_impl( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int sub = 8; // values per sub-chunk (== bits bytes, byte-aligned) + + typedef float U; + + const short k_lane = simd_lid % k_lanes; + const short sg_row = simd_lid / k_lanes; + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bits / 8; // bytes per weight row + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device T* srow = scales + row * in_vec_size_g; + const device T* brow = biases + row * in_vec_size_g; + + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's groups: decode the group in + // 8-value sub-chunks and reuse each chunk across the streamed vectors. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + U scale = srow[g]; + U bias = brow[g]; +#pragma unroll + for (int sc = 0; sc < group_size / sub; sc++) { + const int k0 = g * group_size + sc * sub; + const device uint8_t* wc = wrow + k0 * bits / 8; + U w_dq[sub]; + dequantize(wc, scale, bias, w_dq); +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device T* xc = xv[v] + k0; + U acc = 0; +#pragma unroll + for (int i = 0; i < sub; i++) { + acc += static_cast(xc[i]) * w_dq[i]; + } + result[v] += acc; + } + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); + } + } + } +} + template METAL_FUNC void qvm_impl( const device uint32_t* w, @@ -983,6 +1088,7 @@ METAL_FUNC void qvm_impl( device T* y, const int in_vec_size, const int out_vec_size, + const int in_vec_stride, uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { @@ -1016,7 +1122,7 @@ METAL_FUNC void qvm_impl( ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w; scales += out_col / group_size + simd_lid * out_vec_size_g; biases += out_col / group_size + simd_lid * out_vec_size_g; - x += tid.x * in_vec_size + simd_lid; + x += tid.x * in_vec_stride + simd_lid; y += tid.x * out_vec_size + out_col; if (out_col >= out_vec_size) { @@ -1102,6 +1208,7 @@ METAL_FUNC void qmm_t_impl( const constant int& K, const constant int& N, const constant int& M, + const constant int& K_eff, uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], @@ -1156,7 +1263,7 @@ METAL_FUNC void qmm_t_impl( if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -1166,7 +1273,7 @@ METAL_FUNC void qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -1178,7 +1285,7 @@ METAL_FUNC void qmm_t_impl( } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -1188,7 +1295,7 @@ METAL_FUNC void qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -1491,7 +1598,13 @@ template quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1543,7 +1656,13 @@ template simd_lid); } -template +template < + typename T, + int group_size, + const int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1595,7 +1714,71 @@ template simd_lid); } -template +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void affine_qmv_wide( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + const constant int64_t* b_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + b_strides, + tid); + } + qmv_wide_impl( + w, + scales, + biases, + x, + y, + in_vec_size, + out_vec_size, + M, + tid, + simd_gid, + simd_lid); +} + +template < + typename T, + const int group_size, + const int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void affine_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1642,6 +1825,7 @@ template y, in_vec_size, out_vec_size, + in_vec_size, tid, simd_gid, simd_lid); @@ -1690,6 +1874,9 @@ template int in_vec_size_adj = tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size; + // The in_vec_stride is the full K dimension, not the partition size + int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; + qvm_impl( w, scales, @@ -1698,6 +1885,7 @@ template y, in_vec_size_adj, out_vec_size, + in_vec_stride, tid, simd_gid, simd_lid); @@ -1759,17 +1947,91 @@ template < tid); } qmm_t_impl( - w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < typename T, const int group_size, const int bits, - const bool batched, + const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32> +[[kernel]] void affine_qmm_t_splitk( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* biases [[buffer(2)]], + const device T* x [[buffer(3)]], + device T* y [[buffer(4)]], + const constant int& K [[buffer(5)]], + const constant int& N [[buffer(6)]], + const constant int& M [[buffer(7)]], + const constant int& k_partition_size [[buffer(8)]], + const constant int& split_k_partition_stride [[buffer(9)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + (void)lid; + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + + const int k_start = tid.z * k_partition_size; + x += k_start; + + auto wl = (const device uint8_t*)w; + wl += k_start * bytes_per_pack / pack_factor; + scales += k_start / group_size; + biases += k_start / group_size; + y += tid.z * static_cast(split_k_partition_stride); + + qmm_t_impl( + (const device uint32_t*)wl, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + k_partition_size, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1822,7 +2084,7 @@ template < w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void affine_gather_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1884,7 +2146,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1946,7 +2208,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2003,6 +2265,7 @@ template y, in_vec_size, out_vec_size, + in_vec_size, tid, simd_gid, simd_lid); @@ -2073,16 +2336,31 @@ template < b_strides, tid); qmm_t_impl( - w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_gather_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2340,7 +2618,7 @@ template < } } -template +template [[kernel]] void affine_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], @@ -2445,7 +2723,7 @@ template } } -template +template [[kernel]] void affine_dequantize( const device uint8_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], diff --git a/Source/Cmlx/mlx-generated/metal/quantized_nax.h b/Source/Cmlx/mlx-generated/metal/quantized_nax.h index c26ff646b..ed32eb59a 100644 --- a/Source/Cmlx/mlx-generated/metal/quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/quantized_nax.h @@ -491,17 +491,16 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -510,22 +509,24 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -534,14 +535,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -549,16 +554,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } @@ -612,15 +619,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -628,7 +635,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -641,7 +648,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -671,7 +678,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -752,14 +759,14 @@ struct QuantizedBlockLoader< const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -768,7 +775,7 @@ struct QuantizedBlockLoader< scales(scales_ + bi * src_ld / group_size + group_id), biases(biases_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -781,7 +788,7 @@ struct QuantizedBlockLoader< } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -811,7 +818,7 @@ struct QuantizedBlockLoader< } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { // if (group_steps > 1) { @@ -826,8 +833,8 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; + scales += group_stride; + biases += group_stride; } } }; @@ -986,16 +993,13 @@ METAL_FUNC void qmm_t_nax_tgp_impl( // Make the weight loader loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -1003,22 +1007,17 @@ METAL_FUNC void qmm_t_nax_tgp_impl( constexpr bool transpose_a = false; constexpr bool transpose_b = true; - const short sgp_sm = min(SM, short(M - (y_row + tm))); + const short sgp_sm = min(int(SM), M - (y_row + tm)); const bool is_unaligned_sm = (sgp_sm != SM); - const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn))); + const short sgp_sn = aligned_N ? SN : min(int(SN), N - (y_col + tn)); const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -1037,8 +1036,8 @@ METAL_FUNC void qmm_t_nax_tgp_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -1102,7 +1101,6 @@ METAL_FUNC void qmm_n_nax_tgp_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { (void)lid; - (void)M; static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -1123,38 +1121,34 @@ METAL_FUNC void qmm_n_nax_tgp_impl( bits>; // Set the block - const int K_w = K * bytes_per_pack / pack_factor; - const int K_g = K / group_size; const int y_row = tid.y * BM; const int y_col = tid.x * BN; auto wl = (const device uint8_t*)w; + // Here w is [K, N]: packed and group-quantized along N, with row stride N. x += y_row * static_cast(K); - wl += y_col * K_w; - scales += y_col * K_g; - biases += y_col * K_g; + wl += y_col * bytes_per_pack / pack_factor; + scales += y_col / group_size; + biases += y_col / group_size; y += y_row * static_cast(N) + y_col; - // Make the x loader and mma operation - // const short num_els = min(BM, M - y_row); - // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + // Make the weight loader + loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); + const short sgp_sm = min(int(SM), M - (y_row + tm)); + const short ldb_tgp = BN_padded; constexpr bool transpose_a = false; @@ -1162,12 +1156,7 @@ METAL_FUNC void qmm_n_nax_tgp_impl( using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -1179,12 +1168,17 @@ METAL_FUNC void qmm_n_nax_tgp_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; - Atile.load(x + kk1, K); + if (sgp_sm == SM) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + Btile.template load(Ws + tn + kk1 * ldb_tgp); tile_matmad_nax( @@ -1204,7 +1198,11 @@ METAL_FUNC void qmm_n_nax_tgp_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - Dtile.store(y + tm * N + tn, N); + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm)); + } } template < @@ -1534,16 +1532,13 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); @@ -1561,10 +1556,6 @@ template < using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - // Do as many matmuls as necessary uint32_t index; short offset; @@ -1585,8 +1576,11 @@ template < } threadgroup_barrier(mem_flags::mem_none); - NAXTile Dtile; + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); const device T* xn = x + tm * K; @@ -1616,31 +1610,33 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -1654,36 +1650,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - volatile int compiler_barrier; + volatile int compiler_barrier; - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; + (void)compiler_barrier; + } } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { @@ -1702,4 +1697,4 @@ template < }); }); } -} \ No newline at end of file +} diff --git a/Source/Cmlx/mlx-generated/metal/random.metal b/Source/Cmlx/mlx-generated/metal/random.metal index eb6234d8e..b434b7320 100644 --- a/Source/Cmlx/mlx-generated/metal/random.metal +++ b/Source/Cmlx/mlx-generated/metal/random.metal @@ -35,7 +35,7 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) { device const uint32_t* keys, device char* out, constant const bool& odd, - constant const uint& bytes_per_key, + constant const ulong& bytes_per_key, uint2 grid_dim [[threads_per_grid]], uint2 index [[thread_position_in_grid]]) { auto kidx = 2 * index.x; @@ -68,7 +68,7 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) { device const uint32_t* keys, device char* out, constant const bool& odd, - constant const uint& bytes_per_key, + constant const ulong& bytes_per_key, constant const int& ndim, constant const int* key_shape, constant const int64_t* key_strides, diff --git a/Source/Cmlx/mlx-generated/metal/reduction/ops.h b/Source/Cmlx/mlx-generated/metal/reduction/ops.h index 11d8e83ac..f875e51c9 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/ops.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/ops.h @@ -7,12 +7,12 @@ #define DEFINE_SIMD_REDUCE() \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ return simd_reduce_impl(val); \ } \ \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ for (short i = simd_size / 2; i > 0; i /= 2) { \ val = operator()(val, simd_shuffle_down(val, i)); \ } \ @@ -28,7 +28,8 @@ union bool4_or_uint { struct None { template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_store_explicit(out, val, offset); } }; @@ -37,7 +38,7 @@ template struct And { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_all(val); } @@ -47,7 +48,7 @@ struct And { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (!val) { bool4_or_uint update; update.b = {true, true, true, true}; @@ -56,20 +57,20 @@ struct And { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (!val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out &= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a && b; } }; @@ -78,7 +79,7 @@ template struct Or { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_any(val); } @@ -88,7 +89,7 @@ struct Or { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (val) { bool4_or_uint update; update.b = {false, false, false, false}; @@ -97,20 +98,20 @@ struct Or { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out |= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a || b; } }; @@ -120,19 +121,29 @@ struct Sum { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_sum(val); } static constexpr constant U init = U(0); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_add_explicit(out, val, offset); } + void atomic_update( + device mlx_atomic* out, + complex64_t val, + size_t offset = 0) thread { + auto out_lanes = reinterpret_cast*>(out); + mlx_atomic_fetch_add_explicit(out_lanes, val.real, 2 * offset); + mlx_atomic_fetch_add_explicit(out_lanes, val.imag, 2 * offset + 1); + } + // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a + b; } }; @@ -142,19 +153,20 @@ struct Prod { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_product(val); } static constexpr constant U init = U(1); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_mul_explicit(out, val, offset); } // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a * b; } }; @@ -164,12 +176,14 @@ struct Min { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_min(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -179,18 +193,19 @@ struct Min { static constexpr constant U init = Limits::max; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_min_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a < b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -199,7 +214,7 @@ struct Min { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); @@ -221,12 +236,14 @@ struct Max { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_max(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -236,18 +253,19 @@ struct Max { static constexpr constant U init = Limits::min; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_max_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a > b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -256,7 +274,7 @@ struct Max { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_all.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_all.h index e0d08392c..47ad63fbb 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_all.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_all.h @@ -37,13 +37,13 @@ template < for (IdxT b = 0; b < blocks; b++) { for (int i = 0; i < N_READS; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } in += lsize.x * N_READS; } if (extra > 0) { for (int i = 0; i < extra; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } } diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h index c109faf0b..b1546adb5 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_col.h @@ -43,13 +43,13 @@ template row = in + loop.location(); if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -125,7 +125,7 @@ template for (IdxT r = gid.z * lsize.y + lid.y; r < total_rows; r += lsize.y * gsize.z) { row = in + loop.location(); - total = op(static_cast(*row), total); + total = op(cast_to(*row), total); loop.next(lsize.y * gsize.z, reduce_shape, reduce_strides); } @@ -207,13 +207,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -352,13 +352,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); diff --git a/Source/Cmlx/mlx-generated/metal/reduction/reduce_row.h b/Source/Cmlx/mlx-generated/metal/reduction/reduce_row.h index 936d75bb5..09d1d8834 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/reduce_row.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/reduce_row.h @@ -34,7 +34,7 @@ METAL_FUNC void per_thread_row_reduce( for (int i = 0; i < blocks; i++) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } inputs[j] += lsize_x * N_READS; @@ -46,13 +46,13 @@ METAL_FUNC void per_thread_row_reduce( if (index + N_READS <= extra) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } else { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; index + i < extra; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } @@ -337,7 +337,8 @@ template < // lid.x * N_READS breaks the per_thread_row_reduce interface a bit. Maybe it // needs a small refactor. - in += elem_to_loc(out_idx, shape, strides, ndim) + lid.x * N_READS; + in += + elem_to_loc(out_idx, shape, strides, ndim) + IdxT(lid.x) * N_READS; LoopedElemToLoc 2)> loop(reduce_ndim); const device T* row; diff --git a/Source/Cmlx/mlx-generated/metal/rms_norm.metal b/Source/Cmlx/mlx-generated/metal/rms_norm.metal index 22fae273e..c2ebe6dfc 100644 --- a/Source/Cmlx/mlx-generated/metal/rms_norm.metal +++ b/Source/Cmlx/mlx-generated/metal/rms_norm.metal @@ -27,19 +27,18 @@ template threadgroup float local_sums[SIMD_SIZE]; float acc = 0; + float thread_x[N_READS]; x += gid * size_t(axis_size) + lid * N_READS; w += w_stride * lid * N_READS; if (lid * N_READS + N_READS <= axis_size) { for (int i = 0; i < N_READS; i++) { - float xi = x[i]; - acc += xi * xi; + thread_x[i] = x[i]; + acc += thread_x[i] * thread_x[i]; } } else { for (int i = 0; i < N_READS; i++) { - if ((lid * N_READS + i) < axis_size) { - float xi = x[i]; - acc += xi * xi; - } + thread_x[i] = (lid * N_READS + i < axis_size) ? (float)x[i] : 0; + acc += thread_x[i] * thread_x[i]; } } acc = simd_sum(acc); @@ -64,16 +63,18 @@ template } threadgroup_barrier(mem_flags::mem_threadgroup); - // Write the outputs + // Write the outputs using cached x values out += gid * size_t(axis_size) + lid * N_READS; if (lid * N_READS + N_READS <= axis_size) { for (int i = 0; i < N_READS; i++) { - out[i] = w[w_stride * i] * static_cast(x[i] * local_inv_mean[0]); + out[i] = + w[w_stride * i] * static_cast(thread_x[i] * local_inv_mean[0]); } } else { for (int i = 0; i < N_READS; i++) { if ((lid * N_READS + i) < axis_size) { - out[i] = w[w_stride * i] * static_cast(x[i] * local_inv_mean[0]); + out[i] = + w[w_stride * i] * static_cast(thread_x[i] * local_inv_mean[0]); } } } @@ -165,21 +166,32 @@ template constant float& eps, constant uint& axis_size, constant uint& w_stride, + constant uint& n_rows, + constant uint& rows_per_group, uint gid [[threadgroup_position_in_grid]], uint lid [[thread_position_in_threadgroup]], uint simd_lane_id [[thread_index_in_simdgroup]], uint simd_group_id [[simdgroup_index_in_threadgroup]]) { - // Advance the input pointers - x += gid * size_t(axis_size) + lid * N_READS; - g += gid * size_t(axis_size) + lid * N_READS; w += w_stride * lid * N_READS; + float thread_w[N_READS]; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + thread_w[i] = w[w_stride * i]; + } + } else { + for (int i = 0; i < N_READS; i++) { + thread_w[i] = + (lid * N_READS + i < axis_size) ? (float)w[w_stride * i] : 0; + } + } // Allocate registers for the computation and accumulators float thread_x[N_READS]; - float thread_w[N_READS]; float thread_g[N_READS]; - float sumx2 = 0; - float sumgwx = 0; + float gw_acc[N_READS]; + for (int i = 0; i < N_READS; i++) { + gw_acc[i] = 0; + } // Allocate shared memory to implement the reduction constexpr int SIMD_SIZE = 32; @@ -188,75 +200,99 @@ template threadgroup float local_normalizer[1]; threadgroup float local_meangwx[1]; - // Read and accumulate locally - if (lid * N_READS + N_READS <= axis_size) { - for (int i = 0; i < N_READS; i++) { - thread_x[i] = x[i]; - thread_w[i] = w[w_stride * i]; - thread_g[i] = g[i]; + uint row_end = gid * rows_per_group + rows_per_group; + if (row_end > n_rows) { + row_end = n_rows; + } + for (uint row = gid * rows_per_group; row < row_end; ++row) { + const device T* x_in = x + size_t(row) * axis_size + lid * N_READS; + const device T* g_in = g + size_t(row) * axis_size + lid * N_READS; - sumx2 += thread_x[i] * thread_x[i]; - sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; - } - } else { - for (int i = 0; i < N_READS; i++) { - if ((lid * N_READS + i) < axis_size) { - thread_x[i] = x[i]; - thread_w[i] = w[w_stride * i]; - thread_g[i] = g[i]; + float sumx2 = 0; + float sumgwx = 0; + + // Read and accumulate locally + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + thread_x[i] = x_in[i]; + thread_g[i] = g_in[i]; sumx2 += thread_x[i] * thread_x[i]; sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + thread_x[i] = x_in[i]; + thread_g[i] = g_in[i]; + + sumx2 += thread_x[i] * thread_x[i]; + sumgwx += thread_x[i] * thread_w[i] * thread_g[i]; + } + } } - } - // Accumulate across threads - sumx2 = simd_sum(sumx2); - sumgwx = simd_sum(sumgwx); - if (simd_group_id == 0) { - local_sumx2[simd_lane_id] = 0; - local_sumgwx[simd_lane_id] = 0; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (simd_lane_id == 0) { - local_sumx2[simd_group_id] = sumx2; - local_sumgwx[simd_group_id] = sumgwx; - } - threadgroup_barrier(mem_flags::mem_threadgroup); - if (simd_group_id == 0) { - sumx2 = simd_sum(local_sumx2[simd_lane_id]); - sumgwx = simd_sum(local_sumgwx[simd_lane_id]); + // Accumulate across threads + sumx2 = simd_sum(sumx2); + sumgwx = simd_sum(sumgwx); + if (simd_group_id == 0) { + local_sumx2[simd_lane_id] = 0; + local_sumgwx[simd_lane_id] = 0; + } + threadgroup_barrier(mem_flags::mem_threadgroup); if (simd_lane_id == 0) { - local_meangwx[0] = sumgwx / axis_size; - local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps); + local_sumx2[simd_group_id] = sumx2; + local_sumgwx[simd_group_id] = sumgwx; } - } - threadgroup_barrier(mem_flags::mem_threadgroup); - float meangwx = local_meangwx[0]; - float normalizer = local_normalizer[0]; - float normalizer3 = normalizer * normalizer * normalizer; - - // Write the outputs - gx += gid * size_t(axis_size) + lid * N_READS; - gw += gid * size_t(axis_size) + lid * N_READS; - if (lid * N_READS + N_READS <= axis_size) { - for (int i = 0; i < N_READS; i++) { - gx[i] = static_cast( - thread_g[i] * thread_w[i] * normalizer - - thread_x[i] * meangwx * normalizer3); - if (has_w) { - gw[i] = static_cast(thread_g[i] * thread_x[i] * normalizer); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (simd_group_id == 0) { + sumx2 = simd_sum(local_sumx2[simd_lane_id]); + sumgwx = simd_sum(local_sumgwx[simd_lane_id]); + if (simd_lane_id == 0) { + local_meangwx[0] = sumgwx / axis_size; + local_normalizer[0] = metal::precise::rsqrt(sumx2 / axis_size + eps); } } - } else { - for (int i = 0; i < N_READS; i++) { - if ((lid * N_READS + i) < axis_size) { - gx[i] = static_cast( + threadgroup_barrier(mem_flags::mem_threadgroup); + float meangwx = local_meangwx[0]; + float normalizer = local_normalizer[0]; + float normalizer3 = normalizer * normalizer * normalizer; + + // Write the outputs + device T* gx_out = gx + size_t(row) * axis_size + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + gx_out[i] = static_cast( thread_g[i] * thread_w[i] * normalizer - thread_x[i] * meangwx * normalizer3); if (has_w) { - gw[i] = static_cast(thread_g[i] * thread_x[i] * normalizer); + gw_acc[i] += thread_g[i] * thread_x[i] * normalizer; + } + } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + gx_out[i] = static_cast( + thread_g[i] * thread_w[i] * normalizer - + thread_x[i] * meangwx * normalizer3); + if (has_w) { + gw_acc[i] += thread_g[i] * thread_x[i] * normalizer; + } + } + } + } + } + + if (has_w) { + gw += size_t(gid) * axis_size + lid * N_READS; + if (lid * N_READS + N_READS <= axis_size) { + for (int i = 0; i < N_READS; i++) { + gw[i] = static_cast(gw_acc[i]); + } + } else { + for (int i = 0; i < N_READS; i++) { + if ((lid * N_READS + i) < axis_size) { + gw[i] = static_cast(gw_acc[i]); } } } diff --git a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal index ae04c6bab..8fdd08762 100644 --- a/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/scaled_dot_product_attention.metal @@ -28,14 +28,29 @@ using namespace metal; qk_dim, \ value_dim) +#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt) \ + instantiate_kernel( \ + "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \ + sdpa_vector_2pass_1_gqa, \ + type, \ + qk_dim, \ + value_dim, \ + 8, \ + hpt) + #define instantiate_sdpa_vector_heads(type) \ instantiate_sdpa_vector(type, 64, 64) \ instantiate_sdpa_vector(type, 96, 96) \ instantiate_sdpa_vector(type, 128, 128) \ + instantiate_sdpa_vector(type, 192, 128) \ + instantiate_sdpa_vector(type, 192, 192) \ instantiate_sdpa_vector(type, 256, 256) \ + instantiate_sdpa_vector_gqa(type, 64, 64, 8) \ + instantiate_sdpa_vector_gqa(type, 128, 128, 4) \ instantiate_sdpa_vector_aggregation(type, 64) \ instantiate_sdpa_vector_aggregation(type, 96) \ instantiate_sdpa_vector_aggregation(type, 128) \ + instantiate_sdpa_vector_aggregation(type, 192) \ instantiate_sdpa_vector_aggregation(type, 256) instantiate_sdpa_vector_heads(float) diff --git a/Source/Cmlx/mlx-generated/metal/scan.h b/Source/Cmlx/mlx-generated/metal/scan.h index a1f10340b..9907171b0 100644 --- a/Source/Cmlx/mlx-generated/metal/scan.h +++ b/Source/Cmlx/mlx-generated/metal/scan.h @@ -6,12 +6,12 @@ #define DEFINE_SIMD_SCAN() \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ return simd_scan_impl(val); \ } \ \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ for (int i = 1; i <= 16; i *= 2) { \ val = operator()(val, simd_shuffle_and_fill_up(val, init, i)); \ } \ @@ -20,12 +20,12 @@ #define DEFINE_SIMD_EXCLUSIVE_SCAN() \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ return simd_exclusive_scan_impl(val); \ } \ \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ val = simd_scan(val); \ return simd_shuffle_and_fill_up(val, init, 1); \ } @@ -38,15 +38,15 @@ struct CumSum { static constexpr constant U init = static_cast(0); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a + b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_sum(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_sum(x); } }; @@ -59,15 +59,15 @@ struct CumProd { static constexpr constant U init = static_cast(1.0f); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a * b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_product(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_product(x); } }; @@ -77,11 +77,11 @@ struct CumProd { static constexpr constant bool init = true; template - bool operator()(bool a, T b) { + bool operator()(bool a, T b) thread { return a & static_cast(b); } - bool simd_scan(bool x) { + bool simd_scan(bool x) thread { for (int i = 1; i <= 16; i *= 2) { bool other = simd_shuffle_and_fill_up(x, init, i); x &= other; @@ -89,7 +89,7 @@ struct CumProd { return x; } - bool simd_exclusive_scan(bool x) { + bool simd_exclusive_scan(bool x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -99,20 +99,29 @@ template struct CumMax { static constexpr constant U init = Limits::min; - template - U operator()(U a, T b) { + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } return (a >= b) ? a : b; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + return combine(a, static_cast(b)); + } + + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x >= other) ? x : other; + x = combine(x, other); } return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -122,20 +131,29 @@ template struct CumMin { static constexpr constant U init = Limits::max; - template - U operator()(U a, T b) { + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } return (a <= b) ? a : b; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + return combine(a, static_cast(b)); + } + + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x <= other) ? x : other; + x = combine(x, other); } return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -146,11 +164,11 @@ struct CumLogaddexp { static constexpr constant U init = Limits::min; template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return LogAddExp{}(a, static_cast(b)); } - U simd_scan(U x) { + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); x = LogAddExp{}(x, other); @@ -158,7 +176,7 @@ struct CumLogaddexp { return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h index 1eec72be3..3f40dbd7c 100644 --- a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h +++ b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h @@ -317,6 +317,149 @@ template } } +// Duplication-free variant for high gqa_factor decode: each simdgroup owns a +// contiguous token sub-chunk and computes HPT of its group's query heads, so +// each K/V byte is read G / HPT times instead of G times. Single-token +// queries without mask or sinks only; the partials layout matches +// sdpa_vector_2pass_2. +template +[[kernel]] void sdpa_vector_2pass_1_gqa( + const device T* queries [[buffer(0)]], + const device T* keys [[buffer(1)]], + const device T* values [[buffer(2)]], + device T* out [[buffer(3)]], + device float* sums [[buffer(4)]], + device float* maxs [[buffer(5)]], + const constant int& N [[buffer(7)]], + const constant size_t& k_head_stride [[buffer(8)]], + const constant size_t& k_seq_stride [[buffer(9)]], + const constant size_t& v_head_stride [[buffer(10)]], + const constant size_t& v_seq_stride [[buffer(11)]], + const constant float& scale [[buffer(12)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 tpg [[threadgroups_per_grid]], + uint3 tidtg [[thread_position_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int BD = 32; + constexpr int qk_per_thread = D / BD; + constexpr int v_per_thread = V / BD; + constexpr int NT = G / HPT; + + typedef float U; + + const int kv_head_idx = tid.x; + const int batch_idx = tid.y; + const int block_idx = tid.z; + const int blocks = tpg.z; + const int g = tidtg.y; + const int cchunk = g / NT; + const int h0 = (g % NT) * HPT; + const int num_kv_heads = tpg.x; + const int num_q_heads = num_kv_heads * G; + const int base_head = batch_idx * num_q_heads + kv_head_idx * G; + + const int chunk = (N + blocks - 1) / blocks; + const int kstart = block_idx * chunk; + const int kend = min(N, kstart + chunk); + const int sub = (chunk + HPT - 1) / HPT; + const int s0 = kstart + cchunk * sub; + const int s1 = min(kend, s0 + sub); + + const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride + + simd_lid * qk_per_thread; + const device T* vp = values + kv_head_idx * v_head_stride + + s0 * v_seq_stride + simd_lid * v_per_thread; + + U q[HPT][qk_per_thread]; + for (int j = 0; j < HPT; j++) { + const device T* qp = + queries + (base_head + h0 + j) * D + simd_lid * qk_per_thread; + for (int i = 0; i < qk_per_thread; i++) { + q[j][i] = static_cast(scale) * qp[i]; + } + } + + U max_score[HPT]; + U sum_exp_score[HPT]; + U o[HPT][v_per_thread]; + for (int j = 0; j < HPT; j++) { + max_score[j] = Limits::finite_min; + sum_exp_score[j] = 0; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = 0; + } + } + + for (int t = s0; t < s1; t++) { + U kr[qk_per_thread]; + U vr[v_per_thread]; + for (int i = 0; i < qk_per_thread; i++) { + kr[i] = kp[i]; + } + for (int i = 0; i < v_per_thread; i++) { + vr[i] = vp[i]; + } + kp += k_seq_stride; + vp += v_seq_stride; + for (int j = 0; j < HPT; j++) { + U score = 0; + for (int i = 0; i < qk_per_thread; i++) { + score += q[j][i] * kr[i]; + } + score = simd_sum(score); + U new_max = max(max_score[j], score); + U factor = fast::exp(max_score[j] - new_max); + U exp_score = fast::exp(score - new_max); + max_score[j] = new_max; + sum_exp_score[j] = sum_exp_score[j] * factor + exp_score; + for (int i = 0; i < v_per_thread; i++) { + o[j][i] = o[j][i] * factor + exp_score * vr[i]; + } + } + } + + threadgroup U o_sh[G * HPT * V]; + threadgroup U se_sh[G * HPT]; + threadgroup U mx_sh[G * HPT]; + for (int j = 0; j < HPT; j++) { + int slot = (h0 + j) * HPT + cchunk; + U inv = sum_exp_score[j] > 0 ? 1 / sum_exp_score[j] : 0; + for (int i = 0; i < v_per_thread; i++) { + o_sh[slot * V + simd_lid * v_per_thread + i] = o[j][i] * inv; + } + if (simd_lid == 0) { + se_sh[slot] = sum_exp_score[j]; + mx_sh[slot] = max_score[j]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + U gmax = Limits::finite_min; + for (int s = 0; s < HPT; s++) { + gmax = max(gmax, mx_sh[g * HPT + s]); + } + U denom = 0; + U acc[v_per_thread] = {0}; + for (int s = 0; s < HPT; s++) { + U w = se_sh[g * HPT + s] * fast::exp(mx_sh[g * HPT + s] - gmax); + denom += w; + for (int i = 0; i < v_per_thread; i++) { + acc[i] += w * o_sh[(g * HPT + s) * V + simd_lid * v_per_thread + i]; + } + } + + const int o_offset = base_head + g; + device T* op = + out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread; + for (int i = 0; i < v_per_thread; i++) { + op[i] = static_cast(acc[i]); + } + if (simd_lid == 0) { + sums[o_offset * blocks + block_idx] = denom; + maxs[o_offset * blocks + block_idx] = gmax; + } +} + template [[kernel]] void sdpa_vector_2pass_2( const device T* partials [[buffer(0)]], diff --git a/Source/Cmlx/mlx-generated/metal/searchsorted.h b/Source/Cmlx/mlx-generated/metal/searchsorted.h new file mode 100644 index 000000000..87b7c3f5b --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/searchsorted.h @@ -0,0 +1,34 @@ +// Copyright © 2026 Apple Inc. + +template +METAL_FUNC uint +searchsorted_impl(device const T* a, T v, uint n, int64_t a_stride) { + LessThan lt; + uint lo = 0; + uint hi = n; + while (lo < hi) { + uint mid = lo + (hi - lo) / 2; + T m = a[int64_t(mid) * a_stride]; + bool below = Right ? !lt(v, m) : lt(m, v); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +template +[[kernel]] void searchsorted( + device const T* a [[buffer(0)]], + device const T* v [[buffer(1)]], + device uint* out [[buffer(2)]], + constant const uint& n [[buffer(3)]], + constant const int64_t& a_stride [[buffer(4)]], + uint3 index [[thread_position_in_grid]], + uint3 grid_dim [[threads_per_grid]]) { + auto offset = + index.x + grid_dim.x * (int64_t(index.y) + int64_t(grid_dim.y) * index.z); + out[offset] = searchsorted_impl(a, v[offset], n, a_stride); +} diff --git a/Source/Cmlx/mlx-generated/metal/searchsorted.metal b/Source/Cmlx/mlx-generated/metal/searchsorted.metal new file mode 100644 index 000000000..d041e7257 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/searchsorted.metal @@ -0,0 +1,28 @@ +// Copyright © 2026 Apple Inc. + +#include + +// clang-format off +#include "utils.h" +#include "sort.h" +#include "searchsorted.h" + +#define instantiate_searchsorted(tname, type) \ + instantiate_kernel("searchsorted_" #tname "_left", \ + searchsorted, type, false) \ + instantiate_kernel("searchsorted_" #tname "_right", \ + searchsorted, type, true) + +instantiate_searchsorted(bool_, bool) +instantiate_searchsorted(uint8, uint8_t) +instantiate_searchsorted(uint16, uint16_t) +instantiate_searchsorted(uint32, uint32_t) +instantiate_searchsorted(uint64, uint64_t) +instantiate_searchsorted(int8, int8_t) +instantiate_searchsorted(int16, int16_t) +instantiate_searchsorted(int32, int32_t) +instantiate_searchsorted(int64, int64_t) +instantiate_searchsorted(float16, half) +instantiate_searchsorted(float32, float) +instantiate_searchsorted(bfloat16, bfloat16_t) +instantiate_searchsorted(complex64, complex64_t) // clang-format on diff --git a/Source/Cmlx/mlx-generated/metal/sort.h b/Source/Cmlx/mlx-generated/metal/sort.h index 0d357333a..ea2640bac 100644 --- a/Source/Cmlx/mlx-generated/metal/sort.h +++ b/Source/Cmlx/mlx-generated/metal/sort.h @@ -29,14 +29,26 @@ struct Init>> { static constexpr constant T v = metal::numeric_limits::quiet_NaN(); }; +template <> +struct Init { + static constexpr constant complex64_t v = complex64_t( + metal::numeric_limits::quiet_NaN(), + metal::numeric_limits::quiet_NaN()); +}; + template struct LessThan { static constexpr constant T init = Init::v; - METAL_FUNC bool operator()(T a, T b) const { - if constexpr ( - metal::is_floating_point_v || metal::is_same_v) { - bool an = isnan(a); - bool bn = isnan(b); + METAL_FUNC bool operator()(T a, T b) const thread { + if constexpr (metal::is_floating_point_v) { + bool an = metal::isnan(a); + bool bn = metal::isnan(b); + if (an | bn) { + return (!an) & bn; + } + } else if constexpr (metal::is_same_v) { + bool an = metal::isnan(a.real) || metal::isnan(a.imag); + bool bn = metal::isnan(b.real) || metal::isnan(b.imag); if (an | bn) { return (!an) & bn; } @@ -376,8 +388,11 @@ template < using ValT = typename sort_kernel::ValT; using IdxT = typename sort_kernel::IdxT; - auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); - auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); + // Signed offsets: a non-sorted axis may have a negative stride. + auto in_block_idx = + elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); + auto out_block_idx = + elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); inp += in_block_idx; out += out_block_idx; @@ -520,7 +535,8 @@ template < BLOCK_THREADS, N_PER_THREAD>; - auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); + // Signed offset: a non-sorted axis may have a negative stride. + auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); inp += block_idx; out_vals += tid.y * size_sorted_axis; out_idxs += tid.y * size_sorted_axis; diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h index df891fa3f..1926aa3cd 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.h @@ -234,11 +234,16 @@ template < } int kb_lim = params->NK; + int kb_min_causal = params->NK; if (do_causal) { int q_max = (tid.x + 1) * BQ + params->qL_off; kb_lim = (q_max + BK - 1) / BK; kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); } // Loop over KV seq length @@ -298,7 +303,7 @@ template < } // Mask out if causal - if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) { + if (do_causal && kb >= kb_min_causal) { using stile_t = decltype(Stile); using selem_t = typename stile_t::elem_type; constexpr auto neg_inf = Limits::finite_min; @@ -423,7 +428,7 @@ template < for (short id = 0; id < TD; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } @@ -433,7 +438,7 @@ template < Vtile.template load( &Vs[Vs_offset + kk * LDV_tgp + dd]); - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal index a68dcfc5c..9efa26fe0 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention.metal @@ -12,8 +12,12 @@ attention, dtype, bq, bk, bd, wm, wn, mtype, float) #define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn(iname, itype, 32, 16, 256, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 16, 192, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 16, 128, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 32, 96, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 80, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 32, 32, 72, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 32, 32, 64, 4, 1, mname, mtype) #define instantiate_attn_mask_helper(iname, itype) \ diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h index 4edc17292..5010f1dee 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/kernels/steel_attention_nax.h @@ -21,9 +21,9 @@ constant bool has_sinks [[function_constant(302)]]; template struct TransformScale { T scale; - METAL_FUNC TransformScale(T scale_) : scale(scale_) {} + METAL_FUNC TransformScale(T scale_) thread : scale(scale_) {} - METAL_FUNC T apply(T x) const { + METAL_FUNC T apply(T x) const thread { return scale * x; } }; @@ -125,38 +125,36 @@ template < make_uniform(params->scale) * make_uniform(1.44269504089f); // Prepare MMA tiles - constexpr short UQ = 16; - constexpr short UD = 32; + constexpr short kU = 16; constexpr int kNWarps = WM * WN; static_assert( - BQ >= (kNWarps * UQ) && BQ % (kNWarps * UQ) == 0, + BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); // Q seq frags per warp - constexpr int TQ = BQ / (kNWarps * UQ); + constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) - constexpr int TD = BD / UD; + constexpr int TD = BD / kU; + // KV seq frags per warp + constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - - using OSubTile = NAXSubTile; - NAXTile Otile; + using otile_t = NAXTile; + otile_t Otile; Otile.clear(); // Prepare mma tile offsets - const short2 simd_coord = OSubTile::NAXFrag_t::get_coord(); + const short tm = kU * TQ * simd_group_id; + Q += tm * int(params->Q_strides[2]); + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = UQ * TQ * simd_group_id; - - Q += (tm + sm) * int(params->Q_strides[2]) + sn; - K += sm * int(params->K_strides[2]) + sn; - V += sm * int(params->V_strides[2]) + sn; // Init row reduction variables - constexpr short kRowsPT = decltype(Otile)::kRowsPerThread; + constexpr short kRowsPT = otile_t::kRowsPerThread; metal::vec max_score; metal::vec sum_score{0}; @@ -176,83 +174,76 @@ template < } int kb_lim = params->NK; + int kb_min_causal = params->NK; if (do_causal) { int q_max = (tid.x + 1) * BQ + params->qL_off; kb_lim = (q_max + BK - 1) / BK; kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); } const bool is_last_bq = int(tid.x) == (params->NQ_aligned); // const bool is_last_tq = int(simd_group_id) >= (params->qL_rem / UQ); const bool is_last_q = is_last_bq; - const short lim_rows_q = params->qL_rem - (tm + sm); - const short lim_rows_k = params->kL_rem - sm; + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; // Loop over KV seq length for (int kb = 0; kb < kb_lim; kb++) { const int is_last_k = (kb == (params->NK_aligned)); // Do S = Q @ K.T - constexpr short UDs = 16; - constexpr short UKs = 32; - - constexpr short TDs = BD / UDs; - constexpr short TKs = BK / UKs; - - using SSubTile = NAXSubTile; - using QSubTile = NAXSubTile; - using KSubTile = NAXSubTile; - - NAXTile Stile; + using stile_t = NAXTile; + stile_t Stile; Stile.clear(); STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short ik = 0; ik < TKs; ik++) { - STEEL_PRAGMA_UNROLL - for (short id = 0; id < TDs; id++) { - NAXTile Qtile; - NAXTile Ktile; - - const int Q_load_off = iq * UQ * int(params->Q_strides[2]) + id * UDs; - const int K_load_off = - ik * UKs * int(params->K_strides[2]) + id * UDs; + for (short ik = 0; ik < TK; ik += 2) { + // Unrolling the head-dim loop by 4, rather than fully, potentially + // lets the compiler interleave the next K-tile loads with the running + // mma chain instead of hoisting all TD loads up front and is faster + // for head dim 128. +#pragma clang loop unroll_count(4) + for (short id = 0; id < TD; id++) { + NAXTile Qtile; + NAXTile Ktile; + + const int Q_load_off = iq * kU * int(params->Q_strides[2]) + id * kU; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; if (!align_Q && is_last_q) { - // Qtile.load_rows( - // Q + Q_load_off, - // int(params->Q_strides[2]), - // lim_rows_q - iq * UQ); - Qtile.load_safe( + Qtile.load_rows( Q + Q_load_off, int(params->Q_strides[2]), - short2(BD, lim_rows_q - iq * UQ)); + lim_rows_q - iq * kU); } else { Qtile.load(Q + Q_load_off, int(params->Q_strides[2])); } if (!align_K && is_last_k) { - // Ktile.load_rows( - // K + K_load_off, - // int(params->K_strides[2]), - // lim_rows_k - ik * UKs); - Ktile.load_safe( + Ktile.load_rows( K + K_load_off, int(params->K_strides[2]), - short2(BD, lim_rows_k - ik * UKs)); + lim_rows_k - ik * kU); } else { Ktile.load(K + K_load_off, int(params->K_strides[2])); } - subtile_matmad_nax( - Stile.subtile_at(iq, ik), - Qtile.subtile_at(0, 0), + stile_t::NAXFrag_t::mma( + Stile.frag_at(iq, ik), + Stile.frag_at(iq, ik + 1), + Qtile.frag_at(0, 0), metal::false_type{}, - Ktile.subtile_at(0, 0), + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), metal::true_type{}); } } @@ -260,22 +251,10 @@ template < // Scale S STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { Stile.elems()[ii] *= float(scale2); } - // Scale and Retile S - constexpr short UK = 16; - constexpr short TK = BK / UK; - using PSubTile = NAXSubTile; - - NAXTile Ptile; - - STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { - Ptile.elems()[ii] = Stile.elems()[ii]; - } - // Mask out length sequence if (!align_K && is_last_k) { constexpr auto neg_inf = Limits::finite_min; @@ -284,16 +263,16 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - const short col_pos = sn + ik * UK; + const short col_pos = ik * kU + sn; - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) { + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) { - const auto loc = ii * PSubTile::kFragThrCols + jj; - fg[loc] = ((col_pos + jj) >= params->kL_rem) ? neg_inf : fg[loc]; + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; } } } @@ -301,7 +280,7 @@ template < } // Mask out if causal - if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) { + if (do_causal && kb >= kb_min_causal) { constexpr auto neg_inf = Limits::finite_min; const int base_row = tid.x * BQ + params->qL_off + tm; @@ -311,18 +290,16 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - const short row_pos = base_row + iq * UQ; - const short col_pos = base_col + ik * UK; - - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) { + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) { - const auto r = row_pos + ii * PSubTile::kFragRowsJump + sm; - const auto c = col_pos + jj + sn; - const auto loc = ii * PSubTile::kFragThrCols + jj; + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = + base_row + iq * kU + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; fg[loc] = (r < c) ? neg_inf : fg[loc]; } } @@ -339,33 +316,65 @@ template < constexpr bool is_bool = is_same_v; using melem_t = typename metal::conditional_t; - using MSubTile = NAXSubTile; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; - STEEL_PRAGMA_UNROLL - for (short iq = 0; iq < TQ; iq++) { - STEEL_PRAGMA_UNROLL - for (short ik = 0; ik < TK; ik++) { - const short row_pos = base_row + iq * UQ + sm; - const short col_pos = base_col + ik * UK + sn; + if (base_row + BQ <= params->qL && base_col + BK <= params->kL) { + for (short iq = 0; iq < TQ; iq++) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row + iq * kU; + const int col_pos = base_col + ik * kU; - MSubTile mfrag; - mfrag.load_safe( - mask, - int64_t(mask_params->M_strides[2]), - Int<1>{}, - params->qL, - params->kL, - row_pos, - col_pos); + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < MSubTile::kElemsPerFrag; jj++) { - if constexpr (is_bool) { - fg[jj] = mfrag.elems()[jj] ? fg[jj] : neg_inf; - } else { - fg[jj] += M_LOG2E_F * AccumType(mfrag.elems()[jj]); + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row + iq * kU; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(iq, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } } } } @@ -383,10 +392,10 @@ template < } // Row max - Ptile.template row_reduce(new_max); + Stile.template row_reduce(new_max); // exp(Si - rowmax(Si)) - Ptile.template row_bin_op(new_max); + Stile.template row_bin_op(new_max); // Factor exp(rowmax(Si) - rowmax(Si-1)) STEEL_PRAGMA_UNROLL @@ -401,7 +410,7 @@ template < sum_score[i] = sum_score[i] * factor[i]; } - Ptile.template row_reduce(sum_score); + Stile.template row_reduce(sum_score); // Update O Otile.template row_bin_op(factor); @@ -412,39 +421,36 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TD; id += 2) { if constexpr (BD == 128) { - if (id == 2) { + if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } } STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - using VSubTile = NAXSubTile; - NAXTile Vtile; + NAXTile Vtile; - const int V_load_off = ik * UK * int(params->V_strides[2]) + id * UD; + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; if (!align_K && is_last_k) { - // Vtile.load_rows( - // V + V_load_off, - // int(params->V_strides[2]), - // lim_rows_k - ik * UK); - Vtile.load_safe( + Vtile.load_rows( V + V_load_off, int(params->V_strides[2]), - short2(BD, lim_rows_k - ik * UK)); + lim_rows_k - ik * kU); } else { Vtile.load(V + V_load_off, int(params->V_strides[2])); } - subtile_matmad_nax( - Otile.subtile_at(iq, id), - Ptile.subtile_at(iq, ik), - metal::bool_constant{}, - Vtile.subtile_at(0, 0), - metal::bool_constant{}); + otile_t::NAXFrag_t::mma( + Otile.frag_at(iq, id), + Otile.frag_at(iq, id + 1), + Stile.frag_at(iq, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); } } } @@ -467,14 +473,434 @@ template < Otile.template row_bin_op(rcp); // Store results - O += (tm + sm) * int(params->O_strides[2]) + sn; + O += tm * int(params->O_strides[2]); if (!align_Q && is_last_q) { if (lim_rows_q <= 0) return; - // Otile.store_rows(O, params->O_strides[2], lim_rows_q); - Otile.store_safe(O, params->O_strides[2], short2(BD, lim_rows_q)); + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); + } else { + Otile.store(O, int(params->O_strides[2])); + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Head-dim split attention kernel +/////////////////////////////////////////////////////////////////////////////// + +// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup +// accumulator working set of attention_nax (TD output fragments plus the S +// fragments) is what gates tensor-unit throughput, so this kernel splits the +// head dim across the WN = 2 simdgroups of the second warp dimension: each +// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, +// halving its accumulator set. The pair exchanges its partial Q@K.T sums +// through threadgroup memory, then both simdgroups run softmax redundantly on +// the full S tile (the row statistics are cheap) and each accumulates P@V for +// its own half of Dv. + +// clang-format off +template < + typename T, + int BQ, + int BK, + int BD, + int WM, + int WN, + typename MaskType = float, + typename AccumType = float> +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax_dsplit( + const device T* Q [[buffer(0)]], + const device T* K [[buffer(1)]], + const device T* V [[buffer(2)]], + device T* O [[buffer(3)]], + const constant AttnParams* params [[buffer(4)]], + const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], + const device MaskType* mask [[buffer(6), function_constant(has_mask)]], + const device T* sinks [[buffer(7), function_constant(has_sinks)]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on + + // Pacifying compiler + (void)lid; + + // Move to correct block + ulong3 tidl{tid.x, tid.y, tid.z}; + + Q += tidl.z * params->Q_strides[0] + // Batch + tidl.y * params->Q_strides[1] + // Head + tidl.x * BQ * params->Q_strides[2]; // Sequence + + ulong kv_head_idx = int(tid.y) / params->gqa_factor; + K += tidl.z * params->K_strides[0] + // Batch + kv_head_idx * params->K_strides[1]; // Head + + V += tidl.z * params->V_strides[0] + // Batch + kv_head_idx * params->V_strides[1]; // Head + + O += tidl.z * params->O_strides[0] + // Batch + tidl.y * params->O_strides[1] + // Head + tidl.x * BQ * params->O_strides[2]; // Sequence + + if (has_mask) { + mask += tidl.z * mask_params->M_strides[0] + // Batch + tidl.y * mask_params->M_strides[1]; // Head + } + + const metal::uniform scale2 = + make_uniform(params->scale) * make_uniform(1.44269504089f); + + // Prepare MMA tiles + constexpr short kU = 16; + + // The WM simdgroups along the first warp dimension split the Q sequence; + // the WN simdgroups along the second split the head dim. The exchange + // below reduces exactly one peer, so WN is fixed at 2. + static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + constexpr int kNWarps = WM; + static_assert( + BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, + "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); + + // Q seq frags per warp + constexpr int TQ = BQ / (kNWarps * kU); + // HeadDim frags over the full head dim + constexpr int TD = BD / kU; + // KV seq frags per warp + constexpr short TK = BK / kU; + + static_assert(TQ == 1, "Check TQ"); + static_assert(TD % WN == 0, "The head dim must split evenly across WN"); + + // HeadDim frags / columns owned by each of the WN simdgroups of a row group + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + + static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); + static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + + using otile_t = NAXTile; + otile_t Otile; + Otile.clear(); + + const short tm = kU * TQ * row_group; + Q += tm * int(params->Q_strides[2]) + d_half * BDh; + K += d_half * BDh; + V += d_half * BDh; + O += tm * int(params->O_strides[2]) + d_half * BDh; + + constexpr short kRowsPT = otile_t::kRowsPerThread; + + metal::vec max_score; + metal::vec sum_score{0}; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + } + + if (has_sinks) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = M_LOG2E_F * static_cast(sinks[tidl.y]); + sum_score[i] = 1; + } + } + + int kb_lim = params->NK; + int kb_min_causal = params->NK; + + if (do_causal) { + int q_max = (tid.x + 1) * BQ + params->qL_off; + kb_lim = (q_max + BK - 1) / BK; + kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); + } + + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; + + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + + // One slot per (row group, half): a fragment pair in per-lane-linear + // layout. Both halves share the fragment-to-lane mapping, so the + // exchange needs no coordinate math. + threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + + // Keep the simdgroup's Q half resident in registers for the whole KV + // loop: TDh fragments of T are cheap next to the accumulators. + NAXTile Qtiles[TDh]; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + const int Q_load_off = id * kU; + if (!align_Q && is_last_q) { + Qtiles[id].load_rows( + Q + Q_load_off, int(params->Q_strides[2]), lim_rows_q); + } else { + Qtiles[id].load(Q + Q_load_off, int(params->Q_strides[2])); + } + } + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); + const short sm = simd_coord.y; + const short sn = simd_coord.x; + + // Loop over KV seq length + for (int kb = 0; kb < kb_lim; kb++) { + const int is_last_k = (kb == (params->NK_aligned)); + + stile_t Stile; + Stile.clear(); + + // S = Q @ K.T, this half of D only, exchanged pair by pair. + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Ktile.load_rows( + K + K_load_off, int(params->K_strides[2]), lim_rows_k - ik * kU); + } else { + Ktile.load(K + K_load_off, int(params->K_strides[2])); + } + + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + + // Exchange the partial pair and reduce. + threadgroup AccumType* slot = s_xchg[row_group][d_half]; + thread auto& s0 = Stile.frag_at(0, ik); + thread auto& s1 = Stile.frag_at(0, ik + 1); + const short base = short(simd_lane_id) * (2 * kEPF); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + i] = s0[i]; + slot[base + kEPF + i] = s1[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s0[i] += peer[base + i]; + s1[i] += peer[base + kEPF + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // Scale S + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= float(scale2); + } + + // Mask out length sequence + if (!align_K && is_last_k) { + constexpr auto neg_inf = Limits::finite_min; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const short col_pos = ik * kU + sn; + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; + } + } + } + } + + // Mask out if causal + if (do_causal && kb >= kb_min_causal) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + params->qL_off + tm; + const int base_col = kb * BK; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = base_row + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = (r < c) ? neg_inf : fg[loc]; + } + } + } + } + + // Other masking as needed + if (has_mask) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + tm; + const int base_col = kb * BK; + + constexpr bool is_bool = is_same_v; + using melem_t = typename metal::conditional_t; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; + + if (base_row + kU <= params->qL && base_col + BK <= params->kL) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } + + // Do softmax (redundantly per half; the row statistics are cheap) + metal::vec new_max; + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + new_max[i] = max_score[i]; + } + + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + sum_score[i] = sum_score[i] * factor[i]; + } + + Stile.template row_reduce(sum_score); + + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O = P @ V, this half of Dv only. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Vtile.load_rows( + V + V_load_off, int(params->V_strides[2]), lim_rows_k - ik * kU); + } else { + Vtile.load(V + V_load_off, int(params->V_strides[2])); + } + + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + // Next block + K += BK * int(params->K_strides[2]); + V += BK * int(params->V_strides[2]); + } + + // Normalize output + threadgroup_barrier(mem_flags::mem_none); + + metal::vec rcp; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + rcp[i] = 1.f / sum_score[i]; + } + + Otile.template row_bin_op(rcp); + + if (!align_Q && is_last_q) { + if (lim_rows_q <= 0) + return; + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); } else { Otile.store(O, int(params->O_strides[2])); } diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/loader.h b/Source/Cmlx/mlx-generated/metal/steel/attn/loader.h index 3b7c51666..abc35f153 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/loader.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/loader.h @@ -49,18 +49,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -71,7 +71,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -80,7 +80,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -128,7 +128,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; @@ -173,18 +173,18 @@ struct BlockLoaderT { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * kDstStrRow + bj * kDstStrCol), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -196,7 +196,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -207,7 +207,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -255,7 +255,7 @@ struct BlockLoaderT { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/mma.h b/Source/Cmlx/mlx-generated/metal/steel/attn/mma.h index a735848d6..330f7db4b 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/mma.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/mma.h @@ -24,7 +24,7 @@ struct Shape2D { RInt r; CInt c; - Shape2D(RInt r_, CInt c_) : r(r_), c(c_) {} + Shape2D(RInt r_, CInt c_) thread : r(r_), c(c_) {} }; template @@ -257,24 +257,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -283,16 +284,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const { + METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -304,7 +305,7 @@ struct MMATile { } template - METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) { + METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -316,7 +317,7 @@ struct MMATile { } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -333,7 +334,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -350,7 +351,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -365,7 +366,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -380,8 +381,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -400,8 +403,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -512,7 +517,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -530,7 +535,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -557,7 +562,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -571,7 +576,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -590,7 +595,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -604,7 +610,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -633,7 +639,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -678,7 +684,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -711,7 +717,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/nax.h b/Source/Cmlx/mlx-generated/metal/steel/attn/nax.h index 77f3ee411..a54643175 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/nax.h @@ -72,11 +72,13 @@ struct BaseNAXFrag { StrY str_y, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -109,13 +111,16 @@ struct BaseNAXFrag { LimX lim_x, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -130,7 +135,10 @@ struct BaseNAXFrag { } } else { - dst = dtype_frag_t(0); + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } } } } @@ -153,14 +161,18 @@ struct BaseNAXFrag { LimY lim_y, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if ((r < lx) && ((c + j) < ly)) { dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j) * str_y]); } else { @@ -186,11 +198,13 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -225,13 +239,16 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -268,15 +285,19 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if (r < lx && (c + j) < ly) { dst[r * str_x + (c + j) * str_y] = static_cast(src[i * kElemCols + j]); } @@ -308,7 +329,7 @@ struct BaseNAXFrag { OffY off_y = Int<0>{}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); const_for_loop<0, kElemRows, 1>([&](auto idx_row) { const auto r = off_x + idx_row * Int{}; @@ -362,627 +383,424 @@ struct BaseNAXFrag { } } } -}; - -template < - typename T, - short kRows_, - short kCols_, - typename NAXFrag_ = BaseNAXFrag> -struct NAXSubTile { - using NAXFrag_t = NAXFrag_; - STEEL_CONST short kRows = kRows_; - STEEL_CONST short kCols = kCols_; - - STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; - STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; - STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; - - STEEL_CONST short kSubTileRows = kRows / kFragRows; - STEEL_CONST short kSubTileCols = kCols / kFragCols; - - STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols; - STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag; - STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows; - STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols; - - STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; - STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; - STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - - using frag_type = typename NAXFrag_t::template dtype_frag_t; - - frag_type val_frags[kNumFrags]; - - METAL_FUNC constexpr void clear() { + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers STEEL_PRAGMA_UNROLL - for (short i = 0; i < kNumFrags; ++i) { - val_frags[i] = frag_type(0); + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; } - } - - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC thread T* elems() { - return reinterpret_cast(val_frags); - } - - METAL_FUNC const thread T* elems() const { - return reinterpret_cast(val_frags); - } - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - thread T* vptr = (thread T*)(&vals); + // Load B into right operand registers STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::template row_reduce( - frag_at(i, j), &vptr[i * kFragThrRows]); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; } - } - template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - thread T* vptr = (thread T*)(&vals); + // Load C into output registers (op handles accumulation) STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::template row_bin_op( - frag_at(i, j), &vptr[i * kFragThrRows]); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; } - } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load( - SrcPtrType src, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load( - frag_at(i, j), - src, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } - } + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store( - DstPtrType dst, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) const { + // Copy out results STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store( - frag_at(i, j), - dst, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; } } template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_rows( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) { + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_rows( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; } - } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_safe( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) { + // Load B into right operand registers STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_safe( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; } - } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_rows( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) const { + // Load C into output registers (op handles accumulation) STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; } - } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_safe( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) const { + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; } } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename StartX, - typename StopX, - typename StartY, - typename StopY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_slice( - DstPtrType dst, - StrX str_x, - StrY str_y, - StartX start_x, - StopX stop_x, - StartY start_y, - StopY stop_y, - OffX off_x = Int<0>{}, - OffY off_y = Int<0>{}) const { - const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) { - const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) { - NAXFrag_t::store_slice( - frag_at(), - dst, - str_x, - str_y, - start_x, - stop_x, - start_y, - stop_y, - off_x + idx_row * Int{}, - off_y + idx_col * Int{}); - }); - }); - } }; template < - short RC, - short CC, - short RA, - short CA, - short RB, - short CB, - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - typename NAXFrag_t = BaseNAXFrag> -METAL_FUNC void subtile_matmad_nax( - thread NAXSubTile& C, - thread NAXSubTile& A, - metal::bool_constant, - thread NAXSubTile& B, - metal::bool_constant) { - // Static checks - constexpr short FMa = transpose_a ? CA : RA; - constexpr short FMc = RC; - static_assert(FMa == FMc, "NAX matmul: M dimensions do not match"); - - constexpr short FNb = transpose_b ? RB : CB; - constexpr short FNc = CC; - static_assert(FNb == FNc, "NAX matmul: N dimensions do not match"); - - constexpr short FKa = transpose_a ? RA : CA; - constexpr short FKb = transpose_b ? CB : RB; - static_assert(FKa == FKb, "NAX matmul: N dimensions do not match"); - - constexpr short FM = FMc; - constexpr short FN = FNc; - constexpr short FK = FKa; - - constexpr int TM = FM / 16; - constexpr int TN = FN / 16; - constexpr int TK = FK / 16; - - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - FM, - FN, - FK, - transpose_a, - transpose_b, - true, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); - - mpp::tensor_ops::matmul2d gemm_op; - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_a ? kk : mm; - const short fj = transpose_a ? mm : kk; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i]; - } - } - } - - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_b ? nn : kk; - const short fj = transpose_b ? kk : nn; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i]; - } - } - } - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - ct_c[i] = C.elems()[i]; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - C.elems()[i] = ct_c[i]; - } -} - -template + typename T, + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> struct NAXTile { - using NAXSubTile_t = NAXSubTile_; + using NAXFrag_t = NAXFrag_; using elem_type = T; - STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows; - STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols; - STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile; + + STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; + STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; + STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; STEEL_CONST short kTileRows = kTileRows_; STEEL_CONST short kTileCols = kTileCols_; - STEEL_CONST short kRows = kTileRows * kSubTileRows; - STEEL_CONST short kCols = kTileCols * kSubTileCols; + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - NAXSubTile_t val_subtiles[kSubTiles]; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; METAL_FUNC NAXTile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTiles; ++i) { - val_subtiles[i].clear(); + for (short i = 0; i < kNumFrags; ++i) { + val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread NAXSubTile_t& subtile_at( - const short i, - const short j) { - return val_subtiles[i * kTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( + METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { - return val_subtiles[i * kTileCols + j]; + const short j) const thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const { - return val_subtiles[i * kTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC thread elem_type* elems() { - return reinterpret_cast(val_subtiles[0].elems()); + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC const thread elem_type* elems() const { - return reinterpret_cast(val_subtiles[0].elems()); + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - auto sub_rows = (thread metal::vec*)(&vals); + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_reduce(sub_rows[i]); + NAXFrag_t::template row_reduce( + frag_at(i, j), &vptr[i * kFragThrRows]); } } } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - auto sub_rows = (thread metal::vec*)(&vals); + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_bin_op(sub_rows[i]); + NAXFrag_t::template row_bin_op( + frag_at(i, j), &vptr[i * kFragThrRows]); } } } template - METAL_FUNC void load(const threadgroup U* src) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), src, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(threadgroup U* dst) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), dst, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void load(const device U* src, const int ld) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(device U* dst, const int ld) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_safe( + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_rows( + frag_at(), src, ld, Int<1>{}, - src_tile_dims.y, - src_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_rows( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_safe( + frag_at(), + src, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + src_tile_dims.y, + src_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_safe( + METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), dst, ld, Int<1>{}, - dst_tile_dims.y, - dst_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_rows( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), + dst, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + dst_tile_dims.y, + dst_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template @@ -990,10 +808,11 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { - subtile_at().store_slice( + NAXFrag_t::store_slice( + frag_at(), dst, ld, Int<1>{}, @@ -1001,8 +820,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1022,51 +841,54 @@ METAL_FUNC void tile_matmad_nax( metal::bool_constant) { // Static checks constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; - constexpr short TMc = CTile::kTileRows; - static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match"); - - constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows; - constexpr short FMc = CTile::kSubTileRows; - static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match"); + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; - constexpr short TNc = CTile::kTileCols; - static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match"); - - constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols; - constexpr short FNc = CTile::kSubTileCols; - static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match"); + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; - constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows; - static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match"); - - constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols; - constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows; - static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match"); + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); - constexpr short TM = TMc; - constexpr short TN = TNc; - constexpr short TK = TKa; + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; - // Do matmul here - STEEL_PRAGMA_UNROLL - for (short i = 0; i < TM; ++i) { + if constexpr (TN == 1 && TM % 2 == 0) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < TN; ++j) { + for (short mm = 0; mm < TM; mm += 2) { STEEL_PRAGMA_UNROLL - for (short k = 0; k < TK; ++k) { - const short ra = transpose_a ? k : i; - const short ca = transpose_a ? i : k; - const short rb = transpose_b ? j : k; - const short cb = transpose_b ? k : j; - - subtile_matmad_nax( - C.subtile_at(i, j), - A.subtile_at(ra, ca), - metal::bool_constant{}, - B.subtile_at(rb, cb), - metal::bool_constant{}); + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } } } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/attn/transforms.h b/Source/Cmlx/mlx-generated/metal/steel/attn/transforms.h index 3d8ca0541..d14e6dc9e 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/attn/transforms.h +++ b/Source/Cmlx/mlx-generated/metal/steel/attn/transforms.h @@ -24,7 +24,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -40,14 +40,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast(x * alpha + (beta * c)); } }; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h index 850ec15be..f559596b7 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv.h @@ -135,7 +135,7 @@ implicit_gemm_conv_2d( C += tid.z * N; B += c_col * K; - C += c_row * (N * params->groups) + c_col; + C += static_cast(c_row) * N * params->groups + c_col; const int2 offsets_a(0, c_row); const int2 offsets_b(0, c_col); diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_3d.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_3d.h index d2fbac0fc..f2ccc1c03 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_3d.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_3d.h @@ -94,7 +94,7 @@ implicit_gemm_conv_3d( C += tid.z * N; B += c_col * K; - C += c_row * (N * params->groups) + c_col; + C += static_cast(c_row) * N * params->groups + c_col; const int2 offsets_a(0, c_row); const int2 offsets_b(0, c_col); diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h index b775dd55c..d0797ebea 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/kernels/steel_conv_general.h @@ -200,14 +200,14 @@ implicit_gemm_conv_2d_general( (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + base_ow; if (n < params->N && oh < params->oS[0] && ow < params->oS[1]) { - int offset_cm = n * params->out_strides[0] + + size_t offset_cm = static_cast(n) * params->out_strides[0] + oh * params->out_strides[1] + ow * params->out_strides[2]; STEEL_PRAGMA_UNROLL for (int j = 0; j < mma_t::TN; j++) { // Get accumulated result and associated offset in C thread const auto& accum = mma_op.Ctile.frag_at(i, j); - int offset = offset_cm + (j * mma_t::TN_stride); + size_t offset = offset_cm + (j * mma_t::TN_stride); constexpr short kelems = decltype(mma_op.Ctile)::kElemsPerFrag; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h index a516c1ad8..1284efe4a 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_l.h @@ -64,10 +64,10 @@ struct Conv2DInputBlockLoaderLargeFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -103,7 +103,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -131,7 +131,7 @@ struct Conv2DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -213,10 +213,10 @@ struct Conv2DInputBlockLoaderSmallFilter { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -287,7 +287,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -312,7 +312,7 @@ struct Conv2DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[1]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -394,11 +394,11 @@ struct Conv2DWeightBlockLoader { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -408,7 +408,7 @@ struct Conv2DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -435,7 +435,7 @@ struct Conv2DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_hw < (params->wS[1] * params->wS[0])) { src += weight_step; return; @@ -504,10 +504,10 @@ struct Conv3DInputBlockLoaderLargeFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -559,7 +559,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -588,7 +588,7 @@ struct Conv3DInputBlockLoaderLargeFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -683,10 +683,10 @@ struct Conv3DInputBlockLoaderSmallFilter { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -777,7 +777,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { mask_t d_mask = mask_t(1) << weight_d; mask_t h_mask = mask_t(1) << weight_h; mask_t w_mask = mask_t(1) << weight_w; @@ -804,7 +804,7 @@ struct Conv3DInputBlockLoaderSmallFilter { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_w < params->wS[2]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < n_rows; i++) { @@ -897,11 +897,11 @@ struct Conv3DWeightBlockLoader { const constant MLXConvParams<3>* params_, const constant ImplicitGemmConv3DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -911,7 +911,7 @@ struct Conv3DWeightBlockLoader { do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (BN != 8 || do_read) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BN; i += TROWS) { @@ -938,7 +938,7 @@ struct Conv3DWeightBlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { if (++weight_dhw < (params->wS[0] * params->wS[1] * params->wS[2])) { src += weight_step; return; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h index 1f37fb215..a90e9c886 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_channel_n.h @@ -99,10 +99,10 @@ struct Conv2DInputBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), gemm_params(gemm_params_), @@ -131,7 +131,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (weight_hw >= params->wS[1] * params->wS[0]) { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { @@ -187,7 +187,7 @@ struct Conv2DInputBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; @@ -243,11 +243,11 @@ struct Conv2DWeightBlockLoaderSmallChannels { const constant MLXConvParams<2>* params_, const constant ImplicitGemmConv2DParams* gemm_params_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld), params(params_), @@ -256,7 +256,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { do_read(read_n + BN <= gemm_params_->N) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { if (bi >= BROWS || bj >= BCOLS) return; @@ -310,7 +310,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_hw += TCOLS; } }; diff --git a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h index 9043a3c4d..b44ff341e 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h +++ b/Source/Cmlx/mlx-generated/metal/steel/conv/loaders/loader_general.h @@ -67,10 +67,10 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), jump_params(jump_params_), @@ -101,7 +101,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -137,7 +137,7 @@ struct Conv2DInputBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -184,7 +184,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; @@ -263,11 +263,11 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -279,7 +279,7 @@ struct Conv2DWeightBlockLoaderGeneral { start_row(offsets.y + bi) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -308,7 +308,7 @@ struct Conv2DWeightBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -358,7 +358,7 @@ struct Conv2DWeightBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/gemm_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/gemm_nax.h index 9ccd2a96f..7aafc2e37 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/gemm_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/gemm_nax.h @@ -22,9 +22,6 @@ template < bool kAlignedM, bool kAlignedN, bool kAlignedK, - short UM, - short UN, - short UK, typename AccumType = float> auto gemm_loop( const device T* A, @@ -35,9 +32,9 @@ auto gemm_loop( int gemm_k_iterations_aligned, const short sgp_sm, const short sgp_sn) { - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; constexpr int RA = transpose_a ? TK : TM; constexpr int CA = transpose_a ? TM : TK; @@ -45,25 +42,25 @@ auto gemm_loop( constexpr int RB = transpose_b ? TN : TK; constexpr int CB = transpose_b ? TK : TN; - using DSubTile = NAXSubTile; - using ASubTile = - NAXSubTile; - using BSubTile = - NAXSubTile; - - NAXTile Dtile; + NAXTile Dtile; Dtile.clear(); + const bool has_output = sgp_sm > 0 && sgp_sn > 0; + int gemm_k_iterations_ = gemm_k_iterations_aligned; STEEL_PRAGMA_NO_UNROLL for (int kk0 = 0; kk0 < gemm_k_iterations_; kk0++) { threadgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + continue; + } STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; const int k = kk1; volatile int compiler_barrier; @@ -103,51 +100,38 @@ auto gemm_loop( if constexpr (!kAlignedK) { simdgroup_barrier(mem_flags::mem_none); + if constexpr (!kAlignedM || !kAlignedN) { + if (!has_output) + return Dtile; + } const short rem_bk = K - gemm_k_iterations_ * BK; STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < rem_bk; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - STEEL_PRAGMA_UNROLL - for (int mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (int nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (int kk = 0; kk < TK; kk++) { - const int m = mm * UM; - const int n = nn * UN; - const int k = kk1 + kk * UK; - const short psk = max(0, rem_bk - k); - - const int A_offset = transpose_a ? (m + k * lda) : (m * lda + k); - const int B_offset = transpose_b ? (k + n * ldb) : (k * ldb + n); - - { - const short psm = kAlignedM ? SM : max(0, sgp_sm - m); - const short rmax = transpose_a ? psk : psm; - const short cmax = transpose_a ? psm : psk; - Atile.load_safe(A + A_offset, lda, short2(cmax, rmax)); - } - - { - const short psn = kAlignedN ? SN : max(0, sgp_sn - n); - const short rmax = transpose_b ? psn : psk; - const short cmax = transpose_b ? psk : psn; - Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax)); - } - - subtile_matmad_nax( - Dtile.subtile_at(mm, nn), - Atile.subtile_at(0, 0), - metal::bool_constant{}, - Btile.subtile_at(0, 0), - metal::bool_constant{}); - } - } - } + NAXTile Atile; + NAXTile Btile; + + const int k = kk1; + const short psk = max(0, rem_bk - k); + + const short2 Aklims = + transpose_a ? short2(sgp_sm, psk) : short2(psk, sgp_sm); + const short2 Bklims = + transpose_b ? short2(psk, sgp_sn) : short2(sgp_sn, psk); + + const int A_offset = transpose_a ? k * lda : k; + const int B_offset = transpose_b ? k : k * ldb; + + Atile.load_safe(A + A_offset, lda, Aklims); + Btile.load_safe(B + B_offset, ldb, Bklims); + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_fused_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_fused_nax.h index 4ff926065..8267707ad 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_fused_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_fused_nax.h @@ -15,7 +15,7 @@ constant bool align_K [[function_constant(202)]]; template < bool kAlignedM, bool kAlignedN, - typename NAXTile_t, + class NAXTile_t, typename T> void gemm_epilogue( thread NAXTile_t& Dtile, @@ -27,37 +27,40 @@ void gemm_epilogue( (void)params; - constexpr short UM = NAXTile_t::kSubTileRows; - constexpr short UN = NAXTile_t::kSubTileCols; - using CSubTile = NAXSubTile; - using V = typename NAXTile_t::elem_type; constexpr short TM = NAXTile_t::kTileRows; constexpr short TN = NAXTile_t::kTileCols; - constexpr short kElemsPerSubTile = NAXTile_t::kElemsPerSubTile; + constexpr short kElemsPerFrag = NAXTile_t::kElemsPerFrag; + + using CFrag = typename NAXTile_t::NAXFrag_t; + using cfrag_t = typename CFrag::template dtype_frag_t; - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - const short m = mm * UM; - const short n = nn * UN; + const_for_loop<0, TM, 1>([&](auto mm) { + const_for_loop<0, TN, 1>([&](auto nn) { + auto m = mm * Int{}; + auto n = nn * Int{}; - CSubTile CTile; + cfrag_t celems; if constexpr (kAlignedM && kAlignedN) { - CTile.load(C, addmm_params->ldc, addmm_params->fdc, m, n); + CFrag::load(celems, C, addmm_params->ldc, addmm_params->fdc, m, n); } else { - CTile.load_safe( - C, addmm_params->ldc, addmm_params->fdc, sgp_sm, sgp_sn, m, n); + CFrag::load_safe( + celems, + C, + addmm_params->ldc, + addmm_params->fdc, + sgp_sm, + sgp_sn, + m, + n); } - auto delems = Dtile.subtile_at(mm, nn).elems(); - auto celems = CTile.elems(); + thread auto& delems = Dtile.template frag_at(); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kElemsPerSubTile; i++) { + for (short i = 0; i < kElemsPerFrag; i++) { if (do_axpby) { delems[i] = addmm_params->alpha * delems[i] + addmm_params->beta * static_cast(celems[i]); @@ -65,8 +68,8 @@ void gemm_epilogue( delems[i] += static_cast(celems[i]); } } - } - } + }); + }); } // clang-format off @@ -144,15 +147,12 @@ template < C += c_row_long * addmm_params->ldc + c_col_long * addmm_params->fdc; } - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); @@ -175,8 +175,7 @@ template < C += tm * addmm_params->ldc + tn * addmm_params->fdc; } - using DSubTile = NAXSubTile; - NAXTile Dtile; + NAXTile Dtile; dispatch_bool(align_K, [&](auto kAlignedK) { dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { @@ -192,9 +191,6 @@ template < kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>( A, B, @@ -204,14 +200,17 @@ template < params->gemm_k_iterations_aligned, sgp_sm, sgp_sn); - if (use_out_source) { - gemm_epilogue( - Dtile, C, params, addmm_params, sgp_sm, sgp_sn); - } - if constexpr (kAlignedM && kAlignedN) { - Dtile.store(D, int(params->ldd)); - } else { - Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + if ((kAlignedM.value || sgp_sm > 0) && + (kAlignedN.value || sgp_sn > 0)) { + if (use_out_source) { + gemm_epilogue( + Dtile, C, params, addmm_params, sgp_sm, sgp_sn); + } + if constexpr (kAlignedM && kAlignedN) { + Dtile.store(D, int(params->ldd)); + } else { + Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + } } }); }); diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h index 67cd73784..7fec27060 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_gather_nax.h @@ -25,14 +25,11 @@ gather_mm_rhs_nax( const constant GEMMParams* params [[buffer(4)]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; if (params->tiles_n <= static_cast(tid.x) || params->tiles_m <= static_cast(tid.y)) { @@ -88,13 +85,12 @@ gather_mm_rhs_nax( } threadgroup_barrier(mem_flags::mem_none); - using DSubTile = NAXSubTile; - NAXTile Ctile; + NAXTile Ctile; dispatch_bool(align_K, [&](auto kAlignedK) { dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< + auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K T, SM, SN, @@ -105,9 +101,6 @@ gather_mm_rhs_nax( kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>; Ctile = do_gemm( A, diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_masked.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_masked.h index 6546215e6..f15d702a2 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_masked.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_masked.h @@ -11,7 +11,7 @@ using namespace mlx::steel; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -29,7 +29,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h new file mode 100644 index 000000000..9cfece5c5 --- /dev/null +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_segmented_nax.h @@ -0,0 +1,125 @@ +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool segments_contiguous [[function_constant(199)]]; +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; + +template < + typename T, + int BM, + int BN, + int BK, + int WM, + int WN, + bool transpose_a, + bool transpose_b, + typename AccumType = float> +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] +void segmented_mm_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device uint32_t* segments [[buffer(2)]], + device T* C [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { + const int tid_m = (BK > 64) ? tid.y : tid.z; + const int tid_n = (BK > 64) ? tid.x : tid.y; + const int tid_s = (BK > 64) ? tid.z : tid.x; + + const int c_row = tid_m * BM; + const int c_col = tid_n * BN; + const size_t c_row_long = size_t(c_row); + const size_t c_col_long = size_t(c_col); + + if (params->tiles_n <= static_cast(tid_n) || + params->tiles_m <= static_cast(tid_m)) { + return; + } + + A += transpose_a ? c_row_long : c_row_long * params->lda; + B += transpose_b ? c_col_long * params->ldb : c_col_long; + C += c_row_long * params->ldd + c_col_long; + + uint32_t k_start, k_end; + if (segments_contiguous) { + k_start = segments[2 * tid_s]; + k_end = segments[2 * tid_s + 1]; + } else { + k_start = segments[tid_s]; + k_end = segments[tid_s + 1]; + } + A += transpose_a ? k_start * params->lda : k_start; + B += transpose_b ? k_start : k_start * params->ldb; + C += tid_s * params->batch_stride_d; + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + + const int sgp_sm_int = + align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); + const short sgp_sm = short(sgp_sm_int); + const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + + const int sgp_sn_int = + align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); + const short sgp_sn = short(sgp_sn_int); + const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); + + A += transpose_a ? tm : (tm * params->lda); + B += transpose_b ? (tn * params->ldb) : tn; + C += tm * params->ldd + tn; + + NAXTile Dtile; + Dtile.clear(); + + const int segment_k_size = k_end - k_start; + const int segment_k_iters = segment_k_size / BK; + const bool segment_k_aligned = (segment_k_size % BK) == 0; + + dispatch_bool(segment_k_aligned, [&](auto kAlignedK) { + dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { + Dtile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + segment_k_size, + segment_k_iters, + sgp_sm, + sgp_sn); + }); + }); + }); + + dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { + if constexpr (kAlignedM && kAlignedN) { + Dtile.store(C, int(params->ldd)); + } else { + Dtile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } + }); + }); +} diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_splitk_nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_splitk_nax.h index 1b6b82808..982c770c6 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_splitk_nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/kernels/steel_gemm_splitk_nax.h @@ -72,15 +72,12 @@ template < (c_row_long * params->ldc + c_col_long); // NAX tile configuration - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; // Calculate simdgroup offsets and alignment const short tm = SM * (simd_group_id / WN); @@ -100,8 +97,7 @@ template < B += transpose_b ? (tn * params->ldb) : tn; C += tm * params->ldc + tn; - using DSubTile = NAXSubTile; - NAXTile Dtile; + NAXTile Dtile; // gemm_loop through the partition // Check K-alignment at runtime (partition-specific) @@ -123,9 +119,6 @@ template < kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>( A, B, diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/loader.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/loader.h index cc79de861..7a97c2c36 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/loader.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/loader.h @@ -49,18 +49,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -71,7 +71,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -80,7 +80,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -128,7 +128,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/mma.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/mma.h index 8b9ddb296..e42615a24 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/mma.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/mma.h @@ -89,7 +89,7 @@ struct BaseMMAFrag { for (short j = 0; j < kElemCols; j++) { if ((off_x + i) < lim_x && (off_y + j) < lim_y) { dst[i * kElemCols + j] = - static_cast(src[(off_x + i) * str_x + (off_x + j) * str_y]); + static_cast(src[(off_x + i) * str_x + (off_y + j) * str_y]); } else { dst[i * kElemCols + j] = T(0); } @@ -234,24 +234,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -260,16 +261,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -286,7 +287,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -303,7 +304,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -318,7 +319,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -333,8 +334,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -353,8 +356,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -377,7 +382,7 @@ struct MMATile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -487,7 +492,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -505,7 +510,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -532,7 +537,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -545,8 +550,11 @@ struct BlockMMA { Ctile.template store(D, ldd); } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -566,7 +574,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -585,7 +593,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -599,7 +608,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -628,7 +637,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -673,7 +682,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -706,7 +715,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -819,7 +828,7 @@ struct BlockMMA< /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -839,7 +848,7 @@ struct BlockMMA< /* Karatsuba MMA: 3 real MMAs per K-chunk */ METAL_FUNC void mma( const threadgroup complex64_t* As, - const threadgroup complex64_t* Bs) { + const threadgroup complex64_t* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -897,7 +906,7 @@ struct BlockMMA< } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Adjust for simdgroup and thread location D += sm * ldd + sn; @@ -916,8 +925,11 @@ struct BlockMMA< } } - METAL_FUNC void - store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + METAL_FUNC void store_result_slice( + device U* D, + const int ldd, + short2 start, + short2 stop) thread { D += sm * ldd + sn; start -= short2(sn, sm); stop -= short2(sn, sm); @@ -948,7 +960,7 @@ struct BlockMMA< } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { D += sm * ldd + sn; dst_tile_dims -= short2(sn, sm); if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) @@ -974,7 +986,8 @@ struct BlockMMA< /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) { complex64_t out = epilogue_op.apply( @@ -990,7 +1003,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -1022,7 +1035,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -1071,7 +1084,7 @@ struct BlockMMA< const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1106,7 +1119,7 @@ struct BlockMMA< const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h index 740068be8..a54643175 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h @@ -7,7 +7,6 @@ #include #include "../../steel/defines.h" -#include "../../steel/gemm/transforms.h" #include "../../steel/utils/integral_constant.h" #include @@ -74,10 +73,12 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -111,12 +112,15 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -131,7 +135,10 @@ struct BaseNAXFrag { } } else { - dst = dtype_frag_t(0); + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } } } } @@ -155,13 +162,17 @@ struct BaseNAXFrag { OffX off_x = {}, OffY off_y = {}) { const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if ((r < lx) && ((c + j) < ly)) { dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j) * str_y]); } else { @@ -188,10 +199,12 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -227,12 +240,15 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -270,14 +286,18 @@ struct BaseNAXFrag { using U = pointer_element_t; const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if (r < lx && (c + j) < ly) { dst[r * str_x + (c + j) * str_y] = static_cast(src[i * kElemCols + j]); } @@ -363,634 +383,424 @@ struct BaseNAXFrag { } } } -}; -template < - typename T, - short kRows_, - short kCols_, - typename NAXFrag_t = BaseNAXFrag> -struct NAXSubTile { - STEEL_CONST short kRows = kRows_; - STEEL_CONST short kCols = kCols_; - - STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; - STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; - STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; - - STEEL_CONST short kSubTileRows = kRows / kFragRows; - STEEL_CONST short kSubTileCols = kCols / kFragCols; - - STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols; - STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag; - - STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows; - STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols; - - STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; - STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; - STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - - using frag_type = typename NAXFrag_t::template dtype_frag_t; - - frag_type val_frags[kNumFrags]; - - METAL_FUNC constexpr void clear() { + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers STEEL_PRAGMA_UNROLL - for (short i = 0; i < kNumFrags; ++i) { - val_frags[i] = frag_type(0); + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; } - } - - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC thread T* elems() { - return reinterpret_cast(val_frags); - } - METAL_FUNC const thread T* elems() const { - return reinterpret_cast(val_frags); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { + // Load B into right operand registers STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::template row_reduce( - frag_at(i, j), &vals[i * kFragThrRows]); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; } - } - template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { + // Load C into output registers (op handles accumulation) STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::template row_bin_op( - frag_at(i, j), &vals[i * kFragThrRows]); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; } - } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load( - SrcPtrType src, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load( - frag_at(i, j), - src, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } - } + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store( - DstPtrType dst, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) const { + // Copy out results STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store( - frag_at(i, j), - dst, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; } } template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_rows( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) { + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_rows( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; } - } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_safe( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) { + // Load B into right operand registers STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::load_safe( - frag_at(i, j), - src, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; } - } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_safe( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) const { + // Load C into output registers (op handles accumulation) STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; } - } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_rows( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) const { + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; } } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename StartX, - typename StopX, - typename StartY, - typename StopY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_slice( - DstPtrType dst, - StrX str_x, - StrY str_y, - StartX start_x, - StopX stop_x, - StartY start_y, - StopY stop_y, - OffX off_x = Int<0>{}, - OffY off_y = Int<0>{}) const { - const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) { - const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) { - NAXFrag_t::store_slice( - frag_at(), - dst, - str_x, - str_y, - start_x, - stop_x, - start_y, - stop_y, - off_x + idx_row * Int{}, - off_y + idx_col * Int{}); - }); - }); - } }; template < - short RC, - short CC, - short RA, - short CA, - short RB, - short CB, - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - typename NAXFrag_t = BaseNAXFrag> -METAL_FUNC void subtile_matmad_nax( - thread NAXSubTile& C, - thread NAXSubTile& A, - metal::bool_constant, - thread NAXSubTile& B, - metal::bool_constant) { - // Static checks - constexpr short FMa = transpose_a ? CA : RA; - constexpr short FMc = RC; - static_assert(FMa == FMc, "NAX matmul: M dimensions do not match"); - - constexpr short FNb = transpose_b ? RB : CB; - constexpr short FNc = CC; - static_assert(FNb == FNc, "NAX matmul: N dimensions do not match"); - - constexpr short FKa = transpose_a ? RA : CA; - constexpr short FKb = transpose_b ? CB : RB; - static_assert(FKa == FKb, "NAX matmul: N dimensions do not match"); - - constexpr short FM = FMc; - constexpr short FN = FNc; - constexpr short FK = FKa; - - constexpr int TM = FM / 16; - constexpr int TN = FN / 16; - constexpr int TK = FK / 16; - - // Create Matmul descriptor - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - FM, - FN, - FK, - transpose_a, - transpose_b, - true, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); - - // Create matmul op - mpp::tensor_ops::matmul2d gemm_op; - - // Create matmul operands in registers - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - - // Create matmul output in register - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - // Load A in to left operand registers - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_a ? kk : mm; - const short fj = transpose_a ? mm : kk; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i]; - } - } - } - - // Load B into right operand registers - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_b ? nn : kk; - const short fj = transpose_b ? kk : nn; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i]; - } - } - } - - // Load C into output registers (op handles accumulation) - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - ct_c[i] = C.elems()[i]; - } - - // Do matmul - gemm_op.run(ct_a, ct_b, ct_c); - - // Copy out results - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - C.elems()[i] = ct_c[i]; - } -} - -template + typename T, + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> struct NAXTile { - using NAXSubTile_t = NAXSubTile_; + using NAXFrag_t = NAXFrag_; using elem_type = T; - STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows; - STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols; - STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile; + + STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; + STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; + STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; STEEL_CONST short kTileRows = kTileRows_; STEEL_CONST short kTileCols = kTileCols_; - STEEL_CONST short kRows = kTileRows * kSubTileRows; - STEEL_CONST short kCols = kTileCols * kSubTileCols; + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - NAXSubTile_t val_subtiles[kSubTiles]; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; METAL_FUNC NAXTile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTiles; ++i) { - val_subtiles[i].clear(); + for (short i = 0; i < kNumFrags; ++i) { + val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread NAXSubTile_t& subtile_at( - const short i, - const short j) { - return val_subtiles[i * kTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( + METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { - return val_subtiles[i * kTileCols + j]; + const short j) const thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const { - return val_subtiles[i * kTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC thread elem_type* elems() { - return reinterpret_cast(val_subtiles[0].elems()); + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } } - METAL_FUNC const thread elem_type* elems() const { - return reinterpret_cast(val_subtiles[0].elems()); + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - auto sub_rows = (thread metal::vec*)(&vals); + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_reduce(sub_rows[i]); + NAXFrag_t::template row_reduce( + frag_at(i, j), &vptr[i * kFragThrRows]); } } } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - auto sub_rows = (thread metal::vec*)(&vals); + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_bin_op(sub_rows[i]); + NAXFrag_t::template row_bin_op( + frag_at(i, j), &vptr[i * kFragThrRows]); } } } template - METAL_FUNC void load(const threadgroup U* src) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), src, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(threadgroup U* dst) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), dst, Int{}, Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void load(const device U* src, const int ld) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( - &src[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void store(device U* dst, const int ld) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( - &dst[(i * kSubTileRows * ld + j * kSubTileCols)], ld, Int<1>{}); - } - } + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_rows( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_rows( + frag_at(), + src, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_safe( + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load_safe( + frag_at(), src, ld, Int<1>{}, src_tile_dims.y, src_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_rows( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( + frag_at(), + dst, ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_safe( + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), dst, ld, Int<1>{}, dst_tile_dims.y, dst_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template @@ -998,10 +808,11 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { - subtile_at().store_slice( + NAXFrag_t::store_slice( + frag_at(), dst, ld, Int<1>{}, @@ -1009,8 +820,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1030,51 +841,54 @@ METAL_FUNC void tile_matmad_nax( metal::bool_constant) { // Static checks constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; - constexpr short TMc = CTile::kTileRows; - static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match"); - - constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows; - constexpr short FMc = CTile::kSubTileRows; - static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match"); + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; - constexpr short TNc = CTile::kTileCols; - static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match"); - - constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols; - constexpr short FNc = CTile::kSubTileCols; - static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match"); + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; - constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows; - static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match"); - - constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols; - constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows; - static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match"); + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); - constexpr short TM = TMc; - constexpr short TN = TNc; - constexpr short TK = TKa; + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; - // Do matmul here - STEEL_PRAGMA_UNROLL - for (short i = 0; i < TM; ++i) { + if constexpr (TN == 1 && TM % 2 == 0) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < TN; ++j) { + for (short mm = 0; mm < TM; mm += 2) { STEEL_PRAGMA_UNROLL - for (short k = 0; k < TK; ++k) { - const short ra = transpose_a ? k : i; - const short ca = transpose_a ? i : k; - const short rb = transpose_b ? j : k; - const short cb = transpose_b ? k : j; - - subtile_matmad_nax( - C.subtile_at(i, j), - A.subtile_at(ra, ca), - metal::bool_constant{}, - B.subtile_at(rb, cb), - metal::bool_constant{}); + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } } } } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/transforms.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/transforms.h index 704776baf..980347610 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/gemm/transforms.h +++ b/Source/Cmlx/mlx-generated/metal/steel/gemm/transforms.h @@ -24,7 +24,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -40,14 +40,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast( x * static_cast(alpha) + (static_cast(beta) * c)); } diff --git a/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h b/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h index 40bcff8c0..f85fca32c 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h +++ b/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h @@ -20,13 +20,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -56,7 +52,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -131,4 +128,4 @@ METAL_FUNC constexpr auto sum(T x, Us... us) { } // namespace steel } // namespace mlx -#pragma METAL internals : disable \ No newline at end of file +#pragma METAL internals : disable diff --git a/Source/Cmlx/mlx-generated/metal/ternary_ops.h b/Source/Cmlx/mlx-generated/metal/ternary_ops.h index e0235d9dd..d97e238c6 100644 --- a/Source/Cmlx/mlx-generated/metal/ternary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/ternary_ops.h @@ -4,7 +4,7 @@ struct Select { template - T operator()(bool condition, T x, T y) { + T operator()(bool condition, T x, T y) thread { return condition ? x : y; } }; diff --git a/Source/Cmlx/mlx-generated/metal/unary_ops.h b/Source/Cmlx/mlx-generated/metal/unary_ops.h index 0ec0febc7..a2dc80426 100644 --- a/Source/Cmlx/mlx-generated/metal/unary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/unary_ops.h @@ -16,125 +16,125 @@ constant float inf = metal::numeric_limits::infinity(); struct Abs { template - T operator()(T x) { + T operator()(T x) thread { return metal::abs(x); }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; }; }; struct ArcCos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acos(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcCosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acosh(x); }; }; struct ArcSin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asin(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcSinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asinh(x); }; }; struct ArcTan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atan(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcTanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atanh(x); }; }; struct BitwiseInvert { template - T operator()(T x) { + T operator()(T x) thread { return ~x; }; }; struct Ceil { template - T operator()(T x) { + T operator()(T x) thread { return metal::ceil(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Cos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cos(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cos(x.real) * metal::precise::cosh(x.imag), -metal::precise::sin(x.real) * metal::precise::sinh(x.imag)}; @@ -143,11 +143,11 @@ struct Cos { struct Cosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cosh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cosh(x.real) * metal::precise::cos(x.imag), metal::precise::sinh(x.real) * metal::precise::sin(x.imag)}; @@ -155,89 +155,89 @@ struct Cosh { }; struct Conjugate { - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return complex64_t{x.real, -x.imag}; } }; struct Erf { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erf(static_cast(x))); }; }; struct ErfInv { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erfinv(static_cast(x))); }; }; struct Exp { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::exp(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return cexpf(x); } }; struct Expm1 { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(expm1f(static_cast(x))); }; }; struct Floor { template - T operator()(T x) { + T operator()(T x) thread { return metal::floor(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Imag { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.imag; }; }; struct Log { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto r = metal::precise::log(Abs{}(x).real); auto i = metal::precise::atan2(x.imag, x.real); return {r, i}; @@ -246,11 +246,11 @@ struct Log { struct Log2 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log2(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN2_F, y.imag / M_LN2_F}; }; @@ -258,11 +258,11 @@ struct Log2 { struct Log10 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log10(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN10_F, y.imag / M_LN10_F}; }; @@ -270,44 +270,44 @@ struct Log10 { struct Log1p { template - T operator()(T x) { + T operator()(T x) thread { return log1p(x); }; }; struct LogicalNot { template - T operator()(T x) { + T operator()(T x) thread { return !x; }; }; struct Negative { template - T operator()(T x) { + T operator()(T x) thread { return -x; }; }; struct Real { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.real; }; }; struct Round { template - T operator()(T x) { + T operator()(T x) thread { return metal::rint(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::rint(x.real), metal::rint(x.imag)}; }; }; struct Sigmoid { template - T operator()(T x) { + T operator()(T x) thread { auto y = 1 / (1 + metal::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } @@ -315,13 +315,13 @@ struct Sigmoid { struct Sign { template - T operator()(T x) { + T operator()(T x) thread { return (x > T(0)) - (x < T(0)); }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x != 0; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x == complex64_t(0)) { return x; } @@ -332,11 +332,11 @@ struct Sign { struct Sin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sin(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sin(x.real) * metal::precise::cosh(x.imag), metal::precise::cos(x.real) * metal::precise::sinh(x.imag)}; @@ -345,11 +345,11 @@ struct Sin { struct Sinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sinh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sinh(x.real) * metal::precise::cos(x.imag), metal::precise::cosh(x.real) * metal::precise::sin(x.imag)}; @@ -358,18 +358,18 @@ struct Sinh { struct Square { template - T operator()(T x) { + T operator()(T x) thread { return x * x; }; }; struct Sqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x.real == 0.0 && x.imag == 0.0) { return {0.0, 0.0}; } @@ -383,22 +383,22 @@ struct Sqrt { struct Rsqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::rsqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return 1.0 / Sqrt{}(x); } }; struct Tan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tan(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tan_a = metal::precise::tan(x.real); float tanh_b = metal::precise::tanh(x.imag); float t1 = tan_a * tanh_b; @@ -409,11 +409,11 @@ struct Tan { struct Tanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tanh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tanh_a = metal::precise::tanh(x.real); float tan_b = metal::precise::tan(x.imag); float t1 = tanh_a * tan_b; @@ -422,19 +422,19 @@ struct Tanh { }; }; -complex64_t ArcCos::operator()(complex64_t x) { +complex64_t ArcCos::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(x + i * Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcSin::operator()(complex64_t x) { +complex64_t ArcSin::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(i * x + Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcTan::operator()(complex64_t x) { +complex64_t ArcTan::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto ix = i * x; return (1.0 / complex64_t{0.0, 2.0}) * Log{}((1.0 + ix) / (1.0 - ix)); @@ -442,13 +442,13 @@ complex64_t ArcTan::operator()(complex64_t x) { struct ToFP8 { template - uint8_t operator()(T f) { + uint8_t operator()(T f) thread { return fp8_e4m3(f).bits; } }; struct FromFP8 { - float operator()(uint8_t x) { + float operator()(uint8_t x) thread { return float(*(thread fp8_e4m3*)(&x)); } }; diff --git a/Source/Cmlx/mlx-generated/metal/utils.h b/Source/Cmlx/mlx-generated/metal/utils.h index 9651ef064..cbf5798a1 100644 --- a/Source/Cmlx/mlx-generated/metal/utils.h +++ b/Source/Cmlx/mlx-generated/metal/utils.h @@ -75,14 +75,14 @@ struct Limits { static constexpr constant bool min = false; }; -template <> -struct Limits { - static constexpr constant complex64_t max = complex64_t( - metal::numeric_limits::infinity(), - metal::numeric_limits::infinity()); - static constexpr constant complex64_t min = complex64_t( - -metal::numeric_limits::infinity(), - -metal::numeric_limits::infinity()); +template +struct Limits> { + inline static constexpr constant complex_t max = complex_t( + metal::numeric_limits::infinity(), + metal::numeric_limits::infinity()); + inline static constexpr constant complex_t min = complex_t( + -metal::numeric_limits::infinity(), + -metal::numeric_limits::infinity()); }; /////////////////////////////////////////////////////////////////////////////// @@ -205,9 +205,9 @@ struct LoopedElemToLoc { OffsetT offset{0}; int index{0}; - LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {} + LoopedElemToLoc(int dim) thread : dim(dim), inner_looper(dim - 1) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { if (dim == 0) { return; } @@ -220,7 +220,8 @@ struct LoopedElemToLoc { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { if (dim == 0) { return; } @@ -243,7 +244,7 @@ struct LoopedElemToLoc { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -254,9 +255,9 @@ struct LoopedElemToLoc<1, OffsetT, true> { OffsetT offset{0}; uint index{0}; - LoopedElemToLoc(int dim) : dim(dim) {} + LoopedElemToLoc(int dim) thread : dim(dim) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { index++; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -265,7 +266,8 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { index += n; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -274,7 +276,7 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -283,17 +285,18 @@ template struct LoopedElemToLoc<1, OffsetT, false> { OffsetT offset{0}; - LoopedElemToLoc(int) {} + LoopedElemToLoc(int) thread {} - void next(const constant int*, const constant int64_t* strides) { + void next(const constant int*, const constant int64_t* strides) thread { offset += OffsetT(strides[0]); } - void next(int n, const constant int*, const constant int64_t* strides) { + void next(int n, const constant int*, const constant int64_t* strides) + thread { offset += n * OffsetT(strides[0]); } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -443,3 +446,27 @@ template struct ConditionalType { using type = T; }; + +/////////////////////////////////////////////////////////////////////////////// +// Type casting utils +/////////////////////////////////////////////////////////////////////////////// + +template +inline U cast_to(T val) { + return static_cast(val); +} + +template <> +inline bool cast_to(float val) { + return (as_type(val) & 0x7FFFFFFF) != 0; +} + +template <> +inline bool cast_to(bfloat16_t val) { + return (as_type(val) & 0x7FFF) != 0; +} + +template <> +inline bool cast_to(complex64_t val) { + return cast_to(val.real) || cast_to(val.imag); +} diff --git a/Source/Cmlx/mlx-generated/quantized.cpp b/Source/Cmlx/mlx-generated/quantized.cpp index 8ee194497..583bba8ac 100644 --- a/Source/Cmlx/mlx-generated/quantized.cpp +++ b/Source/Cmlx/mlx-generated/quantized.cpp @@ -493,25 +493,25 @@ qouter(const thread uint8_t* w, U x, U scale, U bias, thread U* result) { } } -template -inline void -dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { +// Decode one quantized block (scale * q + bias) into w_local. W (the output +// pointer type) serves the threadgroup block loader or a thread-local decode. +template +inline void dequantize(const device uint8_t* w, U scale, U bias, W w_local) { static_assert( bits == 2 || bits == 3 || bits == 4 || bits == 5 || bits == 6 || bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -520,22 +520,24 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -544,14 +546,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -559,16 +565,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } @@ -622,15 +630,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -638,7 +646,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -651,7 +659,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -681,7 +689,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -987,6 +995,103 @@ METAL_FUNC void qmv_impl( } } +// Affine analog of fp_qmv_wide. Weights carry a scale and bias per group, so +// each group is decoded in 8-value sub-chunks (scale * q + bias, registers +// bounded for any group_size) and reused across the vecs_per_tg vectors. +template +METAL_FUNC void qmv_wide_impl( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + constexpr int num_simdgroups = 2; + constexpr int results_per_simdgroup = SIMD_SIZE / k_lanes; + constexpr int sub = 8; // values per sub-chunk (== bits bytes, byte-aligned) + + typedef float U; + + const short k_lane = simd_lid % k_lanes; + const short sg_row = simd_lid / k_lanes; + + const int out_row = tid.y * (results_per_simdgroup * num_simdgroups) + + results_per_simdgroup * simd_gid + sg_row; + const int vec0 = tid.x * vecs_per_tg; + + const int row = min(out_row, out_vec_size - 1); + + const int in_vec_size_w = in_vec_size * bits / 8; // bytes per weight row + const int in_vec_size_g = in_vec_size / group_size; + const device uint8_t* wrow = (const device uint8_t*)w + row * in_vec_size_w; + const device T* srow = scales + row * in_vec_size_g; + const device T* brow = biases + row * in_vec_size_g; + + const device T* xv[vecs_per_tg]; + for (int v = 0; v < vecs_per_tg; v++) { + xv[v] = x + min(vec0 + v, M - 1) * in_vec_size; + } + + U result[vecs_per_tg] = {0}; + + // Each lane reduces a strided subset of the row's groups: decode the group in + // 8-value sub-chunks and reuse each chunk across the streamed vectors. + for (int g = k_lane; g < in_vec_size_g; g += k_lanes) { + U scale = srow[g]; + U bias = brow[g]; +#pragma unroll + for (int sc = 0; sc < group_size / sub; sc++) { + const int k0 = g * group_size + sc * sub; + const device uint8_t* wc = wrow + k0 * bits / 8; + U w_dq[sub]; + dequantize(wc, scale, bias, w_dq); +#pragma unroll + for (int v = 0; v < vecs_per_tg; v++) { + const device T* xc = xv[v] + k0; + U acc = 0; +#pragma unroll + for (int i = 0; i < sub; i++) { + acc += static_cast(xc[i]) * w_dq[i]; + } + result[v] += acc; + } + } + } + + // Reduce each vector's partial over its k_lanes with a shuffle ladder: + // simd_sum would mix the results_per_simdgroup rows a simdgroup spans. + for (int v = 0; v < vecs_per_tg; v++) { + if constexpr (k_lanes >= 32) { + result[v] += simd_shuffle_down(result[v], 16); + } + if constexpr (k_lanes >= 16) { + result[v] += simd_shuffle_down(result[v], 8); + } + if constexpr (k_lanes >= 8) { + result[v] += simd_shuffle_down(result[v], 4); + } + if constexpr (k_lanes >= 4) { + result[v] += simd_shuffle_down(result[v], 2); + } + if constexpr (k_lanes >= 2) { + result[v] += simd_shuffle_down(result[v], 1); + } + } + + if (k_lane == 0 && out_row < out_vec_size) { + for (int v = 0; v < vecs_per_tg; v++) { + if (vec0 + v < M) { + y[(vec0 + v) * out_vec_size + out_row] = static_cast(result[v]); + } + } + } +} + template METAL_FUNC void qvm_impl( const device uint32_t* w, @@ -996,6 +1101,7 @@ METAL_FUNC void qvm_impl( device T* y, const int in_vec_size, const int out_vec_size, + const int in_vec_stride, uint3 tid [[threadgroup_position_in_grid]], uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { @@ -1029,7 +1135,7 @@ METAL_FUNC void qvm_impl( ws += out_col * bytes_per_pack / pack_factor + simd_lid * out_vec_size_w; scales += out_col / group_size + simd_lid * out_vec_size_g; biases += out_col / group_size + simd_lid * out_vec_size_g; - x += tid.x * in_vec_size + simd_lid; + x += tid.x * in_vec_stride + simd_lid; y += tid.x * out_vec_size + out_col; if (out_col >= out_vec_size) { @@ -1115,6 +1221,7 @@ METAL_FUNC void qmm_t_impl( const constant int& K, const constant int& N, const constant int& M, + const constant int& K_eff, uint3 tid [[threadgroup_position_in_grid]], uint lid [[thread_index_in_threadgroup]], uint simd_gid [[simdgroup_index_in_threadgroup]], @@ -1169,7 +1276,7 @@ METAL_FUNC void qmm_t_impl( if (num_els < BM) { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_safe(short2(BK, num_outs)); @@ -1179,7 +1286,7 @@ METAL_FUNC void qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_safe(short2(BK, num_els)); loader_w.load_unsafe(); @@ -1191,7 +1298,7 @@ METAL_FUNC void qmm_t_impl( } } else { if (!aligned_N && num_outs < BN) { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_safe(short2(BK, num_outs)); @@ -1201,7 +1308,7 @@ METAL_FUNC void qmm_t_impl( loader_w.next(); } } else { - for (int k = 0; k < K; k += BK) { + for (int k = 0; k < K_eff; k += BK) { threadgroup_barrier(mem_flags::mem_threadgroup); loader_x.load_unsafe(); loader_w.load_unsafe(); @@ -1504,7 +1611,13 @@ template quad_lid); } -template +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1556,7 +1669,13 @@ template simd_lid); } -template +template < + typename T, + int group_size, + const int bits, + bool batched, + bool has_global_scale = false, + int results_per_simdgroup = 4> [[kernel]] void affine_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1608,7 +1727,71 @@ template simd_lid); } -template +template < + typename T, + int group_size, + int bits, + int vecs_per_tg, + int k_lanes, + bool batched> +[[kernel]] void affine_qmv_wide( + const device uint32_t* w, + const device T* scales, + const device T* biases, + const device T* x, + device T* y, + const constant int& in_vec_size, + const constant int& out_vec_size, + const constant int& M, + const constant int& x_batch_ndims, + const constant int* x_shape, + const constant int64_t* x_strides, + const constant int& w_batch_ndims, + const constant int* w_shape, + const constant int64_t* w_strides, + const constant int64_t* s_strides, + const constant int64_t* b_strides, + uint3 tid [[threadgroup_position_in_grid]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + if (batched) { + adjust_matrix_offsets( + x, + w, + scales, + biases, + y, + out_vec_size * M, + x_batch_ndims, + x_shape, + x_strides, + w_batch_ndims, + w_shape, + w_strides, + s_strides, + b_strides, + tid); + } + qmv_wide_impl( + w, + scales, + biases, + x, + y, + in_vec_size, + out_vec_size, + M, + tid, + simd_gid, + simd_lid); +} + +template < + typename T, + const int group_size, + const int bits, + bool batched, + bool has_global_scale = false> [[kernel]] void affine_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1655,6 +1838,7 @@ template y, in_vec_size, out_vec_size, + in_vec_size, tid, simd_gid, simd_lid); @@ -1703,6 +1887,9 @@ template int in_vec_size_adj = tid.z % split_k == split_k - 1 ? final_block_size : in_vec_size; + // The in_vec_stride is the full K dimension, not the partition size + int in_vec_stride = (split_k - 1) * in_vec_size + final_block_size; + qvm_impl( w, scales, @@ -1711,6 +1898,7 @@ template y, in_vec_size_adj, out_vec_size, + in_vec_stride, tid, simd_gid, simd_lid); @@ -1772,17 +1960,91 @@ template < tid); } qmm_t_impl( - w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < typename T, const int group_size, const int bits, - const bool batched, + const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32> +[[kernel]] void affine_qmm_t_splitk( + const device uint32_t* w [[buffer(0)]], + const device T* scales [[buffer(1)]], + const device T* biases [[buffer(2)]], + const device T* x [[buffer(3)]], + device T* y [[buffer(4)]], + const constant int& K [[buffer(5)]], + const constant int& N [[buffer(6)]], + const constant int& M [[buffer(7)]], + const constant int& k_partition_size [[buffer(8)]], + const constant int& split_k_partition_stride [[buffer(9)]], + uint3 tid [[threadgroup_position_in_grid]], + uint lid [[thread_index_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + (void)lid; + + constexpr int BK_padded = (BK + 16 / sizeof(T)); + constexpr int pack_factor = get_pack_factor(); + constexpr int bytes_per_pack = get_bytes_per_pack(); + + threadgroup T Xs[BM * BK_padded]; + threadgroup T Ws[BN * BK_padded]; + + const int k_start = tid.z * k_partition_size; + x += k_start; + + auto wl = (const device uint8_t*)w; + wl += k_start * bytes_per_pack / pack_factor; + scales += k_start / group_size; + biases += k_start / group_size; + y += tid.z * static_cast(split_k_partition_stride); + + qmm_t_impl( + (const device uint32_t*)wl, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + k_partition_size, + tid, + lid, + simd_gid, + simd_lid); +} + +template < + typename T, + int group_size, + int bits, + bool batched, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1835,7 +2097,7 @@ template < w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); } -template +template [[kernel]] void affine_gather_qmv_fast( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1897,7 +2159,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qmv( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -1959,7 +2221,7 @@ template simd_lid); } -template +template [[kernel]] void affine_gather_qvm( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2016,6 +2278,7 @@ template y, in_vec_size, out_vec_size, + in_vec_size, tid, simd_gid, simd_lid); @@ -2086,16 +2349,31 @@ template < b_strides, tid); qmm_t_impl( - w, scales, biases, x, y, Xs, Ws, K, N, M, tid, lid, simd_gid, simd_lid); + w, + scales, + biases, + x, + y, + Xs, + Ws, + K, + N, + M, + K, + tid, + lid, + simd_gid, + simd_lid); } template < typename T, - const int group_size, - const int bits, - const int BM = 32, - const int BK = 32, - const int BN = 32> + int group_size, + int bits, + bool has_global_scale = false, + int BM = 32, + int BK = 32, + int BN = 32> [[kernel]] void affine_gather_qmm_n( const device uint32_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], @@ -2353,7 +2631,7 @@ template < } } -template +template [[kernel]] void affine_quantize( const device T* w [[buffer(0)]], device uint8_t* out [[buffer(1)]], @@ -2458,7 +2736,7 @@ template } } -template +template [[kernel]] void affine_dequantize( const device uint8_t* w [[buffer(0)]], const device T* scales [[buffer(1)]], diff --git a/Source/Cmlx/mlx-generated/quantized_nax.cpp b/Source/Cmlx/mlx-generated/quantized_nax.cpp index 2607fb2ab..123a4a37c 100644 --- a/Source/Cmlx/mlx-generated/quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/quantized_nax.cpp @@ -504,17 +504,16 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { bits == 8, "Template undefined for bits not in {2, 3, 4, 5, 6, 8}"); + const float s = float(scale); + const float b = float(bias); + if (bits == 2) { - U s[4] = { - scale, - scale / static_cast(4.0f), - scale / static_cast(16.0f), - scale / static_cast(64.0f)}; + float sc[4] = {s, s / 4.0f, s / 16.0f, s / 64.0f}; for (int i = 0; i < (N / 4); i++) { - w_local[4 * i] = s[0] * (w[i] & 0x03) + bias; - w_local[4 * i + 1] = s[1] * (w[i] & 0x0c) + bias; - w_local[4 * i + 2] = s[2] * (w[i] & 0x30) + bias; - w_local[4 * i + 3] = s[3] * (w[i] & 0xc0) + bias; + w_local[4 * i] = static_cast(sc[0] * (w[i] & 0x03) + b); + w_local[4 * i + 1] = static_cast(sc[1] * (w[i] & 0x0c) + b); + w_local[4 * i + 2] = static_cast(sc[2] * (w[i] & 0x30) + b); + w_local[4 * i + 3] = static_cast(sc[3] * (w[i] & 0xc0) + b); } } @@ -523,22 +522,24 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 3 * i; - w_local[0] = (w[0] & 0x7) * scale + bias; - w_local[1] = ((w[0] & 0x38) >> 3) * scale + bias; - w_local[2] = (((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * scale + bias; - w_local[3] = ((w[1] & 0xe) >> 1) * scale + bias; - w_local[4] = ((w[1] & 0x70) >> 4) * scale + bias; - w_local[5] = (((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * scale + bias; - w_local[6] = ((w[2] & 0x1c) >> 2) * scale + bias; - w_local[7] = ((w[2] & 0xe0) >> 5) * scale + bias; + w_local[0] = static_cast((w[0] & 0x7) * s + b); + w_local[1] = static_cast(((w[0] & 0x38) >> 3) * s + b); + w_local[2] = + static_cast((((w[0] & 0xc0) >> 6) + ((w[1] & 0x1) << 2)) * s + b); + w_local[3] = static_cast(((w[1] & 0xe) >> 1) * s + b); + w_local[4] = static_cast(((w[1] & 0x70) >> 4) * s + b); + w_local[5] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0x3) << 1)) * s + b); + w_local[6] = static_cast(((w[2] & 0x1c) >> 2) * s + b); + w_local[7] = static_cast(((w[2] & 0xe0) >> 5) * s + b); } } else if (bits == 4) { - U s[2] = {scale, scale / static_cast(16.0f)}; + float sc[2] = {s, s / 16.0f}; for (int i = 0; i < (N / 2); i++) { - w_local[2 * i] = s[0] * (w[i] & 0x0f) + bias; - w_local[2 * i + 1] = s[1] * (w[i] & 0xf0) + bias; + w_local[2 * i] = static_cast(sc[0] * (w[i] & 0x0f) + b); + w_local[2 * i + 1] = static_cast(sc[1] * (w[i] & 0xf0) + b); } } @@ -547,14 +548,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { w_local += 8 * i; w += 5 * i; - w_local[0] = (w[0] & 0x1f) * scale + bias; - w_local[1] = (((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * scale + bias; - w_local[2] = ((w[1] & 0x7c) >> 2) * scale + bias; - w_local[3] = (((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * scale + bias; - w_local[4] = (((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * scale + bias; - w_local[5] = ((w[3] & 0x3e) >> 1) * scale + bias; - w_local[6] = (((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * scale + bias; - w_local[7] = ((w[4] & 0xf8) >> 3) * scale + bias; + w_local[0] = static_cast((w[0] & 0x1f) * s + b); + w_local[1] = + static_cast((((w[0] & 0xe0) >> 5) + ((w[1] & 0x3) << 3)) * s + b); + w_local[2] = static_cast(((w[1] & 0x7c) >> 2) * s + b); + w_local[3] = + static_cast((((w[1] & 0x80) >> 7) + ((w[2] & 0xf) << 1)) * s + b); + w_local[4] = + static_cast((((w[2] & 0xf0) >> 4) + ((w[3] & 0x1) << 4)) * s + b); + w_local[5] = static_cast(((w[3] & 0x3e) >> 1) * s + b); + w_local[6] = + static_cast((((w[3] & 0xc0) >> 6) + ((w[4] & 0x7) << 2)) * s + b); + w_local[7] = static_cast(((w[4] & 0xf8) >> 3) * s + b); } } @@ -562,16 +567,18 @@ dequantize(const device uint8_t* w, U scale, U bias, threadgroup U* w_local) { for (int i = 0; i < (N / 4); i++) { w_local += 4 * i; w += 3 * i; - w_local[0] = (w[0] & 0x3f) * scale + bias; - w_local[1] = (((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * scale + bias; - w_local[2] = (((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * scale + bias; - w_local[3] = ((w[2] >> 2) & 0x3f) * scale + bias; + w_local[0] = static_cast((w[0] & 0x3f) * s + b); + w_local[1] = + static_cast((((w[0] >> 6) & 0x03) + ((w[1] & 0x0f) << 2)) * s + b); + w_local[2] = + static_cast((((w[1] >> 4) & 0x0f) + ((w[2] & 0x03) << 4)) * s + b); + w_local[3] = static_cast(((w[2] >> 2) & 0x3f) * s + b); } } else if (bits == 8) { for (int i = 0; i < N; i++) { - w_local[i] = scale * w[i] + bias; + w_local[i] = static_cast(s * w[i] + b); } } } @@ -625,15 +632,15 @@ struct QuantizedBlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), group_step_cnt(0), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), dst(dst_ + bi * dst_ld + bj * pack_factor), src(src_ + bi * src_ld * bytes_per_pack / pack_factor + @@ -641,7 +648,7 @@ struct QuantizedBlockLoader { scales(scales_ + bi * src_ld / group_size), biases(biases_ + bi * src_ld / group_size) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -654,7 +661,7 @@ struct QuantizedBlockLoader { } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -684,7 +691,7 @@ struct QuantizedBlockLoader { } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { if (group_steps > 1) { @@ -765,14 +772,14 @@ struct QuantizedBlockLoader< const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), tile_stride( - reduction_dim ? BCOLS_PACKED * bytes_per_pack + reduction_dim ? BCOLS_PACKED* bytes_per_pack : BROWS * src_ld * bytes_per_pack / pack_factor), - group_stride(BROWS * src_ld / group_size), + group_stride(BROWS* src_ld / group_size), thread_idx(simd_group_id * 32 + simd_lane_id), - bi(n_reads * thread_idx / BCOLS_PACKED), + bi(n_reads* thread_idx / BCOLS_PACKED), bj((n_reads * thread_idx) % BCOLS_PACKED), group_id((bj * pack_factor) / group_size), dst(dst_ + bi * dst_ld + bj * pack_factor), @@ -781,7 +788,7 @@ struct QuantizedBlockLoader< scales(scales_ + bi * src_ld / group_size + group_id), biases(biases_ + bi * src_ld / group_size + group_id) {} - void load_unsafe() const { + void load_unsafe() const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -794,7 +801,7 @@ struct QuantizedBlockLoader< } } - void load_safe(short2 src_tile_dim) const { + void load_safe(short2 src_tile_dim) const thread { if (BCOLS_PACKED * BROWS < tgp_size && bi >= BROWS) { return; } @@ -824,7 +831,7 @@ struct QuantizedBlockLoader< } } - void next() { + void next() thread { src += tile_stride; if (reduction_dim == 1) { // if (group_steps > 1) { @@ -839,8 +846,8 @@ struct QuantizedBlockLoader< biases += n_groups; // } } else { - scales += n_groups * group_stride; - biases += n_groups * group_stride; + scales += group_stride; + biases += group_stride; } } }; @@ -999,16 +1006,13 @@ METAL_FUNC void qmm_t_nax_tgp_impl( // Make the weight loader loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); @@ -1016,22 +1020,17 @@ METAL_FUNC void qmm_t_nax_tgp_impl( constexpr bool transpose_a = false; constexpr bool transpose_b = true; - const short sgp_sm = min(SM, short(M - (y_row + tm))); + const short sgp_sm = min(int(SM), M - (y_row + tm)); const bool is_unaligned_sm = (sgp_sm != SM); - const short sgp_sn = aligned_N ? SN : min(SN, short(N - (y_col + tn))); + const short sgp_sn = aligned_N ? SN : min(int(SN), N - (y_col + tn)); const short tgp_bn = aligned_N ? BN : min(BN, int(N - (y_col))); const bool is_unaligned_bn = aligned_N ? false : (tgp_bn != BN); using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -1050,8 +1049,8 @@ METAL_FUNC void qmm_t_nax_tgp_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; @@ -1115,7 +1114,6 @@ METAL_FUNC void qmm_n_nax_tgp_impl( uint simd_gid [[simdgroup_index_in_threadgroup]], uint simd_lid [[thread_index_in_simdgroup]]) { (void)lid; - (void)M; static_assert(BK >= SIMD_SIZE, "BK should be larger than SIMD_SIZE"); static_assert(BK % SIMD_SIZE == 0, "BK should be divisible by SIMD_SIZE"); @@ -1136,38 +1134,34 @@ METAL_FUNC void qmm_n_nax_tgp_impl( bits>; // Set the block - const int K_w = K * bytes_per_pack / pack_factor; - const int K_g = K / group_size; const int y_row = tid.y * BM; const int y_col = tid.x * BN; auto wl = (const device uint8_t*)w; + // Here w is [K, N]: packed and group-quantized along N, with row stride N. x += y_row * static_cast(K); - wl += y_col * K_w; - scales += y_col * K_g; - biases += y_col * K_g; + wl += y_col * bytes_per_pack / pack_factor; + scales += y_col / group_size; + biases += y_col / group_size; y += y_row * static_cast(N) + y_col; - // Make the x loader and mma operation - // const short num_els = min(BM, M - y_row); - // const short num_outs = min(BN, N - y_col); - loader_w_t loader_w(wl, scales, biases, K, Ws, simd_gid, simd_lid); + // Make the weight loader + loader_w_t loader_w(wl, scales, biases, N, Ws, simd_gid, simd_lid); - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_gid / WN); const short tn = SN * (simd_gid % WN); + const short sgp_sm = min(int(SM), M - (y_row + tm)); + const short ldb_tgp = BN_padded; constexpr bool transpose_a = false; @@ -1175,12 +1169,7 @@ METAL_FUNC void qmm_n_nax_tgp_impl( using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - - NAXTile Dtile; - + NAXTile Dtile; Dtile.clear(); x += tm * K; @@ -1192,12 +1181,17 @@ METAL_FUNC void qmm_n_nax_tgp_impl( STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + NAXTile Atile; + NAXTile Btile; volatile int compiler_barrier; - Atile.load(x + kk1, K); + if (sgp_sm == SM) { + Atile.load(x + kk1, K); + } else { + Atile.load_safe(x + kk1, K, short2(SK, sgp_sm)); + } + Btile.template load(Ws + tn + kk1 * ldb_tgp); tile_matmad_nax( @@ -1217,7 +1211,11 @@ METAL_FUNC void qmm_n_nax_tgp_impl( // Store results to device memory threadgroup_barrier(mem_flags::mem_threadgroup); - Dtile.store(y + tm * N + tn, N); + if (sgp_sm == SM) { + Dtile.store(y + tm * N + tn, N); + } else { + Dtile.store_safe(y + tm * N + tn, N, short2(SN, sgp_sm)); + } } template < @@ -1547,16 +1545,13 @@ template < scales += transpose ? y_col_long * K_g : y_col / group_size; biases += transpose ? y_col_long * K_g : y_col / group_size; - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; - constexpr short TK = SK / UK; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + constexpr short TK = SK / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); @@ -1574,10 +1569,6 @@ template < using AccumType = float; - using ASubTile = NAXSubTile; - using BSubTile = NAXSubTile; - using DSubTile = NAXSubTile; - // Do as many matmuls as necessary uint32_t index; short offset; @@ -1598,8 +1589,11 @@ template < } threadgroup_barrier(mem_flags::mem_none); - NAXTile Dtile; + const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); + const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); + const bool sg_active = m_hi_lim > m_lo_lim; + NAXTile Dtile; Dtile.clear(); const device T* xn = x + tm * K; @@ -1629,31 +1623,33 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; - - volatile int compiler_barrier; - - if constexpr (kAlignedM.value) { - Atile.load(xn + kk1, K); - } else { - Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + if (sg_active) { + NAXTile Atile; + NAXTile Btile; + + volatile int compiler_barrier; + + if constexpr (kAlignedM.value) { + Atile.load(xn + kk1, K); + } else { + Atile.load_safe(xn + kk1, K, short2(SK, sgp_sm)); + } + + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } + + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); + + (void)compiler_barrier; } - - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } - - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); - - (void)compiler_barrier; } xn += BK; @@ -1667,36 +1663,35 @@ template < STEEL_PRAGMA_NO_UNROLL for (int kk1 = 0; kk1 < BK; kk1 += SK) { - NAXTile Atile; - NAXTile Btile; + if (sg_active) { + NAXTile Atile; + NAXTile Btile; - volatile int compiler_barrier; + volatile int compiler_barrier; - const short psk = min(int(SK), max(0, (BK - kk1))); - Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); + const short psk = min(int(SK), max(0, (BK - kk1))); + Atile.load_safe(xn + kk1, K, short2(psk, sgp_sm)); - if constexpr (transpose) { - Btile.template load(Ws + tn * BK_padded + kk1); - } else { - Btile.template load(Ws + tn + kk1 * BN_padded); - } + if constexpr (transpose) { + Btile.template load(Ws + tn * BK_padded + kk1); + } else { + Btile.template load(Ws + tn + kk1 * BN_padded); + } - tile_matmad_nax( - Dtile, - Atile, - metal::bool_constant{}, - Btile, - metal::bool_constant{}); + tile_matmad_nax( + Dtile, + Atile, + metal::bool_constant{}, + Btile, + metal::bool_constant{}); - (void)compiler_barrier; + (void)compiler_barrier; + } } } threadgroup_barrier(mem_flags::mem_threadgroup); - const short m_lo_lim = min(int(sgp_sm), max(0, offset - tm)); - const short m_hi_lim = min(int(sgp_sm), max(0, offset_next - tm)); - // Store results to device memory if constexpr (kAlignedN.value) { if (m_lo_lim == 0 && m_hi_lim == SM) { diff --git a/Source/Cmlx/mlx-generated/reduce.cpp b/Source/Cmlx/mlx-generated/reduce.cpp index cc29a385e..cac2d8de5 100644 --- a/Source/Cmlx/mlx-generated/reduce.cpp +++ b/Source/Cmlx/mlx-generated/reduce.cpp @@ -50,13 +50,13 @@ template < for (IdxT b = 0; b < blocks; b++) { for (int i = 0; i < N_READS; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } in += lsize.x * N_READS; } if (extra > 0) { for (int i = 0; i < extra; i++) { - total = op(static_cast(in[i]), total); + total = op(cast_to(in[i]), total); } } @@ -128,13 +128,13 @@ template row = in + loop.location(); if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -210,7 +210,7 @@ template for (IdxT r = gid.z * lsize.y + lid.y; r < total_rows; r += lsize.y * gsize.z) { row = in + loop.location(); - total = op(static_cast(*row), total); + total = op(cast_to(*row), total); loop.next(lsize.y * gsize.z, reduce_shape, reduce_strides); } @@ -292,13 +292,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -437,13 +437,13 @@ template < if (safe) { for (int i = 0; i < n_reads; i++) { - totals[i] = op(static_cast(row[i]), totals[i]); + totals[i] = op(cast_to(row[i]), totals[i]); } } else { U vals[n_reads]; for (int i = 0; i < n_reads; i++) { vals[i] = - (column + i < reduction_stride) ? static_cast(row[i]) : op.init; + (column + i < reduction_stride) ? cast_to(row[i]) : op.init; } for (int i = 0; i < n_reads; i++) { totals[i] = op(vals[i], totals[i]); @@ -537,7 +537,7 @@ METAL_FUNC void per_thread_row_reduce( for (int i = 0; i < blocks; i++) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } inputs[j] += lsize_x * N_READS; @@ -549,13 +549,13 @@ METAL_FUNC void per_thread_row_reduce( if (index + N_READS <= extra) { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; i < N_READS; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } else { for (int j = 0; j < N_WRITES; j++) { for (int i = 0; index + i < extra; i++) { - totals[j] = op(static_cast(inputs[j][i]), totals[j]); + totals[j] = op(cast_to(inputs[j][i]), totals[j]); } } } @@ -840,7 +840,8 @@ template < // lid.x * N_READS breaks the per_thread_row_reduce interface a bit. Maybe it // needs a small refactor. - in += elem_to_loc(out_idx, shape, strides, ndim) + lid.x * N_READS; + in += + elem_to_loc(out_idx, shape, strides, ndim) + IdxT(lid.x) * N_READS; LoopedElemToLoc 2)> loop(reduce_ndim); const device T* row; diff --git a/Source/Cmlx/mlx-generated/reduce_utils.cpp b/Source/Cmlx/mlx-generated/reduce_utils.cpp index 75fdf28f0..d277363ba 100644 --- a/Source/Cmlx/mlx-generated/reduce_utils.cpp +++ b/Source/Cmlx/mlx-generated/reduce_utils.cpp @@ -107,6 +107,19 @@ METAL_FUNC void mlx_atomic_fetch_mul_explicit( T expected = mlx_atomic_load_explicit(object, offset); while (!mlx_atomic_compare_exchange_weak_explicit( object, &expected, val * expected, offset)) { + // Workaround: Metal's atomic_compare_exchange_weak_explicit does + // not perform bitwise comparison as required by the C++ atomics spec. + // The compiler lowers the success check to `fcmp fast ueq` under + // no-nans-fp-math, which evaluates to false when either operand is NaN - + // even when the bit patterns are identical. With NaN in memory the CAS + // can never succeed, so the loop spins. Bail out instead: memory is + // already NaN and that is the correct reduction result regardless of + // this thread's update. + if constexpr (metal::is_floating_point_v) { + if (isnan(expected)) { + break; + } + } } } @@ -171,7 +184,7 @@ union uint_or_packed { template struct mlx_atomic_update_helper { - uint operator()(uint_or_packed init, T update, size_t elem_offset) { + uint operator()(uint_or_packed init, T update, size_t elem_offset) thread { Op op; init.val[elem_offset] = op(update, init.val[elem_offset]); return init.bits; @@ -208,7 +221,7 @@ struct __None { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { #pragma unused(b) return a; } @@ -222,7 +235,7 @@ struct __Add { return true; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a + b; } }; @@ -234,7 +247,7 @@ struct __Mul { return b != 0; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return a * b; } }; @@ -245,7 +258,7 @@ struct __Max { return a > b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return max(a, b); } }; @@ -256,7 +269,7 @@ struct __Min { return a < b; } - T operator()(T a, T b) { + T operator()(T a, T b) thread { return min(a, b); } }; @@ -369,12 +382,12 @@ METAL_FUNC bool mlx_atomic_compare_exchange_weak_explicit( #define DEFINE_SIMD_REDUCE() \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ return simd_reduce_impl(val); \ } \ \ template = true> \ - T simd_reduce(T val) { \ + T simd_reduce(T val) thread { \ for (short i = simd_size / 2; i > 0; i /= 2) { \ val = operator()(val, simd_shuffle_down(val, i)); \ } \ @@ -390,7 +403,8 @@ union bool4_or_uint { struct None { template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_store_explicit(out, val, offset); } }; @@ -399,7 +413,7 @@ template struct And { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_all(val); } @@ -409,7 +423,7 @@ struct And { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (!val) { bool4_or_uint update; update.b = {true, true, true, true}; @@ -418,20 +432,20 @@ struct And { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (!val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out &= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a && b; } }; @@ -440,7 +454,7 @@ template struct Or { DEFINE_SIMD_REDUCE() - bool simd_reduce_impl(bool val) { + bool simd_reduce_impl(bool val) thread { return simd_any(val); } @@ -450,7 +464,7 @@ struct Or { device mlx_atomic* out, bool val, int elem_idx, - size_t offset = 0) { + size_t offset = 0) thread { if (val) { bool4_or_uint update; update.b = {false, false, false, false}; @@ -459,20 +473,20 @@ struct Or { } } - void - atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, bool val, size_t offset = 0) + thread { if (val) { mlx_atomic_store_explicit(out, val, offset); } } // Non atomic update - void update(device bool* out, bool val) { + void update(device bool* out, bool val) thread { *out |= val; } // Operator - bool operator()(bool a, bool b) { + bool operator()(bool a, bool b) thread { return a || b; } }; @@ -482,19 +496,29 @@ struct Sum { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_sum(val); } static constexpr constant U init = U(0); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_add_explicit(out, val, offset); } + void atomic_update( + device mlx_atomic* out, + complex64_t val, + size_t offset = 0) thread { + auto out_lanes = reinterpret_cast*>(out); + mlx_atomic_fetch_add_explicit(out_lanes, val.real, 2 * offset); + mlx_atomic_fetch_add_explicit(out_lanes, val.imag, 2 * offset + 1); + } + // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a + b; } }; @@ -504,19 +528,20 @@ struct Prod { DEFINE_SIMD_REDUCE() template - T simd_reduce_impl(T val) { + T simd_reduce_impl(T val) thread { return simd_product(val); } static constexpr constant U init = U(1); template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_mul_explicit(out, val, offset); } // Operator - U operator()(U a, U b) { + U operator()(U a, U b) thread { return a * b; } }; @@ -526,12 +551,14 @@ struct Min { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_min(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -541,18 +568,19 @@ struct Min { static constexpr constant U init = Limits::max; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_min_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a < b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -561,7 +589,7 @@ struct Min { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); @@ -583,12 +611,14 @@ struct Max { DEFINE_SIMD_REDUCE() template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { return simd_max(val); } template - metal::enable_if_t, T> simd_reduce_impl(T val) { + metal::enable_if_t, T> simd_reduce_impl( + T val) thread { if (simd_any(val != val)) { return static_cast(NAN); } @@ -598,18 +628,19 @@ struct Max { static constexpr constant U init = Limits::min; template - void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) { + void atomic_update(device mlx_atomic* out, T val, size_t offset = 0) + thread { mlx_atomic_fetch_max_explicit(out, val, offset); } // Operator template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { return a > b ? a : b; } template - metal::enable_if_t, T> operator()(T a, T b) { + metal::enable_if_t, T> operator()(T a, T b) thread { if (metal::isnan(a) || metal::isnan(b)) { return static_cast(NAN); } else { @@ -618,7 +649,7 @@ struct Max { } template <> - complex64_t operator()(complex64_t a, complex64_t b) { + complex64_t operator()(complex64_t a, complex64_t b) thread { bool real_is_nan = metal::isnan(a.real) || metal::isnan(b.real); bool imag_is_nan = metal::isnan(a.imag) || metal::isnan(b.imag); diff --git a/Source/Cmlx/mlx-generated/scan.cpp b/Source/Cmlx/mlx-generated/scan.cpp index 05e13b2b3..1fd1d74dc 100644 --- a/Source/Cmlx/mlx-generated/scan.cpp +++ b/Source/Cmlx/mlx-generated/scan.cpp @@ -21,33 +21,40 @@ constant mlx::os_log logger("mlx", "binary_ops"); struct Add { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x + y; } }; struct FloorDivide { template - T operator()(T x, T y) { + metal::enable_if_t & !metal::is_signed_v, T> + operator()(T x, T y) thread { return x / y; } - template <> - float operator()(float x, float y) { - return trunc(x / y); + template + metal::enable_if_t & metal::is_signed_v, T> + operator()(T x, T y) thread { + auto q = x / y; + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + return q; } - template <> - half operator()(half x, half y) { - return trunc(x / y); + template + metal::enable_if_t, T> operator()(T x, T y) thread { + return floor(x / y); } template <> - bfloat16_t operator()(bfloat16_t x, bfloat16_t y) { - return trunc(x / y); + complex64_t operator()(complex64_t x, complex64_t y) thread { + // Complex is not supported, simply make compiler happy. + return x / y; } }; struct Divide { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x / y; } }; @@ -55,12 +62,12 @@ struct Divide { struct Remainder { template metal::enable_if_t & !metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { return x % y; } template metal::enable_if_t & metal::is_signed_v, T> - operator()(T x, T y) { + operator()(T x, T y) thread { auto r = x % y; if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -68,7 +75,7 @@ struct Remainder { return r; } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { T r = fmod(x, y); if (r != 0 && (r < 0 != y < 0)) { r += y; @@ -76,25 +83,25 @@ struct Remainder { return r; } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { return x % y; } }; struct Equal { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y; } }; struct NaNEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x == y || (metal::isnan(x) && metal::isnan(y)); } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x == y || (metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) && metal::isnan(y.imag)) || @@ -105,35 +112,35 @@ struct NaNEqual { struct Greater { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x > y; } }; struct GreaterEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x >= y; } }; struct Less { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x < y; } }; struct LessEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x <= y; } }; struct LogAddExp { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { if (metal::isnan(x) || metal::isnan(y)) { return metal::numeric_limits::quiet_NaN(); } @@ -145,7 +152,7 @@ struct LogAddExp { : (maxval + log1p(metal::exp(minval - maxval))); }; - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) || metal::isnan(y.imag)) { return metal::numeric_limits::quiet_NaN(); @@ -166,12 +173,12 @@ struct LogAddExp { struct Maximum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::max(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -179,7 +186,7 @@ struct Maximum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -189,12 +196,12 @@ struct Maximum { struct Minimum { template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { return metal::min(x, y); } template - metal::enable_if_t, T> operator()(T x, T y) { + metal::enable_if_t, T> operator()(T x, T y) thread { if (metal::isnan(x)) { return x; } @@ -202,7 +209,7 @@ struct Minimum { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (metal::isnan(x.real) || metal::isnan(x.imag)) { return x; } @@ -212,30 +219,32 @@ struct Minimum { struct Multiply { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x * y; } }; struct NotEqual { template - bool operator()(T x, T y) { + bool operator()(T x, T y) thread { return x != y; } template <> - bool operator()(complex64_t x, complex64_t y) { + bool operator()(complex64_t x, complex64_t y) thread { return x.real != y.real || x.imag != y.imag; } }; struct Power { template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { return metal::pow(base, exp); } template - metal::enable_if_t, T> operator()(T base, T exp) { + metal::enable_if_t, T> operator()(T base, T exp) + thread { T res = 1; // Undefined to raise integer to negative power if (exp < 0) { @@ -255,7 +264,7 @@ struct Power { } template <> - complex64_t operator()(complex64_t x, complex64_t y) { + complex64_t operator()(complex64_t x, complex64_t y) thread { if (x.real == 0 && x.imag == 0) { if (metal::isnan(y.real) || metal::isnan(y.imag)) { auto nan = metal::numeric_limits::quiet_NaN(); @@ -273,70 +282,70 @@ struct Power { struct Subtract { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x - y; } }; struct LogicalAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x && y; }; }; struct LogicalOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x || y; }; }; struct BitwiseAnd { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x & y; }; }; struct BitwiseOr { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x | y; }; }; struct BitwiseXor { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x ^ y; }; }; struct LeftShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x << y; }; }; struct RightShift { template - T operator()(T x, T y) { + T operator()(T x, T y) thread { return x >> y; }; }; struct ArcTan2 { template - T operator()(T y, T x) { + T operator()(T y, T x) thread { return metal::precise::atan2(y, x); } }; struct DivMod { template - metal::array operator()(T x, T y) { + metal::array operator()(T x, T y) thread { return {FloorDivide{}(x, y), Remainder{}(x, y)}; }; }; @@ -352,12 +361,12 @@ struct DivMod { #define DEFINE_SIMD_SCAN() \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ return simd_scan_impl(val); \ } \ \ template = true> \ - T simd_scan(T val) { \ + T simd_scan(T val) thread { \ for (int i = 1; i <= 16; i *= 2) { \ val = operator()(val, simd_shuffle_and_fill_up(val, init, i)); \ } \ @@ -366,12 +375,12 @@ struct DivMod { #define DEFINE_SIMD_EXCLUSIVE_SCAN() \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ return simd_exclusive_scan_impl(val); \ } \ \ template = true> \ - T simd_exclusive_scan(T val) { \ + T simd_exclusive_scan(T val) thread { \ val = simd_scan(val); \ return simd_shuffle_and_fill_up(val, init, 1); \ } @@ -384,15 +393,15 @@ struct CumSum { static constexpr constant U init = static_cast(0); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a + b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_sum(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_sum(x); } }; @@ -405,15 +414,15 @@ struct CumProd { static constexpr constant U init = static_cast(1.0f); template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return a * b; } - U simd_scan_impl(U x) { + U simd_scan_impl(U x) thread { return simd_prefix_inclusive_product(x); } - U simd_exclusive_scan_impl(U x) { + U simd_exclusive_scan_impl(U x) thread { return simd_prefix_exclusive_product(x); } }; @@ -423,11 +432,11 @@ struct CumProd { static constexpr constant bool init = true; template - bool operator()(bool a, T b) { + bool operator()(bool a, T b) thread { return a & static_cast(b); } - bool simd_scan(bool x) { + bool simd_scan(bool x) thread { for (int i = 1; i <= 16; i *= 2) { bool other = simd_shuffle_and_fill_up(x, init, i); x &= other; @@ -435,7 +444,7 @@ struct CumProd { return x; } - bool simd_exclusive_scan(bool x) { + bool simd_exclusive_scan(bool x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -445,20 +454,29 @@ template struct CumMax { static constexpr constant U init = Limits::min; - template - U operator()(U a, T b) { + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } return (a >= b) ? a : b; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + return combine(a, static_cast(b)); + } + + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x >= other) ? x : other; + x = combine(x, other); } return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -468,20 +486,29 @@ template struct CumMin { static constexpr constant U init = Limits::max; - template - U operator()(U a, T b) { + static U combine(U a, U b) { + if constexpr (metal::is_floating_point_v) { + if (metal::isnan(a) || metal::isnan(b)) { + return metal::numeric_limits::quiet_NaN(); + } + } return (a <= b) ? a : b; } - U simd_scan(U x) { + template + U operator()(U a, T b) thread { + return combine(a, static_cast(b)); + } + + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); - x = (x <= other) ? x : other; + x = combine(x, other); } return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } @@ -492,11 +519,11 @@ struct CumLogaddexp { static constexpr constant U init = Limits::min; template - U operator()(U a, T b) { + U operator()(U a, T b) thread { return LogAddExp{}(a, static_cast(b)); } - U simd_scan(U x) { + U simd_scan(U x) thread { for (int i = 1; i <= 16; i *= 2) { U other = simd_shuffle_and_fill_up(x, init, i); x = LogAddExp{}(x, other); @@ -504,7 +531,7 @@ struct CumLogaddexp { return x; } - U simd_exclusive_scan(U x) { + U simd_exclusive_scan(U x) thread { x = simd_scan(x); return simd_shuffle_and_fill_up(x, init, 1); } diff --git a/Source/Cmlx/mlx-generated/scatter.cpp b/Source/Cmlx/mlx-generated/scatter.cpp index a07727a5f..b8bf7c108 100644 --- a/Source/Cmlx/mlx-generated/scatter.cpp +++ b/Source/Cmlx/mlx-generated/scatter.cpp @@ -97,6 +97,86 @@ METAL_FUNC void scatter_impl( } } +template < + typename T, + typename IdxT, + typename Op, + bool OUT_ROW_CONTIG, + bool UPD_ROW_CONTIG, + bool UPD_SCALAR, + int NWORK, + int NDIM> +[[kernel]] void slice_update_op_impl( + const device T* updates [[buffer(0)]], + device T* out [[buffer(1)]], + const constant int* update_shape [[buffer(2)]], + const constant int64_t* update_strides [[buffer(3)]], + const constant int& update_ndim [[buffer(4)]], + const constant int64_t& update_size [[buffer(5)]], + const constant int64_t* output_strides [[buffer(6)]], + const constant int64_t& output_offset [[buffer(7)]], + uint3 gid [[thread_position_in_grid]], + uint3 gsize [[threads_per_grid]]) { + Op op; + + IdxT idx = + (IdxT(gid.z) * IdxT(gsize.y) + IdxT(gid.y)) * IdxT(gsize.x) * NWORK + + IdxT(gid.x) * NWORK; + IdxT out_idx; + IdxT update_idx; + + if constexpr (OUT_ROW_CONTIG) { + out_idx = idx; + } else if constexpr (NDIM == 1) { + out_idx = NWORK * gid.x * output_strides[0]; + } else if constexpr (NDIM == 2) { + out_idx = gid.y * output_strides[0] + NWORK * gid.x * output_strides[1]; + } else if constexpr (NDIM == 3) { + out_idx = gid.z * output_strides[0] + gid.y * output_strides[1] + + NWORK * gid.x * output_strides[2]; + } else { + out_idx = elem_to_loc(idx, update_shape, output_strides, update_ndim); + } + + if constexpr (UPD_SCALAR) { + update_idx = 0; + } else if constexpr (UPD_ROW_CONTIG) { + update_idx = idx; + } else if constexpr (NDIM == 1) { + update_idx = NWORK * gid.x * update_strides[0]; + } else if constexpr (NDIM == 2) { + update_idx = gid.y * update_strides[0] + NWORK * gid.x * update_strides[1]; + } else if constexpr (NDIM == 3) { + update_idx = gid.z * update_strides[0] + gid.y * update_strides[1] + + NWORK * gid.x * update_strides[2]; + } else { + update_idx = + elem_to_loc(idx, update_shape, update_strides, update_ndim); + } + + out += output_offset; + + if constexpr (OUT_ROW_CONTIG && (UPD_ROW_CONTIG || UPD_SCALAR)) { + for (int j = 0; j < NWORK; j++) { + out[out_idx] = op(out[out_idx], updates[update_idx]); + out_idx++; + if constexpr (!UPD_SCALAR) { + update_idx++; + } + } + } else { + auto out_stride = output_strides[update_ndim - 1]; + auto update_stride = update_strides[update_ndim - 1]; + for (int j = 0; j < NWORK; j++) { + out[out_idx] = op(out[out_idx], updates[update_idx]); + out_idx += out_stride; + if constexpr (!UPD_SCALAR) { + update_idx += update_stride; + } + } + } +} + /////////////////////////////////////////////////////////////////////////////// )preamble"; } diff --git a/Source/Cmlx/mlx-generated/searchsorted.cpp b/Source/Cmlx/mlx-generated/searchsorted.cpp new file mode 100644 index 000000000..1bee9cd1a --- /dev/null +++ b/Source/Cmlx/mlx-generated/searchsorted.cpp @@ -0,0 +1,53 @@ +namespace mlx::core::metal { + +const char* searchsorted() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/searchsorted.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/searchsorted.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/searchsorted.h" +// Copyright © 2026 Apple Inc. + +template +METAL_FUNC uint +searchsorted_impl(device const T* a, T v, uint n, int64_t a_stride) { + LessThan lt; + uint lo = 0; + uint hi = n; + while (lo < hi) { + uint mid = lo + (hi - lo) / 2; + T m = a[int64_t(mid) * a_stride]; + bool below = Right ? !lt(v, m) : lt(m, v); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +template +[[kernel]] void searchsorted( + device const T* a [[buffer(0)]], + device const T* v [[buffer(1)]], + device uint* out [[buffer(2)]], + constant const uint& n [[buffer(3)]], + constant const int64_t& a_stride [[buffer(4)]], + uint3 index [[thread_position_in_grid]], + uint3 grid_dim [[threads_per_grid]]) { + auto offset = + index.x + grid_dim.x * (int64_t(index.y) + int64_t(grid_dim.y) * index.z); + out[offset] = searchsorted_impl(a, v[offset], n, a_stride); +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/sort.cpp b/Source/Cmlx/mlx-generated/sort.cpp index f04797cef..fe4516bca 100644 --- a/Source/Cmlx/mlx-generated/sort.cpp +++ b/Source/Cmlx/mlx-generated/sort.cpp @@ -42,14 +42,26 @@ struct Init>> { static constexpr constant T v = metal::numeric_limits::quiet_NaN(); }; +template <> +struct Init { + static constexpr constant complex64_t v = complex64_t( + metal::numeric_limits::quiet_NaN(), + metal::numeric_limits::quiet_NaN()); +}; + template struct LessThan { static constexpr constant T init = Init::v; - METAL_FUNC bool operator()(T a, T b) const { - if constexpr ( - metal::is_floating_point_v || metal::is_same_v) { - bool an = isnan(a); - bool bn = isnan(b); + METAL_FUNC bool operator()(T a, T b) const thread { + if constexpr (metal::is_floating_point_v) { + bool an = metal::isnan(a); + bool bn = metal::isnan(b); + if (an | bn) { + return (!an) & bn; + } + } else if constexpr (metal::is_same_v) { + bool an = metal::isnan(a.real) || metal::isnan(a.imag); + bool bn = metal::isnan(b.real) || metal::isnan(b.imag); if (an | bn) { return (!an) & bn; } @@ -389,8 +401,11 @@ template < using ValT = typename sort_kernel::ValT; using IdxT = typename sort_kernel::IdxT; - auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); - auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); + // Signed offsets: a non-sorted axis may have a negative stride. + auto in_block_idx = + elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim); + auto out_block_idx = + elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim); inp += in_block_idx; out += out_block_idx; @@ -533,7 +548,8 @@ template < BLOCK_THREADS, N_PER_THREAD>; - auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); + // Signed offset: a non-sorted axis may have a negative stride. + auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim); inp += block_idx; out_vals += tid.y * size_sorted_axis; out_idxs += tid.y * size_sorted_axis; diff --git a/Source/Cmlx/mlx-generated/steel_attention.cpp b/Source/Cmlx/mlx-generated/steel_attention.cpp index f102699af..75c2edbdd 100644 --- a/Source/Cmlx/mlx-generated/steel_attention.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention.cpp @@ -72,18 +72,18 @@ struct BlockLoader { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -94,7 +94,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { *((threadgroup ReadVector*)(&dst[i * dst_ld])) = @@ -103,7 +103,7 @@ struct BlockLoader { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -151,7 +151,7 @@ struct BlockLoader { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; @@ -196,18 +196,18 @@ struct BlockLoaderT { const int src_ld_, threadgroup T* dst_, ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(src_ld_), - tile_stride(reduction_dim ? BCOLS : BROWS * src_ld), + tile_stride(reduction_dim ? BCOLS : BROWS* src_ld), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * kDstStrRow + bj * kDstStrCol), src(src_ + bi * src_ld + bj) {} /* Apply operation to threadgroup without bound checking */ template - METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const { + METAL_FUNC void apply_inplace_op(thread const UnaryOp& op) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -219,7 +219,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < BROWS; i += TROWS) { STEEL_PRAGMA_UNROLL @@ -230,7 +230,7 @@ struct BlockLoaderT { } /* Load from device memory into threadgroup memory - with bound checking */ - METAL_FUNC void load_safe(short2 src_tile_dim) const { + METAL_FUNC void load_safe(short2 src_tile_dim) const thread { src_tile_dim = src_tile_dim - short2(bj, bi); // Skip loading if thread has no valid reads @@ -278,7 +278,7 @@ struct BlockLoaderT { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { src += tile_stride; } }; @@ -362,7 +362,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -378,14 +378,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast(x * alpha + (beta * c)); } }; @@ -493,13 +493,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -529,7 +525,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -633,7 +630,7 @@ struct Shape2D { RInt r; CInt c; - Shape2D(RInt r_, CInt c_) : r(r_), c(c_) {} + Shape2D(RInt r_, CInt c_) thread : r(r_), c(c_) {} }; template @@ -866,24 +863,25 @@ struct MMATile { METAL_FUNC MMATile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { + const short j) const thread { return val_frags[i * kTileCols + j]; } - METAL_FUNC mat_type mat_at(const short i, const short j) { + METAL_FUNC mat_type mat_at(const short i, const short j) thread { mat_type val_mat; STEEL_PRAGMA_UNROLL for (short ii = 0; ii < kElemsPerFrag; ++ii) { @@ -892,16 +890,16 @@ struct MMATile { return val_mat; } - METAL_FUNC thread elem_type* elems() { + METAL_FUNC thread elem_type* elems() thread { return reinterpret_cast(val_frags); } - METAL_FUNC const thread elem_type* elems() const { + METAL_FUNC const thread elem_type* elems() const thread { return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const { + METAL_FUNC void row_reduce(thread T vals[kRowsPerThread]) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -913,7 +911,7 @@ struct MMATile { } template - METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) { + METAL_FUNC void row_bin_op(thread T vals[kRowsPerThread]) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -925,7 +923,7 @@ struct MMATile { } template - METAL_FUNC void load(const threadgroup U* src) { + METAL_FUNC void load(const threadgroup U* src) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -942,7 +940,7 @@ struct MMATile { } template - METAL_FUNC void store(threadgroup U* dst) const { + METAL_FUNC void store(threadgroup U* dst) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -959,7 +957,7 @@ struct MMATile { } template - METAL_FUNC void load(const device U* src, const int ld) { + METAL_FUNC void load(const device U* src, const int ld) thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -974,7 +972,7 @@ struct MMATile { } template - METAL_FUNC void store(device U* dst, const int ld) const { + METAL_FUNC void store(device U* dst, const int ld) const thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -989,8 +987,10 @@ struct MMATile { } template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -1009,8 +1009,10 @@ struct MMATile { } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { STEEL_PRAGMA_UNROLL for (int i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL @@ -1121,7 +1123,7 @@ struct BlockMMA { /* Constructor */ METAL_FUNC BlockMMA( ushort simd_group_id [[simdgroup_index_in_threadgroup]], - ushort simd_lane_id [[thread_index_in_simdgroup]]) { + ushort simd_lane_id [[thread_index_in_simdgroup]]) thread { // Determine thread position in simdgroup matrix short tm = kFragSize * (simd_group_id / WN); short tn = kFragSize * (simd_group_id % WN); @@ -1139,7 +1141,7 @@ struct BlockMMA { } /* (BM, BK) X (BK, BN) multiply accumulate function */ - METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) thread { // Adjust for simdgroup and thread location As += As_offset; Bs += Bs_offset; @@ -1166,7 +1168,7 @@ struct BlockMMA { } /* Store results from simdgroup_matrix results into device memory */ - METAL_FUNC void store_result(device U* D, const int ldd) { + METAL_FUNC void store_result(device U* D, const int ldd) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1180,7 +1182,7 @@ struct BlockMMA { } METAL_FUNC void - store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) thread { // Apply epilogue STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1199,7 +1201,8 @@ struct BlockMMA { /* Apply epilogue */ template - METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + METAL_FUNC void apply_epilogue( + thread const UnaryEpilogue& epilogue_op) thread { // Loop over all simdgroup tiles STEEL_PRAGMA_UNROLL for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { @@ -1213,7 +1216,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; @@ -1242,7 +1245,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const BinaryEpilogue& epilogue_op) { + thread const BinaryEpilogue& epilogue_op) thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; dst_tile_dims -= short2(sn, sm); @@ -1287,7 +1290,7 @@ struct BlockMMA { const device U* C, const int ldc, const int fdc, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -1320,7 +1323,7 @@ struct BlockMMA { const int ldc, const int fdc, short2 dst_tile_dims, - thread const Epilogue& epilogue_op) const { + thread const Epilogue& epilogue_op) const thread { // Adjust for simdgroup and thread location C += (sm)*ldc + (sn)*fdc; D += (sm)*ldd + sn; @@ -2012,11 +2015,16 @@ template < } int kb_lim = params->NK; + int kb_min_causal = params->NK; if (do_causal) { int q_max = (tid.x + 1) * BQ + params->qL_off; kb_lim = (q_max + BK - 1) / BK; kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); } // Loop over KV seq length @@ -2076,7 +2084,7 @@ template < } // Mask out if causal - if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) { + if (do_causal && kb >= kb_min_causal) { using stile_t = decltype(Stile); using selem_t = typename stile_t::elem_type; constexpr auto neg_inf = Limits::finite_min; @@ -2201,7 +2209,7 @@ template < for (short id = 0; id < TD; id++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } @@ -2211,7 +2219,7 @@ template < Vtile.template load( &Vs[Vs_offset + kk * LDV_tgp + dd]); - if constexpr (BD == 128) { + if constexpr (BD >= 128) { simdgroup_barrier(mem_flags::mem_none); } diff --git a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp index e40632ec9..4ca7573a8 100644 --- a/Source/Cmlx/mlx-generated/steel_attention_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention_nax.cpp @@ -18,12 +18,6 @@ const char* steel_attention_nax() { #define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") #define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") -/////////////////////////////////////////////////////////////////////////////// -// Contents from "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.3.525.0.RrIjD2/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.850/include/metal/__exec/units.h" -/////////////////////////////////////////////////////////////////////////////// - -#line 1 "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.3.525.0.RrIjD2/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.850/include/metal/__exec/units.h" - /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/utils/type_traits.h" /////////////////////////////////////////////////////////////////////////////// @@ -109,13 +103,9 @@ struct integral_constant { using value_type = T; using type = integral_constant; - METAL_FUNC constexpr operator value_type() const noexcept { + METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -145,7 +135,8 @@ using Int = integral_constant; METAL_FUNC constexpr auto __operator__( \ integral_constant, integral_constant) { \ constexpr auto res = tv __op__ uv; \ - return integral_constant{}; \ + using res_t = metal::remove_addrspace_t; \ + return integral_constant{}; \ } integral_const_binop(+, operator+); @@ -222,6 +213,72 @@ METAL_FUNC constexpr auto sum(T x, Us... us) { #pragma METAL internals : disable +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsBase.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsBase.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.6.42.0.n257DK/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.883/include/metal/__exec/units.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/private/var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain-v17.6.42.0.n257DK/Metal.xctoolchain/usr/metal/32023/lib/clang/32023.883/include/metal/__exec/units.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTraits.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTraits.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTypes.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsTypes.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsUtility.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsUtility.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsConvolution2dImpl.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsConvolution2dImpl.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsConvolution2d.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsConvolution2d.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsAvailability.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsAvailability.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsMatMul2dImpl.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/__impl/MPPTensorOpsMatMul2dImpl.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsMatMul2d.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MPPTensorOpsMatMul2d.h" + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MetalPerformancePrimitives.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "/Applications/XcodeWonderF.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX26.5.sdk/System/Library/Frameworks/MetalPerformancePrimitives.framework/Headers/MetalPerformancePrimitives.h" + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/steel/attn/nax.h" /////////////////////////////////////////////////////////////////////////////// @@ -298,11 +355,13 @@ struct BaseNAXFrag { StrY str_y, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -335,13 +394,16 @@ struct BaseNAXFrag { LimX lim_x, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -356,7 +418,10 @@ struct BaseNAXFrag { } } else { - dst = dtype_frag_t(0); + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = T(0); + } } } } @@ -379,14 +444,18 @@ struct BaseNAXFrag { LimY lim_y, OffX off_x = {}, OffY off_y = {}) { - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + src += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if ((r < lx) && ((c + j) < ly)) { dst[i * kElemCols + j] = static_cast(src[r * str_x + (c + j) * str_y]); } else { @@ -412,11 +481,13 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL @@ -451,13 +522,16 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; - if (r < lim_x) { + if (r < lx) { if constexpr (metal::is_same_v>) { STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { @@ -494,15 +568,19 @@ struct BaseNAXFrag { OffY off_y = {}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); + dst += sc.y * str_x + sc.x * str_y; + auto lx = lim_x - sc.y; + auto ly = lim_y - sc.x; + STEEL_PRAGMA_UNROLL for (short i = 0; i < kElemRows; i++) { - const auto r = off_x + i * kElemRowsJump + sc.y; - const auto c = off_y + sc.x; + const auto r = off_x + i * kElemRowsJump; + const auto c = off_y; STEEL_PRAGMA_UNROLL for (short j = 0; j < kElemCols; j++) { - if (r < lim_x && (c + j) < lim_y) { + if (r < lx && (c + j) < ly) { dst[r * str_x + (c + j) * str_y] = static_cast(src[i * kElemCols + j]); } @@ -534,7 +612,7 @@ struct BaseNAXFrag { OffY off_y = Int<0>{}) { using U = pointer_element_t; - const short2 sc = short2{0, 0}; // get_coord(); + const short2 sc = get_coord(); const_for_loop<0, kElemRows, 1>([&](auto idx_row) { const auto r = off_x + idx_row * Int{}; @@ -588,81 +666,272 @@ struct BaseNAXFrag { } } } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cn0, + thread dtype_frag_t& Cn1, + const thread dtype_frag_t& A, + metal::bool_constant, + const thread dtype_frag_t& Bn0, + const thread dtype_frag_t& Bn1, + metal::bool_constant) { + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = A[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = Bn0[i]; + ct_b[kElemsPerFrag + i] = Bn1[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cn0[i]; + ct_c[kElemsPerFrag + i] = Cn1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cn0[i] = ct_c[i]; + Cn1[i] = ct_c[kElemsPerFrag + i]; + } + } + + template < + typename CType, + typename AType, + typename BType, + bool transpose_a = false, + bool transpose_b = false> + METAL_FUNC static constexpr void mma( + thread dtype_frag_t& Cm0, + thread dtype_frag_t& Cm1, + const thread dtype_frag_t& Am0, + const thread dtype_frag_t& Am1, + metal::bool_constant, + const thread dtype_frag_t& B, + metal::bool_constant) { + // Create Matmul descriptor + constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( + 16, + 32, + 16, + transpose_a, + transpose_b, + true, + mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); + + // Create matmul op + mpp::tensor_ops::matmul2d gemm_op; + + // Create matmul operands in registers + auto ct_a = + gemm_op + .template get_left_input_cooperative_tensor(); + auto ct_b = + gemm_op + .template get_right_input_cooperative_tensor(); + + // Create matmul output in register + auto ct_c = gemm_op.template get_destination_cooperative_tensor< + metal::remove_addrspace_t, + metal::remove_addrspace_t, + CType>(); + + // Load A in to left operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[i]; + } + + // Load B into right operand registers + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_b[i] = B[i]; + } + + // Load C into output registers (op handles accumulation) + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + ct_c[i] = Cm0[i]; + ct_c[kElemsPerFrag + i] = Cm1[i]; + } + + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); + + // Copy out results + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + i]; + } + } }; template < typename T, - short kRows_, - short kCols_, - typename NAXFrag_ = BaseNAXFrag> -struct NAXSubTile { + short kTileRows_, + short kTileCols_, + class NAXFrag_ = BaseNAXFrag> +struct NAXTile { using NAXFrag_t = NAXFrag_; - STEEL_CONST short kRows = kRows_; - STEEL_CONST short kCols = kCols_; + using elem_type = T; STEEL_CONST short kFragRows = NAXFrag_t::kFragRows; STEEL_CONST short kFragCols = NAXFrag_t::kFragCols; STEEL_CONST short kElemsPerFrag = NAXFrag_t::kElemsPerFrag; - STEEL_CONST short kSubTileRows = kRows / kFragRows; - STEEL_CONST short kSubTileCols = kCols / kFragCols; + STEEL_CONST short kTileRows = kTileRows_; + STEEL_CONST short kTileCols = kTileCols_; - STEEL_CONST short kNumFrags = kSubTileRows * kSubTileCols; - STEEL_CONST short kElemsPerSubTile = kNumFrags * kElemsPerFrag; + STEEL_CONST short kRows = kTileRows * kFragRows; + STEEL_CONST short kCols = kTileCols * kFragCols; - STEEL_CONST int kRowsPerThread = kSubTileRows * NAXFrag_t::kElemRows; - STEEL_CONST int kColsPerThread = kSubTileCols * NAXFrag_t::kElemCols; + STEEL_CONST short kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - using frag_type = typename NAXFrag_t::template dtype_frag_t; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; + + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - frag_type val_frags[kNumFrags]; + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; + + METAL_FUNC NAXTile() thread {} - METAL_FUNC constexpr void clear() { + METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL for (short i = 0; i < kNumFrags; ++i) { val_frags[i] = frag_type(0); } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) + thread { + return val_frags[i * kTileCols + j]; } METAL_FUNC constexpr const thread frag_type& frag_at( const short i, - const short j) const { - return val_frags[i * kSubTileCols + j]; + const short j) const thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr thread frag_type& frag_at() { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread frag_type& frag_at() const { - return val_frags[i * kSubTileCols + j]; + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + return val_frags[i * kTileCols + j]; } - METAL_FUNC thread T* elems() { - return reinterpret_cast(val_frags); + template + METAL_FUNC constexpr thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j, + metal::bool_constant) const thread { + if constexpr (transpose) { + return frag_at(j, i); + } else { + return frag_at(i, j); + } + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + template + METAL_FUNC constexpr const thread frag_type& frag_at() const thread { + if constexpr (transpose) { + return frag_at(); + } else { + return frag_at(); + } + } + + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); } - METAL_FUNC const thread T* elems() const { - return reinterpret_cast(val_frags); + METAL_FUNC const thread elem_type* elems() const thread { + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - thread T* vptr = (thread T*)(&vals); + METAL_FUNC void row_reduce( + thread metal::vec& vals) const thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { + for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + for (short j = 0; j < kTileCols; ++j) { NAXFrag_t::template row_reduce( frag_at(i, j), &vptr[i * kFragThrRows]); } @@ -670,545 +939,151 @@ struct NAXSubTile { } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - thread T* vptr = (thread T*)(&vals); + METAL_FUNC void row_bin_op( + thread metal::vec& vals) thread { + auto vptr = (thread T*)(&vals); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { + for (short i = 0; i < kTileRows; ++i) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + for (short j = 0; j < kTileCols; ++j) { NAXFrag_t::template row_bin_op( frag_at(i, j), &vptr[i * kFragThrRows]); } } } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load( - SrcPtrType src, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + template + METAL_FUNC void load(const threadgroup U* src) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load( - frag_at(i, j), + frag_at(), src, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store( - DstPtrType dst, - StrX str_x, - StrY str_y, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kSubTileCols; ++j) { + template + METAL_FUNC void store(threadgroup U* dst) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::store( - frag_at(i, j), + frag_at(), dst, - str_x, - str_y, - off_x + i * kFragRows, - off_y + j * kFragCols); - } - } + Int{}, + Int{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_rows( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { + template + METAL_FUNC void load(const device U* src, const int ld) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::load( + frag_at(), + src, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void store(device U* dst, const int ld) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store( + frag_at(), + dst, + ld, + Int<1>{}, + idx_row * Int{}, + idx_col * Int{}); + }); + }); + } + + template + METAL_FUNC void + load_rows(const device U* src, const int ld, const short n_rows) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_rows( - frag_at(i, j), + frag_at(), src, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } - template < - typename SrcPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void load_safe( - SrcPtrType src, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { + template + METAL_FUNC void load_safe( + const device U* src, + const int ld, + const short2 src_tile_dims) thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { NAXFrag_t::load_safe( - frag_at(i, j), + frag_at(), src, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } + ld, + Int<1>{}, + src_tile_dims.y, + src_tile_dims.x, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_rows( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename LimX, - typename LimY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_safe( - DstPtrType dst, - StrX str_x, - StrY str_y, - LimX lim_x, - LimY lim_y, - OffX off_x = {}, - OffY off_y = {}) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { - NAXFrag_t::store_safe( - frag_at(i, j), - dst, - str_x, - str_y, - lim_x, - lim_y, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } - } - - template < - typename DstPtrType, - typename StrX, - typename StrY, - typename StartX, - typename StopX, - typename StartY, - typename StopY, - typename OffX = Int<0>, - typename OffY = Int<0>> - METAL_FUNC constexpr void store_slice( - DstPtrType dst, - StrX str_x, - StrY str_y, - StartX start_x, - StopX stop_x, - StartY start_y, - StopY stop_y, - OffX off_x = Int<0>{}, - OffY off_y = Int<0>{}) const { - const_for_loop<0, kSubTileRows, 1>([&](auto idx_row) { - const_for_loop<0, kSubTileCols, 1>([&](auto idx_col) { - NAXFrag_t::store_slice( + template + METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) + const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_rows( frag_at(), dst, - str_x, - str_y, - start_x, - stop_x, - start_y, - stop_y, - off_x + idx_row * Int{}, - off_y + idx_col * Int{}); - }); - }); - } -}; - -template < - short RC, - short CC, - short RA, - short CA, - short RB, - short CB, - typename CType, - typename AType, - typename BType, - bool transpose_a, - bool transpose_b, - typename NAXFrag_t = BaseNAXFrag> -METAL_FUNC void subtile_matmad_nax( - thread NAXSubTile& C, - thread NAXSubTile& A, - metal::bool_constant, - thread NAXSubTile& B, - metal::bool_constant) { - // Static checks - constexpr short FMa = transpose_a ? CA : RA; - constexpr short FMc = RC; - static_assert(FMa == FMc, "NAX matmul: M dimensions do not match"); - - constexpr short FNb = transpose_b ? RB : CB; - constexpr short FNc = CC; - static_assert(FNb == FNc, "NAX matmul: N dimensions do not match"); - - constexpr short FKa = transpose_a ? RA : CA; - constexpr short FKb = transpose_b ? CB : RB; - static_assert(FKa == FKb, "NAX matmul: N dimensions do not match"); - - constexpr short FM = FMc; - constexpr short FN = FNc; - constexpr short FK = FKa; - - constexpr int TM = FM / 16; - constexpr int TN = FN / 16; - constexpr int TK = FK / 16; - - constexpr auto desc = mpp::tensor_ops::matmul2d_descriptor( - FM, - FN, - FK, - transpose_a, - transpose_b, - true, - mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate); - - mpp::tensor_ops::matmul2d gemm_op; - - auto ct_a = - gemm_op.template get_left_input_cooperative_tensor(); - auto ct_b = - gemm_op - .template get_right_input_cooperative_tensor(); - auto ct_c = gemm_op.template get_destination_cooperative_tensor< - decltype(ct_a), - decltype(ct_b), - CType>(); - - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_a ? kk : mm; - const short fj = transpose_a ? mm : kk; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_a[(TK * mm + kk) * 8 + i] = A.frag_at(fi, fj)[i]; - } - } - } - - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - STEEL_PRAGMA_UNROLL - for (short kk = 0; kk < TK; kk++) { - const short fi = transpose_b ? nn : kk; - const short fj = transpose_b ? kk : nn; - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < 8; i++) { - ct_b[(TN * kk + nn) * 8 + i] = B.frag_at(fi, fj)[i]; - } - } - } - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - ct_c[i] = C.elems()[i]; - } - - gemm_op.run(ct_a, ct_b, ct_c); - - STEEL_PRAGMA_UNROLL - for (short i = 0; i < ct_c.get_capacity(); i++) { - C.elems()[i] = ct_c[i]; - } -} - -template -struct NAXTile { - using NAXSubTile_t = NAXSubTile_; - using elem_type = T; - STEEL_CONST short kSubTileRows = NAXSubTile_t::kRows; - STEEL_CONST short kSubTileCols = NAXSubTile_t::kCols; - STEEL_CONST short kElemsPerSubTile = NAXSubTile_t::kElemsPerSubTile; - - STEEL_CONST short kTileRows = kTileRows_; - STEEL_CONST short kTileCols = kTileCols_; - - STEEL_CONST short kRows = kTileRows * kSubTileRows; - STEEL_CONST short kCols = kTileCols * kSubTileCols; - - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; - - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; - - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; - - NAXSubTile_t val_subtiles[kSubTiles]; - - METAL_FUNC NAXTile() thread {} - - METAL_FUNC constexpr void clear() { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kSubTiles; ++i) { - val_subtiles[i].clear(); - } - } - - METAL_FUNC constexpr thread NAXSubTile_t& subtile_at( - const short i, - const short j) { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( - const short i, - const short j) const { - return val_subtiles[i * kTileCols + j]; - } - - template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC thread elem_type* elems() { - return reinterpret_cast(val_subtiles[0].elems()); - } - - METAL_FUNC const thread elem_type* elems() const { - return reinterpret_cast(val_subtiles[0].elems()); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const { - auto sub_rows = (thread metal::vec*)(&vals); - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_reduce(sub_rows[i]); - } - } - } - - template - METAL_FUNC void row_bin_op(thread metal::vec& vals) { - auto sub_rows = (thread metal::vec*)(&vals); - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).template row_bin_op(sub_rows[i]); - } - } - } - - template - METAL_FUNC void load(const threadgroup U* src) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( - src, - Int{}, - Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } - } - - template - METAL_FUNC void store(threadgroup U* dst) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( - dst, - Int{}, - Int{}, - i * kSubTileRows, - j * kSubTileCols); - } - } - } - - template - METAL_FUNC void load(const device U* src, const int ld) { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{}); - } - } - } - - template - METAL_FUNC void store(device U* dst, const int ld) const { - STEEL_PRAGMA_UNROLL - for (short i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (short j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{}); - } - } - } - - template - METAL_FUNC void - load_safe(const device U* src, const int ld, const short2 src_tile_dims) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_safe( - src, - ld, - Int<1>{}, - src_tile_dims.y, - src_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } - } - - template - METAL_FUNC void - load_rows(const device U* src, const int ld, const short n_rows) { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).load_rows( - &src[(i * kSubTileRows) * ld + (j * kSubTileCols)], ld, Int<1>{}, - n_rows - i * kSubTileRows); - } - } + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template - METAL_FUNC void - store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_safe( + METAL_FUNC void store_safe( + device U* dst, + const int ld, + const short2 dst_tile_dims) const thread { + const_for_loop<0, kTileRows, 1>([&](auto idx_row) { + const_for_loop<0, kTileCols, 1>([&](auto idx_col) { + NAXFrag_t::store_safe( + frag_at(), dst, ld, Int<1>{}, dst_tile_dims.y, dst_tile_dims.x, - i * kSubTileRows, - j * kSubTileCols); - } - } - } - - template - METAL_FUNC void store_rows(device U* dst, const int ld, const short n_rows) - const { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_rows( - &dst[(i * kSubTileRows) * ld + (j * kSubTileCols)], - ld, - Int<1>{}, - n_rows - i * kSubTileRows); - } - } + idx_row * Int{}, + idx_col * Int{}); + }); + }); } template @@ -1216,10 +1091,11 @@ struct NAXTile { device U* dst, const int ld, const short2 start, - const short2 stop) const { + const short2 stop) const thread { const_for_loop<0, kTileRows, 1>([&](auto idx_row) { const_for_loop<0, kTileCols, 1>([&](auto idx_col) { - subtile_at().store_slice( + NAXFrag_t::store_slice( + frag_at(), dst, ld, Int<1>{}, @@ -1227,8 +1103,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1248,51 +1124,54 @@ METAL_FUNC void tile_matmad_nax( metal::bool_constant) { // Static checks constexpr short TMa = transpose_a ? ATile::kTileCols : ATile::kTileRows; - constexpr short TMc = CTile::kTileRows; - static_assert(TMa == TMc, "NAX tile matmul: M dimensions do not match"); - - constexpr short FMa = transpose_a ? ATile::kSubTileCols : ATile::kSubTileRows; - constexpr short FMc = CTile::kSubTileRows; - static_assert(FMa == FMc, "NAX subtile matmul: M dimensions do not match"); + constexpr short TM = CTile::kTileRows; + static_assert(TMa == TM, "MXU tile matmul: M dimensions do not match"); constexpr short TNb = transpose_b ? BTile::kTileRows : BTile::kTileCols; - constexpr short TNc = CTile::kTileCols; - static_assert(TNb == TNc, "NAX tile matmul: N dimensions do not match"); - - constexpr short FNb = transpose_b ? BTile::kSubTileRows : BTile::kSubTileCols; - constexpr short FNc = CTile::kSubTileCols; - static_assert(FNb == FNc, "NAX subtile matmul: N dimensions do not match"); + constexpr short TN = CTile::kTileCols; + static_assert(TNb == TN, "MXU tile matmul: N dimensions do not match"); constexpr short TKa = transpose_a ? ATile::kTileRows : ATile::kTileCols; - constexpr short TKb = transpose_b ? BTile::kTileCols : BTile::kTileRows; - static_assert(TKa == TKb, "NAX tile matmul: K dimensions do not match"); + constexpr short TK = transpose_b ? BTile::kTileCols : BTile::kTileRows; + static_assert(TKa == TK, "MXU tile matmul: K dimensions do not match"); - constexpr short FKa = transpose_a ? ATile::kSubTileRows : ATile::kSubTileCols; - constexpr short FKb = transpose_b ? BTile::kSubTileCols : BTile::kSubTileRows; - static_assert(FKa == FKb, "NAX subtile matmul: K dimensions do not match"); + constexpr auto ta = metal::bool_constant{}; + constexpr auto tb = metal::bool_constant{}; - constexpr short TM = TMc; - constexpr short TN = TNc; - constexpr short TK = TKa; - - // Do matmul here - STEEL_PRAGMA_UNROLL - for (short i = 0; i < TM; ++i) { + if constexpr (TN == 1 && TM % 2 == 0) { STEEL_PRAGMA_UNROLL - for (short j = 0; j < TN; ++j) { + for (short mm = 0; mm < TM; mm += 2) { STEEL_PRAGMA_UNROLL - for (short k = 0; k < TK; ++k) { - const short ra = transpose_a ? k : i; - const short ca = transpose_a ? i : k; - const short rb = transpose_b ? j : k; - const short cb = transpose_b ? k : j; - - subtile_matmad_nax( - C.subtile_at(i, j), - A.subtile_at(ra, ca), - metal::bool_constant{}, - B.subtile_at(rb, cb), - metal::bool_constant{}); + for (short nn = 0; nn < TN; ++nn) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm + 1, nn), + A.frag_at(mm, kk, ta), + A.frag_at(mm + 1, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + metal::bool_constant{}); + } + } + } + } else if constexpr (TN % 2 == 0) { + STEEL_PRAGMA_UNROLL + for (short mm = 0; mm < TM; ++mm) { + STEEL_PRAGMA_UNROLL + for (short nn = 0; nn < TN; nn += 2) { + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < TK; ++kk) { + CTile::NAXFrag_t::mma( + C.frag_at(mm, nn), + C.frag_at(mm, nn + 1), + A.frag_at(mm, kk, ta), + metal::bool_constant{}, + B.frag_at(kk, nn, tb), + B.frag_at(kk, nn + 1, tb), + metal::bool_constant{}); + } } } } @@ -1426,7 +1305,7 @@ struct TransformNone { template struct TransformAdd { - TransformAdd(const float, const float) {} + TransformAdd(const float, const float) thread {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); @@ -1442,14 +1321,14 @@ struct TransformAxpby { const float alpha; const float beta; - TransformAxpby(const float alpha_, const float beta_) - : alpha(alpha_), beta(beta_) {} + TransformAxpby(const float alpha_, const float beta_) thread : alpha(alpha_), + beta(beta_) {} static METAL_FUNC OutT apply(InT x) { return static_cast(x); } - METAL_FUNC OutT apply(InT x, OutT c) const { + METAL_FUNC OutT apply(InT x, OutT c) const thread { return static_cast(x * alpha + (beta * c)); } }; @@ -1496,9 +1375,9 @@ constant bool has_sinks [[function_constant(302)]]; template struct TransformScale { T scale; - METAL_FUNC TransformScale(T scale_) : scale(scale_) {} + METAL_FUNC TransformScale(T scale_) thread : scale(scale_) {} - METAL_FUNC T apply(T x) const { + METAL_FUNC T apply(T x) const thread { return scale * x; } }; @@ -1600,38 +1479,36 @@ template < make_uniform(params->scale) * make_uniform(1.44269504089f); // Prepare MMA tiles - constexpr short UQ = 16; - constexpr short UD = 32; + constexpr short kU = 16; constexpr int kNWarps = WM * WN; static_assert( - BQ >= (kNWarps * UQ) && BQ % (kNWarps * UQ) == 0, + BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); // Q seq frags per warp - constexpr int TQ = BQ / (kNWarps * UQ); + constexpr int TQ = BQ / (kNWarps * kU); // HeadDim frags (all warps load the same frags) - constexpr int TD = BD / UD; + constexpr int TD = BD / kU; + // KV seq frags per warp + constexpr short TK = BK / kU; static_assert(TQ == 1, "Check TQ"); - - using OSubTile = NAXSubTile; - NAXTile Otile; + using otile_t = NAXTile; + otile_t Otile; Otile.clear(); // Prepare mma tile offsets - const short2 simd_coord = OSubTile::NAXFrag_t::get_coord(); + const short tm = kU * TQ * simd_group_id; + Q += tm * int(params->Q_strides[2]); + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); const short sm = simd_coord.y; const short sn = simd_coord.x; - const short tm = UQ * TQ * simd_group_id; - - Q += (tm + sm) * int(params->Q_strides[2]) + sn; - K += sm * int(params->K_strides[2]) + sn; - V += sm * int(params->V_strides[2]) + sn; // Init row reduction variables - constexpr short kRowsPT = decltype(Otile)::kRowsPerThread; + constexpr short kRowsPT = otile_t::kRowsPerThread; metal::vec max_score; metal::vec sum_score{0}; @@ -1651,83 +1528,76 @@ template < } int kb_lim = params->NK; + int kb_min_causal = params->NK; if (do_causal) { int q_max = (tid.x + 1) * BQ + params->qL_off; kb_lim = (q_max + BK - 1) / BK; kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); } const bool is_last_bq = int(tid.x) == (params->NQ_aligned); // const bool is_last_tq = int(simd_group_id) >= (params->qL_rem / UQ); const bool is_last_q = is_last_bq; - const short lim_rows_q = params->qL_rem - (tm + sm); - const short lim_rows_k = params->kL_rem - sm; + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; // Loop over KV seq length for (int kb = 0; kb < kb_lim; kb++) { const int is_last_k = (kb == (params->NK_aligned)); // Do S = Q @ K.T - constexpr short UDs = 16; - constexpr short UKs = 32; - - constexpr short TDs = BD / UDs; - constexpr short TKs = BK / UKs; - - using SSubTile = NAXSubTile; - using QSubTile = NAXSubTile; - using KSubTile = NAXSubTile; - - NAXTile Stile; + using stile_t = NAXTile; + stile_t Stile; Stile.clear(); STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short ik = 0; ik < TKs; ik++) { - STEEL_PRAGMA_UNROLL - for (short id = 0; id < TDs; id++) { - NAXTile Qtile; - NAXTile Ktile; - - const int Q_load_off = iq * UQ * int(params->Q_strides[2]) + id * UDs; - const int K_load_off = - ik * UKs * int(params->K_strides[2]) + id * UDs; + for (short ik = 0; ik < TK; ik += 2) { + // Unrolling the head-dim loop by 4, rather than fully, potentially + // lets the compiler interleave the next K-tile loads with the running + // mma chain instead of hoisting all TD loads up front and is faster + // for head dim 128. +#pragma clang loop unroll_count(4) + for (short id = 0; id < TD; id++) { + NAXTile Qtile; + NAXTile Ktile; + + const int Q_load_off = iq * kU * int(params->Q_strides[2]) + id * kU; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; if (!align_Q && is_last_q) { - // Qtile.load_rows( - // Q + Q_load_off, - // int(params->Q_strides[2]), - // lim_rows_q - iq * UQ); - Qtile.load_safe( + Qtile.load_rows( Q + Q_load_off, int(params->Q_strides[2]), - short2(BD, lim_rows_q - iq * UQ)); + lim_rows_q - iq * kU); } else { Qtile.load(Q + Q_load_off, int(params->Q_strides[2])); } if (!align_K && is_last_k) { - // Ktile.load_rows( - // K + K_load_off, - // int(params->K_strides[2]), - // lim_rows_k - ik * UKs); - Ktile.load_safe( + Ktile.load_rows( K + K_load_off, int(params->K_strides[2]), - short2(BD, lim_rows_k - ik * UKs)); + lim_rows_k - ik * kU); } else { Ktile.load(K + K_load_off, int(params->K_strides[2])); } - subtile_matmad_nax( - Stile.subtile_at(iq, ik), - Qtile.subtile_at(0, 0), + stile_t::NAXFrag_t::mma( + Stile.frag_at(iq, ik), + Stile.frag_at(iq, ik + 1), + Qtile.frag_at(0, 0), metal::false_type{}, - Ktile.subtile_at(0, 0), + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), metal::true_type{}); } } @@ -1735,22 +1605,10 @@ template < // Scale S STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { Stile.elems()[ii] *= float(scale2); } - // Scale and Retile S - constexpr short UK = 16; - constexpr short TK = BK / UK; - using PSubTile = NAXSubTile; - - NAXTile Ptile; - - STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { - Ptile.elems()[ii] = Stile.elems()[ii]; - } - // Mask out length sequence if (!align_K && is_last_k) { constexpr auto neg_inf = Limits::finite_min; @@ -1759,16 +1617,16 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - const short col_pos = sn + ik * UK; + const short col_pos = ik * kU + sn; - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) { + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) { - const auto loc = ii * PSubTile::kFragThrCols + jj; - fg[loc] = ((col_pos + jj) >= params->kL_rem) ? neg_inf : fg[loc]; + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; } } } @@ -1776,7 +1634,7 @@ template < } // Mask out if causal - if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) { + if (do_causal && kb >= kb_min_causal) { constexpr auto neg_inf = Limits::finite_min; const int base_row = tid.x * BQ + params->qL_off + tm; @@ -1786,18 +1644,16 @@ template < for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - const short row_pos = base_row + iq * UQ; - const short col_pos = base_col + ik * UK; - - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); STEEL_PRAGMA_UNROLL - for (short ii = 0; ii < PSubTile::kFragThrRows; ii++) { + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < PSubTile::kFragThrCols; jj++) { - const auto r = row_pos + ii * PSubTile::kFragRowsJump + sm; - const auto c = col_pos + jj + sn; - const auto loc = ii * PSubTile::kFragThrCols + jj; + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = + base_row + iq * kU + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; fg[loc] = (r < c) ? neg_inf : fg[loc]; } } @@ -1814,33 +1670,65 @@ template < constexpr bool is_bool = is_same_v; using melem_t = typename metal::conditional_t; - using MSubTile = NAXSubTile; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; - STEEL_PRAGMA_UNROLL - for (short iq = 0; iq < TQ; iq++) { - STEEL_PRAGMA_UNROLL - for (short ik = 0; ik < TK; ik++) { - const short row_pos = base_row + iq * UQ + sm; - const short col_pos = base_col + ik * UK + sn; + if (base_row + BQ <= params->qL && base_col + BK <= params->kL) { + for (short iq = 0; iq < TQ; iq++) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row + iq * kU; + const int col_pos = base_col + ik * kU; - MSubTile mfrag; - mfrag.load_safe( - mask, - int64_t(mask_params->M_strides[2]), - Int<1>{}, - params->qL, - params->kL, - row_pos, - col_pos); + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); - thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0); + thread auto& fg = Stile.frag_at(iq, ik); + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short jj = 0; jj < MSubTile::kElemsPerFrag; jj++) { - if constexpr (is_bool) { - fg[jj] = mfrag.elems()[jj] ? fg[jj] : neg_inf; - } else { - fg[jj] += M_LOG2E_F * AccumType(mfrag.elems()[jj]); + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row + iq * kU; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(iq, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } } } } @@ -1858,10 +1746,10 @@ template < } // Row max - Ptile.template row_reduce(new_max); + Stile.template row_reduce(new_max); // exp(Si - rowmax(Si)) - Ptile.template row_bin_op(new_max); + Stile.template row_bin_op(new_max); // Factor exp(rowmax(Si) - rowmax(Si-1)) STEEL_PRAGMA_UNROLL @@ -1876,7 +1764,7 @@ template < sum_score[i] = sum_score[i] * factor[i]; } - Ptile.template row_reduce(sum_score); + Stile.template row_reduce(sum_score); // Update O Otile.template row_bin_op(factor); @@ -1887,39 +1775,36 @@ template < STEEL_PRAGMA_UNROLL for (short iq = 0; iq < TQ; iq++) { STEEL_PRAGMA_UNROLL - for (short id = 0; id < TD; id++) { + for (short id = 0; id < TD; id += 2) { if constexpr (BD == 128) { - if (id == 2) { + if (id == 4) { threadgroup_barrier(mem_flags::mem_none); } } STEEL_PRAGMA_UNROLL for (short ik = 0; ik < TK; ik++) { - using VSubTile = NAXSubTile; - NAXTile Vtile; + NAXTile Vtile; - const int V_load_off = ik * UK * int(params->V_strides[2]) + id * UD; + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; if (!align_K && is_last_k) { - // Vtile.load_rows( - // V + V_load_off, - // int(params->V_strides[2]), - // lim_rows_k - ik * UK); - Vtile.load_safe( + Vtile.load_rows( V + V_load_off, int(params->V_strides[2]), - short2(BD, lim_rows_k - ik * UK)); + lim_rows_k - ik * kU); } else { Vtile.load(V + V_load_off, int(params->V_strides[2])); } - subtile_matmad_nax( - Otile.subtile_at(iq, id), - Ptile.subtile_at(iq, ik), - metal::bool_constant{}, - Vtile.subtile_at(0, 0), - metal::bool_constant{}); + otile_t::NAXFrag_t::mma( + Otile.frag_at(iq, id), + Otile.frag_at(iq, id + 1), + Stile.frag_at(iq, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); } } } @@ -1942,14 +1827,434 @@ template < Otile.template row_bin_op(rcp); // Store results - O += (tm + sm) * int(params->O_strides[2]) + sn; + O += tm * int(params->O_strides[2]); if (!align_Q && is_last_q) { if (lim_rows_q <= 0) return; - // Otile.store_rows(O, params->O_strides[2], lim_rows_q); - Otile.store_safe(O, params->O_strides[2], short2(BD, lim_rows_q)); + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); + } else { + Otile.store(O, int(params->O_strides[2])); + } +} + +/////////////////////////////////////////////////////////////////////////////// +// Head-dim split attention kernel +/////////////////////////////////////////////////////////////////////////////// + +// Variant of attention_nax for wide heads (bd = 256). There, the per-simdgroup +// accumulator working set of attention_nax (TD output fragments plus the S +// fragments) is what gates tensor-unit throughput, so this kernel splits the +// head dim across the WN = 2 simdgroups of the second warp dimension: each +// simdgroup of a pair owns one half of D for Q@K.T and one half of Dv for P@V, +// halving its accumulator set. The pair exchanges its partial Q@K.T sums +// through threadgroup memory, then both simdgroups run softmax redundantly on +// the full S tile (the row statistics are cheap) and each accumulates P@V for +// its own half of Dv. + +// clang-format off +template < + typename T, + int BQ, + int BK, + int BD, + int WM, + int WN, + typename MaskType = float, + typename AccumType = float> +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void attention_nax_dsplit( + const device T* Q [[buffer(0)]], + const device T* K [[buffer(1)]], + const device T* V [[buffer(2)]], + device T* O [[buffer(3)]], + const constant AttnParams* params [[buffer(4)]], + const constant AttnMaskParams* mask_params [[buffer(5), function_constant(has_mask)]], + const device MaskType* mask [[buffer(6), function_constant(has_mask)]], + const device T* sinks [[buffer(7), function_constant(has_sinks)]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]]) { // clang-format on + + // Pacifying compiler + (void)lid; + + // Move to correct block + ulong3 tidl{tid.x, tid.y, tid.z}; + + Q += tidl.z * params->Q_strides[0] + // Batch + tidl.y * params->Q_strides[1] + // Head + tidl.x * BQ * params->Q_strides[2]; // Sequence + + ulong kv_head_idx = int(tid.y) / params->gqa_factor; + K += tidl.z * params->K_strides[0] + // Batch + kv_head_idx * params->K_strides[1]; // Head + + V += tidl.z * params->V_strides[0] + // Batch + kv_head_idx * params->V_strides[1]; // Head + + O += tidl.z * params->O_strides[0] + // Batch + tidl.y * params->O_strides[1] + // Head + tidl.x * BQ * params->O_strides[2]; // Sequence + + if (has_mask) { + mask += tidl.z * mask_params->M_strides[0] + // Batch + tidl.y * mask_params->M_strides[1]; // Head + } + + const metal::uniform scale2 = + make_uniform(params->scale) * make_uniform(1.44269504089f); + + // Prepare MMA tiles + constexpr short kU = 16; + + // The WM simdgroups along the first warp dimension split the Q sequence; + // the WN simdgroups along the second split the head dim. The exchange + // below reduces exactly one peer, so WN is fixed at 2. + static_assert(WN == 2, "The head-dim split kernel needs WN == 2"); + constexpr int kNWarps = WM; + static_assert( + BQ >= (kNWarps * kU) && BQ % (kNWarps * kU) == 0, + "Each simdgroup must host atleast 1 simdgroup matrix along Q sequence."); + + // Q seq frags per warp + constexpr int TQ = BQ / (kNWarps * kU); + // HeadDim frags over the full head dim + constexpr int TD = BD / kU; + // KV seq frags per warp + constexpr short TK = BK / kU; + + static_assert(TQ == 1, "Check TQ"); + static_assert(TD % WN == 0, "The head dim must split evenly across WN"); + + // HeadDim frags / columns owned by each of the WN simdgroups of a row group + constexpr int TDh = TD / WN; + constexpr int BDh = BD / WN; + + static_assert(TDh % 2 == 0, "P@V accumulates output fragments in pairs"); + static_assert(TK % 2 == 0, "S fragments are exchanged pair by pair"); + + const short row_group = simd_group_id / WN; + const short d_half = simd_group_id % WN; + + using otile_t = NAXTile; + otile_t Otile; + Otile.clear(); + + const short tm = kU * TQ * row_group; + Q += tm * int(params->Q_strides[2]) + d_half * BDh; + K += d_half * BDh; + V += d_half * BDh; + O += tm * int(params->O_strides[2]) + d_half * BDh; + + constexpr short kRowsPT = otile_t::kRowsPerThread; + + metal::vec max_score; + metal::vec sum_score{0}; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = Limits::finite_min; + } + + if (has_sinks) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + max_score[i] = M_LOG2E_F * static_cast(sinks[tidl.y]); + sum_score[i] = 1; + } + } + + int kb_lim = params->NK; + int kb_min_causal = params->NK; + + if (do_causal) { + int q_max = (tid.x + 1) * BQ + params->qL_off; + kb_lim = (q_max + BK - 1) / BK; + kb_lim = min(params->NK, kb_lim); + + int q_min = tid.x * BQ + params->qL_off; + q_min = max(0, q_min); + kb_min_causal = (q_min / BK); + } + + const bool is_last_q = int(tid.x) == (params->NQ_aligned); + const short lim_rows_q = params->qL_rem - tm; + const short lim_rows_k = params->kL_rem; + + using stile_t = NAXTile; + constexpr short kEPF = stile_t::NAXFrag_t::kElemsPerFrag; + + // One slot per (row group, half): a fragment pair in per-lane-linear + // layout. Both halves share the fragment-to-lane mapping, so the + // exchange needs no coordinate math. + threadgroup AccumType s_xchg[WM][WN][2 * kEPF * 32]; + + // Keep the simdgroup's Q half resident in registers for the whole KV + // loop: TDh fragments of T are cheap next to the accumulators. + NAXTile Qtiles[TDh]; + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + const int Q_load_off = id * kU; + if (!align_Q && is_last_q) { + Qtiles[id].load_rows( + Q + Q_load_off, int(params->Q_strides[2]), lim_rows_q); + } else { + Qtiles[id].load(Q + Q_load_off, int(params->Q_strides[2])); + } + } + + const short2 simd_coord = otile_t::NAXFrag_t::get_coord(); + const short sm = simd_coord.y; + const short sn = simd_coord.x; + + // Loop over KV seq length + for (int kb = 0; kb < kb_lim; kb++) { + const int is_last_k = (kb == (params->NK_aligned)); + + stile_t Stile; + Stile.clear(); + + // S = Q @ K.T, this half of D only, exchanged pair by pair. + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik += 2) { + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id++) { + NAXTile Ktile; + const int K_load_off = ik * kU * int(params->K_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Ktile.load_rows( + K + K_load_off, int(params->K_strides[2]), lim_rows_k - ik * kU); + } else { + Ktile.load(K + K_load_off, int(params->K_strides[2])); + } + + stile_t::NAXFrag_t::mma( + Stile.frag_at(0, ik), + Stile.frag_at(0, ik + 1), + Qtiles[id].frag_at(0, 0), + metal::false_type{}, + Ktile.frag_at(0, 0), + Ktile.frag_at(1, 0), + metal::true_type{}); + } + + // Exchange the partial pair and reduce. + threadgroup AccumType* slot = s_xchg[row_group][d_half]; + thread auto& s0 = Stile.frag_at(0, ik); + thread auto& s1 = Stile.frag_at(0, ik + 1); + const short base = short(simd_lane_id) * (2 * kEPF); + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + slot[base + i] = s0[i]; + slot[base + kEPF + i] = s1[i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const threadgroup AccumType* peer = s_xchg[row_group][1 - d_half]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kEPF; i++) { + s0[i] += peer[base + i]; + s1[i] += peer[base + kEPF + i]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + // Scale S + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kElemsPerTile; ii++) { + Stile.elems()[ii] *= float(scale2); + } + + // Mask out length sequence + if (!align_K && is_last_k) { + constexpr auto neg_inf = Limits::finite_min; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const short col_pos = ik * kU + sn; + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = ((col_pos + jj) < params->kL_rem) ? fg[loc] : neg_inf; + } + } + } + } + + // Mask out if causal + if (do_causal && kb >= kb_min_causal) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + params->qL_off + tm; + const int base_col = kb * BK; + + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < stile_t::kFragThrRows; ii++) { + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < stile_t::kFragThrCols; jj++) { + const auto r = base_row + ii * stile_t::kFragRowsJump + sm; + const auto c = base_col + ik * kU + jj + sn; + const auto loc = ii * stile_t::kFragThrCols + jj; + fg[loc] = (r < c) ? neg_inf : fg[loc]; + } + } + } + } + + // Other masking as needed + if (has_mask) { + constexpr auto neg_inf = Limits::finite_min; + + const int base_row = tid.x * BQ + tm; + const int base_col = kb * BK; + + constexpr bool is_bool = is_same_v; + using melem_t = typename metal::conditional_t; + using mtile_t = NAXTile; + using mfrag_t = typename mtile_t::frag_type; + + if (base_row + kU <= params->qL && base_col + BK <= params->kL) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + const int row_pos = base_row; + const int col_pos = base_col + ik * kU; + + mfrag_t mfrag; + mtile_t::NAXFrag_t::load_safe( + mfrag, + mask, + int64_t(mask_params->M_strides[2]), + Int<1>{}, + params->qL, + params->kL, + row_pos, + col_pos); + + thread auto& fg = Stile.frag_at(0, ik); + + STEEL_PRAGMA_UNROLL + for (short jj = 0; jj < mtile_t::kElemsPerFrag; jj++) { + if constexpr (is_bool) { + fg[jj] = mfrag[jj] ? fg[jj] : neg_inf; + } else { + fg[jj] += M_LOG2E_F * AccumType(mfrag[jj]); + } + } + } + } + } + + // Do softmax (redundantly per half; the row statistics are cheap) + metal::vec new_max; + metal::vec factor; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + new_max[i] = max_score[i]; + } + + Stile.template row_reduce(new_max); + Stile.template row_bin_op(new_max); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + factor[i] = fast::exp2(max_score[i] - new_max[i]); + max_score[i] = new_max[i]; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + sum_score[i] = sum_score[i] * factor[i]; + } + + Stile.template row_reduce(sum_score); + + Otile.template row_bin_op(factor); + + simdgroup_barrier(mem_flags::mem_none); + + // O = P @ V, this half of Dv only. + STEEL_PRAGMA_UNROLL + for (short id = 0; id < TDh; id += 2) { + STEEL_PRAGMA_UNROLL + for (short ik = 0; ik < TK; ik++) { + NAXTile Vtile; + + const int V_load_off = ik * kU * int(params->V_strides[2]) + id * kU; + + if (!align_K && is_last_k) { + Vtile.load_rows( + V + V_load_off, int(params->V_strides[2]), lim_rows_k - ik * kU); + } else { + Vtile.load(V + V_load_off, int(params->V_strides[2])); + } + + otile_t::NAXFrag_t::mma( + Otile.frag_at(0, id), + Otile.frag_at(0, id + 1), + Stile.frag_at(0, ik), + metal::false_type{}, + Vtile.frag_at(0, 0), + Vtile.frag_at(0, 1), + metal::false_type{}); + } + } + + // Next block + K += BK * int(params->K_strides[2]); + V += BK * int(params->V_strides[2]); + } + + // Normalize output + threadgroup_barrier(mem_flags::mem_none); + + metal::vec rcp; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kRowsPT; ++i) { + rcp[i] = 1.f / sum_score[i]; + } + + Otile.template row_bin_op(rcp); + + if (!align_Q && is_last_q) { + if (lim_rows_q <= 0) + return; + Otile.store_rows(O, int(params->O_strides[2]), lim_rows_q); } else { Otile.store(O, int(params->O_strides[2])); } diff --git a/Source/Cmlx/mlx-generated/steel_conv.cpp b/Source/Cmlx/mlx-generated/steel_conv.cpp index a7d6b33c2..2af0e5605 100644 --- a/Source/Cmlx/mlx-generated/steel_conv.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv.cpp @@ -148,7 +148,7 @@ implicit_gemm_conv_2d( C += tid.z * N; B += c_col * K; - C += c_row * (N * params->groups) + c_col; + C += static_cast(c_row) * N * params->groups + c_col; const int2 offsets_a(0, c_row); const int2 offsets_b(0, c_col); diff --git a/Source/Cmlx/mlx-generated/steel_conv_3d.cpp b/Source/Cmlx/mlx-generated/steel_conv_3d.cpp index 825825838..db699b376 100644 --- a/Source/Cmlx/mlx-generated/steel_conv_3d.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv_3d.cpp @@ -107,7 +107,7 @@ implicit_gemm_conv_3d( C += tid.z * N; B += c_col * K; - C += c_row * (N * params->groups) + c_col; + C += static_cast(c_row) * N * params->groups + c_col; const int2 offsets_a(0, c_row); const int2 offsets_b(0, c_col); diff --git a/Source/Cmlx/mlx-generated/steel_conv_general.cpp b/Source/Cmlx/mlx-generated/steel_conv_general.cpp index b8ca6de78..0024a77f5 100644 --- a/Source/Cmlx/mlx-generated/steel_conv_general.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv_general.cpp @@ -90,10 +90,10 @@ struct Conv2DInputBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), params(params_), jump_params(jump_params_), @@ -124,7 +124,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -160,7 +160,7 @@ struct Conv2DInputBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { STEEL_PRAGMA_UNROLL for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { // Find bounds @@ -207,7 +207,7 @@ struct Conv2DInputBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; @@ -286,11 +286,11 @@ struct Conv2DWeightBlockLoaderGeneral { const short base_wh_, const short base_ww_, uint simd_group_id [[simdgroup_index_in_threadgroup]], - uint simd_lane_id [[thread_index_in_simdgroup]]) + uint simd_lane_id [[thread_index_in_simdgroup]]) thread : src_ld(params_->wt_strides[0]), thread_idx(simd_group_id * 32 + simd_lane_id), bi(thread_idx / TCOLS), - bj(vec_size * (thread_idx % TCOLS)), + bj(vec_size*(thread_idx % TCOLS)), dst(dst_ + bi * dst_ld + bj), src(src_ + bi * src_ld + bj), params(params_), @@ -302,7 +302,7 @@ struct Conv2DWeightBlockLoaderGeneral { start_row(offsets.y + bi) {} /* Load from device memory into threadgroup memory - without bound checking */ - METAL_FUNC void load_unsafe() const { + METAL_FUNC void load_unsafe() const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -331,7 +331,7 @@ struct Conv2DWeightBlockLoaderGeneral { } } - METAL_FUNC void load_safe(const short remaining_k) const { + METAL_FUNC void load_safe(const short remaining_k) const thread { const device T* curr_src = src + weight_h * params->wt_strides[1] + weight_w * params->wt_strides[2]; @@ -381,7 +381,7 @@ struct Conv2DWeightBlockLoaderGeneral { } /* Iteration helper */ - METAL_FUNC void next() { + METAL_FUNC void next() thread { weight_w += jump_params->f_wgt_jump_w; if (weight_w < params->wS[1]) { return; @@ -609,14 +609,14 @@ implicit_gemm_conv_2d_general( (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + base_ow; if (n < params->N && oh < params->oS[0] && ow < params->oS[1]) { - int offset_cm = n * params->out_strides[0] + + size_t offset_cm = static_cast(n) * params->out_strides[0] + oh * params->out_strides[1] + ow * params->out_strides[2]; STEEL_PRAGMA_UNROLL for (int j = 0; j < mma_t::TN; j++) { // Get accumulated result and associated offset in C thread const auto& accum = mma_op.Ctile.frag_at(i, j); - int offset = offset_cm + (j * mma_t::TN_stride); + size_t offset = offset_cm + (j * mma_t::TN_stride); constexpr short kelems = decltype(mma_op.Ctile)::kElemsPerFrag; diff --git a/Source/Cmlx/mlx-generated/steel_gemm_fused_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_fused_nax.cpp index ccbf8e5e4..394942381 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_fused_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_fused_nax.cpp @@ -28,7 +28,7 @@ constant bool align_K [[function_constant(202)]]; template < bool kAlignedM, bool kAlignedN, - typename NAXTile_t, + class NAXTile_t, typename T> void gemm_epilogue( thread NAXTile_t& Dtile, @@ -40,37 +40,40 @@ void gemm_epilogue( (void)params; - constexpr short UM = NAXTile_t::kSubTileRows; - constexpr short UN = NAXTile_t::kSubTileCols; - using CSubTile = NAXSubTile; - using V = typename NAXTile_t::elem_type; constexpr short TM = NAXTile_t::kTileRows; constexpr short TN = NAXTile_t::kTileCols; - constexpr short kElemsPerSubTile = NAXTile_t::kElemsPerSubTile; + constexpr short kElemsPerFrag = NAXTile_t::kElemsPerFrag; + + using CFrag = typename NAXTile_t::NAXFrag_t; + using cfrag_t = typename CFrag::template dtype_frag_t; - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { - const short m = mm * UM; - const short n = nn * UN; + const_for_loop<0, TM, 1>([&](auto mm) { + const_for_loop<0, TN, 1>([&](auto nn) { + auto m = mm * Int{}; + auto n = nn * Int{}; - CSubTile CTile; + cfrag_t celems; if constexpr (kAlignedM && kAlignedN) { - CTile.load(C, addmm_params->ldc, addmm_params->fdc, m, n); + CFrag::load(celems, C, addmm_params->ldc, addmm_params->fdc, m, n); } else { - CTile.load_safe( - C, addmm_params->ldc, addmm_params->fdc, sgp_sm, sgp_sn, m, n); + CFrag::load_safe( + celems, + C, + addmm_params->ldc, + addmm_params->fdc, + sgp_sm, + sgp_sn, + m, + n); } - auto delems = Dtile.subtile_at(mm, nn).elems(); - auto celems = CTile.elems(); + thread auto& delems = Dtile.template frag_at(); STEEL_PRAGMA_UNROLL - for (short i = 0; i < kElemsPerSubTile; i++) { + for (short i = 0; i < kElemsPerFrag; i++) { if (do_axpby) { delems[i] = addmm_params->alpha * delems[i] + addmm_params->beta * static_cast(celems[i]); @@ -78,8 +81,8 @@ void gemm_epilogue( delems[i] += static_cast(celems[i]); } } - } - } + }); + }); } // clang-format off @@ -157,15 +160,12 @@ template < C += c_row_long * addmm_params->ldc + c_col_long * addmm_params->fdc; } - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; const short tm = SM * (simd_group_id / WN); const short tn = SN * (simd_group_id % WN); @@ -188,8 +188,7 @@ template < C += tm * addmm_params->ldc + tn * addmm_params->fdc; } - using DSubTile = NAXSubTile; - NAXTile Dtile; + NAXTile Dtile; dispatch_bool(align_K, [&](auto kAlignedK) { dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { @@ -205,9 +204,6 @@ template < kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>( A, B, @@ -217,14 +213,17 @@ template < params->gemm_k_iterations_aligned, sgp_sm, sgp_sn); - if (use_out_source) { - gemm_epilogue( - Dtile, C, params, addmm_params, sgp_sm, sgp_sn); - } - if constexpr (kAlignedM && kAlignedN) { - Dtile.store(D, int(params->ldd)); - } else { - Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + if ((kAlignedM.value || sgp_sm > 0) && + (kAlignedN.value || sgp_sn > 0)) { + if (use_out_source) { + gemm_epilogue( + Dtile, C, params, addmm_params, sgp_sm, sgp_sn); + } + if constexpr (kAlignedM && kAlignedN) { + Dtile.store(D, int(params->ldd)); + } else { + Dtile.store_safe(D, int(params->ldd), short2(sgp_sn, sgp_sm)); + } } }); }); diff --git a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp index f4037edb1..7312ba458 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_gather_nax.cpp @@ -38,14 +38,11 @@ gather_mm_rhs_nax( const constant GEMMParams* params [[buffer(4)]], uint simd_group_id [[simdgroup_index_in_threadgroup]], uint3 tid [[threadgroup_position_in_grid]]) { - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; if (params->tiles_n <= static_cast(tid.x) || params->tiles_m <= static_cast(tid.y)) { @@ -101,13 +98,12 @@ gather_mm_rhs_nax( } threadgroup_barrier(mem_flags::mem_none); - using DSubTile = NAXSubTile; - NAXTile Ctile; + NAXTile Ctile; dispatch_bool(align_K, [&](auto kAlignedK) { dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { - auto do_gemm = gemm_loop< + auto do_gemm = gemm_loop< // Matmul for partial BM, full BN and full K T, SM, SN, @@ -118,9 +114,6 @@ gather_mm_rhs_nax( kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>; Ctile = do_gemm( A, diff --git a/Source/Cmlx/mlx-generated/steel_gemm_masked.cpp b/Source/Cmlx/mlx-generated/steel_gemm_masked.cpp index cbb3d8d15..a220e52e4 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_masked.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_masked.cpp @@ -35,7 +35,7 @@ using namespace mlx::steel; struct _NoMask { char x; - constexpr METAL_FUNC operator bool() { + constexpr METAL_FUNC operator bool() thread { return true; } constexpr METAL_FUNC operator bool() const threadgroup { @@ -53,7 +53,7 @@ template struct ScaleOp { OutT scale; - METAL_FUNC OutT apply(InT x) const { + METAL_FUNC OutT apply(InT x) const thread { return static_cast(x) * scale; } }; diff --git a/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp new file mode 100644 index 000000000..519087e1b --- /dev/null +++ b/Source/Cmlx/mlx-generated/steel_gemm_segmented_nax.cpp @@ -0,0 +1,144 @@ +namespace mlx::core::metal { + +const char* steel_gemm_segmented_nax() { + return R"preamble( +// Copyright © 2025 Apple Inc. + +// Auto generated source for mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h + +/////////////////////////////////////////////////////////////////////////////// +// Contents from "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h" +/////////////////////////////////////////////////////////////////////////////// + +#line 1 "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h" +// Copyright © 2026 Apple Inc. + +using namespace mlx::steel; + +constant bool segments_contiguous [[function_constant(199)]]; +constant bool align_M [[function_constant(200)]]; +constant bool align_N [[function_constant(201)]]; + +template < + typename T, + int BM, + int BN, + int BK, + int WM, + int WN, + bool transpose_a, + bool transpose_b, + typename AccumType = float> +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] +void segmented_mm_nax( + const device T* A [[buffer(0)]], + const device T* B [[buffer(1)]], + const device uint32_t* segments [[buffer(2)]], + device T* C [[buffer(3)]], + const constant GEMMParams* params [[buffer(4)]], + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint3 tid [[threadgroup_position_in_grid]]) { + const int tid_m = (BK > 64) ? tid.y : tid.z; + const int tid_n = (BK > 64) ? tid.x : tid.y; + const int tid_s = (BK > 64) ? tid.z : tid.x; + + const int c_row = tid_m * BM; + const int c_col = tid_n * BN; + const size_t c_row_long = size_t(c_row); + const size_t c_col_long = size_t(c_col); + + if (params->tiles_n <= static_cast(tid_n) || + params->tiles_m <= static_cast(tid_m)) { + return; + } + + A += transpose_a ? c_row_long : c_row_long * params->lda; + B += transpose_b ? c_col_long * params->ldb : c_col_long; + C += c_row_long * params->ldd + c_col_long; + + uint32_t k_start, k_end; + if (segments_contiguous) { + k_start = segments[2 * tid_s]; + k_end = segments[2 * tid_s + 1]; + } else { + k_start = segments[tid_s]; + k_end = segments[tid_s + 1]; + } + A += transpose_a ? k_start * params->lda : k_start; + B += transpose_b ? k_start : k_start * params->ldb; + C += tid_s * params->batch_stride_d; + + constexpr short SM = BM / WM; + constexpr short SN = BN / WN; + constexpr short SK = 32; + + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; + + const short tm = SM * (simd_group_id / WN); + const short tn = SN * (simd_group_id % WN); + + const int sgp_sm_int = + align_M ? int(SM) : min(int(SM), params->M - (c_row + tm)); + const short sgp_sm = short(sgp_sm_int); + const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM); + + const int sgp_sn_int = + align_N ? int(SN) : min(int(SN), params->N - (c_col + tn)); + const short sgp_sn = short(sgp_sn_int); + const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN); + + A += transpose_a ? tm : (tm * params->lda); + B += transpose_b ? (tn * params->ldb) : tn; + C += tm * params->ldd + tn; + + NAXTile Dtile; + Dtile.clear(); + + const int segment_k_size = k_end - k_start; + const int segment_k_iters = segment_k_size / BK; + const bool segment_k_aligned = (segment_k_size % BK) == 0; + + dispatch_bool(segment_k_aligned, [&](auto kAlignedK) { + dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { + Dtile = gemm_loop< + T, + SM, + SN, + SK, + BK, + transpose_a, + transpose_b, + kAlignedM.value, + kAlignedN.value, + kAlignedK.value, + AccumType>( + A, + B, + params->lda, + params->ldb, + segment_k_size, + segment_k_iters, + sgp_sm, + sgp_sn); + }); + }); + }); + + dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) { + dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) { + if constexpr (kAlignedM && kAlignedN) { + Dtile.store(C, int(params->ldd)); + } else { + Dtile.store_safe(C, int(params->ldd), short2(sgp_sn, sgp_sm)); + } + }); + }); +} + +/////////////////////////////////////////////////////////////////////////////// +)preamble"; +} + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx-generated/steel_gemm_splitk_nax.cpp b/Source/Cmlx/mlx-generated/steel_gemm_splitk_nax.cpp index 25c87ceb3..ff3127cae 100644 --- a/Source/Cmlx/mlx-generated/steel_gemm_splitk_nax.cpp +++ b/Source/Cmlx/mlx-generated/steel_gemm_splitk_nax.cpp @@ -85,15 +85,12 @@ template < (c_row_long * params->ldc + c_col_long); // NAX tile configuration - constexpr short UM = 16; - constexpr short UN = 32; - constexpr short UK = 16; constexpr short SM = BM / WM; constexpr short SN = BN / WN; constexpr short SK = 32; - constexpr short TM = SM / UM; - constexpr short TN = SN / UN; + constexpr short TM = SM / 16; + constexpr short TN = SN / 16; // Calculate simdgroup offsets and alignment const short tm = SM * (simd_group_id / WN); @@ -113,8 +110,7 @@ template < B += transpose_b ? (tn * params->ldb) : tn; C += tm * params->ldc + tn; - using DSubTile = NAXSubTile; - NAXTile Dtile; + NAXTile Dtile; // gemm_loop through the partition // Check K-alignment at runtime (partition-specific) @@ -136,9 +132,6 @@ template < kAlignedM.value, kAlignedN.value, kAlignedK.value, - UM, - UN, - UK, AccumType>( A, B, diff --git a/Source/Cmlx/mlx-generated/ternary_ops.cpp b/Source/Cmlx/mlx-generated/ternary_ops.cpp index 8a15ee36b..dbcf6a388 100644 --- a/Source/Cmlx/mlx-generated/ternary_ops.cpp +++ b/Source/Cmlx/mlx-generated/ternary_ops.cpp @@ -16,7 +16,7 @@ const char* ternary_ops() { struct Select { template - T operator()(bool condition, T x, T y) { + T operator()(bool condition, T x, T y) thread { return condition ? x : y; } }; diff --git a/Source/Cmlx/mlx-generated/unary_ops.cpp b/Source/Cmlx/mlx-generated/unary_ops.cpp index 3226c35cc..3a680b7f2 100644 --- a/Source/Cmlx/mlx-generated/unary_ops.cpp +++ b/Source/Cmlx/mlx-generated/unary_ops.cpp @@ -321,7 +321,7 @@ float erfinv(float a) { struct fp8_e4m3 { template - fp8_e4m3(T f) { + fp8_e4m3(T f) thread { // From PyTorch // https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148 uint32_t fp8_max = 543 << 21; @@ -348,7 +348,7 @@ struct fp8_e4m3 { bits |= static_cast(sign >> 24); } - operator float16_t() { + operator float16_t() thread { uint16_t v = (bits & 127) << 7; half converted = as_type(v); converted *= 256.0; @@ -356,11 +356,11 @@ struct fp8_e4m3 { return (sign ? -converted : converted); } - operator bfloat16_t() { + operator bfloat16_t() thread { return static_cast(this->operator float16_t()); } - operator float() { + operator float() thread { return static_cast(this->operator float16_t()); } @@ -368,7 +368,7 @@ struct fp8_e4m3 { }; struct fp8_e8m0 { - fp8_e8m0(float x) { + fp8_e8m0(float x) thread { if (!metal::isfinite(x)) { bits = 0xFF; return; @@ -385,12 +385,12 @@ struct fp8_e8m0 { bits = static_cast(n + 127); } - operator bfloat16_t() { + operator bfloat16_t() thread { uint16_t out = (bits == 0 ? 0x40 : (static_cast(bits) << 7)); return as_type(out); } - operator float() { + operator float() thread { uint32_t out = (bits == 0 ? 0x400000 : (static_cast(bits) << 23)); return as_type(out); } @@ -398,6 +398,17 @@ struct fp8_e8m0 { uint8_t bits; }; +// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down +// leaves the block's largest elements outside the element range, where they +// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0. +inline float mx_scale_round_up(float x) { + fp8_e8m0 s(x); + if (s.bits < 0xFE && float(s) < x) { + s.bits += 1; + } + return float(s); +} + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/unary_ops.h" /////////////////////////////////////////////////////////////////////////////// @@ -416,125 +427,125 @@ constant float inf = metal::numeric_limits::infinity(); struct Abs { template - T operator()(T x) { + T operator()(T x) thread { return metal::abs(x); }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0}; }; }; struct ArcCos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acos(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcCosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::acosh(x); }; }; struct ArcSin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asin(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcSinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::asinh(x); }; }; struct ArcTan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atan(x); }; - complex64_t operator()(complex64_t x); + complex64_t operator()(complex64_t x) thread; }; struct ArcTanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::atanh(x); }; }; struct BitwiseInvert { template - T operator()(T x) { + T operator()(T x) thread { return ~x; }; }; struct Ceil { template - T operator()(T x) { + T operator()(T x) thread { return metal::ceil(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Cos { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cos(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cos(x.real) * metal::precise::cosh(x.imag), -metal::precise::sin(x.real) * metal::precise::sinh(x.imag)}; @@ -543,11 +554,11 @@ struct Cos { struct Cosh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::cosh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::cosh(x.real) * metal::precise::cos(x.imag), metal::precise::sinh(x.real) * metal::precise::sin(x.imag)}; @@ -555,89 +566,89 @@ struct Cosh { }; struct Conjugate { - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return complex64_t{x.real, -x.imag}; } }; struct Erf { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erf(static_cast(x))); }; }; struct ErfInv { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(erfinv(static_cast(x))); }; }; struct Exp { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::exp(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return cexpf(x); } }; struct Expm1 { template - T operator()(T x) { + T operator()(T x) thread { return static_cast(expm1f(static_cast(x))); }; }; struct Floor { template - T operator()(T x) { + T operator()(T x) thread { return metal::floor(x); }; - int8_t operator()(int8_t x) { + int8_t operator()(int8_t x) thread { return x; }; - int16_t operator()(int16_t x) { + int16_t operator()(int16_t x) thread { return x; }; - int32_t operator()(int32_t x) { + int32_t operator()(int32_t x) thread { return x; }; - int64_t operator()(int64_t x) { + int64_t operator()(int64_t x) thread { return x; }; - uint8_t operator()(uint8_t x) { + uint8_t operator()(uint8_t x) thread { return x; }; - uint16_t operator()(uint16_t x) { + uint16_t operator()(uint16_t x) thread { return x; }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x; }; - uint64_t operator()(uint64_t x) { + uint64_t operator()(uint64_t x) thread { return x; }; - bool operator()(bool x) { + bool operator()(bool x) thread { return x; }; }; struct Imag { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.imag; }; }; struct Log { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto r = metal::precise::log(Abs{}(x).real); auto i = metal::precise::atan2(x.imag, x.real); return {r, i}; @@ -646,11 +657,11 @@ struct Log { struct Log2 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log2(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN2_F, y.imag / M_LN2_F}; }; @@ -658,11 +669,11 @@ struct Log2 { struct Log10 { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::log10(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { auto y = Log{}(x); return {y.real / M_LN10_F, y.imag / M_LN10_F}; }; @@ -670,44 +681,44 @@ struct Log10 { struct Log1p { template - T operator()(T x) { + T operator()(T x) thread { return log1p(x); }; }; struct LogicalNot { template - T operator()(T x) { + T operator()(T x) thread { return !x; }; }; struct Negative { template - T operator()(T x) { + T operator()(T x) thread { return -x; }; }; struct Real { - float operator()(complex64_t x) { + float operator()(complex64_t x) thread { return x.real; }; }; struct Round { template - T operator()(T x) { + T operator()(T x) thread { return metal::rint(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return {metal::rint(x.real), metal::rint(x.imag)}; }; }; struct Sigmoid { template - T operator()(T x) { + T operator()(T x) thread { auto y = 1 / (1 + metal::exp(metal::abs(x))); return (x < 0) ? y : 1 - y; } @@ -715,13 +726,13 @@ struct Sigmoid { struct Sign { template - T operator()(T x) { + T operator()(T x) thread { return (x > T(0)) - (x < T(0)); }; - uint32_t operator()(uint32_t x) { + uint32_t operator()(uint32_t x) thread { return x != 0; }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x == complex64_t(0)) { return x; } @@ -732,11 +743,11 @@ struct Sign { struct Sin { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sin(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sin(x.real) * metal::precise::cosh(x.imag), metal::precise::cos(x.real) * metal::precise::sinh(x.imag)}; @@ -745,11 +756,11 @@ struct Sin { struct Sinh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sinh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return { metal::precise::sinh(x.real) * metal::precise::cos(x.imag), metal::precise::cosh(x.real) * metal::precise::sin(x.imag)}; @@ -758,18 +769,18 @@ struct Sinh { struct Square { template - T operator()(T x) { + T operator()(T x) thread { return x * x; }; }; struct Sqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::sqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { if (x.real == 0.0 && x.imag == 0.0) { return {0.0, 0.0}; } @@ -783,22 +794,22 @@ struct Sqrt { struct Rsqrt { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::rsqrt(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { return 1.0 / Sqrt{}(x); } }; struct Tan { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tan(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tan_a = metal::precise::tan(x.real); float tanh_b = metal::precise::tanh(x.imag); float t1 = tan_a * tanh_b; @@ -809,11 +820,11 @@ struct Tan { struct Tanh { template - T operator()(T x) { + T operator()(T x) thread { return metal::precise::tanh(x); }; - complex64_t operator()(complex64_t x) { + complex64_t operator()(complex64_t x) thread { float tanh_a = metal::precise::tanh(x.real); float tan_b = metal::precise::tan(x.imag); float t1 = tanh_a * tan_b; @@ -822,19 +833,19 @@ struct Tanh { }; }; -complex64_t ArcCos::operator()(complex64_t x) { +complex64_t ArcCos::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(x + i * Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcSin::operator()(complex64_t x) { +complex64_t ArcSin::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto y = Log{}(i * x + Sqrt{}(1.0 - x * x)); return {y.imag, -y.real}; }; -complex64_t ArcTan::operator()(complex64_t x) { +complex64_t ArcTan::operator()(complex64_t x) thread { auto i = complex64_t{0.0, 1.0}; auto ix = i * x; return (1.0 / complex64_t{0.0, 2.0}) * Log{}((1.0 + ix) / (1.0 - ix)); @@ -842,13 +853,13 @@ complex64_t ArcTan::operator()(complex64_t x) { struct ToFP8 { template - uint8_t operator()(T f) { + uint8_t operator()(T f) thread { return fp8_e4m3(f).bits; } }; struct FromFP8 { - float operator()(uint8_t x) { + float operator()(uint8_t x) thread { return float(*(thread fp8_e4m3*)(&x)); } }; diff --git a/Source/Cmlx/mlx-generated/utils.cpp b/Source/Cmlx/mlx-generated/utils.cpp index deb3d7419..18d9644d0 100644 --- a/Source/Cmlx/mlx-generated/utils.cpp +++ b/Source/Cmlx/mlx-generated/utils.cpp @@ -422,163 +422,285 @@ instantiate_metal_simd_reduction_funcs(bfloat16_t, bfloat16_t, float); #include + using namespace metal; -struct complex64_t; +template +struct complex_t; template -static constexpr constant bool can_convert_to_complex64 = - !is_same_v && is_convertible_v; +static constexpr constant bool is_complex_v = false; template -static constexpr constant bool can_convert_from_complex64 = - !is_same_v && - (is_convertible_v || is_convertible_v); +static constexpr constant bool is_complex_v> = true; -struct complex64_t { - float real; - float imag; +// Metal accepts explicit bfloat casts that is_convertible_v reports as false. +template +static constexpr constant bool is_lane_convertible_v = + is_convertible_v || + (is_same_v && is_convertible_v) || + (is_same_v && is_convertible_v); + +template +struct complex_t { + using value_type = T; + + T real; + T imag; // Constructors - constexpr complex64_t(float real, float imag) : real(real), imag(imag) {}; - constexpr complex64_t() : real(0), imag(0) {}; - constexpr complex64_t() threadgroup : real(0), imag(0) {}; + constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {}; + constexpr complex_t() thread : real(0), imag(0) {}; + constexpr complex_t() threadgroup : real(0), imag(0) {}; + + // Conversions from scalar types + template < + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) thread : real(static_cast(x)), + imag(static_cast(0)) {} - // Conversions to complex64_t template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) thread : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) threadgroup : real(static_cast(x)), + imag(static_cast(0)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) threadgroup : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) device : real(static_cast(x)), + imag(static_cast(0)) {} template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) device : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr complex_t(U x) constant : real(static_cast(x)), + imag(static_cast(0)) {} + // Conversions between complex types template < - typename T, - typename = typename enable_if>::type> - constexpr complex64_t(T x) constant : real(x), imag(0) {} + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) thread : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) threadgroup + : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) device : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + template < + typename U, + typename = typename enable_if< + !is_same_v && is_lane_convertible_v>::type> + constexpr complex_t(complex_t x) constant : real(static_cast(x.real)), + imag(static_cast(x.imag)) {} + + // Conversions to and from two-lane vectors (the FFT lane representation) + constexpr complex_t(vec v) thread : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) threadgroup : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) device : real(v.x), imag(v.y) {}; + constexpr complex_t(vec v) constant : real(v.x), imag(v.y) {}; + + constexpr operator vec() const thread { + return vec(real, imag); + } - // Conversions from complex64_t + constexpr operator vec() const threadgroup { + return vec(real, imag); + } + + constexpr operator vec() const device { + return vec(real, imag); + } + + constexpr operator vec() const constant { + return vec(real, imag); + } + + // Conversions to scalar types template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const thread { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const thread { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const threadgroup { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const threadgroup { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const device { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const device { + return static_cast(real); } template < - typename T, - typename = typename enable_if>::type> - constexpr operator T() const constant { - return static_cast(real); + typename U, + typename = typename enable_if< + !is_complex_v && is_lane_convertible_v>::type> + constexpr operator U() const constant { + return static_cast(real); } }; -constexpr complex64_t operator-(complex64_t x) { +using complex32_t = complex_t; +using complex64_t = complex_t; + +static_assert(sizeof(complex32_t) == 2 * sizeof(half)); +static_assert(sizeof(complex64_t) == 2 * sizeof(float)); +static_assert(sizeof(complex_t) == 2 * sizeof(bfloat16_t)); + +template +constexpr complex_t operator-(complex_t x) { return {-x.real, -x.imag}; } -constexpr bool operator>=(complex64_t a, complex64_t b) { +template +constexpr bool operator>=(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); } -constexpr bool operator>(complex64_t a, complex64_t b) { +template +constexpr bool operator>(complex_t a, complex_t b) { return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); } -constexpr bool operator<=(complex64_t a, complex64_t b) { +template +constexpr bool operator<=(complex_t a, complex_t b) { return operator>=(b, a); } -constexpr bool operator<(complex64_t a, complex64_t b) { +template +constexpr bool operator<(complex_t a, complex_t b) { return operator>(b, a); } -constexpr bool operator==(complex64_t a, complex64_t b) { +template +constexpr bool operator==(complex_t a, complex_t b) { return a.real == b.real && a.imag == b.imag; } -constexpr complex64_t operator+(complex64_t a, complex64_t b) { +template +constexpr complex_t operator+(complex_t a, complex_t b) { return {a.real + b.real, a.imag + b.imag}; } -constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) { +template +constexpr thread complex_t& operator+=( + thread complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr threadgroup complex64_t& operator+=( - threadgroup complex64_t& a, - complex64_t b) { +template +constexpr threadgroup complex_t& operator+=( + threadgroup complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) { +template +constexpr device complex_t& operator+=( + device complex_t& a, + complex_t b) { a.real += b.real; a.imag += b.imag; return a; } -constexpr complex64_t operator+(float a, complex64_t b) { - return {a + b.real, b.imag}; +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(U a, complex_t b) { + return {static_cast(a) + b.real, b.imag}; } -constexpr complex64_t operator+(complex64_t a, float b) { - return {a.real + b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator+(complex_t a, U b) { + return {a.real + static_cast(b), a.imag}; } -constexpr complex64_t operator-(complex64_t a, complex64_t b) { +template +constexpr complex_t operator-(complex_t a, complex_t b) { return {a.real - b.real, a.imag - b.imag}; } -constexpr complex64_t operator-(float a, complex64_t b) { - return {a - b.real, -b.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(U a, complex_t b) { + return {static_cast(a) - b.real, -b.imag}; } -constexpr complex64_t operator-(complex64_t a, float b) { - return {a.real - b, a.imag}; + +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator-(complex_t a, U b) { + return {a.real - static_cast(b), a.imag}; } -constexpr complex64_t operator*(complex64_t a, complex64_t b) { +template +constexpr complex_t operator*(complex_t a, complex_t b) { return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; } -constexpr complex64_t operator/(complex64_t a, complex64_t b) { +template +constexpr complex_t operator/(complex_t a, complex_t b) { auto denom = b.real * b.real + b.imag * b.imag; auto x = a.real * b.real + a.imag * b.imag; auto y = a.imag * b.real - a.real * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator/(float a, complex64_t b) { +template < + typename T, + typename U, + enable_if_t && is_lane_convertible_v, bool> = true> +constexpr complex_t operator/(U a, complex_t b) { + auto scalar = static_cast(a); auto denom = b.real * b.real + b.imag * b.imag; - auto x = a * b.real; - auto y = -a * b.imag; + auto x = scalar * b.real; + auto y = -scalar * b.imag; return {x / denom, y / denom}; } -constexpr complex64_t operator%(complex64_t a, complex64_t b) { +template +constexpr complex_t operator%(complex_t a, complex_t b) { auto real = a.real - (b.real * static_cast(a.real / b.real)); auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); if (real != 0 && (real < 0 != b.real < 0)) { @@ -590,6 +712,13 @@ constexpr complex64_t operator%(complex64_t a, complex64_t b) { return {real, imag}; } +static_assert( + (complex_t{1.0h, 2.0h} * complex_t{3.0h, 4.0h}).real == -5.0h); +static_assert( + (complex_t{bfloat16_t(1.0f), bfloat16_t(2.0f)} * + complex_t{bfloat16_t(3.0f), bfloat16_t(4.0f)}) + .real == bfloat16_t(-5.0f)); + /////////////////////////////////////////////////////////////////////////////// // Contents from "mlx/backend/metal/kernels/defines.h" /////////////////////////////////////////////////////////////////////////////// @@ -641,7 +770,7 @@ struct os_log { constexpr os_log(constant char*, constant char*) constant {} template - void log_debug(constant char*, Args...) const {} + void log_debug(constant char*, Args...) const thread {} template void log_debug(constant char*, Args...) const constant {} @@ -726,14 +855,14 @@ struct Limits { static constexpr constant bool min = false; }; -template <> -struct Limits { - static constexpr constant complex64_t max = complex64_t( - metal::numeric_limits::infinity(), - metal::numeric_limits::infinity()); - static constexpr constant complex64_t min = complex64_t( - -metal::numeric_limits::infinity(), - -metal::numeric_limits::infinity()); +template +struct Limits> { + inline static constexpr constant complex_t max = complex_t( + metal::numeric_limits::infinity(), + metal::numeric_limits::infinity()); + inline static constexpr constant complex_t min = complex_t( + -metal::numeric_limits::infinity(), + -metal::numeric_limits::infinity()); }; /////////////////////////////////////////////////////////////////////////////// @@ -856,9 +985,9 @@ struct LoopedElemToLoc { OffsetT offset{0}; int index{0}; - LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {} + LoopedElemToLoc(int dim) thread : dim(dim), inner_looper(dim - 1) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { if (dim == 0) { return; } @@ -871,7 +1000,8 @@ struct LoopedElemToLoc { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { if (dim == 0) { return; } @@ -894,7 +1024,7 @@ struct LoopedElemToLoc { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -905,9 +1035,9 @@ struct LoopedElemToLoc<1, OffsetT, true> { OffsetT offset{0}; uint index{0}; - LoopedElemToLoc(int dim) : dim(dim) {} + LoopedElemToLoc(int dim) thread : dim(dim) {} - void next(const constant int* shape, const constant int64_t* strides) { + void next(const constant int* shape, const constant int64_t* strides) thread { index++; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -916,7 +1046,8 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - void next(int n, const constant int* shape, const constant int64_t* strides) { + void next(int n, const constant int* shape, const constant int64_t* strides) + thread { index += n; if (dim > 1) { offset = elem_to_loc(index, shape, strides, dim); @@ -925,7 +1056,7 @@ struct LoopedElemToLoc<1, OffsetT, true> { } } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -934,17 +1065,18 @@ template struct LoopedElemToLoc<1, OffsetT, false> { OffsetT offset{0}; - LoopedElemToLoc(int) {} + LoopedElemToLoc(int) thread {} - void next(const constant int*, const constant int64_t* strides) { + void next(const constant int*, const constant int64_t* strides) thread { offset += OffsetT(strides[0]); } - void next(int n, const constant int*, const constant int64_t* strides) { + void next(int n, const constant int*, const constant int64_t* strides) + thread { offset += n * OffsetT(strides[0]); } - OffsetT location() { + OffsetT location() thread { return offset; } }; @@ -1095,6 +1227,30 @@ struct ConditionalType { using type = T; }; +/////////////////////////////////////////////////////////////////////////////// +// Type casting utils +/////////////////////////////////////////////////////////////////////////////// + +template +inline U cast_to(T val) { + return static_cast(val); +} + +template <> +inline bool cast_to(float val) { + return (as_type(val) & 0x7FFFFFFF) != 0; +} + +template <> +inline bool cast_to(bfloat16_t val) { + return (as_type(val) & 0x7FFF) != 0; +} + +template <> +inline bool cast_to(complex64_t val) { + return cast_to(val.real) || cast_to(val.imag); +} + /////////////////////////////////////////////////////////////////////////////// )preamble"; } diff --git a/Source/MLX/ArrayAt.swift b/Source/MLX/ArrayAt.swift index d57747881..5482faeeb 100644 --- a/Source/MLX/ArrayAt.swift +++ b/Source/MLX/ArrayAt.swift @@ -91,6 +91,66 @@ public struct ArrayAtIndices { let indexOperations: [MLXArrayIndexOperation] let stream: StreamOrDevice + /// The `mlx_slice_update_*` functions all share this signature. + private typealias SliceUpdateReduce = ( + UnsafeMutablePointer?, mlx_array, mlx_array, + UnsafePointer?, Int, UnsafePointer?, Int, UnsafePointer?, Int, + mlx_stream + ) -> Int32 + + /// Apply `values` to the indexed region, preferring a slice update over a scatter. + /// + /// Mirrors the structure of `mlx_add_item` and friends: if the indices describe a pure + /// slice, use the corresponding reducing slice update, otherwise fall back to a scatter. + /// + /// - Parameters: + /// - values: the update value + /// - sliceUpdate: the `mlx_slice_update_*` function for this reduction + /// - scatter: the `mlx_scatter_*` function for this reduction + /// - elementwise: the whole-array fallback when the indices select everything + private func update( + _ values: some ScalarOrArray, + sliceUpdate: SliceUpdateReduce, + scatter: ( + UnsafeMutablePointer?, mlx_array, mlx_vector_array, mlx_array, + UnsafePointer?, Int, mlx_stream + ) -> Int32, + elementwise: (MLXArray, MLXArray) -> MLXArray + ) -> MLXArray { + let values = values.asMLXArray(dtype: array.dtype) + + switch sliceUpdateArguments( + src: array, operations: indexOperations, update: values, stream: stream) + { + case .slice(let update, let starts, let ends, let strides): + var result = mlx_array_new() + _ = sliceUpdate( + &result, array.ctx, update.ctx, starts, starts.count, ends, ends.count, strides, + strides.count, stream.ctx) + return MLXArray(result) + + case .broadcast(let update): + return elementwise(array, update) + + case nil: + let (indices, update, axes) = scatterArguments( + src: array, operations: indexOperations, update: values, stream: stream) + + if !indices.isEmpty { + let indices_vector = new_mlx_vector_array(indices) + defer { mlx_vector_array_free(indices_vector) } + + var result = mlx_array_new() + _ = scatter( + &result, array.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx) + + return MLXArray(result) + } else { + return elementwise(array, update) + } + } + } + /// Add values via `at[]` operator. /// /// ```swift @@ -102,22 +162,9 @@ public struct ArrayAtIndices { /// ### See Also /// - ``MLXArray/at`` public func add(_ values: some ScalarOrArray) -> MLXArray { - let values = values.asMLXArray(dtype: array.dtype) - let (indices, update, axes) = scatterArguments( - src: array, operations: indexOperations, update: values, stream: stream) - - if !indices.isEmpty { - let indices_vector = new_mlx_vector_array(indices) - defer { mlx_vector_array_free(indices_vector) } - - var result = mlx_array_new() - mlx_scatter_add( - &result, array.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx) - - return MLXArray(result) - } else { - return array + update - } + update( + values, sliceUpdate: mlx_slice_update_add, scatter: mlx_scatter_add, + elementwise: { $0 + $1 }) } /// Subtract values via `at[]` operator. @@ -145,22 +192,9 @@ public struct ArrayAtIndices { /// ### See Also /// - ``MLXArray/at`` public func multiply(_ values: some ScalarOrArray) -> MLXArray { - let values = values.asMLXArray(dtype: array.dtype) - let (indices, update, axes) = scatterArguments( - src: array, operations: indexOperations, update: values, stream: stream) - - if !indices.isEmpty { - let indices_vector = new_mlx_vector_array(indices) - defer { mlx_vector_array_free(indices_vector) } - - var result = mlx_array_new() - mlx_scatter_prod( - &result, array.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx) - - return MLXArray(result) - } else { - return array * update - } + update( + values, sliceUpdate: mlx_slice_update_prod, scatter: mlx_scatter_prod, + elementwise: { $0 * $1 }) } /// Divide values via `at[]` operator. @@ -188,22 +222,9 @@ public struct ArrayAtIndices { /// ### See Also /// - ``MLXArray/at`` public func minimum(_ values: some ScalarOrArray) -> MLXArray { - let values = values.asMLXArray(dtype: array.dtype) - let (indices, update, axes) = scatterArguments( - src: array, operations: indexOperations, update: values, stream: stream) - - if !indices.isEmpty { - let indices_vector = new_mlx_vector_array(indices) - defer { mlx_vector_array_free(indices_vector) } - - var result = mlx_array_new() - mlx_scatter_min( - &result, array.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx) - - return MLXArray(result) - } else { - return MLX.minimum(array, update) - } + update( + values, sliceUpdate: mlx_slice_update_min, scatter: mlx_scatter_min, + elementwise: { MLX.minimum($0, $1) }) } /// Update to maximum values via `at[]` operator. @@ -217,22 +238,9 @@ public struct ArrayAtIndices { /// ### See Also /// - ``MLXArray/at`` public func maximum(_ values: some ScalarOrArray) -> MLXArray { - let values = values.asMLXArray(dtype: array.dtype) - let (indices, update, axes) = scatterArguments( - src: array, operations: indexOperations, update: values, stream: stream) - - if !indices.isEmpty { - let indices_vector = new_mlx_vector_array(indices) - defer { mlx_vector_array_free(indices_vector) } - - var result = mlx_array_new() - mlx_scatter_max( - &result, array.ctx, indices_vector, update.ctx, axes, axes.count, stream.ctx) - - return MLXArray(result) - } else { - return MLX.maximum(array, update) - } + update( + values, sliceUpdate: mlx_slice_update_max, scatter: mlx_scatter_max, + elementwise: { MLX.maximum($0, $1) }) } } diff --git a/Source/MLX/Cmlx+Util.swift b/Source/MLX/Cmlx+Util.swift index 61854afa0..9279a74fe 100644 --- a/Source/MLX/Cmlx+Util.swift +++ b/Source/MLX/Cmlx+Util.swift @@ -179,3 +179,7 @@ func new_mlx_kwargs_closure(keys: [String], _ f: @escaping ([MLXArray]) -> [MLXA return mlx_closure_kwargs_new_func_payload(trampoline, payload, free) } + +func mlx_optional_dtype(_ dtype: DType?) -> mlx_optional_dtype { + .init(value: dtype?.cmlxDtype ?? MLX_FLOAT16, has_value: dtype != nil) +} diff --git a/Source/MLX/Documentation.docc/Articles/converting-python.md b/Source/MLX/Documentation.docc/Articles/converting-python.md index fa82f3094..6c94c240a 100644 --- a/Source/MLX/Documentation.docc/Articles/converting-python.md +++ b/Source/MLX/Documentation.docc/Articles/converting-python.md @@ -171,8 +171,8 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts. `cosh` | ``MLX/cosh(_:stream:)`` `cummax` | ``MLX/cummax(_:axis:reverse:inclusive:stream:)`` `cummin` | ``MLX/cummin(_:axis:reverse:inclusive:stream:)`` -`cumprod` | ``MLX/cumprod(_:axis:reverse:inclusive:stream:)`` -`cumsum` | ``MLX/cumsum(_:axis:reverse:inclusive:stream:)`` +`cumprod` | ``MLX/cumprod(_:axis:reverse:inclusive:dtype:stream:)`` +`cumsum` | ``MLX/cumsum(_:axis:reverse:inclusive:dtype:stream:)`` `dequantize` | ``MLX/dequantized(_:scales:biases:groupSize:bits:mode:globalScale:dtype:stream:)`` `divide` | ``MLX/divide(_:_:stream:)`` `equal` | ``MLX/equal(_:_:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/arithmetic.md b/Source/MLX/Documentation.docc/Organization/arithmetic.md index 7023b95a6..2a5851293 100644 --- a/Source/MLX/Documentation.docc/Organization/arithmetic.md +++ b/Source/MLX/Documentation.docc/Organization/arithmetic.md @@ -178,6 +178,7 @@ Note: the `-` and `/` operators are not able to be linked here. - ``logicalAnd(_:_:stream:)`` - ``logicalNot(_:stream:)`` - ``logicalOr(_:_:stream:)`` +- ``logicalXor(_:_:stream:)`` - ``maximum(_:_:stream:)`` - ``minimum(_:_:stream:)`` - ``multiply(_:_:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/cumulative.md b/Source/MLX/Documentation.docc/Organization/cumulative.md index 3d20ae6d0..b9c4aeb6c 100644 --- a/Source/MLX/Documentation.docc/Organization/cumulative.md +++ b/Source/MLX/Documentation.docc/Organization/cumulative.md @@ -38,9 +38,9 @@ These are available as both methods on `MLXArray` and free functions. They each - ``cummax(_:reverse:inclusive:stream:)`` - ``cummin(_:axis:reverse:inclusive:stream:)`` - ``cummin(_:reverse:inclusive:stream:)`` -- ``cumprod(_:axis:reverse:inclusive:stream:)`` -- ``cumprod(_:reverse:inclusive:stream:)`` -- ``cumsum(_:axis:reverse:inclusive:stream:)`` -- ``cumsum(_:reverse:inclusive:stream:)`` +- ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` +- ``cumprod(_:reverse:inclusive:dtype:stream:)`` +- ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` +- ``cumsum(_:reverse:inclusive:dtype:stream:)`` - ``logCumsumExp(_:axis:reverse:inclusive:stream:)`` - ``logCumsumExp(_:reverse:inclusive:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/initialization.md b/Source/MLX/Documentation.docc/Organization/initialization.md index 4b6fb0ba7..447ba6ed2 100644 --- a/Source/MLX/Documentation.docc/Organization/initialization.md +++ b/Source/MLX/Documentation.docc/Organization/initialization.md @@ -275,15 +275,15 @@ there are specific initializers to request it: ### MLXArray Factory Free Methods - ``MLX/zeros(_:type:stream:)`` -- ``MLX/zeros(like:stream:)`` +- ``MLX/zeros(like:dtype:stream:)`` - ``MLX/ones(_:type:stream:)`` -- ``MLX/ones(like:stream:)`` +- ``MLX/ones(like:dtype:stream:)`` - ``MLX/eye(_:m:k:type:stream:)`` - ``MLX/full(_:values:type:stream:)`` - ``MLX/full(_:values:stream:)`` - ``MLX/identity(_:type:stream:)`` -- ``MLX/linspace(_:_:count:stream:)-7vj0o`` -- ``MLX/linspace(_:_:count:stream:)-6w959`` +- ``linspace(_:_:count:endpoint:stream:)-47tc1`` +- ``linspace(_:_:count:endpoint:stream:)-38sfd`` - ``MLXArray/repeated(_:count:axis:stream:)`` - ``MLXArray/repeated(_:count:stream:)`` - ``MLX/repeat(_:count:axis:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/logical.md b/Source/MLX/Documentation.docc/Organization/logical.md index ce9b1d4d4..4704cbfc6 100644 --- a/Source/MLX/Documentation.docc/Organization/logical.md +++ b/Source/MLX/Documentation.docc/Organization/logical.md @@ -59,5 +59,6 @@ if (a < b).all().item() { - ``logicalAnd(_:_:stream:)`` - ``logicalNot(_:stream:)`` - ``logicalOr(_:_:stream:)`` +- ``logicalXor(_:_:stream:)`` - ``notEqual(_:_:stream:)`` - ``where(_:_:_:stream:)`` diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index e392e4282..e0592dbec 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -74,10 +74,10 @@ operations as methods for convenience. - ``cummax(_:reverse:inclusive:stream:)`` - ``cummin(_:axis:reverse:inclusive:stream:)`` - ``cummin(_:reverse:inclusive:stream:)`` -- ``cumprod(_:axis:reverse:inclusive:stream:)`` -- ``cumprod(_:reverse:inclusive:stream:)`` -- ``cumsum(_:axis:reverse:inclusive:stream:)`` -- ``cumsum(_:reverse:inclusive:stream:)`` +- ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` +- ``cumprod(_:reverse:inclusive:dtype:stream:)`` +- ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` +- ``cumsum(_:reverse:inclusive:dtype:stream:)`` ### Indexes @@ -99,14 +99,15 @@ operations as methods for convenience. ### Factory - ``MLX/zeros(_:type:stream:)`` -- ``MLX/zeros(like:stream:)`` +- ``MLX/zeros(like:dtype:stream:)`` - ``MLX/ones(_:type:stream:)`` -- ``MLX/ones(like:stream:)`` +- ``MLX/ones(like:dtype:stream:)`` - ``MLX/eye(_:m:k:type:stream:)`` - ``MLX/full(_:values:type:stream:)`` - ``MLX/full(_:values:stream:)`` - ``MLX/identity(_:type:stream:)`` -- ``MLX/linspace(_:_:count:stream:)-6w959`` +- ``linspace(_:_:count:endpoint:stream:)-47tc1`` +- ``linspace(_:_:count:endpoint:stream:)-38sfd`` - ``MLX/repeated(_:count:axis:stream:)`` - ``MLX/repeated(_:count:stream:)`` - ``MLX/repeat(_:count:axis:stream:)`` diff --git a/Source/MLX/FFT.swift b/Source/MLX/FFT.swift index 988385b43..1a8c4477c 100644 --- a/Source/MLX/FFT.swift +++ b/Source/MLX/FFT.swift @@ -3,6 +3,32 @@ import Cmlx import Foundation +/// Normalization mode for discrete Fourier Transforms. +/// +/// ### See Also +/// - +public enum FFTNorm: Sendable { + /// No normalization on the forward transform; the inverse transform is scaled by `1/n`. + /// + /// This is the default and matches the behavior of prior versions. + case backward + + /// Both the forward and inverse transforms are scaled by `1/sqrt(n)`, making the + /// transform unitary. + case ortho + + /// The forward transform is scaled by `1/n`; no normalization on the inverse transform. + case forward + + var ctx: mlx_fft_norm { + switch self { + case .backward: MLX_FFT_NORM_BACKWARD + case .ortho: MLX_FFT_NORM_ORTHO + case .forward: MLX_FFT_NORM_FORWARD + } + } +} + public enum MLXFFT { /// One dimensional discrete Fourier Transform. @@ -12,20 +38,23 @@ public enum MLXFFT { /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axis /// /// ### See Also /// - public static func fft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() mlx_fft_fftn( - &result, array.ctx, [(n ?? array.dim(axis)).int32], 1, [axis.int32], 1, stream.ctx) + &result, array.ctx, [(n ?? array.dim(axis)).int32], 1, [axis.int32], 1, norm.ctx, + stream.ctx) return MLXArray(result) } @@ -36,16 +65,19 @@ public enum MLXFFT { /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axis /// /// ### See Also /// - public static func ifft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_fft_ifft(&result, array.ctx, (n ?? array.dim(axis)).int32, axis.int32, stream.ctx) + mlx_fft_ifft( + &result, array.ctx, (n ?? array.dim(axis)).int32, axis.int32, norm.ctx, stream.ctx) return MLXArray(result) } @@ -56,6 +88,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes /// @@ -63,9 +96,10 @@ public enum MLXFFT { /// - public static func fft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - fftn(array, s: s, axes: axes, stream: stream) + fftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// Two dimensional inverse discrete Fourier Transform. @@ -75,6 +109,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axes /// @@ -82,9 +117,10 @@ public enum MLXFFT { /// - public static func ifft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - ifftn(array, s: s, axes: axes, stream: stream) + ifftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// n-dimensional discrete Fourier Transform. @@ -94,6 +130,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes /// @@ -101,35 +138,40 @@ public enum MLXFFT { /// - public static func fftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() if let s, let axes { // both supplied mlx_fft_fft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let axes { // no n, compute from dim() let n = axes.map { array.dim($0) } mlx_fft_fft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let s { // axes are the rightmost dimensions matching the number of dimensions of n let axes = Array(-s.count ..< 0) mlx_fft_fft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else { let axes = Array(0 ..< array.ndim) let n = axes.map { array.dim($0) } mlx_fft_fft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } } @@ -141,6 +183,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axes /// @@ -148,35 +191,40 @@ public enum MLXFFT { /// - public static func ifftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() if let s, let axes { // both supplied mlx_fft_ifft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let axes { // no n, compute from dim() let n = axes.map { array.dim($0) } mlx_fft_ifft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let s { // axes are the rightmost dimensions matching the number of dimensions of n let axes = Array(-s.count ..< 0) mlx_fft_ifft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else { let axes = Array(0 ..< array.ndim) let n = axes.map { array.dim($0) } mlx_fft_ifft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } } @@ -191,17 +239,20 @@ public enum MLXFFT { /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axis. The output data type will be complex. /// /// ### See Also /// - public static func rfft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() mlx_fft_rfftn( - &result, array.ctx, [(n ?? array.dim(axis)).int32], 1, [axis.int32], 1, stream.ctx) + &result, array.ctx, [(n ?? array.dim(axis)).int32], 1, [axis.int32], 1, norm.ctx, + stream.ctx) return MLXArray(result) } @@ -215,17 +266,19 @@ public enum MLXFFT { /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n / 2 + 1`. If not specified `array.dim(axis) / 2 + 1` will be used. /// - axis: axis along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on - /// - Returns: inverse of ``rfft(_:n:axis:stream:)`` + /// - Returns: inverse of ``rfft(_:n:axis:norm:stream:)`` /// /// ### See Also /// - public static func irfft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { let n = n ?? (array.dim(axis) - 1) * 2 var result = mlx_array_new() - mlx_fft_irfft(&result, array.ctx, n.int32, axis.int32, stream.ctx) + mlx_fft_irfft(&result, array.ctx, n.int32, axis.int32, norm.ctx, stream.ctx) return MLXArray(result) } @@ -240,6 +293,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes. The output data type will be complex. /// @@ -247,9 +301,10 @@ public enum MLXFFT { /// - public static func rfft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - rfftn(array, s: s, axes: axes, stream: stream) + rfftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// Inverse two dimensional discrete Fourier Transform on a real input. @@ -264,16 +319,18 @@ public enum MLXFFT { /// match `n / 2 + 1`. If not specified `array.dim(axis) / 2 + 1` will be used. /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on - /// - Returns: inverse of ``rfft2(_:s:axes:stream:)`` + /// - Returns: inverse of ``rfft2(_:s:axes:norm:stream:)`` /// /// ### See Also /// - public static func irfft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - irfftn(array, s: s, axes: axes, stream: stream) + irfftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// n-dimensional real discrete Fourier Transform. @@ -287,6 +344,7 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes. The output data type will be complex. /// @@ -294,35 +352,40 @@ public enum MLXFFT { /// - public static func rfftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() if let s, let axes { // both supplied mlx_fft_rfft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let axes { // no n, compute from dim() let n = axes.map { array.dim($0) } mlx_fft_rfft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let s { // axes are the rightmost dimensions matching the number of dimensions of n let axes = Array(-s.count ..< 0) mlx_fft_rfft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else { let axes = Array(0 ..< array.ndim) let n = axes.map { array.dim($0) } mlx_fft_rfft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } } @@ -339,21 +402,24 @@ public enum MLXFFT { /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `n / 2 + 1`. If not specified `array.dim(axis) / 2 + 1` will be used. /// - axes: axes along which to perform the FFT + /// - norm: normalization mode /// - stream: stream or device to evaluate on - /// - Returns: inverse of ``rfftn(_:s:axes:stream:)`` + /// - Returns: inverse of ``rfftn(_:s:axes:norm:stream:)`` /// /// ### See Also /// - public static func irfftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() if let s, let axes { // both supplied mlx_fft_irfft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let axes { // no n, compute from dim() @@ -361,14 +427,16 @@ public enum MLXFFT { n[n.count - 1] = (n[n.count - 1] - 1) * 2 mlx_fft_irfft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else if let s { // axes are the rightmost dimensions matching the number of dimensions of n let axes = Array(-s.count ..< 0) mlx_fft_irfft2( - &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, s.asInt32, s.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } else { let axes = Array(0 ..< array.ndim) @@ -376,11 +444,56 @@ public enum MLXFFT { n[n.count - 1] = (n[n.count - 1] - 1) * 2 mlx_fft_irfft2( - &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, stream.ctx) + &result, array.ctx, n.asInt32, n.count, axes.asInt32, axes.count, norm.ctx, + stream.ctx) return MLXArray(result) } } + /// The sample frequencies of a discrete Fourier Transform. + /// + /// Returns `n` frequencies in cycles per unit of the sample spacing `d`, ordered + /// to match the output of ``fft(_:n:axis:norm:stream:)``. + /// + /// - Parameters: + /// - n: window length + /// - d: sample spacing + /// - stream: stream or device to evaluate on + /// - Returns: array of sample frequencies + /// + /// ### See Also + /// - + /// - ``rfftfreq(_:d:stream:)`` + public static func fftfreq(_ n: Int, d: Double = 1.0, stream: StreamOrDevice = .default) + -> MLXArray + { + var result = mlx_array_new() + mlx_fft_fftfreq(&result, n.int32, d, stream.ctx) + return MLXArray(result) + } + + /// The sample frequencies of a discrete Fourier Transform on a real input. + /// + /// Like ``fftfreq(_:d:stream:)`` but returns only the `n / 2 + 1` non-negative + /// frequencies, matching the output of ``rfft(_:n:axis:norm:stream:)``. + /// + /// - Parameters: + /// - n: window length + /// - d: sample spacing + /// - stream: stream or device to evaluate on + /// - Returns: array of non-negative sample frequencies + /// + /// ### See Also + /// - + /// - ``fftfreq(_:d:stream:)`` + public static func rfftfreq(_ n: Int, d: Double = 1.0, stream: StreamOrDevice = .default) + -> MLXArray + { + var result = mlx_array_new() + mlx_fft_rfftfreq(&result, n.int32, d, stream.ctx) + return MLXArray(result) + } + /// Shift the zero-frequency component to the center of the spectrum. /// /// - Parameters: @@ -431,15 +544,19 @@ public enum MLXFFT { /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axis /// /// ### See Also /// - -public func fft(_ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default) +public func fft( + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default +) -> MLXArray { - MLXFFT.fft(array, n: n, axis: axis, stream: stream) + MLXFFT.fft(array, n: n, axis: axis, norm: norm, stream: stream) } /// One dimensional inverse discrete Fourier Transform. @@ -449,15 +566,17 @@ public func fft(_ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: Stream /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axis /// /// ### See Also /// - public func ifft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.ifft(array, n: n, axis: axis, stream: stream) + MLXFFT.ifft(array, n: n, axis: axis, norm: norm, stream: stream) } /// Two dimensional discrete Fourier Transform. @@ -467,6 +586,7 @@ public func ifft( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes /// @@ -474,9 +594,10 @@ public func ifft( /// - public func fft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.fft2(array, s: s, axes: axes, stream: stream) + MLXFFT.fft2(array, s: s, axes: axes, norm: norm, stream: stream) } /// Two dimensional inverse discrete Fourier Transform. @@ -486,6 +607,7 @@ public func fft2( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axes /// @@ -493,9 +615,10 @@ public func fft2( /// - public func ifft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.ifft2(array, s: s, axes: axes, stream: stream) + MLXFFT.ifft2(array, s: s, axes: axes, norm: norm, stream: stream) } /// n-dimensional discrete Fourier Transform. @@ -505,6 +628,7 @@ public func ifft2( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axis)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes /// @@ -512,9 +636,10 @@ public func ifft2( /// - public func fftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.fftn(array, s: s, axes: axes, stream: stream) + MLXFFT.fftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// n-dimensional inverse discrete Fourier Transform. @@ -524,6 +649,7 @@ public func fftn( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: inverse DFT of the input along the given axes /// @@ -531,9 +657,10 @@ public func fftn( /// - public func ifftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.ifftn(array, s: s, axes: axes, stream: stream) + MLXFFT.ifftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// One dimensional discrete Fourier Transform on a real input. @@ -546,15 +673,17 @@ public func ifftn( /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n`. If not specified `array.dim(axis)` will be used. /// - axis: axis along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axis. The output data type will be complex. /// /// ### See Also /// - public func rfft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.rfft(array, n: n, axis: axis, stream: stream) + MLXFFT.rfft(array, n: n, axis: axis, norm: norm, stream: stream) } /// Inverse one dimensional discrete Fourier Transform on a real input. @@ -567,15 +696,17 @@ public func rfft( /// - n: size of the transformed axis. The corresponding axis in the input is truncated or padded with zeros to /// match `n / 2 + 1`. If not specified `array.dim(axis) / 2 + 1` will be used. /// - axis: axis along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on -/// - Returns: inverse of ``rfft(_:n:axis:stream:)`` +/// - Returns: inverse of ``rfft(_:n:axis:norm:stream:)`` /// /// ### See Also /// - public func irfft( - _ array: MLXArray, n: Int? = nil, axis: Int = -1, stream: StreamOrDevice = .default + _ array: MLXArray, n: Int? = nil, axis: Int = -1, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.irfft(array, n: n, axis: axis, stream: stream) + MLXFFT.irfft(array, n: n, axis: axis, norm: norm, stream: stream) } /// Two dimensional real discrete Fourier Transform. @@ -589,6 +720,7 @@ public func irfft( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes. The output data type will be complex. /// @@ -596,9 +728,10 @@ public func irfft( /// - public func rfft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], stream: StreamOrDevice = .default + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.rfft2(array, s: s, axes: axes, stream: stream) + MLXFFT.rfft2(array, s: s, axes: axes, norm: norm, stream: stream) } /// Inverse two dimensional discrete Fourier Transform on a real input. @@ -613,17 +746,18 @@ public func rfft2( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on -/// - Returns: inverse of ``rfft2(_:s:axes:stream:)`` +/// - Returns: inverse of ``rfft2(_:s:axes:norm:stream:)`` /// /// ### See Also /// - public func irfft2( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [-2, -1], + axes: (some Collection)? = [-2, -1], norm: FFTNorm = .backward, stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.irfft2(array, s: s, axes: axes, stream: stream) + MLXFFT.irfft2(array, s: s, axes: axes, norm: norm, stream: stream) } /// n-dimensional real discrete Fourier Transform. @@ -637,6 +771,7 @@ public func irfft2( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on /// - Returns: DFT of the input along the given axes. The output data type will be complex. /// @@ -644,9 +779,10 @@ public func irfft2( /// - public func rfftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.rfftn(array, s: s, axes: axes, stream: stream) + MLXFFT.rfftn(array, s: s, axes: axes, norm: norm, stream: stream) } /// Inverse n-dimensional discrete Fourier Transform on a real input. @@ -661,16 +797,54 @@ public func rfftn( /// - s: sizes of the transformed axis. The corresponding axes in the input are truncated or padded with zeros to /// match `s`. If not specified `array.dim(axes)` will be used. /// - axes: axes along which to perform the FFT +/// - norm: normalization mode /// - stream: stream or device to evaluate on -/// - Returns: inverse of ``rfftn(_:s:axes:stream:)`` +/// - Returns: inverse of ``rfftn(_:s:axes:norm:stream:)`` /// /// ### See Also /// - public func irfftn( _ array: MLXArray, s: (some Collection)? = [Int]?.none, - axes: (some Collection)? = [Int]?.none, stream: StreamOrDevice = .default + axes: (some Collection)? = [Int]?.none, norm: FFTNorm = .backward, + stream: StreamOrDevice = .default ) -> MLXArray { - MLXFFT.irfftn(array, s: s, axes: axes, stream: stream) + MLXFFT.irfftn(array, s: s, axes: axes, norm: norm, stream: stream) +} + +/// The sample frequencies of a discrete Fourier Transform. +/// +/// Returns `n` frequencies in cycles per unit of the sample spacing `d`, ordered +/// to match the output of ``fft(_:n:axis:norm:stream:)``. +/// +/// - Parameters: +/// - n: window length +/// - d: sample spacing +/// - stream: stream or device to evaluate on +/// - Returns: array of sample frequencies +/// +/// ### See Also +/// - +/// - ``rfftfreq(_:d:stream:)`` +public func fftfreq(_ n: Int, d: Double = 1.0, stream: StreamOrDevice = .default) -> MLXArray { + MLXFFT.fftfreq(n, d: d, stream: stream) +} + +/// The sample frequencies of a discrete Fourier Transform on a real input. +/// +/// Like ``fftfreq(_:d:stream:)`` but returns only the `n / 2 + 1` non-negative +/// frequencies, matching the output of ``rfft(_:n:axis:norm:stream:)``. +/// +/// - Parameters: +/// - n: window length +/// - d: sample spacing +/// - stream: stream or device to evaluate on +/// - Returns: array of non-negative sample frequencies +/// +/// ### See Also +/// - +/// - ``fftfreq(_:d:stream:)`` +public func rfftfreq(_ n: Int, d: Double = 1.0, stream: StreamOrDevice = .default) -> MLXArray { + MLXFFT.rfftfreq(n, d: d, stream: stream) } /// Shift the zero-frequency component to the center of the spectrum. diff --git a/Source/MLX/Factory.swift b/Source/MLX/Factory.swift index 913540511..0abbae72d 100644 --- a/Source/MLX/Factory.swift +++ b/Source/MLX/Factory.swift @@ -20,7 +20,7 @@ extension MLXArray { /// /// ### See Also /// - - /// - ``zeros(like:stream:)`` + /// - ``zeros(like:dtype:stream:)`` /// - ``ones(_:type:stream:)`` static public func zeros( _ shape: some Collection, type: (some HasDType).Type = Float.self, @@ -44,7 +44,7 @@ extension MLXArray { /// /// ### See Also /// - - /// - ``zeros(like:stream:)`` + /// - ``zeros(like:dtype:stream:)`` /// - ``ones(_:type:stream:)`` static public func zeros( _ shape: some Collection, dtype: DType, stream: StreamOrDevice = .default @@ -63,14 +63,17 @@ extension MLXArray { /// /// - Parameters: /// - array: array to copy shape and dtype from + /// - dtype: optional type to produce the result in, otherwise use array's dtype /// - stream: stream or device to evaluate on /// /// ### See Also /// - /// - ``zeros(_:type:stream:)`` /// - ``ones(_:type:stream:)`` - static public func zeros(like array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { - MLX.zeros(like: array, stream: stream) + static public func zeros( + like array: MLXArray, dtype: DType? = nil, stream: StreamOrDevice = .default + ) -> MLXArray { + MLX.zeros(like: array, dtype: dtype, stream: stream) } /// Construct an array of ones. @@ -88,7 +91,7 @@ extension MLXArray { /// /// ### See Also /// - - /// - ``ones(like:stream:)`` + /// - ``ones(like:dtype:stream:)`` /// - ``zeros(_:type:stream:)`` static public func ones( _ shape: some Collection, type: (some HasDType).Type = Float.self, @@ -112,7 +115,7 @@ extension MLXArray { /// /// ### See Also /// - - /// - ``ones(like:stream:)`` + /// - ``ones(like:dtype:stream:)`` /// - ``zeros(_:type:stream:)`` static public func ones( _ shape: some Collection, dtype: DType, stream: StreamOrDevice = .default @@ -131,14 +134,17 @@ extension MLXArray { /// /// - Parameters: /// - array: array to copy shape and dtype from + /// - dtype: optional type to produce the result in, otherwise use array's dtype /// - stream: stream or device to evaluate on /// /// ### See Also /// - /// - ``ones(_:type:stream:)`` /// - ``zeros(_:type:stream:)`` - static public func ones(like array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { - MLX.ones(like: array, stream: stream) + static public func ones( + like array: MLXArray, dtype: DType? = nil, stream: StreamOrDevice = .default + ) -> MLXArray { + MLX.ones(like: array, dtype: dtype, stream: stream) } /// Create an identity matrix or a general diagonal matrix. @@ -331,7 +337,7 @@ extension MLXArray { /// /// ```swift /// // Create a 50 element 1-D array with values from 0 to 50 - /// let r = MLXArray.linSpace(0, 50) + /// let r = MLXArray.linspace(0, 50) /// ``` /// /// - Parameters: @@ -355,7 +361,7 @@ extension MLXArray { /// /// ```swift /// // Create a 50 element 1-D array with values from 0 to 1 - /// let r = MLXArray.linSpace(0.0, 1.0) + /// let r = MLXArray.linspace(0.0, 1.0) /// ``` /// /// - Parameters: @@ -659,7 +665,7 @@ extension MLXArray { /// /// ### See Also /// - -/// - ``zeros(like:stream:)`` +/// - ``zeros(like:dtype:stream:)`` /// - ``ones(_:type:stream:)`` public func zeros( _ shape: some Collection, type: (some HasDType).Type = Float.self, @@ -685,7 +691,7 @@ public func zeros( /// /// ### See Also /// - -/// - ``zeros(like:stream:)`` +/// - ``zeros(like:dtype:stream:)`` /// - ``ones(_:type:stream:)`` public func zeros( _ shape: some Collection, dtype: DType, stream: StreamOrDevice = .default @@ -706,15 +712,22 @@ public func zeros( /// /// - Parameters: /// - array: array to copy shape and dtype from +/// - dtype: optional type to produce the result in, otherwise use array's dtype /// - stream: stream or device to evaluate on /// /// ### See Also /// - /// - ``zeros(_:type:stream:)`` /// - ``ones(_:type:stream:)`` -public func zeros(like array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { +public func zeros( + like array: MLXArray, dtype: DType? = nil, stream: StreamOrDevice = .default +) -> MLXArray { var result = mlx_array_new() - mlx_zeros_like(&result, array.ctx, stream.ctx) + if let dtype { + mlx_zeros_like_dtype(&result, array.ctx, dtype.cmlxDtype, stream.ctx) + } else { + mlx_zeros_like(&result, array.ctx, stream.ctx) + } return MLXArray(result) } @@ -733,7 +746,7 @@ public func zeros(like array: MLXArray, stream: StreamOrDevice = .default) -> ML /// /// ### See Also /// - -/// - ``ones(like:stream:)`` +/// - ``ones(like:dtype:stream:)`` /// - ``zeros(_:type:stream:)`` public func ones( _ shape: some Collection, type: (some HasDType).Type = Float.self, @@ -759,7 +772,7 @@ public func ones( /// /// ### See Also /// - -/// - ``zeros(like:stream:)`` +/// - ``zeros(like:dtype:stream:)`` /// - ``ones(_:type:stream:)`` public func ones( _ shape: some Collection, dtype: DType, stream: StreamOrDevice = .default @@ -780,15 +793,22 @@ public func ones( /// /// - Parameters: /// - array: array to copy shape and dtype from +/// - dtype: optional type to produce the result in, otherwise use array's dtype /// - stream: stream or device to evaluate on /// /// ### See Also /// - /// - ``ones(_:type:stream:)`` /// - ``zeros(_:type:stream:)`` -public func ones(like array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { +public func ones( + like array: MLXArray, dtype: DType? = nil, stream: StreamOrDevice = .default +) -> MLXArray { var result = mlx_array_new() - mlx_ones_like(&result, array.ctx, stream.ctx) + if let dtype { + mlx_ones_like_dtype(&result, array.ctx, dtype.cmlxDtype, stream.ctx) + } else { + mlx_ones_like(&result, array.ctx, stream.ctx) + } return MLXArray(result) } @@ -993,23 +1013,27 @@ public func identity(_ n: Int, dtype: DType, stream: StreamOrDevice = .default) /// /// ```swift /// // Create a 50 element 1-D array with values from 0 to 50 -/// let r = MLXArray.linSpace(0, 50) +/// let r = MLXArray.linspace(0, 50) /// ``` /// /// - Parameters: /// - start: start value /// - stop: stop value /// - count: number of samples +/// - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval /// - stream: stream or device to evaluate on /// /// ### See Also /// - -/// - ``linspace(_:_:count:stream:)-7vj0o`` +/// - ``linspace(_:_:count:endpoint:stream:)`` public func linspace( - _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default + _ start: T, _ stop: T, count: Int = 50, + endpoint: Bool = true, + stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryInteger { var result = mlx_array_new() - mlx_linspace(&result, Double(start), Double(stop), count.int32, T.dtype.cmlxDtype, stream.ctx) + mlx_linspace_endpoint( + &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx) return MLXArray(result) } @@ -1019,23 +1043,27 @@ public func linspace( /// /// ```swift /// // Create a 50 element 1-D array with values from 0 to 1 -/// let r = MLXArray.linSpace(0.0, 1.0) +/// let r = MLXArray.linspace(0.0, 1.0) /// ``` /// /// - Parameters: /// - start: start value /// - stop: stop value /// - count: number of samples +/// - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval /// - stream: stream or device to evaluate on /// /// ### See Also /// - -/// - ``linspace(_:_:count:stream:)-6w959`` +/// - ``linspace(_:_:count:endpoint:stream:)`` public func linspace( - _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default + _ start: T, _ stop: T, count: Int = 50, + endpoint: Bool = true, + stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryFloatingPoint { var result = mlx_array_new() - mlx_linspace(&result, Double(start), Double(stop), count.int32, T.dtype.cmlxDtype, stream.ctx) + mlx_linspace_endpoint( + &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx) return MLXArray(result) } diff --git a/Source/MLX/GPU+Metal.swift b/Source/MLX/GPU+Metal.swift index c46440817..cc4e83c89 100644 --- a/Source/MLX/GPU+Metal.swift +++ b/Source/MLX/GPU+Metal.swift @@ -51,6 +51,27 @@ public enum GPU { Memory.peakMemory } + /// The metallib override path. This can be set early in the process to override or provide + /// the metallib load path. + public static var metallib: URL? { + get { + var s = mlx_string_new() + defer { mlx_string_free(s) } + + mlx_metal_get_metallib_path(&s) + + if let path = String(cString: mlx_string_data(s), encoding: .utf8), !path.isEmpty { + return URL(fileURLWithPath: path) + } else { + return nil + } + } + set { + let path = newValue?.path(percentEncoded: false) ?? "" + mlx_metal_set_metallib_path(path.cString(using: .utf8)) + } + } + /// Return a snapshot of memory stats -- see ``Memory/Snapshot`` for more details. /// /// Get the current memory use. This can be used to measure before/after and current memory use: diff --git a/Source/MLX/GraphUtils.swift b/Source/MLX/GraphUtils.swift new file mode 100644 index 000000000..ada2e2134 --- /dev/null +++ b/Source/MLX/GraphUtils.swift @@ -0,0 +1,122 @@ +// Copyright © 2025 Apple Inc. + +import Cmlx +import Foundation + +/// Assigns human readable names to arrays in a compute graph. +/// +/// Arrays that aren't given a name are numbered automatically. Pass a namer to +/// ``exportToDot(_:namer:)`` or ``graphDescription(_:namer:)`` to make their output easier +/// to read: +/// +/// ```swift +/// let x = MLXArray([1, 2, 3]) +/// let y = MLXArray([4, 5, 6]) +/// +/// let namer = NodeNamer() +/// namer.setName("x", for: x) +/// namer.setName("y", for: y) +/// +/// print(exportToDot([x + y], namer: namer)) +/// ``` +/// +/// ### See Also +/// - ``exportToDot(_:namer:)`` +/// - ``graphDescription(_:namer:)`` +public final class NodeNamer { + + let ctx: mlx_node_namer + + /// Create an empty namer. + public init() { + ctx = mlx_node_namer_new() + } + + deinit { + mlx_node_namer_free(ctx) + } + + /// Assign `name` to `array`. + public func setName(_ name: String, for array: MLXArray) { + mlx_node_namer_set_name(ctx, array.ctx, name) + } + + /// The name of `array`. + /// + /// Arrays without an assigned name are given a generated one, which is recorded so that + /// subsequent calls return the same value. + public func name(for array: MLXArray) -> String? { + var name: UnsafePointer? + mlx_node_namer_get_name(&name, ctx, array.ctx) + guard let name else { return nil } + return String(cString: name, encoding: .utf8) + } +} + +/// Capture everything written to a `FILE *` as a `String`. +private func captureFileOutput(_ body: (UnsafeMutablePointer) -> Void) -> String { + var buffer: UnsafeMutablePointer? + var size = 0 + + guard let file = open_memstream(&buffer, &size) else { return "" } + body(file) + + // closing flushes and finalizes buffer/size + fclose(file) + + guard let buffer else { return "" } + defer { free(buffer) } + + return String(cString: buffer, encoding: .utf8) ?? "" +} + +/// A [GraphViz](https://graphviz.org) `DOT` representation of the compute graph that produces +/// `outputs`. +/// +/// The arrays are not evaluated -- this describes the graph as it currently stands. +/// +/// ```swift +/// let x = MLXArray([1, 2, 3]) +/// let dot = exportToDot([x * 2 + 1]) +/// try dot.write(to: url, atomically: true, encoding: .utf8) +/// ``` +/// +/// - Parameters: +/// - outputs: the outputs of the graph to describe +/// - namer: optional namer providing readable names for arrays +/// - Returns: the graph in `DOT` format +/// +/// ### See Also +/// - ``graphDescription(_:namer:)`` +/// - ``NodeNamer`` +public func exportToDot(_ outputs: [MLXArray], namer: NodeNamer? = nil) -> String { + let namer = namer ?? NodeNamer() + let vector = new_mlx_vector_array(outputs) + defer { mlx_vector_array_free(vector) } + + return captureFileOutput { file in + mlx_export_to_dot(file, namer.ctx, vector) + } +} + +/// A textual description of the compute graph that produces `outputs`. +/// +/// The arrays are not evaluated -- this describes the graph as it currently stands. +/// +/// - Parameters: +/// - outputs: the outputs of the graph to describe +/// - namer: optional namer providing readable names for arrays +/// - Returns: the description of the graph +/// +/// ### See Also +/// - ``exportToDot(_:namer:)`` +/// - ``NodeNamer`` +public func graphDescription(_ outputs: [MLXArray], namer: NodeNamer? = nil) -> String { + let namer = namer ?? NodeNamer() + let vector = new_mlx_vector_array(outputs) + defer { mlx_vector_array_free(vector) } + + return captureFileOutput { file in + mlx_print_graph(file, namer.ctx, vector) + } +} diff --git a/Source/MLX/IO.swift b/Source/MLX/IO.swift index fa84d2848..83c897d3c 100644 --- a/Source/MLX/IO.swift +++ b/Source/MLX/IO.swift @@ -214,7 +214,9 @@ private func new_mlx_io_vtable_dataIO() -> mlx_io_vtable { case SEEK_CUR: state.offset += Int(offset) case SEEK_END: - state.offset = state.offset - Int(offset) + // offset is relative to the end of the data, not the current position. + // mlx's load_safetensors uses seek(0, end) + tell() to size the input. + state.offset = state.data.count + Int(offset) default: break } diff --git a/Source/MLX/Linalg.swift b/Source/MLX/Linalg.swift index da889492e..3acc45719 100644 --- a/Source/MLX/Linalg.swift +++ b/Source/MLX/Linalg.swift @@ -528,6 +528,46 @@ public enum MLXLinalg { return MLXArray(result) } + /// Compute the determinant of a square matrix. + /// + /// This function supports arrays with at least 2 dimensions. When the + /// input has more than two dimensions, the determinant is computed for + /// each matrix in the last two dimensions. + /// + /// - Parameters: + /// - array: input array + /// - stream: stream or device to evaluate on + public static func det(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { + var result = mlx_array_new() + mlx_linalg_det(&result, array.ctx, stream.ctx) + return MLXArray(result) + } + + /// Compute the sign and natural log of the absolute value of the + /// determinant of a square matrix. + /// + /// This function supports arrays with at least 2 dimensions. When the + /// input has more than two dimensions, the sign and log-absolute-determinant + /// are computed for each matrix in the last two dimensions. + /// + /// For a singular matrix, `sign` is 0 and `logabsdet` is `-inf`. + /// + /// The determinant can be reconstructed as `det = sign * exp(logabsdet)`. + /// This is more numerically stable than computing the determinant directly + /// for matrices with large or small determinants. + /// + /// - Parameters: + /// - array: input array + /// - stream: stream or device to evaluate on + public static func slogdet( + _ array: MLXArray, stream: StreamOrDevice = .default + ) -> (sign: MLXArray, logabsdet: MLXArray) { + var sign = mlx_array_new() + var logabsdet = mlx_array_new() + mlx_linalg_slogdet(&sign, &logabsdet, array.ctx, stream.ctx) + return (MLXArray(sign), MLXArray(logabsdet)) + } + } // MLXLinalg /// Matrix or vector norm. @@ -925,3 +965,38 @@ public func eigvalsh( public func pinv(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { return MLXLinalg.pinv(array, stream: stream) } + +/// Compute the determinant of a square matrix. +/// +/// This function supports arrays with at least 2 dimensions. When the +/// input has more than two dimensions, the determinant is computed for +/// each matrix in the last two dimensions. +/// +/// - Parameters: +/// - array: input array +/// - stream: stream or device to evaluate on +public func det(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { + MLXLinalg.det(array, stream: stream) +} + +/// Compute the sign and natural log of the absolute value of the +/// determinant of a square matrix. +/// +/// This function supports arrays with at least 2 dimensions. When the +/// input has more than two dimensions, the sign and log-absolute-determinant +/// are computed for each matrix in the last two dimensions. +/// +/// For a singular matrix, `sign` is 0 and `logabsdet` is `-inf`. +/// +/// The determinant can be reconstructed as `det = sign * exp(logabsdet)`. +/// This is more numerically stable than computing the determinant directly +/// for matrices with large or small determinants. +/// +/// - Parameters: +/// - array: input array +/// - stream: stream or device to evaluate on +public func slogdet( + _ array: MLXArray, stream: StreamOrDevice = .default +) -> (sign: MLXArray, logabsdet: MLXArray) { + MLXLinalg.slogdet(array, stream: stream) +} diff --git a/Source/MLX/MLXArray+Indexing.swift b/Source/MLX/MLXArray+Indexing.swift index 9a891b5df..bebaf23d2 100644 --- a/Source/MLX/MLXArray+Indexing.swift +++ b/Source/MLX/MLXArray+Indexing.swift @@ -832,11 +832,32 @@ func gatherND( // MARK: - index set (slice) -func updateSlice( +/// The result of translating index operations into `mlx_slice_update*` arguments. +/// +/// See ``sliceUpdateArguments(src:operations:update:stream:)``. +enum SliceUpdateArguments { + /// The update applies to the slice described by `starts`/`ends`/`strides`. + case slice(update: MLXArray, starts: [Int32], ends: [Int32], strides: [Int32]) + + /// The operations select the entire array -- the update has been broadcast to `src`'s shape. + case broadcast(MLXArray) +} + +/// Translate index operations into arguments for the `mlx_slice_update*` family. +/// +/// Returns `nil` if the operations cannot be expressed as a slice update, e.g. if any of them +/// are arrays. Callers should fall back to a scatter in that case. +/// +/// - Parameters: +/// - src: the input array +/// - operations: array of index operations +/// - update: the update value +/// - Returns: the slice update arguments or `nil` if not applicable +func sliceUpdateArguments( src: MLXArray, operations: [MLXArrayIndexOperation], update: MLXArray, stream: StreamOrDevice = .default -) -> MLXArray? { - // See mlx_slice_update +) -> SliceUpdateArguments? { + // See mlx_compute_slice_update_args let ndim = src.ndim if ndim == 0 || operations.count == 0 { return nil @@ -855,19 +876,14 @@ func updateSlice( var ends = src.shape.asInt32 var strides = [Int32](repeating: 1, count: ndim) - // If it's just a simple slice, just do a slice update and return + // If it's just a simple slice, the params are complete if operations.count == 1, case .slice(let slice) = operations[0] { let size = src.dim(0).int32 starts[0] = slice.start(size) ends[0] = slice.end(size) strides[0] = slice.stride - var result = mlx_array_new() - mlx_slice_update( - &result, - src.ctx, update.ctx, starts, starts.count, ends, ends.count, strides, strides.count, - stream.ctx) - return MLXArray(result) + return .slice(update: update, starts: starts, ends: ends, strides: strides) } // Expand ellipses into a series of ':' (full slice) slices @@ -876,7 +892,7 @@ func updateSlice( // If no non-None indices return the broadcasted update let nonNewAxisOperationCount = countNonNewAxisOperations(operations) if nonNewAxisOperationCount == 0 { - return broadcast(update, to: src.shape) + return .broadcast(broadcast(update, to: src.shape)) } // Process entries @@ -898,7 +914,7 @@ func updateSlice( switch item { case .ellipsis, .array: // these were replaced or rejected earlier - fatalError("unexpected item \(item) in updateSlice") + fatalError("unexpected item \(item) in sliceUpdateArguments") case .index(let index): let size = src.dim(axis).int32 @@ -930,12 +946,31 @@ func updateSlice( update = reshaped(update, updateReshape) - var result = mlx_array_new() - mlx_slice_update( - &result, - src.ctx, update.ctx, starts, starts.count, ends, ends.count, strides, strides.count, - stream.ctx) - return MLXArray(result) + return .slice(update: update, starts: starts, ends: ends, strides: strides) +} + +func updateSlice( + src: MLXArray, operations: [MLXArrayIndexOperation], update: MLXArray, + stream: StreamOrDevice = .default +) -> MLXArray? { + // See mlx_slice_update + switch sliceUpdateArguments( + src: src, operations: operations, update: update, stream: stream) + { + case nil: + return nil + + case .broadcast(let update): + return update + + case .slice(let update, let starts, let ends, let strides): + var result = mlx_array_new() + mlx_slice_update( + &result, + src.ctx, update.ctx, starts, starts.count, ends, ends.count, strides, strides.count, + stream.ctx) + return MLXArray(result) + } } // MARK: - index set (scatter) diff --git a/Source/MLX/MLXArray+Ops.swift b/Source/MLX/MLXArray+Ops.swift index f940d2970..f9927a9f7 100644 --- a/Source/MLX/MLXArray+Ops.swift +++ b/Source/MLX/MLXArray+Ops.swift @@ -1551,7 +1551,7 @@ extension MLXArray { axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cummax(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_cummax_axis(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -1576,9 +1576,8 @@ extension MLXArray { public func cummax( reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = self.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cummax(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_cummax(&result, ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -1605,7 +1604,7 @@ extension MLXArray { axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cummin(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_cummin_axis(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -1630,9 +1629,8 @@ extension MLXArray { public func cummin( reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = self.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cummin(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_cummin(&result, ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -1649,17 +1647,21 @@ extension MLXArray { /// - axis: axis to reduce over /// - reverse: reverse the reduction /// - inclusive: include the initial value + /// - dtype: if specified, convert the input to this type /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``cumprod(reverse:inclusive:stream:)`` - /// - ``cumprod(_:axis:reverse:inclusive:stream:)`` + /// - ``cumprod(reverse:inclusive:dtype:stream:)`` + /// - ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` public func cumprod( - axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default + axis: Int, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cumprod(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumprod_axis(&result, ctx, axis.int32, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -1675,18 +1677,21 @@ extension MLXArray { /// - Parameters: /// - reverse: reverse the reduction /// - inclusive: include the initial value + /// - dtype: if specified, convert the input to this type /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``cumprod(axis:reverse:inclusive:stream:)`` - /// - ``cumprod(_:axis:reverse:inclusive:stream:)`` + /// - ``cumprod(axis:reverse:inclusive:dtype:stream:)`` + /// - ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` public func cumprod( - reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default + reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray { - let flat = self.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cumprod(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumprod(&result, ctx, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -1703,17 +1708,21 @@ extension MLXArray { /// - axis: axis to reduce over /// - reverse: reverse the reduction /// - inclusive: include the initial value + /// - dtype: if specified, convert the input to this type /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``cumsum(reverse:inclusive:stream:)`` - /// - ``cumsum(_:axis:reverse:inclusive:stream:)`` + /// - ``cumsum(reverse:inclusive:dtype:stream:)`` + /// - ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` public func cumsum( - axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default + axis: Int, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cumsum(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumsum_axis(&result, ctx, axis.int32, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -1729,18 +1738,21 @@ extension MLXArray { /// - Parameters: /// - reverse: reverse the reduction /// - inclusive: include the initial value + /// - dtype: if specified, convert the input to this type /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``cumsum(axis:reverse:inclusive:stream:)`` - /// - ``cumsum(_:axis:reverse:inclusive:stream:)`` + /// - ``cumsum(axis:reverse:inclusive:dtype:stream:)`` + /// - ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` public func cumsum( - reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default + reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray { - let flat = self.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cumsum(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumsum(&result, ctx, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -1998,7 +2010,7 @@ extension MLXArray { axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_logcumsumexp(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_logcumsumexp_axis(&result, ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -2014,9 +2026,8 @@ extension MLXArray { public func logCumsumExp( reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = self.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_logcumsumexp(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_logcumsumexp(&result, ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } diff --git a/Source/MLX/MLXFast.swift b/Source/MLX/MLXFast.swift index 92c96da83..47e5c86cf 100644 --- a/Source/MLX/MLXFast.swift +++ b/Source/MLX/MLXFast.swift @@ -114,11 +114,15 @@ public enum MLXFast { /// - mask: mask array /// - sinks: optional array of attention sinks /// - memoryEfficientThreshold: unused + /// - forceFused: If `true` use a fused kernel regardless of the heuristic choice. This + /// can result in slower kernels in some case but can also reduce memory consumption /// - stream: stream to evaluate on public static func scaledDotProductAttention( queries: MLXArray, keys: MLXArray, values: MLXArray, scale: Float, mask: MLXArray?, sinks: MLXArray? = nil, - memoryEfficientThreshold: Int? = nil, stream: StreamOrDevice = .default + memoryEfficientThreshold: Int? = nil, + forceFused: Bool = false, + stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() @@ -127,6 +131,7 @@ public enum MLXFast { queries.ctx, keys.ctx, values.ctx, scale, "", mask?.ctx ?? MLXArray.mlxNone.ctx, (sinks ?? .mlxNone).ctx, + forceFused, stream.ctx) return MLXArray(result) } @@ -199,11 +204,14 @@ public enum MLXFast { /// - scale: scale for queries, typically `1 / sqrt(q.dim(-1))` /// - mask: a ``ScaledDotProductAttentionMaskMode`` /// - sinks: optional array of attention sinks + /// - forceFused: If `true` use a fused kernel regardless of the heuristic choice. This + /// can result in slower kernels in some case but can also reduce memory consumption /// - stream: stream to evaluate on public static func scaledDotProductAttention( queries: MLXArray, keys: MLXArray, values: MLXArray, scale: Float, mask: ScaledDotProductAttentionMaskMode, sinks: MLXArray? = nil, + forceFused: Bool = false, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() @@ -213,6 +221,7 @@ public enum MLXFast { queries.ctx, keys.ctx, values.ctx, scale, mask.mode, mask.mask?.ctx ?? MLXArray.mlxNone.ctx, (sinks ?? .mlxNone).ctx, + forceFused, stream.ctx) return MLXArray(result) } @@ -259,6 +268,22 @@ public enum MLXFast { return MLXArray(result) } + /// Cross entropy loss with class indices as targets. + /// + /// - Parameters: + /// - logits: The unnormalized logits. The loss is computed over the last axis. + /// - targets: Class indices. The shape should match the shape of `logits` with the last axis removed. + /// The indices must be in `0 ..< logits.dim(-1)`. + /// - stream: stream or device to evaluate on + public static func crossEntropy( + logits: MLXArray, targets: MLXArray, + stream: StreamOrDevice = .default + ) -> MLXArray { + var result = mlx_array_new() + mlx_fast_cross_entropy( + &result, logits.ctx, targets.ctx, stream.ctx) + return MLXArray(result) + } } /// Optimized implementation of `NN.RoPE`. @@ -367,3 +392,5 @@ public func layerNorm( ) -> MLXArray { return MLXFast.layerNorm(x, weight: weight, bias: bias, eps: eps, stream: stream) } + +// Note: crossEntropy() requires MLXFast.crossEntropy() to avoid collision with MLXNN diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index 4b90ffbf6..f581d40b2 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -460,7 +460,7 @@ public func cummax( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cummax(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_cummax_axis(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -481,9 +481,8 @@ public func cummax( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = array.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cummax(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_cummax(&result, array.ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -505,7 +504,7 @@ public func cummin( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cummin(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_cummin_axis(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -526,9 +525,8 @@ public func cummin( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = array.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cummin(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_cummin(&result, array.ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -543,14 +541,16 @@ public func cummin( /// /// ### See Also /// - -/// - ``cumprod(_:reverse:inclusive:stream:)`` +/// - ``cumprod(_:reverse:inclusive:dtype:stream:)`` /// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` public func cumprod( _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cumprod(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumprod_axis(&result, array.ctx, axis.int32, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -565,15 +565,16 @@ public func cumprod( /// /// ### See Also /// - -/// - ``cumprod(_:axis:reverse:inclusive:stream:)`` +/// - ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` /// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` public func cumprod( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = array.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cumprod(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumprod(&result, array.ctx, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -588,14 +589,16 @@ public func cumprod( /// /// ### See Also /// - -/// - ``cumsum(_:reverse:inclusive:stream:)`` +/// - ``cumsum(_:reverse:inclusive:dtype:stream:)`` /// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` public func cumsum( _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_cumsum(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumsum_axis(&result, array.ctx, axis.int32, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -610,15 +613,16 @@ public func cumsum( /// /// ### See Also /// - -/// - ``cumsum(_:axis:reverse:inclusive:stream:)`` +/// - ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` /// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` public func cumsum( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = array.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_cumsum(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + let dtype = mlx_optional_dtype(dtype) + mlx_cumsum(&result, array.ctx, reverse, inclusive, dtype, stream.ctx) return MLXArray(result) } @@ -844,7 +848,7 @@ public func logCumsumExp( axis: Int, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_logcumsumexp(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) + mlx_logcumsumexp_axis(&result, array.ctx, axis.int32, reverse, inclusive, stream.ctx) return MLXArray(result) } @@ -862,9 +866,8 @@ public func logCumsumExp( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, stream: StreamOrDevice = .default ) -> MLXArray { - let flat = array.reshaped([-1], stream: stream) var result = mlx_array_new() - mlx_logcumsumexp(&result, flat.ctx, 0, reverse, inclusive, stream.ctx) + mlx_logcumsumexp(&result, array.ctx, reverse, inclusive, stream.ctx) return MLXArray(result) } diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index b7e95a330..362f16de0 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -406,7 +406,7 @@ public func bartlett( _ m: Int, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_bartlett(&result, Int32(m), stream.ctx) + mlx_bartlett(&result, m.int32, stream.ctx) return MLXArray(result) } @@ -423,7 +423,7 @@ public func blackman( _ m: Int, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_blackman(&result, Int32(m), stream.ctx) + mlx_blackman(&result, m.int32, stream.ctx) return MLXArray(result) } @@ -1154,7 +1154,7 @@ public func dequantized( var result = mlx_array_new() let gs = mlx_optional_int(value: Int32(groupSize ?? 0), has_value: groupSize != nil) let bits = mlx_optional_int(value: Int32(bits ?? 0), has_value: bits != nil) - let dtype = mlx_optional_dtype(value: dtype?.cmlxDtype ?? MLX_FLOAT16, has_value: dtype != nil) + let dtype = mlx_optional_dtype(dtype) mlx_dequantize( &result, w.ctx, scales.ctx, (biases ?? .mlxNone).ctx, gs, bits, mode.rawValue, @@ -1164,6 +1164,21 @@ public func dequantized( return MLXArray(result) } +/// The n-th discrete difference along the given axis. +/// +/// - Parameters: +/// - array: Input array +/// - n: The number of times to difference +/// - axis: The axis along which to difference +/// - stream: Stream or device to evaluate on +public func diff( + _ array: MLXArray, n: Int = 1, axis: Int = -1, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_diff(&result, array.ctx, n.int32, axis.int32, stream.ctx) + return MLXArray(result) +} + /// Element-wise division. /// /// Divide two arrays with . @@ -1374,6 +1389,35 @@ public func expm1(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXAr return MLXArray(result) } +/// Reverse the order of elements along the given axes. +/// +/// - Parameters: +/// - array: Input array +/// - axes: The axes along which to flip or None for all +/// - stream: Stream or device to evaluate on +public func flipped( + _ array: MLXArray, axes: [Int]? = nil, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + let axes = axes ?? Array(0 ..< array.ndim) + mlx_flip_axes(&result, array.ctx, axes.asInt32, axes.count, stream.ctx) + return MLXArray(result) +} + +/// Reverse the order of elements along the given axis. +/// +/// - Parameters: +/// - array: Input array +/// - axis: The axis along which to flip +/// - stream: Stream or device to evaluate on +public func flipped( + _ array: MLXArray, axis: Int, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_flip_axis(&result, array.ctx, axis.int32, stream.ctx) + return MLXArray(result) +} + /// Convert an E4M3 FP8 array to the given floating-point ``DType``. /// /// The input array is interpreted as raw FP8 bytes. @@ -1624,7 +1668,7 @@ public func hamming( _ m: Int, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_hamming(&result, Int32(m), stream.ctx) + mlx_hamming(&result, m.int32, stream.ctx) return MLXArray(result) } @@ -1641,7 +1685,7 @@ public func hanning( _ m: Int, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_hanning(&result, Int32(m), stream.ctx) + mlx_hanning(&result, m.int32, stream.ctx) return MLXArray(result) } @@ -1945,8 +1989,27 @@ public func logicalOr( return MLXArray(result) } +/// Element-wise logical exclusive or. +/// +/// - Parameters: +/// - a: input array or scalar +/// - b: input array or scalar +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - +public func logicalXor( + _ a: some ScalarOrArray, _ b: some ScalarOrArray, stream: StreamOrDevice = .default +) -> MLXArray { + let (a, b) = toArrays(a, b) + var result = mlx_array_new() + mlx_logical_xor(&result, a.ctx, b.ctx, stream.ctx) + return MLXArray(result) +} + /// Indexing mode for ``meshGrid(_:sparse:indexing:stream:)``. -public enum MeshGridIndexing: String, Sendable { +public enum MeshGridIndexing: String, Sendable, Codable { /// cartesian indexing case xy @@ -2014,7 +2077,7 @@ public func median( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_median(&result, a.ctx, [axis.int32], 1, keepDims, stream.ctx) + mlx_median_axis(&result, a.ctx, axis.int32, keepDims, stream.ctx) return MLXArray(result) } @@ -2033,7 +2096,7 @@ public func median( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_median(&result, a.ctx, axes.asInt32, axes.count, keepDims, stream.ctx) + mlx_median_axes(&result, a.ctx, axes.asInt32, axes.count, keepDims, stream.ctx) return MLXArray(result) } @@ -2051,7 +2114,7 @@ public func median( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_median(&result, a.ctx, nil, 0, keepDims, stream.ctx) + mlx_median(&result, a.ctx, keepDims, stream.ctx) return MLXArray(result) } @@ -2325,6 +2388,17 @@ public func partitioned(_ array: MLXArray, kth: Int, stream: StreamOrDevice = .d return MLXArray(result) } +/// Element-wise unary plus. Returns a copy of the input. +/// +/// - Parameters: +/// - array: input array +/// - stream: stream or device to evaluate on +public func positive(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { + var result = mlx_array_new() + mlx_positive(&result, array.ctx, stream.ctx) + return MLXArray(result) +} + /// Put values along an axis at the specified indices. /// /// - Parameters: @@ -2621,6 +2695,31 @@ public func roll( return MLXArray(result) } +public enum SearchSortedSide: String, Sendable, Codable { + case left + case right +} + +/// Find the indices that keep `sorted` sorted when inserting `values`. +/// +/// - Parameters: +/// - sorted: 1-D array in ascending order. +/// - values: values to insert in any shape +/// - side: if `.left` the first suitable index is returned, otherwise last +/// - stream: stream or device to evaluate on +/// - Returns: `.uint32` array in the same shape as `values` +public func search( + sorted: MLXArray, + values: MLXArray, + side: SearchSortedSide = .left, + stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_searchsorted( + &result, sorted.ctx, values.ctx, side.rawValue.cString(using: .utf8), stream.ctx) + return MLXArray(result) +} + /// Element-wise logistic sigmoid. /// /// For details, please see @@ -2882,6 +2981,21 @@ public func stacked( return MLXArray(result) } +/// Split an array into a sequence of arrays along the given axis. +/// +/// The inverse of ``stacked(_:axis:stream:)``. The given axis is removed from each of the returned arrays. +/// +/// ### See Also +/// - +public func unstacked( + _ array: MLXArray, axis: Int = 0, stream: StreamOrDevice = .default +) -> [MLXArray] { + var result = mlx_vector_array_new() + defer { mlx_vector_array_free(result) } + mlx_unstack_axis(&result, array.ctx, axis.int32, stream.ctx) + return mlx_vector_array_values(result) +} + /// Stop gradients from being computed. /// ///The operation is the identity but it prevents gradients from flowing @@ -3058,6 +3172,21 @@ public func tensordot( return MLXArray(result) } +/// Compute the vector dot product of two arrays along an axis. +/// +/// - Parameters: +/// - a: first array +/// - b: second array +/// - axis: axis over which to compute the dot product +/// - stream: stream or device to evaluate on +public func vecdot( + _ a: MLXArray, _ b: MLXArray, axis: Int = -1, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_vecdot(&result, a.ctx, b.ctx, axis.int32, stream.ctx) + return MLXArray(result) +} + /// Construct array by repeating given array the number of times given by `repetitions`. /// /// - Parameters: @@ -3170,7 +3299,7 @@ public func trace( stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_trace( + mlx_trace_axes( &result, array.ctx, offset.int32, axis1.int32, axis2.int32, (dtype ?? array.dtype).cmlxDtype, stream.ctx) @@ -3207,6 +3336,17 @@ public func triu(_ array: MLXArray, k: Int = 0, stream: StreamOrDevice = .defaul return MLXArray(result) } +/// Element-wise truncation towards zero. +/// +/// - Parameters: +/// - array: Input array +/// - stream: stream or device to evaluate on +public func trunc(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArray { + var result = mlx_array_new() + mlx_trunc(&result, array.ctx, stream.ctx) + return MLXArray(result) +} + /// Select from `x` or `y` according to `condition`. /// /// The condition and input arrays must be the same shape or @@ -3302,7 +3442,7 @@ public func unflatten( _ a: MLXArray, axis: Int, shape: some Collection, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() - mlx_unflatten(&result, a.ctx, axis.int32, shape.map { Int32($0) }, shape.count, stream.ctx) + mlx_unflatten(&result, a.ctx, axis.int32, shape.map { $0.int32 }, shape.count, stream.ctx) return MLXArray(result) } diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index 8c0e929d2..917d544ad 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -83,8 +83,23 @@ public final class Stream: @unchecked Sendable, Equatable { let ctx: mlx_stream - public static let gpu = Stream(mlx_default_gpu_stream_new()) - public static let cpu = Stream(mlx_default_cpu_stream_new()) + private static func newStreamThreadUnsafe(_ deviceType: DeviceType) -> mlx_stream { + var cDeviceType: mlx_device_type + switch deviceType { + case DeviceType.cpu: + cDeviceType = MLX_CPU + case DeviceType.gpu: + cDeviceType = MLX_GPU + } + return evalLock.withLock { + let ctx = mlx_device_new_type(cDeviceType, 0) + defer { mlx_device_free(ctx) } + return mlx_stream_new_thread_unsafe(ctx) + } + } + + public static let gpu = Stream(newStreamThreadUnsafe(.gpu)) + public static let cpu = Stream(newStreamThreadUnsafe(.cpu)) @TaskLocal static var defaultStream: Stream? @@ -111,15 +126,15 @@ public final class Stream: @unchecked Sendable, Equatable { /// Default stream on the default device. public init() { let device = Device.defaultDevice() - var ctx = mlx_stream_new() - mlx_get_default_stream(&ctx, device.ctx) - self.ctx = ctx + self.ctx = evalLock.withLock { + mlx_stream_new_thread_unsafe(device.ctx) + } } @available(*, deprecated, message: "use init(Device) -- index not supported") public init(index: Int32, _ device: Device) { - self.ctx = withEvalLock { - mlx_stream_new_device(device.ctx) + self.ctx = evalLock.withLock { + mlx_stream_new_thread_unsafe(device.ctx) } } @@ -127,8 +142,8 @@ public final class Stream: @unchecked Sendable, Equatable { /// /// See also ``withNewDefaultStream(device:_:)-5bwc3`` public init(_ device: Device) { - self.ctx = withEvalLock { - mlx_stream_new_device(device.ctx) + self.ctx = evalLock.withLock { + mlx_stream_new_thread_unsafe(device.ctx) } } diff --git a/Source/MLX/Transforms+Compile.swift b/Source/MLX/Transforms+Compile.swift index 0bae50dd3..c9616f753 100644 --- a/Source/MLX/Transforms+Compile.swift +++ b/Source/MLX/Transforms+Compile.swift @@ -97,7 +97,11 @@ final class CompiledFunction: @unchecked (Sendable) { + "see CompiledFunction.deinit") EvalLockOwnership.recordCompileErase() #endif - _ = mlx_detail_compile_erase(functionID) + // remove the compiled structure from the back end + var cache = mlx_compile_cache_new() + mlx_detail_compile_cache(&cache) + defer { mlx_compile_cache_free(cache) } + mlx_detail_compile_erase(cache, functionID) } } diff --git a/Tests/MLXTests/ArrayAtTests.swift b/Tests/MLXTests/ArrayAtTests.swift index 67e6d9eed..57b077e9b 100644 --- a/Tests/MLXTests/ArrayAtTests.swift +++ b/Tests/MLXTests/ArrayAtTests.swift @@ -37,4 +37,65 @@ class ArrayAtTests: XCTestCase { let a3 = MLXArray([0, 0]) assertEqual(a3.at[idx].add(1), MLXArray([2, 2])) } + + // MARK: - slice updates + // + // pure slice indices route through the mlx_slice_update_* family rather than a + // scatter -- these check that path produces the same answers. + + func testSliceAdd() { + let a = MLXArray(0 ..< 6) + assertEqual(a.at[1 ..< 4].add(10), MLXArray([0, 11, 12, 13, 4, 5])) + } + + func testSliceSubtract() { + let a = MLXArray(0 ..< 6) + assertEqual(a.at[1 ..< 4].subtract(1), MLXArray([0, 0, 1, 2, 4, 5])) + } + + func testSliceMultiply() { + let a = MLXArray(1 ..< 7) + assertEqual(a.at[2 ..< 5].multiply(10), MLXArray([1, 2, 30, 40, 50, 6])) + } + + func testSliceDivide() { + let a = MLXArray(converting: [2.0, 4.0, 6.0, 8.0]) + assertEqual( + a.at[1 ..< 3].divide(2), MLXArray(converting: [2.0, 2.0, 3.0, 8.0]), atol: 1e-6) + } + + func testSliceMinimumMaximum() { + let a = MLXArray([5, 5, 5, 5]) + assertEqual(a.at[0 ..< 2].minimum(3), MLXArray([3, 3, 5, 5])) + assertEqual(a.at[0 ..< 2].maximum(7), MLXArray([7, 7, 5, 5])) + } + + func testSliceIndexAdd() { + // a single integer index is also expressible as a slice update + let a = MLXArray(0 ..< 4) + assertEqual(a.at[1].add(100), MLXArray([0, 101, 2, 3])) + } + + func testMultiDimensionalSliceAdd() { + let a = MLXArray(0 ..< 6, [2, 3]) + assertEqual( + a.at[0 ..< 1, 1 ..< 3].add(10), + MLXArray([0, 11, 12, 3, 4, 5], [2, 3])) + } + + func testStridedSliceAdd() { + let a = MLXArray(0 ..< 6) + assertEqual(a.at[.stride(by: 2)].add(10), MLXArray([10, 1, 12, 3, 14, 5])) + } + + func testSliceAndScatterAgree() { + // the slice path and the scatter path should produce the same result -- + // indexing with an array forces the scatter fallback + let a = MLXArray(0 ..< 6) + + let viaSlice = a.at[1 ..< 4].add(10) + let viaScatter = a.at[MLXArray([1, 2, 3])].add(10) + + assertEqual(viaSlice, viaScatter) + } } diff --git a/Tests/MLXTests/FFTTests.swift b/Tests/MLXTests/FFTTests.swift new file mode 100644 index 000000000..e602c5872 --- /dev/null +++ b/Tests/MLXTests/FFTTests.swift @@ -0,0 +1,73 @@ +// Copyright © 2025 Apple Inc. + +import Foundation +import MLX +import XCTest + +class FFTTests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + func testFFTRoundTrip() { + let a = MLXArray(converting: [1.0, 2.0, 3.0, 4.0]) + + for norm in [FFTNorm.backward, .ortho, .forward] { + let r = ifft(fft(a, norm: norm), norm: norm) + assertEqual(r.realPart(), a, atol: 1e-5) + } + } + + func testFFTNormScaling() { + let a = MLXArray(converting: [1.0, 1.0, 1.0, 1.0]) + let n = 4.0 + + // for a constant signal only the DC term is non-zero and its magnitude + // shows how each mode scales the forward transform + let backward = fft(a, norm: .backward)[0].realPart().item(Double.self) + let ortho = fft(a, norm: .ortho)[0].realPart().item(Double.self) + let forward = fft(a, norm: .forward)[0].realPart().item(Double.self) + + XCTAssertEqual(backward, n, accuracy: 1e-5) + XCTAssertEqual(ortho, n / n.squareRoot(), accuracy: 1e-5) + XCTAssertEqual(forward, 1.0, accuracy: 1e-5) + } + + func testFFTNormDefaultIsBackward() { + let a = MLXArray(converting: [1.0, 2.0, 3.0, 4.0]) + assertEqual(fft(a).realPart(), fft(a, norm: .backward).realPart(), atol: 1e-5) + } + + func testFFTNAndAxes() { + let a = MLXArray(0 ..< 24, [2, 3, 4]).asType(.float32) + + // norm should thread through each of the n-dimensional variants + for norm in [FFTNorm.backward, .ortho, .forward] { + assertEqual(ifftn(fftn(a, norm: norm), norm: norm).realPart(), a, atol: 1e-4) + assertEqual( + ifft2(fft2(a, norm: norm), norm: norm).realPart(), a, atol: 1e-4) + assertEqual( + irfftn(rfftn(a, norm: norm), norm: norm), a, atol: 1e-4) + } + } + + func testFFTFrequencies() { + // matches numpy/mlx: [0, 1, 2, 3, -4, -3, -2, -1] / (n * d) + let f = fftfreq(8) + assertEqual( + f, MLXArray(converting: [0, 0.125, 0.25, 0.375, -0.5, -0.375, -0.25, -0.125]), + atol: 1e-6) + + // d scales the result + let scaled = fftfreq(8, d: 0.5) + assertEqual(scaled, f * 2, atol: 1e-6) + } + + func testRFFTFrequencies() { + // only the non-negative frequencies -- n / 2 + 1 of them + let f = rfftfreq(8) + XCTAssertEqual(f.shape, [5]) + assertEqual(f, MLXArray(converting: [0, 0.125, 0.25, 0.375, 0.5]), atol: 1e-6) + } +} diff --git a/Tests/MLXTests/GraphUtilsTests.swift b/Tests/MLXTests/GraphUtilsTests.swift new file mode 100644 index 000000000..7e9c5ca27 --- /dev/null +++ b/Tests/MLXTests/GraphUtilsTests.swift @@ -0,0 +1,69 @@ +// Copyright © 2025 Apple Inc. + +import Foundation +import MLX +import XCTest + +class GraphUtilsTests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + func testExportToDot() { + let x = MLXArray([1, 2, 3]) + let dot = exportToDot([x * 2 + 1]) + + XCTAssertFalse(dot.isEmpty) + XCTAssertTrue(dot.contains("digraph"), "expected a digraph, got:\n\(dot)") + + // the graph should mention the primitives that produce the output + XCTAssertTrue(dot.contains("Add"), "expected Add in:\n\(dot)") + XCTAssertTrue(dot.contains("Multiply"), "expected Multiply in:\n\(dot)") + } + + func testGraphDescription() { + let x = MLXArray([1, 2, 3]) + let description = graphDescription([x + x]) + + XCTAssertFalse(description.isEmpty) + XCTAssertTrue(description.contains("Add"), "expected Add in:\n\(description)") + } + + func testNodeNamer() { + let x = MLXArray([1, 2, 3]) + let y = MLXArray([4, 5, 6]) + + let namer = NodeNamer() + namer.setName("lhs", for: x) + namer.setName("rhs", for: y) + + XCTAssertEqual(namer.name(for: x), "lhs") + XCTAssertEqual(namer.name(for: y), "rhs") + + let dot = exportToDot([x + y], namer: namer) + XCTAssertTrue(dot.contains("lhs"), "expected lhs in:\n\(dot)") + XCTAssertTrue(dot.contains("rhs"), "expected rhs in:\n\(dot)") + } + + func testNodeNamerGeneratesNames() { + // unnamed arrays get a generated name, and it is stable + let x = MLXArray([1, 2, 3]) + let namer = NodeNamer() + + let first = namer.name(for: x) + XCTAssertNotNil(first) + XCTAssertEqual(namer.name(for: x), first) + } + + func testGraphIsNotEvaluated() { + // describing the graph should not force evaluation + let x = MLXArray([1, 2, 3]) + let y = x * 2 + + _ = exportToDot([y]) + + // still usable afterwards + assertEqual(y, MLXArray([2, 4, 6])) + } +} diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 6f700df4e..373ff8ca3 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -96,4 +96,20 @@ final class SaveTests: XCTestCase { XCTAssertEqual(loadedMetadata, metadata) } + /// `loadArrays(data:)` seeks to the end to size its input, so the in-memory IO + /// stream has to report the size of the whole buffer rather than an offset + /// relative to the read position. + public func testLoadFromDataLargerThanHeader() throws { + // enough data that a SEEK_END mistake shows up as a truncated size + let arrays: [String: MLXArray] = [ + "big": MLX.ones([64, 64]) + ] + + let data = try saveToData(arrays: arrays) + XCTAssertGreaterThan(data.count, 8) + + let loaded = try loadArrays(data: data) + assertEqual(try XCTUnwrap(loaded["big"]), try XCTUnwrap(arrays["big"])) + } + } diff --git a/tools/fix-metal-includes.sh b/tools/fix-metal-includes.sh index 622d43114..6606683cf 100755 --- a/tools/fix-metal-includes.sh +++ b/tools/fix-metal-includes.sh @@ -22,11 +22,12 @@ KERNELS_DIR="${CMLX_MLX_DIR}/${KERNELS_INCLUDE_PATH}" KERNEL_LIST=" \ arg_reduce.metal \ conv.metal \ -gemv.metal \ +dot.metal \ layer_norm.metal \ random.metal \ rms_norm.metal \ rope.metal \ +searchsorted.metal \ scaled_dot_product_attention.metal \ steel/attn/kernels/steel_attention.metal" diff --git a/tools/update-mlx-xcodeproj.sh b/tools/update-mlx-xcodeproj.sh index a37d04d9d..74b96cf2d 100755 --- a/tools/update-mlx-xcodeproj.sh +++ b/tools/update-mlx-xcodeproj.sh @@ -45,9 +45,11 @@ for x in \ distributed/distributed.h \ distributed/ops.h \ einsum.h \ + error.h \ export.h \ fast.h \ fft.h \ + graph_utils.h \ io.h \ linalg.h \ memory.h \ @@ -71,6 +73,8 @@ for x in \ scheduler.h \ primitives.h \ backend/metal/device.h \ + backend/metal/resident.h \ + backend/common/metal_kernel.h \ backend/common/utils.h \ backend/cpu/encoder.h \ backend/gpu/eval.h diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index 87b93ba14..c9f0a2e50 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -299,9 +299,13 @@ "examples/example-closure.c", "examples/example-export.c", "examples/example-float64.c", + "examples/example-gguf.c", "examples/example-grad.c", + "examples/example-graph.c", "examples/example-metal-kernel.c", + "examples/example-ops.c", "examples/example-safe-tensors.c", + "examples/example-thread-stream.c", examples/example.c, LICENSE, "mlx-c.pc.in", @@ -317,6 +321,7 @@ mlx/c/export.h, mlx/c/fast.h, mlx/c/fft.h, + mlx/c/graph_utils.h, mlx/c/half.h, mlx/c/io_types.h, mlx/c/io.h, @@ -329,10 +334,13 @@ mlx/c/optional.h, mlx/c/private/array.h, mlx/c/private/closure.h, + mlx/c/private/compile.h, mlx/c/private/device.h, mlx/c/private/distributed_group.h, mlx/c/private/enums.h, mlx/c/private/export.h, + mlx/c/private/gguf.h, + mlx/c/private/graph_utils.h, mlx/c/private/io.h, mlx/c/private/map.h, mlx/c/private/mlx.h, @@ -363,30 +371,32 @@ isa = PBXFileSystemSynchronizedBuildFileExceptionSet; membershipExceptions = ( ".clang-format", - ".github/actions/build-cuda-release/action.yml", ".github/actions/build-docs/action.yml", - ".github/actions/build-linux-release/action.yml", - ".github/actions/build-linux/action.yml", - ".github/actions/build-macos-release/action.yml", ".github/actions/build-macos/action.yml", - ".github/actions/build-windows/action.yml", - ".github/actions/setup-linux/action.yml", - ".github/actions/setup-macos/action.yml", - ".github/actions/setup-windows/action.yml", + ".github/actions/build-wheel/action.yml", + .github/actions/build/action.yml, + .github/actions/setup/action.yml, ".github/actions/test-linux/action.yml", + ".github/actions/test-macos/action.yml", + ".github/actions/test-wheel/action.yml", ".github/actions/test-windows/action.yml", .github/dependabot.yml, .github/ISSUE_TEMPLATE/bug_report.md, + .github/ISSUE_TEMPLATE/config.yml, + .github/ISSUE_TEMPLATE/other.md, .github/pull_request_template.md, ".github/scripts/build-sanitizer-tests.sh", + ".github/scripts/find-eligible-contributors.js", ".github/scripts/setup+build-cpp-linux-fedora-container.sh", + ".github/scripts/update-bypass-list.js", .github/workflows/build_and_test.yml, .github/workflows/documentation.yml, - .github/workflows/nightly.yml, .github/workflows/release.yml, + .github/workflows/update_bypass_list.yml, .gitignore, ".pre-commit-config.yaml", ACKNOWLEDGMENTS.md, + AGENTS.md, benchmarks/cpp/autograd.cpp, benchmarks/cpp/CMakeLists.txt, benchmarks/cpp/compare_devices.cpp, @@ -398,6 +408,7 @@ benchmarks/python/batch_matmul_bench.py, benchmarks/python/blas/bench_gemm.py, benchmarks/python/blas/bench_gemv.py, + benchmarks/python/block_masked_mm_bench.py, benchmarks/python/comparative/bench_mlx.py, benchmarks/python/comparative/bench_torch.py, benchmarks/python/comparative/compare.py, @@ -429,10 +440,13 @@ benchmarks/python/scatter_bench.py, benchmarks/python/sdpa_bench.py, benchmarks/python/sdpa_vector_bench.py, + benchmarks/python/segmented_mm_bench.py, benchmarks/python/single_ops.py, + benchmarks/python/slice_update_bench.py, benchmarks/python/synchronize_bench.py, benchmarks/python/time_utils.py, CITATION.cff, + CLAUDE.md, cmake/extension.cmake, cmake/FindCUDNN.cmake, cmake/FindNCCL.cmake, @@ -445,6 +459,12 @@ docs/.nojekyll, docs/Doxyfile, docs/index.html, + docs/logo/mlx_logo_dark.pdf, + docs/logo/mlx_logo_dark.png, + docs/logo/mlx_logo_dark.svg, + docs/logo/mlx_logo.pdf, + docs/logo/mlx_logo.png, + docs/logo/mlx_logo.svg, docs/Makefile, docs/README.md, docs/requirements.txt, @@ -498,17 +518,21 @@ docs/src/python/optimizers/common_optimizers.rst, docs/src/python/optimizers/optimizer.rst, docs/src/python/optimizers/schedulers.rst, + docs/src/python/printoptions.rst, docs/src/python/random.rst, docs/src/python/transforms.rst, docs/src/python/tree_utils.rst, docs/src/usage/compile.rst, docs/src/usage/distributed.rst, + docs/src/usage/environment_variables.rst, docs/src/usage/export.rst, docs/src/usage/function_transforms.rst, docs/src/usage/indexing.rst, + docs/src/usage/kv_cache.rst, docs/src/usage/launching_distributed.rst, docs/src/usage/lazy_evaluation.rst, docs/src/usage/numpy.rst, + docs/src/usage/precision.rst, docs/src/usage/quick_start.rst, docs/src/usage/saving_and_loading.rst, docs/src/usage/unified_memory.rst, @@ -540,6 +564,8 @@ examples/extensions/requirements.txt, examples/extensions/setup.py, examples/extensions/test.py, + examples/python/distributed_data_parallel.py, + examples/python/distributed_tensor_parallel.py, examples/python/linear_regression.py, examples/python/logistic_regression.py, examples/python/qqmm.py, @@ -559,6 +585,7 @@ mlx/backend/common/copy.h, mlx/backend/common/hadamard.h, mlx/backend/common/matmul.h, + mlx/backend/common/metal_kernel.h, mlx/backend/common/quantized.h, mlx/backend/common/reduce.h, mlx/backend/common/slicing.h, @@ -580,6 +607,7 @@ mlx/backend/cpu/gemms/simd_bf16.cpp, mlx/backend/cpu/gemms/simd_fp16.cpp, mlx/backend/cpu/gemms/simd_gemm.h, + mlx/backend/cpu/jit_compiler.cpp, mlx/backend/cpu/jit_compiler.h, mlx/backend/cpu/lapack.h, mlx/backend/cpu/make_compiled_preamble.ps1, @@ -635,6 +663,7 @@ mlx/backend/cuda/copy/copy_general_input.cu, mlx/backend/cuda/copy/copy_general.cu, mlx/backend/cuda/copy/copy.cuh, + mlx/backend/cuda/cross_entropy.cu, mlx/backend/cuda/cublas_utils.cpp, mlx/backend/cuda/cublas_utils.h, mlx/backend/cuda/cuda_utils.h, @@ -652,29 +681,43 @@ mlx/backend/cuda/device/cast_op.cuh, mlx/backend/cuda/device/complex.cuh, mlx/backend/cuda/device/config.h, + mlx/backend/cuda/device/cute_dequant.cuh, mlx/backend/cuda/device/fp16_math.cuh, mlx/backend/cuda/device/gather_axis.cuh, mlx/backend/cuda/device/gather.cuh, + mlx/backend/cuda/device/gemm_sm70.cuh, + mlx/backend/cuda/device/hadamard.cuh, mlx/backend/cuda/device/indexing.cuh, + mlx/backend/cuda/device/qmm_naive.cuh, + mlx/backend/cuda/device/qmm_sm80.cuh, + mlx/backend/cuda/device/qmm_sm90.cuh, mlx/backend/cuda/device/scatter_axis.cuh, mlx/backend/cuda/device/scatter_ops.cuh, mlx/backend/cuda/device/scatter.cuh, + mlx/backend/cuda/device/slice_update.cuh, mlx/backend/cuda/device/ternary_ops.cuh, mlx/backend/cuda/device/unary_ops.cuh, mlx/backend/cuda/device/utils.cuh, + mlx/backend/cuda/dirs.cpp, mlx/backend/cuda/distributed.cu, mlx/backend/cuda/eval.cpp, mlx/backend/cuda/event.cu, mlx/backend/cuda/event.h, mlx/backend/cuda/fence.cpp, + mlx/backend/cuda/fft.cu, + mlx/backend/cuda/gemms/block_mask.cu, + mlx/backend/cuda/gemms/block_mask.h, mlx/backend/cuda/gemms/cublas_gemm_batched_12_0.cpp, mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu, mlx/backend/cuda/gemms/cublas_gemm.cpp, mlx/backend/cuda/gemms/cublas_gemm.h, + mlx/backend/cuda/gemms/gather_gemm.cu, + mlx/backend/cuda/gemms/gather_gemm.h, mlx/backend/cuda/gemms/gemv.cu, mlx/backend/cuda/gemms/gemv.h, mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu, mlx/backend/cuda/gemms/grouped_gemm.h, + mlx/backend/cuda/hadamard.cu, mlx/backend/cuda/indexing.cpp, mlx/backend/cuda/jit_module.cpp, mlx/backend/cuda/jit_module.h, @@ -686,22 +729,24 @@ mlx/backend/cuda/lru_cache.h, mlx/backend/cuda/matmul.cpp, mlx/backend/cuda/primitives.cpp, + mlx/backend/cuda/ptx.cuh, mlx/backend/cuda/quantized/affine_quantize.cu, mlx/backend/cuda/quantized/convert_fp8.cu, mlx/backend/cuda/quantized/cublas_qqmm.cpp, mlx/backend/cuda/quantized/cublas_qqmm.h, mlx/backend/cuda/quantized/fp_quantize.cu, + mlx/backend/cuda/quantized/fp_quantize.cuh, mlx/backend/cuda/quantized/mxfp8_quantize.cuh, mlx/backend/cuda/quantized/no_qqmm_impl.cpp, mlx/backend/cuda/quantized/nvfp4_quantize.cuh, - mlx/backend/cuda/quantized/qmm/CMakeLists.txt, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu, - mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh, + mlx/backend/cuda/quantized/qmm/fp_qmv.cu, + mlx/backend/cuda/quantized/qmm/qmm_naive.cu, + mlx/backend/cuda/quantized/qmm/qmm_sm80.cu, + mlx/backend/cuda/quantized/qmm/qmm_sm90.cu, + mlx/backend/cuda/quantized/qmm/qmm_utils.h, + mlx/backend/cuda/quantized/qmm/qmm.cu, mlx/backend/cuda/quantized/qmm/qmm.h, + mlx/backend/cuda/quantized/qmm/qmv.cu, mlx/backend/cuda/quantized/qqmm_impl.cpp, mlx/backend/cuda/quantized/qqmm_impl.h, mlx/backend/cuda/quantized/qqmm_utils.cu, @@ -776,11 +821,13 @@ mlx/backend/gpu/copy.h, mlx/backend/gpu/device_info.h, mlx/backend/gpu/eval.h, + mlx/backend/gpu/scan.h, mlx/backend/gpu/slicing.h, mlx/backend/metal/allocator.h, mlx/backend/metal/binary.h, mlx/backend/metal/CMakeLists.txt, mlx/backend/metal/device.h, + mlx/backend/metal/event.h, mlx/backend/metal/jit/includes.h, mlx/backend/metal/jit/indexing.h, mlx/backend/metal/kernels.h, @@ -802,6 +849,8 @@ mlx/backend/metal/kernels/copy.h, mlx/backend/metal/kernels/copy.metal, mlx/backend/metal/kernels/defines.h, + mlx/backend/metal/kernels/dot.h, + mlx/backend/metal/kernels/dot.metal, mlx/backend/metal/kernels/erf.h, mlx/backend/metal/kernels/expm1f.h, mlx/backend/metal/kernels/fence.metal, @@ -817,6 +866,7 @@ mlx/backend/metal/kernels/fp8.h, mlx/backend/metal/kernels/gemv_masked.h, mlx/backend/metal/kernels/gemv_masked.metal, + mlx/backend/metal/kernels/gemv.h, mlx/backend/metal/kernels/gemv.metal, mlx/backend/metal/kernels/hadamard.h, mlx/backend/metal/kernels/indexing/gather_axis.h, @@ -850,6 +900,8 @@ mlx/backend/metal/kernels/scan.h, mlx/backend/metal/kernels/scan.metal, mlx/backend/metal/kernels/sdpa_vector.h, + mlx/backend/metal/kernels/searchsorted.h, + mlx/backend/metal/kernels/searchsorted.metal, mlx/backend/metal/kernels/softmax.h, mlx/backend/metal/kernels/softmax.metal, mlx/backend/metal/kernels/sort.h, @@ -889,6 +941,8 @@ mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.metal, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.metal, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h, + mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.metal, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.metal, mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h, @@ -917,7 +971,6 @@ mlx/backend/metal/nojit_kernels.cpp, mlx/backend/metal/reduce.h, mlx/backend/metal/resident.h, - mlx/backend/metal/scan.h, mlx/backend/metal/ternary.h, mlx/backend/metal/unary.h, mlx/backend/metal/utils.h, @@ -941,24 +994,38 @@ mlx/distributed/CMakeLists.txt, mlx/distributed/distributed_impl.h, mlx/distributed/distributed.h, + mlx/distributed/jaccl/.gitignore, mlx/distributed/jaccl/CMakeLists.txt, mlx/distributed/jaccl/jaccl.cpp, mlx/distributed/jaccl/jaccl.h, - mlx/distributed/jaccl/mesh_impl.h, - mlx/distributed/jaccl/mesh.cpp, - mlx/distributed/jaccl/mesh.h, - mlx/distributed/jaccl/ring_impl.h, - mlx/distributed/jaccl/ring.cpp, - mlx/distributed/jaccl/ring.h, - mlx/distributed/jaccl/utils.cpp, - mlx/distributed/jaccl/utils.h, + mlx/distributed/jaccl/lib/CMakeLists.txt, + mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp, + mlx/distributed/jaccl/lib/examples/CMakeLists.txt, + mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp, + mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp, + mlx/distributed/jaccl/lib/examples/minimal_env.cpp, + mlx/distributed/jaccl/lib/jaccl/group.h, + mlx/distributed/jaccl/lib/jaccl/jaccl.cpp, + mlx/distributed/jaccl/lib/jaccl/jaccl.h, + mlx/distributed/jaccl/lib/jaccl/mesh_impl.h, + mlx/distributed/jaccl/lib/jaccl/mesh.cpp, + mlx/distributed/jaccl/lib/jaccl/mesh.h, + mlx/distributed/jaccl/lib/jaccl/rdma.cpp, + mlx/distributed/jaccl/lib/jaccl/rdma.h, + mlx/distributed/jaccl/lib/jaccl/reduction_ops.h, + mlx/distributed/jaccl/lib/jaccl/ring_impl.h, + mlx/distributed/jaccl/lib/jaccl/ring.cpp, + mlx/distributed/jaccl/lib/jaccl/ring.h, + mlx/distributed/jaccl/lib/jaccl/tcp.cpp, + mlx/distributed/jaccl/lib/jaccl/tcp.h, + mlx/distributed/jaccl/lib/jaccl/threadpool.h, + mlx/distributed/jaccl/lib/jaccl/types.h, + mlx/distributed/jaccl/lib/README.md, mlx/distributed/mpi/CMakeLists.txt, mlx/distributed/mpi/mpi_declarations.h, mlx/distributed/mpi/mpi.cpp, mlx/distributed/mpi/mpi.h, mlx/distributed/nccl/CMakeLists.txt, - mlx/distributed/nccl/nccl_stub/CMakeLists.txt, - mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp, mlx/distributed/nccl/nccl.cpp, mlx/distributed/nccl/nccl.h, mlx/distributed/ops.h, @@ -971,6 +1038,7 @@ mlx/dtype_utils.h, mlx/dtype.h, mlx/einsum.h, + mlx/error.h, mlx/event.h, mlx/export_impl.h, mlx/export.h, @@ -1006,6 +1074,7 @@ mlx/utils.h, mlx/version.h, pyproject.toml, + python/mlx/__array_api_info.py, python/mlx/__main__.py, python/mlx/_distributed_utils/common.py, python/mlx/_distributed_utils/config.py, @@ -1062,7 +1131,9 @@ python/src/mlx_func.h, python/src/mlx.cpp, python/src/ops.cpp, + python/src/print.cpp, python/src/random.cpp, + python/src/random.h, python/src/small_vector.h, python/src/stream.cpp, python/src/transforms.cpp, @@ -1071,7 +1142,6 @@ python/src/utils.cpp, python/src/utils.h, python/tests/__main__.py, - python/tests/cuda_skip.py, python/tests/mlx_distributed_tests.py, python/tests/mlx_tests.py, python/tests/mpi_test_distributed.py, @@ -1105,6 +1175,7 @@ python/tests/test_quantized.py, python/tests/test_random.py, python/tests/test_reduce.py, + python/tests/test_threads.py, python/tests/test_tree.py, python/tests/test_upsample.py, python/tests/test_vmap.py, @@ -1129,6 +1200,7 @@ tests/load_tests.cpp, tests/ops_tests.cpp, tests/random_tests.cpp, + tests/residency_tests.cpp, tests/scheduler_tests.cpp, tests/tests.cpp, tests/utils_tests.cpp, @@ -1144,12 +1216,14 @@ "mlx-allocator.h", "mlx-api.h", "mlx-array.h", + "mlx-backend-common-metal_kernel.h", "mlx-backend-common-utils.h", "mlx-backend-cpu-encoder.h", "mlx-backend-cuda-cuda.h", "mlx-backend-gpu-eval.h", "mlx-backend-metal-device.h", "mlx-backend-metal-metal.h", + "mlx-backend-metal-resident.h", "mlx-c-array.h", "mlx-c-closure.h", "mlx-c-compile.h", @@ -1161,6 +1235,7 @@ "mlx-c-export.h", "mlx-c-fast.h", "mlx-c-fft.h", + "mlx-c-graph_utils.h", "mlx-c-half.h", "mlx-c-io_types.h", "mlx-c-io.h", @@ -1184,11 +1259,13 @@ "mlx-distributed-ops.h", "mlx-dtype.h", "mlx-einsum.h", + "mlx-error.h", "mlx-event.h", "mlx-export_impl.h", "mlx-export.h", "mlx-fast.h", "mlx-fft.h", + "mlx-graph_utils.h", "mlx-io-load.h", "mlx-io.h", "mlx-linalg.h", @@ -1213,6 +1290,7 @@ C3CB32A92EB168CD0029A645 /* Exceptions for "mlx-generated" folder in "Cmlx" target */ = { isa = PBXFileSystemSynchronizedBuildFileExceptionSet; membershipExceptions = ( + cuda/cuda_jit_sources.h, metal/arange.h, metal/atomic.h, metal/bf16_math.h, @@ -1224,6 +1302,7 @@ metal/complex.h, metal/copy.h, metal/defines.h, + metal/dot.h, metal/erf.h, metal/expm1f.h, metal/fft.h, @@ -1234,6 +1313,7 @@ metal/fp4.h, metal/fp8.h, metal/gemv_masked.h, + metal/gemv.h, metal/hadamard.h, metal/indexing/gather_axis.h, metal/indexing/gather_front.h, @@ -1256,6 +1336,7 @@ metal/reduction/reduce_row.h, metal/scan.h, metal/sdpa_vector.h, + metal/searchsorted.h, metal/softmax.h, metal/sort.h, metal/steel/attn/attn.h, @@ -1283,6 +1364,7 @@ metal/steel/gemm/kernels/steel_gemm_gather_nax.h, metal/steel/gemm/kernels/steel_gemm_gather.h, metal/steel/gemm/kernels/steel_gemm_masked.h, + metal/steel/gemm/kernels/steel_gemm_segmented_nax.h, metal/steel/gemm/kernels/steel_gemm_segmented.h, metal/steel/gemm/kernels/steel_gemm_splitk_nax.h, metal/steel/gemm/kernels/steel_gemm_splitk.h, @@ -1310,14 +1392,22 @@ ChangeLog.md, CMakeLists.txt, CONTRIBUTING.md, + doc/api.md, doc/api.rst, doc/build.py, + "doc/ChangeLog-old.md", doc/CMakeLists.txt, doc/conf.py, doc/contents.rst, + doc/fmt.css, + doc/fmt.js, doc/fmt.less, + "doc/get-started.md", + doc/index.md, doc/index.rst, + doc/perf.svg, "doc/python-license.txt", + doc/syntax.md, doc/syntax.rst, doc/usage.rst, include/fmt/args.h, @@ -1343,6 +1433,7 @@ "support/build-docs.py", support/build.gradle, "support/C++.sublime-syntax", + "support/check-commits", support/cmake/FindSetEnv.cmake, "support/cmake/fmt-config.cmake.in", support/cmake/fmt.pc.in, @@ -1350,8 +1441,13 @@ "support/compute-powers.py", support/docopt.py, support/manage.py, + support/mkdocs, + support/mkdocs.yml, support/printable.py, + support/python/mkdocstrings_handlers/cxx/__init__.py, + support/python/mkdocstrings_handlers/cxx/templates/README, support/README, + support/release.py, support/Vagrantfile, "test/add-subdirectory-test/CMakeLists.txt", "test/add-subdirectory-test/main.cc", @@ -1402,6 +1498,7 @@ "test/noexception-test.cc", "test/os-test.cc", "test/ostream-test.cc", + "test/perf-sanity.cc", "test/posix-mock-test.cc", "test/posix-mock.h", "test/printf-test.cc", diff --git a/xcode/xcconfig/Cmlx.xcconfig b/xcode/xcconfig/Cmlx.xcconfig index 37d668607..c6b2b5f4f 100644 --- a/xcode/xcconfig/Cmlx.xcconfig +++ b/xcode/xcconfig/Cmlx.xcconfig @@ -13,5 +13,5 @@ USER_HEADER_SEARCH_PATHS = $(SRCROOT)/../Source/Cmlx/mlx-c $(SRCROOT)/../Source/ MTL_COMPILER_FLAGS = -Wall -Wextra -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions -GCC_PREPROCESSOR_DEFINITIONS = _METAL_=1 SWIFTPM_BUNDLE=\"com.apple.mlx.Cmlx\" MLX_VERSION=\"0.31.1\" MLX_USE_ACCELERATE=1 METAL_PATH=\"default.metallib\" ACCELERATE_NEW_LAPACK=1 +GCC_PREPROCESSOR_DEFINITIONS = _METAL_=1 SWIFTPM_BUNDLE=\"com.apple.mlx.Cmlx\" MLX_VERSION=\"0.32.0\" MLX_USE_ACCELERATE=1 METAL_PATH=\"default.metallib\" ACCELERATE_NEW_LAPACK=1 From ea8a179690170ca891a97bc0473198ab1ecda5f4 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Fri, 11 Sep 2026 22:30:51 +0400 Subject: [PATCH 26/39] Fix wired-memory ticket lifecycle during cancellation (#471) --- Source/MLX/WiredMemory.swift | 88 +++---- .../MLXTests/WiredMemoryLifecycleTests.swift | 244 ++++++++++++++++++ 2 files changed, 281 insertions(+), 51 deletions(-) create mode 100644 Tests/MLXTests/WiredMemoryLifecycleTests.swift diff --git a/Source/MLX/WiredMemory.swift b/Source/MLX/WiredMemory.swift index 455ab0a11..2fee142ca 100644 --- a/Source/MLX/WiredMemory.swift +++ b/Source/MLX/WiredMemory.swift @@ -291,43 +291,23 @@ public struct WiredMemoryTicket: Sendable, Identifiable { } extension WiredMemoryTicket { - /// Guards against calling `end()` more than once when task cancellation - /// races with normal completion inside `withWiredLimit`. - private final class EndOnceGuard: @unchecked Sendable { - private var _ended = false - private let _lock = NSLock() - - /// Returns `true` exactly once; all subsequent calls return `false`. - func tryMark() -> Bool { - _lock.lock() - defer { _lock.unlock() } - if _ended { return false } - _ended = true - return true - } - } - - /// Convenience wrapper that guarantees start/end pairing and is safe under - /// cancellation. This is the recommended pattern for inference. + /// Keeps the ticket active until the body returns or throws, including during cancellation. + /// + /// If admission is cancelled, the body still runs in the cancelled task so it can + /// clean up. Check cancellation before starting work that requires admission. + /// Use a distinct ticket for each overlapping scope. public static func withWiredLimit( _ ticket: WiredMemoryTicket, _ body: () async throws -> R ) async rethrows -> R { _ = await ticket.start() - let guard_ = EndOnceGuard() - return try await withTaskCancellationHandler { - do { - let result = try await body() - if guard_.tryMark() { _ = await ticket.end() } - return result - } catch { - if guard_.tryMark() { _ = await ticket.end() } - throw error - } - } onCancel: { - Task { - if guard_.tryMark() { _ = await ticket.end() } - } + do { + let result = try await body() + _ = await ticket.end() + return result + } catch { + _ = await ticket.end() + throw error } } @@ -366,6 +346,11 @@ public actor WiredMemoryManager { let kind: WiredMemoryTicketKind } + private struct AdmissionWaiter { + let ticketID: UUID + let continuation: CheckedContinuation + } + /// Stable grouping key for policies. private enum PolicyKey: Hashable { case identifier(AnyHashable) @@ -379,8 +364,8 @@ public actor WiredMemoryManager { private var policies: [PolicyKey: any WiredMemoryPolicy] = [:] /// Last limit applied by the manager. private var currentLimit: Int? - /// Admission waiters parked while capacity is unavailable. - private var waiters: [UUID: CheckedContinuation] = [:] + /// Each wait has its own identity so late cancellation cannot affect a later wait. + private var waiters: [UUID: AdmissionWaiter] = [:] /// Timestamp used to enforce shrink cooldown. private var lastLimitChange: Date? /// Hysteresis configuration for shrink behavior. @@ -458,9 +443,6 @@ public actor WiredMemoryManager { var baselineValue = resolveBaselineAndEmit(refresh: baseline == nil || !hasActiveWork()) if tickets[id] != nil { - #if DEBUG - assertionFailure("Ticket already started: \(id)") - #endif emit( kind: .ticketStartIgnored, ticketID: id, @@ -489,18 +471,25 @@ public actor WiredMemoryManager { return currentLimit ?? baselineValue } + let waiterID = UUID() await withTaskCancellationHandler { await withCheckedContinuation { continuation in - waiters[id] = continuation + waiters[waiterID] = AdmissionWaiter(ticketID: id, continuation: continuation) } - } onCancel: { [id] in - Task { await self.cancelWaiter(id: id) } + } onCancel: { + Task { await self.cancelWaiter(id: waiterID) } } baselineValue = resolveBaselineAndEmit(refresh: baseline == nil || !hasActiveWork()) if Task.isCancelled { return currentLimit ?? baselineValue } + if tickets[id] != nil { + emit( + kind: .ticketStartIgnored, ticketID: id, size: normalizedSize, + baseline: baselineValue) + return currentLimit ?? baselineValue + } } emit( @@ -533,8 +522,9 @@ public actor WiredMemoryManager { /// If this was the last ticket, the manager restores the baseline and /// clears internal state. public func end(id: UUID, policy: any WiredMemoryPolicy) -> Int { - if let waiter = waiters.removeValue(forKey: id) { - waiter.resume() + for (waiterID, waiter) in waiters where waiter.ticketID == id { + waiters.removeValue(forKey: waiterID) + waiter.continuation.resume() } guard WiredMemoryBackend.isSupported || policyOnlyMode else { @@ -545,9 +535,6 @@ public actor WiredMemoryManager { } guard let state = tickets.removeValue(forKey: id) else { - #if DEBUG - assertionFailure("Ticket not active: \(id)") - #endif emit(kind: .ticketEndIgnored, ticketID: id) return currentLimit ?? baseline ?? 0 } @@ -774,8 +761,8 @@ public actor WiredMemoryManager { guard !waiters.isEmpty else { return } let pending = waiters waiters.removeAll() - for (_, continuation) in pending { - continuation.resume() + for waiter in pending.values { + waiter.continuation.resume() } } @@ -786,10 +773,9 @@ public actor WiredMemoryManager { /// Cancel an admission waiter and emit a debug event. private func cancelWaiter(id: UUID) { - if let waiter = waiters.removeValue(forKey: id) { - waiter.resume() - } - emit(kind: .admissionCancelled, ticketID: id) + guard let waiter = waiters.removeValue(forKey: id) else { return } + waiter.continuation.resume() + emit(kind: .admissionCancelled, ticketID: waiter.ticketID) } #if DEBUG diff --git a/Tests/MLXTests/WiredMemoryLifecycleTests.swift b/Tests/MLXTests/WiredMemoryLifecycleTests.swift new file mode 100644 index 000000000..3978ca694 --- /dev/null +++ b/Tests/MLXTests/WiredMemoryLifecycleTests.swift @@ -0,0 +1,244 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +final class WiredMemoryLifecycleTests: XCTestCase { + private enum Failure: Error { case expected } + private struct Policy: WiredMemoryPolicy, Hashable { + let capacity: Int + + func limit(baseline: Int, activeSizes: [Int]) -> Int { + baseline + activeSizes.reduce(0, +) + } + + func canAdmit(baseline: Int, activeSizes: [Int], newSize: Int) -> Bool { + activeSizes.reduce(0, +) + newSize <= capacity + } + } + + private actor Gate { + private var isOpen = false + private var continuation: CheckedContinuation? + + func wait() async { + if isOpen { return } + await withCheckedContinuation { continuation = $0 } + } + + func open() { + isOpen = true + continuation?.resume() + continuation = nil + } + } + + private static func manager() -> WiredMemoryManager { + .makeForTesting( + configuration: .init( + policyOnlyWhenUnsupported: true, baselineOverride: 1024, + useRecommendedWorkingSetWhenUnsupported: false)) + } + + func testDuplicateStartAndUnmatchedEndAreIdempotent() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let ticket = WiredMemoryTicket(size: 64, policy: Policy(capacity: 64), manager: manager) + let events = await manager.events() + _ = await ticket.end() + _ = await ticket.start() + _ = await ticket.start() + _ = await ticket.end() + _ = await ticket.end() + + var started = 0 + var ignoredStarts = 0 + var ended = 0 + var ignoredEnds = 0 + for await event in events { + switch event.kind { + case .ticketStarted: started += 1 + case .ticketStartIgnored: ignoredStarts += 1 + case .ticketEnded: ended += 1 + case .ticketEndIgnored: ignoredEnds += 1 + default: break + } + if ignoredEnds == 2 { break } + } + XCTAssertEqual(started, 1) + XCTAssertEqual(ignoredStarts, 1) + XCTAssertEqual(ended, 1) + } + } + + func testCancellationWhileAdmissionIsDeniedStillRunsCleanup() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let ticket = WiredMemoryTicket(size: 1, policy: Policy(capacity: 0), manager: manager) + let events = await manager.events() + let task = Task { + await ticket.withWiredLimit { + XCTAssertTrue(Task.isCancelled) + return 123 + } + } + for await event in events { + if event.kind == .admissionWait { break } + } + task.cancel() + let result = await task.value + XCTAssertEqual(result, 123) + // The cancelled waiter must be removed before the scope returns. + await manager.updateConfiguration { $0.baselineOverride = 2048 } + } + } + + func testCancellationKeepsTicketUntilBodyFinishes() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let ticket = WiredMemoryTicket(size: 64, policy: Policy(capacity: 64), manager: manager) + let events = await manager.events() + let gate = Gate() + let started = expectation(description: "body started") + let endedEarly = expectation(description: "ticket ended before body completed") + endedEarly.isInverted = true + let observer = Task { + for await event in events { + if event.kind == .ticketEnded { + endedEarly.fulfill() + break + } + } + } + let task = Task { + await ticket.withWiredLimit { + started.fulfill() + await gate.wait() + XCTAssertTrue(Task.isCancelled) + } + } + await fulfillment(of: [started], timeout: 5) + task.cancel() + await fulfillment(of: [endedEarly], timeout: 0.2) + observer.cancel() + await observer.value + await gate.open() + await task.value + // End is awaited; no active ticket may remain after task.value. + await manager.updateConfiguration { $0.baselineOverride = 2048 } + } + } + + func testThrowingBodyReleasesTicketBeforeReturning() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let ticket = WiredMemoryTicket(size: 64, policy: Policy(capacity: 64), manager: manager) + do { + try await ticket.withWiredLimit { throw Failure.expected } + XCTFail("Expected the body's error") + } catch { + XCTAssertTrue(error is Failure) + } + await manager.updateConfiguration { $0.baselineOverride = 2048 } + } + } + + func testAlreadyCancelledScopeCanCleanUpWithoutAdmission() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let ticket = WiredMemoryTicket(size: 1, policy: Policy(capacity: 0), manager: manager) + let task = Task { + withUnsafeCurrentTask { $0?.cancel() } + return await ticket.withWiredLimit { + XCTAssertTrue(Task.isCancelled) + return 123 + } + } + let result = await task.value + XCTAssertEqual(result, 123) + await manager.updateConfiguration { $0.baselineOverride = 2048 } + } + } + + func testCancellationRacingCompletionAlwaysFinishesCleanup() async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + for index in 0 ..< 100 { + let ticket = WiredMemoryTicket( + size: 64, policy: Policy(capacity: 64), manager: manager) + let task = Task { + await ticket.withWiredLimit { + await Task.yield() + return index + } + } + if index.isMultiple(of: 2) { await Task.yield() } + task.cancel() + let result = await task.value + XCTAssertEqual(result, index) + await manager.updateConfiguration { $0.baselineOverride = 1024 } + } + } + } + + func testConcurrentStartsOfWaitingTicketDoNotLoseWaiters() async { + await checkConcurrentStarts(cancelFirst: false) + } + + func testCancellingOneStartDoesNotCancelAnotherWaiterForSameTicket() async { + await checkConcurrentStarts(cancelFirst: true) + } + + private func checkConcurrentStarts(cancelFirst: Bool) async { + await Device.withDefaultDevice(.cpu) { + let manager = Self.manager() + let policy = Policy(capacity: 64) + let blocker = WiredMemoryTicket(size: 64, policy: policy, manager: manager) + let ticket = WiredMemoryTicket(size: 32, policy: policy, manager: manager) + _ = await blocker.start() + let events = await manager.events() + let firstFinished = expectation(description: "first start returned") + let secondFinished = expectation(description: "second start returned") + let first = Task { + _ = await ticket.start() + firstFinished.fulfill() + } + // Ensure the first waiter is installed before adding the second. + for await event in events { + if event.kind == .admissionWait { break } + } + let secondEvents = await manager.events() + let second = Task { + _ = await ticket.start() + secondFinished.fulfill() + } + for await event in secondEvents { + if event.kind == .admissionWait { break } + } + if cancelFirst { + first.cancel() + let result = await XCTWaiter.fulfillment(of: [firstFinished], timeout: 5) + guard result == .completed else { + XCTFail("Cancelling the first start did not resume its continuation") + second.cancel() + _ = await blocker.end() + return + } + } + _ = await blocker.end() + let expected = cancelFirst ? [secondFinished] : [firstFinished, secondFinished] + let result = await XCTWaiter.fulfillment(of: expected, timeout: 5) + guard result == .completed else { + XCTFail("A duplicate start lost an admission continuation") + first.cancel() + second.cancel() + return + } + await first.value + await second.value + _ = await ticket.end() + await manager.updateConfiguration { $0.baselineOverride = 2048 } + } + } +} From a4c3b54f3ab80e452364ee912761c0f3a4e88425 Mon Sep 17 00:00:00 2001 From: Lukas Zima <76228598+Silenterc@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:20:38 +0200 Subject: [PATCH 27/39] Fix condition typo (#473) --- Source/MLX/Ops.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index 362f16de0..97091cac1 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -3356,8 +3356,8 @@ public func trunc(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXAr /// /// - Parameters: /// - condition: condition array -/// - a: input selected from where condiiton is non-zero or `true` -/// - b: input selected from where condiiton is zero or `false` +/// - a: input selected from where condition is non-zero or `true` +/// - b: input selected from where condition is zero or `false` /// - stream: stream or device to evaluate on /// /// ### See Also @@ -3380,8 +3380,8 @@ public func `where`( /// /// - Parameters: /// - condition: condition array -/// - a: input selected from where condiiton is non-zero or `true` -/// - b: input selected from where condiiton is zero or `false` +/// - a: input selected from where condition is non-zero or `true` +/// - b: input selected from where condition is zero or `false` /// - stream: stream or device to evaluate on /// /// ### See Also From 555692ae79aba0fdb53ef665537c0ba24492df46 Mon Sep 17 00:00:00 2001 From: Shubhransh Gupta <54713516+shubhransh-gupta@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:50:49 +0530 Subject: [PATCH 28/39] fix(ops): fix convolve padRight for even-sized kernels in same mode (#466) (#467) --- Source/MLX/Ops.swift | 2 +- Tests/MLXTests/OpsTests.swift | 28 ++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index 97091cac1..0c689303b 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -1003,7 +1003,7 @@ public func convolve( padding = weightSize / 2 } else { let padLeft = weightSize / 2 - let padRight = max(0, padLeft / 2 - 1) + let padRight = max(0, padLeft - 1) input = padded(input, widths: [0, [padLeft, padRight], 0], stream: stream) } diff --git a/Tests/MLXTests/OpsTests.swift b/Tests/MLXTests/OpsTests.swift index 46fd06cc2..4301ac048 100644 --- a/Tests/MLXTests/OpsTests.swift +++ b/Tests/MLXTests/OpsTests.swift @@ -113,4 +113,32 @@ class OpsTests: XCTestCase { XCTAssertNil(b2) } + func testConvolve() { + let a = MLXArray([1, 2, 3, 4, 5] as [Float]) + let vEven = MLXArray([1, 2, 3, 4] as [Float]) + let vOdd = MLXArray([1, 2, 3] as [Float]) + + // Even kernel size + let fullEven = convolve(a, vEven, mode: .full) + assertEqual(fullEven, MLXArray([1, 4, 10, 20, 30, 34, 31, 20] as [Float])) + + let validEven = convolve(a, vEven, mode: .valid) + assertEqual(validEven, MLXArray([20, 30] as [Float])) + + let sameEven = convolve(a, vEven, mode: .same) + XCTAssertEqual(sameEven.shape, [5]) + assertEqual(sameEven, MLXArray([4, 10, 20, 30, 34] as [Float])) + + // Odd kernel size + let fullOdd = convolve(a, vOdd, mode: .full) + assertEqual(fullOdd, MLXArray([1, 4, 10, 16, 22, 22, 15] as [Float])) + + let validOdd = convolve(a, vOdd, mode: .valid) + assertEqual(validOdd, MLXArray([10, 16, 22] as [Float])) + + let sameOdd = convolve(a, vOdd, mode: .same) + XCTAssertEqual(sameOdd.shape, [5]) + assertEqual(sameOdd, MLXArray([4, 10, 16, 22, 22] as [Float])) + } + } From 2bebe4e9ad127758ebcd76c6ad45a1740d0d2852 Mon Sep 17 00:00:00 2001 From: Lukas Zima <76228598+Silenterc@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:09:25 +0200 Subject: [PATCH 29/39] Add countNonzero to the Swift API (#479) * Add countNonzero + docs * Add tests --- .../Organization/reduction.md | 6 ++ Source/MLX/MLXArray+Ops.swift | 59 +++++++++++++++++ Source/MLX/Ops+Array.swift | 63 +++++++++++++++++++ Tests/MLXTests/MLXArray+OpsTests.swift | 15 +++++ Tests/MLXTests/OpsTests.swift | 15 +++++ 5 files changed, 158 insertions(+) diff --git a/Source/MLX/Documentation.docc/Organization/reduction.md b/Source/MLX/Documentation.docc/Organization/reduction.md index f34b0eb74..4d7e989d9 100644 --- a/Source/MLX/Documentation.docc/Organization/reduction.md +++ b/Source/MLX/Documentation.docc/Organization/reduction.md @@ -33,6 +33,9 @@ See also and ### MLXArray Aggregating Reduction Functions +- ``MLXArray/countNonzero(keepDims:stream:)`` +- ``MLXArray/countNonzero(axis:keepDims:stream:)`` +- ``MLXArray/countNonzero(axes:keepDims:stream:)`` - ``MLXArray/logSumExp(keepDims:stream:)`` - ``MLXArray/logSumExp(axis:keepDims:stream:)`` - ``MLXArray/logSumExp(axes:keepDims:stream:)`` @@ -66,6 +69,9 @@ See also and ### Aggregating Reduction Free Functions +- ``countNonzero(_:keepDims:stream:)`` +- ``countNonzero(_:axis:keepDims:stream:)`` +- ``countNonzero(_:axes:keepDims:stream:)`` - ``logSumExp(_:keepDims:stream:)`` - ``logSumExp(_:axis:keepDims:stream:)`` - ``logSumExp(_:axes:keepDims:stream:)`` diff --git a/Source/MLX/MLXArray+Ops.swift b/Source/MLX/MLXArray+Ops.swift index f9927a9f7..bb80a16f7 100644 --- a/Source/MLX/MLXArray+Ops.swift +++ b/Source/MLX/MLXArray+Ops.swift @@ -1528,6 +1528,65 @@ extension MLXArray { return MLXArray(result) } + /// Count the number of non-zero elements in the array over the given axes. + /// + /// - Parameters: + /// - axes: axes to reduce over + /// - keepDims: if `true` keep the reduced axes as singleton dimensions + /// - stream: stream or device to evaluate on + /// + /// ### See Also + /// - + /// - ``countNonzero(axis:keepDims:stream:)`` + /// - ``countNonzero(keepDims:stream:)`` + /// - ``countNonzero(_:axes:keepDims:stream:)`` + public func countNonzero( + axes: some Collection, keepDims: Bool = false, stream: StreamOrDevice = .default + ) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero_axes(&result, ctx, axes.asInt32, axes.count, keepDims, stream.ctx) + return MLXArray(result) + } + + /// Count the number of non-zero elements in the array over the given axis. + /// + /// - Parameters: + /// - axis: axis to reduce over + /// - keepDims: if `true` keep reduced axis as singleton dimension + /// - stream: stream or device to evaluate on + /// + /// ### See Also + /// - + /// - ``countNonzero(axes:keepDims:stream:)`` + /// - ``countNonzero(keepDims:stream:)`` + /// - ``countNonzero(_:axes:keepDims:stream:)`` + public func countNonzero( + axis: Int, keepDims: Bool = false, stream: StreamOrDevice = .default + ) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero_axis(&result, ctx, axis.int32, keepDims, stream.ctx) + return MLXArray(result) + } + + /// Count the number of non-zero elements in the array over all axes. + /// + /// - Parameters: + /// - keepDims: if `true` keep the reduced axes as singleton dimensions + /// - stream: stream or device to evaluate on + /// + /// ### See Also + /// - + /// - ``countNonzero(axes:keepDims:stream:)`` + /// - ``countNonzero(axis:keepDims:stream:)`` + /// - ``countNonzero(_:axes:keepDims:stream:)`` + public func countNonzero( + keepDims: Bool = false, stream: StreamOrDevice = .default + ) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero(&result, ctx, keepDims, stream.ctx) + return MLXArray(result) + } + /// Return the cumulative maximum of the elements along the given axis. /// /// ```swift diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index f581d40b2..b21d857e1 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -442,6 +442,69 @@ public func conjugate(_ array: MLXArray, stream: StreamOrDevice = .default) -> M return MLXArray(result) } +/// Count the number of non-zero elements in the array over the given axes. +/// +/// - Parameters: +/// - array: input array +/// - axes: axes to reduce over +/// - keepDims: if `true` keep the reduced axes as singleton dimensions +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - ``countNonzero(_:axis:keepDims:stream:)`` +/// - ``countNonzero(_:keepDims:stream:)`` +/// - ``MLXArray/countNonzero(axes:keepDims:stream:)`` +public func countNonzero( + _ array: MLXArray, axes: some Collection, keepDims: Bool = false, + stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero_axes(&result, array.ctx, axes.asInt32, axes.count, keepDims, stream.ctx) + return MLXArray(result) +} + +/// Count the number of non-zero elements in the array over the given axis. +/// +/// - Parameters: +/// - array: input array +/// - axis: axis to reduce over +/// - keepDims: if `true` keep reduced axis as singleton dimension +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - ``countNonzero(_:axes:keepDims:stream:)`` +/// - ``countNonzero(_:keepDims:stream:)`` +/// - ``MLXArray/countNonzero(axes:keepDims:stream:)`` +public func countNonzero( + _ array: MLXArray, axis: Int, keepDims: Bool = false, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero_axis(&result, array.ctx, axis.int32, keepDims, stream.ctx) + return MLXArray(result) +} + +/// Count the number of non-zero elements in the array over all axes. +/// +/// - Parameters: +/// - array: input array +/// - keepDims: if `true` keep the reduced axes as singleton dimensions +/// - stream: stream or device to evaluate on +/// +/// ### See Also +/// - +/// - ``countNonzero(_:axes:keepDims:stream:)`` +/// - ``countNonzero(_:axis:keepDims:stream:)`` +/// - ``MLXArray/countNonzero(axes:keepDims:stream:)`` +public func countNonzero( + _ array: MLXArray, keepDims: Bool = false, stream: StreamOrDevice = .default +) -> MLXArray { + var result = mlx_array_new() + mlx_count_nonzero(&result, array.ctx, keepDims, stream.ctx) + return MLXArray(result) +} + /// Return the cumulative maximum of the elements along the given axis. /// /// ```swift diff --git a/Tests/MLXTests/MLXArray+OpsTests.swift b/Tests/MLXTests/MLXArray+OpsTests.swift index 94bf78dfb..5797efc20 100644 --- a/Tests/MLXTests/MLXArray+OpsTests.swift +++ b/Tests/MLXTests/MLXArray+OpsTests.swift @@ -181,4 +181,19 @@ class MLXArrayOpsTests: XCTestCase { XCTAssertEqual(r2, [Complex(-4, 24)]) } + func testCountNonzero() { + let a = MLXArray([0, 1, 2, 0, 3, 0], [2, 3]) + let resAxes0 = a.countNonzero(axes: [0]) + let resAxes1 = a.countNonzero(axes: [1]) + let resAxis0 = a.countNonzero(axis: 0) + let resAxis1 = a.countNonzero(axis: 1) + let resAll = a.countNonzero() + + assertEqual(resAxes0, [0, 2, 1]) + assertEqual(resAxes0, resAxis0) + assertEqual(resAxes1, [2, 1]) + assertEqual(resAxes1, resAxis1) + assertEqual(resAll, 3.asMLXArray(dtype: .int32)) + } + } diff --git a/Tests/MLXTests/OpsTests.swift b/Tests/MLXTests/OpsTests.swift index 4301ac048..27f58ce7c 100644 --- a/Tests/MLXTests/OpsTests.swift +++ b/Tests/MLXTests/OpsTests.swift @@ -86,6 +86,21 @@ class OpsTests: XCTestCase { XCTAssertEqual(c.dtype, .float32) } + func testCountNonzero() { + let a = MLXArray([0, 1, 2, 0, 3, 0], [2, 3]) + let resAxes0 = countNonzero(a, axes: [0]) + let resAxes1 = countNonzero(a, axes: [1]) + let resAxis0 = countNonzero(a, axis: 0) + let resAxis1 = countNonzero(a, axis: 1) + let resAll = countNonzero(a) + + assertEqual(resAxes0, [0, 2, 1]) + assertEqual(resAxes0, resAxis0) + assertEqual(resAxes1, [2, 1]) + assertEqual(resAxes1, resAxis1) + assertEqual(resAll, 3.asMLXArray(dtype: .int32)) + } + func testFlatten() { let a = zeros([4, 5, 6, 7]) let b = flatten(a, startAxis: 1, endAxis: 2) From 8541c5c9bf9951efb94d7c44f4df0505865a95c5 Mon Sep 17 00:00:00 2001 From: Ronald Mannak Date: Thu, 17 Sep 2026 09:02:28 -0700 Subject: [PATCH 30/39] Distributed support (#482) Add distributed communication and sharded neural network layers Port mlx.core.distributed and distributed layer primitives to Swift: - Backend & Primitives: - Enable the ring distributed backend (CPU socket-based communication). - Add MLXDistributed.Group with rank/size/split and 8 communication primitives (allSum, allGather, send, recv, etc.). - Default distributed operations to the CPU stream. - Make initialize and split throw on failure. - Layers & Optimization: - Add AllToShardedLinear, ShardedToAllLinear, and quantized counterparts (subclassing Linear / QuantizedLinear for seamless model replacement). - Implement shardLinear, shardInPlace (supporting predicate sharding), sumGradients, averageGradients, and clipGradNormSharded. - Ensure correct boundary checks and integer scaling for quantized layers. - Testing & Infrastructure: - Add a reusable multi-process test harness (DistributedHarness) for spawning local ranks via loopback sockets. - Port Python ring tests (collectives, dtypes, donation, send/recv). - Port layer, predicate sharding, and gradient tests; enable in CI. --- .github/workflows/pull_request.yml | 6 + Package.swift | 8 +- Source/MLX/Distributed.swift | 299 +++++++++ Source/MLX/Documentation.docc/MLX.md | 4 + Source/MLXNN/Distributed.swift | 754 ++++++++++++++++++++++ Source/MLXNN/Quantized.swift | 4 +- Tests/MLXTests/DistributedHarness.swift | 305 +++++++++ Tests/MLXTests/DistributedNNTests.swift | 670 +++++++++++++++++++ Tests/MLXTests/DistributedRingTests.swift | 191 ++++++ Tests/MLXTests/DistributedTests.swift | 136 ++++ xcode/MLX.xcodeproj/project.pbxproj | 2 +- 11 files changed, 2372 insertions(+), 7 deletions(-) create mode 100644 Source/MLX/Distributed.swift create mode 100644 Source/MLXNN/Distributed.swift create mode 100644 Tests/MLXTests/DistributedHarness.swift create mode 100644 Tests/MLXTests/DistributedNNTests.swift create mode 100644 Tests/MLXTests/DistributedRingTests.swift create mode 100644 Tests/MLXTests/DistributedTests.swift diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index 56452955c..d8c76493c 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -116,6 +116,12 @@ jobs: xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/CmlxTests.xctest xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest + - name: Run Distributed Tests SwiftPM (Xcode, macOS) + shell: sh + run: | + MLX_TEST_DISTRIBUTED=1 xcrun xctest -XCTest DistributedRingTests ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest + MLX_TEST_DISTRIBUTED=1 xcrun xctest -XCTest DistributedNNRingTests ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest + - name: Build xcodeproj (Xcode, macOS) shell: sh run: | diff --git a/Package.swift b/Package.swift index 4a7ee0076..1cb1bf138 100644 --- a/Package.swift +++ b/Package.swift @@ -248,10 +248,8 @@ let cmlx = Target.target( // vendor docs "vendor-README.md", - // example code + mlx-c distributed + // example code "mlx-c/examples", - "mlx-c/mlx/c/distributed.cpp", - "mlx-c/mlx/c/distributed_group.cpp", // vendored library, include header only "json", @@ -293,9 +291,9 @@ let cmlx = Target.target( "mlx/mlx/backend/metal/kernels", "mlx/mlx/backend/metal/nojit_kernels.cpp", - // do not build distributed support (yet) + // distributed backends: enable ring, disable MPI + NCCL + JACCL "mlx/mlx/distributed/mpi/mpi.cpp", - "mlx/mlx/distributed/ring/ring.cpp", + "mlx/mlx/distributed/ring/no_ring.cpp", "mlx/mlx/distributed/nccl/nccl.cpp", "mlx/mlx/distributed/jaccl/jaccl.cpp", "mlx/mlx/distributed/jaccl/lib", diff --git a/Source/MLX/Distributed.swift b/Source/MLX/Distributed.swift new file mode 100644 index 000000000..87d8f7659 --- /dev/null +++ b/Source/MLX/Distributed.swift @@ -0,0 +1,299 @@ +// Copyright © 2026 Apple Inc. + +import Cmlx +import Foundation + +/// Communication operations for running MLX across multiple processes. +/// +/// This mirrors `mlx.core.distributed` in the Python API. +/// +/// Forming a group can fail as soon as it is attempted, so +/// ``initialize(backend:strict:)`` and ``Group/split(color:key:)`` throw. +/// +/// The collectives and the point to point operations are lazy like every +/// other MLX operation. Invalid arguments -- e.g. sending in a group of size +/// one -- are reported when the operation is created, and communication +/// failures when its result is evaluated. Use ``withError(_:)-2wfiu`` and +/// ``checkedEval(_:)-(Any...)`` to receive either as a Swift error: +/// +/// ```swift +/// let group = try MLXDistributed.initialize(strict: true) +/// try withError { +/// let sum = MLXDistributed.allSum(x, group: group) +/// try checkedEval(sum) +/// } +/// ``` +/// +/// Every operation evaluates on the CPU stream by default, which is where the +/// ring, JACCL and MPI backends communicate. NCCL, available only in a CUDA +/// build, communicates on the GPU, so pass `stream: .gpu` for an NCCL group. +/// +/// Every process must create *and evaluate* these operations in the same +/// order. They are lazy, so the communication happens where the result is +/// evaluated: evaluating two collectives in one order in one process and in +/// the other order in another makes the processes exchange mismatched buffers +/// rather than report an error. Beware of iterating a `Dictionary`, whose +/// order differs from process to process. +public enum MLXDistributed { + + /// A communication backend. + /// + /// Use ``isAvailable(_:)`` to check whether MLX was built with support for + /// a given backend. Note that availability means MLX can instantiate the + /// backend, not that a communication group can actually be formed -- use + /// ``initialize(backend:strict:)`` with `strict: true` for that. + public enum Backend: String, CaseIterable, Sendable { + /// Select the first available backend. + case any + case ring + case mpi + case nccl + case jaccl + } + + /// Check if a communication backend is available. + /// + /// - Parameter backend: the backend to check for, defaulting to ``Backend/any`` + /// - Returns: whether MLX has the capability of instantiating that backend + public static func isAvailable(_ backend: Backend = .any) -> Bool { + mlx_distributed_is_available(backend.rawValue) + } + + /// A group of independent MLX processes that can communicate. + /// + /// Obtain a group with ``initialize(backend:strict:)``. + public final class Group: @unchecked Sendable { + + let ctx: mlx_distributed_group + + init(_ ctx: mlx_distributed_group) { + self.ctx = ctx + } + + deinit { + mlx_distributed_group_free(ctx) + } + + /// The rank of this process in the group. + public var rank: Int { + Int(mlx_distributed_group_rank(ctx)) + } + + /// The number of processes in the group. + public var size: Int { + Int(mlx_distributed_group_size(ctx)) + } + + /// Split the group into subgroups based on the provided color. + /// + /// Processes that use the same color go to the same group. The `key` + /// argument defines the rank in the new group -- the smaller the key + /// the smaller the rank. If the key is negative the rank in the + /// current group is used. + /// + /// Throws if the group cannot be split. The ring and JACCL backends do + /// not support splitting, and neither does a group of size one. + /// + /// - Parameters: + /// - color: a value to group processes into subgroups + /// - key: a key to optionally change the rank ordering of the processes + /// - Returns: the subgroup this process belongs to + public func split(color: Int, key: Int = -1) throws -> Group { + var result = mlx_distributed_group_new() + do { + try withError { + _ = mlx_distributed_group_split(&result, ctx, Int32(color), Int32(key)) + } + } catch { + mlx_distributed_group_free(result) + throw error + } + return Group(result) + } + } + + /// Initialize the distributed backend and return the global group. + /// + /// If no backend is configured and `strict` is `false` this returns a + /// group of size one, in which the collectives return their input + /// unchanged, so the same code also runs as a single process. Forming a + /// real group waits for the other processes to connect. + /// + /// Throws if the backend cannot be initialized: when `strict` is `true` and + /// no backend is configured, and -- even when `strict` is `false` -- when a + /// backend is configured but cannot form its group, e.g. the ring backend + /// with a malformed `MLX_HOSTFILE`. + /// + /// MLX caches the group for each backend for the life of the process, + /// including the group of size one that a non-strict call falls back to, + /// and a later strict call for that backend returns the cached group + /// instead of throwing. A non-strict ``Backend/any`` that finds no backend + /// caches its fallback under ``Backend/jaccl``. Initialize strictly + /// before anything else initializes the backend. + /// + /// - Parameters: + /// - backend: the backend to use, defaulting to ``Backend/any`` + /// - strict: if `true`, throw rather than fall back to a group of size one + /// when no backend can be initialized + /// - Returns: the global group + @discardableResult + public static func initialize(backend: Backend = .any, strict: Bool = false) throws -> Group { + var result = mlx_distributed_group_new() + do { + try withError { + _ = mlx_distributed_init(&result, strict, backend.rawValue) + } + } catch { + mlx_distributed_group_free(result) + throw error + } + return Group(result) + } + + // MARK: - Collectives + + // Note: mlx-c has no representation for an "unspecified" stream -- + // mlx_stream_get_() rejects a null handle -- so the backend's own default + // (Group::communication_stream) is unreachable from Swift. Python passes + // stream=None and lets ring/mpi/jaccl pick the CPU and nccl pick the GPU. + // The ring and JACCL backends both want the CPU, and the collectives have + // no GPU implementation, so the CPU stream is the default here. A CUDA + // build with nccl has to pass .gpu, which MLXDistributed documents. + + /// All reduce sum. + /// + /// Sum `x` across all processes in the group. + /// + /// - Parameters: + /// - x: input array + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func allSum( + _ x: MLXArray, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_all_sum(&result, x.ctx, group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// All reduce max. + /// + /// - Parameters: + /// - x: input array + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func allMax( + _ x: MLXArray, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_all_max(&result, x.ctx, group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// All reduce min. + /// + /// - Parameters: + /// - x: input array + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func allMin( + _ x: MLXArray, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_all_min(&result, x.ctx, group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// Gather arrays from all processes, concatenating along the first axis. + /// + /// - Parameters: + /// - x: input array + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func allGather( + _ x: MLXArray, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_all_gather(&result, x.ctx, group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// Sum `x` across the group and shard the result along the first axis. + /// + /// `x.dim(0)` must be divisible by the group size. The result is + /// equivalent to `allSum(x)` sliced to this process' chunk, but is + /// performed as a single reduce-scatter collective. + /// + /// Currently implemented by the NCCL and JACCL backends only. + /// + /// - Parameters: + /// - x: input array + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func sumScatter( + _ x: MLXArray, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_sum_scatter(&result, x.ctx, group.groupCtx, stream.ctx) + return MLXArray(result) + } + + // MARK: - Point to point + + /// Send `x` to the process with rank `dst`. + /// + /// - Parameters: + /// - x: input array + /// - dst: rank of the destination process + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + /// - Returns: an array identical to `x` which, when evaluated, performs the send + public static func send( + _ x: MLXArray, to dst: Int, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_send(&result, x.ctx, Int32(dst), group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// Receive an array with the given shape and type from the process with rank `src`. + /// + /// - Parameters: + /// - shape: shape of the array to receive + /// - dtype: type of the array to receive + /// - src: rank of the source process + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func recv( + _ shape: [Int], dtype: DType, from src: Int, group: Group? = nil, + stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + let shape = shape.asInt32 + mlx_distributed_recv( + &result, shape, shape.count, dtype.cmlxDtype, Int32(src), group.groupCtx, stream.ctx) + return MLXArray(result) + } + + /// Receive an array with the same shape and type as `x` from the process with rank `src`. + /// + /// - Parameters: + /// - x: array whose shape and type describe what will be received + /// - src: rank of the source process + /// - group: the group, or `nil` to use the global group + /// - stream: stream to evaluate on + public static func recvLike( + _ x: MLXArray, from src: Int, group: Group? = nil, stream: StreamOrDevice = .cpu + ) -> MLXArray { + var result = mlx_array_new() + mlx_distributed_recv_like(&result, x.ctx, Int32(src), group.groupCtx, stream.ctx) + return MLXArray(result) + } +} + +extension Optional where Wrapped == MLXDistributed.Group { + /// The underlying group handle, or a null handle meaning "the global group". + fileprivate var groupCtx: mlx_distributed_group { + self?.ctx ?? mlx_distributed_group(ctx: nil) + } +} diff --git a/Source/MLX/Documentation.docc/MLX.md b/Source/MLX/Documentation.docc/MLX.md index 5e2f98c5c..fca5773a8 100644 --- a/Source/MLX/Documentation.docc/MLX.md +++ b/Source/MLX/Documentation.docc/MLX.md @@ -130,3 +130,7 @@ See for configuration, best practices, and policy guidance. - ``Device`` - ``DeviceType`` - ``Stream`` + +### Distributed + +- ``MLXDistributed`` diff --git a/Source/MLXNN/Distributed.swift b/Source/MLXNN/Distributed.swift new file mode 100644 index 000000000..477c6088a --- /dev/null +++ b/Source/MLXNN/Distributed.swift @@ -0,0 +1,754 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX + +/// How a linear layer is sharded across a group. +public enum ShardingType: Sendable { + /// A common input becomes a sharded output. + case allToSharded + + /// A sharded input becomes a common output. + case shardedToAll +} + +/// Describes how an unsharded weight is composed. +/// +/// A fused QKV matrix, for example, is three segments stacked together and +/// each has to be sharded separately. +public enum Segments: Sendable, Equatable { + /// The weight is `count` equally sized segments. + case count(Int) + + /// The weight is split at these indices. + case indices([Int]) + + /// The weight is split at these fractions of the axis. + case fractions([Double]) + + /// The positions that separate the segments of an axis of length `dimension`. + /// + /// Throws if `count` equally sized segments do not fit, which + /// `split(parts:)` would report with a `fatalError`. + func boundaries(of dimension: Int) throws -> [Int] { + switch self { + case .count(let count): + guard count >= 1, dimension % count == 0 else { + throw ShardingError.invalidSegments(self, dimension: dimension) + } + return (1 ..< count).map { $0 * dimension / count } + case .indices(let indices): + return indices + case .fractions(let fractions): + return fractions.map { Int($0 * Double(dimension)) } + } + } +} + +/// A layer cannot be sharded across a group. +public enum ShardingError: Error, CustomStringConvertible, Equatable { + /// A dimension is not divisible by the size of the group. + case indivisible(dimension: String, of: Int, across: Int) + + /// A shard would not hold whole quantization groups. + case quantizationGroup(inputDimensions: Int, groupSize: Int, across: Int) + + /// The segments cannot split a dimension, e.g. `.count(3)` of 8. + case invalidSegments(Segments, dimension: Int) + + /// A ``QuantizedLinear`` was given to a float sharded layer. + /// + /// ``shardLinear(_:sharding:segments:group:)`` returns the quantized flavor. + case quantizedLayer + + /// The layer being sharded is missing a parameter. + case missingParameter(String) + + public var description: String { + switch self { + case .indivisible(let dimension, let value, let size): + "Cannot shard the \(dimension) of size \(value) across \(size) processes." + case .quantizationGroup(let inputDimensions, let groupSize, let size): + """ + Sharding \(inputDimensions) inputs across \ + \(size == 1 ? "1 process" : "\(size) processes") splits a quantization group of \ + \(groupSize). + """ + case .invalidSegments(let segments, let dimension): + "The segments \(segments) do not fit a dimension of size \(dimension)." + case .quantizedLayer: + """ + A quantized layer cannot be sharded as a float layer. Use shardLinear, \ + QuantizedAllToShardedLinear or QuantizedShardedToAllLinear. + """ + case .missingParameter(let name): + "The layer being sharded has no \(name)." + } + } +} + +/// Returns a function that is the identity in the forward pass and sums the +/// gradients across the group in the backward pass. +/// +/// - Parameter group: the group to sum across +public func sumGradients(group: MLXDistributed.Group) -> (MLXArray) -> MLXArray { + if group.size == 1 { + return { $0 } + } + + let f = CustomFunction { + Forward { inputs in + inputs + } + VJP { _, cotangents in + cotangents.map { MLXDistributed.allSum($0, group: group) } + } + } + + return { f([$0])[0] } +} + +/// The sharding axis and segments for a parameter, or `nil` to leave it alone. +/// +/// The path is the parameter's path in the module, e.g. `layers.0.conv.weight`. +public typealias ShardingPredicate = (String, MLXArray) -> (axis: Int, segments: Segments)? + +/// Shard the rows -- the output dimensions -- of every parameter. +/// +/// `scales` and `biases` hold one row per output row, so they shard the same +/// way as the weight. `bias`, the affine bias, is one dimensional. +private func allToShardedPredicate(_ segments: Segments) -> ShardingPredicate { + { path, weight in + path.hasSuffix("bias") ? (-1, segments) : (max(weight.ndim - 2, 0), segments) + } +} + +/// Shard the columns -- the input dimensions -- of every parameter. +/// +/// The affine bias applies to the summed result, so every process keeps it +/// whole. +private func shardedToAllPredicate(_ segments: Segments) -> ShardingPredicate { + { path, _ in + path.hasSuffix("bias") ? nil : (-1, segments) + } +} + +private func predicate(for sharding: ShardingType, segments: Segments) -> ShardingPredicate { + switch sharding { + case .allToSharded: allToShardedPredicate(segments) + case .shardedToAll: shardedToAllPredicate(segments) + } +} + +/// Returns a new parameter tree holding this process' shard of the weights. +/// +/// The segments are checked in a group of one as well, where nothing is split, +/// so that segments that cannot work are reported before the model runs on +/// several processes. +private func shard( + _ parameters: ModuleParameters, group: MLXDistributed.Group, + _ sharding: ShardingPredicate +) throws -> ModuleParameters { + let size = group.size + let rank = group.rank + + let sharded = try parameters.flattened().map { path, weight -> (String, MLXArray) in + // a scalar -- the NVFP4 global scale, for example -- is the same in + // every process + guard weight.ndim > 0, let (axis, segments) = sharding(path, weight) else { + return (path, weight) + } + + let boundaries = try segments.boundaries(of: weight.dim(axis)) + guard size > 1 else { + return (path, weight) + } + + let segmented = + boundaries.isEmpty ? [weight] : weight.split(indices: boundaries, axis: axis) + let parts = try segmented.map { part -> MLXArray in + // split(parts:) needs equal sections and calls fatalError rather + // than reporting, so check before asking. A dimension can divide + // by the group size while one of its segments does not. + guard part.dim(axis) % size == 0 else { + throw ShardingError.indivisible( + dimension: "segment of \(path)", of: part.dim(axis), across: size) + } + return part.split(parts: size, axis: axis)[rank] + } + + return (path, concatenated(parts, axis: axis).contiguous()) + } + + // in a group of one every parameter stays whole + return size == 1 ? parameters : ModuleParameters.unflattened(sharded) +} + +/// Shard the input dimensions of a quantized layer. +/// +/// The packed weight, the scales and the biases hold the inputs at different +/// resolutions, so a boundary between two segments lands at a different +/// position in each. Every segment has to hold whole quantization groups on +/// every process. Then every boundary lies on a group edge and scales to each +/// parameter exactly in integer arithmetic, because every supported mode packs +/// a group into whole `uint32` words. Scaling through a `Double` could +/// truncate to the position before. +private func quantizedShardedToAllPredicate( + _ segments: Segments, inputDimensions: Int, groupSize: Int, size: Int +) throws -> ShardingPredicate { + let boundaries = try segments.boundaries(of: inputDimensions) + guard boundaries == boundaries.sorted(), + boundaries.allSatisfy({ (0 ... inputDimensions).contains($0) }) + else { + throw ShardingError.invalidSegments(segments, dimension: inputDimensions) + } + + let edges = [0] + boundaries + [inputDimensions] + for (start, end) in zip(edges, edges.dropFirst()) { + guard (end - start) % (groupSize * size) == 0 else { + throw ShardingError.quantizationGroup( + inputDimensions: end - start, groupSize: groupSize, across: size) + } + } + + return { path, weight in + path.hasSuffix("bias") + ? nil : (-1, .indices(boundaries.map { $0 * weight.dim(-1) / inputDimensions })) + } +} + +/// Shard a module in place by replacing its parameters with sharded ones. +/// +/// The module itself is unchanged, so distributed communication only happens +/// if the module supports it natively. +/// +/// - Parameters: +/// - module: the module whose parameters are sharded in place +/// - sharding: the kind of sharding to apply +/// - segments: the segments that comprise each unsharded weight +/// - group: the group to shard across, or `nil` to use the global group +public func shardInPlace( + _ module: Module, sharding: ShardingType, segments: Segments = .count(1), + group: MLXDistributed.Group? = nil +) throws { + try shardInPlace( + module, predicate: predicate(for: sharding, segments: segments), group: group) +} + +/// Shard a module in place, choosing how each parameter is split. +/// +/// This is the general form of ``shardInPlace(_:sharding:segments:group:)``, and +/// matches what Python's `shard_inplace` accepts. The predicate receives a +/// parameter and its path and returns the axis to shard it along, with the +/// segments that make it up, or `nil` to leave the parameter whole. +/// +/// The module itself is unchanged, so distributed communication only happens +/// if the module supports it natively. +/// +/// - Parameters: +/// - module: the module whose parameters are sharded in place +/// - predicate: chooses how each parameter is sharded +/// - group: the group to shard across, or `nil` to use the global group +public func shardInPlace( + _ module: Module, predicate: ShardingPredicate, group: MLXDistributed.Group? = nil +) throws { + let group = try group ?? MLXDistributed.initialize() + _ = try module.update(parameters: shard(module.parameters(), group: group, predicate)) +} + +/// Create a new linear layer with sharded parameters that also performs the +/// distributed communication, either in the forward or the backward pass. +/// +/// Unlike ``shardInPlace(_:sharding:segments:group:)`` the original layer is +/// not changed. A ``QuantizedLinear`` yields the quantized flavor of the +/// sharded layer, so a quantized model shards like any other. +/// +/// - Parameters: +/// - layer: the linear layer to shard +/// - sharding: the kind of sharding to apply +/// - segments: the segments that comprise each unsharded weight +/// - group: the group to shard across, or `nil` to use the global group +/// - Returns: the sharded layer, which can replace `layer` in a model +public func shardLinear( + _ layer: Linear, sharding: ShardingType, segments: Segments = .count(1), + group: MLXDistributed.Group? = nil +) throws -> Linear { + // QuantizedLinear is a Linear, so it has to be matched first + if let layer = layer as? QuantizedLinear { + return switch sharding { + case .allToSharded: + try QuantizedAllToShardedLinear(layer, segments: segments, group: group) + case .shardedToAll: + try QuantizedShardedToAllLinear(layer, segments: segments, group: group) + } + } + + return switch sharding { + case .allToSharded: + try AllToShardedLinear(layer, segments: segments, group: group) + case .shardedToAll: + try ShardedToAllLinear(layer, segments: segments, group: group) + } +} + +/// Each member of the group applies part of the affine transformation so that +/// the result is sharded across the group. +/// +/// The gradients are automatically aggregated from each member of the group. +/// +/// This is a ``Linear``, so it can replace one in a model. +open class AllToShardedLinear: Linear { + + /// The group the output dimensions are sharded across. + public let group: MLXDistributed.Group + + private let aggregateGradients: (MLXArray) -> MLXArray + + /// - Parameters: + /// - inputDimensions: number of input dimensions + /// - outputDimensions: number of output dimensions, sharded across the group + /// - bias: if `true` this layer will apply a bias + /// - group: the group to shard across, or `nil` to use the global group + public init( + _ inputDimensions: Int, _ outputDimensions: Int, bias: Bool = true, + group: MLXDistributed.Group? = nil + ) throws { + let group = try group ?? MLXDistributed.initialize() + guard outputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "output", of: outputDimensions, across: group.size) + } + + self.group = group + self.aggregateGradients = sumGradients(group: group) + + let scale = sqrt(1.0 / Float(inputDimensions)) + let shardedOutput = outputDimensions / group.size + super.init( + weight: MLXRandom.uniform(-scale ..< scale, [shardedOutput, inputDimensions]), + bias: bias ? MLXRandom.uniform(-scale ..< scale, [shardedOutput]) : nil) + } + + /// Hold parameters that are already this process' shard. + init(shardedWeight weight: MLXArray, bias: MLXArray?, group: MLXDistributed.Group) { + self.group = group + self.aggregateGradients = sumGradients(group: group) + super.init(weight: weight, bias: bias) + } + + /// Create a sharded layer from an existing ``Linear``. + /// + /// Throws for a ``QuantizedLinear``: shard it with + /// ``shardLinear(_:sharding:segments:group:)`` or ``QuantizedAllToShardedLinear``. + public convenience init( + _ other: Linear, segments: Segments = .count(1), group: MLXDistributed.Group? = nil + ) throws { + // a QuantizedLinear is a Linear, but its packed weight is not a float one + guard !(other is QuantizedLinear) else { + throw ShardingError.quantizedLayer + } + + let group = try group ?? MLXDistributed.initialize() + let (outputDimensions, _) = other.shape + guard outputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "output", of: outputDimensions, across: group.size) + } + + let parameters = try Dictionary( + uniqueKeysWithValues: shard( + other.parameters(), group: group, allToShardedPredicate(segments) + ).flattened()) + guard let weight = parameters["weight"] else { + throw ShardingError.missingParameter("weight") + } + + self.init(shardedWeight: weight, bias: parameters["bias"], group: group) + } + + open override func callAsFunction(_ x: MLXArray) -> MLXArray { + // every shard reads the same input, so their gradients are summed + let x = aggregateGradients(x) + + if let bias { + return addMM(bias, x, weight.T) + } else { + return matmul(x, weight.T) + } + } + + /// Quantizing keeps the layer sharded. + /// + /// Without this a quantized model built by ``quantize(model:groupSize:bits:mode:filter:apply:)`` + /// would hold plain ``QuantizedLinear`` layers that no longer communicate. + public override func toQuantized(groupSize: Int, bits: Int, mode: QuantizationMode) -> Module { + QuantizedAllToShardedLinear( + shardedWeight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode, + group: group) + } +} + +/// Each member of the group applies part of the affine transformation and the +/// results are then aggregated. +/// +/// Every member of the group ends up with the same result. +/// +/// This is a ``Linear``, so it can replace one in a model. +open class ShardedToAllLinear: Linear { + + /// The group the input dimensions are sharded across. + public let group: MLXDistributed.Group + + /// - Parameters: + /// - inputDimensions: number of input dimensions, sharded across the group + /// - outputDimensions: number of output dimensions + /// - bias: if `true` this layer will apply a bias + /// - group: the group to shard across, or `nil` to use the global group + public init( + _ inputDimensions: Int, _ outputDimensions: Int, bias: Bool = true, + group: MLXDistributed.Group? = nil + ) throws { + let group = try group ?? MLXDistributed.initialize() + guard inputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "input", of: inputDimensions, across: group.size) + } + + self.group = group + + // Each process holds its own slice of the weight, but the bias is + // added after the reduction, so every process has to hold the same + // one. Random values would differ from process to process -- their + // generators are seeded independently -- and the layer would quietly + // produce a different result in each. Python's quantized flavor uses + // zeros for this reason; its float one does not, and diverges. + let scale = sqrt(1.0 / Float(inputDimensions)) + super.init( + weight: MLXRandom.uniform( + -scale ..< scale, [outputDimensions, inputDimensions / group.size]), + bias: bias ? MLXArray.zeros([outputDimensions]) : nil) + } + + /// Hold parameters that are already this process' shard. + init(shardedWeight weight: MLXArray, bias: MLXArray?, group: MLXDistributed.Group) { + self.group = group + super.init(weight: weight, bias: bias) + } + + /// Create a sharded layer from an existing ``Linear``. + /// + /// Throws for a ``QuantizedLinear``: shard it with + /// ``shardLinear(_:sharding:segments:group:)`` or ``QuantizedShardedToAllLinear``. + public convenience init( + _ other: Linear, segments: Segments = .count(1), group: MLXDistributed.Group? = nil + ) throws { + // a QuantizedLinear is a Linear, but its packed weight is not a float one + guard !(other is QuantizedLinear) else { + throw ShardingError.quantizedLayer + } + + let group = try group ?? MLXDistributed.initialize() + let (_, inputDimensions) = other.shape + guard inputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "input", of: inputDimensions, across: group.size) + } + + let parameters = try Dictionary( + uniqueKeysWithValues: shard( + other.parameters(), group: group, shardedToAllPredicate(segments) + ).flattened()) + guard let weight = parameters["weight"] else { + throw ShardingError.missingParameter("weight") + } + + self.init(shardedWeight: weight, bias: parameters["bias"], group: group) + } + + open override func callAsFunction(_ x: MLXArray) -> MLXArray { + // each process holds part of the sum, and the bias belongs to the + // whole, so it is added after the reduction + var x = matmul(x, weight.T) + x = MLXDistributed.allSum(x, group: group) + + if let bias { + x = x + bias + } + return x + } + + /// Quantizing keeps the layer sharded. + public override func toQuantized(groupSize: Int, bits: Int, mode: QuantizationMode) -> Module { + QuantizedShardedToAllLinear( + shardedWeight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode, + group: group) + } +} + +/// The quantized flavor of ``AllToShardedLinear``. +/// +/// Like ``QuantizedLinear`` its parameters are frozen. +open class QuantizedAllToShardedLinear: QuantizedLinear { + + /// The group the output dimensions are sharded across. + public let group: MLXDistributed.Group + + private let aggregateGradients: (MLXArray) -> MLXArray + + /// Quantize a weight that is already this process' shard. + public init( + shardedWeight weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4, + mode: QuantizationMode = .affine, group: MLXDistributed.Group + ) { + self.group = group + self.aggregateGradients = sumGradients(group: group) + super.init(weight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode) + } + + /// Hold quantized parameters that are already this process' shard. + init( + shardedQuantizedWeight weight: MLXArray, bias: MLXArray?, scales: MLXArray, + biases: MLXArray?, groupSize: Int, bits: Int, mode: QuantizationMode, + globalScale: MLXArray?, group: MLXDistributed.Group + ) { + self.group = group + self.aggregateGradients = sumGradients(group: group) + super.init( + weight: weight, bias: bias, scales: scales, biases: biases, groupSize: groupSize, + bits: bits, mode: mode, globalScale: globalScale) + self.freeze() + } + + /// Create a sharded layer from an existing ``QuantizedLinear``. + /// + /// The packed weight, `scales` and `biases` all carry one row per output, + /// so they shard together along the output dimension. + public convenience init( + _ other: QuantizedLinear, segments: Segments = .count(1), + group: MLXDistributed.Group? = nil + ) throws { + let group = try group ?? MLXDistributed.initialize() + let (outputDimensions, _) = other.shape + guard outputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "output", of: outputDimensions, across: group.size) + } + + let parameters = try Dictionary( + uniqueKeysWithValues: shard( + other.parameters(), group: group, allToShardedPredicate(segments) + ).flattened()) + guard let weight = parameters["weight"] else { + throw ShardingError.missingParameter("weight") + } + guard let scales = parameters["scales"] else { + throw ShardingError.missingParameter("scales") + } + + self.init( + shardedQuantizedWeight: weight, bias: parameters["bias"], scales: scales, + biases: parameters["biases"], groupSize: other.groupSize, bits: other.bits, + mode: other.mode, globalScale: parameters["global_scale"], group: group) + } + + open override func callAsFunction(_ x: MLXArray) -> MLXArray { + // the quantized matmul, the global scale and the bias are all local to + // this shard, so only the gradient aggregation is added + super.callAsFunction(aggregateGradients(x)) + } +} + +/// The quantized flavor of ``ShardedToAllLinear``. +/// +/// Like ``QuantizedLinear`` its parameters are frozen. +open class QuantizedShardedToAllLinear: QuantizedLinear { + + /// The group the input dimensions are sharded across. + public let group: MLXDistributed.Group + + /// Quantize a weight that is already this process' shard. + public init( + shardedWeight weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4, + mode: QuantizationMode = .affine, group: MLXDistributed.Group + ) { + self.group = group + super.init(weight: weight, bias: bias, groupSize: groupSize, bits: bits, mode: mode) + } + + /// Hold quantized parameters that are already this process' shard. + init( + shardedQuantizedWeight weight: MLXArray, bias: MLXArray?, scales: MLXArray, + biases: MLXArray?, groupSize: Int, bits: Int, mode: QuantizationMode, + globalScale: MLXArray?, group: MLXDistributed.Group + ) { + self.group = group + super.init( + weight: weight, bias: bias, scales: scales, biases: biases, groupSize: groupSize, + bits: bits, mode: mode, globalScale: globalScale) + self.freeze() + } + + /// Create a sharded layer from an existing ``QuantizedLinear``. + /// + /// This shards the input dimension, which is the packed one, so every + /// segment has to hold whole quantization groups on every process. + public convenience init( + _ other: QuantizedLinear, segments: Segments = .count(1), + group: MLXDistributed.Group? = nil + ) throws { + let group = try group ?? MLXDistributed.initialize() + let (_, inputDimensions) = other.shape + guard inputDimensions % group.size == 0 else { + throw ShardingError.indivisible( + dimension: "input", of: inputDimensions, across: group.size) + } + + let parameters = try Dictionary( + uniqueKeysWithValues: shard( + other.parameters(), group: group, + quantizedShardedToAllPredicate( + segments, inputDimensions: inputDimensions, groupSize: other.groupSize, + size: group.size) + ).flattened()) + guard let weight = parameters["weight"] else { + throw ShardingError.missingParameter("weight") + } + guard let scales = parameters["scales"] else { + throw ShardingError.missingParameter("scales") + } + + self.init( + shardedQuantizedWeight: weight, bias: parameters["bias"], scales: scales, + biases: parameters["biases"], groupSize: other.groupSize, bits: other.bits, + mode: other.mode, globalScale: parameters["global_scale"], group: group) + } + + open override func callAsFunction(_ x: MLXArray) -> MLXArray { + // this cannot call super: the bias applies to the summed result, not to + // this process' partial product + var x = quantizedMM( + x, weight, scales: scales, biases: biases, transpose: true, groupSize: groupSize, + bits: bits, mode: mode) + x = applyNVFP4GlobalScale(x, globalScale: globalScale) + x = MLXDistributed.allSum(x, group: group) + + if let bias { + x = x + bias + } + return x + } +} + +/// Average the gradients across the processes in the group. +/// +/// Small gradients are concatenated into batches of at least `allReduceSize` +/// bytes so that they are communicated in one step, which is considerably +/// faster than one call per array. +/// +/// - Parameters: +/// - gradients: the gradients, which must have the same structure in every process +/// - group: the group to average across, or `nil` to use the global group +/// - allReduceSize: group arrays until their size in bytes exceeds this, +/// or `0` to disable grouping +/// - stream: stream to evaluate on +public func averageGradients( + _ gradients: ModuleParameters, group: MLXDistributed.Group? = nil, + allReduceSize: Int = 32 * 1024 * 1024, stream: StreamOrDevice = .cpu +) throws -> ModuleParameters { + let group = try group ?? MLXDistributed.initialize() + let size = group.size + if size == 1 { + return gradients + } + + let flat = gradients.flattened() + if flat.isEmpty { + return gradients + } + + // one all reduce per gradient + if allReduceSize <= 0 { + return ModuleParameters.unflattened( + flat.map { + ($0.0, MLXDistributed.allSum($0.1, group: group, stream: stream) / size) + }) + } + + // arrays of mixed types cannot be concatenated + let dtype = flat[0].1.dtype + guard flat.allSatisfy({ $0.1.dtype == dtype }) else { + return try averageGradients( + gradients, group: group, allReduceSize: 0, stream: stream) + } + + // gather the gradients into groups that are at least allReduceSize bytes + let batches = groupBySize(flat.map { $0.1.nbytes }, limit: allReduceSize) + + // concatenate, reduce, split + var result = [(String, MLXArray)]() + for batch in batches { + let sizes = batch.map { flat[$0].1.size } + var big = concatenated(batch.map { flat[$0].1.reshaped([-1]) }) + big = MLXDistributed.allSum(big, group: group, stream: stream) / size + + let indices = sizes.dropLast().reduce(into: [Int]()) { $0.append(($0.last ?? 0) + $1) } + let parts = indices.isEmpty ? [big] : big.split(indices: indices) + + for (part, i) in zip(parts, batch) { + result.append((flat[i].0, part.reshaped(flat[i].1.shape))) + } + } + + return ModuleParameters.unflattened(result) +} + +/// Group consecutive arrays, given their sizes in bytes, into batches of at +/// least `limit` bytes. The last batch holds whatever remains. +/// +/// This is `_group_by_size` from Python's `mlx.nn.utils`. It is a function of +/// its own so that the tests can count the batches `averageGradients` forms, +/// the way the Python tests count its `all_sum` calls. +func groupBySize(_ sizes: [Int], limit: Int) -> [[Int]] { + var batches = [[Int]]() + var batch = [Int]() + var batchBytes = 0 + for (i, size) in sizes.enumerated() { + batch.append(i) + batchBytes += size + if batchBytes >= limit { + batches.append(batch) + batch = [] + batchBytes = 0 + } + } + if !batch.isEmpty { + batches.append(batch) + } + return batches +} + +/// Clip the global norm of gradients that are sharded across a group. +/// +/// This is the sharded counterpart of `clipGradNorm` in MLXOptimizers: no +/// process holds the whole gradient, so the local squared norms are summed +/// across the group before anything is rescaled. +/// +/// - Parameters: +/// - gradients: this process' shard of the gradients +/// - maxNorm: the maximum allowed global norm +/// - group: the group the gradients are sharded across, or `nil` to use the +/// global group +/// - stream: stream to evaluate on +/// - Returns: the rescaled shard and the global gradient norm +public func clipGradNormSharded( + gradients: ModuleParameters, maxNorm: Float, group: MLXDistributed.Group? = nil, + stream: StreamOrDevice = .cpu +) throws -> (ModuleParameters, MLXArray) { + let group = try group ?? MLXDistributed.initialize() + + let localNormSquared = gradients.reduce(MLXArray(0)) { $0 + $1.square().sum() } + let totalNorm = sqrt(MLXDistributed.allSum(localNormSquared, group: group, stream: stream)) + let normalizer = minimum(maxNorm / (totalNorm + 1e-6), 1) + + return (gradients.mapValues { $0 * normalizer }, totalNorm) +} diff --git a/Source/MLXNN/Quantized.swift b/Source/MLXNN/Quantized.swift index ca5913ee7..b1e429794 100644 --- a/Source/MLXNN/Quantized.swift +++ b/Source/MLXNN/Quantized.swift @@ -27,7 +27,9 @@ public protocol Quantized: Module { var mode: QuantizationMode { get } } -private func applyNVFP4GlobalScale(_ value: MLXArray, globalScale: MLXArray?) -> MLXArray { +/// Internal so that the sharded quantized layers, which cannot reuse +/// ``QuantizedLinear``'s forward pass, apply the scale the same way. +func applyNVFP4GlobalScale(_ value: MLXArray, globalScale: MLXArray?) -> MLXArray { guard let globalScale else { return value } // NVFP4 encodes its E4M3 group scales with (448 * 6) / globalScale. return (value * (globalScale / (448 * 6))).asType(value.dtype) diff --git a/Tests/MLXTests/DistributedHarness.swift b/Tests/MLXTests/DistributedHarness.swift new file mode 100644 index 000000000..4b7692dc0 --- /dev/null +++ b/Tests/MLXTests/DistributedHarness.swift @@ -0,0 +1,305 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +#if canImport(Darwin) + import Darwin +#elseif canImport(Glibc) + import Glibc +#endif + +/// Runs a test body in several processes that form a ring, the way the Python +/// distributed tests are run under a launcher. +/// +/// The test process is itself rank 0. It reserves a loopback port for every +/// rank, writes a temporary hostfile, re-executes the test bundle once for each +/// of the other ranks and then runs the same body, so no separate worker +/// executable is needed. +/// +/// MLX caches the communication group per process, so a process can only join +/// one ring. Each test class that needs ranks therefore has a single test +/// method and is run in its own `xctest` invocation; a second body in the same +/// process is skipped rather than left hanging against a ring that the first +/// one built. +enum DistributedHarness { + + /// Set in the environment of the ranks the harness spawns. + /// + /// `MLX_RANK` alone would not do: a whole test run can be launched with it + /// already set, and then every process would wait to be spawned. + static let spawnedRankVariable = "MLX_TEST_DISTRIBUTED_SPAWNED_RANK" + + /// Run `body` in `ranks` processes. + /// + /// - Parameters: + /// - ranks: how many processes take part, including this one + /// - testName: the `Class/method` selector every rank runs + /// - timeout: how long the spawned ranks may run before giving up + /// - body: the test body, which every rank runs with its own group + static func run( + ranks: Int, testName: String, timeout: TimeInterval = 120, + body: (MLXDistributed.Group) throws -> Void + ) throws { + let environment = ProcessInfo.processInfo.environment + + if environment[spawnedRankVariable] == "1" { + // spawned by rank 0 -- the hostfile and rank are already set + try runRank(ranks: ranks, body) + return + } + + guard environment["MLX_TEST_DISTRIBUTED"] == "1" else { + throw XCTSkip("Set MLX_TEST_DISTRIBUTED=1 to run the multi process tests.") + } + + guard state.claim() else { + throw XCTSkip( + """ + Another multi process test already formed a group in this process. \ + Run this class in its own xctest invocation. + """) + } + + guard let runner = testRunner() else { + throw XCTSkip("Could not locate the test bundle to re-execute as a rank.") + } + + guard let ports = reserveFreePorts(count: ranks) else { + throw XCTSkip("Could not reserve \(ranks) loopback ports.") + } + + let hostfile = try writeHostfile(ports: ports) + defer { try? FileManager.default.removeItem(at: hostfile) } + + let spawned = (1 ..< ranks).map { + makeRank($0, runner: runner, testName: testName, hostfile: hostfile) + } + let watchdog = Watchdog(spawned: spawned, timeout: timeout) + defer { + watchdog.disarm() + for process in spawned where process.isRunning { + process.terminate() + } + } + for process in spawned { + try process.run() + } + + setenv("MLX_HOSTFILE", hostfile.path, 1) + setenv("MLX_RANK", "0", 1) + defer { + unsetenv("MLX_HOSTFILE") + unsetenv("MLX_RANK") + } + + try runRank(ranks: ranks, body) + + for (index, process) in spawned.enumerated() { + process.waitUntilExit() + XCTAssertEqual( + process.terminationStatus, 0, + "rank \(index + 1) failed -- see its output above") + } + } + + /// Join the ring and run the body. + /// + /// The body is wrapped in `withError` so that an MLX error becomes a test + /// failure naming the rank. Without an active scope the global handler + /// calls `fatalError()` and the message never reaches XCTest. + private static func runRank( + ranks: Int, _ body: (MLXDistributed.Group) throws -> Void + ) throws { + do { + let group = try MLXDistributed.initialize(backend: .ring, strict: true) + + // a group of one makes every assertion about sharding pass without + // proving anything, so the size is part of the contract + print("[rank \(group.rank)] joined a group of size \(group.size)") + XCTAssertEqual(group.size, ranks, "joined a group of the wrong size") + + try withError { + try body(group) + } + } catch { + let rank = ProcessInfo.processInfo.environment["MLX_RANK"] ?? "?" + XCTFail("rank \(rank) failed: \(error)") + throw error + } + } + + // MARK: - Process helpers + + /// The `xctest` tool running this bundle, plus the bundle itself. + private static func testRunner() -> (executable: URL, bundle: URL)? { + let executable = URL(fileURLWithPath: CommandLine.arguments[0]) + guard executable.lastPathComponent == "xctest" else { return nil } + + let bundle = Bundle(for: DistributedRingTests.self).bundleURL + guard bundle.pathExtension == "xctest" else { return nil } + + return (executable, bundle) + } + + /// Write the JSON hostfile the ring backend expects. + private static func writeHostfile(ports: [Int]) throws -> URL { + let hosts = ports.map { ["127.0.0.1:\($0)"] } + let data = try JSONSerialization.data(withJSONObject: hosts) + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("mlx-ring-hostfile-\(UUID().uuidString).json") + try data.write(to: url) + + return url + } + + /// Prepare the test bundle to be re-executed as the given rank. + /// + /// The environment is inherited without XCTest's own session variables. + /// Under `xcodebuild test` those make the spawned `xctest` wait for Xcode + /// instead of running the test named on its command line. + private static func makeRank( + _ rank: Int, runner: (executable: URL, bundle: URL), testName: String, hostfile: URL + ) -> Process { + let process = Process() + process.executableURL = runner.executable + process.arguments = ["-XCTest", testName, runner.bundle.path] + + var environment = ProcessInfo.processInfo.environment.filter { key, value in + !key.hasPrefix("XCTest") && !key.hasPrefix("XCInject") + && !(key == "DYLD_INSERT_LIBRARIES" && value.contains("XCTest")) + } + environment["MLX_RANK"] = "\(rank)" + environment["MLX_HOSTFILE"] = hostfile.path + environment[spawnedRankVariable] = "1" + process.environment = environment + + return process + } +} + +/// Tracks whether this process has already formed a group. +private final class HarnessState: @unchecked Sendable { + + private let lock = NSLock() + private var used = false + + /// Returns `true` the first time it is called in a process. + func claim() -> Bool { + lock.withLock { + if used { + return false + } + used = true + return true + } + } +} + +private let state = HarnessState() + +/// Ends the test process when a spawned rank fails or the ranks take too long. +/// +/// Forming the ring blocks in `accept()` until every rank has connected, and +/// MLX has no timeout, so a rank that dies on startup would leave rank 0 +/// waiting forever. Nothing can interrupt that thread, so the watchdog +/// terminates the spawned ranks and exits: a failure rather than a hang. +private final class Watchdog: @unchecked Sendable { + + private let lock = NSLock() + private var armed = true + private let spawned: [Process] + + init(spawned: [Process], timeout: TimeInterval) { + self.spawned = spawned + + for (index, process) in spawned.enumerated() { + process.terminationHandler = { [weak self] terminated in + let status = terminated.terminationStatus + guard status != 0 else { return } + + // rank 0 may still report the failure itself, e.g. as an error + // from a collective, so give it a moment first + DispatchQueue.global().asyncAfter(deadline: .now() + 10) { + self?.fire("rank \(index + 1) exited with status \(status)") + } + } + } + + DispatchQueue.global().asyncAfter(deadline: .now() + timeout) { [weak self] in + self?.fire("the ranks did not finish within \(Int(timeout)) seconds") + } + } + + func disarm() { + lock.withLock { armed = false } + } + + private func fire(_ reason: String) { + guard lock.withLock({ armed }) else { return } + + for process in spawned where process.isRunning { + process.terminate() + } + FileHandle.standardError.write(Data("DistributedHarness: \(reason), exiting\n".utf8)) + exit(1) + } +} + +/// Reserve loopback ports by binding them and closing them again. +/// +/// Every socket stays open until all ports are picked, so the same port is +/// not handed out twice. +/// +/// This lives at file scope on purpose: inside an `XCTestCase` the name +/// `bind` resolves to `NSObject.bind(_:to:withKeyPath:options:)` rather than +/// the socket call. +private func reserveFreePorts(count: Int) -> [Int]? { + var ports = [Int]() + var sockets = [Int32]() + defer { sockets.forEach { close($0) } } + + for _ in 0 ..< count { + // Glibc types SOCK_STREAM as an enum rather than the Int32 that + // socket() takes. + #if canImport(Darwin) + let socketType = SOCK_STREAM + #else + let socketType = Int32(SOCK_STREAM.rawValue) + #endif + + let fd = socket(AF_INET, socketType, 0) + guard fd >= 0 else { return nil } + sockets.append(fd) + + var reuse: Int32 = 1 + setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &reuse, socklen_t(MemoryLayout.size)) + + var address = sockaddr_in() + address.sin_family = sa_family_t(AF_INET) + address.sin_port = 0 + address.sin_addr.s_addr = inet_addr("127.0.0.1") + + let bound = withUnsafePointer(to: &address) { + $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { + bind(fd, $0, socklen_t(MemoryLayout.size)) + } + } + guard bound == 0 else { return nil } + + var assigned = sockaddr_in() + var length = socklen_t(MemoryLayout.size) + let named = withUnsafeMutablePointer(to: &assigned) { + $0.withMemoryRebound(to: sockaddr.self, capacity: 1) { + getsockname(fd, $0, &length) + } + } + guard named == 0 else { return nil } + + ports.append(Int(UInt16(bigEndian: assigned.sin_port))) + } + + return ports +} diff --git a/Tests/MLXTests/DistributedNNTests.swift b/Tests/MLXTests/DistributedNNTests.swift new file mode 100644 index 000000000..7cae539e0 --- /dev/null +++ b/Tests/MLXTests/DistributedNNTests.swift @@ -0,0 +1,670 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +@testable import MLXNN + +// Ports of the layer tests from the Python distributed tests +// (`python/tests/mlx_distributed_tests.py`). Python runs them for every group +// size under a launcher; each body here runs twice, once in a single process +// where the sharding degenerates but the layers are still built and run, and +// once across real ranks where every rank holds a different slice. + +/// Port of `test_shard_linear`. +/// +/// - Parameter world: the group to shard across +func shardLinearBody(world: MLXDistributed.Group) throws { + MLXRandom.seed(0xF0F0_F0F0) + + let lower = world.rank * 1024 / world.size + let upper = (world.rank + 1) * 1024 / world.size + + let x = MLXRandom.normal([4, 1024]) + let linear = Linear(1024, 1024, bias: true) + + // MARK: - float + + let sharded1 = try shardLinear(linear, sharding: .allToSharded, group: world) + let sharded2 = try shardLinear(linear, sharding: .shardedToAll, group: world) + + let y = linear(x) + let y1 = sharded1(x) + let y2 = sharded2(x[lower ..< upper, axis: 1]) + try checkedEval(y, y1, y2) + + XCTAssertTrue( + y.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "sharded-to-all must reproduce the unsharded output") + XCTAssertTrue( + y[lower ..< upper, axis: 1].allClose(y1, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "all-to-sharded must reproduce this rank's slice of the output") + + // MARK: - quantized + + // Python: qlin = lin.to_quantized() + let quantized = QuantizedLinear(linear) + let qSharded1 = try shardLinear(quantized, sharding: .allToSharded, group: world) + let qSharded2 = try shardLinear(quantized, sharding: .shardedToAll, group: world) + + XCTAssertTrue( + qSharded1 is QuantizedAllToShardedLinear, + "a QuantizedLinear must shard into the quantized layer, not the float one") + XCTAssertTrue(qSharded2 is QuantizedShardedToAllLinear) + + let qy = quantized(x) + let qy1 = qSharded1(x) + let qy2 = qSharded2(x[lower ..< upper, axis: 1]) + try checkedEval(qy, qy1, qy2) + + XCTAssertTrue( + qy.allClose(qy2, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "quantized sharded-to-all must reproduce the unsharded output") + XCTAssertTrue( + qy[lower ..< upper, axis: 1].allClose(qy1, rtol: 1e-5, atol: 1e-8).item(Bool.self), + "quantized all-to-sharded must reproduce this rank's slice") + + // MARK: - a non affine mode, which carries no quantization biases + + // Python: lin.to_quantized(group_size=32, bits=8, mode="mxfp8") + let mxfp8 = QuantizedLinear(linear, groupSize: 32, bits: 8, mode: .mxfp8) + XCTAssertEqual(mxfp8.mode, .mxfp8) + XCTAssertNil(mxfp8.biases, "mxfp8 quantization has no biases") + + let mxSharded1 = try XCTUnwrap( + try shardLinear(mxfp8, sharding: .allToSharded, group: world) as? QuantizedLinear) + let mxSharded2 = try XCTUnwrap( + try shardLinear(mxfp8, sharding: .shardedToAll, group: world) as? QuantizedLinear) + + XCTAssertEqual(mxSharded1.mode, .mxfp8, "the mode must survive sharding") + XCTAssertEqual(mxSharded2.mode, .mxfp8, "the mode must survive sharding") + XCTAssertNil(mxSharded1.biases) + XCTAssertNil(mxSharded2.biases) + + let my = mxfp8(x) + let my1 = mxSharded1(x) + let my2 = mxSharded2(x[lower ..< upper, axis: 1]) + try checkedEval(my, my1, my2) + + XCTAssertTrue( + my.allClose(my2, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "mxfp8 sharded-to-all must reproduce the unsharded output") + XCTAssertTrue( + my[lower ..< upper, axis: 1].allClose(my1, rtol: 1e-5, atol: 1e-8).item(Bool.self), + "mxfp8 all-to-sharded must reproduce this rank's slice") + + // MARK: - quantizing a sharded layer keeps it sharded + + let quantizedShard = sharded1.toQuantized(groupSize: 64, bits: 4, mode: .affine) + XCTAssertTrue( + quantizedShard is QuantizedAllToShardedLinear, + "quantizing a sharded layer must not drop the communication") + + // MARK: - backward + + try shardLinearBackward(world: world) +} + +/// The gradients of a sharded model must match the slice of the unsharded +/// model's gradients that this rank owns. +private func shardLinearBackward(world: MLXDistributed.Group) throws { + MLXRandom.seed(0xF0F0_F0F0) + + let layers = (0 ..< 4).map { _ in Linear(128, 128) } + let model = Sequential(layers: layers[0], layers[1], layers[2], layers[3]) + let sharded = Sequential( + layers: try shardLinear(layers[0], sharding: .allToSharded, group: world), + try shardLinear(layers[1], sharding: .shardedToAll, group: world), + try shardLinear(layers[2], sharding: .allToSharded, group: world), + try shardLinear(layers[3], sharding: .shardedToAll, group: world)) + + func loss(_ model: Sequential, _ x: MLXArray, _ y: MLXArray) -> MLXArray { + (model(x) * y).sum() + } + + let x = MLXRandom.normal([4, 128]) + let y = MLXRandom.normal([4, 128]) + + let (l1, g1) = valueAndGrad(model: model, loss)(model, x, y) + let (l2, g2) = valueAndGrad(model: sharded, loss)(sharded, x, y) + try checkedEval(l1, g1, l2, g2) + + let lower = world.rank * 128 / world.size + let upper = (world.rank + 1) * 128 / world.size + let full = Dictionary(uniqueKeysWithValues: g1.flattened()) + let shard = Dictionary(uniqueKeysWithValues: g2.flattened()) + + XCTAssertTrue(l1.allClose(l2, rtol: 1e-4, atol: 1e-6).item(Bool.self), "loss mismatch") + + // all-to-sharded splits the output dimension, so rows of the weight + // gradient and the bias gradient + for layer in [0, 2] { + assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 0) + assertGradientShard(full, shard, "layers.\(layer).bias", lower ..< upper, axis: 0) + } + + // sharded-to-all splits the input dimension -- columns of the weight + // gradient -- and leaves the bias whole + for layer in [1, 3] { + assertGradientShard(full, shard, "layers.\(layer).weight", lower ..< upper, axis: 1) + assertGradientWhole(full, shard, "layers.\(layer).bias") + } +} + +/// Port of `test_shard_predicate`. +/// +/// A module that is not a ``Linear`` is sharded by handing ``shardInPlace`` a +/// predicate. The module aggregates where the sharding requires it, which the +/// predicate cannot do on its own. +func shardPredicateBody(world: MLXDistributed.Group) throws { + MLXRandom.seed(0xF0F0_F0F0) + + // even layers shard their output channels; odd layers shard their input + // channels and keep the bias, which applies to the summed result + let sharding: ShardingPredicate = { path, _ in + let parts = path.split(separator: ".") + guard parts.count > 1, let layer = Int(parts[1]) else { return nil } + + if layer % 2 == 0 { + return (0, .count(1)) + } + return parts.last == "bias" ? nil : (-1, .count(1)) + } + + let model = Sequential( + layers: AggregatingConv(3, 128), AggregatingConv(128, 128), + AggregatingConv(128, 128), AggregatingConv(128, 3)) + let sharded = Sequential( + layers: AggregatingConv(3, 128), + AggregatingConv(128, 128, aggregate: world), + AggregatingConv(128, 128), + AggregatingConv(128, 3, aggregate: world)) + + _ = sharded.update(parameters: model.parameters()) + try shardInPlace(sharded, predicate: sharding, group: world) + + let x = MLXRandom.normal([4, 16, 16, 3]) + let y1 = model(x) + let y2 = sharded(x) + try checkedEval(y1, y2) + + XCTAssertTrue( + y1.allClose(y2, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "a predicate sharded model must reproduce the unsharded output") +} + +/// Sharding swaps the layers of a model in place, the way quantization does. +/// +/// Models declare their projections as `Linear`, and a model's shard method +/// installs the sharded layers with `update(modules:)`. That only works +/// because the sharded layers are `Linear` subclasses: a sharded layer that +/// was not a `Linear` could not be installed at all. +func moduleSubstitutionBody(world: MLXDistributed.Group) throws { + MLXRandom.seed(0xF0F0_F0F0) + let x = MLXRandom.normal([4, 64]) + + // the check has to be able to fail: a module that is not a Linear is refused + let control = ProjectionBlock(64, 256) + let notLinear: [(String, Module)] = [("up", Identity()), ("down", control.down)] + XCTAssertThrowsError( + try control.update(modules: ModuleChildren.unflattened(notLinear), verify: .all) + ) { error in + XCTAssertTrue(error is UpdateError, "unexpected error: \(error)") + } + + // a float model, and one quantized first like a loaded quantized model + for quantized in [false, true] { + let block = ProjectionBlock(64, 256) + if quantized { + quantize(model: block) + } + let expected = block(x) + try checkedEval(expected) + + // what a model's shard method does + let up = try shardLinear(block.up, sharding: .allToSharded, group: world) + let down = try shardLinear(block.down, sharding: .shardedToAll, group: world) + let sharded: [(String, Module)] = [("up", up), ("down", down)] + try block.update(modules: ModuleChildren.unflattened(sharded), verify: .all) + + if quantized { + XCTAssertTrue(block.up is QuantizedAllToShardedLinear, "\(type(of: block.up))") + XCTAssertTrue(block.down is QuantizedShardedToAllLinear, "\(type(of: block.down))") + } else { + XCTAssertTrue(block.up is AllToShardedLinear, "\(type(of: block.up))") + XCTAssertTrue(block.down is ShardedToAllLinear, "\(type(of: block.down))") + } + + // every process computes the output of the whole block + let y = block(x) + try checkedEval(y) + XCTAssertTrue( + expected.allClose(y, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "the sharded block must reproduce the unsharded output (quantized: \(quantized))") + } +} + +/// Port of `test_donation`. +/// +/// A collective donates its result to the operation that consumes it, so +/// summing and then scaling must not cost more memory than summing alone. +/// +/// This one has no single process half: in a group of one MLX returns the +/// input unchanged, so there is no buffer to donate and the input is still +/// live. Python only ever runs it under a launcher for the same reason. +func donationBody(world: MLXDistributed.Group) throws { + let cpu = Stream.defaultStream(.cpu) + + let x = MLXRandom.normal([1024]) + try checkedEval(x) + cpu.synchronize() + GPU.resetPeakMemory() + + let scale = MLXArray(2.0) + + // Python rebinds one name, so the first result is released before the + // second is evaluated. Holding on to it here would add its buffer to the + // peak and make the comparison below fail for the wrong reason. + let allSumOnly: Int + do { + let sum = MLXDistributed.allSum(x, group: world) + try checkedEval(sum) + cpu.synchronize() + allSumOnly = Memory.peakMemory + } + + let scaled = MLXDistributed.allSum(x, group: world) * scale + try checkedEval(scaled) + cpu.synchronize() + let allSumWithBinary = Memory.peakMemory + + // the instrument has to move at all, or the comparison proves nothing + XCTAssertGreaterThan(allSumOnly, 0, "peak memory is not tracked for this stream") + XCTAssertEqual( + allSumOnly, allSumWithBinary, + "the multiply must donate the buffer the all sum produced") +} + +/// Port of `test_quantized_sharded_linear_construction` from the Python nn +/// tests. +/// +/// Every bit width packs a different number of values into each `uint32`, so +/// the sharded layers have to come out with the shapes the layer they were +/// built from would have, divided along the sharded axis. +func quantizedShardedConstructionBody(world: MLXDistributed.Group) throws { + for bits in [2, 3, 4, 5, 6, 8] { + let quantized = QuantizedLinear(Linear(1536, 1024), groupSize: 64, bits: bits) + let allToSharded = try QuantizedAllToShardedLinear(quantized, group: world) + let shardedToAll = try QuantizedShardedToAllLinear(quantized, group: world) + + // rows are outputs, columns are the packed inputs + XCTAssertEqual( + allToSharded.weight.dim(0), quantized.weight.dim(0) / world.size, "bits \(bits)") + XCTAssertEqual(allToSharded.weight.dim(1), quantized.weight.dim(1), "bits \(bits)") + + XCTAssertEqual(shardedToAll.weight.dim(0), quantized.weight.dim(0), "bits \(bits)") + XCTAssertEqual( + shardedToAll.weight.dim(1), quantized.weight.dim(1) / world.size, "bits \(bits)") + + // and the scales have to cover exactly the inputs the weight holds + for layer in [allToSharded, shardedToAll] { + XCTAssertEqual( + layer.weight.dim(1) * 32 / bits, layer.scales.dim(1) * layer.groupSize, + "bits \(bits)") + } + } +} + +/// Cases the Python tests do not cover, each one a bug this port had. +func shardingEdgeCasesBody(world: MLXDistributed.Group) throws { + // A sharded-to-all layer adds its bias after the reduction, so every + // process has to hold the same bias. Seed the processes differently: + // nothing keeps their generators in step -- they may have done different + // work before this -- and the bias must not depend on that. + MLXRandom.seed(UInt64(world.rank + 1)) + let layer = try ShardedToAllLinear(8 * world.size, 4, group: world) + let bias = try XCTUnwrap(layer.bias) + let gathered = MLXDistributed.allGather(bias, group: world) + try checkedEval(gathered) + + let width = bias.dim(0) + for rank in 1 ..< world.size { + XCTAssertTrue( + gathered[0 ..< width].allClose(gathered[rank * width ..< (rank + 1) * width]) + .item(Bool.self), + "rank \(rank) holds a different bias than rank 0") + } + + // A dimension can divide by the group size while one of its segments does + // not. That has to be reported rather than trap inside split(parts:). + if world.size > 1 { + XCTAssertThrowsError( + try shardLinear( + Linear(8 * world.size, 8), sharding: .shardedToAll, + segments: .indices([3]), group: world) + ) { error in + XCTAssertEqual( + error as? ShardingError, + .indivisible(dimension: "segment of weight", of: 3, across: world.size)) + } + } + + // Segments that cannot split a dimension at all are reported too, in a + // group of one as well. split(parts:) would call fatalError. + XCTAssertThrowsError( + try shardLinear( + Linear(8 * world.size, 8), sharding: .shardedToAll, segments: .count(3), + group: world) + ) { error in + XCTAssertEqual( + error as? ShardingError, .invalidSegments(.count(3), dimension: 8 * world.size)) + } + + // A QuantizedLinear is a Linear, so the float layers accept one as far as + // the compiler is concerned, but they would keep its packed weight as a + // float weight and drop its scales. + let quantized = QuantizedLinear(Linear(64, 64)) + XCTAssertThrowsError(try AllToShardedLinear(quantized, group: world)) { error in + XCTAssertEqual(error as? ShardingError, .quantizedLayer) + } + XCTAssertThrowsError(try ShardedToAllLinear(quantized, group: world)) { error in + XCTAssertEqual(error as? ShardingError, .quantizedLayer) + } + + // A quantized layer packs its inputs in the weight and groups them in the + // scales, so a boundary between segments lands at a different position in + // each. Scaled through a Double, 3840 of 5632 inputs came to + // 479.99999999999994 packed words and 59.99999999999999 groups, which + // truncate to the wrong positions. The shard has to reproduce the layer. + MLXRandom.seed(0xF0F0_F0F0) + let inputs = 5632 + let boundary = 3840 + let wide = QuantizedLinear(Linear(inputs, 8)) + let segmented = try QuantizedShardedToAllLinear( + wide, segments: .indices([boundary]), group: world) + + // this rank's part of each segment + let (rank, size) = (world.rank, world.size) + let first = (rank * boundary / size) ..< ((rank + 1) * boundary / size) + let second = + (boundary + rank * (inputs - boundary) / size) + ..< (boundary + (rank + 1) * (inputs - boundary) / size) + + let x = MLXRandom.normal([2, inputs]) + let y = wide(x) + let ySharded = segmented(concatenated([x[first, axis: 1], x[second, axis: 1]], axis: 1)) + try checkedEval(y, ySharded) + + XCTAssertTrue( + y.allClose(ySharded, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "a segmented quantized shard must reproduce the unsharded output") + + // a boundary inside a quantization group cannot be sharded at all + XCTAssertThrowsError( + try QuantizedShardedToAllLinear(wide, segments: .indices([272]), group: world) + ) { error in + XCTAssertEqual( + error as? ShardingError, + .quantizationGroup(inputDimensions: 272, groupSize: 64, across: world.size)) + } +} + +/// Port of `test_average_gradients`. +/// +/// Python counts the `all_sum` calls by replacing `mx.distributed.all_sum`, +/// which Swift cannot do, so this counts the batches `averageGradients` forms +/// instead. Every batching has to produce the same average, and the average +/// has to include every rank. +func averageGradientsBody(world: MLXDistributed.Group) throws { + // Ten gradients of ten float32 values, 40 bytes each as in Python, so the + // limits below form different batches. Each rank contributes its own + // rank + 1, so the average is (size + 1) / 2 times the gradient's own + // factor. The values differ per gradient and half are matrices on + // purpose: batching concatenates them, reduces once and splits the result + // again, and identical or flat arrays would hide a bad split. + let mean = Float(world.size + 1) / 2 + let shapes = (0 ..< 10).map { $0.isMultiple(of: 2) ? [10] : [2, 5] } + let gradients = ModuleParameters.unflattened( + (0 ..< 10).map { + ("g\($0)", MLXArray.ones(shapes[$0]) * Float(($0 + 1) * (world.rank + 1))) + }) + let expected = Dictionary( + uniqueKeysWithValues: (0 ..< 10).map { + ("g\($0)", MLXArray.ones(shapes[$0]) * (Float($0 + 1) * mean)) + }) + + // Python makes one all_sum call for 32MiB, two for 4 * 50 bytes and ten + // without batching. A limit of one byte batches every gradient alone. + let bytes = gradients.flattened().map { $0.1.nbytes } + XCTAssertEqual(groupBySize(bytes, limit: 32 * 1024 * 1024), [Array(0 ..< 10)]) + XCTAssertEqual(groupBySize(bytes, limit: 4 * 50), [Array(0 ..< 5), Array(5 ..< 10)]) + XCTAssertEqual(groupBySize(bytes, limit: 1), (0 ..< 10).map { [$0] }) + + for allReduceSize in [32 * 1024 * 1024, 4 * 50, 1, 0] { + let averaged = try averageGradients( + gradients, group: world, allReduceSize: allReduceSize) + let flat = averaged.flattened() + try checkedEval(flat.map { $0.1 }) + + XCTAssertEqual(flat.count, 10, "allReduceSize \(allReduceSize)") + for (key, value) in flat { + let want = try XCTUnwrap(expected[key]) + XCTAssertEqual(value.shape, want.shape, "\(key) with allReduceSize \(allReduceSize)") + XCTAssertTrue( + value.allClose(want).item(Bool.self), + "\(key) with allReduceSize \(allReduceSize)") + } + } + + // arrays of different types cannot be concatenated, so they fall back to + // one reduction per array and keep their dtype + let mixed = try averageGradients( + ModuleParameters.unflattened([ + ("wide", MLXArray.ones([4]) * Float(world.rank + 1)), + ("narrow", (MLXArray.ones([4]) * Float(world.rank + 1)).asType(.float16)), + ]), group: world) + // evaluate in the flattened order, which is sorted by key: a Dictionary + // iterates in a per process order, and evaluating the collectives in + // different orders makes the ranks exchange mismatched buffers + let flat = mixed.flattened() + try checkedEval(flat.map { $0.1 }) + + let flatMixed = Dictionary(uniqueKeysWithValues: flat) + XCTAssertEqual(flatMixed["narrow"]?.dtype, .float16, "the dtype must survive") + for key in ["wide", "narrow"] { + let averaged = try XCTUnwrap(flatMixed[key]) + XCTAssertTrue( + averaged.asType(.float32) + .allClose(MLXArray.ones([4]) * (Float(world.size + 1) / 2), rtol: 1e-2) + .item(Bool.self), key) + } +} + +/// Port of `test_clip_grad_norm_sharded`. +func clipGradNormShardedBody(world: MLXDistributed.Group) throws { + let value: Float = 3 + let gradients = ModuleParameters.unflattened([ + ("a", MLXArray.ones([4, 3]) * value), + ("b", MLXArray.ones([5]) * value), + ]) + + // every rank holds the same number of elements, so the global norm counts + // all of them + let localCount = 4 * 3 + 5 + let expectedNorm = (Float(world.size * localCount)).squareRoot() * value + + // a limit far above the norm leaves the shard alone + let (unclipped, norm) = try clipGradNormSharded( + gradients: gradients, maxNorm: 1e9, group: world) + try checkedEval(norm, unclipped.flattened().map { $0.1 }) + + XCTAssertEqual(norm.item(Float.self), expectedNorm, accuracy: expectedNorm * 1e-4) + for (key, clipped) in unclipped.flattened() { + XCTAssertTrue( + clipped.allClose(MLXArray.ones(clipped.shape) * value).item(Bool.self), key) + } + + // below it every gradient is scaled by maxNorm / norm + let maxNorm: Float = 1 + let (clipped, _) = try clipGradNormSharded( + gradients: gradients, maxNorm: maxNorm, group: world) + try checkedEval(clipped.flattened().map { $0.1 }) + + let scale = maxNorm / (expectedNorm + 1e-6) + for (key, scaled) in clipped.flattened() { + XCTAssertTrue( + scaled.allClose(MLXArray.ones(scaled.shape) * value * scale, rtol: 1e-4) + .item(Bool.self), key) + } +} + +/// A block that declares its projections as `Linear`, the way models do. +private class ProjectionBlock: Module, UnaryLayer { + + @ModuleInfo(key: "up") var up: Linear + @ModuleInfo(key: "down") var down: Linear + + init(_ dimensions: Int, _ hiddenDimensions: Int) { + self._up.wrappedValue = Linear(dimensions, hiddenDimensions) + self._down.wrappedValue = Linear(hiddenDimensions, dimensions) + } + + func callAsFunction(_ x: MLXArray) -> MLXArray { + down(relu(up(x))) + } +} + +/// A convolution that can sum its output across the group, so that a layer fed +/// sharded input channels produces the whole result. +private class AggregatingConv: Module, UnaryLayer { + + @ModuleInfo(key: "conv") var conv: Conv2d + + let aggregate: MLXDistributed.Group? + + init( + _ inputChannels: Int, _ outputChannels: Int, kernelSize: Int = 3, + aggregate: MLXDistributed.Group? = nil + ) { + self._conv.wrappedValue = Conv2d( + inputChannels: inputChannels, outputChannels: outputChannels, + kernelSize: IntOrPair(kernelSize)) + self.aggregate = aggregate + } + + func callAsFunction(_ x: MLXArray) -> MLXArray { + let y = conv(x) + guard let aggregate else { return y } + return MLXDistributed.allSum(y, group: aggregate) + } +} + +private func assertGradientShard( + _ full: [String: MLXArray], _ shard: [String: MLXArray], _ key: String, + _ range: Range, axis: Int, file: StaticString = #filePath, line: UInt = #line +) { + guard let expected = full[key], let actual = shard[key] else { + return XCTFail("missing gradient for \(key)", file: file, line: line) + } + XCTAssertTrue( + expected[range, axis: axis].allClose(actual, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "\(key) gradient shard mismatch", file: file, line: line) +} + +private func assertGradientWhole( + _ full: [String: MLXArray], _ shard: [String: MLXArray], _ key: String, + file: StaticString = #filePath, line: UInt = #line +) { + guard let expected = full[key], let actual = shard[key] else { + return XCTFail("missing gradient for \(key)", file: file, line: line) + } + XCTAssertTrue( + expected.allClose(actual, rtol: 1e-4, atol: 1e-6).item(Bool.self), + "\(key) gradient mismatch", file: file, line: line) +} + +/// The single process half: sharding is the identity in a group of size one, +/// but the layers are still built and run. +/// +/// That is enough to catch a layer that mishandles its input -- a +/// ``QuantizedLinear`` sharded as if it were a float layer, for example -- +/// which is why this runs in CI without a launcher. +class DistributedNNTests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + override func setUpWithError() throws { + try XCTSkipIf( + ProcessInfo.processInfo.environment["MLX_TEST_DISTRIBUTED"] == "1", + "The multi process run covers this; see DistributedNNRingTests.") + } + + func testShardLinear() throws { + try shardLinearBody(world: try MLXDistributed.initialize()) + } + + func testShardPredicate() throws { + try shardPredicateBody(world: try MLXDistributed.initialize()) + } + + func testModuleSubstitution() throws { + try moduleSubstitutionBody(world: try MLXDistributed.initialize()) + } + + func testQuantizedShardedConstruction() throws { + try quantizedShardedConstructionBody(world: try MLXDistributed.initialize()) + } + + func testShardingEdgeCases() throws { + try shardingEdgeCasesBody(world: try MLXDistributed.initialize()) + } + + func testAverageGradients() throws { + try averageGradientsBody(world: try MLXDistributed.initialize()) + } + + func testClipGradNormSharded() throws { + try clipGradNormShardedBody(world: try MLXDistributed.initialize()) + } +} + +/// The multi process half, where every rank holds a different slice. +/// +/// Skipped unless `MLX_TEST_DISTRIBUTED=1`: +/// +/// ``` +/// MLX_TEST_DISTRIBUTED=1 xcrun xctest -XCTest DistributedNNRingTests \ +/// .../MLXTests.xctest +/// ``` +/// +/// Like ``DistributedRingTests`` this has a single test method, since a +/// process can only join one ring. It runs every ported body in turn, the way +/// the Python tests all share the group their launcher formed. +class DistributedNNRingTests: XCTestCase { + + static let testName = "DistributedNNRingTests/testShardedLayers" + + /// Four ranks: the dimensions the Python tests use, 1024 and 128, divide by + /// four, and a shard of 1024 inputs still holds whole quantization groups. + static let rankCount = 4 + + override class func setUp() { + setDefaultDevice() + } + + func testShardedLayers() throws { + try DistributedHarness.run(ranks: Self.rankCount, testName: Self.testName) { group in + try shardLinearBody(world: group) + try shardPredicateBody(world: group) + try moduleSubstitutionBody(world: group) + try donationBody(world: group) + try quantizedShardedConstructionBody(world: group) + try shardingEdgeCasesBody(world: group) + try averageGradientsBody(world: group) + try clipGradNormShardedBody(world: group) + } + } +} diff --git a/Tests/MLXTests/DistributedRingTests.swift b/Tests/MLXTests/DistributedRingTests.swift new file mode 100644 index 000000000..71f619513 --- /dev/null +++ b/Tests/MLXTests/DistributedRingTests.swift @@ -0,0 +1,191 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +/// Multi process tests for the ring backend, ported from the Python +/// distributed tests: `test_groups`, `test_all_reduce`, `test_all_gather`, +/// `test_send_recv` and `test_all_gather_vjp`, together with the ring specific +/// `test_all_reduce_extra` and `test_all_gather_extra`. +/// +/// ``DistributedHarness`` runs the body in several processes. They are +/// skipped unless `MLX_TEST_DISTRIBUTED=1` is set: +/// +/// ``` +/// MLX_TEST_DISTRIBUTED=1 xcrun xctest -XCTest DistributedRingTests \ +/// .../MLXTests.xctest +/// ``` +/// +/// There is deliberately a single test method: MLX caches the group per +/// process, and a single method also guarantees every rank issues the same +/// operations in the same order, which the ring backend requires. +class DistributedRingTests: XCTestCase { + + static let testName = "DistributedRingTests/testRingCollectives" + + /// Three ranks, so that the left and right neighbors of a rank are + /// different processes. + static let rankCount = 3 + + /// The dtypes of the Python tests, including the ring specific ones. + static let dtypes: [DType] = [ + .int8, .uint8, .int16, .uint16, .int32, .uint32, .float32, .float16, .bfloat16, + .complex64, + ] + + /// Sizes from the Python tests. The large ones exercise the ring's chunked + /// transfers, which a handful of elements never reach. + static let shapes = [[7], [10], [1024], [1024, 1024]] + + func testRingCollectives() throws { + try DistributedHarness.run(ranks: Self.rankCount, testName: Self.testName) { group in + try runRankBody(group) + } + } + + private func runRankBody(_ group: MLXDistributed.Group) throws { + try groups(group) + try reductions(group) + try allGather(group) + try sendRecv(group) + try allGatherVJP(group) + } + + /// Port of `test_groups`. + private func groups(_ group: MLXDistributed.Group) throws { + XCTAssertEqual(group.size, Self.rankCount) + XCTAssertTrue(group.rank >= 0 && group.rank < group.size) + + // initializing again yields the same group + let again = try MLXDistributed.initialize() + XCTAssertEqual(again.size, group.size) + XCTAssertEqual(again.rank, group.rank) + + // the ring backend does not support splitting + XCTAssertThrowsError(try group.split(color: group.rank % 2)) + } + + /// Port of `test_all_reduce` and `test_all_reduce_extra`. + /// + /// Every rank builds the same array and contributes its own row, so the + /// reduction over the rows is the expected result in every process. + private func reductions(_ group: MLXDistributed.Group) throws { + let tolerances: [DType: Float] = [ + .float32: 1e-6, .float16: 5e-3, .bfloat16: 1e-1, .complex64: 1e-6, + ] + let key = MLXRandom.key(0) + var combinations = 0 + defer { + if group.rank == 0 { + print("reductions: \(combinations) dtype and shape combinations") + } + } + + for dtype in Self.dtypes { + for shape in Self.shapes { + let rtol = tolerances[dtype] ?? 0 + let x = (MLXRandom.uniform(0 ..< 1, [group.size] + shape, key: key) * 10) + .asType(dtype) + let name = "\(dtype) \(shape)" + combinations += 1 + + let sum = MLXDistributed.allSum(x[group.rank], group: group) + let expected = x.sum(axis: 0) + var error = abs(sum - expected) + if rtol > 0 { + error = error / abs(expected) + } + try checkedEval(error) + XCTAssertLessThanOrEqual( + error.max().asType(.float32).item(Float.self), rtol, "allSum \(name)") + + let maximum = MLXDistributed.allMax(x[group.rank], group: group) + try checkedEval(maximum) + XCTAssertTrue( + (maximum .== x.max(axis: 0)).all().item(Bool.self), "allMax \(name)") + + let minimum = MLXDistributed.allMin(x[group.rank], group: group) + try checkedEval(minimum) + XCTAssertTrue( + (minimum .== x.min(axis: 0)).all().item(Bool.self), "allMin \(name)") + } + } + } + + /// Port of `test_all_gather` and `test_all_gather_extra`. + private func allGather(_ group: MLXDistributed.Group) throws { + for dtype in Self.dtypes { + let x = MLXArray.ones([2, 2, 4]).asType(dtype) + let gathered = MLXDistributed.allGather(x, group: group) + try checkedEval(gathered) + + XCTAssertEqual(gathered.shape, [group.size * 2, 2, 4], "allGather \(dtype)") + XCTAssertTrue( + (gathered .== MLXArray(1).asType(dtype)).all().item(Bool.self), + "allGather \(dtype)") + } + + // the shards are concatenated in rank order rather than reduced + let base = MLXArray([1, 2, 3]).asType(.float32) + let gathered = MLXDistributed.allGather(base * Float(group.rank + 1), group: group) + try checkedEval(gathered) + assertEqual(gathered, concatenated((0 ..< group.size).map { base * Float($0 + 1) })) + } + + /// Port of `test_send_recv`. + /// + /// Every rank sends to its right neighbor and receives from its left one. + /// Even ranks send first and odd ranks receive first, so that neighbors do + /// not wait on each other. + private func sendRecv(_ group: MLXDistributed.Group) throws { + let rank = group.rank + let size = group.size + let right = (rank + 1) % size + let left = (rank + size - 1) % size + let key = MLXRandom.key(0) + var transfers = 0 + defer { + if rank == 0 { + print("send/recv: \(transfers) dtype and shape combinations") + } + } + + for dtype in Self.dtypes { + for shape in Self.shapes { + let x = (MLXRandom.uniform(0 ..< 1, [size] + shape, key: key) * 10) + .asType(dtype) + let name = "\(dtype) \(shape)" + transfers += 1 + + let sent: MLXArray + let received: MLXArray + if rank % 2 == 0 { + sent = MLXDistributed.send(x[rank], to: right, group: group) + received = MLXDistributed.recvLike(sent, from: left, group: group) + try checkedEval(sent, received) + } else { + received = MLXDistributed.recvLike(x[rank], from: left, group: group) + sent = MLXDistributed.send(x[rank], to: right, group: group) + try checkedEval(received, sent) + } + + XCTAssertTrue((sent .== x[rank]).all().item(Bool.self), "send \(name)") + XCTAssertTrue((received .== x[left]).all().item(Bool.self), "recv \(name)") + } + } + } + + /// Port of `test_all_gather_vjp`. + /// + /// The gathered array starts with rank 0's contribution, so only rank 0 + /// sees a gradient. + private func allGatherVJP(_ group: MLXDistributed.Group) throws { + let gradient = grad { x in + MLXDistributed.allGather(x, group: group)[0] + }(MLXArray(1.0)) + try checkedEval(gradient) + + XCTAssertEqual(gradient.item(Float.self), group.rank == 0 ? 1.0 : 0.0) + } +} diff --git a/Tests/MLXTests/DistributedTests.swift b/Tests/MLXTests/DistributedTests.swift new file mode 100644 index 000000000..050c873de --- /dev/null +++ b/Tests/MLXTests/DistributedTests.swift @@ -0,0 +1,136 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +/// Single process tests for ``MLXDistributed``. +/// +/// Without `MLX_HOSTFILE` / `MLX_RANK` no backend can form a communication +/// group, so `initialize()` yields an empty group of size one. At that size +/// MLX short circuits the collectives to the identity and rejects the point to +/// point operations, which is enough to cover the whole API surface in CI. +/// +/// Multi process behaviour is verified by launching several copies of a +/// program with a hostfile, matching how the Python tests are run. +class DistributedTests: XCTestCase { + + override class func setUp() { + setDefaultDevice() + } + + /// The assertions here describe a group of size one, so they are not valid + /// in a multi process run -- see ``DistributedRingTests``. + override func setUpWithError() throws { + try XCTSkipIf( + ProcessInfo.processInfo.environment["MLX_TEST_DISTRIBUTED"] == "1", + "Single process assertions; skipped during a multi process run.") + } + + // MARK: - Availability + + func testIsAvailableLinks() { + for backend in MLXDistributed.Backend.allCases { + _ = MLXDistributed.isAvailable(backend) + } + } + + /// The ring backend is built on all platforms -- it needs only TCP sockets. + func testRingIsAvailable() { + XCTAssertTrue(MLXDistributed.isAvailable(.ring)) + + // `any` is true when at least one backend is available + XCTAssertTrue(MLXDistributed.isAvailable(.any)) + } + + /// Backends that are not compiled in report `false` rather than trapping. + func testUnavailableBackendsReportFalse() { + XCTAssertFalse(MLXDistributed.isAvailable(.mpi)) + XCTAssertFalse(MLXDistributed.isAvailable(.jaccl)) + + // nccl is CUDA-only and is never built on Apple platforms + #if !os(Linux) + XCTAssertFalse(MLXDistributed.isAvailable(.nccl)) + #endif + } + + // MARK: - Group + + func testDefaultGroupIsSingleton() throws { + let group = try MLXDistributed.initialize() + XCTAssertEqual(group.rank, 0) + XCTAssertEqual(group.size, 1) + } + + /// A strict init throws when no backend can form a group. + /// + /// Note this uses `.any`: MLX caches successfully registered groups per + /// backend name, but never caches under "any" when initialization fails, + /// so this is independent of the order the tests run in. + func testStrictInitializeThrows() { + XCTAssertThrowsError(try MLXDistributed.initialize(backend: .any, strict: true)) + } + + /// A backend that is configured but cannot form its group throws even when + /// the init is not strict -- here the ring backend with a missing hostfile. + func testInitializeThrowsWhenConfiguredBackendFails() throws { + let environment = ProcessInfo.processInfo.environment + try XCTSkipIf( + environment["MLX_HOSTFILE"] != nil || environment["MLX_RANK"] != nil, + "MLX_HOSTFILE or MLX_RANK is already set") + + let hostfile = FileManager.default.temporaryDirectory + .appendingPathComponent("missing-hostfile-\(UUID().uuidString).json") + setenv("MLX_HOSTFILE", hostfile.path, 1) + setenv("MLX_RANK", "0", 1) + defer { + unsetenv("MLX_HOSTFILE") + unsetenv("MLX_RANK") + } + + XCTAssertThrowsError(try MLXDistributed.initialize(backend: .any)) + } + + func testSplitSingletonGroupThrows() throws { + let group = try MLXDistributed.initialize() + XCTAssertThrowsError(try group.split(color: 0)) + } + + // MARK: - Collectives + + /// In a group of size one the collectives return their input unchanged. + func testCollectivesAreIdentityInSingletonGroup() { + let x = MLXArray([1, 2, 3, 4], [2, 2]) + + assertEqual(MLXDistributed.allSum(x), x) + assertEqual(MLXDistributed.allMax(x), x) + assertEqual(MLXDistributed.allMin(x), x) + assertEqual(MLXDistributed.allGather(x), x) + assertEqual(MLXDistributed.sumScatter(x), x) + } + + // MARK: - Point to point + + func testSendInSingletonGroupReportsError() { + let x = MLXArray([1, 2, 3]) + XCTAssertThrowsError( + try withError { + MLXDistributed.send(x, to: 0) + }) + } + + func testRecvInSingletonGroupReportsError() { + XCTAssertThrowsError( + try withError { + MLXDistributed.recv([3], dtype: .int32, from: 0) + }) + } + + func testRecvLikeInSingletonGroupReportsError() { + let x = MLXArray([1, 2, 3]) + XCTAssertThrowsError( + try withError { + MLXDistributed.recvLike(x, from: 0) + }) + } +} diff --git a/xcode/MLX.xcodeproj/project.pbxproj b/xcode/MLX.xcodeproj/project.pbxproj index c9f0a2e50..ac8e88396 100644 --- a/xcode/MLX.xcodeproj/project.pbxproj +++ b/xcode/MLX.xcodeproj/project.pbxproj @@ -1032,7 +1032,7 @@ mlx/distributed/primitives.h, mlx/distributed/reduction_ops.h, mlx/distributed/ring/CMakeLists.txt, - mlx/distributed/ring/ring.cpp, + mlx/distributed/ring/no_ring.cpp, mlx/distributed/ring/ring.h, mlx/distributed/utils.h, mlx/dtype_utils.h, From bac95f7eb7a65e033ae025e4b28a0cafd18a8da4 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:03:06 -0700 Subject: [PATCH 31/39] add thin logging shim (#484) * add thin logging shim --- Source/MLX/Logging.swift | 300 ++++++++++++++++++++++++++++++ Tests/MLXTests/LoggingTests.swift | 84 +++++++++ 2 files changed, 384 insertions(+) create mode 100644 Source/MLX/Logging.swift create mode 100644 Tests/MLXTests/LoggingTests.swift diff --git a/Source/MLX/Logging.swift b/Source/MLX/Logging.swift new file mode 100644 index 000000000..b909f42e4 --- /dev/null +++ b/Source/MLX/Logging.swift @@ -0,0 +1,300 @@ +// Copyright © 2026 Apple Inc. + +import Foundation + +public enum LogLevel: Int, Comparable, Sendable, CustomStringConvertible { + case trace = 0 + case debug = 1 + case info = 2 + case warning = 3 + case error = 4 + + public static func < (lhs: LogLevel, rhs: LogLevel) -> Bool { + lhs.rawValue < rhs.rawValue + } + + public var description: String { + switch self { + case .trace: "trace" + case .debug: "debug" + case .info: "info" + case .warning: "warning" + case .error: "error" + } + } +} + +/// Protocol for backing of MLXLogger. +/// +/// Install with ``MLXLogger/factory`` or ``StderrHandler/install()`` +public protocol MLXLogHandler: Sendable { + func log( + level: LogLevel, + message: () -> String, + metadata: () -> [String: String]?, + file: StaticString, + function: StaticString, + line: UInt + ) +} + +public typealias MLXLogHandlerFactory = @Sendable (String) -> any MLXLogHandler + +#if canImport(OSLog) + private let defaultFactory: MLXLogHandlerFactory = { OSLogHandler(label: $0) } +#else + private let defaultFactory: MLXLogHandlerFactory = { StderrHandler(label: $0) } +#endif + +private let factoryLock = NSLock() +#if swift(>=5.10) + nonisolated(unsafe) private var handlerFactory = defaultFactory +#else + private var handlerFactory = defaultFactory +#endif + +/// Thin wrapper for logging in MLX and MLX libraries. This +/// provides a hook for library code if it needs logging without +/// adding a dependency or requiring a certain logging backend. +/// +/// By default it will log to `OSLog` (macOS, iOS) or `stderr` +/// (others). Developers can override the backing by implementing +/// ``MLXLogHandler`` and installing it with ``MLXLogger/factory``, +/// e.g. +/// +/// ```swift +/// MLXLogger.factory = { StderrHandler(label: $0) } +/// ``` +/// +/// The factory is used at the time of the first log for the given +/// logger -- whatever factory is installed at that time will be used. +/// +/// Callers can use this the same way that `swift-log` is used: +/// +/// ```swift +/// private let logger = MLXLogger(label: "KVCache") +/// +/// func ...() { +/// logger.error("Failed to ...") +/// } +/// ``` +public final class MLXLogger: @unchecked (Sendable) { + + private let label: String + + /// backing that is realized at first log time + private let lock = NSLock() + private var backing: MLXLogHandler? + + public static var factory: MLXLogHandlerFactory { + get { + factoryLock.lock() + defer { factoryLock.unlock() } + return handlerFactory + } + set { + factoryLock.lock() + defer { factoryLock.unlock() } + handlerFactory = newValue + } + } + + public init(label: String) { + self.label = label + } + + // note: this avoids `lock.withLock { ... }` returning an existential -- + // that pattern crashes the Swift 6.3 compiler on Linux (SIL + // ClosureLifetimeFixup pass) + private func _existingBacking() -> MLXLogHandler? { + lock.lock() + defer { lock.unlock() } + return backing + } + + private func _setBacking(_ handler: MLXLogHandler) { + lock.lock() + defer { lock.unlock() } + self.backing = handler + } + + private func _backing() -> MLXLogHandler { + if let backing = _existingBacking() { + return backing + } + + let backing = Self.factory(label) + _setBacking(backing) + return backing + } + + fileprivate func _log( + level: LogLevel, + message: () -> String, metadata: () -> [String: String]?, + file: StaticString, function: StaticString, line: UInt + ) { + _backing().log( + level: level, + message: message, metadata: metadata, + file: file, function: function, line: line) + } +} + +extension MLXLogger { + + /// Log with the given level and message. + public func log( + level: LogLevel, + message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]?, + file: StaticString = #fileID, + function: StaticString = #function, line: UInt = #line + ) { + _log( + level: level, message: message, metadata: metadata, + file: file, function: function, line: line) + } + + /// Log a trace message. + public func trace( + _ message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]? = nil, + file: StaticString = #fileID, + function: StaticString = #function, + line: UInt = #line + ) { + _log( + level: .trace, message: message, metadata: metadata, + file: file, function: function, line: line) + } + + /// Log a debug message. + public func debug( + _ message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]? = nil, + file: StaticString = #fileID, + function: StaticString = #function, + line: UInt = #line + ) { + _log( + level: .debug, message: message, metadata: metadata, + file: file, function: function, line: line) + } + + /// Log a info message. + public func info( + _ message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]? = nil, + file: StaticString = #fileID, + function: StaticString = #function, + line: UInt = #line + ) { + _log( + level: .info, message: message, metadata: metadata, + file: file, function: function, line: line) + } + + /// Log a warning message. + public func warning( + _ message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]? = nil, + file: StaticString = #fileID, + function: StaticString = #function, + line: UInt = #line + ) { + _log( + level: .warning, message: message, metadata: metadata, + file: file, function: function, line: line) + } + + /// Log an error message. + public func error( + _ message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]? = nil, + file: StaticString = #fileID, + function: StaticString = #function, + line: UInt = #line + ) { + _log( + level: .error, message: message, metadata: metadata, + file: file, function: function, line: line) + } +} + +#if canImport(OSLog) + import OSLog + + /// On systems that support it, a log handler that is backed by OSLog. + public struct OSLogHandler: MLXLogHandler { + + private let logger: os.Logger + + public static func install() { + MLXLogger.factory = { OSLogHandler(label: $0) } + } + + public init(label: String) { + self.logger = Logger(subsystem: "mlx-swift", category: label) + } + + public func log( + level: LogLevel, + message: () -> String, metadata: () -> [String: String]?, + file: StaticString, function: StaticString, line: UInt + ) { + let level: OSLogType = + switch level { + case .trace: .debug + case .debug: .debug + case .info: .info + case .warning: .error + case .error: .error + } + if logger.isEnabled(type: level) { + let message = message() + logger.log(level: level, "\(message, privacy: .public)") + } + } + } +#endif + +/// A logger that logs to `stderr`. +/// +/// Callers can set `StderrHandler.logLevel` to filter output by level. +public struct StderrHandler: MLXLogHandler { + + private let label: String + + private static let lock = NSLock() + #if swift(>=5.10) + nonisolated(unsafe) private static var _logLevel = LogLevel.info + #else + private static var _logLevel = LogLevel.info + #endif + + /// Get or set the lowest log level. Default is `.info`. + public static var logLevel: LogLevel { + get { lock.withLock { _logLevel } } + set { lock.withLock { _logLevel = newValue } } + } + + public static func install() { + MLXLogger.factory = { StderrHandler(label: $0) } + } + + public init(label: String) { + self.label = label + } + + public func log( + level: LogLevel, message: @autoclosure () -> String, + metadata: @autoclosure () -> [String: String]?, file: StaticString, function: StaticString, + line: UInt + ) { + if level >= Self.lock.withLock({ Self._logLevel }) { + let ts = Date().formatted(date: .numeric, time: .standard) + let message = "\(ts) [\(level)]: \(message())\n" + FileHandle.standardError.write(Data(message.utf8)) + } + } +} diff --git a/Tests/MLXTests/LoggingTests.swift b/Tests/MLXTests/LoggingTests.swift new file mode 100644 index 000000000..ae33167bc --- /dev/null +++ b/Tests/MLXTests/LoggingTests.swift @@ -0,0 +1,84 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import Testing + +private let lock = NSLock() +nonisolated(unsafe) private var logs = [String: [String]]() + +@Suite(.serialized) struct LoggingTests { + + class CollectHandler: MLXLogHandler, @unchecked (Sendable) { + + private let label: String + public init(label: String) { + self.label = label + } + + func log( + level: MLX.LogLevel, message: () -> String, metadata: () -> [String: String]?, + file: StaticString, function: StaticString, line: UInt + ) { + let message = level.description + ": " + message() + lock.withLock { + logs[label, default: []].append(message) + } + } + } + + @Test func testLoggingFactory() async throws { + MLXLogger.factory = { CollectHandler(label: $0) } + let logger = MLXLogger(label: "case1") + + logger.debug("test debug") + logger.info("test info") + logger.error("test error") + + let list = lock.withLock { + logs["case1"] ?? [] + } + + #expect( + list == [ + "debug: test debug", + "info: test info", + "error: test error", + ]) + } + + #if canImport(OSLog) + @Test func testOSLog() async throws { + OSLogHandler.install() + let logger = MLXLogger(label: "case2") + + // no assertions but should run without error + logger.debug("test debug") + logger.info("test info") + logger.error("test error") + + // and not collect + let list = lock.withLock { + logs["case2"] + } + #expect(list == nil) + } + #endif + + @Test func testStderr() async throws { + StderrHandler.install() + let logger = MLXLogger(label: "case3") + + // no assertions but should run without error + logger.debug("test debug") + logger.info("test info") + logger.error("test error") + + // and not collect + let list = lock.withLock { + logs["case3"] + } + #expect(list == nil) + } + +} From 9ba154552c6e55bc822ca03008e4e58a6d22c816 Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:04:36 -0700 Subject: [PATCH 32/39] patch for Cuda builds (#478) * patch for Cuda builds - local apply of https://github.com/ml-explore/mlx/pull/4480 + other shutdown items from https://github.com/ml-explore/mlx/issues/4506 - fixes crash on exit for cuda builds when using cmake --- .github/scripts/build-linux-cuda-cmake.sh | 13 ++++- CMakeLists.txt | 52 +++++++++++++++-- cmake/apply-patch.cmake | 69 +++++++++++++++++++++++ cmake/mlx.patch | 34 +++++++++++ 4 files changed, 160 insertions(+), 8 deletions(-) create mode 100644 cmake/apply-patch.cmake create mode 100644 cmake/mlx.patch diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index 2a571ea6d..2fc81d589 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.github/scripts/build-linux-cuda-cmake.sh @@ -16,8 +16,15 @@ rm -rf build mkdir -p build pushd build cmake -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=ON -DMLX_C_BUILD_EXAMPLES=OFF .. -G Ninja + +# Verify the local mlx patch is actually in place: without it the examples below +# crash on exit, which is easy to misread as an unrelated CUDA failure. +git -C _deps/mlx-src apply --reverse --check "$PWD/../cmake/mlx.patch" || { + echo "error: cmake/mlx.patch is not applied to _deps/mlx-src" + exit 1 +} + ninja -# TODO dkoski -- disabled for now until clear_streams is available -# ./example1 --device gpu -# ./tutorial --device gpu +./example1 --device gpu +./tutorial --device gpu popd diff --git a/CMakeLists.txt b/CMakeLists.txt index 41404a01d..22f4b16ff 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -14,7 +14,9 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # ----------------------------- Configuration ----------------------------- # note: mirrors a subset of MLX options exactly (1:1 mapping) -option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON) +# note: MLX_BUILD_EXAMPLES is owned by mlx itself (see below), so the Swift +# examples use their own option name to avoid colliding with mlx's C++ examples +option(MLX_SWIFT_BUILD_EXAMPLES "Build Swift examples for mlx-swift" ON) option(MLX_BUILD_METAL "Build metal backend" ON) option(MLX_BUILD_CUDA "Build cuda backend" OFF) @@ -26,7 +28,45 @@ if(POLICY CMP0135) cmake_policy(SET CMP0135 NEW) endif() +# mlx + +# note: patches are applied via cmake/apply-patch.cmake so that they are +# idempotent (the patch step re-runs on every update step) while still failing +# the build loudly if a patch does not apply -- a silently unpatched dependency +# is very hard to diagnose (it typically shows up as a runtime failure). +set(apply_patch ${CMAKE_COMMAND} -DREPO=) +set(apply_patch_script -P ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake) + +# two local patches for CUDA: +# +# * https://github.com/ml-explore/mlx/pull/4480: leak the global command encoder +# map rather than synchronizing on process shutdown +# * make ~CudaHandle non-throwing: cu::Worker is detached and holds the last +# reference to itself, so ~Worker -> ~CudaEvent -> cudaEventDestroy runs on +# the worker thread at an arbitrary point during exit, racing the CUDA +# runtime's own teardown. A failing destroy there terminated the process. +set(mlx_patch + ${apply_patch} -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/mlx.patch + ${apply_patch_script}) + +# note: this must be declared (and named `mlx`) before mlx-c is made available +# so that our pinned/patched version wins over the one mlx-c declares. Because +# we now pull mlx in first, we are also responsible for the option defaults that +# mlx-c would normally set for us. +set(MLX_BUILD_TESTS OFF) +set(MLX_BUILD_EXAMPLES OFF) +set(MLX_BUILD_BENCHMARKS OFF) +set(MLX_BUILD_PYTHON_BINDINGS OFF) + +FetchContent_Declare( + mlx + GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" + GIT_TAG "v0.32.2" + PATCH_COMMAND ${mlx_patch}) +FetchContent_MakeAvailable(mlx) + # mlx-c +set(MLX_C_BUILD_EXAMPLES OFF) FetchContent_Declare( mlx-c GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git" @@ -34,13 +74,15 @@ FetchContent_Declare( FetchContent_MakeAvailable(mlx-c) # swift-numerics -set(swift_numerics_patch git apply - ${CMAKE_CURRENT_SOURCE_DIR}/cmake/swift-numerics.patch) +set(swift_numerics_patch + ${apply_patch} + -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/swift-numerics.patch + ${apply_patch_script}) FetchContent_Declare( swift-numerics GIT_REPOSITORY "https://github.com/apple/swift-numerics.git" GIT_TAG "1.0.2" - PATCH_COMMAND ${swift_numerics_patch} || true) + PATCH_COMMAND ${swift_numerics_patch}) FetchContent_MakeAvailable(swift-numerics) # MLX package @@ -107,7 +149,7 @@ add_library(MLXLinalg STATIC ${MLXLinalg-src}) target_link_libraries(MLXLinalg PRIVATE MLX) # Examples -if(MLX_BUILD_EXAMPLES) +if(MLX_SWIFT_BUILD_EXAMPLES) add_executable(example1 ${CMAKE_CURRENT_LIST_DIR}/Source/Examples/Example1.swift) target_link_libraries(example1 PRIVATE MLX) diff --git a/cmake/apply-patch.cmake b/cmake/apply-patch.cmake new file mode 100644 index 000000000..63e04f83f --- /dev/null +++ b/cmake/apply-patch.cmake @@ -0,0 +1,69 @@ +# Idempotent `git apply` for FetchContent PATCH_COMMAND. +# +# Usage (from a FetchContent_Declare PATCH_COMMAND): +# +# PATCH_COMMAND ${CMAKE_COMMAND} -DREPO= +# -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/foo.patch -P +# ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake +# +# Unlike `git apply ... || true` this does not swallow real failures: an +# already-applied patch is a no-op (the patch step re-runs on every update), but +# a patch that fails to apply is a hard error instead of a silently unpatched +# dependency. + +cmake_minimum_required(VERSION 3.16) + +if(NOT DEFINED PATCH OR PATCH STREQUAL "") + message(FATAL_ERROR "apply-patch.cmake: -DPATCH= is required") +endif() + +if(NOT DEFINED REPO + OR REPO STREQUAL "" + OR NOT IS_DIRECTORY "${REPO}") + # the patch step runs with the dependency source dir as its working directory + set(REPO "${CMAKE_CURRENT_BINARY_DIR}") +endif() + +if(NOT EXISTS "${PATCH}") + message(FATAL_ERROR "apply-patch.cmake: no such patch file: ${PATCH}") +endif() + +# note: deliberately not `find_package(Git REQUIRED)` -- that pulls in +# FindPackageHandleStandardArgs in script mode for no benefit here. +if(NOT DEFINED GIT_EXECUTABLE OR GIT_EXECUTABLE STREQUAL "") + find_program(GIT_EXECUTABLE NAMES git git.exe) +endif() + +if(NOT GIT_EXECUTABLE) + message( + FATAL_ERROR "apply-patch.cmake: git not found (pass -DGIT_EXECUTABLE=)") +endif() + +message(STATUS "apply-patch: repo=${REPO} patch=${PATCH} git=${GIT_EXECUTABLE}") + +# already applied? (the patch step re-runs whenever the update step re-runs) +execute_process( + COMMAND ${GIT_EXECUTABLE} apply --reverse --check "${PATCH}" + WORKING_DIRECTORY "${REPO}" + RESULT_VARIABLE reverse_check + OUTPUT_QUIET ERROR_QUIET) + +if(reverse_check EQUAL 0) + message(STATUS "apply-patch: already applied: ${PATCH}") + return() +endif() + +# note: git's stdout/stderr are intentionally *not* captured so that the real +# reason a patch failed always reaches the build log, even if this script dies +# before it can format a message of its own. +execute_process( + COMMAND ${GIT_EXECUTABLE} apply --verbose "${PATCH}" + WORKING_DIRECTORY "${REPO}" + RESULT_VARIABLE apply_result) + +if(NOT apply_result EQUAL 0) + message(FATAL_ERROR "apply-patch: failed to apply ${PATCH} in ${REPO} " + "(git apply exit ${apply_result}, see output above)") +endif() + +message(STATUS "apply-patch: applied ${PATCH}") diff --git a/cmake/mlx.patch b/cmake/mlx.patch new file mode 100644 index 000000000..de99b7ebc --- /dev/null +++ b/cmake/mlx.patch @@ -0,0 +1,34 @@ +diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h +index f8a234ee6..3b7a79f27 100644 +--- a/mlx/backend/cuda/cuda_utils.h ++++ b/mlx/backend/cuda/cuda_utils.h +@@ -30,7 +30,11 @@ class CudaHandle { + if (cudaPeekAtLastError() != cudaSuccess) { + return; + } +- reset(); ++ // Inline reset() without throwing -- handle shutdown cases. ++ if (handle_ != nullptr) { ++ Destroy(handle_); ++ handle_ = nullptr; ++ } + } + + CudaHandle(const CudaHandle&) = delete; +diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp +index 472b3d99f..e7d8f2620 100644 +--- a/mlx/backend/cuda/device.cpp ++++ b/mlx/backend/cuda/device.cpp +@@ -617,8 +617,10 @@ std::unordered_map& get_command_encoders() { + } + + std::unordered_map& get_global_command_encoders() { +- static std::unordered_map encoders; +- return encoders; ++ // encoders are leaked intentionally as they would synchronize on process ++ // shutdown ++ static auto* encoders = new std::unordered_map(); ++ return *encoders; + } + + } // namespace mlx::core::cu From 84d9eae4cb4d0b48d4f1a438a4de62729f7fc82a Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:21:31 -0700 Subject: [PATCH 33/39] pool Streams, fix Device inheritance (#472) * pool Streams, fix Device inheritance - see https://github.com/ml-explore/mlx/issues/2118 - Streams are finite and seem to leak OS/Metal resources - implements a pool to reuse Streams -- withNewDefaultStream is fully usable - much more care is taken with GPU devices - using index != 0 now works in all cases - and inherits properly when creating new streams Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com> --- .../CompileLockRepro/CompileLockRepro.swift | 37 +- Source/MLX/Device.swift | 242 ++++++----- .../Articles/using-streams.md | 125 +++++- Source/MLX/Stream.swift | 393 +++++++++++++++--- Tests/MLXTests/CompileLockOrderingTests.swift | 1 - Tests/MLXTests/FFTTests.swift | 1 - Tests/MLXTests/FP8Tests.swift | 1 - Tests/MLXTests/GraphUtilsTests.swift | 1 - Tests/MLXTests/IntegrationTests.swift | 1 - Tests/MLXTests/LinalgTests.swift | 1 - Tests/MLXTests/LossTests.swift | 1 - Tests/MLXTests/MLXArray+IndexingTests.swift | 1 - Tests/MLXTests/MLXArray+InitTests.swift | 1 - Tests/MLXTests/MLXArray+OpsTests.swift | 1 - Tests/MLXTests/MLXArrayTests.swift | 1 - Tests/MLXTests/MLXRandomTests.swift | 1 - Tests/MLXTests/ModuleTests.swift | 1 - Tests/MLXTests/NestedTests.swift | 1 - Tests/MLXTests/OpsTests.swift | 1 - Tests/MLXTests/OptimizerTests.swift | 1 - Tests/MLXTests/SaveTests.swift | 1 - Tests/MLXTests/StreamTests.swift | 376 +++++++++++++++-- Tests/MLXTests/TransformTests.swift | 1 - Tests/MLXTests/Utils.swift | 4 - 24 files changed, 947 insertions(+), 248 deletions(-) diff --git a/Source/CompileLockRepro/CompileLockRepro.swift b/Source/CompileLockRepro/CompileLockRepro.swift index 2f4577b44..aa8032db4 100644 --- a/Source/CompileLockRepro/CompileLockRepro.swift +++ b/Source/CompileLockRepro/CompileLockRepro.swift @@ -112,8 +112,6 @@ struct CompileLockRepro { default: die(2, "unknown --device \(CommandLine.arguments[index + 1])") } } - Device.setDefault(device: device) - print( "compile-lock repro: device=\(device), \(rounds) rounds, \(Int(deadline))s deadline per round" ) @@ -125,7 +123,8 @@ struct CompileLockRepro { var tracedAtLeastOnce = false for round in 1 ... rounds { - tracedAtLeastOnce = runRound(round, keepAlive: &keepAlive) || tracedAtLeastOnce + tracedAtLeastOnce = + runRound(round, device: device, keepAlive: &keepAlive) || tracedAtLeastOnce } if !tracedAtLeastOnce { @@ -141,7 +140,7 @@ struct CompileLockRepro { } /// Runs one round. Returns whether the traced body actually executed. - static func runRound(_ round: Int, keepAlive: inout [Any]) -> Bool { + static func runRound(_ round: Int, device: Device, keepAlive: inout [Any]) -> Bool { let gate = DispatchSemaphore(value: 0) // H (inside the trace) -> W let hDone = DispatchSemaphore(value: 0) let wDone = DispatchSemaphore(value: 0) @@ -165,23 +164,27 @@ struct CompileLockRepro { keepAlive.append(outer) let h = Thread { - let a = MLXArray([1.0, 2.0, 3.0] as [Float]) - let r = outer(a) - eval(r) - hDone.signal() + Device.withDefaultDevice(device) { + let a = MLXArray([1.0, 2.0, 3.0] as [Float]) + let r = outer(a) + eval(r) + hDone.signal() + } } h.name = "H-nested-compiled-call" let w = Thread { - gate.wait() - let b = MLXArray([4.0, 5.0, 6.0] as [Float]) - var iterations = 0 - repeat { - let r = shared(b) // W takes L_cf(shared), then wants evalLock - eval(r) - iterations += 1 - } while !stop.isSet && iterations < 1000 - wDone.signal() + Device.withDefaultDevice(device) { + gate.wait() + let b = MLXArray([4.0, 5.0, 6.0] as [Float]) + var iterations = 0 + repeat { + let r = shared(b) // W takes L_cf(shared), then wants evalLock + eval(r) + iterations += 1 + } while !stop.isSet && iterations < 1000 + wDone.signal() + } } w.name = "W-plain-compiled-call" diff --git a/Source/MLX/Device.swift b/Source/MLX/Device.swift index da5fc697e..c6a6ddc19 100644 --- a/Source/MLX/Device.swift +++ b/Source/MLX/Device.swift @@ -9,60 +9,77 @@ import Foundation public enum DeviceType: String, Hashable, Sendable { case cpu case gpu + + var cDeviceType: mlx_device_type { + switch self { + case .cpu: return MLX_CPU + case .gpu: return MLX_GPU + } + } + + init(_ cDeviceType: mlx_device_type) { + switch cDeviceType { + case MLX_CPU: self = .cpu + case MLX_GPU: self = .gpu + default: fatalError("Unknown deviceType: \(cDeviceType)") + } + } } /// Representation of a Device in MLX. /// -/// Typically this is used via the `stream: ` parameter on a method with a ``StreamOrDevice``: +/// This is typically used with ``withDefaultDevice(_:_:)-17vjl`` or ``Stream/withNewDefaultStream(device:_:)-5bwc3``. /// /// ```swift -/// let a: MLXArray ... -/// let result = sqrt(a, stream: .gpu) +/// Device.withDefaultDevice(.cpu) { +/// // default device is cpu inside this scope/task +/// } +/// +/// Stream.withNewDefaultStream(device: .cpu) { +/// // default device is cpu inside this scope/task AND there +/// // is a new Stream scoped to the work +/// } /// ``` /// -/// Read more at . +/// Implementation: ``Stream`` hold the current device as a Task local variable. /// /// ### See Also /// - /// - ``StreamOrDevice`` -public final class Device: @unchecked Sendable, Equatable { +/// - ``Stream`` +public final class Device: @unchecked Sendable, Hashable { let ctx: mlx_device - let defaultStream: Stream init(_ ctx: mlx_device) { self.ctx = ctx - - var deviceType = MLX_GPU - mlx_device_get_type(&deviceType, ctx) - self.defaultStream = - switch deviceType { - case MLX_CPU: .cpu - case MLX_GPU: .gpu - default: .gpu - } } + /// Initialize a new `Device` with a type and optional index. + /// + /// If `deviceType` is ``DeviceType/cpu`` then index _must_ be 0. public init(_ deviceType: DeviceType, index: Int32 = 0) { - var cDeviceType: mlx_device_type - switch deviceType { - case DeviceType.cpu: - cDeviceType = MLX_CPU - case DeviceType.gpu: - cDeviceType = MLX_GPU + if deviceType == .cpu { + precondition(index == 0) } + let cDeviceType = deviceType.cDeviceType self.ctx = mlx_device_new_type(cDeviceType, index) - self.defaultStream = - switch deviceType { - case .cpu: .cpu - case .gpu: .gpu - } } @available(*, deprecated, message: "please use defaultDevice()") public convenience init() { + self.init(copying: Stream.defaultDevice) + } + + convenience init(stream: mlx_stream) { var ctx = mlx_device_new() - mlx_get_default_device(&ctx) + mlx_stream_get_device(&ctx, stream) + self.init(ctx) + } + + convenience init(copying device: Device) { + var ctx = mlx_device_new() + mlx_device_set(&ctx, device.ctx) self.init(ctx) } @@ -70,114 +87,131 @@ public final class Device: @unchecked Sendable, Equatable { mlx_device_free(ctx) } - /// static CPU device + /// Current CPU device. /// - /// See ``withDefaultDevice(_:_:)-17vjl`` - static public let cpu: Device = Device(.cpu) + /// Note: previously this returned a `static` CPU device. + /// + /// ### See Also + /// - ``gpu`` + /// - ``Stream/cpu`` + /// - ``withDefaultDevice(_:_:)-17vjl`` + static public var cpu: Device { + Stream.cpu.device + } - /// static GPU device + /// Current GPU device. + /// + /// Note: previously this returned a `static` GPU device -- it would + /// always be index 0, even if there were multiple GPUs in the system. /// - /// See ``withDefaultDevice(_:_:)-17vjl`` - static public let gpu: Device = Device(.gpu) + /// ### See Also + /// - ``cpu`` + /// - ``Stream/gpu`` + /// - ``withDefaultDevice(_:_:)-17vjl`` + static public var gpu: Device { + Stream.gpu.device + } - public var deviceType: DeviceType? { + /// The ``DeviceType`` for the device. + /// + /// Note: previously this returned an Optional ``DeviceType``. Now it is + /// not optional and it will be a `fatalError` if the DeviceType is unknown. + public var deviceType: DeviceType { var cDeviceType = MLX_CPU mlx_device_get_type(&cDeviceType, ctx) - return switch cDeviceType { - case MLX_CPU: DeviceType.cpu - case MLX_GPU: DeviceType.gpu - default: nil - } - } - - // support for global default device - static let _lock = NSLock() - #if swift(>=5.10) - nonisolated(unsafe) static var _defaultDevice: Device? - #else - static var _defaultDevice: Device? - #endif - - @TaskLocal static var _tlDefaultDevice = _resolveGlobalDefaultDevice() - - private static func _resolveGlobalDefaultDevice() -> Device { - _lock.withLock { - if let device = _defaultDevice { - return device - } - // Ask the underlying MLX C++ core for its default device rather - // than hard-coding `.gpu`. On Apple platforms with Metal this - // still resolves to GPU; on a CPU-only host (Linux without - // CUDA / no Metal) it correctly resolves to CPU. Hard-coding GPU - // here meant `defaultDevice()` / `StreamOrDevice.default` - // returned an unavailable device on those hosts. - var ctx = mlx_device_new() - mlx_get_default_device(&ctx) - return Device(ctx) - } + return DeviceType(cDeviceType) } /// Return the current default device. - /// - /// This is used by ``StreamOrDevice/default`` -- the default stream parameter - /// to most functions. static public func defaultDevice() -> Device { - _tlDefaultDevice + Stream.defaultDevice } - /// Use a device scoped to a task. + /// Use a device scoped to a Task. + /// + /// This can be used to set the default device to e.g. the CPU: + /// + /// ```swift + /// Device.withDefaultDevice(.cpu) { + /// // default device is cpu inside this scope/task + /// } + /// ``` + /// + /// If a GPU device is given and it does not match the current GPU it will + /// override the default device and provide a new stream: + /// + /// ```swift + /// Device.withDefaultDevice(.init(.gpu, index: 3)) { + /// // if the enclosing GPU was index 0, this will have + /// // both a new default device and a new stream + /// } + /// ``` + /// + /// See also ``Stream/withNewDefaultStream(device:_:)-5bwc3``. static public func withDefaultDevice( _ device: Device, _ body: () throws -> R ) rethrows -> R { - try $_tlDefaultDevice.withValue(device, operation: body) + try Stream.withNewDefaultDevice(device: device, body) } - /// Use a device scoped to a task. + /// Use a device scoped to a Task. + /// + /// This can be used to set the default device to e.g. the CPU: + /// + /// ```swift + /// Device.withDefaultDevice(.cpu) { + /// // default device is cpu inside this scope/task + /// } + /// ``` + /// + /// If a GPU device is given and it does not match the current GPU it will + /// override the default device and provide a new stream: + /// + /// ```swift + /// Device.withDefaultDevice(.init(.gpu, index: 3)) { + /// // if the enclosing GPU was index 0, this will have + /// // both a new default device and a new stream + /// } + /// ``` + /// + /// See also ``Stream/withNewDefaultStream(device:_:)-5bwc3``. static public func withDefaultDevice( _ device: Device, _ body: () async throws -> R ) async rethrows -> R { - try await $_tlDefaultDevice.withValue(device, operation: body) - } - - /// Return the current default stream. - static func defaultStream() -> Stream { - _tlDefaultDevice.defaultStream + try await Stream.withNewDefaultDevice(device: device, body) } - /// Set the default device globally. Prefer the scoped version, ``withDefaultDevice(_:_:)-17vjl``. + /// Set the default device globally. Use the scoped version, ``withDefaultDevice(_:_:)-17vjl`` + /// -- this is only usable before any mlx resources are created. /// - /// For example: + /// Beware: using the static cpu and gpu values will render this unusable. If required, call like this: /// /// ```swift - /// Device.setDefault(device: Device(.cpu, index: 1)) + /// Device.setDefault(device: .init(.cpu)) /// ``` - /// - /// By default this is ``gpu``. - /// - /// ### See Also - /// - ``withDefaultDevice(_:_:)-17vjl`` - /// - ``StreamOrDevice/default`` - @available(*, deprecated, message: "please use withDefaultDevice()") + @available( + *, deprecated, + message: "use withDefaultDevice() -- this only works before any mlx resources are created" + ) static public func setDefault(device: Device?) { - _lock.withLock { - if let device { - // sets the mlx core default device -- only used - // by the deprecated init(). this isn't thread - // safe or really usable across tasks/threads - // but is kept for backward compatibility - mlx_set_default_device(device.ctx) - } - _defaultDevice = device + if let ctx = device?.ctx { + mlx_set_default_device(ctx) } } - /// Compare two ``Device`` for equality -- this does not compare the index, just the device type. + /// Compare two ``Device`` for equality public static func == (lhs: Device, rhs: Device) -> Bool { - var lhs_type = MLX_CPU - var rhs_type = MLX_CPU - mlx_device_get_type(&lhs_type, lhs.ctx) - mlx_device_get_type(&rhs_type, rhs.ctx) - return lhs_type == rhs_type + mlx_device_equal(lhs.ctx, rhs.ctx) + } + + public func hash(into hasher: inout Hasher) { + var index: Int32 = 0 + mlx_device_get_index(&index, ctx) + hasher.combine(index) + + var type: mlx_device_type = MLX_CPU + mlx_device_get_type(&type, ctx) + hasher.combine(type.rawValue) } } diff --git a/Source/MLX/Documentation.docc/Articles/using-streams.md b/Source/MLX/Documentation.docc/Articles/using-streams.md index 5472f92f3..d90c9cdaa 100644 --- a/Source/MLX/Documentation.docc/Articles/using-streams.md +++ b/Source/MLX/Documentation.docc/Articles/using-streams.md @@ -4,38 +4,121 @@ Controlling where your computations are evaluated. ## Specifying the Stream -All operations (including random number generation) take an optional -argument `stream`. The `stream` specifies which -`Stream` the operation should run on. If the stream is unspecified then -the operation is run on the default stream of the default device: -``Stream/defaultStream(_:)``. The `stream` can also -be a ``Device`` (e.g. `stream: .cpu`) in which case the operation is -run on the default stream of the provided device. +All operations, including random number generation, take an optional `stream:` +argument. It defaults to ``Stream/defaultStream`` -- the default stream of the +default device. For example: ```swift -// produced on cpu -let a = MLXRandom.uniform([100, 100], stream: .cpu) +// produced on the CPU +let a = uniform(0 ..< 1, [100, 100], stream: .cpu) -// produced on gpu -let b = MLXRandom.uniform([100, 100], stream: .gpu) +// produced on the GPU +let b = uniform(0 ..< 1, [100, 100], stream: .gpu) ``` -The parameter is of type ``StreamOrDevice`` and can be initialized with -a stream, device, or some defaults like ``StreamOrDevice/gpu``. +The argument is a ``StreamOrDevice``. ``StreamOrDevice/cpu`` and +``StreamOrDevice/gpu`` select whichever stream is currently in effect for that +device, which the scopes below control. -If you want to execute a block of code on a specific device you can use -``using(device:fn:)``. +## Scoping the Default Device -For example: +To avoid passing `stream:` to every call, change the default device for a +scope. ``Device/withDefaultDevice(_:_:)-17vjl`` does this without creating any +new streams: + +```swift +Device.withDefaultDevice(.cpu) { + // no stream: argument needed -- these run on the CPU + let a = uniform(0 ..< 1, [100, 100]) + let b = uniform(0 ..< 1, [100, 100]) +} +``` + +Inside the scope, ``Device/defaultDevice()`` returns the CPU device and +``StreamOrDevice/default`` resolves to the CPU stream. On exit, the previous +default is restored. + +## Scoping a New Stream + +``Stream/withNewDefaultStream(device:_:)-5bwc3`` sets the default device and +creates *new* streams for the scope: + +```swift +Stream.withNewDefaultStream(device: .cpu) { + // a new stream, private to this scope + let a = uniform(0 ..< 1, [100, 100]) +} +``` + +Use this when work must be independent of the enclosing scope's streams -- for +example so a ``Stream/synchronize()`` waits only on your own work. Use +``Device/withDefaultDevice(_:_:)-17vjl`` when you only want to redirect work to +a different device. + +`withNewDefaultStream` replaces *both* the CPU and GPU streams for the scope, +not only the one named by `device:`. Code inside that explicitly asks for the +other device also gets a new stream: + +```swift +Stream.withNewDefaultStream(device: .cpu) { + // a new CPU stream -- and `stream: .gpu` here is also a new + // stream, not the GPU stream from the enclosing scope + let a = uniform(0 ..< 1, [100, 100], stream: .gpu) +} +``` + +This does not affect correctness: MLX inserts dependencies between streams +automatically, as described in . + +Both calls nest arbitrarily, and both have `async` variants. + +## Propagating Scopes Across Tasks + +These scopes are task-local. They are inherited by `async let` and by child +tasks of a `TaskGroup`: ```swift -using(device: .gpu) { - // this code will run on gpu - let a = MLXRandom.uniform([100, 100]) - let b = MLXRandom.uniform([100, 100]) +await Device.withDefaultDevice(.cpu) { + await withTaskGroup(of: Void.self) { group in + group.addTask { + // still the CPU + let a = uniform(0 ..< 1, [100, 100]) + } + } } ``` -Read more in . +They are *not* inherited by `Task.detached` or by a `Thread`, which see the +process-wide default instead. Work started that way must establish its own +scope: + +```swift +Task.detached { + Device.withDefaultDevice(.cpu) { + // the scope has to be re-established here + } +} +``` + +## Holding a Stream Directly + +When a stream must outlive a single scope, create a ``Stream`` and pass it +explicitly: + +```swift +let stream = Stream(Device(.cpu)) +let a = uniform(0 ..< 1, [100, 100], stream: .stream(stream)) +``` + +Streams are pooled: the underlying resources are recycled when the `Stream` is +released, so creating them in a loop does not exhaust system resources. Prefer +the scoped calls above where they fit. + +## See Also + +- +- ``Stream`` +- ``Device`` +- ``StreamOrDevice`` diff --git a/Source/MLX/Stream.swift b/Source/MLX/Stream.swift index 917d544ad..8b31b4105 100644 --- a/Source/MLX/Stream.swift +++ b/Source/MLX/Stream.swift @@ -32,25 +32,27 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable { /// The default stream on the default device. /// - /// This will be ``Device/gpu`` unless ``Device/setDefault(device:)`` - /// sets it otherwise. + /// See ``Stream/defaultStream``. public static var `default`: StreamOrDevice { - StreamOrDevice(Stream.defaultStream ?? Device.defaultStream()) + StreamOrDevice(Stream.defaultStream) } + /// Evaluate on the given ``Device``. + /// + /// Prefer ``cpu`` and ``gpu`` instead as they will use the current Device for the + /// Task. + @available(*, deprecated, message: "prefer stream(Stream)") public static func device(_ device: Device) -> StreamOrDevice { StreamOrDevice(Stream.defaultStream(device)) } - /// The ``Stream/defaultStream(_:)`` on the ``Device/cpu`` - public static let cpu = device(.cpu) + /// Evaluate on the current CPU stream, ``Stream/cpu``. + public static var cpu: StreamOrDevice { .stream(.cpu) } - /// The ``Stream/defaultStream(_:)`` on the ``Device/gpu`` - /// - /// ### See Also - /// - ``GPU`` - public static let gpu = device(.gpu) + /// Evaluate on the current GPU stream, ``Stream/gpu``. + public static var gpu: StreamOrDevice { .stream(.gpu) } + /// Evaluate on the given stream. public static func stream(_ stream: Stream) -> StreamOrDevice { StreamOrDevice(stream) } @@ -65,9 +67,174 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable { } } +/// Pool for recycling ``Stream`` (`mlx_stream`). +/// +/// See also https://github.com/ml-explore/mlx/issues/2118 -- there are +/// metal/OS resources associated with a GPU stream that may never be collected. +/// Reuse the streams. +/// +/// `Stream.deinit` will return the underlying `mlx_stream` to the pool. +private final class StreamPool: @unchecked Sendable { + private var idle: [Device: [mlx_stream]] = [:] + private let lock = NSLock() + + init() { + } + + func newStream(_ device: Device) -> mlx_stream { + // do not hold the evalLock at the same time as the pool lock + let (valid, stream) = lock.withLock { + if let value = idle[device, default: []].popLast() { + return (true, value) + } + return (false, mlx_stream()) + } + if valid { + return stream + } + + return withEvalLock { + let result = mlx_stream_new_thread_unsafe(device.ctx) + if result.ctx == nil { + fatalError("Unable to create stream: \(device)") + } + return result + } + } + + func releaseStream(_ stream: mlx_stream) { + let device = Device(stream: stream) + + lock.withLock { + idle[device, default: []].append(stream) + } + } +} + +/// global stream pool +private let streamPool = StreamPool() + +private protocol StreamPromise: Sendable { + var current: Stream { get } + var device: Device { get } +} + +private final class LazyStreamPromise: StreamPromise, @unchecked Sendable { + private var stream: Stream? + private let lock = NSLock() + + let device: Device + + init(_ device: Device) { + self.stream = nil + self.device = device + } + + var current: Stream { + if let stream = lock.withLock({ self.stream }) { + return stream + } + + // double checked lock to avoid holding a lock outside evalLock + let stream = Stream(streamPool.newStream(device)) + + return lock.withLock { + if let current = self.stream { + return current + } else { + self.stream = stream + return stream + } + } + } +} + +private final class RealizedStreamPromise: StreamPromise, Sendable { + let current: Stream + let device: Device + + init(_ device: Device) { + self.device = device + self.current = Stream(streamPool.newStream(device)) + } + + init(_ promise: StreamPromise) { + self.device = promise.device + self.current = promise.current + } +} + +/// TaskLocal promise for a pair of Streams -- one for CPU and one for GPU. +/// +/// Implementation note: the default device will be a realized Stream and the +/// none default device will be a lazy promise. +private final class StreamPairPromise: Sendable { + let cpu: StreamPromise + let gpu: StreamPromise + let defaultDeviceType: DeviceType + + var defaultStream: StreamPromise { + switch defaultDeviceType { + case .cpu: cpu + case .gpu: gpu + } + } + + /// Initialize the global promise. + init() { + var default_ctx = mlx_device_new() + mlx_get_default_device(&default_ctx) + + let defaultDevice = Device(default_ctx) + self.defaultDeviceType = defaultDevice.deviceType + if defaultDevice.deviceType == .cpu { + self.cpu = RealizedStreamPromise(defaultDevice) + self.gpu = LazyStreamPromise(Device(.gpu)) + } else { + self.cpu = LazyStreamPromise(Device(.cpu)) + self.gpu = RealizedStreamPromise(defaultDevice) + } + } + + /// Initialize with new streams. + /// + /// See ``Stream/withNewDefaultStream(device:_:)``. + init(cpu: Device, gpu: Device, defaultDeviceType: DeviceType) { + self.defaultDeviceType = defaultDeviceType + if defaultDeviceType == .cpu { + self.cpu = RealizedStreamPromise(cpu) + self.gpu = LazyStreamPromise(gpu) + } else { + self.cpu = LazyStreamPromise(cpu) + self.gpu = RealizedStreamPromise(gpu) + } + } + + /// Initialize with a new default device and a parent. + /// + /// See ``Device/withDefaultDevice(_:_:)``. + init(parent: StreamPairPromise, device: Device) { + self.defaultDeviceType = device.deviceType + + switch self.defaultDeviceType { + case .cpu: + self.cpu = RealizedStreamPromise(parent.cpu) + self.gpu = parent.gpu + + case .gpu: + self.cpu = parent.cpu + if parent.gpu.device == device { + self.gpu = RealizedStreamPromise(parent.gpu) + } else { + self.gpu = RealizedStreamPromise(device) + } + } + } +} + /// A stream of evaluation attached to a particular device. /// -/// Typically this is used via the `stream: ` parameter on a method with a ``StreamOrDevice``: +/// Typically this is used via the `stream:` parameter on a method with a ``StreamOrDevice``: /// /// ```swift /// let a: MLXArray ... @@ -76,81 +243,185 @@ public struct StreamOrDevice: Sendable, CustomStringConvertible, Equatable { /// /// Read more at . /// +/// ## Implementation Notes +/// +/// There is a global pair of CPU/GPU streams. On devices with a GPU (and associated GPU software) +/// the default stream will be a GPU stream. +/// +/// The default streams can be overridden and scoped to a `Task` using: +/// +/// ```swift +/// Stream.withNewDefaultStream { +/// // the default stream is private to this Task +/// } +/// ``` +/// +/// Callers can override the default device by passing a device to that method or using +/// ``Device/withDefaultDevice(_:_:)-17vjl`` to override the default device without +/// getting a new (private) default stream. +/// +/// Also note that the current ``Device/defaultDevice()`` is owned by ``Stream/defaultStream``. +/// /// ### See Also /// - /// - ``StreamOrDevice`` +/// - ``Device`` public final class Stream: @unchecked Sendable, Equatable { + /// reference to the backing C++ stream object let ctx: mlx_stream - private static func newStreamThreadUnsafe(_ deviceType: DeviceType) -> mlx_stream { - var cDeviceType: mlx_device_type - switch deviceType { - case DeviceType.cpu: - cDeviceType = MLX_CPU - case DeviceType.gpu: - cDeviceType = MLX_GPU - } - return evalLock.withLock { - let ctx = mlx_device_new_type(cDeviceType, 0) - defer { mlx_device_free(ctx) } - return mlx_stream_new_thread_unsafe(ctx) - } + /// the global cpu and gpu streams + private static let globalStreams = StreamPairPromise() + + /// the task local override streams, see ``withNewDefaultStream(device:_:)`` + @TaskLocal private static var localStreams: StreamPairPromise? + + /// the current GPU stream + public static var gpu: Stream { + localStreams?.gpu.current ?? globalStreams.gpu.current } - public static let gpu = Stream(newStreamThreadUnsafe(.gpu)) - public static let cpu = Stream(newStreamThreadUnsafe(.cpu)) + /// the current CPU stream + public static var cpu: Stream { + localStreams?.cpu.current ?? globalStreams.cpu.current + } - @TaskLocal static var defaultStream: Stream? + /// the current default stream, see ``StreamOrDevice`` -- if no stream is + /// specified, this will be the stream that all operations run on. + public static var defaultStream: Stream { + localStreams?.defaultStream.current ?? globalStreams.defaultStream.current + } - /// Set the ``StreamOrDevice/default`` scoped to a Task. - public static func withNewDefaultStream(device: Device? = nil, _ body: () throws -> R) - rethrows -> R - { - let device = device ?? Device.defaultDevice() - return try $defaultStream.withValue(Stream(device), operation: body) + static var defaultDevice: Device { + let streams = localStreams ?? globalStreams + return streams.defaultStream.device } - /// Set the ``StreamOrDevice/default`` scoped to a Task. + var device: Device { Device(stream: ctx) } + var index: Int { + var index: Int32 = 0 + mlx_stream_get_index(&index, ctx) + return Int(index) + } + + /// Obtain a new CPU and GPU stream scoped to the block and inherited through `Tasks`. + /// + /// For example you can get a new stream on the current device: + /// + /// ```swift + /// Stream.withNewDefaultStream { + /// // A new Stream on the default Device scoped to the block and Task + /// } + /// ``` + /// + /// Or a new stream on a different device (it will become the default device, scoped to the block + /// and Task): + /// + /// ```swift + /// Stream.withNewDefaultStream(device: .cpu) { + /// // A new Stream on the cpu Device scoped to the block and Task. + /// // The cpu is now the default device. + /// } + /// ``` + /// + /// These calls can be nested arbitrarily. + /// + /// See also ``Device/withDefaultDevice(_:_:)-17vjl`` which can change the default device + /// without creating new scoped Streams. + public static func withNewDefaultStream( + device: Device? = nil, _ body: () throws -> R + ) rethrows -> R { + let streams = localStreams ?? globalStreams + let device = device ?? streams.defaultStream.device + let cpu = device.deviceType == .cpu ? device : streams.cpu.device + let gpu = device.deviceType == .gpu ? device : streams.gpu.device + return try $localStreams.withValue( + StreamPairPromise(cpu: cpu, gpu: gpu, defaultDeviceType: device.deviceType), + operation: body) + } + + /// Obtain a new CPU and GPU stream scoped to the block and inherited through `Tasks`. + /// + /// For example you can get a new stream on the current device: + /// + /// ```swift + /// Stream.withNewDefaultStream { + /// // A new Stream on the default Device scoped to the block and Task + /// } + /// ``` + /// + /// Or a new stream on a different device (it will become the default device, scoped to the block + /// and Task): + /// + /// ```swift + /// Stream.withNewDefaultStream(device: .cpu) { + /// // A new Stream on the cpu Device scoped to the block and Task. + /// // The cpu is now the default device. + /// } + /// ``` + /// + /// These calls can be nested arbitrarily. + /// + /// See also ``Device/withDefaultDevice(_:_:)-17vjl`` which can change the default device + /// without creating new scoped Streams. public static func withNewDefaultStream( device: Device? = nil, _ body: () async throws -> R ) async rethrows -> R { - let device = device ?? Device.defaultDevice() - return try await $defaultStream.withValue(Stream(device), operation: body) + let streams = localStreams ?? globalStreams + let device = device ?? streams.defaultStream.device + let cpu = device.deviceType == .cpu ? device : streams.cpu.device + let gpu = device.deviceType == .gpu ? device : streams.gpu.device + return try await $localStreams.withValue( + StreamPairPromise(cpu: cpu, gpu: gpu, defaultDeviceType: device.deviceType), + operation: body) + } + + /// Implementation for ``Device/withDefaultDevice(_:_:)`` + static func withNewDefaultDevice( + device: Device, _ body: () throws -> R + ) rethrows -> R { + let pair = StreamPairPromise( + parent: localStreams ?? globalStreams, device: device) + return try $localStreams.withValue(pair, operation: body) + } + + /// Implementation for ``Device/withDefaultDevice(_:_:)`` + static func withNewDefaultDevice( + device: Device, _ body: () async throws -> R + ) async rethrows -> R { + let pair = StreamPairPromise( + parent: localStreams ?? globalStreams, device: device) + return try await $localStreams.withValue(pair, operation: body) } init(_ ctx: mlx_stream) { self.ctx = ctx } - /// Default stream on the default device. + /// New stream on the default device. + @available(*, deprecated, message: "use withNewDefaultStream instead") public init() { - let device = Device.defaultDevice() - self.ctx = evalLock.withLock { - mlx_stream_new_thread_unsafe(device.ctx) - } + self.ctx = streamPool.newStream(Device.defaultDevice()) } - @available(*, deprecated, message: "use init(Device) -- index not supported") + /// New stream on the given device. + /// + /// Note: `index` is unused. + @available(*, deprecated, message: "use withNewDefaultStream instead") public init(index: Int32, _ device: Device) { - self.ctx = evalLock.withLock { - mlx_stream_new_thread_unsafe(device.ctx) - } + self.ctx = streamPool.newStream(device) } /// New stream on the given device. /// - /// See also ``withNewDefaultStream(device:_:)-5bwc3`` + /// Prefer ``withNewDefaultStream(device:_:)-5bwc3`` public init(_ device: Device) { - self.ctx = evalLock.withLock { - mlx_stream_new_thread_unsafe(device.ctx) - } + self.ctx = streamPool.newStream(device) } deinit { - _ = withEvalLock { - mlx_stream_free(ctx) - } + streamPool.releaseStream(ctx) } /// Synchronize with the given stream @@ -160,11 +431,27 @@ public final class Stream: @unchecked Sendable, Equatable { } } + @_disfavoredOverload + @available( + *, deprecated, message: "use defaultStream(DeviceType) or withNewDefaultStream(Device)" + ) static public func defaultStream(_ device: Device) -> Stream { switch device.deviceType { + case .cpu: return .cpu + case .gpu: + let streams = localStreams ?? globalStreams + if streams.gpu.device == device { + return streams.gpu.current + } else { + return Stream(device) + } + } + } + + static public func defaultStream(_ deviceType: DeviceType) -> Stream { + switch deviceType { case .cpu: .cpu case .gpu: .gpu - default: fatalError("Unexpected device type: \(device)") } } diff --git a/Tests/MLXTests/CompileLockOrderingTests.swift b/Tests/MLXTests/CompileLockOrderingTests.swift index 745e14b2d..af7c67057 100644 --- a/Tests/MLXTests/CompileLockOrderingTests.swift +++ b/Tests/MLXTests/CompileLockOrderingTests.swift @@ -23,7 +23,6 @@ import XCTest class CompileLockOrderingTests: XCTestCase { override class func setUp() { - setDefaultDevice() } /// The instrumentation must be able to report "not held", otherwise a diff --git a/Tests/MLXTests/FFTTests.swift b/Tests/MLXTests/FFTTests.swift index e602c5872..258b59c59 100644 --- a/Tests/MLXTests/FFTTests.swift +++ b/Tests/MLXTests/FFTTests.swift @@ -7,7 +7,6 @@ import XCTest class FFTTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testFFTRoundTrip() { diff --git a/Tests/MLXTests/FP8Tests.swift b/Tests/MLXTests/FP8Tests.swift index 9ec00c401..4fd0b1f9e 100644 --- a/Tests/MLXTests/FP8Tests.swift +++ b/Tests/MLXTests/FP8Tests.swift @@ -7,7 +7,6 @@ import XCTest class FP8Tests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testFromFP8DecodesE4M3Bytes() { diff --git a/Tests/MLXTests/GraphUtilsTests.swift b/Tests/MLXTests/GraphUtilsTests.swift index 7e9c5ca27..81c7c4975 100644 --- a/Tests/MLXTests/GraphUtilsTests.swift +++ b/Tests/MLXTests/GraphUtilsTests.swift @@ -7,7 +7,6 @@ import XCTest class GraphUtilsTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testExportToDot() { diff --git a/Tests/MLXTests/IntegrationTests.swift b/Tests/MLXTests/IntegrationTests.swift index 242f0fa6d..a8f4a5399 100644 --- a/Tests/MLXTests/IntegrationTests.swift +++ b/Tests/MLXTests/IntegrationTests.swift @@ -16,7 +16,6 @@ import XCTest class MLXIntegrationTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testRandomSeed() { diff --git a/Tests/MLXTests/LinalgTests.swift b/Tests/MLXTests/LinalgTests.swift index 2939988db..d0d0856fd 100644 --- a/Tests/MLXTests/LinalgTests.swift +++ b/Tests/MLXTests/LinalgTests.swift @@ -7,7 +7,6 @@ import XCTest class LinalgTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testNormNoAxes() { diff --git a/Tests/MLXTests/LossTests.swift b/Tests/MLXTests/LossTests.swift index 4d6340979..b0efaeea0 100644 --- a/Tests/MLXTests/LossTests.swift +++ b/Tests/MLXTests/LossTests.swift @@ -7,7 +7,6 @@ import XCTest class LossTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testCrossEntropy() { diff --git a/Tests/MLXTests/MLXArray+IndexingTests.swift b/Tests/MLXTests/MLXArray+IndexingTests.swift index 68673ff51..699183f49 100644 --- a/Tests/MLXTests/MLXArray+IndexingTests.swift +++ b/Tests/MLXTests/MLXArray+IndexingTests.swift @@ -23,7 +23,6 @@ extension MLXArrayIndexOperation: Equatable { class MLXArrayIndexingTests: XCTestCase { override class func setUp() { - setDefaultDevice() } // MARK: - Subscript (get) diff --git a/Tests/MLXTests/MLXArray+InitTests.swift b/Tests/MLXTests/MLXArray+InitTests.swift index 8a6cde37e..8c1017be4 100644 --- a/Tests/MLXTests/MLXArray+InitTests.swift +++ b/Tests/MLXTests/MLXArray+InitTests.swift @@ -13,7 +13,6 @@ import XCTest class MLXArrayInitTests: XCTestCase { override class func setUp() { - setDefaultDevice() } // MARK: - Dtype diff --git a/Tests/MLXTests/MLXArray+OpsTests.swift b/Tests/MLXTests/MLXArray+OpsTests.swift index 5797efc20..2c2cd3a58 100644 --- a/Tests/MLXTests/MLXArray+OpsTests.swift +++ b/Tests/MLXTests/MLXArray+OpsTests.swift @@ -9,7 +9,6 @@ import XCTest class MLXArrayOpsTests: XCTestCase { override class func setUp() { - setDefaultDevice() } // MARK: - Operators diff --git a/Tests/MLXTests/MLXArrayTests.swift b/Tests/MLXTests/MLXArrayTests.swift index 715a77590..1a74fdba2 100644 --- a/Tests/MLXTests/MLXArrayTests.swift +++ b/Tests/MLXTests/MLXArrayTests.swift @@ -8,7 +8,6 @@ import XCTest class MLXArrayTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testArrayProperties() { diff --git a/Tests/MLXTests/MLXRandomTests.swift b/Tests/MLXTests/MLXRandomTests.swift index e3906339d..43fa93c1e 100644 --- a/Tests/MLXTests/MLXRandomTests.swift +++ b/Tests/MLXTests/MLXRandomTests.swift @@ -7,7 +7,6 @@ import XCTest class MLXRandomTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testSplit() { diff --git a/Tests/MLXTests/ModuleTests.swift b/Tests/MLXTests/ModuleTests.swift index 288da7835..07a771b3d 100644 --- a/Tests/MLXTests/ModuleTests.swift +++ b/Tests/MLXTests/ModuleTests.swift @@ -9,7 +9,6 @@ import XCTest class ModuleTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func newTestModule() -> Module { diff --git a/Tests/MLXTests/NestedTests.swift b/Tests/MLXTests/NestedTests.swift index 3e9d83b55..05eb1f0c1 100644 --- a/Tests/MLXTests/NestedTests.swift +++ b/Tests/MLXTests/NestedTests.swift @@ -7,7 +7,6 @@ import XCTest class NestedTests: XCTestCase { override class func setUp() { - setDefaultDevice() } static let defaultValues = [10, 1, 2, 1, 2, 3, 10, 20, 30] diff --git a/Tests/MLXTests/OpsTests.swift b/Tests/MLXTests/OpsTests.swift index 27f58ce7c..b59970c23 100644 --- a/Tests/MLXTests/OpsTests.swift +++ b/Tests/MLXTests/OpsTests.swift @@ -8,7 +8,6 @@ import XCTest class OpsTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testAsStridedReshape() { diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index d773a5563..95e04af7d 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -10,7 +10,6 @@ import XCTest class OptimizerTests: XCTestCase { override class func setUp() { - setDefaultDevice() } class ShapeModule: Module { diff --git a/Tests/MLXTests/SaveTests.swift b/Tests/MLXTests/SaveTests.swift index 373ff8ca3..0ef8e1d55 100644 --- a/Tests/MLXTests/SaveTests.swift +++ b/Tests/MLXTests/SaveTests.swift @@ -16,7 +16,6 @@ final class SaveTests: XCTestCase { ) override func setUpWithError() throws { - setDefaultDevice() try FileManager.default.createDirectory( at: temporaryPath, withIntermediateDirectories: false diff --git a/Tests/MLXTests/StreamTests.swift b/Tests/MLXTests/StreamTests.swift index 32040e9df..e7f107a90 100644 --- a/Tests/MLXTests/StreamTests.swift +++ b/Tests/MLXTests/StreamTests.swift @@ -1,23 +1,62 @@ // Copyright © 2024 Apple Inc. import Foundation -import MLX import XCTest +@testable import MLX + class StreamTests: XCTestCase { + // MARK: - Device: Equatable / Hashable + func testEquatableDevice() { let s1 = Device.gpu let s2 = Device(.gpu, index: 3) let s3 = Device.cpu + let s4 = Device(.gpu) + + // equality does not ignore the index -- this changed: it used to + // compare the device type only + XCTAssertNotEqual(s1, s2) - // equality ignores index - XCTAssertEqual(s1, s2) + XCTAssertEqual(s1, s4) XCTAssertNotEqual(s1, s3) XCTAssertNotEqual(s2, s3) } + func testHashableDevice() { + XCTAssertEqual(Device.gpu.hashValue, Device(.gpu).hashValue) + XCTAssertEqual(Device.cpu.hashValue, Device(.cpu).hashValue) + + // usable as a dictionary key + var counts: [Device: Int] = [:] + counts[Device.gpu, default: 0] += 1 + counts[Device(.gpu), default: 0] += 1 + counts[Device(.gpu, index: 3), default: 0] += 1 + counts[Device.cpu, default: 0] += 1 + + XCTAssertEqual(counts.count, 3) + XCTAssertEqual(counts[Device.gpu], 2) + XCTAssertEqual(counts[Device(.gpu, index: 3)], 1) + XCTAssertEqual(counts[Device.cpu], 1) + } + + @available(*, deprecated) + func testDeprecatedDeviceInitCopiesDefault() { + // Device() produces the same value as defaultDevice + XCTAssertEqual(Device(), Device.defaultDevice()) + + Device.withDefaultDevice(.cpu) { + XCTAssertEqual(Device(), Device.cpu) + } + + let gpu = Device(.gpu, index: 7) + Device.withDefaultDevice(gpu) { + XCTAssertEqual(Device(), gpu) + } + } + func testDeviceType() { let s1 = Device.gpu let s2 = Device(.gpu, index: 3) @@ -28,37 +67,112 @@ class StreamTests: XCTestCase { XCTAssertEqual(s3.deviceType, .cpu) } + // MARK: - withDefaultDevice + func testUsingDevice() { let defaultDevice = Device.defaultDevice() Device.withDefaultDevice(.cpu) { - // these _should_ be the same - XCTAssertTrue(Device.defaultDevice().description.contains("cpu")) - XCTAssertTrue(StreamOrDevice.default.description.contains("cpu")) + // the default device and the default stream must agree + XCTAssertEqual(Device.defaultDevice(), Device.cpu) + XCTAssertEqual(StreamOrDevice.default.stream, MLX.Stream.cpu) } XCTAssertEqual(defaultDevice, Device.defaultDevice()) Device.withDefaultDevice(.gpu) { - XCTAssertTrue(Device.defaultDevice().description.contains("gpu")) - XCTAssertTrue(StreamOrDevice.default.description.contains("gpu")) + XCTAssertEqual(Device.defaultDevice(), Device.gpu) + XCTAssertEqual(StreamOrDevice.default.stream, MLX.Stream.gpu) + } + + let gpu = Device(.gpu, index: 7) + Device.withDefaultDevice(gpu) { + XCTAssertEqual(Device.defaultDevice(), gpu) + XCTAssertEqual(StreamOrDevice.default.stream, MLX.Stream.gpu) } - XCTAssertTrue(StreamOrDevice.default.description.contains("gpu")) + + // restored on scope exit + XCTAssertEqual(defaultDevice, Device.defaultDevice()) } - func testSetUnsetDefaultDevice() { - // Issue #237 -- setting an unsetting the default device in a loop - // exhausts many resources - for _ in 1 ..< 10000 { - let defaultDevice = MLX.Device.defaultDevice() - MLX.Device.setDefault(device: .cpu) - defer { - MLX.Device.setDefault(device: defaultDevice) + func testWithDefaultDeviceNests() { + Device.withDefaultDevice(.cpu) { + XCTAssertEqual(Device.defaultDevice(), Device.cpu) + + Device.withDefaultDevice(.gpu) { + XCTAssertEqual(Device.defaultDevice(), Device.gpu) + + Device.withDefaultDevice(.cpu) { + XCTAssertEqual(Device.defaultDevice(), Device.cpu) + } + + XCTAssertEqual(Device.defaultDevice(), Device.gpu) } - let x = MLXArray(1) - let _ = x * x + XCTAssertEqual(Device.defaultDevice(), Device.cpu) + } + } + + func testNestedStreamsAndDevices() { + let outerCpu = MLX.Stream.cpu + let outerGpu = MLX.Stream.gpu + + Device.withDefaultDevice(.cpu) { + XCTAssertEqual(MLX.Stream.defaultStream, outerCpu) + XCTAssertEqual(MLX.Stream.cpu, outerCpu) + XCTAssertEqual(MLX.Stream.gpu, outerGpu) + } + + let gpu7 = Device(.gpu, index: 7) + Stream.withNewDefaultStream(device: gpu7) { + // default is now gpu/7 + XCTAssertNotEqual(MLX.Stream.defaultStream, outerGpu) + XCTAssertEqual(MLX.Stream.defaultStream, MLX.Stream.gpu) + XCTAssertEqual(MLX.Stream.defaultStream.device, gpu7) + + // not the same stream + XCTAssertNotEqual(MLX.Stream.cpu, outerCpu) + + let gpu7Stream = MLX.Stream.defaultStream + + // this will override the cpu stream but not the gpu stream + Device.withDefaultDevice(.cpu) { + XCTAssertEqual(MLX.Stream.defaultStream.device, outerCpu.device) + + // this is the same stream, we are just changing the default device + XCTAssertEqual(MLX.Stream.gpu, gpu7Stream) + } + + // new scoped stream on GPU index 0 + Stream.withNewDefaultStream(device: .init(.gpu, index: 0)) { + XCTAssertEqual(MLX.Stream.defaultStream.device, outerGpu.device) + XCTAssertNotEqual(MLX.Stream.gpu, gpu7Stream) + XCTAssertNotEqual(MLX.Stream.gpu, outerGpu) + XCTAssertEqual(MLX.Stream.gpu.device, outerGpu.device) + XCTAssertNotEqual(MLX.Stream.gpu.device, gpu7) + } + + // new scoped stream with default cpu but gpu inheriting from parent (gpu7) + Stream.withNewDefaultStream(device: .cpu) { + XCTAssertEqual(MLX.Stream.defaultStream.device, outerCpu.device) + + // not the same stream + XCTAssertNotEqual(MLX.Stream.gpu, gpu7Stream) + + // but the same device + XCTAssertEqual(MLX.Stream.gpu.device, gpu7) + } + + } + } + + func testWithNewDefaultStreamDefaultsToCurrentDevice() { + // `device:` omitted means "a new stream on the current default device". + Device.withDefaultDevice(.cpu) { + MLX.Stream.withNewDefaultStream { + XCTAssertEqual(Device.defaultDevice(), Device.cpu) + XCTAssertEqual(MLX.Stream.defaultStream, MLX.Stream.cpu) + } } - print("here") } func testWithDefaultDevice() { @@ -71,7 +185,179 @@ class StreamTests: XCTestCase { } } } - print("here") + } + + // MARK: - Task-local scoping (async + propagation boundaries) + + func testScopeIsInheritedByChildTasks() async { + // The scope is task-local, so structured child tasks inherit it. + await Device.withDefaultDevice(.cpu) { + async let child = Device.defaultDevice() + let value = await child + XCTAssertEqual(value, Device.cpu) + + await withTaskGroup(of: Device.self) { group in + for _ in 0 ..< 4 { + group.addTask { Device.defaultDevice() } + } + for await device in group { + XCTAssertEqual(device, Device.cpu) + } + } + } + } + + func testScopeIsNotInheritedByDetachedWork() async { + // detached Tasks and Threads do not inherit + let global = Device.defaultDevice() + + // pick a scope device that differs from the process default so the + // assertion below can actually distinguish them + let scoped = Device(.gpu, index: 3) + XCTAssertNotEqual(scoped, global) + + await Device.withDefaultDevice(scoped) { + XCTAssertEqual(Device.defaultDevice(), scoped) + + let detached = await Task.detached { Device.defaultDevice() }.value + XCTAssertEqual(detached, global, "a detached task must see the global default") + + let fromThread = await withCheckedContinuation { continuation in + let thread = Thread { continuation.resume(returning: Device.defaultDevice()) } + thread.start() + } + XCTAssertEqual(fromThread, global, "a raw Thread must see the global default") + } + } + + func testConcurrentFirstTouchYieldsOneStream() async { + // test of the double checked lock in StreamPromise + await MLX.Stream.withNewDefaultStream(device: .cpu) { + let streams = await withTaskGroup(of: ObjectIdentifier.self) { group in + for _ in 0 ..< 16 { + group.addTask { ObjectIdentifier(MLX.Stream.defaultStream) } + } + var seen = Set() + for await id in group { seen.insert(id) } + return seen + } + + XCTAssertEqual( + streams.count, 1, + "a cold promise handed out \(streams.count) distinct streams under contention") + } + } + + // MARK: - Stream / StreamOrDevice resolution + + func testDeviceTypeOverloads() { + XCTAssertEqual(MLX.Stream.defaultStream(.cpu), MLX.Stream.cpu) + XCTAssertEqual(MLX.Stream.defaultStream(.gpu), MLX.Stream.gpu) + + XCTAssertEqual(StreamOrDevice.device(.cpu).stream, MLX.Stream.cpu) + XCTAssertEqual(StreamOrDevice.device(.gpu).stream, MLX.Stream.gpu) + + XCTAssertEqual(StreamOrDevice.cpu.stream, MLX.Stream.cpu) + XCTAssertEqual(StreamOrDevice.gpu.stream, MLX.Stream.gpu) + } + + func testStreamOrDeviceStaticsTrackTheCurrentScope() { + // `StreamOrDevice.cpu`/`.gpu` became computed, so they must track the + // current scope rather than caching whatever was current at first access. + MLX.Stream.withNewDefaultStream(device: .cpu) { + XCTAssertEqual(StreamOrDevice.cpu.stream, MLX.Stream.cpu) + XCTAssertEqual(StreamOrDevice.default.stream, MLX.Stream.defaultStream) + } + } + + @available(*, deprecated) + func testDeprecatedDeviceOverloadHonorsTheDevice() { + XCTAssertEqual(MLX.Stream.defaultStream(Device.cpu), MLX.Stream.cpu) + XCTAssertEqual(MLX.Stream.defaultStream(Device.gpu), MLX.Stream.gpu) + XCTAssertEqual(StreamOrDevice.device(Device.cpu).stream, MLX.Stream.cpu) + } + + // MARK: - Stream pool + + func testConcurrentStreamsAreDistinct() { + // Streams handed out at the same time must never alias -- a double-release + // into the pool would let two owners share one `mlx_stream`. + final class Box: @unchecked Sendable { + private let lock = NSLock() + private var streams: [MLX.Stream] = [] + func append(_ stream: MLX.Stream) { lock.withLock { streams.append(stream) } } + var all: [MLX.Stream] { lock.withLock { streams } } + } + + let n = 32 + let box = Box() + + DispatchQueue.concurrentPerform(iterations: n) { _ in + box.append(MLX.Stream(.cpu)) + } + + let live = box.all + XCTAssertEqual(live.count, n) + for i in 0 ..< live.count { + for j in (i + 1) ..< live.count { + XCTAssertNotEqual( + live[i], live[j], "the pool handed the same stream to two owners") + } + } + } + + func testPoolDoesNotCrossDevices() throws { + // recycling stream innards should never mix types + for _ in 0 ..< 100 { + autoreleasepool { + let device = Device(.gpu, index: Int32.random(in: 0 ..< 10)) + let stream = MLX.Stream(device) + XCTAssertEqual(stream.device, device) + } + } + } + + func testScopeExitRecyclesTheStream() { + // see https://github.com/ml-explore/mlx/issues/2118 -- this used to + // exhaust metal/OS resources; the pool makes it bounded + + // stream pooling should reuse the same stream over and over + var indices = Set() + + for _ in 0 ..< 100 { + MLX.Stream.withNewDefaultStream(device: .gpu) { + indices.insert(MLX.Stream.defaultStream.index) + + let x = MLXArray(1) + let _ = x * x + } + } + + XCTAssertEqual( + indices.count, 1, + "100 sequential scopes used \(indices.count) distinct streams, expected 1") + } + + func testPoolingConcurrent() { + // but if we have two live streams they cannot be the same + let s1 = Stream(.cpu) + let s2 = Stream(.cpu) + + XCTAssertNotEqual(s1.index, s2.index) + } + + // MARK: - Misc + + func testStreamSynchronize() { + let stream = MLX.Stream(.cpu) + stream.synchronize() + MLX.Stream.defaultStream.synchronize() + } + + func testStreamOrDeviceDescriptionMatchesItsStream() { + let stream = MLX.Stream(.cpu) + XCTAssertEqual(StreamOrDevice.stream(stream).description, stream.description) + XCTAssertEqual(StreamOrDevice.default.description, MLX.Stream.defaultStream.description) } func testStreamOrDeviceStreamUsesGivenStream() { @@ -84,22 +370,50 @@ class StreamTests: XCTestCase { XCTAssertEqual(StreamOrDevice.stream(.gpu).stream, Stream.gpu) } - func disabledTestCreateStream() { - // see https://github.com/ml-explore/mlx/issues/2118 - for _ in 1 ..< 10000 { - let _ = Stream(.cpu) + func testStreamPoolRecyclesStreams() { + // Streams must be recycled through the pool rather than allocated without + // bound -- see https://github.com/ml-explore/mlx/issues/2118 and #237. + + let n = 5 + + // hold n at once, drop them, then take n again -- the second batch must + // be the same underlying streams + var first = Set() + var live: [MLX.Stream] = [] + for _ in 0 ..< n { + let stream = Stream(.cpu) + first.insert(stream.index) + live.append(stream) } - print("here") + XCTAssertEqual(first.count, n, "streams held at the same time must be distinct") + live.removeAll() + + var second = Set() + for _ in 0 ..< n { + let stream = Stream(.cpu) + second.insert(stream.index) + live.append(stream) + } + XCTAssertEqual(second, first, "the pool did not recycle the released streams") + live.removeAll() + + // and the set of distinct streams stays bounded over many rounds + var observed = first + for _ in 0 ..< 10000 { + let stream = Stream(.cpu) + observed.insert(stream.index) + } + XCTAssertEqual( + observed, first, + "10k sequential streams allocated \(observed.count) distinct streams, expected \(n)") } - func disabledTestCreateStreamScoped() { + func testCreateStream() { // see https://github.com/ml-explore/mlx/issues/2118 for _ in 1 ..< 10000 { - Stream.withNewDefaultStream(device: .cpu) { - let x = MLXArray(1) - let _ = x * x - } + let _ = Stream(.cpu) } + print("here") } } diff --git a/Tests/MLXTests/TransformTests.swift b/Tests/MLXTests/TransformTests.swift index 1ae105c22..71b58c990 100644 --- a/Tests/MLXTests/TransformTests.swift +++ b/Tests/MLXTests/TransformTests.swift @@ -10,7 +10,6 @@ import XCTest class TransformTests: XCTestCase { override class func setUp() { - setDefaultDevice() } func testEval() { diff --git a/Tests/MLXTests/Utils.swift b/Tests/MLXTests/Utils.swift index 520306bd5..00c829c86 100644 --- a/Tests/MLXTests/Utils.swift +++ b/Tests/MLXTests/Utils.swift @@ -33,7 +33,3 @@ func assertNotEqual( array1.allClose(array2, rtol: rtol, atol: atol).item(Bool.self), "contents same:\n\(array1)\n\(array2)") } - -func setDefaultDevice() { - MLX.Device.setDefault(device: .gpu) -} From 901941965d82e4a216d4d117231d847d194c563d Mon Sep 17 00:00:00 2001 From: David Koski <46639364+davidkoski@users.noreply.github.com> Date: Thu, 17 Sep 2026 11:40:52 -0700 Subject: [PATCH 34/39] replace and improve integration tests (#477) * replace and improve integration tests - previously we used a python script to generate swift code and check output against known values from mlx (python) - this is replaced with a new and more powerful script (claude rewrite) - much more integration test coverage and tools to examine that - fixes per problems found by new integration tests - linspace output type, add dtype override - linspace default output type is float32 (matches python) - convolve .same padding side matches python (padRight) - nanToNum default for +/- infinity honors optionals (mlx defaults) - tensorDot axes default 1 -> 2 (matches python) - MLXNN Pooling padWidths matches python - MLXNN Alibi slop function matches python - MLXNN GRU apply bias to match python - MLXNN Lion optimizer default beta matches python - MLXNN Adafactor moving average matches python - MLXNN Muon zero power newton shulz matches python * fix linspace disambiguation * split integration tests into their own test target -- we need to run them with TF32 disabled so we can guarantee consistent results on NAX vs non NAX hardware Co-authored-by: Nick Henderson <525247+nwh@users.noreply.github.com> --- .github/workflows/pull_request.yml | 1 + MAINTENANCE.md | 7 +- Package.swift | 6 + .../Articles/converting-python.md | 10 +- .../Organization/cumulative.md | 4 +- .../Organization/initialization.md | 8 +- .../MLX/Documentation.docc/free-functions.md | 4 +- Source/MLX/Factory.swift | 63 +- Source/MLX/Ops+Array.swift | 8 +- Source/MLX/Ops.swift | 9 +- Source/MLXNN/Pooling.swift | 7 +- Source/MLXNN/PositionalEncoding.swift | 30 +- Source/MLXNN/Recurrent.swift | 6 +- Source/MLXOptimizers/Optimizers.swift | 46 +- .../Generated/GeneratedActivationsTests.swift | 823 ++ .../Generated/GeneratedBinaryTests.swift | 1462 ++++ .../Generated/GeneratedConvolutionTests.swift | 422 + .../Generated/GeneratedDefaultsTests.swift | 1410 ++++ .../Generated/GeneratedElementwiseTests.swift | 1517 ++++ .../Generated/GeneratedFFTTests.swift | 1513 ++++ .../Generated/GeneratedFactoryTests.swift | 500 ++ .../Generated/GeneratedFastTests.swift | 201 + .../Generated/GeneratedIndexingTests.swift | 166 + .../Generated/GeneratedLinalgTests.swift | 405 + .../Generated/GeneratedLossesTests.swift | 555 ++ .../GeneratedModuleActivationsTests.swift | 965 +++ .../GeneratedModuleAttentionTests.swift | 124 + .../GeneratedModuleConvolutionTests.swift | 306 + .../GeneratedModuleDropoutTests.swift | 109 + .../GeneratedModuleLinearTests.swift | 163 + .../GeneratedModuleNormalizationTests.swift | 308 + .../GeneratedModulePoolingTests.swift | 246 + .../GeneratedModulePositionalTests.swift | 191 + .../GeneratedModuleRecurrentTests.swift | 163 + .../GeneratedModuleUpsampleTests.swift | 162 + .../GeneratedOptimizerFunctionsTests.swift | 297 + .../Generated/GeneratedOptimizersTests.swift | 2113 +++++ .../GeneratedQuantizationTests.swift | 571 ++ .../GeneratedRandomInputsTests.swift | 596 ++ .../Generated/GeneratedRandomTests.swift | 269 + .../Generated/GeneratedReductionTests.swift | 1506 ++++ .../Generated/GeneratedSchedulesTests.swift | 224 + .../Generated/GeneratedShapeTests.swift | 1172 +++ .../IntegrationSupport.swift | 228 + Tests/MLXTests/DistributedNNTests.swift | 8 - Tests/MLXTests/DistributedTests.swift | 4 - Tests/MLXTests/IntegrationTests.swift | 7076 ----------------- Tests/MLXTests/MLXArray+InitTests.swift | 40 + Tests/MLXTests/OpsTests.swift | 79 + Tests/MLXTests/OptimizerTests.swift | 174 +- Tests/MLXTests/PoolingTests.swift | 83 + Tests/MLXTests/PositionalEncodingTests.swift | 167 + scripts/verify-docs.sh | 10 +- tools/audit_defaults.py | 428 + tools/audit_integration_coverage.py | 688 ++ tools/generate_integration_tests.py | 1063 --- tools/integration-coverage-report.md | 648 ++ tools/integration_tests/README.md | 407 + tools/integration_tests/cases.py | 2216 ++++++ tools/integration_tests/check_api.py | 242 + tools/integration_tests/core.py | 944 +++ tools/integration_tests/fake_mlx.py | 401 + tools/integration_tests/generate.py | 257 + tools/integration_tests/modules.py | 502 ++ tools/integration_tests/optimizers.py | 452 ++ xcode/MLX.xcodeproj/project.pbxproj | 283 +- xcode/MLX.xctestplan | 7 + .../MLXIntegrationTests.swift | 19 + 68 files changed, 26867 insertions(+), 8227 deletions(-) create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedActivationsTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedBinaryTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedConvolutionTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedDefaultsTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedElementwiseTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedFFTTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedFactoryTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedIndexingTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedLinalgTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedLossesTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleActivationsTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleAttentionTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleConvolutionTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleDropoutTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleLinearTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleNormalizationTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModulePoolingTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModulePositionalTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleRecurrentTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedModuleUpsampleTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedOptimizerFunctionsTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedOptimizersTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedRandomInputsTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedRandomTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedReductionTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedSchedulesTests.swift create mode 100644 Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift create mode 100644 Tests/MLXIntegrationTests/IntegrationSupport.swift delete mode 100644 Tests/MLXTests/IntegrationTests.swift create mode 100644 Tests/MLXTests/PositionalEncodingTests.swift create mode 100644 tools/audit_defaults.py create mode 100755 tools/audit_integration_coverage.py delete mode 100644 tools/generate_integration_tests.py create mode 100644 tools/integration-coverage-report.md create mode 100644 tools/integration_tests/README.md create mode 100644 tools/integration_tests/cases.py create mode 100644 tools/integration_tests/check_api.py create mode 100644 tools/integration_tests/core.py create mode 100644 tools/integration_tests/fake_mlx.py create mode 100644 tools/integration_tests/generate.py create mode 100644 tools/integration_tests/modules.py create mode 100644 tools/integration_tests/optimizers.py create mode 100644 xcode/MLXIntegrationTests/MLXIntegrationTests.swift diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index d8c76493c..751ff4493 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -115,6 +115,7 @@ jobs: run: | xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/CmlxTests.xctest xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest + xcrun xctest ~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXIntegrationTests.xctest - name: Run Distributed Tests SwiftPM (Xcode, macOS) shell: sh diff --git a/MAINTENANCE.md b/MAINTENANCE.md index f71c5b268..dbf4e1007 100644 --- a/MAINTENANCE.md +++ b/MAINTENANCE.md @@ -91,12 +91,7 @@ dependencies: [.product(name: "MLX", package: "mlx-swift"), .product(name: "MLXOptimizers", package: "mlx-swift")] ``` -10. Update `tools/generate_integration_tests.py` as needed - -``` -import MLXNN -@testable import MLXOptimizers -``` +10. Update `tools/integration_tests/cases.py` as needed, regenerate tests if needed 11. Update tests as needed diff --git a/Package.swift b/Package.swift index 1cb1bf138..f9bde7a86 100644 --- a/Package.swift +++ b/Package.swift @@ -405,6 +405,12 @@ let package = Package( "MLX", "MLXNN", "MLXOptimizers", ] ), + .testTarget( + name: "MLXIntegrationTests", + dependencies: [ + "MLX", "MLXNN", "MLXOptimizers", + ] + ), // ------ // Example programs diff --git a/Source/MLX/Documentation.docc/Articles/converting-python.md b/Source/MLX/Documentation.docc/Articles/converting-python.md index 6c94c240a..ddd45b344 100644 --- a/Source/MLX/Documentation.docc/Articles/converting-python.md +++ b/Source/MLX/Documentation.docc/Articles/converting-python.md @@ -107,8 +107,8 @@ Note: some of the symbols are not linkable. `cos` | ``MLXArray/cos(stream:)`` `cummax` | ``MLXArray/cummax(axis:reverse:inclusive:stream:)`` `cummin` | ``MLXArray/cummin(axis:reverse:inclusive:stream:)`` -`cumprod` | ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` -`cumsum` | ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` +`cumprod` | ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)`` +`cumsum` | ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)`` `exp` | ``MLXArray/exp(stream:)`` `flatten` | ``MLXArray/flattened(start:end:stream:)`` `log` | ``MLXArray/log(stream:)`` @@ -190,7 +190,7 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts. `identity` | ``MLXArray/identity(_:type:stream:)`` `less` | ``MLX/less(_:_:stream:)`` `less_equal` | ``MLX/lessEqual(_:_:stream:)`` -`linspace` | ``MLXArray/linspace(_:_:count:stream:)-(Int,Int,Int,StreamOrDevice)`` +`linspace` | ``MLXArray/linspace(_:_:count:dtype:stream:)-(Int,Int,Int,DType?,StreamOrDevice)`` `load` | ``MLX/loadArray(url:stream:)`` and ``MLX/loadArrays(url:stream:)`` `log` | ``MLX/log(_:stream:)`` `log10` | ``MLX/log10(_:stream:)`` @@ -210,7 +210,7 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts. `negative` | ``MLX/negative(_:stream:)`` `not_equal` | ``MLX/notEqual(_:_:stream:)`` `ones` | ``MLXArray/ones(_:type:stream:)`` -`ones_like` | ``MLXArray/ones(like:stream:)`` +`ones_like` | ``MLXArray/ones(like:dtype:stream:)`` `pad` | ``MLX/padded(_:width:mode:value:stream:)`` `partition` | ``MLX/partitioned(_:kth:axis:stream:)`` `power` | ``MLX/pow(_:_:stream:)-(MLXArray,MLXArray,_)`` @@ -255,4 +255,4 @@ This is a mapping of `mx` free functions to their ``MLX`` counterparts. `var` | ``MLX/variance(_:axes:keepDims:ddof:stream:)`` `where` | ``MLX/which(_:_:_:stream:)`` `zeros` | ``MLXArray/zeros(_:type:stream:)`` -`zeros_like` | ``MLXArray/zeros(like:stream:)`` +`zeros_like` | ``MLXArray/zeros(like:dtype:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/cumulative.md b/Source/MLX/Documentation.docc/Organization/cumulative.md index b9c4aeb6c..541f23d13 100644 --- a/Source/MLX/Documentation.docc/Organization/cumulative.md +++ b/Source/MLX/Documentation.docc/Organization/cumulative.md @@ -25,9 +25,9 @@ These are available as both methods on `MLXArray` and free functions. They each - ``MLXArray/cummax(reverse:inclusive:stream:)`` - ``MLXArray/cummin(axis:reverse:inclusive:stream:)`` - ``MLXArray/cummin(reverse:inclusive:stream:)`` -- ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` +- ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)`` - ``MLXArray/cumprod(reverse:inclusive:stream:)`` -- ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` +- ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)`` - ``MLXArray/cumsum(reverse:inclusive:stream:)`` - ``MLXArray/logCumsumExp(axis:reverse:inclusive:stream:)`` - ``MLXArray/logCumsumExp(reverse:inclusive:stream:)`` diff --git a/Source/MLX/Documentation.docc/Organization/initialization.md b/Source/MLX/Documentation.docc/Organization/initialization.md index 447ba6ed2..dc3e6132f 100644 --- a/Source/MLX/Documentation.docc/Organization/initialization.md +++ b/Source/MLX/Documentation.docc/Organization/initialization.md @@ -264,8 +264,8 @@ there are specific initializers to request it: - ``MLXArray/full(_:values:type:stream:)`` - ``MLXArray/full(_:values:stream:)`` - ``MLXArray/identity(_:type:stream:)`` -- ``MLXArray/linspace(_:_:count:stream:)-(Int,Int,Int,StreamOrDevice)`` -- ``MLXArray/linspace(_:_:count:stream:)-(Double,Double,Int,StreamOrDevice)`` +- ``MLXArray/linspace(_:_:count:dtype:stream:)-(Int,Int,Int,DType?,StreamOrDevice)`` +- ``MLXArray/linspace(_:_:count:dtype:stream:)-(Double,Double,Int,DType?,StreamOrDevice)`` - ``MLXArray/repeated(_:count:axis:stream:)`` - ``MLXArray/repeated(_:count:stream:)`` - ``MLXArray/repeat(_:count:axis:stream:)`` @@ -282,8 +282,8 @@ there are specific initializers to request it: - ``MLX/full(_:values:type:stream:)`` - ``MLX/full(_:values:stream:)`` - ``MLX/identity(_:type:stream:)`` -- ``linspace(_:_:count:endpoint:stream:)-47tc1`` -- ``linspace(_:_:count:endpoint:stream:)-38sfd`` +- ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu`` +- ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2`` - ``MLXArray/repeated(_:count:axis:stream:)`` - ``MLXArray/repeated(_:count:stream:)`` - ``MLX/repeat(_:count:axis:stream:)`` diff --git a/Source/MLX/Documentation.docc/free-functions.md b/Source/MLX/Documentation.docc/free-functions.md index e0592dbec..c54ebddbf 100644 --- a/Source/MLX/Documentation.docc/free-functions.md +++ b/Source/MLX/Documentation.docc/free-functions.md @@ -106,8 +106,8 @@ operations as methods for convenience. - ``MLX/full(_:values:type:stream:)`` - ``MLX/full(_:values:stream:)`` - ``MLX/identity(_:type:stream:)`` -- ``linspace(_:_:count:endpoint:stream:)-47tc1`` -- ``linspace(_:_:count:endpoint:stream:)-38sfd`` +- ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu`` +- ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2`` - ``MLX/repeated(_:count:axis:stream:)`` - ``MLX/repeated(_:count:stream:)`` - ``MLX/repeat(_:count:axis:stream:)`` diff --git a/Source/MLX/Factory.swift b/Source/MLX/Factory.swift index 0abbae72d..0741e0c29 100644 --- a/Source/MLX/Factory.swift +++ b/Source/MLX/Factory.swift @@ -331,31 +331,34 @@ extension MLXArray { MLX.identity(n, dtype: dtype, stream: stream) } - /// Generate `num` evenly spaced numbers over interval `[start, stop]` for `BinaryInteger`. + /// Generate `count` evenly spaced numbers over interval `[start, stop]` for `BinaryInteger`. /// - /// Example: + /// The result is floating point (`float32` by default) even for integer + /// bounds -- see ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu``. /// /// ```swift - /// // Create a 50 element 1-D array with values from 0 to 50 - /// let r = MLXArray.linspace(0, 50) + /// // [0, 0.5, 1] as float32 + /// let r = MLXArray.linspace(0, 1, count: 3) /// ``` /// /// - Parameters: /// - start: start value /// - stop: stop value /// - count: number of samples + /// - dtype: dtype of the result, `float32` if not specified /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``linspace(_:_:count:stream:)-92x6l`` + /// - ``linspace(_:_:count:dtype:stream:)-3fx01`` static public func linspace( - _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default + _ start: T, _ stop: T, count: Int = 50, dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryInteger { - MLX.linspace(start, stop, count: count, stream: stream) + MLX.linspace(start, stop, count: count, dtype: dtype, stream: stream) } - /// Generate `num` evenly spaced numbers over interval `[start, stop]` for `BinaryFloatingPoint`. + /// Generate `count` evenly spaced numbers over interval `[start, stop]` for `BinaryFloatingPoint`. /// /// Example: /// @@ -368,15 +371,17 @@ extension MLXArray { /// - start: start value /// - stop: stop value /// - count: number of samples + /// - dtype: dtype of the result, derived from `T` if not specified /// - stream: stream or device to evaluate on /// /// ### See Also /// - - /// - ``linspace(_:_:count:stream:)-7m7eg`` + /// - ``linspace(_:_:count:dtype:stream:)-9yqai`` static public func linspace( - _ start: T, _ stop: T, count: Int = 50, stream: StreamOrDevice = .default + _ start: T, _ stop: T, count: Int = 50, dtype: DType? = nil, + stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryFloatingPoint { - MLX.linspace(start, stop, count: count, stream: stream) + MLX.linspace(start, stop, count: count, dtype: dtype, stream: stream) } /// Generate values in the half-open interval `[0, stop)`. @@ -1007,13 +1012,18 @@ public func identity(_ n: Int, dtype: DType, stream: StreamOrDevice = .default) return MLXArray(result) } -/// Generate `num` evenly spaced numbers over interval `[start, stop]`. +/// Generate `count` evenly spaced numbers over interval `[start, stop]`. /// -/// Example: +/// The result is floating point (`float32` by default), matching python's +/// `mx.linspace`, even for integer bounds: evenly spaced values between two +/// integers are generally not integers. Pass `dtype:` for anything else: /// /// ```swift -/// // Create a 50 element 1-D array with values from 0 to 50 -/// let r = MLXArray.linspace(0, 50) +/// // [0, 0.5, 1] as float32 -- *not* [0, 0, 1] as int64 +/// let r = MLXArray.linspace(0, 1, count: 3) +/// +/// // opt in to an integer (truncating) result +/// let i = MLXArray.linspace(0, 10, count: 6, dtype: .int32) /// ``` /// /// - Parameters: @@ -1021,25 +1031,31 @@ public func identity(_ n: Int, dtype: DType, stream: StreamOrDevice = .default) /// - stop: stop value /// - count: number of samples /// - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval +/// - dtype: dtype of the result, `float32` if not specified /// - stream: stream or device to evaluate on /// /// ### See Also /// - -/// - ``linspace(_:_:count:endpoint:stream:)`` +/// - ``linspace(_:_:count:endpoint:dtype:stream:)-2b6eu`` public func linspace( _ start: T, _ stop: T, count: Int = 50, endpoint: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryInteger { var result = mlx_array_new() mlx_linspace_endpoint( - &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx) + &result, Double(start), Double(stop), count.int32, endpoint, + (dtype ?? .float32).cmlxDtype, stream.ctx) return MLXArray(result) } -/// Generate `num` evenly spaced numbers over interval `[start, stop]`. +/// Generate `count` evenly spaced numbers over interval `[start, stop]`. /// -/// Example: +/// The result dtype follows the bounds (`Float` -> `float32`, +/// `Float16` -> `float16`) except that `Double` produces `float32`, matching +/// ``MLXArray/init(_:)`` and python's `mx.linspace` -- `float64` is not +/// supported on the GPU. Pass `dtype:` to be explicit. /// /// ```swift /// // Create a 50 element 1-D array with values from 0 to 1 @@ -1051,19 +1067,24 @@ public func linspace( /// - stop: stop value /// - count: number of samples /// - endpoint: if `true` then the endpoint is the last sample, if `false` it is a half-open interval +/// - dtype: dtype of the result, derived from `T` if not specified /// - stream: stream or device to evaluate on /// /// ### See Also /// - -/// - ``linspace(_:_:count:endpoint:stream:)`` +/// - ``linspace(_:_:count:endpoint:dtype:stream:)-8k1d2`` public func linspace( _ start: T, _ stop: T, count: Int = 50, endpoint: Bool = true, + dtype: DType? = nil, stream: StreamOrDevice = .default ) -> MLXArray where T: BinaryFloatingPoint { + // Double.dtype is float64, but we do not automatically promote to float64 + // (see HasDType conformances) and it is not available on the GPU + let resolved = dtype ?? (T.dtype == .float64 ? .float32 : T.dtype) var result = mlx_array_new() mlx_linspace_endpoint( - &result, Double(start), Double(stop), count.int32, endpoint, T.dtype.cmlxDtype, stream.ctx) + &result, Double(start), Double(stop), count.int32, endpoint, resolved.cmlxDtype, stream.ctx) return MLXArray(result) } diff --git a/Source/MLX/Ops+Array.swift b/Source/MLX/Ops+Array.swift index b21d857e1..a1cba744f 100644 --- a/Source/MLX/Ops+Array.swift +++ b/Source/MLX/Ops+Array.swift @@ -605,7 +605,7 @@ public func cummin( /// ### See Also /// - /// - ``cumprod(_:reverse:inclusive:dtype:stream:)`` -/// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` +/// - ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)`` public func cumprod( _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true, dtype: DType? = nil, @@ -629,7 +629,7 @@ public func cumprod( /// ### See Also /// - /// - ``cumprod(_:axis:reverse:inclusive:dtype:stream:)`` -/// - ``MLXArray/cumprod(axis:reverse:inclusive:stream:)`` +/// - ``MLXArray/cumprod(axis:reverse:inclusive:dtype:stream:)`` public func cumprod( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, dtype: DType? = nil, @@ -653,7 +653,7 @@ public func cumprod( /// ### See Also /// - /// - ``cumsum(_:reverse:inclusive:dtype:stream:)`` -/// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` +/// - ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)`` public func cumsum( _ array: MLXArray, axis: Int, reverse: Bool = false, inclusive: Bool = true, dtype: DType? = nil, @@ -677,7 +677,7 @@ public func cumsum( /// ### See Also /// - /// - ``cumsum(_:axis:reverse:inclusive:dtype:stream:)`` -/// - ``MLXArray/cumsum(axis:reverse:inclusive:stream:)`` +/// - ``MLXArray/cumsum(axis:reverse:inclusive:dtype:stream:)`` public func cumsum( _ array: MLXArray, reverse: Bool = false, inclusive: Bool = true, dtype: DType? = nil, diff --git a/Source/MLX/Ops.swift b/Source/MLX/Ops.swift index 0c689303b..5a0da9e1a 100644 --- a/Source/MLX/Ops.swift +++ b/Source/MLX/Ops.swift @@ -1002,6 +1002,9 @@ public func convolve( if weightSize % 2 == 1 { padding = weightSize / 2 } else { + // even sized weights use asymmetric padding -- this must match + // python's `mx.convolve()` so that the result is the centered + // `input.size` window of the full convolution let padLeft = weightSize / 2 let padRight = max(0, padLeft - 1) @@ -2184,7 +2187,7 @@ public func multiply( /// - public func nanToNum( _ array: MLXArray, - nan: Float = 0, posInf: Float? = 0, negInf: Float? = 0, + nan: Float = 0, posInf: Float? = nil, negInf: Float? = nil, stream: StreamOrDevice = .default ) -> MLXArray { let posInf = mlx_optional_float(value: posInf ?? 0, has_value: posInf != nil) @@ -3110,7 +3113,7 @@ public func tanh(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArr /// - Parameters: /// - a: input array /// - b: input array -/// - axes: sum over the last `axes` dimensions +/// - axes: sum over the last `axes` dimensions of `a` and the first `axes` of `b` /// - stream: stream or device to evaluate on /// - Returns: tensor dot product /// @@ -3118,7 +3121,7 @@ public func tanh(_ array: MLXArray, stream: StreamOrDevice = .default) -> MLXArr /// - /// - ``tensordot(_:_:axes:stream:)-(MLXArray,MLXArray,Int,StreamOrDevice)`` public func tensordot( - _ a: MLXArray, _ b: MLXArray, axes: Int = 1, stream: StreamOrDevice = .default + _ a: MLXArray, _ b: MLXArray, axes: Int = 2, stream: StreamOrDevice = .default ) -> MLXArray { var result = mlx_array_new() mlx_tensordot_axis(&result, a.ctx, b.ctx, axes.int32, stream.ctx) diff --git a/Source/MLXNN/Pooling.swift b/Source/MLXNN/Pooling.swift index c978176db..1110f318f 100644 --- a/Source/MLXNN/Pooling.swift +++ b/Source/MLXNN/Pooling.swift @@ -41,9 +41,10 @@ open class Pool: Module, UnaryLayer { // Apply padding if any padding value is greater than 0 if padding.contains(where: { $0 > 0 }) { - // batch and channel dimension get no padding - let padWidths: [IntOrPair] = - [0, 0] + padding.map { .init($0) } + [0, 0] + // one width per dimension: batch and channel get no padding, each + // spatial dimension is padded symmetrically (`IntOrPair` is already + // the (before, after) pair for a single dimension) + let padWidths: [IntOrPair] = [0] + padding.map { IntOrPair($0) } + [0] let paddingValue = paddingValue.asMLXArray(dtype: input.dtype) input = padded(input, widths: padWidths, mode: .constant, value: paddingValue) } diff --git a/Source/MLXNN/PositionalEncoding.swift b/Source/MLXNN/PositionalEncoding.swift index bf690feba..2dca4ea1d 100644 --- a/Source/MLXNN/PositionalEncoding.swift +++ b/Source/MLXNN/PositionalEncoding.swift @@ -156,10 +156,30 @@ final public class ALiBi: Module { public override init() { } + /// The per-head slopes, matching python's `ALiBi.create_alibi_slope()`. + /// + /// For a power of two head count the slopes are `2^(-8i/n)` for `i` in + /// `1...n`. Otherwise python uses the slopes of the next power of two *below* + /// `n` and pads them with every other slope of the next power of two *above*, + /// which is not the same as extending the geometric series. static func alibiSlope(numHeads: Int) -> MLXArray { - let x = pow(pow(2, 8), (1 / Float(numHeads))) - let out = pow(x, -MLXArray(1 ..< (numHeads + 1))) - return out.expandedDimensions(axes: [-1, -2]) + func slopes(_ n: Int) -> [Float] { + let log2n = Foundation.log2(Double(n)) + if log2n == log2n.rounded(.down) { + let start = Foundation.pow(2.0, -Foundation.pow(2.0, 3 - log2n)) + return (1 ... n).map { Float(Foundation.pow(start, Double($0))) } + } + + let closestPowerOf2 = Int(Foundation.pow(2.0, log2n.rounded(.down))) + let interleaved = slopes(2 * closestPowerOf2) + .enumerated() + .filter { $0.offset.isMultiple(of: 2) } + .map { $0.element } + .prefix(n - closestPowerOf2) + return slopes(closestPowerOf2) + interleaved + } + + return MLXArray(slopes(numHeads)).expandedDimensions(axes: [-1, -2]) } static func alibiMatrix(key: Key) -> MLXArray { @@ -167,8 +187,10 @@ final public class ALiBi: Module { return value } + // x1 is a column and x2 a row so that the difference is the (q, k) + // distance matrix -- python: `x1[:, None] - x2[None, :]` let x1 = MLXArray(key.offset ..< key.qSequenceLength).expandedDimensions(axis: 1) - let x2 = MLXArray(0 ..< key.kSequenceLength).expandedDimensions(axis: 1) + let x2 = MLXArray(0 ..< key.kSequenceLength).expandedDimensions(axis: 0) let distanceMatrix = -abs(expandedDimensions((x1 - x2), axes: [0, 1])) let slope = alibiSlope(numHeads: key.numHeads) diff --git a/Source/MLXNN/Recurrent.swift b/Source/MLXNN/Recurrent.swift index 59a87607d..ea3ca857a 100644 --- a/Source/MLXNN/Recurrent.swift +++ b/Source/MLXNN/Recurrent.swift @@ -160,8 +160,12 @@ open class GRU: Module { var n = x_n[.ellipsis, index, 0...] if hidden != nil { - // Note: xProj_n was computed earlier + // Note: hProj_n was computed earlier and already includes bhn n = n + r * hProj_n + } else if let bhn { + // matches python: without an incoming hidden state the hidden + // bias is still applied, gated by r + n = n + r * bhn } n = tanh(n) diff --git a/Source/MLXOptimizers/Optimizers.swift b/Source/MLXOptimizers/Optimizers.swift index a01083b85..42f247517 100644 --- a/Source/MLXOptimizers/Optimizers.swift +++ b/Source/MLXOptimizers/Optimizers.swift @@ -9,6 +9,17 @@ import MLXNN /// ### See Also /// - /// - ``OptimizerBase`` +// Note on scalar precision: these optimizers store their hyperparameters as `Float` +// while python stores them as double, so coefficients derived from them differ in +// the last bits -- `1 - Float(0.95)` is 0.050000012, while python's `1 - 0.95` +// rounds to 0.05. Widening at the point of use (`1 - Double(momentum)`) cannot +// recover python's value: the constant was already rounded when it was stored. +// +// The 2.4e-7 difference is invisible in the smooth optimizers, but Muon's +// Newton-Schulz iteration grows it to ~2e-4 relative over a few steps, which is why +// the generated `Muon` cases compare with a looser tolerance than the rest. See +// `todos/muon-04-pending-p3-float-hyperparameters.md`. + public protocol Optimizer: Updatable, Evaluatable { /// Apply the gradients to the parameters of the model and update the model with the new parameters. @@ -561,17 +572,17 @@ open class Lion: OptimizerBaseArrayState { /// The learning rate public var learningRate: Float - /// The coefficients used for computing running averages of the gradient and its square - public var betas: (Float, Float) = (0.9, 0.999) + /// The coefficients used for computing the gradient momentum and update direction + public var betas: (Float, Float) = (0.9, 0.99) /// The weight decay public var weightDecay: Float = 0.0 /// Initialize the optimizer. /// - Parameters: /// - learningRate: the learning rate - /// - betas: coefficients used for computing running averages of the gradient and its square + /// - betas: coefficients used for computing the gradient momentum and update direction /// - weightDecay:the weight decay - public init(learningRate: Float, betas: (Float, Float) = (0.9, 0.999), weightDecay: Float = 0.0) + public init(learningRate: Float, betas: (Float, Float) = (0.9, 0.99), weightDecay: Float = 0.0) { self.learningRate = learningRate self.betas = betas @@ -710,7 +721,9 @@ open class Adafactor: OptimizerBase { func approvateExpMovingAverage(expAvgSqRow: MLXArray, expAvgSqCol: MLXArray) -> MLXArray { let rFactor = rsqrt(expAvgSqRow / mean(expAvgSqRow, axis: -1, keepDims: true)) let cFactor = rsqrt(expAvgSqCol) - return matmul(rFactor.expandedDimensions(axis: -1), cFactor.expandedDimensions(axis: 0)) + // broadcast rather than matmul so this also works for parameters with more + // than two dimensions + return rFactor.expandedDimensions(axis: -1) * cFactor.expandedDimensions(axis: -2) } override open func applySingle(gradient: MLXArray, parameter: MLXArray, state: State) -> ( @@ -730,8 +743,15 @@ open class Adafactor: OptimizerBase { var update = square(gradient) + eps.0 if factored { - var expAvgSqRow = state.expAvgSqRow! - var expAvgSqCol = state.expAvgSqCol! + // the state is created from the gradient shape so that it always agrees + // with the branch taken here, even if a parameter's rank changed + let rowShape = Array(gradientShape.dropLast()) + let columnShape = Array(gradientShape.dropLast(2)) + [gradientShape.last!] + + var expAvgSqRow = + state.expAvgSqRow ?? MLXArray.zeros(rowShape, dtype: gradient.dtype) + var expAvgSqCol = + state.expAvgSqCol ?? MLXArray.zeros(columnShape, dtype: gradient.dtype) expAvgSqRow = (beta2 * expAvgSqRow) + (1 - beta2) * mean(update, axis: -1) expAvgSqCol = (beta2 * expAvgSqCol) + (1 - beta2) * mean(update, axis: -2) @@ -742,7 +762,7 @@ open class Adafactor: OptimizerBase { update = approvateExpMovingAverage(expAvgSqRow: expAvgSqRow, expAvgSqCol: expAvgSqCol) update = update * gradient } else { - var expAvgSq = state.expAvgSq! + var expAvgSq = state.expAvgSq ?? MLXArray.zeros(like: gradient) expAvgSq = (beta2 * expAvgSq) + (1 - beta2) * update state.expAvgSq = expAvgSq update = rsqrt(expAvgSq) * gradient @@ -752,7 +772,7 @@ open class Adafactor: OptimizerBase { update = learningRate * update if let beta1 { - var expAvg = state.expAvg! + var expAvg = state.expAvg ?? MLXArray.zeros(like: gradient) expAvg = (beta1 * expAvg) + (1 - beta1) * update state.expAvg = expAvg update = expAvg @@ -811,19 +831,19 @@ open class Muon: OptimizerBaseArrayState { } /// Orthogonalize a 2D matrix via a quintic Newton-Schulz iteration. - private func zeropowerViaNewtonSchulz5(_ input: MLXArray, steps: Int) -> MLXArray { + func zeropowerViaNewtonSchulz5(_ input: MLXArray, steps: Int) -> MLXArray { precondition(input.ndim == 2, "Newton-Schulz iteration expects a 2D array") let (a, b, c): (Float, Float, Float) = (3.4445, -4.7750, 2.0315) let transposeNeeded = input.dim(-2) > input.dim(-1) var X = transposeNeeded ? input.transposed(1, 0) : input // Frobenius-normalize so the iteration converges. - X = X / (sqrt((X * X).sum(keepDims: true)) + 1e-7) + X = X / (MLX.norm(X, keepDims: true) + 1e-7) for _ in 0 ..< steps { let A = matmul(X, X.transposed(1, 0)) - let B = b * A + c * matmul(A, A) - X = a * X + matmul(B, X) + let B = addMM(b * A, A, A, alpha: c, beta: 1.0) + X = addMM(a * X, B, X, alpha: 1.0, beta: 1.0) } return transposeNeeded ? X.transposed(1, 0) : X diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedActivationsTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedActivationsTests.swift new file mode 100644 index 000000000..6bb65b882 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedActivationsTests.swift @@ -0,0 +1,823 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 34 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: Activations") +struct GeneratedActivationsTests { + + @Test("relu") + func test_relu() throws { + try withIntegrationState(seed: 67500) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.relu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5200212001800537, + minimum: 0.0, + maximum: 2.232034921646118, + absoluteSum: 6.2402544021606445, + positionChecksum: 3.4608532587687173, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.8907164931297302, 0.0, 0.36071768403053284]), + tolerance: .float32) + } + } + + @Test("reluSquared") + func test_reluSquared() throws { + try withIntegrationState(seed: 41758) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.reluSquared(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.453746497631073, + minimum: 0.0, + maximum: 3.3485403060913086, + absoluteSum: 5.444957733154297, + positionChecksum: 2.6932687759399414, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, 0.5844548344612122, 0.09588181227445602, 0.10871507227420807, 0.0, + 0.499411016702652, + ]), + tolerance: .float32) + } + } + + @Test("relu6") + func test_relu6() throws { + try withIntegrationState(seed: 12149) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.relu6(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.240987628698349, + minimum: 0.0, + maximum: 0.8727357983589172, + absoluteSum: 2.8918514251708984, + positionChecksum: 2.2303806940714517, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, 0.0, 0.0, 0.2970517873764038, 0.8727357983589172, 0.5916443467140198, + ]), + tolerance: .float32) + } + } + + @Test("leakyRelu") + func test_leakyRelu() throws { + try withIntegrationState(seed: 56128) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.leakyRelu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0438215732574463, + minimum: -0.018068933859467506, + maximum: 2.521155834197998, + absoluteSum: 12.575902938842773, + positionChecksum: 7.75894292195638, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7122242450714111, 1.8448302745819092, 0.9383899569511414, + 1.2284408807754517, -0.018068933859467506, 2.521155834197998, + ]), + tolerance: .float32) + } + } + + @Test("leakyRelu/slope") + func test_leakyRelu_slope() throws { + try withIntegrationState(seed: 39880) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.leakyRelu(x, negativeSlope: 0.2) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.4364432692527771, + minimum: -0.2560635805130005, + maximum: 1.8702387809753418, + absoluteSum: 6.699383735656738, + positionChecksum: 4.133757909138997, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6765661239624023, 1.2138452529907227, 0.13987021148204803, + -0.16138756275177002, 1.8702387809753418, 1.159955382347107, + ]), + tolerance: .float32) + } + } + + @Test("elu") + func test_elu() throws { + try withIntegrationState(seed: 86398) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.elu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.25546324253082275, + minimum: -0.8628273606300354, + maximum: 2.130934000015259, + absoluteSum: 7.796996116638184, + positionChecksum: 3.408879597981771, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.1739104986190796, 2.130934000015259, -0.8628273606300354, + -0.36768800020217896, 0.5139061212539673, 0.48894003033638, + ]), + tolerance: .float32) + } + } + + @Test("elu/alpha") + func test_elu_alpha() throws { + try withIntegrationState(seed: 33060) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.elu(x, alpha: 0.5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1415794938802719, + minimum: -0.3538806438446045, + maximum: 1.1850383281707764, + absoluteSum: 4.625144004821777, + positionChecksum: 2.805095672607422, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.09659192711114883, -0.1328100562095642, 0.11225724220275879, + 0.22676324844360352, -0.3352320194244385, -0.26999032497406006, + ]), + tolerance: .float32) + } + } + + @Test("celu") + func test_celu() throws { + try withIntegrationState(seed: 19670) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.celu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.14471694827079773, + minimum: -0.7049131393432617, + maximum: 1.0946556329727173, + absoluteSum: 5.854050159454346, + positionChecksum: 2.7480808893839517, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.045488059520721436, 1.0946556329727173, 0.7886189222335815, + -0.6293155550956726, 0.3050518333911896, 0.0885806754231453, + ]), + tolerance: .float32) + } + } + + @Test("celu/alpha") + func test_celu_alpha() throws { + try withIntegrationState(seed: 5167) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.celu(x, alpha: 0.5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.14974400401115417, + minimum: -0.43982020020484924, + maximum: 1.1591838598251343, + absoluteSum: 5.7410783767700195, + positionChecksum: 3.1327006022135415, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.43982020020484924, -0.4007747173309326, -0.206155925989151, + 0.9863490462303162, -0.006057232618331909, 0.0833292305469513, + ]), + tolerance: .float32) + } + } + + @Test("silu") + func test_silu() throws { + try withIntegrationState(seed: 84525) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.silu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1968996822834015, + minimum: -0.2784386873245239, + maximum: 1.4008158445358276, + absoluteSum: 5.122735500335693, + positionChecksum: 2.5903452237447104, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.2193196713924408, 1.4008158445358276, 0.881088376045227, + -0.2784386873245239, 0.7538488507270813, 0.07409554719924927, + ]), + tolerance: .float32) + } + } + + @Test("mish") + func test_mish() throws { + try withIntegrationState(seed: 7302) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.mish(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.21375694870948792, + minimum: -0.3031258285045624, + maximum: 1.1241786479949951, + absoluteSum: 5.914010047912598, + positionChecksum: 3.54515806833903, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.27689889073371887, -0.29029443860054016, 1.1014254093170166, + 1.1241786479949951, 0.8475918173789978, 0.5795539617538452, + ]), + tolerance: .float32) + } + } + + @Test("selu") + func test_selu() throws { + try withIntegrationState(seed: 76809) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.selu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.6044207811355591, + minimum: -1.5526567697525024, + maximum: 0.91837078332901, + absoluteSum: 10.77687931060791, + positionChecksum: 5.267512957255046, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.5526567697525024, -1.3209490776062012, -0.4850633144378662, + -0.5401492118835449, 0.6920045018196106, -1.2968387603759766, + ]), + tolerance: .float32) + } + } + + @Test("softplus") + func test_softplus() throws { + try withIntegrationState(seed: 6943) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softplus(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.8977645635604858, + minimum: 0.26141318678855896, + maximum: 2.1883463859558105, + absoluteSum: 10.773174285888672, + positionChecksum: 6.323604583740234, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.26141318678855896, 1.2181404829025269, 0.3854416310787201, + 1.010901689529419, 2.1883463859558105, 0.720491886138916, + ]), + tolerance: .float32) + } + } + + @Test("softsign") + func test_softsign() throws { + try withIntegrationState(seed: 40624) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softsign(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.004115648567676544, + minimum: -0.6314003467559814, + maximum: 0.7147443294525146, + absoluteSum: 5.524916648864746, + positionChecksum: 3.1572577158610025, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.6314003467559814, 0.5352221131324768, -0.5489830374717712, + -0.627332329750061, 0.7147443294525146, 0.6036931872367859, + ]), + tolerance: .float32) + } + } + + @Test("softshrink") + func test_softshrink() throws { + try withIntegrationState(seed: 92275) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softshrink(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.22316786646842957, + minimum: -1.119591474533081, + maximum: 0.3092571496963501, + absoluteSum: 3.999713897705078, + positionChecksum: 1.6630719502766926, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.3092571496963501, -0.43543922901153564, -0.818618893623352, + -0.6015747785568237, 0.041445255279541016, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("softshrink/lambda") + func test_softshrink_lambda() throws { + try withIntegrationState(seed: 6070) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softshrink(x, lambda: 0.2) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.05193863809108734, + minimum: -0.8008320927619934, + maximum: 1.2317386865615845, + absoluteSum: 6.1652302742004395, + positionChecksum: 3.3861312866210938, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.768548309803009, -0.7094013690948486, -0.5380290746688843, + -0.8008320927619934, 0.23463128507137299, 1.2317386865615845, + ]), + tolerance: .float32) + } + } + + @Test("softmin") + func test_softmin() throws { + try withIntegrationState(seed: 54376) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softmin(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3333333432674408, + minimum: 0.08368699252605438, + maximum: 0.767266571521759, + absoluteSum: 4.0, + positionChecksum: 2.2911632855733237, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.08813121169805527, 0.7124828100204468, 0.13925348222255707, + 0.4346010982990265, 0.08368699252605438, 0.5157842040061951, + ]), + tolerance: .float32) + } + } + + @Test("softmin/axis") + func test_softmin_axis() throws { + try withIntegrationState(seed: 147) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.softmin(x, axis: 0) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.25, + minimum: 0.055289801210165024, + maximum: 0.43014487624168396, + absoluteSum: 3.0, + positionChecksum: 1.5460936228434246, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.4214111864566803, 0.2248254418373108, 0.15415097773075104, + 0.19324195384979248, 0.33697277307510376, 0.18245795369148254, + ]), + tolerance: .float32) + } + } + + @Test("logSoftmax") + func test_logSoftmax() throws { + try withIntegrationState(seed: 36714) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.logSoftmax(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -1.3329758644104004, + minimum: -2.9164342880249023, + maximum: -0.40699923038482666, + absoluteSum: 15.995709419250488, + positionChecksum: 8.133668263753256, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9382022023200989, -2.536550998687744, -1.00897216796875, + -2.07830810546875, -1.2845144271850586, -0.8200601935386658, + ]), + tolerance: .float32) + } + } + + @Test("logSoftmax/axis") + func test_logSoftmax_axis() throws { + try withIntegrationState(seed: 19083) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.logSoftmax(x, axis: 0) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -2.090636730194092, + minimum: -5.523003101348877, + maximum: -0.09577543288469315, + absoluteSum: 25.08763885498047, + positionChecksum: 13.576393127441406, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9566798806190491, -3.1437976360321045, -1.5162262916564941, + -0.5962467789649963, -2.607832908630371, -0.09577543288469315, + ]), + tolerance: .float32) + } + } + + @Test("logSigmoid") + func test_logSigmoid() throws { + try withIntegrationState(seed: 13659) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.logSigmoid(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.6867945194244385, + minimum: -1.328415870666504, + maximum: -0.23934398591518402, + absoluteSum: 8.241534233093262, + positionChecksum: 4.763319969177246, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.44154447317123413, -1.2001237869262695, -0.7010847926139832, + -0.5049616694450378, -0.43374401330947876, -1.30607008934021, + ]), + tolerance: .float32) + } + } + + @Test("gelu") + func test_gelu() throws { + try withIntegrationState(seed: 56993) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.gelu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.23518109321594238, + minimum: -0.15614983439445496, + maximum: 1.5465519428253174, + absoluteSum: 4.124262809753418, + positionChecksum: 2.6090475718180337, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.29401835799217224, -0.015457117930054665, 1.0918797254562378, + 0.12314192205667496, -0.02627631649374962, 1.5465519428253174, + ]), + tolerance: .float32) + } + } + + @Test("geluApproximate") + func test_geluApproximate() throws { + try withIntegrationState(seed: 59006) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.geluApproximate(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.34252291917800903, + minimum: -0.14033935964107513, + maximum: 1.349687933921814, + absoluteSum: 4.6187591552734375, + positionChecksum: 2.154048442840576, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.14033935964107513, 1.349687933921814, -0.08664838969707489, + 0.11771057546138763, -0.02725459821522236, 0.6522086262702942, + ]), + tolerance: .float32) + } + } + + @Test("geluFastApproximate") + func test_geluFastApproximate() throws { + try withIntegrationState(seed: 45811) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.geluFastApproximate(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.7524160742759705, + minimum: -0.15481270849704742, + maximum: 3.242194414138794, + absoluteSum: 9.772090911865234, + positionChecksum: 4.8879391352335615, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.75855553150177, -0.14043818414211273, 0.5213433504104614, + -0.15481270849704742, 0.8235732316970825, -0.07629833370447159, + ]), + tolerance: .float32) + } + } + + @Test("glu") + func test_glu() throws { + try withIntegrationState(seed: 49808) { + let x = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.glu(x) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: -0.15694786608219147, + minimum: -1.4327794313430786, + maximum: 0.9524335861206055, + absoluteSum: 7.32803201675415, + positionChecksum: 4.967343807220459, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.06294449418783188, 0.1666058599948883, 0.3972969651222229, + -1.0184520483016968, -0.4588458836078644, -1.4327794313430786, + ]), + tolerance: .float32) + } + } + + @Test("glu/axis") + func test_glu_axis() throws { + try withIntegrationState(seed: 29911) { + let x = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.glu(x, axis: 0) + expectSummary( + result, + ArraySummary( + shape: [2, 8], + dtype: .float32, + mean: 0.322890043258667, + minimum: -0.396503210067749, + maximum: 1.6453715562820435, + absoluteSum: 7.251438617706299, + positionChecksum: 4.466485977172852, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.4837617576122284, 0.14164136350154877, -0.20633184909820557, + 0.618834376335144, -0.396503210067749, 0.431894451379776, + ]), + tolerance: .float32) + } + } + + @Test("step") + func test_step() throws { + try withIntegrationState(seed: 79720) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.step(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 6.0, + positionChecksum: 3.6666666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 0.0, 0.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("step/threshold") + func test_step_threshold() throws { + try withIntegrationState(seed: 47231) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.step(x, threshold: 0.5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.3333333333333335, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("hardSwish") + func test_hardSwish() throws { + try withIntegrationState(seed: 54778) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.hardSwish(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.6743363738059998, + minimum: -0.36365342140197754, + maximum: 2.4174962043762207, + absoluteSum: 10.894828796386719, + positionChecksum: 5.196015357971191, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.9947257041931152, 1.401398777961731, 1.191110610961914, + -0.35835549235343933, 0.4512682259082794, -0.31963229179382324, + ]), + tolerance: .float32) + } + } + + @Test("hardTanH") + func test_hardTanH() throws { + try withIntegrationState(seed: 64404) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.hardTanH(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.1564551740884781, + minimum: -1.0, + maximum: 1.0, + absoluteSum: 7.426859378814697, + positionChecksum: 3.8016363779703775, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0, 1.0, -0.22063256800174713, -0.4806191623210907, 0.5961291193962097, + -1.0, + ]), + tolerance: .float32) + } + } + + @Test("hardTanH/range") + func test_hardTanH_range() throws { + try withIntegrationState(seed: 12317) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.hardTanH(x, min: -0.5, max: 0.5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.24953168630599976, + minimum: -0.5, + maximum: 0.5, + absoluteSum: 5.2974853515625, + positionChecksum: 2.741623560587565, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-0.5, 0.5, 0.38924914598464966, 0.5, 0.4365149736404419, 0.5]), + tolerance: .float32) + } + } + + @Test("hardShrink") + func test_hardShrink() throws { + try withIntegrationState(seed: 98167) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.hardShrink(x) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.34890827536582947, + minimum: -1.311409831047058, + maximum: 1.3795651197433472, + absoluteSum: 6.809719085693359, + positionChecksum: 3.9342784881591797, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 1.3795651197433472, 0.0, 1.1961815357208252]), + tolerance: .float32) + } + } + + @Test("hardShrink/lambda") + func test_hardShrink_lambda() throws { + try withIntegrationState(seed: 17487) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.hardShrink(x, lambda: 0.2) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1098986268043518, + minimum: -0.9246480464935303, + maximum: 0.6220173239707947, + absoluteSum: 3.96665096282959, + positionChecksum: 1.78379487991333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9246480464935303, 0.4655166268348694, 0.0, 0.5469300150871277, 0.0, + 0.4056254029273987, + ]), + tolerance: .float32) + } + } + + @Test("prelu") + func test_prelu() throws { + try withIntegrationState(seed: 57294) { + let x = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let alpha = MLXRandom.uniform(low: 0.1, high: 0.5, [3], dtype: .float32) + let result = MLXNN.prelu(x, alpha: alpha) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3869323432445526, + minimum: -0.9682585597038269, + maximum: 1.647214412689209, + absoluteSum: 9.230525970458984, + positionChecksum: 5.1564515431722, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.647214412689209, 1.027597427368164, -0.08731298893690109, + 0.215767964720726, -0.9682585597038269, 1.128369688987732, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedBinaryTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedBinaryTests.swift new file mode 100644 index 000000000..72978092f --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedBinaryTests.swift @@ -0,0 +1,1462 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 61 + +import Foundation +import MLX +import Testing + +@Suite("generated: Binary") +struct GeneratedBinaryTests { + + @Test("add") + func test_add() throws { + try withIntegrationState(seed: 31173) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.add(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.6918743848800659, + minimum: -4.831423282623291, + maximum: 1.7133017778396606, + absoluteSum: 15.499208450317383, + positionChecksum: 8.161814371744791, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.7133017778396606, 1.263370156288147, -0.09468770027160645, + -2.1408917903900146, -1.1286695003509521, -0.5542188882827759, + ]), + tolerance: .float32) + } + } + + @Test("subtract") + func test_subtract() throws { + try withIntegrationState(seed: 70937) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.subtract(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.1385997086763382, + minimum: -3.5849361419677734, + maximum: 2.4345836639404297, + absoluteSum: 16.477638244628906, + positionChecksum: 10.1959228515625, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.21980679035186768, 2.4345836639404297, -1.0636004209518433, + 0.7404163479804993, -3.5849361419677734, 2.1673178672790527, + ]), + tolerance: .float32) + } + } + + @Test("multiply") + func test_multiply() throws { + try withIntegrationState(seed: 763) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.multiply(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.15953750908374786, + minimum: -0.6936066150665283, + maximum: 1.4158040285110474, + absoluteSum: 3.8533198833465576, + positionChecksum: 1.7182307243347168, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.3082238733768463, 0.13711540400981903, 1.4158040285110474, + -0.004460824187844992, 0.039156656712293625, -0.10085374861955643, + ]), + tolerance: .float32) + } + } + + @Test("divide") + func test_divide() throws { + try withIntegrationState(seed: 52989) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.divide(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 13.941320419311523, + minimum: -3.5286707878112793, + maximum: 129.2261199951172, + absoluteSum: 190.4403076171875, + positionChecksum: 118.93863932291667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.3581501543521881, -2.9087166786193848, 0.7051337957382202, + -2.402435779571533, 0.1394377052783966, -0.26220136880874634, + ]), + tolerance: .float32) + } + } + + @Test("remainder") + func test_remainder() throws { + try withIntegrationState(seed: 90342) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.remainder(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.006593585014343262, + minimum: -0.6975749731063843, + maximum: 1.1502771377563477, + absoluteSum: 5.422736167907715, + positionChecksum: 3.33150323232015, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.02850806713104248, -0.6895581483840942, -0.6510635614395142, + 0.3228166699409485, 0.7109462022781372, 1.1502771377563477, + ]), + tolerance: .float32) + } + } + + @Test("maximum") + func test_maximum() throws { + try withIntegrationState(seed: 93290) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.maximum(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.325461745262146, + minimum: -0.5284785032272339, + maximum: 1.6095561981201172, + absoluteSum: 7.125810623168945, + positionChecksum: 3.293111801147461, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.08404814451932907, 1.6095561981201172, -0.43784892559051514, + -0.5284785032272339, 0.484352171421051, 0.12096934765577316, + ]), + tolerance: .float32) + } + } + + @Test("minimum") + func test_minimum() throws { + try withIntegrationState(seed: 90053) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.minimum(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.22598682343959808, + minimum: -1.4718194007873535, + maximum: 1.3531242609024048, + absoluteSum: 7.797120094299316, + positionChecksum: 3.343695322672526, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.8319122791290283, -1.4718194007873535, 0.052562080323696136, + -0.15580348670482635, -0.6957205533981323, 0.44752341508865356, + ]), + tolerance: .float32) + } + } + + @Test("logAddExp") + func test_logAddExp() throws { + try withIntegrationState(seed: 53415) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logAddExp(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.949323832988739, + minimum: -0.1343468427658081, + maximum: 1.801002025604248, + absoluteSum: 11.660578727722168, + positionChecksum: 6.266410827636719, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.801002025604248, 0.9974111318588257, 1.0150268077850342, + 1.2767889499664307, 0.6966230273246765, 1.4964959621429443, + ]), + tolerance: .float32) + } + } + + @Test("atan2") + func test_atan2() throws { + try withIntegrationState(seed: 42823) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.atan2(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.7183784246444702, + minimum: -2.9523682594299316, + maximum: 2.919264316558838, + absoluteSum: 23.06583023071289, + positionChecksum: 13.557057698567709, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.835385322570801, 0.07711521536111832, -2.8996195793151855, + -1.5876164436340332, 2.8714325428009033, -1.5721657276153564, + ]), + tolerance: .float32) + } + } + + @Test("floorDivide") + func test_floorDivide() throws { + try withIntegrationState(seed: 33967) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.floorDivide(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 2.9166667461395264, + minimum: -3.0, + maximum: 37.0, + absoluteSum: 51.0, + positionChecksum: 28.416666666666668, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 0.0, -3.0, 1.0, -3.0]), + tolerance: .float32) + } + } + + @Test("pow") + func test_pow() throws { + try withIntegrationState(seed: 1911) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let b = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.pow(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0408190488815308, + minimum: 0.028537485748529434, + maximum: 2.14806866645813, + absoluteSum: 12.489828109741211, + positionChecksum: 6.593130747477214, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.673996090888977, 0.39860963821411133, 1.0955862998962402, + 1.98665189743042, 0.6895277500152588, 0.877358078956604, + ]), + tolerance: .float32) + } + } + + @Test("equal") + func test_equal() throws { + try withIntegrationState(seed: 99685) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.equal(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.9166666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 1.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("notEqual") + func test_notEqual() throws { + try withIntegrationState(seed: 8645) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.notEqual(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.75, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 9.0, + positionChecksum: 4.25, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("less") + func test_less() throws { + try withIntegrationState(seed: 85782) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.less(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 6.0, + positionChecksum: 3.9166666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 0.0, 0.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("lessEqual") + func test_lessEqual() throws { + try withIntegrationState(seed: 77066) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.lessEqual(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5833333730697632, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 7.0, + positionChecksum: 4.333333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 1.0, 1.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("greater") + func test_greater() throws { + try withIntegrationState(seed: 90710) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.greater(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.9166666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 1.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("greaterEqual") + func test_greaterEqual() throws { + try withIntegrationState(seed: 85655) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.greaterEqual(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5833333730697632, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 7.0, + positionChecksum: 3.6666666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 0.0, 1.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("operator/add") + func test_operator_add() throws { + try withIntegrationState(seed: 77360) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a + b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.27254313230514526, + minimum: -1.8845726251602173, + maximum: 3.533946990966797, + absoluteSum: 11.555704116821289, + positionChecksum: 7.568548838297526, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0339982509613037, 0.43266522884368896, 0.7669238448143005, + 0.01592520996928215, 1.4233887195587158, 3.533946990966797, + ]), + tolerance: .float32) + } + } + + @Test("operator/add/scalarRHS") + func test_operator_add_scalarRHS() throws { + try withIntegrationState(seed: 96149) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a + 1.3 + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0331014394760132, + minimum: -0.7769765853881836, + maximum: 2.1503243446350098, + absoluteSum: 13.951169967651367, + positionChecksum: 8.28534189860026, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.6562983989715576, -0.7769765853881836, 1.057753086090088, + 0.4848131537437439, 0.8572401404380798, 2.1503243446350098, + ]), + tolerance: .float32) + } + } + + @Test("operator/add/scalarLHS") + func test_operator_add_scalarLHS() throws { + try withIntegrationState(seed: 12303) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = 0.5 + a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.53053218126297, + minimum: -0.9197744131088257, + maximum: 2.8648841381073, + absoluteSum: 10.549100875854492, + positionChecksum: 4.844537417093913, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.174974262714386, -0.5325756072998047, 1.936285376548767, + 0.10657274723052979, -0.9197744131088257, 0.23554858565330505, + ]), + tolerance: .float32) + } + } + + @Test("operator/subtract") + func test_operator_subtract() throws { + try withIntegrationState(seed: 60591) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a - b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.27403706312179565, + minimum: -0.7568876147270203, + maximum: 2.6599271297454834, + absoluteSum: 8.948841094970703, + positionChecksum: 3.8554064432779946, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.2547779083251953, 1.4729726314544678, 2.6599271297454834, + -0.31046876311302185, 0.3260926604270935, 0.31457972526550293, + ]), + tolerance: .float32) + } + } + + @Test("operator/subtract/scalarRHS") + func test_operator_subtract_scalarRHS() throws { + try withIntegrationState(seed: 28905) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a - 1.3 + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.9529048800468445, + minimum: -2.7124485969543457, + maximum: 0.6620986461639404, + absoluteSum: 12.759056091308594, + positionChecksum: 8.024199803670248, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.5383754968643188, -1.0335915088653564, -1.7431178092956543, + -2.7124485969543457, -1.635654091835022, -0.5485337972640991, + ]), + tolerance: .float32) + } + } + + @Test("operator/subtract/scalarLHS") + func test_operator_subtract_scalarLHS() throws { + try withIntegrationState(seed: 24115) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = 0.5 - a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.36531418561935425, + minimum: -1.047974944114685, + maximum: 2.1034936904907227, + absoluteSum: 9.299480438232422, + positionChecksum: 5.230242093404134, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.08025974035263062, 0.06461766362190247, -0.45031672716140747, + 0.8556313514709473, -0.4064660668373108, -1.047974944114685, + ]), + tolerance: .float32) + } + } + + @Test("operator/multiply") + func test_operator_multiply() throws { + try withIntegrationState(seed: 27085) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a * b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.4524042308330536, + minimum: -2.1970603466033936, + maximum: 0.457118958234787, + absoluteSum: 6.941291809082031, + positionChecksum: 3.4856224060058594, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.1970603466033936, 0.11475061625242233, 0.457118958234787, + -1.40084969997406, -0.4138683080673218, -1.1669623851776123, + ]), + tolerance: .float32) + } + } + + @Test("operator/multiply/scalarRHS") + func test_operator_multiply_scalarRHS() throws { + try withIntegrationState(seed: 63195) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a * 1.3 + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.30748769640922546, + minimum: -2.6694865226745605, + maximum: 1.9966527223587036, + absoluteSum: 14.375575065612793, + positionChecksum: 7.712694803873698, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.6694865226745605, -1.180345058441162, -0.7085902690887451, + -0.6169244647026062, 1.6982543468475342, -1.4327366352081299, + ]), + tolerance: .float32) + } + } + + @Test("operator/multiply/scalarLHS") + func test_operator_multiply_scalarLHS() throws { + try withIntegrationState(seed: 35649) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = 0.5 * a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.09928220510482788, + minimum: -0.9127927422523499, + maximum: 0.8308178186416626, + absoluteSum: 4.5851569175720215, + positionChecksum: 2.197216033935547, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9127927422523499, 0.20200762152671814, -0.110417440533638, + -0.5224007964134216, 0.8308178186416626, 0.3921858072280884, + ]), + tolerance: .float32) + } + } + + @Test("operator/divide") + func test_operator_divide() throws { + try withIntegrationState(seed: 60155) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a / b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.8858397006988525, + minimum: -4.235711097717285, + maximum: 17.544374465942383, + absoluteSum: 42.191917419433594, + positionChecksum: 24.856480916341145, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.19931618869304657, -0.5708990693092346, 0.011940893717110157, + 0.45701560378074646, -0.8684737086296082, -4.235711097717285, + ]), + tolerance: .float32) + } + } + + @Test("operator/divide/scalarRHS") + func test_operator_divide_scalarRHS() throws { + try withIntegrationState(seed: 40534) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a / 1.3 + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1341741383075714, + minimum: -0.765953540802002, + maximum: 1.2086530923843384, + absoluteSum: 4.993619918823242, + positionChecksum: 2.275047938028971, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.0909951850771904, 1.2086530923843384, 0.4493958055973053, + 0.1280558556318283, 0.15634208917617798, -0.08264032751321793, + ]), + tolerance: .float32) + } + } + + @Test("operator/divide/scalarLHS") + func test_operator_divide_scalarLHS() throws { + try withIntegrationState(seed: 24204) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = 0.5 / a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -10.988043785095215, + minimum: -135.88998413085938, + maximum: 4.149685382843018, + absoluteSum: 146.1427764892578, + positionChecksum: 62.945159912109375, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8811553716659546, 0.2845892608165741, -135.88998413085938, + 4.149685382843018, 0.4647601842880249, -0.7214249968528748, + ]), + tolerance: .float32) + } + } + + @Test("operator/remainder") + func test_operator_remainder() throws { + try withIntegrationState(seed: 71211) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a % b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.05340499430894852, + minimum: -0.9634010791778564, + maximum: 1.3972878456115723, + absoluteSum: 3.9187231063842773, + positionChecksum: 1.8389933904012044, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9634010791778564, 0.14139050245285034, -0.22529059648513794, + 0.050104208290576935, 0.3663478493690491, 0.0027291178703308105, + ]), + tolerance: .float32) + } + } + + @Test("operator/remainder/scalarRHS") + func test_operator_remainder_scalarRHS() throws { + try withIntegrationState(seed: 35295) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a % 1.3 + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5924494862556458, + minimum: 0.051908016204833984, + maximum: 1.2263563871383667, + absoluteSum: 7.10939359664917, + positionChecksum: 4.6820723215738935, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.293434739112854, 0.06582850962877274, 1.0644726753234863, + 0.3147084712982178, 1.1621785163879395, 0.9618409872055054, + ]), + tolerance: .float32) + } + } + + @Test("operator/remainder/scalarLHS") + func test_operator_remainder_scalarLHS() throws { + try withIntegrationState(seed: 2437) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = 0.5 % a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.20201009511947632, + minimum: -0.4174903631210327, + maximum: 0.5, + absoluteSum: 3.881199598312378, + positionChecksum: 2.143406550089518, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.07111084461212158, 0.5, -0.06144183874130249, 0.5, 0.5, + -0.4174903631210327, + ]), + tolerance: .float32) + } + } + + @Test("operator/pow") + func test_operator_pow() throws { + try withIntegrationState(seed: 55330) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let b = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = a ** b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.8067885637283325, + minimum: 0.041118938475847244, + maximum: 1.7501370906829834, + absoluteSum: 9.681462287902832, + positionChecksum: 4.114236831665039, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.7501370906829834, 1.371971845626831, 0.3343576490879059, + 0.29617589712142944, 0.8186249136924744, 0.041118938475847244, + ]), + tolerance: .float32) + } + } + + @Test("operator/negate") + func test_operator_negate() throws { + try withIntegrationState(seed: 27407) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = -a + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.36295005679130554, + minimum: -1.8532614707946777, + maximum: 0.9766664505004883, + absoluteSum: 9.706311225891113, + positionChecksum: 6.226847966512044, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.22385217249393463, 0.6678629517555237, -1.2412484884262085, + -1.682862401008606, -0.23273582756519318, -1.8532614707946777, + ]), + tolerance: .float32) + } + } + + @Test("operator/equal") + func test_operator_equal() throws { + try withIntegrationState(seed: 5151) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .== b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.3333333432674408, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.0, + positionChecksum: 1.8333333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("operator/notEqual") + func test_operator_notEqual() throws { + try withIntegrationState(seed: 22291) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .!= b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.6666666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 8.0, + positionChecksum: 3.8333333333333335, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 1.0, 1.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("operator/less") + func test_operator_less() throws { + try withIntegrationState(seed: 59592) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .< b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.25, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 2.25, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("operator/lessEqual") + func test_operator_lessEqual() throws { + try withIntegrationState(seed: 52583) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .<= b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5833333730697632, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 7.0, + positionChecksum: 3.75, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 0.0, 1.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("operator/greater") + func test_operator_greater() throws { + try withIntegrationState(seed: 67466) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .> b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 3.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("operator/greaterEqual") + func test_operator_greaterEqual() throws { + try withIntegrationState(seed: 84422) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = a .>= b + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 6.0, + positionChecksum: 3.1666666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 0.0, 1.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("logicalAnd") + func test_logicalAnd() throws { + try withIntegrationState(seed: 77847) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let b = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.logicalAnd(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.3333333432674408, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.0, + positionChecksum: 1.5833333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("logicalOr") + func test_logicalOr() throws { + try withIntegrationState(seed: 96553) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let b = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.logicalOr(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.75, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 9.0, + positionChecksum: 4.75, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 1.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("logicalXor") + func test_logicalXor() throws { + try withIntegrationState(seed: 20568) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let b = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.logicalXor(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 6.0, + positionChecksum: 3.3333333333333335, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("which") + func test_which() throws { + try withIntegrationState(seed: 39197) { + let mask = MLXRandom.bernoulli(0.5, [4, 3]) + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.which(mask, a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.08233775198459625, + minimum: -0.9875338673591614, + maximum: 2.0078234672546387, + absoluteSum: 10.938934326171875, + positionChecksum: 5.835131963094075, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.5632016062736511, 1.5790736675262451, 0.684048056602478, + 2.0078234672546387, 1.3224455118179321, -0.5925998687744141, + ]), + tolerance: .float32) + } + } + + @Test("bitwiseAnd") + func test_bitwiseAnd() throws { + try withIntegrationState(seed: 18363) { + let a = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let result = MLX.bitwiseAnd(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 3.8333334922790527, + minimum: 0.0, + maximum: 12.0, + absoluteSum: 46.0, + positionChecksum: 25.333333333333332, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [8.0, 5.0, 0.0, 10.0, 0.0, 12.0]), + tolerance: .exact) + } + } + + @Test("bitwiseOr") + func test_bitwiseOr() throws { + try withIntegrationState(seed: 78878) { + let a = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let result = MLX.bitwiseOr(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 10.166666984558105, + minimum: 6.0, + maximum: 15.0, + absoluteSum: 122.0, + positionChecksum: 63.166666666666664, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [11.0, 11.0, 6.0, 14.0, 8.0, 6.0]), + tolerance: .exact) + } + } + + @Test("bitwiseXOr") + func test_bitwiseXOr() throws { + try withIntegrationState(seed: 66230) { + let a = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let result = MLX.bitwiseXOr(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 7.0, + minimum: 0.0, + maximum: 14.0, + absoluteSum: 84.0, + positionChecksum: 39.916666666666664, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [12.0, 10.0, 11.0, 4.0, 4.0, 9.0]), + tolerance: .exact) + } + } + + @Test("bitwiseInvert") + func test_bitwiseInvert() throws { + try withIntegrationState(seed: 53813) { + let a = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let result = MLX.bitwiseInvert(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: -7.9166669845581055, + minimum: -16.0, + maximum: -1.0, + absoluteSum: 95.0, + positionChecksum: 61.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-8.0, -1.0, -4.0, -10.0, -16.0, -16.0]), + tolerance: .exact) + } + } + + @Test("leftShift") + func test_leftShift() throws { + try withIntegrationState(seed: 83785) { + let a = MLXRandom.randInt(low: 0, high: 16, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.leftShift(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 32.41666793823242, + minimum: 0.0, + maximum: 120.0, + absoluteSum: 389.0, + positionChecksum: 275.0833333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [16.0, 40.0, 9.0, 0.0, 16.0, 120.0]), + tolerance: .exact) + } + } + + @Test("rightShift") + func test_rightShift() throws { + try withIntegrationState(seed: 20116) { + let a = MLXRandom.randInt(low: 0, high: 1024, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.rightShift(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 284.5, + minimum: 24.0, + maximum: 856.0, + absoluteSum: 3414.0, + positionChecksum: 1920.3333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [201.0, 287.0, 512.0, 51.0, 224.0, 302.0]), + tolerance: .exact) + } + } + + @Test("matmul") + func test_matmul() throws { + try withIntegrationState(seed: 69847) { + let a = MLXRandom.normal([10, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([8, 13], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matmul(a, b) + expectSummary( + result, + ArraySummary( + shape: [10, 13], + dtype: .float32, + mean: 0.013026426546275616, + minimum: -7.829680442810059, + maximum: 5.517763614654541, + absoluteSum: 251.1970977783203, + positionChecksum: 140.54157151442308, + sampleIndices: [0, 26, 52, 77, 103, 129], + samples: [ + -1.5777220726013184, -2.814626455307007, -0.9406680464744568, + 1.512544870376587, 1.1238691806793213, 2.8315815925598145, + ]), + tolerance: .float32) + } + } + + @Test("matmul/batched") + func test_matmul_batched() throws { + try withIntegrationState(seed: 58020) { + let a = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([2, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.matmul(a, b) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 3], + dtype: .float32, + mean: 0.560451865196228, + minimum: -4.739752292633057, + maximum: 9.252123832702637, + absoluteSum: 51.17353057861328, + positionChecksum: 27.05833689371745, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 0.6251500844955444, 0.8671976327896118, 9.252123832702637, + 0.1872236281633377, 1.1555585861206055, -2.105651617050171, + ]), + tolerance: .float32) + } + } + + @Test("inner") + func test_inner() throws { + try withIntegrationState(seed: 36443) { + let a = MLXRandom.normal([12], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([12], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.inner(a, b) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -2.8012592792510986, + minimum: -2.8012592792510986, + maximum: -2.8012592792510986, + absoluteSum: 2.8012592792510986, + positionChecksum: 2.8012592792510986, + sampleIndices: [0], + samples: [-2.8012592792510986]), + tolerance: .float32) + } + } + + @Test("outer") + func test_outer() throws { + try withIntegrationState(seed: 99139) { + let a = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.outer(a, b) + expectSummary( + result, + ArraySummary( + shape: [5, 4], + dtype: .float32, + mean: 0.2968282103538513, + minimum: -1.7715340852737427, + maximum: 3.274460792541504, + absoluteSum: 17.70854377746582, + positionChecksum: 9.23790283203125, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [ + -0.5336867570877075, -1.4807971715927124, -0.4454834759235382, + 0.15995460748672485, 0.636084258556366, -0.22449640929698944, + ]), + tolerance: .float32) + } + } + + @Test("vecdot") + func test_vecdot() throws { + try withIntegrationState(seed: 57547) { + let a = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.vecdot(a, b) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.39991307258605957, + minimum: 0.07599055767059326, + maximum: 1.119821548461914, + absoluteSum: 1.5996522903442383, + positionChecksum: 1.2174617052078247, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.08565317094326019, 0.07599055767059326, 1.119821548461914, + 0.31818699836730957, + ]), + tolerance: .float32) + } + } + + @Test("tensordot") + func test_tensordot() throws { + try withIntegrationState(seed: 41209) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tensordot(a, b, axes: ([1, 2], [1, 0])) + expectSummary( + result, + ArraySummary( + shape: [2, 2], + dtype: .float32, + mean: 2.566420793533325, + minimum: -2.453749656677246, + maximum: 8.781932830810547, + absoluteSum: 15.173182487487793, + positionChecksum: 11.776962280273438, + sampleIndices: [0, 1, 2, 3], + samples: [ + -2.453749656677246, 2.286130666732788, 1.6513692140579224, + 8.781932830810547, + ]), + tolerance: .float32) + } + } + + @Test("addMM") + func test_addMM() throws { + try withIntegrationState(seed: 78623) { + let c = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let a = MLXRandom.normal([4, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.addMM(c, a, b, alpha: 0.5, beta: 2.0) + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .float32, + mean: 0.18949465453624725, + minimum: -2.853945732116699, + maximum: 4.330568313598633, + absoluteSum: 26.563777923583984, + positionChecksum: 16.041049194335937, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [ + -1.6629900932312012, 0.6326870918273926, -0.28035566210746765, + 4.330568313598633, 0.5440598726272583, -1.714925765991211, + ]), + tolerance: .float32) + } + } + + @Test("kron") + func test_kron() throws { + try withIntegrationState(seed: 42558) { + let a = MLXRandom.normal([2, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([3, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.kron(a, b) + expectSummary( + result, + ArraySummary( + shape: [6, 6], + dtype: .float32, + mean: -0.23324330151081085, + minimum: -1.7474167346954346, + maximum: 1.0404057502746582, + absoluteSum: 22.564327239990234, + positionChecksum: 11.563710530598959, + sampleIndices: [0, 7, 14, 21, 28, 35], + samples: [ + 0.4093790054321289, 0.27812692523002625, -1.7474167346954346, + -0.03940131887793541, -0.668895959854126, 0.8408368229866028, + ]), + tolerance: .float32) + } + } + + @Test("norm/frobeniusKeepDims") + func test_norm_frobeniusKeepDims() throws { + try withIntegrationState(seed: 80992) { + let a = MLXRandom.normal([3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 1], + dtype: .float32, + mean: 3.2592105865478516, + minimum: 3.2592105865478516, + maximum: 3.2592105865478516, + absoluteSum: 3.2592105865478516, + positionChecksum: 3.2592105865478516, + sampleIndices: [0], + samples: [3.2592105865478516]), + tolerance: .float32) + } + } + + @Test("addMM/newtonSchulzForm") + func test_addMM_newtonSchulzForm() throws { + try withIntegrationState(seed: 34985) { + let a = MLXRandom.normal([3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.addMM(-4.7750 * a, a, a, alpha: 2.0315, beta: 1.0) + expectSummary( + result, + ArraySummary( + shape: [3, 3], + dtype: .float32, + mean: -0.7809702157974243, + minimum: -4.185337066650391, + maximum: 4.572365760803223, + absoluteSum: 24.614940643310547, + positionChecksum: 14.121938069661459, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [ + -0.531244158744812, 4.220737934112549, -2.525264263153076, + -0.26388388872146606, -1.9338984489440918, -4.185337066650391, + ]), + tolerance: .float32) + } + } + + @Test("einsum") + func test_einsum() throws { + try withIntegrationState(seed: 27180) { + let a = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([5, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.einsum("ij,jk->ik", a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 6], + dtype: .float32, + mean: 0.5131218433380127, + minimum: -6.676657199859619, + maximum: 8.000497817993164, + absoluteSum: 71.23182678222656, + positionChecksum: 36.53941090901693, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 2.163822889328003, -0.3911140263080597, 1.3212900161743164, + 8.000497817993164, -1.8788937330245972, -1.330705165863037, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedConvolutionTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedConvolutionTests.swift new file mode 100644 index 000000000..078667b15 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedConvolutionTests.swift @@ -0,0 +1,422 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 16 + +import Foundation +import MLX +import Testing + +@Suite("generated: Convolution") +struct GeneratedConvolutionTests { + + @Test("convolve/full") + func test_convolve_full() throws { + try withIntegrationState(seed: 4719) { + let a = MLXRandom.normal([20], dtype: .float32, loc: 0.0, scale: 1.0) + let v = MLXRandom.normal([4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convolve(a, v) + expectSummary( + result, + ArraySummary( + shape: [23], + dtype: .float32, + mean: -0.010141642764210701, + minimum: -2.3259332180023193, + maximum: 1.7424554824829102, + absoluteSum: 19.981769561767578, + positionChecksum: 10.247398044752037, + sampleIndices: [0, 4, 9, 13, 18, 22], + samples: [ + -0.008841410279273987, 1.1080988645553589, 0.9621633291244507, + 1.7424554824829102, -0.30217692255973816, 0.002387824235484004, + ]), + tolerance: .float32) + } + } + + @Test("convolve/same") + func test_convolve_same() throws { + try withIntegrationState(seed: 68267) { + let a = MLXRandom.normal([20], dtype: .float32, loc: 0.0, scale: 1.0) + let v = MLXRandom.normal([4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convolve(a, v, mode: .same) + expectSummary( + result, + ArraySummary( + shape: [20], + dtype: .float32, + mean: 0.4013185203075409, + minimum: -3.1303353309631348, + maximum: 4.640427112579346, + absoluteSum: 29.873620986938477, + positionChecksum: 15.805455017089844, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [ + 0.08655143529176712, 0.08424151688814163, -3.1303353309631348, + 2.3075802326202393, -1.9558295011520386, -0.9330341219902039, + ]), + tolerance: .float32) + } + } + + @Test("convolve/valid") + func test_convolve_valid() throws { + try withIntegrationState(seed: 33328) { + let a = MLXRandom.normal([20], dtype: .float32, loc: 0.0, scale: 1.0) + let v = MLXRandom.normal([4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convolve(a, v, mode: .valid) + expectSummary( + result, + ArraySummary( + shape: [17], + dtype: .float32, + mean: 0.18561282753944397, + minimum: -3.5847108364105225, + maximum: 4.357660293579102, + absoluteSum: 23.426000595092773, + positionChecksum: 14.26570488424862, + sampleIndices: [0, 3, 6, 10, 13, 16], + samples: [ + 1.5556355714797974, 0.9753752946853638, 0.8864786624908447, + 0.2802148759365082, 2.622657299041748, 2.514362677175086e-05, + ]), + tolerance: .float32) + } + } + + @Test("conv1d") + func test_conv1d() throws { + try withIntegrationState(seed: 33875) { + let input = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv1d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 3], + dtype: .float32, + mean: -0.23364683985710144, + minimum: -7.760387897491455, + maximum: 6.223398685455322, + absoluteSum: 115.83242797851562, + positionChecksum: 60.60389709472656, + sampleIndices: [0, 9, 19, 28, 38, 47], + samples: [ + -2.8744194507598877, -2.188898801803589, -2.1922268867492676, + 3.848215341567993, 3.4207470417022705, -3.5831971168518066, + ]), + tolerance: .float32) + } + } + + @Test("conv1d/stridePadding") + func test_conv1d_stridePadding() throws { + try withIntegrationState(seed: 94979) { + let input = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv1d(input, weight, stride: 2, padding: 1) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: 0.45748910307884216, + minimum: -7.211019039154053, + maximum: 8.565747261047363, + absoluteSum: 72.53246307373047, + positionChecksum: 44.807047526041664, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.2506880760192871, -1.96531081199646, -0.8448853492736816, + -1.1593079566955566, -1.14354407787323, 4.385194778442383, + ]), + tolerance: .float32) + } + } + + @Test("conv1d/dilation") + func test_conv1d_dilation() throws { + try withIntegrationState(seed: 17243) { + let input = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv1d(input, weight, dilation: 2) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 3], + dtype: .float32, + mean: -0.8824627995491028, + minimum: -6.105772972106934, + maximum: 6.817114353179932, + absoluteSum: 88.25929260253906, + positionChecksum: 49.157989501953125, + sampleIndices: [0, 7, 14, 21, 28, 35], + samples: [ + -1.936726450920105, -1.7665845155715942, -0.4554101824760437, + -2.9440197944641113, 0.016120949760079384, 5.547804355621338, + ]), + tolerance: .float32) + } + } + + @Test("conv1d/groups") + func test_conv1d_groups() throws { + try withIntegrationState(seed: 44073) { + let input = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv1d(input, weight, groups: 2) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 4], + dtype: .float32, + mean: -0.29602718353271484, + minimum: -6.234485149383545, + maximum: 7.6336469650268555, + absoluteSum: 137.21255493164062, + positionChecksum: 70.91708374023438, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.6050371527671814, 2.1364376544952393, 1.7200886011123657, + -1.4107182025909424, -2.151890754699707, -5.122685432434082, + ]), + tolerance: .float32) + } + } + + @Test("conv2d") + func test_conv2d() throws { + try withIntegrationState(seed: 54064) { + let input = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv2d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 6, 4], + dtype: .float32, + mean: 0.11624614149332047, + minimum: -11.582426071166992, + maximum: 12.33521842956543, + absoluteSum: 929.9988403320312, + positionChecksum: 462.14453125, + sampleIndices: [0, 57, 115, 172, 230, 287], + samples: [ + -2.1473069190979004, 2.8220767974853516, 4.751214027404785, + -3.128023624420166, -0.647574782371521, -7.356082916259766, + ]), + tolerance: .float32) + } + } + + @Test("conv2d/stridePadding") + func test_conv2d_stridePadding() throws { + try withIntegrationState(seed: 57937) { + let input = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv2d(input, weight, stride: [2, 1], padding: [1, 0]) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 6, 4], + dtype: .float32, + mean: 0.16354301571846008, + minimum: -11.90718936920166, + maximum: 14.7648286819458, + absoluteSum: 820.2205200195312, + positionChecksum: 408.4796549479167, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + -6.419956207275391, 3.057616710662842, -3.531296968460083, + 7.2112555503845215, -1.983046531677246, 2.6352462768554688, + ]), + tolerance: .float32) + } + } + + @Test("conv3d") + func test_conv3d() throws { + try withIntegrationState(seed: 25745) { + let input = MLXRandom.normal([1, 4, 6, 6, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 2, 3, 3, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv3d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 4, 3], + dtype: .float32, + mean: -0.19438402354717255, + minimum: -21.18064308166504, + maximum: 13.98067569732666, + absoluteSum: 585.3264770507812, + positionChecksum: 298.1985134548611, + sampleIndices: [0, 29, 57, 86, 114, 143], + samples: [ + -1.989408254623413, -4.56206750869751, -4.980064392089844, + 0.4489351511001587, 1.2831811904907227, -0.33633241057395935, + ]), + tolerance: .float32) + } + } + + @Test("convTransposed1d") + func test_convTransposed1d() throws { + try withIntegrationState(seed: 56758) { + let input = MLXRandom.normal([2, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convTransposed1d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [2, 10, 3], + dtype: .float32, + mean: -0.08723472058773041, + minimum: -7.693572521209717, + maximum: 7.150848865509033, + absoluteSum: 127.65473937988281, + positionChecksum: 62.71107177734375, + sampleIndices: [0, 12, 24, 35, 47, 59], + samples: [ + 0.22783879935741425, -2.8541452884674072, 1.6900525093078613, + -1.3801987171173096, -0.30918294191360474, 0.5386923551559448, + ]), + tolerance: .float32) + } + } + + @Test("convTransposed1d/stride") + func test_convTransposed1d_stride() throws { + try withIntegrationState(seed: 7340) { + let input = MLXRandom.normal([2, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convTransposed1d( + input, weight, stride: 2, padding: 1, outputPadding: 1) + expectSummary( + result, + ArraySummary( + shape: [2, 16, 3], + dtype: .float32, + mean: 0.20248651504516602, + minimum: -6.641228199005127, + maximum: 7.307250499725342, + absoluteSum: 178.95211791992188, + positionChecksum: 95.32716878255208, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 2.1703476905822754, 0.1148005798459053, 1.0188913345336914, + 1.8733576536178589, -2.3983073234558105, 1.2071207761764526, + ]), + tolerance: .float32) + } + } + + @Test("convTransposed2d") + func test_convTransposed2d() throws { + try withIntegrationState(seed: 60757) { + let input = MLXRandom.normal([2, 6, 6, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convTransposed2d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 4], + dtype: .float32, + mean: -0.15208807587623596, + minimum: -13.422320365905762, + maximum: 12.339730262756348, + absoluteSum: 1462.8289794921875, + positionChecksum: 722.030029296875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -1.8898229598999023, 6.999209403991699, 3.9277143478393555, + -0.8446962237358093, 0.1216912716627121, -0.5072753429412842, + ]), + tolerance: .float32) + } + } + + @Test("convTransposed3d") + func test_convTransposed3d() throws { + try withIntegrationState(seed: 37364) { + let input = MLXRandom.normal([1, 4, 4, 4, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 2, 2, 2, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convTransposed3d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [1, 5, 5, 5, 3], + dtype: .float32, + mean: 0.038671448826789856, + minimum: -10.21535873413086, + maximum: 12.558921813964844, + absoluteSum: 831.7537841796875, + positionChecksum: 421.350625, + sampleIndices: [0, 75, 150, 224, 299, 374], + samples: [ + 1.1298370361328125, -0.9043198823928833, 1.8892734050750732, + -0.6862713098526001, 0.47907447814941406, -0.20345965027809143, + ]), + tolerance: .float32) + } + } + + @Test("convGeneral") + func test_convGeneral() throws { + try withIntegrationState(seed: 91113) { + let input = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convGeneral(input, weight, strides: 2, padding: 1) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: 0.563105046749115, + minimum: -9.814929008483887, + maximum: 15.443587303161621, + absoluteSum: 483.46710205078125, + positionChecksum: 256.685791015625, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 4.08807897567749, 0.49885979294776917, 5.748270511627197, + -0.7634090781211853, -2.1041526794433594, 2.326464891433716, + ]), + tolerance: .float32) + } + } + + @Test("convGeneral/flip") + func test_convGeneral_flip() throws { + try withIntegrationState(seed: 61813) { + let input = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([4, 3, 3, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convGeneral(input, weight, kernelDilation: 2, flip: true) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: 1.1288611888885498, + minimum: -10.111289978027344, + maximum: 13.187017440795898, + absoluteSum: 500.216796875, + positionChecksum: 253.00144958496094, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + -0.9194361567497253, 0.8226830959320068, 13.187017440795898, + -2.3197848796844482, -3.078174352645874, -1.5299072265625, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedDefaultsTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedDefaultsTests.swift new file mode 100644 index 000000000..b915c6b39 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedDefaultsTests.swift @@ -0,0 +1,1410 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 59 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: Defaults") +struct GeneratedDefaultsTests { + + @Test("nanToNum") + func test_nanToNum() throws { + // python replaces +/-inf with the dtype max/min unless told otherwise + try withIntegrationState(seed: 12047) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.nanToNum(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .float32, + mean: 0.0, + minimum: -3.4028234663852886e+38, + maximum: 3.4028234663852886e+38, + absoluteSum: Double.infinity, + positionChecksum: Double.infinity, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [1.0, -2.5, 3.4028234663852886e+38, -3.4028234663852886e+38, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("tensordot") + func test_tensordot() throws { + // axes defaults to 2 (as in numpy), not 1 + try withIntegrationState(seed: 43747) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([3, 4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tensordot(a, b) + expectSummary( + result, + ArraySummary( + shape: [2, 5], + dtype: .float32, + mean: -0.1383628100156784, + minimum: -6.713634967803955, + maximum: 6.5990705490112305, + absoluteSum: 30.51371955871582, + positionChecksum: 18.927197265625, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [ + 0.017482735216617584, -2.9454421997070312, 6.5990705490112305, + 0.03260062262415886, 0.44432538747787476, 2.5363638401031494, + ]), + tolerance: .float32) + } + } + + @Test("isClose/within") + func test_isClose_within() throws { + // 1e-6 apart: inside the default rtol 1e-5 / atol 1e-8 + try withIntegrationState(seed: 9065) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = a + 1e-6 + let result = MLX.isClose(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.8333333730697632, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 10.0, + positionChecksum: 5.333333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 1.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("isClose/outside") + func test_isClose_outside() throws { + try withIntegrationState(seed: 50854) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = a + 1e-3 + let result = MLX.isClose(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("allClose") + func test_allClose() throws { + try withIntegrationState(seed: 52631) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = a + 1e-6 + let result = MLX.allClose(a, b) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0], + samples: [0.0]), + tolerance: .exact) + } + } + + @Test("arrayEqual") + func test_arrayEqual() throws { + try withIntegrationState(seed: 12896) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.arrayEqual(a, a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0], + samples: [1.0]), + tolerance: .exact) + } + } + + @Test("diff") + func test_diff() throws { + try withIntegrationState(seed: 78249) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diff(a) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: 0.16116583347320557, + minimum: -1.240124225616455, + maximum: 1.5920885801315308, + absoluteSum: 7.024305820465088, + positionChecksum: 3.702437162399292, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + 0.5840742588043213, -1.240124225616455, -0.21392442286014557, + 1.5920885801315308, 0.6998119950294495, -0.25859102606773376, + ]), + tolerance: .float32) + } + } + + @Test("trace") + func test_trace() throws { + try withIntegrationState(seed: 60095) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.trace(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2.147444009780884, + minimum: 2.147444009780884, + maximum: 2.147444009780884, + absoluteSum: 2.147444009780884, + positionChecksum: 2.147444009780884, + sampleIndices: [0], + samples: [2.147444009780884]), + tolerance: .float32) + } + } + + @Test("diag") + func test_diag() throws { + try withIntegrationState(seed: 94808) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diag(a) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.4410656690597534, + minimum: -1.6250548362731934, + maximum: 0.6272233128547668, + absoluteSum: 3.3691916465759277, + positionChecksum: 1.827361822128296, + sampleIndices: [0, 1, 2, 3], + samples: [ + -1.6250548362731934, 0.17524121701717377, -0.9416722655296326, + 0.6272233128547668, + ]), + tolerance: .float32) + } + } + + @Test("diagonal") + func test_diagonal() throws { + try withIntegrationState(seed: 71392) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diagonal(a) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.10765820741653442, + minimum: -1.8548674583435059, + maximum: 1.9325915575027466, + absoluteSum: 4.152509689331055, + positionChecksum: 3.0919158458709717, + sampleIndices: [0, 1, 2, 3], + samples: [ + -0.0060711028054356575, 1.9325915575027466, 0.35897988080978394, + -1.8548674583435059, + ]), + tolerance: .float32) + } + } + + @Test("flipped") + func test_flipped() throws { + try withIntegrationState(seed: 2426) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.flipped(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3812035024166107, + minimum: -0.9314361214637756, + maximum: 2.043067693710327, + absoluteSum: 8.247232437133789, + positionChecksum: 5.1659698486328125, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.5396879315376282, -0.9314361214637756, 0.1392711102962494, + 1.5242443084716797, 2.043067693710327, 0.7740973234176636, + ]), + tolerance: .float32) + } + } + + @Test("round") + func test_round() throws { + try withIntegrationState(seed: 4970) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 10.0) + let result = MLX.round(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 2.3333334922790527, + minimum: -13.0, + maximum: 17.0, + absoluteSum: 88.0, + positionChecksum: 49.166666666666664, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-13.0, 9.0, -3.0, 17.0, 10.0, 12.0]), + tolerance: .float32) + } + } + + @Test("softmax") + func test_softmax() throws { + try withIntegrationState(seed: 65331) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.softmax(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.0833333358168602, + minimum: 0.019660301506519318, + maximum: 0.1961202472448349, + absoluteSum: 1.0, + positionChecksum: 0.5366625785827637, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.056501682847738266, 0.023698141798377037, 0.14301876723766327, + 0.019660301506519318, 0.09527640044689178, 0.10355699062347412, + ]), + tolerance: .float32) + } + } + + @Test("cumsum") + func test_cumsum() throws { + try withIntegrationState(seed: 57443) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumsum(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.0727943405508995, + minimum: -2.4656498432159424, + maximum: 1.888854742050171, + absoluteSum: 12.072671890258789, + positionChecksum: 7.7668304443359375, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.4540281295776367, 1.1351778507232666, 0.9642502069473267, + -2.4656498432159424, -0.36840271949768066, 1.888854742050171, + ]), + tolerance: .float32) + } + } + + @Test("cumprod") + func test_cumprod() throws { + try withIntegrationState(seed: 75562) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumprod(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.08276982605457306, + minimum: -0.6033686399459839, + maximum: 0.7661811709403992, + absoluteSum: 2.5519943237304688, + positionChecksum: 0.796099583307902, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.1745307743549347, -0.6033686399459839, 0.7661811709403992, + -0.0003647314733825624, -0.0005849742447026074, -0.0001368979865219444, + ]), + tolerance: .float32) + } + } + + @Test("cummax") + func test_cummax() throws { + try withIntegrationState(seed: 27207) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummax(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 1.3060426712036133, + minimum: 0.6872203350067139, + maximum: 1.429807186126709, + absoluteSum: 15.67251205444336, + positionChecksum: 9.108100255330404, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6872203350067139, 1.429807186126709, 1.429807186126709, 1.429807186126709, + 1.429807186126709, 1.429807186126709, + ]), + tolerance: .float32) + } + } + + @Test("cummin") + func test_cummin() throws { + try withIntegrationState(seed: 84926) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummin(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: -1.6473878622055054, + minimum: -2.1973390579223633, + maximum: -0.8774563074111938, + absoluteSum: 19.768653869628906, + positionChecksum: 12.632850646972656, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.8774563074111938, -0.8774563074111938, -0.8774563074111938, + -2.1973390579223633, -2.1973390579223633, -2.1973390579223633, + ]), + tolerance: .float32) + } + } + + @Test("logCumsumExp") + func test_logCumsumExp() throws { + try withIntegrationState(seed: 43574) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logCumsumExp(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 2.4215335845947266, + minimum: -0.11741061508655548, + maximum: 3.1633148193359375, + absoluteSum: 29.29322052001953, + positionChecksum: 18.49832534790039, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.11741061508655548, 2.4433906078338623, 2.7169978618621826, + 2.9348068237304688, 3.073072671890259, 3.1633148193359375, + ]), + tolerance: .float32) + } + } + + @Test("median") + func test_median() throws { + try withIntegrationState(seed: 92274) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.median(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.07507183402776718, + minimum: -0.07507183402776718, + maximum: -0.07507183402776718, + absoluteSum: 0.07507183402776718, + positionChecksum: 0.07507183402776718, + sampleIndices: [0], + samples: [-0.07507183402776718]), + tolerance: .float32) + } + } + + @Test("variance") + func test_variance() throws { + try withIntegrationState(seed: 5488) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.variance(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.4012840986251831, + minimum: 0.4012840986251831, + maximum: 0.4012840986251831, + absoluteSum: 0.4012840986251831, + positionChecksum: 0.4012840986251831, + sampleIndices: [0], + samples: [0.4012840986251831]), + tolerance: .float32) + } + } + + @Test("std") + func test_std() throws { + try withIntegrationState(seed: 52117) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.std(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.7837796211242676, + minimum: 0.7837796211242676, + maximum: 0.7837796211242676, + absoluteSum: 0.7837796211242676, + positionChecksum: 0.7837796211242676, + sampleIndices: [0], + samples: [0.7837796211242676]), + tolerance: .float32) + } + } + + @Test("hadamardTransform") + func test_hadamardTransform() throws { + try withIntegrationState(seed: 17021) { + let a = MLXRandom.normal([4, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.hadamardTransform(a) + expectSummary( + result, + ArraySummary( + shape: [4, 16], + dtype: .float32, + mean: -0.1470964550971985, + minimum: -2.6109731197357178, + maximum: 1.7692186832427979, + absoluteSum: 53.5502815246582, + positionChecksum: 25.162355422973633, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 1.0269194841384888, -1.6071994304656982, -1.6351640224456787, + -1.1893904209136963, 1.2207064628601074, 1.7337404489517212, + ]), + tolerance: .float32) + } + } + + @Test("linspace") + func test_linspace() throws { + try withIntegrationState(seed: 63792) { + let result = MLX.linspace(0.0, 1.0) + expectSummary( + result, + ArraySummary( + shape: [50], + dtype: .float32, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 25.0, + positionChecksum: 17.0, + sampleIndices: [0, 10, 20, 29, 39, 49], + samples: [ + 0.0, 0.20408163964748383, 0.40816327929496765, 0.5918367505073547, + 0.795918345451355, 1.0, + ]), + tolerance: .float32) + } + } + + @Test("eye") + func test_eye() throws { + try withIntegrationState(seed: 94476) { + let result = MLX.eye(4) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.25, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.0, + positionChecksum: 2.125, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [1.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .float32) + } + } + + @Test("tri") + func test_tri() throws { + // python requires m and k; Swift defaults them to n and 0 + try withIntegrationState(seed: 81451) { + let result = MLX.tri(4) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.625, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 10.0, + positionChecksum: 6.25, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [1.0, 0.0, 0.0, 1.0, 1.0, 1.0]), + tolerance: .float32) + } + } + + @Test("convolve") + func test_convolve() throws { + try withIntegrationState(seed: 62456) { + let a = MLXRandom.normal([20], dtype: .float32, loc: 0.0, scale: 1.0) + let v = MLXRandom.normal([4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.convolve(a, v) + expectSummary( + result, + ArraySummary( + shape: [23], + dtype: .float32, + mean: 0.08946806937456131, + minimum: -1.803636908531189, + maximum: 1.8454562425613403, + absoluteSum: 19.04773712158203, + positionChecksum: 8.943277110224185, + sampleIndices: [0, 4, 9, 13, 18, 22], + samples: [ + 1.0646476745605469, 1.7893160581588745, -0.8670757412910461, + -0.4499398171901703, 1.8454562425613403, -0.021289121359586716, + ]), + tolerance: .float32) + } + } + + @Test("conv1d") + func test_conv1d() throws { + try withIntegrationState(seed: 88096) { + let input = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([3, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.conv1d(input, weight) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 3], + dtype: .float32, + mean: 0.5933355093002319, + minimum: -7.674775123596191, + maximum: 7.810962677001953, + absoluteSum: 154.97207641601562, + positionChecksum: 74.69058227539062, + sampleIndices: [0, 9, 19, 28, 38, 47], + samples: [ + 1.575302243232727, -0.19398178160190582, 4.896422863006592, + 7.810962677001953, 2.7325944900512695, 3.176520586013794, + ]), + tolerance: .float32) + } + } + + @Test("quantized") + func test_quantized() throws { + // group size and bit count default in the C++ core (64 / 4) + try withIntegrationState(seed: 5234) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w).wq + expectSummary( + result, + ArraySummary( + shape: [64, 16], + dtype: .uint32, + mean: 2255552512.0, + minimum: 10032196.0, + maximum: 4293439744.0, + absoluteSum: 2309685772288.0, + positionChecksum: 1166251655168.0, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + 1967645312.0, 1215345152.0, 1610332928.0, 788708928.0, 3749942784.0, + 3667559680.0, + ]), + tolerance: .exact) + } + } + + @Test("dequantized") + func test_dequantized() throws { + try withIntegrationState(seed: 20669) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let q = MLX.quantized(w) + let result = MLX.dequantized(q.wq, scales: q.scales, biases: q.biases) + expectSummary( + result, + ArraySummary( + shape: [64, 128], + dtype: .float32, + mean: -0.0056404173374176025, + minimum: -0.9999048709869385, + maximum: 0.9997067451477051, + absoluteSum: 4079.2216796875, + positionChecksum: 2048.1728515625, + sampleIndices: [0, 1638, 3276, 4915, 6553, 8191], + samples: [ + 0.7425611615180969, -0.24511511623859406, -0.9868167042732239, + -0.12466387450695038, -0.6139729022979736, -0.24921728670597076, + ]), + tolerance: .float32) + } + } + + @Test("quantizedMM") + func test_quantizedMM() throws { + try withIntegrationState(seed: 17188) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w) + let result = MLX.quantizedMM(x, q.wq, scales: q.scales, biases: q.biases) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: -0.008240580558776855, + minimum: -17.94013214111328, + maximum: 18.256729125976562, + absoluteSum: 2450.4580078125, + positionChecksum: 1212.8326416015625, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 4.712850093841553, 6.662078857421875, 6.3404364585876465, 5.768905162811279, + -6.76831579208374, -5.105240821838379, + ]), + tolerance: .float32) + } + } + + @Test("norm") + func test_norm() throws { + try withIntegrationState(seed: 51906) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2.6667308807373047, + minimum: 2.6667308807373047, + maximum: 2.6667308807373047, + absoluteSum: 2.6667308807373047, + positionChecksum: 2.6667308807373047, + sampleIndices: [0], + samples: [2.6667308807373047]), + tolerance: .float32) + } + } + + @Test("cholesky") + func test_cholesky() throws { + // upper defaults to false on both sides + try withIntegrationState(seed: 65482) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.cholesky(spd, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.6790149211883545, + minimum: -1.1485399007797241, + maximum: 3.0164363384246826, + absoluteSum: 13.845975875854492, + positionChecksum: 8.125425338745117, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 3.0164363384246826, 0.0, 0.0, -0.2918112277984619, 0.553765058517456, + 2.652482271194458, + ]), + tolerance: .float32) + } + } + + @Test("triInv") + func test_triInv() throws { + try withIntegrationState(seed: 32119) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let lower = MLX.tril(spd) + let result = MLX.triInv(lower, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.021889742463827133, + minimum: -0.040400367230176926, + maximum: 0.14794215559959412, + absoluteSum: 0.6479233503341675, + positionChecksum: 0.3984556496143341, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.09358751028776169, 0.0, 0.0, -0.03134113550186157, 0.005859235301613808, + 0.14794215559959412, + ]), + tolerance: .float32) + } + } + + @Test("solveTriangular") + func test_solveTriangular() throws { + try withIntegrationState(seed: 82366) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let lower = MLX.tril(spd) + let b = MLXRandom.normal([4, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.solveTriangular(lower, b, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: 0.007373856380581856, + minimum: -0.15618708729743958, + maximum: 0.1315477341413498, + absoluteSum: 0.5120813250541687, + positionChecksum: 0.23803149163722992, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + -0.15618708729743958, -0.065582275390625, 0.1315477341413498, + 0.03207559138536453, 0.0042139156721532345, 0.09015661478042603, + ]), + tolerance: .float32) + } + } + + @Test("fft") + func test_fft() throws { + try withIntegrationState(seed: 13721) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: 0.5256693959236145, + minimum: -20.80100440979004, + maximum: 19.87451171875, + absoluteSum: 795.9764404296875, + positionChecksum: 385.171640625, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -1.5068063735961914, -8.961091041564941, -7.156008243560791, + 8.837318420410156, 0.48036813735961914, 0.8878612518310547, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -0.07993728667497635, + minimum: -23.762022018432617, + maximum: 17.19780921936035, + absoluteSum: 733.4420776367188, + positionChecksum: 353.8171875, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + 8.099128723144531, -16.915706634521484, 0.5834121704101562, + 5.633106231689453, 4.721193313598633, -21.837533950805664, + ]), + tolerance: .float32) + } + } + + @Test("fftfreq") + func test_fftfreq() throws { + try withIntegrationState(seed: 87337) { + let result = MLX.fftfreq(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: -0.03125, + minimum: -0.5, + maximum: 0.4375, + absoluteSum: 4.0, + positionChecksum: 2.25, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [0.0, 0.1875, 0.375, -0.4375, -0.25, -0.0625]), + tolerance: .float32) + } + } + + @Test("gumbel") + func test_gumbel() throws { + try withIntegrationState(seed: 53461) { + let result = MLXRandom.gumbel() + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.186912178993225, + minimum: 1.186912178993225, + maximum: 1.186912178993225, + absoluteSum: 1.186912178993225, + positionChecksum: 1.186912178993225, + sampleIndices: [0], + samples: [1.186912178993225]), + tolerance: .float32) + } + } + + @Test("uniform") + func test_uniform() throws { + try withIntegrationState(seed: 83902) { + let result = MLXRandom.uniform() + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.9644234776496887, + minimum: 0.9644234776496887, + maximum: 0.9644234776496887, + absoluteSum: 0.9644234776496887, + positionChecksum: 0.9644234776496887, + sampleIndices: [0], + samples: [0.9644234776496887]), + tolerance: .float32) + } + } + + @Test("normal") + func test_normal() throws { + try withIntegrationState(seed: 67758) { + let result = MLXRandom.normal() + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.3701243996620178, + minimum: -0.3701243996620178, + maximum: -0.3701243996620178, + absoluteSum: 0.3701243996620178, + positionChecksum: 0.3701243996620178, + sampleIndices: [0], + samples: [-0.3701243996620178]), + tolerance: .float32) + } + } + + @Test("bernoulli") + func test_bernoulli() throws { + try withIntegrationState(seed: 34396) { + let result = MLXRandom.bernoulli() + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0], + samples: [1.0]), + tolerance: .exact) + } + } + + @Test("sorted/flat") + func test_sorted_flat() throws { + try withIntegrationState(seed: 45340) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sorted(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.7008785009384155, + minimum: -0.6916800737380981, + maximum: 1.6399438381195068, + absoluteSum: 9.793901443481445, + positionChecksum: 6.6583404541015625, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.6916800737380981, 0.2463919222354889, 0.49870675802230835, + 0.9339690804481506, 1.1501116752624512, 1.6399438381195068, + ]), + tolerance: .float32) + } + } + + @Test("argSort/flat") + func test_argSort_flat() throws { + try withIntegrationState(seed: 29628) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argSort(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .uint32, + mean: 5.5, + minimum: 0.0, + maximum: 11.0, + absoluteSum: 66.0, + positionChecksum: 31.833333333333332, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [9.0, 6.0, 11.0, 1.0, 0.0, 4.0]), + tolerance: .exact) + } + } + + @Test("partitioned/flat") + func test_partitioned_flat() throws { + try withIntegrationState(seed: 94638) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.partitioned(a, kth: 3) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .float32, + mean: -0.2733670473098755, + minimum: -1.770876169204712, + maximum: 1.3596465587615967, + absoluteSum: 10.031713485717773, + positionChecksum: 4.404526074727376, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.770876169204712, -1.5603525638580322, -0.5617737770080566, + 0.21052542328834534, 0.6110134124755859, 1.3596465587615967, + ]), + tolerance: .float32) + } + } + + @Test("argPartition/flat") + func test_argPartition_flat() throws { + try withIntegrationState(seed: 12369) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argPartition(a, kth: 3) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .uint32, + mean: 5.5, + minimum: 0.0, + maximum: 11.0, + absoluteSum: 66.0, + positionChecksum: 35.166666666666664, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [11.0, 9.0, 4.0, 7.0, 10.0, 8.0]), + tolerance: .exact) + } + } + + @Test("top/flat") + func test_top_flat() throws { + try withIntegrationState(seed: 32885) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.top(a, k: 3) + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .float32, + mean: 0.8019466400146484, + minimum: 0.5089910626411438, + maximum: 1.005195140838623, + absoluteSum: 2.4058399200439453, + positionChecksum: 1.7692945798238118, + sampleIndices: [0, 1, 2], + samples: [0.5089910626411438, 0.8916536569595337, 1.005195140838623]), + tolerance: .float32) + } + } + + @Test("roll/flat") + func test_roll_flat() throws { + try withIntegrationState(seed: 90193) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.roll(a, shift: 2) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5598515272140503, + minimum: -1.5043145418167114, + maximum: 2.463324546813965, + absoluteSum: 10.23849868774414, + positionChecksum: 4.870436032613118, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.2558262050151825, 1.3071829080581665, 0.6110653281211853, + 1.2561404705047607, 0.022526530548930168, 0.05659160390496254, + ]), + tolerance: .float32) + } + } + + @Test("takeAlong/flat") + func test_takeAlong_flat() throws { + try withIntegrationState(seed: 42978) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 12, [4], type: Int32.self) + let result = MLX.takeAlong(a, indices) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.19436131417751312, + minimum: -0.8136857151985168, + maximum: 0.2381412237882614, + absoluteSum: 1.560187578201294, + positionChecksum: 1.216184139251709, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.2381412237882614, 0.1532299816608429, -0.35513073205947876, + -0.8136857151985168, + ]), + tolerance: .float32) + } + } + + @Test("repeated/flat") + func test_repeated_flat() throws { + try withIntegrationState(seed: 68684) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.repeated(a, count: 2) + expectSummary( + result, + ArraySummary( + shape: [24], + dtype: .float32, + mean: -0.0766555666923523, + minimum: -0.9837857484817505, + maximum: 0.9884175062179565, + absoluteSum: 12.13071060180664, + positionChecksum: 6.469824473063151, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.3888154923915863, 0.35114219784736633, -0.4268342852592468, + -0.750239372253418, -0.3761099576950073, 0.05179966613650322, + ]), + tolerance: .float32) + } + } + + @Test("crossEntropy") + func test_crossEntropy() throws { + try withIntegrationState(seed: 58839) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.crossEntropy(logits: logits, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.5683512687683105, + minimum: 0.43761518597602844, + maximum: 3.6751675605773926, + absoluteSum: 6.273405075073242, + positionChecksum: 4.924343585968018, + sampleIndices: [0, 1, 2, 3], + samples: [ + 1.1801968812942505, 0.43761518597602844, 0.9804257154464722, + 3.6751675605773926, + ]), + tolerance: .float32) + } + } + + @Test("binaryCrossEntropy") + func test_binaryCrossEntropy() throws { + try withIntegrationState(seed: 37138) { + let logits = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLXNN.binaryCrossEntropy(logits: logits, targets: targets.asType(.float32)) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.8657888174057007, + minimum: 0.8657888174057007, + maximum: 0.8657888174057007, + absoluteSum: 0.8657888174057007, + positionChecksum: 0.8657888174057007, + sampleIndices: [0], + samples: [0.8657888174057007]), + tolerance: .float32) + } + } + + @Test("l1Loss") + func test_l1Loss() throws { + try withIntegrationState(seed: 94773) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.l1Loss(predictions: predictions, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.5470814108848572, + minimum: 0.5470814108848572, + maximum: 0.5470814108848572, + absoluteSum: 0.5470814108848572, + positionChecksum: 0.5470814108848572, + sampleIndices: [0], + samples: [0.5470814108848572]), + tolerance: .float32) + } + } + + @Test("mseLoss") + func test_mseLoss() throws { + try withIntegrationState(seed: 57835) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.mseLoss(predictions: predictions, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.724809169769287, + minimum: 1.724809169769287, + maximum: 1.724809169769287, + absoluteSum: 1.724809169769287, + positionChecksum: 1.724809169769287, + sampleIndices: [0], + samples: [1.724809169769287]), + tolerance: .float32) + } + } + + @Test("smoothL1Loss") + func test_smoothL1Loss() throws { + try withIntegrationState(seed: 44045) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.smoothL1Loss(predictions: predictions, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.5390201807022095, + minimum: 0.5390201807022095, + maximum: 0.5390201807022095, + absoluteSum: 0.5390201807022095, + positionChecksum: 0.5390201807022095, + sampleIndices: [0], + samples: [0.5390201807022095]), + tolerance: .float32) + } + } + + @Test("huberLoss") + func test_huberLoss() throws { + try withIntegrationState(seed: 20289) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.huberLoss(inputs: inputs, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.933381199836731, + minimum: 0.017336202785372734, + maximum: 2.534982681274414, + absoluteSum: 11.200573921203613, + positionChecksum: 5.508398691813151, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6122961044311523, 0.14958451688289642, 2.534982681274414, + 0.14333206415176392, 1.719494104385376, 0.608657717704773, + ]), + tolerance: .float32) + } + } + + @Test("logCoshLoss") + func test_logCoshLoss() throws { + try withIntegrationState(seed: 65708) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.logCoshLoss(inputs: inputs, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.6543585062026978, + minimum: 0.025879621505737305, + maximum: 1.854630708694458, + absoluteSum: 7.852302074432373, + positionChecksum: 4.758116404215495, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.09128165245056152, 1.542405366897583, 0.11331796646118164, + 0.13976502418518066, 1.8090453147888184, 0.6400725841522217, + ]), + tolerance: .float32) + } + } + + @Test("hingeLoss") + func test_hingeLoss() throws { + try withIntegrationState(seed: 32168) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let signs = MLXRandom.bernoulli(0.5, [4, 3]) + let targets = MLX.which(signs, 1.0, -1.0) + let result = MLXNN.hingeLoss(inputs: inputs, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.990768313407898, + minimum: 0.0, + maximum: 2.4934892654418945, + absoluteSum: 11.889219284057617, + positionChecksum: 6.786125183105469, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.3253191709518433, 1.9193181991577148, 0.0, 0.9395096898078918, + 0.37627631425857544, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("nllLoss") + func test_nllLoss() throws { + try withIntegrationState(seed: 78313) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let inputs = MLX.log(MLX.softmax(logits, axis: -1)) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.nllLoss(inputs: inputs, targets: targets) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.7705587148666382, + minimum: 0.7432368993759155, + maximum: 3.2677173614501953, + absoluteSum: 7.082234859466553, + positionChecksum: 5.156996726989746, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.8151545524597168, 2.2561261653900146, 0.7432368993759155, + 3.2677173614501953, + ]), + tolerance: .float32) + } + } + + @Test("tripletLoss") + func test_tripletLoss() throws { + try withIntegrationState(seed: 18126) { + let anchors = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let positives = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let negatives = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.tripletLoss( + anchors: anchors, positives: positives, negatives: negatives) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.5873920917510986, + minimum: 1.2556252479553223, + maximum: 1.9852931499481201, + absoluteSum: 6.3495683670043945, + positionChecksum: 3.7734334468841553, + sampleIndices: [0, 1, 2, 3], + samples: [ + 1.9852931499481201, 1.5465178489685059, 1.2556252479553223, + 1.5621321201324463, + ]), + tolerance: .float32) + } + } + + @Test("cosineSimilarityLoss") + func test_cosineSimilarityLoss() throws { + try withIntegrationState(seed: 76691) { + let x1 = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let x2 = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.cosineSimilarityLoss(x1: x1, x2: x2) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.3369714617729187, + minimum: 0.12991397082805634, + maximum: 0.6953044533729553, + absoluteSum: 1.3478858470916748, + positionChecksum: 0.698710560798645, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.3417125642299652, 0.6953044533729553, 0.1809549331665039, + 0.12991397082805634, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedElementwiseTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedElementwiseTests.swift new file mode 100644 index 000000000..8dd22458a --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedElementwiseTests.swift @@ -0,0 +1,1517 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 63 + +import Foundation +import MLX +import Testing + +@Suite("generated: Elementwise") +struct GeneratedElementwiseTests { + + @Test("exp") + func test_exp() throws { + try withIntegrationState(seed: 47954) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.exp(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.510894775390625, + minimum: 0.2611827552318573, + maximum: 7.178291320800781, + absoluteSum: 18.1307373046875, + positionChecksum: 7.237987518310547, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 7.178291320800781, 1.276309847831726, 0.298198401927948, 0.2611827552318573, + 0.6586772799491882, 0.5715979933738708, + ]), + tolerance: .float32) + } + } + + @Test("expm1") + func test_expm1() throws { + try withIntegrationState(seed: 78136) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.expm1(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3363969922065735, + minimum: -0.7551273703575134, + maximum: 1.9503211975097656, + absoluteSum: 8.781084060668945, + positionChecksum: 4.586121241251628, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.3306465744972229, 1.373231053352356, -0.3045533299446106, + -0.6832998991012573, -0.6291795969009399, 0.5119333863258362, + ]), + tolerance: .float32) + } + } + + @Test("log") + func test_log() throws { + try withIntegrationState(seed: 57228) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.log(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.13923785090446472, + minimum: -0.6309344172477722, + maximum: 0.6703048944473267, + absoluteSum: 4.917095184326172, + positionChecksum: 2.5226616859436035, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.6309344172477722, 0.15799346566200256, 0.2860988676548004, + -0.1460823118686676, -0.28112390637397766, 0.6090521812438965, + ]), + tolerance: .float32) + } + } + + @Test("log2") + func test_log2() throws { + try withIntegrationState(seed: 96094) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.log2(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.3897838592529297, + minimum: -3.062103271484375, + maximum: 0.6787765026092529, + absoluteSum: 11.070028305053711, + positionChecksum: 5.862569173177083, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.607052743434906, -1.283917784690857, -1.0272319316864014, + 0.2580295503139496, -3.062103271484375, 0.6499488949775696, + ]), + tolerance: .float32) + } + } + + @Test("log10") + func test_log10() throws { + try withIntegrationState(seed: 4788) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.log10(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.07758500427007675, + minimum: -0.7380918264389038, + maximum: 0.2524550259113312, + absoluteSum: 3.0936691761016846, + positionChecksum: 1.3356029192606609, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.24614253640174866, -0.3498530387878418, 0.16428828239440918, + 0.1663184016942978, 0.16893142461776733, 0.05258811637759209, + ]), + tolerance: .float32) + } + } + + @Test("log1p") + func test_log1p() throws { + try withIntegrationState(seed: 75204) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.log1p(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.7237419486045837, + minimum: 0.18861305713653564, + maximum: 1.0836082696914673, + absoluteSum: 8.684903144836426, + positionChecksum: 4.419689814249675, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8431658744812012, 0.45729416608810425, 1.0836082696914673, + 1.0338786840438843, 0.3002285659313202, 0.9232745170593262, + ]), + tolerance: .float32) + } + } + + @Test("sqrt") + func test_sqrt() throws { + try withIntegrationState(seed: 65948) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.sqrt(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0848369598388672, + minimum: 0.5140633583068848, + maximum: 1.4113861322402954, + absoluteSum: 13.018043518066406, + positionChecksum: 6.730450948079427, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.4113861322402954, 1.3208208084106445, 0.5140633583068848, + 0.8169410824775696, 0.6033309102058411, 1.0939743518829346, + ]), + tolerance: .float32) + } + } + + @Test("rsqrt") + func test_rsqrt() throws { + try withIntegrationState(seed: 35677) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.rsqrt(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.1557154655456543, + minimum: 0.7132951021194458, + maximum: 3.140223264694214, + absoluteSum: 13.868584632873535, + positionChecksum: 7.229054133097331, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7132951021194458, 2.035323143005371, 0.964241623878479, 0.896981954574585, + 0.7307946681976318, 0.9336780309677124, + ]), + tolerance: .float32) + } + } + + @Test("square") + func test_square() throws { + try withIntegrationState(seed: 72464) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.square(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.695485532283783, + minimum: 0.014689148403704166, + maximum: 4.136717319488525, + absoluteSum: 8.345826148986816, + positionChecksum: 6.028196334838867, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.020748626440763474, 0.4422823488712311, 0.02497928962111473, + 0.11787982285022736, 4.136717319488525, 0.20451615750789642, + ]), + tolerance: .float32) + } + } + + @Test("abs") + func test_abs() throws { + try withIntegrationState(seed: 54927) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.abs(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.7679557800292969, + minimum: 0.05212230235338211, + maximum: 1.779400110244751, + absoluteSum: 9.215469360351562, + positionChecksum: 5.223146438598633, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.2723303735256195, 0.05212230235338211, 0.7042833566665649, + 0.4646765887737274, 0.6566354632377625, 1.0102406740188599, + ]), + tolerance: .float32) + } + } + + @Test("negative") + func test_negative() throws { + try withIntegrationState(seed: 65451) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.negative(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.40208250284194946, + minimum: -2.357377290725708, + maximum: 1.9783403873443604, + absoluteSum: 11.139066696166992, + positionChecksum: 6.509375254313151, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.32914111018180847, 0.6736426949501038, -0.43585315346717834, + -0.034667275846004486, -2.357377290725708, 0.7370271682739258, + ]), + tolerance: .float32) + } + } + + @Test("sign") + func test_sign() throws { + try withIntegrationState(seed: 14106) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sign(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.1666666716337204, + minimum: -1.0, + maximum: 1.0, + absoluteSum: 12.0, + positionChecksum: 6.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-1.0, -1.0, 1.0, -1.0, 1.0, -1.0]), + tolerance: .float32) + } + } + + @Test("floor") + func test_floor() throws { + try withIntegrationState(seed: 48869) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.floor(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.4166666865348816, + minimum: -2.0, + maximum: 1.0, + absoluteSum: 11.0, + positionChecksum: 7.083333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-1.0, 1.0, 0.0, 1.0, -1.0, -1.0]), + tolerance: .float32) + } + } + + @Test("ceil") + func test_ceil() throws { + try withIntegrationState(seed: 65987) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.ceil(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.8333333730697632, + minimum: -1.0, + maximum: 3.0, + absoluteSum: 14.0, + positionChecksum: 7.25, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [3.0, -0.0, 1.0, -1.0, 1.0, -1.0]), + tolerance: .float32) + } + } + + @Test("trunc") + func test_trunc() throws { + try withIntegrationState(seed: 30046) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.trunc(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.25, + minimum: -0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 1.0833333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("reciprocal") + func test_reciprocal() throws { + try withIntegrationState(seed: 56536) { + let a = MLXRandom.uniform(low: 0.5, high: 2.0, [4, 3], dtype: .float32) + let result = MLX.reciprocal(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0375549793243408, + minimum: 0.5049313306808472, + maximum: 1.9941542148590088, + absoluteSum: 12.450658798217773, + positionChecksum: 6.7019303639729815, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8576971888542175, 1.7296605110168457, 1.9941542148590088, + 0.6068732738494873, 1.437481164932251, 1.684537410736084, + ]), + tolerance: .float32) + } + } + + @Test("sin") + func test_sin() throws { + try withIntegrationState(seed: 22723) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sin(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.026744266971945763, + minimum: -0.9563843011856079, + maximum: 0.8906667828559875, + absoluteSum: 7.560203552246094, + positionChecksum: 4.530291557312012, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.3366437554359436, -0.10839032381772995, -0.41767552495002747, + -0.9191461205482483, -0.8813959360122681, 0.8294884562492371, + ]), + tolerance: .float32) + } + } + + @Test("cos") + func test_cos() throws { + try withIntegrationState(seed: 79724) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cos(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.6566232442855835, + minimum: -0.3807412087917328, + maximum: 0.9957085251808167, + absoluteSum: 8.786455154418945, + positionChecksum: 4.586637496948242, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8214612603187561, -0.3807412087917328, 0.9957085251808167, + 0.9071047306060791, 0.9621671438217163, -0.07274708896875381, + ]), + tolerance: .float32) + } + } + + @Test("tan") + func test_tan() throws { + try withIntegrationState(seed: 53934) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tan(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5176577568054199, + minimum: -0.9707298874855042, + maximum: 3.0964441299438477, + absoluteSum: 11.18362045288086, + positionChecksum: 4.061483383178711, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.44313281774520874, 3.0964441299438477, -0.4185173213481903, + -0.9707298874855042, -0.04070381820201874, -0.15154463052749634, + ]), + tolerance: .float32) + } + } + + @Test("sinh") + func test_sinh() throws { + try withIntegrationState(seed: 99750) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sinh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.2838282585144043, + minimum: -3.0325629711151123, + maximum: 1.4805724620819092, + absoluteSum: 14.353582382202148, + positionChecksum: 5.958340962727864, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.3117921352386475, -2.4681921005249023, 1.1744270324707031, + -0.6643975973129272, -0.37228044867515564, -0.08928694576025009, + ]), + tolerance: .float32) + } + } + + @Test("cosh") + func test_cosh() throws { + try withIntegrationState(seed: 31959) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cosh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.7174079418182373, + minimum: 1.0034161806106567, + maximum: 5.466761589050293, + absoluteSum: 20.60889434814453, + positionChecksum: 12.778172810872396, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.0733435153961182, 1.0399240255355835, 1.1779524087905884, + 2.703171968460083, 1.1084223985671997, 1.0640565156936646, + ]), + tolerance: .float32) + } + } + + @Test("tanh") + func test_tanh() throws { + try withIntegrationState(seed: 35559) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tanh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.09210336208343506, + minimum: -0.8818567395210266, + maximum: 0.9250870943069458, + absoluteSum: 6.192976474761963, + positionChecksum: 3.220587412516276, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.581309974193573, 0.3747791647911072, 0.3662618398666382, + 0.20108503103256226, -0.43759405612945557, 0.9250870943069458, + ]), + tolerance: .float32) + } + } + + @Test("asin") + func test_asin() throws { + try withIntegrationState(seed: 21542) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + let result = MLX.asin(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.08627458661794662, + minimum: -0.46629178524017334, + maximum: 0.8269193768501282, + absoluteSum: 4.095224380493164, + positionChecksum: 2.3354612986246743, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6206415295600891, 0.10426806658506393, -0.29076746106147766, + -0.33669838309288025, 0.8269193768501282, 0.12828460335731506, + ]), + tolerance: .float32) + } + } + + @Test("acos") + func test_acos() throws { + try withIntegrationState(seed: 78281) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + let result = MLX.acos(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.745958685874939, + minimum: 0.736549437046051, + maximum: 2.466407060623169, + absoluteSum: 20.95150375366211, + positionChecksum: 12.173113505045572, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.836718201637268, 1.2751185894012451, 2.083163261413574, 1.961542010307312, + 2.466407060623169, 2.0229902267456055, + ]), + tolerance: .float32) + } + } + + @Test("atan") + func test_atan() throws { + try withIntegrationState(seed: 90539) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + let result = MLX.atan(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.046083711087703705, + minimum: -0.7164856195449829, + maximum: 0.5717471241950989, + absoluteSum: 4.243737697601318, + positionChecksum: 2.1660304069519043, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.33865731954574585, 0.058168165385723114, 0.495836466550827, + 0.28327685594558716, -0.7164856195449829, 0.06804904341697693, + ]), + tolerance: .float32) + } + } + + @Test("asinh") + func test_asinh() throws { + try withIntegrationState(seed: 18773) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.asinh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.17652754485607147, + minimum: -1.0286109447479248, + maximum: 1.2197173833847046, + absoluteSum: 6.706778049468994, + positionChecksum: 3.318317731221517, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.5257062315940857, -0.8774291276931763, 0.18421567976474762, + 0.0898388996720314, 0.9146624207496643, 0.02543237805366516, + ]), + tolerance: .float32) + } + } + + @Test("acosh") + func test_acosh() throws { + try withIntegrationState(seed: 95620) { + let a = MLXRandom.uniform(low: 1.0, high: 3.0, [4, 3], dtype: .float32) + let result = MLX.acosh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.4068176746368408, + minimum: 0.5282204747200012, + maximum: 1.7533022165298462, + absoluteSum: 16.881811141967773, + positionChecksum: 9.038970947265625, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.674250841140747, 1.7533022165298462, 1.330007791519165, 1.561944842338562, + 1.1493288278579712, 1.735342264175415, + ]), + tolerance: .float32) + } + } + + @Test("atanh") + func test_atanh() throws { + try withIntegrationState(seed: 17268) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + let result = MLX.atanh(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.08335167169570923, + minimum: -0.8625828623771667, + maximum: 1.104719638824463, + absoluteSum: 5.097476005554199, + positionChecksum: 2.553055922190348, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.558688223361969, -0.10614049434661865, -0.8625828623771667, + 0.07864232361316681, -0.030338648706674576, 0.7645785808563232, + ]), + tolerance: .float32) + } + } + + @Test("erf") + func test_erf() throws { + try withIntegrationState(seed: 19849) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.erf(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1678692102432251, + minimum: -0.9197378754615784, + maximum: 0.9991385340690613, + absoluteSum: 6.304012298583984, + positionChecksum: 3.7706054051717124, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7421068549156189, 0.12956653535366058, 0.9657787680625916, + 0.40262550115585327, 0.33637773990631104, 0.9991385340690613, + ]), + tolerance: .float32) + } + } + + @Test("erfInverse") + func test_erfInverse() throws { + try withIntegrationState(seed: 83926) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + let result = MLX.erfInverse(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.016824109479784966, + minimum: -0.750474750995636, + maximum: 1.074080467224121, + absoluteSum: 5.0765275955200195, + positionChecksum: 2.2263415654500327, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.750474750995636, 0.003604303812608123, -0.47791653871536255, + -0.009403870441019535, -0.17815247178077698, 0.3282304108142853, + ]), + tolerance: .float32) + } + } + + @Test("sigmoid") + func test_sigmoid() throws { + try withIntegrationState(seed: 83793) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sigmoid(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5078827142715454, + minimum: 0.2062499076128006, + maximum: 0.8442580103874207, + absoluteSum: 6.094592094421387, + positionChecksum: 3.197172164916992, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.5074858069419861, 0.6258241534233093, 0.20864970982074738, + 0.3561042547225952, 0.2062499076128006, 0.8442580103874207, + ]), + tolerance: .float32) + } + } + + @Test("degrees") + func test_degrees() throws { + try withIntegrationState(seed: 24094) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.degrees(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -2.4433441162109375, + minimum: -85.32889556884766, + maximum: 128.0925750732422, + absoluteSum: 596.0240478515625, + positionChecksum: 281.8431396484375, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 128.0925750732422, 36.44943618774414, 30.74724006652832, -70.09713745117188, + 55.25877380371094, -10.830322265625, + ]), + tolerance: .float32) + } + } + + @Test("radians") + func test_radians() throws { + try withIntegrationState(seed: 7077) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.radians(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.0018685739487409592, + minimum: -0.031736548990011215, + maximum: 0.03308562934398651, + absoluteSum: 0.18592886626720428, + positionChecksum: 0.08721178770065308, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.017862409353256226, -0.009540251456201077, 0.03308562934398651, + -0.013581288047134876, -0.011545875109732151, -0.007317093200981617, + ]), + tolerance: .float32) + } + } + + @Test("stopGradient") + func test_stopGradient() throws { + try withIntegrationState(seed: 91151) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.stopGradient(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.26480633020401, + minimum: -0.7478859424591064, + maximum: 1.6692116260528564, + absoluteSum: 6.601988792419434, + positionChecksum: 2.663097381591797, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8156958222389221, 1.6692116260528564, 0.24522577226161957, + -0.047961488366127014, 0.6883101463317871, -0.17975953221321106, + ]), + tolerance: .float32) + } + } + + @Test("round") + func test_round() throws { + try withIntegrationState(seed: 63999) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 10.0) + let result = MLX.round(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.5, + minimum: -17.0, + maximum: 23.0, + absoluteSum: 98.0, + positionChecksum: 39.166666666666664, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-14.0, 23.0, -1.0, 2.0, 7.0, 0.0]), + tolerance: .float32) + } + } + + @Test("round/decimals") + func test_round_decimals() throws { + try withIntegrationState(seed: 2520) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 10.0) + let result = MLX.round(a, decimals: 2) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.6791664361953735, + minimum: -16.440000534057617, + maximum: 15.34999942779541, + absoluteSum: 110.8699951171875, + positionChecksum: 59.361663818359375, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -5.789999961853027, -5.799999713897705, -10.619999885559082, + 10.90999984741211, -11.369999885559082, 5.349999904632568, + ]), + tolerance: .float32) + } + } + + @Test("clip") + func test_clip() throws { + try withIntegrationState(seed: 84323) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.clip(a, min: -0.5, max: 0.5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.20147940516471863, + minimum: -0.5, + maximum: 0.5, + absoluteSum: 4.310338497161865, + positionChecksum: 2.2437593142191568, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.5, -0.5, -0.15267881751060486, -0.22263969480991364, 0.5, + -0.2244129180908203, + ]), + tolerance: .float32) + } + } + + @Test("logicalNot") + func test_logicalNot() throws { + try withIntegrationState(seed: 65221) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.logicalNot(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.0833333358168602, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 0.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("isNaN") + func test_isNaN() throws { + try withIntegrationState(seed: 94051) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.isNaN(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .bool, + mean: 0.1666666716337204, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 0.8333333333333334, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("isInf") + func test_isInf() throws { + try withIntegrationState(seed: 17171) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.isInf(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .bool, + mean: 0.3333333432674408, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 2.0, + positionChecksum: 1.1666666666666667, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 0.0, 1.0, 1.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("isFinite") + func test_isFinite() throws { + try withIntegrationState(seed: 93517) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.isFinite(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .bool, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 1.5, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [1.0, 1.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("isPosInf") + func test_isPosInf() throws { + try withIntegrationState(seed: 67816) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.isPosInf(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .bool, + mean: 0.1666666716337204, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 0.5, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 0.0, 1.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("isNegInf") + func test_isNegInf() throws { + try withIntegrationState(seed: 31786) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.isNegInf(a) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .bool, + mean: 0.1666666716337204, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 0.6666666666666666, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 0.0, 0.0, 1.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("nanToNum") + func test_nanToNum() throws { + // arguments are explicit: python defaults posinf/neginf to the dtype max/min while Swift defaults them to 0 + try withIntegrationState(seed: 94828) { + let a = MLXArray([1.0, -2.5, Float.infinity, -Float.infinity, Float.nan, 0.0]) + let result = MLX.nanToNum(a, nan: 1.0, posInf: 100.0, negInf: -100.0) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .float32, + mean: -0.0833333358168602, + minimum: -100.0, + maximum: 100.0, + absoluteSum: 204.5, + positionChecksum: 118.5, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [1.0, -2.5, 100.0, -100.0, 1.0, 0.0]), + tolerance: .float32) + } + } + + @Test("abs/method") + func test_abs_method() throws { + try withIntegrationState(seed: 59334) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.abs() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0287009477615356, + minimum: 0.007374108768999577, + maximum: 2.067824363708496, + absoluteSum: 12.34441089630127, + positionChecksum: 7.004905700683594, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6036763191223145, 1.1487308740615845, 0.4937532842159271, + 2.067824363708496, 1.0368303060531616, 2.0537874698638916, + ]), + tolerance: .float32) + } + } + + @Test("exp/method") + func test_exp_method() throws { + try withIntegrationState(seed: 65744) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.exp() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.2003825902938843, + minimum: 0.19767798483371735, + maximum: 3.021728992462158, + absoluteSum: 14.404590606689453, + positionChecksum: 7.480013529459636, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.19767798483371735, 1.8948596715927124, 1.4384976625442505, + 3.021728992462158, 0.4999668002128601, 0.38851505517959595, + ]), + tolerance: .float32) + } + } + + @Test("sqrt/method") + func test_sqrt_method() throws { + try withIntegrationState(seed: 3785) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + let result = a.sqrt() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.1024643182754517, + minimum: 0.7357668280601501, + maximum: 1.4061843156814575, + absoluteSum: 13.229571342468262, + positionChecksum: 7.240408579508464, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.4061843156814575, 1.0316053628921509, 1.038780689239502, + 1.0228774547576904, 1.3849389553070068, 0.9444799423217773, + ]), + tolerance: .float32) + } + } + + @Test("round/method") + func test_round_method() throws { + try withIntegrationState(seed: 25776) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 10.0) + let result = a.round() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.0, + minimum: -20.0, + maximum: 12.0, + absoluteSum: 74.0, + positionChecksum: 39.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-3.0, -4.0, 4.0, 12.0, -0.0, -2.0]), + tolerance: .float32) + } + } + + @Test("asType/float16") + func test_asType_float16() throws { + try withIntegrationState(seed: 396) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.asType(.float16) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float16, + mean: 0.3122355341911316, + minimum: -0.763671875, + maximum: 1.630859375, + absoluteSum: 7.469482421875, + positionChecksum: 4.177174886067708, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.335693359375, 1.630859375, 0.72021484375, 0.63427734375, -0.72607421875, + 1.4150390625, + ]), + tolerance: .float16) + } + } + + @Test("asType/int32") + func test_asType_int32() throws { + try withIntegrationState(seed: 68400) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 10.0) + let result = a.asType(.int32) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: -4.6666669845581055, + minimum: -23.0, + maximum: 11.0, + absoluteSum: 106.0, + positionChecksum: 49.083333333333336, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [11.0, -10.0, -5.0, -3.0, -11.0, 6.0]), + tolerance: .exact) + } + } + + @Test("exp/float16") + func test_exp_float16() throws { + try withIntegrationState(seed: 43848) { + let a = MLXRandom.normal([4, 3], dtype: .float16, loc: 0.0, scale: 1.0) + let result = MLX.exp(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float16, + mean: 1.9106242656707764, + minimum: 0.174072265625, + maximum: 6.24609375, + absoluteSum: 22.927490234375, + positionChecksum: 12.168294270833334, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.255859375, 2.857421875, 1.10546875, 0.388671875, 1.515625, 1.5947265625, + ]), + tolerance: .float16) + } + } + + @Test("exp/bfloat16") + func test_exp_bfloat16() throws { + try withIntegrationState(seed: 82777) { + let a = MLXRandom.normal([4, 3], dtype: .bfloat16, loc: 0.0, scale: 1.0) + let result = MLX.exp(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bfloat16, + mean: 2.4685873985290527, + minimum: 0.21484375, + maximum: 9.625, + absoluteSum: 29.623046875, + positionChecksum: 22.324381510416668, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.494140625, 2.140625, 0.6796875, 0.84375, 0.7890625, 8.9375]), + tolerance: .float16) + } + } + + @Test("positive") + func test_positive() throws { + try withIntegrationState(seed: 39984) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.positive(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.24678385257720947, + minimum: -1.4747556447982788, + maximum: 1.6195077896118164, + absoluteSum: 8.60529613494873, + positionChecksum: 4.061902681986491, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.7380889058113098, -1.2995355129241943, 0.5283272862434387, + 0.3148530423641205, 0.2088206708431244, -0.6653445363044739, + ]), + tolerance: .float32) + } + } + + @Test("view/int16") + func test_view_int16() throws { + try withIntegrationState(seed: 63133) { + let a = MLXRandom.randInt(low: 0, high: 1024, [4, 3], type: Int32.self) + let result = MLX.view(a, dtype: .int16) + expectSummary( + result, + ArraySummary( + shape: [4, 6], + dtype: .int16, + mean: 213.625, + minimum: 0.0, + maximum: 919.0, + absoluteSum: 5127.0, + positionChecksum: 2602.625, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [22.0, 0.0, 0.0, 863.0, 255.0, 0.0]), + tolerance: .exact) + } + } + + @Test("contiguous") + func test_contiguous() throws { + try withIntegrationState(seed: 69937) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.contiguous(a.T) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: -0.5083404779434204, + minimum: -2.389263391494751, + maximum: 0.8913607597351074, + absoluteSum: 9.336750984191895, + positionChecksum: 4.510031382242839, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.004110712558031082, -2.389263391494751, -0.3058277368545532, + -0.40270015597343445, -0.49763554334640503, 0.01963740773499012, + ]), + tolerance: .float32) + } + } + + @Test("isClose") + func test_isClose() throws { + try withIntegrationState(seed: 90080) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.isClose(a, b, rtol: 0.5, atol: 0.1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.25, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 1.6666666666666667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 0.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("allClose") + func test_allClose() throws { + try withIntegrationState(seed: 65428) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.allClose(a, a + 0.001, rtol: 0.01) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0], + samples: [0.0]), + tolerance: .exact) + } + } + + @Test("arrayEqual") + func test_arrayEqual() throws { + try withIntegrationState(seed: 84233) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.arrayEqual(a, a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0], + samples: [1.0]), + tolerance: .exact) + } + } + + @Test("arrayEqual/false") + func test_arrayEqual_false() throws { + try withIntegrationState(seed: 90639) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let b = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + let result = MLX.arrayEqual(a, b) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0], + samples: [0.0]), + tolerance: .exact) + } + } + + @Test("realPart") + func test_realPart() throws { + // Swift exposes realPart() as a method, not a free function + try withIntegrationState(seed: 28164) { + let r = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = c.realPart() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.27432000637054443, + minimum: -2.3572521209716797, + maximum: 0.864957869052887, + absoluteSum: 8.50927734375, + positionChecksum: 3.2805039087931314, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.3572521209716797, -0.278932124376297, 0.36380791664123535, + 0.6124622225761414, 0.09972750395536423, -0.07274171710014343, + ]), + tolerance: .float32) + } + } + + @Test("imaginaryPart") + func test_imaginaryPart() throws { + try withIntegrationState(seed: 58944) { + let r = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = c.imaginaryPart() + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.4124636650085449, + minimum: -1.5858224630355835, + maximum: 1.1684366464614868, + absoluteSum: 9.83724308013916, + positionChecksum: 6.4064591725667315, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.04229888319969177, 0.016773255541920662, -0.5229106545448303, + -1.2519052028656006, -1.4781920909881592, 0.7881496548652649, + ]), + tolerance: .float32) + } + } + + @Test("conjugate") + func test_conjugate() throws { + try withIntegrationState(seed: 29664) { + let r = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.conjugate(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.23909756541252136, + minimum: -0.8671482801437378, + maximum: 1.3631097078323364, + absoluteSum: 9.518288612365723, + positionChecksum: 5.316549301147461, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7143032550811768, -0.7656220197677612, 0.36717352271080017, + 1.3309167623519897, -0.14781005680561066, -0.8671482801437378, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.04297981783747673, + minimum: -1.5997480154037476, + maximum: 1.4610912799835205, + absoluteSum: 10.585673332214355, + positionChecksum: 5.1464080810546875, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.5997480154037476, -0.7909235954284668, -1.143326997756958, + 1.0923233032226562, 1.4610912799835205, 1.02470064163208, + ]), + tolerance: .float32) + } + } + + @Test("asImaginary") + func test_asImaginary() throws { + // verifies the complex value the other complex cases are built from + try withIntegrationState(seed: 52131) { + let r = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = c + expectSummary( + result.realPart(), + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.4877490997314453, + minimum: -1.8236697912216187, + maximum: 0.596900999546051, + absoluteSum: 7.777519226074219, + positionChecksum: 3.4768377939860025, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.25812309980392456, -1.8236697912216187, -1.0013736486434937, + -0.42360997200012207, 0.10724081844091415, -0.371466726064682, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.037692904472351074, + minimum: -1.3910119533538818, + maximum: 1.7863659858703613, + absoluteSum: 9.36073112487793, + positionChecksum: 4.2958984375, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.7863659858703613, 0.4858953356742859, 1.0185210704803467, + 1.220772624015808, 0.1476796418428421, 0.20525570213794708, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedFFTTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedFFTTests.swift new file mode 100644 index 000000000..32f6eafeb --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedFFTTests.swift @@ -0,0 +1,1513 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 41 + +import Foundation +import MLX +import Testing + +@Suite("generated: FFT") +struct GeneratedFFTTests { + + @Test("fft") + func test_fft() throws { + try withIntegrationState(seed: 3161) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -0.30193811655044556, + minimum: -26.834150314331055, + maximum: 25.471298217773438, + absoluteSum: 748.5565795898438, + positionChecksum: 364.7257421875, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -17.864389419555664, 0.47785067558288574, -7.920133590698242, + 1.171248197555542, 5.3682661056518555, -11.829191207885742, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -2.2110860347747803, + minimum: -33.87453842163086, + maximum: 27.859336853027344, + absoluteSum: 759.38427734375, + positionChecksum: 398.85734375, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -10.048587799072266, 0.7990264892578125, 2.4799821376800537, + 8.100057601928711, -5.986302375793457, -4.987302303314209, + ]), + tolerance: .float32) + } + } + + @Test("fft/nShort") + func test_fft_nShort() throws { + try withIntegrationState(seed: 2147) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c, n: 80) + expectSummary( + result.realPart(), + ArraySummary( + shape: [80], + dtype: .float32, + mean: -0.30215585231781006, + minimum: -19.026569366455078, + maximum: 23.793720245361328, + absoluteSum: 603.0272216796875, + positionChecksum: 314.046435546875, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + 8.233772277832031, -4.0249857902526855, 2.5544846057891846, + -9.478172302246094, 3.6788711547851562, 5.409932613372803, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [80], + dtype: .float32, + mean: -0.330033540725708, + minimum: -17.620588302612305, + maximum: 18.184288024902344, + absoluteSum: 535.7166748046875, + positionChecksum: 275.080224609375, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + -4.2741594314575195, -5.628664493560791, -7.816094875335693, + 7.278494358062744, -11.146451950073242, -5.636445999145508, + ]), + tolerance: .float32) + } + } + + @Test("fft/nLong") + func test_fft_nLong() throws { + try withIntegrationState(seed: 27280) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c, n: 120) + expectSummary( + result.realPart(), + ArraySummary( + shape: [120], + dtype: .float32, + mean: 1.245316982269287, + minimum: -24.74895668029785, + maximum: 26.73784065246582, + absoluteSum: 952.5072021484375, + positionChecksum: 484.21611328125, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + 17.85637664794922, 7.068179607391357, 3.323228120803833, 18.420766830444336, + 10.24211311340332, -7.67423152923584, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [120], + dtype: .float32, + mean: 0.865074634552002, + minimum: -23.82689666748047, + maximum: 21.837039947509766, + absoluteSum: 956.6898803710938, + positionChecksum: 484.05237630208336, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + 9.100713729858398, -9.670844078063965, -3.221973419189453, + -2.571913719177246, -14.796401023864746, 15.20323371887207, + ]), + tolerance: .float32) + } + } + + @Test("fft/axis") + func test_fft_axis() throws { + try withIntegrationState(seed: 67284) { + let r = MLXRandom.normal([10, 10], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([10, 10], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c, axis: 0) + expectSummary( + result.realPart(), + ArraySummary( + shape: [10, 10], + dtype: .float32, + mean: -0.465768426656723, + minimum: -9.097143173217773, + maximum: 7.025123119354248, + absoluteSum: 236.0032501220703, + positionChecksum: 112.51224609375, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -5.653952598571777, -1.792625904083252, -3.0560266971588135, + -1.7548167705535889, 3.503523111343384, -1.5613975524902344, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [10, 10], + dtype: .float32, + mean: -0.3379652798175812, + minimum: -7.193725109100342, + maximum: 7.352320671081543, + absoluteSum: 243.55081176757812, + positionChecksum: 124.103173828125, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -6.474360466003418, -1.6081452369689941, 3.647308588027954, + -0.7049122452735901, 0.582835853099823, 1.1408429145812988, + ]), + tolerance: .float32) + } + } + + @Test("fft/ortho") + func test_fft_ortho() throws { + try withIntegrationState(seed: 41079) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft(c, norm: .ortho) + expectSummary( + result.realPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -0.1377011239528656, + minimum: -2.420893907546997, + maximum: 2.4660956859588623, + absoluteSum: 89.8805160522461, + positionChecksum: 44.410625, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + 0.9301566481590271, 1.4401391744613647, -1.4074681997299194, + -2.0820672512054443, -0.47943899035453796, -0.2989818751811981, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: 0.0205818060785532, + minimum: -2.0081610679626465, + maximum: 2.822983503341675, + absoluteSum: 78.08838653564453, + positionChecksum: 37.138740234375, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -1.1429274082183838, 1.4222431182861328, 0.8448073267936707, + 0.45329347252845764, 1.7841476202011108, 0.44647422432899475, + ]), + tolerance: .float32) + } + } + + @Test("ifft") + func test_ifft() throws { + try withIntegrationState(seed: 92555) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: 0.016596531495451927, + minimum: -0.23650917410850525, + maximum: 0.2672997713088989, + absoluteSum: 8.624792098999023, + positionChecksum: 4.234874267578125, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -0.010019684210419655, 0.14668859541416168, 0.11429940909147263, + -0.045612435787916183, 0.14252117276191711, 0.036860380321741104, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: 0.018229888752102852, + minimum: -0.23968389630317688, + maximum: 0.2401280254125595, + absoluteSum: 8.7319917678833, + positionChecksum: 4.291281127929688, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -0.0015521335881203413, -0.17396560311317444, 0.0605366975069046, + -0.03737438842654228, 0.04746906831860542, -0.12103308737277985, + ]), + tolerance: .float32) + } + } + + @Test("ifft/nShort") + func test_ifft_nShort() throws { + try withIntegrationState(seed: 12331) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft(c, n: 80) + expectSummary( + result.realPart(), + ArraySummary( + shape: [80], + dtype: .float32, + mean: -0.019246213138103485, + minimum: -0.3084016740322113, + maximum: 0.20837044715881348, + absoluteSum: 7.586906909942627, + positionChecksum: 3.7272315979003907, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + -0.06569941341876984, 0.08424773812294006, 0.13474620878696442, + 0.1767660677433014, -0.20869246125221252, 0.027556825429201126, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [80], + dtype: .float32, + mean: 0.016126060858368874, + minimum: -0.20641577243804932, + maximum: 0.277831107378006, + absoluteSum: 7.127358436584473, + positionChecksum: 3.7531387329101564, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + -0.05102470517158508, -0.11271238327026367, 0.026768634095788002, + 0.09892851114273071, -0.12387256324291229, 0.01915312185883522, + ]), + tolerance: .float32) + } + } + + @Test("ifft/nLong") + func test_ifft_nLong() throws { + try withIntegrationState(seed: 38177) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft(c, n: 120) + expectSummary( + result.realPart(), + ArraySummary( + shape: [120], + dtype: .float32, + mean: -0.004438468720763922, + minimum: -0.2090398520231247, + maximum: 0.2671051621437073, + absoluteSum: 8.520634651184082, + positionChecksum: 4.005149841308594, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + 0.07100244611501694, 0.07003705203533173, 0.2671051621437073, + -0.2090398520231247, -0.09596047550439835, 0.02532544918358326, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [120], + dtype: .float32, + mean: -0.0033197151497006416, + minimum: -0.2966475784778595, + maximum: 0.22516706585884094, + absoluteSum: 8.153952598571777, + positionChecksum: 3.686529286702474, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + -0.12483212351799011, 0.09503603726625443, 0.016529258340597153, + 0.10363555699586868, 0.10037048161029816, -0.031290605664253235, + ]), + tolerance: .float32) + } + } + + @Test("ifft/axis") + func test_ifft_axis() throws { + try withIntegrationState(seed: 92391) { + let r = MLXRandom.normal([10, 10], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([10, 10], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft(c, axis: 0) + expectSummary( + result.realPart(), + ArraySummary( + shape: [10, 10], + dtype: .float32, + mean: -0.03038243018090725, + minimum: -0.9769951701164246, + maximum: 0.7989615797996521, + absoluteSum: 27.44314193725586, + positionChecksum: 14.012554931640626, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + 0.37563377618789673, 0.14058302342891693, 0.17051009833812714, + 0.10858341306447983, 0.29239600896835327, -0.06135018542408943, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [10, 10], + dtype: .float32, + mean: -0.07071053236722946, + minimum: -0.8794716596603394, + maximum: 0.5949147343635559, + absoluteSum: 24.313308715820312, + positionChecksum: 13.0312890625, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -0.008079779334366322, 0.014237630181014538, -0.20407943427562714, + 0.25475987792015076, -0.5532649755477905, -0.37195757031440735, + ]), + tolerance: .float32) + } + } + + @Test("ifft/ortho") + func test_ifft_ortho() throws { + try withIntegrationState(seed: 8262) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft(c, norm: .ortho) + expectSummary( + result.realPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -0.14974907040596008, + minimum: -2.6609697341918945, + maximum: 2.1586873531341553, + absoluteSum: 76.76493835449219, + positionChecksum: 35.6903564453125, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + -0.7524006366729736, -0.7424914240837097, -1.2931209802627563, + -0.3934279680252075, 0.7683846354484558, 0.5575725436210632, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [100], + dtype: .float32, + mean: -0.1610265076160431, + minimum: -2.565251588821411, + maximum: 2.1245665550231934, + absoluteSum: 91.36912536621094, + positionChecksum: 48.544853515625, + sampleIndices: [0, 20, 40, 59, 79, 99], + samples: [ + 0.05022697523236275, -0.2160770297050476, -2.0412232875823975, + -0.921349287033081, -1.5685405731201172, -2.192484140396118, + ]), + tolerance: .float32) + } + } + + @Test("rfft") + func test_rfft() throws { + try withIntegrationState(seed: 9717) { + let c = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfft(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [51], + dtype: .float32, + mean: 0.08805207908153534, + minimum: -14.910188674926758, + maximum: 16.072185516357422, + absoluteSum: 306.84942626953125, + positionChecksum: 164.0440793504902, + sampleIndices: [0, 10, 20, 30, 40, 50], + samples: [ + -3.9297866821289062, -7.619717121124268, -0.037673234939575195, + 4.216407775878906, 1.6698544025421143, 9.09442138671875, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [51], + dtype: .float32, + mean: 0.4946885406970978, + minimum: -14.891788482666016, + maximum: 12.950984954833984, + absoluteSum: 284.84588623046875, + positionChecksum: 144.68315333946077, + sampleIndices: [0, 10, 20, 30, 40, 50], + samples: [ + 0.0, 3.381943702697754, -3.114729404449463, 2.588118076324463, + -1.9571096897125244, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("rfft/n") + func test_rfft_n() throws { + try withIntegrationState(seed: 39379) { + let c = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfft(c, n: 80) + expectSummary( + result.realPart(), + ArraySummary( + shape: [41], + dtype: .float32, + mean: 1.0768029689788818, + minimum: -9.284590721130371, + maximum: 11.904441833496094, + absoluteSum: 186.73614501953125, + positionChecksum: 107.89291158536585, + sampleIndices: [0, 8, 16, 24, 32, 40], + samples: [ + 5.147792816162109, -0.7290818691253662, 2.711458206176758, + -9.033034324645996, -4.209567070007324, 8.454961776733398, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [41], + dtype: .float32, + mean: -0.010239856317639351, + minimum: -14.581356048583984, + maximum: 16.252479553222656, + absoluteSum: 198.701171875, + positionChecksum: 104.59896627286585, + sampleIndices: [0, 8, 16, 24, 32, 40], + samples: [ + 0.0, 1.6902306079864502, 3.0469067096710205, 4.965663433074951, + -2.1530351638793945, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("rfft/forward") + func test_rfft_forward() throws { + try withIntegrationState(seed: 29470) { + let c = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfft(c, norm: .forward) + expectSummary( + result.realPart(), + ArraySummary( + shape: [51], + dtype: .float32, + mean: -0.008272930979728699, + minimum: -0.18479548394680023, + maximum: 0.1636444628238678, + absoluteSum: 2.7358570098876953, + positionChecksum: 1.4796739466050093, + sampleIndices: [0, 10, 20, 30, 40, 50], + samples: [ + -0.01737872324883938, -0.0004667544271796942, -0.0013617968652397394, + -0.09146632254123688, -0.08097280561923981, 0.1343451887369156, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [51], + dtype: .float32, + mean: 0.022197062149643898, + minimum: -0.13841067254543304, + maximum: 0.19061492383480072, + absoluteSum: 2.6014790534973145, + positionChecksum: 1.4728336708218444, + sampleIndices: [0, 10, 20, 30, 40, 50], + samples: [ + 0.0, 0.011355984024703503, 0.03543727844953537, -0.13841067254543304, + 0.0947970375418663, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("irfft") + func test_irfft() throws { + try withIntegrationState(seed: 45651) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfft(c) + expectSummary( + result, + ArraySummary( + shape: [198], + dtype: .float32, + mean: -0.009573065675795078, + minimum: -0.2509799599647522, + maximum: 0.19757866859436035, + absoluteSum: 14.547061920166016, + positionChecksum: 7.517489346590909, + sampleIndices: [0, 39, 79, 118, 158, 197], + samples: [ + -0.11132393777370453, 0.011129573918879032, 0.02290239743888378, + -0.10837771743535995, -0.03477989509701729, -0.05752992257475853, + ]), + tolerance: .float32) + } + } + + @Test("irfft/n") + func test_irfft_n() throws { + try withIntegrationState(seed: 16439) { + let r = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([100], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfft(c, n: 120) + expectSummary( + result, + ArraySummary( + shape: [120], + dtype: .float32, + mean: 0.014314486645162106, + minimum: -0.28694266080856323, + maximum: 0.346214234828949, + absoluteSum: 13.746326446533203, + positionChecksum: 6.66098378499349, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + 0.09791383892297745, -0.020166436210274696, -0.1698228269815445, + 0.21948356926441193, -0.07427943497896194, -0.10716888308525085, + ]), + tolerance: .float32) + } + } + + @Test("fft2") + func test_fft2() throws { + try withIntegrationState(seed: 35659) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft2(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.7178468704223633, + minimum: -30.726062774658203, + maximum: 22.191179275512695, + absoluteSum: 3245.3369140625, + positionChecksum: 1631.6517333984375, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 3.754521369934082, -6.490096569061279, 4.651754856109619, + 14.567571640014648, 8.419650077819824, -0.5989406108856201, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.3027850389480591, + minimum: -23.579471588134766, + maximum: 25.952545166015625, + absoluteSum: 3262.20458984375, + positionChecksum: 1579.3980712890625, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.2110457420349121, 5.290388107299805, -2.667041301727295, + 1.3738980293273926, 12.329010963439941, -5.085733413696289, + ]), + tolerance: .float32) + } + } + + @Test("fft2/s") + func test_fft2_s() throws { + try withIntegrationState(seed: 49089) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft2(c, s: [3, 4]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 3, 4], + dtype: .float32, + mean: -0.44498777389526367, + minimum: -7.310737609863281, + maximum: 7.252446174621582, + absoluteSum: 287.83111572265625, + positionChecksum: 149.67510986328125, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + -2.878303050994873, -1.4361096620559692, 0.4035429358482361, + 6.749013900756836, 1.0702213048934937, -3.8870134353637695, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 3, 4], + dtype: .float32, + mean: 0.6117333173751831, + minimum: -7.615391254425049, + maximum: 8.767561912536621, + absoluteSum: 268.90863037109375, + positionChecksum: 147.26704915364584, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + -1.313608169555664, 3.5226364135742188, -4.22772216796875, + -0.37668323516845703, 8.767561912536621, -5.13897180557251, + ]), + tolerance: .float32) + } + } + + @Test("fft2/axes") + func test_fft2_axes() throws { + try withIntegrationState(seed: 55820) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fft2(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.14448444545269012, + minimum: -32.98771667480469, + maximum: 22.09111785888672, + absoluteSum: 3448.20849609375, + positionChecksum: 1765.2010498046875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.04199075698852539, 2.22198486328125, -0.7248520851135254, + -4.36610221862793, -14.677776336669922, -9.698610305786133, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.6306827068328857, + minimum: -23.172414779663086, + maximum: 25.345382690429688, + absoluteSum: 3380.444580078125, + positionChecksum: 1648.940185546875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 10.400022506713867, -4.030159950256348, 11.986205101013184, + 3.302541732788086, -2.0049829483032227, 1.248794436454773, + ]), + tolerance: .float32) + } + } + + @Test("ifft2") + func test_ifft2() throws { + try withIntegrationState(seed: 56653) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft2(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.0019193659536540508, + minimum: -0.36529168486595154, + maximum: 0.41034168004989624, + absoluteSum: 50.309425354003906, + positionChecksum: 25.473011016845703, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.004118995741009712, -0.17179054021835327, -0.09692619740962982, + -0.009604329243302345, -0.15581375360488892, 0.12122435122728348, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.004310143645852804, + minimum: -0.3949809670448303, + maximum: 0.33320337533950806, + absoluteSum: 51.721160888671875, + positionChecksum: 26.324140548706055, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.010328251868486404, 0.20638486742973328, -0.15939217805862427, + 0.16267183423042297, 0.06081516668200493, 0.11304222047328949, + ]), + tolerance: .float32) + } + } + + @Test("ifft2/s") + func test_ifft2_s() throws { + try withIntegrationState(seed: 62821) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft2(c, s: [3, 4]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 3, 4], + dtype: .float32, + mean: 0.011248579248785973, + minimum: -0.7063571214675903, + maximum: 0.7444247007369995, + absoluteSum: 23.008529663085938, + positionChecksum: 11.561017354329428, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 0.21029286086559296, -0.2403487265110016, -0.21712619066238403, + -0.17577317357063293, -0.4854602813720703, -0.17871509492397308, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 3, 4], + dtype: .float32, + mean: 0.046339940279722214, + minimum: -0.6162391304969788, + maximum: 0.8353030681610107, + absoluteSum: 18.841552734375, + positionChecksum: 9.221977233886719, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 0.17300370335578918, 0.613670825958252, 0.1690458208322525, + -0.17819438874721527, -0.15327095985412598, -0.05432131886482239, + ]), + tolerance: .float32) + } + } + + @Test("ifft2/axes") + func test_ifft2_axes() throws { + try withIntegrationState(seed: 15965) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifft2(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.001953849568963051, + minimum: -0.4720636308193207, + maximum: 0.5106112360954285, + absoluteSum: 53.05870819091797, + positionChecksum: 27.151283264160156, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.0033838478848338127, -0.10563762485980988, -0.022997520864009857, + 0.1938212811946869, -0.20600074529647827, -0.08471836894750595, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.00011229666415601969, + minimum: -0.2937026619911194, + maximum: 0.40563082695007324, + absoluteSum: 47.61387634277344, + positionChecksum: 24.217378616333008, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.09508480131626129, -0.14579325914382935, -0.18434379994869232, + 0.11046306043863297, 0.020605478435754776, 0.0009845122694969177, + ]), + tolerance: .float32) + } + } + + @Test("fftn") + func test_fftn() throws { + try withIntegrationState(seed: 20276) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fftn(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -1.9362115859985352, + minimum: -71.32183837890625, + maximum: 69.47334289550781, + absoluteSum: 9444.30859375, + positionChecksum: 4618.19921875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -40.14142608642578, -46.13766860961914, -42.42741775512695, + -28.554052352905273, 51.425296783447266, -16.503108978271484, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.8104066252708435, + minimum: -74.49031066894531, + maximum: 64.66445922851562, + absoluteSum: 8921.556640625, + positionChecksum: 4586.01806640625, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 26.956165313720703, 5.077007293701172, -26.426931381225586, + -0.8038463592529297, -10.05626392364502, 9.05034351348877, + ]), + tolerance: .float32) + } + } + + @Test("fftn/s") + func test_fftn_s() throws { + try withIntegrationState(seed: 56309) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fftn(c, s: [3, 4], axes: [0, 1]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [3, 4, 8], + dtype: .float32, + mean: -0.4631645083427429, + minimum: -9.234552383422852, + maximum: 7.494760513305664, + absoluteSum: 269.39691162109375, + positionChecksum: 135.80084228515625, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 2.372249126434326, -0.43843603134155273, 2.477461814880371, + 0.17259609699249268, 5.412355422973633, 1.3921151161193848, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [3, 4, 8], + dtype: .float32, + mean: 0.6277634501457214, + minimum: -6.913886070251465, + maximum: 11.678964614868164, + absoluteSum: 276.4273986816406, + positionChecksum: 139.78584798177084, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 1.672297716140747, 3.777714252471924, 3.101365566253662, + -1.5718486309051514, 1.980665922164917, -5.670261859893799, + ]), + tolerance: .float32) + } + } + + @Test("fftn/axes") + func test_fftn_axes() throws { + try withIntegrationState(seed: 33152) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.fftn(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.069861501455307, + minimum: -26.00143814086914, + maximum: 24.77533531188965, + absoluteSum: 3271.591796875, + positionChecksum: 1632.01904296875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 3.444851875305176, 15.690192222595215, -6.4840569496154785, + 2.230067729949951, 1.1066653728485107, -12.76463794708252, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.1090853214263916, + minimum: -22.648542404174805, + maximum: 21.17329978942871, + absoluteSum: 3264.510986328125, + positionChecksum: 1672.77587890625, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.9768218994140625, 5.668068885803223, -2.5782880783081055, + 5.189621925354004, 7.344736099243164, -5.029444694519043, + ]), + tolerance: .float32) + } + } + + @Test("ifftn") + func test_ifftn() throws { + try withIntegrationState(seed: 41746) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifftn(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.0007427402306348085, + minimum: -0.15388083457946777, + maximum: 0.1925613135099411, + absoluteSum: 17.775423049926758, + positionChecksum: 9.001312255859375, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.049047283828258514, -0.009180162101984024, 0.03828829899430275, + 0.038404885679483414, 0.07229334115982056, 0.006673771888017654, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: 0.0006254386971704662, + minimum: -0.18057334423065186, + maximum: 0.12072519958019257, + absoluteSum: 18.66337776184082, + positionChecksum: 9.401142120361328, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.050497107207775116, -0.03860977664589882, -0.06357460469007492, + -0.02733970433473587, 0.05312860757112503, -0.020095955580472946, + ]), + tolerance: .float32) + } + } + + @Test("ifftn/s") + func test_ifftn_s() throws { + try withIntegrationState(seed: 2929) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifftn(c, s: [3, 4], axes: [0, 1]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [3, 4, 8], + dtype: .float32, + mean: -0.03799237683415413, + minimum: -0.7618478536605835, + maximum: 0.8563539981842041, + absoluteSum: 22.170364379882812, + positionChecksum: 10.169722239176432, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 0.49846258759498596, -0.5838869214057922, -0.09089171886444092, + 0.23879596590995789, 0.07031029462814331, -0.007234610617160797, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [3, 4, 8], + dtype: .float32, + mean: 0.03390979766845703, + minimum: -0.6404374241828918, + maximum: 0.5621622800827026, + absoluteSum: 21.58087730407715, + positionChecksum: 11.274425506591797, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 0.42583590745925903, 0.26459410786628723, -0.007343888282775879, + 0.465649276971817, -0.23035968840122223, 0.15064139664173126, + ]), + tolerance: .float32) + } + } + + @Test("ifftn/axes") + func test_ifftn_axes() throws { + try withIntegrationState(seed: 23441) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.ifftn(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.008047770708799362, + minimum: -0.4037286043167114, + maximum: 0.4154580235481262, + absoluteSum: 49.243099212646484, + positionChecksum: 24.31148910522461, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.015621490776538849, -0.12264824658632278, -0.04659213125705719, + 0.0034754127264022827, -0.09425322711467743, 0.11336925625801086, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 8], + dtype: .float32, + mean: -0.002788180485367775, + minimum: -0.3810046315193176, + maximum: 0.4174731373786926, + absoluteSum: 50.17457962036133, + positionChecksum: 25.259380340576172, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.035300590097904205, 0.08818861842155457, -0.04737217724323273, + -0.02006063610315323, 0.08536103367805481, -0.05704180896282196, + ]), + tolerance: .float32) + } + } + + @Test("rfft2") + func test_rfft2() throws { + try withIntegrationState(seed: 10302) { + let c = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfft2(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.06260216236114502, + minimum: -15.983535766601562, + maximum: 15.228923797607422, + absoluteSum: 1476.531494140625, + positionChecksum: 682.576953125, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -15.983535766601562, 2.8445022106170654, -1.0613014698028564, + -3.684018135070801, 2.3670763969421387, 0.32685017585754395, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.055629778653383255, + minimum: -17.28950309753418, + maximum: 16.831018447875977, + absoluteSum: 1299.924560546875, + positionChecksum: 631.76669921875, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 0.0, 0.0, -0.3817894458770752, -2.158597230911255, -4.935375690460205, + -2.6903560161590576, + ]), + tolerance: .float32) + } + } + + @Test("rfft2/axes") + func test_rfft2_axes() throws { + try withIntegrationState(seed: 47042) { + let c = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfft2(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.7185362577438354, + minimum: -14.52419662475586, + maximum: 16.742055892944336, + absoluteSum: 1495.4569091796875, + positionChecksum: 726.011865234375, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 5.749485015869141, -7.658758640289307, -0.8341364860534668, + 2.9679250717163086, 6.413942813873291, 7.591013431549072, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.23355670273303986, + minimum: -15.147994995117188, + maximum: 14.86296272277832, + absoluteSum: 1337.6533203125, + positionChecksum: 668.728515625, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 0.0, -5.069750785827637, -0.016709327697753906, 2.846343517303467, + 13.277387619018555, 4.648657321929932, + ]), + tolerance: .float32) + } + } + + @Test("rfftn") + func test_rfftn() throws { + try withIntegrationState(seed: 87041) { + let c = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfftn(c) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: -1.1790698766708374, + minimum: -43.51289367675781, + maximum: 40.62009048461914, + absoluteSum: 4045.7822265625, + positionChecksum: 2089.2171875, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 2.826408863067627, 18.93130874633789, -16.50474739074707, 5.28663969039917, + -1.7076082229614258, -12.364112854003906, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: -0.06564798206090927, + minimum: -53.53523254394531, + maximum: 53.53523254394531, + absoluteSum: 4218.8720703125, + positionChecksum: 2227.225, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 0.0, -14.105966567993164, -23.85370635986328, -16.85041046142578, + -1.7126522064208984, 8.338759422302246, + ]), + tolerance: .float32) + } + } + + @Test("rfftn/axes") + func test_rfftn_axes() throws { + try withIntegrationState(seed: 46670) { + let c = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rfftn(c, axes: [0, 2]) + expectSummary( + result.realPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.07727640122175217, + minimum: -16.674985885620117, + maximum: 17.34377098083496, + absoluteSum: 1540.9918212890625, + positionChecksum: 760.8494140625, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + -15.779212951660156, -1.9917529821395874, -7.319093704223633, + 1.1237940788269043, 17.34377098083496, 6.796459197998047, + ]), + tolerance: .float32) + expectSummary( + result.imaginaryPart(), + ArraySummary( + shape: [8, 8, 5], + dtype: .float32, + mean: 0.20542661845684052, + minimum: -16.39087677001953, + maximum: 16.39087677001953, + absoluteSum: 1281.357666015625, + positionChecksum: 672.893212890625, + sampleIndices: [0, 64, 128, 191, 255, 319], + samples: [ + 0.0, -0.4302929639816284, 3.292222499847412, -1.093297004699707, + 16.39087677001953, 9.523144721984863, + ]), + tolerance: .float32) + } + } + + @Test("irfft2") + func test_irfft2() throws { + try withIntegrationState(seed: 34359) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfft2(c) + expectSummary( + result, + ArraySummary( + shape: [8, 8, 14], + dtype: .float32, + mean: -0.004153141751885414, + minimum: -0.4831690192222595, + maximum: 0.39062052965164185, + absoluteSum: 90.77721405029297, + positionChecksum: 45.66217476981027, + sampleIndices: [0, 179, 358, 537, 716, 895], + samples: [ + -0.04847263544797897, 0.20513692498207092, -0.229575976729393, + -0.044226326048374176, -0.07635796815156937, -0.04189996048808098, + ]), + tolerance: .float32) + } + } + + @Test("irfft2/axes") + func test_irfft2_axes() throws { + try withIntegrationState(seed: 42730) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfft2(c, axes: [0, 2]) + expectSummary( + result, + ArraySummary( + shape: [8, 8, 14], + dtype: .float32, + mean: -0.0023053386248648167, + minimum: -0.3896021842956543, + maximum: 0.4151984751224518, + absoluteSum: 91.33279418945312, + positionChecksum: 45.873779296875, + sampleIndices: [0, 179, 358, 537, 716, 895], + samples: [ + -0.11438766866922379, 0.0355401448905468, 0.13291791081428528, + -0.1825760304927826, 0.04413542523980141, -0.03667726740241051, + ]), + tolerance: .float32) + } + } + + @Test("irfftn") + func test_irfftn() throws { + try withIntegrationState(seed: 58696) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfftn(c) + expectSummary( + result, + ArraySummary( + shape: [8, 8, 14], + dtype: .float32, + mean: 0.00038374235737137496, + minimum: -0.13770824670791626, + maximum: 0.16200695931911469, + absoluteSum: 32.69968795776367, + positionChecksum: 16.683267865862167, + sampleIndices: [0, 179, 358, 537, 716, 895], + samples: [ + 0.05379920452833176, 0.045248210430145264, 0.008416865952312946, + 0.0006176760070957243, 0.025639304891228676, -0.06305667012929916, + ]), + tolerance: .float32) + } + } + + @Test("irfftn/axes") + func test_irfftn_axes() throws { + try withIntegrationState(seed: 9062) { + let r = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let i = MLXRandom.normal([8, 8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let c = r + i.asImaginary() + let result = MLX.irfftn(c, axes: [0, 2]) + expectSummary( + result, + ArraySummary( + shape: [8, 8, 14], + dtype: .float32, + mean: 0.0023820779751986265, + minimum: -0.34605321288108826, + maximum: 0.40482693910598755, + absoluteSum: 86.85626220703125, + positionChecksum: 44.459620884486604, + sampleIndices: [0, 179, 358, 537, 716, 895], + samples: [ + -0.09112237393856049, 0.04324769228696823, 0.007158036343753338, + 0.15657760202884674, 0.08600138127803802, -0.22598956525325775, + ]), + tolerance: .float32) + } + } + + @Test("fftfreq") + func test_fftfreq() throws { + try withIntegrationState(seed: 29017) { + let result = MLX.fftfreq(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: -0.03125, + minimum: -0.5, + maximum: 0.4375, + absoluteSum: 4.0, + positionChecksum: 2.25, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [0.0, 0.1875, 0.375, -0.4375, -0.25, -0.0625]), + tolerance: .float32) + } + } + + @Test("fftfreq/d") + func test_fftfreq_d() throws { + try withIntegrationState(seed: 91367) { + let result = MLX.fftfreq(16, d: 0.25) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: -0.125, + minimum: -2.0, + maximum: 1.75, + absoluteSum: 16.0, + positionChecksum: 9.0, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [0.0, 0.75, 1.5, -1.75, -1.0, -0.25]), + tolerance: .float32) + } + } + + @Test("rfftfreq") + func test_rfftfreq() throws { + try withIntegrationState(seed: 94336) { + let result = MLX.rfftfreq(16) + expectSummary( + result, + ArraySummary( + shape: [9], + dtype: .float32, + mean: 0.25, + minimum: 0.0, + maximum: 0.5, + absoluteSum: 2.25, + positionChecksum: 1.6666666666666667, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [0.0, 0.125, 0.1875, 0.3125, 0.375, 0.5]), + tolerance: .float32) + } + } + + @Test("fftshift") + func test_fftshift() throws { + try withIntegrationState(seed: 60895) { + let a = MLXRandom.normal([8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.fftshift(a) + expectSummary( + result, + ArraySummary( + shape: [8, 8], + dtype: .float32, + mean: 0.16273652017116547, + minimum: -2.3805737495422363, + maximum: 2.3832385540008545, + absoluteSum: 52.84351348876953, + positionChecksum: 24.834022521972656, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 1.3552110195159912, 0.08066458255052567, 1.1147830486297607, + 0.0029193360824137926, 0.69813072681427, -0.2748071253299713, + ]), + tolerance: .float32) + } + } + + @Test("fftshift/axes") + func test_fftshift_axes() throws { + try withIntegrationState(seed: 3962) { + let a = MLXRandom.normal([8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.fftshift(a, axes: [0]) + expectSummary( + result, + ArraySummary( + shape: [8, 8], + dtype: .float32, + mean: -0.012623829767107964, + minimum: -1.9678735733032227, + maximum: 1.761767864227295, + absoluteSum: 51.23778533935547, + positionChecksum: 25.559467315673828, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.9273226857185364, 1.414719820022583, 0.45456352829933167, + 1.4807580709457397, -0.10517635941505432, 0.7928916811943054, + ]), + tolerance: .float32) + } + } + + @Test("ifftshift") + func test_ifftshift() throws { + try withIntegrationState(seed: 43884) { + let a = MLXRandom.normal([8, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.ifftshift(a) + expectSummary( + result, + ArraySummary( + shape: [8, 8], + dtype: .float32, + mean: 0.12466083467006683, + minimum: -1.6268000602722168, + maximum: 2.804321765899658, + absoluteSum: 48.001792907714844, + positionChecksum: 23.223478317260742, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.6544477939605713, 0.8364695906639099, 1.0058847665786743, + 2.804321765899658, -0.050703976303339005, -0.33208170533180237, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedFactoryTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedFactoryTests.swift new file mode 100644 index 000000000..27bb07d78 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedFactoryTests.swift @@ -0,0 +1,500 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 23 + +import Foundation +import MLX +import Testing + +@Suite("generated: Factory") +struct GeneratedFactoryTests { + + @Test("arange/stop") + func test_arange_stop() throws { + try withIntegrationState(seed: 76748) { + let result = MLX.arange(10) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 4.5, + minimum: 0.0, + maximum: 9.0, + absoluteSum: 45.0, + positionChecksum: 33.0, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [0.0, 2.0, 4.0, 5.0, 7.0, 9.0]), + tolerance: .exact) + } + } + + @Test("arange/range") + func test_arange_range() throws { + try withIntegrationState(seed: 98523) { + let result = MLX.arange(2, 12) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .int32, + mean: 6.5, + minimum: 2.0, + maximum: 11.0, + absoluteSum: 65.0, + positionChecksum: 44.0, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [2.0, 4.0, 6.0, 7.0, 9.0, 11.0]), + tolerance: .exact) + } + } + + @Test("arange/step") + func test_arange_step() throws { + try withIntegrationState(seed: 7526) { + let result = MLX.arange(0.0, 2.0, step: 0.25) + expectSummary( + result, + ArraySummary( + shape: [8], + dtype: .float32, + mean: 0.875, + minimum: 0.0, + maximum: 1.75, + absoluteSum: 7.0, + positionChecksum: 5.25, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [0.0, 0.25, 0.75, 1.0, 1.5, 1.75]), + tolerance: .float32) + } + } + + @Test("arange/dtype") + func test_arange_dtype() throws { + try withIntegrationState(seed: 62487) { + let result = MLX.arange(0, 10, step: 2, dtype: .int32) + expectSummary( + result, + ArraySummary( + shape: [5], + dtype: .int32, + mean: 4.0, + minimum: 0.0, + maximum: 8.0, + absoluteSum: 20.0, + positionChecksum: 16.0, + sampleIndices: [0, 1, 2, 3, 4], + samples: [0.0, 2.0, 4.0, 6.0, 8.0]), + tolerance: .exact) + } + } + + @Test("linspace") + func test_linspace() throws { + try withIntegrationState(seed: 87908) { + let result = MLX.linspace(0.0, 1.0, count: 9) + expectSummary( + result, + ArraySummary( + shape: [9], + dtype: .float32, + mean: 0.5, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.5, + positionChecksum: 3.3333333333333335, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [0.0, 0.25, 0.375, 0.625, 0.75, 1.0]), + tolerance: .float32) + } + } + + @Test("linspace/int") + func test_linspace_int() throws { + // integer bounds still produce float32, matching python + try withIntegrationState(seed: 96709) { + let result = MLX.linspace(0, 10, count: 6) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .float32, + mean: 5.0, + minimum: 0.0, + maximum: 10.0, + absoluteSum: 30.0, + positionChecksum: 23.333333333333332, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 2.0, 4.0, 6.0, 8.0, 10.0]), + tolerance: .float32) + } + } + + @Test("linspace/dtype") + func test_linspace_dtype() throws { + try withIntegrationState(seed: 61299) { + let result = MLX.linspace(0, 10, count: 6, dtype: .int32) + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .int32, + mean: 5.0, + minimum: 0.0, + maximum: 10.0, + absoluteSum: 30.0, + positionChecksum: 23.333333333333332, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [0.0, 2.0, 4.0, 6.0, 8.0, 10.0]), + tolerance: .exact) + } + } + + @Test("linspace/endpoint") + func test_linspace_endpoint() throws { + // python has no endpoint parameter: the half-open interval is the inclusive one with the last sample dropped + try withIntegrationState(seed: 29128) { + let result = MLX.linspace(0.0, 1.0, count: 5, endpoint: false) + expectSummary( + result, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.4000000059604645, + minimum: 0.0, + maximum: 0.800000011920929, + absoluteSum: 2.0, + positionChecksum: 1.6, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.0, 0.20000000298023224, 0.4000000059604645, 0.6000000238418579, + 0.800000011920929, + ]), + tolerance: .float32) + } + } + + @Test("zeros") + func test_zeros() throws { + try withIntegrationState(seed: 87556) { + let result = MLX.zeros([3, 4]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("zeros/dtype") + func test_zeros_dtype() throws { + try withIntegrationState(seed: 92392) { + let result = MLX.zeros([3, 4], dtype: .int32) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .int32, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("ones") + func test_ones() throws { + try withIntegrationState(seed: 95935) { + let result = MLX.ones([3, 4]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 12.0, + positionChecksum: 6.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), + tolerance: .float32) + } + } + + @Test("ones/dtype") + func test_ones_dtype() throws { + try withIntegrationState(seed: 51426) { + let result = MLX.ones([3, 4], dtype: .int16) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .int16, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 12.0, + positionChecksum: 6.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("zeros/like") + func test_zeros_like() throws { + try withIntegrationState(seed: 79637) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.zeros(like: a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("ones/like") + func test_ones_like() throws { + try withIntegrationState(seed: 45715) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.ones(like: a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 12.0, + positionChecksum: 6.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), + tolerance: .float32) + } + } + + @Test("full") + func test_full() throws { + try withIntegrationState(seed: 96888) { + let result = MLX.full([2, 3], values: Float(2.5)) + expectSummary( + result, + ArraySummary( + shape: [2, 3], + dtype: .float32, + mean: 2.5, + minimum: 2.5, + maximum: 2.5, + absoluteSum: 15.0, + positionChecksum: 8.75, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [2.5, 2.5, 2.5, 2.5, 2.5, 2.5]), + tolerance: .float32) + } + } + + @Test("full/dtype") + func test_full_dtype() throws { + try withIntegrationState(seed: 47656) { + let result = MLX.full([2, 3], values: MLXArray(7), dtype: .int32) + expectSummary( + result, + ArraySummary( + shape: [2, 3], + dtype: .int32, + mean: 7.0, + minimum: 7.0, + maximum: 7.0, + absoluteSum: 42.0, + positionChecksum: 24.5, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [7.0, 7.0, 7.0, 7.0, 7.0, 7.0]), + tolerance: .exact) + } + } + + @Test("eye") + func test_eye() throws { + try withIntegrationState(seed: 78454) { + let result = MLX.eye(4) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.25, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.0, + positionChecksum: 2.125, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [1.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .float32) + } + } + + @Test("eye/rectangular") + func test_eye_rectangular() throws { + try withIntegrationState(seed: 51823) { + let result = MLX.eye(4, m: 6, k: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 6], + dtype: .float32, + mean: 0.1666666716337204, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 4.0, + positionChecksum: 2.0833333333333335, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("identity") + func test_identity() throws { + try withIntegrationState(seed: 81094) { + let result = MLX.identity(5) + expectSummary( + result, + ArraySummary( + shape: [5, 5], + dtype: .float32, + mean: 0.19999998807907104, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.6, + sampleIndices: [0, 5, 10, 14, 19, 24], + samples: [1.0, 0.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .float32) + } + } + + @Test("bartlett") + func test_bartlett() throws { + try withIntegrationState(seed: 66814) { + let result = MLX.bartlett(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: 0.46666663885116577, + minimum: 0.0, + maximum: 0.9333333969116211, + absoluteSum: 7.466666221618652, + positionChecksum: 3.9666662216186523, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.0, 0.40000003576278687, 0.8000000715255737, 0.7999999523162842, + 0.39999985694885254, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("hanning") + func test_hanning() throws { + try withIntegrationState(seed: 34335) { + let result = MLX.hanning(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: 0.4687499701976776, + minimum: 0.0, + maximum: 0.9890736937522888, + absoluteSum: 7.499999523162842, + positionChecksum: 3.984375, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.0, 0.34549155831336975, 0.904508650302887, 0.9045085310935974, + 0.3454914093017578, 7.64274186065882e-15, + ]), + tolerance: .float32) + } + } + + @Test("hamming") + func test_hamming() throws { + try withIntegrationState(seed: 86817) { + let result = MLX.hamming(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: 0.5112500190734863, + minimum: 0.08000001311302185, + maximum: 0.9899479746818542, + absoluteSum: 8.180000305175781, + positionChecksum: 4.345624923706055, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.08000001311302185, 0.39785224199295044, 0.9121478796005249, + 0.9121477603912354, 0.3978521227836609, 0.08000001311302185, + ]), + tolerance: .float32) + } + } + + @Test("blackman") + func test_blackman() throws { + try withIntegrationState(seed: 98518) { + let result = MLX.blackman(16) + expectSummary( + result, + ArraySummary( + shape: [16], + dtype: .float32, + mean: 0.39375001192092896, + minimum: 0.0, + maximum: 0.9821575284004211, + absoluteSum: 6.300000190734863, + positionChecksum: 3.346874713897705, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.0, 0.20077016949653625, 0.8492298722267151, 0.8492297530174255, + 0.20077009499073029, 0.0, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift new file mode 100644 index 000000000..6a0102a4a --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedFastTests.swift @@ -0,0 +1,201 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 7 + +import Foundation +import MLX +import Testing + +@Suite("generated: Fast") +struct GeneratedFastTests { + + @Test("rmsNorm") + func test_rmsNorm() throws { + try withIntegrationState(seed: 66852) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.rmsNorm(x, weight: weight, eps: 1e-5) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.06485586613416672, + minimum: -3.9406793117523193, + maximum: 3.7056970596313477, + absoluteSum: 191.23782348632812, + positionChecksum: 95.85537719726562, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -1.0069217681884766, -0.8062147498130798, 1.2844102382659912, + 0.10886163264513016, 1.0201539993286133, -0.6644749045372009, + ]), + tolerance: .float32) + } + } + + @Test("layerNorm") + func test_layerNorm() throws { + try withIntegrationState(seed: 45114) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let weight = MLXRandom.normal([16], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.layerNorm(x, weight: weight, bias: bias, eps: 1e-5) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.23579367995262146, + minimum: -3.412724733352661, + maximum: 3.5934901237487793, + absoluteSum: 225.9371337890625, + positionChecksum: 108.25701904296875, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.18614619970321655, -0.702523946762085, -0.5655195713043213, + 1.331803560256958, -1.6021459102630615, -0.23616120219230652, + ]), + tolerance: .float32) + } + } + + @Test("layerNorm/noAffine") + func test_layerNorm_noAffine() throws { + try withIntegrationState(seed: 56298) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.layerNorm(x, eps: 1e-5) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -6.51925802230835e-09, + minimum: -2.750777006149292, + maximum: 2.616267681121826, + absoluteSum: 200.528564453125, + positionChecksum: 99.69488525390625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.5727211833000183, 0.3872288167476654, 2.390326976776123, + 1.077345848083496, -1.7498226165771484, -0.0715562179684639, + ]), + tolerance: .float32) + } + } + + @Test("RoPE") + func test_RoPE() throws { + try withIntegrationState(seed: 66733) { + let x = MLXRandom.normal([2, 4, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.RoPE( + x, dimensions: 16, traditional: false, base: 10000.0, scale: 1.0, offset: 0) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 16], + dtype: .float32, + mean: -0.00884411484003067, + minimum: -3.2420690059661865, + maximum: 3.197585344314575, + absoluteSum: 821.5716552734375, + positionChecksum: 414.88031005859375, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + -0.6869831681251526, 0.25631067156791687, 0.8028337955474854, + -0.9851393699645996, -0.39124467968940735, 2.371519088745117, + ]), + tolerance: .float32) + } + } + + @Test("RoPE/traditional") + func test_RoPE_traditional() throws { + try withIntegrationState(seed: 74377) { + let x = MLXRandom.normal([2, 4, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.RoPE( + x, dimensions: 8, traditional: true, base: 10000.0, scale: 0.5, offset: 2) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 16], + dtype: .float32, + mean: -0.028798004612326622, + minimum: -2.746978521347046, + maximum: 3.72855281829834, + absoluteSum: 807.062744140625, + positionChecksum: 392.5193786621094, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + -1.4806770086288452, -0.9948027729988098, 1.3070459365844727, + -0.7166029214859009, 0.09865694493055344, -2.09395694732666, + ]), + tolerance: .float32) + } + } + + @Test("scaledDotProductAttention") + func test_scaledDotProductAttention() throws { + try withIntegrationState(seed: 82359) { + let queries = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let values = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.scaledDotProductAttention( + queries: queries, keys: keys, values: values, scale: 0.35, mask: nil) + expectSummary( + result, + ArraySummary( + shape: [1, 2, 4, 8], + dtype: .float32, + mean: 0.197175532579422, + minimum: -1.386329174041748, + maximum: 2.5051798820495605, + absoluteSum: 37.204132080078125, + positionChecksum: 20.866201400756836, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 0.5489547848701477, 0.5017918944358826, 0.6151727437973022, + -0.010526720434427261, 2.5051798820495605, 0.639729380607605, + ]), + tolerance: .float32) + } + } + + @Test("scaledDotProductAttention/mask") + func test_scaledDotProductAttention_mask() throws { + try withIntegrationState(seed: 63535) { + let queries = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let keys = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let values = MLXRandom.normal([1, 2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MLX.tril(MLX.ones([4, 4])) .== 1 + let result = MLX.scaledDotProductAttention( + queries: queries, keys: keys, values: values, scale: 0.35, mask: mask) + expectSummary( + result, + ArraySummary( + shape: [1, 2, 4, 8], + dtype: .float32, + mean: 0.0719122439622879, + minimum: -1.5812186002731323, + maximum: 1.1895414590835571, + absoluteSum: 33.8242301940918, + positionChecksum: 17.089391708374023, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.4037952125072479, 0.0268552303314209, -0.7593691945075989, + -0.8534163236618042, -0.503488302230835, 0.6569595336914062, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedIndexingTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedIndexingTests.swift new file mode 100644 index 000000000..e5c2b106c --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedIndexingTests.swift @@ -0,0 +1,166 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 6 + +import Foundation +import MLX +import Testing + +@Suite("generated: Indexing") +struct GeneratedIndexingTests { + + @Test("index/range") + func test_index_range() throws { + try withIntegrationState(seed: 56799) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a[1 ..< 4, 0 ..< 3] + expectSummary( + result, + ArraySummary( + shape: [3, 3], + dtype: .float32, + mean: 0.20556111633777618, + minimum: -1.9838948249816895, + maximum: 2.7164838314056396, + absoluteSum: 10.062643051147461, + positionChecksum: 5.407042609320746, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [ + -0.7313419580459595, 1.4610944986343384, -1.9838948249816895, + 2.7164838314056396, -1.3910598754882812, 0.4201461672782898, + ]), + tolerance: .float32) + } + } + + @Test("index/lastColumn") + func test_index_lastColumn() throws { + try withIntegrationState(seed: 67368) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a[0..., -1] + expectSummary( + result, + ArraySummary( + shape: [6], + dtype: .float32, + mean: -0.07870198786258698, + minimum: -1.1848795413970947, + maximum: 1.5605968236923218, + absoluteSum: 3.762835741043091, + positionChecksum: 1.838436762491862, + sampleIndices: [0, 1, 2, 3, 4, 5], + samples: [ + -0.4085389971733093, -1.1848795413970947, -0.5241052508354187, + 1.5605968236923218, 0.07067099958658218, 0.014044167473912239, + ]), + tolerance: .float32) + } + } + + @Test("index/newAxis") + func test_index_newAxis() throws { + try withIntegrationState(seed: 7679) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a[.ellipsis, .newAxis] + expectSummary( + result, + ArraySummary( + shape: [6, 5, 1], + dtype: .float32, + mean: -0.09582757204771042, + minimum: -1.8003718852996826, + maximum: 1.4870085716247559, + absoluteSum: 22.161052703857422, + positionChecksum: 11.798158772786458, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.08150071650743484, 0.26980727910995483, 0.2373223602771759, + -0.4347092807292938, 1.3970370292663574, -0.6056225299835205, + ]), + tolerance: .float32) + } + } + + @Test("index/stride") + func test_index_stride() throws { + try withIntegrationState(seed: 33648) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a[.stride(by: 2)] + expectSummary( + result, + ArraySummary( + shape: [3, 5], + dtype: .float32, + mean: 0.07586007565259933, + minimum: -1.3099104166030884, + maximum: 2.5639853477478027, + absoluteSum: 12.466233253479004, + positionChecksum: 6.5820460001627605, + sampleIndices: [0, 3, 6, 8, 11, 14], + samples: [ + -0.0001792133116396144, -0.5972116589546204, -0.8127403259277344, + 2.5639853477478027, 0.8785732388496399, -0.3988332748413086, + ]), + tolerance: .float32) + } + } + + @Test("index/reverse") + func test_index_reverse() throws { + try withIntegrationState(seed: 42592) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a[.ellipsis, .stride(by: -1)] + expectSummary( + result, + ArraySummary( + shape: [6, 5], + dtype: .float32, + mean: -0.13014544546604156, + minimum: -2.513247489929199, + maximum: 1.7819768190383911, + absoluteSum: 19.7058162689209, + positionChecksum: 8.921414184570313, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 1.695084571838379, -0.2929709255695343, 0.3612818419933319, + 0.3690113425254822, -0.16253980994224548, 0.26952359080314636, + ]), + tolerance: .float32) + } + } + + @Test("index/array") + func test_index_array() throws { + try withIntegrationState(seed: 61633) { + let a = MLXRandom.normal([6, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 6, [4], type: Int32.self) + let result = a[indices] + expectSummary( + result, + ArraySummary( + shape: [4, 5], + dtype: .float32, + mean: 0.04731393977999687, + minimum: -1.259320855140686, + maximum: 2.2643637657165527, + absoluteSum: 13.387282371520996, + positionChecksum: 7.656300354003906, + sampleIndices: [0, 4, 8, 11, 15, 19], + samples: [ + 0.24186666309833527, -0.5770583152770996, 0.4117497503757477, + -1.259320855140686, -0.39659786224365234, 2.2643637657165527, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedLinalgTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedLinalgTests.swift new file mode 100644 index 000000000..2be56a9e6 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedLinalgTests.swift @@ -0,0 +1,405 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 16 + +import Foundation +import MLX +import Testing + +@Suite("generated: Linalg") +struct GeneratedLinalgTests { + + @Test("norm") + func test_norm() throws { + try withIntegrationState(seed: 61303) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2.4654452800750732, + minimum: 2.4654452800750732, + maximum: 2.4654452800750732, + absoluteSum: 2.4654452800750732, + positionChecksum: 2.4654452800750732, + sampleIndices: [0], + samples: [2.4654452800750732]), + tolerance: .float32) + } + } + + @Test("norm/fro") + func test_norm_fro() throws { + try withIntegrationState(seed: 13358) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, ord: .fro, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 3.750117301940918, + minimum: 3.750117301940918, + maximum: 3.750117301940918, + absoluteSum: 3.750117301940918, + positionChecksum: 3.750117301940918, + sampleIndices: [0], + samples: [3.750117301940918]), + tolerance: .float32) + } + } + + @Test("norm/ord1/axis") + func test_norm_ord1_axis() throws { + try withIntegrationState(seed: 50165) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, ord: 1.0, axis: 0, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .float32, + mean: 2.603952407836914, + minimum: 1.1081544160842896, + maximum: 4.008610248565674, + absoluteSum: 7.811856746673584, + positionChecksum: 4.678925196329753, + sampleIndices: [0, 1, 2], + samples: [2.69509220123291, 4.008610248565674, 1.1081544160842896]), + tolerance: .float32) + } + } + + @Test("norm/axes") + func test_norm_axes() throws { + try withIntegrationState(seed: 97742) { + let a = MLXRandom.normal([2, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, axes: [1, 2], stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [2], + dtype: .float32, + mean: 3.1824984550476074, + minimum: 2.9801127910614014, + maximum: 3.3848843574523926, + absoluteSum: 6.364996910095215, + positionChecksum: 4.672554969787598, + sampleIndices: [0, 1], + samples: [3.3848843574523926, 2.9801127910614014]), + tolerance: .float32) + } + } + + @Test("norm/keepDims") + func test_norm_keepDims() throws { + try withIntegrationState(seed: 52042) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.norm(a, axis: -1, keepDims: true, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 1], + dtype: .float32, + mean: 1.7273744344711304, + minimum: 0.8218560218811035, + maximum: 2.5514132976531982, + absoluteSum: 6.9094977378845215, + positionChecksum: 3.605181932449341, + sampleIndices: [0, 1, 2, 3], + samples: [ + 2.5514132976531982, 2.026794910430908, 1.5094335079193115, + 0.8218560218811035, + ]), + tolerance: .float32) + } + } + + @Test("inv") + func test_inv() throws { + try withIntegrationState(seed: 45668) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.inv(spd, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.03864220157265663, + minimum: -0.040189146995544434, + maximum: 0.21930460631847382, + absoluteSum: 0.9170589447021484, + positionChecksum: 0.4542868137359619, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.18467867374420166, -0.01917562633752823, 0.0063242255710065365, + 0.006324225105345249, -0.01917562447488308, 0.15314576029777527, + ]), + tolerance: .float32) + } + } + + @Test("triInv") + func test_triInv() throws { + try withIntegrationState(seed: 99938) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let lower = MLX.tril(spd) + let result = MLX.triInv(lower, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.04744876176118851, + minimum: -0.017073122784495354, + maximum: 0.20120765268802643, + absoluteSum: 0.8115707635879517, + positionChecksum: 0.4696817994117737, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.16164816915988922, 0.0, 0.0, 0.03999091684818268, 0.006804905831813812, + 0.20120765268802643, + ]), + tolerance: .float32) + } + } + + @Test("cholesky") + func test_cholesky() throws { + try withIntegrationState(seed: 62767) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.cholesky(spd, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.8446391820907593, + minimum: -0.9616007208824158, + maximum: 3.0558056831359863, + absoluteSum: 17.262596130371094, + positionChecksum: 10.090898513793945, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 3.0558056831359863, 0.0, 0.0, -0.9125845432281494, 0.9799381494522095, + 2.471771001815796, + ]), + tolerance: .float32) + } + } + + @Test("cholesky/upper") + func test_cholesky_upper() throws { + try withIntegrationState(seed: 38847) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.cholesky(spd, upper: true, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.7331788539886475, + minimum: -0.7065603137016296, + maximum: 3.136165142059326, + absoluteSum: 13.893322944641113, + positionChecksum: 6.400030612945557, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 3.0352115631103516, -0.1763664186000824, -0.03923870995640755, 0.0, 0.0, + 2.291856050491333, + ]), + tolerance: .float32) + } + } + + @Test("choleskyInv") + func test_choleskyInv() throws { + try withIntegrationState(seed: 92684) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let l = MLX.cholesky(spd, stream: .cpu) + let result = MLX.choleskyInv(l, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.037270255386829376, + minimum: -0.026440046727657318, + maximum: 0.20521625876426697, + absoluteSum: 0.795662522315979, + positionChecksum: 0.3411255478858948, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 0.20521625876426697, -0.01557123102247715, -0.026440046727657318, + -0.026440046727657318, -0.01557123102247715, 0.07566997408866882, + ]), + tolerance: .float32) + } + } + + @Test("pinv") + func test_pinv() throws { + try withIntegrationState(seed: 49272) { + let a = MLXRandom.normal([5, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.pinv(a, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [3, 5], + dtype: .float32, + mean: 0.14109282195568085, + minimum: -0.7837748527526855, + maximum: 0.786260187625885, + absoluteSum: 5.669179916381836, + positionChecksum: 2.6510897318522137, + sampleIndices: [0, 3, 6, 8, 11, 14], + samples: [ + -0.7837748527526855, -0.39348694682121277, 0.19790665805339813, + 0.3927387297153473, 0.528360903263092, 0.377410352230072, + ]), + tolerance: .float32) + } + } + + @Test("cross") + func test_cross() throws { + try withIntegrationState(seed: 83359) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cross(a, b) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.19259056448936462, + minimum: -0.550257682800293, + maximum: 3.921427011489868, + absoluteSum: 6.5804123878479, + positionChecksum: 4.06615416208903, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.028709368780255318, -0.5001630187034607, -0.550257682800293, + 3.921427011489868, 0.14121165871620178, 0.2556290924549103, + ]), + tolerance: .float32) + } + } + + @Test("solve") + func test_solve() throws { + try withIntegrationState(seed: 29554) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let b = MLXRandom.normal([4, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.solve(spd, b, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: -0.011842835694551468, + minimum: -0.191814124584198, + maximum: 0.21963128447532654, + absoluteSum: 1.128840684890747, + positionChecksum: 0.6843239068984985, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + -0.18137195706367493, -0.09362111240625381, -0.10529350489377975, + -0.191814124584198, 0.09047023206949234, 0.2069474756717682, + ]), + tolerance: .float32) + } + } + + @Test("solveTriangular") + func test_solveTriangular() throws { + try withIntegrationState(seed: 13290) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let lower = MLX.tril(spd) + let b = MLXRandom.normal([4, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.solveTriangular(lower, b, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: 0.061786651611328125, + minimum: -0.15068136155605316, + maximum: 0.28889700770378113, + absoluteSum: 0.938929557800293, + positionChecksum: 0.47943881154060364, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + 0.06748572736978531, 0.28889700770378113, -0.03723353147506714, + 0.017042802646756172, -0.034403301775455475, 0.1020500585436821, + ]), + tolerance: .float32) + } + } + + @Test("det") + func test_det() throws { + try withIntegrationState(seed: 74576) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.det(spd, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2897.23046875, + minimum: 2897.23046875, + maximum: 2897.23046875, + absoluteSum: 2897.23046875, + positionChecksum: 2897.23046875, + sampleIndices: [0], + samples: [2897.23046875]), + tolerance: .float32) + } + } + + @Test("eigvalsh") + func test_eigvalsh() throws { + // symmetric input: eigenvalues are returned in ascending order + try withIntegrationState(seed: 56419) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let spd = MLX.matmul(a, a.T) + 4.0 * MLX.eye(4) + let result = MLX.eigvalsh(spd, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 8.057674407958984, + minimum: 4.122847080230713, + maximum: 15.702493667602539, + absoluteSum: 32.23069763183594, + positionChecksum: 24.50469970703125, + sampleIndices: [0, 1, 2, 3], + samples: [ + 4.122847080230713, 6.13009786605835, 6.275260925292969, 15.702493667602539, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedLossesTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedLossesTests.swift new file mode 100644 index 000000000..8555b75f2 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedLossesTests.swift @@ -0,0 +1,555 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 23 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: Losses") +struct GeneratedLossesTests { + + @Test("crossEntropy/none") + func test_crossEntropy_none() throws { + try withIntegrationState(seed: 30964) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.crossEntropy(logits: logits, targets: targets, reduction: .none) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.681290864944458, + minimum: 1.0067064762115479, + maximum: 2.6064746379852295, + absoluteSum: 6.725163459777832, + positionChecksum: 3.676652669906616, + sampleIndices: [0, 1, 2, 3], + samples: [ + 2.6064746379852295, 1.2626385688781738, 1.8493441343307495, + 1.0067064762115479, + ]), + tolerance: .float32) + } + } + + @Test("mseLoss/none") + func test_mseLoss_none() throws { + try withIntegrationState(seed: 14505) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.mseLoss(predictions: predictions, targets: targets, reduction: .none) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 2.298888921737671, + minimum: 0.011387436650693417, + maximum: 6.817906379699707, + absoluteSum: 27.586666107177734, + positionChecksum: 15.784774780273438, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.08654762804508209, 4.789333820343018, 0.49890097975730896, + 1.4658656120300293, 4.84807825088501, 1.1505988836288452, + ]), + tolerance: .float32) + } + } + + @Test("l1Loss/none") + func test_l1Loss_none() throws { + try withIntegrationState(seed: 96391) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.l1Loss(predictions: predictions, targets: targets, reduction: .none) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.1529018878936768, + minimum: 0.06354738771915436, + maximum: 2.881420135498047, + absoluteSum: 13.834822654724121, + positionChecksum: 6.185560862223308, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.8485625982284546, 1.470560073852539, 0.10926105082035065, + 0.631935179233551, 0.40655890107154846, 1.2638256549835205, + ]), + tolerance: .float32) + } + } + + @Test("crossEntropy/mean") + func test_crossEntropy_mean() throws { + try withIntegrationState(seed: 20203) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.crossEntropy(logits: logits, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.336897850036621, + minimum: 1.336897850036621, + maximum: 1.336897850036621, + absoluteSum: 1.336897850036621, + positionChecksum: 1.336897850036621, + sampleIndices: [0], + samples: [1.336897850036621]), + tolerance: .float32) + } + } + + @Test("mseLoss/mean") + func test_mseLoss_mean() throws { + try withIntegrationState(seed: 21174) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.mseLoss(predictions: predictions, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.5020225048065186, + minimum: 1.5020225048065186, + maximum: 1.5020225048065186, + absoluteSum: 1.5020225048065186, + positionChecksum: 1.5020225048065186, + sampleIndices: [0], + samples: [1.5020225048065186]), + tolerance: .float32) + } + } + + @Test("l1Loss/mean") + func test_l1Loss_mean() throws { + try withIntegrationState(seed: 72088) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.l1Loss(predictions: predictions, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.33334481716156, + minimum: 1.33334481716156, + maximum: 1.33334481716156, + absoluteSum: 1.33334481716156, + positionChecksum: 1.33334481716156, + sampleIndices: [0], + samples: [1.33334481716156]), + tolerance: .float32) + } + } + + @Test("crossEntropy/sum") + func test_crossEntropy_sum() throws { + try withIntegrationState(seed: 79221) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.crossEntropy(logits: logits, targets: targets, reduction: .sum) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 7.240591049194336, + minimum: 7.240591049194336, + maximum: 7.240591049194336, + absoluteSum: 7.240591049194336, + positionChecksum: 7.240591049194336, + sampleIndices: [0], + samples: [7.240591049194336]), + tolerance: .float32) + } + } + + @Test("mseLoss/sum") + func test_mseLoss_sum() throws { + try withIntegrationState(seed: 59747) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.mseLoss(predictions: predictions, targets: targets, reduction: .sum) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 35.32448959350586, + minimum: 35.32448959350586, + maximum: 35.32448959350586, + absoluteSum: 35.32448959350586, + positionChecksum: 35.32448959350586, + sampleIndices: [0], + samples: [35.32448959350586]), + tolerance: .float32) + } + } + + @Test("l1Loss/sum") + func test_l1Loss_sum() throws { + try withIntegrationState(seed: 61040) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.l1Loss(predictions: predictions, targets: targets, reduction: .sum) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 11.155969619750977, + minimum: 11.155969619750977, + maximum: 11.155969619750977, + absoluteSum: 11.155969619750977, + positionChecksum: 11.155969619750977, + sampleIndices: [0], + samples: [11.155969619750977]), + tolerance: .float32) + } + } + + @Test("crossEntropy/weights") + func test_crossEntropy_weights() throws { + try withIntegrationState(seed: 88185) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let weights = MLXRandom.uniform(low: 0.5, high: 1.5, [4], dtype: .float32) + let result = MLXNN.crossEntropy( + logits: logits, targets: targets, weights: weights, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2.4085261821746826, + minimum: 2.4085261821746826, + maximum: 2.4085261821746826, + absoluteSum: 2.4085261821746826, + positionChecksum: 2.4085261821746826, + sampleIndices: [0], + samples: [2.4085261821746826]), + tolerance: .float32) + } + } + + @Test("crossEntropy/labelSmoothing") + func test_crossEntropy_labelSmoothing() throws { + try withIntegrationState(seed: 27380) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.crossEntropy( + logits: logits, targets: targets, labelSmoothing: 0.1, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.5042340755462646, + minimum: 1.5042340755462646, + maximum: 1.5042340755462646, + absoluteSum: 1.5042340755462646, + positionChecksum: 1.5042340755462646, + sampleIndices: [0], + samples: [1.5042340755462646]), + tolerance: .float32) + } + } + + @Test("crossEntropy/probabilities") + func test_crossEntropy_probabilities() throws { + try withIntegrationState(seed: 36) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let t = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLX.softmax(t, axis: -1) + let result = MLXNN.crossEntropy(logits: logits, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 2.351641893386841, + minimum: 2.351641893386841, + maximum: 2.351641893386841, + absoluteSum: 2.351641893386841, + positionChecksum: 2.351641893386841, + sampleIndices: [0], + samples: [2.351641893386841]), + tolerance: .float32) + } + } + + @Test("binaryCrossEntropy/logits") + func test_binaryCrossEntropy_logits() throws { + try withIntegrationState(seed: 63990) { + let logits = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLXNN.binaryCrossEntropy( + logits: logits, targets: targets.asType(.float32), reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.7426967620849609, + minimum: 0.7426967620849609, + maximum: 0.7426967620849609, + absoluteSum: 0.7426967620849609, + positionChecksum: 0.7426967620849609, + sampleIndices: [0], + samples: [0.7426967620849609]), + tolerance: .float32) + } + } + + @Test("binaryCrossEntropy/probabilities") + func test_binaryCrossEntropy_probabilities() throws { + try withIntegrationState(seed: 37781) { + let p = MLXRandom.uniform(low: 0.1, high: 0.9, [4, 3], dtype: .float32) + let targets = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLXNN.binaryCrossEntropy( + logits: p, targets: targets.asType(.float32), withLogits: false, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.8758541941642761, + minimum: 0.8758541941642761, + maximum: 0.8758541941642761, + absoluteSum: 0.8758541941642761, + positionChecksum: 0.8758541941642761, + sampleIndices: [0], + samples: [0.8758541941642761]), + tolerance: .float32) + } + } + + @Test("nllLoss") + func test_nllLoss() throws { + try withIntegrationState(seed: 70502) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let inputs = MLX.log(MLX.softmax(logits, axis: -1)) + let targets = MLXRandom.randInt(low: 0, high: 5, [4], type: Int32.self) + let result = MLXNN.nllLoss(inputs: inputs, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.6425096988677979, + minimum: 1.6425096988677979, + maximum: 1.6425096988677979, + absoluteSum: 1.6425096988677979, + positionChecksum: 1.6425096988677979, + sampleIndices: [0], + samples: [1.6425096988677979]), + tolerance: .float32) + } + } + + @Test("klDivLoss") + func test_klDivLoss() throws { + try withIntegrationState(seed: 90379) { + let a = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let inputs = MLX.log(MLX.softmax(a, axis: -1)) + let targets = MLX.log(MLX.softmax(b, axis: -1)) + let result = MLXNN.klDivLoss(inputs: inputs, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.4696533679962158, + minimum: 0.4696533679962158, + maximum: 0.4696533679962158, + absoluteSum: 0.4696533679962158, + positionChecksum: 0.4696533679962158, + sampleIndices: [0], + samples: [0.4696533679962158]), + tolerance: .float32) + } + } + + @Test("smoothL1Loss") + func test_smoothL1Loss() throws { + try withIntegrationState(seed: 16579) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.smoothL1Loss( + predictions: predictions, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.2783452272415161, + minimum: 1.2783452272415161, + maximum: 1.2783452272415161, + absoluteSum: 1.2783452272415161, + positionChecksum: 1.2783452272415161, + sampleIndices: [0], + samples: [1.2783452272415161]), + tolerance: .float32) + } + } + + @Test("smoothL1Loss/beta") + func test_smoothL1Loss_beta() throws { + try withIntegrationState(seed: 63527) { + let predictions = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.smoothL1Loss( + predictions: predictions, targets: targets, beta: 0.5, reduction: .none) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0599005222320557, + minimum: 0.011740408837795258, + maximum: 3.1852192878723145, + absoluteSum: 12.718805313110352, + positionChecksum: 6.934684753417969, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.0957175493240356, 0.011740408837795258, 1.24065363407135, + 0.014730868861079216, 1.9761757850646973, 0.6690434217453003, + ]), + tolerance: .float32) + } + } + + @Test("tripletLoss") + func test_tripletLoss() throws { + try withIntegrationState(seed: 98850) { + let anchors = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let positives = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let negatives = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.tripletLoss( + anchors: anchors, positives: positives, negatives: negatives, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.7783023118972778, + minimum: 1.7783023118972778, + maximum: 1.7783023118972778, + absoluteSum: 1.7783023118972778, + positionChecksum: 1.7783023118972778, + sampleIndices: [0], + samples: [1.7783023118972778]), + tolerance: .float32) + } + } + + @Test("hingeLoss") + func test_hingeLoss() throws { + try withIntegrationState(seed: 90714) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let signs = MLXRandom.bernoulli(0.5, [4, 3]) + let targets = MLX.which(signs, 1.0, -1.0) + let result = MLXNN.hingeLoss(inputs: inputs, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.9594718813896179, + minimum: 0.9594718813896179, + maximum: 0.9594718813896179, + absoluteSum: 0.9594718813896179, + positionChecksum: 0.9594718813896179, + sampleIndices: [0], + samples: [0.9594718813896179]), + tolerance: .float32) + } + } + + @Test("huberLoss") + func test_huberLoss() throws { + try withIntegrationState(seed: 48851) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.huberLoss( + inputs: inputs, targets: targets, delta: 0.5, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.3303348124027252, + minimum: 0.3303348124027252, + maximum: 0.3303348124027252, + absoluteSum: 0.3303348124027252, + positionChecksum: 0.3303348124027252, + sampleIndices: [0], + samples: [0.3303348124027252]), + tolerance: .float32) + } + } + + @Test("logCoshLoss") + func test_logCoshLoss() throws { + try withIntegrationState(seed: 71520) { + let inputs = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let targets = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.logCoshLoss(inputs: inputs, targets: targets, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.8429071307182312, + minimum: 0.8429071307182312, + maximum: 0.8429071307182312, + absoluteSum: 0.8429071307182312, + positionChecksum: 0.8429071307182312, + sampleIndices: [0], + samples: [0.8429071307182312]), + tolerance: .float32) + } + } + + @Test("cosineSimilarityLoss") + func test_cosineSimilarityLoss() throws { + try withIntegrationState(seed: 8460) { + let x1 = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let x2 = MLXRandom.normal([4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXNN.cosineSimilarityLoss(x1: x1, x2: x2, reduction: .mean) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.043844811618328094, + minimum: -0.043844811618328094, + maximum: -0.043844811618328094, + absoluteSum: 0.043844811618328094, + positionChecksum: 0.043844811618328094, + sampleIndices: [0], + samples: [-0.043844811618328094]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleActivationsTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleActivationsTests.swift new file mode 100644 index 000000000..ad9dbb228 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleActivationsTests.swift @@ -0,0 +1,965 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 34 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleActivations") +struct GeneratedModuleActivationsTests { + + @Test("Identity") + func test_Identity() throws { + try withIntegrationState(seed: 34070) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Identity() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.017332343384623528, + minimum: -3.227935791015625, + maximum: 2.948580741882324, + absoluteSum: 204.18060302734375, + positionChecksum: 105.40927124023438, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.3334098160266876, 0.5547546148300171, -0.570229172706604, + -0.4087222218513489, -1.4258067607879639, -0.2681485116481781, + ]), + tolerance: .float32) + } + } + + @Test("Sigmoid") + func test_Sigmoid() throws { + try withIntegrationState(seed: 9913) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Sigmoid() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.5089752078056335, + minimum: 0.04935688152909279, + maximum: 0.9119895696640015, + absoluteSum: 130.2976531982422, + positionChecksum: 67.42926788330078, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.18784688413143158, 0.2168743908405304, 0.7701733708381653, + 0.1441628485918045, 0.6048126220703125, 0.3422277867794037, + ]), + tolerance: .float32) + } + } + + @Test("ReLU") + func test_ReLU() throws { + try withIntegrationState(seed: 16541) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ReLU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.3870805501937866, + minimum: 0.0, + maximum: 2.8602685928344727, + absoluteSum: 99.09262084960938, + positionChecksum: 47.08649444580078, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.8868525624275208, 1.1041529178619385, 0.0, 1.1260745525360107, 0.0, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("ReLU6") + func test_ReLU6() throws { + try withIntegrationState(seed: 54790) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ReLU6() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.38837987184524536, + minimum: 0.0, + maximum: 2.9087584018707275, + absoluteSum: 99.42524719238281, + positionChecksum: 49.03082275390625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.0, 0.0, 0.3142980933189392, 0.0, 0.22747474908828735, 0.24886935949325562, + ]), + tolerance: .float32) + } + } + + @Test("ReLUSquared") + func test_ReLUSquared() throws { + try withIntegrationState(seed: 86178) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ReLUSquared() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.5489815473556519, + minimum: 0.0, + maximum: 6.659243106842041, + absoluteSum: 140.53927612304688, + positionChecksum: 72.461181640625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.0, 0.13198506832122803, 2.2052338123321533, 0.0, 0.0, 0.19578227400779724, + ]), + tolerance: .float32) + } + } + + @Test("LeakyReLU") + func test_LeakyReLU() throws { + try withIntegrationState(seed: 42095) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LeakyReLU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.4470902681350708, + minimum: -0.03602779656648636, + maximum: 2.8485965728759766, + absoluteSum: 116.6517333984375, + positionChecksum: 62.598758697509766, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.01762707717716694, 0.8074039816856384, 0.26561081409454346, + 0.8008984327316284, -0.017988381907343864, 1.6079342365264893, + ]), + tolerance: .float32) + } + } + + @Test("LeakyReLU/slope") + func test_LeakyReLU_slope() throws { + try withIntegrationState(seed: 76511) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LeakyReLU(negativeSlope: 0.2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.3090943694114685, + minimum: -0.5242452025413513, + maximum: 3.7356669902801514, + absoluteSum: 117.27056884765625, + positionChecksum: 58.45794677734375, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.061162084341049194, -0.01622505858540535, 0.19621963798999786, + 0.9918569326400757, -0.08134179562330246, 2.142888307571411, + ]), + tolerance: .float32) + } + } + + @Test("ELU") + func test_ELU() throws { + try withIntegrationState(seed: 85369) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ELU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.16009706258773804, + minimum: -0.9679020643234253, + maximum: 2.54449462890625, + absoluteSum: 162.6907958984375, + positionChecksum: 79.58113098144531, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.8480066657066345, -0.45555174350738525, 0.7897208333015442, + 0.7451042532920837, -0.7147349119186401, 0.16544604301452637, + ]), + tolerance: .float32) + } + } + + @Test("ELU/alpha") + func test_ELU_alpha() throws { + try withIntegrationState(seed: 48075) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ELU(alpha: 0.5) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.2684254050254822, + minimum: -0.46672412753105164, + maximum: 3.140700578689575, + absoluteSum: 134.91976928710938, + positionChecksum: 70.10220336914062, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.214599609375, 0.44013354182243347, 1.0033553838729858, 1.466883897781372, + 0.8385943174362183, 0.34053468704223633, + ]), + tolerance: .float32) + } + } + + @Test("CELU") + func test_CELU() throws { + try withIntegrationState(seed: 40999) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = CELU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.27527904510498047, + minimum: -0.9589483737945557, + maximum: 2.492225408554077, + absoluteSum: 168.98414611816406, + positionChecksum: 87.66098022460938, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.8132992386817932, 0.05264721438288689, 0.7257114052772522, + -0.36488646268844604, -0.6876139640808105, -0.2170543074607849, + ]), + tolerance: .float32) + } + } + + @Test("SiLU") + func test_SiLU() throws { + try withIntegrationState(seed: 94108) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = SiLU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.29252636432647705, + minimum: -0.2782213091850281, + maximum: 3.081599473953247, + absoluteSum: 117.98585510253906, + positionChecksum: 61.460487365722656, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.104568600654602, -0.07563341408967972, -0.049893204122781754, + -0.24418365955352783, -0.2484523504972458, 0.3073877692222595, + ]), + tolerance: .float32) + } + } + + @Test("SELU") + func test_SELU() throws { + try withIntegrationState(seed: 97176) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = SELU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.12662386894226074, + minimum: -1.6931291818618774, + maximum: 2.9934651851654053, + absoluteSum: 210.03921508789062, + positionChecksum: 101.77021026611328, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.8677928447723389, 0.6797854900360107, 0.36840230226516724, + -0.4589764475822449, -0.5336998105049133, 0.24484601616859436, + ]), + tolerance: .float32) + } + } + + @Test("Mish") + func test_Mish() throws { + try withIntegrationState(seed: 86013) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Mish() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.1439376026391983, + minimum: -0.3088423013687134, + maximum: 2.430266857147217, + absoluteSum: 103.19464111328125, + positionChecksum: 48.830989837646484, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.398360013961792, -0.014234477654099464, 0.8261752128601074, + -0.3047178089618683, -0.29602304100990295, -0.1958984136581421, + ]), + tolerance: .float32) + } + } + + @Test("Tanh") + func test_Tanh() throws { + try withIntegrationState(seed: 65516) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Tanh() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.049693141132593155, + minimum: -0.9932478070259094, + maximum: 0.9843816161155701, + absoluteSum: 147.73788452148438, + positionChecksum: 76.96502685546875, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.2730482816696167, 0.02616955153644085, -0.3087116479873657, + 0.5383175611495972, 0.22366760671138763, -0.2514197528362274, + ]), + tolerance: .float32) + } + } + + @Test("GELU") + func test_GELU() throws { + try withIntegrationState(seed: 56464) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GELU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.279835045337677, + minimum: -0.1699712574481964, + maximum: 2.9361231327056885, + absoluteSum: 100.90007781982422, + positionChecksum: 48.360328674316406, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.3581196069717407, 0.6479476094245911, 2.9361231327056885, + 0.5618335008621216, 0.5602741837501526, 1.0405502319335938, + ]), + tolerance: .float32) + } + } + + @Test("GELU/precise") + func test_GELU_precise() throws { + try withIntegrationState(seed: 5645) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GELU(approximation: .precise) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.26011642813682556, + minimum: -0.17004051804542542, + maximum: 2.3098645210266113, + absoluteSum: 98.03465270996094, + positionChecksum: 50.40411376953125, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.5097497701644897, -0.05096469447016716, -0.16896995902061462, + 1.1236003637313843, -0.16408808529376984, 0.2334529012441635, + ]), + tolerance: .float32) + } + } + + @Test("GELU/fast") + func test_GELU_fast() throws { + try withIntegrationState(seed: 11942) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GELU(approximation: .fast) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.2423984706401825, + minimum: -0.16360986232757568, + maximum: 2.6349453926086426, + absoluteSum: 93.47364807128906, + positionChecksum: 45.76264953613281, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.0383901596069336, -0.09115244448184967, 1.3228586912155151, + -0.1617099493741989, -0.16003209352493286, 0.08449243754148483, + ]), + tolerance: .float32) + } + } + + @Test("HardSwish") + func test_HardSwish() throws { + try withIntegrationState(seed: 53183) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = HardSwish() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.1723780632019043, + minimum: -0.3749813735485077, + maximum: 3.0541574954986572, + absoluteSum: 100.03382873535156, + positionChecksum: 49.33647155761719, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.3716834783554077, 0.10859835892915726, 0.44419172406196594, + 0.8135273456573486, -0.215162992477417, -0.1706112176179886, + ]), + tolerance: .float32) + } + } + + @Test("HardTanh") + func test_HardTanh() throws { + try withIntegrationState(seed: 35106) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = HardTanh() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.05235737934708595, + minimum: -1.0, + maximum: 1.0, + absoluteSum: 155.99127197265625, + positionChecksum: 75.52446746826172, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.07812279462814331, 0.43642038106918335, -0.6854920983314514, -1.0, + -0.8840487003326416, -1.0, + ]), + tolerance: .float32) + } + } + + @Test("HardShrink") + func test_HardShrink() throws { + try withIntegrationState(seed: 10671) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = HardShrink() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.1441853642463684, + minimum: -2.248732566833496, + maximum: 2.803218364715576, + absoluteSum: 192.95077514648438, + positionChecksum: 102.37713623046875, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.5038650631904602, 0.6754856109619141, 0.0, 0.0, -0.7869282364845276, + -1.0194823741912842, + ]), + tolerance: .float32) + } + } + + @Test("HardShrink/lambda") + func test_HardShrink_lambda() throws { + try withIntegrationState(seed: 29385) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = HardShrink(lambda: 0.2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.06390944868326187, + minimum: -2.980595827102661, + maximum: 2.9201619625091553, + absoluteSum: 205.49563598632812, + positionChecksum: 104.30592346191406, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.8073967695236206, 2.180393934249878, 0.3150389492511749, + -0.29229095578193665, -1.2221306562423706, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("Softplus") + func test_Softplus() throws { + try withIntegrationState(seed: 89825) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softplus() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.7824997305870056, + minimum: 0.0757046714425087, + maximum: 3.211472511291504, + absoluteSum: 200.31993103027344, + positionChecksum: 98.73367309570312, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.0939661264419556, 0.4173877239227295, 0.9904651641845703, + 0.8673735857009888, 1.2018206119537354, 0.3981037139892578, + ]), + tolerance: .float32) + } + } + + @Test("Softsign") + func test_Softsign() throws { + try withIntegrationState(seed: 69518) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softsign() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.007850009948015213, + minimum: -0.686312198638916, + maximum: 0.7605715990066528, + absoluteSum: 97.87689208984375, + positionChecksum: 49.9131965637207, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.2170916199684143, -0.6419945359230042, -0.6414955258369446, + -0.34234535694122314, 0.4860936999320984, 0.4500749111175537, + ]), + tolerance: .float32) + } + } + + @Test("Softshrink") + func test_Softshrink() throws { + try withIntegrationState(seed: 1067) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softshrink() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.00816175527870655, + minimum: -2.2436203956604004, + maximum: 2.848755121231079, + absoluteSum: 98.35267639160156, + positionChecksum: 51.70696258544922, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.0, 0.14205098152160645, 0.0, 0.6208071708679199, 0.1955741047859192, + 0.4562690854072571, + ]), + tolerance: .float32) + } + } + + @Test("Softshrink/lambda") + func test_Softshrink_lambda() throws { + try withIntegrationState(seed: 65904) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softshrink(lambda: 0.2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.07054832577705383, + minimum: -3.6172547340393066, + maximum: 2.4144792556762695, + absoluteSum: 156.82997131347656, + positionChecksum: 75.9085693359375, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.08811737596988678, 0.0, 0.861285388469696, -1.2284928560256958, 0.0, + -0.251159131526947, + ]), + tolerance: .float32) + } + } + + @Test("Softmax") + func test_Softmax() throws { + try withIntegrationState(seed: 76552) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softmax() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.0625, + minimum: 0.002436474896967411, + maximum: 0.45796260237693787, + absoluteSum: 16.0, + positionChecksum: 8.038972854614258, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.011925941333174706, 0.015899190679192543, 0.038694579154253006, + 0.03524196892976761, 0.015219993889331818, 0.12464060634374619, + ]), + tolerance: .float32) + } + } + + @Test("Softmin") + func test_Softmin() throws { + try withIntegrationState(seed: 7761) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Softmin() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.0625, + minimum: 0.0016566670965403318, + maximum: 0.6610931754112244, + absoluteSum: 16.0, + positionChecksum: 8.019933700561523, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.2532690465450287, 0.0048106457106769085, 0.008263001218438148, + 0.09788679331541061, 0.13839659094810486, 0.05248511955142021, + ]), + tolerance: .float32) + } + } + + @Test("LogSoftmax") + func test_LogSoftmax() throws { + try withIntegrationState(seed: 82322) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LogSoftmax() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -3.238546848297119, + minimum: -5.579071521759033, + maximum: -0.37829434871673584, + absoluteSum: 829.0679931640625, + positionChecksum: 420.2214050292969, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -2.0407981872558594, -2.1493773460388184, -4.613822937011719, + -2.9762957096099854, -4.447968482971191, -1.2499594688415527, + ]), + tolerance: .float32) + } + } + + @Test("LogSigmoid") + func test_LogSigmoid() throws { + try withIntegrationState(seed: 79811) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LogSigmoid() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.8042079210281372, + minimum: -2.9104275703430176, + maximum: -0.05430770292878151, + absoluteSum: 205.87722778320312, + positionChecksum: 103.44607543945312, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.646495521068573, -0.8605636358261108, -0.6894464492797852, + -1.5663222074508667, -0.37833893299102783, -0.36881721019744873, + ]), + tolerance: .float32) + } + } + + @Test("Step") + func test_Step() throws { + try withIntegrationState(seed: 68019) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Step() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .int32, + mean: 0.484375, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 124.0, + positionChecksum: 63.3359375, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [1.0, 0.0, 1.0, 1.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("Step/threshold") + func test_Step_threshold() throws { + try withIntegrationState(seed: 39538) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Step(threshold: 0.5) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .int32, + mean: 0.2890625, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 74.0, + positionChecksum: 35.3515625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [0.0, 1.0, 0.0, 0.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("GLU") + func test_GLU() throws { + try withIntegrationState(seed: 70327) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GLU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8], + dtype: .float32, + mean: -0.019226614385843277, + minimum: -2.286162853240967, + maximum: 1.3630620241165161, + absoluteSum: 49.290138244628906, + positionChecksum: 22.783916473388672, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.8543118834495544, -0.3675357699394226, -0.8153498768806458, + 0.032235924154520035, -0.4751351773738861, 0.2292122095823288, + ]), + tolerance: .float32) + } + } + + @Test("PReLU") + func test_PReLU() throws { + try withIntegrationState(seed: 7115) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = PReLU() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [1])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.6037850379943848, + minimum: 0.00014645590272266418, + maximum: 3.3934574127197266, + absoluteSum: 154.5689697265625, + positionChecksum: 72.80538940429688, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.044510796666145325, 0.5477976202964783, 0.10613550990819931, + 0.43876010179519653, 0.5276330709457397, 0.3869277536869049, + ]), + tolerance: .float32) + } + } + + @Test("PReLU/count") + func test_PReLU_count() throws { + try withIntegrationState(seed: 70876) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = PReLU(count: 16, value: 0.3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.4174256920814514, + minimum: -0.8470731973648071, + maximum: 2.8875372409820557, + absoluteSum: 128.29067993164062, + positionChecksum: 61.88407897949219, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.7127238512039185, 0.9431889653205872, 0.14541825652122498, + -0.006025301292538643, 0.03635843098163605, 0.8640216588973999, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleAttentionTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleAttentionTests.swift new file mode 100644 index 000000000..b6e4522b9 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleAttentionTests.swift @@ -0,0 +1,124 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 3 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleAttention") +struct GeneratedModuleAttentionTests { + + @Test("MultiHeadAttention") + func test_MultiHeadAttention() throws { + try withIntegrationState(seed: 24603) { + let x = MLXRandom.normal([2, 6, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MultiHeadAttention(dimensions: 16, numHeads: 4) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters( + module, + [ + ("key_proj.weight", [16, 16]), ("out_proj.weight", [16, 16]), + ("query_proj.weight", [16, 16]), ("value_proj.weight", [16, 16]), + ]) + let result = module(x, keys: x, values: x) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 16], + dtype: .float32, + mean: 0.05091768503189087, + minimum: -0.7064140439033508, + maximum: 1.1835263967514038, + absoluteSum: 40.15238952636719, + positionChecksum: 21.73735809326172, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + 0.27396339178085327, 0.2314339280128479, 0.05882430821657181, + -0.01647859998047352, -0.12509478628635406, -0.09450595080852509, + ]), + tolerance: .float32) + } + } + + @Test("MultiHeadAttention/bias") + func test_MultiHeadAttention_bias() throws { + try withIntegrationState(seed: 27712) { + let x = MLXRandom.normal([2, 6, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MultiHeadAttention(dimensions: 16, numHeads: 4, bias: true) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters( + module, + [ + ("key_proj.bias", [16]), ("key_proj.weight", [16, 16]), ("out_proj.bias", [16]), + ("out_proj.weight", [16, 16]), ("query_proj.bias", [16]), + ("query_proj.weight", [16, 16]), ("value_proj.bias", [16]), + ("value_proj.weight", [16, 16]), + ]) + let result = module(x, keys: x, values: x) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 16], + dtype: .float32, + mean: 0.26262208819389343, + minimum: -1.393288493156433, + maximum: 2.2614102363586426, + absoluteSum: 78.51240539550781, + positionChecksum: 30.15442403157552, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + 2.2614102363586426, 0.005398362874984741, 0.10388334095478058, + 0.21886420249938965, 0.3588539659976959, 0.5437532663345337, + ]), + tolerance: .float32) + } + } + + @Test("MultiHeadAttention/mask") + func test_MultiHeadAttention_mask() throws { + try withIntegrationState(seed: 65643) { + let x = MLXRandom.normal([2, 6, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let mask = MultiHeadAttention.createAdditiveCausalMask(6) + let module = MultiHeadAttention(dimensions: 16, numHeads: 4) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters( + module, + [ + ("key_proj.weight", [16, 16]), ("out_proj.weight", [16, 16]), + ("query_proj.weight", [16, 16]), ("value_proj.weight", [16, 16]), + ]) + let result = module(x, keys: x, values: x, mask: mask) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 16], + dtype: .float32, + mean: 0.024126103147864342, + minimum: -1.0875264406204224, + maximum: 0.8838620781898499, + absoluteSum: 60.221214294433594, + positionChecksum: 30.494303385416668, + sampleIndices: [0, 38, 76, 115, 153, 191], + samples: [ + 0.47034403681755066, -0.05138539895415306, 0.000970873050391674, + -0.5258967280387878, -0.12062773108482361, -1.0875264406204224, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleConvolutionTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleConvolutionTests.swift new file mode 100644 index 000000000..72d10f2bb --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleConvolutionTests.swift @@ -0,0 +1,306 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 10 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleConvolution") +struct GeneratedModuleConvolutionTests { + + @Test("Conv1d") + func test_Conv1d() throws { + try withIntegrationState(seed: 80399) { + let x = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv1d(inputChannels: 4, outputChannels: 3, kernelSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 3], + dtype: .float32, + mean: -0.11053398251533508, + minimum: -2.7783169746398926, + maximum: 3.0954372882843018, + absoluteSum: 45.54235076904297, + positionChecksum: 22.108502705891926, + sampleIndices: [0, 9, 19, 28, 38, 47], + samples: [ + 3.0954372882843018, 0.5464149713516235, 0.0334894061088562, + -0.09258026629686356, 1.676621913909912, -0.2100052833557129, + ]), + tolerance: .float32) + } + } + + @Test("Conv1d/stridePadding") + func test_Conv1d_stridePadding() throws { + try withIntegrationState(seed: 36805) { + let x = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv1d( + inputChannels: 4, outputChannels: 3, kernelSize: 3, stride: 2, padding: 1) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: -0.17636795341968536, + minimum: -2.4992713928222656, + maximum: 2.399353504180908, + absoluteSum: 23.765769958496094, + positionChecksum: 10.79921162923177, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + -1.8267059326171875, 2.399353504180908, -1.4456043243408203, + -0.010623306035995483, -1.917959451675415, 0.3439697027206421, + ]), + tolerance: .float32) + } + } + + @Test("Conv1d/noBias") + func test_Conv1d_noBias() throws { + try withIntegrationState(seed: 91755) { + let x = MLXRandom.normal([2, 10, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv1d(inputChannels: 4, outputChannels: 3, kernelSize: 3, bias: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [3, 3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 3], + dtype: .float32, + mean: 0.02788696438074112, + minimum: -1.910292148590088, + maximum: 1.8673605918884277, + absoluteSum: 34.24660110473633, + positionChecksum: 15.799858093261719, + sampleIndices: [0, 9, 19, 28, 38, 47], + samples: [ + -0.6793568730354309, -1.4656273126602173, -0.2600403130054474, + 0.4456599950790405, -0.5511552691459656, -0.8705923557281494, + ]), + tolerance: .float32) + } + } + + @Test("Conv2d") + func test_Conv2d() throws { + try withIntegrationState(seed: 77676) { + let x = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv2d(inputChannels: 3, outputChannels: 4, kernelSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [4]), ("weight", [4, 3, 3, 3])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 6, 4], + dtype: .float32, + mean: -0.19274936616420746, + minimum: -5.951562881469727, + maximum: 5.163274765014648, + absoluteSum: 359.1477355957031, + positionChecksum: 194.03301323784723, + sampleIndices: [0, 57, 115, 172, 230, 287], + samples: [ + -3.717362642288208, -0.37188804149627686, 3.1527411937713623, + -1.7658523321151733, -0.7540384531021118, -1.4457536935806274, + ]), + tolerance: .float32) + } + } + + @Test("Conv2d/stridePadding") + func test_Conv2d_stridePadding() throws { + try withIntegrationState(seed: 70327) { + let x = MLXRandom.normal([2, 8, 8, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv2d( + inputChannels: 3, outputChannels: 4, kernelSize: 3, stride: 2, padding: 1) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [4]), ("weight", [4, 3, 3, 3])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: -0.16016966104507446, + minimum: -4.5320258140563965, + maximum: 3.6546685695648193, + absoluteSum: 128.35308837890625, + positionChecksum: 58.57609558105469, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + -0.4038444757461548, -1.0283045768737793, 2.045264720916748, + 0.5583839416503906, -0.3402917981147766, -1.4707754850387573, + ]), + tolerance: .float32) + } + } + + @Test("Conv3d") + func test_Conv3d() throws { + try withIntegrationState(seed: 74637) { + let x = MLXRandom.normal([1, 4, 6, 6, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Conv3d(inputChannels: 2, outputChannels: 3, kernelSize: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 2, 2, 2, 2])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 5, 5, 3], + dtype: .float32, + mean: -0.23182404041290283, + minimum: -3.550316095352173, + maximum: 3.392254590988159, + absoluteSum: 165.49330139160156, + positionChecksum: 86.74056423611111, + sampleIndices: [0, 45, 90, 134, 179, 224], + samples: [ + -0.2590126395225525, -1.4220865964889526, 1.4350488185882568, + -1.436377763748169, -0.26314473152160645, 0.30482929944992065, + ]), + tolerance: .float32) + } + } + + @Test("ConvTransposed1d") + func test_ConvTransposed1d() throws { + try withIntegrationState(seed: 94091) { + let x = MLXRandom.normal([2, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ConvTransposed1d(inputChannels: 4, outputChannels: 3, kernelSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 10, 3], + dtype: .float32, + mean: -0.16663511097431183, + minimum: -1.8646354675292969, + maximum: 1.3452529907226562, + absoluteSum: 25.450639724731445, + positionChecksum: 12.595225016276041, + sampleIndices: [0, 12, 24, 35, 47, 59], + samples: [ + -0.7107505798339844, 0.8472874164581299, -0.5283459424972534, + -0.22366145253181458, 0.5042819380760193, -0.5187655091285706, + ]), + tolerance: .float32) + } + } + + @Test("ConvTransposed1d/stride") + func test_ConvTransposed1d_stride() throws { + try withIntegrationState(seed: 96722) { + let x = MLXRandom.normal([2, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ConvTransposed1d( + inputChannels: 4, outputChannels: 3, kernelSize: 3, stride: 2, padding: 1, + outputPadding: 1) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 16, 3], + dtype: .float32, + mean: -0.16872534155845642, + minimum: -2.399252414703369, + maximum: 1.8922030925750732, + absoluteSum: 51.313255310058594, + positionChecksum: 26.30189259847005, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + 0.008320808410644531, -0.12555092573165894, 0.2811652421951294, + -1.5479183197021484, -0.2482323944568634, 0.49897417426109314, + ]), + tolerance: .float32) + } + } + + @Test("ConvTransposed2d") + func test_ConvTransposed2d() throws { + try withIntegrationState(seed: 17256) { + let x = MLXRandom.normal([2, 6, 6, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ConvTransposed2d(inputChannels: 3, outputChannels: 4, kernelSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [4]), ("weight", [4, 3, 3, 3])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 4], + dtype: .float32, + mean: -0.13198184967041016, + minimum: -4.098501682281494, + maximum: 3.9713425636291504, + absoluteSum: 394.7406005859375, + positionChecksum: 200.25326538085938, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.45735645294189453, 0.35417571663856506, 0.3488615155220032, + 1.8767280578613281, -0.24553531408309937, -0.6553887128829956, + ]), + tolerance: .float32) + } + } + + @Test("ConvTransposed3d") + func test_ConvTransposed3d() throws { + try withIntegrationState(seed: 18025) { + let x = MLXRandom.normal([1, 4, 4, 4, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ConvTransposed3d(inputChannels: 2, outputChannels: 3, kernelSize: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [3]), ("weight", [3, 2, 2, 2, 2])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [1, 5, 5, 5, 3], + dtype: .float32, + mean: -0.18090620636940002, + minimum: -4.558497428894043, + maximum: 3.8166122436523438, + absoluteSum: 291.21270751953125, + positionChecksum: 155.24589583333332, + sampleIndices: [0, 75, 150, 224, 299, 374], + samples: [ + -1.108421802520752, 0.058435797691345215, -0.2240082025527954, + 1.144876480102539, 0.4049822986125946, 0.36139535903930664, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleDropoutTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleDropoutTests.swift new file mode 100644 index 000000000..53ff30898 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleDropoutTests.swift @@ -0,0 +1,109 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 3 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleDropout") +struct GeneratedModuleDropoutTests { + + @Test("Dropout/eval") + func test_Dropout_eval() throws { + // eval mode is the identity; training mode is stochastic and cannot be compared value-by-value + try withIntegrationState(seed: 81768) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Dropout() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 0.05213969200849533, + minimum: -2.930783987045288, + maximum: 3.0198965072631836, + absoluteSum: 203.62631225585938, + positionChecksum: 104.69136810302734, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.10624595731496811, -0.14009599387645721, 0.12543074786663055, + -1.2116183042526245, -0.6216195225715637, 0.6007198095321655, + ]), + tolerance: .float32) + } + } + + @Test("Dropout2d/eval") + func test_Dropout2d_eval() throws { + // eval mode is the identity; training mode is stochastic and cannot be compared value-by-value + try withIntegrationState(seed: 24737) { + let x = MLXRandom.normal([2, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Dropout2d() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 4], + dtype: .float32, + mean: -0.05712714046239853, + minimum: -3.243155002593994, + maximum: 2.784644842147827, + absoluteSum: 398.7391357421875, + positionChecksum: 194.79966735839844, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 2.463557481765747, 0.979790449142456, -1.3742802143096924, + -1.0140464305877686, 0.9424789547920227, 0.03446979075670242, + ]), + tolerance: .float32) + } + } + + @Test("Dropout3d/eval") + func test_Dropout3d_eval() throws { + // eval mode is the identity; training mode is stochastic and cannot be compared value-by-value + try withIntegrationState(seed: 78645) { + let x = MLXRandom.normal([2, 4, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Dropout3d() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8, 4], + dtype: .float32, + mean: 0.00032923324033617973, + minimum: -3.3692824840545654, + maximum: 3.9168732166290283, + absoluteSum: 1644.861328125, + positionChecksum: 813.3851318359375, + sampleIndices: [0, 409, 819, 1228, 1638, 2047], + samples: [ + 1.438454508781433, 0.07942906767129898, -1.1821718215942383, + -0.8269960880279541, 0.7035877704620361, 0.6526941061019897, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleLinearTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleLinearTests.swift new file mode 100644 index 000000000..4bd4c35d4 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleLinearTests.swift @@ -0,0 +1,163 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 5 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleLinear") +struct GeneratedModuleLinearTests { + + @Test("Linear") + func test_Linear() throws { + try withIntegrationState(seed: 45463) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Linear(16, 5) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [5]), ("weight", [5, 16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 5], + dtype: .float32, + mean: -0.030131885781884193, + minimum: -3.8416616916656494, + maximum: 3.5228404998779297, + absoluteSum: 76.65104675292969, + positionChecksum: 48.569488525390625, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + -0.2741870582103729, 0.2716590166091919, -0.3759540915489197, + -0.07041969895362854, -0.8177027106285095, 3.5228404998779297, + ]), + tolerance: .float32) + } + } + + @Test("Linear/noBias") + func test_Linear_noBias() throws { + try withIntegrationState(seed: 88724) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Linear(16, 5, bias: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [5, 16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 5], + dtype: .float32, + mean: 0.11451786011457443, + minimum: -2.7055249214172363, + maximum: 3.1095452308654785, + absoluteSum: 64.63607025146484, + positionChecksum: 34.513467407226564, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + -0.16995857656002045, 1.6557778120040894, -0.1264149248600006, + 0.21584510803222656, -0.18836407363414764, -1.5945661067962646, + ]), + tolerance: .float32) + } + } + + @Test("Bilinear") + func test_Bilinear() throws { + try withIntegrationState(seed: 3925) { + let x = MLXRandom.normal([2, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let y = MLXRandom.normal([2, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Bilinear(16, 8, 5) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [5]), ("weight", [5, 8, 16])]) + let result = module(x, y) + expectSummary( + result, + ArraySummary( + shape: [2, 5], + dtype: .float32, + mean: 0.6467210650444031, + minimum: -1.8467845916748047, + maximum: 4.493323802947998, + absoluteSum: 13.118966102600098, + positionChecksum: 4.266597747802734, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [ + 4.493323802947998, 1.3232699632644653, -1.8467845916748047, + -0.7817299365997314, -0.029828175902366638, 0.7220736145973206, + ]), + tolerance: .float32) + } + } + + @Test("Embedding") + func test_Embedding() throws { + try withIntegrationState(seed: 44353) { + let x = MLXRandom.randInt(low: 0, high: 10, [2, 6], type: Int32.self) + let module = Embedding(embeddingCount: 10, dimensions: 8) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [10, 8])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 8], + dtype: .float32, + mean: 0.07708332687616348, + minimum: -0.4000000059604645, + maximum: 0.48750001192092896, + absoluteSum: 27.900001525878906, + positionChecksum: 15.243748982747396, + sampleIndices: [0, 19, 38, 57, 76, 95], + samples: [ + -0.09999999403953552, 0.4375, -0.32499998807907104, 0.3125, + -0.3499999940395355, 0.2875000238418579, + ]), + tolerance: .float32) + } + } + + @Test("Embedding/asLinear") + func test_Embedding_asLinear() throws { + try withIntegrationState(seed: 8970) { + let x = MLXRandom.normal([2, 6, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Embedding(embeddingCount: 10, dimensions: 8) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [10, 8])]) + let result = module.asLinear(x) + expectSummary( + result, + ArraySummary( + shape: [2, 6, 10], + dtype: .float32, + mean: -0.015905074775218964, + minimum: -2.8021512031555176, + maximum: 2.801405429840088, + absoluteSum: 62.93570327758789, + positionChecksum: 31.226318359375, + sampleIndices: [0, 24, 48, 71, 95, 119], + samples: [ + -2.8021512031555176, 0.007820483297109604, -0.06263187527656555, + 0.44406285881996155, -0.010466308332979679, -1.489911437034607, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleNormalizationTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleNormalizationTests.swift new file mode 100644 index 000000000..c36813f8b --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleNormalizationTests.swift @@ -0,0 +1,308 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 10 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleNormalization") +struct GeneratedModuleNormalizationTests { + + @Test("LayerNorm") + func test_LayerNorm() throws { + try withIntegrationState(seed: 42737) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LayerNorm(dimensions: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.0198993980884552, + minimum: -1.304665446281433, + maximum: 1.182969570159912, + absoluteSum: 74.0040054321289, + positionChecksum: 37.30052185058594, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.03280201554298401, -0.5021070241928101, -0.24673715233802795, + 0.10571826994419098, 0.23079484701156616, 0.9026331305503845, + ]), + tolerance: .float32) + } + } + + @Test("LayerNorm/noAffine") + func test_LayerNorm_noAffine() throws { + try withIntegrationState(seed: 71715) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = LayerNorm(dimensions: 16, affine: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: 1.862645149230957e-09, + minimum: -2.7004194259643555, + maximum: 2.4202475547790527, + absoluteSum: 202.51004028320312, + positionChecksum: 101.72293090820312, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.364133596420288, 0.1137239933013916, -0.1254626214504242, + 0.008553391322493553, -0.20014025270938873, -1.4844169616699219, + ]), + tolerance: .float32) + } + } + + @Test("RMSNorm") + func test_RMSNorm() throws { + try withIntegrationState(seed: 74037) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RMSNorm(dimensions: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.02211989276111126, + minimum: -0.9623683094978333, + maximum: 1.1941437721252441, + absoluteSum: 51.25483322143555, + positionChecksum: 26.368240356445312, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.029898809269070625, -0.015912216156721115, -0.06328903138637543, + -0.10268687456846237, -0.041741251945495605, -0.15935736894607544, + ]), + tolerance: .float32) + } + } + + @Test("GroupNorm") + func test_GroupNorm() throws { + try withIntegrationState(seed: 30539) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GroupNorm(groupCount: 4, dimensions: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.010309945791959763, + minimum: -1.3219797611236572, + maximum: 1.3990989923477173, + absoluteSum: 73.74345397949219, + positionChecksum: 36.934391021728516, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.17130780220031738, -0.5502849221229553, -0.1820463389158249, + 0.13417498767375946, -0.11009815335273743, 1.3455901145935059, + ]), + tolerance: .float32) + } + } + + @Test("GroupNorm/pytorchCompatible") + func test_GroupNorm_pytorchCompatible() throws { + try withIntegrationState(seed: 49381) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GroupNorm(groupCount: 4, dimensions: 16, pytorchCompatible: true) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.030907047912478447, + minimum: -1.2020220756530762, + maximum: 1.2239670753479004, + absoluteSum: 73.9453125, + positionChecksum: 37.40062713623047, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.6553022861480713, -0.08870765566825867, -0.05848407745361328, + 0.06563569605350494, 0.47302502393722534, 0.6089176535606384, + ]), + tolerance: .float32) + } + } + + @Test("InstanceNorm") + func test_InstanceNorm() throws { + try withIntegrationState(seed: 7588) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = InstanceNorm(dimensions: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -9.313225746154785e-09, + minimum: -2.293649435043335, + maximum: 2.4967687129974365, + absoluteSum: 211.96717834472656, + positionChecksum: 104.72128295898438, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 1.8193515539169312, -0.2942982316017151, -0.3463905453681946, + 0.1246948391199112, 2.166072130203247, -1.5744894742965698, + ]), + tolerance: .float32) + } + } + + @Test("InstanceNorm/affine") + func test_InstanceNorm_affine() throws { + try withIntegrationState(seed: 86086) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = InstanceNorm(dimensions: 16, affine: true) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.03125, + minimum: -1.5313172340393066, + maximum: 1.369132399559021, + absoluteSum: 72.45367431640625, + positionChecksum: 35.79180145263672, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.10569164156913757, -0.18218281865119934, -0.2373899221420288, + 0.01955053210258484, 0.29726508259773254, 0.3344022035598755, + ]), + tolerance: .float32) + } + } + + @Test("BatchNorm/eval") + func test_BatchNorm_eval() throws { + // eval mode uses the (rewritten) running statistics + try withIntegrationState(seed: 32289) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = BatchNorm(featureCount: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters( + module, + [("bias", [16]), ("running_mean", [16]), ("running_var", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: Double.nan, + minimum: Double.nan, + maximum: Double.nan, + absoluteSum: Double.nan, + positionChecksum: Double.nan, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + Double.nan, Double.nan, Double.nan, 0.19975754618644714, + 0.17086565494537354, 0.7938545942306519, + ]), + tolerance: .float32) + } + } + + @Test("BatchNorm/training") + func test_BatchNorm_training() throws { + // training mode normalizes with the batch statistics; deterministic, unlike Dropout + try withIntegrationState(seed: 8639) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = BatchNorm(featureCount: 16) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(true) + expectParameters( + module, + [("bias", [16]), ("running_mean", [16]), ("running_var", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.0312499962747097, + minimum: -1.5580859184265137, + maximum: 1.4559376239776611, + absoluteSum: 74.08992004394531, + positionChecksum: 36.43730926513672, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.5708073377609253, -0.2267388105392456, -0.17827975749969482, + 0.17880982160568237, 0.019031628966331482, 0.21674005687236786, + ]), + tolerance: .float32) + } + } + + @Test("BatchNorm/noTrackRunningStats") + func test_BatchNorm_noTrackRunningStats() throws { + try withIntegrationState(seed: 91221) { + let x = MLXRandom.normal([2, 8, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let module = BatchNorm(featureCount: 16, trackRunningStats: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("bias", [16]), ("weight", [16])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 16], + dtype: .float32, + mean: -0.03125, + minimum: -1.1940438747406006, + maximum: 1.2612149715423584, + absoluteSum: 74.57074737548828, + positionChecksum: 37.47613525390625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.3706926703453064, -0.5099658370018005, -0.023155465722084045, + 0.10782907903194427, 0.10105209052562714, 1.017714262008667, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModulePoolingTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModulePoolingTests.swift new file mode 100644 index 000000000..1208a34ec --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModulePoolingTests.swift @@ -0,0 +1,246 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 8 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModulePooling") +struct GeneratedModulePoolingTests { + + @Test("MaxPool1d") + func test_MaxPool1d() throws { + try withIntegrationState(seed: 29744) { + let x = MLXRandom.normal([2, 16, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MaxPool1d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 4], + dtype: .float32, + mean: 0.5671930313110352, + minimum: -1.6935234069824219, + maximum: 2.3494832515716553, + absoluteSum: 54.285911560058594, + positionChecksum: 28.31031036376953, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -1.6935234069824219, 0.8108446002006531, 0.4986465275287628, + 0.9882473945617676, 0.5640528202056885, 0.9553364515304565, + ]), + tolerance: .float32) + } + } + + @Test("MaxPool2d") + func test_MaxPool2d() throws { + try withIntegrationState(seed: 98643) { + let x = MLXRandom.normal([2, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MaxPool2d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: 1.061241626739502, + minimum: -0.757491409778595, + maximum: 3.0933570861816406, + absoluteSum: 142.98167419433594, + positionChecksum: 71.34737396240234, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 1.6198300123214722, 0.06948284059762955, -0.12833164632320404, + 1.8321651220321655, 1.6542015075683594, 0.931220293045044, + ]), + tolerance: .float32) + } + } + + @Test("MaxPool3d") + func test_MaxPool3d() throws { + try withIntegrationState(seed: 32850) { + let x = MLXRandom.normal([2, 4, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MaxPool3d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 2, 4, 4, 4], + dtype: .float32, + mean: 1.3603229522705078, + minimum: -0.38012799620628357, + maximum: 3.5154623985290527, + absoluteSum: 349.03045654296875, + positionChecksum: 171.13186645507812, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + 0.4982339143753052, 1.8628464937210083, 0.7978113293647766, + 1.2853648662567139, 1.1705279350280762, 1.8667374849319458, + ]), + tolerance: .float32) + } + } + + @Test("MaxPool2d/padding") + func test_MaxPool2d_padding() throws { + try withIntegrationState(seed: 60585) { + let x = MLXRandom.normal([2, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = MaxPool2d(kernelSize: 3, stride: 2, padding: 1) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: 1.378739356994629, + minimum: 0.026348749175667763, + maximum: 2.6323399543762207, + absoluteSum: 176.4786376953125, + positionChecksum: 89.82207489013672, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 1.1064631938934326, 1.1356632709503174, 1.3230632543563843, + 1.220381259918213, 2.2429463863372803, 1.0948508977890015, + ]), + tolerance: .float32) + } + } + + @Test("AvgPool1d") + func test_AvgPool1d() throws { + try withIntegrationState(seed: 69169) { + let x = MLXRandom.normal([2, 16, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = AvgPool1d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 4], + dtype: .float32, + mean: 0.07272765040397644, + minimum: -2.300410270690918, + maximum: 1.7861367464065552, + absoluteSum: 36.99386978149414, + positionChecksum: 19.82023811340332, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 0.10242053866386414, 0.11492392420768738, 0.6828688383102417, + 0.174363374710083, -0.2848045229911804, -1.0003540515899658, + ]), + tolerance: .float32) + } + } + + @Test("AvgPool2d") + func test_AvgPool2d() throws { + try withIntegrationState(seed: 9458) { + let x = MLXRandom.normal([2, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = AvgPool2d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: -0.08202866464853287, + minimum: -1.2255580425262451, + maximum: 1.491541862487793, + absoluteSum: 48.739749908447266, + positionChecksum: 24.986846923828125, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + -0.4838985204696655, 0.0945008397102356, -0.11756622791290283, + 0.09166160225868225, -0.45479485392570496, 0.6930731534957886, + ]), + tolerance: .float32) + } + } + + @Test("AvgPool3d") + func test_AvgPool3d() throws { + try withIntegrationState(seed: 62547) { + let x = MLXRandom.normal([2, 4, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = AvgPool3d(kernelSize: 2, stride: 2) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 2, 4, 4, 4], + dtype: .float32, + mean: -0.001542090903967619, + minimum: -1.1280696392059326, + maximum: 0.8791224360466003, + absoluteSum: 73.494384765625, + positionChecksum: 36.46441650390625, + sampleIndices: [0, 51, 102, 153, 204, 255], + samples: [ + -0.4229593873023987, 0.4040684401988983, 0.1515209674835205, + 0.12673257291316986, 0.5045558214187622, 0.24828201532363892, + ]), + tolerance: .float32) + } + } + + @Test("AvgPool2d/padding") + func test_AvgPool2d_padding() throws { + try withIntegrationState(seed: 60609) { + let x = MLXRandom.normal([2, 8, 8, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = AvgPool2d(kernelSize: 3, stride: 2, padding: 1) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 4, 4], + dtype: .float32, + mean: 0.038703806698322296, + minimum: -0.7088121771812439, + maximum: 0.7457450032234192, + absoluteSum: 24.897192001342773, + positionChecksum: 12.805156707763672, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.28262028098106384, -0.18140751123428345, 0.34491124749183655, + 0.2511276602745056, -0.11475422978401184, -0.3191181719303131, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModulePositionalTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModulePositionalTests.swift new file mode 100644 index 000000000..2ca306dee --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModulePositionalTests.swift @@ -0,0 +1,191 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 6 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModulePositional") +struct GeneratedModulePositionalTests { + + @Test("RoPE") + func test_RoPE() throws { + try withIntegrationState(seed: 70361) { + let x = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RoPE(dimensions: 8) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8], + dtype: .float32, + mean: -0.17599579691886902, + minimum: -2.1909637451171875, + maximum: 1.6637766361236572, + absoluteSum: 48.165287017822266, + positionChecksum: 21.670984268188477, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 1.1579456329345703, -1.2916109561920166, -0.03200756013393402, + 0.1322382688522339, 0.4934418797492981, 0.43062347173690796, + ]), + tolerance: .float32) + } + } + + @Test("RoPE/traditional") + func test_RoPE_traditional() throws { + try withIntegrationState(seed: 84217) { + let x = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RoPE(dimensions: 8, traditional: true, base: 500.0, scale: 0.5) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8], + dtype: .float32, + mean: 0.0022530164569616318, + minimum: -2.0894505977630615, + maximum: 2.9366118907928467, + absoluteSum: 46.10230255126953, + positionChecksum: 20.703702926635742, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + 1.162760615348816, 0.5243428945541382, -0.13487887382507324, + -0.6033228039741516, 1.8898133039474487, -0.35043108463287354, + ]), + tolerance: .float32) + } + } + + @Test("RoPE/offset") + func test_RoPE_offset() throws { + try withIntegrationState(seed: 44446) { + let x = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RoPE(dimensions: 8) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x, offset: 2) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8], + dtype: .float32, + mean: -0.13316422700881958, + minimum: -2.70853853225708, + maximum: 2.8416025638580322, + absoluteSum: 53.91521453857422, + positionChecksum: 30.396228790283203, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.6359602212905884, -0.9670241475105286, -0.32228633761405945, + -1.538472056388855, -0.9030501842498779, 0.39927321672439575, + ]), + tolerance: .float32) + } + } + + @Test("SinusoidalPositionalEncoding") + func test_SinusoidalPositionalEncoding() throws { + try withIntegrationState(seed: 72004) { + let x = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = SinusoidalPositionalEncoding(dimensions: 8) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: 0.23017042875289917, + minimum: -0.4999251067638397, + maximum: 0.5, + absoluteSum: 137.57159423828125, + positionChecksum: 69.34072875976562, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 0.46070006489753723, 0.4999995231628418, 0.1861143559217453, + 8.04811788839288e-05, -0.020217712968587875, 0.5, + ]), + tolerance: .float32) + } + } + + @Test("SinusoidalPositionalEncoding/cosineFirst") + func test_SinusoidalPositionalEncoding_cosineFirst() throws { + try withIntegrationState(seed: 28369) { + let x = MLXRandom.normal([2, 4, 8], dtype: .float32, loc: 0.0, scale: 1.0) + let module = SinusoidalPositionalEncoding( + dimensions: 8, cosineFirst: true, fullTurns: true) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 8, 8], + dtype: .float32, + mean: 0.17157743871212006, + minimum: -0.4993824362754822, + maximum: 0.5, + absoluteSum: 142.95547485351562, + positionChecksum: 71.41983795166016, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -0.4852396547794342, 0.007850918918848038, 0.09497486054897308, + 0.4999999701976776, 0.4368760883808136, 0.00014115065278019756, + ]), + tolerance: .float32) + } + } + + @Test("ALiBi") + func test_ALiBi() throws { + try withIntegrationState(seed: 51200) { + let scores = MLXRandom.normal([1, 4, 6, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let module = ALiBi() + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(attentionScores: scores) + expectSummary( + result, + ArraySummary( + shape: [1, 4, 6, 6], + dtype: .float32, + mean: -0.1866210252046585, + minimum: -2.942293643951416, + maximum: 2.480872631072998, + absoluteSum: 124.28718566894531, + positionChecksum: 60.13202582465278, + sampleIndices: [0, 29, 57, 86, 114, 143], + samples: [ + -0.519200325012207, -1.532111644744873, 1.0923478603363037, + 0.1874494105577469, -1.2987788915634155, -1.300144910812378, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleRecurrentTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleRecurrentTests.swift new file mode 100644 index 000000000..8ebaa14d1 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleRecurrentTests.swift @@ -0,0 +1,163 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 5 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleRecurrent") +struct GeneratedModuleRecurrentTests { + + @Test("RNN") + func test_RNN() throws { + try withIntegrationState(seed: 26873) { + let x = MLXRandom.normal([2, 5, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RNN(inputSize: 4, hiddenSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("Whh", [3, 3]), ("Wxh", [3, 4]), ("bias", [3])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: -0.1807376742362976, + minimum: -0.9347691535949707, + maximum: 0.7658460736274719, + absoluteSum: 11.239614486694336, + positionChecksum: 6.376473999023437, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + -0.35051560401916504, -0.33143526315689087, -0.8694400787353516, + -0.299630343914032, -0.5201353430747986, 0.19474464654922485, + ]), + tolerance: .float32) + } + } + + @Test("RNN/noBias") + func test_RNN_noBias() throws { + try withIntegrationState(seed: 54659) { + let x = MLXRandom.normal([2, 5, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RNN(inputSize: 4, hiddenSize: 3, bias: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("Whh", [3, 3]), ("Wxh", [3, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: -0.04676102474331856, + minimum: -0.7166611552238464, + maximum: 0.7711976170539856, + absoluteSum: 7.510426044464111, + positionChecksum: 3.6819297790527346, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.40871885418891907, -0.7166611552238464, 0.23077711462974548, + 0.0683915838599205, 0.21350324153900146, -0.5236168503761292, + ]), + tolerance: .float32) + } + } + + @Test("RNN/hidden") + func test_RNN_hidden() throws { + try withIntegrationState(seed: 52626) { + let x = MLXRandom.normal([2, 5, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let hidden = MLXRandom.normal([2, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = RNN(inputSize: 4, hiddenSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("Whh", [3, 3]), ("Wxh", [3, 4]), ("bias", [3])]) + let result = module(x, hidden: hidden) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: -0.0772470086812973, + minimum: -0.911475419998169, + maximum: 0.95676189661026, + absoluteSum: 12.752267837524414, + positionChecksum: 6.333803304036459, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.6317100524902344, -0.3342420160770416, -0.4475969076156616, + -0.8060013651847839, 0.5694852471351624, -0.5043428540229797, + ]), + tolerance: .float32) + } + } + + @Test("GRU") + func test_GRU() throws { + try withIntegrationState(seed: 26515) { + let x = MLXRandom.normal([2, 5, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GRU(inputSize: 4, hiddenSize: 3) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("Wh", [9, 3]), ("Wx", [9, 4]), ("b", [9]), ("bhn", [3])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: 0.3049412667751312, + minimum: -0.1211482584476471, + maximum: 0.6732051968574524, + absoluteSum: 9.443852424621582, + positionChecksum: 5.356678263346354, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.00961192138493061, 0.2725975513458252, -0.1211482584476471, + 0.378728985786438, 0.33235982060432434, 0.5703566074371338, + ]), + tolerance: .float32) + } + } + + @Test("GRU/noBias") + func test_GRU_noBias() throws { + try withIntegrationState(seed: 84535) { + let x = MLXRandom.normal([2, 5, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let module = GRU(inputSize: 4, hiddenSize: 3, bias: false) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, [("Wh", [9, 3]), ("Wx", [9, 4])]) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 5, 3], + dtype: .float32, + mean: 0.028713371604681015, + minimum: -0.3385675251483917, + maximum: 0.5721112489700317, + absoluteSum: 5.291589260101318, + positionChecksum: 2.6151153564453127, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 0.03434830531477928, 0.13748443126678467, 0.4411085546016693, + -0.020659416913986206, -0.03192802146077156, -0.3385675251483917, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedModuleUpsampleTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedModuleUpsampleTests.swift new file mode 100644 index 000000000..51808a0ea --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedModuleUpsampleTests.swift @@ -0,0 +1,162 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 5 + +import Foundation +import MLX +import MLXNN +import Testing + +@Suite("generated: ModuleUpsample") +struct GeneratedModuleUpsampleTests { + + @Test("Upsample/nearest") + func test_Upsample_nearest() throws { + try withIntegrationState(seed: 39945) { + let x = MLXRandom.normal([2, 4, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Upsample(scaleFactor: 2.0) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 3], + dtype: .float32, + mean: -0.08313984423875809, + minimum: -1.9472978115081787, + maximum: 1.9504168033599854, + absoluteSum: 274.71282958984375, + positionChecksum: 145.42476399739584, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -0.9178804755210876, -0.5090406537055969, 0.6381695866584778, + 0.4494686424732208, 0.5923683643341064, -1.504755973815918, + ]), + tolerance: .float32) + } + } + + @Test("Upsample/linear") + func test_Upsample_linear() throws { + try withIntegrationState(seed: 2223) { + let x = MLXRandom.normal([2, 4, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Upsample(scaleFactor: 2.0, mode: .linear()) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 3], + dtype: .float32, + mean: 0.07992034405469894, + minimum: -2.166339159011841, + maximum: 2.1690256595611572, + absoluteSum: 198.35586547851562, + positionChecksum: 104.31976318359375, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -0.6088932156562805, 0.43297699093818665, -0.9983523488044739, + 0.5581990480422974, 0.9666178822517395, 0.052967239171266556, + ]), + tolerance: .float32) + } + } + + @Test("Upsample/linear/alignCorners") + func test_Upsample_linear_alignCorners() throws { + try withIntegrationState(seed: 15253) { + let x = MLXRandom.normal([2, 4, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Upsample(scaleFactor: 2.0, mode: .linear(alignCorners: true)) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 3], + dtype: .float32, + mean: -0.15478087961673737, + minimum: -2.823728561401367, + maximum: 1.8036024570465088, + absoluteSum: 241.51821899414062, + positionChecksum: 117.62404378255208, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -2.823728561401367, -0.9696316719055176, 0.006423458456993103, + -0.19877995550632477, -0.32860827445983887, 0.10723376274108887, + ]), + tolerance: .float32) + } + } + + @Test("Upsample/cubic") + func test_Upsample_cubic() throws { + try withIntegrationState(seed: 98182) { + let x = MLXRandom.normal([2, 4, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Upsample(scaleFactor: 2.0, mode: .cubic()) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 8, 3], + dtype: .float32, + mean: -0.05168304964900017, + minimum: -2.2958436012268066, + maximum: 2.3127639293670654, + absoluteSum: 248.95880126953125, + positionChecksum: 116.72739664713542, + sampleIndices: [0, 77, 153, 230, 306, 383], + samples: [ + -1.6836540699005127, 0.061760637909173965, -0.49417844414711, + -0.23918810486793518, 0.038516491651535034, 0.0978599488735199, + ]), + tolerance: .float32) + } + } + + @Test("Upsample/scaleFactors") + func test_Upsample_scaleFactors() throws { + try withIntegrationState(seed: 48748) { + let x = MLXRandom.normal([2, 4, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let module = Upsample(scaleFactor: [2.0, 3.0]) + module.update(parameters: module.mapParameters { deterministicParameter($0) }) + module.train(false) + expectParameters(module, []) + let result = module(x) + expectSummary( + result, + ArraySummary( + shape: [2, 8, 12, 3], + dtype: .float32, + mean: -0.11570308357477188, + minimum: -1.899631142616272, + maximum: 2.335793972015381, + absoluteSum: 417.2052001953125, + positionChecksum: 212.24254014756946, + sampleIndices: [0, 115, 230, 345, 460, 575], + samples: [ + -0.7018576264381409, 0.24986159801483154, -1.3673666715621948, + 0.004559833090752363, 0.339202344417572, 1.4709136486053467, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedOptimizerFunctionsTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedOptimizerFunctionsTests.swift new file mode 100644 index 000000000..ab8ee659d --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedOptimizerFunctionsTests.swift @@ -0,0 +1,297 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 11 + +import Foundation +import MLX +import MLXNN +import Testing + +@testable import MLXOptimizers + +@Suite("generated: OptimizerFunctions") +struct GeneratedOptimizerFunctionsTests { + + @Test("clipGradNorm/norm") + func test_clipGradNorm_norm() throws { + // the returned norm is the norm *before* clipping + try withIntegrationState(seed: 74913) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = clipGradNorm( + gradients: ModuleParameters.unflattened([("a", a), ("b", b)]), maxNorm: 1.0 + ).1 + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 4.2577900886535645, + minimum: 4.2577900886535645, + maximum: 4.2577900886535645, + absoluteSum: 4.2577900886535645, + positionChecksum: 4.2577900886535645, + sampleIndices: [0], + samples: [4.2577900886535645]), + tolerance: .float32) + } + } + + @Test("clipGradNorm/clipped") + func test_clipGradNorm_clipped() throws { + // the norm of a normal [4, 3] + [5] is well above 1, so this clips + try withIntegrationState(seed: 32730) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = clipGradNorm( + gradients: ModuleParameters.unflattened([("a", a), ("b", b)]), maxNorm: 1.0 + ).0[unwrapping: "a"]! + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.141945943236351, + minimum: -0.134745791554451, + maximum: 0.4862045347690582, + absoluteSum: 2.1906065940856934, + positionChecksum: 1.2294964790344238, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.13243113458156586, 0.17849785089492798, 0.1051086038351059, + 0.24804984033107758, 0.32634231448173523, -0.0005411746096797287, + ]), + tolerance: .float32) + } + } + + @Test("clipGradNorm/underTheLimit") + func test_clipGradNorm_underTheLimit() throws { + // already inside the limit, so the gradients pass through unchanged + try withIntegrationState(seed: 41005) { + let a = MLXRandom.uniform(low: -0.02, high: 0.02, [4, 3], dtype: .float32) + let result = clipGradNorm( + gradients: ModuleParameters.unflattened([("a", a)]), maxNorm: 1.0 + ).0[unwrapping: "a"]! + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.0004688164626713842, + minimum: -0.016733653843402863, + maximum: 0.018089659512043, + absoluteSum: 0.12719818949699402, + positionChecksum: 0.06582144896189372, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0020372532308101654, 0.01661718264222145, 0.008448302745819092, + -0.016733653843402863, -0.01313621737062931, -0.006226498633623123, + ]), + tolerance: .float32) + } + } + + @Test("clipGradNorm/array") + func test_clipGradNorm_array() throws { + // the Collection overload should agree with the ModuleParameters one + try withIntegrationState(seed: 30787) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = clipGradNorm(gradients: [a, b], maxNorm: 0.5).1 + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 4.24692964553833, + minimum: 4.24692964553833, + maximum: 4.24692964553833, + absoluteSum: 4.24692964553833, + positionChecksum: 4.24692964553833, + sampleIndices: [0], + samples: [4.24692964553833]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/tall/steps1") + func test_newtonSchulz_tall_steps1() throws { + try withIntegrationState(seed: 35922) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.14183175563812256, + minimum: -0.7796086072921753, + maximum: 0.9717819690704346, + absoluteSum: 5.605131149291992, + positionChecksum: 3.1666199366251626, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.34602782130241394, 0.6664655804634094, 0.5670061111450195, + 0.9717819690704346, 0.24358342587947845, -0.7796086072921753, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/tall/steps5") + func test_newtonSchulz_tall_steps5() throws { + try withIntegrationState(seed: 97931) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 5) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.055028900504112244, + minimum: -0.9158430099487305, + maximum: 0.7314736247062683, + absoluteSum: 5.023744106292725, + positionChecksum: 2.7080958684285483, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7314736247062683, -0.21115955710411072, -0.19221490621566772, + 0.6600404977798462, 0.5243890881538391, -0.13377171754837036, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/wide/steps1") + func test_newtonSchulz_wide_steps1() throws { + try withIntegrationState(seed: 50292) { + let a = MLXRandom.normal([3, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 1) + expectSummary( + result, + ArraySummary( + shape: [3, 6], + dtype: .float32, + mean: 0.029883000999689102, + minimum: -0.7599508762359619, + maximum: 0.7960574626922607, + absoluteSum: 6.28695821762085, + positionChecksum: 3.4281453026665583, + sampleIndices: [0, 3, 7, 10, 14, 17], + samples: [ + -0.05090735852718353, 0.15048637986183167, 0.5751805305480957, + 0.04742121696472168, -0.30298084020614624, 0.017879322171211243, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/wide/steps5") + func test_newtonSchulz_wide_steps5() throws { + try withIntegrationState(seed: 47885) { + let a = MLXRandom.normal([3, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 5) + expectSummary( + result, + ArraySummary( + shape: [3, 6], + dtype: .float32, + mean: 0.11007862538099289, + minimum: -0.3809030055999756, + maximum: 0.8513897657394409, + absoluteSum: 5.038133144378662, + positionChecksum: 2.663041008843316, + sampleIndices: [0, 3, 7, 10, 14, 17], + samples: [ + -0.2114892601966858, -0.3663442134857178, 0.32633036375045776, + 0.10934531688690186, 0.2657369375228882, 0.4214053153991699, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/square/steps1") + func test_newtonSchulz_square_steps1() throws { + try withIntegrationState(seed: 54544) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: -0.01559077575802803, + minimum: -0.7211068868637085, + maximum: 0.7834442853927612, + absoluteSum: 5.679539680480957, + positionChecksum: 3.048487663269043, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + -0.13931389153003693, -0.3704991936683655, -0.3554251492023468, + 0.6534769535064697, -0.014799904078245163, -0.7211068868637085, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/square/steps5") + func test_newtonSchulz_square_steps5() throws { + try withIntegrationState(seed: 28105) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(a, steps: 5) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: -0.1895444095134735, + minimum: -0.8144998550415039, + maximum: 0.6692322492599487, + absoluteSum: 6.853260040283203, + positionChecksum: 3.5909621715545654, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + -0.8144998550415039, -0.253656804561615, -0.3669908046722412, + -0.5989289283752441, 0.261578232049942, -0.6159399747848511, + ]), + tolerance: .float32) + } + } + + @Test("newtonSchulz/nearlyOrthogonal") + func test_newtonSchulz_nearlyOrthogonal() throws { + // an already orthogonal input is the fixed point of the iteration + try withIntegrationState(seed: 28031) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.qr(a, stream: .cpu).0 + let result = Muon(learningRate: 0.1).zeropowerViaNewtonSchulz5(q, steps: 5) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.11143779754638672, + minimum: -0.5158382654190063, + maximum: 0.572691798210144, + absoluteSum: 5.199581146240234, + positionChecksum: 2.7055320739746094, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + -0.16971194744110107, 0.50611412525177, -0.3549635410308838, + 0.5651206970214844, 0.20236217975616455, 0.572691798210144, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedOptimizersTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedOptimizersTests.swift new file mode 100644 index 000000000..a64b903fb --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedOptimizersTests.swift @@ -0,0 +1,2113 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 42 + +import Foundation +import MLX +import MLXNN +import Testing + +@testable import MLXOptimizers + +@Suite("generated: Optimizers") +struct GeneratedOptimizersTests { + + @Test("SGD") + func test_SGD() throws { + try withIntegrationState(seed: 52904) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.1919216811656952, + minimum: -1.3976359367370605, + maximum: 1.300575613975525, + absoluteSum: 8.8001127243042, + positionChecksum: 4.43962033589681, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10629726201295853, -0.980926513671875, -1.0590994358062744, + -0.5857878923416138, -0.16218313574790955, -0.2747551202774048, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.2919028103351593, + minimum: -0.7022152543067932, + maximum: 0.2861098051071167, + absoluteSum: 2.031733512878418, + positionChecksum: 0.9501703262329102, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.5892770290374756, -0.7022152543067932, 0.2861098051071167, + -0.3718429207801819, -0.08228857815265656, + ]), + tolerance: .float32) + } + } + + @Test("SGD/momentum") + func test_SGD_momentum() throws { + try withIntegrationState(seed: 68261) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1, momentum: 0.9) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.20715034008026123, + minimum: -2.3862385749816895, + maximum: 1.9399197101593018, + absoluteSum: 13.50024700164795, + positionChecksum: 7.357227325439453, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.9399197101593018, -1.117470622062683, -0.3328934907913208, + 1.0880004167556763, 0.503385066986084, 0.4228379726409912, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.07929762452840805, + minimum: -1.334601640701294, + maximum: 1.1424953937530518, + absoluteSum: 3.4348740577697754, + positionChecksum: 2.365402412414551, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.3766975402832031, -0.07365846633911133, 1.1424953937530518, + -1.334601640701294, -0.5074209570884705, + ]), + tolerance: .float32) + } + } + + @Test("SGD/dampening") + func test_SGD_dampening() throws { + try withIntegrationState(seed: 79936) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1, momentum: 0.9, dampening: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.12276717275381088, + minimum: -1.7722910642623901, + maximum: 2.0786659717559814, + absoluteSum: 11.036983489990234, + positionChecksum: 5.762930552164714, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.274435818195343, 1.7900145053863525, -1.7722910642623901, + -1.1252398490905762, -1.0930812358856201, -0.4445021152496338, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.7147542834281921, + minimum: -0.8840757012367249, + maximum: -0.5355558395385742, + absoluteSum: 3.5737712383270264, + positionChecksum: 2.2165500640869142, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.6281787157058716, -0.6859133839607239, -0.8400475978851318, + -0.5355558395385742, -0.8840757012367249, + ]), + tolerance: .float32) + } + } + + @Test("SGD/weightDecay") + func test_SGD_weightDecay() throws { + try withIntegrationState(seed: 54772) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1, momentum: 0.9, weightDecay: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.14776979386806488, + minimum: -1.3177533149719238, + maximum: 1.8211688995361328, + absoluteSum: 9.807472229003906, + positionChecksum: 5.342832565307617, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.6273982524871826, -0.31063804030418396, -0.7584969401359558, + -0.4017985761165619, -0.48426440358161926, 1.0547175407409668, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.48087412118911743, + minimum: -0.46386152505874634, + maximum: 2.3275372982025146, + absoluteSum: 3.509000539779663, + positionChecksum: 2.3551183700561524, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.46386152505874634, 0.32813137769699097, 0.3010169267654419, + 2.3275372982025146, -0.08845347166061401, + ]), + tolerance: .float32) + } + } + + @Test("SGD/nesterov") + func test_SGD_nesterov() throws { + try withIntegrationState(seed: 34485) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1, momentum: 0.9, nesterov: true) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5082465410232544, + minimum: -1.6663074493408203, + maximum: 1.9156103134155273, + absoluteSum: 11.580924987792969, + positionChecksum: 6.680131276448567, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.5632967352867126, 0.5782409310340881, 1.9156103134155273, + -1.0746759176254272, 0.9145888686180115, -1.6663074493408203, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.14118611812591553, + minimum: -0.5578364133834839, + maximum: 1.4228016138076782, + absoluteSum: 2.7750954627990723, + positionChecksum: 1.6474948883056642, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.476746141910553, 0.2838525176048279, 1.4228016138076782, + 0.03385895490646362, -0.5578364133834839, + ]), + tolerance: .float32) + } + } + + @Test("RMSprop") + func test_RMSprop() throws { + try withIntegrationState(seed: 83905) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = RMSprop(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.2316233515739441, + minimum: -1.397747278213501, + maximum: 1.4456602334976196, + absoluteSum: 6.713398456573486, + positionChecksum: 4.262948671976726, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.13057762384414673, 0.1284933090209961, 0.578632652759552, + 0.40558111667633057, 1.171437382698059, -0.2595949172973633, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.2978183627128601, + minimum: -0.9918203353881836, + maximum: 1.4725697040557861, + absoluteSum: 4.434231281280518, + positionChecksum: 2.373613166809082, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.9918203353881836, 1.4725697040557861, -0.9049973487854004, + -0.10810491442680359, -0.9567388296127319, + ]), + tolerance: .float32) + } + } + + @Test("RMSprop/alpha") + func test_RMSprop_alpha() throws { + try withIntegrationState(seed: 78027) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = RMSprop(learningRate: 0.1, alpha: 0.5, eps: 1e-6) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.48947882652282715, + minimum: -2.4106788635253906, + maximum: 1.5219806432724, + absoluteSum: 11.106720924377441, + positionChecksum: 4.809648513793945, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.4106788635253906, 0.7164310812950134, -1.4400393962860107, + 0.128189355134964, 1.5219806432724, -0.10650129616260529, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.14489522576332092, + minimum: -0.25065889954566956, + maximum: 0.5770130753517151, + absoluteSum: 1.6556437015533447, + positionChecksum: 1.1530281066894532, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.2149249017238617, 0.08289781957864761, 0.5770130753517151, + -0.25065889954566956, 0.530148983001709, + ]), + tolerance: .float32) + } + } + + @Test("AdaGrad") + func test_AdaGrad() throws { + try withIntegrationState(seed: 47873) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = AdaGrad(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5292344093322754, + minimum: -0.2918749749660492, + maximum: 1.540786623954773, + absoluteSum: 7.637607574462891, + positionChecksum: 4.699300448099772, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7778294086456299, -0.2918749749660492, 0.20912456512451172, + -0.27893269062042236, 0.9240051507949829, 0.8926200270652771, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.32091256976127625, + minimum: -2.3398289680480957, + maximum: 1.4234572649002075, + absoluteSum: 4.687746047973633, + positionChecksum: 3.28100700378418, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 1.4234572649002075, 0.1181345283985138, -0.17913809418678284, + -0.6271874904632568, -2.3398289680480957, + ]), + tolerance: .float32) + } + } + + @Test("AdaDelta") + func test_AdaDelta() throws { + try withIntegrationState(seed: 4025) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = AdaDelta(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.22225290536880493, + minimum: -1.8476370573043823, + maximum: 1.6021405458450317, + absoluteSum: 11.606254577636719, + positionChecksum: 5.977849960327148, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.8476370573043823, -1.1789054870605469, 1.6021405458450317, + 0.3608337938785553, 0.5647247433662415, -1.187432050704956, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.5898528099060059, + minimum: -1.4875011444091797, + maximum: 0.8329288363456726, + absoluteSum: 4.615121841430664, + positionChecksum: 3.0648353576660154, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.6340699791908264, -0.4064069092273712, -1.254214882850647, + -1.4875011444091797, 0.8329288363456726, + ]), + tolerance: .float32) + } + } + + @Test("AdaDelta/rho") + func test_AdaDelta_rho() throws { + try withIntegrationState(seed: 56156) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = AdaDelta(learningRate: 0.1, rho: 0.5, eps: 1e-5) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.6433796882629395, + minimum: -1.97148597240448, + maximum: 3.078878402709961, + absoluteSum: 11.844547271728516, + positionChecksum: 6.222557703653972, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.07173069566488266, 0.7374876737594604, 0.23356150090694427, + -0.018778571859002113, 0.8372526168823242, 0.6959485411643982, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.2940554618835449, + minimum: -0.3743702173233032, + maximum: 0.8940013647079468, + absoluteSum: 2.21901798248291, + positionChecksum: 1.5835657119750977, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.320523202419281, -0.3743702173233032, 0.14193440973758698, + 0.48818862438201904, 0.8940013647079468, + ]), + tolerance: .float32) + } + } + + @Test("Adam") + func test_Adam() throws { + try withIntegrationState(seed: 2855) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adam(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.20924310386180878, + minimum: -1.4468461275100708, + maximum: 1.958080768585205, + absoluteSum: 11.411166191101074, + positionChecksum: 5.717811584472656, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.958080768585205, -0.2732287645339966, 1.9326785802841187, + 0.7665133476257324, -1.0778547525405884, -0.7260398268699646, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.43796560168266296, + minimum: -0.16814911365509033, + maximum: 0.8766471147537231, + absoluteSum: 2.5261263847351074, + positionChecksum: 1.248355770111084, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.8766471147537231, 0.6090426445007324, 0.18284839391708374, + 0.689439058303833, -0.16814911365509033, + ]), + tolerance: .float32) + } + } + + @Test("Adam/betas") + func test_Adam_betas() throws { + try withIntegrationState(seed: 26177) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adam(learningRate: 0.1, betas: (0.8, 0.9)) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.13137023150920868, + minimum: -2.1283836364746094, + maximum: 1.7582262754440308, + absoluteSum: 9.839652061462402, + positionChecksum: 5.175012588500977, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0437983274459839, 0.7374716997146606, -2.1283836364746094, + 0.43236860632896423, 1.7582262754440308, 0.6263670325279236, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.47237634658813477, + minimum: -2.5258002281188965, + maximum: 0.687131404876709, + absoluteSum: 3.736144542694092, + positionChecksum: 3.006717872619629, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.687131404876709, -0.10542242228984833, -0.1645507514476776, + -0.25323984026908875, -2.5258002281188965, + ]), + tolerance: .float32) + } + } + + @Test("Adam/biasCorrection") + func test_Adam_biasCorrection() throws { + // bias correction is step dependent: identical at step 1, not later + try withIntegrationState(seed: 75773) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adam(learningRate: 0.1, biasCorrection: true) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.38266444206237793, + minimum: -0.3790029287338257, + maximum: 1.114660620689392, + absoluteSum: 5.856780052185059, + positionChecksum: 3.123239199320475, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.1374974548816681, 0.47027307748794556, 0.7321585416793823, + 1.114660620689392, -0.11590324342250824, 0.2759057581424713, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.7057530879974365, + minimum: -1.7636712789535522, + maximum: 1.1074949502944946, + absoluteSum: 5.743755340576172, + positionChecksum: 2.8100845336914064, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -1.7636712789535522, 1.1074949502944946, -1.479508638381958, + -1.3321667909622192, -0.060913801193237305, + ]), + tolerance: .float32) + } + } + + @Test("AdamW") + func test_AdamW() throws { + try withIntegrationState(seed: 75439) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = AdamW(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5645228028297424, + minimum: -0.8113094568252563, + maximum: 1.9090187549591064, + absoluteSum: 9.240545272827148, + positionChecksum: 4.90969975789388, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.8444397449493408, 0.6620670557022095, 0.06166234612464905, + -0.19734594225883484, -0.22448092699050903, -0.8113094568252563, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.11740782111883163, + minimum: -1.336305022239685, + maximum: 1.2092561721801758, + absoluteSum: 4.171524524688721, + positionChecksum: 2.539629364013672, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -1.336305022239685, -0.45593777298927307, 0.11859375238418579, + 1.2092561721801758, 1.0514320135116577, + ]), + tolerance: .float32) + } + } + + @Test("AdamW/weightDecay") + func test_AdamW_weightDecay() throws { + try withIntegrationState(seed: 45988) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = AdamW(learningRate: 0.1, weightDecay: 0.1, biasCorrection: true) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.020839249715209007, + minimum: -0.9159260392189026, + maximum: 0.964216411113739, + absoluteSum: 5.242161750793457, + positionChecksum: 2.8174708684285483, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.24218687415122986, -0.09974238276481628, -0.7899527549743652, + 0.964216411113739, 0.6594827175140381, 0.3614733815193176, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.12586472928524017, + minimum: -0.33768603205680847, + maximum: 0.8633996248245239, + absoluteSum: 2.1387057304382324, + positionChecksum: 1.5357303619384766, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.33768603205680847, -0.10294844210147858, 0.5206149816513062, + -0.3140564560890198, 0.8633996248245239, + ]), + tolerance: .float32) + } + } + + @Test("Adamax") + func test_Adamax() throws { + try withIntegrationState(seed: 15957) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adamax(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.09984465688467026, + minimum: -0.9413069486618042, + maximum: 1.3880528211593628, + absoluteSum: 9.504159927368164, + positionChecksum: 4.476319630940755, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.34331476688385, -0.785368800163269, -0.5558609962463379, + 0.9454747438430786, 0.5932865738868713, -0.9413069486618042, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.4311213493347168, + minimum: -0.5994178056716919, + maximum: 1.8316594362258911, + absoluteSum: 3.5902762413024902, + positionChecksum: 2.48956184387207, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.6177985668182373, -0.11791691929101944, 0.4234835207462311, + 1.8316594362258911, -0.5994178056716919, + ]), + tolerance: .float32) + } + } + + @Test("Lion") + func test_Lion() throws { + try withIntegrationState(seed: 36413) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Lion(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.13024353981018066, + minimum: -1.2807341814041138, + maximum: 0.7626199722290039, + absoluteSum: 6.7521820068359375, + positionChecksum: 3.5021870930989585, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7626199722290039, -0.6869360208511353, 0.7073222398757935, + 0.5880645513534546, 0.3303239941596985, -0.37053045630455017, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.6855365633964539, + minimum: -0.7185075879096985, + maximum: 2.3334498405456543, + absoluteSum: 4.864697456359863, + positionChecksum: 2.822071075439453, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.7441189289093018, -0.7185075879096985, 2.3334498405456543, + 0.4142351746559143, 0.6543862223625183, + ]), + tolerance: .float32) + } + } + + @Test("Lion/weightDecay") + func test_Lion_weightDecay() throws { + try withIntegrationState(seed: 7472) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Lion(learningRate: 0.1, betas: (0.8, 0.9), weightDecay: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3012966513633728, + minimum: -1.460852026939392, + maximum: 1.3098098039627075, + absoluteSum: 6.7537384033203125, + positionChecksum: 3.571783701578776, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.10823702067136765, 0.02829088270664215, 0.49554669857025146, + 0.34063977003097534, 0.1592463254928589, 0.7434209585189819, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.19140489399433136, + minimum: -1.1581501960754395, + maximum: 0.6329450011253357, + absoluteSum: 3.3792479038238525, + positionChecksum: 1.5446331024169921, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -1.1581501960754395, -0.917064905166626, 0.5781667828559875, + 0.6329450011253357, -0.09292108565568924, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor") + func test_Adafactor() throws { + // the 2-D parameter is factored and the 1-D one is not + try withIntegrationState(seed: 1795) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor(learningRate: 0.1, relativeStep: false) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.25985509157180786, + minimum: -1.102706789970398, + maximum: 1.0535223484039307, + absoluteSum: 6.386711120605469, + positionChecksum: 3.502171516418457, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.13309255242347717, 0.447610080242157, -1.102706789970398, + -0.7704825401306152, -0.013223405927419662, -0.8421300053596497, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.07308216392993927, + minimum: -1.8284789323806763, + maximum: 1.3594447374343872, + absoluteSum: 4.476809978485107, + positionChecksum: 3.4291236877441404, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.1359250843524933, 0.5603299140930176, -0.5926315784454346, + -1.8284789323806763, 1.3594447374343872, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor/relativeStep") + func test_Adafactor_relativeStep() throws { + // without a learning rate the step size comes from the step count + try withIntegrationState(seed: 77705) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor() + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.11738912761211395, + minimum: -0.6956945657730103, + maximum: 1.4269239902496338, + absoluteSum: 6.289972305297852, + positionChecksum: 4.015622456868489, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.15101845562458038, -0.4634033441543579, -0.4472257196903229, + -0.2334236353635788, 1.4269239902496338, -0.6956945657730103, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.5099145174026489, + minimum: -2.1994707584381104, + maximum: 1.0201630592346191, + absoluteSum: 4.589898586273193, + positionChecksum: 2.507386016845703, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.40986695885658264, -2.1994707584381104, -0.5772594213485718, + 1.0201630592346191, -0.3831382393836975, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor/beta1") + func test_Adafactor_beta1() throws { + try withIntegrationState(seed: 48053) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor(learningRate: 0.1, beta1: 0.9, relativeStep: false) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.16388410329818726, + minimum: -0.9698700308799744, + maximum: 1.4167040586471558, + absoluteSum: 8.533632278442383, + positionChecksum: 4.301989237467448, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.9218981862068176, 1.1756681203842163, 1.4167040586471558, + -0.41768917441368103, -0.8146899342536926, -0.12940138578414917, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.5856785774230957, + minimum: -1.6530791521072388, + maximum: 0.16239207983016968, + absoluteSum: 3.2531771659851074, + positionChecksum: 1.8036436080932616, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.9667976498603821, -0.4675377607345581, 0.16239207983016968, + -1.6530791521072388, -0.0033703986555337906, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor/noScaleParameter") + func test_Adafactor_noScaleParameter() throws { + try withIntegrationState(seed: 81017) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor( + learningRate: 0.1, weightDecay: 0.1, scaleParameter: false, relativeStep: false) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.14465275406837463, + minimum: -1.1899194717407227, + maximum: 1.5379257202148438, + absoluteSum: 8.267154693603516, + positionChecksum: 5.085835138956706, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.026530228555202484, -0.9377098083496094, 0.4836030900478363, + 0.6172863245010376, 0.016763836145401, -1.1380316019058228, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.34424692392349243, + minimum: -0.3117382228374481, + maximum: 1.2032966613769531, + absoluteSum: 2.567047357559204, + positionChecksum: 1.3552351951599122, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.42866355180740356, 1.2032966613769531, -0.11116818338632584, + 0.5121809244155884, -0.3117382228374481, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor/warmupInit") + func test_Adafactor_warmupInit() throws { + try withIntegrationState(seed: 48640) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor(warmupInit: true) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.3013237416744232, + minimum: -3.0539402961730957, + maximum: 1.6986980438232422, + absoluteSum: 13.95006275177002, + positionChecksum: 7.426012674967448, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.6986980438232422, -2.1536543369293213, -0.3216652274131775, + 0.16456238925457, 0.7074159383773804, -1.3977259397506714, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.565890908241272, + minimum: -1.6157656908035278, + maximum: 0.5499652028083801, + absoluteSum: 4.476120471954346, + positionChecksum: 2.3359392166137694, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -1.2608189582824707, 0.5499652028083801, -1.6157656908035278, + -0.7762027978897095, 0.2733677625656128, + ]), + tolerance: .float32) + } + } + + @Test("Muon") + func test_Muon() throws { + // loose: Muon compounds the Float-vs-double hyperparameter difference (todos/muon-04) to ~1e-4; momentumZero/oneStep/newtonSchulz are the controls + try withIntegrationState(seed: 20620) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.03163754940032959, + minimum: -2.5139663219451904, + maximum: 2.240089178085327, + absoluteSum: 11.950052261352539, + positionChecksum: 6.561954498291016, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.03799111396074295, 2.12414288520813, -1.5080902576446533, + 2.240089178085327, -2.5139663219451904, -0.06351476162672043, + ]), + tolerance: .loose) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.4560411870479584, + minimum: -2.053658962249756, + maximum: 0.9637644290924072, + absoluteSum: 4.570568561553955, + positionChecksum: 2.6469955444335938, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.18141686916351318, -1.2737054824829102, -2.053658962249756, + 0.9637644290924072, -0.09802297502756119, + ]), + tolerance: .loose) + } + } + + @Test("Muon/vectorOnly") + func test_Muon_vectorOnly() throws { + // rank 1: plain momentum update, no orthogonalization + try withIntegrationState(seed: 69306) { + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.20210061967372894, + minimum: -1.4918893575668335, + maximum: 2.1059176921844482, + absoluteSum: 6.641729354858398, + positionChecksum: 3.698236846923828, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 2.1059176921844482, -0.45383501052856445, 1.720198631286621, + -1.4918893575668335, -0.8698888421058655, + ]), + tolerance: .float32) + } + } + + @Test("Muon/matrixOnly") + func test_Muon_matrixOnly() throws { + // tall matrix: the Newton-Schulz iteration transposes first + try withIntegrationState(seed: 75197) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.08751313388347626, + minimum: -1.5594749450683594, + maximum: 0.9060735702514648, + absoluteSum: 6.130833625793457, + positionChecksum: 3.4926980336507163, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.7616370916366577, 0.08012694865465164, 0.3566833436489105, + -1.5594749450683594, 0.5005760788917542, 0.8296913504600525, + ]), + tolerance: .loose) + } + } + + @Test("Muon/wide") + func test_Muon_wide() throws { + // wide matrix: no transpose in the iteration + try withIntegrationState(seed: 50653) { + let weight = MLXRandom.normal([3, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([3, 6], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [3, 6], + dtype: .float32, + mean: 0.018045898526906967, + minimum: -1.1826281547546387, + maximum: 2.2200303077697754, + absoluteSum: 13.78824234008789, + positionChecksum: 5.91921149359809, + sampleIndices: [0, 3, 7, 10, 14, 17], + samples: [ + 2.2200303077697754, 1.1188658475875854, -0.8831031918525696, + -0.7170143723487854, -0.5960718989372253, -0.6579134464263916, + ]), + tolerance: .loose) + } + } + + @Test("Muon/oneStep") + func test_Muon_oneStep() throws { + // first update only, so the momentum state does not carry + try withIntegrationState(seed: 18547) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 1 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.04324345290660858, + minimum: -2.483652353286743, + maximum: 1.6655775308609009, + absoluteSum: 6.838813781738281, + positionChecksum: 2.857765515645345, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.38655349612236023, 0.19337239861488342, -0.26349616050720215, + 0.9585155248641968, 0.10933685302734375, -0.05299711227416992, + ]), + tolerance: .float32) + } + } + + @Test("Muon/oneNewtonSchulzStep") + func test_Muon_oneNewtonSchulzStep() throws { + // a single iteration: if this agrees and the 5 step default does not, the difference is amplification rather than a formula error + try withIntegrationState(seed: 64243) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1, nsSteps: 1) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.06221400946378708, + minimum: -2.0203006267547607, + maximum: 2.590712785720825, + absoluteSum: 10.942645072937012, + positionChecksum: 6.706094741821289, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.2853350043296814, -0.027788713574409485, 1.4993284940719604, + -0.9181387424468994, -1.357496976852417, -0.14571359753608704, + ]), + tolerance: .float32) + } + } + + @Test("Muon/twoSteps") + func test_Muon_twoSteps() throws { + // with oneStep and threeSteps this shows how fast a difference grows + try withIntegrationState(seed: 45980) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 2 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.028408613055944443, + minimum: -0.846992552280426, + maximum: 1.4628791809082031, + absoluteSum: 7.093116760253906, + positionChecksum: 3.6778895060221353, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.5931099057197571, 1.4628791809082031, -0.8363688588142395, + 0.7522042393684387, -0.7285213470458984, 0.362490713596344, + ]), + tolerance: .loose) + } + } + + @Test("Muon/momentumZero") + func test_Muon_momentumZero() throws { + // momentum 0 removes the (1 - momentum) coefficient, which python computes in double and Swift in Float; if this agrees and the default does not, that coefficient is the seed of the difference + try withIntegrationState(seed: 18459) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1, momentum: 0.0, nesterov: false) + var parameters = ModuleParameters.unflattened([("weight", weight)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)) + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.597805380821228, + minimum: -2.283677816390991, + maximum: 0.8412001132965088, + absoluteSum: 10.496516227722168, + positionChecksum: 5.905821482340495, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0269737243652344, 0.06046082079410553, -0.9181715846061707, + -2.283677816390991, 0.8412001132965088, -0.649782121181488, + ]), + tolerance: .float32) + } + } + + @Test("Muon/noWeightDecay") + func test_Muon_noWeightDecay() throws { + // isolates the weight decay term (python defaults it to 0.01) + try withIntegrationState(seed: 15172) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon(learningRate: 0.1, weightDecay: 0.0) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.041766803711652756, + minimum: -1.3217413425445557, + maximum: 1.0777087211608887, + absoluteSum: 8.2870512008667, + positionChecksum: 4.308595657348633, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8723462224006653, -0.9842115640640259, -1.3217413425445557, + -0.7982552647590637, -1.110180377960205, 0.4663331210613251, + ]), + tolerance: .loose) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.23301398754119873, + minimum: -0.6206360459327698, + maximum: 1.185815691947937, + absoluteSum: 3.5016415119171143, + positionChecksum: 2.000499153137207, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.7082110643386841, -0.6206360459327698, 1.185815691947937, + 0.43932899832725525, -0.5476498603820801, + ]), + tolerance: .loose) + } + } + + @Test("Lion/oneStep") + func test_Lion_oneStep() throws { + // Lion's update is lr * sign(c), so a mismatched beta shows up as a difference of exactly 2 * lr per element + try withIntegrationState(seed: 34768) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Lion(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 1 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.19433651864528656, + minimum: -1.3650480508804321, + maximum: 0.5978026390075684, + absoluteSum: 5.644954681396484, + positionChecksum: 2.9201577504475913, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.2685225009918213, -0.7982114553451538, -1.3650480508804321, + -0.06641935557126999, -0.21257930994033813, 0.5896715521812439, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.07899947464466095, + minimum: -1.4680547714233398, + maximum: 0.9950604438781738, + absoluteSum: 3.8233253955841064, + positionChecksum: 2.4469171524047852, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.6411066055297852, 0.9950604438781738, 0.6133296489715576, + 0.10577392578125, -1.4680547714233398, + ]), + tolerance: .float32) + } + } + + @Test("Muon/noNesterov") + func test_Muon_noNesterov() throws { + try withIntegrationState(seed: 59801) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Muon( + learningRate: 0.1, momentum: 0.8, weightDecay: 0.0, nesterov: false, nsSteps: 3) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.2554815411567688, + minimum: -0.9771368503570557, + maximum: 2.7309908866882324, + absoluteSum: 8.80552864074707, + positionChecksum: 4.27720578511556, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.04743572324514389, -0.7059667706489563, 0.5713173151016235, + 0.7252044677734375, -0.2173776626586914, -0.25553134083747864, + ]), + tolerance: .loose) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.5632168650627136, + minimum: -1.19088876247406, + maximum: -0.19167444109916687, + absoluteSum: 2.816084146499634, + positionChecksum: 2.11552734375, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.19167444109916687, -0.280126690864563, -0.35240861773490906, + -1.19088876247406, -0.80098557472229, + ]), + tolerance: .loose) + } + } + + @Test("Adam/tenSteps") + func test_Adam_tenSteps() throws { + try withIntegrationState(seed: 36115) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adam(learningRate: 0.05, biasCorrection: true) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 10 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.02142399549484253, + minimum: -0.9353051781654358, + maximum: 0.8552088737487793, + absoluteSum: 5.777521133422852, + positionChecksum: 2.9126103719075522, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.24671006202697754, 0.8552088737487793, -0.9353051781654358, + 0.46673938632011414, -0.26257777214050293, -0.3439560532569885, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.20665812492370605, + minimum: -1.7892528772354126, + maximum: 0.8451172709465027, + absoluteSum: 3.4166171550750732, + positionChecksum: 1.5547866821289062, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -1.7892528772354126, -0.013102374039590359, 0.8451172709465027, + -0.42259863018989563, 0.3465459942817688, + ]), + tolerance: .float32) + } + } + + @Test("SGD/singleStep") + func test_SGD_singleStep() throws { + // the first step alone, for comparison with the multi-step cases + try withIntegrationState(seed: 38685) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = SGD(learningRate: 0.1, momentum: 0.9) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 1 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.0076949745416641235, + minimum: -0.7979379296302795, + maximum: 1.260810375213623, + absoluteSum: 6.1943678855896, + positionChecksum: 3.0758225123087564, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.127612590789795, -0.1752968430519104, 0.3452015519142151, + -0.7125493288040161, 0.18754345178604126, -0.3442728519439697, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.0020482302643358707, + minimum: -0.7621648907661438, + maximum: 0.3773851990699768, + absoluteSum: 1.540561318397522, + positionChecksum: 0.814434814453125, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + -0.7621648907661438, -0.013236373662948608, 0.08243922144174576, + 0.3773851990699768, 0.3053356409072876, + ]), + tolerance: .float32) + } + } + + @Test("Adam/shapes") + func test_Adam_shapes() throws { + try withIntegrationState(seed: 98800) { + let scalar = MLXRandom.normal([Int](), dtype: .float32, loc: 0.0, scale: 1.0) + let vector = MLXRandom.normal([7], dtype: .float32, loc: 0.0, scale: 1.0) + let matrix = MLXRandom.normal([3, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let tensor = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let scalarTarget = MLXRandom.normal([Int](), dtype: .float32, loc: 0.0, scale: 1.0) + let vectorTarget = MLXRandom.normal([7], dtype: .float32, loc: 0.0, scale: 1.0) + let matrixTarget = MLXRandom.normal([3, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let tensorTarget = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adam(learningRate: 0.1) + var parameters = ModuleParameters.unflattened([ + ("scalar", scalar), ("vector", vector), ("matrix", matrix), ("tensor", tensor), + ]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("scalar", 2 * (parameters[unwrapping: "scalar"]! - scalarTarget)), + ("vector", 2 * (parameters[unwrapping: "vector"]! - vectorTarget)), + ("matrix", 2 * (parameters[unwrapping: "matrix"]! - matrixTarget)), + ("tensor", 2 * (parameters[unwrapping: "tensor"]! - tensorTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "scalar"]!, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.04785709083080292, + minimum: -0.04785709083080292, + maximum: -0.04785709083080292, + absoluteSum: 0.04785709083080292, + positionChecksum: 0.04785709083080292, + sampleIndices: [0], + samples: [-0.04785709083080292]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "vector"]!, + ArraySummary( + shape: [7], + dtype: .float32, + mean: -0.5064614415168762, + minimum: -2.039198875427246, + maximum: 0.8779197931289673, + absoluteSum: 6.674189567565918, + positionChecksum: 2.9789466857910156, + sampleIndices: [0, 1, 2, 4, 5, 6], + samples: [ + -1.4115877151489258, -2.039198875427246, -0.9141958951950073, + 0.8779197931289673, -0.7447269558906555, 0.33861857652664185, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "matrix"]!, + ArraySummary( + shape: [3, 5], + dtype: .float32, + mean: -0.028072740882635117, + minimum: -1.7022106647491455, + maximum: 1.6618009805679321, + absoluteSum: 14.534392356872559, + positionChecksum: 7.248495483398438, + sampleIndices: [0, 3, 6, 8, 11, 14], + samples: [ + -1.1987392902374268, 0.766374945640564, -1.7022106647491455, + 0.8694322109222412, -0.6111583113670349, -0.4536212980747223, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "tensor"]!, + ArraySummary( + shape: [2, 3, 4], + dtype: .float32, + mean: -0.10296154022216797, + minimum: -1.9619927406311035, + maximum: 1.3071478605270386, + absoluteSum: 13.969619750976562, + positionChecksum: 7.8157094319661455, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.4101582169532776, 1.3071478605270386, 0.08434449136257172, + 0.14656518399715424, -1.9619927406311035, -1.6114184856414795, + ]), + tolerance: .float32) + } + } + + @Test("Adafactor/shapes") + func test_Adafactor_shapes() throws { + try withIntegrationState(seed: 41330) { + let vector = MLXRandom.normal([7], dtype: .float32, loc: 0.0, scale: 1.0) + let matrix = MLXRandom.normal([3, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let tensor = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let vectorTarget = MLXRandom.normal([7], dtype: .float32, loc: 0.0, scale: 1.0) + let matrixTarget = MLXRandom.normal([3, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let tensorTarget = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let optimizer = Adafactor(learningRate: 0.1, relativeStep: false) + var parameters = ModuleParameters.unflattened([ + ("vector", vector), ("matrix", matrix), ("tensor", tensor), + ]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("vector", 2 * (parameters[unwrapping: "vector"]! - vectorTarget)), + ("matrix", 2 * (parameters[unwrapping: "matrix"]! - matrixTarget)), + ("tensor", 2 * (parameters[unwrapping: "tensor"]! - tensorTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "vector"]!, + ArraySummary( + shape: [7], + dtype: .float32, + mean: 0.24444937705993652, + minimum: -0.7275797724723816, + maximum: 1.3433971405029297, + absoluteSum: 3.5701165199279785, + positionChecksum: 2.086233820234026, + sampleIndices: [0, 1, 2, 4, 5, 6], + samples: [ + 0.7943839430809021, 0.034648530185222626, -0.20190563797950745, + 0.14161652326583862, 0.3265848159790039, -0.7275797724723816, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "matrix"]!, + ArraySummary( + shape: [3, 5], + dtype: .float32, + mean: -0.022888900712132454, + minimum: -1.4636974334716797, + maximum: 1.163638949394226, + absoluteSum: 8.799198150634766, + positionChecksum: 4.993682861328125, + sampleIndices: [0, 3, 6, 8, 11, 14], + samples: [ + 0.019102375954389572, 0.8244732618331909, -0.2120092213153839, + 0.47228384017944336, 0.7905675172805786, -0.47224366664886475, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "tensor"]!, + ArraySummary( + shape: [2, 3, 4], + dtype: .float32, + mean: 0.1986042559146881, + minimum: -0.7255576252937317, + maximum: 1.7102957963943481, + absoluteSum: 12.307641983032227, + positionChecksum: 6.778558095296224, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.2963857650756836, 0.9113672375679016, -0.4616476595401764, + 0.8774430155754089, 0.8562024831771851, 0.14194099605083466, + ]), + tolerance: .float32) + } + } + + @Test("SGD/positive") + func test_SGD_positive() throws { + try withIntegrationState(seed: 1980) { + let weight = MLXRandom.uniform(low: 0.5, high: 1.5, [4, 3], dtype: .float32) + let bias = MLXRandom.uniform(low: 0.5, high: 1.5, [5], dtype: .float32) + let weightTarget = MLXRandom.uniform(low: 0.5, high: 1.5, [4, 3], dtype: .float32) + let biasTarget = MLXRandom.uniform(low: 0.5, high: 1.5, [5], dtype: .float32) + let optimizer = SGD(learningRate: 0.05, momentum: 0.9, weightDecay: 0.2) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for _ in 0 ..< 3 { + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.0045102834701538, + minimum: 0.7203009724617004, + maximum: 1.2829526662826538, + absoluteSum: 12.054122924804688, + positionChecksum: 6.408329010009766, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.152671456336975, 0.759495735168457, 0.9144626259803772, + 0.8906338214874268, 1.2829526662826538, 0.7203009724617004, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.9440234303474426, + minimum: 0.7767097353935242, + maximum: 1.3086878061294556, + absoluteSum: 4.720117092132568, + positionChecksum: 2.8720470428466798, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.8757483959197998, 0.9107287526130676, 0.8482428789138794, + 1.3086878061294556, 0.7767097353935242, + ]), + tolerance: .float32) + } + } + + @Test("SGD/cosineDecay") + func test_SGD_cosineDecay() throws { + try withIntegrationState(seed: 70093) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let schedule = cosineDecay(0.1, decaySteps: 10) + let optimizer = SGD(learningRate: schedule(0), momentum: 0.9) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for step in 0 ..< 5 { + optimizer.learningRate = schedule(step) + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5896511077880859, + minimum: -1.0688788890838623, + maximum: 2.913727045059204, + absoluteSum: 14.172250747680664, + positionChecksum: 6.361111958821614, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0688788890838623, 1.9024821519851685, 0.9821493029594421, + -0.4725768268108368, -0.4211561679840088, 0.9433878660202026, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.006274223327636719, + minimum: -2.107053756713867, + maximum: 1.7379882335662842, + absoluteSum: 4.676799774169922, + positionChecksum: 2.9456262588500977, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 1.7379882335662842, -0.21566040813922882, 0.44083574414253235, + 0.1752614676952362, -2.107053756713867, + ]), + tolerance: .float32) + } + } + + @Test("Adam/exponentialDecay") + func test_Adam_exponentialDecay() throws { + try withIntegrationState(seed: 60243) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let schedule = exponentialDecay(0.1, decayRate: 0.9) + let optimizer = Adam(learningRate: schedule(0)) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for step in 0 ..< 5 { + optimizer.learningRate = schedule(step) + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.3028128147125244, + minimum: -1.2761456966400146, + maximum: 0.6678332686424255, + absoluteSum: 7.528482913970947, + positionChecksum: 4.288489977518718, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.0958678722381592, 0.2696917951107025, -0.22582115232944489, + -1.2761456966400146, 0.34309184551239014, -1.0258384943008423, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: 0.3606184124946594, + minimum: -0.35021162033081055, + maximum: 1.2646141052246094, + absoluteSum: 2.5035152435302734, + positionChecksum: 1.0792693138122558, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 1.2646141052246094, 0.3743942081928253, 0.425295889377594, + 0.08899958431720734, -0.35021162033081055, + ]), + tolerance: .float32) + } + } + + @Test("SGD/stepDecay") + func test_SGD_stepDecay() throws { + try withIntegrationState(seed: 56543) { + let weight = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let bias = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let weightTarget = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let biasTarget = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let schedule = stepDecay(0.1, decayRate: 0.5, stepSize: 2) + let optimizer = SGD(learningRate: schedule(0)) + var parameters = ModuleParameters.unflattened([("weight", weight), ("bias", bias)]) + for step in 0 ..< 6 { + optimizer.learningRate = schedule(step) + let gradients = ModuleParameters.unflattened([ + ("weight", 2 * (parameters[unwrapping: "weight"]! - weightTarget)), + ("bias", 2 * (parameters[unwrapping: "bias"]! - biasTarget)), + ]) + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + } + expectSummary( + parameters[unwrapping: "weight"]!, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.45779770612716675, + minimum: -0.9003390669822693, + maximum: 2.1724798679351807, + absoluteSum: 9.183296203613281, + positionChecksum: 5.28756841023763, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.4551224708557129, 0.24549242854118347, 2.1724798679351807, + 1.0964405536651611, -0.5553304553031921, 0.46668076515197754, + ]), + tolerance: .float32) + expectSummary( + parameters[unwrapping: "bias"]!, + ArraySummary( + shape: [5], + dtype: .float32, + mean: -0.051898252218961716, + minimum: -0.9068479537963867, + maximum: 1.0226328372955322, + absoluteSum: 3.2517483234405518, + positionChecksum: 1.760629653930664, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 1.0226328372955322, -0.9068479537963867, -0.085511215031147, + 0.47349560260772705, -0.763260543346405, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift new file mode 100644 index 000000000..11ce7fffa --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedQuantizationTests.swift @@ -0,0 +1,571 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 22 + +import Foundation +import MLX +import Testing + +@Suite("generated: Quantization") +struct GeneratedQuantizationTests { + + @Test("quantized/bits2/wq") + func test_quantized_bits2_wq() throws { + try withIntegrationState(seed: 44380) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 2).wq + expectSummary( + result, + ArraySummary( + shape: [64, 8], + dtype: .uint32, + mean: 2706155520.0, + minimum: 55539704.0, + maximum: 4294343168.0, + absoluteSum: 1385551626240.0, + positionChecksum: 690314412032.0, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 1849313664.0, 1907358848.0, 2144858240.0, 328123136.0, 3683949312.0, + 1981644544.0, + ]), + tolerance: .exact) + } + } + + @Test("quantized/bits2/scales") + func test_quantized_bits2_scales() throws { + try withIntegrationState(seed: 66112) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 2).scales + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: -0.006201982498168945, + minimum: -0.49998176097869873, + maximum: 0.4999878704547882, + absoluteSum: 62.84138107299805, + positionChecksum: 31.624141693115234, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.4971696436405182, 0.4996141493320465, 0.49712875485420227, + 0.4998661279678345, 0.49900156259536743, -0.48680394887924194, + ]), + tolerance: .float32) + } + } + + @Test("quantized/bits2/biases") + func test_quantized_bits2_biases() throws { + try withIntegrationState(seed: 88914) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 2).biases! + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: -0.04578951746225357, + minimum: -0.9995979070663452, + maximum: 0.9998750686645508, + absoluteSum: 126.25718688964844, + positionChecksum: 63.589569091796875, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.9953732490539551, -0.9974033832550049, 0.9865076541900635, + -0.9958232045173645, -0.9863418936729431, -0.9527243375778198, + ]), + tolerance: .float32) + } + } + + @Test("dequantized/bits2") + func test_dequantized_bits2() throws { + try withIntegrationState(seed: 42894) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let q = MLX.quantized(w, groupSize: 64, bits: 2) + let result = MLX.dequantized( + q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 2) + expectSummary( + result, + ArraySummary( + shape: [64, 128], + dtype: .float32, + mean: 0.02056698314845562, + minimum: -0.9999780654907227, + maximum: 0.9999353885650635, + absoluteSum: 3605.364501953125, + positionChecksum: 1786.4384765625, + sampleIndices: [0, 1638, 3276, 4915, 6553, 8191], + samples: [ + -0.4797353148460388, 0.0, 0.9931421279907227, 0.49951058626174927, + -0.4997556805610657, -0.4949713349342346, + ]), + tolerance: .float32) + } + } + + @Test("quantizedMM/bits2") + func test_quantizedMM_bits2() throws { + try withIntegrationState(seed: 40462) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w, groupSize: 64, bits: 2) + let result = MLX.quantizedMM( + x, q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 2) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: -0.2884669303894043, + minimum: -21.208293914794922, + maximum: 16.252973556518555, + absoluteSum: 2394.50634765625, + positionChecksum: 1203.15673828125, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 5.748183727264404, -3.221513271331787, -0.056482791900634766, + -0.08748602867126465, -2.5631000995635986, 7.229959487915039, + ]), + tolerance: .float32) + } + } + + @Test("quantized/bits4/wq") + func test_quantized_bits4_wq() throws { + try withIntegrationState(seed: 44608) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 4).wq + expectSummary( + result, + ArraySummary( + shape: [64, 16], + dtype: .uint32, + mean: 2183851520.0, + minimum: 21900200.0, + maximum: 4294455040.0, + absoluteSum: 2236263956480.0, + positionChecksum: 1127263502336.0, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + 1132585600.0, 573891328.0, 1776925312.0, 2403129856.0, 2522929152.0, + 1568619904.0, + ]), + tolerance: .exact) + } + } + + @Test("quantized/bits4/scales") + func test_quantized_bits4_scales() throws { + try withIntegrationState(seed: 46023) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 4).scales + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: -0.003454235615208745, + minimum: -0.12493161857128143, + maximum: 0.1249997615814209, + absoluteSum: 15.703900337219238, + positionChecksum: 7.915512561798096, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + -0.12353004515171051, 0.12238717824220657, -0.11830480396747589, + 0.12461955100297928, -0.1237143725156784, 0.12481806427240372, + ]), + tolerance: .float32) + } + } + + @Test("quantized/bits4/biases") + func test_quantized_bits4_biases() throws { + try withIntegrationState(seed: 34133) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 4).biases! + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: -0.09647571295499802, + minimum: -0.9999547600746155, + maximum: 0.9996299743652344, + absoluteSum: 126.08390808105469, + positionChecksum: 63.495304107666016, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + -0.9987098574638367, 0.9933186769485474, -0.9994155168533325, + 0.9901624917984009, -0.9994679093360901, 0.9919745922088623, + ]), + tolerance: .float32) + } + } + + @Test("dequantized/bits4") + func test_dequantized_bits4() throws { + try withIntegrationState(seed: 62267) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let q = MLX.quantized(w, groupSize: 64, bits: 4) + let result = MLX.dequantized( + q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 4) + expectSummary( + result, + ArraySummary( + shape: [64, 128], + dtype: .float32, + mean: 0.004197360016405582, + minimum: -0.9999241828918457, + maximum: 0.9990442991256714, + absoluteSum: 4095.062255859375, + positionChecksum: 2047.1114501953125, + sampleIndices: [0, 1638, 3276, 4915, 6553, 8191], + samples: [ + 0.24011309444904327, 0.12464109063148499, 0.6011228561401367, + 0.6184362769126892, 0.12476050853729248, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("quantizedMM/bits4") + func test_quantizedMM_bits4() throws { + try withIntegrationState(seed: 15163) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w, groupSize: 64, bits: 4) + let result = MLX.quantizedMM( + x, q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 4) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: -0.029651537537574768, + minimum: -16.354642868041992, + maximum: 19.714792251586914, + absoluteSum: 2415.8642578125, + positionChecksum: 1229.1107177734375, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 3.3405404090881348, 1.769765019416809, -1.461564540863037, + -1.8853719234466553, -13.134235382080078, 0.8762807846069336, + ]), + tolerance: .float32) + } + } + + @Test("quantized/bits8/wq") + func test_quantized_bits8_wq() throws { + try withIntegrationState(seed: 11800) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 8).wq + expectSummary( + result, + ArraySummary( + shape: [64, 32], + dtype: .uint32, + mean: 2115291136.0, + minimum: 169010.0, + maximum: 4294641408.0, + absoluteSum: 4332116246528.0, + positionChecksum: 2162359140352.0, + sampleIndices: [0, 409, 819, 1228, 1638, 2047], + samples: [ + 3760455424.0, 1464697344.0, 4021152000.0, 3835848704.0, 36983752.0, + 3903625984.0, + ]), + tolerance: .exact) + } + } + + @Test("quantized/bits8/scales") + func test_quantized_bits8_scales() throws { + try withIntegrationState(seed: 81192) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 8).scales + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: -0.00018952885875478387, + minimum: -0.007808024063706398, + maximum: 0.007807521149516106, + absoluteSum: 0.9724078178405762, + positionChecksum: 0.490849107503891, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.007772293407469988, 0.007766560185700655, -0.007544763386249542, + -0.007312706205993891, -0.00779489241540432, 0.007582233287394047, + ]), + tolerance: .float32) + } + } + + @Test("quantized/bits8/biases") + func test_quantized_bits8_biases() throws { + try withIntegrationState(seed: 59258) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 64, bits: 8).biases! + expectSummary( + result, + ArraySummary( + shape: [64, 2], + dtype: .float32, + mean: 0.09264960139989853, + minimum: -0.999893069267273, + maximum: 0.9999345541000366, + absoluteSum: 125.88665771484375, + positionChecksum: 63.497344970703125, + sampleIndices: [0, 25, 51, 76, 102, 127], + samples: [ + 0.9969533681869507, 0.9536901712417603, 0.9930459260940552, + -0.9810078740119934, 0.9552137851715088, -0.996850848197937, + ]), + tolerance: .float32) + } + } + + @Test("dequantized/bits8") + func test_dequantized_bits8() throws { + try withIntegrationState(seed: 81104) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let q = MLX.quantized(w, groupSize: 64, bits: 8) + let result = MLX.dequantized( + q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 8) + expectSummary( + result, + ArraySummary( + shape: [64, 128], + dtype: .float32, + mean: 0.004389249254018068, + minimum: -0.9991215467453003, + maximum: 0.9998486042022705, + absoluteSum: 4120.2646484375, + positionChecksum: 2047.842041015625, + sampleIndices: [0, 1638, 3276, 4915, 6553, 8191], + samples: [ + -0.3541724383831024, -0.645488440990448, 0.8101414442062378, + 0.7884412407875061, 0.3327997326850891, 0.5721215605735779, + ]), + tolerance: .float32) + } + } + + @Test("quantizedMM/bits8") + func test_quantizedMM_bits8() throws { + try withIntegrationState(seed: 75568) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w, groupSize: 64, bits: 8) + let result = MLX.quantizedMM( + x, q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 8) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: 0.16322675347328186, + minimum: -19.405149459838867, + maximum: 15.277058601379395, + absoluteSum: 2676.939697265625, + positionChecksum: 1373.523681640625, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + 5.699150085449219, 4.081239223480225, -7.194596290588379, 5.197780609130859, + 1.424774169921875, 0.47141361236572266, + ]), + tolerance: .float32) + } + } + + @Test("quantized/groupSize32") + func test_quantized_groupSize32() throws { + try withIntegrationState(seed: 81977) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let result = MLX.quantized(w, groupSize: 32, bits: 4).wq + expectSummary( + result, + ArraySummary( + shape: [64, 16], + dtype: .uint32, + mean: 2162841344.0, + minimum: 1575843.0, + maximum: 4294780416.0, + absoluteSum: 2214749536256.0, + positionChecksum: 1112428249088.0, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + 1404200832.0, 4060153600.0, 2282566144.0, 787873920.0, 2414542848.0, + 3172955136.0, + ]), + tolerance: .exact) + } + } + + @Test("quantizedMM/noTranspose") + func test_quantizedMM_noTranspose() throws { + try withIntegrationState(seed: 6989) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [128, 64], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w, groupSize: 64, bits: 4) + let result = MLX.quantizedMM( + x, q.wq, scales: q.scales, biases: q.biases, transpose: false, groupSize: 64, + bits: 4) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: 0.04205028712749481, + minimum: -21.537073135375977, + maximum: 19.365976333618164, + absoluteSum: 2747.372314453125, + positionChecksum: 1449.293701171875, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -1.0757472515106201, 7.758847713470459, 5.602324485778809, + -0.715898871421814, -0.12151119112968445, -1.0324546098709106, + ]), + tolerance: .float32) + } + } + + @Test("gatherQuantizedMM") + func test_gatherQuantizedMM() throws { + // without indices this is quantizedMM; the indexed form needs the multi-input generator + try withIntegrationState(seed: 93943) { + let w = MLXRandom.uniform(low: -1.0, high: 1.0, [64, 128], dtype: .float32) + let x = MLXRandom.normal([8, 128], dtype: .float32, loc: 0.0, scale: 1.0) + let q = MLX.quantized(w, groupSize: 64, bits: 4) + let result = MLX.gatherQuantizedMM( + x, q.wq, scales: q.scales, biases: q.biases, groupSize: 64, bits: 4) + expectSummary( + result, + ArraySummary( + shape: [8, 64], + dtype: .float32, + mean: -0.27394434809684753, + minimum: -17.840972900390625, + maximum: 17.40959358215332, + absoluteSum: 2546.3671875, + positionChecksum: 1242.66064453125, + sampleIndices: [0, 102, 204, 307, 409, 511], + samples: [ + -3.621098041534424, -0.9873456954956055, 0.8476672172546387, + 3.559670925140381, -2.1285338401794434, 4.59012508392334, + ]), + tolerance: .float32) + } + } + + @Test("blockMaskedMM") + func test_blockMaskedMM() throws { + try withIntegrationState(seed: 73180) { + let a = MLXRandom.normal([32, 64], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([64, 32], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.blockMaskedMM(a, b, blockSize: 32) + expectSummary( + result, + ArraySummary( + shape: [32, 32], + dtype: .float32, + mean: 0.3000313639640808, + minimum: -27.04250717163086, + maximum: 25.29352569580078, + absoluteSum: 6362.05419921875, + positionChecksum: 3143.78466796875, + sampleIndices: [0, 205, 409, 614, 818, 1023], + samples: [ + -11.411656379699707, 16.519243240356445, -3.2464497089385986, + -3.709771156311035, 6.376284122467041, -3.328993558883667, + ]), + tolerance: .float32) + } + } + + @Test("gatherMM") + func test_gatherMM() throws { + try withIntegrationState(seed: 60084) { + let a = MLXRandom.normal([4, 3, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([6, 5, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let lhsIndices = MLXRandom.randInt(low: 0, high: 4, [2], type: UInt32.self) + let rhsIndices = MLXRandom.randInt(low: 0, high: 6, [2], type: UInt32.self) + let result = MLX.gatherMM(a, b, lhsIndices: lhsIndices, rhsIndices: rhsIndices) + expectSummary( + result, + ArraySummary( + shape: [2, 3, 2], + dtype: .float32, + mean: 1.3322674036026, + minimum: -0.8716369867324829, + maximum: 4.639623641967773, + absoluteSum: 19.419553756713867, + positionChecksum: 11.04852040608724, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 2.538201093673706, 0.6413965225219727, 1.8107318878173828, + -0.8716369867324829, 0.9088269472122192, 4.639623641967773, + ]), + tolerance: .float32) + } + } + + @Test("toFP8") + func test_toFP8() throws { + try withIntegrationState(seed: 19517) { + let a = MLXRandom.uniform(low: -4.0, high: 4.0, [4, 3], dtype: .float32) + let result = MLX.toFP8(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .uint8, + mean: 64.33333587646484, + minimum: 2.0, + maximum: 198.0, + absoluteSum: 772.0, + positionChecksum: 462.4166666666667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [72.0, 10.0, 58.0, 56.0, 50.0, 198.0]), + tolerance: .exact) + } + } + + @Test("fromFP8") + func test_fromFP8() throws { + try withIntegrationState(seed: 87362) { + let a = MLXRandom.uniform(low: -4.0, high: 4.0, [4, 3], dtype: .float32) + let encoded = MLX.toFP8(a) + let result = MLX.fromFP8(encoded, dtype: .float32) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1614583432674408, + minimum: -3.75, + maximum: 4.0, + absoluteSum: 26.3125, + positionChecksum: 14.755208333333334, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [-0.625, 2.75, 3.25, -1.125, -3.75, 4.0]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedRandomInputsTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedRandomInputsTests.swift new file mode 100644 index 000000000..fa4f56c43 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedRandomInputsTests.swift @@ -0,0 +1,596 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 21 + +import Foundation +import MLX +import Testing + +@Suite("generated: RandomInputs") +struct GeneratedRandomInputsTests { + + @Test("normal/2d") + func test_normal_2d() throws { + try withIntegrationState(seed: 75481) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5905870795249939, + minimum: -1.3774175643920898, + maximum: 2.2758142948150635, + absoluteSum: 12.038883209228516, + positionChecksum: 5.29741636912028, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 2.2758142948150635, 1.4398239850997925, 0.3749207854270935, + 1.3119254112243652, -0.48573410511016846, 0.871476948261261, + ]), + tolerance: .float32) + } + } + + @Test("normal/2d/seed2") + func test_normal_2d_seed2() throws { + try withIntegrationState(seed: 89536) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.13405632972717285, + minimum: -3.1497466564178467, + maximum: 1.7531814575195312, + absoluteSum: 12.751688957214355, + positionChecksum: 6.7695051829020185, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.06993529200553894, 0.6983082294464111, -3.1497466564178467, + 1.1703077554702759, 0.8777164220809937, -1.2233566045761108, + ]), + tolerance: .float32) + } + } + + @Test("normal/1d") + func test_normal_1d() throws { + try withIntegrationState(seed: 80954) { + let a = MLXRandom.normal([17], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [17], + dtype: .float32, + mean: 0.21538208425045013, + minimum: -1.4181712865829468, + maximum: 2.0330264568328857, + absoluteSum: 14.34508991241455, + positionChecksum: 7.750569063074448, + sampleIndices: [0, 3, 6, 10, 13, 16], + samples: [ + -0.5277025103569031, 1.0268818140029907, -0.07764638960361481, + 0.72825026512146, -1.4181712865829468, 0.6145756840705872, + ]), + tolerance: .float32) + } + } + + @Test("normal/4d") + func test_normal_4d() throws { + try withIntegrationState(seed: 41791) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [2, 3, 4, 3], + dtype: .float32, + mean: -0.12563791871070862, + minimum: -2.510462522506714, + maximum: 2.0979740619659424, + absoluteSum: 60.029022216796875, + positionChecksum: 28.137715657552082, + sampleIndices: [0, 14, 28, 43, 57, 71], + samples: [ + -0.9431903958320618, -2.1896116733551025, 0.5952138900756836, + -0.40866464376449585, -0.6583279967308044, 0.3055444657802582, + ]), + tolerance: .float32) + } + } + + @Test("normal/scalar") + func test_normal_scalar() throws { + try withIntegrationState(seed: 67446) { + let a = MLXRandom.normal([Int](), dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.35509586334228516, + minimum: 0.35509586334228516, + maximum: 0.35509586334228516, + absoluteSum: 0.35509586334228516, + positionChecksum: 0.35509586334228516, + sampleIndices: [0], + samples: [0.35509586334228516]), + tolerance: .float32) + } + } + + @Test("normal/locScale") + func test_normal_locScale() throws { + try withIntegrationState(seed: 74498) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 2.5, scale: 0.25) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 2.4573264122009277, + minimum: 2.1506776809692383, + maximum: 2.6786084175109863, + absoluteSum: 29.4879150390625, + positionChecksum: 15.964505513509115, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 2.6650614738464355, 2.2944018840789795, 2.3229405879974365, + 2.448261260986328, 2.595139741897583, 2.4941651821136475, + ]), + tolerance: .float32) + } + } + + @Test("normal/float16") + func test_normal_float16() throws { + try withIntegrationState(seed: 44715) { + let a = MLXRandom.normal([4, 3], dtype: .float16, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float16, + mean: 0.323333740234375, + minimum: -1.005859375, + maximum: 1.4765625, + absoluteSum: 8.0045166015625, + positionChecksum: 4.085835774739583, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.005859375, 0.306396484375, -0.6796875, 0.25390625, 0.66552734375, + 0.1114501953125, + ]), + tolerance: .float16) + } + } + + @Test("normal/bfloat16") + func test_normal_bfloat16() throws { + try withIntegrationState(seed: 69686) { + let a = MLXRandom.normal([4, 3], dtype: .bfloat16, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .bfloat16, + mean: 0.1966959685087204, + minimum: -2.875, + maximum: 2.546875, + absoluteSum: 14.2568359375, + positionChecksum: 9.154378255208334, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.64453125, -1.6484375, -0.0439453125, -1.140625, 0.53125, -2.875]), + tolerance: .float16) + } + } + + @Test("uniform/unit") + func test_uniform_unit() throws { + try withIntegrationState(seed: 62798) { + let a = MLXRandom.uniform(low: 0.0, high: 1.0, [4, 3], dtype: .float32) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.4539196491241455, + minimum: 0.06912913173437119, + maximum: 0.986918032169342, + absoluteSum: 5.447035789489746, + positionChecksum: 2.776052474975586, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.6409809589385986, 0.14968015253543854, 0.986918032169342, + 0.7626613974571228, 0.16480962932109833, 0.13554799556732178, + ]), + tolerance: .float32) + } + } + + @Test("uniform/unit/seed2") + func test_uniform_unit_seed2() throws { + try withIntegrationState(seed: 41931) { + let a = MLXRandom.uniform(low: 0.0, high: 1.0, [4, 3], dtype: .float32) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.4274345636367798, + minimum: 0.010180529206991196, + maximum: 0.9993993639945984, + absoluteSum: 5.129214763641357, + positionChecksum: 2.8637450536092124, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.41757380962371826, 0.4386035203933716, 0.2980853021144867, + 0.9993993639945984, 0.8430754542350769, 0.0188215933740139, + ]), + tolerance: .float32) + } + } + + @Test("uniform/range") + func test_uniform_range() throws { + try withIntegrationState(seed: 23538) { + let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.189245581626892, + minimum: 0.2704201340675354, + maximum: 1.9868327379226685, + absoluteSum: 14.270946502685547, + positionChecksum: 7.686832427978516, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.8097621202468872, 1.754885196685791, 0.6641294956207275, + 1.9868327379226685, 1.5307133197784424, 1.0133975744247437, + ]), + tolerance: .float32) + } + } + + @Test("uniform/negative") + func test_uniform_negative() throws { + try withIntegrationState(seed: 40772) { + let a = MLXRandom.uniform(low: -0.9, high: 0.9, [4, 3], dtype: .float32) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.019281506538391113, + minimum: -0.8162036538124084, + maximum: 0.7615352869033813, + absoluteSum: 5.340817451477051, + positionChecksum: 3.0421508153279624, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.31404179334640503, 0.5715023279190063, -0.10399436950683594, + -0.7621757388114929, 0.17554616928100586, 0.7615352869033813, + ]), + tolerance: .float32) + } + } + + @Test("uniform/float16") + func test_uniform_float16() throws { + try withIntegrationState(seed: 82054) { + let a = MLXRandom.uniform(low: 0.0, high: 1.0, [4, 3], dtype: .float16) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float16, + mean: 0.43940991163253784, + minimum: 0.038726806640625, + maximum: 0.875, + absoluteSum: 5.272918701171875, + positionChecksum: 2.8272933959960938, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.23974609375, 0.08013916015625, 0.29248046875, 0.37353515625, + 0.72900390625, 0.177978515625, + ]), + tolerance: .float16) + } + } + + @Test("randInt/small") + func test_randInt_small() throws { + try withIntegrationState(seed: 21039) { + let a = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 1.75, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 21.0, + positionChecksum: 13.083333333333334, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 3.0, 3.0, 3.0, 1.0]), + tolerance: .exact) + } + } + + @Test("randInt/large") + func test_randInt_large() throws { + try withIntegrationState(seed: 4514) { + let a = MLXRandom.randInt(low: -100, high: 1024, [9], type: Int32.self) + expectSummary( + a, + ArraySummary( + shape: [9], + dtype: .int32, + mean: 518.4444580078125, + minimum: -83.0, + maximum: 998.0, + absoluteSum: 4832.0, + positionChecksum: 2413.1111111111113, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [998.0, 463.0, 975.0, 870.0, 361.0, 140.0]), + tolerance: .exact) + } + } + + @Test("randInt/uint32") + func test_randInt_uint32() throws { + try withIntegrationState(seed: 4708) { + let a = MLXRandom.randInt(low: 0, high: 6, [4], type: UInt32.self) + expectSummary( + a, + ArraySummary( + shape: [4], + dtype: .uint32, + mean: 1.5, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 6.0, + positionChecksum: 4.25, + sampleIndices: [0, 1, 2, 3], + samples: [2.0, 0.0, 1.0, 3.0]), + tolerance: .exact) + } + } + + @Test("bernoulli/half") + func test_bernoulli_half() throws { + try withIntegrationState(seed: 17534) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.6666666666666665, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 1.0, 1.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("bernoulli/quarter") + func test_bernoulli_quarter() throws { + try withIntegrationState(seed: 82624) { + let a = MLXRandom.bernoulli(0.25, [4, 3]) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.25, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 1.9166666666666667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]), + tolerance: .exact) + } + } + + @Test("sequence/normal") + func test_sequence_normal() throws { + try withIntegrationState(seed: 91718) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.6608368158340454, + minimum: -3.117097854614258, + maximum: 0.4458395838737488, + absoluteSum: 9.928923606872559, + positionChecksum: 4.727120081583659, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -3.117097854614258, -1.0726861953735352, -0.45238417387008667, + -0.7030599117279053, -0.81221604347229, -0.6846268773078918, + ]), + tolerance: .float32) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + b, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.05415819212794304, + minimum: -2.0514211654663086, + maximum: 2.580822229385376, + absoluteSum: 12.921215057373047, + positionChecksum: 8.349390029907227, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.852959156036377, 0.06281861662864685, 1.4253885746002197, + 0.7632550597190857, 2.580822229385376, -1.264439344406128, + ]), + tolerance: .float32) + let c = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + c, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.2736375033855438, + minimum: -1.6628878116607666, + maximum: 1.6362630128860474, + absoluteSum: 10.95200252532959, + positionChecksum: 5.341580708821614, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.2123521566390991, 1.0446020364761353, 0.8530498147010803, + 0.12905484437942505, -0.4385685622692108, -0.6690441370010376, + ]), + tolerance: .float32) + } + } + + @Test("sequence/mixed") + func test_sequence_mixed() throws { + try withIntegrationState(seed: 1510) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.33849453926086426, + minimum: -2.251725673675537, + maximum: 1.393682599067688, + absoluteSum: 12.362427711486816, + positionChecksum: 6.745829264322917, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.7234569787979126, -0.23261785507202148, -0.3641354739665985, + -2.251725673675537, -0.11936888098716736, 0.1819031834602356, + ]), + tolerance: .float32) + let b = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3], dtype: .float32) + expectSummary( + b, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.373868703842163, + minimum: 0.1270606517791748, + maximum: 1.9089659452438354, + absoluteSum: 16.48642349243164, + positionChecksum: 10.1610959370931, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.1270606517791748, 1.5031565427780151, 1.8709282875061035, + 1.6418914794921875, 1.796919345855713, 1.742146611213684, + ]), + tolerance: .float32) + let c = MLXRandom.randInt(low: 0, high: 4, [4, 3], type: Int32.self) + expectSummary( + c, + ArraySummary( + shape: [4, 3], + dtype: .int32, + mean: 1.4166667461395264, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 17.0, + positionChecksum: 10.0, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 3.0, 1.0, 2.0, 0.0, 3.0]), + tolerance: .exact) + let d = MLXRandom.bernoulli(0.5, [4, 3]) + expectSummary( + d, + ArraySummary( + shape: [4, 3], + dtype: .bool, + mean: 0.4166666865348816, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 5.0, + positionChecksum: 2.8333333333333335, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 0.0, 0.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("sequence/shapes") + func test_sequence_shapes() throws { + try withIntegrationState(seed: 74672) { + let a = MLXRandom.normal([10, 8], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + a, + ArraySummary( + shape: [10, 8], + dtype: .float32, + mean: 0.131831094622612, + minimum: -2.4943573474884033, + maximum: 1.8212450742721558, + absoluteSum: 63.549659729003906, + positionChecksum: 34.66279296875, + sampleIndices: [0, 16, 32, 47, 63, 79], + samples: [ + 0.22473129630088806, -0.6478517651557922, -0.24332140386104584, + -1.8043724298477173, 0.17549262940883636, -2.4943573474884033, + ]), + tolerance: .float32) + let b = MLXRandom.normal([8, 13], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + b, + ArraySummary( + shape: [8, 13], + dtype: .float32, + mean: 0.02560247667133808, + minimum: -2.513389825820923, + maximum: 2.4972832202911377, + absoluteSum: 82.86776733398438, + positionChecksum: 42.1690673828125, + sampleIndices: [0, 21, 41, 62, 82, 103], + samples: [ + -1.0536662340164185, -0.4492167830467224, -0.45552271604537964, + 1.7218302488327026, 1.0314418077468872, -1.0060495138168335, + ]), + tolerance: .float32) + let c = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + expectSummary( + c, + ArraySummary( + shape: [2, 3, 4], + dtype: .float32, + mean: -0.17977163195610046, + minimum: -2.1054158210754395, + maximum: 1.6299241781234741, + absoluteSum: 17.145761489868164, + positionChecksum: 8.302991231282553, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -1.1952673196792603, -1.2896186113357544, 0.006688464432954788, + -0.3580016493797302, -1.3585586547851562, 0.4489697217941284, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedRandomTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedRandomTests.swift new file mode 100644 index 000000000..c12688c85 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedRandomTests.swift @@ -0,0 +1,269 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 11 + +import Foundation +import MLX +import Testing + +@Suite("generated: Random") +struct GeneratedRandomTests { + + @Test("gumbel") + func test_gumbel() throws { + try withIntegrationState(seed: 65021) { + let result = MLXRandom.gumbel([4, 3]) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.846875786781311, + minimum: -0.9293802976608276, + maximum: 3.7100210189819336, + absoluteSum: 15.268047332763672, + positionChecksum: 6.564333597819011, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.0602260828018188, 0.20688922703266144, 3.7100210189819336, + 1.311169147491455, 0.7526293396949768, -0.8609374761581421, + ]), + tolerance: .float32) + } + } + + @Test("gumbel/dtype") + func test_gumbel_dtype() throws { + try withIntegrationState(seed: 90296) { + let result = MLXRandom.gumbel([4, 3], dtype: .float16) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float16, + mean: 0.8638967275619507, + minimum: -1.18359375, + maximum: 5.22265625, + absoluteSum: 17.61090087890625, + positionChecksum: 8.492075602213541, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 5.22265625, 0.91845703125, -0.77490234375, 3.275390625, 1.8720703125, + -1.18359375, + ]), + tolerance: .float16) + } + } + + @Test("laplace") + func test_laplace() throws { + // the Swift dtype: parameter has no default + try withIntegrationState(seed: 62794) { + let result = MLXRandom.laplace([4, 3], dtype: .float32) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.37702733278274536, + minimum: -2.6754982471466064, + maximum: 2.45698618888855, + absoluteSum: 12.961724281311035, + positionChecksum: 8.14331309000651, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.23149043321609497, -0.45272037386894226, 0.8567035794258118, + -2.6754982471466064, 1.9753601551055908, 0.19610698521137238, + ]), + tolerance: .float32) + } + } + + @Test("laplace/locScale") + func test_laplace_locScale() throws { + try withIntegrationState(seed: 37443) { + let result = MLXRandom.laplace([4, 3], dtype: .float32, loc: 1.0, scale: 2.0) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.0226888470351696, + minimum: -4.823460578918457, + maximum: 4.867783546447754, + absoluteSum: 30.37203025817871, + positionChecksum: 17.33179982503255, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.4936223030090332, 0.259247362613678, 4.175117492675781, + -4.638853549957275, -1.1265568733215332, 4.867783546447754, + ]), + tolerance: .float32) + } + } + + @Test("truncatedNormal") + func test_truncatedNormal() throws { + try withIntegrationState(seed: 37638) { + let result = MLXRandom.truncatedNormal(low: -1.0, high: 1.0, [4, 3]) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.05009012669324875, + minimum: -0.8407586216926575, + maximum: 0.7136946320533752, + absoluteSum: 6.225048542022705, + positionChecksum: 3.2496751149495444, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.3011457920074463, 0.6435555815696716, -0.3013858199119568, + -0.6681084036827087, 0.2138742208480835, -0.4170948266983032, + ]), + tolerance: .float32) + } + } + + @Test("categorical") + func test_categorical() throws { + try withIntegrationState(seed: 30654) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXRandom.categorical(logits) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .uint32, + mean: 2.75, + minimum: 0.0, + maximum: 4.0, + absoluteSum: 11.0, + positionChecksum: 8.25, + sampleIndices: [0, 1, 2, 3], + samples: [0.0, 4.0, 3.0, 4.0]), + tolerance: .exact) + } + } + + @Test("categorical/count") + func test_categorical_count() throws { + try withIntegrationState(seed: 71811) { + let logits = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXRandom.categorical(logits, count: 3) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .uint32, + mean: 1.75, + minimum: 0.0, + maximum: 3.0, + absoluteSum: 21.0, + positionChecksum: 12.333333333333334, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 2.0, 0.0, 3.0, 0.0, 3.0]), + tolerance: .exact) + } + } + + @Test("permutation/count") + func test_permutation_count() throws { + try withIntegrationState(seed: 53709) { + let result = MLXRandom.permutation(10) + expectSummary( + result, + ArraySummary( + shape: [10], + dtype: .uint32, + mean: 4.5, + minimum: 0.0, + maximum: 9.0, + absoluteSum: 45.0, + positionChecksum: 24.7, + sampleIndices: [0, 2, 4, 5, 7, 9], + samples: [9.0, 3.0, 2.0, 6.0, 0.0, 7.0]), + tolerance: .exact) + } + } + + @Test("permutation/array") + func test_permutation_array() throws { + try withIntegrationState(seed: 9976) { + let a = MLXRandom.normal([6, 2], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLXRandom.permutation(a) + expectSummary( + result, + ArraySummary( + shape: [6, 2], + dtype: .float32, + mean: 0.30458083748817444, + minimum: -1.2631405591964722, + maximum: 3.427476167678833, + absoluteSum: 9.221382141113281, + positionChecksum: 5.262217839558919, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.9780406951904297, -1.2631405591964722, -0.6314418315887451, + 0.45676153898239136, 1.034820795059204, -0.07250352948904037, + ]), + tolerance: .float32) + } + } + + @Test("multivariateNormal") + func test_multivariateNormal() throws { + try withIntegrationState(seed: 597) { + let mean = MLX.zeros([3]) + let covariance = MLX.eye(3) * 2.0 + let result = MLXRandom.multivariateNormal( + mean: mean, covariance: covariance, shape: [4], dtype: .float32, stream: .cpu) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.37988409399986267, + minimum: -3.102618932723999, + maximum: 2.8767926692962646, + absoluteSum: 18.955251693725586, + positionChecksum: 9.738911946614584, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.462942361831665, 0.007976667955517769, 2.2923548221588135, + -3.102618932723999, -0.5047874450683594, 1.6479119062423706, + ]), + tolerance: .float32) + } + } + + @Test("key") + func test_key() throws { + try withIntegrationState(seed: 20155) { + let result = MLXRandom.key(42) + expectSummary( + result, + ArraySummary( + shape: [2], + dtype: .uint32, + mean: 21.0, + minimum: 0.0, + maximum: 42.0, + absoluteSum: 42.0, + positionChecksum: 42.0, + sampleIndices: [0, 1], + samples: [0.0, 42.0]), + tolerance: .exact) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedReductionTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedReductionTests.swift new file mode 100644 index 000000000..ea27372e7 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedReductionTests.swift @@ -0,0 +1,1506 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 64 + +import Foundation +import MLX +import Testing + +@Suite("generated: Reduction") +struct GeneratedReductionTests { + + @Test("sum") + func test_sum() throws { + try withIntegrationState(seed: 22575) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sum(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.7057861685752869, + minimum: 0.7057861685752869, + maximum: 0.7057861685752869, + absoluteSum: 0.7057861685752869, + positionChecksum: 0.7057861685752869, + sampleIndices: [0], + samples: [0.7057861685752869]), + tolerance: .float32) + } + } + + @Test("sum/axis") + func test_sum_axis() throws { + try withIntegrationState(seed: 37581) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sum(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.5632065534591675, + minimum: -2.874988555908203, + maximum: 0.6334655284881592, + absoluteSum: 4.267636775970459, + positionChecksum: 2.3913214206695557, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.37393978238105774, -2.874988555908203, 0.6334655284881592, + -0.38524290919303894, + ]), + tolerance: .float32) + } + } + + @Test("sum/axes") + func test_sum_axes() throws { + try withIntegrationState(seed: 15457) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sum(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.44520968198776245, + minimum: -6.6333723068237305, + maximum: 7.999377727508545, + absoluteSum: 31.566116333007812, + positionChecksum: 17.56298573811849, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.32674166560173035, 7.999377727508545, -0.22247016429901123, + -6.6333723068237305, 1.6439025402069092, -2.4121196269989014, + ]), + tolerance: .float32) + } + } + + @Test("sum/keepDims") + func test_sum_keepDims() throws { + try withIntegrationState(seed: 30341) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sum(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: 0.3402286171913147, + minimum: -5.441669464111328, + maximum: 6.115833282470703, + absoluteSum: 33.430877685546875, + positionChecksum: 17.420809427897137, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -3.12631893157959, 1.0557702779769897, -2.0913119316101074, + 2.757711410522461, -5.441669464111328, -0.5352753400802612, + ]), + tolerance: .float32) + } + } + + @Test("sum/method") + func test_sum_method() throws { + try withIntegrationState(seed: 92671) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.sum(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.891804575920105, + minimum: -1.5811764001846313, + maximum: -0.49883174896240234, + absoluteSum: 3.56721830368042, + positionChecksum: 2.2104756832122803, + sampleIndices: [0, 1, 2, 3], + samples: [ + -0.8713735938072205, -0.6158364415168762, -1.5811764001846313, + -0.49883174896240234, + ]), + tolerance: .float32) + } + } + + @Test("mean") + func test_mean() throws { + try withIntegrationState(seed: 78073) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.mean(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.23587274551391602, + minimum: -0.23587274551391602, + maximum: -0.23587274551391602, + absoluteSum: 0.23587274551391602, + positionChecksum: 0.23587274551391602, + sampleIndices: [0], + samples: [-0.23587274551391602]), + tolerance: .float32) + } + } + + @Test("mean/axis") + func test_mean_axis() throws { + try withIntegrationState(seed: 14185) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.mean(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.2951846718788147, + minimum: -1.14504075050354, + maximum: 0.3134152293205261, + absoluteSum: 2.174551010131836, + positionChecksum: 1.1131901741027832, + sampleIndices: [0, 1, 2, 3], + samples: [ + -1.14504075050354, 0.3134152293205261, 0.18349099159240723, + -0.5326040983200073, + ]), + tolerance: .float32) + } + } + + @Test("mean/axes") + func test_mean_axes() throws { + try withIntegrationState(seed: 90885) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.mean(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.09846246242523193, + minimum: -0.46042436361312866, + maximum: 0.7632273435592651, + absoluteSum: 3.7009592056274414, + positionChecksum: 2.2039335568745932, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.17073218524456024, -0.06515144556760788, 0.4433436095714569, + -0.46042436361312866, -0.4089309573173523, 0.05317772924900055, + ]), + tolerance: .float32) + } + } + + @Test("mean/keepDims") + func test_mean_keepDims() throws { + try withIntegrationState(seed: 64348) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.mean(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: -0.10993402451276779, + minimum: -0.8585461378097534, + maximum: 0.3680957555770874, + absoluteSum: 3.5669546127319336, + positionChecksum: 2.2769004503885903, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.138699010014534, 0.04041095823049545, -0.17920581996440887, + 0.15857860445976257, -0.15389952063560486, -0.5889500975608826, + ]), + tolerance: .float32) + } + } + + @Test("mean/method") + func test_mean_method() throws { + try withIntegrationState(seed: 64427) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.mean(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.2153686285018921, + minimum: -0.5716754794120789, + maximum: 0.36203593015670776, + absoluteSum: 1.5855462551116943, + positionChecksum: 1.1746270656585693, + sampleIndices: [0, 1, 2, 3], + samples: [ + -0.13388535380363464, 0.36203593015670776, -0.5179495811462402, + -0.5716754794120789, + ]), + tolerance: .float32) + } + } + + @Test("min") + func test_min() throws { + try withIntegrationState(seed: 27506) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.min(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -1.6234827041625977, + minimum: -1.6234827041625977, + maximum: -1.6234827041625977, + absoluteSum: 1.6234827041625977, + positionChecksum: 1.6234827041625977, + sampleIndices: [0], + samples: [-1.6234827041625977]), + tolerance: .float32) + } + } + + @Test("min/axis") + func test_min_axis() throws { + try withIntegrationState(seed: 17376) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.min(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -1.527522087097168, + minimum: -3.202768564224243, + maximum: -0.5903887152671814, + absoluteSum: 6.110088348388672, + positionChecksum: 3.523585319519043, + sampleIndices: [0, 1, 2, 3], + samples: [ + -1.1166950464248657, -3.202768564224243, -0.5903887152671814, + -1.2002359628677368, + ]), + tolerance: .float32) + } + } + + @Test("min/axes") + func test_min_axes() throws { + try withIntegrationState(seed: 1356) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.min(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: -1.2078721523284912, + minimum: -2.452378988265991, + maximum: -0.4843156933784485, + absoluteSum: 14.494464874267578, + positionChecksum: 8.333949406941732, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.393852710723877, -0.9224222302436829, -1.487610936164856, + -1.094190001487732, -0.6395735144615173, -2.452378988265991, + ]), + tolerance: .float32) + } + } + + @Test("min/keepDims") + func test_min_keepDims() throws { + try withIntegrationState(seed: 78536) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.min(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: -1.075392484664917, + minimum: -2.119365692138672, + maximum: -0.35634124279022217, + absoluteSum: 12.904708862304688, + positionChecksum: 7.181203842163086, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.1877634525299072, -1.355299949645996, -0.35634124279022217, + -0.6733994483947754, -0.7651041150093079, -1.3274004459381104, + ]), + tolerance: .float32) + } + } + + @Test("min/method") + func test_min_method() throws { + try withIntegrationState(seed: 56994) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.min(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.8903401494026184, + minimum: -2.030451536178589, + maximum: 0.40401411056518555, + absoluteSum: 4.369388580322266, + positionChecksum: 2.9881086349487305, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.40401411056518555, -2.030451536178589, -0.25217559933662415, + -1.6827476024627686, + ]), + tolerance: .float32) + } + } + + @Test("max") + func test_max() throws { + try withIntegrationState(seed: 63179) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.max(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.1173732280731201, + minimum: 1.1173732280731201, + maximum: 1.1173732280731201, + absoluteSum: 1.1173732280731201, + positionChecksum: 1.1173732280731201, + sampleIndices: [0], + samples: [1.1173732280731201]), + tolerance: .float32) + } + } + + @Test("max/axis") + func test_max_axis() throws { + try withIntegrationState(seed: 6753) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.max(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.34814202785491943, + minimum: -0.42853012681007385, + maximum: 1.4546650648117065, + absoluteSum: 2.6034679412841797, + positionChecksum: 1.8354352712631226, + sampleIndices: [0, 1, 2, 3], + samples: [ + -0.17691978812217712, 0.5433529615402222, 1.4546650648117065, + -0.42853012681007385, + ]), + tolerance: .float32) + } + } + + @Test("max/axes") + func test_max_axes() throws { + try withIntegrationState(seed: 90061) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.max(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 1.2743048667907715, + minimum: 0.052372973412275314, + maximum: 2.9989659786224365, + absoluteSum: 15.291658401489258, + positionChecksum: 8.274401346842447, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.4796899557113647, 0.4484626054763794, 2.9989659786224365, + 0.052372973412275314, 1.8464957475662231, 0.9091650247573853, + ]), + tolerance: .float32) + } + } + + @Test("max/keepDims") + func test_max_keepDims() throws { + try withIntegrationState(seed: 80866) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.max(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: 1.187777042388916, + minimum: 0.4238319396972656, + maximum: 2.3939080238342285, + absoluteSum: 14.253324508666992, + positionChecksum: 7.477830251057942, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.456592321395874, 0.9078659415245056, 0.4238319396972656, + 0.6892192959785461, 0.9878159761428833, 0.4668642282485962, + ]), + tolerance: .float32) + } + } + + @Test("max/method") + func test_max_method() throws { + try withIntegrationState(seed: 52822) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.max(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.5650159120559692, + minimum: 0.33188390731811523, + maximum: 1.0610710382461548, + absoluteSum: 2.260063648223877, + positionChecksum: 1.1534028053283691, + sampleIndices: [0, 1, 2, 3], + samples: [ + 1.0610710382461548, 0.37632203102111816, 0.49078676104545593, + 0.33188390731811523, + ]), + tolerance: .float32) + } + } + + @Test("product") + func test_product() throws { + try withIntegrationState(seed: 1646) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.product(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.09494685381650925, + minimum: -0.09494685381650925, + maximum: -0.09494685381650925, + absoluteSum: 0.09494685381650925, + positionChecksum: 0.09494685381650925, + sampleIndices: [0], + samples: [-0.09494685381650925]), + tolerance: .float32) + } + } + + @Test("product/axis") + func test_product_axis() throws { + try withIntegrationState(seed: 51321) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.product(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.2498457133769989, + minimum: -0.070404052734375, + maximum: 0.6425824761390686, + absoluteSum: 1.1401909589767456, + positionChecksum: 0.6402937173843384, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.6425824761390686, 0.0007186831790022552, -0.070404052734375, + 0.42648571729660034, + ]), + tolerance: .float32) + } + } + + @Test("product/axes") + func test_product_axes() throws { + try withIntegrationState(seed: 89717) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.product(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: -0.003022553399205208, + minimum: -0.1403750628232956, + maximum: 0.047358572483062744, + absoluteSum: 0.3018551468849182, + positionChecksum: 0.08646868666013081, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.1403750628232956, -0.023114101961255074, 0.047358572483062744, + -0.0006650111172348261, 0.007170831318944693, 0.02495034784078598, + ]), + tolerance: .float32) + } + } + + @Test("product/keepDims") + func test_product_keepDims() throws { + try withIntegrationState(seed: 24620) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.product(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: -0.01248251087963581, + minimum: -0.30200156569480896, + maximum: 0.234603613615036, + absoluteSum: 0.7817379832267761, + positionChecksum: 0.4315507411956787, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.0029180373530834913, 0.0015591266565024853, 0.01657288521528244, + 0.031356602907180786, -0.10858920961618423, 0.029305333271622658, + ]), + tolerance: .float32) + } + } + + @Test("product/method") + func test_product_method() throws { + try withIntegrationState(seed: 43022) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.product(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.7211005687713623, + minimum: -2.585625171661377, + maximum: 0.2048765867948532, + absoluteSum: 3.434208393096924, + positionChecksum: 1.9454559087753296, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.07002640515565872, -2.585625171661377, -0.5736801624298096, + 0.2048765867948532, + ]), + tolerance: .float32) + } + } + + @Test("logSumExp") + func test_logSumExp() throws { + try withIntegrationState(seed: 10288) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logSumExp(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 3.6624972820281982, + minimum: 3.6624972820281982, + maximum: 3.6624972820281982, + absoluteSum: 3.6624972820281982, + positionChecksum: 3.6624972820281982, + sampleIndices: [0], + samples: [3.6624972820281982]), + tolerance: .float32) + } + } + + @Test("logSumExp/axis") + func test_logSumExp_axis() throws { + try withIntegrationState(seed: 78003) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logSumExp(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.3402810096740723, + minimum: 0.5028432607650757, + maximum: 1.6911296844482422, + absoluteSum: 5.361124038696289, + positionChecksum: 3.1376407146453857, + sampleIndices: [0, 1, 2, 3], + samples: [ + 1.6911296844482422, 1.6588503122329712, 0.5028432607650757, 1.50830078125, + ]), + tolerance: .float32) + } + } + + @Test("logSumExp/axes") + func test_logSumExp_axes() throws { + try withIntegrationState(seed: 33247) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logSumExp(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 2.3641576766967773, + minimum: 1.0966796875, + maximum: 3.5836031436920166, + absoluteSum: 28.369890213012695, + positionChecksum: 15.860939025878906, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 2.2214109897613525, 2.3977997303009033, 2.9473743438720703, + 2.1123032569885254, 2.5586776733398438, 2.292973041534424, + ]), + tolerance: .float32) + } + } + + @Test("logSumExp/keepDims") + func test_logSumExp_keepDims() throws { + try withIntegrationState(seed: 5045) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logSumExp(a, axes: [0, -1], keepDims: true) + expectSummary( + result, + ArraySummary( + shape: [1, 3, 4, 1], + dtype: .float32, + mean: 1.9824609756469727, + minimum: 0.9807406663894653, + maximum: 3.0540881156921387, + absoluteSum: 23.789531707763672, + positionChecksum: 12.198070526123047, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.818997859954834, 2.2178637981414795, 1.7742036581039429, + 1.2414398193359375, 0.9807406663894653, 1.9903795719146729, + ]), + tolerance: .float32) + } + } + + @Test("logSumExp/method") + func test_logSumExp_method() throws { + try withIntegrationState(seed: 98817) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = a.logSumExp(axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.328906536102295, + minimum: 0.47593313455581665, + maximum: 1.9842123985290527, + absoluteSum: 5.31562614440918, + positionChecksum: 3.2436769008636475, + sampleIndices: [0, 1, 2, 3], + samples: [ + 1.28114652633667, 1.9842123985290527, 0.47593313455581665, + 1.5743341445922852, + ]), + tolerance: .float32) + } + } + + @Test("variance") + func test_variance() throws { + try withIntegrationState(seed: 24490) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.variance(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 0.3176921606063843, + minimum: 0.3176921606063843, + maximum: 0.3176921606063843, + absoluteSum: 0.3176921606063843, + positionChecksum: 0.3176921606063843, + sampleIndices: [0], + samples: [0.3176921606063843]), + tolerance: .float32) + } + } + + @Test("variance/axis") + func test_variance_axis() throws { + try withIntegrationState(seed: 94775) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.variance(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.5114014148712158, + minimum: 0.12645000219345093, + maximum: 1.533726453781128, + absoluteSum: 2.0456056594848633, + positionChecksum: 1.4190025329589844, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.2399284839630127, 0.12645000219345093, 1.533726453781128, + 0.14550063014030457, + ]), + tolerance: .float32) + } + } + + @Test("variance/ddof") + func test_variance_ddof() throws { + try withIntegrationState(seed: 73532) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.variance(a, axis: -1, ddof: 1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 1.3407645225524902, + minimum: 0.09674856066703796, + maximum: 2.520488739013672, + absoluteSum: 5.363058090209961, + positionChecksum: 3.5370237827301025, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.09674856066703796, 2.520488739013672, 1.9729126691818237, + 0.7729080319404602, + ]), + tolerance: .float32) + } + } + + @Test("std") + func test_std() throws { + try withIntegrationState(seed: 24170) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.std(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 1.0559873580932617, + minimum: 1.0559873580932617, + maximum: 1.0559873580932617, + absoluteSum: 1.0559873580932617, + positionChecksum: 1.0559873580932617, + sampleIndices: [0], + samples: [1.0559873580932617]), + tolerance: .float32) + } + } + + @Test("std/axis") + func test_std_axis() throws { + try withIntegrationState(seed: 15089) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.std(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.6319688558578491, + minimum: 0.5627264380455017, + maximum: 0.6639640927314758, + absoluteSum: 2.5278754234313965, + positionChecksum: 1.5501139163970947, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.6447125673294067, 0.6564723253250122, 0.6639640927314758, + 0.5627264380455017, + ]), + tolerance: .float32) + } + } + + @Test("std/ddof") + func test_std_ddof() throws { + try withIntegrationState(seed: 23642) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.std(a, axis: -1, ddof: 1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.6409159898757935, + minimum: 0.15130622684955597, + maximum: 1.0482271909713745, + absoluteSum: 2.563663959503174, + positionChecksum: 1.684018611907959, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.6390142440795898, 0.7251163721084595, 0.15130622684955597, + 1.0482271909713745, + ]), + tolerance: .float32) + } + } + + @Test("median") + func test_median() throws { + try withIntegrationState(seed: 76670) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.median(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -0.48609113693237305, + minimum: -0.48609113693237305, + maximum: -0.48609113693237305, + absoluteSum: 0.48609113693237305, + positionChecksum: 0.48609113693237305, + sampleIndices: [0], + samples: [-0.48609113693237305]), + tolerance: .float32) + } + } + + @Test("median/axis") + func test_median_axis() throws { + try withIntegrationState(seed: 54563) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.median(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: 0.21008887887001038, + minimum: -0.07639828324317932, + maximum: 0.35667484998703003, + absoluteSum: 0.9931520223617554, + positionChecksum: 0.6930854320526123, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.2302650660276413, -0.07639828324317932, 0.35667484998703003, + 0.3298138678073883, + ]), + tolerance: .float32) + } + } + + @Test("all") + func test_all() throws { + try withIntegrationState(seed: 78751) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.all(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0], + samples: [0.0]), + tolerance: .exact) + } + } + + @Test("all/axis") + func test_all_axis() throws { + try withIntegrationState(seed: 95363) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.all(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .bool, + mean: 0.0, + minimum: 0.0, + maximum: 0.0, + absoluteSum: 0.0, + positionChecksum: 0.0, + sampleIndices: [0, 1, 2, 3], + samples: [0.0, 0.0, 0.0, 0.0]), + tolerance: .exact) + } + } + + @Test("any") + func test_any() throws { + try withIntegrationState(seed: 12694) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.any(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .bool, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0], + samples: [1.0]), + tolerance: .exact) + } + } + + @Test("any/axis") + func test_any_axis() throws { + try withIntegrationState(seed: 52226) { + let a = MLXRandom.bernoulli(0.5, [4, 3]) + let result = MLX.any(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .bool, + mean: 0.75, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 3.0, + positionChecksum: 2.0, + sampleIndices: [0, 1, 2, 3], + samples: [1.0, 0.0, 1.0, 1.0]), + tolerance: .exact) + } + } + + @Test("argMin") + func test_argMin() throws { + try withIntegrationState(seed: 22413) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argMin(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .uint32, + mean: 1.0, + minimum: 1.0, + maximum: 1.0, + absoluteSum: 1.0, + positionChecksum: 1.0, + sampleIndices: [0], + samples: [1.0]), + tolerance: .exact) + } + } + + @Test("argMin/axis") + func test_argMin_axis() throws { + try withIntegrationState(seed: 96731) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argMin(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .uint32, + mean: 1.25, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 5.0, + positionChecksum: 2.75, + sampleIndices: [0, 1, 2, 3], + samples: [1.0, 2.0, 2.0, 0.0]), + tolerance: .exact) + } + } + + @Test("argMax") + func test_argMax() throws { + try withIntegrationState(seed: 30164) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argMax(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .uint32, + mean: 8.0, + minimum: 8.0, + maximum: 8.0, + absoluteSum: 8.0, + positionChecksum: 8.0, + sampleIndices: [0], + samples: [8.0]), + tolerance: .exact) + } + } + + @Test("argMax/axis") + func test_argMax_axis() throws { + try withIntegrationState(seed: 36698) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argMax(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .uint32, + mean: 1.75, + minimum: 1.0, + maximum: 2.0, + absoluteSum: 7.0, + positionChecksum: 4.5, + sampleIndices: [0, 1, 2, 3], + samples: [2.0, 1.0, 2.0, 2.0]), + tolerance: .exact) + } + } + + @Test("cumsum/axis") + func test_cumsum_axis() throws { + try withIntegrationState(seed: 2662) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumsum(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.10426246374845505, + minimum: -2.4091546535491943, + maximum: 2.405540943145752, + absoluteSum: 10.7410888671875, + positionChecksum: 8.23110834757487, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.11669416725635529, -0.1358266919851303, 0.07263723015785217, + 1.0716676712036133, -1.257238745689392, -2.4091546535491943, + ]), + tolerance: .float32) + } + } + + @Test("cumsum/reverse") + func test_cumsum_reverse() throws { + try withIntegrationState(seed: 60267) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumsum(a, axis: -1, reverse: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -1.1823921203613281, + minimum: -3.4265153408050537, + maximum: 1.0207629203796387, + absoluteSum: 17.93625259399414, + positionChecksum: 7.194801330566406, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -3.4265153408050537, -1.890081763267517, -2.832381248474121, + -0.3406890034675598, -1.968650460243225, 0.3049541711807251, + ]), + tolerance: .float32) + } + } + + @Test("cumsum/exclusive") + func test_cumsum_exclusive() throws { + try withIntegrationState(seed: 81804) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumsum(a, axis: -1, inclusive: false) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.3773741126060486, + minimum: -3.1300058364868164, + maximum: 1.900256633758545, + absoluteSum: 11.56186294555664, + positionChecksum: 7.195838928222656, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, -1.3729093074798584, 1.349768877029419, -2.0088613033294678, 0.0, + 1.900256633758545, + ]), + tolerance: .float32) + } + } + + @Test("cumprod/axis") + func test_cumprod_axis() throws { + try withIntegrationState(seed: 92128) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumprod(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.6749899387359619, + minimum: -3.232989549636841, + maximum: 1.2675939798355103, + absoluteSum: 13.808900833129883, + positionChecksum: 6.150314966837565, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -2.8824124336242676, 0.42125698924064636, 0.020262163132429123, + -2.222736358642578, 0.544766902923584, 1.2675939798355103, + ]), + tolerance: .float32) + } + } + + @Test("cumprod/reverse") + func test_cumprod_reverse() throws { + try withIntegrationState(seed: 88787) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumprod(a, axis: -1, reverse: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.013112984597682953, + minimum: -1.2184795141220093, + maximum: 1.3128567934036255, + absoluteSum: 7.380435943603516, + positionChecksum: 4.74104372660319, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.35103660821914673, 0.676842451095581, 0.7622838020324707, + 0.04095830023288727, -1.0784687995910645, -1.2184795141220093, + ]), + tolerance: .float32) + } + } + + @Test("cumprod/exclusive") + func test_cumprod_exclusive() throws { + try withIntegrationState(seed: 35949) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cumprod(a, axis: -1, inclusive: false) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.18877795338630676, + minimum: -1.8461512327194214, + maximum: 1.2107491493225098, + absoluteSum: 10.803426742553711, + positionChecksum: 4.96627680460612, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.0, -1.8461512327194214, -0.90610271692276, 0.6101259589195251, 1.0, + 0.3891559839248657, + ]), + tolerance: .float32) + } + } + + @Test("cummax/axis") + func test_cummax_axis() throws { + try withIntegrationState(seed: 48106) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummax(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3044331669807434, + minimum: -1.4819244146347046, + maximum: 1.6274610757827759, + absoluteSum: 8.384651184082031, + positionChecksum: 5.544460296630859, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7419719696044922, 0.7419719696044922, -0.04854029417037964, + -0.3050103485584259, -1.4819244146347046, 1.6274610757827759, + ]), + tolerance: .float32) + } + } + + @Test("cummax/reverse") + func test_cummax_reverse() throws { + try withIntegrationState(seed: 62634) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummax(a, axis: -1, reverse: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.5222153663635254, + minimum: -1.3533363342285156, + maximum: 1.8465720415115356, + absoluteSum: 10.072877883911133, + positionChecksum: 5.809856414794922, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.820504367351532, 0.17221690714359283, -0.27490541338920593, + 1.8465720415115356, 0.864747166633606, 0.24279813468456268, + ]), + tolerance: .float32) + } + } + + @Test("cummax/exclusive") + func test_cummax_exclusive() throws { + try withIntegrationState(seed: 25900) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummax(a, axis: -1, inclusive: false) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -Double.infinity, + minimum: -Double.infinity, + maximum: 0.9814871549606323, + absoluteSum: Double.infinity, + positionChecksum: Double.infinity, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -Double.infinity, 0.9814871549606323, 0.16378529369831085, + -0.5856958031654358, -Double.infinity, -1.9102404117584229, + ]), + tolerance: .float32) + } + } + + @Test("cummin/axis") + func test_cummin_axis() throws { + try withIntegrationState(seed: 94795) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummin(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.7058196067810059, + minimum: -1.1635726690292358, + maximum: 0.8255786299705505, + absoluteSum: 10.120992660522461, + positionChecksum: 5.628684997558594, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.8448988199234009, -0.8448988199234009, -1.0783499479293823, + -0.9089512825012207, 0.8255786299705505, -0.8575506806373596, + ]), + tolerance: .float32) + } + } + + @Test("cummin/reverse") + func test_cummin_reverse() throws { + try withIntegrationState(seed: 80709) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummin(a, axis: -1, reverse: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.37917232513427734, + minimum: -2.6307315826416016, + maximum: 2.1301822662353516, + absoluteSum: 12.440810203552246, + positionChecksum: 6.774248123168945, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.9820541143417358, 0.7694716453552246, -0.08122467249631882, + -2.6307315826416016, 0.2754596769809723, 0.44595789909362793, + ]), + tolerance: .float32) + } + } + + @Test("cummin/exclusive") + func test_cummin_exclusive() throws { + try withIntegrationState(seed: 21818) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.cummin(a, axis: -1, inclusive: false) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: Double.infinity, + minimum: -0.9863429665565491, + maximum: Double.infinity, + absoluteSum: Double.infinity, + positionChecksum: Double.infinity, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + Double.infinity, -0.716964066028595, 0.8032096028327942, + -0.37494421005249023, Double.infinity, -0.9684954285621643, + ]), + tolerance: .float32) + } + } + + @Test("logCumsumExp/axis") + func test_logCumsumExp_axis() throws { + try withIntegrationState(seed: 68907) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logCumsumExp(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 1.124143123626709, + minimum: -0.18204832077026367, + maximum: 2.4726126194000244, + absoluteSum: 14.127202987670898, + positionChecksum: 9.762095133463541, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.18204832077026367, 0.888251006603241, 0.21538907289505005, + 2.3542518615722656, 1.4005087614059448, 1.8368728160858154, + ]), + tolerance: .float32) + } + } + + @Test("logCumsumExp/reverse") + func test_logCumsumExp_reverse() throws { + try withIntegrationState(seed: 2059) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logCumsumExp(a, axis: -1, reverse: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.058185141533613205, + minimum: -1.9204189777374268, + maximum: 1.5514507293701172, + absoluteSum: 13.139307022094727, + positionChecksum: 6.322711944580078, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.4820340871810913, -1.6690078973770142, 1.4352582693099976, + -1.3616853952407837, 0.6133323311805725, -1.0922249555587769, + ]), + tolerance: .float32) + } + } + + @Test("logCumsumExp/exclusive") + func test_logCumsumExp_exclusive() throws { + try withIntegrationState(seed: 27959) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.logCumsumExp(a, axis: -1, inclusive: false) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -Double.infinity, + minimum: -Double.infinity, + maximum: 1.805098056793213, + absoluteSum: Double.infinity, + positionChecksum: Double.infinity, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -Double.infinity, 0.3031471371650696, -0.4027698040008545, + 1.8003222942352295, -Double.infinity, 0.24152666330337524, + ]), + tolerance: .float32) + } + } + + @Test("softmax/axis") + func test_softmax_axis() throws { + try withIntegrationState(seed: 96786) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.softmax(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3333333432674408, + minimum: 0.02120232954621315, + maximum: 0.927946150302887, + absoluteSum: 4.0, + positionChecksum: 2.1860249837239585, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.060398709028959274, 0.887180745601654, 0.02120232954621315, + 0.5514998435974121, 0.10449923574924469, 0.4300844967365265, + ]), + tolerance: .float32) + } + } + + @Test("softmax/axes") + func test_softmax_axes() throws { + try withIntegrationState(seed: 10846) { + let a = MLXRandom.normal([2, 3, 4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.softmax(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [2, 3, 4, 3], + dtype: .float32, + mean: 0.1666666716337204, + minimum: 0.02015802077949047, + maximum: 0.5400499105453491, + absoluteSum: 12.0, + positionChecksum: 6.345125834147136, + sampleIndices: [0, 14, 28, 43, 57, 71], + samples: [ + 0.22917267680168152, 0.04581531509757042, 0.4089527130126953, + 0.3178185224533081, 0.030428849160671234, 0.3574194312095642, + ]), + tolerance: .float32) + } + } + + @Test("softmax/precise") + func test_softmax_precise() throws { + // precise: computes in float32; the python side is already float32 + try withIntegrationState(seed: 68889) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.softmax(a, axis: -1, precise: true) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.3333333432674408, + minimum: 0.016829494386911392, + maximum: 0.6464029550552368, + absoluteSum: 4.0, + positionChecksum: 2.197746435801188, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10401526838541031, 0.24958176910877228, 0.016829494386911392, + 0.5826361179351807, 0.26546579599380493, 0.6221785545349121, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedSchedulesTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedSchedulesTests.swift new file mode 100644 index 000000000..b774efd8c --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedSchedulesTests.swift @@ -0,0 +1,224 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 8 + +import Foundation +import MLX +import Testing + +@testable import MLXOptimizers + +@Suite("generated: Schedules") +struct GeneratedSchedulesTests { + + @Test("exponentialDecay") + func test_exponentialDecay() throws { + try withIntegrationState(seed: 61845) { + let schedule = exponentialDecay(0.1, decayRate: 0.9) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.059797536581754684, + minimum: 0.03138105943799019, + maximum: 0.10000000149011612, + absoluteSum: 0.717570424079895, + positionChecksum: 0.31554583708445233, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10000000149011612, 0.08100000023841858, 0.06560999900102615, + 0.04782969132065773, 0.038742050528526306, 0.03138105943799019, + ]), + tolerance: .float32) + } + } + + @Test("stepDecay") + func test_stepDecay() throws { + try withIntegrationState(seed: 34315) { + let schedule = stepDecay(0.1, decayRate: 0.5, stepSize: 3) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.046875, + minimum: 0.012500000186264515, + maximum: 0.10000000149011612, + absoluteSum: 0.5625, + positionChecksum: 0.19687501589457193, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10000000149011612, 0.10000000149011612, 0.05000000074505806, + 0.02500000037252903, 0.012500000186264515, 0.012500000186264515, + ]), + tolerance: .float32) + } + } + + @Test("cosineDecay") + func test_cosineDecay() throws { + // constant at the end value beyond decaySteps, which the 12 samples cover + try withIntegrationState(seed: 58536) { + let schedule = cosineDecay(0.1, decaySteps: 8) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.03750000149011612, + minimum: 0.0, + maximum: 0.10000000149011612, + absoluteSum: 0.45000001788139343, + positionChecksum: 0.11609552303949992, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10000000149011612, 0.08535534143447876, 0.04999999701976776, + 0.0038060189690440893, 0.0, 0.0, + ]), + tolerance: .float32) + } + } + + @Test("cosineDecay/end") + func test_cosineDecay_end() throws { + try withIntegrationState(seed: 5429) { + let schedule = cosineDecay(0.1, decaySteps: 8, end: 0.01) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.04374999925494194, + minimum: 0.009999999776482582, + maximum: 0.10000000149011612, + absoluteSum: 0.5249999761581421, + positionChecksum: 0.16948598623275757, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10000000149011612, 0.08681980520486832, 0.054999999701976776, + 0.013425417244434357, 0.009999999776482582, 0.009999999776482582, + ]), + tolerance: .float32) + } + } + + @Test("linearSchedule") + func test_linearSchedule() throws { + try withIntegrationState(seed: 81538) { + let schedule = linearSchedule(0.0, end: 0.1, steps: 8) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.0625, + minimum: 0.0, + maximum: 0.10000000149011612, + absoluteSum: 0.75, + positionChecksum: 0.5250000158945719, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, 0.02500000037252903, 0.05000000074505806, 0.08749999850988388, + 0.10000000149011612, 0.10000000149011612, + ]), + tolerance: .float32) + } + } + + @Test("linearSchedule/down") + func test_linearSchedule_down() throws { + try withIntegrationState(seed: 61284) { + let schedule = linearSchedule(0.1, end: 0.0, steps: 5) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.02500000409781933, + minimum: 7.450580596923828e-09, + maximum: 0.10000000149011612, + absoluteSum: 0.30000004172325134, + positionChecksum: 0.05833337207635244, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.10000000149011612, 0.06000000238418579, 0.020000003278255463, + 7.450580596923828e-09, 7.450580596923828e-09, 7.450580596923828e-09, + ]), + tolerance: .float32) + } + } + + @Test("joinSchedules") + func test_joinSchedules() throws { + // warmup then decay -- the classic use, and the boundary handling is the part that is easy to get wrong + try withIntegrationState(seed: 21828) { + let schedule = joinSchedules( + [linearSchedule(0.0, end: 0.1, steps: 4), cosineDecay(0.1, decaySteps: 8)], + boundaries: [4]) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.05000000447034836, + minimum: 0.0, + maximum: 0.10000000149011612, + absoluteSum: 0.6000000238418579, + positionChecksum: 0.3077622056007385, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, 0.05000000074505806, 0.10000000149011612, 0.0691341683268547, + 0.030865823850035667, 0.0038060189690440893, + ]), + tolerance: .float32) + } + } + + @Test("joinSchedules/three") + func test_joinSchedules_three() throws { + try withIntegrationState(seed: 92857) { + let schedule = joinSchedules( + [ + linearSchedule(0.0, end: 0.1, steps: 3), + stepDecay(0.1, decayRate: 0.5, stepSize: 2), + linearSchedule(0.05, end: 0.0, steps: 4), + ], boundaries: [3, 7]) + let values = MLXArray((0 ..< 12).map { schedule($0) }) + expectSummary( + values, + ArraySummary( + shape: [12], + dtype: .float32, + mean: 0.04375000298023224, + minimum: 0.0, + maximum: 0.10000000149011612, + absoluteSum: 0.5250000357627869, + positionChecksum: 0.24513888359069824, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.0, 0.06666667014360428, 0.10000000149011612, 0.05000000074505806, + 0.02500000037252903, 0.0, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift new file mode 100644 index 000000000..bed682645 --- /dev/null +++ b/Tests/MLXIntegrationTests/Generated/GeneratedShapeTests.swift @@ -0,0 +1,1172 @@ +// Copyright © 2026 Apple Inc. +// +// GENERATED by tools/integration_tests -- DO NOT EDIT. +// +// Values in this file were produced by python mlx and are compared against +// the Swift API. Regenerate with: +// +// python3 tools/integration_tests/generate.py +// +// python mlx: 0.32.2.dev20260910+1f8e74e3f +// vendored mlx (Cmlx): 0.32.2 +// generator revision: 2 +// cases: 49 + +import Foundation +import MLX +import Testing + +@Suite("generated: Shape") +struct GeneratedShapeTests { + + @Test("concatenated") + func test_concatenated() throws { + try withIntegrationState(seed: 31820) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.concatenated([a, b]) + expectSummary( + result, + ArraySummary( + shape: [8, 3], + dtype: .float32, + mean: 0.08590242266654968, + minimum: -1.2395046949386597, + maximum: 1.4470280408859253, + absoluteSum: 16.01652717590332, + positionChecksum: 7.683911005655925, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.6949948668479919, 1.133658766746521, -0.05453184247016907, + -0.12073975056409836, 0.3892839252948761, 0.24856531620025635, + ]), + tolerance: .float32) + } + } + + @Test("concatenated/axis") + func test_concatenated_axis() throws { + try withIntegrationState(seed: 20890) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.concatenated([a, b], axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 6], + dtype: .float32, + mean: -0.08199743926525116, + minimum: -2.4462075233459473, + maximum: 1.5454013347625732, + absoluteSum: 17.703125, + positionChecksum: 8.136674880981445, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 1.5267335176467896, -1.4755029678344727, 0.2989487051963806, + -0.7911329865455627, -0.8274667263031006, -0.6284770965576172, + ]), + tolerance: .float32) + } + } + + @Test("stacked") + func test_stacked() throws { + try withIntegrationState(seed: 94057) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.stacked([a, b]) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 3], + dtype: .float32, + mean: -0.2751633822917938, + minimum: -1.3530937433242798, + maximum: 1.9417942762374878, + absoluteSum: 18.261486053466797, + positionChecksum: 9.054578145345053, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 0.9709651470184326, -0.3056058883666992, -0.9813033938407898, + -1.3027057647705078, -0.97268146276474, 0.43900448083877563, + ]), + tolerance: .float32) + } + } + + @Test("stacked/axis") + func test_stacked_axis() throws { + try withIntegrationState(seed: 63795) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let b = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.stacked([a, b], axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3, 2], + dtype: .float32, + mean: -0.2341117560863495, + minimum: -2.2961699962615967, + maximum: 1.4487446546554565, + absoluteSum: 20.772422790527344, + positionChecksum: 11.166987101236979, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 0.1765039563179016, -0.06140340119600296, 1.031309962272644, + 1.1500154733657837, -2.2923014163970947, -1.327276587486267, + ]), + tolerance: .float32) + } + } + + @Test("flipped") + func test_flipped() throws { + try withIntegrationState(seed: 15319) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.flipped(a) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.2945154309272766, + minimum: -1.2866392135620117, + maximum: 1.67825448513031, + absoluteSum: 8.736289978027344, + positionChecksum: 5.237860679626465, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.4458805322647095, -0.5673931837081909, 0.14704769849777222, + 0.32986971735954285, -1.144845962524414, -0.7193918228149414, + ]), + tolerance: .float32) + } + } + + @Test("flipped/axis") + func test_flipped_axis() throws { + try withIntegrationState(seed: 45420) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.flipped(a, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.1344703733921051, + minimum: -1.337192416191101, + maximum: 1.5279310941696167, + absoluteSum: 11.764471054077148, + positionChecksum: 6.069742838541667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.5263196229934692, 0.5950191020965576, -1.337192416191101, + -1.1254405975341797, 0.7738415002822876, -0.5284457206726074, + ]), + tolerance: .float32) + } + } + + @Test("tiled") + func test_tiled() throws { + try withIntegrationState(seed: 25082) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tiled(a, repetitions: [2, 1]) + expectSummary( + result, + ArraySummary( + shape: [8, 3], + dtype: .float32, + mean: 0.13411381840705872, + minimum: -0.9933890700340271, + maximum: 1.0480691194534302, + absoluteSum: 13.16995620727539, + positionChecksum: 6.219711939493815, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.9933890700340271, -0.25454410910606384, -0.058666568249464035, + 0.957951545715332, 0.1598086655139923, 0.7156504392623901, + ]), + tolerance: .float32) + } + } + + @Test("repeated") + func test_repeated() throws { + try withIntegrationState(seed: 11706) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.repeated(a, count: 2, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 6], + dtype: .float32, + mean: -0.0626811534166336, + minimum: -1.3356250524520874, + maximum: 1.4082560539245605, + absoluteSum: 11.977014541625977, + positionChecksum: 7.969771067301433, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 0.03741425648331642, 0.7571091055870056, 0.26252982020378113, + -1.3356250524520874, 0.012510073371231556, 1.4082560539245605, + ]), + tolerance: .float32) + } + } + + @Test("padded") + func test_padded() throws { + try withIntegrationState(seed: 8104) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.padded(a, width: 2) + expectSummary( + result, + ArraySummary( + shape: [8, 7], + dtype: .float32, + mean: -0.016057996079325676, + minimum: -1.319933295249939, + maximum: 1.7083896398544312, + absoluteSum: 9.278562545776367, + positionChecksum: 4.887723650251116, + sampleIndices: [0, 11, 22, 33, 44, 55], + samples: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("padded/edge") + func test_padded_edge() throws { + try withIntegrationState(seed: 78439) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.padded(a, width: 2, mode: .edge) + expectSummary( + result, + ArraySummary( + shape: [8, 7], + dtype: .float32, + mean: 0.650275707244873, + minimum: -1.475324273109436, + maximum: 1.795180082321167, + absoluteSum: 59.184226989746094, + positionChecksum: 35.910391671316965, + sampleIndices: [0, 11, 22, 33, 44, 55], + samples: [ + -0.6165251731872559, 0.5478887557983398, 0.3702576458454132, + 1.1758822202682495, 1.7887344360351562, 1.2311887741088867, + ]), + tolerance: .float32) + } + } + + @Test("padded/widths") + func test_padded_widths() throws { + try withIntegrationState(seed: 15378) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.padded(a, widths: [IntOrPair([1, 2]), IntOrPair([0, 1])]) + expectSummary( + result, + ArraySummary( + shape: [7, 4], + dtype: .float32, + mean: 0.03199724853038788, + minimum: -1.4607915878295898, + maximum: 1.7407435178756714, + absoluteSum: 8.945435523986816, + positionChecksum: 4.13947514125279, + sampleIndices: [0, 5, 11, 16, 22, 27], + samples: [0.0, -0.9670328497886658, 0.0, -0.6681327223777771, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("padded/value") + func test_padded_value() throws { + try withIntegrationState(seed: 16948) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.padded(a, width: 1, value: MLXArray(Float(2.5))) + expectSummary( + result, + ArraySummary( + shape: [6, 5], + dtype: .float32, + mean: 1.8162819147109985, + minimum: -0.7915933132171631, + maximum: 2.8892898559570312, + absoluteSum: 56.07164001464844, + positionChecksum: 28.535982259114583, + sampleIndices: [0, 6, 12, 17, 23, 29], + samples: [ + 2.5, -0.7915933132171631, 2.8892898559570312, 0.3641357123851776, + 0.6989915370941162, 2.5, + ]), + tolerance: .float32) + } + } + + @Test("roll") + func test_roll() throws { + try withIntegrationState(seed: 19553) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.roll(a, shift: 2, axis: 0) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.26138874888420105, + minimum: -2.099207639694214, + maximum: 1.4279882907867432, + absoluteSum: 10.050413131713867, + positionChecksum: 5.608279546101888, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.016545826569199562, 1.4216101169586182, -0.5516468286514282, + 1.170422077178955, 0.0015867743641138077, 1.4279882907867432, + ]), + tolerance: .float32) + } + } + + @Test("transposed") + func test_transposed() throws { + try withIntegrationState(seed: 60390) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.transposed(a) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: 0.0771273821592331, + minimum: -1.0791972875595093, + maximum: 2.472217082977295, + absoluteSum: 10.713860511779785, + positionChecksum: 5.417025248209636, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.5892202854156494, -1.0791972875595093, 0.30885958671569824, + -1.0246267318725586, 1.2981765270233154, -0.5083723068237305, + ]), + tolerance: .float32) + } + } + + @Test("transposed/axes") + func test_transposed_axes() throws { + try withIntegrationState(seed: 83812) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.transposed(a, axes: [2, 0, 1]) + expectSummary( + result, + ArraySummary( + shape: [4, 2, 3], + dtype: .float32, + mean: -0.02823314070701599, + minimum: -1.502172827720642, + maximum: 1.6366482973098755, + absoluteSum: 19.033903121948242, + positionChecksum: 10.281106313069662, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.9478650689125061, 1.2468520402908325, -0.5932059288024902, + -0.07760278135538101, -0.11842215806245804, -1.2484960556030273, + ]), + tolerance: .float32) + } + } + + @Test("reshaped") + func test_reshaped() throws { + try withIntegrationState(seed: 68729) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.reshaped(a, [3, 4]) + expectSummary( + result, + ArraySummary( + shape: [3, 4], + dtype: .float32, + mean: -0.26473572850227356, + minimum: -1.2564424276351929, + maximum: 1.1341665983200073, + absoluteSum: 10.037672996520996, + positionChecksum: 5.030018170674642, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.2564424276351929, -0.8097322583198547, 0.3184874653816223, + -1.0139796733856201, -0.5813552141189575, -0.4897230267524719, + ]), + tolerance: .float32) + } + } + + @Test("squeezed") + func test_squeezed() throws { + try withIntegrationState(seed: 41015) { + let a = MLXRandom.normal([4, 1, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.squeezed(a, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.01943860575556755, + minimum: -1.4817900657653809, + maximum: 1.1566555500030518, + absoluteSum: 7.114997863769531, + positionChecksum: 3.67331600189209, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.22527667880058289, -1.4817900657653809, -0.4570019245147705, + -0.14241276681423187, -0.3114512860774994, 0.7334579229354858, + ]), + tolerance: .float32) + } + } + + @Test("expandedDimensions") + func test_expandedDimensions() throws { + try withIntegrationState(seed: 19999) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.expandedDimensions(a, axes: [0, -1]) + expectSummary( + result, + ArraySummary( + shape: [1, 4, 3, 1], + dtype: .float32, + mean: 0.09923312067985535, + minimum: -1.8916503190994263, + maximum: 1.8866270780563354, + absoluteSum: 12.624441146850586, + positionChecksum: 7.273377736409505, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 1.5140305757522583, 1.8866270780563354, -0.9565528631210327, + 1.1745210886001587, -1.0232282876968384, 1.3436799049377441, + ]), + tolerance: .float32) + } + } + + @Test("movedAxis") + func test_movedAxis() throws { + try withIntegrationState(seed: 64658) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.movedAxis(a, source: 0, destination: 2) + expectSummary( + result, + ArraySummary( + shape: [3, 4, 2], + dtype: .float32, + mean: -0.14943131804466248, + minimum: -2.045391082763672, + maximum: 1.752381682395935, + absoluteSum: 20.457523345947266, + positionChecksum: 11.733778635660807, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 0.7180314660072327, -1.3107750415802002, 0.3325449526309967, + -1.3910497426986694, -0.08477483689785004, -2.045391082763672, + ]), + tolerance: .float32) + } + } + + @Test("swappedAxes") + func test_swappedAxes() throws { + try withIntegrationState(seed: 48013) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.swappedAxes(a, 0, 2) + expectSummary( + result, + ArraySummary( + shape: [4, 3, 2], + dtype: .float32, + mean: -0.1168517917394638, + minimum: -1.750725269317627, + maximum: 1.7686270475387573, + absoluteSum: 21.129409790039062, + positionChecksum: 10.469156901041666, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -1.2027404308319092, -1.109092354774475, -0.47926244139671326, + 1.6138956546783447, 0.576898992061615, 0.4427987039089203, + ]), + tolerance: .float32) + } + } + + @Test("flattened") + func test_flattened() throws { + try withIntegrationState(seed: 34896) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.flattened(a) + expectSummary( + result, + ArraySummary( + shape: [24], + dtype: .float32, + mean: 0.09383618831634521, + minimum: -2.187122344970703, + maximum: 2.043957471847534, + absoluteSum: 21.715774536132812, + positionChecksum: 11.356913248697916, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -2.187122344970703, -1.3866022825241089, -0.039398424327373505, + 1.7639521360397339, -1.4106767177581787, -0.17367228865623474, + ]), + tolerance: .float32) + } + } + + @Test("flattened/range") + func test_flattened_range() throws { + try withIntegrationState(seed: 32731) { + let a = MLXRandom.normal([2, 3, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.flattened(a, start: 1, end: 2) + expectSummary( + result, + ArraySummary( + shape: [2, 12], + dtype: .float32, + mean: 0.08807133883237839, + minimum: -1.6179569959640503, + maximum: 1.552121877670288, + absoluteSum: 17.646623611450195, + positionChecksum: 8.385078430175781, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.7485483288764954, -1.4518442153930664, -1.6179569959640503, + 0.10551830381155014, 1.3873385190963745, -0.1756194829940796, + ]), + tolerance: .float32) + } + } + + @Test("unflatten") + func test_unflatten() throws { + try withIntegrationState(seed: 91953) { + let a = MLXRandom.normal([2, 12], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.unflatten(a, axis: 1, shape: [3, 4]) + expectSummary( + result, + ArraySummary( + shape: [2, 3, 4], + dtype: .float32, + mean: 0.24676339328289032, + minimum: -1.1256226301193237, + maximum: 1.6207804679870605, + absoluteSum: 13.571483612060547, + positionChecksum: 6.478286107381185, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + 1.2019108533859253, 0.9293265342712402, -0.8174373507499695, + 0.5026535987854004, 0.7566536068916321, 1.0857075452804565, + ]), + tolerance: .float32) + } + } + + @Test("broadcast") + func test_broadcast() throws { + try withIntegrationState(seed: 88800) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.broadcast(a, to: [2, 4, 3]) + expectSummary( + result, + ArraySummary( + shape: [2, 4, 3], + dtype: .float32, + mean: -0.19611158967018127, + minimum: -0.9691779613494873, + maximum: 1.3313171863555908, + absoluteSum: 14.432416915893555, + positionChecksum: 6.544670104980469, + sampleIndices: [0, 5, 9, 14, 18, 23], + samples: [ + -0.9691779613494873, -0.34518253803253174, -0.6659537553787231, + 1.3313171863555908, -0.8420485854148865, 0.033554743975400925, + ]), + tolerance: .float32) + } + } + + @Test("atLeast2D") + func test_atLeast2D() throws { + try withIntegrationState(seed: 76585) { + let a = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.atLeast2D(a) + expectSummary( + result, + ArraySummary( + shape: [1, 5], + dtype: .float32, + mean: 0.9482474327087402, + minimum: -0.18628670275211334, + maximum: 2.5243351459503174, + absoluteSum: 5.1138105392456055, + positionChecksum: 3.568856048583984, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 0.47301429510116577, -0.18628670275211334, 2.5243351459503174, + 0.2251853346824646, 1.7049888372421265, + ]), + tolerance: .float32) + } + } + + @Test("take") + func test_take() throws { + try withIntegrationState(seed: 96628) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 4, [5], type: Int32.self) + let result = MLX.take(a, indices, axis: 0) + expectSummary( + result, + ArraySummary( + shape: [5, 3], + dtype: .float32, + mean: -0.13128289580345154, + minimum: -1.813438057899475, + maximum: 1.2486591339111328, + absoluteSum: 13.089598655700684, + positionChecksum: 6.21908213297526, + sampleIndices: [0, 3, 6, 8, 11, 14], + samples: [ + -1.4467936754226685, -1.4467936754226685, -1.813438057899475, + -1.3840922117233276, 0.2269812375307083, -1.2374604940414429, + ]), + tolerance: .float32) + } + } + + @Test("take/flat") + func test_take_flat() throws { + try withIntegrationState(seed: 20188) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 12, [4], type: Int32.self) + let result = MLX.take(a, indices) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.20158261060714722, + minimum: -1.8244850635528564, + maximum: 0.7776317596435547, + absoluteSum: 2.9791224002838135, + positionChecksum: 1.8212910890579224, + sampleIndices: [0, 1, 2, 3], + samples: [ + -0.06824146211147308, -1.8244850635528564, 0.7776317596435547, + 0.3087643086910248, + ]), + tolerance: .float32) + } + } + + @Test("takeAlong") + func test_takeAlong() throws { + try withIntegrationState(seed: 75580) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 3, [4, 1], type: Int32.self) + let result = MLX.takeAlong(a, indices, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 1], + dtype: .float32, + mean: 0.9248526692390442, + minimum: -0.45212575793266296, + maximum: 2.0080971717834473, + absoluteSum: 4.603662490844727, + positionChecksum: 3.3495242595672607, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.18274369835853577, 2.0080971717834473, -0.45212575793266296, + 1.960695505142212, + ]), + tolerance: .float32) + } + } + + @Test("sorted") + func test_sorted() throws { + try withIntegrationState(seed: 9140) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.sorted(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.03630068898200989, + minimum: -1.7828558683395386, + maximum: 1.2648265361785889, + absoluteSum: 11.236763954162598, + positionChecksum: 5.297865549723308, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -1.7828558683395386, 1.2601059675216675, -1.2201143503189087, + 0.1945764720439911, -0.847751796245575, 1.2648265361785889, + ]), + tolerance: .float32) + } + } + + @Test("top") + func test_top() throws { + try withIntegrationState(seed: 29038) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.top(a, k: 2, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: 0.7949435710906982, + minimum: -0.42735984921455383, + maximum: 1.7273166179656982, + absoluteSum: 7.214267730712891, + positionChecksum: 3.4773030281066895, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + 0.8981940746307373, 1.7273166179656982, 1.2254910469055176, + -0.42735984921455383, 0.1055675819516182, 0.7501816153526306, + ]), + tolerance: .float32) + } + } + + @Test("tril") + func test_tril() throws { + try withIntegrationState(seed: 23707) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.tril(a) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.0325717031955719, + minimum: -1.7094670534133911, + maximum: 1.558538556098938, + absoluteSum: 10.388333320617676, + positionChecksum: 6.133814334869385, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [ + 1.5056425333023071, 0.0, 0.0, -0.28900396823883057, 1.558538556098938, + 0.8308115601539612, + ]), + tolerance: .float32) + } + } + + @Test("triu/k") + func test_triu_k() throws { + try withIntegrationState(seed: 57960) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.triu(a, k: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 4], + dtype: .float32, + mean: 0.08439695835113525, + minimum: -2.722848653793335, + maximum: 1.6600544452667236, + absoluteSum: 7.721507549285889, + positionChecksum: 2.9740920066833496, + sampleIndices: [0, 3, 6, 9, 12, 15], + samples: [0.0, 1.6600544452667236, -2.722848653793335, 0.0, 0.0, 0.0]), + tolerance: .float32) + } + } + + @Test("tri") + func test_tri() throws { + try withIntegrationState(seed: 39346) { + let result = MLX.tri(4, m: 3, k: 0, dtype: .float32) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.75, + minimum: 0.0, + maximum: 1.0, + absoluteSum: 9.0, + positionChecksum: 5.583333333333333, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [1.0, 0.0, 1.0, 1.0, 1.0, 1.0]), + tolerance: .float32) + } + } + + @Test("diag") + func test_diag() throws { + try withIntegrationState(seed: 26965) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diag(a) + expectSummary( + result, + ArraySummary( + shape: [4], + dtype: .float32, + mean: -0.21551969647407532, + minimum: -2.043041706085205, + maximum: 0.5067105889320374, + absoluteSum: 3.2240047454833984, + positionChecksum: 2.5592713356018066, + sampleIndices: [0, 1, 2, 3], + samples: [ + 0.5067105889320374, 0.4645487666130066, 0.20970351994037628, + -2.043041706085205, + ]), + tolerance: .float32) + } + } + + @Test("diagonal/offset") + func test_diagonal_offset() throws { + try withIntegrationState(seed: 21806) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diagonal(a, offset: 1) + expectSummary( + result, + ArraySummary( + shape: [3], + dtype: .float32, + mean: 1.1711373329162598, + minimum: 0.8744285702705383, + maximum: 1.4721418619155884, + absoluteSum: 3.5134119987487793, + positionChecksum: 2.2448037465413413, + sampleIndices: [0, 1, 2], + samples: [1.1668416261672974, 1.4721418619155884, 0.8744285702705383]), + tolerance: .float32) + } + } + + @Test("putAlong") + func test_putAlong() throws { + try withIntegrationState(seed: 81478) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let indices = MLXRandom.randInt(low: 0, high: 3, [4, 1], type: Int32.self) + let values = MLXRandom.normal([4, 1], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.putAlong(a, indices, values: values, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.2807766795158386, + minimum: -1.0164457559585571, + maximum: 2.2925961017608643, + absoluteSum: 8.769731521606445, + positionChecksum: 4.279629707336426, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + 0.7069619297981262, -0.7591624855995178, -0.08726170659065247, + -0.21965672075748444, 2.0261762142181396, -0.06557861715555191, + ]), + tolerance: .float32) + } + } + + @Test("argSort") + func test_argSort() throws { + try withIntegrationState(seed: 68988) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argSort(a, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .uint32, + mean: 1.0, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 12.0, + positionChecksum: 6.5, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [2.0, 1.0, 2.0, 0.0, 2.0, 1.0]), + tolerance: .exact) + } + } + + @Test("argSort/flat") + func test_argSort_flat() throws { + try withIntegrationState(seed: 69849) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argSort(a) + expectSummary( + result, + ArraySummary( + shape: [12], + dtype: .uint32, + mean: 5.5, + minimum: 0.0, + maximum: 11.0, + absoluteSum: 66.0, + positionChecksum: 28.916666666666668, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [8.0, 9.0, 2.0, 11.0, 0.0, 1.0]), + tolerance: .exact) + } + } + + @Test("partitioned") + func test_partitioned() throws { + try withIntegrationState(seed: 30615) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.partitioned(a, kth: 1, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: -0.32990625500679016, + minimum: -1.1124590635299683, + maximum: 0.795062243938446, + absoluteSum: 7.30184268951416, + positionChecksum: 4.664909362792969, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.405587762594223, 0.795062243938446, -0.04451719671487808, + -0.6547178030014038, -1.1124590635299683, -0.9105625152587891, + ]), + tolerance: .float32) + } + } + + @Test("argPartition") + func test_argPartition() throws { + try withIntegrationState(seed: 22252) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.argPartition(a, kth: 1, axis: -1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .uint32, + mean: 1.0, + minimum: 0.0, + maximum: 2.0, + absoluteSum: 12.0, + positionChecksum: 6.666666666666667, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [0.0, 1.0, 1.0, 0.0, 2.0, 0.0]), + tolerance: .exact) + } + } + + @Test("diff") + func test_diff() throws { + try withIntegrationState(seed: 47492) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diff(a) + expectSummary( + result, + ArraySummary( + shape: [4, 2], + dtype: .float32, + mean: -0.19391421973705292, + minimum: -1.7796225547790527, + maximum: 1.2921435832977295, + absoluteSum: 8.09886360168457, + positionChecksum: 4.752481460571289, + sampleIndices: [0, 1, 3, 4, 6, 7], + samples: [ + -0.9081496000289917, 0.06589078903198242, 1.2921435832977295, + 0.8175460696220398, -0.7752882242202759, 1.0981947183609009, + ]), + tolerance: .float32) + } + } + + @Test("diff/n2") + func test_diff_n2() throws { + try withIntegrationState(seed: 36184) { + let a = MLXRandom.normal([4, 5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.diff(a, n: 2, axis: 1) + expectSummary( + result, + ArraySummary( + shape: [4, 3], + dtype: .float32, + mean: 0.04767768830060959, + minimum: -2.7871458530426025, + maximum: 3.4202797412872314, + absoluteSum: 16.604129791259766, + positionChecksum: 8.50756581624349, + sampleIndices: [0, 2, 4, 7, 9, 11], + samples: [ + -0.49136456847190857, 3.4202797412872314, -0.6970906853675842, + 0.4519670903682709, -0.7846707105636597, 2.357470750808716, + ]), + tolerance: .float32) + } + } + + @Test("asStrided") + func test_asStrided() throws { + try withIntegrationState(seed: 78763) { + let a = MLXRandom.normal([4, 3], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.asStrided(a, [3, 3], strides: [1, 2], offset: 0) + expectSummary( + result, + ArraySummary( + shape: [3, 3], + dtype: .float32, + mean: -0.5163639187812805, + minimum: -2.0306308269500732, + maximum: 0.5134000182151794, + absoluteSum: 6.464855194091797, + positionChecksum: 4.407473246256511, + sampleIndices: [0, 2, 3, 5, 6, 8], + samples: [ + 0.5134000182151794, -1.3329793214797974, -0.3853185772895813, + -0.47415676712989807, 0.1477542519569397, -2.0306308269500732, + ]), + tolerance: .float32) + } + } + + @Test("atLeast1D") + func test_atLeast1D() throws { + try withIntegrationState(seed: 18570) { + let a = MLXRandom.normal([Int](), dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.atLeast1D(a) + expectSummary( + result, + ArraySummary( + shape: [1], + dtype: .float32, + mean: 0.623765230178833, + minimum: 0.623765230178833, + maximum: 0.623765230178833, + absoluteSum: 0.623765230178833, + positionChecksum: 0.623765230178833, + sampleIndices: [0], + samples: [0.623765230178833]), + tolerance: .float32) + } + } + + @Test("atLeast3D") + func test_atLeast3D() throws { + try withIntegrationState(seed: 92904) { + let a = MLXRandom.normal([5], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.atLeast3D(a) + expectSummary( + result, + ArraySummary( + shape: [1, 5, 1], + dtype: .float32, + mean: 0.3729848563671112, + minimum: -1.1092168092727661, + maximum: 1.8132435083389282, + absoluteSum: 6.191127777099609, + positionChecksum: 3.4658935546875, + sampleIndices: [0, 1, 2, 3, 4], + samples: [ + 1.8132435083389282, -1.1092168092727661, 0.9958317875862122, + -1.0538851022720337, 1.2189509868621826, + ]), + tolerance: .float32) + } + } + + @Test("trace") + func test_trace() throws { + try withIntegrationState(seed: 40135) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.trace(a) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: 3.576962471008301, + minimum: 3.576962471008301, + maximum: 3.576962471008301, + absoluteSum: 3.576962471008301, + positionChecksum: 3.576962471008301, + sampleIndices: [0], + samples: [3.576962471008301]), + tolerance: .float32) + } + } + + @Test("trace/offset") + func test_trace_offset() throws { + try withIntegrationState(seed: 47702) { + let a = MLXRandom.normal([4, 4], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.trace(a, offset: 1) + expectSummary( + result, + ArraySummary( + shape: [Int](), + dtype: .float32, + mean: -1.7535713911056519, + minimum: -1.7535713911056519, + maximum: -1.7535713911056519, + absoluteSum: 1.7535713911056519, + positionChecksum: 1.7535713911056519, + sampleIndices: [0], + samples: [-1.7535713911056519]), + tolerance: .float32) + } + } + + @Test("hadamardTransform") + func test_hadamardTransform() throws { + try withIntegrationState(seed: 82262) { + let a = MLXRandom.normal([4, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.hadamardTransform(a) + expectSummary( + result, + ArraySummary( + shape: [4, 16], + dtype: .float32, + mean: -0.09357999265193939, + minimum: -2.850245952606201, + maximum: 2.2235944271087646, + absoluteSum: 52.54167175292969, + positionChecksum: 26.58550262451172, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.9944025278091431, 0.9899498820304871, 1.330984354019165, + -1.6608858108520508, -1.4785182476043701, 0.39248645305633545, + ]), + tolerance: .float32) + } + } + + @Test("hadamardTransform/scale") + func test_hadamardTransform_scale() throws { + try withIntegrationState(seed: 50260) { + let a = MLXRandom.normal([4, 16], dtype: .float32, loc: 0.0, scale: 1.0) + let result = MLX.hadamardTransform(a, scale: 0.25) + expectSummary( + result, + ArraySummary( + shape: [4, 16], + dtype: .float32, + mean: -0.18685631453990936, + minimum: -3.034515380859375, + maximum: 2.81048583984375, + absoluteSum: 58.34088897705078, + positionChecksum: 28.606985092163086, + sampleIndices: [0, 13, 25, 38, 50, 63], + samples: [ + -0.47638797760009766, -0.6138679385185242, -0.5108878016471863, + -1.6855876445770264, -1.0864219665527344, 0.4750545620918274, + ]), + tolerance: .float32) + } + } +} diff --git a/Tests/MLXIntegrationTests/IntegrationSupport.swift b/Tests/MLXIntegrationTests/IntegrationSupport.swift new file mode 100644 index 000000000..65fab9512 --- /dev/null +++ b/Tests/MLXIntegrationTests/IntegrationSupport.swift @@ -0,0 +1,228 @@ +// Copyright © 2025 Apple Inc. + +// Hand written support for the generated integration tests -- NOT generated. +// +// The generated tests (see tools/integration_tests) declare their inputs with a +// seeded random state, evaluate one expression, and compare the result against a +// summary of the value python produced. All of the comparison policy lives here +// so it can be tuned without regenerating. + +import Foundation +import MLX +import MLXNN +import Testing + +#if canImport(Darwin) + import Darwin +#elseif canImport(Glibc) + import Glibc +#endif + +/// Statistics of an `MLXArray`, computed by python `mlx` in +/// `tools/integration_tests/core.py` (`summarize`) and recomputed here. +/// +/// The definitions must match the generator exactly: +/// +/// ``` +/// values = array.astype(float32).reshape(-1) +/// absoluteSum = sum(|values|) +/// positionChecksum = sum(|values| * arange(1, n + 1)) / n +/// samples = values[evenly spaced indices] +/// ``` +/// +/// `positionChecksum` and `samples` are what make this order sensitive: `mean`, +/// `min`, `max` and `absoluteSum` are all invariant under a permutation of the +/// elements, so a wrong axis or a transposed result would otherwise pass. +struct ArraySummary: Sendable { + let shape: [Int] + let dtype: DType + let mean: Double + let minimum: Double + let maximum: Double + let absoluteSum: Double + let positionChecksum: Double + let sampleIndices: [Int] + let samples: [Double] +} + +/// Comparison tolerance: `|actual - expected| <= absolute + relative * |expected|`. +struct Tolerance: Sendable { + let relative: Double + let absolute: Double + + /// integer and bool results must match exactly + static let exact = Tolerance(relative: 0, absolute: 0) + + /// float32 computed on GPU vs python: relative, with a floor for values near zero + static let float32 = Tolerance(relative: 1e-4, absolute: 1e-6) + + /// float16 / bfloat16 + static let float16 = Tolerance(relative: 5e-3, absolute: 1e-3) + + /// for cases where reduction order legitimately differs + static let loose = Tolerance(relative: 1e-2, absolute: 1e-4) +} + +/// Pin float32 matmuls (rather than TF32) for this process, before any MLX call. +/// +/// `MLX_ENABLE_TF32` defaults to *enabled*, and on hardware with neural accelerators +/// (M5 and later) that runs float32 matmuls on the accelerators in TF32. TF32 +/// results differ from float32 by ~1e-4 relative -- far outside the tolerances here +/// -- so every matmul-shaped case (`Linear`, `Conv*`, attention, `GRU`, `Muon`, +/// quantized matmul, `einsum`, ...) would fail unless the machine running the tests +/// matches the machine that generated the values. +/// +/// mlx reads the variable **once**, at its first use, so it has to be set before any +/// MLX work happens. That is why these tests are their own target +/// (`MLXIntegrationTests` in `Package.swift`): nothing else runs in this process, and +/// every generated case enters through `withIntegrationState(seed:)`, so setting it +/// here is enough no matter how the tests are launched. The test plan and CI also +/// set it in the environment; this covers everyone else. +private let matmulPrecisionIsFloat32: Bool = { + setenv("MLX_ENABLE_TF32", "0", 1) + return ProcessInfo.processInfo.environment["MLX_ENABLE_TF32"] == "0" +}() + +/// Run a generated case. +/// +/// - the default device is scoped to the block rather than set globally, so the +/// generated tests can run in parallel with tests that want another device +/// - MLX errors are converted into Swift errors instead of ending the process, +/// so a bad case fails its own test with the mlx message +/// - the random state is task-local and equivalent to python's +/// `mx.random.seed(seed)` +/// - the matmul precision is checked, because the values assume float32 matmuls +func withIntegrationState( + seed: UInt64, sourceLocation: SourceLocation = #_sourceLocation, _ body: () throws -> R +) throws -> R { + #expect( + matmulPrecisionIsFloat32, + """ + could not set MLX_ENABLE_TF32=0: the generated values assume float32 \ + matmuls, not TF32. Set it in the environment, or regenerate with \ + --enable-tf32 on this machine. + """, + sourceLocation: sourceLocation) + + return try Device.withDefaultDevice(.gpu) { + try withError { + try withRandomState(MLXRandom.RandomState(seed: seed), body: body) + } + } +} + +/// Assert that `array` matches the summary python produced. +func expectSummary( + _ array: MLXArray, _ expected: ArraySummary, tolerance: Tolerance = .float32, + sourceLocation: SourceLocation = #_sourceLocation +) { + #expect(array.shape == expected.shape, "shape", sourceLocation: sourceLocation) + #expect(array.dtype == expected.dtype, "dtype", sourceLocation: sourceLocation) + + guard array.shape == expected.shape else { + // the remaining comparisons would be meaningless (and may trap) + return + } + + // `reshaped([-1])` rather than `flattened()`: matches python's `reshape(-1)` + // and also works for a 0-d (scalar) result + let values = array.asType(.float32).reshaped([-1]) + let count = values.size + let absolute = MLX.abs(values) + let weights = MLX.arange(1, count + 1, dtype: .float32) + + expectClose( + Double(values.mean().item(Float.self)), expected.mean, tolerance, "mean", + sourceLocation: sourceLocation) + expectClose( + Double(values.min().item(Float.self)), expected.minimum, tolerance, "minimum", + sourceLocation: sourceLocation) + expectClose( + Double(values.max().item(Float.self)), expected.maximum, tolerance, "maximum", + sourceLocation: sourceLocation) + expectClose( + Double(absolute.sum().item(Float.self)), expected.absoluteSum, tolerance, "absoluteSum", + sourceLocation: sourceLocation) + expectClose( + Double((absolute * weights).sum().item(Float.self)) / Double(count), + expected.positionChecksum, tolerance, "positionChecksum", + sourceLocation: sourceLocation) + + for (index, expectedValue) in zip(expected.sampleIndices, expected.samples) { + expectClose( + Double(values[index].item(Float.self)), expectedValue, tolerance, + "element[\(index)]", sourceLocation: sourceLocation) + } +} + +private func expectClose( + _ actual: Double, _ expected: Double, _ tolerance: Tolerance, _ label: String, + sourceLocation: SourceLocation +) { + if expected.isNaN { + #expect( + actual.isNaN, "\(label): expected nan, got \(actual)", sourceLocation: sourceLocation) + return + } + if expected.isInfinite { + #expect( + actual == expected, "\(label): expected \(expected), got \(actual)", + sourceLocation: sourceLocation) + return + } + + let limit = tolerance.absolute + tolerance.relative * Swift.abs(expected) + let delta = Swift.abs(actual - expected) + #expect( + delta <= limit, + "\(label): expected \(expected), got \(actual) (delta \(delta) > tolerance \(limit))", + sourceLocation: sourceLocation) +} + +// MARK: - Modules + +/// The value a generated module test writes into every parameter. +/// +/// Module tests do **not** rely on python and Swift drawing the same random +/// initialization (they do not: the two implementations draw different numbers of +/// keys in different orders). Instead every parameter is replaced with a +/// deterministic function of its own shape, which both sides can produce exactly. +/// +/// This must match `PARAMETER_PY` in `tools/integration_tests/core.py`: +/// +/// ``` +/// (mx.arange(v.size, dtype=mx.float32).reshape(v.shape) / v.size - 0.5).astype(v.dtype) +/// ``` +func deterministicParameter(_ parameter: MLXArray) -> MLXArray { + let size = parameter.size + let values = MLX.arange(size, dtype: .float32).reshaped(parameter.shape) / Float(size) - 0.5 + return values.asType(parameter.dtype) +} + +/// Assert that a module's parameters are named and shaped exactly as python's. +/// +/// The keys matter beyond this test: mlx-swift uses them to load python +/// checkpoints, so a rename here is a compatibility break (this is why +/// `BatchNorm` spells its keys `running_mean` / `running_var`). The shapes catch +/// a transposed weight, which would otherwise show up as a confusing value +/// mismatch (or a broadcast error) further down. +func expectParameters( + _ module: Module, _ expected: [(String, [Int])], + sourceLocation: SourceLocation = #_sourceLocation +) { + let actual = module.parameters().flattened() + .map { ($0.0, $0.1.shape) } + .sorted { $0.0 < $1.0 } + + #expect( + actual.map { $0.0 } == expected.map { $0.0 }, "parameter keys", + sourceLocation: sourceLocation) + + for (actualParameter, expectedParameter) in zip(actual, expected) + where actualParameter.0 == expectedParameter.0 { + #expect( + actualParameter.1 == expectedParameter.1, + "parameter \(actualParameter.0) shape", + sourceLocation: sourceLocation) + } +} diff --git a/Tests/MLXTests/DistributedNNTests.swift b/Tests/MLXTests/DistributedNNTests.swift index 7cae539e0..b57848678 100644 --- a/Tests/MLXTests/DistributedNNTests.swift +++ b/Tests/MLXTests/DistributedNNTests.swift @@ -592,10 +592,6 @@ private func assertGradientWhole( /// which is why this runs in CI without a launcher. class DistributedNNTests: XCTestCase { - override class func setUp() { - setDefaultDevice() - } - override func setUpWithError() throws { try XCTSkipIf( ProcessInfo.processInfo.environment["MLX_TEST_DISTRIBUTED"] == "1", @@ -651,10 +647,6 @@ class DistributedNNRingTests: XCTestCase { /// four, and a shard of 1024 inputs still holds whole quantization groups. static let rankCount = 4 - override class func setUp() { - setDefaultDevice() - } - func testShardedLayers() throws { try DistributedHarness.run(ranks: Self.rankCount, testName: Self.testName) { group in try shardLinearBody(world: group) diff --git a/Tests/MLXTests/DistributedTests.swift b/Tests/MLXTests/DistributedTests.swift index 050c873de..f9defc304 100644 --- a/Tests/MLXTests/DistributedTests.swift +++ b/Tests/MLXTests/DistributedTests.swift @@ -15,10 +15,6 @@ import XCTest /// program with a hostfile, matching how the Python tests are run. class DistributedTests: XCTestCase { - override class func setUp() { - setDefaultDevice() - } - /// The assertions here describe a group of size one, so they are not valid /// in a multi process run -- see ``DistributedRingTests``. override func setUpWithError() throws { diff --git a/Tests/MLXTests/IntegrationTests.swift b/Tests/MLXTests/IntegrationTests.swift deleted file mode 100644 index a8f4a5399..000000000 --- a/Tests/MLXTests/IntegrationTests.swift +++ /dev/null @@ -1,7076 +0,0 @@ -// Copyright © 2024 Apple Inc. - -import Foundation -import MLX -import MLXNN -import XCTest - -@testable import MLXOptimizers - -/// Integration tests comparing results vs known results from python -/// integration. Generated by `tools/generate_integration_tests.py`. -/// -/// Note: this is not meant to be complete coverage, merely a sanity -/// check that the wrapping of the c++ core matches python (e.g. calls -/// the same functions). -class MLXIntegrationTests: XCTestCase { - - override class func setUp() { - } - - func testRandomSeed() { - MLXRandom.seed(864) - let r = MLXRandom.normal() - XCTAssertEqual( - r.item(Float.self), 1.3235496282577515, - accuracy: 0.001) - - } - - func testAddOp() { - MLXRandom.seed(394) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.17949937283992767, - accuracy: -0.0035899874567985536) - XCTAssertEqual( - a.sum().item(Float.self), -2.1539924144744873, - accuracy: -0.04307984828948975) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.12489726394414902, - accuracy: 0.00249794527888298) - XCTAssertEqual( - b.sum().item(Float.self), 1.4987671375274658, - accuracy: 0.029975342750549316) - let result = a + b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.05460209771990776, - accuracy: -0.0010920419543981553) - XCTAssertEqual( - result.sum().item(Float.self), -0.6552251577377319, - accuracy: -0.01310450315475464) - } - - func testAddOp1() { - MLXRandom.seed(776) - let a = 0.5 - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.1590166687965393, - accuracy: 0.003180333375930786) - XCTAssertEqual( - b.sum().item(Float.self), 1.9082000255584717, - accuracy: 0.03816400051116944) - let result = a + b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6590167284011841, - accuracy: 0.013180334568023682) - XCTAssertEqual( - result.sum().item(Float.self), 7.908200263977051, - accuracy: 0.15816400527954103) - } - - func testAddOp2() { - MLXRandom.seed(911) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.523838222026825, - accuracy: 0.0104767644405365) - XCTAssertEqual( - a.sum().item(Float.self), 6.28605842590332, - accuracy: 0.12572116851806642) - let b = 1.3 - let result = a + b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.8238381147384644, - accuracy: 0.03647676229476929) - XCTAssertEqual( - result.sum().item(Float.self), 21.886056900024414, - accuracy: 0.4377211380004883) - } - - func testSubOp() { - MLXRandom.seed(430) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.05095359683036804, - accuracy: 0.0010190719366073608) - XCTAssertEqual( - a.sum().item(Float.self), 0.6114431619644165, - accuracy: 0.01222886323928833) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.13153076171875, - accuracy: 0.002630615234375) - XCTAssertEqual( - b.sum().item(Float.self), 1.5783690214157104, - accuracy: 0.03156738042831421) - let result = a - b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.08057717978954315, - accuracy: -0.001611543595790863) - XCTAssertEqual( - result.sum().item(Float.self), -0.966926097869873, - accuracy: -0.01933852195739746) - } - - func testSubOp1() { - MLXRandom.seed(41) - let a = 0.5 - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.27060821652412415, - accuracy: -0.005412164330482483) - XCTAssertEqual( - b.sum().item(Float.self), -3.2472984790802, - accuracy: -0.06494596958160401) - let result = a - b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7706083059310913, - accuracy: 0.015412166118621827) - XCTAssertEqual( - result.sum().item(Float.self), 9.247299194335938, - accuracy: 0.18494598388671876) - } - - func testSubOp2() { - MLXRandom.seed(265) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.18188762664794922, - accuracy: 0.0036377525329589844) - XCTAssertEqual( - a.sum().item(Float.self), 2.1826515197753906, - accuracy: 0.043653030395507816) - let b = 1.3 - let result = a - b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -1.118112325668335, - accuracy: -0.0223622465133667) - XCTAssertEqual( - result.sum().item(Float.self), -13.417346954345703, - accuracy: -0.2683469390869141) - } - - func testMulOp() { - MLXRandom.seed(988) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.03418738767504692, - accuracy: 0.0006837477535009385) - XCTAssertEqual( - a.sum().item(Float.self), 0.41024863719940186, - accuracy: 0.008204972743988037) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.841292142868042, - accuracy: 0.01682584285736084) - XCTAssertEqual( - b.sum().item(Float.self), 10.095505714416504, - accuracy: 0.20191011428833008) - let result = a * b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.44235628843307495, - accuracy: 0.0088471257686615) - XCTAssertEqual( - result.sum().item(Float.self), 5.30827522277832, - accuracy: 0.10616550445556641) - } - - func testMulOp1() { - MLXRandom.seed(523) - let a = 0.5 - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.112852081656456, - accuracy: 0.00225704163312912) - XCTAssertEqual( - b.sum().item(Float.self), 1.3542249202728271, - accuracy: 0.027084498405456542) - let result = a * b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.056426040828228, - accuracy: 0.00112852081656456) - XCTAssertEqual( - result.sum().item(Float.self), 0.6771124601364136, - accuracy: 0.013542249202728271) - } - - func testMulOp2() { - MLXRandom.seed(497) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.24078607559204102, - accuracy: -0.004815721511840821) - XCTAssertEqual( - a.sum().item(Float.self), -2.889432907104492, - accuracy: -0.057788658142089847) - let b = 1.3 - let result = a * b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.3130219280719757, - accuracy: -0.006260438561439514) - XCTAssertEqual( - result.sum().item(Float.self), -3.756263017654419, - accuracy: -0.07512526035308838) - } - - func testDivOp() { - MLXRandom.seed(414) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.12723450362682343, - accuracy: -0.0025446900725364686) - XCTAssertEqual( - a.sum().item(Float.self), -1.5268139839172363, - accuracy: -0.030536279678344727) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.47808408737182617, - accuracy: -0.009561681747436523) - XCTAssertEqual( - b.sum().item(Float.self), -5.737009048461914, - accuracy: -0.11474018096923828) - let result = a / b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.7393910884857178, - accuracy: 0.054787821769714355) - XCTAssertEqual( - result.sum().item(Float.self), 32.8726921081543, - accuracy: 0.657453842163086) - } - - func testDivOp1() { - MLXRandom.seed(940) - let a = 0.5 - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.3181919455528259, - accuracy: -0.0063638389110565186) - XCTAssertEqual( - b.sum().item(Float.self), -3.818303346633911, - accuracy: -0.07636606693267822) - let result = a / b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.33945927023887634, - accuracy: -0.006789185404777527) - XCTAssertEqual( - result.sum().item(Float.self), -4.073511123657227, - accuracy: -0.08147022247314453) - } - - func testDivOp2() { - MLXRandom.seed(802) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.08001989126205444, - accuracy: 0.001600397825241089) - XCTAssertEqual( - a.sum().item(Float.self), 0.9602386951446533, - accuracy: 0.019204773902893067) - let b = 1.3 - let result = a / b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.06155376881361008, - accuracy: 0.0012310753762722016) - XCTAssertEqual( - result.sum().item(Float.self), 0.7386451959609985, - accuracy: 0.014772903919219971) - } - - func testModOp() { - MLXRandom.seed(849) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.5959362983703613, - accuracy: -0.011918725967407227) - XCTAssertEqual( - a.sum().item(Float.self), -7.151235103607178, - accuracy: -0.14302470207214354) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.42154282331466675, - accuracy: 0.008430856466293336) - XCTAssertEqual( - b.sum().item(Float.self), 5.058513641357422, - accuracy: 0.10117027282714844) - let result = a % b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.28858500719070435, - accuracy: 0.005771700143814087) - XCTAssertEqual( - result.sum().item(Float.self), 3.463019847869873, - accuracy: 0.06926039695739747) - } - - func testModOp1() { - MLXRandom.seed(310) - let a = 0.5 - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.15466859936714172, - accuracy: 0.0030933719873428344) - XCTAssertEqual( - b.sum().item(Float.self), 1.8560230731964111, - accuracy: 0.03712046146392822) - let result = a % b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.029992982745170593, - accuracy: -0.0005998596549034119) - XCTAssertEqual( - result.sum().item(Float.self), -0.3599157929420471, - accuracy: -0.007198315858840942) - } - - func testModOp2() { - MLXRandom.seed(991) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.34309279918670654, - accuracy: -0.006861855983734131) - XCTAssertEqual( - a.sum().item(Float.self), -4.1171135902404785, - accuracy: -0.08234227180480957) - let b = 1.3 - let result = a % b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6319072246551514, - accuracy: 0.012638144493103028) - XCTAssertEqual( - result.sum().item(Float.self), 7.582886219024658, - accuracy: 0.15165772438049316) - } - - func testPowOp() { - MLXRandom.seed(488) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.8261654376983643, - accuracy: 0.016523308753967285) - XCTAssertEqual( - a.sum().item(Float.self), 9.913985252380371, - accuracy: 0.19827970504760742) - let b = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 1.3262829780578613, - accuracy: 0.02652565956115723) - XCTAssertEqual( - b.sum().item(Float.self), 15.915395736694336, - accuracy: 0.3183079147338867) - let result = a ** b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.8045400381088257, - accuracy: 0.016090800762176515) - XCTAssertEqual( - result.sum().item(Float.self), 9.65447998046875, - accuracy: 0.193089599609375) - } - - func testPowOp1() { - MLXRandom.seed(366) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.9292365908622742, - accuracy: 0.018584731817245483) - XCTAssertEqual( - a.sum().item(Float.self), 11.150838851928711, - accuracy: 0.22301677703857423) - let b = 1.3 - let result = a ** b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.9557135105133057, - accuracy: 0.019114270210266113) - XCTAssertEqual( - result.sum().item(Float.self), 11.468562126159668, - accuracy: 0.22937124252319335) - } - - func testEqualOp() { - MLXRandom.seed(597) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.268618643283844, - accuracy: -0.00537237286567688) - XCTAssertEqual( - a.sum().item(Float.self), -3.223423480987549, - accuracy: -0.06446846961975097) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.002711281180381775, - accuracy: -5.42256236076355e-05) - XCTAssertEqual( - b.sum().item(Float.self), -0.0325353741645813, - accuracy: -0.000650707483291626) - let result = a .== b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testEqualOp1() { - MLXRandom.seed(913) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.040915559977293015, - accuracy: 0.0008183111995458603) - XCTAssertEqual( - a.sum().item(Float.self), 0.490986704826355, - accuracy: 0.0098197340965271) - let b = 1.3 - let result = a .== b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testNotEqualOp() { - MLXRandom.seed(929) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2947184443473816, - accuracy: -0.005894368886947632) - XCTAssertEqual( - a.sum().item(Float.self), -3.536621332168579, - accuracy: -0.07073242664337158) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.009182363748550415, - accuracy: -0.0001836472749710083) - XCTAssertEqual( - b.sum().item(Float.self), -0.11018836498260498, - accuracy: -0.0022037672996520997) - let result = a .!= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testNotEqualOp1() { - MLXRandom.seed(223) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.009602058678865433, - accuracy: -0.00019204117357730867) - XCTAssertEqual( - a.sum().item(Float.self), -0.11522470414638519, - accuracy: -0.002304494082927704) - let b = 1.3 - let result = a .!= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testLessThanOp() { - MLXRandom.seed(516) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.671758770942688, - accuracy: 0.01343517541885376) - XCTAssertEqual( - a.sum().item(Float.self), 8.061104774475098, - accuracy: 0.16122209548950195) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.04753169044852257, - accuracy: -0.0009506338089704514) - XCTAssertEqual( - b.sum().item(Float.self), -0.5703802704811096, - accuracy: -0.011407605409622193) - let result = a .< b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLessThanOp1() { - MLXRandom.seed(142) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.21625718474388123, - accuracy: 0.004325143694877624) - XCTAssertEqual( - a.sum().item(Float.self), 2.595086097717285, - accuracy: 0.051901721954345705) - let b = 1.3 - let result = a .< b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLessThanEqualOp() { - MLXRandom.seed(288) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.19039803743362427, - accuracy: -0.0038079607486724855) - XCTAssertEqual( - a.sum().item(Float.self), -2.284776449203491, - accuracy: -0.045695528984069825) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.10338123887777328, - accuracy: 0.002067624777555466) - XCTAssertEqual( - b.sum().item(Float.self), 1.240574836730957, - accuracy: 0.024811496734619142) - let result = a .<= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLessThanEqualOp1() { - MLXRandom.seed(143) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.025643646717071533, - accuracy: -0.0005128729343414307) - XCTAssertEqual( - a.sum().item(Float.self), -0.3077237606048584, - accuracy: -0.006154475212097168) - let b = 1.3 - let result = a .<= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testGreaterThanOp() { - MLXRandom.seed(773) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.18657740950584412, - accuracy: -0.003731548190116882) - XCTAssertEqual( - a.sum().item(Float.self), -2.23892879486084, - accuracy: -0.044778575897216795) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.0327937975525856, - accuracy: 0.0006558759510517121) - XCTAssertEqual( - b.sum().item(Float.self), 0.39352554082870483, - accuracy: 0.007870510816574097) - let result = a .> b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testGreaterThanOp1() { - MLXRandom.seed(97) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.04777970910072327, - accuracy: -0.0009555941820144654) - XCTAssertEqual( - a.sum().item(Float.self), -0.5733565092086792, - accuracy: -0.011467130184173583) - let b = 1.3 - let result = a .> b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testGreaterThanEqualOp() { - MLXRandom.seed(633) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.17345473170280457, - accuracy: 0.0034690946340560913) - XCTAssertEqual( - a.sum().item(Float.self), 2.0814566612243652, - accuracy: 0.0416291332244873) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.2481841742992401, - accuracy: 0.004963683485984803) - XCTAssertEqual( - b.sum().item(Float.self), 2.978209972381592, - accuracy: 0.05956419944763184) - let result = a .>= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testGreaterThanEqualOp1() { - MLXRandom.seed(818) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.01743374764919281, - accuracy: -0.0003486749529838562) - XCTAssertEqual( - a.sum().item(Float.self), -0.20920497179031372, - accuracy: -0.004184099435806275) - let b = 1.3 - let result = a .>= b - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testAbs() { - MLXRandom.seed(256) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.15743832290172577, - accuracy: -0.0031487664580345156) - XCTAssertEqual( - a.sum().item(Float.self), -1.8892598152160645, - accuracy: -0.03778519630432129) - let result = a.abs() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.8130466341972351, - accuracy: 0.0162609326839447) - XCTAssertEqual( - result.sum().item(Float.self), 9.756559371948242, - accuracy: 0.19513118743896485) - } - - func testAbs1() { - MLXRandom.seed(931) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.22450898587703705, - accuracy: -0.004490179717540741) - XCTAssertEqual( - a.sum().item(Float.self), -2.6941077709198, - accuracy: -0.053882155418396) - let result = abs(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6981315016746521, - accuracy: 0.013962630033493042) - XCTAssertEqual( - result.sum().item(Float.self), 8.377577781677246, - accuracy: 0.16755155563354493) - } - - func testAll() { - MLXRandom.seed(545) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.05837365239858627, - accuracy: 0.0011674730479717256) - XCTAssertEqual( - a.sum().item(Float.self), 0.7004837989807129, - accuracy: 0.014009675979614259) - let result = a.all() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAll1() { - MLXRandom.seed(722) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.01978858932852745, - accuracy: -0.00039577178657054904) - XCTAssertEqual( - a.sum().item(Float.self), -0.2374630719423294, - accuracy: -0.004749261438846588) - let result = all(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAll2() { - MLXRandom.seed(829) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.03525291010737419, - accuracy: 0.0007050582021474838) - XCTAssertEqual( - a.sum().item(Float.self), 0.4230349063873291, - accuracy: 0.008460698127746582) - let result = a.all(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAll3() { - MLXRandom.seed(616) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.1122930571436882, - accuracy: 0.002245861142873764) - XCTAssertEqual( - a.sum().item(Float.self), 1.347516655921936, - accuracy: 0.02695033311843872) - let result = all(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAll4() { - MLXRandom.seed(923) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4720103442668915, - accuracy: 0.00944020688533783) - XCTAssertEqual( - a.sum().item(Float.self), 33.984745025634766, - accuracy: 0.6796949005126953) - let result = a.all(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAll5() { - MLXRandom.seed(150) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4883745014667511, - accuracy: 0.009767490029335022) - XCTAssertEqual( - a.sum().item(Float.self), 35.1629638671875, - accuracy: 0.70325927734375) - let result = all(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny() { - MLXRandom.seed(317) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.1289883852005005, - accuracy: -0.0025797677040100097) - XCTAssertEqual( - a.sum().item(Float.self), -1.5478605031967163, - accuracy: -0.030957210063934325) - let result = a.any() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny1() { - MLXRandom.seed(101) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.08150328695774078, - accuracy: -0.0016300657391548157) - XCTAssertEqual( - a.sum().item(Float.self), -0.9780394434928894, - accuracy: -0.01956078886985779) - let result = any(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny2() { - MLXRandom.seed(747) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.23745165765285492, - accuracy: 0.004749033153057099) - XCTAssertEqual( - a.sum().item(Float.self), 2.8494198322296143, - accuracy: 0.056988396644592286) - let result = a.any(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny3() { - MLXRandom.seed(75) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3039236068725586, - accuracy: -0.006078472137451172) - XCTAssertEqual( - a.sum().item(Float.self), -3.647083282470703, - accuracy: -0.07294166564941407) - let result = any(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny4() { - MLXRandom.seed(920) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.41749411821365356, - accuracy: 0.008349882364273071) - XCTAssertEqual( - a.sum().item(Float.self), 30.0595760345459, - accuracy: 0.601191520690918) - let result = a.any(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testAny5() { - MLXRandom.seed(870) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5518582463264465, - accuracy: 0.01103716492652893) - XCTAssertEqual( - a.sum().item(Float.self), 39.73379135131836, - accuracy: 0.7946758270263672) - let result = any(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testArgMax() { - MLXRandom.seed(700) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.20903855562210083, - accuracy: -0.004180771112442016) - XCTAssertEqual( - a.sum().item(Float.self), -2.50846266746521, - accuracy: -0.0501692533493042) - let result = a.argMax() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 8.0, - accuracy: 0.16) - XCTAssertEqual( - result.sum().item(Float.self), 8, - accuracy: 0.16) - } - - func testArgMax1() { - MLXRandom.seed(338) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.07185819745063782, - accuracy: 0.0014371639490127564) - XCTAssertEqual( - a.sum().item(Float.self), 0.8622983694076538, - accuracy: 0.017245967388153077) - let result = argMax(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 7.0, - accuracy: 0.14) - XCTAssertEqual( - result.sum().item(Float.self), 7, - accuracy: 0.14) - } - - func testArgMax2() { - MLXRandom.seed(483) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.0843842476606369, - accuracy: 0.001687684953212738) - XCTAssertEqual( - a.sum().item(Float.self), 1.012610912322998, - accuracy: 0.020252218246459962) - let result = a.argMax(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 0.75, - accuracy: 0.015) - XCTAssertEqual( - result.sum().item(Float.self), 3, - accuracy: 0.06) - } - - func testArgMax3() { - MLXRandom.seed(573) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5405035018920898, - accuracy: 0.010810070037841797) - XCTAssertEqual( - a.sum().item(Float.self), 6.48604154586792, - accuracy: 0.1297208309173584) - let result = argMax(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 1.25, - accuracy: 0.025) - XCTAssertEqual( - result.sum().item(Float.self), 5, - accuracy: 0.1) - } - - func testArgMin() { - MLXRandom.seed(103) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.2983229458332062, - accuracy: 0.005966458916664124) - XCTAssertEqual( - a.sum().item(Float.self), 3.5798752307891846, - accuracy: 0.0715975046157837) - let result = a.argMin() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 11.0, - accuracy: 0.22) - XCTAssertEqual( - result.sum().item(Float.self), 11, - accuracy: 0.22) - } - - func testArgMin1() { - MLXRandom.seed(362) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.057320237159729004, - accuracy: -0.00114640474319458) - XCTAssertEqual( - a.sum().item(Float.self), -0.687842845916748, - accuracy: -0.013756856918334961) - let result = argMin(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 8.0, - accuracy: 0.16) - XCTAssertEqual( - result.sum().item(Float.self), 8, - accuracy: 0.16) - } - - func testArgMin2() { - MLXRandom.seed(444) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.22760924696922302, - accuracy: 0.004552184939384461) - XCTAssertEqual( - a.sum().item(Float.self), 2.7313108444213867, - accuracy: 0.05462621688842773) - let result = a.argMin(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 0.25, - accuracy: 0.005) - XCTAssertEqual( - result.sum().item(Float.self), 1, - accuracy: 0.02) - } - - func testArgMin3() { - MLXRandom.seed(323) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.4485395550727844, - accuracy: -0.008970791101455688) - XCTAssertEqual( - a.sum().item(Float.self), -5.382474422454834, - accuracy: -0.10764948844909668) - let result = argMin(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .uint32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5, - accuracy: 0.01) - XCTAssertEqual( - result.sum().item(Float.self), 2, - accuracy: 0.04) - } - - func testCummax() { - MLXRandom.seed(625) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2803424596786499, - accuracy: -0.0056068491935729985) - XCTAssertEqual( - a.sum().item(Float.self), -3.3641092777252197, - accuracy: -0.0672821855545044) - let result = a.cummax() - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3141178786754608, - accuracy: 0.006282357573509216) - XCTAssertEqual( - result.sum().item(Float.self), 3.7694144248962402, - accuracy: 0.07538828849792481) - } - - func testCummax1() { - MLXRandom.seed(655) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.37031763792037964, - accuracy: 0.0074063527584075925) - XCTAssertEqual( - a.sum().item(Float.self), 4.443811416625977, - accuracy: 0.08887622833251953) - let result = cummax(a) - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.3793113231658936, - accuracy: 0.027586226463317872) - XCTAssertEqual( - result.sum().item(Float.self), 16.551734924316406, - accuracy: 0.3310346984863281) - } - - func testCummax2() { - MLXRandom.seed(934) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.05676526948809624, - accuracy: -0.0011353053897619249) - XCTAssertEqual( - a.sum().item(Float.self), -0.6811832189559937, - accuracy: -0.013623664379119873) - let result = a.cummax(axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.018919458612799644, - accuracy: 0.0003783891722559929) - XCTAssertEqual( - result.sum().item(Float.self), 0.22703349590301514, - accuracy: 0.004540669918060303) - } - - func testCummax3() { - MLXRandom.seed(209) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.07204179465770721, - accuracy: 0.0014408358931541443) - XCTAssertEqual( - a.sum().item(Float.self), 0.8645014762878418, - accuracy: 0.017290029525756836) - let result = cummax(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.061327576637268, - accuracy: 0.02122655153274536) - XCTAssertEqual( - result.sum().item(Float.self), 12.735930442810059, - accuracy: 0.2547186088562012) - } - - func testCummin() { - MLXRandom.seed(989) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.16920256614685059, - accuracy: -0.0033840513229370117) - XCTAssertEqual( - a.sum().item(Float.self), -2.030430793762207, - accuracy: -0.04060861587524414) - let result = a.cummin() - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.9566025733947754, - accuracy: -0.019132051467895508) - XCTAssertEqual( - result.sum().item(Float.self), -11.479230880737305, - accuracy: -0.2295846176147461) - } - - func testCummin1() { - MLXRandom.seed(565) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.644216775894165, - accuracy: -0.0128843355178833) - XCTAssertEqual( - a.sum().item(Float.self), -7.7306013107299805, - accuracy: -0.1546120262145996) - let result = cummin(a) - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -1.7292345762252808, - accuracy: -0.03458469152450561) - XCTAssertEqual( - result.sum().item(Float.self), -20.75081443786621, - accuracy: -0.4150162887573242) - } - - func testCummin2() { - MLXRandom.seed(488) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.398279070854187, - accuracy: -0.00796558141708374) - XCTAssertEqual( - a.sum().item(Float.self), -4.779348850250244, - accuracy: -0.09558697700500489) - let result = a.cummin(axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.9816936254501343, - accuracy: -0.019633872509002687) - XCTAssertEqual( - result.sum().item(Float.self), -11.780323028564453, - accuracy: -0.23560646057128906) - } - - func testCummin3() { - MLXRandom.seed(453) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.0004353722033556551, - accuracy: 8.707444067113101e-06) - XCTAssertEqual( - a.sum().item(Float.self), 0.005224466323852539, - accuracy: 0.00010448932647705078) - let result = cummin(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.4506932497024536, - accuracy: -0.009013864994049072) - XCTAssertEqual( - result.sum().item(Float.self), -5.408318996429443, - accuracy: -0.10816637992858887) - } - - func testCumprod() { - MLXRandom.seed(886) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.0051885247230529785, - accuracy: -0.00010377049446105957) - XCTAssertEqual( - a.sum().item(Float.self), -0.06226229667663574, - accuracy: -0.001245245933532715) - let result = a.cumprod() - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.14253053069114685, - accuracy: 0.002850610613822937) - XCTAssertEqual( - result.sum().item(Float.self), 1.7103662490844727, - accuracy: 0.034207324981689456) - } - - func testCumprod1() { - MLXRandom.seed(533) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.12974336743354797, - accuracy: -0.0025948673486709596) - XCTAssertEqual( - a.sum().item(Float.self), -1.5569202899932861, - accuracy: -0.031138405799865723) - let result = cumprod(a) - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.020257528871297836, - accuracy: 0.0004051505774259567) - XCTAssertEqual( - result.sum().item(Float.self), 0.24309033155441284, - accuracy: 0.004861806631088257) - } - - func testCumprod2() { - MLXRandom.seed(266) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.20902606844902039, - accuracy: 0.004180521368980407) - XCTAssertEqual( - a.sum().item(Float.self), 2.508312702178955, - accuracy: 0.050166254043579106) - let result = a.cumprod(axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7459201812744141, - accuracy: 0.014918403625488281) - XCTAssertEqual( - result.sum().item(Float.self), 8.951042175292969, - accuracy: 0.17902084350585937) - } - - func testCumprod3() { - MLXRandom.seed(63) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.20442306995391846, - accuracy: 0.004088461399078369) - XCTAssertEqual( - a.sum().item(Float.self), 2.4530768394470215, - accuracy: 0.04906153678894043) - let result = cumprod(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.20277641713619232, - accuracy: 0.004055528342723847) - XCTAssertEqual( - result.sum().item(Float.self), 2.433316946029663, - accuracy: 0.048666338920593265) - } - - func testCumsum() { - MLXRandom.seed(824) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.07800808548927307, - accuracy: 0.0015601617097854615) - XCTAssertEqual( - a.sum().item(Float.self), 0.9360969662666321, - accuracy: 0.018721939325332643) - let result = a.cumsum() - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.16304373741149902, - accuracy: -0.0032608747482299806) - XCTAssertEqual( - result.sum().item(Float.self), -1.9565248489379883, - accuracy: -0.03913049697875977) - } - - func testCumsum1() { - MLXRandom.seed(940) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3181919455528259, - accuracy: -0.0063638389110565186) - XCTAssertEqual( - a.sum().item(Float.self), -3.818303346633911, - accuracy: -0.07636606693267822) - let result = cumsum(a) - XCTAssertEqual(result.shape, [12]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -3.2638518810272217, - accuracy: -0.06527703762054443) - XCTAssertEqual( - result.sum().item(Float.self), -39.166221618652344, - accuracy: -0.7833244323730469) - } - - func testCumsum2() { - MLXRandom.seed(561) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.19832412898540497, - accuracy: 0.0039664825797081) - XCTAssertEqual( - a.sum().item(Float.self), 2.379889488220215, - accuracy: 0.0475977897644043) - let result = a.cumsum(axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7345349788665771, - accuracy: 0.014690699577331543) - XCTAssertEqual( - result.sum().item(Float.self), 8.814419746398926, - accuracy: 0.17628839492797851) - } - - func testCumsum3() { - MLXRandom.seed(937) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.38423240184783936, - accuracy: -0.007684648036956787) - XCTAssertEqual( - a.sum().item(Float.self), -4.610788822174072, - accuracy: -0.09221577644348145) - let result = cumsum(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.6559594869613647, - accuracy: -0.013119189739227296) - XCTAssertEqual( - result.sum().item(Float.self), -7.871513366699219, - accuracy: -0.15743026733398438) - } - - func testExpandedDimensions() { - MLXRandom.seed(14) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.20228910446166992, - accuracy: -0.004045782089233399) - XCTAssertEqual( - a.sum().item(Float.self), -2.427469253540039, - accuracy: -0.04854938507080078) - let result = expandedDimensions(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3, 1]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.20228910446166992, - accuracy: -0.004045782089233399) - XCTAssertEqual( - result.sum().item(Float.self), -2.427469253540039, - accuracy: -0.04854938507080078) - } - - func testExpandedDimensions1() { - MLXRandom.seed(95) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5254887342453003, - accuracy: 0.010509774684906006) - XCTAssertEqual( - a.sum().item(Float.self), 37.83518981933594, - accuracy: 0.7567037963867188) - let result = expandedDimensions(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [1, 2, 3, 4, 3, 1]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5254887342453003, - accuracy: 0.010509774684906006) - XCTAssertEqual( - result.sum().item(Float.self), 37.83518981933594, - accuracy: 0.7567037963867188) - } - - func testFloor() { - MLXRandom.seed(736) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.035255465656518936, - accuracy: -0.0007051093131303787) - XCTAssertEqual( - a.sum().item(Float.self), -0.42306557297706604, - accuracy: -0.00846131145954132) - let result = floor(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.4166666865348816, - accuracy: -0.008333333730697633) - XCTAssertEqual( - result.sum().item(Float.self), -5.0, - accuracy: -0.1) - } - - func testLog() { - MLXRandom.seed(860) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.9451824426651001, - accuracy: 0.018903648853302004) - XCTAssertEqual( - a.sum().item(Float.self), 11.342188835144043, - accuracy: 0.22684377670288086) - let result = a.log() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.26455360651016235, - accuracy: -0.0052910721302032475) - XCTAssertEqual( - result.sum().item(Float.self), -3.174643039703369, - accuracy: -0.06349286079406738) - } - - func testLog1() { - MLXRandom.seed(408) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.1787240505218506, - accuracy: 0.023574481010437014) - XCTAssertEqual( - a.sum().item(Float.self), 14.144688606262207, - accuracy: 0.28289377212524414) - let result = log(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.08071314543485641, - accuracy: 0.0016142629086971284) - XCTAssertEqual( - result.sum().item(Float.self), 0.9685577154159546, - accuracy: 0.01937115430831909) - } - - func testLog2() { - MLXRandom.seed(727) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.0156004428863525, - accuracy: 0.020312008857727052) - XCTAssertEqual( - a.sum().item(Float.self), 12.18720531463623, - accuracy: 0.2437441062927246) - let result = a.log2() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.2411154955625534, - accuracy: -0.004822309911251068) - XCTAssertEqual( - result.sum().item(Float.self), -2.893385887145996, - accuracy: -0.057867717742919926) - } - - func testLog21() { - MLXRandom.seed(844) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.0754908323287964, - accuracy: 0.02150981664657593) - XCTAssertEqual( - a.sum().item(Float.self), 12.905889511108398, - accuracy: 0.25811779022216796) - let result = log2(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.04823940992355347, - accuracy: -0.0009647881984710694) - XCTAssertEqual( - result.sum().item(Float.self), -0.5788729190826416, - accuracy: -0.011577458381652831) - } - - func testLog10() { - MLXRandom.seed(803) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.1359423398971558, - accuracy: 0.022718846797943115) - XCTAssertEqual( - a.sum().item(Float.self), 13.631307601928711, - accuracy: 0.2726261520385742) - let result = a.log10() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.07397204637527466, - accuracy: -0.0014794409275054933) - XCTAssertEqual( - result.sum().item(Float.self), -0.8876644968986511, - accuracy: -0.017753289937973024) - } - - func testLog101() { - MLXRandom.seed(684) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.3902796506881714, - accuracy: 0.027805593013763428) - XCTAssertEqual( - a.sum().item(Float.self), 16.6833553314209, - accuracy: 0.333667106628418) - let result = log10(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.11499349772930145, - accuracy: 0.002299869954586029) - XCTAssertEqual( - result.sum().item(Float.self), 1.3799219131469727, - accuracy: 0.027598438262939454) - } - - func testLog1p() { - MLXRandom.seed(640) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.2758265733718872, - accuracy: 0.025516531467437743) - XCTAssertEqual( - a.sum().item(Float.self), 15.309918403625488, - accuracy: 0.3061983680725098) - let result = a.log1p() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7979601621627808, - accuracy: 0.015959203243255615) - XCTAssertEqual( - result.sum().item(Float.self), 9.575521469116211, - accuracy: 0.1915104293823242) - } - - func testLog1p1() { - MLXRandom.seed(1) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.045818567276001, - accuracy: 0.02091637134552002) - XCTAssertEqual( - a.sum().item(Float.self), 12.549821853637695, - accuracy: 0.25099643707275393) - let result = log1p(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.686416506767273, - accuracy: 0.01372833013534546) - XCTAssertEqual( - result.sum().item(Float.self), 8.236997604370117, - accuracy: 0.16473995208740236) - } - - func testLogSumExp() { - MLXRandom.seed(626) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.09981077909469604, - accuracy: -0.001996215581893921) - XCTAssertEqual( - a.sum().item(Float.self), -1.1977293491363525, - accuracy: -0.023954586982727052) - let result = a.logSumExp() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.6805179119110107, - accuracy: 0.05361035823822022) - XCTAssertEqual( - result.sum().item(Float.self), 2.6805179119110107, - accuracy: 0.05361035823822022) - } - - func testLogSumExp1() { - MLXRandom.seed(505) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.16046583652496338, - accuracy: 0.0032093167304992677) - XCTAssertEqual( - a.sum().item(Float.self), 1.9255900382995605, - accuracy: 0.038511800765991214) - let result = logSumExp(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 3.126417636871338, - accuracy: 0.06252835273742675) - XCTAssertEqual( - result.sum().item(Float.self), 3.126417636871338, - accuracy: 0.06252835273742675) - } - - func testLogSumExp2() { - MLXRandom.seed(847) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2282017469406128, - accuracy: -0.004564034938812256) - XCTAssertEqual( - a.sum().item(Float.self), -2.7384209632873535, - accuracy: -0.05476841926574707) - let result = a.logSumExp(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.163524866104126, - accuracy: 0.02327049732208252) - XCTAssertEqual( - result.sum().item(Float.self), 4.654099464416504, - accuracy: 0.09308198928833009) - } - - func testLogSumExp3() { - MLXRandom.seed(888) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.05087994784116745, - accuracy: 0.001017598956823349) - XCTAssertEqual( - a.sum().item(Float.self), 0.610559344291687, - accuracy: 0.01221118688583374) - let result = logSumExp(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.5178632736206055, - accuracy: 0.03035726547241211) - XCTAssertEqual( - result.sum().item(Float.self), 6.071453094482422, - accuracy: 0.12142906188964844) - } - - func testLogSumExp4() { - MLXRandom.seed(341) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5185534358024597, - accuracy: 0.010371068716049195) - XCTAssertEqual( - a.sum().item(Float.self), 37.335845947265625, - accuracy: 0.7467169189453126) - let result = a.logSumExp(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.3395986557006836, - accuracy: 0.046791973114013674) - XCTAssertEqual( - result.sum().item(Float.self), 28.07518196105957, - accuracy: 0.5615036392211914) - } - - func testLogSumExp5() { - MLXRandom.seed(249) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.48312702775001526, - accuracy: 0.009662540555000305) - XCTAssertEqual( - a.sum().item(Float.self), 34.7851448059082, - accuracy: 0.695702896118164) - let result = logSumExp(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.314382553100586, - accuracy: 0.04628765106201172) - XCTAssertEqual( - result.sum().item(Float.self), 27.7725887298584, - accuracy: 0.555451774597168) - } - - func testMax() { - MLXRandom.seed(747) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.23745165765285492, - accuracy: 0.004749033153057099) - XCTAssertEqual( - a.sum().item(Float.self), 2.8494198322296143, - accuracy: 0.056988396644592286) - let result = a.max() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.8713812828063965, - accuracy: 0.03742762565612793) - XCTAssertEqual( - result.sum().item(Float.self), 1.8713812828063965, - accuracy: 0.03742762565612793) - } - - func testMax1() { - MLXRandom.seed(333) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.08884426206350327, - accuracy: 0.0017768852412700653) - XCTAssertEqual( - a.sum().item(Float.self), 1.0661311149597168, - accuracy: 0.021322622299194335) - let result = max(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.4511317014694214, - accuracy: 0.029022634029388428) - XCTAssertEqual( - result.sum().item(Float.self), 1.4511317014694214, - accuracy: 0.029022634029388428) - } - - func testMax2() { - MLXRandom.seed(720) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.03598502278327942, - accuracy: -0.0007197004556655884) - XCTAssertEqual( - a.sum().item(Float.self), -0.431820273399353, - accuracy: -0.00863640546798706) - let result = a.max(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.0814929008483887, - accuracy: 0.021629858016967773) - XCTAssertEqual( - result.sum().item(Float.self), 4.325971603393555, - accuracy: 0.08651943206787109) - } - - func testMax3() { - MLXRandom.seed(891) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2537704408168793, - accuracy: -0.005075408816337585) - XCTAssertEqual( - a.sum().item(Float.self), -3.0452451705932617, - accuracy: -0.06090490341186523) - let result = max(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.30163291096687317, - accuracy: 0.006032658219337464) - XCTAssertEqual( - result.sum().item(Float.self), 1.2065316438674927, - accuracy: 0.024130632877349855) - } - - func testMax4() { - MLXRandom.seed(64) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5222265124320984, - accuracy: 0.010444530248641969) - XCTAssertEqual( - a.sum().item(Float.self), 37.60030746459961, - accuracy: 0.7520061492919922) - let result = a.max(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.8983621597290039, - accuracy: 0.01796724319458008) - XCTAssertEqual( - result.sum().item(Float.self), 10.780345916748047, - accuracy: 0.21560691833496093) - } - - func testMax5() { - MLXRandom.seed(195) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4561886787414551, - accuracy: 0.009123773574829101) - XCTAssertEqual( - a.sum().item(Float.self), 32.845584869384766, - accuracy: 0.6569116973876953) - let result = max(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7991632223129272, - accuracy: 0.015983264446258545) - XCTAssertEqual( - result.sum().item(Float.self), 9.589958190917969, - accuracy: 0.19179916381835938) - } - - func testMean() { - MLXRandom.seed(939) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.1672295778989792, - accuracy: -0.0033445915579795836) - XCTAssertEqual( - a.sum().item(Float.self), -2.0067548751831055, - accuracy: -0.040135097503662114) - let result = a.mean() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.1672295778989792, - accuracy: -0.0033445915579795836) - XCTAssertEqual( - result.sum().item(Float.self), -0.1672295778989792, - accuracy: -0.0033445915579795836) - } - - func testMean1() { - MLXRandom.seed(581) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3291994333267212, - accuracy: -0.006583988666534424) - XCTAssertEqual( - a.sum().item(Float.self), -3.950392961502075, - accuracy: -0.0790078592300415) - let result = mean(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.3291994333267212, - accuracy: -0.006583988666534424) - XCTAssertEqual( - result.sum().item(Float.self), -0.3291994333267212, - accuracy: -0.006583988666534424) - } - - func testMean2() { - MLXRandom.seed(227) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.28339889645576477, - accuracy: -0.005667977929115295) - XCTAssertEqual( - a.sum().item(Float.self), -3.4007866382598877, - accuracy: -0.06801573276519776) - let result = a.mean(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.2833988666534424, - accuracy: -0.005667977333068848) - XCTAssertEqual( - result.sum().item(Float.self), -1.1335954666137695, - accuracy: -0.022671909332275392) - } - - func testMean3() { - MLXRandom.seed(244) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5779480934143066, - accuracy: 0.011558961868286134) - XCTAssertEqual( - a.sum().item(Float.self), 6.93537712097168, - accuracy: 0.1387075424194336) - let result = mean(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5779481530189514, - accuracy: 0.011558963060379028) - XCTAssertEqual( - result.sum().item(Float.self), 2.3117926120758057, - accuracy: 0.046235852241516114) - } - - func testMean4() { - MLXRandom.seed(822) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4916161596775055, - accuracy: 0.00983232319355011) - XCTAssertEqual( - a.sum().item(Float.self), 35.3963623046875, - accuracy: 0.70792724609375) - let result = a.mean(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4916161894798279, - accuracy: 0.009832323789596557) - XCTAssertEqual( - result.sum().item(Float.self), 5.8993940353393555, - accuracy: 0.11798788070678712) - } - - func testMean5() { - MLXRandom.seed(990) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4603706896305084, - accuracy: 0.009207413792610168) - XCTAssertEqual( - a.sum().item(Float.self), 33.146690368652344, - accuracy: 0.6629338073730469) - let result = mean(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4603707790374756, - accuracy: 0.009207415580749513) - XCTAssertEqual( - result.sum().item(Float.self), 5.524449348449707, - accuracy: 0.11048898696899415) - } - - func testMin() { - MLXRandom.seed(145) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.16449235379695892, - accuracy: -0.0032898470759391784) - XCTAssertEqual( - a.sum().item(Float.self), -1.9739081859588623, - accuracy: -0.03947816371917725) - let result = a.min() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -1.3645747900009155, - accuracy: -0.02729149580001831) - XCTAssertEqual( - result.sum().item(Float.self), -1.3645747900009155, - accuracy: -0.02729149580001831) - } - - func testMin1() { - MLXRandom.seed(822) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.025983989238739014, - accuracy: 0.0005196797847747803) - XCTAssertEqual( - a.sum().item(Float.self), 0.31180787086486816, - accuracy: 0.006236157417297363) - let result = min(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -1.8071335554122925, - accuracy: -0.03614267110824585) - XCTAssertEqual( - result.sum().item(Float.self), -1.8071335554122925, - accuracy: -0.03614267110824585) - } - - func testMin2() { - MLXRandom.seed(556) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.47528061270713806, - accuracy: 0.00950561225414276) - XCTAssertEqual( - a.sum().item(Float.self), 5.703367233276367, - accuracy: 0.11406734466552734) - let result = a.min(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.6413009166717529, - accuracy: -0.012826018333435059) - XCTAssertEqual( - result.sum().item(Float.self), -2.5652036666870117, - accuracy: -0.051304073333740235) - } - - func testMin3() { - MLXRandom.seed(458) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.1867717206478119, - accuracy: 0.003735434412956238) - XCTAssertEqual( - a.sum().item(Float.self), 2.241260528564453, - accuracy: 0.04482521057128906) - let result = min(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.8374209403991699, - accuracy: -0.0167484188079834) - XCTAssertEqual( - result.sum().item(Float.self), -3.3496837615966797, - accuracy: -0.0669936752319336) - } - - func testMin4() { - MLXRandom.seed(93) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4830058515071869, - accuracy: 0.009660117030143738) - XCTAssertEqual( - a.sum().item(Float.self), 34.77642059326172, - accuracy: 0.6955284118652344) - let result = a.min(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.11048994958400726, - accuracy: 0.0022097989916801454) - XCTAssertEqual( - result.sum().item(Float.self), 1.3258793354034424, - accuracy: 0.02651758670806885) - } - - func testMin5() { - MLXRandom.seed(82) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5077146887779236, - accuracy: 0.010154293775558472) - XCTAssertEqual( - a.sum().item(Float.self), 36.555458068847656, - accuracy: 0.7311091613769531) - let result = min(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.12779314815998077, - accuracy: 0.0025558629631996154) - XCTAssertEqual( - result.sum().item(Float.self), 1.5335177183151245, - accuracy: 0.030670354366302492) - } - - func testProduct() { - MLXRandom.seed(327) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.296897828578949, - accuracy: -0.00593795657157898) - XCTAssertEqual( - a.sum().item(Float.self), -3.5627739429473877, - accuracy: -0.07125547885894776) - let result = a.product() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 3.555632019924815e-06, - accuracy: 7.11126403984963e-08) - XCTAssertEqual( - result.sum().item(Float.self), 3.555632019924815e-06, - accuracy: 7.11126403984963e-08) - } - - func testProduct1() { - MLXRandom.seed(896) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.05823778361082077, - accuracy: -0.0011647556722164155) - XCTAssertEqual( - a.sum().item(Float.self), -0.6988533735275269, - accuracy: -0.013977067470550537) - let result = product(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0047314404509961605, - accuracy: 9.462880901992321e-05) - XCTAssertEqual( - result.sum().item(Float.self), 0.0047314404509961605, - accuracy: 9.462880901992321e-05) - } - - func testProduct2() { - MLXRandom.seed(520) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.19801780581474304, - accuracy: 0.003960356116294861) - XCTAssertEqual( - a.sum().item(Float.self), 2.376213550567627, - accuracy: 0.04752427101135254) - let result = a.product(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0764155238866806, - accuracy: 0.0015283104777336122) - XCTAssertEqual( - result.sum().item(Float.self), 0.3056620955467224, - accuracy: 0.006113241910934449) - } - - func testProduct3() { - MLXRandom.seed(955) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.36796003580093384, - accuracy: 0.007359200716018677) - XCTAssertEqual( - a.sum().item(Float.self), 4.415520191192627, - accuracy: 0.08831040382385254) - let result = product(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0716419368982315, - accuracy: 0.0014328387379646302) - XCTAssertEqual( - result.sum().item(Float.self), 0.286567747592926, - accuracy: 0.005731354951858521) - } - - func testProduct4() { - MLXRandom.seed(501) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5341216921806335, - accuracy: 0.010682433843612671) - XCTAssertEqual( - a.sum().item(Float.self), 38.45676040649414, - accuracy: 0.7691352081298828) - let result = a.product(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.02406497858464718, - accuracy: 0.0004812995716929436) - XCTAssertEqual( - result.sum().item(Float.self), 0.28877973556518555, - accuracy: 0.0057755947113037115) - } - - func testProduct5() { - MLXRandom.seed(111) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4650188088417053, - accuracy: 0.009300376176834107) - XCTAssertEqual( - a.sum().item(Float.self), 33.481353759765625, - accuracy: 0.6696270751953125) - let result = product(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.007884619757533073, - accuracy: 0.00015769239515066148) - XCTAssertEqual( - result.sum().item(Float.self), 0.09461543709039688, - accuracy: 0.0018923087418079377) - } - - func testReciprocal() { - MLXRandom.seed(308) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.12339428067207336, - accuracy: -0.0024678856134414673) - XCTAssertEqual( - a.sum().item(Float.self), -1.4807313680648804, - accuracy: -0.029614627361297607) - let result = a.reciprocal() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.3978341817855835, - accuracy: 0.02795668363571167) - XCTAssertEqual( - result.sum().item(Float.self), 16.774009704589844, - accuracy: 0.33548019409179686) - } - - func testReciprocal1() { - MLXRandom.seed(564) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.20050053298473358, - accuracy: -0.004010010659694672) - XCTAssertEqual( - a.sum().item(Float.self), -2.406006336212158, - accuracy: -0.048120126724243165) - let result = reciprocal(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.753860592842102, - accuracy: -0.015077211856842042) - XCTAssertEqual( - result.sum().item(Float.self), -9.046326637268066, - accuracy: -0.18092653274536133) - } - - func testRound() { - MLXRandom.seed(298) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.23528259992599487, - accuracy: 0.004705651998519898) - XCTAssertEqual( - a.sum().item(Float.self), 2.8233911991119385, - accuracy: 0.05646782398223877) - let result = a.round() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.25, - accuracy: 0.005) - XCTAssertEqual( - result.sum().item(Float.self), 3.0, - accuracy: 0.06) - } - - func testRound1() { - MLXRandom.seed(723) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.16675782203674316, - accuracy: 0.0033351564407348632) - XCTAssertEqual( - a.sum().item(Float.self), 2.001093864440918, - accuracy: 0.04002187728881836) - let result = round(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.25, - accuracy: 0.005) - XCTAssertEqual( - result.sum().item(Float.self), 3.0, - accuracy: 0.06) - } - - func testSin() { - MLXRandom.seed(127) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.18461060523986816, - accuracy: -0.0036922121047973633) - XCTAssertEqual( - a.sum().item(Float.self), -2.215327262878418, - accuracy: -0.04430654525756836) - let result = a.sin() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.15534619987010956, - accuracy: -0.0031069239974021914) - XCTAssertEqual( - result.sum().item(Float.self), -1.86415433883667, - accuracy: -0.0372830867767334) - } - - func testSin1() { - MLXRandom.seed(560) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.15535306930541992, - accuracy: -0.0031070613861083987) - XCTAssertEqual( - a.sum().item(Float.self), -1.864236831665039, - accuracy: -0.03728473663330078) - let result = sin(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.13987457752227783, - accuracy: -0.002797491550445557) - XCTAssertEqual( - result.sum().item(Float.self), -1.6784948110580444, - accuracy: -0.03356989622116089) - } - - func testCos() { - MLXRandom.seed(340) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.18920934200286865, - accuracy: 0.003784186840057373) - XCTAssertEqual( - a.sum().item(Float.self), 2.270512104034424, - accuracy: 0.04541024208068848) - let result = a.cos() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7257080078125, - accuracy: 0.01451416015625) - XCTAssertEqual( - result.sum().item(Float.self), 8.70849609375, - accuracy: 0.174169921875) - } - - func testCos1() { - MLXRandom.seed(834) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.18475177884101868, - accuracy: 0.0036950355768203737) - XCTAssertEqual( - a.sum().item(Float.self), 2.2170212268829346, - accuracy: 0.04434042453765869) - let result = cos(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4963728189468384, - accuracy: 0.009927456378936769) - XCTAssertEqual( - result.sum().item(Float.self), 5.9564738273620605, - accuracy: 0.11912947654724121) - } - - func testSqrt() { - MLXRandom.seed(944) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.1813615560531616, - accuracy: 0.023627231121063234) - XCTAssertEqual( - a.sum().item(Float.self), 14.176338195800781, - accuracy: 0.2835267639160156) - let result = a.sqrt() - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.040137767791748, - accuracy: 0.02080275535583496) - XCTAssertEqual( - result.sum().item(Float.self), 12.481653213500977, - accuracy: 0.24963306427001955) - } - - func testSqrt1() { - MLXRandom.seed(553) - let a = MLXRandom.uniform(low: 0.1, high: 2.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.2980049848556519, - accuracy: 0.025960099697113038) - XCTAssertEqual( - a.sum().item(Float.self), 15.576059341430664, - accuracy: 0.3115211868286133) - let result = sqrt(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.118424415588379, - accuracy: 0.02236848831176758) - XCTAssertEqual( - result.sum().item(Float.self), 13.421092987060547, - accuracy: 0.26842185974121097) - } - - func testSum() { - MLXRandom.seed(208) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.013726795092225075, - accuracy: 0.0002745359018445015) - XCTAssertEqual( - a.sum().item(Float.self), 0.1647215336561203, - accuracy: 0.003294430673122406) - let result = a.sum() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.1647215336561203, - accuracy: 0.003294430673122406) - XCTAssertEqual( - result.sum().item(Float.self), 0.1647215336561203, - accuracy: 0.003294430673122406) - } - - func testSum1() { - MLXRandom.seed(986) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.08680947870016098, - accuracy: 0.0017361895740032197) - XCTAssertEqual( - a.sum().item(Float.self), 1.0417137145996094, - accuracy: 0.020834274291992187) - let result = sum(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.0417137145996094, - accuracy: 0.020834274291992187) - XCTAssertEqual( - result.sum().item(Float.self), 1.0417137145996094, - accuracy: 0.020834274291992187) - } - - func testSum2() { - MLXRandom.seed(818) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.01743374764919281, - accuracy: -0.0003486749529838562) - XCTAssertEqual( - a.sum().item(Float.self), -0.20920497179031372, - accuracy: -0.004184099435806275) - let result = a.sum(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.05230128765106201, - accuracy: -0.0010460257530212403) - XCTAssertEqual( - result.sum().item(Float.self), -0.20920515060424805, - accuracy: -0.004184103012084961) - } - - func testSum3() { - MLXRandom.seed(617) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2992609739303589, - accuracy: -0.005985219478607178) - XCTAssertEqual( - a.sum().item(Float.self), -3.5911316871643066, - accuracy: -0.07182263374328614) - let result = sum(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.8977828621864319, - accuracy: -0.017955657243728638) - XCTAssertEqual( - result.sum().item(Float.self), -3.5911314487457275, - accuracy: -0.07182262897491455) - } - - func testSum4() { - MLXRandom.seed(560) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4772148132324219, - accuracy: 0.009544296264648438) - XCTAssertEqual( - a.sum().item(Float.self), 34.359466552734375, - accuracy: 0.6871893310546875) - let result = a.sum(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.8632893562316895, - accuracy: 0.05726578712463379) - XCTAssertEqual( - result.sum().item(Float.self), 34.35947036743164, - accuracy: 0.6871894073486329) - } - - func testSum5() { - MLXRandom.seed(601) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5380033850669861, - accuracy: 0.010760067701339722) - XCTAssertEqual( - a.sum().item(Float.self), 38.736244201660156, - accuracy: 0.7747248840332032) - let result = sum(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 3.228020429611206, - accuracy: 0.06456040859222412) - XCTAssertEqual( - result.sum().item(Float.self), 38.736244201660156, - accuracy: 0.7747248840332032) - } - - func testVariance() { - MLXRandom.seed(294) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.047887593507766724, - accuracy: 0.0009577518701553345) - XCTAssertEqual( - a.sum().item(Float.self), 0.5746511220932007, - accuracy: 0.011493022441864014) - let result = a.variance() - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6768754124641418, - accuracy: 0.013537508249282838) - XCTAssertEqual( - result.sum().item(Float.self), 0.6768754124641418, - accuracy: 0.013537508249282838) - } - - func testVariance1() { - MLXRandom.seed(455) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.07983233034610748, - accuracy: 0.0015966466069221496) - XCTAssertEqual( - a.sum().item(Float.self), 0.957987904548645, - accuracy: 0.019159758090972902) - let result = variance(a) - XCTAssertEqual(result.shape, []) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5957978367805481, - accuracy: 0.011915956735610963) - XCTAssertEqual( - result.sum().item(Float.self), 0.5957978367805481, - accuracy: 0.011915956735610963) - } - - func testVariance2() { - MLXRandom.seed(93) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.028267383575439453, - accuracy: -0.0005653476715087891) - XCTAssertEqual( - a.sum().item(Float.self), -0.33920860290527344, - accuracy: -0.006784172058105469) - let result = a.variance(axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.0322922468185425, - accuracy: 0.02064584493637085) - XCTAssertEqual( - result.sum().item(Float.self), 4.12916898727417, - accuracy: 0.0825833797454834) - } - - func testVariance3() { - MLXRandom.seed(610) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.01769038289785385, - accuracy: -0.00035380765795707704) - XCTAssertEqual( - a.sum().item(Float.self), -0.21228459477424622, - accuracy: -0.004245691895484924) - let result = variance(a, axis: -1) - XCTAssertEqual(result.shape, [4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.591558039188385, - accuracy: 0.0118311607837677) - XCTAssertEqual( - result.sum().item(Float.self), 2.36623215675354, - accuracy: 0.0473246431350708) - } - - func testVariance4() { - MLXRandom.seed(817) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4460234045982361, - accuracy: 0.008920468091964721) - XCTAssertEqual( - a.sum().item(Float.self), 32.113685607910156, - accuracy: 0.6422737121582032) - let result = a.variance(axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0666377991437912, - accuracy: 0.001332755982875824) - XCTAssertEqual( - result.sum().item(Float.self), 0.7996535301208496, - accuracy: 0.015993070602416993) - } - - func testVariance5() { - MLXRandom.seed(394) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.56817227602005, - accuracy: 0.011363445520401002) - XCTAssertEqual( - a.sum().item(Float.self), 40.90840530395508, - accuracy: 0.8181681060791016) - let result = variance(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.07289446890354156, - accuracy: 0.0014578893780708313) - XCTAssertEqual( - result.sum().item(Float.self), 0.8747336268424988, - accuracy: 0.017494672536849977) - } - - func testAcos() { - MLXRandom.seed(324) - let a = MLXRandom.uniform(low: 0.1, high: 1.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5365500450134277, - accuracy: 0.010731000900268555) - XCTAssertEqual( - a.sum().item(Float.self), 6.438600063323975, - accuracy: 0.1287720012664795) - let result = acos(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.9425711631774902, - accuracy: 0.018851423263549806) - XCTAssertEqual( - result.sum().item(Float.self), 11.310853958129883, - accuracy: 0.22621707916259767) - } - - func testAcosh() { - MLXRandom.seed(589) - let a = MLXRandom.uniform(low: 1, high: 3, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.9692752361297607, - accuracy: 0.039385504722595215) - XCTAssertEqual( - a.sum().item(Float.self), 23.631301879882812, - accuracy: 0.47262603759765626) - let result = acosh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.259089469909668, - accuracy: 0.02518178939819336) - XCTAssertEqual( - result.sum().item(Float.self), 15.109073638916016, - accuracy: 0.3021814727783203) - } - - func testAsin() { - MLXRandom.seed(247) - let a = MLXRandom.uniform(low: 0.1, high: 1.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4699937105178833, - accuracy: 0.009399874210357666) - XCTAssertEqual( - a.sum().item(Float.self), 5.6399245262146, - accuracy: 0.112798490524292) - let result = asin(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5069084763526917, - accuracy: 0.010138169527053834) - XCTAssertEqual( - result.sum().item(Float.self), 6.082901477813721, - accuracy: 0.12165802955627442) - } - - func testAsinh() { - MLXRandom.seed(297) - let a = MLXRandom.uniform(low: 1, high: 3, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 1.7906322479248047, - accuracy: 0.03581264495849609) - XCTAssertEqual( - a.sum().item(Float.self), 21.487586975097656, - accuracy: 0.4297517395019531) - let result = asinh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.3248964548110962, - accuracy: 0.026497929096221923) - XCTAssertEqual( - result.sum().item(Float.self), 15.898756980895996, - accuracy: 0.31797513961791996) - } - - func testAtan() { - MLXRandom.seed(188) - let a = MLXRandom.uniform(low: 0.1, high: 1.0, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.43250617384910583, - accuracy: 0.008650123476982116) - XCTAssertEqual( - a.sum().item(Float.self), 5.1900739669799805, - accuracy: 0.10380147933959961) - let result = atan(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3987109661102295, - accuracy: 0.00797421932220459) - XCTAssertEqual( - result.sum().item(Float.self), 4.784531593322754, - accuracy: 0.09569063186645509) - } - - func testAtanh() { - MLXRandom.seed(193) - let a = MLXRandom.uniform(low: 0.1, high: 0.9, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.593014657497406, - accuracy: 0.01186029314994812) - XCTAssertEqual( - a.sum().item(Float.self), 7.116175651550293, - accuracy: 0.14232351303100585) - let result = atanh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.759239137172699, - accuracy: 0.01518478274345398) - XCTAssertEqual( - result.sum().item(Float.self), 9.110869407653809, - accuracy: 0.18221738815307617) - } - - func testCeil() { - MLXRandom.seed(841) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3292633295059204, - accuracy: -0.006585266590118408) - XCTAssertEqual( - a.sum().item(Float.self), -3.951159715652466, - accuracy: -0.07902319431304931) - let result = ceil(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3333333432674408, - accuracy: 0.006666666865348816) - XCTAssertEqual( - result.sum().item(Float.self), 4.0, - accuracy: 0.08) - } - - func testCosh() { - MLXRandom.seed(191) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.5465152263641357, - accuracy: -0.010930304527282716) - XCTAssertEqual( - a.sum().item(Float.self), -6.558182239532471, - accuracy: -0.1311636447906494) - let result = cosh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 1.735574722290039, - accuracy: 0.03471149444580078) - XCTAssertEqual( - result.sum().item(Float.self), 20.82689666748047, - accuracy: 0.41653793334960937) - } - - func testErf() { - MLXRandom.seed(33) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.15268580615520477, - accuracy: 0.0030537161231040956) - XCTAssertEqual( - a.sum().item(Float.self), 1.8322296142578125, - accuracy: 0.03664459228515625) - let result = erf(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.16305577754974365, - accuracy: 0.003261115550994873) - XCTAssertEqual( - result.sum().item(Float.self), 1.9566693305969238, - accuracy: 0.03913338661193848) - } - - func testErfInverse() { - MLXRandom.seed(627) - let a = MLXRandom.uniform(low: 0.1, high: 0.9, [4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4853326678276062, - accuracy: 0.009706653356552124) - XCTAssertEqual( - a.sum().item(Float.self), 5.823991775512695, - accuracy: 0.11647983551025391) - let result = erfInverse(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5043282508850098, - accuracy: 0.010086565017700196) - XCTAssertEqual( - result.sum().item(Float.self), 6.051938533782959, - accuracy: 0.12103877067565919) - } - - func testLogicalNot() { - MLXRandom.seed(672) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.16328130662441254, - accuracy: 0.0032656261324882506) - XCTAssertEqual( - a.sum().item(Float.self), 1.9593756198883057, - accuracy: 0.03918751239776611) - let result = logicalNot(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testNegative() { - MLXRandom.seed(266) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.20902606844902039, - accuracy: 0.004180521368980407) - XCTAssertEqual( - a.sum().item(Float.self), 2.508312702178955, - accuracy: 0.050166254043579106) - let result = negative(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.20902606844902039, - accuracy: -0.004180521368980407) - XCTAssertEqual( - result.sum().item(Float.self), -2.508312702178955, - accuracy: -0.050166254043579106) - } - - func testSigmoid() { - MLXRandom.seed(487) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.4737564027309418, - accuracy: -0.009475128054618835) - XCTAssertEqual( - a.sum().item(Float.self), -5.685076713562012, - accuracy: -0.11370153427124023) - let result = sigmoid(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.40112757682800293, - accuracy: 0.00802255153656006) - XCTAssertEqual( - result.sum().item(Float.self), 4.813530921936035, - accuracy: 0.0962706184387207) - } - - func testSign() { - MLXRandom.seed(70) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.19362865388393402, - accuracy: -0.0038725730776786806) - XCTAssertEqual( - a.sum().item(Float.self), -2.3235437870025635, - accuracy: -0.04647087574005127) - let result = sign(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.3333333432674408, - accuracy: -0.006666666865348816) - XCTAssertEqual( - result.sum().item(Float.self), -4.0, - accuracy: -0.08) - } - - func testSinh() { - MLXRandom.seed(91) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.37823864817619324, - accuracy: 0.007564772963523865) - XCTAssertEqual( - a.sum().item(Float.self), 4.538863658905029, - accuracy: 0.09077727317810058) - let result = sinh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.8675357103347778, - accuracy: 0.017350714206695556) - XCTAssertEqual( - result.sum().item(Float.self), 10.410428047180176, - accuracy: 0.20820856094360352) - } - - func testSoftMax() { - MLXRandom.seed(695) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.14673520624637604, - accuracy: -0.002934704124927521) - XCTAssertEqual( - a.sum().item(Float.self), -1.7608224153518677, - accuracy: -0.03521644830703735) - let result = softMax(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0833333283662796, - accuracy: 0.0016666665673255921) - XCTAssertEqual( - result.sum().item(Float.self), 0.9999998807907104, - accuracy: 0.01999999761581421) - } - - func testSoftMax1() { - MLXRandom.seed(775) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.27563148736953735, - accuracy: 0.005512629747390747) - XCTAssertEqual( - a.sum().item(Float.self), 3.307577610015869, - accuracy: 0.06615155220031739) - let result = softMax(a, axis: -1) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3333333432674408, - accuracy: 0.006666666865348816) - XCTAssertEqual( - result.sum().item(Float.self), 4.0, - accuracy: 0.08) - } - - func testSoftMax2() { - MLXRandom.seed(133) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 3, 4, 3]) - XCTAssertEqual(a.shape, [2, 3, 4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5304720401763916, - accuracy: 0.010609440803527832) - XCTAssertEqual( - a.sum().item(Float.self), 38.19398498535156, - accuracy: 0.7638796997070313) - let result = softMax(a, axes: [0, -1]) - XCTAssertEqual(result.shape, [2, 3, 4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.1666666716337204, - accuracy: 0.003333333432674408) - XCTAssertEqual( - result.sum().item(Float.self), 12.0, - accuracy: 0.24) - } - - func testTan() { - MLXRandom.seed(897) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3048655390739441, - accuracy: -0.006097310781478882) - XCTAssertEqual( - a.sum().item(Float.self), -3.65838623046875, - accuracy: -0.073167724609375) - let result = tan(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0694185122847557, - accuracy: 0.001388370245695114) - XCTAssertEqual( - result.sum().item(Float.self), 0.8330221176147461, - accuracy: 0.016660442352294923) - } - - func testTanh() { - MLXRandom.seed(153) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.45169347524642944, - accuracy: -0.009033869504928588) - XCTAssertEqual( - a.sum().item(Float.self), -5.420321464538574, - accuracy: -0.10840642929077149) - let result = tanh(a) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.3471390902996063, - accuracy: -0.006942781805992127) - XCTAssertEqual( - result.sum().item(Float.self), -4.165668964385986, - accuracy: -0.08331337928771973) - } - - func testMLXadd() { - MLXRandom.seed(945) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.3744015097618103, - accuracy: 0.007488030195236206) - XCTAssertEqual( - a.sum().item(Float.self), 4.4928178787231445, - accuracy: 0.08985635757446289) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.20527128875255585, - accuracy: -0.004105425775051117) - XCTAssertEqual( - b.sum().item(Float.self), -2.4632554054260254, - accuracy: -0.04926510810852051) - let result = MLX.add(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.16913020610809326, - accuracy: 0.0033826041221618653) - XCTAssertEqual( - result.sum().item(Float.self), 2.029562473297119, - accuracy: 0.040591249465942385) - } - - func testConv1d() { - MLXRandom.seed(39) - let a = MLXRandom.uniform(0.0 ..< 1.0, [4, 10, 4]) - XCTAssertEqual(a.shape, [4, 10, 4]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5018469095230103, - accuracy: 0.010036938190460205) - XCTAssertEqual( - a.sum().item(Float.self), 80.29550170898438, - accuracy: 1.6059100341796875) - let b = MLXRandom.uniform(0.0 ..< 1.0, [2, 10, 4]) - XCTAssertEqual(b.shape, [2, 10, 4]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.5621036291122437, - accuracy: 0.011242072582244873) - XCTAssertEqual( - b.sum().item(Float.self), 44.96828842163086, - accuracy: 0.8993657684326172) - let result = conv1d(a, b) - XCTAssertEqual(result.shape, [4, 1, 2]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 11.004511833190918, - accuracy: 0.22009023666381836) - XCTAssertEqual( - result.sum().item(Float.self), 88.03609466552734, - accuracy: 1.7607218933105468) - } - - func testConv2d() { - MLXRandom.seed(862) - let a = MLXRandom.uniform(0.0 ..< 1.0, [4, 10, 12, 4]) - XCTAssertEqual(a.shape, [4, 10, 12, 4]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5062730312347412, - accuracy: 0.010125460624694825) - XCTAssertEqual( - a.sum().item(Float.self), 972.044189453125, - accuracy: 19.4408837890625) - let b = MLXRandom.uniform(0.0 ..< 1.0, [2, 10, 12, 4]) - XCTAssertEqual(b.shape, [2, 10, 12, 4]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.4991826117038727, - accuracy: 0.009983652234077454) - XCTAssertEqual( - b.sum().item(Float.self), 479.21527099609375, - accuracy: 9.584305419921876) - let result = conv2d(a, b) - XCTAssertEqual(result.shape, [4, 1, 1, 2]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 121.8356704711914, - accuracy: 2.4367134094238283) - XCTAssertEqual( - result.sum().item(Float.self), 974.6853637695312, - accuracy: 19.493707275390626) - } - - func testConvolve() { - MLXRandom.seed(82) - let a = MLXRandom.uniform(0.0 ..< 1.0, [20]) - XCTAssertEqual(a.shape, [20]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4547460675239563, - accuracy: 0.009094921350479125) - XCTAssertEqual( - a.sum().item(Float.self), 9.094921112060547, - accuracy: 0.18189842224121094) - let b = MLXRandom.uniform(0.0 ..< 1.0, [4]) - XCTAssertEqual(b.shape, [4]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.36941707134246826, - accuracy: 0.007388341426849365) - XCTAssertEqual( - b.sum().item(Float.self), 1.477668285369873, - accuracy: 0.02955336570739746) - let result = convolve(a, b) - XCTAssertEqual(result.shape, [23]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5843163132667542, - accuracy: 0.011686326265335084) - XCTAssertEqual( - result.sum().item(Float.self), 13.439274787902832, - accuracy: 0.26878549575805666) - } - - func testDivide() { - MLXRandom.seed(919) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.2972916066646576, - accuracy: -0.005945832133293152) - XCTAssertEqual( - a.sum().item(Float.self), -3.5674991607666016, - accuracy: -0.07134998321533204) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.04941798374056816, - accuracy: 0.0009883596748113633) - XCTAssertEqual( - b.sum().item(Float.self), 0.5930157899856567, - accuracy: 0.011860315799713136) - let result = divide(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -3.7239060401916504, - accuracy: -0.07447812080383301) - XCTAssertEqual( - result.sum().item(Float.self), -44.68687057495117, - accuracy: -0.8937374114990234) - } - - func testEqual() { - MLXRandom.seed(716) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.3898356854915619, - accuracy: 0.007796713709831238) - XCTAssertEqual( - a.sum().item(Float.self), 4.678028106689453, - accuracy: 0.09356056213378906) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.30894142389297485, - accuracy: -0.0061788284778594976) - XCTAssertEqual( - b.sum().item(Float.self), -3.707296848297119, - accuracy: -0.07414593696594239) - let result = equal(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), false) - } - - func testGreater() { - MLXRandom.seed(945) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.3744015097618103, - accuracy: 0.007488030195236206) - XCTAssertEqual( - a.sum().item(Float.self), 4.4928178787231445, - accuracy: 0.08985635757446289) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.20527128875255585, - accuracy: -0.004105425775051117) - XCTAssertEqual( - b.sum().item(Float.self), -2.4632554054260254, - accuracy: -0.04926510810852051) - let result = greater(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testGreaterEqual() { - MLXRandom.seed(849) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.5959362983703613, - accuracy: -0.011918725967407227) - XCTAssertEqual( - a.sum().item(Float.self), -7.151235103607178, - accuracy: -0.14302470207214354) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.42154282331466675, - accuracy: 0.008430856466293336) - XCTAssertEqual( - b.sum().item(Float.self), 5.058513641357422, - accuracy: 0.10117027282714844) - let result = greaterEqual(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLess() { - MLXRandom.seed(553) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.41794854402542114, - accuracy: 0.008358970880508423) - XCTAssertEqual( - a.sum().item(Float.self), 5.015382289886475, - accuracy: 0.1003076457977295) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.19319520890712738, - accuracy: -0.0038639041781425476) - XCTAssertEqual( - b.sum().item(Float.self), -2.318342447280884, - accuracy: -0.04636684894561768) - let result = less(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLessEqual() { - MLXRandom.seed(699) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.06386314332485199, - accuracy: 0.0012772628664970398) - XCTAssertEqual( - a.sum().item(Float.self), 0.7663577198982239, - accuracy: 0.015327154397964478) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.44176310300827026, - accuracy: -0.008835262060165406) - XCTAssertEqual( - b.sum().item(Float.self), -5.301156997680664, - accuracy: -0.10602313995361329) - let result = lessEqual(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), false) - XCTAssertEqual(result.any().item(), true) - } - - func testLogAddExp() { - MLXRandom.seed(400) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.19093476235866547, - accuracy: -0.0038186952471733096) - XCTAssertEqual( - a.sum().item(Float.self), -2.291217088699341, - accuracy: -0.04582434177398682) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.09774001687765121, - accuracy: 0.001954800337553024) - XCTAssertEqual( - b.sum().item(Float.self), 1.1728801727294922, - accuracy: 0.023457603454589845) - let result = logAddExp(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.7312957048416138, - accuracy: 0.014625914096832275) - XCTAssertEqual( - result.sum().item(Float.self), 8.775547981262207, - accuracy: 0.17551095962524416) - } - - func testMatmul() { - MLXRandom.seed(857) - let a = MLXRandom.uniform(0.0 ..< 1.0, [10, 8]) - XCTAssertEqual(a.shape, [10, 8]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.517481803894043, - accuracy: 0.01034963607788086) - XCTAssertEqual( - a.sum().item(Float.self), 41.39854431152344, - accuracy: 0.8279708862304688) - let b = MLXRandom.uniform(0.0 ..< 1.0, [8, 13]) - XCTAssertEqual(b.shape, [8, 13]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.4867754876613617, - accuracy: 0.009735509753227234) - XCTAssertEqual( - b.sum().item(Float.self), 50.62464904785156, - accuracy: 1.0124929809570313) - let result = matmul(a, b) - XCTAssertEqual(result.shape, [10, 13]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 2.032482862472534, - accuracy: 0.04064965724945068) - XCTAssertEqual( - result.sum().item(Float.self), 264.2227783203125, - accuracy: 5.28445556640625) - } - - func testMaximum() { - MLXRandom.seed(722) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.01978858932852745, - accuracy: -0.00039577178657054904) - XCTAssertEqual( - a.sum().item(Float.self), -0.2374630719423294, - accuracy: -0.004749261438846588) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.34027594327926636, - accuracy: 0.006805518865585327) - XCTAssertEqual( - b.sum().item(Float.self), 4.083311080932617, - accuracy: 0.08166622161865235) - let result = maximum(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6268421411514282, - accuracy: 0.012536842823028565) - XCTAssertEqual( - result.sum().item(Float.self), 7.522105693817139, - accuracy: 0.15044211387634276) - } - - func testMinimum() { - MLXRandom.seed(537) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.15274646878242493, - accuracy: 0.0030549293756484985) - XCTAssertEqual( - a.sum().item(Float.self), 1.8329575061798096, - accuracy: 0.03665915012359619) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.16062960028648376, - accuracy: 0.003212592005729675) - XCTAssertEqual( - b.sum().item(Float.self), 1.9275552034378052, - accuracy: 0.03855110406875611) - let result = minimum(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.4159923493862152, - accuracy: -0.008319846987724304) - XCTAssertEqual( - result.sum().item(Float.self), -4.991908073425293, - accuracy: -0.09983816146850587) - } - - func testMultiply() { - MLXRandom.seed(282) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.016050420701503754, - accuracy: -0.0003210084140300751) - XCTAssertEqual( - a.sum().item(Float.self), -0.19260503351688385, - accuracy: -0.003852100670337677) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.2504279613494873, - accuracy: 0.0050085592269897465) - XCTAssertEqual( - b.sum().item(Float.self), 3.0051355361938477, - accuracy: 0.060102710723876955) - let result = multiply(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.2712330222129822, - accuracy: -0.005424660444259643) - XCTAssertEqual( - result.sum().item(Float.self), -3.254796266555786, - accuracy: -0.06509592533111572) - } - - func testNotEqual() { - MLXRandom.seed(534) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.5144619941711426, - accuracy: -0.010289239883422853) - XCTAssertEqual( - a.sum().item(Float.self), -6.173543453216553, - accuracy: -0.12347086906433105) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), -0.15448352694511414, - accuracy: -0.0030896705389022827) - XCTAssertEqual( - b.sum().item(Float.self), -1.8538023233413696, - accuracy: -0.03707604646682739) - let result = notEqual(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .bool) - XCTAssertEqual(result.all().item(), true) - XCTAssertEqual(result.any().item(), true) - } - - func testRemainder() { - MLXRandom.seed(831) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.27154868841171265, - accuracy: 0.005430973768234253) - XCTAssertEqual( - a.sum().item(Float.self), 3.2585840225219727, - accuracy: 0.06517168045043946) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.1420227438211441, - accuracy: 0.0028404548764228823) - XCTAssertEqual( - b.sum().item(Float.self), 1.7042728662490845, - accuracy: 0.03408545732498169) - let result = remainder(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.007391604594886303, - accuracy: -0.00014783209189772606) - XCTAssertEqual( - result.sum().item(Float.self), -0.08869925141334534, - accuracy: -0.0017739850282669067) - } - - func testSubtract() { - MLXRandom.seed(241) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.1329968273639679, - accuracy: -0.002659936547279358) - XCTAssertEqual( - a.sum().item(Float.self), -1.5959618091583252, - accuracy: -0.031919236183166506) - let b = MLXRandom.normal([4, 3]) - XCTAssertEqual(b.shape, [4, 3]) - XCTAssertEqual(b.dtype, .float32) - XCTAssertEqual( - b.mean().item(Float.self), 0.03048303723335266, - accuracy: 0.0006096607446670532) - XCTAssertEqual( - b.sum().item(Float.self), 0.36579644680023193, - accuracy: 0.007315928936004639) - let result = subtract(a, b) - XCTAssertEqual(result.shape, [4, 3]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.16347989439964294, - accuracy: -0.003269597887992859) - XCTAssertEqual( - result.sum().item(Float.self), -1.9617586135864258, - accuracy: -0.03923517227172851) - } - - func testQuantize() { - MLXRandom.seed(869) - let w = MLXRandom.uniform(0.0 ..< 1.0, [32, 256]) - let (wq, scales, biases) = quantized(w, bits: 8) - XCTAssertEqual(wq.shape, [32, 64]) - XCTAssertEqual(wq.dtype, .uint32) - XCTAssertEqual( - wq.mean().item(Float.self), 732984.1875, - accuracy: 14659.68375) - XCTAssertEqual( - wq.sum().item(Float.self), 1_501_151_616, - accuracy: 30023031.62) - - XCTAssertEqual(scales.shape, [32, 4]) - XCTAssertEqual(scales.dtype, .float32) - XCTAssertEqual( - scales.mean().item(Float.self), -0.0037985900416970253, - accuracy: -7.597180083394051e-05) - XCTAssertEqual( - scales.sum().item(Float.self), -0.48621952533721924, - accuracy: -0.009724390506744385) - - if let biases { - XCTAssertEqual(biases.shape, [32, 4]) - XCTAssertEqual(biases.dtype, .float32) - XCTAssertEqual( - biases.mean().item(Float.self), 0.9862217307090759, - accuracy: 0.01972443461418152) - XCTAssertEqual( - biases.sum().item(Float.self), 126.23638153076172, - accuracy: 2.5247276306152346) - } else { - XCTFail("biases should not be nil") - } - - } - - func testFft_() { - MLXRandom.seed(220) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(r.shape, [100, 100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.49703681468963623, - accuracy: 0.009940736293792725) - XCTAssertEqual( - r.sum().item(Float.self), 4970.3681640625, - accuracy: 99.40736328125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(i.shape, [100, 100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49617284536361694, - accuracy: 0.00992345690727234) - XCTAssertEqual( - i.sum().item(Float.self), 4961.728515625, - accuracy: 99.2345703125) - let c = r + i.asImaginary() - let result = fft(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100, 100]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5186547636985779, - accuracy: 0.010373095273971558) - XCTAssertEqual( - resultReal.sum().item(Float.self), 5186.5478515625, - accuracy: 103.73095703125) - XCTAssertEqual(resultImaginary.shape, [100, 100]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.4938635528087616, - accuracy: 0.009877271056175233) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 4938.6357421875, - accuracy: 98.77271484375001) - } - - func testFft_1() { - MLXRandom.seed(916) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(r.shape, [100, 100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4961331784725189, - accuracy: 0.00992266356945038) - XCTAssertEqual( - r.sum().item(Float.self), 4961.33203125, - accuracy: 99.226640625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(i.shape, [100, 100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49821561574935913, - accuracy: 0.009964312314987183) - XCTAssertEqual( - i.sum().item(Float.self), 4982.15625, - accuracy: 99.643125) - let c = r + i.asImaginary() - let result = fft(c, n: 80, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100, 80]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.4616168737411499, - accuracy: 0.009232337474822999) - XCTAssertEqual( - resultReal.sum().item(Float.self), 3692.934814453125, - accuracy: 73.85869628906251) - XCTAssertEqual(resultImaginary.shape, [100, 80]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.4507213234901428, - accuracy: 0.009014426469802857) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 3605.7705078125, - accuracy: 72.11541015625001) - } - - func testFft_2() { - MLXRandom.seed(695) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(r.shape, [100, 100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.497734934091568, - accuracy: 0.00995469868183136) - XCTAssertEqual( - r.sum().item(Float.self), 4977.349609375, - accuracy: 99.5469921875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(i.shape, [100, 100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5017335414886475, - accuracy: 0.010034670829772949) - XCTAssertEqual( - i.sum().item(Float.self), 5017.33544921875, - accuracy: 100.346708984375) - let c = r + i.asImaginary() - let result = fft(c, n: 120, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100, 120]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5616999268531799, - accuracy: 0.011233998537063599) - XCTAssertEqual( - resultReal.sum().item(Float.self), 6740.3994140625, - accuracy: 134.80798828125) - XCTAssertEqual(resultImaginary.shape, [100, 120]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.5033472180366516, - accuracy: 0.010066944360733033) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 6040.1669921875, - accuracy: 120.80333984375) - } - - func testFft_3() { - MLXRandom.seed(603) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(r.shape, [100, 100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4938477873802185, - accuracy: 0.00987695574760437) - XCTAssertEqual( - r.sum().item(Float.self), 4938.47802734375, - accuracy: 98.769560546875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100, 100]) - XCTAssertEqual(i.shape, [100, 100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5019298791885376, - accuracy: 0.010038597583770752) - XCTAssertEqual( - i.sum().item(Float.self), 5019.298828125, - accuracy: 100.3859765625) - let c = r + i.asImaginary() - let result = fft(c, axis: 0, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100, 100]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5286791324615479, - accuracy: 0.010573582649230957) - XCTAssertEqual( - resultReal.sum().item(Float.self), 5286.79150390625, - accuracy: 105.735830078125) - XCTAssertEqual(resultImaginary.shape, [100, 100]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.5054363012313843, - accuracy: 0.010108726024627685) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 5054.36328125, - accuracy: 101.087265625) - } - - func testIfft_() { - MLXRandom.seed(845) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4823684096336365, - accuracy: 0.00964736819267273) - XCTAssertEqual( - r.sum().item(Float.self), 48.23684310913086, - accuracy: 0.9647368621826172) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.47583281993865967, - accuracy: 0.009516656398773193) - XCTAssertEqual( - i.sum().item(Float.self), 47.583282470703125, - accuracy: 0.9516656494140625) - let c = r + i.asImaginary() - let result = ifft(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.006908764597028494, - accuracy: 0.00013817529194056987) - XCTAssertEqual( - resultReal.sum().item(Float.self), 0.6908764839172363, - accuracy: 0.013817529678344726) - XCTAssertEqual(resultImaginary.shape, [100]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.005689740646630526, - accuracy: 0.00011379481293261051) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 0.5689740777015686, - accuracy: 0.011379481554031371) - } - - func testIfft_1() { - MLXRandom.seed(972) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.46860700845718384, - accuracy: 0.009372140169143678) - XCTAssertEqual( - r.sum().item(Float.self), 46.86070251464844, - accuracy: 0.9372140502929688) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4761311113834381, - accuracy: 0.009522622227668762) - XCTAssertEqual( - i.sum().item(Float.self), 47.61311340332031, - accuracy: 0.9522622680664062) - let c = r + i.asImaginary() - let result = ifft(c, n: 80, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [80]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.005004124250262976, - accuracy: 0.00010008248500525951) - XCTAssertEqual( - resultReal.sum().item(Float.self), 0.40032994747161865, - accuracy: 0.008006598949432373) - XCTAssertEqual(resultImaginary.shape, [80]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.0003413451777305454, - accuracy: 6.826903554610908e-06) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 0.02730761468410492, - accuracy: 0.0005461522936820984) - } - - func testIfft_2() { - MLXRandom.seed(429) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4427134692668915, - accuracy: 0.008854269385337829) - XCTAssertEqual( - r.sum().item(Float.self), 44.27134704589844, - accuracy: 0.8854269409179688) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.46247541904449463, - accuracy: 0.009249508380889893) - XCTAssertEqual( - i.sum().item(Float.self), 46.24754333496094, - accuracy: 0.9249508666992188) - let c = r + i.asImaginary() - let result = ifft(c, n: 120, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [120]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.0033431013580411673, - accuracy: 6.686202716082335e-05) - XCTAssertEqual( - resultReal.sum().item(Float.self), 0.40117213129997253, - accuracy: 0.00802344262599945) - XCTAssertEqual(resultImaginary.shape, [120]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.0048831733874976635, - accuracy: 9.766346774995328e-05) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 0.5859807729721069, - accuracy: 0.01171961545944214) - } - - func testIfft_3() { - MLXRandom.seed(593) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4806605279445648, - accuracy: 0.009613210558891297) - XCTAssertEqual( - r.sum().item(Float.self), 48.06605529785156, - accuracy: 0.9613211059570312) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.48936721682548523, - accuracy: 0.009787344336509705) - XCTAssertEqual( - i.sum().item(Float.self), 48.93672180175781, - accuracy: 0.9787344360351563) - let c = r + i.asImaginary() - let result = ifft(c, axis: 0, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [100]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.009695236571133137, - accuracy: 0.00019390473142266273) - XCTAssertEqual( - resultReal.sum().item(Float.self), 0.9695236682891846, - accuracy: 0.019390473365783693) - XCTAssertEqual(resultImaginary.shape, [100]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.008965172804892063, - accuracy: 0.00017930345609784126) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 0.896517276763916, - accuracy: 0.017930345535278322) - } - - func testRfft_() { - MLXRandom.seed(281) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.512880265712738, - accuracy: 0.01025760531425476) - XCTAssertEqual( - r.sum().item(Float.self), 51.28802490234375, - accuracy: 1.025760498046875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5312841534614563, - accuracy: 0.010625683069229126) - XCTAssertEqual( - i.sum().item(Float.self), 53.12841796875, - accuracy: 1.062568359375) - let c = r + i.asImaginary() - let result = rfft(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [51]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 1.0831178426742554, - accuracy: 0.021662356853485106) - XCTAssertEqual( - resultReal.sum().item(Float.self), 55.23900604248047, - accuracy: 1.1047801208496093) - XCTAssertEqual(resultImaginary.shape, [51]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.31754177808761597, - accuracy: 0.006350835561752319) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 16.194629669189453, - accuracy: 0.32389259338378906) - } - - func testRfft_1() { - MLXRandom.seed(461) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4627799987792969, - accuracy: 0.009255599975585938) - XCTAssertEqual( - r.sum().item(Float.self), 46.27799987792969, - accuracy: 0.9255599975585938) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5184644460678101, - accuracy: 0.010369288921356202) - XCTAssertEqual( - i.sum().item(Float.self), 51.84644317626953, - accuracy: 1.0369288635253906) - let c = r + i.asImaginary() - let result = rfft(c, n: 80, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [41]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 1.3340529203414917, - accuracy: 0.026681058406829834) - XCTAssertEqual( - resultReal.sum().item(Float.self), 54.696170806884766, - accuracy: 1.0939234161376954) - XCTAssertEqual(resultImaginary.shape, [41]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.09739136695861816, - accuracy: 0.0019478273391723632) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 3.993046283721924, - accuracy: 0.07986092567443848) - } - - func testRfft_2() { - MLXRandom.seed(504) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.47210997343063354, - accuracy: 0.009442199468612671) - XCTAssertEqual( - r.sum().item(Float.self), 47.21099853515625, - accuracy: 0.944219970703125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4813127815723419, - accuracy: 0.009626255631446838) - XCTAssertEqual( - i.sum().item(Float.self), 48.13127899169922, - accuracy: 0.9626255798339844) - let c = r + i.asImaginary() - let result = rfft(c, n: 120, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [61]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.421972393989563, - accuracy: 0.00843944787979126) - XCTAssertEqual( - resultReal.sum().item(Float.self), 25.740318298339844, - accuracy: 0.5148063659667969) - XCTAssertEqual(resultImaginary.shape, [61]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), -0.24379387497901917, - accuracy: -0.004875877499580384) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), -14.871427536010742, - accuracy: -0.29742855072021485) - } - - func testRfft_3() { - MLXRandom.seed(676) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4711613357067108, - accuracy: 0.009423226714134217) - XCTAssertEqual( - r.sum().item(Float.self), 47.11613464355469, - accuracy: 0.9423226928710937) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.47868087887763977, - accuracy: 0.009573617577552795) - XCTAssertEqual( - i.sum().item(Float.self), 47.86808776855469, - accuracy: 0.9573617553710938) - let c = r + i.asImaginary() - let result = rfft(c, axis: 0, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [51]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5287001729011536, - accuracy: 0.010574003458023071) - XCTAssertEqual( - resultReal.sum().item(Float.self), 26.963706970214844, - accuracy: 0.5392741394042969) - XCTAssertEqual(resultImaginary.shape, [51]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.25333887338638306, - accuracy: 0.005066777467727661) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 12.920282363891602, - accuracy: 0.25840564727783205) - } - - func testIrfft_() { - MLXRandom.seed(656) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.476201593875885, - accuracy: 0.0095240318775177) - XCTAssertEqual( - r.sum().item(Float.self), 47.62015914916992, - accuracy: 0.9524031829833984) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5490368008613586, - accuracy: 0.010980736017227172) - XCTAssertEqual( - i.sum().item(Float.self), 54.90367889404297, - accuracy: 1.0980735778808595) - let c = r + i.asImaginary() - let result = irfft(c, stream: .cpu) - XCTAssertEqual(result.shape, [198]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0028988695703446865, - accuracy: 5.797739140689373e-05) - XCTAssertEqual( - result.sum().item(Float.self), 0.5739761590957642, - accuracy: 0.011479523181915283) - } - - func testIrfft_1() { - MLXRandom.seed(717) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5386900305747986, - accuracy: 0.010773800611495972) - XCTAssertEqual( - r.sum().item(Float.self), 53.86900329589844, - accuracy: 1.0773800659179689) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4603566527366638, - accuracy: 0.009207133054733276) - XCTAssertEqual( - i.sum().item(Float.self), 46.035667419433594, - accuracy: 0.9207133483886719) - let c = r + i.asImaginary() - let result = irfft(c, n: 80, stream: .cpu) - XCTAssertEqual(result.shape, [80]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0038530919700860977, - accuracy: 7.706183940172196e-05) - XCTAssertEqual( - result.sum().item(Float.self), 0.3082473576068878, - accuracy: 0.006164947152137757) - } - - func testIrfft_2() { - MLXRandom.seed(938) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5083585381507874, - accuracy: 0.010167170763015747) - XCTAssertEqual( - r.sum().item(Float.self), 50.835853576660156, - accuracy: 1.0167170715332032) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4999184012413025, - accuracy: 0.00999836802482605) - XCTAssertEqual( - i.sum().item(Float.self), 49.99184036254883, - accuracy: 0.9998368072509766) - let c = r + i.asImaginary() - let result = irfft(c, n: 120, stream: .cpu) - XCTAssertEqual(result.shape, [120]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0034326675813645124, - accuracy: 6.865335162729025e-05) - XCTAssertEqual( - result.sum().item(Float.self), 0.41192010045051575, - accuracy: 0.008238402009010316) - } - - func testIrfft_3() { - MLXRandom.seed(812) - let r = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(r.shape, [100]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4437080919742584, - accuracy: 0.008874161839485169) - XCTAssertEqual( - r.sum().item(Float.self), 44.370811462402344, - accuracy: 0.8874162292480469) - let i = MLXRandom.uniform(0.0 ..< 1.0, [100]) - XCTAssertEqual(i.shape, [100]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5141764879226685, - accuracy: 0.01028352975845337) - XCTAssertEqual( - i.sum().item(Float.self), 51.41764831542969, - accuracy: 1.0283529663085937) - let c = r + i.asImaginary() - let result = irfft(c, axis: 0, stream: .cpu) - XCTAssertEqual(result.shape, [198]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.004781102295964956, - accuracy: 9.562204591929913e-05) - XCTAssertEqual( - result.sum().item(Float.self), 0.9466582536697388, - accuracy: 0.018933165073394775) - } - - func testFft2_() { - MLXRandom.seed(365) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5127536058425903, - accuracy: 0.010255072116851807) - XCTAssertEqual( - r.sum().item(Float.self), 262.52984619140625, - accuracy: 5.250596923828125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5092339515686035, - accuracy: 0.01018467903137207) - XCTAssertEqual( - i.sum().item(Float.self), 260.727783203125, - accuracy: 5.2145556640625) - let c = r + i.asImaginary() - let result = fft2(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.3378489017486572, - accuracy: 0.0067569780349731445) - XCTAssertEqual( - resultReal.sum().item(Float.self), 172.9786376953125, - accuracy: 3.45957275390625) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.4062993824481964, - accuracy: 0.008125987648963929) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 208.02528381347656, - accuracy: 4.160505676269532) - } - - func testFft2_1() { - MLXRandom.seed(84) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5123671293258667, - accuracy: 0.010247342586517334) - XCTAssertEqual( - r.sum().item(Float.self), 262.33197021484375, - accuracy: 5.246639404296875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.502106249332428, - accuracy: 0.01004212498664856) - XCTAssertEqual( - i.sum().item(Float.self), 257.0783996582031, - accuracy: 5.141567993164062) - let c = r + i.asImaginary() - let result = fft2(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 4]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5198299884796143, - accuracy: 0.010396599769592285) - XCTAssertEqual( - resultReal.sum().item(Float.self), 49.9036750793457, - accuracy: 0.9980735015869141) - XCTAssertEqual(resultImaginary.shape, [8, 3, 4]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.5818836688995361, - accuracy: 0.011637673377990723) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 55.8608283996582, - accuracy: 1.117216567993164) - } - - func testFft2_2() { - MLXRandom.seed(332) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4954897463321686, - accuracy: 0.009909794926643371) - XCTAssertEqual( - r.sum().item(Float.self), 253.6907501220703, - accuracy: 5.073815002441406) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5005261898040771, - accuracy: 0.010010523796081543) - XCTAssertEqual( - i.sum().item(Float.self), 256.2694091796875, - accuracy: 5.12538818359375) - let c = r + i.asImaginary() - let result = fft2(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.3796372711658478, - accuracy: 0.007592745423316956) - XCTAssertEqual( - resultReal.sum().item(Float.self), 194.37428283691406, - accuracy: 3.8874856567382814) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.6848205327987671, - accuracy: 0.013696410655975343) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 350.62811279296875, - accuracy: 7.012562255859375) - } - - func testFft2_3() { - MLXRandom.seed(627) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.48098403215408325, - accuracy: 0.009619680643081665) - XCTAssertEqual( - r.sum().item(Float.self), 246.26382446289062, - accuracy: 4.925276489257812) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49524593353271484, - accuracy: 0.009904918670654296) - XCTAssertEqual( - i.sum().item(Float.self), 253.56591796875, - accuracy: 5.071318359375) - let c = r + i.asImaginary() - let result = fft2(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 5, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.39667367935180664, - accuracy: 0.007933473587036133) - XCTAssertEqual( - resultReal.sum().item(Float.self), 158.6694793701172, - accuracy: 3.173389587402344) - XCTAssertEqual(resultImaginary.shape, [8, 5, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.46095579862594604, - accuracy: 0.009219115972518921) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 184.38232421875, - accuracy: 3.687646484375) - } - - func testIfft2_() { - MLXRandom.seed(118) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4939275085926056, - accuracy: 0.009878550171852112) - XCTAssertEqual( - r.sum().item(Float.self), 252.89088439941406, - accuracy: 5.0578176879882815) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5031070709228516, - accuracy: 0.010062141418457031) - XCTAssertEqual( - i.sum().item(Float.self), 257.5908203125, - accuracy: 5.15181640625) - let c = r + i.asImaginary() - let result = ifft2(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.009310859255492687, - accuracy: 0.00018621718510985374) - XCTAssertEqual( - resultReal.sum().item(Float.self), 4.767159938812256, - accuracy: 0.09534319877624511) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.006652233190834522, - accuracy: 0.00013304466381669044) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 3.4059433937072754, - accuracy: 0.0681188678741455) - } - - func testIfft2_1() { - MLXRandom.seed(498) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5009815096855164, - accuracy: 0.010019630193710327) - XCTAssertEqual( - r.sum().item(Float.self), 256.5025329589844, - accuracy: 5.1300506591796875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49243441224098206, - accuracy: 0.009848688244819642) - XCTAssertEqual( - i.sum().item(Float.self), 252.1264190673828, - accuracy: 5.0425283813476565) - let c = r + i.asImaginary() - let result = ifft2(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 4]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.055086322128772736, - accuracy: 0.0011017264425754547) - XCTAssertEqual( - resultReal.sum().item(Float.self), 5.2882866859436035, - accuracy: 0.10576573371887207) - XCTAssertEqual(resultImaginary.shape, [8, 3, 4]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.052069030702114105, - accuracy: 0.001041380614042282) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 4.998626708984375, - accuracy: 0.0999725341796875) - } - - func testIfft2_2() { - MLXRandom.seed(601) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5043255686759949, - accuracy: 0.010086511373519898) - XCTAssertEqual( - r.sum().item(Float.self), 258.2146911621094, - accuracy: 5.164293823242188) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5282406806945801, - accuracy: 0.010564813613891602) - XCTAssertEqual( - i.sum().item(Float.self), 270.459228515625, - accuracy: 5.4091845703125) - let c = r + i.asImaginary() - let result = ifft2(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.005900253541767597, - accuracy: 0.00011800507083535194) - XCTAssertEqual( - resultReal.sum().item(Float.self), 3.0209298133850098, - accuracy: 0.06041859626770019) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.009091435000300407, - accuracy: 0.00018182870000600816) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 4.654814720153809, - accuracy: 0.09309629440307618) - } - - func testIfft2_3() { - MLXRandom.seed(645) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5046442151069641, - accuracy: 0.010092884302139282) - XCTAssertEqual( - r.sum().item(Float.self), 258.3778381347656, - accuracy: 5.1675567626953125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4897768199443817, - accuracy: 0.009795536398887635) - XCTAssertEqual( - i.sum().item(Float.self), 250.76573181152344, - accuracy: 5.015314636230469) - let c = r + i.asImaginary() - let result = ifft2(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 5, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.011821601539850235, - accuracy: 0.0002364320307970047) - XCTAssertEqual( - resultReal.sum().item(Float.self), 4.728640556335449, - accuracy: 0.09457281112670898) - XCTAssertEqual(resultImaginary.shape, [8, 5, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.007806990761309862, - accuracy: 0.00015613981522619725) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 3.1227962970733643, - accuracy: 0.06245592594146729) - } - - func testFftn_() { - MLXRandom.seed(343) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4879269003868103, - accuracy: 0.009758538007736206) - XCTAssertEqual( - r.sum().item(Float.self), 249.81857299804688, - accuracy: 4.996371459960938) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49864500761032104, - accuracy: 0.009972900152206421) - XCTAssertEqual( - i.sum().item(Float.self), 255.30624389648438, - accuracy: 5.1061248779296875) - let c = r + i.asImaginary() - let result = fftn(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.8467151522636414, - accuracy: 0.01693430304527283) - XCTAssertEqual( - resultReal.sum().item(Float.self), 433.5181579589844, - accuracy: 8.670363159179688) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.28999343514442444, - accuracy: 0.005799868702888489) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 148.4766387939453, - accuracy: 2.9695327758789065) - } - - func testFftn_1() { - MLXRandom.seed(865) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.506011962890625, - accuracy: 0.0101202392578125) - XCTAssertEqual( - r.sum().item(Float.self), 259.078125, - accuracy: 5.1815625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4955616891384125, - accuracy: 0.00991123378276825) - XCTAssertEqual( - i.sum().item(Float.self), 253.7275848388672, - accuracy: 5.074551696777344) - let c = r + i.asImaginary() - let result = fftn(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 4]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.3270348012447357, - accuracy: 0.006540696024894714) - XCTAssertEqual( - resultReal.sum().item(Float.self), 31.395339965820312, - accuracy: 0.6279067993164062) - XCTAssertEqual(resultImaginary.shape, [8, 3, 4]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.4940086305141449, - accuracy: 0.009880172610282898) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 47.424827575683594, - accuracy: 0.9484965515136718) - } - - func testFftn_2() { - MLXRandom.seed(194) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.48457071185112, - accuracy: 0.0096914142370224) - XCTAssertEqual( - r.sum().item(Float.self), 248.10020446777344, - accuracy: 4.962004089355469) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4813538193702698, - accuracy: 0.009627076387405396) - XCTAssertEqual( - i.sum().item(Float.self), 246.45315551757812, - accuracy: 4.929063110351563) - let c = r + i.asImaginary() - let result = fftn(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.6415871381759644, - accuracy: 0.012831742763519288) - XCTAssertEqual( - resultReal.sum().item(Float.self), 328.49261474609375, - accuracy: 6.569852294921875) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.5985268950462341, - accuracy: 0.011970537900924684) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 306.4457702636719, - accuracy: 6.128915405273438) - } - - func testFftn_3() { - MLXRandom.seed(248) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.48458826541900635, - accuracy: 0.009691765308380127) - XCTAssertEqual( - r.sum().item(Float.self), 248.10919189453125, - accuracy: 4.962183837890625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5036082863807678, - accuracy: 0.010072165727615356) - XCTAssertEqual( - i.sum().item(Float.self), 257.8474426269531, - accuracy: 5.156948852539062) - let c = r + i.asImaginary() - let result = fftn(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 5, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.44054391980171204, - accuracy: 0.008810878396034241) - XCTAssertEqual( - resultReal.sum().item(Float.self), 176.2175750732422, - accuracy: 3.5243515014648437) - XCTAssertEqual(resultImaginary.shape, [8, 5, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.5384612679481506, - accuracy: 0.010769225358963012) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 215.384521484375, - accuracy: 4.3076904296875) - } - - func testIfftn_() { - MLXRandom.seed(16) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4910537600517273, - accuracy: 0.009821075201034547) - XCTAssertEqual( - r.sum().item(Float.self), 251.41952514648438, - accuracy: 5.028390502929688) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4993511438369751, - accuracy: 0.009987022876739502) - XCTAssertEqual( - i.sum().item(Float.self), 255.66778564453125, - accuracy: 5.113355712890625) - let c = r + i.asImaginary() - let result = ifftn(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.000745800556614995, - accuracy: 1.49160111322999e-05) - XCTAssertEqual( - resultReal.sum().item(Float.self), 0.38184988498687744, - accuracy: 0.007636997699737549) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.00013762176968157291, - accuracy: 2.7524353936314585e-06) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 0.07046234607696533, - accuracy: 0.0014092469215393067) - } - - func testIfftn_1() { - MLXRandom.seed(749) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5013963580131531, - accuracy: 0.010027927160263062) - XCTAssertEqual( - r.sum().item(Float.self), 256.7149353027344, - accuracy: 5.134298706054688) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49108564853668213, - accuracy: 0.009821712970733643) - XCTAssertEqual( - i.sum().item(Float.self), 251.43585205078125, - accuracy: 5.028717041015625) - let c = r + i.asImaginary() - let result = ifftn(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 4]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.029997922480106354, - accuracy: 0.0005999584496021271) - XCTAssertEqual( - resultReal.sum().item(Float.self), 2.87980055809021, - accuracy: 0.0575960111618042) - XCTAssertEqual(resultImaginary.shape, [8, 3, 4]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.04324401170015335, - accuracy: 0.0008648802340030671) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 4.151424884796143, - accuracy: 0.08302849769592285) - } - - func testIfftn_2() { - MLXRandom.seed(277) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5160338878631592, - accuracy: 0.010320677757263183) - XCTAssertEqual( - r.sum().item(Float.self), 264.2093505859375, - accuracy: 5.28418701171875) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.48932191729545593, - accuracy: 0.009786438345909119) - XCTAssertEqual( - i.sum().item(Float.self), 250.53282165527344, - accuracy: 5.010656433105469) - let c = r + i.asImaginary() - let result = ifftn(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 8]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.007633390370756388, - accuracy: 0.00015266780741512775) - XCTAssertEqual( - resultReal.sum().item(Float.self), 3.9082958698272705, - accuracy: 0.07816591739654541) - XCTAssertEqual(resultImaginary.shape, [8, 8, 8]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.010958677157759666, - accuracy: 0.00021917354315519335) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 5.610842704772949, - accuracy: 0.11221685409545899) - } - - func testIfftn_3() { - MLXRandom.seed(119) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5159510970115662, - accuracy: 0.010319021940231323) - XCTAssertEqual( - r.sum().item(Float.self), 264.1669616699219, - accuracy: 5.283339233398437) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5054599046707153, - accuracy: 0.010109198093414307) - XCTAssertEqual( - i.sum().item(Float.self), 258.79547119140625, - accuracy: 5.175909423828125) - let c = r + i.asImaginary() - let result = ifftn(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 5, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.008541149087250233, - accuracy: 0.00017082298174500465) - XCTAssertEqual( - resultReal.sum().item(Float.self), 3.416459560394287, - accuracy: 0.06832919120788575) - XCTAssertEqual(resultImaginary.shape, [8, 5, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.013507889583706856, - accuracy: 0.00027015779167413714) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 5.40315580368042, - accuracy: 0.1080631160736084) - } - - func testRfft2_() { - MLXRandom.seed(722) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4949501156806946, - accuracy: 0.009899002313613892) - XCTAssertEqual( - r.sum().item(Float.self), 253.41445922851562, - accuracy: 5.068289184570313) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4984145760536194, - accuracy: 0.009968291521072387) - XCTAssertEqual( - i.sum().item(Float.self), 255.18826293945312, - accuracy: 5.103765258789062) - let c = r + i.asImaginary() - let result = rfft2(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 5]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.8082265853881836, - accuracy: 0.016164531707763673) - XCTAssertEqual( - resultReal.sum().item(Float.self), 258.63250732421875, - accuracy: 5.172650146484375) - XCTAssertEqual(resultImaginary.shape, [8, 8, 5]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.03572067618370056, - accuracy: 0.0007144135236740113) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 11.43061637878418, - accuracy: 0.2286123275756836) - } - - func testRfft2_1() { - MLXRandom.seed(225) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5058765411376953, - accuracy: 0.010117530822753906) - XCTAssertEqual( - r.sum().item(Float.self), 259.0087890625, - accuracy: 5.18017578125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4878406226634979, - accuracy: 0.009756812453269958) - XCTAssertEqual( - i.sum().item(Float.self), 249.77439880371094, - accuracy: 4.995487976074219) - let c = r + i.asImaginary() - let result = rfft2(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 3]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.7214547991752625, - accuracy: 0.014429095983505249) - XCTAssertEqual( - resultReal.sum().item(Float.self), 51.94474411010742, - accuracy: 1.0388948822021484) - XCTAssertEqual(resultImaginary.shape, [8, 3, 3]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), -0.03948277235031128, - accuracy: -0.0007896554470062257) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), -2.842759609222412, - accuracy: -0.05685519218444824) - } - - func testRfft2_2() { - MLXRandom.seed(380) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4951004981994629, - accuracy: 0.009902009963989258) - XCTAssertEqual( - r.sum().item(Float.self), 253.491455078125, - accuracy: 5.0698291015625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4919356107711792, - accuracy: 0.009838712215423585) - XCTAssertEqual( - i.sum().item(Float.self), 251.87103271484375, - accuracy: 5.037420654296875) - let c = r + i.asImaginary() - let result = rfft2(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 5]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.86641925573349, - accuracy: 0.0173283851146698) - XCTAssertEqual( - resultReal.sum().item(Float.self), 277.254150390625, - accuracy: 5.5450830078125) - XCTAssertEqual(resultImaginary.shape, [8, 8, 5]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.006711974740028381, - accuracy: 0.00013423949480056762) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 2.147831916809082, - accuracy: 0.04295663833618164) - } - - func testRfft2_3() { - MLXRandom.seed(813) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.4799543023109436, - accuracy: 0.009599086046218872) - XCTAssertEqual( - r.sum().item(Float.self), 245.73660278320312, - accuracy: 4.914732055664063) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4987284541130066, - accuracy: 0.009974569082260132) - XCTAssertEqual( - i.sum().item(Float.self), 255.34896850585938, - accuracy: 5.1069793701171875) - let c = r + i.asImaginary() - let result = rfft2(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.7862035036087036, - accuracy: 0.015724070072174072) - XCTAssertEqual( - resultReal.sum().item(Float.self), 188.68882751464844, - accuracy: 3.7737765502929688) - XCTAssertEqual(resultImaginary.shape, [8, 3, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.11010777205228806, - accuracy: 0.002202155441045761) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 26.42586326599121, - accuracy: 0.5285172653198242) - } - - func testIrfft2_() { - MLXRandom.seed(174) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.514319121837616, - accuracy: 0.01028638243675232) - XCTAssertEqual( - r.sum().item(Float.self), 263.3313903808594, - accuracy: 5.266627807617188) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5133156776428223, - accuracy: 0.010266313552856446) - XCTAssertEqual( - i.sum().item(Float.self), 262.817626953125, - accuracy: 5.2563525390625) - let c = r + i.asImaginary() - let result = irfft2(c, stream: .cpu) - XCTAssertEqual(result.shape, [8, 8, 14]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.005966233555227518, - accuracy: 0.00011932467110455036) - XCTAssertEqual( - result.sum().item(Float.self), 5.345745086669922, - accuracy: 0.10691490173339845) - } - - func testIrfft2_1() { - MLXRandom.seed(340) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5179769992828369, - accuracy: 0.010359539985656738) - XCTAssertEqual( - r.sum().item(Float.self), 265.2042236328125, - accuracy: 5.30408447265625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5126595497131348, - accuracy: 0.010253190994262695) - XCTAssertEqual( - i.sum().item(Float.self), 262.481689453125, - accuracy: 5.2496337890625) - let c = r + i.asImaginary() - let result = irfft2(c, s: [3, 4], stream: .cpu) - XCTAssertEqual(result.shape, [8, 3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.047055989503860474, - accuracy: 0.0009411197900772095) - XCTAssertEqual( - result.sum().item(Float.self), 4.5173749923706055, - accuracy: 0.09034749984741211) - } - - func testIrfft2_2() { - MLXRandom.seed(436) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5037882328033447, - accuracy: 0.010075764656066894) - XCTAssertEqual( - r.sum().item(Float.self), 257.9395751953125, - accuracy: 5.15879150390625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5085217356681824, - accuracy: 0.010170434713363648) - XCTAssertEqual( - i.sum().item(Float.self), 260.3631286621094, - accuracy: 5.207262573242188) - let c = r + i.asImaginary() - let result = irfft2(c, axes: [0, 2], stream: .cpu) - XCTAssertEqual(result.shape, [8, 8, 14]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.004838266409933567, - accuracy: 9.676532819867134e-05) - XCTAssertEqual( - result.sum().item(Float.self), 4.335086345672607, - accuracy: 0.08670172691345215) - } - - func testIrfft2_3() { - MLXRandom.seed(835) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.507753849029541, - accuracy: 0.010155076980590821) - XCTAssertEqual( - r.sum().item(Float.self), 259.969970703125, - accuracy: 5.1993994140625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49466848373413086, - accuracy: 0.009893369674682618) - XCTAssertEqual( - i.sum().item(Float.self), 253.270263671875, - accuracy: 5.0654052734375) - let c = r + i.asImaginary() - let result = irfft2(c, s: [10, 5], axes: [2, 1], stream: .cpu) - XCTAssertEqual(result.shape, [8, 5, 10]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.013026298955082893, - accuracy: 0.0002605259791016579) - XCTAssertEqual( - result.sum().item(Float.self), 5.210519790649414, - accuracy: 0.10421039581298829) - } - - func testRfftn_() { - MLXRandom.seed(63) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5128694772720337, - accuracy: 0.010257389545440674) - XCTAssertEqual( - r.sum().item(Float.self), 262.58917236328125, - accuracy: 5.251783447265625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4861868619918823, - accuracy: 0.009723737239837646) - XCTAssertEqual( - i.sum().item(Float.self), 248.92767333984375, - accuracy: 4.978553466796875) - let c = r + i.asImaginary() - let result = rfftn(c, stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 5]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5379745960235596, - accuracy: 0.010759491920471192) - XCTAssertEqual( - resultReal.sum().item(Float.self), 172.15187072753906, - accuracy: 3.443037414550781) - XCTAssertEqual(resultImaginary.shape, [8, 8, 5]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.12080971151590347, - accuracy: 0.0024161942303180697) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 38.65910720825195, - accuracy: 0.773182144165039) - } - - func testRfftn_1() { - MLXRandom.seed(103) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5262448787689209, - accuracy: 0.010524897575378419) - XCTAssertEqual( - r.sum().item(Float.self), 269.4373779296875, - accuracy: 5.38874755859375) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4829203188419342, - accuracy: 0.009658406376838685) - XCTAssertEqual( - i.sum().item(Float.self), 247.2552032470703, - accuracy: 4.945104064941407) - let c = r + i.asImaginary() - let result = rfftn(c, s: [3, 4], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 3]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.5502752065658569, - accuracy: 0.011005504131317139) - XCTAssertEqual( - resultReal.sum().item(Float.self), 39.619815826416016, - accuracy: 0.7923963165283203) - XCTAssertEqual(resultImaginary.shape, [8, 3, 3]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), 0.025733524933457375, - accuracy: 0.0005146704986691475) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), 1.852813720703125, - accuracy: 0.0370562744140625) - } - - func testRfftn_2() { - MLXRandom.seed(801) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.50965815782547, - accuracy: 0.010193163156509399) - XCTAssertEqual( - r.sum().item(Float.self), 260.9449768066406, - accuracy: 5.218899536132812) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5044752955436707, - accuracy: 0.010089505910873414) - XCTAssertEqual( - i.sum().item(Float.self), 258.2913513183594, - accuracy: 5.165827026367188) - let c = r + i.asImaginary() - let result = rfftn(c, axes: [0, 2], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 8, 5]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.8939148187637329, - accuracy: 0.017878296375274657) - XCTAssertEqual( - resultReal.sum().item(Float.self), 286.052734375, - accuracy: 5.7210546875) - XCTAssertEqual(resultImaginary.shape, [8, 8, 5]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), -0.05308229848742485, - accuracy: -0.001061645969748497) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), -16.98633575439453, - accuracy: -0.33972671508789065) - } - - func testRfftn_3() { - MLXRandom.seed(149) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5064899921417236, - accuracy: 0.010129799842834472) - XCTAssertEqual( - r.sum().item(Float.self), 259.3228759765625, - accuracy: 5.18645751953125) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.4949530363082886, - accuracy: 0.009899060726165771) - XCTAssertEqual( - i.sum().item(Float.self), 253.41595458984375, - accuracy: 5.068319091796875) - let c = r + i.asImaginary() - let result = rfftn(c, s: [10, 5], axes: [2, 1], stream: .cpu) - let resultReal = result.realPart() - let resultImaginary = result.imaginaryPart() - XCTAssertEqual(resultReal.shape, [8, 3, 10]) - XCTAssertEqual(resultReal.dtype, .float32) - XCTAssertEqual( - resultReal.mean().item(Float.self), 0.7501934170722961, - accuracy: 0.015003868341445924) - XCTAssertEqual( - resultReal.sum().item(Float.self), 180.04641723632812, - accuracy: 3.6009283447265625) - XCTAssertEqual(resultImaginary.shape, [8, 3, 10]) - XCTAssertEqual(resultImaginary.dtype, .float32) - XCTAssertEqual( - resultImaginary.mean().item(Float.self), -0.039641283452510834, - accuracy: -0.0007928256690502167) - XCTAssertEqual( - resultImaginary.sum().item(Float.self), -9.513907432556152, - accuracy: -0.19027814865112305) - } - - func testIrfftn_() { - MLXRandom.seed(875) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5080322623252869, - accuracy: 0.010160645246505737) - XCTAssertEqual( - r.sum().item(Float.self), 260.1125183105469, - accuracy: 5.202250366210937) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5002621412277222, - accuracy: 0.010005242824554443) - XCTAssertEqual( - i.sum().item(Float.self), 256.13421630859375, - accuracy: 5.122684326171875) - let c = r + i.asImaginary() - let result = irfftn(c, stream: .cpu) - XCTAssertEqual(result.shape, [8, 8, 14]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 9.236017649527639e-05, - accuracy: 1.8472035299055278e-06) - XCTAssertEqual( - result.sum().item(Float.self), 0.0827547162771225, - accuracy: 0.00165509432554245) - } - - func testIrfftn_1() { - MLXRandom.seed(714) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.49393296241760254, - accuracy: 0.00987865924835205) - XCTAssertEqual( - r.sum().item(Float.self), 252.8936767578125, - accuracy: 5.05787353515625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5138311386108398, - accuracy: 0.010276622772216797) - XCTAssertEqual( - i.sum().item(Float.self), 263.08154296875, - accuracy: 5.261630859375) - let c = r + i.asImaginary() - let result = irfftn(c, s: [3, 4], stream: .cpu) - XCTAssertEqual(result.shape, [8, 3, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.041039690375328064, - accuracy: 0.0008207938075065613) - XCTAssertEqual( - result.sum().item(Float.self), 3.939810276031494, - accuracy: 0.07879620552062988) - } - - func testIrfftn_2() { - MLXRandom.seed(224) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5111634731292725, - accuracy: 0.01022326946258545) - XCTAssertEqual( - r.sum().item(Float.self), 261.7156982421875, - accuracy: 5.23431396484375) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.49477437138557434, - accuracy: 0.009895487427711487) - XCTAssertEqual( - i.sum().item(Float.self), 253.32447814941406, - accuracy: 5.066489562988282) - let c = r + i.asImaginary() - let result = irfftn(c, axes: [0, 2], stream: .cpu) - XCTAssertEqual(result.shape, [8, 8, 14]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.00553929852321744, - accuracy: 0.00011078597046434879) - XCTAssertEqual( - result.sum().item(Float.self), 4.9632110595703125, - accuracy: 0.09926422119140625) - } - - func testIrfftn_3() { - MLXRandom.seed(46) - let r = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(r.shape, [8, 8, 8]) - XCTAssertEqual(r.dtype, .float32) - XCTAssertEqual( - r.mean().item(Float.self), 0.5066219568252563, - accuracy: 0.010132439136505127) - XCTAssertEqual( - r.sum().item(Float.self), 259.39044189453125, - accuracy: 5.187808837890625) - let i = MLXRandom.uniform(0.0 ..< 1.0, [8, 8, 8]) - XCTAssertEqual(i.shape, [8, 8, 8]) - XCTAssertEqual(i.dtype, .float32) - XCTAssertEqual( - i.mean().item(Float.self), 0.5006622076034546, - accuracy: 0.010013244152069093) - XCTAssertEqual( - i.sum().item(Float.self), 256.33905029296875, - accuracy: 5.1267810058593755) - let c = r + i.asImaginary() - let result = irfftn(c, s: [10, 5], axes: [2, 1], stream: .cpu) - XCTAssertEqual(result.shape, [8, 5, 10]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.010128638707101345, - accuracy: 0.00020257277414202692) - XCTAssertEqual( - result.sum().item(Float.self), 4.051455497741699, - accuracy: 0.08102910995483399) - } - - func testSGD() { - MLXRandom.seed(836) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.07495629787445068, - accuracy: 0.0014991259574890137) - XCTAssertEqual( - a.sum().item(Float.self), 0.8994755148887634, - accuracy: 0.01798951029777527) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.09326401352882385, - accuracy: 0.001865280270576477) - XCTAssertEqual( - aGrad.sum().item(Float.self), 1.1191681623458862, - accuracy: 0.022383363246917726) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = SGD(learningRate: 0.1).apply(gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.06562991440296173, - accuracy: 0.0013125982880592346) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 0.787558913230896, - accuracy: 0.01575117826461792) - } - - func testSGD1() { - MLXRandom.seed(587) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.10747849941253662, - accuracy: -0.0021495699882507326) - XCTAssertEqual( - a.sum().item(Float.self), -1.2897419929504395, - accuracy: -0.02579483985900879) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), -0.32190611958503723, - accuracy: -0.006438122391700745) - XCTAssertEqual( - aGrad.sum().item(Float.self), -3.8628733158111572, - accuracy: -0.07725746631622314) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = SGD(learningRate: 0.1, momentum: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.07528789341449738, - accuracy: -0.0015057578682899475) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -0.9034546613693237, - accuracy: -0.018069093227386476) - } - - func testSGD2() { - MLXRandom.seed(649) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.7766839861869812, - accuracy: 0.015533679723739624) - XCTAssertEqual( - a.sum().item(Float.self), 9.320207595825195, - accuracy: 0.18640415191650392) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.06737084686756134, - accuracy: 0.0013474169373512267) - XCTAssertEqual( - aGrad.sum().item(Float.self), 0.8084501028060913, - accuracy: 0.016169002056121828) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = SGD(learningRate: 0.1, momentum: 0.1, dampening: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.7706205248832703, - accuracy: 0.015412410497665405) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 9.247446060180664, - accuracy: 0.18494892120361328) - } - - func testRMSprop() { - MLXRandom.seed(931) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.22450898587703705, - accuracy: -0.004490179717540741) - XCTAssertEqual( - a.sum().item(Float.self), -2.6941077709198, - accuracy: -0.053882155418396) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.24865484237670898, - accuracy: 0.00497309684753418) - XCTAssertEqual( - aGrad.sum().item(Float.self), 2.983858108520508, - accuracy: 0.05967716217041016) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = RMSprop(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.22450922429561615, - accuracy: -0.004490184485912323) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -2.694110631942749, - accuracy: -0.05388221263885498) - } - - func testAdagrad() { - MLXRandom.seed(958) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.04584333300590515, - accuracy: -0.0009168666601181031) - XCTAssertEqual( - a.sum().item(Float.self), -0.5501199960708618, - accuracy: -0.011002399921417237) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.23250393569469452, - accuracy: 0.00465007871389389) - XCTAssertEqual( - aGrad.sum().item(Float.self), 2.7900471687316895, - accuracy: 0.05580094337463379) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = AdaGrad(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.06250998377799988, - accuracy: -0.0012501996755599975) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -0.7501198053359985, - accuracy: -0.01500239610671997) - } - - func testAdaDelta() { - MLXRandom.seed(547) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3483370244503021, - accuracy: -0.006966740489006043) - XCTAssertEqual( - a.sum().item(Float.self), -4.180044174194336, - accuracy: -0.08360088348388672) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.5226783752441406, - accuracy: 0.010453567504882813) - XCTAssertEqual( - aGrad.sum().item(Float.self), 6.272140026092529, - accuracy: 0.12544280052185058) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = AdaDelta(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.348442405462265, - accuracy: -0.0069688481092453) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -4.181308746337891, - accuracy: -0.08362617492675782) - } - - func testAdam() { - MLXRandom.seed(616) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.1122930571436882, - accuracy: 0.002245861142873764) - XCTAssertEqual( - a.sum().item(Float.self), 1.347516655921936, - accuracy: 0.02695033311843872) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.305597722530365, - accuracy: 0.0061119544506073) - XCTAssertEqual( - aGrad.sum().item(Float.self), 3.66717267036438, - accuracy: 0.0733434534072876) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Adam(learningRate: 0.1).apply(gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.11229278147220612, - accuracy: 0.0022458556294441224) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 1.3475133180618286, - accuracy: 0.026950266361236572) - } - - func testAdamW() { - MLXRandom.seed(696) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3633918762207031, - accuracy: -0.007267837524414063) - XCTAssertEqual( - a.sum().item(Float.self), -4.3607025146484375, - accuracy: -0.08721405029296875) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.22175447642803192, - accuracy: 0.0044350895285606385) - XCTAssertEqual( - aGrad.sum().item(Float.self), 2.6610536575317383, - accuracy: 0.05322107315063477) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = AdamW(learningRate: 0.1).apply(gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.4684376120567322, - accuracy: -0.009368752241134645) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -5.621251106262207, - accuracy: -0.11242502212524415) - } - - func testAdamax() { - MLXRandom.seed(75) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.3039236068725586, - accuracy: -0.006078472137451172) - XCTAssertEqual( - a.sum().item(Float.self), -3.647083282470703, - accuracy: -0.07294166564941407) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), -0.24271723628044128, - accuracy: -0.004854344725608826) - XCTAssertEqual( - aGrad.sum().item(Float.self), -2.912606716156006, - accuracy: -0.05825213432312012) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Adamax(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.3039236068725586, - accuracy: -0.006078472137451172) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -3.647083282470703, - accuracy: -0.07294166564941407) - } - - func testLion() { - MLXRandom.seed(27) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.1776922345161438, - accuracy: 0.003553844690322876) - XCTAssertEqual( - a.sum().item(Float.self), 2.1323068141937256, - accuracy: 0.042646136283874515) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), -0.02118723653256893, - accuracy: -0.00042374473065137863) - XCTAssertEqual( - aGrad.sum().item(Float.self), -0.2542468309402466, - accuracy: -0.005084936618804932) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Lion(learningRate: 0.1).apply(gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.21102556586265564, - accuracy: 0.004220511317253113) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 2.532306671142578, - accuracy: 0.05064613342285156) - } - - func testLion1() { - MLXRandom.seed(127) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.18461060523986816, - accuracy: -0.0036922121047973633) - XCTAssertEqual( - a.sum().item(Float.self), -2.215327262878418, - accuracy: -0.04430654525756836) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), -0.03600400686264038, - accuracy: -0.0007200801372528076) - XCTAssertEqual( - aGrad.sum().item(Float.self), -0.43204808235168457, - accuracy: -0.008640961647033691) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Lion(learningRate: 0.1, weightDecay: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.18276450037956238, - accuracy: -0.003655290007591248) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -2.193173885345459, - accuracy: -0.04386347770690918) - } - - func testAdafactor() { - MLXRandom.seed(650) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), -0.5207136869430542, - accuracy: -0.010414273738861083) - XCTAssertEqual( - a.sum().item(Float.self), -6.248563766479492, - accuracy: -0.12497127532958985) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.4333036541938782, - accuracy: 0.008666073083877564) - XCTAssertEqual( - aGrad.sum().item(Float.self), 5.199643611907959, - accuracy: 0.10399287223815919) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Adafactor(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), -0.5268284678459167, - accuracy: -0.010536569356918336) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), -6.321941375732422, - accuracy: -0.12643882751464844) - } - - func testAdafactor1() { - MLXRandom.seed(193) - let a = MLXRandom.normal([4, 3]) - XCTAssertEqual(a.shape, [4, 3]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4008181691169739, - accuracy: 0.008016363382339478) - XCTAssertEqual( - a.sum().item(Float.self), 4.809817790985107, - accuracy: 0.09619635581970215) - let aGrad = MLXRandom.normal([4, 3]) - XCTAssertEqual(aGrad.shape, [4, 3]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.21447472274303436, - accuracy: 0.004289494454860688) - XCTAssertEqual( - aGrad.sum().item(Float.self), 2.5736966133117676, - accuracy: 0.05147393226623535) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Adafactor(learningRate: 0.1, beta1: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [4, 3]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.39943069219589233, - accuracy: 0.007988613843917847) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 4.793168067932129, - accuracy: 0.09586336135864258) - } - - func testAdafactor2() { - MLXRandom.seed(620) - let a = MLXRandom.uniform(0.0 ..< 1.0, [10]) - XCTAssertEqual(a.shape, [10]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4890245497226715, - accuracy: 0.00978049099445343) - XCTAssertEqual( - a.sum().item(Float.self), 4.89024543762207, - accuracy: 0.09780490875244141) - let aGrad = MLXRandom.uniform(0.0 ..< 1.0, [10]) - XCTAssertEqual(aGrad.shape, [10]) - XCTAssertEqual(aGrad.dtype, .float32) - XCTAssertEqual( - aGrad.mean().item(Float.self), 0.6818901896476746, - accuracy: 0.013637803792953491) - XCTAssertEqual( - aGrad.sum().item(Float.self), 6.818902015686035, - accuracy: 0.1363780403137207) - let aModel = ModuleParameters(values: ["a": .value(a)]) - let aGradParams = ModuleParameters(values: ["a": .value(aGrad)]) - let result = Adafactor(learningRate: 0.1).apply( - gradients: aGradParams, modelParameters: aModel) - XCTAssertEqual(result[unwrapping: "a"]!.shape, [10]) - XCTAssertEqual(result[unwrapping: "a"]!.dtype, .float32) - XCTAssertEqual( - result[unwrapping: "a"]!.mean().item(Float.self), 0.4835330545902252, - accuracy: 0.009670661091804504) - XCTAssertEqual( - result[unwrapping: "a"]!.sum().item(Float.self), 4.835330486297607, - accuracy: 0.09670660972595214) - } - - func testGLU() { - MLXRandom.seed(850) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5472526550292969, - accuracy: 0.010945053100585937) - XCTAssertEqual( - a.sum().item(Float.self), 140.0966796875, - accuracy: 2.80193359375) - let result = GLU()(a) - XCTAssertEqual(result.shape, [2, 8, 8]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.33327674865722656, - accuracy: 0.006665534973144531) - XCTAssertEqual( - result.sum().item(Float.self), 42.659423828125, - accuracy: 0.8531884765625) - } - - func testSigmoid1() { - MLXRandom.seed(589) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5296978950500488, - accuracy: 0.010593957901000976) - XCTAssertEqual( - a.sum().item(Float.self), 135.6026611328125, - accuracy: 2.71205322265625) - let result = Sigmoid()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.6270139813423157, - accuracy: 0.012540279626846314) - XCTAssertEqual( - result.sum().item(Float.self), 160.5155792236328, - accuracy: 3.2103115844726564) - } - - func testMish() { - MLXRandom.seed(122) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5017197728157043, - accuracy: 0.010034395456314087) - XCTAssertEqual( - a.sum().item(Float.self), 128.4402618408203, - accuracy: 2.5688052368164063) - let result = Mish()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.39537572860717773, - accuracy: 0.007907514572143556) - XCTAssertEqual( - result.sum().item(Float.self), 101.2161865234375, - accuracy: 2.0243237304687502) - } - - func testReLU() { - MLXRandom.seed(400) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.47832274436950684, - accuracy: 0.009566454887390137) - XCTAssertEqual( - a.sum().item(Float.self), 122.45062255859375, - accuracy: 2.449012451171875) - let result = ReLU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.47832274436950684, - accuracy: 0.009566454887390137) - XCTAssertEqual( - result.sum().item(Float.self), 122.45062255859375, - accuracy: 2.449012451171875) - } - - func testLeakyReLU() { - MLXRandom.seed(93) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4999306797981262, - accuracy: 0.009998613595962524) - XCTAssertEqual( - a.sum().item(Float.self), 127.98225402832031, - accuracy: 2.559645080566406) - let result = LeakyReLU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4999306797981262, - accuracy: 0.009998613595962524) - XCTAssertEqual( - result.sum().item(Float.self), 127.98225402832031, - accuracy: 2.559645080566406) - } - - func testReLU6() { - MLXRandom.seed(379) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.49325865507125854, - accuracy: 0.009865173101425172) - XCTAssertEqual( - a.sum().item(Float.self), 126.27421569824219, - accuracy: 2.525484313964844) - let result = ReLU6()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.49325865507125854, - accuracy: 0.009865173101425172) - XCTAssertEqual( - result.sum().item(Float.self), 126.27421569824219, - accuracy: 2.525484313964844) - } - - func testSoftmax() { - MLXRandom.seed(853) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5143963098526001, - accuracy: 0.010287926197052003) - XCTAssertEqual( - a.sum().item(Float.self), 131.68545532226562, - accuracy: 2.6337091064453126) - let result = Softmax()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.0624999962747097, - accuracy: 0.001249999925494194) - XCTAssertEqual( - result.sum().item(Float.self), 15.999999046325684, - accuracy: 0.31999998092651366) - } - - func testSoftplus() { - MLXRandom.seed(118) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.49898141622543335, - accuracy: 0.009979628324508667) - XCTAssertEqual( - a.sum().item(Float.self), 127.73924255371094, - accuracy: 2.554784851074219) - let result = SoftPlus()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.9828577637672424, - accuracy: 0.01965715527534485) - XCTAssertEqual( - result.sum().item(Float.self), 251.61158752441406, - accuracy: 5.032231750488282) - } - - func testSoftsign() { - MLXRandom.seed(37) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5065512657165527, - accuracy: 0.010131025314331054) - XCTAssertEqual( - a.sum().item(Float.self), 129.6771240234375, - accuracy: 2.59354248046875) - let result = SoftSign()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.314089834690094, - accuracy: 0.00628179669380188) - XCTAssertEqual( - result.sum().item(Float.self), 80.40699768066406, - accuracy: 1.6081399536132812) - } - - func testCELU() { - MLXRandom.seed(620) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4667481780052185, - accuracy: 0.00933496356010437) - XCTAssertEqual( - a.sum().item(Float.self), 119.48753356933594, - accuracy: 2.389750671386719) - let result = CELU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4667481780052185, - accuracy: 0.00933496356010437) - XCTAssertEqual( - result.sum().item(Float.self), 119.48753356933594, - accuracy: 2.389750671386719) - } - - func testSiLU() { - MLXRandom.seed(22) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5029705762863159, - accuracy: 0.010059411525726319) - XCTAssertEqual( - a.sum().item(Float.self), 128.76046752929688, - accuracy: 2.5752093505859377) - let result = SiLU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3319709300994873, - accuracy: 0.0066394186019897465) - XCTAssertEqual( - result.sum().item(Float.self), 84.98455810546875, - accuracy: 1.699691162109375) - } - - func testLogSoftmax() { - MLXRandom.seed(199) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.527843713760376, - accuracy: 0.010556874275207519) - XCTAssertEqual( - a.sum().item(Float.self), 135.12799072265625, - accuracy: 2.702559814453125) - let result = LogSoftMax()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -2.8109545707702637, - accuracy: -0.05621909141540527) - XCTAssertEqual( - result.sum().item(Float.self), -719.6043701171875, - accuracy: -14.39208740234375) - } - - func testLogSigmoid() { - MLXRandom.seed(984) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5109776854515076, - accuracy: 0.010219553709030152) - XCTAssertEqual( - a.sum().item(Float.self), 130.81028747558594, - accuracy: 2.616205749511719) - let result = LogSigmoid()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.4795985519886017, - accuracy: -0.009591971039772034) - XCTAssertEqual( - result.sum().item(Float.self), -122.77722930908203, - accuracy: -2.4555445861816407) - } - - func testPReLU() { - MLXRandom.seed(993) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.49665144085884094, - accuracy: 0.00993302881717682) - XCTAssertEqual( - a.sum().item(Float.self), 127.14276885986328, - accuracy: 2.5428553771972657) - let result = PReLU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.49665144085884094, - accuracy: 0.00993302881717682) - XCTAssertEqual( - result.sum().item(Float.self), 127.14276885986328, - accuracy: 2.5428553771972657) - } - - func testGELU() { - MLXRandom.seed(189) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.49295032024383545, - accuracy: 0.009859006404876709) - XCTAssertEqual( - a.sum().item(Float.self), 126.19528198242188, - accuracy: 2.5239056396484374) - let result = GELU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.3656383752822876, - accuracy: 0.007312767505645752) - XCTAssertEqual( - result.sum().item(Float.self), 93.60342407226562, - accuracy: 1.8720684814453126) - } - - func testTanh1() { - MLXRandom.seed(735) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.47412270307540894, - accuracy: 0.009482454061508178) - XCTAssertEqual( - a.sum().item(Float.self), 121.37541198730469, - accuracy: 2.4275082397460936) - let result = Tanh()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4130796790122986, - accuracy: 0.008261593580245972) - XCTAssertEqual( - result.sum().item(Float.self), 105.74839782714844, - accuracy: 2.114967956542969) - } - - func testHardswish() { - MLXRandom.seed(126) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4918924570083618, - accuracy: 0.009837849140167236) - XCTAssertEqual( - a.sum().item(Float.self), 125.92446899414062, - accuracy: 2.5184893798828125) - let result = HardSwish()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.2996022403240204, - accuracy: 0.005992044806480408) - XCTAssertEqual( - result.sum().item(Float.self), 76.69817352294922, - accuracy: 1.5339634704589844) - } - - func testStep() { - MLXRandom.seed(490) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4793606400489807, - accuracy: 0.009587212800979614) - XCTAssertEqual( - a.sum().item(Float.self), 122.71632385253906, - accuracy: 2.454326477050781) - let result = Step()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .int32) - XCTAssertEqual( - result.mean().item(Float.self), 1.0, - accuracy: 0.02) - XCTAssertEqual( - result.sum().item(Float.self), 256, - accuracy: 5.12) - } - - func testSELU() { - MLXRandom.seed(215) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4930267930030823, - accuracy: 0.009860535860061645) - XCTAssertEqual( - a.sum().item(Float.self), 126.21485900878906, - accuracy: 2.524297180175781) - let result = SELU()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.5180231928825378, - accuracy: 0.010360463857650756) - XCTAssertEqual( - result.sum().item(Float.self), 132.6139373779297, - accuracy: 2.6522787475585936) - } - - func testLinear() { - MLXRandom.seed(744) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5086885690689087, - accuracy: 0.010173771381378174) - XCTAssertEqual( - a.sum().item(Float.self), 130.22427368164062, - accuracy: 2.6044854736328125) - let result = Linear(16, 5)(a) - XCTAssertEqual(result.shape, [2, 8, 5]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.10419309139251709, - accuracy: 0.002083861827850342) - XCTAssertEqual( - result.sum().item(Float.self), 8.335447311401367, - accuracy: 0.16670894622802734) - } - - func testConv1d1() { - MLXRandom.seed(819) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.512987494468689, - accuracy: 0.01025974988937378) - XCTAssertEqual( - a.sum().item(Float.self), 131.32479858398438, - accuracy: 2.6264959716796876) - let result = Conv1d(inputChannels: 16, outputChannels: 2, kernelSize: 8)(a) - XCTAssertEqual(result.shape, [2, 1, 2]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.2648651897907257, - accuracy: 0.005297303795814514) - XCTAssertEqual( - result.sum().item(Float.self), 1.0594607591629028, - accuracy: 0.021189215183258055) - } - - func testConv2d1() { - MLXRandom.seed(62) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 8, 4]) - XCTAssertEqual(a.shape, [2, 8, 8, 4]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5225042700767517, - accuracy: 0.010450085401535034) - XCTAssertEqual( - a.sum().item(Float.self), 267.5221862792969, - accuracy: 5.350443725585937) - let result = Conv2d(inputChannels: 4, outputChannels: 2, kernelSize: 8)(a) - XCTAssertEqual(result.shape, [2, 1, 1, 2]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.27932149171829224, - accuracy: -0.005586429834365845) - XCTAssertEqual( - result.sum().item(Float.self), -1.117285966873169, - accuracy: -0.02234571933746338) - } - - func testDropout() { - MLXRandom.seed(959) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5114291906356812, - accuracy: 0.010228583812713623) - XCTAssertEqual( - a.sum().item(Float.self), 130.92587280273438, - accuracy: 2.6185174560546876) - let result = Dropout()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.47791361808776855, - accuracy: 0.009558272361755372) - XCTAssertEqual( - result.sum().item(Float.self), 122.34588623046875, - accuracy: 2.446917724609375) - } - - func testDropout2d() { - MLXRandom.seed(695) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4578399062156677, - accuracy: 0.009156798124313355) - XCTAssertEqual( - a.sum().item(Float.self), 117.20701599121094, - accuracy: 2.344140319824219) - let result = Dropout2d()(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.36828434467315674, - accuracy: 0.007365686893463135) - XCTAssertEqual( - result.sum().item(Float.self), 94.28079223632812, - accuracy: 1.8856158447265625) - } - - func testDropout3d() { - MLXRandom.seed(23) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 8, 4]) - XCTAssertEqual(a.shape, [2, 8, 8, 4]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5006061792373657, - accuracy: 0.010012123584747314) - XCTAssertEqual( - a.sum().item(Float.self), 256.31036376953125, - accuracy: 5.126207275390625) - let result = Dropout3d()(a) - XCTAssertEqual(result.shape, [2, 8, 8, 4]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.23728415369987488, - accuracy: 0.004745683073997498) - XCTAssertEqual( - result.sum().item(Float.self), 121.48948669433594, - accuracy: 2.429789733886719) - } - - func testEmbedding() { - MLXRandom.seed(557) - let a = MLXRandom.randInt(low: 0, high: 10, [2, 8, 8, 4]) - XCTAssertEqual(a.shape, [2, 8, 8, 4]) - XCTAssertEqual(a.dtype, .int32) - XCTAssertEqual( - a.mean().item(Float.self), 4.60546875, - accuracy: 0.09210937500000001) - XCTAssertEqual( - a.sum().item(Float.self), 2358, - accuracy: 47.160000000000004) - let result = Embedding(embeddingCount: 10, dimensions: 8)(a) - XCTAssertEqual(result.shape, [2, 8, 8, 4, 8]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.0011973462533205748, - accuracy: -2.3946925066411497e-05) - XCTAssertEqual( - result.sum().item(Float.self), -4.904330253601074, - accuracy: -0.0980866050720215) - } - - func testInstanceNorm() { - MLXRandom.seed(435) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5000646114349365, - accuracy: 0.01000129222869873) - XCTAssertEqual( - a.sum().item(Float.self), 128.01654052734375, - accuracy: 2.560330810546875) - let result = InstanceNorm(dimensions: 8)(a)[0, 0] - XCTAssertEqual(result.shape, [16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.10645411163568497, - accuracy: 0.0021290822327136995) - XCTAssertEqual( - result.sum().item(Float.self), 1.7032657861709595, - accuracy: 0.03406531572341919) - } - - func testLayerNorm() { - MLXRandom.seed(635) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.4926903247833252, - accuracy: 0.009853806495666504) - XCTAssertEqual( - a.sum().item(Float.self), 126.12872314453125, - accuracy: 2.522574462890625) - let result = LayerNorm(dimensions: 16)(a)[.ellipsis, 0] - XCTAssertEqual(result.shape, [2, 8]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.2909903824329376, - accuracy: 0.005819807648658752) - XCTAssertEqual( - result.sum().item(Float.self), 4.655846118927002, - accuracy: 0.09311692237854004) - } - - func testRMSNorm() { - MLXRandom.seed(103) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5054763555526733, - accuracy: 0.010109527111053467) - XCTAssertEqual( - a.sum().item(Float.self), 129.40194702148438, - accuracy: 2.5880389404296875) - let result = RMSNorm(dimensions: 16)(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.8729387521743774, - accuracy: 0.01745877504348755) - XCTAssertEqual( - result.sum().item(Float.self), 223.47232055664062, - accuracy: 4.469446411132813) - } - - func testGroupNorm() { - MLXRandom.seed(855) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.48666587471961975, - accuracy: 0.009733317494392395) - XCTAssertEqual( - a.sum().item(Float.self), 124.58646392822266, - accuracy: 2.491729278564453) - let result = GroupNorm(groupCount: 4, dimensions: 16)(a)[0, 0] - XCTAssertEqual(result.shape, [16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), -0.054606519639492035, - accuracy: -0.0010921303927898408) - XCTAssertEqual( - result.sum().item(Float.self), -0.8737043142318726, - accuracy: -0.017474086284637452) - } - - func testBatchNorm() { - MLXRandom.seed(266) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5058146715164185, - accuracy: 0.010116293430328369) - XCTAssertEqual( - a.sum().item(Float.self), 129.48855590820312, - accuracy: 2.5897711181640624) - let result = BatchNorm(featureCount: 16)(a)[0, 0] - XCTAssertEqual(result.shape, [16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4397852420806885, - accuracy: 0.00879570484161377) - XCTAssertEqual( - result.sum().item(Float.self), 7.036563873291016, - accuracy: 0.14073127746582031) - } - - func testRoPE() { - MLXRandom.seed(71) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5082664489746094, - accuracy: 0.010165328979492188) - XCTAssertEqual( - a.sum().item(Float.self), 130.1162109375, - accuracy: 2.60232421875) - let result = RoPE(dimensions: 8)(a) - XCTAssertEqual(result.shape, [2, 8, 16]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.4562537670135498, - accuracy: 0.009125075340270997) - XCTAssertEqual( - result.sum().item(Float.self), 116.80096435546875, - accuracy: 2.3360192871093752) - } - - func testRoPEArrayOffset() { - MLXRandom.seed(42) - let batch = MLXRandom.uniform(0.0 ..< 1.0, [3, 8, 16]) - XCTAssertEqual(batch.shape, [3, 8, 16]) - - let offsets = MLXArray([50, 20, 0]) - - // Test MLXFast.RoPE with array offset - let result = MLXFast.RoPE( - batch, dimensions: 8, traditional: false, - base: 10000, scale: 1.0, offset: offsets) - XCTAssertEqual(result.shape, [3, 8, 16]) - - // Verify against individual scalar offset calls - for i in 0 ..< 3 { - let single = batch[i].expandedDimensions(axis: 0) - let offsetValue = [50, 20, 0][i] - let expected = MLXFast.RoPE( - single, dimensions: 8, traditional: false, - base: 10000, scale: 1.0, offset: offsetValue) - XCTAssert(allClose(result[i], expected[0]).all().item()) - } - } - - func testRoPEArrayOffsetModule() { - MLXRandom.seed(123) - let batch = MLXRandom.uniform(0.0 ..< 1.0, [3, 8, 16]) - let offsets = MLXArray([10, 5, 0]) - - let rope = RoPE(dimensions: 8) - - // Test MLXNN RoPE module with array offset - let result = rope(batch, offset: offsets) - XCTAssertEqual(result.shape, [3, 8, 16]) - - // Verify shape and dtype preserved - XCTAssertEqual(result.dtype, batch.dtype) - } - - func testRoPEScalarArrayOffset() { - MLXRandom.seed(99) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - - // Scalar array offset should work the same as int offset - let scalarOffset = MLXArray(5) - let resultArray = MLXFast.RoPE( - a, dimensions: 8, traditional: false, - base: 10000, scale: 1.0, offset: scalarOffset) - - let resultInt = MLXFast.RoPE( - a, dimensions: 8, traditional: false, - base: 10000, scale: 1.0, offset: 5) - - XCTAssert(allClose(resultArray, resultInt).all().item()) - } - - func testSinusoidalPositionalEncoding() { - MLXRandom.seed(226) - let a = MLXRandom.uniform(0.0 ..< 1.0, [2, 8, 16]) - XCTAssertEqual(a.shape, [2, 8, 16]) - XCTAssertEqual(a.dtype, .float32) - XCTAssertEqual( - a.mean().item(Float.self), 0.5026599168777466, - accuracy: 0.010053198337554931) - XCTAssertEqual( - a.sum().item(Float.self), 128.68093872070312, - accuracy: 2.5736187744140624) - let result = SinusoidalPositionalEncoding(dimensions: 8)(a) - XCTAssertEqual(result.shape, [2, 8, 16, 8]) - XCTAssertEqual(result.dtype, .float32) - XCTAssertEqual( - result.mean().item(Float.self), 0.2705308198928833, - accuracy: 0.005410616397857666) - XCTAssertEqual( - result.sum().item(Float.self), 554.047119140625, - accuracy: 11.0809423828125) - } - -} diff --git a/Tests/MLXTests/MLXArray+InitTests.swift b/Tests/MLXTests/MLXArray+InitTests.swift index 8c1017be4..35e344718 100644 --- a/Tests/MLXTests/MLXArray+InitTests.swift +++ b/Tests/MLXTests/MLXArray+InitTests.swift @@ -438,6 +438,46 @@ class MLXArrayInitTests: XCTestCase { + "buffer permanently pins whatever the closure held") } + // MARK: - linspace + + func testLinspaceIntegerBoundsAreFloat() { + // matches python: mx.linspace(0, 1, 3) -> float32 [0, 0.5, 1]. + // an integer result would truncate to [0, 0, 1] + let a = MLXArray.linspace(0, 1, count: 3) + XCTAssertEqual(a.dtype, .float32) + assertEqual(a, MLXArray(converting: [0.0, 0.5, 1.0]), atol: 1e-6) + } + + func testLinspaceIntegerDType() { + // an integer (truncating) result is available, but has to be asked for + let a = MLXArray.linspace(0, 10, count: 6, dtype: .int32) + XCTAssertEqual(a.dtype, .int32) + XCTAssertEqual(a.asArray(Int32.self), [0, 2, 4, 6, 8, 10]) + } + + func testLinspaceDoubleBoundsAreFloat32() { + // Double.dtype is float64, but like MLXArray(1.0) we do not promote to + // float64 -- it is not available on the GPU + let a = MLXArray.linspace(0.0, 1.0, count: 5) + XCTAssertEqual(a.dtype, .float32) + assertEqual(a, MLXArray(converting: [0.0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6) + } + + func testLinspaceFloatBoundsKeepTheirDType() { + XCTAssertEqual(MLXArray.linspace(Float(0), Float(1), count: 5).dtype, .float32) + #if !arch(x86_64) + XCTAssertEqual(MLXArray.linspace(Float16(0), Float16(1), count: 5).dtype, .float16) + #endif + } + + func testLinspaceEndpoint() { + let inclusive = MLX.linspace(0, 1, count: 5) + assertEqual(inclusive, MLXArray(converting: [0, 0.25, 0.5, 0.75, 1.0]), atol: 1e-6) + + let halfOpen = MLX.linspace(0, 1, count: 5, endpoint: false) + assertEqual(halfOpen, MLXArray(converting: [0, 0.2, 0.4, 0.6, 0.8]), atol: 1e-6) + } + #if canImport(IOSurface) func testIOSurface() { let height = 100 diff --git a/Tests/MLXTests/OpsTests.swift b/Tests/MLXTests/OpsTests.swift index b59970c23..41bde1dcb 100644 --- a/Tests/MLXTests/OpsTests.swift +++ b/Tests/MLXTests/OpsTests.swift @@ -53,6 +53,85 @@ class OpsTests: XCTestCase { assertEqual(c, expected) } + func testTensordotDefaultAxes() { + // axes defaults to 2 (as in numpy and python mlx): sum over the last two + // dimensions of a and the first two of b + let a = MLXArray(0 ..< 24, [2, 3, 4]).asType(.float32) + let b = MLXArray(0 ..< 60, [3, 4, 5]).asType(.float32) + + assertEqual(tensordot(a, b), tensordot(a, b, axes: 2)) + XCTAssertEqual(tensordot(a, b).shape, [2, 5]) + } + + func testNanToNumDefaults() { + // by default infinities become the largest finite value for the dtype + // (matching python mlx) and NaN becomes 0 + let a = MLXArray( + [1.5, Float.nan, Float.infinity, -Float.infinity] as [Float]) + + let result = nanToNum(a) + XCTAssertEqual(result.dtype, .float32) + assertEqual( + result, + MLXArray( + [ + 1.5, 0, Float.greatestFiniteMagnitude, -Float.greatestFiniteMagnitude, + ] as [Float])) + + // and they can be replaced explicitly + assertEqual( + nanToNum(a, nan: -1, posInf: 100, negInf: -100), + MLXArray([1.5, -1, 100, -100] as [Float])) + } + + func testNanToNumDefaultsFloat16() { + // the replacement follows the dtype + let a = MLXArray([Float.infinity, -Float.infinity] as [Float]).asType(.float16) + let limit = Float(DType.float16.finfo!.max) + + let result = nanToNum(a) + XCTAssertEqual(result.dtype, .float16) + XCTAssertEqual(result[0].item(Float.self), limit) + XCTAssertEqual(result[1].item(Float.self), -limit) + } + + func testConvolveModeShapes() { + // matches numpy/python mlx: full is M + K - 1, same is M, valid is M - K + 1 + let a = MLXArray(0 ..< 20).asType(.float32) + + for kernelSize in [3, 4, 5, 6] { + let v = MLXArray(1 ..< (kernelSize + 1)).asType(.float32) + + XCTAssertEqual( + convolve(a, v, mode: .full).shape, [a.size + kernelSize - 1], + "full, kernel \(kernelSize)") + XCTAssertEqual( + convolve(a, v, mode: .same).shape, [a.size], "same, kernel \(kernelSize)") + XCTAssertEqual( + convolve(a, v, mode: .valid).shape, [a.size - kernelSize + 1], + "valid, kernel \(kernelSize)") + } + } + + func testConvolveSameAndValidAreWindowsOfFull() { + // `same` is the centered `a.size` window of the full convolution and + // `valid` is the window with no zero padding at all -- true for both odd + // and even sized weights (even sizes need asymmetric padding) + let a = MLXArray(0 ..< 20).asType(.float32) + + for kernelSize in [3, 4, 5, 6] { + let v = MLXArray(1 ..< (kernelSize + 1)).asType(.float32) + let full = convolve(a, v, mode: .full) + + let sameStart = (kernelSize - 1) / 2 + assertEqual( + convolve(a, v, mode: .same), full[sameStart ..< (sameStart + a.size)]) + + assertEqual( + convolve(a, v, mode: .valid), full[(kernelSize - 1) ..< a.size]) + } + } + func testConvertScalarInt() { let a = MLXArray(0 ..< 10) let b = a .< (a + 1) diff --git a/Tests/MLXTests/OptimizerTests.swift b/Tests/MLXTests/OptimizerTests.swift index 95e04af7d..5d1e3727a 100644 --- a/Tests/MLXTests/OptimizerTests.swift +++ b/Tests/MLXTests/OptimizerTests.swift @@ -13,17 +13,24 @@ class OptimizerTests: XCTestCase { } class ShapeModule: Module { + // ranks 1 through 3: some optimizers (Adafactor, Muon) treat them differently let first = [MLXArray.zeros([10]), MLXArray.zeros([1])] let second = MLXArray.zeros([1]) + let matrix = MLXArray.zeros([3, 5]) + let tensor = MLXArray.zeros([2, 3, 4]) } - func checkShape(optimizer: OptimizerBase) { + func checkShape(optimizer: OptimizerBase, steps: Int = 2) { let model = ShapeModule() let params = model.parameters() let grads = params.mapValues { MLXArray.ones(like: $0) } - let optimizer = SGD(learningRate: 0.1) - let update = optimizer.apply(gradients: grads, modelParameters: model.parameters()) + // note: more than one step, and using the optimizer that was passed in -- + // this used to build its own SGD, so every caller was really testing SGD + var update = params + for _ in 0 ..< steps { + update = optimizer.apply(gradients: grads, modelParameters: update) + } eval(update) let shapesEqual = params.mapValues(update) { (e1, e2) -> Bool in @@ -212,6 +219,167 @@ class OptimizerTests: XCTestCase { checkTrain(optimizer: Adafactor(learningRate: 0.1), compile: true) } + // MARK: - Muon + + func testMuonUpdateIsConditioned() { + // the Newton-Schulz iteration used by Muon is a *crude* polar approximation: + // it pulls the singular values of the update into a band around 1 (it does + // not converge to an exactly orthogonal matrix). With no momentum and + // nesterov off the update direction is the orthogonalized gradient, so the + // singular values of the update should be far better conditioned than the + // gradient's. + let optimizer = Muon( + learningRate: 1.0, momentum: 0.0, weightDecay: 0.0, nesterov: false, nsSteps: 5) + + // full rank and deterministic: `arange` reshaped is rank 2, which the + // iteration cannot orthogonalize + let values = (0 ..< 24).map { Float(sin(Double($0) * 1.7) + 0.1 * Double($0 % 5)) } + let parameter = MLXArray(values, [6, 4]) + let parameters = ModuleParameters.unflattened([("w", parameter)]) + + let updated = optimizer.apply(gradients: parameters, modelParameters: parameters) + + // parameter - learningRate * scale * update, with scale = sqrt(max(1, 6/4)) + let scale = Float((6.0 / 4.0).squareRoot()) + let update = (parameter - updated[unwrapping: "w"]!) / scale + + // singular values via the gram matrix, which avoids picking between the svd + // overloads: eigvalsh returns them ascending + func singularValues(_ array: MLXArray) -> MLXArray { + MLX.sqrt( + MLX.maximum(0, MLX.eigvalsh(matmul(array.T, array), stream: .cpu))) + } + + let inputSingular = singularValues(parameter) + let updateSingular = singularValues(update) + + let inputCondition = + inputSingular.max().item(Float.self) / inputSingular.min().item(Float.self) + let updateCondition = + updateSingular.max().item(Float.self) / updateSingular.min().item(Float.self) + + XCTAssertGreaterThan(inputCondition, 10) + XCTAssertLessThan(updateCondition, 2) + // and the values are pulled toward 1 rather than being rescaled arbitrarily + XCTAssertGreaterThan(updateSingular.min().item(Float.self), 0.5) + XCTAssertLessThan(updateSingular.max().item(Float.self), 1.5) + } + + func testMuonLeavesLowRankParametersAlone() { + // rank 0/1 parameters skip the orthogonalization: the update is the plain + // momentum direction, so a single step moves by learningRate * gradient + let optimizer = Muon( + learningRate: 0.1, momentum: 0.0, weightDecay: 0.0, nesterov: false, nsSteps: 5) + + let parameter = MLXArray(converting: [1.0, 2.0, 3.0]) + let gradients = ModuleParameters.unflattened([("b", MLXArray.ones([3]))]) + let updated = optimizer.apply( + gradients: gradients, + modelParameters: ModuleParameters.unflattened([("b", parameter)])) + + assertEqual( + updated[unwrapping: "b"]!, MLXArray(converting: [0.9, 1.9, 2.9]), atol: 1e-6) + } + + // MARK: - defaults + // + // these have to match python: `tools/audit_defaults.py` compares them + // automatically, and this is the regression test for the one that was wrong + + func testLionDefaults() { + let lion = Lion(learningRate: 0.1) + XCTAssertEqual(lion.betas.0, 0.9) + // python's Lion uses 0.99 here -- Adam is the one with 0.999 + XCTAssertEqual(lion.betas.1, 0.99) + XCTAssertEqual(lion.weightDecay, 0) + } + + func testAdamAndAdamWDefaults() { + XCTAssertEqual(Adam(learningRate: 0.1).betas.1, 0.999) + XCTAssertEqual(AdamW(learningRate: 0.1).betas.1, 0.999) + XCTAssertEqual(AdamW(learningRate: 0.1).weightDecay, 0.01) + } + + func testMuonDefaults() { + let muon = Muon(learningRate: 0.1) + XCTAssertEqual(muon.momentum, 0.95) + XCTAssertEqual(muon.weightDecay, 0.01) + XCTAssertTrue(muon.nesterov) + XCTAssertEqual(muon.nsSteps, 5) + } + + // MARK: - Adafactor + // + // Adafactor keeps *factored* state for parameters of rank 2 and above and + // unfactored state below that, so its state and its update path have to agree + // about the rank -- rank 3+ was broken (an outer product via matmul, which only + // works for rank 2) and the mismatch then crashed on a nil state. + + func testAdafactorRanks() throws { + let optimizer = Adafactor(learningRate: 0.1, relativeStep: false) + + var parameters = ModuleParameters.unflattened([ + ("vector", MLXArray.ones([7])), + ("matrix", MLXArray.ones([3, 5])), + ("tensor", MLXArray.ones([2, 3, 4])), + ("big", MLXArray.ones([2, 3, 4, 5])), + ]) + + for _ in 0 ..< 3 { + let gradients = parameters.mapValues { 0.5 * $0 } + parameters = optimizer.apply(gradients: gradients, modelParameters: parameters) + eval(parameters) + } + + for (key, shape) in [ + ("vector", [7]), ("matrix", [3, 5]), ("tensor", [2, 3, 4]), ("big", [2, 3, 4, 5]), + ] { + let value = parameters[unwrapping: key]! + XCTAssertEqual(value.shape, shape, key) + XCTAssertTrue(isFinite(value).all().item(Bool.self), "\(key) is not finite") + // the update should have moved the parameter off its starting value + XCTAssertFalse( + value.allClose(MLXArray.ones(like: value), rtol: 1e-6).item(Bool.self), + "\(key) did not change") + } + } + + func testAdafactorFactoredUpdateIsBatched() { + // the row/column factors combine per leading index: slice `i` of a rank 3 + // update must equal the rank 2 update of slice `i` of the factors + let optimizer = Adafactor(learningRate: 0.1, relativeStep: false) + + let row = MLXArray(converting: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]) + let column = MLXArray(converting: [0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], [2, 4]) + + let combined = optimizer.approvateExpMovingAverage( + expAvgSqRow: row, expAvgSqCol: column) + XCTAssertEqual(combined.shape, [2, 3, 4]) + + for i in 0 ..< 2 { + let slice = optimizer.approvateExpMovingAverage( + expAvgSqRow: row[i], expAvgSqCol: column[i]) + assertEqual(combined[i], slice, atol: 1e-6) + } + } + + func testAdafactorRankChangeDoesNotCrash() { + // the state is created for the rank seen on the first step; a later call with + // a different rank for the same key used to hit a nil state + let optimizer = Adafactor(learningRate: 0.1, relativeStep: false) + + var matrix = ModuleParameters.unflattened([("w", MLXArray.ones([3, 4]))]) + matrix = optimizer.apply( + gradients: matrix.mapValues { 0.5 * $0 }, modelParameters: matrix) + + var vector = ModuleParameters.unflattened([("w", MLXArray.ones([5]))]) + vector = optimizer.apply( + gradients: vector.mapValues { 0.5 * $0 }, modelParameters: vector) + eval(vector) + + XCTAssertEqual(vector[unwrapping: "w"]!.shape, [5]) + } + class TwoParameterModel: Module { let weight = MLXArray.zeros([3]) let bias = MLXArray.zeros([3]) diff --git a/Tests/MLXTests/PoolingTests.swift b/Tests/MLXTests/PoolingTests.swift index 90534fca9..2363b21ef 100644 --- a/Tests/MLXTests/PoolingTests.swift +++ b/Tests/MLXTests/PoolingTests.swift @@ -67,4 +67,87 @@ class MLXNNPoolingTests: XCTestCase { let output = pool.callAsFunction(input) assertEqual(output, MLXArray(converting: [2.5, 4.5, 10.5, 12.5], [1, 2, 2, 1])) } + + // MARK: - padding + // + // padding shifts every spatial dimension and must leave the batch and channel + // dimensions alone: the output is + // floor((size + 2 * padding - kernel) / stride) + 1 + + func testMaxPooling1dPadding() { + let input = MLXArray(0 ..< 4, [1, 4, 1]) + let pool = MaxPool1d(kernelSize: 2, stride: 2, padding: 1) + let output = pool(input) + + // padded with -inf: [-inf, 0, 1, 2, 3, -inf] + XCTAssertEqual(output.shape, [1, 3, 1]) + assertEqual(output, MLXArray([0, 2, 3], [1, 3, 1])) + } + + func testAvgPooling1dPadding() { + let input = MLXArray(0 ..< 4, [1, 4, 1]).asType(.float32) + let pool = AvgPool1d(kernelSize: 2, stride: 2, padding: 1) + let output = pool(input) + + // padded with 0, which is included in the average: [0, 0, 1, 2, 3, 0] + XCTAssertEqual(output.shape, [1, 3, 1]) + assertEqual(output, MLXArray(converting: [0, 1.5, 1.5], [1, 3, 1])) + } + + func testMaxPooling2dPadding() { + let input = MLXArray(0 ..< 16, [1, 4, 4, 1]) + let pool = MaxPool2d(kernelSize: 3, stride: 2, padding: 1) + let output = pool(input) + + XCTAssertEqual(output.shape, [1, 2, 2, 1]) + assertEqual(output, MLXArray([5, 7, 13, 15], [1, 2, 2, 1])) + } + + func testAvgPooling2dPadding() { + let input = MLXArray(0 ..< 4, [1, 2, 2, 1]).asType(.float32) + let pool = AvgPool2d(kernelSize: 2, stride: 2, padding: 1) + let output = pool(input) + + XCTAssertEqual(output.shape, [1, 2, 2, 1]) + assertEqual(output, MLXArray(converting: [0, 0.25, 0.5, 0.75], [1, 2, 2, 1])) + } + + func testPoolingPaddingKeepsBatchAndChannels() { + // the padding must not touch the batch or channel dimensions: this shape + // regressed to [2, 3, 4, 6] when the pad widths were built per pair + // instead of per dimension + let input = MLXArray(0 ..< (2 * 8 * 8 * 4), [2, 8, 8, 4]).asType(.float32) + + XCTAssertEqual( + MaxPool2d(kernelSize: 3, stride: 2, padding: 1)(input).shape, [2, 4, 4, 4]) + XCTAssertEqual( + AvgPool2d(kernelSize: 3, stride: 2, padding: 1)(input).shape, [2, 4, 4, 4]) + XCTAssertEqual( + MaxPool2d(kernelSize: 2, stride: 2, padding: 0)(input).shape, [2, 4, 4, 4]) + } + + func testPooling1dAnd3dPaddingShapes() { + let input1d = MLXArray(0 ..< (2 * 16 * 4), [2, 16, 4]).asType(.float32) + XCTAssertEqual( + MaxPool1d(kernelSize: 3, stride: 2, padding: 1)(input1d).shape, [2, 8, 4]) + XCTAssertEqual( + AvgPool1d(kernelSize: 3, stride: 2, padding: 1)(input1d).shape, [2, 8, 4]) + + let input3d = MLXArray(0 ..< (2 * 4 * 8 * 8 * 4), [2, 4, 8, 8, 4]).asType(.float32) + XCTAssertEqual( + MaxPool3d(kernelSize: 3, stride: 2, padding: 1)(input3d).shape, [2, 2, 4, 4, 4]) + XCTAssertEqual( + AvgPool3d(kernelSize: 3, stride: 2, padding: 1)(input3d).shape, [2, 2, 4, 4, 4]) + } + + func testPoolingAsymmetricPaddingPerAxis() { + let input = MLXArray(0 ..< (1 * 8 * 6 * 1), [1, 8, 6, 1]).asType(.float32) + + // padding only the width axis + XCTAssertEqual( + MaxPool2d(kernelSize: 2, stride: 2, padding: [0, 1])(input).shape, [1, 4, 4, 1]) + // and only the height axis + XCTAssertEqual( + MaxPool2d(kernelSize: 2, stride: 2, padding: [1, 0])(input).shape, [1, 5, 3, 1]) + } } diff --git a/Tests/MLXTests/PositionalEncodingTests.swift b/Tests/MLXTests/PositionalEncodingTests.swift new file mode 100644 index 000000000..9d5ac1502 --- /dev/null +++ b/Tests/MLXTests/PositionalEncodingTests.swift @@ -0,0 +1,167 @@ +// Copyright © 2026 Apple Inc. + +import Foundation +import MLX +import XCTest + +@testable import MLXNN + +/// Tests for layers whose formulas are easy to get subtly wrong -- these pin down +/// the parts that a shape check would not catch. +class PositionalEncodingTests: XCTestCase { + + // MARK: - ALiBi + + func testALiBiSlopesPowerOfTwo() { + // matches python `ALiBi.create_alibi_slope`: 2^(-8i/n) for i in 1...n + let slopes = ALiBi.alibiSlope(numHeads: 4).reshaped([-1]) + assertEqual( + slopes, MLXArray(converting: [0.25, 0.0625, 0.015625, 0.00390625]), atol: 1e-7) + } + + func testALiBiSlopesNotPowerOfTwo() { + // python does *not* extend the geometric series: it takes the slopes of the + // power of two below and pads with every other slope of the one above + assertEqual( + ALiBi.alibiSlope(numHeads: 3).reshaped([-1]), + MLXArray(converting: [0.0625, 0.00390625, 0.25]), atol: 1e-7) + assertEqual( + ALiBi.alibiSlope(numHeads: 6).reshaped([-1]), + MLXArray(converting: [0.25, 0.0625, 0.015625, 0.00390625, 0.5, 0.125]), + atol: 1e-7) + } + + func testALiBiMaskIsADistanceMatrix() { + // the mask must be -|i - j| * slope: a (q, k) matrix, not a broadcast of a + // single column (which collapsed to all zeros when q == k) + let scores = MLXArray.zeros([1, 2, 4, 4]) + let mask = ALiBi()(attentionScores: scores) + + XCTAssertEqual(mask.shape, [1, 2, 4, 4]) + + let slopes: [Float] = [0.0625, 0.00390625] // numHeads == 2 + for head in 0 ..< 2 { + for q in 0 ..< 4 { + for k in 0 ..< 4 { + let expected = -Float(abs(q - k)) * slopes[head] + XCTAssertEqual( + mask[0, head, q, k].item(Float.self), expected, accuracy: 1e-6, + "head \(head), q \(q), k \(k)") + } + } + } + } + + func testALiBiOffsetShiftsTheQueryPositions() { + // q positions start at `offset`, so with offset 1 and 3 queries the + // distances are |[1, 2, 3] - [0, 1, 2, 3]| + let scores = MLXArray.zeros([1, 1, 3, 4]) + let mask = ALiBi()(attentionScores: scores, offset: 1) + + XCTAssertEqual(mask.shape, [1, 1, 3, 4]) + + let slope: Float = 0.00390625 // numHeads == 1 + for q in 0 ..< 3 { + for k in 0 ..< 4 { + let expected = -Float(abs((q + 1) - k)) * slope + XCTAssertEqual( + mask[0, 0, q, k].item(Float.self), expected, accuracy: 1e-6, + "q \(q), k \(k)") + } + } + } + + func testALiBiAddsToTheScores() { + let scores = MLXArray.ones([1, 1, 3, 3]) + let mask = ALiBi()(attentionScores: MLXArray.zeros([1, 1, 3, 3])) + assertEqual(ALiBi()(attentionScores: scores), scores + mask) + } +} + +/// Recurrent layers: the formulas have branches that only run for some argument +/// combinations, which is where they drift from python. +class RecurrentTests: XCTestCase { + + private func parameters(_ module: Module) -> ModuleParameters { + // deterministic values so the tests do not depend on the random init + module.mapParameters { parameter in + let size = parameter.size + return MLX.arange(size, dtype: .float32).reshaped(parameter.shape) / Float(size) + - 0.5 + } + } + + func testGRUAppliesHiddenBiasWithoutAnIncomingHiddenState() { + // python applies `r * bhn` even when no hidden state is passed in; when + // this was missing, `bhn` had no effect on the first call at all + let gru = GRU(inputSize: 4, hiddenSize: 3) + gru.update(parameters: parameters(gru)) + + let x = MLX.arange(2 * 5 * 4, dtype: .float32).reshaped([2, 5, 4]) / 40 + + let before = gru(x) + + gru.update(parameters: ModuleParameters.unflattened([("bhn", MLXArray.ones([3]))])) + let after = gru(x) + + assertNotEqual(before, after) + } + + func testGRUMatchesAReferenceImplementation() { + let hiddenSize = 3 + let gru = GRU(inputSize: 4, hiddenSize: hiddenSize) + gru.update(parameters: parameters(gru)) + + let x = MLX.arange(1 * 2 * 4, dtype: .float32).reshaped([1, 2, 4]) / 8 + let result = gru(x) + + // the same formula written out, with no incoming hidden state + let wx = gru.wx + let wh = gru.wh + let b = gru.b! + let bhn = gru.bhn! + + let projected = addMM(b, x, wx.T) + var hidden: MLXArray? = nil + var steps = [MLXArray]() + + for step in 0 ..< x.dim(-2) { + var rz = projected[.ellipsis, step, ..<(2 * hiddenSize)] + var hiddenN: MLXArray? = nil + + if let hidden { + let projectedHidden = matmul(hidden, wh.T) + rz = rz + projectedHidden[.ellipsis, ..<(2 * hiddenSize)] + hiddenN = projectedHidden[.ellipsis, (2 * hiddenSize)...] + bhn + } + + rz = sigmoid(rz) + let r = rz[.ellipsis, ..