diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index c84e02f27..2fc81d589 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.github/scripts/build-linux-cuda-cmake.sh @@ -16,6 +16,14 @@ 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 ./example1 --device gpu ./tutorial --device gpu diff --git a/.github/scripts/setup+build-linux-container-cmake.sh b/.github/scripts/setup+build-linux-container-cmake.sh index 8b365bc42..93f8ce870 100755 --- a/.github/scripts/setup+build-linux-container-cmake.sh +++ b/.github/scripts/setup+build-linux-container-cmake.sh @@ -2,6 +2,10 @@ set -ex # [Setup] Install dependencies inside the container. +# Note: mlx-c requires a C++20 toolchain with support: +# - libstdc++ >= 13 (Ubuntu noble, Fedora 39+) +# - or libc++ >= 17 +# Debian bookworm (libstdc++ 12) is *not* sufficient. if command -v apt-get >/dev/null 2>&1; then apt-get update -y apt-get install -y \ @@ -12,6 +16,9 @@ if command -v apt-get >/dev/null 2>&1; then liblapacke-dev \ libopenblas-dev + export CC=/usr/bin/gcc + export CXX=/usr/bin/g++ + elif command -v dnf >/dev/null 2>&1; then dnf update -y dnf install -y \ @@ -23,6 +30,11 @@ elif command -v dnf >/dev/null 2>&1; then clang \ ninja-build + # Fedora's default gcc (13+) has ; keep it as the compiler so this + # build keeps matching what has been passing in CI. + export CC=/usr/bin/gcc + export CXX=/usr/bin/g++ + else echo "No supported package manager found (apt-get, dnf)" exit 1 @@ -31,8 +43,8 @@ fi # [CMake] CI Build Sanity Check: Verifies code compilation, not for release. export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" export DEBUG=1 -export CMAKE_C_COMPILER=/usr/bin/clang -export CMAKE_CXX_COMPILER=/usr/bin/clang++ + +"$CXX" --version rm -rf build mkdir -p build diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index b46745c75..7dcd7ae49 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -69,6 +69,7 @@ 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) + mac_build_and_test: needs: lint if: github.repository == 'ml-explore/mlx-swift' @@ -82,46 +83,54 @@ 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: | 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 - env: - # workaround for CI failure - DEVELOPER_DIR: /Applications/Xcode-latest.app 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' - name: Run Tests SwiftPM (Xcode, macOS) shell: sh - env: - # workaround for CI failure - DEVELOPER_DIR: /Applications/Xcode-latest.app 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 + 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 - env: - # workaround for CI failure - DEVELOPER_DIR: /Applications/Xcode-latest.app run: | rm -rf ~/Library/Developer/Xcode/DerivedData/* cd xcode xcodebuild build-for-testing -scheme MLX -destination 'platform=macOS' - name: Run Tests xcodeproj (Xcode, macOS) - env: - # workaround for CI failure - DEVELOPER_DIR: /Applications/Xcode-latest.app working-directory: xcode run: ../.github/scripts/run-xcode-tests.sh @@ -153,15 +162,20 @@ jobs: matrix: include: # See https://github.com/swiftlang/swift-docker/tree/main/ + # + # Note: mlx-c uses std::format, which requires libstdc++ >= 13 (or + # libc++ >= 17). Debian bookworm only ships gcc/libstdc++ 12 (and + # gcc-13 is not available in bookworm-backports), so use Ubuntu + # noble (24.04, gcc/libstdc++ 13) for the Debian-family build. - host: ubuntu-22.04 arch: x86_64 - container: swift:bookworm + container: swift:noble - host: ubuntu-22.04 arch: x86_64 container: swift:fedora39 - host: ubuntu-22.04-arm arch: aarch64 - container: swift:bookworm + container: swift:noble - host: ubuntu-22.04-arm arch: aarch64 container: swift:fedora39 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7c5a75282..7f8e7e61a 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -26,6 +26,11 @@ jobs: shell: bash run: xcodebuild -showComponent MetalToolchain + - name: Setup gh + shell: sh + run: | + brew install gh + - name: Build XCFrameworks shell: bash env: diff --git a/CMakeLists.txt b/CMakeLists.txt index 18f350bf5..f75ba6c9c 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,21 +28,61 @@ 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" - GIT_TAG "v0.6.0") + GIT_TAG "ebc88f10caa1b625e6b581437a8dea6df8a70085") 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 @@ -57,6 +99,21 @@ if(MLX_BUILD_CUDA) CUDNN::cudnn_all # Promoted from PRIVATE, required by mlx ) endif() + +# cLogSupport package small C shim that exposes os_log_type_enabled() to Swift +# (see #491) +set(cLogSupport_include ${CMAKE_CURRENT_LIST_DIR}/Source/cLogSupport/include) +set(cLogSupport_module_dir ${CMAKE_CURRENT_BINARY_DIR}/cLogSupport-module) +file( + WRITE ${cLogSupport_module_dir}/module.modulemap + "module cLogSupport {\n header \"${cLogSupport_include}/cLogSupport.h\"\n export *\n}\n" +) + +add_library(cLogSupport STATIC + ${CMAKE_CURRENT_LIST_DIR}/Source/cLogSupport/cLogSupport.c) +target_include_directories(cLogSupport PUBLIC ${cLogSupport_include} + ${cLogSupport_module_dir}) + file(GLOB MLX-src ${CMAKE_CURRENT_LIST_DIR}/Source/MLX/*.swift) if(NOT MLX_BUILD_METAL) @@ -64,11 +121,15 @@ 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) # Note: mlxc is target name exported by mlx-c -target_link_libraries(MLX PUBLIC mlxc Numerics) +target_link_libraries(MLX PUBLIC mlxc Numerics cLogSupport) target_compile_options(MLX PRIVATE -package-name MLX) # MLXRandom package @@ -103,7 +164,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/MAINTENANCE.md b/MAINTENANCE.md index f71c5b268..9e603b0a4 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 @@ -120,9 +115,15 @@ pre-commit run --all-files SwiftPM is able to fetch repositories from github and build them _if_ they have a `Package.swift` at the top level. It is unable to do this for repositories -that do not have a `Package.swift`. For this reason `mlx-swift` uses +that do not have a `Package.swift`. For this reason upstream `mlx-swift` uses git submodules to include the `mlx` and `mlx-c` repositories. +> **SharpAI fork:** this fork vendors `Source/Cmlx/mlx` and `Source/Cmlx/mlx-c` as plain tracked +> directories (upstream release + the SharpAI patch series) instead of submodules. Follow +> [UPSTREAM_MERGE_PLAN.md](UPSTREAM_MERGE_PLAN.md) for the fork's procedure; the steps below +> describe the upstream process, which the fork follows with step 1 replaced by "re-vendor the +> tag/commit upstream's submodules point to, then re-apply the carried patches". + When a new version of `mlx` and its equivalent `mlx-c` are to be used, there is a process to go through to update `mlx-swift`. diff --git a/Package.swift b/Package.swift index ff9c0801f..7390b8f74 100644 --- a/Package.swift +++ b/Package.swift @@ -4,86 +4,150 @@ import PackageDescription +// SharpAI fork: upstream builds CUDA on Linux through a SwiftPM build-tool plugin +// (CudaBuild/encuda), which requires swift-tools-version 6.3 (experimentalCGen). The +// Apple toolchains this fork supports are older, and the fork ships no CUDA backend, so +// the plugin/targets are intentionally not wired up here (Source/Encuda and +// Plugins/CudaBuild stay in the tree, unused, to keep future upstream merges clean). +let cudaBuildPlugins: [Target.PluginUsage] = [] +let cudaPackageDependencies: [Package.Dependency] = [] +let cudaTargets: [Target] = [] + +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/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", + "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", +] + #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 platformExcludes: [String] + let cxxSettings: [CXXSetting] + let linkerSettings: [LinkerSetting] + let mlxSwiftExcludes: [String] - let cxxSettings: [CXXSetting] = [] + // Linux without CUDA (CPU only) - let linkerSettings: [LinkerSetting] = [ + 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])), ] - let mlxSwiftExcludes: [String] = [ + 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] = + [ + // 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", - // 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 +165,9 @@ import PackageDescription .linkedFramework("Accelerate"), ] - let mlxSwiftExcludes: [String] = [] + let mlxSwiftExcludes: [String] = [ + "GPU+CUDA.swift" + ] #endif let cmlx = Target.target( @@ -111,10 +177,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", @@ -146,40 +210,7 @@ let cmlx = Target.target( "mlx/python", "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.cpp", - "mlx/mlx/backend/cuda/device_info.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", + "mlx/build", // build variants (we are opting _out_ of these) "mlx/mlx/io/no_safetensors.cpp", @@ -190,29 +221,27 @@ 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/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"), .headerSearchPath("mlx-c"), + .headerSearchPath("mlx-generated/cuda"), ], cxxSettings: cxxSettings + [ .headerSearchPath("mlx"), .headerSearchPath("mlx-c"), .headerSearchPath("json/single_include/nlohmann"), .headerSearchPath("fmt/include"), - .define("MLX_VERSION", to: "\"0.24.2\""), - .define("MLX_ENABLE_NAX", to: "1"), + .define("MLX_VERSION", to: "\"0.32.0\""), ], - linkerSettings: linkerSettings + linkerSettings: linkerSettings, + plugins: cudaBuildPlugins, ) let package = Package( @@ -238,7 +267,7 @@ let package = Package( dependencies: [ // for Complex type .package(url: "https://github.com/apple/swift-numerics", from: "1.0.0") - ], + ] + cudaPackageDependencies, targets: [ cmlx, .testTarget( @@ -246,10 +275,17 @@ let package = Package( dependencies: ["Cmlx"] ), + .target( + name: "cLogSupport", + path: "Source/cLogSupport", + publicHeadersPath: "include" + ), + .target( name: "MLX", dependencies: [ "Cmlx", + "cLogSupport", .product(name: "Numerics", package: "swift-numerics"), ], exclude: mlxSwiftExcludes, @@ -306,6 +342,12 @@ let package = Package( "MLX", "MLXNN", "MLXOptimizers", ] ), + .testTarget( + name: "MLXIntegrationTests", + dependencies: [ + "MLX", "MLXNN", "MLXOptimizers", + ] + ), // ------ // Example programs @@ -334,8 +376,20 @@ let package = Package( path: "Source/Examples", sources: ["CustomFunctionExampleSimple.swift"] ), - ], - cxxLanguageStandard: .gnucxx17 + // 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 ) if Context.environment["MLX_SWIFT_BUILD_DOC"] == "1" diff --git a/Plugins/CudaBuild/plugin.swift b/Plugins/CudaBuild/plugin.swift new file mode 100644 index 000000000..15ecd2e88 --- /dev/null +++ b/Plugins/CudaBuild/plugin.swift @@ -0,0 +1,248 @@ +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: 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] + ) + ) + } + + // 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: linkArguments( + clangpp: clangUrl.url.path, + verboseFlag: verboseFlag, + stdArgs: stdArgs, + incrementalFlag: incrementalFlag, + inputs: outputCpps.map { $0.path }, + output: linkOutput.path + ), + inputFiles: outputCpps, + outputFiles: [linkOutput] + ) + ) + + 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 + "/") + 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/.clang-format b/Source/.clang-format deleted file mode 100644 index eab4576fe..000000000 --- a/Source/.clang-format +++ /dev/null @@ -1,87 +0,0 @@ ---- -AccessModifierOffset: -1 -AlignAfterOpenBracket: AlwaysBreak -AlignConsecutiveAssignments: false -AlignConsecutiveDeclarations: false -AlignEscapedNewlinesLeft: true -AlignOperands: false -AlignTrailingComments: false -AllowAllParametersOfDeclarationOnNextLine: false -AllowShortBlocksOnASingleLine: false -AllowShortCaseLabelsOnASingleLine: false -AllowShortFunctionsOnASingleLine: Empty -AllowShortIfStatementsOnASingleLine: false -AllowShortLoopsOnASingleLine: false -AlwaysBreakAfterReturnType: None -AlwaysBreakBeforeMultilineStrings: true -AlwaysBreakTemplateDeclarations: true -BinPackArguments: false -BinPackParameters: false -BraceWrapping: - AfterClass: false - AfterControlStatement: false - AfterEnum: false - AfterFunction: false - AfterNamespace: false - AfterObjCDeclaration: false - AfterStruct: false - AfterUnion: false - BeforeCatch: false - BeforeElse: false - IndentBraces: false -BreakBeforeBinaryOperators: None -BreakBeforeBraces: Attach -BreakBeforeTernaryOperators: true -BreakConstructorInitializersBeforeComma: false -BreakAfterJavaFieldAnnotations: false -BreakStringLiterals: false -ColumnLimit: 80 -CommentPragmas: '^ IWYU pragma:' -ConstructorInitializerAllOnOneLineOrOnePerLine: true -ConstructorInitializerIndentWidth: 4 -ContinuationIndentWidth: 4 -Cpp11BracedListStyle: true -DerivePointerAlignment: false -DisableFormat: false -ForEachMacros: [ FOR_EACH, FOR_EACH_R, FOR_EACH_RANGE, ] -IncludeCategories: - - Regex: '^<.*\.h(pp)?>' - Priority: 1 - - Regex: '^<.*' - Priority: 2 - - Regex: '.*' - Priority: 3 -IndentCaseLabels: true -IndentWidth: 2 -IndentWrappedFunctionNames: false -KeepEmptyLinesAtTheStartOfBlocks: false -MacroBlockBegin: '' -MacroBlockEnd: '' -MaxEmptyLinesToKeep: 1 -NamespaceIndentation: None -ObjCBlockIndentWidth: 2 -ObjCSpaceAfterProperty: false -ObjCSpaceBeforeProtocolList: false -PenaltyBreakBeforeFirstCallParameter: 1 -PenaltyBreakComment: 300 -PenaltyBreakFirstLessLess: 120 -PenaltyBreakString: 1000 -PenaltyExcessCharacter: 1000000 -PenaltyReturnTypeOnItsOwnLine: 200 -PointerAlignment: Left -ReflowComments: true -SortIncludes: true -SpaceAfterCStyleCast: false -SpaceBeforeAssignmentOperators: true -SpaceBeforeParens: ControlStatements -SpaceInEmptyParentheses: false -SpacesBeforeTrailingComments: 1 -SpacesInAngles: false -SpacesInContainerLiterals: true -SpacesInCStyleCastParentheses: false -SpacesInParentheses: false -SpacesInSquareBrackets: false -Standard: Cpp11 -TabWidth: 8 -UseTab: Never -... diff --git a/Source/.github/ISSUE_TEMPLATE/bug_report.md b/Source/.github/ISSUE_TEMPLATE/bug_report.md deleted file mode 100644 index 22a085792..000000000 --- a/Source/.github/ISSUE_TEMPLATE/bug_report.md +++ /dev/null @@ -1,28 +0,0 @@ ---- -name: Bug report -about: Create a report about an issue you've encountered -title: "[BUG] " -labels: '' -assignees: '' - ---- - -**Describe the bug** -A clear and concise description of what the bug is. - -**To Reproduce** - -Include code snippet -```python - -``` - -**Expected behavior** -A clear and concise description of what you expected to happen. - -**Desktop (please complete the following information):** - - OS Version: [e.g. MacOS 14.1.2] - - Version [e.g. 0.7.0] - -**Additional context** -Add any other context about the problem here. diff --git a/Source/.github/actions/build-cuda-release/action.yml b/Source/.github/actions/build-cuda-release/action.yml deleted file mode 100644 index 0ed7719d5..000000000 --- a/Source/.github/actions/build-cuda-release/action.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: 'Build CUDA wheel' -description: 'Build CUDA wheel' - -inputs: - arch: - description: 'Platform architecture tag' - required: true - type: choice - options: - - x86_64 - - aarch64 - -runs: - using: "composite" - steps: - - name: Build package - shell: bash - env: - CMAKE_ARGS: -DMLX_BUILD_CUDA=ON - run: | - pip install auditwheel build patchelf setuptools - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w - - auditwheel repair dist/mlx_cuda*.whl \ - --plat manylinux_2_35_${{ inputs.arch }} \ - --exclude libcublas* \ - --exclude libcuda* \ - --exclude libcudnn* \ - --exclude libnccl* \ - --exclude libnvrtc* diff --git a/Source/.github/actions/build-docs/action.yml b/Source/.github/actions/build-docs/action.yml deleted file mode 100644 index 411f6be8d..000000000 --- a/Source/.github/actions/build-docs/action.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: 'Build Documentation' -description: 'Build documentation' - -runs: - using: "composite" - steps: - - name: Setup machine - uses: ./.github/actions/setup-linux - - - name: Install dependencies - shell: bash - run: | - sudo apt-get install -y doxygen - source .venv/bin/activate - pip install -r docs/requirements.txt - pip install . -v - - - name: Build documentation - shell: bash - run: | - source .venv/bin/activate - cd docs - doxygen - make html O=-W - - - name: Create artifact tar - shell: bash - run: tar -cf artifact.tar -C docs --dereference build/html index.html - - # Do it manually because upload-pages-artifact requires gtar - - name: Upload artifact - id: upload-artifact - uses: actions/upload-artifact@v5 - with: - name: github-pages - path: artifact.tar - retention-days: 1 - if-no-files-found: error diff --git a/Source/.github/actions/build-linux-release/action.yml b/Source/.github/actions/build-linux-release/action.yml deleted file mode 100644 index c8dec00b4..000000000 --- a/Source/.github/actions/build-linux-release/action.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: 'Build Linux wheel' -description: 'Build Linux wheel' - -inputs: - build-backend: - description: 'Build the backend mlx-cpu package' - type: boolean - required: false - default: false - arch: - description: 'Platform architecture tag' - required: true - type: choice - options: - - x86_64 - - aarch64 - -runs: - using: "composite" - steps: - - name: Build MLX - shell: bash - run: pip install -e . -v - - - name: Build Python package - shell: bash - run: | - pip install auditwheel patchelf build - python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w - auditwheel repair dist/mlx-*.whl \ - --plat manylinux_2_35_${{ inputs.arch }} \ - --exclude libmlx.so* \ - --only-plat - - - name: Build backend package - if: ${{ inputs.build-backend }} - shell: bash - run: | - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w - auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_${{ inputs.arch }} diff --git a/Source/.github/actions/build-linux/action.yml b/Source/.github/actions/build-linux/action.yml deleted file mode 100644 index 166d7b8a6..000000000 --- a/Source/.github/actions/build-linux/action.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: 'Build and Test on Linux' - -inputs: - toolkit: - description: 'The toolkit to build with' - required: false - default: 'cpu' - -runs: - using: "composite" - steps: - - - name: Install Python package - id: python_build - shell: sh - env: - DEBUG: 1 - CMAKE_ARGS: >- - -DCMAKE_COMPILE_WARNING_AS_ERROR=ON - -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }} - run: | - if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then - # There is no GPU in arm64 runner, use a common arch. - CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a" - # Can not build tests and stubs when the built executables can not run. - CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF" - fi - # Install cpu-only torch to save space - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install --no-build-isolation -e ".[dev]" -v - # Pass the CMAKE_ARGS to following steps. - echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT - - - name: Build CPP only - shell: bash - run: | - cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }} - cmake --build build -j $(nproc) diff --git a/Source/.github/actions/build-macos-release/action.yml b/Source/.github/actions/build-macos-release/action.yml deleted file mode 100644 index 93e6dd943..000000000 --- a/Source/.github/actions/build-macos-release/action.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: 'Build macOS release' -description: 'Build MLX releases macOS' - -inputs: - macos-target: - description: 'macOS build target' - required: false - default: '15.0' - build-backend: - description: 'Build the backend mlx-metal package' - type: boolean - required: false - default: false - -runs: - using: "composite" - steps: - - name: Build Python package - shell: bash -l {0} - env: - DEVELOPER_DIR: /Applications/Xcode-latest.app - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - pip install build - python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w - - - name: Build backend package - if: ${{ inputs.build-backend }} - shell: bash -l {0} - env: - DEVELOPER_DIR: /Applications/Xcode-latest.app - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w diff --git a/Source/.github/actions/build-macos/action.yml b/Source/.github/actions/build-macos/action.yml deleted file mode 100644 index f04ed3d63..000000000 --- a/Source/.github/actions/build-macos/action.yml +++ /dev/null @@ -1,80 +0,0 @@ -name: 'Build and Test on macOS' -description: 'Build and test MLX on macOS' - -runs: - using: "composite" - steps: - - name: Install dependencies - env: - DEBUG: 1 - CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" - shell: bash -l {0} - run: | - pip install --upgrade pip - pip install cmake setuptools typing_extensions - pip install -e . -v - - - name: Install tests dependencies - shell: bash -l {0} - run: | - pip install numpy torch tensorflow - - - name: Run Python tests - shell: bash -l {0} - env: - LOW_MEMORY: 1 - run: | - DEVICE=cpu python -m unittest discover -v python/tests - DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests - mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py - mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) - if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi - - - name: Build example extension - shell: bash -l {0} - run: | - cd examples/extensions - pip install -r requirements.txt - python setup.py build_ext --inplace - python test.py - - - name: Build CPP only - shell: bash -l {0} - run: | - mkdir -p build - cd build - cmake .. - make -j $(sysctl -n hw.ncpu) - - - name: Run CPP tests - shell: bash -l {0} - env: - DEVICE: gpu - METAL_DEVICE_WRAPPER_TYPE: 1 - METAL_DEBUG_ERROR_MODE: 0 - run: ./build/tests/tests - - - name: Build small binary with JIT - shell: bash -l {0} - run: | - mkdir -p build - cd build - cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \ - -DBUILD_SHARED_LIBS=ON \ - -DMLX_BUILD_CPU=OFF \ - -DMLX_BUILD_SAFETENSORS=OFF \ - -DMLX_BUILD_GGUF=OFF \ - -DMLX_METAL_JIT=ON - make -j $(sysctl -n hw.ncpu) - - - name: Run Python tests with JIT - shell: bash -l {0} - env: - LOW_MEMORY: 1 - DEVICE: gpu - METAL_DEVICE_WRAPPER_TYPE: 1 - METAL_DEBUG_ERROR_MODE: 0 - run: | - CMAKE_ARGS="-DMLX_METAL_JIT=ON" \ - pip install -e . -v - python -m unittest discover -v python/tests diff --git a/Source/.github/actions/build-windows/action.yml b/Source/.github/actions/build-windows/action.yml deleted file mode 100644 index 372f84815..000000000 --- a/Source/.github/actions/build-windows/action.yml +++ /dev/null @@ -1,26 +0,0 @@ -name: 'Build on Windows' - -runs: - using: 'composite' - steps: - - name: Install Python package - id: python-build - shell: cmd - env: - # For MSVC, Ninja/Release is the only config supported by ccache. - CMAKE_ARGS: >- - -G Ninja - -DCMAKE_BUILD_TYPE=Release - -DCMAKE_C_COMPILER=cl - -DCMAKE_CXX_COMPILER=cl - -DCMAKE_RC_COMPILER=rc - run: | - uv pip install ".[dev]" -v - :: Pass the CMAKE_ARGS to following steps. - >>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS% - - - name: Build CPP only - shell: cmd - run: | - cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }} - cmake --build build -j %NUMBER_OF_PROCESSORS% diff --git a/Source/.github/actions/setup-linux/action.yml b/Source/.github/actions/setup-linux/action.yml deleted file mode 100644 index 08bb528f5..000000000 --- a/Source/.github/actions/setup-linux/action.yml +++ /dev/null @@ -1,93 +0,0 @@ -name: 'Setup Linux Environment' -description: 'Install dependencies for Linux builds' - -inputs: - toolkit: - description: 'Which toolkit to install' - required: false - default: 'cpu' - python-version: - description: 'Version of python to set up' - required: false - default: '3.14' - use-ccache: - description: 'Whether to enable ccache' - required: false - default: 'true' - -runs: - using: "composite" - steps: - - name: Install common dependencies - shell: bash - run: | - echo "::group::Install common dependencies" - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - zip \ - libblas-dev liblapack-dev liblapacke-dev \ - openmpi-bin openmpi-common libopenmpi-dev - echo "::endgroup::" - - - name: Use ccache - if: ${{ inputs.use-ccache == 'true' }} - uses: hendrikmuhs/ccache-action@v1.2 - with: - key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }} - max-size: 1GB - # ccache-action bug: running "apt-get update" fails on large arm runner. - update-package-index: false - - - uses: actions/setup-python@v6 - with: - python-version: ${{ inputs.python-version }} - - - name: Setup Python venv - shell: bash - run: | - echo "::group::Setup Python venv" - python -m venv .venv - source .venv/bin/activate - pip install setuptools cmake typing_extensions - echo PATH=$PATH >> $GITHUB_ENV - # Search python packages in .venv - echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV - echo "::endgroup::" - - - name: Install CUDA toolkit - if: ${{ startsWith(inputs.toolkit, 'cuda') }} - shell: bash - env: - # Note: the CI machine does not meet CUDA 13's driver requirement. - # Compatibility matrix: - # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html - PACKAGES: | - { - "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6", - "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9", - "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0" - } - run: | - echo "::group::Install CUDA toolkit" - # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is - # Jetson specific. SBSA means Arm Server Base System Architecture. - ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }} - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - libnccl2 libnccl-dev \ - ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} - echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH - echo "::endgroup::" - - - name: CUDA packages and driver report - if: ${{ startsWith(inputs.toolkit, 'cuda') }} - shell: bash - run: | - echo "::group::Installed NVIDIA and CUDA packages" - dpkg -l | egrep "cuda|nvidia" -i - echo "::endgroup::" - echo "::group::NVIDIA-SMI Status" - nvidia-smi || true - echo "::endgroup::" diff --git a/Source/.github/actions/setup-macos/action.yml b/Source/.github/actions/setup-macos/action.yml deleted file mode 100644 index a108e686f..000000000 --- a/Source/.github/actions/setup-macos/action.yml +++ /dev/null @@ -1,24 +0,0 @@ -name: 'Setup macOS Environment' -description: 'Install dependencies for macOS builds' - -inputs: - python-version: - description: 'Python version to use' - required: false - default: '3.10' - -runs: - using: "composite" - steps: - - name: Install Homebrew packages - shell: sh - run: /opt/homebrew/bin/brew install openmpi - - - name: Verify MetalToolchain installed - shell: bash - run: xcodebuild -showComponent MetalToolchain - - - uses: conda-incubator/setup-miniconda@v3 - with: - miniconda-version: "latest" - python-version: ${{ inputs.python-version }} diff --git a/Source/.github/actions/setup-windows/action.yml b/Source/.github/actions/setup-windows/action.yml deleted file mode 100644 index 83afd4de1..000000000 --- a/Source/.github/actions/setup-windows/action.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: 'Setup Windows environment' - -inputs: - python-version: - description: 'Version of python to set up' - required: false - default: '3.14' - use-ccache: - description: 'Whether to enable ccache' - required: false - default: 'true' - -runs: - using: 'composite' - steps: - - name: Use ccache - if: ${{ inputs.use-ccache == 'true' }} - uses: hendrikmuhs/ccache-action@v1.2 - with: - key: ccache-${{ runner.os }}-${{ runner.arch }}-cpu - max-size: 1GB - - - name: Setup Visual Studio cmd - shell: cmd - run: | - :: Find out path to VS. - pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\" - for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x - popd - :: Import VS vars. - call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" x64 - :: Export to all steps. - >>%GITHUB_ENV% set - - - uses: astral-sh/setup-uv@v7 - - - name: Setup Python venv - shell: cmd - run: | - uv venv --python ${{ inputs.python-version }} - call ".venv/Scripts/activate.bat" - >>%GITHUB_ENV% set diff --git a/Source/.github/actions/test-linux/action.yml b/Source/.github/actions/test-linux/action.yml deleted file mode 100644 index 3258670ef..000000000 --- a/Source/.github/actions/test-linux/action.yml +++ /dev/null @@ -1,69 +0,0 @@ -name: 'Run Linux tests' - -inputs: - has-gpu: - description: 'Run GPU tests' - required: false - default: false - -runs: - using: "composite" - steps: - - name: Run MPI tests - shell: bash - run: | - echo "::group::MPI tests" - mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - echo "::endgroup::" - - - name: Run distributed tests - if: ${{ inputs.has-gpu == 'false' }} - shell: bash - run: | - echo "::group::Distributed tests" - mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) - if grep -Fq '[WARN]' stderr.log ; then - grep -F '[WARN]' stderr.log - echo "Distributed ring test failed"; - exit 1; - fi - echo "::endgroup::" - - - name: Run Python tests - CPU - if: ${{ inputs.has-gpu == 'false' }} - shell: bash - env: - DEVICE: cpu - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" - - - name: Run Python tests - GPU - if: ${{ inputs.has-gpu == 'true' }} - shell: bash - env: - DEVICE: gpu - run: | - echo "::group::Python tests - GPU" - python -m tests discover python/tests -v - echo "::endgroup::" - - - name: Run CPP tests - CPU - shell: bash - env: - DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/tests/tests - echo "::endgroup::" - - - name: Run CPP tests - GPU - if: ${{ inputs.has-gpu == 'true' }} - shell: bash - env: - DEVICE: gpu - run: | - echo "::group::CPP tests - GPU" - ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp" - echo "::endgroup::" diff --git a/Source/.github/actions/test-windows/action.yml b/Source/.github/actions/test-windows/action.yml deleted file mode 100644 index ebb886e7b..000000000 --- a/Source/.github/actions/test-windows/action.yml +++ /dev/null @@ -1,20 +0,0 @@ -name: 'Run tests on Windows' - -runs: - using: 'composite' - steps: - - name: Run Python tests - CPU - shell: bash - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" - - - name: Run CPP tests - CPU - shell: bash - env: - DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/tests.exe -tce="*gguf*,test random uniform" - echo "::endgroup::" diff --git a/Source/.github/dependabot.yml b/Source/.github/dependabot.yml deleted file mode 100644 index 5ace4600a..000000000 --- a/Source/.github/dependabot.yml +++ /dev/null @@ -1,6 +0,0 @@ -version: 2 -updates: - - package-ecosystem: "github-actions" - directory: "/" - schedule: - interval: "weekly" diff --git a/Source/.github/pull_request_template.md b/Source/.github/pull_request_template.md deleted file mode 100644 index 02bb9b79a..000000000 --- a/Source/.github/pull_request_template.md +++ /dev/null @@ -1,12 +0,0 @@ -## Proposed changes - -Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #. - -## Checklist - -Put an `x` in the boxes that apply. - -- [ ] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document -- [ ] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes -- [ ] I have added tests that prove my fix is effective or that my feature works -- [ ] I have updated the necessary documentation (if needed) diff --git a/Source/.github/scripts/build-sanitizer-tests.sh b/Source/.github/scripts/build-sanitizer-tests.sh deleted file mode 100755 index 9c911890b..000000000 --- a/Source/.github/scripts/build-sanitizer-tests.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/bin/bash -set -ex - -export CMAKE_C_COMPILER=/usr/bin/clang -export CMAKE_CXX_COMPILER=/usr/bin/clang++ -BASE_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=DEBUG -DCMAKE_COMPILE_WARNING_AS_ERROR=ON" -if [[ "$(uname -s)" != "Darwin" ]]; then - BASE_CMAKE_ARGS+=" -DMLX_BUILD_METAL=OFF" -fi - -run_test() { - local sanitizer_name=$1 - local cmake_sanitizer_flag="-DUSE_${sanitizer_name}=ON" - echo " Running tests with: ${sanitizer_name}" - - case "$sanitizer_name" in - ASAN) - export ASAN_OPTIONS="detect_leaks=0" - ;; - UBSAN) - export UBSAN_OPTIONS="halt_on_error=0:print_stacktrace=1" - ;; - TSAN) - export TSAN_OPTIONS="" - ;; - esac - - rm -rf build - mkdir -p build - pushd build > /dev/null - - cmake .. ${BASE_CMAKE_ARGS} ${cmake_sanitizer_flag} - make -j $(nproc) - ./tests/tests - - popd > /dev/null - unset ${sanitizer_name}_OPTIONS -} - -sanitizer_arg=$(echo "$1" | tr '[:lower:]' '[:upper:]') - -if [[ "$sanitizer_arg" == "ASAN" || "$sanitizer_arg" == "UBSAN" || "$sanitizer_arg" == "TSAN" ]]; then - run_test "$sanitizer_arg" - echo " ${sanitizer_arg} test run completed successfully." -else - echo "Error: Invalid sanitizer '$1'. Please use one of: ASAN, UBSAN, TSAN." - exit 1 -fi diff --git a/Source/.github/scripts/setup+build-cpp-linux-fedora-container.sh b/Source/.github/scripts/setup+build-cpp-linux-fedora-container.sh deleted file mode 100755 index e93757d47..000000000 --- a/Source/.github/scripts/setup+build-cpp-linux-fedora-container.sh +++ /dev/null @@ -1,27 +0,0 @@ -#!/bin/bash -set -ex - -# [Setup] Install dependencies inside the container. -dnf update -y -dnf install -y \ - blas-devel \ - lapack-devel \ - openblas-devel \ - make \ - cmake \ - clang \ - git -dnf clean all - -# [C++] CI Build Sanity Check: Verifies code compilation, not for release. -export CMAKE_ARGS="-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" -export DEBUG=1 -export CMAKE_C_COMPILER=/usr/bin/clang -export CMAKE_CXX_COMPILER=/usr/bin/clang++ - -mkdir -p build -pushd build -cmake .. -DMLX_BUILD_METAL=OFF -DCMAKE_BUILD_TYPE=DEBUG -make -j $(nproc) -./tests/tests -popd diff --git a/Source/.github/workflows/build_and_test.yml b/Source/.github/workflows/build_and_test.yml deleted file mode 100644 index ba0ccfa68..000000000 --- a/Source/.github/workflows/build_and_test.yml +++ /dev/null @@ -1,152 +0,0 @@ -name: Build and Test - -on: - pull_request: - push: - branches: - - main - # For testing CI without starting a pull request: - - test/* - -permissions: - contents: read - -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: ${{ github.ref != 'refs/heads/main' }} - -jobs: - check_lint: - name: Check Lint - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v6 - - uses: pre-commit/action@v3.0.1 - - linux_build_and_test: - name: Linux (cpu, ${{ matrix.arch }}) - needs: check_lint - strategy: - fail-fast: false - matrix: - arch: ['x86_64', 'aarch64'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - - uses: ./.github/actions/build-linux - - uses: ./.github/actions/test-linux - - run: df -h - - cuda_build_and_test: - name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }}) - if: github.repository == 'ml-explore/mlx' - needs: check_lint - strategy: - fail-fast: false - matrix: - arch: ['x86_64', 'aarch64'] - toolkit: ['cuda-12.6', 'cuda-12.9'] - runs-on: ${{ matrix.arch == 'x86_64' && 'gpu-t4-4-core' || 'ubuntu-22.04-arm' }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - toolkit: ${{ matrix.toolkit }} - - uses: ./.github/actions/build-linux - with: - toolkit: ${{ matrix.toolkit }} - - uses: ./.github/actions/test-linux - if: matrix.arch == 'x86_64' - with: - has-gpu: true - - mac_build_and_test: - name: macOS (${{ matrix.macos-target }}) - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - macos-target: ["14.0", "15.0", "26.0"] - runs-on: [self-hosted, macos] - env: - MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }} - needs: check_lint - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos - - uses: ./.github/actions/build-macos - - windows_build_and_test: - name: Windows (cpu, x86_64) - needs: check_lint - runs-on: windows-2025 - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-windows - - uses: ./.github/actions/build-windows - - uses: ./.github/actions/test-windows - - build_documentation: - name: Build Documentation - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22.04 - needs: check_lint - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/build-docs - - linux_sanitizer_build_and_test: - name: Linux Sanitizer Tests (${{ matrix.sanitizer }}) - needs: check_lint - strategy: - fail-fast: false - matrix: - sanitizer: [ASAN, UBSAN] - # todo 12/16/2025: enable TSAN later + consider enabling ASAN for GPU backend tests. - # sanitizer: [ASAN, UBSAN, TSAN] - runs-on: ubuntu-22.04-arm - steps: - - name: Checkout code - uses: actions/checkout@v6 - - - name: Install Dependencies - run: | - export DEBIAN_FRONTEND=noninteractive - sudo apt-get update -y - sudo apt-get install -y \ - build-essential \ - libblas-dev \ - liblapacke-dev \ - libopenblas-dev \ - cmake \ - clang \ - git - sudo apt-get clean - sudo rm -rf /var/lib/apt/lists/* - - - name: Linux Build and Test with ${{ matrix.sanitizer }} - run: | - bash .github/scripts/build-sanitizer-tests.sh ${{ matrix.sanitizer }} - - linux_fedora_build_cpp: - name: Linux Fedora (${{ matrix.arch }}) - needs: check_lint - strategy: - fail-fast: false - matrix: - include: - - host: ubuntu-22.04 - arch: x86_64 - - host: ubuntu-22.04-arm - arch: aarch64 - - runs-on: ${{ matrix.host }} - container: - image: fedora:42 - steps: - - name: Checkout code - uses: actions/checkout@v6 - - - name: CPP Build Test - No Release - run: | - bash ./.github/scripts/setup+build-cpp-linux-fedora-container.sh diff --git a/Source/.github/workflows/documentation.yml b/Source/.github/workflows/documentation.yml deleted file mode 100644 index a9af245fc..000000000 --- a/Source/.github/workflows/documentation.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: Documentation - -on: - workflow_dispatch: - -permissions: - contents: read - -jobs: - build: - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/build-docs - - deploy: - needs: build - permissions: - pages: write - id-token: write - runs-on: ubuntu-latest - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 diff --git a/Source/.github/workflows/nightly.yml b/Source/.github/workflows/nightly.yml deleted file mode 100644 index 84962e14b..000000000 --- a/Source/.github/workflows/nightly.yml +++ /dev/null @@ -1,104 +0,0 @@ -name: Nightly Build - -on: - schedule: - - cron: 33 6 * * 1-5 - workflow_dispatch: - -permissions: - contents: read - -jobs: - build_linux_release: - strategy: - fail-fast: false - matrix: - python_version: ["3.10", "3.14"] - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - - uses: ./.github/actions/build-linux-release - with: - build-backend: ${{ matrix.python-version == '3.10' }} - arch: "x86_64" - - name: Upload mlx artifacts - uses: actions/upload-artifact@v6 - with: - name: linux-wheels-${{ matrix.python_version }} - path: wheelhouse/mlx-*.whl - retention-days: 7 - - name: Upload mlx-cpu artifacts - if: matrix.python_version == '3.10' - uses: actions/upload-artifact@v6 - with: - name: mlx-cpu - path: wheelhouse/mlx_cpu-*.whl - retention-days: 7 - - run: df -h - - build_linux_with_tests: - strategy: - fail-fast: false - matrix: - python_version: ["3.11", "3.12", "3.13", "3.14"] - runner: - - ubuntu-22.04 - - ubuntu-22.04-arm - runs-on: ${{ matrix.runner }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - python-version: ${{ matrix.python_version }} - - uses: ./.github/actions/build-linux - - uses: ./.github/actions/test-linux - - run: df -h - - build_mac_release: - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - python-version: ["3.10", "3.13"] - runs-on: [self-hosted, macos] - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos - with: - python-version: ${{ matrix.python-version }} - - uses: ./.github/actions/build-macos - - name: Build macOS 26 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 26.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 15 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 15.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 14 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 14.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - build_cuda_release: - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22-large - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - toolkit: 'cuda-12.9' - - name: Build Python package - uses: ./.github/actions/build-cuda-release - with: - toolkit: 'cuda-12.9' - arch: 'x86_64' - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - name: mlx-cuda - path: wheelhouse/mlx_cuda_*.whl - retention-days: 7 diff --git a/Source/.github/workflows/release.yml b/Source/.github/workflows/release.yml deleted file mode 100644 index 9b9420a50..000000000 --- a/Source/.github/workflows/release.yml +++ /dev/null @@ -1,256 +0,0 @@ -name: PyPI Release - -on: - push: - tags: - - 'v*' - branches: - - 'test-publish/*' - workflow_dispatch: - inputs: - dry_run: - description: 'Dry run (do not publish to PyPi)' - required: false - type: boolean - dev_release: - description: 'Development release (DEV_RELEASE=1)' - required: false - type: boolean - -permissions: - contents: read - -jobs: - build_documentation: - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/build-docs - - deploy_documentation: - if: ${{ !inputs.dry_run }} - needs: build_documentation - permissions: - pages: write - id-token: write - runs-on: ubuntu-latest - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 - - build_linux_release: - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - arch: ['x86_64', 'aarch64'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} - env: - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - python-version: ${{ matrix.python_version }} - use-ccache: false - - uses: ./.github/actions/build-linux-release - with: - build-backend: ${{ matrix.python_version == '3.10' }} - arch: ${{ matrix.arch }} - - name: Upload MLX artifacts - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }} - path: wheelhouse/mlx-*.whl - if-no-files-found: error - - name: Upload CPU artifacts - if: matrix.python_version == '3.10' - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: mlx-cpu-${{ matrix.arch }} - path: wheelhouse/mlx_cpu-*.whl - if-no-files-found: error - - build_mac_release: - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - runs-on: [self-hosted, macos] - env: - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos - with: - python-version: ${{ matrix.python-version }} - - - name: Install dependencies - shell: bash -l {0} - run: | - pip install --upgrade pip - pip install cmake setuptools typing_extensions - pip install -e . -v - - name: Build macOS 14 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 14.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 15 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 15.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 26 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 26.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Upload MLX artifacts - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: mac-wheels-${{ matrix.python-version }} - path: dist/mlx-*.whl - if-no-files-found: error - - name: Upload Metal artifacts - if: matrix.python-version == '3.10' - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: mlx-metal - path: dist/mlx_metal-*.whl - if-no-files-found: error - - build_cuda_release: - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - arch: ['x86_64', 'aarch64'] - toolkit: ['cuda-12.9', 'cuda-13.0'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }} - env: - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - toolkit: ${{ matrix.toolkit }} - use-ccache: false - - name: Build Python package - uses: ./.github/actions/build-cuda-release - with: - arch: ${{ matrix.arch }} - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }} - path: wheelhouse/mlx_cuda_*.whl - if-no-files-found: error - - pypi-publish: - name: Upload release to PyPI - runs-on: ubuntu-latest - needs: [build_linux_release, build_mac_release] - permissions: - id-token: write - environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} - url: https://pypi.org/p/mlx - steps: - - uses: actions/download-artifact@v7 - with: - pattern: linux-wheels-* - merge-multiple: true - path: dist - - uses: actions/download-artifact@v7 - with: - pattern: mac-wheels-* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-cuda: - name: Upload CUDA release to PyPI - runs-on: ubuntu-latest - needs: [build_cuda_release] - permissions: - id-token: write - environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} - url: https://pypi.org/p/mlx-cuda - steps: - - uses: actions/download-artifact@v7 - with: - pattern: mlx-cuda-* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-cpu: - name: Upload CPU release to PyPI - runs-on: ubuntu-latest - needs: [build_linux_release] - permissions: - id-token: write - environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} - url: https://pypi.org/p/mlx-cpu - steps: - - uses: actions/download-artifact@v7 - with: - pattern: mlx-cpu-* - merge-multiple: true - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ - - pypi-publish-metal: - name: Upload Metal release to PyPI - runs-on: ubuntu-latest - needs: [build_mac_release] - permissions: - id-token: write - environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} - url: https://pypi.org/p/mlx-metal - steps: - - uses: actions/download-artifact@v7 - with: - name: mlx-metal - path: dist - - name: Display structure of downloaded files - run: du -ah dist - - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} - uses: pypa/gh-action-pypi-publish@release/v1 - with: - repository-url: https://upload.pypi.org/legacy/ diff --git a/Source/.gitignore b/Source/.gitignore deleted file mode 100644 index 1daaa46d1..000000000 --- a/Source/.gitignore +++ /dev/null @@ -1,81 +0,0 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# tensor files -*.safe -*.safetensors - -# Metal libraries -*.metallib - -# Distribution / packaging -python/mlx/core -python/mlx/share -python/mlx/include -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -venv/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST -uv.lock -.DS_Store - -# Prerequisites -*.d - -# Compiled Object files -*.slo -*.lo -*.o -*.obj -*.ilk - -# Precompiled Headers -*.gch -*.pch - -# Compiled Dynamic libraries -*.so -*.dylib -*.dll - -# Fortran module files -*.mod -*.smod - -# Compiled Static libraries -*.lai -*.la -*.a -*.lib - -# Executables -*.exe -*.out -*.app - -# Debug symbols -*.pdb - -# VSCode -.vscode/ -# Jetbrains -.cache/ -# vim -*.swp diff --git a/Source/.pre-commit-config.yaml b/Source/.pre-commit-config.yaml deleted file mode 100644 index 0345848e9..000000000 --- a/Source/.pre-commit-config.yaml +++ /dev/null @@ -1,27 +0,0 @@ -repos: -- repo: https://github.com/pre-commit/pre-commit-hooks - rev: v6.0.0 - hooks: - - id: check-yaml - # - id: end-of-file-fixer - # - id: trailing-whitespace -- repo: https://github.com/pre-commit/mirrors-clang-format - rev: v21.1.8 - hooks: - - id: clang-format -# Using this mirror lets us use mypyc-compiled black, which is about 2x faster -- repo: https://github.com/psf/black-pre-commit-mirror - rev: 26.1.0 - hooks: - - id: black - -- repo: https://github.com/pycqa/isort - rev: 7.0.0 - hooks: - - id: isort - args: - - --profile=black -- repo: https://github.com/cheshirekow/cmake-format-precommit - rev: v0.6.13 - hooks: - - id: cmake-format diff --git a/Source/ACKNOWLEDGMENTS.md b/Source/ACKNOWLEDGMENTS.md deleted file mode 100644 index 186908f09..000000000 --- a/Source/ACKNOWLEDGMENTS.md +++ /dev/null @@ -1,268 +0,0 @@ -# Individual Contributors - -If you wish to be acknowledged for your contributions, please list your name -with a short description of your contribution(s) below. For example: - -- Jane Smith: Added the `foo` and `bar` ops. - -MLX was developed with contributions from the following individuals: - -- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops. Added `clip_grad_norm` along with `tree_reduce`. Added `cross`. Added `orthogonal` initializer. -- Juarez Bochi: Fixed bug in cross attention. -- Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example. -- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`. -- Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented pooling layers and ``Upsample``. -- Hinrik Snær Guðmundsson: Added `atleast_1d`, `atleast_2d`, `atleast_3d` ops. -- Luca Arnaboldi: Added `Ceil` and `Floor` ops; implemented pickling, copy and deepcopy for mlx arrays. -- Brian Keene & Atila Orhon, with Argmax Inc.: Added `fast.scaled_dot_product_attention` -- AmirHossein Razlighi: Added chaining support for some of the ops in `nn.Module`. Comparison works for non array objects in `mlx.core.array`. Exception handling for invalid operations in `mlx.core.array`. -- Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions. -- Paul Paczuski: Improved stability of BCE loss calculation -- Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops. -- Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. - - - - - -# Organizations - -MLX has received contributions from the following companies: -- NVIDIA Corporation & Affiliates - -# Third-Party Software - -MLX leverages several third-party software, listed here together with -their license copied verbatim. - -## PocketFFT - -Copyright (C) 2010-2018 Max-Planck-Society -All rights reserved. - -Redistribution and use in source and binary forms, with or without modification, -are permitted provided that the following conditions are met: - -* Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. -* Redistributions in binary form must reproduce the above copyright notice, this - list of conditions and the following disclaimer in the documentation and/or - other materials provided with the distribution. -* Neither the name of the copyright holder nor the names of its contributors may - be used to endorse or promote products derived from this software without - specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED -WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON -ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - -## metal-cpp - - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - -TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - -1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - -2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - -3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - -4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - -5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - -6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - -7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - -8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - -9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - -END OF TERMS AND CONDITIONS - -APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - -Copyright © 2023 Apple Inc. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. diff --git a/Source/CITATION.cff b/Source/CITATION.cff deleted file mode 100644 index ef4168f43..000000000 --- a/Source/CITATION.cff +++ /dev/null @@ -1,24 +0,0 @@ -cff-version: 1.2.0 -title: mlx -message: >- - If you use this software, please cite it using the - metadata from this file. -type: software -authors: - - given-names: Awni - family-names: Hannun - affiliation: Apple - - given-names: Jagrit - family-names: Digani - affiliation: Apple - - given-names: Angelos - family-names: Katharopoulos - affiliation: Apple - - given-names: Ronan - family-names: Collobert - affiliation: Apple -repository-code: 'https://github.com/ml-explore' -abstract: >- - MLX: efficient and flexible machine learning on Apple - silicon -license: MIT diff --git a/Source/CMakeLists.txt b/Source/CMakeLists.txt deleted file mode 100644 index 041a476cf..000000000 --- a/Source/CMakeLists.txt +++ /dev/null @@ -1,437 +0,0 @@ -cmake_minimum_required(VERSION 3.25) - -if(NOT MLX_VERSION) - file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$") - string(REGEX MATCH "#define MLX_VERSION_MAJOR ([0-9]+)" _ "${_mlx_h_version}") - set(_major ${CMAKE_MATCH_1}) - string(REGEX MATCH "#define MLX_VERSION_MINOR ([0-9]+)" _ "${_mlx_h_version}") - set(_minor ${CMAKE_MATCH_1}) - string(REGEX MATCH "#define MLX_VERSION_PATCH ([0-9]+)" _ "${_mlx_h_version}") - set(_patch ${CMAKE_MATCH_1}) - set(MLX_PROJECT_VERSION "${_major}.${_minor}.${_patch}") - set(MLX_VERSION ${MLX_PROJECT_VERSION}) -else() - string(REGEX REPLACE "^([0-9]+\.[0-9]+\.[0-9]+).*" "\\1" MLX_PROJECT_VERSION - ${MLX_VERSION}) -endif() - -project( - mlx - LANGUAGES C CXX - VERSION ${MLX_PROJECT_VERSION}) - -# ----------------------------- Setup ----------------------------- -set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake") -set(CMAKE_CXX_STANDARD 20) -set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_POSITION_INDEPENDENT_CODE ON) -set(CMAKE_INSTALL_MESSAGE NEVER) -set(CMAKE_EXPORT_COMPILE_COMMANDS ON) - -# ----------------------------- Configuration ----------------------------- -option(MLX_BUILD_TESTS "Build tests for mlx" ON) -option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON) -option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF) -option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF) -option(MLX_BUILD_METAL "Build metal backend" ON) -option(MLX_BUILD_CPU "Build cpu backend" ON) -option(MLX_BUILD_CUDA "Build cuda backend" OFF) -option(MLX_METAL_DEBUG "Enhance metal debug workflow" OFF) -option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF) -option(MLX_BUILD_GGUF "Include support for GGUF format" ON) -option(MLX_BUILD_SAFETENSORS "Include support for safetensors format" ON) -option(MLX_BUILD_PYTHON_STUBS "Build stub files for python bindings" ON) -option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF) -option(MLX_USE_CCACHE "Use CCache for compilation cache when available" ON) -option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF) -option(USE_SYSTEM_FMT "Use system's provided fmt library" OFF) -option(USE_ASAN "Enable AddressSanitizer (ASan)" OFF) -option(USE_UBSAN "Enable UndefinedBehaviorSanitizer (UBSan)" OFF) -option(USE_TSAN "Enable ThreadSanitizer (TSan)" OFF) - -# --------------------- Processor tests ------------------------- -message( - STATUS - "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}" -) - -if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin") - if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64") - if(NOT MLX_ENABLE_X64_MAC) - message( - FATAL_ERROR - "Building for x86_64 on macOS is not supported." - " If you are on an Apple silicon system, check the build" - " documentation for possible fixes: " - "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source" - ) - else() - set(MLX_BUILD_METAL OFF) - message(WARNING "Building for x86_64 arch is not officially supported.") - endif() - endif() -else() - set(MLX_BUILD_METAL OFF) -endif() - -if(MLX_USE_CCACHE) - find_program(CCACHE_PROGRAM ccache) - if(CCACHE_PROGRAM) - message(STATUS "Found CCache: ${CCACHE_PROGRAM}") - set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}") - set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}") - set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}") - endif() -endif() - -if(USE_ASAN AND USE_TSAN) - message( - FATAL_ERROR - "AddressSanitizer (ASan) and ThreadSanitizer (TSan) are mutually exclusive and cannot be enabled at the same time." - ) -endif() - -set(SANITIZER_COMPILE_FLAGS "") -set(SANITIZER_LINK_FLAGS "") - -if(USE_ASAN) - if(WIN32 AND MSVC) - list(APPEND SANITIZER_COMPILE_FLAGS /fsanitize=address) - list(APPEND SANITIZER_LINK_FLAGS /fsanitize=address) - else() - list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=address) - list(APPEND SANITIZER_LINK_FLAGS -fsanitize=address) - if(CMAKE_SYSTEM_NAME STREQUAL "Linux") - list(APPEND SANITIZER_LINK_FLAGS -lpthread) - endif() - endif() -endif() - -if(USE_UBSAN) - if(WIN32 AND MSVC) - if(CMAKE_CXX_COMPILER_ID STREQUAL "Clang") - list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined) - list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined) - else() - message( - WARNING - "UndefinedBehaviorSanitizer (UBSan) is not directly supported via a simple flag in MSVC." - ) - endif() - else() - list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=undefined) - list(APPEND SANITIZER_LINK_FLAGS -fsanitize=undefined) - endif() -endif() - -if(USE_TSAN) - if(WIN32 AND MSVC) - message( - FATAL_ERROR - "ThreadSanitizer (TSan) is not supported by the MSVC compiler. Please use Clang or GCC." - ) - elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") - message(FATAL_ERROR "ThreadSanitizer (TSan) is not supported on macOS.") - else() - list(APPEND SANITIZER_COMPILE_FLAGS -fsanitize=thread) - list(APPEND SANITIZER_LINK_FLAGS -fsanitize=thread) - if(CMAKE_SYSTEM_NAME STREQUAL "Linux") - list(APPEND SANITIZER_LINK_FLAGS -lpthread) - endif() - endif() -endif() - -# ----------------------------- Lib ----------------------------- - -include(FetchContent) -# Avoid warning about DOWNLOAD_EXTRACT_TIMESTAMP in CMake 3.24: -cmake_policy(SET CMP0135 NEW) - -add_library(mlx) - -target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS}) -target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS}) - -if(MLX_BUILD_CUDA) - enable_language(CUDA) - find_package(CUDAToolkit REQUIRED) - find_package(CUDNN REQUIRED) -endif() - -if(MLX_BUILD_METAL) - find_library(METAL_LIB Metal) - find_library(FOUNDATION_LIB Foundation) - find_library(QUARTZ_LIB QuartzCore) - if(METAL_LIB) - message(STATUS "Metal found ${METAL_LIB}") - else() - message( - FATAL_ERROR - "Metal not found. Set MLX_BUILD_METAL=OFF to build without GPU") - endif() - - if(MLX_METAL_DEBUG) - add_compile_definitions(MLX_METAL_DEBUG) - endif() - - # Throw an error if xcrun not found - execute_process( - COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version" - OUTPUT_VARIABLE MACOS_SDK_VERSION - OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) - - if(${MACOS_SDK_VERSION} LESS 14.0) - message( - FATAL_ERROR - "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON") - endif() - message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}") - - set(METAL_CPP_URL - https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip) - - if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") - if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0) - message(FATAL_ERROR "MLX requires macOS >= 14.0") - endif() - set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") - endif() - execute_process( - COMMAND - zsh "-c" - "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'" - OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY) - FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL}) - FetchContent_MakeAvailable(metal_cpp) - target_include_directories( - mlx PUBLIC $ - $) - target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB}) -endif() - -if(CMAKE_SYSTEM_NAME STREQUAL "Linux") - # With newer clang/gcc versions following libs are implicitly linked, but when - # building on old distributions they need to be explicitly listed. - target_link_libraries(mlx PRIVATE dl pthread) -endif() - -if(WIN32) - if(MSVC) - # GGUF does not build with MSVC. - set(MLX_BUILD_GGUF OFF) - endif() - # Generate DLL and EXE in the same dir, otherwise EXE will not be able to run. - # This is only done when MLX is built as the top project. - if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR) - set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) - endif() - # Windows implementation of dlfcn.h APIs. - FetchContent_Declare( - dlfcn-win32 - GIT_REPOSITORY https://github.com/dlfcn-win32/dlfcn-win32.git - GIT_TAG v1.4.2 - EXCLUDE_FROM_ALL) - block() - set(BUILD_SHARED_LIBS OFF) - FetchContent_MakeAvailable(dlfcn-win32) - endblock() - target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src") - target_link_libraries(mlx PRIVATE dl) -endif() - -if(MLX_BUILD_CPU) - find_library(ACCELERATE_LIBRARY Accelerate) - if(ACCELERATE_LIBRARY) - message(STATUS "Accelerate found ${ACCELERATE_LIBRARY}") - set(MLX_BUILD_ACCELERATE ON) - else() - message(STATUS "Accelerate not found, using default backend.") - set(MLX_BUILD_ACCELERATE OFF) - endif() - - if(MLX_BUILD_ACCELERATE) - target_link_libraries(mlx PUBLIC ${ACCELERATE_LIBRARY}) - add_compile_definitions(MLX_USE_ACCELERATE) - add_compile_definitions(ACCELERATE_NEW_LAPACK) - elseif(WIN32) - # Download and link prebuilt binaries of OpenBLAS. Note that we can only - # link with the dynamic library, the prebuilt binaries were built with MinGW - # so static-linking would require linking with MinGW's runtime. - FetchContent_Declare( - openblas - URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip" - ) - FetchContent_MakeAvailable(openblas) - target_link_libraries(mlx - PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib") - target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include") - # Make sure the DLL file is placed in the same dir with executables. - set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll") - add_custom_command( - TARGET mlx - POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different ${OPENBLAS_DLL_FILE} - ${CMAKE_BINARY_DIR}) - else() - if(${CMAKE_HOST_APPLE}) - # The blas shipped in macOS SDK is not supported, search homebrew for - # openblas instead. - set(BLA_VENDOR OpenBLAS) - set(LAPACK_ROOT - "${LAPACK_ROOT};$ENV{LAPACK_ROOT};/usr/local/opt/openblas") - endif() - # Search and link with lapack. - find_package(LAPACK REQUIRED) - if(NOT LAPACK_FOUND) - message(FATAL_ERROR "Must have LAPACK installed") - endif() - find_path(LAPACK_INCLUDE_DIRS lapacke.h /usr/include /usr/local/include - /usr/local/opt/openblas/include) - message(STATUS "Lapack lib " ${LAPACK_LIBRARIES}) - message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS}) - target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS}) - target_link_libraries(mlx PRIVATE ${LAPACK_LIBRARIES}) - # List blas after lapack otherwise we may accidentally incldue an old - # version of lapack.h from the include dirs of blas. - find_package(BLAS REQUIRED) - if(NOT BLAS_FOUND) - message(FATAL_ERROR "Must have BLAS installed") - endif() - # TODO find a cleaner way to do this - find_path(BLAS_INCLUDE_DIRS cblas.h /usr/include /usr/local/include - $ENV{BLAS_HOME}/include) - message(STATUS "Blas lib " ${BLAS_LIBRARIES}) - message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS}) - target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS}) - target_link_libraries(mlx PRIVATE ${BLAS_LIBRARIES}) - endif() -else() - set(MLX_BUILD_ACCELERATE OFF) -endif() - -message(STATUS "Downloading json") -FetchContent_Declare( - json - URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz) -FetchContent_MakeAvailable(json) -target_include_directories( - mlx PRIVATE $) - -add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx) - -target_include_directories( - mlx PUBLIC $ - $) - -if(USE_SYSTEM_FMT) - find_package(fmt REQUIRED) -else() - FetchContent_Declare( - fmt - GIT_REPOSITORY https://github.com/fmtlib/fmt.git - GIT_TAG 12.1.0 - EXCLUDE_FROM_ALL) - FetchContent_MakeAvailable(fmt) -endif() -target_link_libraries(mlx PRIVATE $) - -if(MLX_BUILD_PYTHON_BINDINGS) - message(STATUS "Building Python bindings.") - find_package( - Python 3.10 - COMPONENTS Interpreter Development.Module - REQUIRED) - FetchContent_Declare( - nanobind - GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v2.10.2 - GIT_SHALLOW TRUE - EXCLUDE_FROM_ALL) - FetchContent_MakeAvailable(nanobind) - add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/python/src) -endif() - -if(MLX_BUILD_TESTS) - include(CTest) - add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/tests) -endif() - -if(MLX_BUILD_EXAMPLES) - add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/examples/cpp) -endif() - -if(MLX_BUILD_BENCHMARKS) - add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/benchmarks/cpp) -endif() - -# ----------------------------- Installation ----------------------------- -include(GNUInstallDirs) - -if(WIN32) - # Install DLLs to the same dir with extension file (core.pyd) on Windows. - set(CMAKE_INSTALL_BINDIR ".") - if(MLX_BUILD_CPU) - # Install OpenBLAS. - install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN) - endif() -endif() - -# Install library -install( - TARGETS mlx - EXPORT MLXTargets - LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} - ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} - RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR} - INCLUDES - DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) - -# Install headers -install( - DIRECTORY ${CMAKE_CURRENT_LIST_DIR}/mlx - DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} - COMPONENT headers - FILES_MATCHING - PATTERN "*.h" - PATTERN "backend/metal/kernels.h" EXCLUDE) - -# Install metal dependencies -if(MLX_BUILD_METAL) - - # Install metal cpp - install( - DIRECTORY ${metal_cpp_SOURCE_DIR}/ - DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/metal_cpp - COMPONENT metal_cpp_source) - -endif() - -# Install cmake config -set(MLX_CMAKE_BUILD_CONFIG ${CMAKE_BINARY_DIR}/MLXConfig.cmake) -set(MLX_CMAKE_BUILD_VERSION_CONFIG ${CMAKE_BINARY_DIR}/MLXConfigVersion.cmake) -set(MLX_CMAKE_INSTALL_MODULE_DIR share/cmake/MLX) - -install( - EXPORT MLXTargets - FILE MLXTargets.cmake - DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}) - -include(CMakePackageConfigHelpers) - -write_basic_package_version_file( - ${MLX_CMAKE_BUILD_VERSION_CONFIG} - COMPATIBILITY SameMajorVersion - VERSION ${MLX_VERSION}) - -configure_package_config_file( - ${CMAKE_CURRENT_LIST_DIR}/mlx.pc.in ${MLX_CMAKE_BUILD_CONFIG} - INSTALL_DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR} - NO_CHECK_REQUIRED_COMPONENTS_MACRO - PATH_VARS CMAKE_INSTALL_LIBDIR CMAKE_INSTALL_INCLUDEDIR - MLX_CMAKE_INSTALL_MODULE_DIR) - -install(FILES ${MLX_CMAKE_BUILD_CONFIG} ${MLX_CMAKE_BUILD_VERSION_CONFIG} - DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}) - -install(DIRECTORY ${CMAKE_MODULE_PATH}/ - DESTINATION ${MLX_CMAKE_INSTALL_MODULE_DIR}) diff --git a/Source/CODE_OF_CONDUCT.md b/Source/CODE_OF_CONDUCT.md deleted file mode 100644 index e8d213c7b..000000000 --- a/Source/CODE_OF_CONDUCT.md +++ /dev/null @@ -1,132 +0,0 @@ -# Contributor Covenant Code of Conduct - -## Our Pledge - -We as members, contributors, and leaders pledge to make participation in our -community a harassment-free experience for everyone, regardless of age, body -size, visible or invisible disability, ethnicity, sex characteristics, gender -identity and expression, level of experience, education, socio-economic status, -nationality, personal appearance, race, caste, color, religion, or sexual -identity and orientation. - -We pledge to act and interact in ways that contribute to an open, welcoming, -diverse, inclusive, and healthy community. - -## Our Standards - -Examples of behavior that contributes to a positive environment for our -community include: - -* Demonstrating empathy and kindness toward other people -* Being respectful of differing opinions, viewpoints, and experiences -* Giving and gracefully accepting constructive feedback -* Accepting responsibility and apologizing to those affected by our mistakes, - and learning from the experience -* Focusing on what is best not just for us as individuals, but for the overall - community - -Examples of unacceptable behavior include: - -* The use of sexualized language or imagery, and sexual attention or advances of - any kind -* Trolling, insulting or derogatory comments, and personal or political attacks -* Public or private harassment -* Publishing others' private information, such as a physical or email address, - without their explicit permission -* Other conduct which could reasonably be considered inappropriate in a - professional setting - -## Enforcement Responsibilities - -Community leaders are responsible for clarifying and enforcing our standards of -acceptable behavior and will take appropriate and fair corrective action in -response to any behavior that they deem inappropriate, threatening, offensive, -or harmful. - -Community leaders have the right and responsibility to remove, edit, or reject -comments, commits, code, wiki edits, issues, and other contributions that are -not aligned to this Code of Conduct, and will communicate reasons for moderation -decisions when appropriate. - -## Scope - -This Code of Conduct applies within all community spaces, and also applies when -an individual is officially representing the community in public spaces. -Examples of representing our community include using an official e-mail address, -posting via an official social media account, or acting as an appointed -representative at an online or offline event. - -## Enforcement - -Instances of abusive, harassing, or otherwise unacceptable behavior may be -reported to the community leaders responsible for enforcement at -[opensource-conduct@group.apple.com](mailto:opensource-conduct@group.apple.com). -All complaints will be reviewed and investigated promptly and fairly. - -All community leaders are obligated to respect the privacy and security of the -reporter of any incident. - -## Enforcement Guidelines - -Community leaders will follow these Community Impact Guidelines in determining -the consequences for any action they deem in violation of this Code of Conduct: - -### 1. Correction - -**Community Impact**: Use of inappropriate language or other behavior deemed -unprofessional or unwelcome in the community. - -**Consequence**: A private, written warning from community leaders, providing -clarity around the nature of the violation and an explanation of why the -behavior was inappropriate. A public apology may be requested. - -### 2. Warning - -**Community Impact**: A violation through a single incident or series of -actions. - -**Consequence**: A warning with consequences for continued behavior. No -interaction with the people involved, including unsolicited interaction with -those enforcing the Code of Conduct, for a specified period of time. This -includes avoiding interactions in community spaces as well as external channels -like social media. Violating these terms may lead to a temporary or permanent -ban. - -### 3. Temporary Ban - -**Community Impact**: A serious violation of community standards, including -sustained inappropriate behavior. - -**Consequence**: A temporary ban from any sort of interaction or public -communication with the community for a specified period of time. No public or -private interaction with the people involved, including unsolicited interaction -with those enforcing the Code of Conduct, is allowed during this period. -Violating these terms may lead to a permanent ban. - -### 4. Permanent Ban - -**Community Impact**: Demonstrating a pattern of violation of community -standards, including sustained inappropriate behavior, harassment of an -individual, or aggression toward or disparagement of classes of individuals. - -**Consequence**: A permanent ban from any sort of public interaction within the -community. - -## Attribution - -This Code of Conduct is adapted from the [Contributor Covenant][homepage], -version 2.1, available at -[https://www.contributor-covenant.org/version/2/1/code_of_conduct.html][v2.1]. - -Community Impact Guidelines were inspired by -[Mozilla's code of conduct enforcement ladder][Mozilla CoC]. - -For answers to common questions about this code of conduct, see the FAQ at -[https://www.contributor-covenant.org/faq][FAQ]. Translations are available at -[https://www.contributor-covenant.org/translations][translations]. - -[homepage]: https://www.contributor-covenant.org -[v2.1]: https://www.contributor-covenant.org/version/2/1/code_of_conduct.html -[Mozilla CoC]: https://github.com/mozilla/diversity -[FAQ]: https://www.contributor-covenant.org/faq -[translations]: https://www.contributor-covenant.org/translations diff --git a/Source/CONTRIBUTING.md b/Source/CONTRIBUTING.md deleted file mode 100644 index fddb2a974..000000000 --- a/Source/CONTRIBUTING.md +++ /dev/null @@ -1,38 +0,0 @@ -# Contributing to MLX - -We want to make contributing to this project as easy and transparent as -possible. - -## Pull Requests - -1. Fork and submit pull requests to the repo. -2. If you've added code that should be tested, add tests. -3. If a change is likely to impact efficiency, run some of the benchmarks before - and after the change. Examples of benchmarks can be found in `benchmarks/python/`. -4. If you've changed APIs, update the documentation. -5. Every PR should have passing tests and at least one review. -6. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`. - This should install hooks for running `black` and `clang-format` to ensure - consistent style for C++ and python code. - - You can also run the formatters manually as follows: - - ```shell - clang-format -i file.cpp - ``` - - ```shell - black file.py - ``` - - or run `pre-commit run --all-files` to check all files in the repo. - -## Issues - -We use GitHub issues to track public bugs. Please ensure your description is -clear and has sufficient instructions to be able to reproduce the issue. - -## License - -By contributing to MLX, you agree that your contributions will be licensed -under the LICENSE file in the root directory of this source tree. 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/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..a44ec5666 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 8f784840de21cacd0c5c237ea860179e47dcec91. // // 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-device.h b/Source/Cmlx/include-framework/mlx-c-device.h index 0a41ed3ab..09e799610 100644 --- a/Source/Cmlx/include-framework/mlx-c-device.h +++ b/Source/Cmlx/include-framework/mlx-c-device.h @@ -7,7 +7,6 @@ #include #include -#include #ifdef __cplusplus extern "C" { @@ -26,6 +25,12 @@ typedef struct mlx_device_ { void* ctx; } mlx_device; +/** + * A MLX vector of strings. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_string_ mlx_vector_string; + /** * Device type. */ 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..a69b5b2ea 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,10 +201,95 @@ 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); + +// ── SharpAI custom ops (not part of upstream mlx-c) ────────────────────────── + +int mlx_fast_streamed_gather_mm( + mlx_array* res, + const mlx_array x, + const mlx_array w_shape, + uint32_t active_expert, + const char* safetensors_path, + const char* tensor_name, const mlx_stream s); +int mlx_fast_turbo_encode( + mlx_array* res_polar_k, + mlx_array* res_polar_v, + mlx_array* res_residual_k, + mlx_array* res_residual_v, + const mlx_array keys, + const mlx_array values, + int k_bits, + const mlx_stream s); + +int mlx_fast_turbo_decode_k( + mlx_array* res, + const mlx_array packed, + const mlx_stream s); + +int mlx_fast_turbo_decode_v( + mlx_array* res, + const mlx_array packed, + const mlx_stream s); + +int mlx_fast_prefault(mlx_array x); + +// pread() directly into the already-evaluated MLX array's unified memory buffer. +// This gives full NVMe sequential throughput without OS page-fault overhead. +// The array MUST already be evaluated (concrete pointer exists). +// safetensors_path: full path to .safetensors file +// tensor_name: e.g. "model.layers.0.mlp.experts.gate_proj.weight" +// expert_index: 0-based index of the expert to read +int mlx_fast_pread_into( + mlx_array dst, + const char* safetensors_path, + const char* tensor_name, + uint32_t expert_index); + +// Like mlx_fast_pread_into, but writes the expert's bytes into the dst buffer +// starting at byte offset `dst_offset`. Reads exactly `bytes_per_expert` bytes +// (NOT the whole dst array). Use this to populate one slot of a stacked +// `[N_slots, ..., ...]` buffer, where `dst_offset = slot * bytes_per_expert`. +// Bounds check: dst_offset + bytes_per_expert <= dst.nbytes. +int mlx_fast_pread_into_offset( + mlx_array dst, + const char* safetensors_path, + const char* tensor_name, + uint32_t expert_index, + size_t dst_offset); + +// mlx_fast_submit_prefetch (PAPPS Background Worker) +int mlx_fast_submit_prefetch( + const char* safetensors_path, + const char* tensor_name, + uint32_t expert_index); + +void mlx_fast_set_prefetch_enabled(bool enabled); + /**@}*/ +// ── SSD Flash-Stream metrics snapshot ──────────────────────────────────────── +// Cumulative NVMe throughput stats since process start. +// Call mlx_ssd_metrics_snapshot() from any thread; never resets any counter. + +typedef struct MlxSSDMetricsSnapshot { + double throughput_mb_per_s; /* 10-s rolling window average (0 before first window) */ + uint64_t total_bytes_read; /* Lifetime bytes read from SSD */ + uint64_t total_chunks; /* Lifetime expert chunks loaded */ + double avg_chunk_latency_ms; /* Lifetime average per-chunk latency (ms) */ +} MlxSSDMetricsSnapshot; + +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). +void mlx_ssd_metrics_snapshot(MlxSSDMetricsSnapshot* out); + +// TurboKV telemetry — call from Swift on each compression event to accumulate +// stats that appear in the 10-second SSD stream log. +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). +void mlx_turbo_kv_record(uint64_t tokens, uint64_t orig_bytes, uint64_t packed_bytes); + #ifdef __cplusplus } #endif 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..7b6ae483a 100644 --- a/Source/Cmlx/include-framework/mlx-c-io_types.h +++ b/Source/Cmlx/include-framework/mlx-c-io_types.h @@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ { bool (*is_open)(void*); bool (*good)(void*); size_t (*tell)(void*); - void (*seek)(void*, int64_t off, int whence); - void (*read)(void*, char* data, size_t n); - void (*read_at_offset)(void*, char* data, size_t n, size_t off); - void (*write)(void*, const char* data, size_t n); + int (*seek)(void*, int64_t off, int whence); + size_t (*read)(void*, char* data, size_t n); + size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); + size_t (*write)(void*, const char* data, size_t n); const char* (*label)(void*); void (*free)(void*); } mlx_io_vtable; @@ -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..4115e7e91 100644 --- a/Source/Cmlx/include-framework/mlx-c-stream.h +++ b/Source/Cmlx/include-framework/mlx-c-stream.h @@ -4,6 +4,7 @@ #define MLX_STREAM_H #include +#include #include @@ -24,6 +25,22 @@ typedef struct mlx_stream_ { void* ctx; } mlx_stream; +/** + * A per-thread stream object. + * + * The same thread-local stream can be used from any number of threads; the + * stream's GPU command encoder is registered once per thread. + */ +typedef struct mlx_stream_thread_local_ { + void* ctx; +} mlx_stream_thread_local; + +/** + * A vector of mlx_stream. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_stream_ mlx_vector_stream; + /** * Returns a new empty stream. */ @@ -33,6 +50,48 @@ 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); + +/** + * Returns a new per-thread stream on a device. + * + * Streams created this way are unique per thread: evaluating the same + * `mlx_stream_thread_local` from different threads yields distinct streams. + */ +mlx_stream_thread_local mlx_stream_thread_local_new(mlx_device dev); +/** + * Set a per-thread stream to the provided src stream. + */ +int mlx_stream_thread_local_set( + mlx_stream_thread_local* tls, + const mlx_stream_thread_local src); +/** + * Free a per-thread stream. + */ +int mlx_stream_thread_local_free(mlx_stream_thread_local tls); +/** + * Returns the concrete stream backing a per-thread stream on the current + * thread. + */ +int mlx_stream_from_thread_local( + mlx_stream* res, + const mlx_stream_thread_local tls); + +/** + * Get the vector of available streams. + */ +int mlx_get_streams(mlx_vector_stream* res); /** * Set stream to provided src stream. */ @@ -61,6 +120,18 @@ int mlx_stream_get_index(int* index, mlx_stream stream); * Synchronize with the provided stream. */ int mlx_synchronize(mlx_stream stream); +/** + * Synchronize with the default stream. + */ +int mlx_synchronize_default(void); +/** + * Synchronize with the stream corresponding to the current thread. + */ +int mlx_synchronize_thread_local(mlx_stream_thread_local tls); +/** + * Destroy all streams created in the current thread. + */ +int mlx_clear_streams(void); /** * Returns the default stream on the given device. */ diff --git a/Source/Cmlx/include-framework/mlx-c-vector.h b/Source/Cmlx/include-framework/mlx-c-vector.h index 1954c6d0b..b290edc92 100644 --- a/Source/Cmlx/include-framework/mlx-c-vector.h +++ b/Source/Cmlx/include-framework/mlx-c-vector.h @@ -7,6 +7,7 @@ #define MLX_VECTOR_H #include +#include #include #ifdef __cplusplus @@ -124,6 +125,35 @@ int mlx_vector_string_append_value(mlx_vector_string vec, const char* val); size_t mlx_vector_string_size(mlx_vector_string vec); int mlx_vector_string_get(char** res, const mlx_vector_string vec, size_t idx); +/** + * A vector of stream. + */ +typedef struct mlx_vector_stream_ { + void* ctx; +} mlx_vector_stream; +mlx_vector_stream mlx_vector_stream_new(void); +int mlx_vector_stream_set(mlx_vector_stream* vec, const mlx_vector_stream src); +int mlx_vector_stream_free(mlx_vector_stream vec); +mlx_vector_stream mlx_vector_stream_new_data( + const mlx_stream* data, + size_t size); +mlx_vector_stream mlx_vector_stream_new_value(const mlx_stream val); +int mlx_vector_stream_set_data( + mlx_vector_stream* vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_set_value(mlx_vector_stream* vec, const mlx_stream val); +int mlx_vector_stream_append_data( + mlx_vector_stream vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_append_value(mlx_vector_stream vec, const mlx_stream val); +size_t mlx_vector_stream_size(mlx_vector_stream vec); +int mlx_vector_stream_get( + mlx_stream* res, + const mlx_vector_stream vec, + size_t idx); + /**@}*/ #ifdef __cplusplus 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..9e8786d53 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, @@ -101,5 +108,39 @@ MLX_API std::vector precompiled_cuda_kernel( bool ensure_row_contiguous = false, StreamOrDevice s = {}); +/** + * Compress a K-cache tensor to TurboQuant format (3-bit PolarQuant + 1-bit QJL). + * + * keys: [batch, heads, seq, 128] — fp16 / bf16 / fp32 + * returns: uint8 array with the same leading dims and last dim = 68 + * Layout per token: indices[48] | qjl_signs[16] | norm_fp16[2] | rnorm_fp16[2] + */ +MLX_API array turbo_encode_k(const array& keys, StreamOrDevice s = {}); + +/** + * Compress a V-cache tensor to TurboQuant format (3-bit PolarQuant only). + * + * values: [batch, heads, seq, 128] — fp16 / bf16 / fp32 + * returns: uint8 array with the same leading dims and last dim = 50 + * Layout per token: indices[48] | norm_fp16[2] + */ +MLX_API array turbo_encode_v(const array& values, StreamOrDevice s = {}); + +/** + * Decode TurboKV compressed K-cache back to float32. + * + * packed: uint8 with last dim 68 (D=128) or 136 (D=256) + * returns: float32 array with last dim = head_dim (128 or 256) + */ +MLX_API array turbo_decode_k(const array& packed, StreamOrDevice s = {}); + +/** + * Decode TurboKV compressed V-cache back to float32. + * + * packed: uint8 with last dim 50 (D=128) or 100 (D=256) + * returns: float32 array with last dim = head_dim (128 or 256) + */ +MLX_API array turbo_decode_v(const array& packed, StreamOrDevice s = {}); + } // namespace mlx::core::fast #endif 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 1970181be..c5e4fc7b4 100644 --- a/Source/Cmlx/include-framework/mlx-io-load.h +++ b/Source/Cmlx/include-framework/mlx-io-load.h @@ -15,7 +15,9 @@ #include #endif -// Detect iOS to force sequential I/O (prevents EXC_RESOURCE WAKEUPS). +// Detect iOS/tvOS/watchOS to select safe sequential I/O. +// On Apple mobile OSes, EXC_RESOURCE (RESOURCE_TYPE_WAKEUPS) is triggered when +// std::async thread spawning causes too many wakeups per second. #if defined(TARGET_OS_IPHONE) && TARGET_OS_IPHONE #define MLX_IOS_SEQUENTIAL_IO 1 #elif defined(__APPLE__) @@ -95,7 +97,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); } @@ -112,10 +116,14 @@ class ParallelFileReader : public Reader { } private: + // On iOS, always use sequential pread() — no std::async thread spawning. + // std::async wakeup bursts during model loading exceed iOS's per-process + // wakeup rate limit → EXC_RESOURCE (RESOURCE_TYPE_WAKEUPS) crash. + // On macOS, 32MB parallel batches maximize NVMe throughput. #ifdef MLX_IOS_SEQUENTIAL_IO - static constexpr size_t batch_size_ = SIZE_MAX; // iOS: sequential only, no wakeups + static constexpr size_t batch_size_ = SIZE_MAX; #else - static constexpr size_t batch_size_ = 1 << 25; // macOS: 32MB parallel batches + static constexpr size_t batch_size_ = 1 << 25; // 32MB #endif static ThreadPool& thread_pool(); int fd_; @@ -159,7 +167,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/device.h b/Source/Cmlx/include/mlx/c/device.h index 4b74e39d3..705aea9db 100644 --- a/Source/Cmlx/include/mlx/c/device.h +++ b/Source/Cmlx/include/mlx/c/device.h @@ -7,7 +7,6 @@ #include #include "mlx/c/string.h" -#include "mlx/c/vector.h" #ifdef __cplusplus extern "C" { @@ -26,6 +25,12 @@ typedef struct mlx_device_ { void* ctx; } mlx_device; +/** + * A MLX vector of strings. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_string_ mlx_vector_string; + /** * Device type. */ 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 1b2f2f67c..2786c1ce8 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,8 +201,11 @@ 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); +// ── SharpAI custom ops (not part of upstream mlx-c) ────────────────────────── + int mlx_fast_streamed_gather_mm( mlx_array* res, const mlx_array x, @@ -228,6 +237,12 @@ int mlx_fast_turbo_decode_v( int mlx_fast_prefault(mlx_array x); +// pread() directly into the already-evaluated MLX array's unified memory buffer. +// This gives full NVMe sequential throughput without OS page-fault overhead. +// The array MUST already be evaluated (concrete pointer exists). +// safetensors_path: full path to .safetensors file +// tensor_name: e.g. "model.layers.0.mlp.experts.gate_proj.weight" +// expert_index: 0-based index of the expert to read int mlx_fast_pread_into( mlx_array dst, const char* safetensors_path, @@ -267,10 +282,12 @@ typedef struct MlxSSDMetricsSnapshot { double avg_chunk_latency_ms; /* Lifetime average per-chunk latency (ms) */ } MlxSSDMetricsSnapshot; +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). void mlx_ssd_metrics_snapshot(MlxSSDMetricsSnapshot* out); // TurboKV telemetry — call from Swift on each compression event to accumulate // stats that appear in the 10-second SSD stream log. +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). void mlx_turbo_kv_record(uint64_t tokens, uint64_t orig_bytes, uint64_t packed_bytes); #ifdef __cplusplus 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..3f15873e3 100644 --- a/Source/Cmlx/include/mlx/c/io_types.h +++ b/Source/Cmlx/include/mlx/c/io_types.h @@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ { bool (*is_open)(void*); bool (*good)(void*); size_t (*tell)(void*); - void (*seek)(void*, int64_t off, int whence); - void (*read)(void*, char* data, size_t n); - void (*read_at_offset)(void*, char* data, size_t n, size_t off); - void (*write)(void*, const char* data, size_t n); + int (*seek)(void*, int64_t off, int whence); + size_t (*read)(void*, char* data, size_t n); + size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); + size_t (*write)(void*, const char* data, size_t n); const char* (*label)(void*); void (*free)(void*); } mlx_io_vtable; @@ -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..6df1a63c9 100644 --- a/Source/Cmlx/include/mlx/c/stream.h +++ b/Source/Cmlx/include/mlx/c/stream.h @@ -4,6 +4,7 @@ #define MLX_STREAM_H #include +#include #include "mlx/c/device.h" @@ -24,6 +25,22 @@ typedef struct mlx_stream_ { void* ctx; } mlx_stream; +/** + * A per-thread stream object. + * + * The same thread-local stream can be used from any number of threads; the + * stream's GPU command encoder is registered once per thread. + */ +typedef struct mlx_stream_thread_local_ { + void* ctx; +} mlx_stream_thread_local; + +/** + * A vector of mlx_stream. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_stream_ mlx_vector_stream; + /** * Returns a new empty stream. */ @@ -33,6 +50,48 @@ 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); + +/** + * Returns a new per-thread stream on a device. + * + * Streams created this way are unique per thread: evaluating the same + * `mlx_stream_thread_local` from different threads yields distinct streams. + */ +mlx_stream_thread_local mlx_stream_thread_local_new(mlx_device dev); +/** + * Set a per-thread stream to the provided src stream. + */ +int mlx_stream_thread_local_set( + mlx_stream_thread_local* tls, + const mlx_stream_thread_local src); +/** + * Free a per-thread stream. + */ +int mlx_stream_thread_local_free(mlx_stream_thread_local tls); +/** + * Returns the concrete stream backing a per-thread stream on the current + * thread. + */ +int mlx_stream_from_thread_local( + mlx_stream* res, + const mlx_stream_thread_local tls); + +/** + * Get the vector of available streams. + */ +int mlx_get_streams(mlx_vector_stream* res); /** * Set stream to provided src stream. */ @@ -61,6 +120,18 @@ int mlx_stream_get_index(int* index, mlx_stream stream); * Synchronize with the provided stream. */ int mlx_synchronize(mlx_stream stream); +/** + * Synchronize with the default stream. + */ +int mlx_synchronize_default(void); +/** + * Synchronize with the stream corresponding to the current thread. + */ +int mlx_synchronize_thread_local(mlx_stream_thread_local tls); +/** + * Destroy all streams created in the current thread. + */ +int mlx_clear_streams(void); /** * Returns the default stream on the given device. */ diff --git a/Source/Cmlx/include/mlx/c/vector.h b/Source/Cmlx/include/mlx/c/vector.h index 81bcf7495..30d1b1838 100644 --- a/Source/Cmlx/include/mlx/c/vector.h +++ b/Source/Cmlx/include/mlx/c/vector.h @@ -7,6 +7,7 @@ #define MLX_VECTOR_H #include "mlx/c/array.h" +#include "mlx/c/stream.h" #include "mlx/c/string.h" #ifdef __cplusplus @@ -124,6 +125,35 @@ int mlx_vector_string_append_value(mlx_vector_string vec, const char* val); size_t mlx_vector_string_size(mlx_vector_string vec); int mlx_vector_string_get(char** res, const mlx_vector_string vec, size_t idx); +/** + * A vector of stream. + */ +typedef struct mlx_vector_stream_ { + void* ctx; +} mlx_vector_stream; +mlx_vector_stream mlx_vector_stream_new(void); +int mlx_vector_stream_set(mlx_vector_stream* vec, const mlx_vector_stream src); +int mlx_vector_stream_free(mlx_vector_stream vec); +mlx_vector_stream mlx_vector_stream_new_data( + const mlx_stream* data, + size_t size); +mlx_vector_stream mlx_vector_stream_new_value(const mlx_stream val); +int mlx_vector_stream_set_data( + mlx_vector_stream* vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_set_value(mlx_vector_stream* vec, const mlx_stream val); +int mlx_vector_stream_append_data( + mlx_vector_stream vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_append_value(mlx_vector_stream vec, const mlx_stream val); +size_t mlx_vector_stream_size(mlx_vector_stream vec); +int mlx_vector_stream_get( + mlx_stream* res, + const mlx_vector_stream vec, + size_t idx); + /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/mlx-c/CMakeLists.txt b/Source/Cmlx/mlx-c/CMakeLists.txt index 9e5254a22..f96c72a79 100644 --- a/Source/Cmlx/mlx-c/CMakeLists.txt +++ b/Source/Cmlx/mlx-c/CMakeLists.txt @@ -4,11 +4,11 @@ cmake_minimum_required(VERSION 3.16) # ----------------------------- Setup ----------------------------- project(mlx.data LANGUAGES CXX C) -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) # ----------------------------- Configuration ----------------------------- if(NOT MLX_C_VERSION) - set(MLX_C_VERSION 0.5.0) + set(MLX_C_VERSION 0.6.0) endif() option(BUILD_SHARED_LIBS "Build mlx C as a shared library" OFF) @@ -35,7 +35,7 @@ else() FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG v0.30.6) + GIT_TAG v0.32.2) FetchContent_MakeAvailable(mlx) endif() @@ -53,6 +53,7 @@ set(mlxc-src ${CMAKE_CURRENT_LIST_DIR}/mlx/c/export.cpp ${CMAKE_CURRENT_LIST_DIR}/mlx/c/fast.cpp ${CMAKE_CURRENT_LIST_DIR}/mlx/c/fft.cpp + ${CMAKE_CURRENT_LIST_DIR}/mlx/c/graph_utils.cpp ${CMAKE_CURRENT_LIST_DIR}/mlx/c/io.cpp ${CMAKE_CURRENT_LIST_DIR}/mlx/c/io_types.cpp ${CMAKE_CURRENT_LIST_DIR}/mlx/c/linalg.cpp diff --git a/Source/Cmlx/mlx-c/README.md b/Source/Cmlx/mlx-c/README.md index 4c42e9c4b..d57abc96e 100644 --- a/Source/Cmlx/mlx-c/README.md +++ b/Source/Cmlx/mlx-c/README.md @@ -23,15 +23,14 @@ brew install cmake To build, run the following commands: ```shell -mkdir build && cd build -cmake .. -DCMAKE_BUILD_TYPE=Release -make -j +cmake -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build -j ``` From the `build/` directory, you can run an [example](examples/example.c) that uses MLX C with `./example`. -## Contributing +## Contributing Check out the [contribution guidelines](CONTRIBUTING.md) for more information on contributing to MLX C. See the diff --git a/Source/Cmlx/mlx-c/docs/src/conf.py b/Source/Cmlx/mlx-c/docs/src/conf.py index b34aa19c2..660fbb7df 100644 --- a/Source/Cmlx/mlx-c/docs/src/conf.py +++ b/Source/Cmlx/mlx-c/docs/src/conf.py @@ -10,7 +10,7 @@ project = "MLX C" copyright = "2023-2026, MLX Contributors" author = "MLX Contributors" -version = "0.5.0" +version = "0.6.0" release = version # -- General configuration --------------------------------------------------- diff --git a/Source/Cmlx/mlx-c/docs/src/install.rst b/Source/Cmlx/mlx-c/docs/src/install.rst index 799496588..25842ed7a 100644 --- a/Source/Cmlx/mlx-c/docs/src/install.rst +++ b/Source/Cmlx/mlx-c/docs/src/install.rst @@ -11,9 +11,8 @@ To build MLX C, run the following commands: .. code-block:: shell - mkdir build && cd build/ - cmake .. -DCMAKE_BUILD_TYPE=Release - make -j + cmake -B build -DCMAKE_BUILD_TYPE=Release + cmake --build build -j MLX C will fetch `MLX `_ under the hood, compile it, and then compile the C API. diff --git a/Source/Cmlx/mlx-c/examples/CMakeLists.txt b/Source/Cmlx/mlx-c/examples/CMakeLists.txt index ba83e5bc3..61b889bcd 100644 --- a/Source/Cmlx/mlx-c/examples/CMakeLists.txt +++ b/Source/Cmlx/mlx-c/examples/CMakeLists.txt @@ -20,3 +20,22 @@ target_link_libraries(example-closure PUBLIC mlxc) add_executable(example-export ${CMAKE_CURRENT_LIST_DIR}/example-export.c) target_link_libraries(example-export PUBLIC mlxc) + +add_executable(example-gguf ${CMAKE_CURRENT_LIST_DIR}/example-gguf.c) +target_link_libraries(example-gguf PUBLIC mlxc) + +add_executable(example-graph ${CMAKE_CURRENT_LIST_DIR}/example-graph.c) +target_link_libraries(example-graph PUBLIC mlxc) + +add_executable(example-thread-stream + ${CMAKE_CURRENT_LIST_DIR}/example-thread-stream.c) +target_link_libraries(example-thread-stream PUBLIC mlxc) +find_package(Threads REQUIRED) +target_link_libraries(example-thread-stream PUBLIC Threads::Threads) + +add_executable(example-ops ${CMAKE_CURRENT_LIST_DIR}/example-ops.c) +target_link_libraries(example-ops PUBLIC mlxc) + +add_executable(example-stream-tls + ${CMAKE_CURRENT_LIST_DIR}/example-stream-tls.c) +target_link_libraries(example-stream-tls PUBLIC mlxc) diff --git a/Source/Cmlx/mlx-c/examples/arrays.safetensors b/Source/Cmlx/mlx-c/examples/arrays.safetensors new file mode 100644 index 000000000..408049e4a Binary files /dev/null and b/Source/Cmlx/mlx-c/examples/arrays.safetensors differ diff --git a/Source/Cmlx/mlx-c/examples/example-gguf.c b/Source/Cmlx/mlx-c/examples/example-gguf.c new file mode 100644 index 000000000..5193f016a --- /dev/null +++ b/Source/Cmlx/mlx-c/examples/example-gguf.c @@ -0,0 +1,106 @@ +/* Copyright © 2023-2024 Apple Inc. */ + +#include +#include "mlx/c/mlx.h" + +char default_filename[] = "test.gguf"; +const char* dtype_strs[] = { + "MLX_BOOL", + "MLX_UINT8", + "MLX_UINT16", + "MLX_UINT32", + "MLX_UINT64", + "MLX_INT8", + "MLX_INT16", + "MLX_INT32", + "MLX_INT64", + "MLX_FLOAT16", + "MLX_FLOAT32", + "MLX_FLOAT64", + "MLX_BFLOAT16", + "MLX_COMPLEX64", +}; + +void add_zero_array_to_gguf( + mlx_io_gguf gguf, + const char* key, + const int* shape, + int ndim, + mlx_dtype dtype, + const char* metadata, + mlx_stream s) { + mlx_array arr = mlx_array_new(); + mlx_zeros(&arr, shape, ndim, dtype, s); + mlx_io_gguf_set_array(gguf, key, arr); + if (metadata) { + mlx_io_gguf_set_metadata_string(gguf, key, metadata); + } + mlx_array_free(arr); +} + +int main(int argc, char* argv[]) { + mlx_stream stream = mlx_default_cpu_stream_new(); + + char* filename = default_filename; + if (argc > 1) { + filename = argv[1]; + } else { + mlx_io_gguf gguf = mlx_io_gguf_new(); + add_zero_array_to_gguf( + gguf, "array3D", (int[]){3, 4, 5}, 3, MLX_FLOAT16, NULL, stream); + add_zero_array_to_gguf( + gguf, + "array2D", + (int[]){6, 7}, + 2, + MLX_FLOAT32, + "a 6x7 zero array", + stream); + mlx_save_gguf(filename, gguf); + mlx_io_gguf_free(gguf); + } + + mlx_io_gguf gguf = mlx_io_gguf_new(); + mlx_load_gguf(&gguf, filename, stream); + mlx_vector_string keys = mlx_vector_string_new(); + mlx_io_gguf_get_keys(&keys, gguf); + + char* key; + bool flag; + mlx_array value = mlx_array_new(); + for (int i = 0; i < mlx_vector_string_size(keys); i++) { + mlx_vector_string_get(&key, keys, i); + printf("%s ", key); + mlx_io_gguf_get_array(&value, gguf, key); + const int* shape = mlx_array_shape(value); + for (int d = 0; d < mlx_array_ndim(value); d++) { + if (d != 0) { + printf("x"); + } + printf("%d", shape[d]); + } + printf(" %s", dtype_strs[mlx_array_dtype(value)]); + + mlx_io_gguf_has_metadata_array(&flag, gguf, key); + if (flag) { + printf(" [array]"); + } + mlx_io_gguf_has_metadata_string(&flag, gguf, key); + if (flag) { + printf(" [string]"); + } + mlx_io_gguf_has_metadata_vector_string(&flag, gguf, key); + if (flag) { + printf(" [vector string]"); + } + + printf("\n"); + } + + mlx_array_free(value); + mlx_vector_string_free(keys); + mlx_io_gguf_free(gguf); + mlx_stream_free(stream); + + return 0; +} diff --git a/Source/Cmlx/mlx-c/examples/example-graph.c b/Source/Cmlx/mlx-c/examples/example-graph.c new file mode 100644 index 000000000..bda0e8592 --- /dev/null +++ b/Source/Cmlx/mlx-c/examples/example-graph.c @@ -0,0 +1,34 @@ +/* Copyright © 2023-2024 Apple Inc. */ + +#include +#include "mlx/c/mlx.h" + +int main(void) { + mlx_stream stream = mlx_default_cpu_stream_new(); + float data[] = {1, 2, 3, 4, 5, 6}; + int shape[] = {2, 3}; + mlx_array res = mlx_array_new(); + mlx_array val = mlx_array_new_data(data, shape, 2, MLX_FLOAT32); + mlx_array two = mlx_array_new_int(2); + mlx_divide(&res, val, two, stream); + mlx_log(&res, res, stream); + + mlx_node_namer namer = mlx_node_namer_new(); + mlx_node_namer_set_name(namer, val, "inputs"); + mlx_node_namer_set_name(namer, res, "result"); + mlx_vector_array vec = mlx_vector_array_new(); + mlx_vector_array_append_value(vec, val); + mlx_vector_array_append_value(vec, two); + mlx_vector_array_append_value(vec, res); + + mlx_export_to_dot(stdout, namer, vec); + + mlx_array_free(val); + mlx_array_free(two); + mlx_array_free(res); + mlx_vector_array_free(vec); + mlx_node_namer_free(namer); + mlx_stream_free(stream); + + return 0; +} diff --git a/Source/Cmlx/mlx-c/examples/example-ops.c b/Source/Cmlx/mlx-c/examples/example-ops.c new file mode 100644 index 000000000..8a408e855 --- /dev/null +++ b/Source/Cmlx/mlx-c/examples/example-ops.c @@ -0,0 +1,76 @@ +/* Copyright © 2023-2024 Apple Inc. */ + +#include +#include "mlx/c/mlx.h" + +void print_array(const char* msg, mlx_array arr) { + mlx_string str = mlx_string_new(); + mlx_array_tostring(&str, arr); + printf("%s\n%s\n", msg, mlx_string_data(str)); + mlx_string_free(str); +} + +/* Exercises a range of array, elementwise, and linalg ops. */ +int main(void) { + mlx_stream stream = mlx_default_cpu_stream_new(); + + float data[] = {1, -2, 0, 4, -5, 6}; + int shape[] = {2, 3}; + mlx_array arr = mlx_array_new_data(data, shape, 2, MLX_FLOAT32); + mlx_array res = mlx_array_new(); + + mlx_positive(&res, arr, stream); + print_array("positive", res); + + mlx_trunc(&res, arr, stream); + print_array("trunc", res); + + mlx_count_nonzero_axis(&res, arr, 1, false, stream); + print_array("count_nonzero (axis=1)", res); + + mlx_diff(&res, arr, 1, 1, stream); + print_array("diff (n=1, axis=1)", res); + + mlx_flip_axis(&res, arr, 1, stream); + print_array("flip (axis=1)", res); + + float rhs_data[] = {6, 5, 4, 3, 2, 1}; + mlx_array rhs = mlx_array_new_data(rhs_data, shape, 2, MLX_FLOAT32); + mlx_vecdot(&res, arr, rhs, 1, stream); + print_array("vecdot (axis=1)", res); + + mlx_array lbool = mlx_array_new_bool(true); + mlx_array rbool = mlx_array_new_bool(false); + mlx_logical_xor(&res, lbool, rbool, stream); + print_array("logical_xor", res); + + mlx_vector_array parts = mlx_vector_array_new(); + mlx_unstack_axis(&parts, arr, 0, stream); + printf("unstack (axis=0): %zu arrays\n", mlx_vector_array_size(parts)); + + float square_data[] = {4, 3, 6, 3}; + int square_shape[] = {2, 2}; + mlx_array square = + mlx_array_new_data(square_data, square_shape, 2, MLX_FLOAT32); + mlx_linalg_det(&res, square, stream); + print_array("linalg_det", res); + + mlx_array sign = mlx_array_new(); + mlx_array logabsdet = mlx_array_new(); + mlx_linalg_slogdet(&sign, &logabsdet, square, stream); + print_array("slogdet sign", sign); + print_array("slogdet logabsdet", logabsdet); + + mlx_array_free(arr); + mlx_array_free(rhs); + mlx_array_free(res); + mlx_array_free(lbool); + mlx_array_free(rbool); + mlx_array_free(square); + mlx_array_free(sign); + mlx_array_free(logabsdet); + mlx_vector_array_free(parts); + mlx_stream_free(stream); + + return 0; +} diff --git a/Source/Cmlx/mlx-c/examples/example-safe-tensors.c b/Source/Cmlx/mlx-c/examples/example-safe-tensors.c index 635e2ffc6..ab2d49a27 100644 --- a/Source/Cmlx/mlx-c/examples/example-safe-tensors.c +++ b/Source/Cmlx/mlx-c/examples/example-safe-tensors.c @@ -32,7 +32,7 @@ size_t mem_tell(void* desc) { mlx_mem_stream* m_desc = desc; return m_desc->pos; } -void mem_seek(void* desc, int64_t off, int whence) { +int mem_seek(void* desc, int64_t off, int whence) { printf("SEEK\n"); mlx_mem_stream* m_desc = desc; size_t new_pos; @@ -48,43 +48,48 @@ void mem_seek(void* desc, int64_t off, int whence) { break; default: m_desc->err = true; - return; + return -1; } if (new_pos > m_desc->size) { m_desc->err = true; + return -1; } else { m_desc->pos = new_pos; + return 0; } } -void mem_read(void* desc, char* data, size_t n) { +size_t mem_read(void* desc, char* data, size_t n) { printf("READ %ld\n", n); mlx_mem_stream* m_desc = desc; if (n + m_desc->pos > m_desc->size) { m_desc->err = true; - return; + return 0; } memcpy(data, m_desc->data + m_desc->pos, n); m_desc->pos += n; + return n; } -void mem_read_at_offset(void* desc, char* data, size_t n, size_t off) { +size_t mem_read_at_offset(void* desc, char* data, size_t n, size_t off) { printf("READ@OFFSET %ld @ %ld\n", n, off); mlx_mem_stream* m_desc = desc; if (off + n > m_desc->size) { m_desc->err = true; - return; + return 0; } memcpy(data, m_desc->data + off, n); m_desc->pos = off; + return n; } -void mem_write(void* desc, const char* data, size_t n) { +size_t mem_write(void* desc, const char* data, size_t n) { printf("WRITE %ld\n", n); mlx_mem_stream* m_desc = desc; if (n + m_desc->pos > m_desc->size) { m_desc->err = true; - return; + return 0; } memcpy(m_desc->data + m_desc->pos, data, n); m_desc->pos += n; + return n; } const char* mem_label(void* desc) { printf("LABEL\n"); diff --git a/Source/Cmlx/mlx-c/examples/example-stream-tls.c b/Source/Cmlx/mlx-c/examples/example-stream-tls.c new file mode 100644 index 000000000..966f7824a --- /dev/null +++ b/Source/Cmlx/mlx-c/examples/example-stream-tls.c @@ -0,0 +1,56 @@ +/* Copyright © 2026 Apple Inc. */ +/* Exercises the thread-local / vector-stream C bindings. */ + +#include + +#include "mlx/c/mlx.h" + +int main(void) { + mlx_device dev = mlx_device_new(); + mlx_get_default_device(&dev); + + mlx_stream_thread_local tls = mlx_stream_thread_local_new(dev); + if (!tls.ctx) + mlx_error("mlx_stream_thread_local_new returned empty"); + printf("thread-local stream created\n"); + + mlx_stream cur = mlx_stream_new(); + if (mlx_stream_from_thread_local(&cur, tls)) + mlx_error("mlx_stream_from_thread_local failed"); + printf("stream_from_thread_local ok\n"); + + mlx_vector_stream streams = mlx_vector_stream_new(); + if (mlx_get_streams(&streams)) + mlx_error("mlx_get_streams failed"); + + size_t n = mlx_vector_stream_size(streams); + printf("get_streams returned %zu streams\n", n); + if (n == 0) + mlx_error("expected at least one stream"); + + mlx_stream first = mlx_stream_new(); + if (mlx_vector_stream_get(&first, streams, 0)) + mlx_error("mlx_vector_stream_get failed"); + mlx_vector_stream_free(streams); + + if (mlx_synchronize_default()) + mlx_error("mlx_synchronize_default failed"); + if (mlx_synchronize_thread_local(tls)) + mlx_error("mlx_synchronize_thread_local failed"); + + mlx_stream_thread_local tls2 = mlx_stream_thread_local_new(dev); + if (mlx_stream_thread_local_set(&tls2, tls)) + mlx_error("mlx_stream_thread_local_set failed"); + if (mlx_stream_thread_local_free(tls2)) + mlx_error("mlx_stream_thread_local_free failed"); + mlx_stream_thread_local_free(tls); + + if (mlx_clear_streams()) + mlx_error("mlx_clear_streams failed"); + + mlx_stream_free(cur); + mlx_stream_free(first); + mlx_device_free(dev); + printf("ok\n"); + return 0; +} diff --git a/Source/Cmlx/mlx-c/examples/example-thread-stream.c b/Source/Cmlx/mlx-c/examples/example-thread-stream.c new file mode 100644 index 000000000..fe846c31c --- /dev/null +++ b/Source/Cmlx/mlx-c/examples/example-thread-stream.c @@ -0,0 +1,73 @@ +/* Copyright © 2025 Apple Inc. */ + +/* Streams are thread affine: a stream created on one thread cannot be used to + * evaluate on another. `mlx_stream_new_thread_unsafe` opts out of that, which + * is what language bindings that do their own locking need. + * + * This creates a stream on the main thread and evaluates on it from a worker + * thread. Note the mutex: MLX applies no synchronization to these streams, so + * serializing the work is the caller's job. + */ + +#include +#include + +#include "mlx/c/mlx.h" + +static pthread_mutex_t eval_mutex = PTHREAD_MUTEX_INITIALIZER; + +typedef struct { + mlx_stream stream; + float result; + int status; +} work; + +static void* run(void* arg) { + work* w = (work*)arg; + + pthread_mutex_lock(&eval_mutex); + + mlx_array a = mlx_array_new_float(2.0); + mlx_array b = mlx_array_new_float(3.0); + mlx_array sum = mlx_array_new(); + + w->status = mlx_add(&sum, a, b, w->stream); + if (w->status == 0) { + w->status = mlx_array_item_float32(&w->result, sum); + } + + mlx_array_free(a); + mlx_array_free(b); + mlx_array_free(sum); + + pthread_mutex_unlock(&eval_mutex); + return NULL; +} + +int main(void) { + mlx_device dev = mlx_device_new(); + mlx_get_default_device(&dev); + + /* created here, used on the worker thread below */ + work w = {mlx_stream_new_thread_unsafe(dev), 0.0f, 0}; + + pthread_t thread; + if (pthread_create(&thread, NULL, run, &w) != 0) { + printf("failed to spawn thread\n"); + return 1; + } + pthread_join(thread, NULL); + + if (w.status != 0) { + /* the default error handler has already reported the reason */ + printf("evaluation failed\n"); + return 1; + } + + printf("2 + 3 evaluated on another thread = %g\n", w.result); + + mlx_stream_free(w.stream); + mlx_device_free(dev); + + return 0; +} diff --git a/Source/Cmlx/mlx-c/mlx/c/compile.cpp b/Source/Cmlx/mlx-c/mlx/c/compile.cpp index bc7b86493..52ebc3519 100644 --- a/Source/Cmlx/mlx-c/mlx/c/compile.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/compile.cpp @@ -8,6 +8,24 @@ #include "mlx/c/private/mlx.h" #include "mlx/compile_impl.h" +extern "C" mlx_compile_cache mlx_compile_cache_new() { + try { + return mlx_compile_cache_new_(); + } catch (std::exception& e) { + mlx_error(e.what()); + } + return {nullptr}; +} +extern "C" int mlx_compile_cache_free(mlx_compile_cache cache) { + try { + mlx_compile_cache_free_(cache); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + extern "C" int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless) { try { @@ -40,18 +58,29 @@ extern "C" int mlx_detail_compile( } return 0; } -extern "C" int mlx_detail_compile_clear_cache(void) { +extern "C" int mlx_detail_compile_cache(mlx_compile_cache* res) { + try { + mlx_compile_cache_set_(*res, mlx::core::detail::compile_cache()); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_detail_compile_clear_cache(const mlx_compile_cache cache) { try { - mlx::core::detail::compile_clear_cache(); + mlx::core::detail::compile_clear_cache(mlx_compile_cache_get_(cache)); } catch (std::exception& e) { mlx_error(e.what()); return 1; } return 0; } -extern "C" int mlx_detail_compile_erase(uintptr_t fun_id) { +extern "C" int mlx_detail_compile_erase( + const mlx_compile_cache cache, + uintptr_t fun_id) { try { - mlx::core::detail::compile_erase(fun_id); + mlx::core::detail::compile_erase(mlx_compile_cache_get_(cache), fun_id); } catch (std::exception& e) { mlx_error(e.what()); return 1; diff --git a/Source/Cmlx/mlx-c/mlx/c/compile.h b/Source/Cmlx/mlx-c/mlx/c/compile.h index 04567fb3a..68605ce5f 100644 --- a/Source/Cmlx/mlx-c/mlx/c/compile.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/device.h b/Source/Cmlx/mlx-c/mlx/c/device.h index 4b74e39d3..705aea9db 100644 --- a/Source/Cmlx/mlx-c/mlx/c/device.h +++ b/Source/Cmlx/mlx-c/mlx/c/device.h @@ -7,7 +7,6 @@ #include #include "mlx/c/string.h" -#include "mlx/c/vector.h" #ifdef __cplusplus extern "C" { @@ -26,6 +25,12 @@ typedef struct mlx_device_ { void* ctx; } mlx_device; +/** + * A MLX vector of strings. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_string_ mlx_vector_string; + /** * Device type. */ diff --git a/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp b/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp index 47f84f70b..d0db5d5ae 100644 --- a/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/distributed_group.cpp @@ -6,49 +6,83 @@ #include "mlx/c/error.h" #include "mlx/c/private/mlx.h" -extern "C" int mlx_distributed_group_rank(mlx_distributed_group group) { +extern "C" mlx_distributed_group mlx_distributed_group_new(void) { try { - return mlx_distributed_group_get_(group).rank(); + return mlx_distributed_group_new_(); + } catch (std::exception& e) { + mlx_error(e.what()); + return {nullptr}; + } +} + +extern "C" int mlx_distributed_group_free(mlx_distributed_group group) { + try { + mlx_distributed_group_free_(group); } catch (std::exception& e) { mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int +mlx_distributed_init(mlx_distributed_group* res, bool strict, const char* bk) { + try { + if (bk) { + mlx_distributed_group_set_( + *res, mlx::core::distributed::init(strict, bk)); + } else { + mlx_distributed_group_set_(*res, mlx::core::distributed::init(strict)); + } return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; } } -extern "C" int mlx_distributed_group_size(mlx_distributed_group group) { +extern "C" int mlx_distributed_group_rank(mlx_distributed_group group) { try { - return mlx_distributed_group_get_(group).size(); + return mlx_distributed_group_get_(group).rank(); } catch (std::exception& e) { mlx_error(e.what()); return 0; } } -extern "C" mlx_distributed_group -mlx_distributed_group_split(mlx_distributed_group group, int color, int key) { +extern "C" int mlx_distributed_group_size(mlx_distributed_group group) { try { - return mlx_distributed_group_new_( - mlx_distributed_group_get_(group).split(color, key)); + return mlx_distributed_group_get_(group).size(); } catch (std::exception& e) { mlx_error(e.what()); - return mlx_distributed_group_new_(); + return 0; } } -extern "C" bool mlx_distributed_is_available(void) { +extern "C" int mlx_distributed_group_split( + mlx_distributed_group* res, + mlx_distributed_group group, + int color, + int key) { try { - return mlx::core::distributed::is_available(); + mlx_distributed_group_set_( + *res, mlx_distributed_group_get_(group).split(color, key)); + return 0; } catch (std::exception& e) { mlx_error(e.what()); - return false; + return 1; } } -extern "C" mlx_distributed_group mlx_distributed_init(bool strict) { +extern "C" bool mlx_distributed_is_available(const char* bk) { try { - return mlx_distributed_group_new_(mlx::core::distributed::init(strict)); + if (bk) { + return mlx::core::distributed::is_available(bk); + } else { + return mlx::core::distributed::is_available(); + } } catch (std::exception& e) { mlx_error(e.what()); - return mlx_distributed_group_new_(); + return false; } } diff --git a/Source/Cmlx/mlx-c/mlx/c/distributed_group.h b/Source/Cmlx/mlx-c/mlx/c/distributed_group.h index 3cfccc806..bfbaa80eb 100644 --- a/Source/Cmlx/mlx-c/mlx/c/distributed_group.h +++ b/Source/Cmlx/mlx-c/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,18 +54,16 @@ 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(void); - -/** - * Initialize distributed. - */ -mlx_distributed_group mlx_distributed_init(bool strict); +bool mlx_distributed_is_available(const char* bk /* may be null */); /**@}*/ diff --git a/Source/Cmlx/mlx-c/mlx/c/fast.cpp b/Source/Cmlx/mlx-c/mlx/c/fast.cpp index 1e07f26f8..79a788d4c 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fast.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/fast.cpp @@ -29,6 +29,25 @@ void prefault(const array& a) { } } +extern "C" int mlx_fast_cross_entropy( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::fast::cross_entropy( + mlx_array_get_(logits), + mlx_array_get_(targets), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + struct mlx_fast_cuda_kernel_config_cpp_ { std::vector output_shapes; std::vector output_dtypes; @@ -631,6 +650,7 @@ extern "C" 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) { try { mlx_array_set_( @@ -645,6 +665,7 @@ extern "C" int mlx_fast_scaled_dot_product_attention( : std::nullopt), (sinks.ctx ? std::make_optional(mlx_array_get_(sinks)) : std::nullopt), + force_fused, mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -653,10 +674,21 @@ extern "C" int mlx_fast_scaled_dot_product_attention( return 0; } +#include +#include #include +#include +#include +#include #include +#include #include +#include +#include +#include #include +#include +#include #include "mlx/backend/metal/ssd_streamer.h" struct SSDStreamEntry { @@ -770,8 +802,7 @@ extern "C" int mlx_fast_streamed_gather_mm( active_expert, streamer, eo, - mlx_stream_get_(s).device - )); + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -868,11 +899,6 @@ extern "C" int mlx_fast_prefault( // This gives full NVMe sequential throughput (~5 GB/s) while preserving all // MLX tensor metadata (shape, strides, dtype) on the dst array. // ───────────────────────────────────────────────────────────────────────────── -#include -#include -#include -#include - struct STPReadEntry { int fd = -1; size_t data_start = 0; @@ -918,7 +944,6 @@ static STPReadEntry get_safetensors_entry(const std::string& path, const std::st return entry; } -#include // --- PAPPS Async Background Worker (16-Thread Pool) --- struct PAPPSJob { std::string cache_id; @@ -1009,18 +1034,26 @@ class PAPPSQueue { bool try_take(const std::string& cache_id, void* dst, size_t length) { void* src = nullptr; + size_t stored = 0; { std::lock_guard lock(cache_mutex_); auto it = cache_.find(cache_id); if (it != cache_.end()) { src = it->second.first; + stored = it->second.second; cache_.erase(it); } } if (src) { - std::memcpy(dst, src, length); + // Only absorb the prefetched slice if it is exactly the size being + // requested; otherwise treat it as a miss so the caller preads + // synchronously instead of over-reading the cached buffer. + bool exact = (stored == length); + if (exact) { + std::memcpy(dst, src, length); + } free(src); - return true; + return exact; } return false; } diff --git a/Source/Cmlx/mlx-c/mlx/c/fast.h b/Source/Cmlx/mlx-c/mlx/c/fast.h index a25d14d80..2786c1ce8 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fast.h +++ b/Source/Cmlx/mlx-c/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; @@ -194,13 +200,17 @@ int mlx_fast_scaled_dot_product_attention( float scale, 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); +// ── SharpAI custom ops (not part of upstream mlx-c) ────────────────────────── + int mlx_fast_streamed_gather_mm( mlx_array* res, const mlx_array x, const mlx_array w_shape, - const mlx_array expert_indices, + uint32_t active_expert, const char* safetensors_path, const char* tensor_name, const mlx_stream s); @@ -227,16 +237,6 @@ int mlx_fast_turbo_decode_v( int mlx_fast_prefault(mlx_array x); - -// mlx_fast_submit_prefetch (PAPPS Background Worker) -// Instantly queues an asynchronous NVMe read into a background Thread Pool. -int mlx_fast_submit_prefetch( - const char* safetensors_path, - const char* tensor_name, - uint32_t expert_index); - -void mlx_fast_set_prefetch_enabled(bool enabled); - // pread() directly into the already-evaluated MLX array's unified memory buffer. // This gives full NVMe sequential throughput without OS page-fault overhead. // The array MUST already be evaluated (concrete pointer exists). @@ -261,11 +261,19 @@ int mlx_fast_pread_into_offset( uint32_t expert_index, size_t dst_offset); +// mlx_fast_submit_prefetch (PAPPS Background Worker) +int mlx_fast_submit_prefetch( + const char* safetensors_path, + const char* tensor_name, + uint32_t expert_index); + +void mlx_fast_set_prefetch_enabled(bool enabled); + /**@}*/ // ── SSD Flash-Stream metrics snapshot ──────────────────────────────────────── // Cumulative NVMe throughput stats since process start. -// Call mlx_ssd_metrics_snapshot() from any thread to read without resetting counters. +// Call mlx_ssd_metrics_snapshot() from any thread; never resets any counter. typedef struct MlxSSDMetricsSnapshot { double throughput_mb_per_s; /* 10-s rolling window average (0 before first window) */ @@ -274,8 +282,14 @@ typedef struct MlxSSDMetricsSnapshot { double avg_chunk_latency_ms; /* Lifetime average per-chunk latency (ms) */ } MlxSSDMetricsSnapshot; +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). void mlx_ssd_metrics_snapshot(MlxSSDMetricsSnapshot* out); +// TurboKV telemetry — call from Swift on each compression event to accumulate +// stats that appear in the 10-second SSD stream log. +// Implemented in the SharpAI mlx fork (mlx/core/moe_stream_op.cpp). +void mlx_turbo_kv_record(uint64_t tokens, uint64_t orig_bytes, uint64_t packed_bytes); + #ifdef __cplusplus } #endif diff --git a/Source/Cmlx/mlx-c/mlx/c/fft.cpp b/Source/Cmlx/mlx-c/mlx/c/fft.cpp index cd63ccbc6..c65cd436a 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fft.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/fft.cpp @@ -13,11 +13,17 @@ extern "C" int mlx_fft_fft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( *res, - mlx::core::fft::fft(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + mlx::core::fft::fft( + mlx_array_get_(a), + n, + axis, + mlx_fft_norm_to_cpp(norm), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -31,6 +37,7 @@ extern "C" int mlx_fft_fft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -39,6 +46,7 @@ extern "C" int mlx_fft_fft2( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -46,6 +54,16 @@ extern "C" int mlx_fft_fft2( } return 0; } +extern "C" int +mlx_fft_fftfreq(mlx_array* res, int n, double d, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::fft::fftfreq(n, d, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_fft_fftn( mlx_array* res, const mlx_array a, @@ -53,6 +71,7 @@ extern "C" int mlx_fft_fftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -61,6 +80,7 @@ extern "C" int mlx_fft_fftn( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -92,11 +112,17 @@ extern "C" int mlx_fft_ifft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( *res, - mlx::core::fft::ifft(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + mlx::core::fft::ifft( + mlx_array_get_(a), + n, + axis, + mlx_fft_norm_to_cpp(norm), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -110,6 +136,7 @@ extern "C" int mlx_fft_ifft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -118,6 +145,7 @@ extern "C" int mlx_fft_ifft2( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -132,6 +160,7 @@ extern "C" int mlx_fft_ifftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -140,6 +169,7 @@ extern "C" int mlx_fft_ifftn( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -171,11 +201,17 @@ extern "C" int mlx_fft_irfft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( *res, - mlx::core::fft::irfft(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + mlx::core::fft::irfft( + mlx_array_get_(a), + n, + axis, + mlx_fft_norm_to_cpp(norm), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -189,6 +225,7 @@ extern "C" int mlx_fft_irfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -197,6 +234,7 @@ extern "C" int mlx_fft_irfft2( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -211,6 +249,7 @@ extern "C" int mlx_fft_irfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -219,6 +258,7 @@ extern "C" int mlx_fft_irfftn( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -231,11 +271,17 @@ extern "C" int mlx_fft_rfft( const mlx_array a, int n, int axis, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( *res, - mlx::core::fft::rfft(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + mlx::core::fft::rfft( + mlx_array_get_(a), + n, + axis, + mlx_fft_norm_to_cpp(norm), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -249,6 +295,7 @@ extern "C" int mlx_fft_rfft2( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -257,6 +304,7 @@ extern "C" int mlx_fft_rfft2( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -264,6 +312,16 @@ extern "C" int mlx_fft_rfft2( } return 0; } +extern "C" int +mlx_fft_rfftfreq(mlx_array* res, int n, double d, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::fft::rfftfreq(n, d, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_fft_rfftn( mlx_array* res, const mlx_array a, @@ -271,6 +329,7 @@ extern "C" int mlx_fft_rfftn( size_t n_num, const int* axes, size_t axes_num, + mlx_fft_norm norm, const mlx_stream s) { try { mlx_array_set_( @@ -279,6 +338,7 @@ extern "C" int mlx_fft_rfftn( mlx_array_get_(a), mlx::core::Shape(n, n + n_num), std::vector(axes, axes + axes_num), + mlx_fft_norm_to_cpp(norm), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); diff --git a/Source/Cmlx/mlx-c/mlx/c/fft.h b/Source/Cmlx/mlx-c/mlx/c/fft.h index 779803e9b..7140b601f 100644 --- a/Source/Cmlx/mlx-c/mlx/c/fft.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/graph_utils.cpp b/Source/Cmlx/mlx-c/mlx/c/graph_utils.cpp new file mode 100644 index 000000000..e8cb80764 --- /dev/null +++ b/Source/Cmlx/mlx-c/mlx/c/graph_utils.cpp @@ -0,0 +1,81 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#include "mlx/c/graph_utils.h" +#include "mlx/c/error.h" +#include "mlx/c/private/mlx.h" +#include "mlx/graph_utils.h" + +extern "C" mlx_node_namer mlx_node_namer_new() { + try { + return mlx_node_namer_new_(mlx::core::NodeNamer()); + } catch (std::exception& e) { + mlx_error(e.what()); + } + return {nullptr}; +} +extern "C" int mlx_node_namer_free(mlx_node_namer namer) { + try { + mlx_node_namer_free_(namer); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} +extern "C" int mlx_node_namer_set_name( + mlx_node_namer namer, + const mlx_array arr, + const char* name) { + try { + mlx_node_namer_get_(namer).set_name(mlx_array_get_(arr), name); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} +extern "C" int mlx_node_namer_get_name( + const char** name, + mlx_node_namer namer, + const mlx_array arr) { + try { + *name = mlx_node_namer_get_(namer).get_name(mlx_array_get_(arr)).c_str(); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} +extern "C" int mlx_export_to_dot( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs) { + try { + mlx::core::export_to_dot( + CFileOutputStream::as_lvalue(CFileOutputStream(os)), + mlx_node_namer_get_(namer), + mlx_vector_array_get_(outputs)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_print_graph( + FILE* os, + const mlx_node_namer namer, + const mlx_vector_array outputs) { + try { + mlx::core::print_graph( + CFileOutputStream::as_lvalue(CFileOutputStream(os)), + mlx_node_namer_get_(namer), + mlx_vector_array_get_(outputs)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} diff --git a/Source/Cmlx/mlx-c/mlx/c/graph_utils.h b/Source/Cmlx/mlx-c/mlx/c/graph_utils.h new file mode 100644 index 000000000..81eec7a87 --- /dev/null +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/io.cpp b/Source/Cmlx/mlx-c/mlx/c/io.cpp index 9ba706398..24e8093f1 100644 --- a/Source/Cmlx/mlx-c/mlx/c/io.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/io.cpp @@ -30,6 +30,18 @@ extern "C" int mlx_load(mlx_array* res, const char* file, const mlx_stream s) { } return 0; } +extern "C" int +mlx_load_gguf(mlx_io_gguf* gguf, const char* file, const mlx_stream s) { + try { + auto cpp_gguf = mlx::core::load_gguf(file, mlx_stream_get_(s)); + mlx_io_gguf_set_(*gguf, std::move(cpp_gguf)); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + extern "C" int mlx_load_safetensors_reader( mlx_map_string_to_array* res_0, mlx_map_string_to_string* res_1, @@ -84,6 +96,17 @@ extern "C" int mlx_save(const char* file, const mlx_array a) { } return 0; } +extern "C" int mlx_save_gguf(const char* file, mlx_io_gguf gguf) { + try { + auto cpp_gguf = mlx_io_gguf_get_(gguf); + mlx::core::save_gguf(file, cpp_gguf.first, cpp_gguf.second); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + extern "C" int mlx_save_safetensors_writer( mlx_io_writer in_stream, const mlx_map_string_to_array param, diff --git a/Source/Cmlx/mlx-c/mlx/c/io.h b/Source/Cmlx/mlx-c/mlx/c/io.h index 6eb205c9a..e1a0c0a7d 100644 --- a/Source/Cmlx/mlx-c/mlx/c/io.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/io_types.cpp b/Source/Cmlx/mlx-c/mlx/c/io_types.cpp index 6a6668e6f..1a054a6ac 100644 --- a/Source/Cmlx/mlx-c/mlx/c/io_types.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/io_types.cpp @@ -83,3 +83,165 @@ extern "C" int mlx_io_writer_tostring(mlx_string* str_, mlx_io_writer io) { return 1; } } + +extern "C" mlx_io_gguf mlx_io_gguf_new(void) { + try { + return mlx_io_gguf_new_(mlx::core::GGUFLoad()); + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_io_gguf({nullptr}); + } +} + +extern "C" int mlx_io_gguf_free(mlx_io_gguf io) { + try { + mlx_io_gguf_free_(io); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_io_gguf_get_keys(mlx_vector_string* keys, mlx_io_gguf io) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).first; + std::vector cpp_keys; + cpp_keys.reserve(cpp_map.size()); + for (const auto& [key, value] : cpp_map) { + cpp_keys.push_back(key); + } + mlx_vector_string_set_(*keys, cpp_keys); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +extern "C" int +mlx_io_gguf_get_array(mlx_array* arr, mlx_io_gguf io, const char* key) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).first; + auto it = cpp_map.find(key); + if (it != cpp_map.end()) { + mlx_array_set_(*arr, it->second); + } else { + return 2; // not found + } + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +#define IMPLEMENT_GGUF_GET_METADATA(CNAME, CPPTYPE) \ + extern "C" int mlx_io_gguf_get_metadata_##CNAME( \ + mlx_##CNAME* res, mlx_io_gguf io, const char* key) { \ + try { \ + auto& cpp_map = mlx_io_gguf_get_(io).second; \ + auto it = cpp_map.find(key); \ + if (it != cpp_map.end()) { \ + if (std::holds_alternative(it->second)) { \ + mlx_##CNAME##_set_(*res, std::get(it->second)); \ + } else { \ + return 3; /* found key, wrong value type */ \ + } \ + } else { \ + return 2; /* not found */ \ + } \ + return 0; \ + } catch (std::exception & e) { \ + mlx_error(e.what()); \ + return 1; \ + } \ + } + +IMPLEMENT_GGUF_GET_METADATA(array, mlx::core::array); +IMPLEMENT_GGUF_GET_METADATA(string, std::string); +IMPLEMENT_GGUF_GET_METADATA(vector_string, std::vector); + +#define IMPLEMENT_GGUF_HAS_METADATA(CNAME, CPPTYPE) \ + extern "C" int mlx_io_gguf_has_metadata_##CNAME( \ + bool* flag, mlx_io_gguf io, const char* key) { \ + try { \ + auto& cpp_map = mlx_io_gguf_get_(io).second; \ + auto it = cpp_map.find(key); \ + if (it != cpp_map.end()) { \ + if (std::holds_alternative(it->second)) { \ + *flag = true; \ + } else { \ + *flag = false; \ + } \ + return 0; \ + } else { \ + *flag = false; \ + return 2; /* not found */ \ + } \ + return 0; \ + } catch (std::exception & e) { \ + *flag = false; \ + mlx_error(e.what()); \ + return 1; \ + } \ + } + +IMPLEMENT_GGUF_HAS_METADATA(array, mlx::core::array); +IMPLEMENT_GGUF_HAS_METADATA(string, std::string); +IMPLEMENT_GGUF_HAS_METADATA(vector_string, std::vector); + +extern "C" int +mlx_io_gguf_set_array(mlx_io_gguf io, const char* key, const mlx_array arr) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).first; + cpp_map.insert(std::make_pair(std::string(key), mlx_array_get_(arr))); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_io_gguf_set_metadata_array( + mlx_io_gguf io, + const char* key, + const mlx_array marr) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).second; + cpp_map.insert(std::make_pair(std::string(key), mlx_array_get_(marr))); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_io_gguf_set_metadata_string( + mlx_io_gguf io, + const char* key, + const char* mstr) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).second; + cpp_map.insert(std::make_pair(std::string(key), std::string(mstr))); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} + +extern "C" int mlx_io_gguf_set_metadata_vector_string( + mlx_io_gguf io, + const char* key, + const mlx_vector_string mvstr) { + try { + auto& cpp_map = mlx_io_gguf_get_(io).second; + cpp_map.insert( + std::make_pair(std::string(key), mlx_vector_string_get_(mvstr))); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} diff --git a/Source/Cmlx/mlx-c/mlx/c/io_types.h b/Source/Cmlx/mlx-c/mlx/c/io_types.h index 88349b57c..3f15873e3 100644 --- a/Source/Cmlx/mlx-c/mlx/c/io_types.h +++ b/Source/Cmlx/mlx-c/mlx/c/io_types.h @@ -37,10 +37,10 @@ typedef struct mlx_io_vtable_ { bool (*is_open)(void*); bool (*good)(void*); size_t (*tell)(void*); - void (*seek)(void*, int64_t off, int whence); - void (*read)(void*, char* data, size_t n); - void (*read_at_offset)(void*, char* data, size_t n, size_t off); - void (*write)(void*, const char* data, size_t n); + int (*seek)(void*, int64_t off, int whence); + size_t (*read)(void*, char* data, size_t n); + size_t (*read_at_offset)(void*, char* data, size_t n, size_t off); + size_t (*write)(void*, const char* data, size_t n); const char* (*label)(void*); void (*free)(void*); } mlx_io_vtable; @@ -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/mlx-c/mlx/c/linalg.cpp b/Source/Cmlx/mlx-c/mlx/c/linalg.cpp index a9a3c43c1..4ea04a8b0 100644 --- a/Source/Cmlx/mlx-c/mlx/c/linalg.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/linalg.cpp @@ -57,6 +57,17 @@ extern "C" int mlx_linalg_cross( } return 0; } +extern "C" int +mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::linalg::det(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -261,6 +272,24 @@ extern "C" int mlx_linalg_qr( } return 0; } +extern "C" int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s) { + try { + { + auto [tpl_0, tpl_1] = + mlx::core::linalg::slogdet(mlx_array_get_(a), mlx_stream_get_(s)); + mlx_array_set_(*res_0, tpl_0); + mlx_array_set_(*res_1, tpl_1); + }; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/Source/Cmlx/mlx-c/mlx/c/linalg.h b/Source/Cmlx/mlx-c/mlx/c/linalg.h index 91d5d661e..6d5b5ae15 100644 --- a/Source/Cmlx/mlx-c/mlx/c/linalg.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/metal.cpp b/Source/Cmlx/mlx-c/mlx/c/metal.cpp index 46d8034df..005685eae 100644 --- a/Source/Cmlx/mlx-c/mlx/c/metal.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/metal.cpp @@ -8,6 +8,15 @@ #include "mlx/c/error.h" #include "mlx/c/private/mlx.h" +extern "C" int mlx_metal_get_metallib_path(mlx_string* res) { + try { + mlx_string_set_(*res, mlx::core::metal::get_metallib_path()); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_metal_is_available(bool* res) { try { *res = mlx::core::metal::is_available(); @@ -17,6 +26,15 @@ extern "C" int mlx_metal_is_available(bool* res) { } return 0; } +extern "C" int mlx_metal_set_metallib_path(const char* path) { + try { + mlx::core::metal::set_metallib_path(std::string(path)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_metal_start_capture(const char* path) { try { mlx::core::metal::start_capture(std::string(path)); diff --git a/Source/Cmlx/mlx-c/mlx/c/metal.h b/Source/Cmlx/mlx-c/mlx/c/metal.h index 5877b224b..f65d15b62 100644 --- a/Source/Cmlx/mlx-c/mlx/c/metal.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/mlx.h b/Source/Cmlx/mlx-c/mlx/c/mlx.h index ffadac89a..2aa9077c6 100644 --- a/Source/Cmlx/mlx-c/mlx/c/mlx.h +++ b/Source/Cmlx/mlx-c/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/mlx-c/mlx/c/ops.cpp b/Source/Cmlx/mlx-c/mlx/c/ops.cpp index d054b63ab..bf565d1e3 100644 --- a/Source/Cmlx/mlx-c/mlx/c/ops.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/ops.cpp @@ -491,6 +491,15 @@ mlx_atleast_3d(mlx_array* res, const mlx_array a, const mlx_stream s) { } return 0; } +extern "C" int mlx_bartlett(mlx_array* res, int M, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::bartlett(M, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_bitwise_and( mlx_array* res, const mlx_array a, @@ -550,6 +559,15 @@ extern "C" int mlx_bitwise_xor( } return 0; } +extern "C" int mlx_blackman(mlx_array* res, int M, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::blackman(M, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_block_masked_mm( mlx_array* res, const mlx_array a, @@ -955,7 +973,61 @@ extern "C" int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s) { } return 0; } -extern "C" int mlx_cummax( +extern "C" int mlx_count_nonzero_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::count_nonzero( + mlx_array_get_(a), axis, keepdims, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::count_nonzero( + mlx_array_get_(a), + std::vector(axes, axes + axes_num), + keepdims, + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::count_nonzero( + mlx_array_get_(a), keepdims, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_cummax_axis( mlx_array* res, const mlx_array a, int axis, @@ -973,7 +1045,24 @@ extern "C" int mlx_cummax( } return 0; } -extern "C" int mlx_cummin( +extern "C" int mlx_cummax( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::cummax( + mlx_array_get_(a), reverse, inclusive, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_cummin_axis( mlx_array* res, const mlx_array a, int axis, @@ -991,36 +1080,117 @@ extern "C" int mlx_cummin( } return 0; } -extern "C" int mlx_cumprod( +extern "C" int mlx_cummin( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::cummin( + mlx_array_get_(a), reverse, inclusive, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { try { mlx_array_set_( *res, mlx::core::cumprod( - mlx_array_get_(a), axis, reverse, inclusive, mlx_stream_get_(s))); + mlx_array_get_(a), + axis, + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; } return 0; } -extern "C" int mlx_cumsum( +extern "C" int mlx_cumprod( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::cumprod( + mlx_array_get_(a), + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { try { mlx_array_set_( *res, mlx::core::cumsum( - mlx_array_get_(a), axis, reverse, inclusive, mlx_stream_get_(s))); + mlx_array_get_(a), + axis, + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_cumsum( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::cumsum( + mlx_array_get_(a), + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -1078,6 +1248,8 @@ extern "C" int mlx_dequantize( (bits.has_value ? std::make_optional(bits.value) : std::nullopt), std::string(mode), + (global_scale.ctx ? std::make_optional(mlx_array_get_(global_scale)) + : std::nullopt), (dtype.has_value ? std::make_optional( mlx_dtype_to_cpp(dtype.value)) : std::nullopt), @@ -1117,6 +1289,21 @@ extern "C" int mlx_diagonal( } return 0; } +extern "C" int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::diff(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_divide( mlx_array* res, const mlx_array a, @@ -1291,6 +1478,46 @@ extern "C" int mlx_flatten( } return 0; } +extern "C" int mlx_flip_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::flip( + mlx_array_get_(a), + std::vector(axes, axes + axes_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int +mlx_flip_axis(mlx_array* res, const mlx_array a, int axis, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::flip(mlx_array_get_(a), axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_flip(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::flip(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s) { try { @@ -1489,6 +1716,51 @@ extern "C" int mlx_gather_qmm( } return 0; } +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::gather_qqmm( + mlx_array_get_(x), + mlx_array_get_(w), + (scales_w.ctx ? std::make_optional(mlx_array_get_(scales_w)) + : std::nullopt), + (lhs_indices.ctx ? std::make_optional(mlx_array_get_(lhs_indices)) + : std::nullopt), + (rhs_indices.ctx ? std::make_optional(mlx_array_get_(rhs_indices)) + : std::nullopt), + (group_size.has_value ? std::make_optional(group_size.value) + : std::nullopt), + (bits.has_value ? std::make_optional(bits.value) + : std::nullopt), + std::string(mode), + (global_scale_x.ctx + ? std::make_optional(mlx_array_get_(global_scale_x)) + : std::nullopt), + (global_scale_w.ctx + ? std::make_optional(mlx_array_get_(global_scale_w)) + : std::nullopt), + sorted_indices, + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_greater( mlx_array* res, const mlx_array a, @@ -1540,6 +1812,24 @@ extern "C" int mlx_hadamard_transform( } return 0; } +extern "C" int mlx_hamming(mlx_array* res, int M, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::hamming(M, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_hanning(mlx_array* res, int M, const mlx_stream s) { + try { + mlx_array_set_(*res, mlx::core::hanning(M, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_identity(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s) { try { @@ -1721,6 +2011,30 @@ extern "C" int mlx_less_equal( } return 0; } +extern "C" int mlx_linspace_endpoint( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::linspace( + start, + stop, + num, + endpoint, + mlx_dtype_to_cpp(dtype), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_linspace( mlx_array* res, double start, @@ -1796,7 +2110,7 @@ extern "C" int mlx_logaddexp( } return 0; } -extern "C" int mlx_logcumsumexp( +extern "C" int mlx_logcumsumexp_axis( mlx_array* res, const mlx_array a, int axis, @@ -1814,6 +2128,23 @@ extern "C" int mlx_logcumsumexp( } return 0; } +extern "C" int mlx_logcumsumexp( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::logcumsumexp( + mlx_array_get_(a), reverse, inclusive, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_logical_and( mlx_array* res, const mlx_array a, @@ -1834,14 +2165,30 @@ extern "C" int mlx_logical_not(mlx_array* res, const mlx_array a, const mlx_stream s) { try { mlx_array_set_( - *res, mlx::core::logical_not(mlx_array_get_(a), mlx_stream_get_(s))); + *res, mlx::core::logical_not(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_logical_or( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::logical_or( + mlx_array_get_(a), mlx_array_get_(b), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; } return 0; } -extern "C" int mlx_logical_or( +extern "C" int mlx_logical_xor( mlx_array* res, const mlx_array a, const mlx_array b, @@ -1849,7 +2196,7 @@ extern "C" int mlx_logical_or( try { mlx_array_set_( *res, - mlx::core::logical_or( + mlx::core::logical_xor( mlx_array_get_(a), mlx_array_get_(b), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -2058,7 +2405,7 @@ mlx_mean(mlx_array* res, const mlx_array a, bool keepdims, const mlx_stream s) { } return 0; } -extern "C" int mlx_median( +extern "C" int mlx_median_axes( mlx_array* res, const mlx_array a, const int* axes, @@ -2079,6 +2426,38 @@ extern "C" int mlx_median( } return 0; } +extern "C" int mlx_median_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::median( + mlx_array_get_(a), axis, keepdims, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_median( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::median(mlx_array_get_(a), keepdims, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_meshgrid( mlx_vector_array* res, const mlx_vector_array arrays, @@ -2289,6 +2668,22 @@ extern "C" int mlx_ones( } return 0; } +extern "C" int mlx_ones_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::ones_like( + mlx_array_get_(a), mlx_dtype_to_cpp(dtype), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s) { try { @@ -2394,6 +2789,17 @@ mlx_partition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) { } return 0; } +extern "C" int +mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::positive(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_power( mlx_array* res, const mlx_array a, @@ -2504,6 +2910,12 @@ extern "C" int mlx_qqmm( (bits.has_value ? std::make_optional(bits.value) : std::nullopt), std::string(mode), + (global_scale_x.ctx + ? std::make_optional(mlx_array_get_(global_scale_x)) + : std::nullopt), + (global_scale_w.ctx + ? std::make_optional(mlx_array_get_(global_scale_w)) + : std::nullopt), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -2529,6 +2941,8 @@ extern "C" int mlx_quantize( (bits.has_value ? std::make_optional(bits.value) : std::nullopt), std::string(mode), + (global_scale.ctx ? std::make_optional(mlx_array_get_(global_scale)) + : std::nullopt), mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); @@ -3016,6 +3430,26 @@ extern "C" int mlx_scatter_prod_single( } return 0; } +extern "C" int mlx_searchsorted( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::searchsorted( + mlx_array_get_(sorted_sequence), + mlx_array_get_(values), + std::string(side), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_segmented_mm( mlx_array* res, const mlx_array a, @@ -3175,6 +3609,114 @@ extern "C" int mlx_slice_update_dynamic( } return 0; } +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::slice_update_add( + mlx_array_get_(src), + mlx_array_get_(update), + mlx::core::Shape(start, start + start_num), + mlx::core::Shape(stop, stop + stop_num), + mlx::core::Shape(strides, strides + strides_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::slice_update_max( + mlx_array_get_(src), + mlx_array_get_(update), + mlx::core::Shape(start, start + start_num), + mlx::core::Shape(stop, stop + stop_num), + mlx::core::Shape(strides, strides + strides_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::slice_update_min( + mlx_array_get_(src), + mlx_array_get_(update), + mlx::core::Shape(start, start + start_num), + mlx::core::Shape(stop, stop + stop_num), + mlx::core::Shape(strides, strides + strides_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" 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) { + try { + mlx_array_set_( + *res, + mlx::core::slice_update_prod( + mlx_array_get_(src), + mlx_array_get_(update), + mlx::core::Shape(start, start + start_num), + mlx::core::Shape(stop, stop + stop_num), + mlx::core::Shape(strides, strides + strides_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_softmax_axes( mlx_array* res, const mlx_array a, @@ -3701,7 +4243,7 @@ mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s) { } return 0; } -extern "C" int mlx_trace( +extern "C" int mlx_trace_axes( mlx_array* res, const mlx_array a, int offset, @@ -3725,6 +4267,17 @@ extern "C" int mlx_trace( } return 0; } +extern "C" int +mlx_trace(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::trace(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_transpose_axes( mlx_array* res, const mlx_array a, @@ -3794,6 +4347,17 @@ mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s) { } return 0; } +extern "C" int +mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::trunc(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -3815,6 +4379,31 @@ extern "C" int mlx_unflatten( } return 0; } +extern "C" int mlx_unstack_axis( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s) { + try { + mlx_vector_array_set_( + *res, mlx::core::unstack(mlx_array_get_(a), axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int +mlx_unstack(mlx_vector_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_vector_array_set_( + *res, mlx::core::unstack(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -3872,6 +4461,23 @@ extern "C" int mlx_var( } return 0; } +extern "C" int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::vecdot( + mlx_array_get_(a), mlx_array_get_(b), axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_view( mlx_array* res, const mlx_array a, @@ -3927,6 +4533,22 @@ extern "C" int mlx_zeros( } return 0; } +extern "C" int mlx_zeros_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::zeros_like( + mlx_array_get_(a), mlx_dtype_to_cpp(dtype), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s) { try { @@ -3938,23 +4560,3 @@ mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s) { } return 0; } - -extern "C" int mlx_bartlett(mlx_array* res, int M, const mlx_stream s) { - mlx_error("mlx_bartlett not implemented in SharpAI MLX fork"); - return 1; -} - -extern "C" int mlx_blackman(mlx_array* res, int M, const mlx_stream s) { - mlx_error("mlx_blackman not implemented in SharpAI MLX fork"); - return 1; -} - -extern "C" int mlx_hamming(mlx_array* res, int M, const mlx_stream s) { - mlx_error("mlx_hamming not implemented in SharpAI MLX fork"); - return 1; -} - -extern "C" int mlx_hanning(mlx_array* res, int M, const mlx_stream s) { - mlx_error("mlx_hanning not implemented in SharpAI MLX fork"); - return 1; -} diff --git a/Source/Cmlx/mlx-c/mlx/c/ops.h b/Source/Cmlx/mlx-c/mlx/c/ops.h index a1446fb9e..741778570 100644 --- a/Source/Cmlx/mlx-c/mlx/c/ops.h +++ b/Source/Cmlx/mlx-c/mlx/c/ops.h @@ -166,6 +166,7 @@ int mlx_astype( int mlx_atleast_1d(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_atleast_2d(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_atleast_3d(mlx_array* res, const mlx_array a, const mlx_stream s); +int mlx_bartlett(mlx_array* res, int M, const mlx_stream s); int mlx_bitwise_and( mlx_array* res, const mlx_array a, @@ -182,6 +183,7 @@ int mlx_bitwise_xor( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_blackman(mlx_array* res, int M, const mlx_stream s); int mlx_block_masked_mm( mlx_array* res, const mlx_array a, @@ -321,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, @@ -329,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( @@ -362,6 +410,7 @@ int mlx_dequantize( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, mlx_optional_dtype dtype, const mlx_stream s); int mlx_diag(mlx_array* res, const mlx_array a, int k, const mlx_stream s); @@ -372,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, @@ -420,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, @@ -483,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, @@ -498,6 +579,8 @@ int mlx_hadamard_transform( const mlx_array a, mlx_optional_float scale, const mlx_stream s); +int mlx_hamming(mlx_array* res, int M, const mlx_stream s); +int mlx_hanning(mlx_array* res, int M, const mlx_stream s); int mlx_identity(mlx_array* res, int n, mlx_dtype dtype, const mlx_stream s); int mlx_imag(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_inner( @@ -538,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, @@ -554,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, @@ -572,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, @@ -642,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, @@ -716,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, @@ -752,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, @@ -790,6 +909,8 @@ int mlx_qqmm( 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 */, const mlx_stream s); int mlx_quantize( mlx_vector_array* res, @@ -797,6 +918,7 @@ int mlx_quantize( mlx_optional_int group_size, mlx_optional_int bits, const char* mode, + const mlx_array global_scale /* may be null */, const mlx_stream s); int mlx_quantized_matmul( mlx_array* res, @@ -948,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, @@ -996,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, @@ -1155,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, @@ -1163,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, @@ -1179,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, @@ -1186,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, @@ -1207,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, @@ -1224,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/mlx-c/mlx/c/private/compile.h b/Source/Cmlx/mlx-c/mlx/c/private/compile.h new file mode 100644 index 000000000..bf3cef2f9 --- /dev/null +++ b/Source/Cmlx/mlx-c/mlx/c/private/compile.h @@ -0,0 +1,63 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#ifndef MLX_COMPILE_PRIVATE_H +#define MLX_COMPILE_PRIVATE_H + +#include "mlx/c/compile.h" +#include "mlx/compile_impl.h" + +inline mlx_compile_cache mlx_compile_cache_new_() { + return mlx_compile_cache({nullptr}); +} + +inline mlx_compile_cache mlx_compile_cache_new_( + const mlx::core::detail::CompileCacheWeakPtr& s) { + return mlx_compile_cache({new mlx::core::detail::CompileCacheWeakPtr(s)}); +} + +inline mlx_compile_cache mlx_compile_cache_new_( + mlx::core::detail::CompileCacheWeakPtr&& s) { + return mlx_compile_cache( + {new mlx::core::detail::CompileCacheWeakPtr(std::move(s))}); +} + +inline mlx_compile_cache& mlx_compile_cache_set_( + mlx_compile_cache& d, + const mlx::core::detail::CompileCacheWeakPtr& s) { + if (d.ctx) { + *static_cast(d.ctx) = s; + } else { + d.ctx = new mlx::core::detail::CompileCacheWeakPtr(s); + } + return d; +} + +inline mlx_compile_cache& mlx_compile_cache_set_( + mlx_compile_cache& d, + mlx::core::detail::CompileCacheWeakPtr&& s) { + if (d.ctx) { + *static_cast(d.ctx) = std::move(s); + } else { + d.ctx = new mlx::core::detail::CompileCacheWeakPtr(std::move(s)); + } + return d; +} + +inline mlx::core::detail::CompileCacheWeakPtr& mlx_compile_cache_get_( + mlx_compile_cache d) { + if (!d.ctx) { + throw std::runtime_error("expected a non-empty mlx_compile_cache"); + } + return *static_cast(d.ctx); +} + +inline void mlx_compile_cache_free_(mlx_compile_cache d) { + if (d.ctx) { + delete static_cast(d.ctx); + } +} + +#endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/enums.h b/Source/Cmlx/mlx-c/mlx/c/private/enums.h index 098709c0c..59a1e7878 100644 --- a/Source/Cmlx/mlx-c/mlx/c/private/enums.h +++ b/Source/Cmlx/mlx-c/mlx/c/private/enums.h @@ -5,6 +5,7 @@ #include "mlx/c/array.h" #include "mlx/c/compile.h" +#include "mlx/c/fft.h" #include "mlx/mlx.h" namespace { @@ -73,6 +74,18 @@ inline mlx::core::Device::DeviceType mlx_device_type_to_cpp( mlx::core::Device::DeviceType::cpu, mlx::core::Device::DeviceType::gpu}; return map[(int)type]; } +inline mlx_fft_norm mlx_fft_norm_to_c(mlx::core::fft::FFTNorm norm) { + static mlx_fft_norm map[] = { + MLX_FFT_NORM_BACKWARD, MLX_FFT_NORM_ORTHO, MLX_FFT_NORM_FORWARD}; + return map[(int)norm]; +} +inline mlx::core::fft::FFTNorm mlx_fft_norm_to_cpp(mlx_fft_norm norm) { + static mlx::core::fft::FFTNorm map[] = { + mlx::core::fft::FFTNorm::Backward, + mlx::core::fft::FFTNorm::Ortho, + mlx::core::fft::FFTNorm::Forward}; + return map[(int)norm]; +} } // namespace #endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/gguf.h b/Source/Cmlx/mlx-c/mlx/c/private/gguf.h new file mode 100644 index 000000000..58b087c09 --- /dev/null +++ b/Source/Cmlx/mlx-c/mlx/c/private/gguf.h @@ -0,0 +1,41 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#ifndef MLX_IO_TYPES_PRIVATE_H +#define MLX_IO_TYPES_PRIVATE_H + +#include "mlx/c/io_types.h" +#include "mlx/mlx.h" + +inline mlx_io_gguf mlx_io_gguf_new_() { + return mlx_io_gguf({nullptr}); +} + +inline mlx_io_gguf mlx_io_gguf_new_(mlx::core::GGUFLoad&& s) { + return mlx_io_gguf({new mlx::core::GGUFLoad(std::move(s))}); +} + +inline mlx_io_gguf& mlx_io_gguf_set_(mlx_io_gguf& d, mlx::core::GGUFLoad&& s) { + if (d.ctx) { + delete static_cast(d.ctx); + } + d.ctx = new mlx::core::GGUFLoad(std::move(s)); + return d; +} + +inline mlx::core::GGUFLoad& mlx_io_gguf_get_(mlx_io_gguf d) { + if (!d.ctx) { + throw std::runtime_error("expected a non-empty mlx_io_gguf"); + } + return *static_cast(d.ctx); +} + +inline void mlx_io_gguf_free_(mlx_io_gguf d) { + if (d.ctx) { + delete static_cast(d.ctx); + } +} + +#endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/graph_utils.h b/Source/Cmlx/mlx-c/mlx/c/private/graph_utils.h new file mode 100644 index 000000000..a9def0738 --- /dev/null +++ b/Source/Cmlx/mlx-c/mlx/c/private/graph_utils.h @@ -0,0 +1,43 @@ +/* Copyright © 2023-2024 Apple Inc. */ +/* */ +/* This file is auto-generated. Do not edit manually. */ +/* */ + +#ifndef MLX_GRAPH_UTILS_PRIVATE_H +#define MLX_GRAPH_UTILS_PRIVATE_H + +#include "mlx/c/graph_utils.h" +#include "mlx/graph_utils.h" + +inline mlx_node_namer mlx_node_namer_new_() { + return mlx_node_namer({nullptr}); +} + +inline mlx_node_namer mlx_node_namer_new_(mlx::core::NodeNamer&& s) { + return mlx_node_namer({new mlx::core::NodeNamer(std::move(s))}); +} + +inline mlx_node_namer& mlx_node_namer_set_( + mlx_node_namer& d, + mlx::core::NodeNamer&& s) { + if (d.ctx) { + delete static_cast(d.ctx); + } + d.ctx = new mlx::core::NodeNamer(std::move(s)); + return d; +} + +inline mlx::core::NodeNamer& mlx_node_namer_get_(mlx_node_namer d) { + if (!d.ctx) { + throw std::runtime_error("expected a non-empty mlx_node_namer"); + } + return *static_cast(d.ctx); +} + +inline void mlx_node_namer_free_(mlx_node_namer d) { + if (d.ctx) { + delete static_cast(d.ctx); + } +} + +#endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/io.h b/Source/Cmlx/mlx-c/mlx/c/private/io.h index fc99f893b..f17ceb499 100644 --- a/Source/Cmlx/mlx-c/mlx/c/private/io.h +++ b/Source/Cmlx/mlx-c/mlx/c/private/io.h @@ -2,6 +2,8 @@ #define MLX_IO_PRIVATE_H #include +#include + #include "mlx/mlx.h" namespace { @@ -24,29 +26,45 @@ class CReader : public mlx::core::io::Reader { virtual void seek( int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int status = 0; switch (way) { case std::ios_base::beg: - return vtable.seek(desc, off, SEEK_SET); + status = vtable.seek(desc, off, SEEK_SET); break; case std::ios_base::cur: - return vtable.seek(desc, off, SEEK_CUR); + status = vtable.seek(desc, off, SEEK_CUR); break; case std::ios_base::end: - return vtable.seek(desc, off, SEEK_END); + status = vtable.seek(desc, off, SEEK_END); break; default: - throw std::runtime_error("mlx_io_reader: invalid seek way"); + reader_error("invalid seek way"); + } + if (status < 0) { + reader_error("unable to seek"); } } virtual void read(char* data, size_t n) override { - return vtable.read(desc, data, n); + auto read_n = vtable.read(desc, data, n); + if (read_n != n) { + reader_error( + std::format("unable to read {} bytes (read {} instead)", n, read_n)); + } }; virtual void read(char* data, size_t n, size_t offset) override { - return vtable.read_at_offset(desc, data, n, offset); + auto read_n = vtable.read_at_offset(desc, data, n, offset); + if (read_n != n) { + reader_error( + std::format("unable to read {} bytes (read {} instead)", n, read_n)); + } }; virtual std::string label() const override { return vtable.label(desc); }; + void reader_error(const std::string& msg) { + throw std::runtime_error( + std::format("[mlx_io_reader] {} in {}", msg, label())); + } virtual ~CReader() { vtable.free(desc); } @@ -70,26 +88,39 @@ class CWriter : public mlx::core::io::Writer { virtual void seek( int64_t off, std::ios_base::seekdir way = std::ios_base::beg) override { + int status = 0; switch (way) { case std::ios_base::beg: - return vtable.seek(desc, off, SEEK_SET); + status = vtable.seek(desc, off, SEEK_SET); break; case std::ios_base::cur: - return vtable.seek(desc, off, SEEK_CUR); + status = vtable.seek(desc, off, SEEK_CUR); break; case std::ios_base::end: - return vtable.seek(desc, off, SEEK_END); + status = vtable.seek(desc, off, SEEK_END); break; default: - throw std::runtime_error("mlx_io_writer: invalid seek way"); + writer_error("invalid seek way"); + } + if (status < 0) { + writer_error("unable to seek"); } } virtual void write(const char* data, size_t n) override { - return vtable.write(desc, data, n); + auto wrote_n = vtable.write(desc, data, n); + if (wrote_n != n) { + writer_error( + std::format( + "unable to write {} bytes (wrote {} instead)", n, wrote_n)); + } }; virtual std::string label() const override { return vtable.label(desc); }; + void writer_error(const std::string& msg) { + throw std::runtime_error( + std::format("[mlx_io_writer] {} in {}", msg, label())); + } virtual ~CWriter() { vtable.free(desc); } @@ -139,6 +170,43 @@ inline void mlx_io_writer_free_(mlx_io_writer io) { } } +class CFileStreamBuf : public std::streambuf { + public: + explicit CFileStreamBuf(FILE* file) : file_(file) {} + + protected: + int_type overflow(int_type c) override { + if (c != traits_type::eof()) { + if (std::fputc(c, file_) == EOF) { + return traits_type::eof(); + } + } + return c; + } + std::streamsize xsputn(const char* s, std::streamsize n) override { + return std::fwrite(s, 1, n, file_); + } + int sync() override { + return std::fflush(file_) == 0 ? 0 : -1; + } + + private: + FILE* file_; +}; + +class CFileOutputStream : public std::ostream { + public: + explicit CFileOutputStream(FILE* file) : std::ostream(&buf_), buf_(file) {} + + template + static T& as_lvalue(T&& t) { + return t; + } + + private: + CFileStreamBuf buf_; +}; + } // namespace #endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/mlx.h b/Source/Cmlx/mlx-c/mlx/c/private/mlx.h index 496ddd0f3..fef37616c 100644 --- a/Source/Cmlx/mlx-c/mlx/c/private/mlx.h +++ b/Source/Cmlx/mlx-c/mlx/c/private/mlx.h @@ -3,10 +3,13 @@ #include "mlx/c/private/array.h" #include "mlx/c/private/closure.h" +#include "mlx/c/private/compile.h" #include "mlx/c/private/device.h" #include "mlx/c/private/distributed_group.h" #include "mlx/c/private/enums.h" #include "mlx/c/private/export.h" +#include "mlx/c/private/gguf.h" +#include "mlx/c/private/graph_utils.h" #include "mlx/c/private/io.h" #include "mlx/c/private/map.h" #include "mlx/c/private/stream.h" diff --git a/Source/Cmlx/mlx-c/mlx/c/private/stream.h b/Source/Cmlx/mlx-c/mlx/c/private/stream.h index 4fec359db..55d7b4d4a 100644 --- a/Source/Cmlx/mlx-c/mlx/c/private/stream.h +++ b/Source/Cmlx/mlx-c/mlx/c/private/stream.h @@ -52,4 +52,55 @@ inline void mlx_stream_free_(mlx_stream d) { } } +inline mlx_stream_thread_local mlx_stream_thread_local_new_() { + return mlx_stream_thread_local({nullptr}); +} + +inline mlx_stream_thread_local mlx_stream_thread_local_new_( + const mlx::core::ThreadLocalStream& s) { + return mlx_stream_thread_local({new mlx::core::ThreadLocalStream(s)}); +} + +inline mlx_stream_thread_local mlx_stream_thread_local_new_( + mlx::core::ThreadLocalStream&& s) { + return mlx_stream_thread_local( + {new mlx::core::ThreadLocalStream(std::move(s))}); +} + +inline mlx_stream_thread_local& mlx_stream_thread_local_set_( + mlx_stream_thread_local& d, + const mlx::core::ThreadLocalStream& s) { + if (d.ctx) { + *static_cast(d.ctx) = s; + } else { + d.ctx = new mlx::core::ThreadLocalStream(s); + } + return d; +} + +inline mlx_stream_thread_local& mlx_stream_thread_local_set_( + mlx_stream_thread_local& d, + mlx::core::ThreadLocalStream&& s) { + if (d.ctx) { + *static_cast(d.ctx) = std::move(s); + } else { + d.ctx = new mlx::core::ThreadLocalStream(std::move(s)); + } + return d; +} + +inline mlx::core::ThreadLocalStream& mlx_stream_thread_local_get_( + mlx_stream_thread_local d) { + if (!d.ctx) { + throw std::runtime_error("expected a non-empty mlx_stream_thread_local"); + } + return *static_cast(d.ctx); +} + +inline void mlx_stream_thread_local_free_(mlx_stream_thread_local d) { + if (d.ctx) { + delete static_cast(d.ctx); + } +} + #endif diff --git a/Source/Cmlx/mlx-c/mlx/c/private/vector.h b/Source/Cmlx/mlx-c/mlx/c/private/vector.h index 0d98042f1..09d8d0ff6 100644 --- a/Source/Cmlx/mlx-c/mlx/c/private/vector.h +++ b/Source/Cmlx/mlx-c/mlx/c/private/vector.h @@ -207,4 +207,54 @@ inline void mlx_vector_string_free_(mlx_vector_string d) { } } +inline mlx_vector_stream mlx_vector_stream_new_() { + return mlx_vector_stream({nullptr}); +} + +inline mlx_vector_stream mlx_vector_stream_new_( + const std::vector& s) { + return mlx_vector_stream({new std::vector(s)}); +} + +inline mlx_vector_stream mlx_vector_stream_new_( + std::vector&& s) { + return mlx_vector_stream({new std::vector(std::move(s))}); +} + +inline mlx_vector_stream& mlx_vector_stream_set_( + mlx_vector_stream& d, + const std::vector& s) { + if (d.ctx) { + *static_cast*>(d.ctx) = s; + } else { + d.ctx = new std::vector(s); + } + return d; +} + +inline mlx_vector_stream& mlx_vector_stream_set_( + mlx_vector_stream& d, + std::vector&& s) { + if (d.ctx) { + *static_cast*>(d.ctx) = std::move(s); + } else { + d.ctx = new std::vector(std::move(s)); + } + return d; +} + +inline std::vector& mlx_vector_stream_get_( + mlx_vector_stream d) { + if (!d.ctx) { + throw std::runtime_error("expected a non-empty mlx_vector_stream"); + } + return *static_cast*>(d.ctx); +} + +inline void mlx_vector_stream_free_(mlx_vector_stream d) { + if (d.ctx) { + delete static_cast*>(d.ctx); + } +} + #endif diff --git a/Source/Cmlx/mlx-c/mlx/c/stream.cpp b/Source/Cmlx/mlx-c/mlx/c/stream.cpp index 2d17997b3..d55254e83 100644 --- a/Source/Cmlx/mlx-c/mlx/c/stream.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/stream.cpp @@ -32,6 +32,67 @@ extern "C" mlx_stream mlx_stream_new_device(mlx_device dev) { return mlx_stream_new_(); } } +extern "C" mlx_stream mlx_stream_new_thread_unsafe(mlx_device dev) { + try { + return mlx_stream_new_( + mlx::core::new_thread_unsafe_stream(mlx_device_get_(dev))); + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_stream_new_(); + } +} +extern "C" mlx_stream_thread_local mlx_stream_thread_local_new(mlx_device dev) { + try { + return mlx_stream_thread_local_new_( + mlx::core::new_thread_local_stream(mlx_device_get_(dev))); + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_stream_thread_local_new_(); + } +} +extern "C" int mlx_stream_thread_local_set( + mlx_stream_thread_local* tls, + const mlx_stream_thread_local src) { + try { + mlx_stream_thread_local_set_(*tls, mlx_stream_thread_local_get_(src)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_stream_thread_local_free(mlx_stream_thread_local tls) { + try { + mlx_stream_thread_local_free_(tls); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_stream_from_thread_local( + mlx_stream* res, + const mlx_stream_thread_local tls) { + try { + mlx_stream_set_( + *res, + mlx::core::stream_from_thread_local_stream( + mlx_stream_thread_local_get_(tls))); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} +extern "C" int mlx_get_streams(mlx_vector_stream* res) { + try { + mlx_vector_stream_set_(*res, mlx::core::get_streams()); + return 0; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } +} extern "C" int mlx_stream_set(mlx_stream* stream, const mlx_stream src) { try { mlx_stream_set_(*stream, mlx_stream_get_(src)); @@ -80,6 +141,33 @@ extern "C" int mlx_synchronize(mlx_stream stream) { } return 0; } +extern "C" int mlx_synchronize_default(void) { + try { + mlx::core::synchronize(); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_synchronize_thread_local(mlx_stream_thread_local tls) { + try { + mlx::core::synchronize(mlx_stream_thread_local_get_(tls)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} +extern "C" int mlx_clear_streams(void) { + try { + mlx::core::clear_streams(); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_get_default_stream(mlx_stream* stream, mlx_device dev) { try { mlx_stream_set_(*stream, mlx::core::default_stream(mlx_device_get_(dev))); diff --git a/Source/Cmlx/mlx-c/mlx/c/stream.h b/Source/Cmlx/mlx-c/mlx/c/stream.h index d5865b806..6df1a63c9 100644 --- a/Source/Cmlx/mlx-c/mlx/c/stream.h +++ b/Source/Cmlx/mlx-c/mlx/c/stream.h @@ -4,6 +4,7 @@ #define MLX_STREAM_H #include +#include #include "mlx/c/device.h" @@ -24,6 +25,22 @@ typedef struct mlx_stream_ { void* ctx; } mlx_stream; +/** + * A per-thread stream object. + * + * The same thread-local stream can be used from any number of threads; the + * stream's GPU command encoder is registered once per thread. + */ +typedef struct mlx_stream_thread_local_ { + void* ctx; +} mlx_stream_thread_local; + +/** + * A vector of mlx_stream. + * Forward declaration to avoid circular references. + */ +typedef struct mlx_vector_stream_ mlx_vector_stream; + /** * Returns a new empty stream. */ @@ -33,6 +50,48 @@ 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); + +/** + * Returns a new per-thread stream on a device. + * + * Streams created this way are unique per thread: evaluating the same + * `mlx_stream_thread_local` from different threads yields distinct streams. + */ +mlx_stream_thread_local mlx_stream_thread_local_new(mlx_device dev); +/** + * Set a per-thread stream to the provided src stream. + */ +int mlx_stream_thread_local_set( + mlx_stream_thread_local* tls, + const mlx_stream_thread_local src); +/** + * Free a per-thread stream. + */ +int mlx_stream_thread_local_free(mlx_stream_thread_local tls); +/** + * Returns the concrete stream backing a per-thread stream on the current + * thread. + */ +int mlx_stream_from_thread_local( + mlx_stream* res, + const mlx_stream_thread_local tls); + +/** + * Get the vector of available streams. + */ +int mlx_get_streams(mlx_vector_stream* res); /** * Set stream to provided src stream. */ @@ -61,6 +120,18 @@ int mlx_stream_get_index(int* index, mlx_stream stream); * Synchronize with the provided stream. */ int mlx_synchronize(mlx_stream stream); +/** + * Synchronize with the default stream. + */ +int mlx_synchronize_default(void); +/** + * Synchronize with the stream corresponding to the current thread. + */ +int mlx_synchronize_thread_local(mlx_stream_thread_local tls); +/** + * Destroy all streams created in the current thread. + */ +int mlx_clear_streams(void); /** * Returns the default stream on the given device. */ diff --git a/Source/Cmlx/mlx-c/mlx/c/vector.cpp b/Source/Cmlx/mlx-c/mlx/c/vector.cpp index b1d83aaa7..63f1ea29d 100644 --- a/Source/Cmlx/mlx-c/mlx/c/vector.cpp +++ b/Source/Cmlx/mlx-c/mlx/c/vector.cpp @@ -529,3 +529,137 @@ extern "C" size_t mlx_vector_string_size(mlx_vector_string vec) { return 0; } } + +extern "C" mlx_vector_stream mlx_vector_stream_new(void) { + try { + return mlx_vector_stream_new_({}); + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_vector_stream_new_(); + } +} + +extern "C" int mlx_vector_stream_set( + mlx_vector_stream* vec, + const mlx_vector_stream src) { + try { + mlx_vector_stream_set_(*vec, mlx_vector_stream_get_(src)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int mlx_vector_stream_free(mlx_vector_stream vec) { + try { + mlx_vector_stream_free_(vec); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" mlx_vector_stream mlx_vector_stream_new_data( + const mlx_stream* data, + size_t size) { + try { + auto vec = mlx_vector_stream_new(); + for (size_t i = 0; i < size; i++) { + mlx_vector_stream_get_(vec).push_back(mlx_stream_get_(data[i])); + } + return vec; + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_vector_stream_new_(); + } +} + +extern "C" mlx_vector_stream mlx_vector_stream_new_value(const mlx_stream val) { + try { + return mlx_vector_stream_new_({mlx_stream_get_(val)}); + } catch (std::exception& e) { + mlx_error(e.what()); + return mlx_vector_stream_new_(); + } +} + +extern "C" int mlx_vector_stream_set_data( + mlx_vector_stream* vec_, + const mlx_stream* data, + size_t size) { + try { + std::vector cpp_arrs; + for (size_t i = 0; i < size; i++) { + cpp_arrs.push_back(mlx_stream_get_(data[i])); + } + mlx_vector_stream_set_(*vec_, cpp_arrs); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int mlx_vector_stream_set_value( + mlx_vector_stream* vec_, + const mlx_stream val) { + try { + mlx_vector_stream_set_( + *vec_, std::vector({mlx_stream_get_(val)})); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int mlx_vector_stream_append_data( + mlx_vector_stream vec, + const mlx_stream* data, + size_t size) { + try { + for (size_t i = 0; i < size; i++) { + mlx_vector_stream_get_(vec).push_back(mlx_stream_get_(data[i])); + } + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int mlx_vector_stream_append_value( + mlx_vector_stream vec, + const mlx_stream value) { + try { + mlx_vector_stream_get_(vec).push_back(mlx_stream_get_(value)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" int mlx_vector_stream_get( + mlx_stream* res, + const mlx_vector_stream vec, + size_t index) { + try { + mlx_stream_set_(*res, mlx_vector_stream_get_(vec).at(index)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} + +extern "C" size_t mlx_vector_stream_size(mlx_vector_stream vec) { + try { + return mlx_vector_stream_get_(vec).size(); + } catch (std::exception& e) { + mlx_error(e.what()); + return 0; + } +} diff --git a/Source/Cmlx/mlx-c/mlx/c/vector.h b/Source/Cmlx/mlx-c/mlx/c/vector.h index 81bcf7495..30d1b1838 100644 --- a/Source/Cmlx/mlx-c/mlx/c/vector.h +++ b/Source/Cmlx/mlx-c/mlx/c/vector.h @@ -7,6 +7,7 @@ #define MLX_VECTOR_H #include "mlx/c/array.h" +#include "mlx/c/stream.h" #include "mlx/c/string.h" #ifdef __cplusplus @@ -124,6 +125,35 @@ int mlx_vector_string_append_value(mlx_vector_string vec, const char* val); size_t mlx_vector_string_size(mlx_vector_string vec); int mlx_vector_string_get(char** res, const mlx_vector_string vec, size_t idx); +/** + * A vector of stream. + */ +typedef struct mlx_vector_stream_ { + void* ctx; +} mlx_vector_stream; +mlx_vector_stream mlx_vector_stream_new(void); +int mlx_vector_stream_set(mlx_vector_stream* vec, const mlx_vector_stream src); +int mlx_vector_stream_free(mlx_vector_stream vec); +mlx_vector_stream mlx_vector_stream_new_data( + const mlx_stream* data, + size_t size); +mlx_vector_stream mlx_vector_stream_new_value(const mlx_stream val); +int mlx_vector_stream_set_data( + mlx_vector_stream* vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_set_value(mlx_vector_stream* vec, const mlx_stream val); +int mlx_vector_stream_append_data( + mlx_vector_stream vec, + const mlx_stream* data, + size_t size); +int mlx_vector_stream_append_value(mlx_vector_stream vec, const mlx_stream val); +size_t mlx_vector_stream_size(mlx_vector_stream vec); +int mlx_vector_stream_get( + mlx_stream* res, + const mlx_vector_stream vec, + size_t idx); + /**@}*/ #ifdef __cplusplus diff --git a/Source/Cmlx/mlx-c/python/c.py b/Source/Cmlx/mlx-c/python/c.py index 079dfa3a5..ece443b23 100644 --- a/Source/Cmlx/mlx-c/python/c.py +++ b/Source/Cmlx/mlx-c/python/c.py @@ -9,7 +9,9 @@ def to_snake_letters(name): - name = re.sub(r"(?(std::vector)>", @@ -55,6 +60,11 @@ "std::function, std::vector>(std::vector,std::vector)>", "std::function, std::vector>(std::vector,std::vector)>", ], + [ + "mlx_compile_cache", + "std::weak_ptr", + "CompileCacheWeakPtr", + ], ]: if len(t) == 2: ctype, cpptype = t @@ -411,16 +421,32 @@ def register_return_tuple_type(cpp_types, alts=[]): } ) +types.append( + { + "cpp": "mlx::core::fft::FFTNorm", + "alt": "FFTNorm", + "c": "mlx_fft_norm", + "c_to_cpp": lambda s: "mlx_fft_norm_to_cpp(" + s + ")", + "c_arg": lambda s, untyped=False: s if untyped else "mlx_fft_norm " + s, + "c_return_arg": lambda s, untyped=False: ( + s if untyped else "mlx_fft_norm* " + s + ), + "c_new": lambda s: "mlx_fft_norm " + s, + "free": lambda s: "", + "c_assign_from_cpp": lambda d, s: d + " = " + "mlx_fft_norm_to_c(" + s + ")", + } +) + types.append( { "cpp": "std::string", "alt": "std::string", "c_to_cpp": lambda s: "std::string(" + s + ")", "c_arg": lambda s, untyped=False: s if untyped else "const char* " + s, - "c_return_arg": lambda s, untyped=False: s if untyped else "char** " + s, + "c_return_arg": lambda s, untyped=False: s if untyped else "mlx_string* " + s, # "c_new": lambda s: "char* " + s, # "free": lambda s: "", - "c_assign_from_cpp": lambda d, s: d + " = " + s + ".c_str()", + "c_assign_from_cpp": lambda d, s: "mlx_string_set_(*" + d + ", " + s + ")", } ) @@ -440,6 +466,16 @@ def register_return_tuple_type(cpp_types, alts=[]): } ) +types.append( + { + "cpp": "std::ostream", + "c_to_cpp": lambda s: "CFileOutputStream::as_lvalue(CFileOutputStream(" + + s + + "))", + "c_arg": lambda s, untyped=False: s if untyped else "FILE* " + s, + } +) + for ctype in ["int", "size_t", "float", "double", "bool", "uint64_t", "uintptr_t"]: types.append( { diff --git a/Source/Cmlx/mlx-c/python/mlxvariants.py b/Source/Cmlx/mlx-c/python/mlxvariants.py index 561e05f64..a1d5ed15a 100644 --- a/Source/Cmlx/mlx-c/python/mlxvariants.py +++ b/Source/Cmlx/mlx-c/python/mlxvariants.py @@ -103,6 +103,19 @@ def mlx_core(name, defs): "round": ["", None], "trace": ["", None, None], "export_function": [None, "", "kwargs"], + "zeros_like": ["dtype", ""], + "ones_like": ["dtype", ""], + "count_nonzero": ["axis", "axes", ""], + "flip": ["axes", "axis", ""], + "unstack": ["axis", ""], + "median": ["axes", "axis", "", None], + "logcumsumexp": ["axis", ""], + "cumsum": ["axis", "", None, None], + "cumprod": ["axis", "", None, None], + "cummax": ["axis", ""], + "cummin": ["axis", ""], + "trace": ["axes", None, ""], + "linspace": ["endpoint", ""], } return _make_variant_suffixes(name, defs, variants) @@ -133,6 +146,7 @@ def mlx_core_fast(name, defs): def mlx_core_detail(name, defs): if name not in [ "compile", + "compile_cache", "compile_clear_cache", "compile_erase", "vmap_replace", diff --git a/Source/Cmlx/mlx-c/python/type_private_generator.py b/Source/Cmlx/mlx-c/python/type_private_generator.py index 5f271b145..6df496c12 100644 --- a/Source/Cmlx/mlx-c/python/type_private_generator.py +++ b/Source/Cmlx/mlx-c/python/type_private_generator.py @@ -90,6 +90,7 @@ def generate( parser.add_argument("--cpptype", type=str) parser.add_argument("--no-copy", default=False, action="store_true") parser.add_argument("--include", default="", type=str) + parser.add_argument("--mlx-include", default="mlx/mlx.h", type=str) parser.add_argument("--using", default="", type=str) args = parser.parse_args() @@ -106,7 +107,7 @@ def generate( print("#define MLX_" + short_ctype.upper() + "_PRIVATE_H") print() print('#include "mlx/c/' + short_ctype + '.h"') - print('#include "mlx/mlx.h"') + print('#include "' + args.mlx_include + '"') ctypes = args.ctype.split(";") cpptypes = args.cpptype.split(";") usings = args.using.split(";") diff --git a/Source/Cmlx/mlx-c/python/vector_generator.py b/Source/Cmlx/mlx-c/python/vector_generator.py index 2c2bf4294..8462ff16e 100644 --- a/Source/Cmlx/mlx-c/python/vector_generator.py +++ b/Source/Cmlx/mlx-c/python/vector_generator.py @@ -217,6 +217,7 @@ def c_assign_wrap(s): #define MLX_VECTOR_H #include "mlx/c/array.h" +#include "mlx/c/stream.h" #include "mlx/c/string.h" #ifdef __cplusplus @@ -335,4 +336,15 @@ def c_assign_wrap(s): lambda d, s: "*" + d + " = " + s + ".data()", ) ) +print( + generate( + code, + "mlx::core::Stream", + "const mlx_stream", + "stream", + "mlx_stream*", + lambda s: "mlx_stream_get_(" + s + ")", + lambda d, s: "mlx_stream_set_(*" + d + ", " + s + ")", + ) +) print(end) 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 771b8e938..2c8157725 100644 --- a/Source/Cmlx/mlx-generated/binary_ops.cpp +++ b/Source/Cmlx/mlx-generated/binary_ops.cpp @@ -28,20 +28,27 @@ struct Add { struct FloorDivide { template - T operator()(T x, T y) thread { + 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) thread { - 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) thread { - 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) thread { - 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; } }; 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 85f18ed19..25c6d5433 100644 --- a/Source/Cmlx/mlx-generated/conv.cpp +++ b/Source/Cmlx/mlx-generated/conv.cpp @@ -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; @@ -1626,10 +1626,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -1829,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); } 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 new file mode 100644 index 000000000..3ed14c37c --- /dev/null +++ b/Source/Cmlx/mlx-generated/cuda/cuda_jit_sources.h @@ -0,0 +1,8366 @@ +#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, 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, + 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_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, + 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_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, + 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_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[] = { + 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_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_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, 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[] = { + 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, 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, 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/Cmlx/mlx-generated/fft.cpp b/Source/Cmlx/mlx-generated/fft.cpp index 0c7e08968..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,7 +792,7 @@ 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_, @@ -418,21 +815,35 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const thread { - 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 thread { - 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 thread { - 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 thread { - 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 thread { @@ -441,83 +852,96 @@ struct ReadWriter { return grid_index >= batch_size; } - METAL_FUNC void load() const thread { + 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 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) + 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) + 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); } } } @@ -544,53 +968,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step + // Four-step FFT I/O METAL_FUNC void load_strided(int stride, int overall_n) thread { - 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]); + 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) thread { - 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); + 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) thread { - // 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) thread { - 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 @@ -655,7 +1098,7 @@ METAL_FUNC void ReadWriter::write() const thread { 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); } } } @@ -680,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; } @@ -713,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); } } } @@ -771,10 +1214,10 @@ METAL_FUNC void ReadWriter::load() const thread { 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; } } @@ -832,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); @@ -865,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; } @@ -1001,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 @@ -1030,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. @@ -1060,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); } } @@ -1081,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. @@ -1111,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); } } @@ -1131,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_); @@ -1169,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, @@ -1193,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(); } @@ -1204,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, @@ -1212,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 @@ -1235,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, @@ -1260,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_; @@ -1272,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); @@ -1294,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)]); } @@ -1315,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); @@ -1350,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(); } @@ -1359,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: @@ -1375,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, @@ -1399,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); } @@ -1433,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; @@ -1442,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( @@ -1465,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 7bf24dfd5..1c56cdb56 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized.cpp @@ -12,6 +12,9 @@ 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) thread { if (metal::isnan(x)) { @@ -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 @@ -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)); } } @@ -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 03c0558db..6b8473d85 100644 --- a/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp +++ b/Source/Cmlx/mlx-generated/fp_quantized_nax.cpp @@ -12,6 +12,9 @@ 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) thread { if (metal::isnan(x)) { @@ -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" /////////////////////////////////////////////////////////////////////////////// @@ -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)); } } @@ -327,7 +342,7 @@ struct QuantizedBlockLoader { 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 0aa7e992c..ee4e24cc4 100644 --- a/Source/Cmlx/mlx-generated/gemm.cpp +++ b/Source/Cmlx/mlx-generated/gemm.cpp @@ -370,10 +370,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -573,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); } diff --git a/Source/Cmlx/mlx-generated/gemm_nax.cpp b/Source/Cmlx/mlx-generated/gemm_nax.cpp index b9069028e..290dea4ca 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) 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 "/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" /////////////////////////////////////////////////////////////////////////////// @@ -235,10 +106,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -423,10 +290,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 @@ -460,12 +329,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++) { @@ -480,7 +352,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); + } } } } @@ -504,13 +379,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 { @@ -537,10 +416,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 @@ -576,12 +457,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++) { @@ -619,14 +503,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]); } @@ -712,37 +600,185 @@ 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; + + typedef typename NAXFrag_t::template dtype_frag_t frag_type; + + frag_type val_frags[kNumFrags]; // = {frag_type(0)}; - frag_type val_frags[kNumFrags]; + METAL_FUNC NAXTile() thread {} METAL_FUNC constexpr void clear() thread { STEEL_PRAGMA_UNROLL @@ -751,595 +787,237 @@ struct NAXSubTile { } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) thread { - 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 thread { - return val_frags[i * kSubTileCols + j]; + return val_frags[i * kTileCols + j]; } template METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; + return val_frags[i * kTileCols + j]; } template METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; + 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 thread T* elems() thread { - return reinterpret_cast(val_frags); + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); } - METAL_FUNC const thread T* elems() const 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 thread { + 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) thread { + 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 = {}) thread { - 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 thread { - 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 = {}) thread { - 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 = {}) thread { - 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 thread { - 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 thread { - 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 thread { - 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 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(), - dst, - str_x, - str_y, - start_x, - stop_x, - start_y, - stop_y, - off_x + idx_row * Int{}, - off_y + idx_col * Int{}); + src, + ld, + Int<1>{}, + idx_row * Int{}, + 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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - 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() thread { - 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) thread { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( - const short i, - const short j) const thread { - return val_subtiles[i * kTileCols + j]; - } - - template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC thread elem_type* elems() thread { - return reinterpret_cast(val_subtiles[0].elems()); - } - - METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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 thread { - 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) thread { - 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 thread { - 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>{}); - } - } + 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 { - 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)], + 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) thread { - 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 thread { - 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_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 thread { - 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 @@ -1350,7 +1028,8 @@ struct NAXTile { 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>{}, @@ -1358,8 +1037,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1379,51 +1058,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 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{}); + } } } } @@ -1502,6 +1184,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" /////////////////////////////////////////////////////////////////////////////// @@ -1526,9 +1331,6 @@ template < bool kAlignedM, bool kAlignedN, bool kAlignedK, - short UM, - short UN, - short UK, typename AccumType = float> auto gemm_loop( const device T* A, @@ -1539,9 +1341,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; @@ -1549,25 +1351,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; @@ -1607,51 +1409,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/metal/atomic.h b/Source/Cmlx/mlx-generated/metal/atomic.h index dacfa939a..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; + } + } } } diff --git a/Source/Cmlx/mlx-generated/metal/binary_ops.h b/Source/Cmlx/mlx-generated/metal/binary_ops.h index 863d6369e..37650c7d9 100644 --- a/Source/Cmlx/mlx-generated/metal/binary_ops.h +++ b/Source/Cmlx/mlx-generated/metal/binary_ops.h @@ -16,20 +16,27 @@ struct Add { struct FloorDivide { template - T operator()(T x, T y) thread { + 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) thread { - 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) thread { - 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) thread { - 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; } }; diff --git a/Source/Cmlx/mlx-generated/metal/complex.h b/Source/Cmlx/mlx-generated/metal/complex.h index 5f654332b..629b6289b 100644 --- a/Source/Cmlx/mlx-generated/metal/complex.h +++ b/Source/Cmlx/mlx-generated/metal/complex.h @@ -4,164 +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) thread : real(real), - imag(imag) {}; - constexpr complex64_t() thread : 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)) { @@ -172,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 4730253ea..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,7 +85,7 @@ 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_, @@ -73,21 +108,35 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const thread { - 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 thread { - 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 thread { - 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 thread { - 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 thread { @@ -96,83 +145,96 @@ struct ReadWriter { return grid_index >= batch_size; } - METAL_FUNC void load() const thread { + 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 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) + 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) + 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); } } } @@ -199,53 +261,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step + // Four-step FFT I/O METAL_FUNC void load_strided(int stride, int overall_n) thread { - 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]); + 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) thread { - 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); + 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) thread { - // 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) thread { - 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 @@ -310,7 +391,7 @@ METAL_FUNC void ReadWriter::write() const thread { 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); } } } @@ -335,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; } @@ -368,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); } } } @@ -426,10 +507,10 @@ METAL_FUNC void ReadWriter::load() const thread { 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; } } @@ -487,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); @@ -520,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; } diff --git a/Source/Cmlx/mlx-generated/metal/fp4.h b/Source/Cmlx/mlx-generated/metal/fp4.h index 47bf4dda6..af4cbb02c 100644 --- a/Source/Cmlx/mlx-generated/metal/fp4.h +++ b/Source/Cmlx/mlx-generated/metal/fp4.h @@ -1,5 +1,8 @@ #pragma once +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { fp4_e2m1(float x) thread { if (metal::isnan(x)) { diff --git a/Source/Cmlx/mlx-generated/metal/fp8.h b/Source/Cmlx/mlx-generated/metal/fp8.h index 796dd2163..42c5ec128 100644 --- a/Source/Cmlx/mlx-generated/metal/fp8.h +++ b/Source/Cmlx/mlx-generated/metal/fp8.h @@ -78,3 +78,14 @@ 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); +} diff --git a/Source/Cmlx/mlx-generated/metal/fp_quantized.h b/Source/Cmlx/mlx-generated/metal/fp_quantized.h index e2bc0cea7..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 @@ -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)); } } @@ -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 0c25eec5d..35c61c2fa 100644 --- a/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h +++ b/Source/Cmlx/mlx-generated/metal/fp_quantized_nax.h @@ -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)); } } @@ -178,7 +179,7 @@ struct QuantizedBlockLoader { 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/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/quantized.h b/Source/Cmlx/mlx-generated/metal/quantized.h index d60bc882e..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); } } } @@ -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 22238112d..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); } } } @@ -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 b7a9cacb3..f875e51c9 100644 --- a/Source/Cmlx/mlx-generated/metal/reduction/ops.h +++ b/Source/Cmlx/mlx-generated/metal/reduction/ops.h @@ -133,6 +133,15 @@ struct Sum { 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) thread { return a + b; 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 740fd0cc7..9907171b0 100644 --- a/Source/Cmlx/mlx-generated/metal/scan.h +++ b/Source/Cmlx/mlx-generated/metal/scan.h @@ -99,15 +99,24 @@ template struct CumMax { static constexpr constant U init = Limits::min; + 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; + } + template U operator()(U a, T b) thread { - return (a >= b) ? a : b; + 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; } @@ -122,15 +131,24 @@ template struct CumMin { static constexpr constant U init = Limits::max; + 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; + } + template U operator()(U a, T b) thread { - return (a <= b) ? a : b; + 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; } diff --git a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h index 1eec72be3..2f466a567 100644 --- a/Source/Cmlx/mlx-generated/metal/sdpa_vector.h +++ b/Source/Cmlx/mlx-generated/metal/sdpa_vector.h @@ -4,6 +4,161 @@ using namespace metal; +// ---------------------------------------------------------------------------- +// TurboQuant decompression for Metal — ported from TheTom/llama-cpp-turboquant +// Sources: turbo-wht.h, ggml-turbo-quant.c (feature/turboquant-kv-cache branch) +// Paper: Zandieh et al., TurboQuant, AISTATS/ICLR 2026 +// +// Decompression path (used during SDPA attention): +// 1. Look up 3-bit index → centroid value +// 2. Collect all d=128 centroid values into local register array +// 3. Apply inverse WHT rotation: D2 * FWHT * D1 +// 4. Scale by stored corrected norm + +// --------------------------------------------------------------------------- +// WHT sign arrays — seed=42, must match CPU turbo_quant.h exactly +// --------------------------------------------------------------------------- +constant float turbo_wht_signs1[128] = { + -1,1,1,-1,-1,1,-1,1,-1,-1,1,1,1,1,1,1,1,-1,1,-1,1,-1,-1,1,1,1,-1,1,1,-1,-1,-1, + -1,1,1,-1,1,1,-1,1,-1,1,1,-1,-1,1,-1,1,1,1,1,-1,-1,-1,-1,-1,1,-1,1,1,1,1,-1,1, + -1,-1,1,-1,-1,-1,1,-1,-1,-1,1,-1,-1,-1,1,1,1,-1,-1,1,1,1,-1,-1,1,1,-1,1,1,-1,1,-1, + -1,1,1,-1,1,-1,1,-1,1,1,1,1,-1,1,-1,1,1,-1,1,1,-1,-1,-1,-1,-1,1,1,-1,1,1,-1,1}; +constant float turbo_wht_signs2[128] = { + 1,1,1,1,-1,1,1,-1,1,-1,-1,-1,1,-1,-1,-1,1,1,-1,-1,1,-1,1,-1,1,-1,-1,1,-1,1,1,1, + 1,1,-1,-1,-1,1,-1,-1,-1,-1,-1,-1,1,1,1,-1,1,-1,1,1,1,-1,-1,1,-1,-1,-1,-1,-1,-1,1,1, + 1,-1,1,-1,-1,-1,-1,1,-1,1,-1,1,-1,-1,1,1,-1,1,-1,1,1,-1,1,-1,-1,-1,-1,1,-1,-1,1,-1, + 1,-1,1,1,1,-1,-1,1,-1,1,-1,1,1,-1,-1,1,-1,1,-1,1,1,-1,1,-1,1,-1,-1,-1,-1,-1,1,-1}; + +// QJL sign arrays — seed=1042 +constant float turbo_qjl_signs1[128] = { + 1,-1,-1,-1,-1,1,-1,1,1,-1,-1,1,-1,1,-1,1,1,-1,1,-1,-1,-1,1,1,-1,1,1,-1,1,-1,-1,1, + 1,1,1,1,-1,-1,1,1,-1,1,-1,-1,1,-1,1,1,1,-1,1,1,1,-1,-1,1,-1,1,-1,1,1,-1,1,1, + -1,-1,-1,1,1,1,1,1,1,-1,-1,1,1,-1,-1,-1,-1,-1,1,1,1,1,-1,1,1,-1,1,1,1,1,1,1, + 1,-1,1,-1,-1,1,-1,-1,-1,-1,1,-1,1,1,1,-1,-1,1,-1,1,1,1,-1,-1,1,-1,-1,-1,-1,-1,-1,-1}; +constant float turbo_qjl_signs2[128] = { + 1,1,-1,1,1,-1,1,1,-1,-1,1,1,1,-1,1,1,-1,-1,-1,1,-1,1,1,1,-1,1,-1,-1,-1,-1,1,1, + -1,-1,1,-1,1,1,-1,-1,-1,-1,-1,1,1,1,1,1,1,1,1,1,-1,-1,1,1,1,1,1,1,1,-1,1,1, + -1,-1,1,-1,1,1,-1,1,-1,-1,1,1,1,-1,1,-1,1,1,1,1,1,1,-1,1,-1,1,-1,1,-1,1,1,-1, + 1,-1,-1,1,1,-1,1,1,-1,1,1,1,-1,1,1,1,-1,-1,1,-1,1,-1,-1,1,-1,1,-1,1,1,1,1,-1}; + +// --------------------------------------------------------------------------- +// 3-bit Lloyd-Max centroids for N(0, 1/128) — matches ggml-turbo-quant.c +// --------------------------------------------------------------------------- +constant float turbo_centroids_3bit[8] = { + -0.190685f, -0.117832f, -0.065717f, -0.021460f, + 0.021460f, 0.065717f, 0.117832f, 0.190685f +}; + +// --------------------------------------------------------------------------- +// Fast Walsh-Hadamard Transform (in-place, n=128, normalized by 1/sqrt(128)) +// --------------------------------------------------------------------------- +static void turbo_fwht_128(thread float * x) { + for (int h = 1; h < 128; h *= 2) { + for (int i = 0; i < 128; i += h * 2) { + for (int j = i; j < i + h; j++) { + float a = x[j], b = x[j + h]; + x[j] = a + b; + x[j + h] = a - b; + } + } + } + const float inv_sqrt_128 = 0.08838834764831845f; + for (int i = 0; i < 128; i++) x[i] *= inv_sqrt_128; +} + +// Inverse WHT rotation: D2 * FWHT * D1 +static void turbo_rotate_inverse(thread float * x, + constant float * s1, + constant float * s2) { + for (int i = 0; i < 128; i++) x[i] *= s2[i]; + turbo_fwht_128(x); + for (int i = 0; i < 128; i++) x[i] *= s1[i]; +} + +// QJL inverse rotation for residual (same WHT, different seeds) +static void turbo_qjl_rotate(thread float * x, + constant float * qs1, + constant float * qs2) { + for (int i = 0; i < 128; i++) x[i] *= qs2[i]; + turbo_fwht_128(x); + for (int i = 0; i < 128; i++) x[i] *= qs1[i]; +} + +// --------------------------------------------------------------------------- +// Bit-unpack helpers +// --------------------------------------------------------------------------- + +// Extract 3-bit index for coordinate i from packed byte array +static inline uchar unpack_3bit(const device uchar * packed, int i) { + int bit_offset = i * 3; + int byte_idx = bit_offset / 8; + int bit_pos = bit_offset % 8; + uint raw = (uint)packed[byte_idx]; + if (byte_idx + 1 < 48) raw |= ((uint)packed[byte_idx + 1]) << 8; + return (uchar)((raw >> bit_pos) & 0x7); +} + +// Extract QJL sign bit for coordinate i from packed byte array +static inline float unpack_sign_bit(const device uchar * signs, int i) { + return (signs[i / 8] & (1u << (i % 8))) ? 1.0f : -1.0f; +} + +constant float TURBO_QJL_CONST = 1.2533141373155003f; // sqrt(pi/2) + +// --------------------------------------------------------------------------- +// Dequantize a full TurboQuantK vector (3-bit PolarQuant + 1-bit QJL) +// Used for K-cache during SDPA dot product. +// +// packed_indices[48]: 3 bits × 128 coords +// qjl_signs[16]: 1 bit × 128 coords +// norm: corrected L2 scale +// rnorm: residual L2 scale +// out[128]: reconstructed float vector +// --------------------------------------------------------------------------- +static void turbo_dequant_k( + const device uchar * packed_indices, + const device uchar * qjl_signs, + float norm, + float rnorm, + thread float * out) +{ + // Stage 1: PolarQuant centroid lookup + thread float mse[128]; + for (int i = 0; i < 128; i++) { + mse[i] = turbo_centroids_3bit[unpack_3bit(packed_indices, i)]; + } + // Inverse rotation to get back to original space + turbo_rotate_inverse(mse, turbo_wht_signs1, turbo_wht_signs2); + + // Stage 2: QJL residual + thread float signs_f[128]; + for (int i = 0; i < 128; i++) signs_f[i] = unpack_sign_bit(qjl_signs, i); + // Inverse QJL rotation + turbo_qjl_rotate(signs_f, turbo_qjl_signs1, turbo_qjl_signs2); + const float qjl_scale = TURBO_QJL_CONST / 128.0f * rnorm; + + // Combine and scale + for (int i = 0; i < 128; i++) { + out[i] = (mse[i] + signs_f[i] * qjl_scale) * norm; + } +} + +// --------------------------------------------------------------------------- +// Dequantize a TurboQuantV vector (3-bit PolarQuant only — V-cache) +// --------------------------------------------------------------------------- +static void turbo_dequant_v( + const device uchar * packed_indices, + float norm, + thread float * out) +{ + thread float buf[128]; + for (int i = 0; i < 128; i++) { + buf[i] = turbo_centroids_3bit[unpack_3bit(packed_indices, i)]; + } + turbo_rotate_inverse(buf, turbo_wht_signs1, turbo_wht_signs2); + for (int i = 0; i < 128; i++) out[i] = buf[i] * norm; +} + constant bool has_mask [[function_constant(20)]]; constant bool query_transposed [[function_constant(21)]]; constant bool do_causal [[function_constant(22)]]; @@ -317,6 +472,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 7662707af..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 @@ -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/nax.h b/Source/Cmlx/mlx-generated/metal/steel/attn/nax.h index 0997eebd2..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]; + 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>(); - METAL_FUNC constexpr void clear() thread { + // 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) thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; - } - METAL_FUNC thread T* elems() thread { - return reinterpret_cast(val_frags); - } - - METAL_FUNC const thread T* elems() const thread { - return reinterpret_cast(val_frags); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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 { - 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 = {}) thread { - 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 thread { + // 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 = {}) thread { - 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)); - } - } - } + 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>(); - 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 = {}) thread { + // 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_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_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[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 thread { + // 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::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_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 thread { + // 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 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 thread { - 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{}); - }); - }); - } -}; + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); -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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - CType>(); - - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { + // Copy out results 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]; - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + 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 +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 kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - NAXSubTile_t val_subtiles[kSubTiles]; + 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 < 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) thread { - 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 thread { - return val_subtiles[i * kTileCols + j]; + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + 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); + } + } + + 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_subtiles[0].elems()); + return reinterpret_cast(val_frags); } METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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( + 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 thread { - 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( + 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) thread { - 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>{}); - } - } + 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 { - 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>{}); - } - } + 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) thread { - 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) thread { - 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 thread { - 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 thread { - 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 @@ -993,7 +811,8 @@ struct NAXTile { 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 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 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{}); + } } } } 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 81faa53df..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 @@ -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/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_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/mma.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/mma.h index 49a4d44cf..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); } diff --git a/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h b/Source/Cmlx/mlx-generated/metal/steel/gemm/nax.h index b2d877098..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]; + 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>(); - METAL_FUNC constexpr void clear() thread { + // 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) thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC thread T* elems() thread { - return reinterpret_cast(val_frags); - } - - METAL_FUNC const thread T* elems() const thread { - return reinterpret_cast(val_frags); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { + // 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) thread { + // 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 = {}) thread { - 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 thread { + // 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 = {}) thread { - 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)); - } - } - } + 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>(); - 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 = {}) thread { + // 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_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_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[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 thread { + // 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::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_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 thread { + // 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 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 thread { - 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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - 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]; - } - } - } + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); - // Load B into right operand registers - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { + // Copy out results 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]; - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + 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 +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 kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - NAXSubTile_t val_subtiles[kSubTiles]; + 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 < 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) thread { - 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 thread { - return val_subtiles[i * kTileCols + j]; + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + 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); + } + } + + 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_subtiles[0].elems()); + return reinterpret_cast(val_frags); } METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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( + 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 thread { - 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( + 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) thread { - 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>{}); - } - } + 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 { - 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>{}); - } - } + 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 { - 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)], + 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) thread { - 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 thread { - 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_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 thread { - 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 @@ -1001,7 +811,8 @@ struct NAXTile { 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/utils/integral_constant.h b/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h index 60f1591f7..f85fca32c 100644 --- a/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h +++ b/Source/Cmlx/mlx-generated/metal/steel/utils/integral_constant.h @@ -23,10 +23,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -132,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/utils.h b/Source/Cmlx/mlx-generated/metal/utils.h index 0faf26386..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()); }; /////////////////////////////////////////////////////////////////////////////// @@ -446,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 6b321c8e6..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); } } } @@ -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 07d60c19f..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); } } } @@ -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 0c921a3af..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; + } + } } } @@ -495,6 +508,15 @@ struct Sum { 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) thread { return a + b; diff --git a/Source/Cmlx/mlx-generated/scan.cpp b/Source/Cmlx/mlx-generated/scan.cpp index 7e9d7bae6..1fd1d74dc 100644 --- a/Source/Cmlx/mlx-generated/scan.cpp +++ b/Source/Cmlx/mlx-generated/scan.cpp @@ -28,20 +28,27 @@ struct Add { struct FloorDivide { template - T operator()(T x, T y) thread { + 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) thread { - 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) thread { - 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) thread { - 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; } }; @@ -447,15 +454,24 @@ template struct CumMax { static constexpr constant U init = Limits::min; + 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; + } + template U operator()(U a, T b) thread { - return (a >= b) ? a : b; + 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; } @@ -470,15 +486,24 @@ template struct CumMin { static constexpr constant U init = Limits::max; + 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; + } + template U operator()(U a, T b) thread { - return (a <= b) ? a : b; + 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; } 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 95b4039d6..75c2edbdd 100644 --- a/Source/Cmlx/mlx-generated/steel_attention.cpp +++ b/Source/Cmlx/mlx-generated/steel_attention.cpp @@ -496,10 +496,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -2019,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 @@ -2083,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; @@ -2208,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); } @@ -2218,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 1784524aa..df0831701 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" /////////////////////////////////////////////////////////////////////////////// @@ -112,10 +106,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -299,11 +289,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 @@ -336,13 +328,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++) { @@ -357,7 +352,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); + } } } } @@ -380,14 +378,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 { @@ -413,11 +415,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 @@ -452,13 +456,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++) { @@ -495,15 +502,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]); } @@ -535,7 +546,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{}; @@ -589,38 +600,185 @@ 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() thread { STEEL_PRAGMA_UNROLL @@ -629,41 +787,85 @@ struct NAXSubTile { } } - METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) thread { - 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 thread { - return val_frags[i * kSubTileCols + j]; + return val_frags[i * kTileCols + j]; } template METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; + return val_frags[i * kTileCols + j]; } template METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; + 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 thread T* elems() thread { - return reinterpret_cast(val_frags); + METAL_FUNC thread elem_type* elems() thread { + return reinterpret_cast(val_frags); } - METAL_FUNC const thread T* elems() const 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 thread { - 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]); } @@ -671,545 +873,151 @@ struct NAXSubTile { } template - METAL_FUNC void row_bin_op(thread metal::vec& vals) thread { - 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 = {}) thread { - 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 thread { - 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 = {}) thread { - 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 = {}) thread { - 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 thread { - 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), + 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, - lim_x, - off_x + (i * kFragRows), - off_y + (j * kFragCols)); - } - } + ld, + Int<1>{}, + n_rows, + idx_row * Int{}, + idx_col * Int{}); + }); + }); } - 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 thread { - STEEL_PRAGMA_UNROLL - for (int i = 0; i < kSubTileRows; ++i) { - STEEL_PRAGMA_UNROLL - for (int j = 0; j < kSubTileCols; ++j) { + template + 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(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 thread { - 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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - 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() thread { - 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) thread { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at( - const short i, - const short j) const thread { - return val_subtiles[i * kTileCols + j]; - } - - template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + j]; - } - - METAL_FUNC thread elem_type* elems() thread { - return reinterpret_cast(val_subtiles[0].elems()); - } - - METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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 thread { - 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) thread { - 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 thread { - 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) thread { - 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) thread { - 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); - } - } - } - - template - 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 - for (int j = 0; j < kTileCols; ++j) { - subtile_at(i, j).store_safe( 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 thread { - 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 @@ -1220,7 +1028,8 @@ struct NAXTile { 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>{}, @@ -1228,8 +1037,8 @@ struct NAXTile { stop.y, start.x, stop.x, - idx_row * Int{}, - idx_col * Int{}); + idx_row * Int{}, + idx_col * Int{}); }); }); } @@ -1249,51 +1058,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{}); + } } } } @@ -1601,38 +1413,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}; @@ -1652,83 +1462,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{}); } } @@ -1736,22 +1539,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; @@ -1760,16 +1551,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; } } } @@ -1777,7 +1568,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; @@ -1787,18 +1578,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]; } } @@ -1815,33 +1604,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]); + } } } } @@ -1859,10 +1680,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 @@ -1877,7 +1698,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); @@ -1888,39 +1709,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{}); } } } @@ -1943,14 +1761,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 97863594d..0024a77f5 100644 --- a/Source/Cmlx/mlx-generated/steel_conv_general.cpp +++ b/Source/Cmlx/mlx-generated/steel_conv_general.cpp @@ -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_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/unary_ops.cpp b/Source/Cmlx/mlx-generated/unary_ops.cpp index c78a9f088..3a680b7f2 100644 --- a/Source/Cmlx/mlx-generated/unary_ops.cpp +++ b/Source/Cmlx/mlx-generated/unary_ops.cpp @@ -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" /////////////////////////////////////////////////////////////////////////////// diff --git a/Source/Cmlx/mlx-generated/utils.cpp b/Source/Cmlx/mlx-generated/utils.cpp index dd35b66aa..18d9644d0 100644 --- a/Source/Cmlx/mlx-generated/utils.cpp +++ b/Source/Cmlx/mlx-generated/utils.cpp @@ -422,164 +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) thread : real(real), - imag(imag) {}; - constexpr complex64_t() thread : 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)) { @@ -591,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" /////////////////////////////////////////////////////////////////////////////// @@ -727,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()); }; /////////////////////////////////////////////////////////////////////////////// @@ -1099,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/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md index 22a085792..98620ba94 100644 --- a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md +++ b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/bug_report.md @@ -1,12 +1,14 @@ --- name: Bug report -about: Create a report about an issue you've encountered +about: Create a report about a bug you've encountered title: "[BUG] " labels: '' assignees: '' --- +☑️ I understand it is strictly prohibited to use AI to write issues. + **Describe the bug** A clear and concise description of what the bug is. diff --git a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 000000000..3ba13e0ce --- /dev/null +++ b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1 @@ +blank_issues_enabled: false diff --git a/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md new file mode 100644 index 000000000..eb410efbb --- /dev/null +++ b/Source/Cmlx/mlx/.github/ISSUE_TEMPLATE/other.md @@ -0,0 +1,10 @@ +--- +name: Other +about: Any other issue +title: '' +labels: '' +assignees: '' + +--- + +☑️ I understand it is strictly prohibited to use AI to write issues. diff --git a/Source/Cmlx/mlx/.github/actions/build-cuda-release/action.yml b/Source/Cmlx/mlx/.github/actions/build-cuda-release/action.yml deleted file mode 100644 index 0ed7719d5..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-cuda-release/action.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: 'Build CUDA wheel' -description: 'Build CUDA wheel' - -inputs: - arch: - description: 'Platform architecture tag' - required: true - type: choice - options: - - x86_64 - - aarch64 - -runs: - using: "composite" - steps: - - name: Build package - shell: bash - env: - CMAKE_ARGS: -DMLX_BUILD_CUDA=ON - run: | - pip install auditwheel build patchelf setuptools - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w - - auditwheel repair dist/mlx_cuda*.whl \ - --plat manylinux_2_35_${{ inputs.arch }} \ - --exclude libcublas* \ - --exclude libcuda* \ - --exclude libcudnn* \ - --exclude libnccl* \ - --exclude libnvrtc* diff --git a/Source/Cmlx/mlx/.github/actions/build-docs/action.yml b/Source/Cmlx/mlx/.github/actions/build-docs/action.yml index 411f6be8d..4d4286e3c 100644 --- a/Source/Cmlx/mlx/.github/actions/build-docs/action.yml +++ b/Source/Cmlx/mlx/.github/actions/build-docs/action.yml @@ -5,24 +5,28 @@ runs: using: "composite" steps: - name: Setup machine - uses: ./.github/actions/setup-linux + id: setup + uses: ./.github/actions/setup + with: + ccache-key: 'release' + ccache-save: false - name: Install dependencies shell: bash + env: + CMAKE_ARGS: ${{ steps.setup.outputs.cmake-args }} run: | sudo apt-get install -y doxygen - source .venv/bin/activate - pip install -r docs/requirements.txt - pip install . -v - + uv pip install -r docs/requirements.txt + uv pip install . -v + - name: Build documentation shell: bash run: | - source .venv/bin/activate cd docs doxygen make html O=-W - + - name: Create artifact tar shell: bash run: tar -cf artifact.tar -C docs --dereference build/html index.html @@ -30,7 +34,7 @@ runs: # Do it manually because upload-pages-artifact requires gtar - name: Upload artifact id: upload-artifact - uses: actions/upload-artifact@v5 + uses: actions/upload-artifact@v7 with: name: github-pages path: artifact.tar diff --git a/Source/Cmlx/mlx/.github/actions/build-linux-release/action.yml b/Source/Cmlx/mlx/.github/actions/build-linux-release/action.yml deleted file mode 100644 index c8dec00b4..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-linux-release/action.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: 'Build Linux wheel' -description: 'Build Linux wheel' - -inputs: - build-backend: - description: 'Build the backend mlx-cpu package' - type: boolean - required: false - default: false - arch: - description: 'Platform architecture tag' - required: true - type: choice - options: - - x86_64 - - aarch64 - -runs: - using: "composite" - steps: - - name: Build MLX - shell: bash - run: pip install -e . -v - - - name: Build Python package - shell: bash - run: | - pip install auditwheel patchelf build - python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w - auditwheel repair dist/mlx-*.whl \ - --plat manylinux_2_35_${{ inputs.arch }} \ - --exclude libmlx.so* \ - --only-plat - - - name: Build backend package - if: ${{ inputs.build-backend }} - shell: bash - run: | - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w - auditwheel repair dist/mlx_cpu*.whl --plat manylinux_2_35_${{ inputs.arch }} diff --git a/Source/Cmlx/mlx/.github/actions/build-linux/action.yml b/Source/Cmlx/mlx/.github/actions/build-linux/action.yml deleted file mode 100644 index 166d7b8a6..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-linux/action.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: 'Build and Test on Linux' - -inputs: - toolkit: - description: 'The toolkit to build with' - required: false - default: 'cpu' - -runs: - using: "composite" - steps: - - - name: Install Python package - id: python_build - shell: sh - env: - DEBUG: 1 - CMAKE_ARGS: >- - -DCMAKE_COMPILE_WARNING_AS_ERROR=ON - -DMLX_BUILD_CUDA=${{ startsWith(inputs.toolkit, 'cuda') && 'ON' || 'OFF' }} - run: | - if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then - # There is no GPU in arm64 runner, use a common arch. - CMAKE_ARGS="$CMAKE_ARGS -DMLX_CUDA_ARCHITECTURES=90a" - # Can not build tests and stubs when the built executables can not run. - CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF -DMLX_BUILD_PYTHON_STUBS=OFF" - fi - # Install cpu-only torch to save space - pip install torch --index-url https://download.pytorch.org/whl/cpu - pip install --no-build-isolation -e ".[dev]" -v - # Pass the CMAKE_ARGS to following steps. - echo CMAKE_ARGS="$CMAKE_ARGS" >> $GITHUB_OUTPUT - - - name: Build CPP only - shell: bash - run: | - cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ steps.python_build.outputs.CMAKE_ARGS }} - cmake --build build -j $(nproc) diff --git a/Source/Cmlx/mlx/.github/actions/build-macos-release/action.yml b/Source/Cmlx/mlx/.github/actions/build-macos-release/action.yml deleted file mode 100644 index 93e6dd943..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-macos-release/action.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: 'Build macOS release' -description: 'Build MLX releases macOS' - -inputs: - macos-target: - description: 'macOS build target' - required: false - default: '15.0' - build-backend: - description: 'Build the backend mlx-metal package' - type: boolean - required: false - default: false - -runs: - using: "composite" - steps: - - name: Build Python package - shell: bash -l {0} - env: - DEVELOPER_DIR: /Applications/Xcode-latest.app - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - pip install build - python setup.py clean --all - MLX_BUILD_STAGE=1 python -m build -w - - - name: Build backend package - if: ${{ inputs.build-backend }} - shell: bash -l {0} - env: - DEVELOPER_DIR: /Applications/Xcode-latest.app - MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} - run: | - python setup.py clean --all - MLX_BUILD_STAGE=2 python -m build -w diff --git a/Source/Cmlx/mlx/.github/actions/build-macos/action.yml b/Source/Cmlx/mlx/.github/actions/build-macos/action.yml index f04ed3d63..7bcd4e0d0 100644 --- a/Source/Cmlx/mlx/.github/actions/build-macos/action.yml +++ b/Source/Cmlx/mlx/.github/actions/build-macos/action.yml @@ -1,80 +1,39 @@ -name: 'Build and Test on macOS' -description: 'Build and test MLX on macOS' +name: 'Build macOS' +description: 'Build C++ and Python binaries for testing on macOS' + +inputs: + cmake-args: + description: 'The args for generating CMake project' + required: true + macos-target: + description: 'The target macOS version to build for' + required: false + default: '26.2' runs: - using: "composite" + using: 'composite' steps: - name: Install dependencies + shell: bash + run: uv pip install build setuptools + + - name: Build wheel + shell: bash env: DEBUG: 1 - CMAKE_ARGS: "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" - shell: bash -l {0} - run: | - pip install --upgrade pip - pip install cmake setuptools typing_extensions - pip install -e . -v + CMAKE_ARGS: ${{ inputs.cmake-args }} + MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} + run: python -m build -w - - name: Install tests dependencies - shell: bash -l {0} - run: | - pip install numpy torch tensorflow - - - name: Run Python tests - shell: bash -l {0} - env: - LOW_MEMORY: 1 - run: | - DEVICE=cpu python -m unittest discover -v python/tests - DEVICE=gpu METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 python -m unittest discover -v python/tests - mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib/ python python/tests/mpi_test_distributed.py - mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) - if $(grep "\[WARN\]" stderr.log); then echo "Distributed ring test failed"; exit 1; fi - - - name: Build example extension - shell: bash -l {0} - run: | - cd examples/extensions - pip install -r requirements.txt - python setup.py build_ext --inplace - python test.py - - name: Build CPP only - shell: bash -l {0} - run: | - mkdir -p build - cd build - cmake .. - make -j $(sysctl -n hw.ncpu) - - - name: Run CPP tests - shell: bash -l {0} - env: - DEVICE: gpu - METAL_DEVICE_WRAPPER_TYPE: 1 - METAL_DEBUG_ERROR_MODE: 0 - run: ./build/tests/tests - - - name: Build small binary with JIT - shell: bash -l {0} - run: | - mkdir -p build - cd build - cmake .. -DCMAKE_BUILD_TYPE=MinSizeRel \ - -DBUILD_SHARED_LIBS=ON \ - -DMLX_BUILD_CPU=OFF \ - -DMLX_BUILD_SAFETENSORS=OFF \ - -DMLX_BUILD_GGUF=OFF \ - -DMLX_METAL_JIT=ON - make -j $(sysctl -n hw.ncpu) - - - name: Run Python tests with JIT - shell: bash -l {0} + shell: bash env: - LOW_MEMORY: 1 - DEVICE: gpu - METAL_DEVICE_WRAPPER_TYPE: 1 - METAL_DEBUG_ERROR_MODE: 0 + MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} run: | - CMAKE_ARGS="-DMLX_METAL_JIT=ON" \ - pip install -e . -v - python -m unittest discover -v python/tests + if ${{ contains(inputs.cmake-args, 'CMAKE_BUILD_TYPE') }} ; then + cmake . -B build ${{ inputs.cmake-args }} + else + cmake . -B build ${{ inputs.cmake-args }} \ + -DCMAKE_BUILD_TYPE=Debug + fi + cmake --build build -j $(sysctl -n hw.physicalcpu) diff --git a/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml new file mode 100644 index 000000000..2e15efd93 --- /dev/null +++ b/Source/Cmlx/mlx/.github/actions/build-wheel/action.yml @@ -0,0 +1,105 @@ +name: 'Build wheel' +description: 'Build the Python wheels for release on all platforms' + +inputs: + cmake-args: + description: 'The args for generating CMake project' + required: true + build-frontend: + description: 'Build the frontend mlx package' + required: false + default: 'true' + build-backend: + description: 'Build the backend mlx-cpu/mlx-cuda/mlx-metal packages' + required: false + default: 'true' + macos-target: + description: 'The target macOS version to build for' + required: false + default: '26.2' + arch-tag: + description: 'Platform architecture tag' + required: false + default: |- + ${{ case(runner.arch == 'x64', 'x86_64', + runner.arch == 'x86', 'i686', + runner.arch == 'arm', 'armv7l', + runner.arch == 'arm64', 'aarch64', + 'unknown') + }} + +runs: + using: 'composite' + steps: + - name: Install dependencies + shell: bash + run: | + uv pip install build setuptools + if ${{ runner.os == 'Linux' }} ; then + uv pip install auditwheel patchelf + fi + mkdir -p wheelhouse + + - name: Build frontend package + if: inputs.build-frontend == 'true' + shell: bash + env: + CMAKE_ARGS: ${{ inputs.cmake-args }} + MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} + run: | + python setup.py clean --all + MLX_BUILD_FRONTEND_PACKAGE=1 python -m build -w + + - name: Post-process frontend package + if: inputs.build-frontend == 'true' + shell: bash + run: | + if ${{ runner.os == 'Linux' }} ; then + auditwheel repair dist/mlx-*.whl \ + --plat manylinux_2_35_${{ inputs.arch-tag }} \ + --exclude libmlx.so* \ + --only-plat + else + mv dist/mlx-*.whl wheelhouse/ + fi + + - name: Build backend package + if: inputs.build-backend == 'true' + shell: bash + env: + CMAKE_ARGS: ${{ inputs.cmake-args }} + MACOSX_DEPLOYMENT_TARGET: ${{ inputs.macos-target }} + run: | + python setup.py clean --all + MLX_BUILD_BACKEND_PACKAGE=1 python -m build -w + + - name: Post-process backend package + if: inputs.build-backend == 'true' + shell: bash + run: | + if ${{ runner.os == 'Linux' }} ; then + if [ -f dist/mlx_cpu*.whl ]; then + auditwheel repair dist/mlx_cpu*.whl \ + --plat manylinux_2_35_${{ inputs.arch-tag }} + fi + if [ -f dist/mlx_cuda*.whl ]; then + auditwheel repair dist/mlx_cuda*.whl \ + --plat manylinux_2_35_${{ inputs.arch-tag }} \ + --exclude libcublas* \ + --exclude libcuda* \ + --exclude libcudnn* \ + --exclude libcufft* \ + --exclude libnccl* \ + --exclude libnvrtc* + fi + else + if [ -f dist/mlx_cpu*.whl ]; then + mv dist/mlx_cpu*.whl wheelhouse/ + fi + if [ -f dist/mlx_cuda*.whl ]; then + mv dist/mlx_cuda*.whl wheelhouse/ + fi + if [ -f dist/mlx_metal*.whl ]; then + mv dist/mlx_metal*.whl wheelhouse/ + fi + fi diff --git a/Source/Cmlx/mlx/.github/actions/build-windows/action.yml b/Source/Cmlx/mlx/.github/actions/build-windows/action.yml deleted file mode 100644 index 372f84815..000000000 --- a/Source/Cmlx/mlx/.github/actions/build-windows/action.yml +++ /dev/null @@ -1,26 +0,0 @@ -name: 'Build on Windows' - -runs: - using: 'composite' - steps: - - name: Install Python package - id: python-build - shell: cmd - env: - # For MSVC, Ninja/Release is the only config supported by ccache. - CMAKE_ARGS: >- - -G Ninja - -DCMAKE_BUILD_TYPE=Release - -DCMAKE_C_COMPILER=cl - -DCMAKE_CXX_COMPILER=cl - -DCMAKE_RC_COMPILER=rc - run: | - uv pip install ".[dev]" -v - :: Pass the CMAKE_ARGS to following steps. - >>%GITHUB_OUTPUT% ECHO CMAKE_ARGS=%CMAKE_ARGS% - - - name: Build CPP only - shell: cmd - run: | - cmake . -B build ${{ steps.python-build.outputs.CMAKE_ARGS }} - cmake --build build -j %NUMBER_OF_PROCESSORS% diff --git a/Source/Cmlx/mlx/.github/actions/setup-linux/action.yml b/Source/Cmlx/mlx/.github/actions/setup-linux/action.yml deleted file mode 100644 index 08bb528f5..000000000 --- a/Source/Cmlx/mlx/.github/actions/setup-linux/action.yml +++ /dev/null @@ -1,93 +0,0 @@ -name: 'Setup Linux Environment' -description: 'Install dependencies for Linux builds' - -inputs: - toolkit: - description: 'Which toolkit to install' - required: false - default: 'cpu' - python-version: - description: 'Version of python to set up' - required: false - default: '3.14' - use-ccache: - description: 'Whether to enable ccache' - required: false - default: 'true' - -runs: - using: "composite" - steps: - - name: Install common dependencies - shell: bash - run: | - echo "::group::Install common dependencies" - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - zip \ - libblas-dev liblapack-dev liblapacke-dev \ - openmpi-bin openmpi-common libopenmpi-dev - echo "::endgroup::" - - - name: Use ccache - if: ${{ inputs.use-ccache == 'true' }} - uses: hendrikmuhs/ccache-action@v1.2 - with: - key: ccache-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }} - max-size: 1GB - # ccache-action bug: running "apt-get update" fails on large arm runner. - update-package-index: false - - - uses: actions/setup-python@v6 - with: - python-version: ${{ inputs.python-version }} - - - name: Setup Python venv - shell: bash - run: | - echo "::group::Setup Python venv" - python -m venv .venv - source .venv/bin/activate - pip install setuptools cmake typing_extensions - echo PATH=$PATH >> $GITHUB_ENV - # Search python packages in .venv - echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV - echo "::endgroup::" - - - name: Install CUDA toolkit - if: ${{ startsWith(inputs.toolkit, 'cuda') }} - shell: bash - env: - # Note: the CI machine does not meet CUDA 13's driver requirement. - # Compatibility matrix: - # https://docs.nvidia.com/deeplearning/cudnn/backend/latest/reference/support-matrix.html - PACKAGES: | - { - "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6", - "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9", - "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0" - } - run: | - echo "::group::Install CUDA toolkit" - # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is - # Jetson specific. SBSA means Arm Server Base System Architecture. - ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }} - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb - sudo dpkg -i cuda-keyring_1.1-1_all.deb - sudo apt-get update - sudo apt-get install -y --no-install-recommends \ - libnccl2 libnccl-dev \ - ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} - echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH - echo "::endgroup::" - - - name: CUDA packages and driver report - if: ${{ startsWith(inputs.toolkit, 'cuda') }} - shell: bash - run: | - echo "::group::Installed NVIDIA and CUDA packages" - dpkg -l | egrep "cuda|nvidia" -i - echo "::endgroup::" - echo "::group::NVIDIA-SMI Status" - nvidia-smi || true - echo "::endgroup::" diff --git a/Source/Cmlx/mlx/.github/actions/setup-macos/action.yml b/Source/Cmlx/mlx/.github/actions/setup-macos/action.yml deleted file mode 100644 index a108e686f..000000000 --- a/Source/Cmlx/mlx/.github/actions/setup-macos/action.yml +++ /dev/null @@ -1,24 +0,0 @@ -name: 'Setup macOS Environment' -description: 'Install dependencies for macOS builds' - -inputs: - python-version: - description: 'Python version to use' - required: false - default: '3.10' - -runs: - using: "composite" - steps: - - name: Install Homebrew packages - shell: sh - run: /opt/homebrew/bin/brew install openmpi - - - name: Verify MetalToolchain installed - shell: bash - run: xcodebuild -showComponent MetalToolchain - - - uses: conda-incubator/setup-miniconda@v3 - with: - miniconda-version: "latest" - python-version: ${{ inputs.python-version }} diff --git a/Source/Cmlx/mlx/.github/actions/setup-windows/action.yml b/Source/Cmlx/mlx/.github/actions/setup-windows/action.yml deleted file mode 100644 index 83afd4de1..000000000 --- a/Source/Cmlx/mlx/.github/actions/setup-windows/action.yml +++ /dev/null @@ -1,42 +0,0 @@ -name: 'Setup Windows environment' - -inputs: - python-version: - description: 'Version of python to set up' - required: false - default: '3.14' - use-ccache: - description: 'Whether to enable ccache' - required: false - default: 'true' - -runs: - using: 'composite' - steps: - - name: Use ccache - if: ${{ inputs.use-ccache == 'true' }} - uses: hendrikmuhs/ccache-action@v1.2 - with: - key: ccache-${{ runner.os }}-${{ runner.arch }}-cpu - max-size: 1GB - - - name: Setup Visual Studio cmd - shell: cmd - run: | - :: Find out path to VS. - pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\" - for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x - popd - :: Import VS vars. - call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" x64 - :: Export to all steps. - >>%GITHUB_ENV% set - - - uses: astral-sh/setup-uv@v7 - - - name: Setup Python venv - shell: cmd - run: | - uv venv --python ${{ inputs.python-version }} - call ".venv/Scripts/activate.bat" - >>%GITHUB_ENV% set diff --git a/Source/Cmlx/mlx/.github/actions/setup/action.yml b/Source/Cmlx/mlx/.github/actions/setup/action.yml new file mode 100644 index 000000000..510086390 --- /dev/null +++ b/Source/Cmlx/mlx/.github/actions/setup/action.yml @@ -0,0 +1,271 @@ +name: 'Setup environment' +description: 'Install dependencies for Linux and Windows builds' + +inputs: + toolkit: + description: 'Which toolkit to install' + required: false + default: 'cpu' + python-version: + description: 'Version of python to set up' + required: false + default: '3.12' + use-ccache: + description: 'Whether to enable ccache' + required: false + default: 'true' + ccache-key: + description: 'Extra key to use in the key of ccache' + required: false + default: 'test' + ccache-save: + description: 'Whether the ccache in this workflow will be saved' + required: false + default: 'auto' + ccache-toolkit: + description: 'Use this instead of inputs.toolkit in the ccache key' + required: false + default: '' + +outputs: + cmake-args: + description: 'The args for generating CMake project' + value: ${{ steps.cmake-args.outputs.cmakeArgs }} + +runs: + using: 'composite' + steps: + - name: Install Linux dependencies + if: runner.os == 'Linux' + shell: bash + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + gdb g++ ninja-build unzip \ + libblas-dev liblapack-dev liblapacke-dev \ + openmpi-bin openmpi-common libopenmpi-dev + + - name: Install macOS dependencies + if: runner.os == 'macOS' + shell: bash + run: | + brew update + brew trust aws/tap # suppress warning in github actions + brew install openmpi + xcodebuild -showComponent MetalToolchain + sysctl -a | grep machdep.cpu + + - name: Setup Windows environment + if: runner.os == 'Windows' + shell: cmd + run: | + :: Find out path to Visual Studio. + pushd "C:\Program Files (x86)\Microsoft Visual Studio\Installer\" + for /f "delims=" %%x in ('.\vswhere.exe -latest -property InstallationPath') do set VSPATH=%%x + popd + :: Import Visual stubs vars. + call "%VSPATH%\VC\Auxiliary\Build\vcvarsall.bat" ${{ runner.arch }} + :: Avoid using default TMP which uses short path and causes mismatch of CCACHE_BASEDIR. + set TMP=%RUNNER_TEMP% + set TEMP=%RUNNER_TEMP% + set UV_CACHE_DIR=%RUNNER_TEMP% + :: The cuda headers are downloaded. + set CCACHE_COMPILERCHECK=content + set CCACHE_SLOPPINESS=include_file_ctime,include_file_mtime + :: Export to all steps. + >>%GITHUB_ENV% set + + - uses: astral-sh/setup-uv@v8.2.0 + with: + enable-cache: false + quiet: true + + - name: Use ccache + if: inputs.use-ccache == 'true' + uses: hendrikmuhs/ccache-action@v1.2.23 + with: + key: v7-${{ inputs.ccache-key }}-${{ runner.os }}-${{ runner.arch }}-${{ inputs.ccache-toolkit || inputs.toolkit }} + max-size: |- + ${{ case(inputs.ccache-key == 'release', + case(startsWith(inputs.toolkit, 'cuda'), + case(runner.os == 'Linux', '300MB', + '160MB'), + '60MB'), + case(startsWith(inputs.toolkit, 'cuda'), + case(runner.os == 'Linux' && runner.arch == 'x64', '600MB', + runner.os == 'Linux', '400MB', + '320MB'), + runner.os == 'macOS' && inputs.toolkit == 'metal', '300MB', + '200MB')) + }} + save: ${{ !startsWith(github.ref, 'refs/pull/') && (inputs.ccache-save != 'false') }} + # ccache-action bug: running "apt-get update" fails on large arm runner. + update-package-index: false + + - name: Cache JIT-compiled CUDA kernels + if: runner.os == 'Linux' && startsWith(inputs.toolkit, 'cuda') + uses: actions/cache@v5 + with: + path: /tmp/mlx-ptx-cache + key: >- + ptx-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}- + ${{ hashFiles('mlx/backend/cuda/**') }} + + - name: Setup Python venv + if: runner.os != 'Windows' + shell: bash + run: | + uv venv --python ${{ inputs.python-version }} --managed-python + # Make sure all builds use the same cmake binary. + uv pip install cmake + source .venv/bin/activate + echo PATH=$PATH >> $GITHUB_ENV + # Search python packages in .venv + echo PYTHONPATH=`python -c 'import sys; print(sys.path[-1])'` >> $GITHUB_ENV + # Use PTX ccache for CUDA. + if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then + echo MLX_PTX_CACHE_DIR=/tmp/mlx-ptx-cache >> $GITHUB_ENV + fi + + - name: Setup Python venv (Windows) + if: runner.os == 'Windows' + shell: cmd + run: | + uv venv --python ${{ inputs.python-version }}${{ runner.arch == 'arm64' && '-arm64' || ''}} || exit /b + uv pip install cmake + call ".venv/Scripts/activate.bat" + >>%GITHUB_ENV% set + + - name: Install CUDA toolkit (Linux) + if: runner.os == 'Linux' && startsWith(inputs.toolkit, 'cuda') + shell: bash + env: + PACKAGES: | + { + "cuda-12.6": "libcudnn9-dev-cuda-12 cuda-compiler-12-6 cuda-libraries-dev-12-6", + "cuda-12.9": "libcudnn9-dev-cuda-12 cuda-compiler-12-9 cuda-libraries-dev-12-9", + "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-compiler-13-0 cuda-libraries-dev-13-0" + } + run: | + # The CUDA binaries are hosted in the "sbsa" repo, the "arm64" repo is + # Jetson specific. SBSA means Arm Server Base System Architecture. + ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }} + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb + sudo dpkg -i cuda-keyring_1.1-1_all.deb + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libnccl2 libnccl-dev \ + ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} + echo "/usr/local/${{ inputs.toolkit }}/bin" >> $GITHUB_PATH + + - name: Install CUDA Toolkit (Windows) + if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') + shell: powershell + env: + INSTALLERS: | + { + "cuda-12.6": "https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_561.17_windows.exe", + "cuda-12.9": "https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda_12.9.1_576.57_windows.exe", + "cuda-13.0": "https://developer.download.nvidia.com/compute/cuda/13.0.2/local_installers/cuda_13.0.2_windows.exe" + } + PACKAGES: | + { + "cuda-12.6": ["cudart_12.6", "nvcc_12.6", "cublas_12.6", "cublas_dev_12.6", "cufft_12.6", "cufft_dev_12.6", "nvrtc_12.6", "nvrtc_dev_12.6"], + "cuda-12.9": ["cudart_12.9", "nvcc_12.9", "cublas_12.9", "cublas_dev_12.9", "cufft_12.9", "cufft_dev_12.9", "nvrtc_12.9", "nvrtc_dev_12.9"], + "cuda-13.0": ["cudart_13.0", "nvcc_13.0", "cublas_13.0", "cublas_dev_13.0", "cufft_13.0", "cufft_dev_13.0", "nvrtc_13.0", "nvrtc_dev_13.0", "crt_13.0", "nvvm_13.0", "nvptxcompiler_13.0"], + } + run: | + $ErrorActionPreference = "Stop" + $cudaUrl = "${{ fromJson(env.INSTALLERS)[inputs.toolkit] }}" + $cudaInstaller = "./install.exe" + + echo "Downloading '$cudaUrl'..." + $webClient = New-Object System.Net.WebClient + $webClient.DownloadFile($cudaUrl, $cudaInstaller) + echo "Downloaded: $cudaInstaller ($(([math]::Round((Get-Item $cudaInstaller).Length / 1MB, 2))) MB)" + + $args = "-s ${{ join(fromJson(env.PACKAGES)[inputs.toolkit], ' ') }}" + echo "Running '$cudaInstaller $args'..." + Start-Process -FilePath $cudaInstaller -ArgumentList "$args" -NoNewWindow -Wait + $cudaPath = (Resolve-Path "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\*").path + echo "$cudaPath\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append + + - name: Install cuDNN (Windows) + if: runner.os == 'Windows' && startsWith(inputs.toolkit, 'cuda') + id: cudnn + shell: powershell + env: + ARCHIVES: | + { + "cuda-12.6": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda12-archive.zip", + "cuda-12.9": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda12-archive.zip", + "cuda-13.0": "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/windows-x86_64/cudnn-windows-x86_64-9.23.2.1_cuda13-archive.zip" + } + run: | + $ErrorActionPreference = "Stop" + $cudnnUrl = "${{ fromJson(env.ARCHIVES)[inputs.toolkit] }}" + $cudnnZip = "cudnn.zip" + + echo "Downloading '$cudnnUrl'..." + $webClient = New-Object System.Net.WebClient + $webClient.DownloadFile($cudnnUrl, $cudnnZip) + echo "Downloaded: $cudnnZip ($(([math]::Round((Get-Item $cudnnZip).Length / 1MB, 2))) MB)" + + echo "Extracing..." + Expand-Archive -Path $cudnnZip -DestinationPath cudnn-extracted + $cudnnDir = (Get-ChildItem -Path cudnn-extracted -Directory)[0].FullName + echo "cudnnDir=$($cudnnDir -replace '\\', '/')" | Out-File -FilePath $env:GITHUB_OUTPUT + + - name: Generate CMake args + id: cmake-args + shell: bash + run: | + cmakeArgs=( + "-G Ninja" + ) + if ${{ runner.os == 'macOS' }} ; then + if ${{ inputs.toolkit == 'cpu' }} ; then + cmakeArgs+=("-DMLX_BUILD_METAL=OFF") + else + cmakeArgs+=("-DMLX_BUILD_METAL=ON") + if ${{ inputs.toolkit == 'jit' }} ; then + cmakeArgs+=("-DMLX_METAL_JIT=ON") + fi + fi + fi + if ${{ runner.os == 'Windows' }} ; then + cmakeArgs+=( + "-DCMAKE_C_COMPILER=cl" + "-DCMAKE_CXX_COMPILER=cl" + "-DCMAKE_RC_COMPILER=rc" + ) + else + cmakeArgs+=( + "-DCMAKE_COMPILE_WARNING_AS_ERROR=ON" + ) + fi + if ${{ startsWith(inputs.toolkit, 'cuda') }} ; then + cmakeArgs+=("-DMLX_BUILD_CUDA=ON") + # Some machines have no GPU. + if ! __nvcc_device_query ; then + # Use a fallback arch for testing. + cmakeArgs+=("-DMLX_CUDA_ARCHITECTURES=80") + # Can not build tests and stubs when the built executables can not run. + cmakeArgs+=("-DMLX_BUILD_TESTS=OFF") + cmakeArgs+=("-DMLX_BUILD_PYTHON_STUBS=OFF") + fi + echo + # Set cuDNN paths. + if ${{ runner.os == 'Windows' }} ; then + cmakeArgs+=( + "-DCUDNN_INCLUDE_PATH=${{ steps.cudnn.outputs.cudnnDir }}/include" + "-DCUDNN_LIBRARY_PATH=${{ steps.cudnn.outputs.cudnnDir }}/lib/x64" + ) + fi + else + cmakeArgs+=("-DMLX_BUILD_CUDA=OFF") + fi + # Pass to following steps. + IFS=" " + echo ${cmakeArgs[*]} + echo "cmakeArgs=${cmakeArgs[*]}" >> $GITHUB_OUTPUT diff --git a/Source/Cmlx/mlx/.github/actions/test-linux/action.yml b/Source/Cmlx/mlx/.github/actions/test-linux/action.yml index 3258670ef..7c5f43221 100644 --- a/Source/Cmlx/mlx/.github/actions/test-linux/action.yml +++ b/Source/Cmlx/mlx/.github/actions/test-linux/action.yml @@ -1,69 +1,69 @@ -name: 'Run Linux tests' - -inputs: - has-gpu: - description: 'Run GPU tests' - required: false - default: false +name: 'Run tests' +description: 'Run Python and C++ tests on Linux' runs: - using: "composite" + using: 'composite' steps: - - name: Run MPI tests + - name: Check GPU support + id: gpu-check shell: bash run: | - echo "::group::MPI tests" - mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - echo "::endgroup::" + if __nvcc_device_query ; then + echo "good=true" >> $GITHUB_OUTPUT + else + echo "good=false" >> $GITHUB_OUTPUT + fi + echo + + - name: Run MPI tests + if: steps.gpu-check.outputs.good == 'false' + shell: bash + run: mpirun --bind-to none --allow-run-as-root -host localhost:8 -np 8 python python/tests/mpi_test_distributed.py - name: Run distributed tests - if: ${{ inputs.has-gpu == 'false' }} + if: steps.gpu-check.outputs.good == 'false' shell: bash run: | - echo "::group::Distributed tests" - mlx.launch --verbose -n 8 python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) + mlx.launch --verbose -n 8 python python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) if grep -Fq '[WARN]' stderr.log ; then grep -F '[WARN]' stderr.log echo "Distributed ring test failed"; exit 1; fi - echo "::endgroup::" - name: Run Python tests - CPU - if: ${{ inputs.has-gpu == 'false' }} + if: steps.gpu-check.outputs.good == 'false' shell: bash env: DEVICE: cpu - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" + run: python python/tests/run.py -v - name: Run Python tests - GPU - if: ${{ inputs.has-gpu == 'true' }} + if: steps.gpu-check.outputs.good == 'true' shell: bash env: DEVICE: gpu - run: | - echo "::group::Python tests - GPU" - python -m tests discover python/tests -v - echo "::endgroup::" + run: python python/tests/run.py -v - name: Run CPP tests - CPU shell: bash env: DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/tests/tests - echo "::endgroup::" + run: ./build/cpp/mlx/tests/tests - name: Run CPP tests - GPU - if: ${{ inputs.has-gpu == 'true' }} + if: steps.gpu-check.outputs.good == 'true' shell: bash env: DEVICE: gpu + run: ./build/cpp/mlx/tests/tests -sfe="*linalg_tests.cpp" + + - name: Show stack trace on crash + if: failure() + shell: bash run: | - echo "::group::CPP tests - GPU" - ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp" - echo "::endgroup::" + set +e + sleep 10 + if coredumpctl list; then + coredumpctl debug --debugger-arguments="-batch -ex 'thread apply all bt'" + fi diff --git a/Source/Cmlx/mlx/.github/actions/test-macos/action.yml b/Source/Cmlx/mlx/.github/actions/test-macos/action.yml new file mode 100644 index 000000000..55e5acc5c --- /dev/null +++ b/Source/Cmlx/mlx/.github/actions/test-macos/action.yml @@ -0,0 +1,76 @@ +name: 'Run tests' +description: 'Run Python and C++ tests on macOS' + +inputs: + toolkit: + description: 'Which toolkit to test' + required: false + default: 'cpu' + +runs: + using: 'composite' + steps: + - name: Install tests dependencies + shell: bash + run: uv pip install tensorflow + + - name: Run tests + shell: bash + env: + METAL_DEBUG_ERROR_MODE: 0 + METAL_DEVICE_WRAPPER_TYPE: 1 + run: | + run_tests() { + for WHEEL in dist/*.whl; do + break + done + + echo "::group::Install $WHEEL" + uv pip install $WHEEL + echo "::endgroup::" + + echo "::group::Run CPP tests" + ./build/tests/tests + echo "::endgroup::" + + echo "::group::Run distributed tests" + # TODO: Re-enable for GPU after openmpi updating their libprrte + # dependency: https://github.com/openpmix/prrte/issues/2416 + if ${{ inputs.toolkit == 'cpu' }} ; then + mpirun --bind-to none -host localhost:8 -np 8 -x DYLD_LIBRARY_PATH=/opt/homebrew/lib uv run python/tests/mpi_test_distributed.py + fi + mlx.launch --verbose -n 8 uv run python/tests/ring_test_distributed.py -v 2> >(tee -a stderr.log >&2) + if $(grep "\[WARN\]" stderr.log); then + echo "Distributed ring test failed" + exit 1 + fi + echo "::endgroup::" + + echo "::group::Run Python tests" + uv run python/tests/run.py -v + echo "::endgroup::" + + if ${{ inputs.toolkit != 'cpu' }} ; then + echo "::group::Build example extension" + pushd examples/extensions + uv pip install -r requirements.txt + uv run setup.py build_ext --inplace + uv run test.py + popd + echo "::endgroup::" + fi + } + + if [[ -d artifact ]] ; then + for SUBDIR in artifact/*; do + rm -rf build dist + mv $SUBDIR/* . + # Downloaded binaries all lose executable permission. + find build -type f -exec chmod 755 {} \; + # Path to metallib is hard-coded so move it besides executable. + mv build/mlx/backend/metal/kernels/mlx.metallib build/tests + run_tests + done + else + run_tests + fi diff --git a/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml b/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml new file mode 100644 index 000000000..a29e3c1dc --- /dev/null +++ b/Source/Cmlx/mlx/.github/actions/test-wheel/action.yml @@ -0,0 +1,54 @@ +name: 'Test wheel' +description: 'Run tests with the built wheels' + +inputs: + toolkit: + description: 'Which toolkit to test' + required: false + default: 'cpu' + +runs: + using: 'composite' + steps: + - name: Get Python version + id: python + shell: bash + run: python -c "import sys; print(f'version={sys.version_info.major}.{sys.version_info.minor}')" >> $GITHUB_OUTPUT + + - name: Download frontend packages + uses: actions/download-artifact@v8 + with: + pattern: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ steps.python.outputs.version }} + path: wheelhouse + + - name: Download backend packages + uses: actions/download-artifact@v8 + with: + pattern: backend-${{ inputs.toolkit}}-${{ runner.os }}-${{ runner.arch }} + path: wheelhouse + + - name: Test local packages + shell: bash + run: | + # Fall back to PyPI if PyTorch index has issues. + uv pip install torch --torch-backend=cpu + uv pip install numpy + if ${{ inputs.toolkit == 'cpu' }} ; then + uv pip install wheelhouse/mlx_cpu*.whl + uv pip install wheelhouse/mlx-*.whl + elif ${{ startsWith(inputs.toolkit, 'cuda') }} ; then + uv pip install wheelhouse/mlx_cuda*.whl + uv pip install wheelhouse/mlx-*.whl + elif ${{ inputs.toolkit == 'metal' }} ; then + uv pip install wheelhouse/mlx_metal-*-macosx_26_0_arm64.whl + uv pip install wheelhouse/mlx-*-macosx_26_0_arm64.whl + else + echo "No matching backend wheel to install" + exit 1 + fi + if ${{ runner.os == 'macOS' }} ; then + # FIXME: run.py does not quit in macOS CI. + python -m unittest discover -v python/tests + else + uv run python/tests/run.py -v --failfast + fi diff --git a/Source/Cmlx/mlx/.github/actions/test-windows/action.yml b/Source/Cmlx/mlx/.github/actions/test-windows/action.yml index ebb886e7b..6d89c4e5b 100644 --- a/Source/Cmlx/mlx/.github/actions/test-windows/action.yml +++ b/Source/Cmlx/mlx/.github/actions/test-windows/action.yml @@ -1,20 +1,17 @@ -name: 'Run tests on Windows' +name: 'Run tests' +description: 'Run Python and C++ tests on Windows' runs: using: 'composite' steps: - name: Run Python tests - CPU shell: bash - run: | - echo "::group::Python tests - CPU" - python -m unittest discover python/tests -v - echo "::endgroup::" + env: + DEVICE: cpu + run: uv run python/tests/run.py -v - name: Run CPP tests - CPU shell: bash env: DEVICE: cpu - run: | - echo "::group::CPP tests - CPU" - ./build/tests.exe -tce="*gguf*,test random uniform" - echo "::endgroup::" + run: ./build/cpp/mlx/tests.exe -tce="*gguf*" diff --git a/Source/Cmlx/mlx/.github/pull_request_template.md b/Source/Cmlx/mlx/.github/pull_request_template.md index 02bb9b79a..a7f928e8b 100644 --- a/Source/Cmlx/mlx/.github/pull_request_template.md +++ b/Source/Cmlx/mlx/.github/pull_request_template.md @@ -1,12 +1,2 @@ -## Proposed changes - -Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #. - -## Checklist - -Put an `x` in the boxes that apply. - -- [ ] I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document -- [ ] I have run `pre-commit run --all-files` to format my code / installed pre-commit prior to committing changes -- [ ] I have added tests that prove my fix is effective or that my feature works -- [ ] I have updated the necessary documentation (if needed) +- ☑️ I understand it is strictly prohibited to use AI to write PR description +- AI usage disclosure: diff --git a/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js b/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js new file mode 100644 index 000000000..f5a05396c --- /dev/null +++ b/Source/Cmlx/mlx/.github/scripts/find-eligible-contributors.js @@ -0,0 +1,220 @@ +#!/usr/bin/env node + +/** + * find-eligible-contributors.js + * + * Finds users who have opened pull requests on a GitHub repo, + * are NOT collaborators on that repo, and match: + * - fewer than 5 currently open PRs + * - more than 2 merged PRs + * - merge rate (merged / closed) > 50% + * + * Usage: + * GITHUB_TOKEN=xxxx node find-eligible-contributors.js + */ + +const GITHUB_API = 'https://api.github.com'; + +function authHeaders(token) { + return { + Authorization: `Bearer ${token}`, + Accept: 'application/vnd.github+json', + 'X-GitHub-Api-Version': '2022-11-28', + 'User-Agent': 'find-eligible-contributors-script', + }; +} + +async function fetchAllPages(url, token, { stopWhen } = {}) { + const results = []; + let page = 1; + const perPage = 100; + const concurrentRequests = 5; + + while (true) { + const responses = await Promise.all( + Array.from({ length: concurrentRequests }, async (_, index) => { + const pageUrl = new URL(url); + pageUrl.searchParams.set('per_page', String(perPage)); + pageUrl.searchParams.set('page', String(page + index)); + + const res = await fetch(pageUrl, { headers: authHeaders(token) }); + if (!res.ok) { + const body = await res.text(); + throw new Error(`GitHub API error ${res.status} for ${pageUrl}: ${body}`); + } + return res.json(); + }) + ); + + let end = false; + for (const data of responses) { + if (!Array.isArray(data) || data.length === 0) { + end = true; + continue; + } + if (data.length < perPage) { + end = true; + } + + for (const item of data) { + if (stopWhen && stopWhen(item)) break; + results.push(item); + } + } + + if (end) { + break; + } + + page += concurrentRequests; + } + + return results; +} + +async function fetchPullRequests(owner, repo, token) { + const url = `${GITHUB_API}/repos/${owner}/${repo}/pulls?state=all&sort=created&direction=desc`; + + const oneYearAgo = new Date(); + oneYearAgo.setFullYear(oneYearAgo.getFullYear() - 1); + + return fetchAllPages(url, token, { + stopWhen: (pr) => new Date(pr.created_at) < oneYearAgo, + }); +} + +async function fetchCollaborators(owner, repo, token) { + const url = `${GITHUB_API}/repos/${owner}/${repo}/collaborators`; + const collaborators = await fetchAllPages(url, token); + return new Set(collaborators.map((c) => c.login)); +} + +/** + * Aggregate per-user PR stats from the list of PRs. + * Returns Map + */ +function aggregateUserStats(pullRequests) { + const stats = new Map(); + + for (const pr of pullRequests) { + const login = pr.user?.login; + if (!login) continue; + + if (!stats.has(login)) { + stats.set(login, { + open: 0, + closed: 0, + merged: 0, + mostRecentCreatedAt: null, + }); + } + const s = stats.get(login); + + if (pr.state === 'open') { + s.open += 1; + } else if (pr.state === 'closed') { + s.closed += 1; + if (pr.merged_at) { + s.merged += 1; + } + } + + const createdAt = new Date(pr.created_at); + if (!s.mostRecentCreatedAt || createdAt > s.mostRecentCreatedAt) { + s.mostRecentCreatedAt = createdAt; + } + } + + return stats; +} + +/** + * Apply eligibility filters: + * - fewer than 5 currently open PRs + * - more than 2 merged PRs + * - merge rate (merged / closed) > 50% + * - not a collaborator + * - has an open PR or at least one PR within the past 6 weeks + */ +function filterEligibleUsers(stats, collaboratorLogins) { + const SIX_WEEKS_MS = 6 * 7 * 24 * 60 * 60 * 1000; + const cutoff = new Date((new Date).getTime() - SIX_WEEKS_MS); + const eligible = []; + + for (const [login, s] of stats.entries()) { + if (s.open >= 5) continue; + if (s.merged <= 2) continue; + if (collaboratorLogins.has(login)) continue; + if (s.open == 0 && s.mostRecentCreatedAt < cutoff) continue; + if ((s.merged / s.closed) < 0.5) continue; + + eligible.push(login); + } + return eligible; +} + +/** + * Given a repo (owner/repo) and a GitHub token, returns the login names of + * eligible external contributors matching all filters. + * + * @param {Object} params + * @param {string} params.owner - Repo owner (user or org). + * @param {string} params.repo - Repo name. + * @param {string} params.token - GitHub token with repo read (and ideally push) access. + * @returns {Promise} + */ +async function findEligibleContributors({ owner, repo, token } = {}) { + if (!owner || !repo) { + throw new Error('findEligibleContributors requires both "owner" and "repo".'); + } + if (!token) { + throw new Error('findEligibleContributors requires a "token".'); + } + + const [pullRequests, collaboratorLogins] = await Promise.all([ + fetchPullRequests(owner, repo, token), + fetchCollaborators(owner, repo, token), + ]); + + const stats = aggregateUserStats(pullRequests); + return filterEligibleUsers(stats, collaboratorLogins); +} + +// ---- CLI entry point ---- + +function isRunAsCLI() { + return import.meta.url === `file://${process.argv[1]}`; +} + +async function runCLI() { + const [, , owner, repo] = process.argv; + if (!owner || !repo) { + console.error('Usage: node find-eligible-contributors.js '); + process.exit(1); + } + const token = process.env.GITHUB_TOKEN; + if (!token) { + console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); + process.exit(1); + } + + console.error(`Fetching pull requests and collaborators for ${owner}/${repo} in parallel...`); + const eligibleLogins = await findEligibleContributors({ owner, repo, token }); + + console.log(JSON.stringify(eligibleLogins, null, 2)); + console.error(`\n${eligibleLogins.length} user(s) match the criteria.`); +} + +if (isRunAsCLI()) { + runCLI().catch((err) => { + console.error('Fatal error:', err); + process.exit(1); + }); +} + +export { findEligibleContributors }; diff --git a/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js b/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js new file mode 100644 index 000000000..0f0f0817a --- /dev/null +++ b/Source/Cmlx/mlx/.github/scripts/update-bypass-list.js @@ -0,0 +1,161 @@ +#!/usr/bin/env node + +/** + * update-bypass-list.js + * + * Syncs a repo's PR interaction-limits bypass list with the current + * set of eligible external contributors. + * + * Usage: + * GITHUB_TOKEN=xxxx node update-bypass-list.js [--dry-run] + */ + +import { findEligibleContributors } from './find-eligible-contributors.js'; + +const GITHUB_API = 'https://api.github.com'; + +function authHeaders(token) { + return { + Authorization: `Bearer ${token}`, + Accept: 'application/vnd.github+json', + 'X-GitHub-Api-Version': '2026-03-10', + 'User-Agent': 'update-bypass-list-script', + 'Content-Type': 'application/json', + }; +} + +function bypassListUrl(owner, repo) { + return `${GITHUB_API}/repos/${owner}/${repo}/interaction-limits/pulls/bypass-list`; +} + +async function fetchBypassList(owner, repo, token) { + const url = bypassListUrl(owner, repo); + const res = await fetch(url, { headers: authHeaders(token) }); + + if (res.status === 404) { + return new Set(); + } + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to fetch bypass list (${res.status}): ${body}`); + } + + const data = await res.json(); + const logins = data.map((user) => user?.login).filter(Boolean); + + return new Set(logins); +} + +async function addUsersToBypassList(owner, repo, token, usernames) { + if (usernames.length === 0) return; + + const res = await fetch(bypassListUrl(owner, repo), { + method: 'PUT', + headers: authHeaders(token), + body: JSON.stringify({ users: usernames }), + }); + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to add users [${usernames.join(', ')}] (${res.status}): ${body}`); + } +} + +async function removeUsersFromBypassList(owner, repo, token, usernames) { + if (usernames.length === 0) return; + + const res = await fetch(bypassListUrl(owner, repo), { + method: 'DELETE', + headers: authHeaders(token), + body: JSON.stringify({ users: usernames }), + }); + + if (!res.ok) { + const body = await res.text(); + throw new Error(`Failed to remove users [${usernames.join(', ')}] (${res.status}): ${body}`); + } +} + +function diffLists(eligibleLogins, currentBypassLogins) { + const toAdd = [...eligibleLogins].filter((login) => !currentBypassLogins.has(login)); + const toRemove = [...currentBypassLogins].filter((login) => !eligibleLogins.has(login)); + return { toAdd, toRemove }; +} + +/** + * Update the bypass list to eligible contributors. + * + * @param {Object} params + * @param {string} params.owner + * @param {string} params.repo + * @param {string} params.token + * @param {boolean} [params.dryRun] - If true, only computes the diff, makes no API writes. + */ +async function updateBypassList({ owner, repo, token, dryRun = false }) { + console.error(`Getting eligible external contributors for ${owner}/${repo}...`); + const eligibleLogins = new Set(await findEligibleContributors({ owner, repo, token })); + console.error(`Found ${eligibleLogins.size} eligible user(s).`); + + console.error(`Fetching current bypass list...`); + const currentBypassLogins = await fetchBypassList(owner, repo, token); + console.error(`Current bypass list has ${currentBypassLogins.size} user(s).`); + + const { toAdd, toRemove } = diffLists(eligibleLogins, currentBypassLogins); + + console.error(`\nUsers to add (${toAdd.length}): ${toAdd.join(', ') || '(none)'}`); + console.error(`Users to remove (${toRemove.length}): ${toRemove.join(', ') || '(none)'}`); + + if (dryRun) { + console.error('\nDry run mode: no changes will be made.'); + return; + } + + try { + await addUsersToBypassList(owner, repo, token, toAdd); + if (toAdd.length > 0) console.error(`Added ${toAdd.length} user(s) to bypass list.`); + } catch (err) { + console.error(`Error adding users: ${err.message}`); + } + + try { + await removeUsersFromBypassList(owner, repo, token, toRemove); + if (toRemove.length > 0) console.error(`Removed ${toRemove.length} user(s) from bypass list.`); + } catch (err) { + console.error(`Error removing users: ${err.message}`); + } +} + +// ---- CLI entry point ---- + +function isRunAsCLI() { + return import.meta.url === `file://${process.argv[1]}`; +} + +async function runCLI() { + const args = process.argv.slice(2); + const dryRun = args.includes('--dry-run'); + const [owner, repo] = args.filter((a) => !a.startsWith('--')); + + if (!owner || !repo) { + console.error('Usage: node update-bypass-list.js [--dry-run]'); + process.exit(1); + } + + const token = process.env.GITHUB_TOKEN; + if (!token) { + console.error('Error: set GITHUB_TOKEN environment variable with a valid GitHub token.'); + process.exit(1); + } + + await updateBypassList({ owner, repo, token, dryRun }); +} + +if (isRunAsCLI()) { + runCLI().catch((err) => { + console.error('Fatal error:', err); + process.exit(1); + }); +} + +export { updateBypassList }; diff --git a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml b/Source/Cmlx/mlx/.github/workflows/build_and_test.yml index ba0ccfa68..d590d10b1 100644 --- a/Source/Cmlx/mlx/.github/workflows/build_and_test.yml +++ b/Source/Cmlx/mlx/.github/workflows/build_and_test.yml @@ -1,6 +1,7 @@ name: Build and Test on: + workflow_dispatch: pull_request: push: branches: @@ -20,71 +21,123 @@ jobs: name: Check Lint runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: pre-commit/action@v3.0.1 - linux_build_and_test: - name: Linux (cpu, ${{ matrix.arch }}) + build_and_test: + name: ${{ matrix.os }} (${{ matrix.toolkit }}, ${{ matrix.arch }}) + if: github.repository == 'ml-explore/mlx' needs: check_lint strategy: fail-fast: false matrix: + os: ['Linux', 'Windows'] arch: ['x86_64', 'aarch64'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} + toolkit: ['cpu', 'cuda-12.6', 'cuda-12.9', 'cuda-13.0'] + exclude: + # CUDA does not support Windows on arm. + - os: 'Windows' + arch: 'aarch64' + toolkit: 'cuda-12.6' + - os: 'Windows' + arch: 'aarch64' + toolkit: 'cuda-12.9' + - os: 'Windows' + arch: 'aarch64' + toolkit: 'cuda-13.0' + # CUDA 12.6 does not compile with CUTLASS on Windows. + - os: 'Windows' + arch: 'x86_64' + toolkit: 'cuda-12.6' + runs-on: |- + ${{ case(matrix.os == 'Windows', + case(matrix.arch == 'aarch64', 'windows-11-arm', + startsWith(matrix.toolkit, 'cuda'), 'windows-2022-large', + 'windows-2022'), + case(matrix.arch == 'x86_64' && startsWith(matrix.toolkit, 'cuda'), 'gpu-t4-4-core', + matrix.arch == 'aarch64', 'ubuntu-22.04-arm', + 'ubuntu-22.04')) + }} steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - - uses: ./.github/actions/build-linux + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + id: setup + with: + toolkit: ${{ matrix.toolkit }} + - uses: ./.github/actions/build + with: + cmake-args: ${{ steps.setup.outputs.cmake-args }} + # For MSVC, Ninja/Release is the only config supported by ccache. + debug: ${{ matrix.os != 'Windows' }} + - name: Check generated Python stubs with ty + if: matrix.os == 'Linux' && matrix.arch == 'x86_64' && matrix.toolkit == 'cpu' + run: uvx ty check python/mlx/core - uses: ./.github/actions/test-linux - - run: df -h + if: matrix.os == 'Linux' && (matrix.toolkit == 'cpu' || matrix.arch == 'x86_64') + - uses: ./.github/actions/test-windows + if: matrix.os == 'Windows' && matrix.toolkit == 'cpu' - cuda_build_and_test: - name: Linux (${{ matrix.toolkit }}, ${{ matrix.arch }}) + mac_build: + name: macOS (${{ matrix.macos-target }}, ${{ matrix.toolkit }}) if: github.repository == 'ml-explore/mlx' - needs: check_lint strategy: fail-fast: false matrix: - arch: ['x86_64', 'aarch64'] - toolkit: ['cuda-12.6', 'cuda-12.9'] - runs-on: ${{ matrix.arch == 'x86_64' && 'gpu-t4-4-core' || 'ubuntu-22.04-arm' }} + macos-target: ['14.0', '15.0', '26.2'] + toolkit: ['cpu', 'metal', 'jit'] + runs-on: 'macos-26' + needs: check_lint steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + id: setup with: toolkit: ${{ matrix.toolkit }} - - uses: ./.github/actions/build-linux + ccache-key: 'test-${{ matrix.macos-target }}' + # Use the compile cache of metal build for all builds. + ccache-save: ${{ matrix.toolkit == 'metal' }} + ccache-toolkit: 'metal' + - uses: ./.github/actions/build-macos with: - toolkit: ${{ matrix.toolkit }} - - uses: ./.github/actions/test-linux - if: matrix.arch == 'x86_64' + cmake-args: ${{ steps.setup.outputs.cmake-args }} + macos-target: ${{ matrix.macos-target }} + - uses: ./.github/actions/test-macos + if: matrix.toolkit == 'cpu' + with: + toolkit: 'cpu' + - uses: actions/upload-artifact@v7 + if: matrix.toolkit != 'cpu' with: - has-gpu: true + name: mlx-${{ matrix.toolkit }}-macos${{ matrix.macos-target }} + path: | + dist/mlx-*.whl + build/mlx/backend/metal/kernels/mlx.metallib + build/tests/tests + if-no-files-found: error - mac_build_and_test: - name: macOS (${{ matrix.macos-target }}) + mac_test: + name: Test macOS (${{ matrix.toolkit }}) if: github.repository == 'ml-explore/mlx' + runs-on: [self-hosted, macos] + needs: mac_build strategy: + fail-fast: false matrix: - macos-target: ["14.0", "15.0", "26.0"] - runs-on: [self-hosted, macos] - env: - MACOSX_DEPLOYMENT_TARGET: ${{ matrix.macos-target }} - needs: check_lint - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos - - uses: ./.github/actions/build-macos - - windows_build_and_test: - name: Windows (cpu, x86_64) - needs: check_lint - runs-on: windows-2025 + toolkit: ['metal', 'jit'] steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-windows - - uses: ./.github/actions/build-windows - - uses: ./.github/actions/test-windows + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + with: + toolkit: ${{ matrix.toolkit }} + use-ccache: false + - uses: actions/download-artifact@v8 + with: + path: artifact + pattern: mlx-${{ matrix.toolkit }}-* + - run: ls -lhR artifact + - uses: ./.github/actions/test-macos + with: + toolkit: ${{ matrix.toolkit }} build_documentation: name: Build Documentation @@ -92,7 +145,7 @@ jobs: runs-on: ubuntu-22.04 needs: check_lint steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: ./.github/actions/build-docs linux_sanitizer_build_and_test: @@ -107,7 +160,7 @@ jobs: runs-on: ubuntu-22.04-arm steps: - name: Checkout code - uses: actions/checkout@v6 + uses: actions/checkout@v7 - name: Install Dependencies run: | @@ -145,7 +198,7 @@ jobs: image: fedora:42 steps: - name: Checkout code - uses: actions/checkout@v6 + uses: actions/checkout@v7 - name: CPP Build Test - No Release run: | diff --git a/Source/Cmlx/mlx/.github/workflows/documentation.yml b/Source/Cmlx/mlx/.github/workflows/documentation.yml index a9af245fc..fad164799 100644 --- a/Source/Cmlx/mlx/.github/workflows/documentation.yml +++ b/Source/Cmlx/mlx/.github/workflows/documentation.yml @@ -10,7 +10,7 @@ jobs: build: runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: ./.github/actions/build-docs deploy: @@ -25,4 +25,4 @@ jobs: steps: - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v4 + uses: actions/deploy-pages@v5 diff --git a/Source/Cmlx/mlx/.github/workflows/nightly.yml b/Source/Cmlx/mlx/.github/workflows/nightly.yml deleted file mode 100644 index 84962e14b..000000000 --- a/Source/Cmlx/mlx/.github/workflows/nightly.yml +++ /dev/null @@ -1,104 +0,0 @@ -name: Nightly Build - -on: - schedule: - - cron: 33 6 * * 1-5 - workflow_dispatch: - -permissions: - contents: read - -jobs: - build_linux_release: - strategy: - fail-fast: false - matrix: - python_version: ["3.10", "3.14"] - runs-on: ubuntu-22.04 - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - - uses: ./.github/actions/build-linux-release - with: - build-backend: ${{ matrix.python-version == '3.10' }} - arch: "x86_64" - - name: Upload mlx artifacts - uses: actions/upload-artifact@v6 - with: - name: linux-wheels-${{ matrix.python_version }} - path: wheelhouse/mlx-*.whl - retention-days: 7 - - name: Upload mlx-cpu artifacts - if: matrix.python_version == '3.10' - uses: actions/upload-artifact@v6 - with: - name: mlx-cpu - path: wheelhouse/mlx_cpu-*.whl - retention-days: 7 - - run: df -h - - build_linux_with_tests: - strategy: - fail-fast: false - matrix: - python_version: ["3.11", "3.12", "3.13", "3.14"] - runner: - - ubuntu-22.04 - - ubuntu-22.04-arm - runs-on: ${{ matrix.runner }} - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - python-version: ${{ matrix.python_version }} - - uses: ./.github/actions/build-linux - - uses: ./.github/actions/test-linux - - run: df -h - - build_mac_release: - if: github.repository == 'ml-explore/mlx' - strategy: - matrix: - python-version: ["3.10", "3.13"] - runs-on: [self-hosted, macos] - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos - with: - python-version: ${{ matrix.python-version }} - - uses: ./.github/actions/build-macos - - name: Build macOS 26 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 26.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 15 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 15.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - name: Build macOS 14 package - uses: ./.github/actions/build-macos-release - with: - macos-target: 14.0 - build-backend: ${{ matrix.python-version == '3.10' }} - - build_cuda_release: - if: github.repository == 'ml-explore/mlx' - runs-on: ubuntu-22-large - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux - with: - toolkit: 'cuda-12.9' - - name: Build Python package - uses: ./.github/actions/build-cuda-release - with: - toolkit: 'cuda-12.9' - arch: 'x86_64' - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - name: mlx-cuda - path: wheelhouse/mlx_cuda_*.whl - retention-days: 7 diff --git a/Source/Cmlx/mlx/.github/workflows/release.yml b/Source/Cmlx/mlx/.github/workflows/release.yml index 9b9420a50..b53f003ff 100644 --- a/Source/Cmlx/mlx/.github/workflows/release.yml +++ b/Source/Cmlx/mlx/.github/workflows/release.yml @@ -1,35 +1,44 @@ -name: PyPI Release +name: 'Release build' +description: 'Build Python wheels for nightly or offical releases' on: push: tags: - 'v*' - branches: - - 'test-publish/*' workflow_dispatch: inputs: - dry_run: - description: 'Dry run (do not publish to PyPi)' + publish: + description: 'Publish to PyPI' required: false type: boolean - dev_release: + dev-release: description: 'Development release (DEV_RELEASE=1)' required: false type: boolean + schedule: + - cron: 33 6 * * * + +# In jobs we must use |*publish| instead of |inputs.publish| because we can not +# set default value for workflow_dispatch inputs reliably. +env: + publish: &publish ${{ inputs.publish || github.event_name == 'push' }} + pypi-env: &pypi-env ${{ (inputs.publish || github.event_name == 'push') && 'pypi' || 'dry-run' }} permissions: contents: read jobs: build_documentation: + name: Build documentation if: github.repository == 'ml-explore/mlx' runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: ./.github/actions/build-docs deploy_documentation: - if: ${{ !inputs.dry_run }} + name: Deploy documentation + if: *publish needs: build_documentation permissions: pages: write @@ -41,216 +50,272 @@ jobs: steps: - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v4 + uses: actions/deploy-pages@v5 - build_linux_release: - if: github.repository == 'ml-explore/mlx' + build_frontend: + name: ${{ matrix.os }} (python-${{ matrix.python-version }}, ${{ matrix.arch }}) strategy: matrix: - python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"] + os: ['Linux', 'Windows'] arch: ['x86_64', 'aarch64'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }} - env: + python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t'] + # There is no cp310 binary for Windows on arm. + exclude: + - os: 'Windows' + arch: 'aarch64' + python-version: '3.10' + runs-on: |- + ${{ case(matrix.os == 'Windows', + case(matrix.arch == 'aarch64', 'windows-11-arm', + 'windows-2022'), + case(matrix.arch == 'aarch64', 'ubuntu-22.04-arm', + 'ubuntu-22.04')) + }} + env: &build-env PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} + DEV_RELEASE: ${{ inputs.dev-release && 1 || 0 }} steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + id: setup with: - python-version: ${{ matrix.python_version }} - use-ccache: false - - uses: ./.github/actions/build-linux-release + python-version: ${{ matrix.python-version }} + ccache-key: 'release' + ccache-save: false + - uses: ./.github/actions/build-wheel with: - build-backend: ${{ matrix.python_version == '3.10' }} - arch: ${{ matrix.arch }} - - name: Upload MLX artifacts - uses: actions/upload-artifact@v6 + cmake-args: ${{ steps.setup.outputs.cmake-args }} + build-backend: false + - run: unzip -l wheelhouse/*.whl + - uses: actions/upload-artifact@v7 with: - overwrite: true - name: linux-wheels-${{ matrix.python_version }}-${{ matrix.arch }} + name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} path: wheelhouse/mlx-*.whl if-no-files-found: error - - name: Upload CPU artifacts - if: matrix.python_version == '3.10' - uses: actions/upload-artifact@v6 + + build_backend: + name: ${{ matrix.os }} (${{ matrix.toolkit }}, ${{ matrix.arch }}) + if: github.repository == 'ml-explore/mlx' + strategy: + matrix: + os: ['Linux', 'Windows'] + arch: ['x86_64', 'aarch64'] + toolkit: ['cpu', 'cuda-12.9', 'cuda-13.0'] + exclude: + # CUDA does not support Windows on arm. + - os: 'Windows' + arch: 'aarch64' + toolkit: 'cuda-12.9' + - os: 'Windows' + arch: 'aarch64' + toolkit: 'cuda-13.0' + runs-on: |- + ${{ case(matrix.os == 'Windows', + case(matrix.arch == 'aarch64', 'windows-11-arm', + 'windows-2022-large'), + case(matrix.arch == 'aarch64', 'ubuntu-22-large-arm', + 'ubuntu-22-large')) + }} + env: *build-env + steps: + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + id: setup with: - overwrite: true - name: mlx-cpu-${{ matrix.arch }} - path: wheelhouse/mlx_cpu-*.whl - if-no-files-found: error + toolkit: ${{ matrix.toolkit }} + ccache-key: 'release' + - uses: ./.github/actions/build-wheel + with: + cmake-args: ${{ steps.setup.outputs.cmake-args }} + build-frontend: false + - run: unzip -l wheelhouse/*.whl + - uses: actions/upload-artifact@v7 + with: + name: backend-${{ matrix.toolkit }}-${{ runner.os }}-${{ runner.arch }} + path: wheelhouse/*.whl + if-no-files-found: ignore - build_mac_release: + build_mac_wheels: + name: macOS (python-${{ matrix.python-version }}) if: github.repository == 'ml-explore/mlx' strategy: matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - runs-on: [self-hosted, macos] - env: - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} + python-version: *pyver + runs-on: 'macos-26' + env: *build-env steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-macos + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup + id: setup with: + toolkit: 'metal' python-version: ${{ matrix.python-version }} - - - name: Install dependencies - shell: bash -l {0} - run: | - pip install --upgrade pip - pip install cmake setuptools typing_extensions - pip install -e . -v + ccache-key: 'release' + ccache-save: ${{ matrix.python-version == '3.10' }} - name: Build macOS 14 package - uses: ./.github/actions/build-macos-release + uses: ./.github/actions/build-wheel with: - macos-target: 14.0 + macos-target: '14.0' + cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: ${{ matrix.python-version == '3.10' }} - name: Build macOS 15 package - uses: ./.github/actions/build-macos-release + uses: ./.github/actions/build-wheel with: - macos-target: 15.0 + macos-target: '15.0' + cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: ${{ matrix.python-version == '3.10' }} - name: Build macOS 26 package - uses: ./.github/actions/build-macos-release + uses: ./.github/actions/build-wheel with: - macos-target: 26.0 + macos-target: '26.2' + cmake-args: ${{ steps.setup.outputs.cmake-args }} build-backend: ${{ matrix.python-version == '3.10' }} - - name: Upload MLX artifacts - uses: actions/upload-artifact@v6 + - name: Display content of the wheels + run: for WHEEL in wheelhouse/*.whl; do unzip -l $WHEEL; echo; done + - name: Upload frontend packages + uses: actions/upload-artifact@v7 with: - overwrite: true - name: mac-wheels-${{ matrix.python-version }} - path: dist/mlx-*.whl + name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} + path: wheelhouse/mlx-*.whl if-no-files-found: error - - name: Upload Metal artifacts + - name: Upload backend packages if: matrix.python-version == '3.10' - uses: actions/upload-artifact@v6 + uses: actions/upload-artifact@v7 with: - overwrite: true - name: mlx-metal - path: dist/mlx_metal-*.whl + name: backend-metal-${{ runner.os }}-${{ runner.arch }} + path: wheelhouse/mlx_metal-*.whl if-no-files-found: error - build_cuda_release: + test_wheel: + name: Test (${{ matrix.os }}, ${{ matrix.toolkit }}, ${{ matrix.arch }}) if: github.repository == 'ml-explore/mlx' + needs: [build_frontend, build_backend, build_mac_wheels] strategy: matrix: - arch: ['x86_64', 'aarch64'] - toolkit: ['cuda-12.9', 'cuda-13.0'] - runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22-large' || 'ubuntu-22-large-arm' }} - env: - PYPI_RELEASE: 1 - DEV_RELEASE: ${{ inputs.dev_release && 1 || 0 }} + os: ['Linux', 'Windows'] + arch: ['aarch64'] + toolkit: ['cpu'] + include: + - os: 'Linux' + arch: 'x86_64' + toolkit: 'cpu' + - os: 'Linux' + arch: 'x86_64' + toolkit: 'cuda-12.9' + - os: 'Linux' + arch: 'x86_64' + toolkit: 'cuda-13.0' + - os: 'Windows' + arch: 'x86_64' + toolkit: 'cpu' + - os: 'macOS' + arch: 'aarch64' + toolkit: 'metal' + runs-on: |- + ${{ case(matrix.os == 'Windows', case(matrix.arch == 'aarch64', 'windows-11-arm', + 'windows-2022'), + matrix.os == 'macOS', fromJson('["self-hosted","macos"]'), + case(matrix.arch == 'x86_64' && startsWith(matrix.toolkit, 'cuda'), 'gpu-t4-4-core', + matrix.arch == 'aarch64', 'ubuntu-22.04-arm', + 'ubuntu-22.04')) + }} steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/setup-linux + - uses: actions/checkout@v7 + - uses: ./.github/actions/setup with: toolkit: ${{ matrix.toolkit }} use-ccache: false - - name: Build Python package - uses: ./.github/actions/build-cuda-release + - uses: ./.github/actions/test-wheel with: - arch: ${{ matrix.arch }} - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - overwrite: true - name: mlx-${{ matrix.toolkit }}-${{ matrix.arch }} - path: wheelhouse/mlx_cuda_*.whl - if-no-files-found: error + toolkit: ${{ matrix.toolkit }} - pypi-publish: - name: Upload release to PyPI + pypi-publish-frontend: + name: Publish mlx to PyPI runs-on: ubuntu-latest - needs: [build_linux_release, build_mac_release] + needs: [test_wheel] permissions: id-token: write environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} + name: *pypi-env url: https://pypi.org/p/mlx steps: - - uses: actions/download-artifact@v7 - with: - pattern: linux-wheels-* - merge-multiple: true - path: dist - - uses: actions/download-artifact@v7 + - uses: actions/download-artifact@v8 with: - pattern: mac-wheels-* + pattern: frontend-* merge-multiple: true path: dist - name: Display structure of downloaded files run: du -ah dist - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} + if: *publish uses: pypa/gh-action-pypi-publish@release/v1 with: repository-url: https://upload.pypi.org/legacy/ pypi-publish-cuda: - name: Upload CUDA release to PyPI + name: Publish mlx-cuda to PyPI runs-on: ubuntu-latest - needs: [build_cuda_release] + needs: [test_wheel] permissions: id-token: write environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} + name: *pypi-env url: https://pypi.org/p/mlx-cuda steps: - - uses: actions/download-artifact@v7 + - uses: actions/download-artifact@v8 with: - pattern: mlx-cuda-* + pattern: backend-cuda* merge-multiple: true path: dist - name: Display structure of downloaded files run: du -ah dist - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} + if: *publish uses: pypa/gh-action-pypi-publish@release/v1 with: repository-url: https://upload.pypi.org/legacy/ pypi-publish-cpu: - name: Upload CPU release to PyPI + name: Publish mlx-cpu to PyPI runs-on: ubuntu-latest - needs: [build_linux_release] + needs: [test_wheel] permissions: id-token: write environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} + name: *pypi-env url: https://pypi.org/p/mlx-cpu steps: - - uses: actions/download-artifact@v7 + - uses: actions/download-artifact@v8 with: - pattern: mlx-cpu-* + pattern: backend-cpu-* merge-multiple: true path: dist - name: Display structure of downloaded files run: du -ah dist - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} + if: *publish uses: pypa/gh-action-pypi-publish@release/v1 with: repository-url: https://upload.pypi.org/legacy/ pypi-publish-metal: - name: Upload Metal release to PyPI + name: Publish mlx-metal to PyPI runs-on: ubuntu-latest - needs: [build_mac_release] + needs: [test_wheel] permissions: id-token: write environment: - name: ${{ inputs.dry_run && 'dry-run' || 'pypi' }} + name: *pypi-env url: https://pypi.org/p/mlx-metal steps: - - uses: actions/download-artifact@v7 + - uses: actions/download-artifact@v8 with: - name: mlx-metal + pattern: backend-metal-* path: dist - name: Display structure of downloaded files run: du -ah dist - name: Publish package distributions to PyPI - if: ${{ !inputs.dry_run }} + if: *publish uses: pypa/gh-action-pypi-publish@release/v1 with: repository-url: https://upload.pypi.org/legacy/ diff --git a/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml b/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml new file mode 100644 index 000000000..b9b71eee8 --- /dev/null +++ b/Source/Cmlx/mlx/.github/workflows/update_bypass_list.yml @@ -0,0 +1,24 @@ +name: 'Update bypass list' +description: 'Update interaction-limits bypass list with eligible contributors' + +on: + workflow_dispatch: + pull_request_target: + types: + - closed + schedule: + - cron: 33 6 * * * + +permissions: + contents: write + +jobs: + update_bypass_list: + name: 'Update bypass list' + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v7 + - env: + # Can only use PAT with permissions: admin:org_hook, public_repo + GITHUB_TOKEN: ${{ secrets.ZCBENZ_TOKEN_UPDATE_BYPASS_LIST }} + run: node .github/scripts/update-bypass-list.js ml-explore mlx diff --git a/Source/Cmlx/mlx/.gitignore b/Source/Cmlx/mlx/.gitignore index 1daaa46d1..161d8d67e 100644 --- a/Source/Cmlx/mlx/.gitignore +++ b/Source/Cmlx/mlx/.gitignore @@ -11,11 +11,11 @@ __pycache__/ *.metallib # Distribution / packaging +/build/ python/mlx/core python/mlx/share python/mlx/include .Python -build/ develop-eggs/ dist/ downloads/ diff --git a/Source/Cmlx/mlx/.pre-commit-config.yaml b/Source/Cmlx/mlx/.pre-commit-config.yaml index 0345848e9..859d0f660 100644 --- a/Source/Cmlx/mlx/.pre-commit-config.yaml +++ b/Source/Cmlx/mlx/.pre-commit-config.yaml @@ -9,6 +9,7 @@ repos: rev: v21.1.8 hooks: - id: clang-format + files: \.(h|cpp)$ # Using this mirror lets us use mypyc-compiled black, which is about 2x faster - repo: https://github.com/psf/black-pre-commit-mirror rev: 26.1.0 diff --git a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md index 186908f09..7005396e7 100644 --- a/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md +++ b/Source/Cmlx/mlx/ACKNOWLEDGMENTS.md @@ -20,6 +20,8 @@ MLX was developed with contributions from the following individuals: - Paul Paczuski: Improved stability of BCE loss calculation - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops. - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. +- katlun-lgtm: Added `reflect` and `symmetric` padding modes. +- Erwin Zhang: Added `searchsorted`. Fixed the ring backend hanging when a peer disconnects. Improved JACCL error reporting. diff --git a/Source/Cmlx/mlx/CMakeLists.txt b/Source/Cmlx/mlx/CMakeLists.txt index 041a476cf..feb7ce5eb 100644 --- a/Source/Cmlx/mlx/CMakeLists.txt +++ b/Source/Cmlx/mlx/CMakeLists.txt @@ -55,7 +55,7 @@ message( "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}" ) -if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin") +if(CMAKE_SYSTEM_NAME STREQUAL "Darwin") if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64") if(NOT MLX_ENABLE_X64_MAC) message( @@ -70,6 +70,8 @@ if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin") message(WARNING "Building for x86_64 arch is not officially supported.") endif() endif() +elseif(CMAKE_SYSTEM_NAME STREQUAL "iOS") + # Enable the Metal backend by default on iOS; the caller selects the SDK. else() set(MLX_BUILD_METAL OFF) endif() @@ -149,6 +151,7 @@ cmake_policy(SET CMP0135 NEW) add_library(mlx) +target_compile_features(mlx INTERFACE cxx_std_20) target_compile_options(mlx PUBLIC ${SANITIZER_COMPILE_FLAGS}) target_link_options(mlx PUBLIC ${SANITIZER_LINK_FLAGS}) @@ -156,6 +159,10 @@ if(MLX_BUILD_CUDA) enable_language(CUDA) find_package(CUDAToolkit REQUIRED) find_package(CUDNN REQUIRED) + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.1" AND CUDAToolkit_VERSION + VERSION_LESS "13.2") + message(FATAL_ERROR "CUDA Toolkit 13.1 is not supported.") + endif() endif() if(MLX_BUILD_METAL) @@ -174,33 +181,52 @@ if(MLX_BUILD_METAL) add_compile_definitions(MLX_METAL_DEBUG) endif() - # Throw an error if xcrun not found - execute_process( - COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version" - OUTPUT_VARIABLE MACOS_SDK_VERSION - OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) - - if(${MACOS_SDK_VERSION} LESS 14.0) - message( - FATAL_ERROR - "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON") - endif() - message(STATUS "Building with macOS SDK version ${MACOS_SDK_VERSION}") - - set(METAL_CPP_URL - https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip) + # Probe the macOS SDK only on Darwin. On iOS the caller selects the SDK, so + # version detection is deferred and the baseline Metal feature set is used. + if(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + # Get the host macOS version. + execute_process( + COMMAND sw_vers -productVersion + OUTPUT_VARIABLE MACOS_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) + message(STATUS "Building on macOS ${MACOS_VERSION}") + + # Get macOS SDK version. + execute_process( + COMMAND xcrun -sdk macosx --show-sdk-version + OUTPUT_VARIABLE MACOS_SDK_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) + if(MACOS_SDK_VERSION VERSION_LESS 14.0) + message( + FATAL_ERROR + "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON") + endif() + message(STATUS "Building with macOS ${MACOS_SDK_VERSION} SDK") - if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") - if(${CMAKE_OSX_DEPLOYMENT_TARGET} LESS 14.0) + # Build for the host system if CMAKE_OSX_DEPLOYMENT_TARGET is not set. + if(CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") + set(CMAKE_OSX_DEPLOYMENT_TARGET "${MACOS_VERSION}") + endif() + if(CMAKE_OSX_DEPLOYMENT_TARGET VERSION_LESS 14.0) message(FATAL_ERROR "MLX requires macOS >= 14.0") endif() + message(STATUS "Building for macOS ${CMAKE_OSX_DEPLOYMENT_TARGET}") + + # Get Metal version. set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") + execute_process( + COMMAND + zsh "-c" + "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1" + OUTPUT_VARIABLE MLX_METAL_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY) + message(STATUS "Metal version: ${MLX_METAL_VERSION}") + else() + set(MLX_METAL_VERSION 0) endif() - execute_process( - COMMAND - zsh "-c" - "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'" - OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY) + + set(METAL_CPP_URL + https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip) FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL}) FetchContent_MakeAvailable(metal_cpp) target_include_directories( @@ -236,7 +262,8 @@ if(WIN32) FetchContent_MakeAvailable(dlfcn-win32) endblock() target_include_directories(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src") - target_link_libraries(mlx PRIVATE dl) + target_sources(mlx PRIVATE "${dlfcn-win32_SOURCE_DIR}/src/dlfcn.c") + target_link_libraries(mlx PRIVATE Psapi.lib) endif() if(MLX_BUILD_CPU) @@ -254,19 +281,34 @@ if(MLX_BUILD_CPU) add_compile_definitions(MLX_USE_ACCELERATE) add_compile_definitions(ACCELERATE_NEW_LAPACK) elseif(WIN32) + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "ARM64") + set(OPENBLAS_ZIP "OpenBLAS-0.3.33-woa64-dll.zip") + set(OPENBLAS_LIB "openblas") + set(OPENBLAS_INC "include/openblas") + else() + set(OPENBLAS_ZIP "OpenBLAS-0.3.33-x64.zip") + set(OPENBLAS_LIB "libopenblas") + set(OPENBLAS_INC "include") + endif() # Download and link prebuilt binaries of OpenBLAS. Note that we can only # link with the dynamic library, the prebuilt binaries were built with MinGW # so static-linking would require linking with MinGW's runtime. FetchContent_Declare( openblas - URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.31/OpenBLAS-0.3.31-x64.zip" + URL "https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.33/${OPENBLAS_ZIP}" ) FetchContent_MakeAvailable(openblas) - target_link_libraries(mlx - PRIVATE "${openblas_SOURCE_DIR}/lib/libopenblas.lib") - target_include_directories(mlx PRIVATE "${openblas_SOURCE_DIR}/include") + target_include_directories(mlx + PRIVATE "${openblas_SOURCE_DIR}/${OPENBLAS_INC}") + # Make openblas importable by dependencies when built as static library. + set(OPENBLAS_IMPORT_LIBRARY + "${openblas_SOURCE_DIR}/lib/${OPENBLAS_LIB}.lib") + add_library(MLX::OpenBLAS UNKNOWN IMPORTED GLOBAL) + set_target_properties(MLX::OpenBLAS PROPERTIES IMPORTED_LOCATION + "${OPENBLAS_IMPORT_LIBRARY}") + target_link_libraries(mlx PRIVATE MLX::OpenBLAS) # Make sure the DLL file is placed in the same dir with executables. - set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/libopenblas.dll") + set(OPENBLAS_DLL_FILE "${openblas_SOURCE_DIR}/bin/${OPENBLAS_LIB}.dll") add_custom_command( TARGET mlx POST_BUILD @@ -317,6 +359,15 @@ FetchContent_MakeAvailable(json) target_include_directories( mlx PRIVATE $) +# Add standalone JACCL library (RDMA over Thunderbolt distributed backend) +if(MLX_BUILD_CPU + AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" + AND DEFINED MACOS_SDK_VERSION + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2) + add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx/distributed/jaccl/lib + ${CMAKE_BINARY_DIR}/jaccl) +endif() + add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx) target_include_directories( @@ -344,7 +395,7 @@ if(MLX_BUILD_PYTHON_BINDINGS) FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v2.10.2 + GIT_TAG v2.15.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) @@ -373,6 +424,8 @@ if(WIN32) if(MLX_BUILD_CPU) # Install OpenBLAS. install(FILES ${OPENBLAS_DLL_FILE} TYPE BIN) + install(FILES ${OPENBLAS_IMPORT_LIBRARY} + DESTINATION ${CMAKE_INSTALL_LIBDIR}) endif() endif() diff --git a/Source/Cmlx/mlx/CONTRIBUTING.md b/Source/Cmlx/mlx/CONTRIBUTING.md index fddb2a974..eaccfec88 100644 --- a/Source/Cmlx/mlx/CONTRIBUTING.md +++ b/Source/Cmlx/mlx/CONTRIBUTING.md @@ -3,29 +3,30 @@ We want to make contributing to this project as easy and transparent as possible. -## Pull Requests - -1. Fork and submit pull requests to the repo. -2. If you've added code that should be tested, add tests. -3. If a change is likely to impact efficiency, run some of the benchmarks before - and after the change. Examples of benchmarks can be found in `benchmarks/python/`. -4. If you've changed APIs, update the documentation. -5. Every PR should have passing tests and at least one review. -6. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`. - This should install hooks for running `black` and `clang-format` to ensure - consistent style for C++ and python code. +## AI Usage Policy - You can also run the formatters manually as follows: +AI-generated code is allowed. What is not allowed is submitting code you do not +understand. You are 100% responsible for every line, however it was produced, +and must explicitly disclose the manner in which AI was employed. - ```shell - clang-format -i file.cpp - ``` +It is strictly prohibited to use AI to write your posts for you (bug reports, +feature requests, pull request descriptions, Github discussions, responding to +humans, ...). - ```shell - black file.py - ``` +## Pull Requests - or run `pre-commit run --all-files` to check all files in the repo. +- Make sure new code is covered by tests. Add new tests if not, and confirm + the new tests fail in the main branch. +- If performance may be impacted, run benchmarks for both the main branch and + the pull request. +- When providing benchmarking results, include scripts and reproduction steps. +- Format the code with `uvx pre-commit run --all` before submitting a pull + request. You can also install git hooks to run it automatically: + + ```shell + pip install pre-commit + pre-commit install + ``` ## Issues diff --git a/Source/Cmlx/mlx/MANIFEST.in b/Source/Cmlx/mlx/MANIFEST.in index d0daeb7ae..632fae78c 100644 --- a/Source/Cmlx/mlx/MANIFEST.in +++ b/Source/Cmlx/mlx/MANIFEST.in @@ -1,6 +1,6 @@ include CMakeLists.txt include mlx.pc.in -recursive-include mlx/ * +recursive-include mlx * include cmake/* include python/src/* include python/mlx/py.typed # support type hinting as in PEP-561 diff --git a/Source/Cmlx/mlx/README.md b/Source/Cmlx/mlx/README.md index b02061dd4..fb30a24f2 100644 --- a/Source/Cmlx/mlx/README.md +++ b/Source/Cmlx/mlx/README.md @@ -113,9 +113,7 @@ BibTex entry: ```text @software{mlx2023, author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert}, -{% raw %} title = {{MLX}: Efficient and flexible machine learning on Apple silicon}, -{% endraw %} url = {https://github.com/ml-explore}, version = {0.0}, year = {2023}, diff --git a/Source/Cmlx/mlx/benchmarks/cpp/irregular_strides.cpp b/Source/Cmlx/mlx/benchmarks/cpp/irregular_strides.cpp index cc4e975c9..6da1e82af 100644 --- a/Source/Cmlx/mlx/benchmarks/cpp/irregular_strides.cpp +++ b/Source/Cmlx/mlx/benchmarks/cpp/irregular_strides.cpp @@ -9,6 +9,12 @@ namespace mx = mlx::core; +// mx::astype is overloaded, so it cannot be passed directly to the timing +// helpers. Wrap the three argument form instead. +auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); +}; + void time_irregular_binary_ops_1D() { auto device = mx::default_device(); int size = 1000000; @@ -164,7 +170,7 @@ void time_irregular_astype_1D() { int step = 2; auto a = mx::random::uniform({size}); a = slice(a, {0}, {size}, {step}); - TIMEM("1D strided", mx::astype, a, mx::int32, device); + TIMEM("1D strided", astype, a, mx::int32, device); } void time_irregular_astype_2D() { @@ -173,16 +179,16 @@ void time_irregular_astype_2D() { mx::Shape shape = {size, size}; auto a = mx::random::uniform(shape); - TIMEM("2D regular", mx::astype, a, mx::int32, device); + TIMEM("2D regular", astype, a, mx::int32, device); a = mx::transpose(a); - TIMEM("2D mx::transpose", mx::astype, a, mx::int32, device); + TIMEM("2D mx::transpose", astype, a, mx::int32, device); a = mx::broadcast_to(mx::random::uniform({size}), shape); - TIMEM("2D broadcast dim 0", mx::astype, a, mx::int32, device); + TIMEM("2D broadcast dim 0", astype, a, mx::int32, device); a = mx::broadcast_to(mx::random::uniform({size, 1}), shape); - TIMEM("2D broadcast dim 1", mx::astype, a, mx::int32, device); + TIMEM("2D broadcast dim 1", astype, a, mx::int32, device); } int main(int argc, char** argv) { diff --git a/Source/Cmlx/mlx/benchmarks/cpp/single_ops.cpp b/Source/Cmlx/mlx/benchmarks/cpp/single_ops.cpp index 1f93a78d7..05578269d 100644 --- a/Source/Cmlx/mlx/benchmarks/cpp/single_ops.cpp +++ b/Source/Cmlx/mlx/benchmarks/cpp/single_ops.cpp @@ -5,6 +5,12 @@ namespace mx = mlx::core; +// mx::astype is overloaded, so it cannot be passed directly to the timing +// helpers. Wrap the three argument form instead. +auto astype = [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); +}; + void time_creation_ops() { int M = 2000; int N = 500; @@ -28,18 +34,18 @@ void time_type_conversions() { auto a = mx::zeros(shape, mx::float32); mx::eval(a); - TIMEM("mx::float32 to mx::int32", mx::astype, a, mx::int32, device); - TIMEM("mx::float32 to mx::uint32", mx::astype, a, mx::uint32, device); + TIMEM("mx::float32 to mx::int32", astype, a, mx::int32, device); + TIMEM("mx::float32 to mx::uint32", astype, a, mx::uint32, device); a = mx::zeros(shape, mx::int32); mx::eval(a); - TIMEM("mx::int32 to mx::float32", mx::astype, a, mx::float32, device); + TIMEM("mx::int32 to mx::float32", astype, a, mx::float32, device); a = mx::zeros(shape, mx::bool_); mx::eval(a); - TIMEM("bool to mx::float32", mx::astype, a, mx::float32, device); - TIMEM("bool to mx::int32", mx::astype, a, mx::int32, device); - TIMEM("bool to mx::uint32", mx::astype, a, mx::uint32, device); + TIMEM("bool to mx::float32", astype, a, mx::float32, device); + TIMEM("bool to mx::int32", astype, a, mx::int32, device); + TIMEM("bool to mx::uint32", astype, a, mx::uint32, device); } void time_random_generation() { diff --git a/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py b/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py index 68b4a5bd3..aa81369fe 100644 --- a/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py +++ b/Source/Cmlx/mlx/benchmarks/python/comparative/compare.py @@ -4,6 +4,7 @@ import argparse import re +import sys from pathlib import Path from subprocess import run @@ -22,15 +23,15 @@ def run_or_raise(*args, **kwargs): def compare(args): - t_mlx = run_or_raise(["python", BENCH_MLX] + args) - t_torch = run_or_raise(["python", BENCH_TORCH] + args) + t_mlx = run_or_raise([sys.executable, BENCH_MLX] + args) + t_torch = run_or_raise([sys.executable, BENCH_TORCH] + args) print((t_torch - t_mlx) / t_torch, " ".join(args), sep="\t") def compare_mlx_dtypes(args, dt1, dt2): - t_mlx_dt1 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt1]) - t_mlx_dt2 = run_or_raise(["python", BENCH_MLX] + args + ["--dtype", dt2]) + t_mlx_dt1 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt1]) + t_mlx_dt2 = run_or_raise([sys.executable, BENCH_MLX] + args + ["--dtype", dt2]) print((t_mlx_dt2 - t_mlx_dt1) / t_mlx_dt2, " ".join(args), sep="\t") diff --git a/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py b/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py index 17c06d57d..c05db9199 100644 --- a/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/gather_qmm_bench.py @@ -80,5 +80,26 @@ def equivalent_matmul(x, w1, w2): time_fn(equivalent_matmul, x, w1, w2) +def time_gather_qmm_short_runs(): + # Many experts and few tokens, so each expert gets N * I / E = 16 rows. + N, E, I = 512, 256, 8 + x = mx.random.normal((N, 1, 1, D)) / 1024**0.5 + w1 = mx.random.normal((E, M, D)) / 1024**0.5 + w2 = mx.random.normal((E, D, M)) / 1024**0.5 + w1 = mx.quantize(w1) + w2 = mx.quantize(w2) + indices = (mx.random.uniform(shape=(N, I)) * E).astype(mx.uint32) + mx.eval(x, w1, w2, indices) + + def gather_mm(x, w1, w2, indices): + x, idx, inv_order = gather_sort(x, indices) + x = mx.gather_qmm(x, *w1, transpose=True, rhs_indices=idx, sorted_indices=True) + x = mx.gather_qmm(x, *w2, transpose=True, rhs_indices=idx, sorted_indices=True) + return scatter_unsort(x, inv_order, indices.shape) + + time_fn(gather_mm, x, w1, w2, indices) + + if __name__ == "__main__": time_gather_qmm() + time_gather_qmm_short_runs() diff --git a/Source/Cmlx/mlx/benchmarks/python/masked_scatter.py b/Source/Cmlx/mlx/benchmarks/python/masked_scatter.py index 71857c543..e1c84ee6b 100644 --- a/Source/Cmlx/mlx/benchmarks/python/masked_scatter.py +++ b/Source/Cmlx/mlx/benchmarks/python/masked_scatter.py @@ -1,5 +1,6 @@ import math import os +import platform import subprocess import time from copy import copy @@ -17,9 +18,6 @@ if not os.path.isdir(RESULTS_DIR): os.mkdir(RESULTS_DIR) -DEVICE_NAME = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"]) -DEVICE_NAME = DEVICE_NAME.decode("utf-8").strip("\n") - TORCH_DEVICE = torch.device( "mps" if torch.backends.mps.is_available() @@ -27,11 +25,36 @@ ) +def get_device_name(): + if TORCH_DEVICE.type == "cuda": + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + stderr=subprocess.DEVNULL, + ) + return out.decode("utf-8").splitlines()[0].strip() + except Exception: + return "CUDA_GPU" + if TORCH_DEVICE.type == "mps": + try: + out = subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.brand_string"], + stderr=subprocess.DEVNULL, + ) + return out.decode("utf-8").strip() + except Exception: + return "Apple_Silicon" + return platform.processor() or platform.machine() or "CPU" + + +DEVICE_NAME = get_device_name() + + N_WARMUP = 5 N_ITER_BENCH = 50 N_ITER_FUNC = 20 -VECTOR_LENGTHS = [4096 * (2**i) for i in range(10)] +VECTOR_LENGTHS = [4096 * (2**i) for i in range(12)] MASK_DENSITIES = [0.01, 0.1, 0.25, 0.5] D_TYPES = ("float32", "float16") @@ -202,9 +225,10 @@ def main(): ) output_path = os.path.join( RESULTS_DIR, - f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.pdf", + f"{DEVICE_NAME.replace(' ', '_')}_masked_scatter_{dtype}.png", ) fig.savefig(output_path) + print(f"Saved benchmark image: {output_path}") plt.close(fig) diff --git a/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py b/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py index 5eb789de0..7dfc7e0d1 100644 --- a/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/sdpa_bench.py @@ -176,6 +176,17 @@ def get_gflop_count(B, M, N, K): ( 1, 1024, 1024, 64, 32, 8), ( 1, 2048, 2048, 64, 32, 8), ( 1, 4096, 4096, 64, 32, 8), + ( 1, 4096, 5000, 64, 32, 8), + ( 1, 2048, 32121, 64, 32, 8), + ) + + shapes_72 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 72, 32, 8), + ( 1, 2048, 2048, 72, 32, 8), + ( 1, 4096, 4096, 72, 32, 8), + ( 1, 4096, 5000, 72, 32, 8), + ( 1, 2048, 32121, 72, 32, 8), ) shapes_80 = ( @@ -183,6 +194,17 @@ def get_gflop_count(B, M, N, K): ( 1, 1024, 1024, 80, 32, 8), ( 1, 2048, 2048, 80, 32, 8), ( 1, 4096, 4096, 80, 32, 8), + ( 1, 4096, 5000, 80, 32, 8), + ( 1, 2048, 32121, 80, 32, 8), + ) + + shapes_96 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 96, 32, 8), + ( 1, 2048, 2048, 96, 32, 8), + ( 1, 4096, 4096, 96, 32, 8), + ( 1, 4096, 5000, 96, 32, 8), + ( 1, 2048, 32121, 96, 32, 8), ) shapes_128 = ( @@ -190,10 +212,21 @@ def get_gflop_count(B, M, N, K): ( 1, 1024, 1024, 128, 32, 8), ( 1, 2048, 2048, 128, 32, 8), ( 1, 4096, 4096, 128, 32, 8), + ( 1, 4096, 5000, 128, 32, 8), + ( 1, 2048, 32121, 128, 32, 8), + ) + + shapes_256 = ( + # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 1024, 1024, 256, 24, 4), + ( 1, 2048, 2048, 256, 24, 4), + ( 1, 4096, 4096, 256, 24, 4), + ( 1, 4096, 5000, 256, 24, 4), + ( 1, 2048, 32121, 256, 24, 4), ) # fmt: on - shapes = shapes_64 + shapes_80 + shapes_128 + shapes = shapes_64 + shapes_72 + shapes_80 + shapes_96 + shapes_128 + shapes_256 masks = [None, "bool", "causal"] diff --git a/Source/Cmlx/mlx/benchmarks/python/synchronize_bench.py b/Source/Cmlx/mlx/benchmarks/python/synchronize_bench.py index db006d70b..2007ef12c 100644 --- a/Source/Cmlx/mlx/benchmarks/python/synchronize_bench.py +++ b/Source/Cmlx/mlx/benchmarks/python/synchronize_bench.py @@ -42,7 +42,7 @@ def all_gather_benchmark(): def fn(x): for _ in range(its_per_eval): - x = mx.distributed.all_gather(x)[0] + x = mx.distributed.all_gather(x)[: a.shape[0]] return x ms = timeit(fn, a) / its_per_eval diff --git a/Source/Cmlx/mlx/docs/src/conf.py b/Source/Cmlx/mlx/docs/src/conf.py index 446d95cd2..b95a3fe74 100644 --- a/Source/Cmlx/mlx/docs/src/conf.py +++ b/Source/Cmlx/mlx/docs/src/conf.py @@ -28,7 +28,10 @@ python_use_unqualified_type_names = True autosummary_generate = True -autosummary_filename_map = {"mlx.core.Stream": "stream_class"} +autosummary_filename_map = { + "mlx.core.Stream": "stream_class", + "mlx.core.PrintOptions": "printoptions_class", +} intersphinx_mapping = { "python": ("https://docs.python.org/3", None), diff --git a/Source/Cmlx/mlx/docs/src/dev/custom_metal_kernels.rst b/Source/Cmlx/mlx/docs/src/dev/custom_metal_kernels.rst index 4c4ce65ae..6a9471fd5 100644 --- a/Source/Cmlx/mlx/docs/src/dev/custom_metal_kernels.rst +++ b/Source/Cmlx/mlx/docs/src/dev/custom_metal_kernels.rst @@ -62,7 +62,7 @@ The full function signature will be generated using: so we add ``device float16_t* out``. * Template parameters passed using ``template`` In the above, ``template=[("T", mx.float32)]`` adds a template of ``template `` to the function - and instantiates the template with ``custom_kernel_myexp_float``. + and instantiates the template with ``custom_kernel_myexp_float_float16_t_float16_t``. Template parameters can be ``mx.core.Dtype``, ``int`` or ``bool``. * Metal attributes used in ``source`` such as ``[[thread_position_in_grid]]`` These will be added as function arguments. @@ -73,7 +73,7 @@ Putting this all together, the generated function signature for ``myexp`` is as .. code-block:: cpp template - [[kernel]] void custom_kernel_myexp_float( + [[kernel]] void custom_kernel_myexp_float_float16_t_float16_t( const device float16_t* inp [[buffer(0)]], device float16_t* out [[buffer(1)]], uint3 thread_position_in_grid [[thread_position_in_grid]]) { @@ -84,7 +84,7 @@ Putting this all together, the generated function signature for ``myexp`` is as } - template [[host_name("custom_kernel_myexp_float")]] [[kernel]] decltype(custom_kernel_myexp_float) custom_kernel_myexp_float; + template [[host_name("custom_kernel_myexp_float_float16_t_float16_t")]] [[kernel]] decltype(custom_kernel_myexp_float_float16_t_float16_t) custom_kernel_myexp_float_float16_t_float16_t; Note: ``grid`` and ``threadgroup`` are parameters to the Metal `dispatchThreads `_ @@ -92,8 +92,31 @@ function. This means we will launch ``mx.prod(grid)`` threads, subdivided into ``threadgroup`` size threadgroups. For optimal performance, each thread group dimension should be less than or equal to the corresponding grid dimension. -Passing ``verbose=True`` to :func:`ast.metal_kernel.__call__` will print the -generated code for debugging purposes. +Passing ``verbose=True`` when calling the kernel returned by +:func:`fast.metal_kernel` will print the generated code for debugging purposes. + +Math Mode +--------- + +By default :func:`fast.metal_kernel` compiles kernels with +``compile_options={"math_mode": "safe"}`` so special values follow IEEE +behavior, for example ``exp(-inf) == 0``. This is important for kernels such as +masked softmax where causal or sliding-window masks depend on exponentiating +``-inf``. + +If your kernel does not rely on these edge cases, you can opt in to less strict +math with ``compile_options={"math_mode": "relaxed"}`` or +``compile_options={"math_mode": "fast"}``: + +.. code-block:: python + + kernel = mx.fast.metal_kernel( + name="my_kernel", + input_names=["x"], + output_names=["y"], + source=source, + compile_options={"math_mode": "relaxed"}, + ) Using Shape/Strides ------------------- diff --git a/Source/Cmlx/mlx/docs/src/dev/extensions.rst b/Source/Cmlx/mlx/docs/src/dev/extensions.rst index aee1b10b1..89c8991e5 100644 --- a/Source/Cmlx/mlx/docs/src/dev/extensions.rst +++ b/Source/Cmlx/mlx/docs/src/dev/extensions.rst @@ -430,7 +430,7 @@ below. auto kernel = d.get_kernel(kname, lib); // Prepare to encode kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = mx::metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); // Kernel parameters are registered with buffer indices corresponding to @@ -474,7 +474,7 @@ We can now call the :meth:`axpby` operation on both the CPU and the GPU! A few things to note about MLX and Metal before moving on. MLX keeps track of the active ``command_buffer`` and the ``MTLCommandBuffer`` to which it is -associated. We rely on :meth:`d.get_command_encoder` to give us the active +associated. We rely on :meth:`metal::get_command_encoder` to give us the active metal compute command encoder instead of building a new one and calling :meth:`compute_encoder->end_encoding` at the end. MLX adds kernels (compute pipelines) to the active command buffer until some specified limit is hit or diff --git a/Source/Cmlx/mlx/docs/src/examples/data_parallelism.rst b/Source/Cmlx/mlx/docs/src/examples/data_parallelism.rst index b1592611d..9fb4ceba5 100644 --- a/Source/Cmlx/mlx/docs/src/examples/data_parallelism.rst +++ b/Source/Cmlx/mlx/docs/src/examples/data_parallelism.rst @@ -34,7 +34,8 @@ dataset, and optimizer initialization. mx.eval(loss, model.parameters()) All we have to do to average the gradients across machines is perform an -:func:`all_sum` and divide by the size of the :class:`Group`. Namely we +:func:`mlx.core.distributed.all_sum` and divide by the size of the +:class:`mlx.core.distributed.Group`. Namely we have to :func:`mlx.utils.tree_map` the gradients with following function. .. code:: python diff --git a/Source/Cmlx/mlx/docs/src/examples/llama-inference.rst b/Source/Cmlx/mlx/docs/src/examples/llama-inference.rst index 7e06895e3..cd6e0b561 100644 --- a/Source/Cmlx/mlx/docs/src/examples/llama-inference.rst +++ b/Source/Cmlx/mlx/docs/src/examples/llama-inference.rst @@ -322,14 +322,14 @@ several unnecessary copies from disk to numpy and then from numpy to MLX. It will be replaced in the future with direct loading to MLX. You can download the full example code in `mlx-examples`_. Assuming, the -existence of ``weights.pth`` and ``tokenizer.model`` in the current working -directory we can play around with our inference script as follows (the timings -are representative of an M1 Ultra and the 7B parameter Llama model): +existence of the PyTorch Llama weights in ``llama-7B/`` we can play around +with our inference script as follows (the timings are representative of an M1 +Ultra and the 7B parameter Llama model): .. code-block:: bash - $ python convert.py weights.pth llama-7B.mlx.npz - $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely' + $ python convert.py --torch-path llama-7B/ + $ python llama.py --prompt 'Call me Ishmael. Some years ago never mind how long precisely' [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ @@ -347,7 +347,7 @@ time as well as the prompt processing time remains almost constant. .. code-block:: bash - $ python llama.py llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' + $ python llama.py --prompt 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' [INFO] Loading model from disk: 5.247 s Press enter to start generation ------ @@ -355,7 +355,7 @@ time as well as the prompt processing time remains almost constant. ------ [INFO] Prompt processing: 0.579 s [INFO] Full generation: 4.690 s - $ python llama.py --num-tokens 500 llama-7B.mlx.npz tokenizer.model 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' + $ python llama.py --max-tokens 500 --prompt 'Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' [INFO] Loading model from disk: 5.628 s Press enter to start generation ------ diff --git a/Source/Cmlx/mlx/docs/src/index.rst b/Source/Cmlx/mlx/docs/src/index.rst index 74c52aaa2..8b3028be7 100644 --- a/Source/Cmlx/mlx/docs/src/index.rst +++ b/Source/Cmlx/mlx/docs/src/index.rst @@ -32,17 +32,20 @@ are the CPU and GPU. install .. toctree:: - :caption: Usage + :caption: Usage :maxdepth: 1 usage/quick_start usage/lazy_evaluation usage/unified_memory usage/indexing + usage/kv_cache usage/saving_and_loading usage/function_transforms usage/compile usage/numpy + usage/precision + usage/environment_variables usage/distributed usage/using_streams usage/export @@ -78,6 +81,7 @@ are the CPU and GPU. python/optimizers python/distributed python/tree_utils + python/printoptions .. toctree:: :caption: C++ API Reference diff --git a/Source/Cmlx/mlx/docs/src/install.rst b/Source/Cmlx/mlx/docs/src/install.rst index 7b448e874..f9d02205b 100644 --- a/Source/Cmlx/mlx/docs/src/install.rst +++ b/Source/Cmlx/mlx/docs/src/install.rst @@ -15,7 +15,7 @@ silicon computer is To install from PyPI your system must meet the following requirements: -- Using an M series chip (Apple silicon) +- Using `Apple silicon `_ - Using a native Python >= 3.10 - macOS >= 14.0 @@ -83,6 +83,7 @@ Build from source Build Requirements ^^^^^^^^^^^^^^^^^^ +- ``libblas-dev``, ``liblapack-dev``, and ``liblapacke-dev`` (Linux) - A C++ compiler with C++20 support (e.g. Clang >= 15.0) - `cmake `_ -- version 3.25 or later, and ``make`` - Xcode >= 15.0 and macOS SDK >= 14.0 @@ -120,13 +121,13 @@ Once the development dependencies are installed, you can build faster with: .. code-block:: shell - python setup.py build_ext --inplace + python setup.py build_ext --inplace Run the tests with: .. code-block:: shell - python -m unittest discover python/tests + python python/tests/run.py C++ API ^^^^^^^ diff --git a/Source/Cmlx/mlx/docs/src/python/data_types.rst b/Source/Cmlx/mlx/docs/src/python/data_types.rst index 18abd235f..9793648d8 100644 --- a/Source/Cmlx/mlx/docs/src/python/data_types.rst +++ b/Source/Cmlx/mlx/docs/src/python/data_types.rst @@ -76,3 +76,4 @@ documentation for more information. Use :func:`issubdtype` to determine if one DtypeCategory issubdtype finfo + iinfo diff --git a/Source/Cmlx/mlx/docs/src/python/devices_and_streams.rst b/Source/Cmlx/mlx/docs/src/python/devices_and_streams.rst index 78dd3e568..7743cd5d0 100644 --- a/Source/Cmlx/mlx/docs/src/python/devices_and_streams.rst +++ b/Source/Cmlx/mlx/docs/src/python/devices_and_streams.rst @@ -10,12 +10,16 @@ Devices and Streams Device Stream + ThreadLocalStream default_device set_default_device default_stream new_stream + new_thread_local_stream + new_thread_unsafe_stream set_default_stream stream synchronize + clear_streams device_count device_info diff --git a/Source/Cmlx/mlx/docs/src/python/distributed.rst b/Source/Cmlx/mlx/docs/src/python/distributed.rst index 8b48d727e..e743a6178 100644 --- a/Source/Cmlx/mlx/docs/src/python/distributed.rst +++ b/Source/Cmlx/mlx/docs/src/python/distributed.rst @@ -15,8 +15,11 @@ made available. Group is_available init + all_max + all_min all_sum all_gather send recv recv_like + sum_scatter diff --git a/Source/Cmlx/mlx/docs/src/python/fast.rst b/Source/Cmlx/mlx/docs/src/python/fast.rst index b250dcb18..c930c7bb4 100644 --- a/Source/Cmlx/mlx/docs/src/python/fast.rst +++ b/Source/Cmlx/mlx/docs/src/python/fast.rst @@ -10,7 +10,9 @@ Fast rms_norm layer_norm + cross_entropy rope scaled_dot_product_attention metal_kernel cuda_kernel + precompiled_cuda_kernel diff --git a/Source/Cmlx/mlx/docs/src/python/fft.rst b/Source/Cmlx/mlx/docs/src/python/fft.rst index 36d9d7838..78bfe7f8f 100644 --- a/Source/Cmlx/mlx/docs/src/python/fft.rst +++ b/Source/Cmlx/mlx/docs/src/python/fft.rst @@ -20,5 +20,7 @@ FFT irfft2 rfftn irfftn + fftfreq + rfftfreq fftshift ifftshift diff --git a/Source/Cmlx/mlx/docs/src/python/linalg.rst b/Source/Cmlx/mlx/docs/src/python/linalg.rst index 495380c46..1bee822bc 100644 --- a/Source/Cmlx/mlx/docs/src/python/linalg.rst +++ b/Source/Cmlx/mlx/docs/src/python/linalg.rst @@ -14,6 +14,7 @@ Linear Algebra cholesky cholesky_inv cross + det qr svd eigvals @@ -23,5 +24,6 @@ Linear Algebra lu lu_factor pinv + slogdet solve solve_triangular diff --git a/Source/Cmlx/mlx/docs/src/python/nn/distributed.rst b/Source/Cmlx/mlx/docs/src/python/nn/distributed.rst index 07dd8e230..23998b71a 100644 --- a/Source/Cmlx/mlx/docs/src/python/nn/distributed.rst +++ b/Source/Cmlx/mlx/docs/src/python/nn/distributed.rst @@ -15,6 +15,7 @@ create sharded layers from existing :class:`Modules `. shard_linear shard_inplace + fully_shard Layers ^^^^^^ @@ -28,3 +29,4 @@ Layers ShardedToAllLinear QuantizedAllToShardedLinear QuantizedShardedToAllLinear + FullyShardedModule diff --git a/Source/Cmlx/mlx/docs/src/python/nn/functions.rst b/Source/Cmlx/mlx/docs/src/python/nn/functions.rst index d3c533d0c..68849e220 100644 --- a/Source/Cmlx/mlx/docs/src/python/nn/functions.rst +++ b/Source/Cmlx/mlx/docs/src/python/nn/functions.rst @@ -36,5 +36,6 @@ simple functions. softmin softplus softshrink + softsign step tanh diff --git a/Source/Cmlx/mlx/docs/src/python/nn/init.rst b/Source/Cmlx/mlx/docs/src/python/nn/init.rst index 610d767d4..63fb621c9 100644 --- a/Source/Cmlx/mlx/docs/src/python/nn/init.rst +++ b/Source/Cmlx/mlx/docs/src/python/nn/init.rst @@ -43,3 +43,5 @@ distribution, you can do: glorot_uniform he_normal he_uniform + sparse + orthogonal diff --git a/Source/Cmlx/mlx/docs/src/python/nn/layers.rst b/Source/Cmlx/mlx/docs/src/python/nn/layers.rst index b9544bae5..c57bddc46 100644 --- a/Source/Cmlx/mlx/docs/src/python/nn/layers.rst +++ b/Source/Cmlx/mlx/docs/src/python/nn/layers.rst @@ -15,6 +15,7 @@ Layers AvgPool2d AvgPool3d BatchNorm + Bilinear CELU Conv1d Conv2d @@ -34,6 +35,7 @@ Layers HardShrink HardTanh Hardswish + Identity InstanceNorm LayerNorm LeakyReLU @@ -47,6 +49,7 @@ Layers Mish MultiHeadAttention PReLU + QQLinear QuantizedAllToShardedLinear QuantizedEmbedding QuantizedLinear @@ -71,4 +74,8 @@ Layers Step Tanh Transformer + TransformerDecoder + TransformerDecoderLayer + TransformerEncoder + TransformerEncoderLayer Upsample diff --git a/Source/Cmlx/mlx/docs/src/python/ops.rst b/Source/Cmlx/mlx/docs/src/python/ops.rst index 6aabfc695..f55ba1b4d 100644 --- a/Source/Cmlx/mlx/docs/src/python/ops.rst +++ b/Source/Cmlx/mlx/docs/src/python/ops.rst @@ -27,19 +27,26 @@ Operations argpartition argsort array_equal + asarray as_strided + astype atleast_1d atleast_2d atleast_3d + bartlett bitwise_and bitwise_invert bitwise_or bitwise_xor + blackman block_masked_mm broadcast_arrays + broadcast_shapes broadcast_to + can_cast ceil clip + concat concatenate contiguous conj @@ -58,10 +65,13 @@ Operations cummin cumprod cumsum + count_nonzero degrees + depends dequantize diag diagonal + diff divide divmod einsum @@ -74,17 +84,24 @@ Operations expand_dims eye flatten + flip floor floor_divide full + from_dlpack + full_like + from_fp8 gather_mm gather_qmm greater greater_equal hadamard_transform + hamming + hanning identity imag inner + isdtype isfinite isclose isinf @@ -107,6 +124,7 @@ Operations logical_not logical_and logical_or + logical_xor logsumexp matmul max @@ -126,17 +144,21 @@ Operations outer partition pad + permute_dims + positive power prod put_along_axis quantize quantized_matmul + qqmm radians real reciprocal remainder repeat reshape + result_type right_shift roll round @@ -146,12 +168,14 @@ Operations savez_compressed save_gguf save_safetensors + searchsorted sigmoid sign sin sinh slice slice_update + segmented_mm softmax sort split @@ -171,12 +195,16 @@ Operations tensordot tile topk + to_fp8 trace transpose tri tril triu + trunc unflatten + unstack + vecdot var view where diff --git a/Source/Cmlx/mlx/docs/src/python/tree_utils.rst b/Source/Cmlx/mlx/docs/src/python/tree_utils.rst index 6dc60b47d..7c5937017 100644 --- a/Source/Cmlx/mlx/docs/src/python/tree_utils.rst +++ b/Source/Cmlx/mlx/docs/src/python/tree_utils.rst @@ -21,3 +21,4 @@ return python trees will be using the default python ``dict``, ``list`` and tree_map tree_map_with_path tree_reduce + tree_merge diff --git a/Source/Cmlx/mlx/docs/src/usage/compile.rst b/Source/Cmlx/mlx/docs/src/usage/compile.rst index bfe7f0966..f26ac6673 100644 --- a/Source/Cmlx/mlx/docs/src/usage/compile.rst +++ b/Source/Cmlx/mlx/docs/src/usage/compile.rst @@ -155,7 +155,7 @@ contents) inside compiled functions. For debugging, inspecting arrays can be helpful. One way to do that is to globally disable compilation using the :func:`disable_compile` function or -``MLX_DISABLE_COMPILE`` flag. For example the following is okay even though +:envvar:`MLX_DISABLE_COMPILE` flag. For example the following is okay even though ``fun`` is compiled: .. code-block:: python @@ -207,9 +207,9 @@ You have two options to deal with this. The first option is to simply return state.append(z) return mx.exp(z), state - _, state = fun(mx.array(1.0), mx.array(2.0)) - # Prints [array(3, dtype=float32)] - print(state) + _, state = fun(mx.array(1.0), mx.array(2.0)) + # Prints [array(3, dtype=float32)] + print(state) In some cases returning updated state can be pretty inconvenient. Hence, :func:`compile` has a parameter to capture implicit outputs: @@ -464,7 +464,7 @@ recompiled. x = mx.array(1.0) y = mx.array(-2.0) - # Firt call compiles the function + # First call compiles the function print(compiled_fun(x, y)) # Second call with different shapes diff --git a/Source/Cmlx/mlx/docs/src/usage/distributed.rst b/Source/Cmlx/mlx/docs/src/usage/distributed.rst index 1f271835c..431aef710 100644 --- a/Source/Cmlx/mlx/docs/src/usage/distributed.rst +++ b/Source/Cmlx/mlx/docs/src/usage/distributed.rst @@ -335,17 +335,48 @@ of a gigantic model using MLX LM. .. code-block:: - mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json \ - --env MLX_METAL_FAST_SYNCH=1 -- \ # <--- important + mlx.launch --verbose --backend jaccl --hostfile m3-ultra-jaccl.json -- \ /path/to/remote/python -m mlx_lm chat --model mlx-community/DeepSeek-R1-0528-4bit .. note:: - Defining the environment variable ``MLX_METAL_FAST_SYNCH=1`` enables a - different, faster way of synchronizing between the GPU and the CPU. It is - not specific to the JACCL backend and can be used in all cases where the CPU - and GPU need to collaborate for some computation and is pretty critical for - low-latency communication since the communication is done by the CPU. + Defining the environment variable :envvar:`MLX_METAL_FAST_SYNCH` to ``1`` + by passing ``--env MLX_METAL_FAST_SYNCH=1`` enables a different, faster way + of synchronizing between the GPU and the CPU. It is not specific to the + JACCL backend and can be used in all cases where the CPU and GPU need to + collaborate for some computation, and it matters for low-latency + communication since the communication is done by the CPU. + + It is however not reliable that can lead to deadlock and leave the GPU wedged + (see `#3142 `_), so it is off + by default and best left unset. + +Custom side channel +^^^^^^^^^^^^^^^^^^^ + +During initialization JACCL exchanges RDMA connection metadata between ranks +over a side channel. By default this is a simple TCP star all-gather, but you +can provide your own side-channel all-gather via the +``all_gather_factory`` argument of :func:`mlx.core.distributed.init` when using +``backend='jaccl'``. + +The factory is called once per rank with ``(rank, size)`` and must return a +callable with signature ``f(src: bytes, n_bytes: int) -> bytes``. The returned +bytes must have length ``size * n_bytes`` and contain the inputs from all ranks +concatenated in rank order. + +.. code-block:: python + + def make_side_channel(rank, size): + def all_gather(src: bytes, n_bytes: int) -> bytes: + # Exchange src with all ranks and return size * n_bytes bytes + ... + return all_gather + + world = mx.distributed.init( + backend="jaccl", + all_gather_factory=make_side_channel, + ) .. _nccl_section: @@ -479,10 +510,10 @@ Below we list the environment variables required to use each backend. Ring ^^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_HOSTFILE** should contain the path to a json file that contains IPs and +:envvar:`MLX_HOSTFILE` should contain the path to a json file that contains IPs and ports for each rank to listen to, something like the following: .. code-block:: json @@ -494,19 +525,19 @@ ports for each rank to listen to, something like the following: ["123.123.4.1:5000", "123.123.4.2:5000"] ] -**MLX_RING_VERBOSE** is optional and if set to 1 it enables some more logging +:envvar:`MLX_RING_VERBOSE` is optional and if set to 1 it enables some more logging from the distributed backend. JACCL ^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_JACCL_COORDINATOR** should contain the IP and port that rank 0 can listen +:envvar:`MLX_JACCL_COORDINATOR` should contain the IP and port that rank 0 can listen to all the other ranks connect to in order to establish the RDMA connections. -**MLX_IBV_DEVICES** should contain the path to a json file that contains the +:envvar:`MLX_IBV_DEVICES` should contain the path to a json file that contains the ibverbs device names that connect each node to each other node, something like the following: @@ -523,10 +554,10 @@ the following: NCCL ^^^^^ -**MLX_RANK** should contain a single 0-based integer that defines the rank of +:envvar:`MLX_RANK` should contain a single 0-based integer that defines the rank of the process. -**MLX_WORLD_SIZE** should contain the total number of processes that will be +:envvar:`MLX_WORLD_SIZE` should contain the total number of processes that will be launched. **NCCL_HOST_IP** and **NCCL_PORT** should contain the IP and port that all diff --git a/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst b/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst new file mode 100644 index 000000000..3584f4c22 --- /dev/null +++ b/Source/Cmlx/mlx/docs/src/usage/environment_variables.rst @@ -0,0 +1,184 @@ +.. _environment_variables: + +Environment Variables +===================== + +MLX uses environment variables to configure compilation, numerical precision, +backend behavior, and distributed execution. Set them before starting the +process. Many variables are read when the corresponding subsystem is first +initialized and changing them later may have no effect. + +Boolean variables use ``0`` to disable and a nonzero integer to enable unless +otherwise noted. + +General +------- + +.. envvar:: MLX_DISABLE_COMPILE + + Disable compilation globally. This variable is enabled by its presence, so + setting it to ``0`` also disables compilation. Calling + :func:`mlx.core.enable_compile` overrides it. + +.. envvar:: MLX_ENABLE_TF32 + + Allow reduced-precision ``float32`` matrix-multiplication family operations + on supported hardware. The default is ``1``. Set it to ``0`` to keep these + operations in full ``float32`` precision. See :doc:`precision`. + +Distributed +----------- + +The variables required to initialize a distributed process depend on the +backend. :doc:`distributed` describes their formats and how ``mlx.launch`` +sets them. + +.. envvar:: MLX_RANK + + The zero-based rank of the current process. This is used by the Ring, + JACCL, and NCCL backends. ``JACCL_RANK`` is accepted as a higher-priority + alias by JACCL. + +.. envvar:: MLX_HOSTFILE + + The path to the JSON host file used by the Ring backend. + +.. envvar:: MLX_RING_VERBOSE + + Enable verbose logging for the Ring backend. This variable is enabled by + its presence. + +.. envvar:: MLX_IBV_DEVICES + + The path to the JSON device-connectivity file used by JACCL. + ``JACCL_IBV_DEVICES`` is accepted as a higher-priority alias. + +.. envvar:: MLX_JACCL_COORDINATOR + + The coordinator address in ``IP:port`` form used to establish JACCL + connections. ``JACCL_COORDINATOR`` is accepted as a higher-priority alias. + +.. envvar:: MLX_JACCL_RING + + Prefer a ring topology for JACCL. This variable is enabled by its presence. + ``JACCL_RING`` is accepted as a higher-priority alias. + +.. envvar:: MLX_WORLD_SIZE + + The total number of processes in an NCCL group. + +.. envvar:: MLX_NCCL_TIMEOUT + + The timeout in milliseconds for establishing NCCL bootstrap connections. + The default is ``300000``. + +.. envvar:: MLX_MPI_LIBNAME + + Override the MPI dynamic-library name. The default is ``libmpi.dylib`` on + macOS and ``libmpi.so`` on other platforms. + +The NCCL backend also requires ``NCCL_HOST_IP`` and ``NCCL_PORT``. Setting +``NCCL_DEBUG=INFO`` enables additional logging while MLX establishes the +bootstrap connection. NCCL itself recognizes additional `NCCL environment +variables `_. +``CUDA_VISIBLE_DEVICES`` selects the local CUDA device for each process and is +handled by the CUDA runtime. + +Metal +----- + +.. envvar:: MLX_METAL_FAST_SYNCH + + Enable the faster Metal CPU/GPU synchronization path. The default is ``0``. + This requires Metal 3.2 or later (macOS 15 or later, or iOS 18 or later). + +Advanced tuning +--------------- + +These variables tune MLX implementation details. They are primarily intended +for development, diagnostics, and performance experiments. The defaults are +selected automatically for the current hardware and are appropriate for most +users. Their behavior may change as the implementation evolves. + +.. envvar:: MLX_BFS_MAX_WIDTH + + Set the breadth-first-search width limit used when constructing an + evaluation tape. The default is ``20``. + +.. envvar:: MLX_MAX_OPS_PER_BUFFER + + Override the maximum number of operations encoded in one Metal command + buffer or CUDA graph. The default depends on the device. + +.. envvar:: MLX_MAX_MB_PER_BUFFER + + Override the approximate memory limit, in megabytes, for one Metal command + buffer or CUDA graph. The default depends on the device. + +.. envvar:: MLX_METAL_GPU_ARCH + + Override the Metal GPU architecture string reported to MLX. This affects + architecture-specific kernel and scheduling choices, but does not change + the capabilities of the physical GPU. Forcing an architecture that does not + match the GPU can select incompatible kernels and produce incorrect results. + +.. envvar:: MLX_SDPA_BLOCKS + + Override the number of reduction blocks used by the Metal scaled + dot-product attention kernel. Positive values are rounded up to a multiple + of ``32``. + +CUDA +---- + +The MLX-prefixed variables in this section are advanced CUDA backend controls. + +.. envvar:: MLX_USE_CUDA_GRAPHS + + Enable CUDA graph capture and replay. The default is ``1``. + +.. envvar:: MLX_SAVE_CUDA_GRAPHS_DOT_FILE + + Use the specified value as the filename prefix when writing captured CUDA + graphs to numbered DOT files. An unset or empty value disables the output. + +.. envvar:: MLX_PTX_CACHE_DIR + + Override the directory used to cache runtime-compiled PTX. By default MLX + uses an ``mlx//ptx`` directory under the system temporary + directory. + +.. envvar:: MLX_CUDA_USE_CUDNN_SDPA + + Allow the CUDA backend to use cuDNN scaled dot-product attention when the + inputs and device are supported. The default is ``1``. + +.. envvar:: MLX_CUDA_CONV_CACHE_SIZE + + Set the CUDA convolution cache capacity. The default is ``128``. + +.. envvar:: MLX_CUDA_FFT_CACHE_SIZE + + Set the CUDA FFT plan cache capacity. The default is ``128``. + +.. envvar:: MLX_CUDA_GRAPH_CACHE_SIZE + + Set the CUDA graph cache capacity. The default is ``400``. + +.. envvar:: MLX_CUDA_SDPA_CACHE_SIZE + + Set the CUDA forward scaled dot-product attention cache capacity. The + default is ``256``. + +.. envvar:: MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE + + Set the CUDA backward scaled dot-product attention cache capacity. The + default is ``64``. + +.. envvar:: MLX_ENABLE_CACHE_THRASHING_CHECK + + Detect repeated CUDA cache misses and raise an error suggesting a larger + cache capacity. The default is ``1``. + +MLX also uses ``CUDA_HOME`` or ``CUDA_PATH`` to locate CUDA headers for runtime +kernel compilation when they cannot be found in the Python environment. diff --git a/Source/Cmlx/mlx/docs/src/usage/export.rst b/Source/Cmlx/mlx/docs/src/usage/export.rst index e908b85fc..9900ed6d2 100644 --- a/Source/Cmlx/mlx/docs/src/usage/export.rst +++ b/Source/Cmlx/mlx/docs/src/usage/export.rst @@ -109,6 +109,35 @@ keyword arguments when calling the imported function. out, = imported_fun(x, z=y) +Saving Metadata +--------------- + +You can save metadata, such as a model configuration, alongside an exported +function. The metadata is a string, so structured data can be encoded with +JSON: + +.. code-block:: python + + import json + + def fun(x, y): + return x + y + + x = mx.array(1.0) + y = mx.array(1.0) + config = {"description": "adds two arrays", "version": 1} + mx.export_function("add.mlxfn", fun, x, y, metadata=json.dumps(config)) + +Pass ``return_metadata=True`` to read the metadata back when importing: + +.. code-block:: python + + imported_fun, metadata = mx.import_function("add.mlxfn", return_metadata=True) + + # Prints: adds two arrays + print(json.loads(metadata)["description"]) + + Exporting Modules ----------------- @@ -155,6 +184,34 @@ parameters, pass them as inputs to the ``call`` wrapper: mx.export_function("model.mlxfn", call, (mx.zeros(4),), params) +Exporting with a Callback +------------------------- + +To inspect the exported graph, you can pass a callback instead of a file path +to :func:`export_function`. + +.. code-block:: python + + def fun(x): + return x.astype(mx.int32) + + def callback(args): + print(args) + + mx.export_function(callback, fun, mx.array([1.0, 2.0])) + +The argument to the callback (``args``) is a dictionary which includes a +``type`` field. The possible types are: + +* ``"inputs"``: The ordered positional inputs to the exported function +* ``"keyword_inputs"``: The keyword specified inputs to the exported function +* ``"outputs"``: The ordered outputs of the exported function +* ``"constants"``: Any graph constants +* ``"primitives"``: Inner graph nodes representating the operations + +Each type has additional fields in the ``args`` dictionary. + + Shapeless Exports ----------------- diff --git a/Source/Cmlx/mlx/docs/src/usage/indexing.rst b/Source/Cmlx/mlx/docs/src/usage/indexing.rst index 8dc01cfc8..9bce59a6f 100644 --- a/Source/Cmlx/mlx/docs/src/usage/indexing.rst +++ b/Source/Cmlx/mlx/docs/src/usage/indexing.rst @@ -28,12 +28,11 @@ For multi-dimensional arrays, the ``...`` or :obj:`Ellipsis` syntax works as in >>> arr = mx.arange(8).reshape(2, 2, 2) >>> arr[:, :, 0] - array(3, dtype=int32) array([[0, 2], - [4, 6]], dtype=int32 + [4, 6]], dtype=int32) >>> arr[..., 0] array([[0, 2], - [4, 6]], dtype=int32 + [4, 6]], dtype=int32) You can index with ``None`` to create a new axis: @@ -41,9 +40,9 @@ You can index with ``None`` to create a new axis: >>> arr = mx.arange(8) >>> arr.shape - [8] + (8,) >>> arr[None].shape - [1, 8] + (1, 8) You can also use an :obj:`array` to index another :obj:`array`: @@ -161,7 +160,7 @@ Other index types are routed through the standard scatter code. >>> updates = mx.array([5.0, 6.0]) >>> a[mask] = updates >>> a - array([5.0, 2.0, 6.0], dtype=float32) + array([5, 2, 6], dtype=float32) Scalar assignments broadcast to every ``True`` entry in ``mask``. For non-scalar assignments, ``updates`` must provide at least as many elements as there are @@ -174,8 +173,8 @@ assignments, ``updates`` must provide at least as many elements as there are [False, False, True]]) >>> a[mask] = 1.0 >>> a - array([[1.0, 0.0, 1.0], - [0.0, 0.0, 1.0]], dtype=float32) + array([[1, 0, 1], + [0, 0, 1]], dtype=float32) Boolean masks follow NumPy semantics: diff --git a/Source/Cmlx/mlx/docs/src/usage/kv_cache.rst b/Source/Cmlx/mlx/docs/src/usage/kv_cache.rst new file mode 100644 index 000000000..f24666cf4 --- /dev/null +++ b/Source/Cmlx/mlx/docs/src/usage/kv_cache.rst @@ -0,0 +1,91 @@ +.. _kv_cache: + +Writing a Fast KV Cache +======================= + +Autoregressive generation appends one position to the key and value arrays for +each generated token. The append strategy can dominate KV-cache performance. + +Avoid appending naively with :func:`concatenate`: + +.. code-block:: python + + # Avoid this + cache = mx.zeros((1, 0, d)) + for x in steps: + cache = mx.concatenate([cache, x], axis=1) + mx.eval(cache) + +Instead, preallocate fixed-size chunks and update the cache in place: + +.. code-block:: python + + chunk = 256 + cache = mx.zeros((1, chunk, d)) + offset = 0 + for x in steps: + if offset == cache.shape[1]: + cache = mx.concatenate([cache, mx.zeros((1, chunk, d))], axis=1) + cache = mx.slice_update(cache, x, mx.array(offset), (1,)) + offset += 1 + mx.eval(cache) + + keys = cache[:, :offset] + +You can also use indexed assignment, which may read more naturally: +``cache[:, offset : offset + 1, :] = x``. + +The following measurements append one position per step to 20 ``bfloat16`` +caches of shape ``[1, 4, N, 512]`` on an M4 Max: + +.. list-table:: + :widths: 20 30 30 + :header-rows: 1 + + * - Context + - Concatenate + - Preallocate + update + * - 512 + - 0.90 ms / step + - 0.24 ms / step + * - 1024 + - 1.11 ms / step + - 0.21 ms / step + * - 4096 + - 3.73 ms / step + - 0.22 ms / step + +With preallocation, step time remains nearly constant as context length grows. +Concatenation becomes progressively slower. + +Why Concatenating Is Slow +------------------------- + +Concatenation has two costs: copying data and preventing buffer reuse. + +:func:`concatenate` creates a new array and copies the existing cache at every +step. Appending ``n`` positions therefore copies on the order of ``n^2`` +elements. + +Growing the cache also prevents buffer reuse. MLX pools freed device buffers, +but reuses a buffer only for a similarly sized request. It does not split a +larger buffer for a smaller request or combine smaller buffers for a larger +one. Because a growing cache has a new size at every step, each allocation +typically comes from the driver while freed buffers remain unused. + +This allocation work occurs on the CPU. In profiles, it appears as GPU idle +time between kernels rather than as a slow kernel, which can make the model +appear to be the bottleneck. + +Preallocation avoids both costs. The cache shape remains fixed, eliminating +both repeated copies and per-step allocations. + +Choosing a Chunk Size +--------------------- + +Use a chunk size that is a multiple of 256. + +This both amortizes growth and enables the fused cuDNN attention kernel on +CUDA. For single-token attention, the key and value arrays must be slices of a +contiguous cache with a capacity that is a multiple of 256, and at least 256 +positions must be in use. Other chunk sizes silently use a slower path. diff --git a/Source/Cmlx/mlx/docs/src/usage/launching_distributed.rst b/Source/Cmlx/mlx/docs/src/usage/launching_distributed.rst index e1e3128bb..7c3a24492 100644 --- a/Source/Cmlx/mlx/docs/src/usage/launching_distributed.rst +++ b/Source/Cmlx/mlx/docs/src/usage/launching_distributed.rst @@ -160,11 +160,6 @@ host and on the same path. A good checklist to debug errors is the following: ``mlx.launch --print-python`` to see what that path is. * the script you want to run is available on all hosts at the same path -If you are launching from a node with a completely different setup than the -nodes that the program will run on, you can specify ``--no-verify-script`` so -that ``mlx.launch`` does not attempt to verify that the executable and script -exist locally before launching the distributed job. - .. _ring_specifics: Ring Specifics @@ -207,7 +202,7 @@ multi-gpu jobs. For instance .. code-block:: - mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 --no-verify-script -- ./my-job.sh + mlx.launch --backend nccl --hosts linux-1,linux-2 -n 8 -- ./my-job.sh will attempt to launch 16 processes, 8 on each node that will all run ``my-job.sh``. diff --git a/Source/Cmlx/mlx/docs/src/usage/numpy.rst b/Source/Cmlx/mlx/docs/src/usage/numpy.rst index 09b36ddbd..88b820294 100644 --- a/Source/Cmlx/mlx/docs/src/usage/numpy.rst +++ b/Source/Cmlx/mlx/docs/src/usage/numpy.rst @@ -75,25 +75,62 @@ even though no in-place operations on MLX memory are executed. PyTorch ------- -.. warning:: - - PyTorch Support for :obj:`memoryview` is experimental and can break for - multi-dimensional arrays. Casting to NumPy first is advised for now. - -PyTorch supports the buffer protocol, but it requires an explicit -:obj:`memoryview`. +PyTorch supports DLPack inputs and can import MLX arrays directly. +MLX can also import PyTorch tensors through DLPack with ``mx.asarray`` or +``mx.from_dlpack``. Use ``torch.as_tensor`` to import an MLX array with +DLPack; ``torch.tensor`` copies the data instead. Similarly, ``mx.asarray`` +can share DLPack inputs when possible, while ``mx.array`` copies: .. code-block:: python import mlx.core as mx import torch - a = mx.arange(3) - b = torch.tensor(memoryview(a)) - c = mx.array(b.numpy()) + a = mx.arange(3, dtype=mx.float32) + mx.eval(a) + + shared = torch.as_tensor(a) + copied = torch.tensor(a) + +Creating an MLX array from a CPU tensor copies the data into MLX-owned storage. +The arrays do not share memory: + +.. code-block:: python + + b = torch.arange(3) + c = mx.array(b) + + b += 10 + print(c.tolist()) # [0, 1, 2] + +Metal DLPack inputs are different. If a PyTorch MPS tensor is passed to +``mx.asarray`` or to ``mx.from_dlpack`` with ``copy=None``, MLX imports it +without a copy when the underlying Metal buffer is not private. Private Metal +buffers are copied into MLX-managed storage instead. Passing ``copy=False`` +requires zero-copy import and raises an error if a copy would be needed. +Passing ``copy=True`` asks MLX to create a new array instead of reusing the +Metal buffer. Zero-copy imports preserve the DLPack strides. ``mx.array`` also +creates a new array instead of reusing the Metal buffer. MLX arrays exported to +PyTorch with DLPack are exported without a copy on Metal. + +In particular, PyTorch 2.12 and later use shared storage for ordinary MPS +tensors on Apple silicon, while older PyTorch versions may use private storage +and require a copy on import. DLPack conversion does not synchronize pending +Metal work; synchronize or evaluate the producing framework before reading the +converted array. + +.. code-block:: python + + b = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + c = mx.asarray(b) # zero-copy if the Metal buffer can be reused + d = mx.from_dlpack(b, copy=True) # explicit copy + +.. code-block:: python -Conversion from PyTorch tensors back to arrays must be done via intermediate -NumPy arrays with ``numpy()``. + a = mx.arange(3, dtype=mx.float32) + mx.eval(a) + b = torch.as_tensor(a) # zero-copy DLPack import on Metal JAX --- diff --git a/Source/Cmlx/mlx/docs/src/usage/precision.rst b/Source/Cmlx/mlx/docs/src/usage/precision.rst new file mode 100644 index 000000000..440a1f705 --- /dev/null +++ b/Source/Cmlx/mlx/docs/src/usage/precision.rst @@ -0,0 +1,21 @@ +.. _precision: + +Numerical Precision +=================== + +By default, MLX may run ``float32`` matrix-multiplication family +operations (matmul, quantized matmul, grouped matmul, convolution and +attention) at reduced precision on hardware with dedicated +matrix-multiplication units. Inputs and outputs stay ``float32``, but +results can differ from a full-precision reference by several orders of +magnitude more than ``float32`` rounding alone would explain. + +To keep these operations in full ``float32``, set +:envvar:`MLX_ENABLE_TF32` to ``0`` when launching the process: + +.. code-block:: shell + + MLX_ENABLE_TF32=0 python my_script.py + +Which operations take the reduced-precision path, and how large the +difference is, depends on the backend and the hardware. diff --git a/Source/Cmlx/mlx/docs/src/usage/quick_start.rst b/Source/Cmlx/mlx/docs/src/usage/quick_start.rst index bc1d92ad6..8c952a6f7 100644 --- a/Source/Cmlx/mlx/docs/src/usage/quick_start.rst +++ b/Source/Cmlx/mlx/docs/src/usage/quick_start.rst @@ -14,7 +14,7 @@ Import ``mlx.core`` and make an :class:`array`: >> import mlx.core as mx >> a = mx.array([1, 2, 3, 4]) >> a.shape - [4] + (4,) >> a.dtype int32 >> b = mx.array([1.0, 2.0, 3.0, 4.0]) diff --git a/Source/Cmlx/mlx/examples/cmake_project/CMakeLists.txt b/Source/Cmlx/mlx/examples/cmake_project/CMakeLists.txt index 7a73ce1d9..a95558478 100644 --- a/Source/Cmlx/mlx/examples/cmake_project/CMakeLists.txt +++ b/Source/Cmlx/mlx/examples/cmake_project/CMakeLists.txt @@ -2,13 +2,13 @@ cmake_minimum_required(VERSION 3.27) project(example LANGUAGES CXX) -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) # Comment the following two commands only the MLX C++ library is installed and # set(MLX_ROOT "/path/to/mlx") directly if needed. find_package( - Python 3.9 + Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) execute_process( diff --git a/Source/Cmlx/mlx/examples/export/CMakeLists.txt b/Source/Cmlx/mlx/examples/export/CMakeLists.txt index aefcdf59f..abe6e7b5d 100644 --- a/Source/Cmlx/mlx/examples/export/CMakeLists.txt +++ b/Source/Cmlx/mlx/examples/export/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.27) project(import_mlx LANGUAGES CXX) -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) find_package( diff --git a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt b/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt index 0f7018720..eedc3a7da 100644 --- a/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt +++ b/Source/Cmlx/mlx/examples/extensions/CMakeLists.txt @@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.27) project(_ext LANGUAGES CXX) # ----------------------------- Setup ----------------------------- -set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_POSITION_INDEPENDENT_CODE ON) diff --git a/Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp b/Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp index 31badbbda..9ade5c483 100644 --- a/Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp +++ b/Source/Cmlx/mlx/examples/extensions/axpby/axpby.cpp @@ -192,7 +192,7 @@ void Axpby::eval_gpu( auto kernel = d.get_kernel(kname, lib); // Prepare to encode kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = mx::metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); // Kernel parameters are registered with buffer indices corresponding to diff --git a/Source/Cmlx/mlx/examples/extensions/pyproject.toml b/Source/Cmlx/mlx/examples/extensions/pyproject.toml index 405234078..560a58bc2 100644 --- a/Source/Cmlx/mlx/examples/extensions/pyproject.toml +++ b/Source/Cmlx/mlx/examples/extensions/pyproject.toml @@ -3,6 +3,6 @@ requires = [ "setuptools>=42", "cmake>=3.25", "mlx>=0.18.0", - "nanobind==2.10.2", + "nanobind==2.15.0", ] build-backend = "setuptools.build_meta" diff --git a/Source/Cmlx/mlx/examples/extensions/requirements.txt b/Source/Cmlx/mlx/examples/extensions/requirements.txt index 2fa00ac5c..917d125ee 100644 --- a/Source/Cmlx/mlx/examples/extensions/requirements.txt +++ b/Source/Cmlx/mlx/examples/extensions/requirements.txt @@ -1,4 +1,4 @@ setuptools>=42 cmake>=3.25 -mlx>=0.21.0 -nanobind==2.10.2 +mlx>=0.31.2 +nanobind==2.15.0 diff --git a/Source/Cmlx/mlx/examples/extensions/setup.py b/Source/Cmlx/mlx/examples/extensions/setup.py index 8dd1186ab..be990a02b 100644 --- a/Source/Cmlx/mlx/examples/extensions/setup.py +++ b/Source/Cmlx/mlx/examples/extensions/setup.py @@ -14,5 +14,5 @@ packages=["mlx_sample_extensions"], package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, zip_safe=False, - python_requires=">=3.8", + python_requires=">=3.10", ) diff --git a/Source/Cmlx/mlx/examples/python/distributed_data_parallel.py b/Source/Cmlx/mlx/examples/python/distributed_data_parallel.py new file mode 100644 index 000000000..6f19766dd --- /dev/null +++ b/Source/Cmlx/mlx/examples/python/distributed_data_parallel.py @@ -0,0 +1,102 @@ +# Copyright © 2026 Apple Inc. + +""" +Data parallel training of a small MLP. + +Every rank holds a different slice of the same dataset and averages the +gradients each step, so the training is equivalent to running on one process +and the final loss is the same at any number of ranks, up to the order the +floating point additions happen in: + + python examples/python/distributed_data_parallel.py + mlx.launch -n 2 python examples/python/distributed_data_parallel.py + mlx.launch -n 4 python examples/python/distributed_data_parallel.py + +The model is replicated and the batch is split, which is what you reach for +when the model fits on one machine but the data is large. +""" + +import time + +import mlx.core as mx +import mlx.nn as nn +import mlx.optimizers as optim + +num_features = 100 +num_examples = 1_000 +hidden = 64 +num_iters = 200 +lr = 0.05 + +world = mx.distributed.init() + +if num_examples % world.size() != 0: + raise ValueError( + f"Cannot split {num_examples} examples evenly over {world.size()} ranks." + ) + +# Fixed keys, so every rank draws the same dataset rather than one of its own. +w_star = mx.random.normal((num_features,), key=mx.random.key(0)) +X = mx.random.normal((num_examples, num_features), key=mx.random.key(1)) +y = X @ w_star + 1e-2 * mx.random.normal((num_examples,), key=mx.random.key(2)) + +# Keep this rank's slice and drop the rest. +examples_per_rank = num_examples // world.size() +start = world.rank() * examples_per_rank +X = X[start : start + examples_per_rank] +y = y[start : start + examples_per_rank] + + +class MLP(nn.Module): + def __init__(self, dims: int, hidden: int): + super().__init__() + self.layers = [nn.Linear(dims, hidden), nn.Linear(hidden, 1)] + + def __call__(self, x): + return self.layers[1](nn.relu(self.layers[0](x))).squeeze(-1) + + +# Seeding the global rng starts every rank from the same weights, which the +# averaged gradient then keeps in step. +mx.random.seed(0) +model = MLP(num_features, hidden) +mx.eval(model.parameters()) + +optimizer = optim.SGD(learning_rate=lr) + + +def loss_fn(model, X, y): + return 0.5 * mx.mean(mx.square(model(X) - y)) + + +loss_and_grad_fn = nn.value_and_grad(model, loss_fn) + +tic = time.perf_counter() +for _ in range(num_iters): + loss, grads = loss_and_grad_fn(model, X, y) + + # Each rank has gradients for its own slice. Averaging them gives the + # gradients of the whole dataset, which is what keeps this equivalent to + # single process training. One call handles the whole parameter tree. + grads = nn.average_gradients(grads, group=world) + + optimizer.update(model, grads) + mx.eval(model.parameters(), optimizer.state) +toc = time.perf_counter() + +# Every slice is the same size, so averaging the per rank losses gives the loss +# over the whole dataset. +loss = mx.distributed.all_sum(loss_fn(model, X, y), group=world) / world.size() + +# Only rank 0 prints the loss, but every rank has to evaluate it. Arrays are +# lazy, so leaving this to the print below would mean the other ranks never +# join the all_sum and everyone waits forever. +mx.eval(loss) + +throughput = num_iters / (toc - tic) + +if world.rank() == 0: + print( + f"Loss {loss.item():.6f}, Throughput {throughput:.2f} (it/s) " + f"over {world.size()} rank(s), {examples_per_rank} examples per rank" + ) diff --git a/Source/Cmlx/mlx/examples/python/distributed_tensor_parallel.py b/Source/Cmlx/mlx/examples/python/distributed_tensor_parallel.py new file mode 100644 index 000000000..db66c4c65 --- /dev/null +++ b/Source/Cmlx/mlx/examples/python/distributed_tensor_parallel.py @@ -0,0 +1,72 @@ +# Copyright © 2026 Apple Inc. + +""" +Tensor parallel inference for a small MLP. + +The two linear layers are sharded across ranks: the first splits its output +features so each rank computes part of the hidden state, the second splits its +input features and sums the partial results. That costs one all reduce per +block, and the sharded model returns what the full model returns: + + python examples/python/distributed_tensor_parallel.py + mlx.launch -n 2 python examples/python/distributed_tensor_parallel.py + mlx.launch -n 4 python examples/python/distributed_tensor_parallel.py + +Unlike data parallelism this splits the model rather than the batch, so it is +what you reach for when the weights are too big for one machine. +""" + +import mlx.core as mx +import mlx.nn as nn + +dims = 256 +hidden = 1024 +num_tokens = 8 + +world = mx.distributed.init() + +if hidden % world.size() != 0: + raise ValueError( + f"Cannot split {hidden} hidden features evenly over {world.size()} ranks." + ) + + +class MLP(nn.Module): + def __init__(self, dims: int, hidden: int): + super().__init__() + self.up = nn.Linear(dims, hidden) + self.down = nn.Linear(hidden, dims) + + def __call__(self, x): + return self.down(nn.silu(self.up(x))) + + +# Seeding the global rng gives every rank the same weights to shard. +mx.random.seed(0) +model = MLP(dims, hidden) +mx.eval(model.parameters()) + +x = mx.random.normal((num_tokens, dims), key=mx.random.key(1)) +expected = model(x) +mx.eval(expected) + +# Each rank keeps a slice of each weight. from_linear takes the slice belonging +# to this rank, so the layers never hold the full weight afterwards. +model.up = nn.AllToShardedLinear.from_linear(model.up, group=world) +model.down = nn.ShardedToAllLinear.from_linear(model.down, group=world) +mx.eval(model.parameters()) + +y = model(x) + +# Every rank evaluates: the down projection ends in an all reduce, so a rank +# that skipped this would leave the others waiting. +mx.eval(y) + +difference = mx.abs(y - expected).max().item() +hidden_per_rank = hidden // world.size() + +if world.rank() == 0: + print( + f"Max |sharded - full| = {difference:.3e} over {world.size()} rank(s), " + f"{hidden_per_rank} of {hidden} hidden features per rank" + ) diff --git a/Source/Cmlx/mlx/examples/python/qqmm.py b/Source/Cmlx/mlx/examples/python/qqmm.py index 5be7eae2f..dc46fa29a 100644 --- a/Source/Cmlx/mlx/examples/python/qqmm.py +++ b/Source/Cmlx/mlx/examples/python/qqmm.py @@ -33,48 +33,71 @@ def test_qqmm(): [64, 128, 256, 1024, 1024 * 8], # N [64, 128, 256, 1024, 1024 * 8], # K ) + layouts = ["TN", "NT", "TT", "NN"] for group_size, mode, bits in tests: for M, N, K in product(*shapes): for dtype in dtypes: - x = mx.random.normal(shape=(M, K), key=k1, dtype=dtype) - w = mx.random.normal(shape=(N, K), key=k2, dtype=dtype) - w_q, scales_w = mx.quantize(w, group_size, bits, mode=mode) - w_dq = mx.dequantize( - w_q, - scales_w, - group_size=group_size, - bits=bits, - mode=mode, - dtype=dtype, - ) - y_q = mx.qqmm( - x, - w_q, - scales_w, - group_size=group_size, - bits=bits, - mode=mode, - ) - x_q, scales_x = mx.quantize( - x, group_size=group_size, bits=bits, mode=mode - ) - x_dq = mx.dequantize( - x_q, - scales_x, - group_size=group_size, - bits=bits, - mode=mode, - dtype=dtype, - ) - y_hat = mx.matmul(x_dq, mx.transpose(w_dq)) - ulp = ulp_bf16_at(y_hat) - error = (y_q - y_hat).abs() - if not (mx.logical_or(error < 1e-3, error <= ulp).all()): - raise AssertionError( - f"qqmm test failed for shape {(M, N, K)}, " - f"group_size={group_size}, bits={bits}, " - f"mode={mode}, dtype={dtype}" + for layout in layouts: + if layout == "NT": + x_shape = (M, K) + w_shape = (N, K) + elif layout == "TN": + x_shape = (K, M) + w_shape = (K, N) + elif layout == "TT": + x_shape = (K, M) + w_shape = (N, K) + else: # "NN" + x_shape = (M, K) + w_shape = (K, N) + + x = mx.random.normal(shape=x_shape, key=k1, dtype=dtype) + w = mx.random.normal(shape=w_shape, key=k2, dtype=dtype) + + if layout == "TT": + x = mx.transpose(x) + elif layout == "TN": + w = mx.transpose(w) + x = mx.transpose(x) + elif layout == "NN": + w = mx.transpose(w) + + y_q = mx.qqmm( + x, + w, + group_size=group_size, + bits=bits, + mode=mode, + ) + w_q, scales_w = mx.quantize(w, group_size, bits, mode=mode) + w_dq = mx.dequantize( + w_q, + scales_w, + group_size=group_size, + bits=bits, + mode=mode, + dtype=dtype, + ) + x_q, scales_x = mx.quantize( + x, group_size=group_size, bits=bits, mode=mode + ) + x_dq = mx.dequantize( + x_q, + scales_x, + group_size=group_size, + bits=bits, + mode=mode, + dtype=dtype, ) + y_hat = mx.matmul(x_dq, mx.transpose(w_dq)) + ulp = ulp_bf16_at(y_hat) + error = (y_q - y_hat).abs() + if not (mx.logical_or(error < 1e-3, error <= ulp).all()): + raise AssertionError( + f"qqmm test failed for shape {(M, N, K)}, " + f"group_size={group_size}, bits={bits}, " + f"mode={mode}, dtype={dtype}, layout={layout}" + ) def test_qqmm_vjp(): diff --git a/Source/Cmlx/mlx/mlx.pc.in b/Source/Cmlx/mlx/mlx.pc.in index c4e2515d7..a202028b1 100644 --- a/Source/Cmlx/mlx/mlx.pc.in +++ b/Source/Cmlx/mlx/mlx.pc.in @@ -11,6 +11,25 @@ @PACKAGE_INIT@ +include(CMakeFindDependencyMacro) + +if(WIN32 AND @MLX_BUILD_CPU@ AND NOT TARGET MLX::OpenBLAS) + add_library(MLX::OpenBLAS UNKNOWN IMPORTED) + set_target_properties(MLX::OpenBLAS PROPERTIES + IMPORTED_LOCATION "@PACKAGE_CMAKE_INSTALL_LIBDIR@/@OPENBLAS_LIB@.lib" + ) +endif() + +if(@MLX_BUILD_CUDA@) + find_dependency(CUDAToolkit REQUIRED) + + set(_MLX_SAVED_CMAKE_MODULE_PATH "${CMAKE_MODULE_PATH}") + list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}") + find_dependency(CUDNN REQUIRED) + set(CMAKE_MODULE_PATH "${_MLX_SAVED_CMAKE_MODULE_PATH}") + unset(_MLX_SAVED_CMAKE_MODULE_PATH) +endif() + include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/MLXTargets.cmake) include(@PACKAGE_MLX_CMAKE_INSTALL_MODULE_DIR@/extension.cmake) @@ -44,7 +63,6 @@ if (@MLX_BUILD_METAL@) endif() set_target_properties(mlx PROPERTIES - CXX_STANDARD 17 INTERFACE_COMPILE_OPTIONS "${MLX_CXX_FLAGS}" ) diff --git a/Source/Cmlx/mlx/mlx/CMakeLists.txt b/Source/Cmlx/mlx/mlx/CMakeLists.txt index 9c0fd3889..0df4f4234 100644 --- a/Source/Cmlx/mlx/mlx/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/CMakeLists.txt @@ -14,6 +14,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp ${CMAKE_CURRENT_SOURCE_DIR}/random.cpp ${CMAKE_CURRENT_SOURCE_DIR}/scheduler.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/stream.cpp ${CMAKE_CURRENT_SOURCE_DIR}/transforms.cpp ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp @@ -48,9 +49,11 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") target_compile_options(mlx PRIVATE -Wno-psabi) endif() -if(MSVC) - # Some of CUDA's headers include windows.h, which defines min/max macros. +if(WIN32) + # windows.h defines min/max macros that conflict with std::min/std::max. target_compile_definitions(mlx PRIVATE NOMINMAX WIN32_LEAN_AND_MEAN) +endif() +if(MSVC) # Unicode support in fmt does not compile in .cu files. target_compile_definitions(mlx PRIVATE FMT_UNICODE=0) # Disable some MSVC warnings to speed up compilation. diff --git a/Source/Cmlx/mlx/mlx/allocator.h b/Source/Cmlx/mlx/mlx/allocator.h index 824deac2c..c831b7247 100644 --- a/Source/Cmlx/mlx/mlx/allocator.h +++ b/Source/Cmlx/mlx/mlx/allocator.h @@ -72,4 +72,6 @@ inline void release(Buffer buffer) { allocator().release(buffer); } +MLX_API bool can_reuse_alien_buffer(void* ptr); + } // namespace mlx::core::allocator diff --git a/Source/Cmlx/mlx/mlx/array.cpp b/Source/Cmlx/mlx/mlx/array.cpp index 1769f0f84..07d3acd61 100644 --- a/Source/Cmlx/mlx/mlx/array.cpp +++ b/Source/Cmlx/mlx/mlx/array.cpp @@ -64,6 +64,7 @@ array array::unsafe_weak_copy(const array& other) { other.data_size(), other.strides(), other.flags(), + 0, [](auto) {}); cpy.array_desc_->offset = other.array_desc_->offset; return cpy; @@ -134,6 +135,7 @@ bool array::is_available() const { } else if ( status() == Status::evaluated && (!event().valid() || event().is_signaled())) { + detach_event(); set_status(Status::available); return true; } @@ -179,9 +181,10 @@ void array::set_data( size_t data_size, Strides strides, Flags flags, + int64_t offset, Deleter d) { array_desc_->data = std::make_shared(buffer, d); - array_desc_->offset = 0; + array_desc_->offset = offset; array_desc_->data_size = data_size; array_desc_->strides = std::move(strides); array_desc_->flags = flags; diff --git a/Source/Cmlx/mlx/mlx/array.h b/Source/Cmlx/mlx/mlx/array.h index 60d5e50bb..3f45e9cb9 100644 --- a/Source/Cmlx/mlx/mlx/array.h +++ b/Source/Cmlx/mlx/mlx/array.h @@ -426,6 +426,7 @@ class MLX_API array { } void detach_event() const { + array_desc_->event.check_error(); array_desc_->event = Event{}; } @@ -443,6 +444,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/mlx/mlx/backend/common/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/common/CMakeLists.txt index 6c4e25067..2b96fbee4 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/common/CMakeLists.txt @@ -4,6 +4,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/metal_kernel.cpp ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp) diff --git a/Source/Cmlx/mlx/mlx/backend/common/common.cpp b/Source/Cmlx/mlx/mlx/backend/common/common.cpp index 2cda88a31..5426301a3 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/common.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/common.cpp @@ -3,6 +3,7 @@ #include "mlx/backend/common/broadcasting.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -19,27 +20,28 @@ void AsStrided::eval(const std::vector& inputs, array& out) { "AsStrided must be used with row contiguous arrays only."); } - // Compute the flags given the shape and strides - bool row_contiguous = true, col_contiguous = true; - size_t r = 1, c = 1; - for (int i = strides_.size() - 1, j = 0; i >= 0; i--, j++) { - row_contiguous &= (r == strides_[i]) || (shape_[i] == 1); - col_contiguous &= (c == strides_[j]) || (shape_[j] == 1); - r *= shape_[i]; - c *= shape_[j]; + auto [no_bsx_size, row_contiguous, col_contiguous] = + check_contiguity(shape_, strides_); + + int64_t l = 0, h = 0; + bool has_negative_stride = false; + for (int i = 0; i < strides_.size(); i++) { + auto delta = strides_[i] * (shape_[i] - 1); + if (strides_[i] >= 0) { + h += delta; + } else { + l += delta; + has_negative_stride |= shape_[i] > 1; + } } + size_t data_size = out.size() == 0 ? 0 : (h - l) + 1; + auto flags = in.flags(); - // TODO: Compute the contiguous flag in a better way cause now we are - // unnecessarily strict. - flags.contiguous = row_contiguous || col_contiguous; + flags.contiguous = + out.size() == 0 || (!has_negative_stride && no_bsx_size == data_size); flags.row_contiguous = row_contiguous; flags.col_contiguous = col_contiguous; - // There is no easy way to compute the actual data size so we use out.size(). - // The contiguous flag will almost certainly not be set so no code should - // rely on data_size anyway. - size_t data_size = out.size(); - return out.copy_shared_buffer(in, strides_, flags, data_size, offset_); } @@ -98,50 +100,10 @@ void NumberOfElements::eval(const std::vector& inputs, array& out) { numel = 1.0 / numel; } - switch (out.dtype()) { - case bool_: - *out.data() = static_cast(numel); - break; - case uint8: - *out.data() = static_cast(numel); - break; - case uint16: - *out.data() = static_cast(numel); - break; - case uint32: - *out.data() = static_cast(numel); - break; - case uint64: - *out.data() = static_cast(numel); - break; - case int8: - *out.data() = static_cast(numel); - break; - case int16: - *out.data() = static_cast(numel); - break; - case int32: - *out.data() = static_cast(numel); - break; - case int64: - *out.data() = static_cast(numel); - break; - case float16: - *out.data() = static_cast(numel); - break; - case float32: - *out.data() = static_cast(numel); - break; - case bfloat16: - *out.data() = static_cast(numel); - break; - case float64: - *out.data() = static_cast(numel); - break; - case complex64: - *out.data() = static_cast(numel); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + *out.data() = static_cast(numel); + }); } std::pair prepare_reshape(const array& in, const array& out) { diff --git a/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp index aceeb1f7f..9dc5a4c54 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/compiled.cpp @@ -2,46 +2,30 @@ #include "mlx/backend/common/compiled.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/utils.h" namespace mlx::core { void print_constant(std::ostream& os, const array& x) { - switch (x.dtype()) { - case float32: - return print_float_constant(os, x); - case float16: - return print_float_constant(os, x); - case bfloat16: - return print_float_constant(os, x); - case float64: - return print_float_constant(os, x); - case complex64: - return print_complex_constant(os, x); - case int8: + dispatch_all_types(x.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { + os << std::boolalpha << x.item(); + } else if constexpr (std::is_same_v) { os << static_cast(x.item()); - return; - case int16: - return print_int_constant(os, x); - case int32: - return print_int_constant(os, x); - case int64: - return print_int_constant(os, x); - case uint8: + } else if constexpr (std::is_same_v) { os << static_cast(x.item()); - return; - case uint16: - return print_int_constant(os, x); - case uint32: - return print_int_constant(os, x); - case uint64: - return print_int_constant(os, x); - case bool_: - os << std::boolalpha << x.item(); - return; - default: + } else if constexpr (is_complex) { + print_complex_constant(os, x); + } else if constexpr (is_floating_point_v) { + print_float_constant(os, x); + } else if constexpr (std::is_integral_v) { + print_int_constant(os, x); + } else { throw std::runtime_error("Unsupported constant type"); - } + } + }); } std::string get_type_string(Dtype d) { @@ -82,7 +66,16 @@ std::string get_type_string(Dtype d) { } } -bool compiled_check_contiguity( +bool has_negative_strides(const array& x) { + for (auto s : x.strides()) { + if (s < 0) { + return true; + } + } + return false; +} + +std::pair compiled_check_contiguity( const std::vector& inputs, const Shape& shape) { bool contiguous = true; @@ -90,6 +83,7 @@ bool compiled_check_contiguity( bool all_row_contig = true; bool all_col_contig = true; int non_scalar_inputs = 0; + bool negative_strides = false; for (const auto& x : inputs) { if (is_scalar(x)) { continue; @@ -99,6 +93,7 @@ bool compiled_check_contiguity( all_contig &= (x.flags().contiguous && shape_eq); all_row_contig &= (x.flags().row_contiguous && shape_eq); all_col_contig &= (x.flags().col_contiguous && shape_eq); + negative_strides |= has_negative_strides(x); } if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) { contiguous = false; @@ -107,7 +102,7 @@ bool compiled_check_contiguity( } else if (non_scalar_inputs == 0 && !shape.empty()) { contiguous = false; } - return contiguous; + return {contiguous, negative_strides}; } void compiled_allocate_outputs( @@ -170,14 +165,16 @@ void compiled_allocate_outputs( } } -std::tuple> compiled_collapse_contiguous_dims( +std::tuple> +compiled_collapse_contiguous_dims( const std::vector& inputs, const array& out, const std::function& is_constant) { const Shape& shape = out.shape(); - bool contiguous = compiled_check_contiguity(inputs, shape); - if (contiguous) { - return {true, shape, {}}; + auto [contiguous, negative_strides] = + compiled_check_contiguity(inputs, shape); + if (contiguous && !negative_strides) { + return {true, false, shape, {}}; } std::vector strides_vec{out.strides()}; @@ -218,7 +215,11 @@ std::tuple> compiled_collapse_contiguous_dims( } auto tup = collapse_contiguous_dims(shape, strides_vec, INT32_MAX); - return {false, std::move(std::get<0>(tup)), std::move(std::get<1>(tup))}; + return { + false, + negative_strides, + std::move(std::get<0>(tup)), + std::move(std::get<1>(tup))}; } bool compiled_use_large_index( diff --git a/Source/Cmlx/mlx/mlx/backend/common/compiled.h b/Source/Cmlx/mlx/mlx/backend/common/compiled.h index 3be371333..8c6466da0 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/compiled.h +++ b/Source/Cmlx/mlx/mlx/backend/common/compiled.h @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #pragma once +#include #include #include @@ -17,13 +18,30 @@ std::string get_type_string(Dtype d); template void print_float_constant(std::ostream& os, const array& x) { + auto value = x.item(); + + // Non-finite values stream as bare tokens like `nan`/`inf`, which are not + // valid identifiers in the generated kernel source and fail to compile. + // Emit the INFINITY/NAN macros instead, which are provided by both the + // Metal () and CPU () kernel toolchains. Widen to + // double first so the finite check is exact for every float type. + double dvalue = static_cast(value); + if (std::isnan(dvalue)) { + os << "NAN"; + return; + } + if (std::isinf(dvalue)) { + os << (dvalue < 0 ? "-INFINITY" : "INFINITY"); + return; + } + auto old_precision = os.precision(); if constexpr (std::is_same_v) { os << std::setprecision(std::numeric_limits::digits10 + 1); } else { os << std::setprecision(std::numeric_limits::digits10 + 1); } - os << x.item() << std::setprecision(old_precision); + os << value << std::setprecision(old_precision); } template @@ -48,8 +66,9 @@ inline bool is_scalar(const array& x) { return x.ndim() == 0; } -// Check if we can use a contiguous operation given inputs and the output shape -bool compiled_check_contiguity( +// Check if we can use a contiguous operation given inputs and the output shape. +// Also returns if any input has negative strides. +std::pair compiled_check_contiguity( const std::vector& inputs, const Shape& shape); @@ -63,7 +82,8 @@ void compiled_allocate_outputs( allocator::malloc); // Collapse contiguous dims ignoring scalars and constants. -std::tuple> compiled_collapse_contiguous_dims( +std::tuple> +compiled_collapse_contiguous_dims( const std::vector& inputs, const array& out, const std::function& is_constant); diff --git a/Source/Cmlx/mlx/mlx/backend/common/load.cpp b/Source/Cmlx/mlx/mlx/backend/common/load.cpp index ce41963de..b53c92483 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/load.cpp +++ b/Source/Cmlx/mlx/mlx/backend/common/load.cpp @@ -51,7 +51,7 @@ void Load::eval_cpu(const std::vector& inputs, array& out) { } }; auto fut = io::thread_pool().enqueue(std::move(read_task)).share(); - scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.wait(); }); + scheduler::enqueue(stream(), [fut = std::move(fut)]() { fut.get(); }); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp new file mode 100644 index 000000000..7a2eef51f --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.cpp @@ -0,0 +1,378 @@ +// Copyright © 2024 Apple Inc. + +#include +#include + +#include "mlx/backend/common/compiled.h" +#include "mlx/backend/metal/metal.h" +#include "mlx/fast.h" +#include "mlx/fast_primitives.h" +#include "mlx/transforms_impl.h" +#include "mlx/utils.h" + +namespace mlx::core::fast { + +namespace { + +// Inputs with fewer elements are passed in the constant address space +constexpr int max_constant_array_size = 8; + +Stream resolve_metal_kernel_stream(StreamOrDevice s) { + if (metal::is_available()) { + // Default an unspecified stream to the GPU rather than the global default + // device. Other inputs go through to_stream so explicit streams, devices + // and ThreadLocalStream are handled as before. + auto stream = std::holds_alternative(s) + ? default_stream(Device::gpu) + : to_stream(s); + if (stream.device != Device::gpu) { + throw std::invalid_argument("[metal_kernel] Only supports the GPU."); + } + return stream; + } + if (!detail::in_export_tracing()) { + throw std::runtime_error("[metal_kernel] No Metal back-end."); + } + // Exporting without Metal: the kernel cannot run here but it can still be + // recorded in the graph on a placeholder GPU stream. The importing process + // remaps it to one of its own streams. + auto* device = std::get_if(&s); + auto* device_type = std::get_if(&s); + auto* stream = std::get_if(&s); + auto* tl_stream = std::get_if(&s); + if ((device && *device != Device::gpu) || + (device_type && *device_type != Device::gpu) || + (stream && stream->device != Device::gpu) || + (tl_stream && tl_stream->device != Device::gpu)) { + throw std::invalid_argument("[metal_kernel] Only supports the GPU."); + } + return Stream(-1, Device::gpu); +} + +std::string write_signature( + std::string func_name, + const std::string& header, + const std::string& source, + const std::vector& input_names, + const std::vector& inputs, + const std::vector& output_names, + const std::vector& output_dtypes, + const std::vector>& template_args, + const std::vector& attributes, + const std::vector>& shape_infos, + bool atomic_outputs) { + std::string kernel_source; + kernel_source.reserve(header.size() + source.size() + 16384); + kernel_source += header; + // Auto-generate a function signature based on `template_args` + // and the dtype/shape of the arrays passed as `inputs`. + if (!template_args.empty()) { + kernel_source += "template <"; + int i = 0; + for (const auto& [name, arg] : template_args) { + std::string param_type; + if (std::holds_alternative(arg)) { + param_type = "int"; + } else if (std::holds_alternative(arg)) { + param_type = "bool"; + } else if (std::holds_alternative(arg)) { + param_type = "typename"; + } + if (i > 0) { + kernel_source += ", "; + } + kernel_source += param_type; + kernel_source += " "; + kernel_source += name; + i++; + } + kernel_source += ">\n"; + } + kernel_source += "[[kernel]] void "; + kernel_source += func_name; + kernel_source += "(\n"; + + int index = 0; + // Add inputs + for (int i = 0; i < inputs.size(); ++i) { + const auto& name = input_names[i]; + const auto& arr = inputs[i]; + auto dtype = get_type_string(arr.dtype()); + std::string location = + arr.size() < max_constant_array_size ? "constant" : "device"; + std::string ref = arr.ndim() == 0 ? "&" : "*"; + kernel_source += " const "; + kernel_source += location; + kernel_source += " "; + kernel_source += dtype; + kernel_source += ref; + kernel_source += " "; + kernel_source += name; + kernel_source += " [[buffer("; + kernel_source += std::to_string(index); + kernel_source += ")]],\n"; + index++; + // Add input shape, strides and ndim if present in the source + if (arr.ndim() > 0) { + if (std::get<0>(shape_infos[i])) { + kernel_source += + (" const constant int* " + name + "_shape [[buffer(" + + std::to_string(index) + ")]],\n"); + index++; + } + if (std::get<1>(shape_infos[i])) { + kernel_source += + (" const constant int64_t* " + name + "_strides [[buffer(" + + std::to_string(index) + ")]],\n"); + index++; + } + if (std::get<2>(shape_infos[i])) { + kernel_source += + (" const constant int& " + name + "_ndim [[buffer(" + + std::to_string(index) + ")]],\n"); + index++; + } + } + } + // Add outputs + for (int i = 0; i < output_names.size(); ++i) { + const auto& name = output_names[i]; + const auto& dtype = output_dtypes[i]; + kernel_source += " device "; + auto type_string = get_type_string(dtype); + if (atomic_outputs) { + kernel_source += "atomic<"; + } + kernel_source += type_string; + if (atomic_outputs) { + kernel_source += ">"; + } + kernel_source += "* "; + kernel_source += name; + kernel_source += " [[buffer("; + kernel_source += std::to_string(index); + kernel_source += ")]]"; + if (index < inputs.size() + output_names.size() - 1 || + attributes.size() > 0) { + kernel_source += ",\n"; + } else { + kernel_source += ") {\n"; + } + index++; + } + + index = 0; + for (const auto& attr : attributes) { + kernel_source += attr; + if (index < attributes.size() - 1) { + kernel_source += ",\n"; + } else { + kernel_source += ") {\n"; + } + index++; + } + kernel_source += source; + kernel_source += "\n}\n"; + return kernel_source; +} + +std::string write_template( + const std::vector>& template_args) { + std::ostringstream template_def; + template_def << "<"; + int i = 0; + for (const auto& [name, arg] : template_args) { + if (i > 0) { + template_def << ", "; + } + if (std::holds_alternative(arg)) { + template_def << std::get(arg); + } else if (std::holds_alternative(arg)) { + template_def << std::get(arg); + } else if (std::holds_alternative(arg)) { + template_def << get_type_string(std::get(arg)); + } + i++; + } + template_def << ">"; + return template_def.str(); +} + +std::string make_template_hash(const std::string& template_def) { + std::string template_hash; + template_hash.reserve(template_def.size()); + for (size_t i = 0; i < template_def.size(); ++i) { + auto c = template_def[i]; + if (c == '<' || c == '>') { + template_hash += '_'; + } else if ( + c == ',' && i + 1 < template_def.size() && template_def[i + 1] == ' ') { + template_hash += '_'; + ++i; + } else { + template_hash += c; + } + } + template_hash.pop_back(); + return template_hash; +} + +} // namespace + +CustomKernelFunction metal_kernel( + const std::string& name, + const std::vector& input_names, + const std::vector& output_names, + const std::string& source, + const std::string& header /* = "" */, + bool ensure_row_contiguous /* = true */, + bool atomic_outputs /* = false */, + const CompileOptions& compile_options /* = {} */) { + if (output_names.empty()) { + throw std::invalid_argument( + "[metal_kernel] Must specify at least one output."); + } + std::vector> shape_infos; + for (auto& n : input_names) { + std::tuple shape_info; + std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos; + std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos; + std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos; + shape_infos.push_back(shape_info); + } + const std::vector> metal_attributes = { + {"dispatch_quadgroups_per_threadgroup", "uint"}, + {"dispatch_simdgroups_per_threadgroup", "uint"}, + {"dispatch_threads_per_threadgroup", "uint3"}, + {"grid_origin", "uint3"}, + {"grid_size", "uint3"}, + {"quadgroup_index_in_threadgroup", "uint"}, + {"quadgroups_per_threadgroup", "uint"}, + {"simdgroup_index_in_threadgroup", "uint"}, + {"simdgroups_per_threadgroup", "uint"}, + {"thread_execution_width", "uint"}, + {"thread_index_in_quadgroup", "uint"}, + {"thread_index_in_simdgroup", "uint"}, + {"thread_index_in_threadgroup", "uint"}, + {"thread_position_in_grid", "uint3"}, + {"thread_position_in_threadgroup", "uint3"}, + {"threadgroup_position_in_grid", "uint3"}, + {"threadgroups_per_grid", "uint3"}, + {"threads_per_grid", "uint3"}, + {"threads_per_simdgroup", "uint"}, + {"threads_per_threadgroup", "uint3"}, + }; + + std::vector attributes; + for (const auto& [attr, dtype] : metal_attributes) { + if (source.find(attr) != std::string::npos) { + attributes.push_back(" " + dtype + " " + attr + " [[" + attr + "]]"); + } + } + + return [=, + shape_infos = std::move(shape_infos), + attributes = std::move(attributes)]( + const std::vector& inputs, + const std::vector& output_shapes, + const std::vector& output_dtypes, + std::tuple grid, + std::tuple threadgroup, + const std::vector>& + template_args = {}, + std::optional init_value = std::nullopt, + bool verbose = false, + StreamOrDevice s_ = {}) { + auto check_size = [](size_t actual, size_t expected, const char* name) { + if (actual == expected) { + return; + } + std::ostringstream msg; + msg << "[metal_kernel] Expected `" << name << "` to have size " + << expected << " but got size " << actual << "." << std::endl; + throw std::invalid_argument(msg.str()); + }; + check_size(inputs.size(), input_names.size(), "inputs"); + check_size(output_shapes.size(), output_names.size(), "output_shapes"); + check_size(output_dtypes.size(), output_names.size(), "output_dtypes"); + + auto s = resolve_metal_kernel_stream(s_); + + std::string kernel_name = "custom_kernel_" + name; + std::string template_def = ""; + if (!template_args.empty()) { + template_def = write_template(template_args); + auto template_hash = make_template_hash(template_def); + kernel_name += "_"; + kernel_name += template_hash; + } + + // The generated source depends on the dtypes of the inputs and outputs + // and on how each input is passed (see `write_signature`). Include them + // in the kernel name so that a given name always maps to the same source. + for (const auto& arr : inputs) { + kernel_name += "_"; + kernel_name += get_type_string(arr.dtype()); + if (arr.ndim() == 0) { + kernel_name += "s"; + } else if (arr.size() < max_constant_array_size) { + kernel_name += "c"; + } + } + for (const auto& dtype : output_dtypes) { + kernel_name += "_"; + kernel_name += get_type_string(dtype); + } + + std::string kernel_source = write_signature( + kernel_name, + header, + source, + input_names, + inputs, + output_names, + output_dtypes, + template_args, + attributes, + shape_infos, + atomic_outputs); + + if (!template_args.empty()) { + template_def = kernel_name + template_def; + kernel_source += "\ntemplate [[host_name(\""; + kernel_source += kernel_name; + kernel_source += "\")]] [[kernel]] decltype("; + kernel_source += template_def; + kernel_source += ") "; + kernel_source += template_def; + kernel_source += ";\n"; + } + + if (verbose) { + std::cout << "Generated source code for `" << name << "`:" << std::endl + << "```" << std::endl + << kernel_source << std::endl + << "```" << std::endl; + } + + return array::make_arrays( + std::move(output_shapes), + std::move(output_dtypes), + std::make_shared( + s, + std::move(kernel_name), + std::move(kernel_source), + grid, + threadgroup, + shape_infos, + ensure_row_contiguous, + init_value, + std::vector{}, + false, + 0, + compile_options.serialize()), + std::move(inputs)); + }; +} + +} // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.h b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.h new file mode 100644 index 000000000..6ac0f4fdd --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/common/metal_kernel.h @@ -0,0 +1,27 @@ +// 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 diff --git a/Source/Cmlx/mlx/mlx/backend/common/utils.h b/Source/Cmlx/mlx/mlx/backend/common/utils.h index 1b6902ff3..c6d782061 100644 --- a/Source/Cmlx/mlx/mlx/backend/common/utils.h +++ b/Source/Cmlx/mlx/mlx/backend/common/utils.h @@ -116,6 +116,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/mlx/mlx/backend/cpu/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/cpu/CMakeLists.txt index f2c96949f..6192a189a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/cpu/CMakeLists.txt @@ -3,6 +3,7 @@ if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin") set(CLANG TRUE) else() set(COMPILER ${CMAKE_CXX_COMPILER}) + set(CLANG FALSE) endif() set(COMPILE_DEPS @@ -17,7 +18,7 @@ set(COMPILE_DEPS unary_ops.h binary_ops.h) -if(MSVC) +if(WIN32) set(SHELL_EXT ps1) set(SHELL_CMD powershell -ExecutionPolicy Bypass -File) else() @@ -31,6 +32,7 @@ add_custom_command( ${SHELL_CMD} ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.${SHELL_EXT} ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp ${COMPILER} ${PROJECT_SOURCE_DIR} ${CLANG} ${CMAKE_SYSTEM_PROCESSOR} + "${PREAMBLE_SIMD_FLAGS}" DEPENDS make_compiled_preamble.${SHELL_EXT} compiled_preamble.h ${COMPILE_DEPS}) diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp index 66468912d..d93f74cf0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/arg_reduce.cpp @@ -4,6 +4,7 @@ #include "mlx/backend/common/utils.h" #include "mlx/backend/cpu/encoder.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -74,50 +75,10 @@ void ArgReduce::eval_cpu(const std::vector& inputs, array& out) { out = array::unsafe_weak_copy(out), reduce_type_ = reduce_type_, axis_ = axis_]() mutable { - switch (in.dtype()) { - case bool_: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint8: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case uint64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int8: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case int64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float32: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case bfloat16: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case float64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - case complex64: - arg_reduce_dispatch(in, out, reduce_type_, axis_); - break; - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + arg_reduce_dispatch(in, out, reduce_type_, axis_); + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp index c1db2e118..90b0378f6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary.cpp @@ -3,12 +3,14 @@ #include #include #include +#include #include "mlx/allocator.h" #include "mlx/backend/cpu/binary.h" #include "mlx/backend/cpu/binary_ops.h" #include "mlx/backend/cpu/binary_two.h" #include "mlx/backend/cpu/encoder.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -45,56 +47,34 @@ void DivMod::eval_cpu( out_b = array::unsafe_weak_copy(out_b), bopt]() mutable { auto integral_op = [](auto x, auto y) { - return std::make_pair(x / y, x % y); + auto q = x / y; + auto r = x % y; + if constexpr (std::is_signed_v) { + if (r != 0 && (r < 0) != (y < 0)) { + q -= 1; + r += y; + } + } + return std::make_pair(q, r); }; auto float_op = [](auto x, auto y) { - return std::make_pair(std::trunc(x / y), std::fmod(x, y)); + auto r = std::fmod(x, y); + if (r != 0 && (r < 0) != (y < 0)) { + r += y; + } + return std::make_pair(std::floor(x / y), r); }; - switch (out_a.dtype()) { - case bool_: - binary_op(a, b, out_a, out_b, integral_op, bopt); - case uint8: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint16: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint32: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case uint64: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int8: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int16: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int32: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case int64: - binary_op(a, b, out_a, out_b, integral_op, bopt); - break; - case float16: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case float32: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case float64: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case bfloat16: - binary_op(a, b, out_a, out_b, float_op, bopt); - break; - case complex64: - // Should never get here + dispatch_all_types(out_a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("[DivMod] Complex type not supported"); - break; - } + } else if constexpr (std::is_integral_v) { + binary_op(a, b, out_a, out_b, integral_op, bopt); + } else { + binary_op(a, b, out_a, out_b, float_op, bopt); + } + }); }); } @@ -128,26 +108,11 @@ void Equal::eval_cpu(const std::vector& inputs, array& out) { b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (a.dtype()) { - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - default: - throw std::runtime_error( - "[NanEqual::eval_cpu] Only for floating point types."); - } + dispatch_inexact_types( + a.dtype(), "[NanEqual::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } else { comparison_op_cpu(a, b, out, detail::Equal(), stream()); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h index acaca50e1..5333aff9a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/binary.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/binary.h @@ -6,6 +6,7 @@ #include "mlx/array.h" #include "mlx/backend/common/binary.h" #include "mlx/backend/common/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" @@ -309,50 +310,10 @@ void binary_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (out.dtype()) { - case bool_: - binary_op(a, b, out, bopt); - break; - case uint8: - binary_op(a, b, out, bopt); - break; - case uint16: - binary_op(a, b, out, bopt); - break; - case uint32: - binary_op(a, b, out, bopt); - break; - case uint64: - binary_op(a, b, out, bopt); - break; - case int8: - binary_op(a, b, out, bopt); - break; - case int16: - binary_op(a, b, out, bopt); - break; - case int32: - binary_op(a, b, out, bopt); - break; - case int64: - binary_op(a, b, out, bopt); - break; - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } @@ -374,50 +335,10 @@ void comparison_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (a.dtype()) { - case bool_: - binary_op(a, b, out, bopt); - break; - case uint8: - binary_op(a, b, out, bopt); - break; - case uint16: - binary_op(a, b, out, bopt); - break; - case uint32: - binary_op(a, b, out, bopt); - break; - case uint64: - binary_op(a, b, out, bopt); - break; - case int8: - binary_op(a, b, out, bopt); - break; - case int16: - binary_op(a, b, out, bopt); - break; - case int32: - binary_op(a, b, out, bopt); - break; - case int64: - binary_op(a, b, out, bopt); - break; - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - } + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + binary_op(a, b, out, bopt); + }); }); } @@ -439,26 +360,9 @@ void binary_float_op_cpu( b = array::unsafe_weak_copy(b), out = array::unsafe_weak_copy(out), bopt]() mutable { - switch (out.dtype()) { - case float16: - binary_op(a, b, out, bopt); - break; - case float32: - binary_op(a, b, out, bopt); - break; - case float64: - binary_op(a, b, out, bopt); - break; - case bfloat16: - binary_op(a, b, out, bopt); - break; - case complex64: - binary_op(a, b, out, bopt); - break; - default: - throw std::runtime_error( - "[binary_float] Only supports floating point types."); - } + dispatch_inexact_types(out.dtype(), "[binary_float]", [&](auto type_tag) { + binary_op(a, b, out, bopt); + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/cholesky.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/cholesky.cpp index 3c5bbbc93..28cb87dd6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/cholesky.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/cholesky.cpp @@ -26,6 +26,9 @@ void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) { a.flags().row_contiguous ? CopyType::Vector : CopyType::General, stream); + if (a.shape(-1) == 0) { + return; + } auto& encoder = cpu::get_command_encoder(stream); encoder.set_output_array(factor); encoder.dispatch([matrix = factor.data(), diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp index 0806368e0..f5b516e44 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/compiled.cpp @@ -10,7 +10,6 @@ #include #include "mlx/backend/common/compiled.h" -#include "mlx/backend/cpu/compiled_preamble.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/jit_compiler.h" #include "mlx/device.h" @@ -42,15 +41,18 @@ struct CompilerCache { }; static CompilerCache& cache() { - static CompilerCache cache_; - return cache_; + // Leak - see Scheduler singleton comment in scheduler.cpp. + // DLib destructors call dlclose() which unmaps JIT .so files; + // StreamThreads may still be executing that code at exit. + static CompilerCache* cache_ = new CompilerCache; + return *cache_; }; -// GPU compile is always available if the GPU is available and since we are in -// this file CPU compile is also available. +// GPU compile is available through the GPU backend. CPU compile requires a +// usable C++ compiler (MSVC or clang-cl on Windows). namespace detail { bool compile_available_for_device(const Device& device) { - return true; + return device == Device::gpu || JitCompiler::available(); } } // namespace detail @@ -101,7 +103,11 @@ void* compile( std::filesystem::create_directories(output_dir); } +#ifdef _WIN32 + std::string shared_lib_name = kernel_file_name + ".dll"; +#else std::string shared_lib_name = "lib" + kernel_file_name + ".so"; +#endif auto shared_lib_path = (output_dir / shared_lib_name).string(); bool lib_exists = false; { @@ -292,7 +298,7 @@ void Compiled::eval_cpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides] = + auto [contiguous, negative_strides, shape, strides] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); // Collect function input arguments. @@ -316,7 +322,21 @@ void Compiled::eval_cpu( // Get the function auto fn_ptr = compile(kernel_name, [&, contiguous = contiguous]() { std::ostringstream kernel; - kernel << get_kernel_preamble() << std::endl; + kernel << std::get<2>(JitCompiler::get_preamble()) << std::endl; + // The prebuilt preamble is produced by running the preprocessor over + // compiled_preamble.h, which consumes the NAN / INFINITY macro definitions + // (macros vanish during preprocessing). print_float_constant emits those + // macros for non-finite constants, so define them here when missing. The + // #ifndef guards avoid clashing when is already in scope (e.g. the + // include-based preamble path). + kernel << "#ifndef INFINITY\n" + "#define INFINITY (std::numeric_limits::infinity())\n" + "#endif\n"; + kernel << "#ifndef NAN\n" + "#define NAN (std::numeric_limits::quiet_NaN())\n" + "#endif\n"; + kernel << "using namespace mlx::core;" << std::endl; + kernel << "using namespace mlx::core::detail;" << std::endl; kernel << "extern \"C\" {" << std::endl; build_kernel( kernel, diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/compiled_preamble.h b/Source/Cmlx/mlx/mlx/backend/cpu/compiled_preamble.h index 31ca1b468..02e3db8e7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/compiled_preamble.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/compiled_preamble.h @@ -9,4 +9,4 @@ #include "mlx/backend/cpu/binary_ops.h" // clang-format on -const char* get_kernel_preamble(); +const char* get_prebuilt_preamble(); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp index 5d4638ade..17bc6cb9f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/conv.cpp @@ -670,6 +670,24 @@ void slow_conv_3D( }); } +template +void dispatch_slow_conv_type(Dtype dtype, F&& f) { + switch (dtype) { + case float32: + f.template operator()(); + break; + case float16: + f.template operator()(); + break; + case bfloat16: + f.template operator()(); + break; + default: + throw std::invalid_argument( + "[Convolution::eval] got unsupported data type."); + } +} + void dispatch_slow_conv_1D( const array& in, const array& wt, @@ -681,20 +699,8 @@ void dispatch_slow_conv_1D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_1D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_1D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_1D( in, wt, out, @@ -705,22 +711,7 @@ void dispatch_slow_conv_1D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_1D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } void dispatch_slow_conv_2D( @@ -734,20 +725,8 @@ void dispatch_slow_conv_2D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_2D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_2D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_2D( in, wt, out, @@ -758,22 +737,7 @@ void dispatch_slow_conv_2D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_2D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } void dispatch_slow_conv_3D( @@ -787,20 +751,8 @@ void dispatch_slow_conv_3D( const std::vector& in_dilation, bool flip, Stream stream) { - if (in.dtype() == float32) { - return slow_conv_3D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else if (in.dtype() == float16) { - return slow_conv_3D( + dispatch_slow_conv_type(in.dtype(), [&]() { + slow_conv_3D( in, wt, out, @@ -811,22 +763,7 @@ void dispatch_slow_conv_3D( in_dilation, flip, stream); - } else if (in.dtype() == bfloat16) { - return slow_conv_3D( - in, - wt, - out, - padding_lo, - padding_hi, - wt_strides, - wt_dilation, - in_dilation, - flip, - stream); - } else { - throw std::invalid_argument( - "[Convolution::eval] got unsupported data type."); - } + }); } /////////////////////////////////////////////////////////////////////////////// @@ -877,7 +814,11 @@ void explicit_gemm_conv_1D_cpu( auto& encoder = cpu::get_command_encoder(stream); // Pad input - Shape padded_shape = {N, iH + padding_lo[0] + padding_hi[0], C}; + Shape padded_shape = { + N, + safe_cast( + static_cast(iH) + padding_lo[0] + padding_hi[0], "conv"), + C}; array in_padded(padded_shape, conv_dtype, nullptr, {}); // Fill with zeros @@ -1024,7 +965,8 @@ void explicit_gemm_conv_ND_cpu( Shape padded_shape(in.shape().size()); padded_shape.front() = N; for (size_t i = 0; i < iDim.size(); i++) { - padded_shape[i + 1] = iDim[i] + padding_lo[i] + padding_hi[i]; + padded_shape[i + 1] = safe_cast( + static_cast(iDim[i]) + padding_lo[i] + padding_hi[i], "conv"); } padded_shape.back() = C; array in_padded(padded_shape, conv_dtype, nullptr, {}); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp index f9ff22677..381719eb0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/copy.cpp @@ -7,6 +7,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" namespace mlx::core { @@ -70,7 +71,6 @@ void copy_general_general( dynamic_i_offset ? dynamic_i_offset->data() : nullptr; auto o_offset_ptr = dynamic_o_offset ? dynamic_o_offset->data() : nullptr; - auto size = src.size(); if (data_shape.empty()) { auto val = static_cast(*src_ptr); *dst_ptr = val; @@ -107,6 +107,8 @@ void copy_general_general( dst_ptr += o_offset_ptr[0]; } + auto size = std::accumulate( + shape.begin(), shape.end(), int64_t{1}, std::multiplies()); ContiguousIterator in(shape, strides[0], ndim - 3); ContiguousIterator out(shape, strides[1], ndim - 3); auto stride = std::accumulate( @@ -195,50 +197,10 @@ void copy(const array& src, array& dst, CopyType ctype, Args&&... args) { template void copy(const array& src, array& dst, CopyType ctype, Args&&... args) { - switch (dst.dtype()) { - case bool_: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint8: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint16: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint32: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint64: - copy(src, dst, ctype, std::forward(args)...); - break; - case int8: - copy(src, dst, ctype, std::forward(args)...); - break; - case int16: - copy(src, dst, ctype, std::forward(args)...); - break; - case int32: - copy(src, dst, ctype, std::forward(args)...); - break; - case int64: - copy(src, dst, ctype, std::forward(args)...); - break; - case float16: - copy(src, dst, ctype, std::forward(args)...); - break; - case float32: - copy(src, dst, ctype, std::forward(args)...); - break; - case float64: - copy(src, dst, ctype, std::forward(args)...); - break; - case bfloat16: - copy(src, dst, ctype, std::forward(args)...); - break; - case complex64: - copy(src, dst, ctype, std::forward(args)...); - break; - } + dispatch_all_types(dst.dtype(), [&](auto type_tag) { + using DstT = MLX_GET_TYPE(type_tag); + copy(src, dst, ctype, std::forward(args)...); + }); } template @@ -247,50 +209,10 @@ inline void copy_inplace_dispatch( array& dst, CopyType ctype, Args&&... args) { - switch (src.dtype()) { - case bool_: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint8: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint16: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint32: - copy(src, dst, ctype, std::forward(args)...); - break; - case uint64: - copy(src, dst, ctype, std::forward(args)...); - break; - case int8: - copy(src, dst, ctype, std::forward(args)...); - break; - case int16: - copy(src, dst, ctype, std::forward(args)...); - break; - case int32: - copy(src, dst, ctype, std::forward(args)...); - break; - case int64: - copy(src, dst, ctype, std::forward(args)...); - break; - case float16: - copy(src, dst, ctype, std::forward(args)...); - break; - case float32: - copy(src, dst, ctype, std::forward(args)...); - break; - case float64: - copy(src, dst, ctype, std::forward(args)...); - break; - case bfloat16: - copy(src, dst, ctype, std::forward(args)...); - break; - case complex64: - copy(src, dst, ctype, std::forward(args)...); - break; - } + dispatch_all_types(src.dtype(), [&](auto type_tag) { + using SrcT = MLX_GET_TYPE(type_tag); + copy(src, dst, ctype, std::forward(args)...); + }); } } // namespace diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp index 22dc4b4cc..6b788cde2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/distributed.cpp @@ -98,6 +98,23 @@ void Recv::eval_cpu( void ReduceScatter::eval_cpu( const std::vector& inputs, std::vector& outputs) { - throw std::runtime_error("[ReduceScatter] Not implemented yet."); + assert(inputs.size() == 0); + assert(outputs.size() == 1); + + auto [in, copied] = ensure_row_contiguous(inputs[0], stream()); + outputs[0].set_data(allocator::malloc(outputs[0].nbytes())); + + switch (reduce_type_) { + case Sum: + distributed::detail::sum_scatter(group(), in, outputs[0], stream()); + break; + default: + throw std::runtime_error("Only scatter sum is supported for now"); + } + if (copied) { + auto& enc = cpu::get_command_encoder(stream()); + enc.add_temporary(in); + } } + } // namespace mlx::core::distributed diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp index fb63eefdb..7bf674cfd 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eig.cpp @@ -228,6 +228,16 @@ void Eig::eval_cpu( const auto& a = inputs[0]; auto& values = outputs[0]; + values.set_data(allocator::malloc(values.nbytes())); + + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. + if (a.shape(-1) == 0) { + if (compute_eigenvectors_) { + outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); + } + return; + } + auto vectors = compute_eigenvectors_ ? outputs[1] : array(a.shape(), complex64, nullptr, {}); @@ -239,8 +249,6 @@ void Eig::eval_cpu( a.flags().row_contiguous ? CopyType::Vector : CopyType::General, stream()); - values.set_data(allocator::malloc(values.nbytes())); - if (compute_eigenvectors_) { // Set the strides and flags so the eigenvectors // are in the columns of the output diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp index d457c1fd9..ff5a4805a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eigh.cpp @@ -155,10 +155,12 @@ void eigh_impl( auto& encoder = cpu::get_command_encoder(stream); encoder.set_output_array(vectors); encoder.set_output_array(values); + // LAPACK reads the row-major input as its (conjugate) transpose, so the + // requested triangle is the opposite one in LAPACK's view. encoder.dispatch([vec_ptr, eig_ptr, jobz, - uplo = uplo[0], + uplo = uplo[0] == 'L' ? 'U' : 'L', N = vectors.shape(-1), size = vectors.size()]() mutable { // Work query @@ -190,12 +192,20 @@ void Eigh::eval_cpu( const auto& a = inputs[0]; auto& values = outputs[0]; + values.set_data(allocator::malloc(values.nbytes())); + + // Nothing to decompose for n = 0; LAPACK rejects lda = 0. + if (a.shape(-1) == 0) { + if (compute_eigenvectors_) { + outputs[1].set_data(allocator::malloc(outputs[1].nbytes())); + } + return; + } + auto vectors = compute_eigenvectors_ ? outputs[1] : array(a.shape(), a.dtype(), nullptr, {}); - values.set_data(allocator::malloc(values.nbytes())); - copy_cpu( a, vectors, @@ -234,7 +244,7 @@ void Eigh::eval_cpu( break; default: throw std::runtime_error( - "[Eigh::eval_cpu] only supports float32 or float64."); + "[Eigh::eval_cpu] only supports float32, float64, or complex64."); } } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp index f3fedefc3..bb2c06ea0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.cpp @@ -2,15 +2,33 @@ #include "mlx/backend/cpu/encoder.h" +#include + namespace mlx::core::cpu { -CommandEncoder& get_command_encoder(Stream stream) { - static std::unordered_map encoder_map; - auto it = encoder_map.find(stream.index); - if (it == encoder_map.end()) { - it = encoder_map.emplace(stream.index, stream).first; +CommandEncoder& get_command_encoder(Stream s) { + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { + auto& global_encoders = get_global_command_encoders(); + it = global_encoders.find(s.index); + if (it == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(cpu, {}) in current thread.", s.index)); + } } return it->second; } +std::unordered_map& get_command_encoders() { + static thread_local std::unordered_map encoders; + return encoders; +} + +std::unordered_map& get_global_command_encoders() { + static std::unordered_map encoders; + return encoders; +} + } // namespace mlx::core::cpu diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.h b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.h index e04179e5e..eb45d64ca 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/encoder.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/encoder.h @@ -46,11 +46,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)); } @@ -62,6 +61,9 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eval.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/eval.cpp index b23c8d561..354820f0f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eval.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eval.cpp @@ -7,6 +7,20 @@ namespace mlx::core::cpu { +void new_stream(Stream s) { + auto& encoders = get_command_encoders(); + encoders.try_emplace(s.index, s); +} + +void new_thread_unsafe_stream(Stream s) { + auto& encoders = get_global_command_encoders(); + encoders.try_emplace(s.index, s); +} + +void clear_streams() { + get_command_encoders().clear(); +} + void eval(array& arr) { auto s = arr.primitive().stream(); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/eval.h b/Source/Cmlx/mlx/mlx/backend/cpu/eval.h index 20156d61b..775f3e46f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/eval.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/eval.h @@ -7,6 +7,9 @@ namespace mlx::core::cpu { +void new_stream(Stream s); +void new_thread_unsafe_stream(Stream s); void eval(array& arr); +void clear_streams(); } // namespace mlx::core::cpu diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/indexing.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/indexing.cpp index a1b412c7d..15a2af3cb 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/indexing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/indexing.cpp @@ -4,11 +4,14 @@ #include #include "mlx/allocator.h" -#include "mlx/primitives.h" - #include "mlx/backend/common/utils.h" +#include "mlx/backend/cpu/binary.h" +#include "mlx/backend/cpu/binary_ops.h" #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" +#include "mlx/backend/cpu/slicing.h" +#include "mlx/dtype_utils.h" +#include "mlx/primitives.h" namespace mlx::core { @@ -63,43 +66,18 @@ void gather( array& out, const std::vector& axes, const Shape& slice_sizes) { - // If the array is row contiguous then we can do a contiguous copy given - // two conditions on the slice size: - // - Any number of leading ones in the slice sizes are allowed - // - All other slice sizes match the corresponding dimension except the - // first non-singleton slice size - // If the array is col contiguous then the reverse is the case: - // - Any number of trailing ones in the slice sizes are allowed - // - All other slice sizes match the corresponding dimension except the - // first non-singleton slice size from the end - - bool can_copy = false; - if (src.flags().row_contiguous) { - can_copy = true; - - // Ignore leading 1s - int i = 0; - for (; i < slice_sizes.size() && slice_sizes[i] == 1; ++i) - ; - - // Check the remaining - i++; - for (; i < src.ndim() && can_copy; ++i) { - can_copy = (src.shape(i) == slice_sizes[i]); - } - } else if (src.flags().col_contiguous) { - can_copy = true; - - // Ignore trailing 1s - int i = slice_sizes.size() - 1; - for (; i >= 0 && slice_sizes[i] == 1; --i) - ; - - // Skip the next slice size and check the remaining - i--; - for (; i >= 0 && can_copy; --i) { - can_copy = (src.shape(i) == slice_sizes[i]); + // Each gathered slice is written into the (row-contiguous) output as a + // sequential block. We can therefore replace the per-element strided read + // with a single contiguous copy only when the slice is itself laid out + // contiguously in row-major order within the source. + bool can_copy = true; + int64_t expected_stride = 1; + for (int i = src.ndim() - 1; i >= 0 && can_copy; --i) { + if (slice_sizes[i] == 1) { + continue; } + can_copy = (src.strides()[i] == expected_stride); + expected_stride *= slice_sizes[i]; } size_t slice_size = 1; for (auto s : slice_sizes) { @@ -150,50 +128,10 @@ void dispatch_gather( array& out, const std::vector& axes, const Shape& size) { - switch (out.dtype()) { - case bool_: - gather(src, inds, out, axes, size); - break; - case uint8: - gather(src, inds, out, axes, size); - break; - case uint16: - gather(src, inds, out, axes, size); - break; - case uint32: - gather(src, inds, out, axes, size); - break; - case uint64: - gather(src, inds, out, axes, size); - break; - case int8: - gather(src, inds, out, axes, size); - break; - case int16: - gather(src, inds, out, axes, size); - break; - case int32: - gather(src, inds, out, axes, size); - break; - case int64: - gather(src, inds, out, axes, size); - break; - case float16: - gather(src, inds, out, axes, size); - break; - case float32: - gather(src, inds, out, axes, size); - break; - case float64: - gather(src, inds, out, axes, size); - break; - case bfloat16: - gather(src, inds, out, axes, size); - break; - case complex64: - gather(src, inds, out, axes, size); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + gather(src, inds, out, axes, size); + }); } void Gather::eval_cpu(const std::vector& inputs, array& out) { @@ -303,50 +241,10 @@ void dispatch_gather_axis( const array& inds, array& out, const int axis) { - switch (out.dtype()) { - case bool_: - gather_axis(src, inds, out, axis); - break; - case uint8: - gather_axis(src, inds, out, axis); - break; - case uint16: - gather_axis(src, inds, out, axis); - break; - case uint32: - gather_axis(src, inds, out, axis); - break; - case uint64: - gather_axis(src, inds, out, axis); - break; - case int8: - gather_axis(src, inds, out, axis); - break; - case int16: - gather_axis(src, inds, out, axis); - break; - case int32: - gather_axis(src, inds, out, axis); - break; - case int64: - gather_axis(src, inds, out, axis); - break; - case float16: - gather_axis(src, inds, out, axis); - break; - case float32: - gather_axis(src, inds, out, axis); - break; - case float64: - gather_axis(src, inds, out, axis); - break; - case bfloat16: - gather_axis(src, inds, out, axis); - break; - case complex64: - gather_axis(src, inds, out, axis); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + gather_axis(src, inds, out, axis); + }); } void GatherAxis::eval_cpu(const std::vector& inputs, array& out) { @@ -532,50 +430,10 @@ void Scatter::eval_cpu(const std::vector& inputs, array& out) { updates = array::unsafe_weak_copy(updates), inds = std::move(inds), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint8: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case uint64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int8: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case int64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float32: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case float64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case bfloat16: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - case complex64: - dispatch_scatter(out, inds, updates, axes_, reduce_type_); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + dispatch_scatter(out, inds, updates, axes_, reduce_type_); + }); }); } @@ -697,53 +555,10 @@ void ScatterAxis::eval_cpu(const std::vector& inputs, array& out) { idx = array::unsafe_weak_copy(idx), updates = array::unsafe_weak_copy(updates), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint8: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint16: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case uint64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int8: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int16: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case int64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case float16: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - case float32: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case float64: - dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); - break; - case bfloat16: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - case complex64: - dispatch_scatter_axis( - out, idx, updates, axis_, reduce_type_); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + dispatch_scatter_axis(out, idx, updates, axis_, reduce_type_); + }); }); } @@ -788,7 +603,7 @@ void MaskedScatter::eval_cpu(const std::vector& inputs, array& out) { auto& mask = inputs[1]; auto& src = inputs[2]; - // Copy src into out (copy allocates memory for out) + // Copy dst into out (copy allocates memory for out) auto ctype = dst.flags().row_contiguous ? CopyType::Vector : CopyType::General; copy_cpu(dst, out, ctype, stream()); @@ -804,50 +619,134 @@ void MaskedScatter::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([mask = array::unsafe_weak_copy(mask), src = array::unsafe_weak_copy(src), out = array::unsafe_weak_copy(out)]() mutable { - switch (out.dtype()) { - case bool_: - masked_scatter_impl(mask, src, out); - break; - case uint8: - masked_scatter_impl(mask, src, out); - break; - case uint16: - masked_scatter_impl(mask, src, out); - break; - case uint32: - masked_scatter_impl(mask, src, out); - break; - case uint64: - masked_scatter_impl(mask, src, out); - break; - case int8: - masked_scatter_impl(mask, src, out); - break; - case int16: - masked_scatter_impl(mask, src, out); - break; - case int32: - masked_scatter_impl(mask, src, out); - break; - case int64: - masked_scatter_impl(mask, src, out); - break; - case float16: - masked_scatter_impl(mask, src, out); - break; - case float32: - masked_scatter_impl(mask, src, out); - break; - case float64: - masked_scatter_impl(mask, src, out); - break; - case bfloat16: - masked_scatter_impl(mask, src, out); - break; - case complex64: - masked_scatter_impl(mask, src, out); - break; + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + masked_scatter_impl(mask, src, out); + }); + }); +} + +template +void slice_update_impl( + array& out, + const array& upd, + int64_t data_offset, + const Strides& out_strides) { + ContiguousIterator out_it(upd.shape(), out_strides, upd.ndim()); + ContiguousIterator upd_it(upd); + Op op; + + constexpr int SIMD_START = 32; + + T* out_ptr = out.data() + data_offset; + const T* upd_ptr = upd.data(); + int64_t size = upd.size(); + int64_t suffix = out_it.contiguous_suffix(); + + if (upd.data_size() == 1) { + if (suffix >= SIMD_START) { + for (int64_t i = 0; i < size; i += suffix) { + VectorScalar{}( + out_ptr + out_it.loc, upd_ptr, out_ptr + out_it.loc, suffix); + out_it.step(suffix); + } + } else { + T update = upd_ptr[0]; + for (int64_t i = 0; i < size; i++) { + out_ptr[out_it.loc] = op(out_ptr[out_it.loc], update); + out_it.step(); + } } + } else if (suffix == upd_it.contiguous_suffix() && suffix >= SIMD_START) { + for (int64_t i = 0; i < size; i += suffix) { + VectorVector{}( + out_ptr + out_it.loc, + upd_ptr + upd_it.loc, + out_ptr + out_it.loc, + suffix); + out_it.step(suffix); + upd_it.step(suffix); + } + } else { + for (int64_t i = 0; i < size; i++) { + out_ptr[out_it.loc] = op(out_ptr[out_it.loc], upd_ptr[upd_it.loc]); + out_it.step(); + upd_it.step(); + } + } +} + +void SliceUpdate::eval_cpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + if (out.size() == 0) { + out.set_data(allocator::malloc(0)); + return; + } + + auto& in = inputs[0]; + auto& upd = inputs[1]; + + if (upd.size() == 0) { + out.copy_shared_buffer(in); + return; + } + + // Check if materialization is needed + auto ctype = in.flags().contiguous && in.size() == in.data_size() + ? CopyType::Vector + : CopyType::General; + copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream()); + + // Calculate out strides, initial offset and if copy needs to be made + auto [data_offset, out_strides] = + prepare_slice(out, start_indices_, strides_); + + // Do copy + if (reduce_type_ == SliceUpdate::None) { + copy_cpu_inplace( + /* const array& src = */ upd, + /* array& dst = */ out, + /* const std::vector& data_shape = */ upd.shape(), + /* const std::vector& i_strides = */ upd.strides(), + /* const std::vector& o_strides = */ out_strides, + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ data_offset, + /* CopyType ctype = */ CopyType::GeneralGeneral, + stream()); + return; + } + + auto& encoder = cpu::get_command_encoder(stream()); + encoder.set_input_array(upd); + encoder.set_output_array(out); + encoder.dispatch([upd = array::unsafe_weak_copy(upd), + out = array::unsafe_weak_copy(out), + data_offset = data_offset, + out_strides = std::move(out_strides), + reduce_type = reduce_type_]() mutable { + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + switch (reduce_type) { + case SliceUpdate::Sum: + slice_update_impl(out, upd, data_offset, out_strides); + break; + case SliceUpdate::Prod: + slice_update_impl( + out, upd, data_offset, out_strides); + break; + case SliceUpdate::Max: + slice_update_impl( + out, upd, data_offset, out_strides); + break; + case SliceUpdate::Min: + slice_update_impl( + out, upd, data_offset, out_strides); + break; + case SliceUpdate::None: + // Should never be here + break; + } + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.cpp index 267e8fd86..530637e55 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.cpp @@ -1,19 +1,27 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #include "mlx/backend/cpu/jit_compiler.h" +#include "mlx/backend/common/utils.h" +#include "mlx/backend/cpu/compiled_preamble.h" #include +#include #include #include #include -namespace mlx::core { +#if defined(_MSC_VER) && \ + (defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64)) +#include +#endif -#ifdef _MSC_VER +namespace mlx::core { namespace { +#if defined(_MSC_VER) + // Split string into array. std::vector str_split(const std::string& str, char delimiter) { std::vector tokens; @@ -36,11 +44,16 @@ struct VisualStudioInfo { // Get path of Visual Studio. // Use -latest to get only the most recent installation when multiple // versions are installed, avoiding path concatenation issues. + auto pf86 = std::getenv("ProgramFiles(x86)"); + if (!pf86) { + throw std::runtime_error( + "ProgramFiles(x86) environment variable not set."); + } std::string vs_path = JitCompiler::exec( fmt::format( "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\"" - " -latest -property installationPath", - std::getenv("ProgramFiles(x86)"))); + " -latest -property installationPath 2>&1", + pf86)); if (vs_path.empty()) { throw std::runtime_error("Can not find Visual Studio."); } @@ -55,7 +68,7 @@ struct VisualStudioInfo { // Read the envs from vcvarsall. std::string envs = JitCompiler::exec( fmt::format( - "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set", + "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL 2>&1 && set", vs_path, arch)); for (const std::string& line : str_split(envs, '\n')) { @@ -65,15 +78,45 @@ struct VisualStudioInfo { continue; std::string name = line.substr(0, pos); std::string value = line.substr(pos + 1); - if (name == "LIB") { + if (name == "INCLUDE") { + includepaths = str_split(value, ';'); + } else if (name == "LIB") { libpaths = str_split(value, ';'); } else if (name == "VCToolsInstallDir" || name == "VCTOOLSINSTALLDIR") { - cl_exe = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch); + msvc_cl = fmt::format("{0}\\bin\\Host{1}\\{1}\\cl.exe", value, arch); } } + + // Check for clang-cl bundled with Visual Studio. + std::string clang_cl_path = fmt::format( + "{0}\\VC\\Tools\\Llvm\\{1}\\bin\\clang-cl.exe", vs_path, arch); + if (std::filesystem::exists(clang_cl_path)) { + clang_cl = clang_cl_path; + } } + + std::string compiler(bool use_include) const { + if (use_include) { + // With installed headers, either cl.exe or clang-cl can parse the source. + // Prefer cl.exe since it is part of the standard Visual Studio C++ + // toolchain. + return !msvc_cl.empty() ? msvc_cl : clang_cl; + } + + // Without installed headers, fall back to the embedded prebuilt preamble. + // It was generated by the build compiler and may contain compiler-specific + // builtins, so prefer the same compiler family when possible. +#ifdef __clang__ + return !clang_cl.empty() ? clang_cl : msvc_cl; +#else + return !msvc_cl.empty() ? msvc_cl : clang_cl; +#endif + } + std::string arch; - std::string cl_exe; + std::string msvc_cl; + std::string clang_cl; + std::vector includepaths; std::vector libpaths; }; @@ -82,34 +125,154 @@ const VisualStudioInfo& GetVisualStudioInfo() { return info; } +#endif // defined(_MSC_VER) + +bool supports_avx2_target() { +#if defined(_MSC_VER) + +#if defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64) + int info[4]; + __cpuid(info, 0); + if (info[0] < 7) { + return false; + } + __cpuid(info, 1); + bool os_avx = (info[2] & (1 << 27)) && (info[2] & (1 << 28)); + bool fma = info[2] & (1 << 12); + bool f16c = info[2] & (1 << 29); + if (!os_avx || !fma || !f16c) { + return false; + } + unsigned long long xcr0 = _xgetbv(0); + if ((xcr0 & 0x6) != 0x6) { + return false; + } + __cpuidex(info, 7, 0); + return info[1] & (1 << 5); +#else + return false; +#endif // defined(_M_X64) || defined(_M_IX86) || defined(_M_AMD64) + +#elif defined(__GNUC__) || defined(__clang__) + +#if defined(__x86_64__) || defined(__i386__) || defined(__amd64__) + __builtin_cpu_init(); + return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma") && + __builtin_cpu_supports("f16c"); +#else + return false; +#endif + +#else + return false; +#endif // defined(_MSC_VER) +} + } // namespace -#endif // _MSC_VER +const std::tuple& JitCompiler::get_preamble() { + static auto preamble = []() -> std::tuple { + // Check whether the headers are shipped with the binary, if so use the + // preamble from the headers, otherwise use the prebuilt one embedded in + // binary, which may not work with all compilers. + auto root_dir = current_binary_dir(); +#if !defined(_WIN32) + root_dir = root_dir.parent_path(); +#endif + auto include_dir = root_dir / "include"; + if (std::filesystem::exists( + include_dir / "mlx/backend/cpu/compiled_preamble.h")) { + return std::make_tuple( + true, + include_dir.string(), + "#include \"mlx/backend/cpu/compiled_preamble.h\"\n"); + } else { + return std::make_tuple(false, "", get_prebuilt_preamble()); + } + }(); + return preamble; +} + +bool JitCompiler::available() { +#ifdef _MSC_VER + static bool result = [] { + try { + bool use_include = std::get<0>(get_preamble()); + const auto& info = GetVisualStudioInfo(); + return !info.compiler(use_include).empty(); + } catch (...) { + return false; + } + }(); + return result; +#else +#ifdef _WIN32 + static int result = std::system("g++ --version > NUL 2>&1"); +#else + static int result = std::system("g++ --version > /dev/null 2>&1"); +#endif + return result == 0; +#endif +} std::string JitCompiler::build_command( const std::filesystem::path& dir, const std::string& source_file_name, const std::string& shared_lib_name) { + auto& [use_include, include_dir, preamble] = get_preamble(); #ifdef _MSC_VER + std::string compiler_flags; + if (use_include) { + compiler_flags += fmt::format(" /I \"{}\"", include_dir); + } const VisualStudioInfo& info = GetVisualStudioInfo(); + std::string compiler = info.compiler(use_include); + for (const std::string& include : info.includepaths) { + compiler_flags += fmt::format(" /I \"{}\"", include); + } std::string libpaths; for (const std::string& lib : info.libpaths) { - libpaths += fmt::format(" /libpath:\"{0}\"", lib); + libpaths += fmt::format(" /libpath:\"{}\"", lib); + } + // clang-cl accepts the same flags as cl.exe (/LD, /EHsc, etc.) + // but we add -Wno-everything to suppress warnings from system and shipped + // MLX headers. + if (!info.clang_cl.empty() && compiler == info.clang_cl) { + compiler_flags += " -Wno-everything"; } +#ifdef __AVX2__ + if (supports_avx2_target()) { + compiler_flags += " /arch:AVX2"; + } +#endif return fmt::format( "\"" "cd /D \"{0}\" && " - "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17 \"{2}\" " - "/link /out:\"{3}\" {4} 2>&1" + "\"{1}\" /LD /EHsc /MD /Ox /nologo /std:c++17{2} \"{3}\" " + "/link /out:\"{4}\" {5} 2>&1" "\"", dir.string(), - info.cl_exe, + compiler, + compiler_flags, source_file_name, shared_lib_name, libpaths); #else + std::string extra_flags; + if (use_include) { + extra_flags = fmt::format("-I \"{}\"", include_dir); + } +#ifdef __AVX2__ + if (supports_avx2_target()) { + if (!extra_flags.empty()) { + extra_flags += " "; + } + extra_flags += "-mavx2 -mfma -mf16c"; + } +#endif return fmt::format( - "g++ -std=c++17 -O3 -Wall -fPIC -shared \"{0}\" -o \"{1}\" 2>&1", + "g++ -std=c++17 -O3 -Wall -fPIC -shared {} \"{}\" -o \"{}\" 2>&1", + extra_flags, (dir / source_file_name).string(), (dir / shared_lib_name).string()); #endif @@ -129,7 +292,7 @@ std::string JitCompiler::exec(const std::string& cmd) { while (fgets(buffer, sizeof(buffer), pipe)) { ret += buffer; } - // Trim trailing spaces. + // Trim trailing whitespace. ret.erase( std::find_if( ret.rbegin(), diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.h b/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.h index 3a9e988da..c2dac51ad 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/jit_compiler.h @@ -1,4 +1,4 @@ -// Copyright © 2024 Apple Inc. +// Copyright © 2024-2026 Apple Inc. #pragma once #include @@ -7,6 +7,15 @@ namespace mlx::core { class JitCompiler { public: + // Return the includes that should be prepended to the source code. + static const std::tuple& get_preamble(); + + // Check if a JIT compiler is available on this system. + // On Windows, this probes for Visual Studio and a usable C++ compiler + // (MSVC cl.exe or clang-cl). On Linux/macOS, checks for g++ in PATH. + // Returns false (rather than throwing) if no compiler is found. + static bool available(); + // Build a shell command that compiles a source code file to a shared library. static std::string build_command( const std::filesystem::path& dir, diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/luf.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/luf.cpp index 5f1507e18..2e396f4b7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/luf.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/luf.cpp @@ -67,11 +67,10 @@ void luf_impl( /* ipiv */ reinterpret_cast(pivots_ptr), /* info */ &info); - if (info != 0) { + if (info < 0) { std::stringstream ss; - ss << "[LUF::eval_cpu] sgetrf_ failed with code " << info - << ((info > 0) ? " because matrix is singular" - : " because argument had an illegal value"); + ss << "[LUF::eval_cpu] LU factorization failed with error code " + << info << " because argument had an illegal value"; throw std::runtime_error(ss.str()); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.ps1 b/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.ps1 index 0cd2d1f17..e9fa7afba 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.ps1 +++ b/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.ps1 @@ -1,27 +1,52 @@ # This script generates a C++ function that provides the CPU # code for use with kernel generation. # -# Copyright © 2024 Apple Inc. +# Copyright © 2024-2026 Apple Inc. $OUTPUT_FILE = $args[0] $CL = $args[1] $SRCDIR = $args[2] +# args[3] = CLANG -- whether the build compiler is Apple Clang (unused on Windows; +# on macOS it controls -nobuiltininc and -arch flags for the preprocessor. +# MSVC and clang-cl both accept the same preprocessor flags so no distinction needed.) +# args[4] = CMAKE_SYSTEM_PROCESSOR (unused on Windows) +$SIMD_FLAGS = $args[5] # Optional, e.g. "/arch:AVX2" or "-mavx2 -mfma -mf16c" -# Get command result as array. -$CONTENT = & $CL /std:c++17 /EP "/I$SRCDIR" /Tp "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" +# Detect compiler type: MSVC/clang-cl use /EP, GCC/clang++ use -E -P +$CL_NAME = [System.IO.Path]::GetFileNameWithoutExtension($CL) +$IS_MSVC_LIKE = ($CL_NAME -eq 'cl') -or ($CL_NAME -eq 'clang-cl') + +# Build the full argument list to avoid PowerShell splatting issues +# (splatting a single-element array can degrade to char-by-char expansion). +$CL_ARGS = [System.Collections.ArrayList]::new() +if ($IS_MSVC_LIKE) { + [void]$CL_ARGS.Add('/std:c++17') + [void]$CL_ARGS.Add('/EP') +} else { + [void]$CL_ARGS.Add('-std=c++17') + [void]$CL_ARGS.Add('-E') + [void]$CL_ARGS.Add('-P') +} +if ($SIMD_FLAGS) { + foreach ($f in ($SIMD_FLAGS -split ' ')) { [void]$CL_ARGS.Add($f) } +} +if ($IS_MSVC_LIKE) { + [void]$CL_ARGS.Add("/I$SRCDIR") + [void]$CL_ARGS.Add('/Tp') +} else { + [void]$CL_ARGS.Add('-I') + [void]$CL_ARGS.Add("$SRCDIR") +} +[void]$CL_ARGS.Add("$SRCDIR/mlx/backend/cpu/compiled_preamble.h") + +# Get command result as array. Redirect stderr to null to suppress warnings. +$CONTENT = & $CL @CL_ARGS 2>$null # Remove empty lines. # Otherwise there will be too much empty lines making the result unreadable. $CONTENT = $CONTENT | Where-Object { $_.Trim() -ne '' } # Concatenate to string. $CONTENT = $CONTENT -join "`n" -# Append extra content. -$CONTENT = @" -$($CONTENT) -using namespace mlx::core; -using namespace mlx::core::detail; -"@ - # Convert each char to ASCII code. # Unlike the unix script that outputs string literal directly, the output from # MSVC is way too large to be embedded as string and compilation will fail, so @@ -29,7 +54,7 @@ using namespace mlx::core::detail; $CHARCODES = ([System.Text.Encoding]::ASCII.GetBytes($CONTENT) -join ', ') + ', 0' $OUTPUT = @" -const char* get_kernel_preamble() { +const char* get_prebuilt_preamble() { static char preamble[] = { $CHARCODES }; return preamble; } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.sh b/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.sh index 88b4c4615..dc53bb8b5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.sh +++ b/Source/Cmlx/mlx/mlx/backend/cpu/make_compiled_preamble.sh @@ -3,7 +3,7 @@ # This script generates a C++ function that provides the CPU # code for use with kernel generation. # -# Copyright © 2023-24 Apple Inc. +# Copyright © 2023-2026 Apple Inc. OUTPUT_FILE=$1 @@ -11,6 +11,7 @@ GCC=$2 SRCDIR=$3 CLANG=$4 ARCH=$5 +SIMD_FLAGS=$6 # Optional, e.g. "-mavx2 -mfma -mf16c" if [ "$CLANG" = "TRUE" ]; then read -r -d '' INCLUDES <<- EOM @@ -27,15 +28,13 @@ else CC_FLAGS="-std=c++17" fi -CONTENT=$($GCC $CC_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null) +CONTENT=$("$GCC" $CC_FLAGS $SIMD_FLAGS -I "$SRCDIR" -E -P "$SRCDIR/mlx/backend/cpu/compiled_preamble.h" 2>/dev/null) cat << EOF > "$OUTPUT_FILE" -const char* get_kernel_preamble() { +const char* get_prebuilt_preamble() { return R"preamble( $INCLUDES $CONTENT -using namespace mlx::core; -using namespace mlx::core::detail; )preamble"; } EOF diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/masked_mm.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/masked_mm.cpp index 688479c60..db38f3942 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/masked_mm.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/masked_mm.cpp @@ -8,6 +8,7 @@ #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/gemm.h" #include "mlx/backend/cpu/lapack.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -525,83 +526,26 @@ void SegmentedMM::eval_cpu(const std::vector& inputs, array& out) { b_transposed = b_transposed, lda = lda, ldb = ldb]() { - switch (a.dtype()) { - case float64: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case float32: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case float16: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - case bfloat16: - segmented_mm( - a.data(), - b.data(), - segments.data(), - static_cast(out_ptr), - a_transposed, - b_transposed, - lda, - ldb, - a.shape(), - a.strides(), - b.shape(), - b.strides(), - segments.size() / 2, - segments.shape(), - segments.strides()); - break; - default: - throw std::invalid_argument( - "Segmented mm supports only real float types."); - } + dispatch_float_types( + a.dtype(), "[SegmentedMM::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + segmented_mm( + a.data(), + b.data(), + segments.data(), + static_cast(out_ptr), + a_transposed, + b_transposed, + lda, + ldb, + a.shape(), + a.strides(), + b.shape(), + b.strides(), + segments.size() / 2, + segments.shape(), + segments.strides()); + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/matmul.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/matmul.cpp index 7df331671..12781dc1a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/matmul.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/matmul.cpp @@ -7,6 +7,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/gemm.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -97,24 +98,11 @@ void matmul_general( return; } - if (out.dtype() == float32) { - matmul_dispatch( + dispatch_inexact_types(out.dtype(), "[Matmul::eval_cpu]", [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + matmul_dispatch( a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == float16) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == bfloat16) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == float64) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else if (out.dtype() == complex64) { - matmul_dispatch( - a, b, out, a_transposed, b_transposed, lda, ldb, alpha, beta, stream); - } else { - throw std::runtime_error("[Matmul::eval_cpu] Invalid type."); - } + }); cpu::get_command_encoder(stream).add_temporaries(std::move(temps)); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp index 4e59b1ebe..3299b7471 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/primitives.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/allocator.h" #include "mlx/backend/common/slicing.h" @@ -13,6 +14,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/threefry.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -125,50 +127,14 @@ void Transpose::eval_cpu(const std::vector& inputs, array& out) { void Arange::eval_cpu(const std::vector& inputs, array& out) { assert(inputs.size() == 0); out.set_data(allocator::malloc(out.nbytes())); - switch (out.dtype()) { - case bool_: + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("Bool type unsupported for arange."); - break; - case uint8: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case uint64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int8: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case int64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float32: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case float64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case bfloat16: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - case complex64: - arange(start_, start_ + step_, out, out.size(), stream()); - break; - } + } else { + arange(start_, start_ + step_, out, out.size(), stream()); + } + }); } void AsType::eval_cpu(const std::vector& inputs, array& out) { @@ -398,44 +364,6 @@ void DynamicSliceUpdate::eval_cpu( } } -void SliceUpdate::eval_cpu(const std::vector& inputs, array& out) { - assert(inputs.size() == 2); - if (out.size() == 0) { - out.set_data(allocator::malloc(0)); - return; - } - - auto& in = inputs[0]; - auto& upd = inputs[1]; - - if (upd.size() == 0) { - out.copy_shared_buffer(in); - return; - } - - // Check if materialization is needed - auto ctype = in.flags().contiguous && in.size() == in.data_size() - ? CopyType::Vector - : CopyType::General; - copy_cpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream()); - - // Calculate out strides, initial offset and if copy needs to be made - auto [data_offset, out_strides] = - prepare_slice(out, start_indices_, strides_); - - // Do copy - copy_cpu_inplace( - /* const array& src = */ upd, - /* array& dst = */ out, - /* const std::vector& data_shape = */ upd.shape(), - /* const std::vector& i_strides = */ upd.strides(), - /* const std::vector& o_strides = */ out_strides, - /* int64_t i_offset = */ 0, - /* int64_t o_offset = */ data_offset, - /* CopyType ctype = */ CopyType::GeneralGeneral, - stream()); -} - void View::eval_cpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); auto& in = inputs[0]; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp index 13c7e1132..f438d492a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/qrf.cpp @@ -12,6 +12,14 @@ template void qrf_impl(const array& a, array& q, array& r, Stream stream) { const int M = a.shape(-2); const int N = a.shape(-1); + + // Nothing to factorize when either dimension is zero; LAPACK rejects lda = 0. + if (M == 0 || N == 0) { + q.set_data(allocator::malloc(q.nbytes())); + r.set_data(allocator::malloc(r.nbytes())); + return; + } + const int lda = M; size_t num_matrices = a.size() / (M * N); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp index 3c3643a32..15e00cd91 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/quantized.cpp @@ -1,5 +1,6 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2026 Apple Inc. +#include "mlx/backend/common/quantized.h" #include "mlx/backend/common/unary.h" #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" @@ -60,15 +61,6 @@ static inline T dequantize_scale(uint8_t s) { } } -inline constexpr short get_pack_factor(int bits, int wsize = 8) { - return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits); -} - -inline constexpr short get_bytes_per_pack(int bits, int wsize = 8) { - auto power_of_2_bits = (bits & (bits - 1)) == 0; - return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3); -} - template void extract_bits(const uint8_t* w_in, T* w_out) { static_assert(bits == 3 || bits == 5 || bits == 6); @@ -1069,6 +1061,15 @@ uint8_t to_fp8_e8m0(float x) { return static_cast(n + 127); } +// Smallest E8M0 >= x, so a block's largest elements do not saturate. +uint8_t to_fp8_e8m0_round_up(float x) { + uint8_t bits = to_fp8_e8m0(x); + if (bits < 0xFE && dequantize_scale(bits) < x) { + bits += 1; + } + return bits; +} + uint8_t to_fp4_e2m1(float x) { if (std::isnan(x)) { return 0x7; @@ -1120,7 +1121,7 @@ void fp_quantize_dequantize( if (group_size == 16) { scale = dequantize_scale(detail::ToFP8()(scale)); } else { - scale = dequantize_scale(to_fp8_e8m0(scale)); + scale = dequantize_scale(to_fp8_e8m0_round_up(scale)); } for (int j = 0; j < group_size; ++j) { @@ -1367,4 +1368,8 @@ void QQMatmul::eval_cpu(const std::vector& inputs, array& out) { } } +void GatherQQMM::eval_cpu(const std::vector& inputs, array& out) { + throw std::runtime_error("[GatherQQMM] NYI"); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp index 41764f4c8..64b3fab33 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/reduce.cpp @@ -460,6 +460,9 @@ void Reduce::eval_cpu(const std::vector& inputs, array& out) { switch (reduce_type_) { case Reduce::And: case Reduce::Or: { + // Integers can be reduced as whatever type has the same width, since + // only their bits matter. Floats cannot: -0.0 compares equal to zero + // but has a bit set, so it has to be tested as a float. switch (in.dtype()) { case bool_: case uint8: @@ -468,21 +471,32 @@ void Reduce::eval_cpu(const std::vector& inputs, array& out) { break; case int16: case uint16: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case float16: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case bfloat16: - reduce_dispatch_and_or(in, out, reduce_type_, axes_); + reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; case uint32: case int32: - case float32: reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; + case float32: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; case uint64: case int64: - case float64: + // complex64 stays on the integer path. Testing it as a complex + // would go through complex64_t's conversion to float and only look + // at the real part, which would miss 1j. case complex64: reduce_dispatch_and_or(in, out, reduce_type_, axes_); break; + case float64: + reduce_dispatch_and_or(in, out, reduce_type_, axes_); + break; } break; } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp index 4dc2f50ea..ab2ca1436 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/scan.cpp @@ -163,7 +163,8 @@ void scan_op( const Op& op, U init) { if (in.flags().row_contiguous) { - if (in.strides()[axis] == 1) { + // A size-one axis can carry any stride and still be row contiguous. + if (in.strides()[axis] == 1 || in.shape(axis) == 1) { contiguous_scan( in.data(), out.data(), @@ -190,6 +191,19 @@ void scan_op( } } +template +U scan_init(const Dtype& dtype, bool maximum) { + constexpr auto inf = std::numeric_limits::infinity(); + if constexpr (std::is_same_v) { + return maximum ? complex64_t{inf, inf} : complex64_t{-inf, -inf}; + } else if (issubdtype(dtype, floating)) { + return maximum ? static_cast(inf) : static_cast(-inf); + } else { + return maximum ? std::numeric_limits::max() + : std::numeric_limits::min(); + } +} + template void scan_dispatch( Scan::ReduceType rtype, @@ -212,18 +226,28 @@ void scan_dispatch( break; } case Scan::Min: { - auto op = [](U y, T x) { return x < y ? x : y; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(std::numeric_limits::infinity()) - : std::numeric_limits::max(); + auto op = [](U y, T x) { + if constexpr (is_floating_point_v) { + if (std::isnan(y) || std::isnan(static_cast(x))) { + return std::numeric_limits::quiet_NaN(); + } + } + return x < y ? x : y; + }; + auto init = scan_init(in.dtype(), /* maximum = */ true); scan_op(in, out, axis, reverse, inclusive, op, init); break; } case Scan::Max: { - auto op = [](U y, T x) { return x < y ? y : x; }; - auto init = (issubdtype(in.dtype(), floating)) - ? static_cast(-std::numeric_limits::infinity()) - : std::numeric_limits::min(); + auto op = [](U y, T x) { + if constexpr (is_floating_point_v) { + if (std::isnan(y) || std::isnan(static_cast(x))) { + return std::numeric_limits::quiet_NaN(); + } + } + return x < y ? y : x; + }; + auto init = scan_init(in.dtype(), /* maximum = */ false); scan_op(in, out, axis, reverse, inclusive, op, init); break; } @@ -231,7 +255,7 @@ void scan_dispatch( auto op = [](U a, T b) { return detail::LogAddExp{}(a, static_cast(b)); }; - auto init = (issubdtype(in.dtype(), floating)) + auto init = (issubdtype(in.dtype(), inexact)) ? static_cast(-std::numeric_limits::infinity()) : std::numeric_limits::min(); scan_op(in, out, axis, reverse, inclusive, op, init); diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/select.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/select.cpp index bf6a9b825..89164f9dd 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/select.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/select.cpp @@ -4,6 +4,7 @@ #include "mlx/backend/cpu/binary_ops.h" #include "mlx/backend/cpu/ternary.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" namespace mlx::core { @@ -32,53 +33,10 @@ void select_op( out = array::unsafe_weak_copy(out), op, topt]() mutable { - switch (out.dtype()) { - case bool_: - ternary_op(a, b, c, out, op, topt); - break; - case uint8: - ternary_op(a, b, c, out, op, topt); - break; - case uint16: - ternary_op(a, b, c, out, op, topt); - break; - case uint32: - ternary_op(a, b, c, out, op, topt); - break; - case uint64: - ternary_op(a, b, c, out, op, topt); - break; - case int8: - ternary_op(a, b, c, out, op, topt); - break; - case int16: - ternary_op(a, b, c, out, op, topt); - break; - case int32: - ternary_op(a, b, c, out, op, topt); - break; - case int64: - ternary_op(a, b, c, out, op, topt); - break; - case float16: - ternary_op( - a, b, c, out, op, topt); - break; - case float32: - ternary_op(a, b, c, out, op, topt); - break; - case float64: - ternary_op(a, b, c, out, op, topt); - break; - case bfloat16: - ternary_op( - a, b, c, out, op, topt); - break; - case complex64: - ternary_op( - a, b, c, out, op, topt); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + ternary_op(a, b, c, out, op, topt); + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h index f62c67d38..e11cc7f5f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/accelerate_simd.h @@ -241,7 +241,7 @@ Simd remainder(Simd a, Simd b) { } else { r = a - b * (a / b); } - if constexpr (std::is_signed_v) { + if constexpr (is_signed_v) { auto mask = r != 0 && (r < 0 != b < 0); r = select(mask, r + b, r); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h index 775f5dfd1..1ae883f7e 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/base_simd.h @@ -13,6 +13,8 @@ #include // For _BitScanReverse #endif +#include "mlx/types/half_types.h" + namespace mlx::core::simd { template struct Simd; @@ -54,13 +56,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) { @@ -89,7 +84,6 @@ Simd recip(Simd in) { DEFAULT_UNARY(operator-, std::negate{}) DEFAULT_UNARY(operator!, std::logical_not{}) -DEFAULT_UNARY(abs, std::abs) DEFAULT_UNARY(acos, std::acos) DEFAULT_UNARY(acosh, std::acosh) DEFAULT_UNARY(asin, std::asin) @@ -108,6 +102,15 @@ DEFAULT_UNARY(sqrt, std::sqrt) DEFAULT_UNARY(tan, std::tan) DEFAULT_UNARY(tanh, std::tanh) +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) { @@ -210,7 +213,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; } @@ -250,6 +253,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; diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/simd/math.h b/Source/Cmlx/mlx/mlx/backend/cpu/simd/math.h index f9fc8317a..88676323a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/simd/math.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/simd/math.h @@ -115,12 +115,30 @@ Simd sincos(Simd in) { } } +// sincos reduces the argument by rounding x * 4/pi into a uint32, which is +// only exact while that product fits in a float's mantissa. Past that the +// result degrades and then leaves [-1, 1] altogether, so send those arguments +// to libm instead. 2^23 keeps x * 4/pi under 2^24 with room to spare. +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); } } @@ -129,7 +147,7 @@ Simd cos(Simd x) { if constexpr (is_complex) { return std::cos(x.value); } else { - return sincos(x); + return sincos_checked(x); } } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/softmax.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/softmax.cpp index 4c2941e96..d62e86a9c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/softmax.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/softmax.cpp @@ -6,6 +6,7 @@ #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" #include "mlx/primitives.h" #include "mlx/types/limits.h" @@ -139,32 +140,23 @@ void Softmax::eval_cpu(const std::vector& inputs, array& out) { auto in = set_output(inputs[0]); - switch (in.dtype()) { - case float32: - softmax(in, out, stream()); - break; - case float16: + dispatch_all_types(in.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr ( + std::is_same_v || std::is_same_v) { if (precise_) { - softmax(in, out, stream()); + softmax(in, out, stream()); } else { - softmax(in, out, stream()); + softmax(in, out, stream()); } - break; - case bfloat16: - if (precise_) { - softmax(in, out, stream()); - } else { - softmax(in, out, stream()); - } - break; - case float64: - softmax(in, out, stream()); - break; - default: + } else if constexpr ( + std::is_same_v || std::is_same_v) { + softmax(in, out, stream()); + } else { throw std::runtime_error( "[softmax] Only defined for floating point types."); - break; - } + } + }); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp index fcf12d7ad..767e7ef1d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/sort.cpp @@ -18,7 +18,7 @@ namespace { // NaN-aware comparator that places NaNs at the end template bool nan_aware_less(T a, T b) { - if constexpr (std::is_floating_point_v || std::is_same_v) { + if constexpr (is_floating_point_v || std::is_same_v) { if (std::isnan(a)) return false; if (std::isnan(b)) @@ -103,11 +103,11 @@ struct StridedIterator { return *this; } - StridedIterator operator+(difference_type diff) { + StridedIterator operator+(difference_type diff) const { return StridedIterator(ptr_, stride_, diff); } - StridedIterator operator-(difference_type diff) { + StridedIterator operator-(difference_type diff) const { return StridedIterator(ptr_, stride_, -diff); } @@ -198,7 +198,7 @@ void argsort(const array& in, array& out, int axis) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if (std::is_floating_point::value) { + if constexpr (is_floating_point_v) { if (std::isnan(v1)) return false; if (std::isnan(v2)) @@ -299,7 +299,7 @@ void argpartition(const array& in, array& out, int axis, int kth) { auto v2 = data_ptr[b * in_stride]; // Handle NaNs (place them at the end) - if (std::is_floating_point::value) { + if constexpr (is_floating_point_v) { if (std::isnan(v1)) return false; if (std::isnan(v2)) @@ -311,6 +311,43 @@ void argpartition(const array& in, array& out, int axis, int kth) { } } +template +void searchsorted(const array& a, const array& v, array& out) { + auto n = static_cast(a.size()); + auto a_stride = a.strides()[0]; // sequence is 1D + const T* a_ptr = a.data(); + const T* v_ptr = v.data(); + uint32_t* out_ptr = out.data(); + + auto bound = [a_ptr, a_stride, n](T x) { + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + (hi - lo) / 2; + T m = a_ptr[static_cast(mid) * a_stride]; + bool below = Right ? !nan_aware_less(x, m) : nan_aware_less(m, x); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + + if (v.flags().row_contiguous) { + for (size_t i = 0; i < v.size(); ++i) { + out_ptr[i] = bound(v_ptr[i]); + } + } else { + ContiguousIterator it(v); + for (size_t i = 0; i < v.size(); ++i) { + out_ptr[i] = bound(v_ptr[it.loc]); + it.step(); + } + } +} + } // namespace void ArgSort::eval_cpu(const std::vector& inputs, array& out) { @@ -326,36 +363,9 @@ void ArgSort::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([in = array::unsafe_weak_copy(in), out = array::unsafe_weak_copy(out), axis_ = axis_]() mutable { - switch (in.dtype()) { - case bool_: - return argsort(in, out, axis_); - case uint8: - return argsort(in, out, axis_); - case uint16: - return argsort(in, out, axis_); - case uint32: - return argsort(in, out, axis_); - case uint64: - return argsort(in, out, axis_); - case int8: - return argsort(in, out, axis_); - case int16: - return argsort(in, out, axis_); - case int32: - return argsort(in, out, axis_); - case int64: - return argsort(in, out, axis_); - case float32: - return argsort(in, out, axis_); - case float64: - return argsort(in, out, axis_); - case float16: - return argsort(in, out, axis_); - case bfloat16: - return argsort(in, out, axis_); - case complex64: - return argsort(in, out, axis_); - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + argsort(in, out, axis_); + }); }); } @@ -397,36 +407,9 @@ void ArgPartition::eval_cpu(const std::vector& inputs, array& out) { out = array::unsafe_weak_copy(out), axis_ = axis_, kth_ = kth_]() mutable { - switch (in.dtype()) { - case bool_: - return argpartition(in, out, axis_, kth_); - case uint8: - return argpartition(in, out, axis_, kth_); - case uint16: - return argpartition(in, out, axis_, kth_); - case uint32: - return argpartition(in, out, axis_, kth_); - case uint64: - return argpartition(in, out, axis_, kth_); - case int8: - return argpartition(in, out, axis_, kth_); - case int16: - return argpartition(in, out, axis_, kth_); - case int32: - return argpartition(in, out, axis_, kth_); - case int64: - return argpartition(in, out, axis_, kth_); - case float32: - return argpartition(in, out, axis_, kth_); - case float64: - return argpartition(in, out, axis_, kth_); - case float16: - return argpartition(in, out, axis_, kth_); - case bfloat16: - return argpartition(in, out, axis_, kth_); - case complex64: - return argpartition(in, out, axis_, kth_); - } + dispatch_all_types(in.dtype(), [&](auto type_tag) { + argpartition(in, out, axis_, kth_); + }); }); } @@ -445,36 +428,38 @@ void Partition::eval_cpu(const std::vector& inputs, array& out) { encoder.dispatch([out = array::unsafe_weak_copy(out), axis_ = axis_, kth_ = kth_]() mutable { - switch (out.dtype()) { - case bool_: - return partition(out, axis_, kth_); - case uint8: - return partition(out, axis_, kth_); - case uint16: - return partition(out, axis_, kth_); - case uint32: - return partition(out, axis_, kth_); - case uint64: - return partition(out, axis_, kth_); - case int8: - return partition(out, axis_, kth_); - case int16: - return partition(out, axis_, kth_); - case int32: - return partition(out, axis_, kth_); - case int64: - return partition(out, axis_, kth_); - case float32: - return partition(out, axis_, kth_); - case float64: - return partition(out, axis_, kth_); - case float16: - return partition(out, axis_, kth_); - case bfloat16: - return partition(out, axis_, kth_); - case complex64: - return partition(out, axis_, kth_); - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + partition(out, axis_, kth_); + }); + }); +} + +void SearchSorted::eval_cpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + auto& a = inputs[0]; + auto& v = inputs[1]; + + out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + auto& encoder = cpu::get_command_encoder(stream()); + encoder.set_input_array(a); + encoder.set_input_array(v); + encoder.set_output_array(out); + encoder.dispatch([a = array::unsafe_weak_copy(a), + v = array::unsafe_weak_copy(v), + out = array::unsafe_weak_copy(out), + right = right_]() mutable { + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if (right) { + searchsorted(a, v, out); + } else { + searchsorted(a, v, out); + } + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp b/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp index ca01a0a65..80101e907 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cpu/svd.cpp @@ -1,5 +1,7 @@ // Copyright © 2024 Apple Inc. +#include + #include "mlx/allocator.h" #include "mlx/backend/cpu/copy.h" #include "mlx/backend/cpu/encoder.h" @@ -89,7 +91,7 @@ struct SVDWork< if (info != 0) { std::stringstream ss; - ss << "svd_impl: sgesvdx_ failed with code " << info; + ss << "[SVD::eval_cpu] SVD failed with error code " << info; throw std::runtime_error(ss.str()); } } @@ -179,7 +181,7 @@ struct SVDWork> { if (info != 0) { std::stringstream ss; - ss << "svd_impl: sgesvdx_ failed with code " << info; + ss << "[SVD::eval_cpu] SVD failed with error code " << info; throw std::runtime_error(ss.str()); } } @@ -206,6 +208,33 @@ void svd_impl( using R = typename SVDWork::R; + // Nothing to decompose when either dimension is zero; LAPACK rejects + // lda = 0. The singular values are empty but the factors are not when only + // one of the dimensions is zero, so fill them with the identity like LAPACK + // does, which keeps them orthonormal. + if (M == 0 || N == 0) { + auto& encoder = cpu::get_command_encoder(stream); + for (auto& o : outputs) { + o.set_data(allocator::malloc(o.nbytes())); + if (o.size() == 0) { + continue; + } + encoder.set_output_array(o); + encoder.dispatch([o = array::unsafe_weak_copy(o)]() mutable { + auto ptr = o.data(); + const size_t size = o.size(); + const int n = o.shape(-1); + std::fill_n(ptr, size, T(0)); + for (size_t i = 0; i < size; i += static_cast(n) * n) { + for (int j = 0; j < n; ++j) { + ptr[i + static_cast(j) * n + j] = T(1); + } + } + }); + } + return; + } + size_t num_matrices = a.size() / (M * N); // lapack clobbers the input, so we have to make a copy. diff --git a/Source/Cmlx/mlx/mlx/backend/cpu/unary.h b/Source/Cmlx/mlx/mlx/backend/cpu/unary.h index 4fab6a754..096521df8 100644 --- a/Source/Cmlx/mlx/mlx/backend/cpu/unary.h +++ b/Source/Cmlx/mlx/mlx/backend/cpu/unary.h @@ -5,6 +5,7 @@ #include "mlx/backend/common/unary.h" #include "mlx/backend/cpu/encoder.h" #include "mlx/backend/cpu/simd/simd.h" +#include "mlx/dtype_utils.h" #include "mlx/utils.h" namespace mlx::core { @@ -61,50 +62,9 @@ void unary(const array& a, array& out, Op op, Stream stream) { encoder.dispatch([a = array::unsafe_weak_copy(a), out = array::unsafe_weak_copy(out), op = op]() mutable { - switch (out.dtype()) { - case bool_: - unary_op(a, out, op); - break; - case uint8: - unary_op(a, out, op); - break; - case uint16: - unary_op(a, out, op); - break; - case uint32: - unary_op(a, out, op); - break; - case uint64: - unary_op(a, out, op); - break; - case int8: - unary_op(a, out, op); - break; - case int16: - unary_op(a, out, op); - break; - case int32: - unary_op(a, out, op); - break; - case int64: - unary_op(a, out, op); - break; - case float16: - unary_op(a, out, op); - break; - case float32: - unary_op(a, out, op); - break; - case float64: - unary_op(a, out, op); - break; - case bfloat16: - unary_op(a, out, op); - break; - case complex64: - unary_op(a, out, op); - break; - } + dispatch_all_types(out.dtype(), [&](auto type_tag) { + unary_op(a, out, op); + }); }); } @@ -146,27 +106,9 @@ void unary_fp(const array& a, array& out, Op op, Stream stream) { encoder.dispatch([a = array::unsafe_weak_copy(a), out = array::unsafe_weak_copy(out), op = op]() mutable { - switch (out.dtype()) { - case bfloat16: - unary_op(a, out, op); - break; - case float16: - unary_op(a, out, op); - break; - case float32: - unary_op(a, out, op); - break; - case float64: - unary_op(a, out, op); - break; - case complex64: - unary_op(a, out, op); - break; - default: - std::ostringstream err; - err << "[unary_fp] Does not support " << out.dtype(); - throw std::runtime_error(err.str()); - } + dispatch_inexact_types(out.dtype(), "[unary_fp]", [&](auto type_tag) { + unary_op(a, out, op); + }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt index 013b24b2f..9c8d3174c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/cuda/CMakeLists.txt @@ -19,6 +19,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy.cu ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp @@ -26,10 +27,14 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp ${CMAKE_CURRENT_SOURCE_DIR}/event.cu + ${CMAKE_CURRENT_SOURCE_DIR}/fft.cu ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gemms/block_mask.cu ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gather_gemm.cu ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu + ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu @@ -54,15 +59,25 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/sort.cu ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cu ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/qmm.cu + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/qmm_naive.cu + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/qmm_sm80.cu + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/qmm_sm90.cu + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/qmv.cu + ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmm/fp_qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/affine_quantize.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/fp_quantize.cu - ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qmv.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/quantized.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm.cpp ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_utils.cu ${CMAKE_CURRENT_SOURCE_DIR}/quantized/convert_fp8.cu ${CMAKE_CURRENT_SOURCE_DIR}/worker.cpp) +# Put dynamic defines in the dirs.cpp file. +add_library(mlx_dirs OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/dirs.cpp) +target_include_directories(mlx_dirs PRIVATE "${PROJECT_SOURCE_DIR}") +target_link_libraries(mlx PRIVATE $) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/binary) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary) @@ -116,6 +131,17 @@ target_compile_options(mlx target_compile_options( mlx PRIVATE "$<$:--expt-relaxed-constexpr>") +# Ignore warnings from CUTLASS. +target_compile_options( + mlx + PRIVATE $<$:-Xcudafe="--diag_suppress=186,2908,2361">) + +if(NOT MSVC) + # Required for generating optimized CUTLASS code. + target_compile_options( + mlx PRIVATE "$<$:-Xcompiler=-fno-strict-aliasing>") +endif() + # Suppress nvcc warnings on C++ headers. target_compile_options( mlx @@ -140,12 +166,11 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES) COMMAND __nvcc_device_query OUTPUT_VARIABLE MLX_CUDA_ARCHITECTURES OUTPUT_STRIP_TRAILING_WHITESPACE) - set(UPGRADABLE_ARCHITECTURES "90;100;121") if(MLX_CUDA_ARCHITECTURES STREQUAL "") message( FATAL_ERROR "Can not get native CUDA arch, must set MLX_CUDA_ARCHITECTURES") - elseif(MLX_CUDA_ARCHITECTURES IN_LIST UPGRADABLE_ARCHITECTURES) + elseif(MLX_CUDA_ARCHITECTURES GREATER_EQUAL 90) # Use arch-specific compute capability whenever possible. set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a") endif() @@ -154,8 +179,32 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") -# Search CUDA libs from installed python packages. +# Configure Windows CUDA DLL loading. if(WIN32) + set(MLX_CUDA_BIN_DIR + "" + CACHE STRING "Directory containing CUDA DLLs for Windows delay-loading") + set(MLX_CUDNN_BIN_DIR + "" + CACHE STRING "Directory containing cuDNN DLLs for Windows delay-loading") + + # With MLX_LOAD_CUDA_LIBS_FROM_PYTHON, unset dirs use the wheel layout. + # Relative dirs are resolved from the MLX binary. + if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON) + if("${MLX_CUDA_BIN_DIR}" STREQUAL "") + set(MLX_CUDA_BIN_DIR "${CUDAToolkit_BIN_DIR}/x64") + endif() + if("${MLX_CUDNN_BIN_DIR}" STREQUAL "") + set(MLX_CUDNN_BIN_DIR "${CUDNN_BIN_DIR}") + endif() + endif() + + function(mlx_add_dir_definition name) + if(NOT "${${name}}" STREQUAL "") + target_compile_definitions(mlx_dirs PRIVATE ${name}="${${name}}") + endif() + endfunction() + # Resolve paths of unfound DLL at runtime. if(BUILD_SHARED_LIBS) target_link_libraries(mlx PRIVATE "delayimp.lib") @@ -176,11 +225,9 @@ if(WIN32) target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}") endforeach() # Pass the locations where CUDA DLLs are placed. - if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON) - target_compile_definitions( - mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64" - MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}") - endif() + foreach(dir_var MLX_CUDA_BIN_DIR MLX_CUDNN_BIN_DIR) + mlx_add_dir_definition(${dir_var}) + endforeach() else() # For POSIX we rely on RPATH to search for CUDA libs. if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON) @@ -189,9 +236,11 @@ else() APPEND PROPERTY INSTALL_RPATH # The paths here should match the install_requires in setup.py. + "$ORIGIN/../../nvidia/cu13/lib" "$ORIGIN/../../nvidia/cublas/lib" "$ORIGIN/../../nvidia/cuda_nvrtc/lib" "$ORIGIN/../../nvidia/cudnn/lib" + "$ORIGIN/../../nvidia/cufft/lib" "$ORIGIN/../../nvidia/nccl/lib") endif() endif() @@ -214,7 +263,7 @@ install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv # The binary of C++ tests will not be installed so it can not find the CCCL # headers, and we have to hard-code the path. if(MLX_BUILD_TESTS) - target_compile_definitions(mlx + target_compile_definitions(mlx_dirs PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include") endif() @@ -234,6 +283,9 @@ target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) # Use cublasLt. target_link_libraries(mlx PRIVATE CUDA::cublasLt) +# Use cuFFT. +target_link_libraries(mlx PRIVATE CUDA::cufft) + # Use NVRTC and driver APIs. target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver) @@ -249,7 +301,7 @@ set(CUDNN_FRONTEND_BUILD_SAMPLES OFF) set(CUDNN_FRONTEND_BUILD_TESTS OFF) set(CUDNN_FRONTEND_BUILD_PYTHON_BINDINGS OFF) FetchContent_MakeAvailable(cudnn) -target_link_libraries(mlx PRIVATE cudnn_frontend) +target_link_libraries(mlx PRIVATE $) # Link with the actual cuDNN libraries. target_link_libraries(mlx PRIVATE CUDNN::cudnn_all) @@ -257,9 +309,15 @@ target_link_libraries(mlx PRIVATE CUDNN::cudnn_all) FetchContent_Declare( cutlass GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git - GIT_TAG v4.3.5 + GIT_TAG v4.4.2 GIT_SHALLOW TRUE SOURCE_SUBDIR include EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(cutlass) target_include_directories( mlx SYSTEM PRIVATE $) + +# Install CUTLASS headers for JIT. +install(DIRECTORY ${cutlass_SOURCE_DIR}/include/cute + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) +install(DIRECTORY ${cutlass_SOURCE_DIR}/include/cutlass + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp index 0337a27b9..b04d72da0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/allocator.cpp @@ -28,40 +28,22 @@ constexpr int small_block_size = 8; // size and small_block_size. constexpr int small_pool_size = 4 * page_size; -// Check if running on Windows or Windows Subsystem for Linux -bool is_windows() { -#if defined(_WIN32) - return true; -#elif defined(__linux__) - // WSL kernels contain "microsoft" or "WSL" in /proc/version - static bool is_wsl = []() { - std::ifstream version("/proc/version"); - if (version.is_open()) { - std::string line; - std::getline(version, line); - return line.find("microsoft") != std::string::npos || - line.find("Microsoft") != std::string::npos || - line.find("WSL") != std::string::npos; - } - return false; - }(); - return is_wsl; -#else - return false; -#endif -} - bool supports_managed_memory() { static bool managed_memory = []() { int device_count = gpu::device_count(); for (int i = 0; i < device_count; ++i) { auto& d = cu::device(i); - if (!d.managed_memory()) { - return false; - } // Empirically on Windows (and WSL) if there is no concurrentManagedAccess // the managed memory also does not work. - if (is_windows() && !d.concurrent_managed_access()) { + // The same has been observed on NVIDIA Tegra, typically on Jetson Orin + // Nano boards. + // NVIDIA documentation describes Windows, WSL, and Tegra as having a + // specific unified memory paradigm called "Limited unified memory + // support", which corresponds to concurrentManagedAccess = 0. + // See: + // https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/understanding-memory.html#table-unified-memory-levels + // https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/unified-memory.html#um-legacy-devices + if (!d.concurrent_managed_access()) { return false; } } @@ -168,9 +150,6 @@ CudaAllocator::CudaAllocator() free_limit_ = total_memory_ - memory_limit_; max_pool_size_ = memory_limit_; - int curr; - CHECK_CUDA_ERROR(cudaGetDevice(&curr)); - int device_count = gpu::device_count(); free_streams_.resize(device_count); mem_pools_.resize(device_count); @@ -181,7 +160,6 @@ CudaAllocator::CudaAllocator() CHECK_CUDA_ERROR(cudaDeviceGetDefaultMemPool(&mem_pools_[i], i)); } } - CHECK_CUDA_ERROR(cudaSetDevice(curr)); } Buffer @@ -223,6 +201,7 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) { if (device == -1) { data = unified_malloc(size); } else { + cu::device(device).make_current(); if (mem_pools_[device]) { // supports memory pools CHECK_CUDA_ERROR(cudaMallocAsync(&data, size, stream)); } else { @@ -321,7 +300,7 @@ void CudaAllocator::move_to_unified_memory( buf.device = -1; } -// This must be called with mutex_ aquired +// This must be called with mutex_ acquired void CudaAllocator::free_cuda_buffer(CudaBuffer* buf) { if (scalar_pool_.in_pool(buf)) { scalar_pool_.free(buf); @@ -419,6 +398,10 @@ void* Buffer::raw_ptr() { return cbuf.data; } +bool can_reuse_alien_buffer(void* ptr) { + return true; +} + } // namespace allocator size_t get_active_memory() { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/arange.cu b/Source/Cmlx/mlx/mlx/backend/cuda/arange.cu index f4f4326a3..a7986c7ef 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/arange.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/arange.cu @@ -56,7 +56,6 @@ void Arange::eval_gpu(const std::vector& inputs, array& out) { cu::arange, num_blocks, block_dims, - 0, gpu_ptr(out), out.data_size(), static_cast(start_), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu index bc891f5aa..9609b8597 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/arg_reduce.cu @@ -172,7 +172,6 @@ void ArgReduce::eval_gpu(const std::vector& inputs, array& out) { kernel, num_blocks, block_dim(), - 0, gpu_ptr(in), gpu_ptr(out), out.size(), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh index d43027c13..6130f63c7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/binary/binary.cuh @@ -16,8 +16,14 @@ namespace cu { namespace cg = cooperative_groups; +constexpr int BINARY_MAX_BLOCK_DIM = 1024; + template -__global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) { +__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_ss( + const In* a, + const In* b, + Out* out, + IdxT size) { IdxT index = cg::this_grid().thread_rank(); if ((index + 1) * N_READS > size) { @@ -36,7 +42,11 @@ __global__ void binary_ss(const In* a, const In* b, Out* out, IdxT size) { } template -__global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) { +__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_sv( + const In* a, + const In* b, + Out* out, + IdxT size) { IdxT index = cg::this_grid().thread_rank(); if ((index + 1) * N_READS > size) { @@ -57,7 +67,11 @@ __global__ void binary_sv(const In* a, const In* b, Out* out, IdxT size) { } template -__global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) { +__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vs( + const In* a, + const In* b, + Out* out, + IdxT size) { IdxT index = cg::this_grid().thread_rank(); if ((index + 1) * N_READS > size) { @@ -78,7 +92,11 @@ __global__ void binary_vs(const In* a, const In* b, Out* out, IdxT size) { } template -__global__ void binary_vv(const In* a, const In* b, Out* out, IdxT size) { +__global__ __launch_bounds__(BINARY_MAX_BLOCK_DIM) void binary_vv( + const In* a, + const In* b, + Out* out, + IdxT size) { IdxT index = cg::this_grid().thread_rank(); if ((index + 1) * N_READS > size) { @@ -117,7 +135,9 @@ __global__ void binary_g_nd( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -155,7 +175,9 @@ __global__ void binary_g( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -265,10 +287,8 @@ void binary_op_gpu_inplace( if (dim0 >= 4) { work_per_thread = 4; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); if (ndim <= 3) { dispatch_1_2_3(ndim, [&](auto dims_constant) { auto kernel = cu::binary_g_nd< @@ -289,9 +309,8 @@ void binary_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out), @@ -307,9 +326,8 @@ void binary_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out), @@ -333,12 +351,16 @@ void binary_op_gpu_inplace( kernel = cu::binary_vv; } auto [num_blocks, block_dims] = get_launch_args( - out.data_size(), out.shape(), out.strides(), large(), N_READS); + out.data_size(), + out.shape(), + out.strides(), + large(), + N_READS, + cu::BINARY_MAX_BLOCK_DIM); encoder.add_kernel_node( kernel, num_blocks, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/binary_two.cu b/Source/Cmlx/mlx/mlx/backend/cuda/binary_two.cu index 4acd30485..6374376fe 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/binary_two.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/binary_two.cu @@ -146,7 +146,9 @@ __global__ void binary_two_g_nd( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -189,7 +191,9 @@ __global__ void binary_two_g( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -287,10 +291,8 @@ void binary_two_op_gpu_inplace( if (dim0 >= 4) { work_per_thread = 4; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); if (ndim <= 3) { dispatch_1_2_3(ndim, [&](auto dims_constant) { @@ -312,9 +314,8 @@ void binary_two_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out_a), @@ -331,9 +332,8 @@ void binary_two_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out_a), @@ -367,7 +367,6 @@ void binary_two_op_gpu_inplace( kernel, num_blocks, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(out_a), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/compiled.cpp index 29cf2e4fe..a363cf9c3 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/compiled.cpp @@ -220,8 +220,6 @@ constexpr const char* g_jit_includes = R"( #include "mlx/backend/cuda/device/utils.cuh" #include - -#define inf cuda::std::numeric_limits::infinity() )"; void Compiled::eval_gpu( @@ -239,7 +237,9 @@ void Compiled::eval_gpu( } int work_per_thread = 16 / max_size; - cu::JitModule& mod = cu::get_jit_module(s.device, lib_name(), [&]() { + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), lib_name(), [&]() { // Build source code. cu::FusedKernelBuilder builder{ g_jit_includes, lib_name(), inputs_, outputs_, tape_, is_constant_}; @@ -285,11 +285,12 @@ void Compiled::eval_gpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides_vec] = + auto [contiguous, negative_strides, shape, strides_vec] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); - // Whether to use large index. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = + negative_strides || compiled_use_large_index(inputs, outputs, contiguous); cu::KernelArgs args; // Put inputs. @@ -305,8 +306,6 @@ void Compiled::eval_gpu( } } - auto& encoder = cu::get_command_encoder(s); - // Put outputs. compiled_allocate_outputs( inputs, outputs, is_constant_, contiguous, [&](auto n) { @@ -351,7 +350,8 @@ void Compiled::eval_gpu( auto [kernel, max_block_dims] = mod.get_kernel_and_dims(kernel_name); auto [num_blocks, block_dims] = get_launch_args(outputs[0], large, work_per_thread, max_block_dims); - encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args()); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/conv.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/conv.cpp index 3b86bf28d..9e471d235 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/conv.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/conv.cpp @@ -39,7 +39,7 @@ struct ConvCacheKey { }; auto& conv_cache() { - static LRUBytesKeyCache< + static thread_local LRUBytesKeyCache< ConvCacheKey, std::pair>> cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128); @@ -103,7 +103,7 @@ std::optional build_conv_graph( const std::vector& dilation) { auto compute_dtype = (dtype == float16 || dtype == bfloat16) ? float32 : dtype; - DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype); + DnnGraph graph(get_cudnn_handle(encoder.device()), dtype, compute_dtype); auto x_ = graph.tensor_nchw("X", 'x', x); auto w_ = graph.tensor_nchw("W", 'w', w); @@ -139,7 +139,7 @@ std::optional build_conv_graph( if (dtype == float32 && !env::enable_tf32()) { graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE}); } - CHECK_CUDNN_FE_ERROR(graph.build()); + CHECK_CUDNN_ERROR(graph.build()); return graph; } @@ -252,6 +252,10 @@ void register_args( } // namespace +void init_cudnn_conv_cache() { + conv_cache(); +} + void Convolution::eval_gpu(const std::vector& inputs, array& out_) { nvtx3::scoped_range r("Convolution::eval_gpu"); if (out_.size() == 0) { @@ -269,20 +273,19 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out_) { // Search cache. BytesKey cache_key; - cache_key.pod = { - encoder.device().cuda_device(), - dtype_to_cudnn_type(dtype), - vector_key(in.shape()), - vector_key(wt.shape()), - vector_key(kernel_strides_), - vector_key(padding_lo_), - vector_key(padding_hi_), - vector_key(kernel_dilation_), - groups_, - flip_, - get_alignment(in), - get_alignment(wt), - get_alignment(out)}; + cache_key.pod.device_id = encoder.device().cuda_device(); + cache_key.pod.cudnn_dtype = dtype_to_cudnn_type(dtype); + cache_key.pod.input_shape = vector_key(in.shape()); + cache_key.pod.weight_shape = vector_key(wt.shape()); + cache_key.pod.stride = vector_key(kernel_strides_); + cache_key.pod.padding_lo = vector_key(padding_lo_); + cache_key.pod.padding_hi = vector_key(padding_hi_); + cache_key.pod.dilation = vector_key(kernel_dilation_); + cache_key.pod.groups = groups_; + cache_key.pod.flip = flip_; + cache_key.pod.input_alignment = get_alignment(in); + cache_key.pod.weight_alignment = get_alignment(wt); + cache_key.pod.output_alignment = get_alignment(out); if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) { auto& [backend_type, graph] = it->second; if (graph) { @@ -290,7 +293,7 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out_) { std::tie(in, wt, out) = prepare_args(encoder, backend_type, in, wt, out, groups_, s); register_args(encoder, backend_type, in, wt, out, out_); - CHECK_CUDNN_FE_ERROR(graph->encode_capturing( + CHECK_CUDNN_ERROR(graph->encode_capturing( encoder, { {'x', gpu_ptr(in)}, @@ -372,7 +375,7 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out_) { if (graph) { register_args(encoder, backend_type, in, wt, out, out_); - CHECK_CUDNN_FE_ERROR(graph->encode_capturing( + CHECK_CUDNN_ERROR(graph->encode_capturing( encoder, { {'x', gpu_ptr(in)}, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_conv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_conv.cu index a869542c9..af15ace1a 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_conv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_conv.cu @@ -4,6 +4,7 @@ #include "mlx/backend/cuda/gemms/cublas_gemm.h" #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/dtype_utils.h" +#include "mlx/utils.h" #include @@ -24,17 +25,19 @@ __global__ void naive_unfold_nd( auto tid = block.group_index(); auto lid = block.thread_index(); - int index_batch = tid.z / out_pixels; // [0, N) - int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out) + // The matrix-row index (N * out_pixels values) rides grid.x, which allows + // up to 2^31-1 blocks; grid.y and grid.z are limited to 65,535. + int index_batch = tid.x / out_pixels; // [0, N) + int index_out_spatial = tid.x % out_pixels; // [0, H_out * W_out) int index_wt_spatial = - tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) + tid.z * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) if (index_wt_spatial >= filter_size / params.C) { return; } in += tid.y; // [0, C) - out += tid.z * filter_size + index_wt_spatial * params.C + tid.y; + out += tid.x * filter_size + index_wt_spatial * params.C + tid.y; bool valid = index_batch < params.N; @@ -105,9 +108,12 @@ array unfold_inputs_nd( dim3 block_dims; block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024); dim3 num_blocks; - num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x); + // mat_M = N * out_pixels can exceed 65,535 (a 256x256 output already has + // 65,536 positions), so it must ride grid.x; the handful of filter-spatial + // blocks go on grid.z. grid.y and grid.z are limited to 65,535. + num_blocks.x = mat_M; num_blocks.y = params.C; - num_blocks.z = mat_M; + num_blocks.z = cuda::ceil_div(wt_spatial_size, block_dims.x); encoder.set_input_array(in); encoder.set_output_array(unfolded); @@ -117,7 +123,6 @@ array unfold_inputs_nd( cu::naive_unfold_nd, num_blocks, block_dims, - 0, gpu_ptr(in), gpu_ptr(unfolded), filter_size, @@ -137,8 +142,8 @@ void gemm_conv_nd( ConvParams& params, Stream s) { // Get gemm shapes. - int mat_M = out.size() / params.O; // N * H_out * W_out - int mat_K = wt.size() / params.O; // C * H_wt * W_wt + int mat_M = safe_cast(out.size() / params.O, "conv"); // N * H_out * W_out + int mat_K = safe_cast(wt.size() / params.O, "conv"); // C * H_wt * W_wt int mat_N = params.O; // O // Unfold input to (N * H_out * W_out, C * H_wt * W_wt) for gemm. diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu index 44e067c2e..82d32aae7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu @@ -4,6 +4,7 @@ #include "mlx/backend/cuda/gemms/cublas_gemm.h" #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/dtype_utils.h" +#include "mlx/utils.h" #include @@ -24,17 +25,19 @@ __global__ void naive_grouped_unfold_transpose_nd( auto tid = block.group_index(); auto lid = block.thread_index(); - int index_batch = tid.z / out_pixels; // [0, N) - int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out) + // The matrix-row index (N * out_pixels values) rides grid.x, which allows + // up to 2^31-1 blocks; grid.y and grid.z are limited to 65,535. + int index_batch = tid.x / out_pixels; // [0, N) + int index_out_spatial = tid.x % out_pixels; // [0, H_out * W_out) int index_wt_spatial = - tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) + tid.z * block.dim_threads().x + lid.x; // [0, H_wt * W_wt) if (index_wt_spatial >= filter_size / params.C) { return; } in += tid.y; // [0, C) - out += tid.z * filter_size + tid.y * (filter_size / params.C); + out += tid.x * filter_size + tid.y * (filter_size / params.C); bool valid = index_batch < params.N; @@ -108,9 +111,12 @@ array grouped_unfold_transpose_inputs_nd( dim3 block_dims; block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024); dim3 num_blocks; - num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x); + // mat_M = N * out_pixels can exceed 65,535 (a 256x256 output already has + // 65,536 positions), so it must ride grid.x; the handful of filter-spatial + // blocks go on grid.z. grid.y and grid.z are limited to 65,535. + num_blocks.x = mat_M; num_blocks.y = params.C; - num_blocks.z = mat_M; + num_blocks.z = cuda::ceil_div(wt_spatial_size, block_dims.x); encoder.set_input_array(in); encoder.set_output_array(unfolded); @@ -120,7 +126,6 @@ array grouped_unfold_transpose_inputs_nd( cu::naive_grouped_unfold_transpose_nd, num_blocks, block_dims, - 0, gpu_ptr(in), gpu_ptr(unfolded), filter_size, @@ -142,8 +147,9 @@ void gemm_grouped_conv_nd( // Get gemm shapes. int C_per_group = params.C / params.groups; int O_per_group = params.O / params.groups; - int mat_M = out.size() / params.O; // N * H_out * W_out - int mat_K = wt.size() / params.O; // C_per_group * H_wt * W_wt + int mat_M = safe_cast(out.size() / params.O, "conv"); // N * H_out * W_out + int mat_K = + safe_cast(wt.size() / params.O, "conv"); // C_per_group * H_wt * W_wt int mat_N = O_per_group; // O_per_group // Unfold input to (N * H_out * W_out, C * H_wt * W_wt) for gemm. diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_contiguous.cu b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_contiguous.cu index 6f51734b1..b26b1ce1b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_contiguous.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_contiguous.cu @@ -76,7 +76,6 @@ void copy_contiguous( kernel, num_blocks, block_dims, - 0, gpu_ptr(in) + in_offset, gpu_ptr(out) + out_offset, out.data_size()); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general.cu b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general.cu index fec110a86..1277ba028 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general.cu @@ -21,7 +21,9 @@ __global__ void copy_gg_nd( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -59,7 +61,9 @@ __global__ void copy_gg( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -120,10 +124,8 @@ void copy_general( work_per_thread = 4; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); if (ndim <= 3) { dispatch_1_2_3(ndim, [&](auto ndim_constant) { @@ -135,9 +137,8 @@ void copy_general( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, in_ptr, out_ptr, rest, @@ -152,9 +153,8 @@ void copy_general( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, in_ptr, out_ptr, rest, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu index 96cf6fb5c..26c104739 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_dynamic.cu @@ -83,7 +83,6 @@ void copy_general_dynamic( dims_constant()>, num_blocks, block_dims, - 0, in_ptr, out_ptr, out.size(), @@ -99,7 +98,6 @@ void copy_general_dynamic( cu::copy_gg_dynamic, num_blocks, block_dims, - 0, in_ptr, out_ptr, out.size(), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_input.cu b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_input.cu index b1b793970..66e2454d9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_input.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/copy/copy_general_input.cu @@ -21,7 +21,9 @@ __global__ void copy_g_nd( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -53,7 +55,9 @@ __global__ void copy_g( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -154,7 +158,6 @@ void copy_general_input( kernel, {num_blocks_x, num_blocks_y}, block_dims, - 0, in_ptr, out_ptr, int64_t(shape[0]), @@ -175,10 +178,8 @@ void copy_general_input( } else if (dim0 < 4) { work_per_thread = 1; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); if (ndim <= 3) { dispatch_1_2_3(ndim, [&](auto dims_constant) { @@ -193,9 +194,8 @@ void copy_general_input( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, in_ptr, out_ptr, rest, @@ -211,9 +211,8 @@ void copy_general_input( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, in_ptr, out_ptr, rest, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cross_entropy.cu b/Source/Cmlx/mlx/mlx/backend/cuda/cross_entropy.cu new file mode 100644 index 000000000..23b0689be --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cross_entropy.cu @@ -0,0 +1,245 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/device/cast_op.cuh" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/gpu/copy.h" +#include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" + +#include +#include +#include + +#include + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +// fused together logsumexp + gather +// cast to float32 inside the kernel +// to avoid logits.astype(mx.float32) +// for each row: loss = logsumexp(x) - x_t +// first we accumulate logsumexp, then we do a gather +template +__global__ void cross_entropy( + const T* x, // [M, N] + const int* y, // [M,] + float* loss, // [M,] <- will be always in fp32 lse - x + int axis_size // N +) { + cg::greater max_op; + cg::plus plus_op; + + float prevmax; + float curmax = Limits::finite_min(); + float normalizer = 0; + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + x += grid.block_rank() * axis_size; // offset input + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); + auto vals = load_vector(x, index, axis_size, Limits::min()); + prevmax = curmax; +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + curmax = max_op(curmax, static_cast(vals[i])); + } + // scale already accumulated normiliser + normalizer = normalizer * __expf(prevmax - curmax); + // add vals scaled by curmax +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + normalizer += __expf(static_cast(vals[i]) - curmax); + } + } + prevmax = curmax; + curmax = cg::reduce(warp, curmax, max_op); + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // second reduce in a block + __shared__ float warp_max[WARP_SIZE]; + __shared__ float warp_normaliser[WARP_SIZE]; + + if (warp.thread_rank() == 0) { + warp_max[warp.meta_group_rank()] = curmax; + warp_normaliser[warp.meta_group_rank()] = normalizer; + } + block.sync(); + bool is_valid = warp.thread_rank() < warp.meta_group_size(); + curmax = + is_valid ? warp_max[warp.thread_rank()] : Limits::finite_min(); + prevmax = curmax; + curmax = + cg::reduce(warp, curmax, max_op); // max within a block (global row max) + normalizer = is_valid ? warp_normaliser[warp.thread_rank()] : 0.0f; + normalizer = normalizer * __expf(prevmax - curmax); + normalizer = cg::reduce(warp, normalizer, plus_op); + // gather and writing the output: + auto row = grid.block_rank(); + if (block.thread_rank() == 0) { + float gap = curmax - static_cast(x[y[row]]); + loss[row] = isinf(curmax) ? gap : log(normalizer) + gap; + } +} + +// get loss from the forward +template +__global__ void cross_entropy_vjp( + const T* x, // [M, N] + const int* y, // [M,] + const float* loss, // [M,] + const float* gy, // cotangent [M,] + T* grads, // [M, N] lse is accumulated in float, x is casted to float + int axis_size // N +) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto row = grid.block_rank(); + + x += row * axis_size; // offset input + grads += row * axis_size; // offset output + auto y_n = y[row]; // target index [0, N) + auto g = gy[row]; // cotangent + auto loss_n = loss[row]; + auto x_t = static_cast(x[y_n]); + block.sync(); + for (int r = 0; r < cuda::ceil_div(axis_size, BLOCK_DIM * N_READS); r++) { + auto index = r * BLOCK_DIM + block.thread_rank(); // [0, N) + auto vals = load_vector(x, index, axis_size, T{}); +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + int col = index * N_READS + i; + float val = __expf((static_cast(vals[i]) - x_t) - loss_n); + vals[i] = static_cast(g * (val - (col == y_n ? 1.0f : 0.0f))); + } + store_vector(grads, index, vals, axis_size); + } +} +} // namespace cu + +namespace fast { + +bool CrossEntropy::use_fallback(Stream s) { + return s.device == Device::cpu; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropy::eval_gpu"); + assert(inputs.size() == 2); // logits and target + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + encoder.add_temporary(x_copy); + return x_copy; + } + }; + auto in = ensure_row_contiguous(inputs[0]); // [n_rows, V] + auto target = ensure_row_contiguous(inputs[1]); // [n_rows,] + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows] in fp32 + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(out), + axis_size); + }); + }); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + nvtx3::scoped_range r("CrossEntropyVJP::eval_gpu"); + assert(inputs.size() == 4); // logits, target, loss, cotangent + auto& s = stream(); + auto& out = outputs[0]; + auto& encoder = cu::get_command_encoder(s); + auto ensure_row_contiguous = [&s, &encoder](const array& x) { + if (x.flags().row_contiguous) { + return x; + } else { + array x_copy = contiguous_copy_gpu(x, s); + encoder.add_temporary(x_copy); + return x_copy; + } + }; + + auto check_input = [&s](const array& x, bool& copied) { + if (x.flags().row_contiguous) { + copied = false; + return x; + } + copied = true; + return contiguous_copy_gpu(x, s); + }; + bool donate_x = inputs[0].is_donatable(); + bool copied; + auto in = check_input(inputs[0], copied); // [n_rows, V] + donate_x |= copied; + auto target = ensure_row_contiguous(inputs[1]); // [n_rows,] + auto loss = ensure_row_contiguous(inputs[2]); // [n_rows,] fp32 + auto cotan = ensure_row_contiguous(inputs[3]); // [n_rows,] fp32 + if (donate_x) { + out.copy_shared_buffer(in); + } else { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); // [n_rows, V] + } + + int axis_size = in.shape().back(); + int n_rows = in.data_size() / axis_size; + + encoder.set_input_array(in); + encoder.set_input_array(target); + encoder.set_input_array(loss); + encoder.set_input_array(cotan); + encoder.set_output_array(out); + dispatch_float_types(in.dtype(), "cross_entropy_vjp", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) { + auto kernel = cu::cross_entropy_vjp; + encoder.add_kernel_node( + kernel, + n_rows, + block_dim(), + gpu_ptr(in), + gpu_ptr(target), + gpu_ptr(loss), + gpu_ptr(cotan), + gpu_ptr(out), + axis_size); + }); + }); +} + +} // namespace fast + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.cpp index 1176bd49d..8a03eccde 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.cpp @@ -2,43 +2,12 @@ #include "mlx/backend/cuda/cublas_utils.h" #include "mlx/backend/cuda/cuda.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/utils.h" namespace mlx::core { -namespace cublas_utils { - -namespace { - -struct CublasPreference { - CublasPreference(cu::Device& device) { - // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 MiB - // for Hopper+: - // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace - uint64_t MiB = 1024 * 1024; - uint64_t workspace_size = - device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB; - - CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&pref_)); - CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute( - pref_, - CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, - &workspace_size, - sizeof(uint64_t))); - } - - ~CublasPreference() { - CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref_)); - } - - cublasLtMatmulPreference_t pref_{nullptr}; -}; -} // namespace - -cublasLtMatmulPreference_t get_preference(cu::Device& device) { - static CublasPreference pref(device); - return pref.pref_; -} +namespace cublas_utils { cublasLtMatrixLayout_t create_matrix_layout( cudaDataType_t type, @@ -70,6 +39,59 @@ cublasLtMatrixLayout_t create_matrix_layout( } // namespace cublas_utils +namespace { + +auto& cublas_handles_cache() { + struct CublasHandles { + ~CublasHandles() { + if (handle) { + CHECK_CUBLAS_ERROR(cublasLtDestroy(handle)); + CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceDestroy(pref)); + } + } + cublasLtHandle_t handle{nullptr}; + cublasLtMatmulPreference_t pref{nullptr}; + }; + static thread_local std::vector cache(gpu::device_count()); + return cache; +} + +auto get_cublas_handles(cu::Device& device) { + auto& storage = cublas_handles_cache().at(device.cuda_device()); + if (!storage.handle) { + // Create cublasLt handle. + device.make_current(); + CHECK_CUBLAS_ERROR(cublasLtCreate(&storage.handle)); + // The recommended cublas workspace size is 4 MiB for pre-Hopper and 32 + // MiB for Hopper+: + // https://docs.nvidia.com/cuda/cublas/#cublassetworkspace + uint64_t MiB = 1024 * 1024; + uint64_t workspace_size = + device.compute_capability_major() >= 9 ? 32 * MiB : 4 * MiB; + CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&storage.pref)); + CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceSetAttribute( + storage.pref, + CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, + sizeof(uint64_t))); + } + return std::make_tuple(storage.handle, storage.pref); +} + +} // namespace + +void check_cublas_error(const char* name, cublasStatus_t err) { + if (err != CUBLAS_STATUS_SUCCESS) { + // TODO: Use cublasGetStatusString when it is widely available. + throw std::runtime_error( + fmt::format("{} failed with code: {}.", name, static_cast(err))); + } +} + +void init_cublas_handles_cache() { + cublas_handles_cache(); +} + CublasMatmulBase::~CublasMatmulBase() { CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(a_desc_)); CHECK_CUBLAS_ERROR(cublasLtMatrixLayoutDestroy(b_desc_)); @@ -98,20 +120,12 @@ void CublasMatmulBase::init_base( M_ = a_rows; N_ = b_cols; scale_type_ = scale_type; - handle_ = device.get_cublaslt_handle(); - pref_ = cublas_utils::get_preference(device); + std::tie(handle_, pref_) = get_cublas_handles(device); heuristic_.state = CUBLAS_STATUS_NOT_INITIALIZED; CHECK_CUBLAS_ERROR( cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type)); - int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST; - CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( - matmul_desc_, - CUBLASLT_MATMUL_DESC_POINTER_MODE, - &pointer_mode, - sizeof(int32_t))); - // In cublasLt matrices use column-major layout, while it is possible to use // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias // epilogue does not work with the option. So instead we swap A and B to make diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.h index 56702b121..053fa6b08 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cublas_utils.h @@ -1,17 +1,15 @@ // Copyright © 2025 Apple Inc. #pragma once -#include #include "mlx/array.h" #include "mlx/backend/cuda/device.h" #include "mlx/dtype_utils.h" +#include + namespace mlx::core { namespace cublas_utils { -// Get the shared cublas preference for a device -cublasLtMatmulPreference_t get_preference(cu::Device& device); - cublasLtMatrixLayout_t create_matrix_layout( cudaDataType_t type, uint64_t rows, @@ -42,6 +40,12 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) { } // namespace cublas_utils +void check_cublas_error(const char* name, cublasStatus_t err); + +#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd)) + +void init_cublas_handles_cache(); + class CublasMatmulBase { public: virtual ~CublasMatmulBase(); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h index 4c60fec29..f8a234ee6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cuda_utils.h @@ -2,23 +2,17 @@ #pragma once -#include #include #include -#include namespace mlx::core { // Throw exception if the cuda API does not succeed. -void check_cublas_error(const char* name, cublasStatus_t err); void check_cuda_error(const char* name, cudaError_t err); void check_cuda_error(const char* name, CUresult err); -void check_cudnn_error(const char* name, cudnnStatus_t err); // The macro version that prints the command that failed. -#define CHECK_CUBLAS_ERROR(cmd) check_cublas_error(#cmd, (cmd)) #define CHECK_CUDA_ERROR(cmd) check_cuda_error(#cmd, (cmd)) -#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd)) // Base class for RAII managed CUDA resources. template @@ -56,6 +50,12 @@ class CudaHandle { } } + Handle release() { + Handle handle = handle_; + handle_ = nullptr; + return handle; + } + operator Handle() const { return handle_; } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.cpp index af2647f83..eee33068c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.cpp @@ -2,6 +2,7 @@ #include "mlx/backend/cuda/cudnn_utils.h" #include "mlx/backend/cuda/device.h" +#include "mlx/backend/gpu/device_info.h" namespace mlx::core { @@ -47,8 +48,48 @@ inline auto nhwc_to_nchw(const array& x) { return std::make_tuple(std::move(shape), std::move(strides)); } +auto& cudnn_handles_cache() { + struct CudnnHandle { + ~CudnnHandle() { + if (handle) { + CHECK_CUDNN_ERROR(cudnnDestroy(handle)); + } + } + cudnnHandle_t handle{nullptr}; + }; + static thread_local std::vector cache(gpu::device_count()); + return cache; +} + } // namespace +void check_cudnn_error(const char* name, cudnnStatus_t err) { + if (err != CUDNN_STATUS_SUCCESS) { + throw std::runtime_error( + fmt::format("{} failed: {}.", name, cudnnGetErrorString(err))); + } +} + +void check_cudnn_error(const char* name, fe::error_t err) { + if (!err.is_good()) { + throw std::runtime_error( + fmt::format("{} failed: {}.", name, err.get_message())); + } +} + +cudnnHandle_t get_cudnn_handle(cu::Device& device) { + auto& storage = cudnn_handles_cache().at(device.cuda_device()); + if (!storage.handle) { + device.make_current(); + CHECK_CUDNN_ERROR(cudnnCreate(&storage.handle)); + } + return storage.handle; +} + +void init_cudnn_handles_cache() { + cudnn_handles_cache(); +} + fe::error_t DnnGraph::prepare() { RETURN_IF_ERROR(validate()); try { @@ -71,10 +112,26 @@ fe::error_t DnnGraph::encode_graph( cu::CommandEncoder& encoder, std::unordered_map variant_pack) { cudnnSetStream(handle_, encoder.stream()); - CudaGraph cuda_graph(encoder.device()); - RETURN_IF_ERROR(populate_cuda_graph( - handle_, variant_pack, prepare_workspace(encoder), cuda_graph)); - encoder.add_graph_node(cuda_graph); + auto* workspace_ptr = prepare_workspace(encoder); + if (!cached_cuda_graph_) { + // First call: populate the CUDA graph from the cuDNN execution plan. + // Also compute and cache the subgraph key to avoid calling + // cudaGraphKernelNodeGetAttribute on every subsequent call (expensive + // on WDDM where each driver API call has ~40-400us overhead). + cached_cuda_graph_.emplace(encoder.device()); + RETURN_IF_ERROR(populate_cuda_graph( + handle_, variant_pack, workspace_ptr, *cached_cuda_graph_)); + std::tie(cached_subgraph_key_, cached_is_updatable_) = + cu::subgraph_to_key(*cached_cuda_graph_); + } else { + // Subsequent calls: patch data pointers without re-running kernel setup. + RETURN_IF_ERROR(update_cuda_graph( + handle_, variant_pack, workspace_ptr, *cached_cuda_graph_)); + } + // Add the cuDNN child graph to the parent CUDA graph for batched launch. + // The pre-computed subgraph key avoids expensive per-node attribute queries. + encoder.add_graph_node( + *cached_cuda_graph_, cached_subgraph_key_, cached_is_updatable_); return {}; } @@ -93,7 +150,7 @@ fe::error_t DnnGraph::encode_capturing( void* DnnGraph::prepare_workspace(cu::CommandEncoder& encoder) { int64_t workspace_size = 0; - CHECK_CUDNN_FE_ERROR(get_workspace_size(workspace_size)); + CHECK_CUDNN_ERROR(get_workspace_size(workspace_size)); return allocate_workspace(encoder, workspace_size); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.h index ef9302132..40820e592 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cudnn_utils.h @@ -2,6 +2,10 @@ #pragma once +#include +#include + +#include "mlx/backend/cuda/cuda_utils.h" #include "mlx/backend/cuda/device/config.h" #include "mlx/backend/cuda/utils.h" #include "mlx/dtype_utils.h" @@ -17,14 +21,16 @@ class CommandEncoder; namespace fe = cudnn_frontend; -#define CHECK_CUDNN_FE_ERROR(cmd) \ - do { \ - auto error = cmd; \ - if (!error.is_good()) { \ - throw std::runtime_error( \ - fmt::format("{} failed: {}.", #cmd, error.get_message())); \ - } \ - } while (0) +void check_cudnn_error(const char* name, cudnnStatus_t err); +void check_cudnn_error(const char* name, fe::error_t err); + +#define CHECK_CUDNN_ERROR(cmd) check_cudnn_error(#cmd, (cmd)) + +cudnnHandle_t get_cudnn_handle(cu::Device& device); + +void init_cudnn_handles_cache(); +void init_cudnn_conv_cache(); +void init_cudnn_sdpa_cache(); // Return pointer alignment of |x|'s data. inline uint8_t get_alignment(const array& x) { @@ -123,6 +129,20 @@ class DnnGraph : public fe::graph::Graph { return attrs; } + // Create a 4D cuDNN tensor from 1D array, with |axis| being contiguous dim. + auto tensor_4d(const char* name, int64_t uid, const array& x, int axis) { + assert(x.ndim() == 1); + auto attrs = Graph::tensor(fe::graph::Tensor_attributes().set_name(name)); + std::vector shape(4, 1); + std::vector strides(4, 1); + shape.at(axis) = x.size(); + if (axis > 0) { + strides.at(axis - 1) = x.size(); + } + set_tensor_attrs(attrs, uid, x, shape, strides); + return attrs; + } + // Create a cuDNN tensor for scalar. auto scalar(const char* name, int64_t uid, Dtype dtype) { return Graph::tensor( @@ -168,6 +188,9 @@ class DnnGraph : public fe::graph::Graph { const array& x); cudnnHandle_t handle_; + std::optional cached_cuda_graph_; + std::string cached_subgraph_key_; + bool cached_is_updatable_{true}; }; } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/custom_kernel.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/custom_kernel.cpp index 7d2d0d533..c230656d8 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/custom_kernel.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/custom_kernel.cpp @@ -21,8 +21,6 @@ constexpr const char* default_header = R"( #include -#define inf cuda::std::numeric_limits::infinity() - )"; std::string template_arguments_hash( @@ -312,9 +310,11 @@ void CustomKernel::eval_gpu( // Compile the custom kernel std::string kernel_name = (is_precompiled_) ? name_ : "mlx::core::cu::" + name_; + std::string module_name = + fmt::format("{}_{:x}", name_, std::hash{}(source_)); cu::JitModule& mod = cu::get_jit_module( - s.device, - name_, + encoder.device(), + module_name, [&]() { return std::make_tuple( is_precompiled_, source_, std::vector{kernel_name}); @@ -373,7 +373,8 @@ void CustomKernel::eval_gpu( kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem); } }); - encoder.add_kernel_node(kernel, grid, block, shared_memory_, args.args()); + encoder.add_kernel_node_raw( + kernel, grid, block, {}, shared_memory_, args.args()); } } // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/cutlass_utils.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/cutlass_utils.cuh index c9bbf0d4f..ee92de74b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/cutlass_utils.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/cutlass_utils.cuh @@ -2,8 +2,10 @@ #pragma once +#include "mlx/backend/cuda/utils.h" #include "mlx/dtype.h" +#include #include #include #include @@ -22,7 +24,7 @@ inline void check_cutlass_error(const char* name, cutlass::Status status) { } // The macro version that prints the command that failed. -#define CHECK_CUTLASS_ERROR(cmd) check_cutlass_error(#cmd, (cmd)) +#define CHECK_CUTLASS_ERROR(cmd) ::mlx::core::check_cutlass_error(#cmd, (cmd)) // Maps CPU types to CUTLASS types. template @@ -43,4 +45,24 @@ struct CTypeToCutlassType { template using cutlass_type_t = typename CTypeToCutlassType::type; +// Convert Dtype to CUTLASS C++ types. +inline const char* dtype_to_cutlass_type(const Dtype& dtype) { + if (dtype == float16) { + return "cutlass::half_t"; + } + if (dtype == bfloat16) { + return "cutlass::bfloat16_t"; + } + return dtype_to_cuda_type(dtype); +} + +// Convert cute shape to string. +inline auto cta_tiler_to_string(auto cta_tiler) { + return fmt::format( + "cute::Shape, cute::Int<{}>, cute::Int<{}>>", + int(cute::size<0>(cta_tiler)), + int(cute::size<1>(cta_tiler)), + int(cute::size<2>(cta_tiler))); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp index 0ac1cc5e4..7092d4497 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/delayload.cpp @@ -7,7 +7,11 @@ #include // clang-format on -namespace mlx::core { +namespace mlx::core::cu { + +// Defined in dirs.cpp to avoid invalidating compile cache. +const char* cuda_bin_dir(); +const char* cudnn_bin_dir(); namespace fs = std::filesystem; @@ -15,35 +19,35 @@ inline fs::path relative_to_current_binary(const char* relative) { return fs::absolute(current_binary_dir() / relative); } -inline fs::path cublas_bin_dir() { -#if defined(MLX_CUDA_BIN_DIR) - return MLX_CUDA_BIN_DIR; -#else +inline fs::path cublas_dir() { + if (const char* dir = cuda_bin_dir()) { + return fs::path(dir); + } return relative_to_current_binary("../nvidia/cublas/bin"); -#endif } fs::path load_nvrtc() { -#if defined(MLX_CUDA_BIN_DIR) - fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR; -#else - fs::path nvrtc_bin_dir = - relative_to_current_binary("../nvidia/cuda_nvrtc/bin"); -#endif + fs::path nvrtc_dir; + if (const char* dir = cuda_bin_dir()) { + nvrtc_dir = fs::path(dir); + } else { + nvrtc_dir = relative_to_current_binary("../nvidia/cuda_nvrtc/bin"); + } // Internally nvrtc loads some libs dynamically, add to search dirs. - ::AddDllDirectory(nvrtc_bin_dir.c_str()); - return nvrtc_bin_dir; + ::AddDllDirectory(nvrtc_dir.c_str()); + return nvrtc_dir; } fs::path load_cudnn() { -#if defined(MLX_CUDNN_BIN_DIR) - fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR; -#else - fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin"); -#endif + fs::path cudnn_dir; + if (const char* dir = cudnn_bin_dir()) { + cudnn_dir = fs::path(dir); + } else { + cudnn_dir = relative_to_current_binary("../nvidia/cudnn/bin"); + } // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would // get errors like "Invalid handle. Cannot load symbol cudnnCreate". - for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) { + for (const auto& dll : fs::directory_iterator(cudnn_dir)) { if (dll.path().filename().string().starts_with("cudnn_graph") && dll.path().extension() == ".dll") { ::LoadLibraryW(dll.path().c_str()); @@ -52,9 +56,9 @@ fs::path load_cudnn() { } // Internally cuDNN loads some libs dynamically, add to search dirs. load_nvrtc(); - ::AddDllDirectory(cudnn_bin_dir.c_str()); - ::AddDllDirectory(cublas_bin_dir().c_str()); - return cudnn_bin_dir; + ::AddDllDirectory(cudnn_dir.c_str()); + ::AddDllDirectory(cublas_dir().c_str()); + return cudnn_dir; } // Called by system when failed to locate a lazy-loaded DLL. @@ -63,18 +67,21 @@ FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) { if (dliNotify == dliNotePreLoadLibrary) { std::string dll = pdli->szDll; if (dll.starts_with("cudnn")) { - static auto cudnn_bin_dir = load_cudnn(); - mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str()); + static auto cudnn_dir = load_cudnn(); + mod = ::LoadLibraryW((cudnn_dir / dll).c_str()); } else if (dll.starts_with("cublas")) { - mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str()); + mod = ::LoadLibraryW((cublas_dir() / dll).c_str()); } else if (dll.starts_with("nvrtc")) { - static auto nvrtc_bin_dir = load_nvrtc(); - mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str()); + static auto nvrtc_dir = load_nvrtc(); + mod = ::LoadLibraryW((nvrtc_dir / dll).c_str()); + } else if (const char* dir = cuda_bin_dir()) { + mod = ::LoadLibraryW((fs::path(dir) / dll).c_str()); } } return reinterpret_cast(mod); } -} // namespace mlx::core +} // namespace mlx::core::cu -extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper; +extern "C" const PfnDliHook __pfnDliNotifyHook2 = + mlx::core::cu::delayload_helper; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp index 1d7ad7dc6..472b3d99f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device.cpp @@ -1,7 +1,6 @@ // Copyright © 2025 Apple Inc. #include "mlx/backend/cuda/device.h" -#include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/worker.h" #include "mlx/backend/gpu/device_info.h" #include "mlx/utils.h" @@ -31,6 +30,11 @@ const char* save_cuda_graphs_dot_file() { return filename; } +inline bool is_empty_dim(dim3 dim) { + return (dim.x == 0 && dim.y == 0 && dim.z == 0) || + (dim.x == 1 && dim.y == 1 && dim.z == 1); +} + } // namespace Device::Device(int device) : device_(device) { @@ -50,49 +54,20 @@ Device::Device(int device) : device_(device) { &memory_pools_, cudaDevAttrMemoryPoolsSupported, device_)); } -Device::~Device() { - if (cudnn_handle_) { - CHECK_CUDNN_ERROR(cudnnDestroy(cudnn_handle_)); - } - if (cublaslt_handle_) { - CHECK_CUBLAS_ERROR(cublasLtDestroy(cublaslt_handle_)); - } -} +Device::~Device() = default; void Device::make_current() { // We need to set/get current CUDA device very frequently, cache it to reduce - // actual calls of CUDA APIs. - static thread_local int current = 0; + // actual calls of CUDA APIs. Use -1 as sentinel so the first call on each + // new thread always calls cudaSetDevice (which establishes the CUDA primary + // context). Without this, device 0 would never get set on a new thread. + static thread_local int current = -1; if (current != device_) { CHECK_CUDA_ERROR(cudaSetDevice(device_)); current = device_; } } -CommandEncoder& Device::get_command_encoder(Stream s) { - auto it = encoders_.find(s.index); - if (it == encoders_.end()) { - it = encoders_.try_emplace(s.index, *this).first; - } - return it->second; -} - -cublasLtHandle_t Device::get_cublaslt_handle() { - if (!cublaslt_handle_) { - make_current(); - CHECK_CUBLAS_ERROR(cublasLtCreate(&cublaslt_handle_)); - } - return cublaslt_handle_; -} - -cudnnHandle_t Device::get_cudnn_handle() { - if (!cudnn_handle_) { - make_current(); - CHECK_CUDNN_ERROR(cudnnCreate(&cudnn_handle_)); - } - return cudnn_handle_; -} - CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) { enc.device().make_current(); if (!use_cuda_graphs()) { @@ -231,13 +206,19 @@ CommandEncoder::CommandEncoder(Device& d) : device_(d), stream_(d), graph_(d), - worker_(d), + worker_(std::make_shared(d)), graph_cache_("MLX_CUDA_GRAPH_CACHE_SIZE", /* default_capacity */ 400) { std::tie(max_ops_per_graph_, max_mb_per_graph_) = get_graph_limits(d); + worker_->start(); +} + +CommandEncoder::~CommandEncoder() { + synchronize(); + worker_->stop(); } void CommandEncoder::add_completed_handler(std::function task) { - worker_.add_task(std::move(task)); + worker_->add_task(std::move(task)); } void CommandEncoder::set_input_array(const array& arr) { @@ -259,51 +240,88 @@ void CommandEncoder::set_output_array(const array& arr) { active_outputs_.push_back(id); } -void CommandEncoder::add_kernel_node( +void CommandEncoder::add_kernel_node_raw( void* func, dim3 grid_dim, dim3 block_dim, + dim3 cluster_dim, uint32_t smem_bytes, void** params) { + bool use_cluster = !is_empty_dim(cluster_dim); + assert(!use_cluster || device_.compute_capability_major() >= 9); + if (!use_cuda_graphs()) { node_count_++; - CHECK_CUDA_ERROR(cudaLaunchKernel( - func, grid_dim, block_dim, params, smem_bytes, stream())); + cudaLaunchConfig_t config = {}; + config.gridDim = grid_dim; + config.blockDim = block_dim; + config.dynamicSmemBytes = smem_bytes; + config.stream = stream(); + cudaLaunchAttribute attr = {}; + if (use_cluster) { + attr.id = cudaLaunchAttributeClusterDimension; + attr.val.clusterDim.x = cluster_dim.x; + attr.val.clusterDim.y = cluster_dim.y; + attr.val.clusterDim.z = cluster_dim.z; + config.attrs = &attr; + config.numAttrs = 1; + } + CHECK_CUDA_ERROR(cudaLaunchKernelExC(&config, func, params)); return; } + cudaKernelNodeParams kernel_params = {0}; kernel_params.func = func; kernel_params.gridDim = grid_dim; kernel_params.blockDim = block_dim; kernel_params.kernelParams = params; kernel_params.sharedMemBytes = smem_bytes; - add_kernel_node(kernel_params); + cudaGraphNode_t node = add_kernel_node_raw(kernel_params); + if (use_cluster) { + cudaKernelNodeAttrValue attr = {}; + attr.clusterDim.x = cluster_dim.x; + attr.clusterDim.y = cluster_dim.y; + attr.clusterDim.z = cluster_dim.z; + CHECK_CUDA_ERROR(cudaGraphKernelNodeSetAttribute( + node, cudaLaunchAttributeClusterDimension, &attr)); + } } -void CommandEncoder::add_kernel_node( +void CommandEncoder::add_kernel_node_raw( CUfunction func, dim3 grid_dim, dim3 block_dim, + dim3 cluster_dim, uint32_t smem_bytes, void** params) { + bool use_cluster = !is_empty_dim(cluster_dim); + assert(!use_cluster || device_.compute_capability_major() >= 9); + if (!use_cuda_graphs()) { node_count_++; - CHECK_CUDA_ERROR(cuLaunchKernel( - func, - grid_dim.x, - grid_dim.y, - grid_dim.z, - block_dim.x, - block_dim.y, - block_dim.z, - smem_bytes, - stream(), - params, - nullptr)); + CUlaunchConfig config = {}; + config.gridDimX = grid_dim.x; + config.gridDimY = grid_dim.y; + config.gridDimZ = grid_dim.z; + config.blockDimX = block_dim.x; + config.blockDimY = block_dim.y; + config.blockDimZ = block_dim.z; + config.sharedMemBytes = smem_bytes; + config.hStream = stream(); + CUlaunchAttribute attr = {}; + if (use_cluster) { + attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; + attr.value.clusterDim.x = cluster_dim.x; + attr.value.clusterDim.y = cluster_dim.y; + attr.value.clusterDim.z = cluster_dim.z; + config.attrs = &attr; + config.numAttrs = 1; + } + CHECK_CUDA_ERROR(cuLaunchKernelEx(&config, func, params, nullptr)); return; } - CUDA_KERNEL_NODE_PARAMS kernel_params = {0}; + CUDA_KERNEL_NODE_PARAMS kernel_params = {}; kernel_params.func = func; kernel_params.gridDimX = grid_dim.x; kernel_params.gridDimY = grid_dim.y; @@ -313,19 +331,31 @@ void CommandEncoder::add_kernel_node( kernel_params.blockDimZ = block_dim.z; kernel_params.kernelParams = params; kernel_params.sharedMemBytes = smem_bytes; - add_kernel_node(kernel_params); + CUgraphNode node = add_kernel_node_raw(kernel_params); + if (use_cluster) { + CUlaunchAttributeValue attr = {}; + attr.clusterDim.x = cluster_dim.x; + attr.clusterDim.y = cluster_dim.y; + attr.clusterDim.z = cluster_dim.z; + CHECK_CUDA_ERROR(cuGraphKernelNodeSetAttribute( + node, CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION, &attr)); + } } -void CommandEncoder::add_kernel_node(const cudaKernelNodeParams& params) { +cudaGraphNode_t CommandEncoder::add_kernel_node_raw( + const cudaKernelNodeParams& params) { cudaGraphNode_t node; CHECK_CUDA_ERROR(cudaGraphAddKernelNode(&node, graph_, NULL, 0, ¶ms)); insert_graph_dependencies(GraphNode{node, "K"}); + return node; } -void CommandEncoder::add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params) { +CUgraphNode CommandEncoder::add_kernel_node_raw( + const CUDA_KERNEL_NODE_PARAMS& params) { CUgraphNode node; CHECK_CUDA_ERROR(cuGraphAddKernelNode(&node, graph_, NULL, 0, ¶ms)); insert_graph_dependencies(GraphNode{node, "K"}); + return node; } std::pair subgraph_to_key(cudaGraph_t graph) { @@ -407,6 +437,24 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) { insert_graph_dependencies(GraphNode{node, sub_graph_key}); } +void CommandEncoder::add_graph_node( + cudaGraph_t child, + const std::string& subgraph_key, + bool is_updatable) { + if (!use_cuda_graphs()) { + node_count_++; + CudaGraphExec graph_exec; + graph_exec.instantiate(child); + device_.make_current(); + CHECK_CUDA_ERROR(cudaGraphLaunch(graph_exec, stream())); + return; + } + is_graph_updatable_ &= is_updatable; + cudaGraphNode_t node; + CHECK_CUDA_ERROR(cudaGraphAddChildGraphNode(&node, graph_, NULL, 0, child)); + insert_graph_dependencies(GraphNode{node, subgraph_key}); +} + bool CommandEncoder::needs_commit() { return (node_count_ > max_ops_per_graph_) || ((bytes_in_graph_ >> 20) > max_mb_per_graph_); @@ -414,6 +462,42 @@ bool CommandEncoder::needs_commit() { void CommandEncoder::commit() { nvtx3::scoped_range r("CommandEncoder::commit"); + try { + commit_impl(); + } catch (...) { + // Clear pending CUDA error first. + cudaGetLastError(); + // Clear states. + clear_graph_state(); + node_count_ = 0; + bytes_in_graph_ = 0; + // Clear graph. + try { + graph_.reset(); + } catch (...) { + // Destroying could fail. + graph_.release(); + } + try { + graph_ = CudaGraph(device_); + } catch (...) { + // Keep the original error. + } + // Re-throw the error. + throw; + } +} + +void CommandEncoder::synchronize() { + CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); + auto p = std::make_shared>(); + std::future f = p->get_future(); + add_completed_handler([p = std::move(p)]() { p->set_value(); }); + commit(); + f.wait(); +} + +void CommandEncoder::commit_impl() { if (!temporaries_.empty()) { add_completed_handler([temporaries = std::move(temporaries_)]() {}); } @@ -472,43 +556,40 @@ void CommandEncoder::commit() { } // Reset state - from_nodes_.clear(); - to_nodes_.clear(); - graph_deps_key_.clear(); - graph_nodes_key_.clear(); - node_map_.clear(); + clear_graph_state(); graph_ = CudaGraph(device_); - is_graph_updatable_ = true; } // Put completion handlers in a batch. - worker_.commit(stream_); + worker_->commit(stream_); node_count_ = 0; bytes_in_graph_ = 0; } -void CommandEncoder::synchronize() { - CHECK_CUDA_ERROR(cudaStreamSynchronize(stream_)); - auto p = std::make_shared>(); - std::future f = p->get_future(); - add_completed_handler([p = std::move(p)]() { p->set_value(); }); - commit(); - f.wait(); +void CommandEncoder::clear_graph_state() { + from_nodes_.clear(); + to_nodes_.clear(); + graph_deps_key_.clear(); + graph_nodes_key_.clear(); + node_map_.clear(); + active_deps_.clear(); + active_outputs_.clear(); + concurrent_nodes_.clear(); + is_graph_updatable_ = true; } Device& device(int cuda_device) { - static auto devices = []() { - std::vector devices; + // The devices are leak intentionally as user code may still be accessing + // device after main thread teardown. + static auto* devices = []() { + auto* devices = new std::vector; int device_count = gpu::device_count(); for (int i = 0; i < device_count; ++i) { - devices.emplace_back(i); + devices->emplace_back(i); } - // Initialize the jit module cache here ensures it is not unloaded before - // any evaluation is done. - get_jit_module_cache(); return devices; }(); - return devices.at(cuda_device); + return devices->at(cuda_device); } Device& device(mlx::core::Device d) { @@ -516,7 +597,28 @@ Device& device(mlx::core::Device d) { } CommandEncoder& get_command_encoder(Stream s) { - return device(s.device).get_command_encoder(s); + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { + auto& global_encoders = get_global_command_encoders(); + it = global_encoders.find(s.index); + if (it == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(gpu, {}) in current thread.", s.index)); + } + } + return it->second; +} + +std::unordered_map& get_command_encoders() { + static thread_local std::unordered_map encoders; + return encoders; +} + +std::unordered_map& get_global_command_encoders() { + static std::unordered_map encoders; + return encoders; } } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device.h b/Source/Cmlx/mlx/mlx/backend/cuda/device.h index 527827afc..198f0b5ad 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device.h @@ -5,17 +5,19 @@ #include "mlx/array.h" #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/lru_cache.h" -#include "mlx/backend/cuda/worker.h" +#include "mlx/backend/cuda/utils.h" #include "mlx/stream.h" -#include -#include -#include - +#include #include namespace mlx::core::cu { +// Compute a key and updatability flag for a CUDA graph by walking its nodes. +std::pair subgraph_to_key(cudaGraph_t graph); + +class Worker; + class CommandEncoder { public: struct CaptureContext { @@ -32,6 +34,7 @@ class CommandEncoder { }; explicit CommandEncoder(Device& d); + ~CommandEncoder(); CommandEncoder(const CommandEncoder&) = delete; CommandEncoder& operator=(const CommandEncoder&) = delete; @@ -47,10 +50,17 @@ class CommandEncoder { void set_output_array(const array& arr); template - void add_kernel_node( - F* func, + void + add_kernel_node(F* func, dim3 grid_dim, dim3 block_dim, Params&&... params) { + add_kernel_node_ex(func, grid_dim, block_dim, {}, 0, params...); + } + + template + void add_kernel_node_ex( + F func, dim3 grid_dim, dim3 block_dim, + dim3 cluster_dim, uint32_t smem_bytes, Params&&... params) { constexpr size_t num = sizeof...(Params); @@ -59,24 +69,41 @@ class CommandEncoder { ([&](auto&& p) { ptrs[i++] = static_cast(&p); }( std::forward(params)), ...); - add_kernel_node((void*)func, grid_dim, block_dim, smem_bytes, ptrs); + if constexpr (std::is_same_v) { + add_kernel_node_raw( + func, grid_dim, block_dim, cluster_dim, smem_bytes, ptrs); + } else { + add_kernel_node_raw( + reinterpret_cast(func), + grid_dim, + block_dim, + cluster_dim, + smem_bytes, + ptrs); + } } - void add_kernel_node( - CUfunction func, + void add_kernel_node_raw( + void* func, dim3 grid_dim, dim3 block_dim, + dim3 cluster_dim, uint32_t smem_bytes, void** params); - void add_kernel_node( - void* func, + void add_kernel_node_raw( + CUfunction func, dim3 grid_dim, dim3 block_dim, + dim3 cluster_dim, uint32_t smem_bytes, void** params); void add_graph_node(cudaGraph_t child); + void add_graph_node( + cudaGraph_t child, + const std::string& subgraph_key, + bool is_updatable); void add_temporary(const array& arr) { temporaries_.push_back(arr.data_shared_ptr()); @@ -98,8 +125,8 @@ class CommandEncoder { void synchronize(); private: - void add_kernel_node(const cudaKernelNodeParams& params); - void add_kernel_node(const CUDA_KERNEL_NODE_PARAMS& params); + cudaGraphNode_t add_kernel_node_raw(const cudaKernelNodeParams& params); + CUgraphNode add_kernel_node_raw(const CUDA_KERNEL_NODE_PARAMS& params); struct GraphNode { cudaGraphNode_t node; @@ -111,13 +138,15 @@ class CommandEncoder { std::string id; }; + void commit_impl(); + void clear_graph_state(); void insert_graph_dependencies(GraphNode node); void insert_graph_dependencies(std::vector nodes); Device& device_; CudaStream stream_; CudaGraph graph_; - Worker worker_; + std::shared_ptr worker_; int node_count_{0}; bool in_concurrent_{false}; std::vector from_nodes_; @@ -148,10 +177,6 @@ class Device { // Make this device the current cuda device, this method is thread-safe. void make_current(); - CommandEncoder& get_command_encoder(Stream s); - cublasLtHandle_t get_cublaslt_handle(); - cudnnHandle_t get_cudnn_handle(); - int cuda_device() const { return device_; } @@ -183,13 +208,13 @@ class Device { int managed_memory_; int memory_pools_; std::string device_name_; - cublasLtHandle_t cublaslt_handle_{nullptr}; - cudnnHandle_t cudnn_handle_{nullptr}; - std::unordered_map encoders_; }; -Device& device(int cuda_device); -Device& device(mlx::core::Device d); -CommandEncoder& get_command_encoder(Stream s); +MLX_API Device& device(int cuda_device); +MLX_API Device& device(mlx::core::Device d); +MLX_API CommandEncoder& get_command_encoder(Stream s); + +std::unordered_map& get_command_encoders(); +std::unordered_map& get_global_command_encoders(); } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/binary_ops.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/binary_ops.cuh index b0b796280..436886446 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/binary_ops.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/binary_ops.cuh @@ -17,9 +17,18 @@ struct FloorDivide { template __device__ T operator()(T x, T y) { if constexpr (cuda::std::is_integral_v) { + auto q = x / y; + if constexpr (cuda::std::is_signed_v) { + if (x % y != 0 && (x < 0) != (y < 0)) { + q -= 1; + } + } + return q; + } else if constexpr (is_complex_v) { + // Complex is not supported, simply make compiler happy. return x / y; } else { - return cuda::std::trunc(x / y); + return cuda::std::floor(x / y); } } }; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/cute_dequant.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/cute_dequant.cuh similarity index 77% rename from Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/cute_dequant.cuh rename to Source/Cmlx/mlx/mlx/backend/cuda/device/cute_dequant.cuh index 7a1c6b106..6416c5b87 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/cute_dequant.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/cute_dequant.cuh @@ -5,6 +5,7 @@ #include #include #include +#include namespace cutlass { @@ -109,13 +110,13 @@ namespace cute { // Required by tiled copy for 3/5/6-bit weights. struct uint24_t { - std::array bytes; + cuda::std::array bytes; }; struct uint40_t { - std::array bytes; + cuda::std::array bytes; }; struct uint48_t { - std::array bytes; + cuda::std::array bytes; }; template <> @@ -133,15 +134,22 @@ struct uint_bit<48> { } // namespace cute -namespace cutlass_gemm { +namespace mlx::core::cu { + +using namespace cute; // Whether the quant type is affine quantization. template constexpr bool quant_has_bias_v = !cutlass::has_negative_zero_v; // Dequantize CuTe tensors with out = w * s + z. -__device__ __forceinline__ void -cute_vectorized_dequant(auto w, auto s, auto z, auto out) { +template < + typename TensorW, + typename TensorS, + typename TensorZ, + typename TensorO> +CUTE_DEVICE void +cute_vectorized_dequant(TensorW w, TensorS s, TensorZ z, TensorO out) { using namespace cute; using Element = typename decltype(out)::value_type; using Quant = typename decltype(w)::value_type; @@ -165,4 +173,36 @@ cute_vectorized_dequant(auto w, auto s, auto z, auto out) { copy(make_tensor(make_rmem_ptr(&w_dq), out.layout()), out); } -} // namespace cutlass_gemm +template < + typename TensorW, + typename TensorS, + typename TensorZ, + typename TensorO> +CUTE_DEVICE void +cute_naive_dequant(TensorW w, TensorS s, TensorZ z, TensorO out) { + using Element = typename decltype(out)::value_type; + using Quant = typename decltype(w)::value_type; + using Scale = typename decltype(s)::value_type; + transform(w, out, [](Quant q) { return Element(q); }); + transform(out, s, out, [](Element e, Scale s) { return e * Element(s); }); + if constexpr (quant_has_bias_v) { + transform(out, z, out, plus{}); + } +} + +template < + typename TensorW, + typename TensorS, + typename TensorZ, + typename TensorO> +CUTE_DEVICE void cute_dequant(TensorW w, TensorS s, TensorZ z, TensorO out) { + if constexpr ( + stride(coalesce(w.layout())) == Int<1>{} && + is_static_v) { + cute_vectorized_dequant(w, s, z, out); + } else { + cute_naive_dequant(w, s, z, out); + } +} + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/gemm_sm70.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/gemm_sm70.cuh new file mode 100644 index 000000000..f66b89075 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/gemm_sm70.cuh @@ -0,0 +1,285 @@ +// Copyright © 2026 Apple Inc. + +#include + +// clang-format off + +namespace mlx::core::cu { + +using namespace cute; + +template +struct SharedStorage { + ArrayEngine> A; + ArrayEngine> B; +}; + +template +inline constexpr auto make_smem_layout(TileM bM, TileK bK) { + // TODO: Calculate swizzle based on tile shape. + if constexpr (KMajor) { + auto swizzle = composition(Swizzle<3,3,3>{}, + Layout>, + Stride<_8,Stride<_1,_64>>>{}); + return tile_to_shape(swizzle, make_shape(bM, bK)); + } else { + auto swizzle = composition(Swizzle<3,3,3>{}, + Layout, Stride<_8,_1>>{}); + return tile_to_shape(swizzle, make_shape(bM, bK)); + } +} + +template +inline constexpr auto make_smem_layouts(CtaTiler cta_tiler) { + // Note: Kernel launcher assumes num_threads being same for all parameters. + auto [bM, bN, bK] = cta_tiler; + auto sA_layout = make_smem_layout(bM, bK); + auto sB_layout = make_smem_layout(bN, bK); + return cute::make_tuple(sA_layout, sB_layout); +} + +template +inline constexpr auto make_tiled_mma(CtaTiler cta_tiler) { + // Note: Kernel launcher assumes num_threads being same for all parameters. + using Atom = cuda::std::conditional_t< + SM80, + cuda::std::conditional_t< + cuda::std::is_same_v, + SM80_16x8x16_F32F16F16F32_TN, + cuda::std::conditional_t< + cuda::std::is_same_v, + SM80_16x8x16_F32BF16BF16F32_TN, + UniversalFMA + > + >, + UniversalFMA>; + if constexpr (!SM80 || cuda::std::is_same_v) { + return make_tiled_mma(Atom{}, Layout>{}); + } else { + if constexpr (size<0>(cta_tiler) >= 32) { + return make_tiled_mma(Atom{}, Layout>{}, Tile<_32,_32,_16>{}); + } else { + return make_tiled_mma(Atom{}, Layout>{}, Tile<_16,_32,_16>{}); + } + } +} + +template +inline constexpr auto make_tiled_copy(NumThreads num_threads, TileM bM, TileK bK) { + // TODO: Only do 1-element read for the tile of residue. + auto n_read = Int{}; + auto atom = Copy_Atom>>, T>{}; + if constexpr (KMajor) { + auto k_threads = bK / n_read; + return make_tiled_copy( + atom, + make_layout(make_shape(Int{}, k_threads), LayoutRight{}), + make_layout(make_shape(Int<1>{}, n_read))); + } else { + auto m_threads = bM / n_read; + return make_tiled_copy( + atom, + make_layout(make_shape(m_threads, Int{}), LayoutLeft{}), + make_layout(make_shape(n_read, Int<1>{}))); + } +} + +template +CUTE_DEVICE void gemm_sm70_mainloop( + CtaTiler cta_tiler, + TensorA gA, + TensorB gB, + TensorC gC, + int m_max_coord, + int n_max_coord, + int k_residue, + int thread_idx) { + // Get the types of operands. + using Element = typename decltype(gA)::value_type; + + // Shift tensor so we handle residue of K in the 0th tile. + gA = domain_offset(make_coord(0, k_residue, 0), gA); + gB = domain_offset(make_coord(0, k_residue, 0), gB); + + // Define smem layouts. + auto [sA_layout, sB_layout] = make_smem_layouts(cta_tiler); + + // Shared memory buffer. + extern __shared__ char smem_buf[]; + using SharedStorage = SharedStorage; + SharedStorage& smem = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K) + + // Define MMA. + auto mma = make_tiled_mma(cta_tiler); + auto num_threads = size(mma); + + // Define copy atoms. + auto [bM, bN, bK] = cta_tiler; + TiledCopy copy_a = make_tiled_copy(num_threads, bM, bK); + TiledCopy copy_b = make_tiled_copy(num_threads, bN, bK); + + // Partition the copying of A/B/C tiles across the threads. + ThrCopy thr_copy_a = copy_a.get_slice(thread_idx); + Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K) + Tensor tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K) + + ThrCopy thr_copy_b = copy_b.get_slice(thread_idx); + Tensor tBgB = thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K) + Tensor tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_M,BCPY_K) + + // MMA. + ThrMMA thr_mma = mma.get_slice(thread_idx); + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K) + Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) + + // Predicates for m/n bounds. + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) + Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K) + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) + Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) + Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N) + Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K) + Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K) + Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) + CUTE_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord; + } + CUTE_UNROLL + for (int n = 0; n < size<0>(tBpB); ++n) { + tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord; + } + + // GMEM => RMEM. + auto fetch_gmem = [&](int tile) { + copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA); + copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB); + }; + // RMEM => SMEM. + auto store_smem = [&]() { + __syncthreads(); + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); + }; + + // Clear the rmem tiles to account for predicated off loads. + clear(tArA); + clear(tBrB); + + // Prefetch first tile. + Tensor tAgA_k = tAgA(_,_,_,0); + Tensor tBgB_k = tBgB(_,_,_,0); + CUTE_UNROLL + for (int k = 0; k < size<2>(tArA); ++k) { + if (get<1>(tAcA(0,0,k)) >= -k_residue) { + copy_if(copy_a, tApA(_,k), tAgA_k(_,_,k), tArA(_,_,k)); + } + } + CUTE_UNROLL + for (int k = 0; k < size<2>(tBrB); ++k) { + if (get<1>(tBcB(0,0,k)) >= -k_residue) { + copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k)); + } + } + + // Clear accumulators. + clear(tCrC); + + // Loop over CTA tiles. + auto K_TILE_MAX = size<3>(tAgA); + for (int tile = 0; tile < K_TILE_MAX; ++tile) { + store_smem(); + // Avoid fetching full 0th-tile when there is residue. + if (K_TILE_MAX > 1) { + fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); + } + gemm(mma, tCsA, tCsB, tCrC); + } + + // Epilogue. + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) { + tCgC(i) = Element(tCrC(i)); + } + } +} + +template +inline constexpr auto make_matrix_stride(int m, int k) { + if constexpr (KMajor) { + return cute::make_stride(k, cute::Int<1>{}, m * k); + } else { + return cute::make_stride(cute::Int<1>{}, m, m * k); + } +} + +template +__global__ +__launch_bounds__(decltype(size(make_tiled_mma(CtaTiler{})))::value) +void gemm_sm70_kernel( + const Element* A, + const Element* B, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, + Element* C, + int m, int n, int k, int l) { + int thread_idx = int(threadIdx.x); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); + + // Define layouts (mixed). + auto dA = make_matrix_stride(m, k); // (dM,dK,dL) + auto dB = make_matrix_stride(n, k); // (dN,dK,dL) + auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) + + // Represent the full tensors. + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, l), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, l), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, l), dC); // (M,N,L) + + // For gather, use index lookup for input batch slicing. + uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; + uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord; + + // Get batch slice. + Tensor mA = mA_mkl(_,_,a_batch); // (M,K) + Tensor mB = mB_nkl(_,_,b_batch); // (N,K) + Tensor mC = mC_mnl(_,_,l_coord); // (M,N) + + // Get the appropriate blocks for this thread block. + auto cta_tiler = CtaTiler{}; + auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + // Compute tile residues for predication. + int m_max_coord = m - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord + int n_max_coord = n - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord + int k_residue = k - size<1>(gA) * size<2>(gA); + + gemm_sm70_mainloop( + cta_tiler, + gA, + gB, + gC, + m_max_coord, n_max_coord, k_residue, + thread_idx); +} + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh new file mode 100644 index 000000000..e2ad855b9 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_naive.cuh @@ -0,0 +1,316 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device/cute_dequant.cuh" +#include "mlx/backend/cuda/device/gemm_sm70.cuh" +#include "mlx/backend/cuda/device/utils.cuh" + +#include + +// clang-format off + +namespace mlx::core::cu { + +using namespace cute; + +template +CUTE_DEVICE void qmm_naive_mainloop( + CtaTiler cta_tiler, + TensorA gA, + TensorB gB, + TensorS gS, + TensorZ gZ, + TensorC gC, + const float* global_scale, + int m_max_coord, + int n_max_coord, + int k_residue, + int thread_idx) { + // Get the types of operands. + using Element = typename decltype(gA)::value_type; + using Quant = typename decltype(gB)::value_type; + using Scale = typename decltype(gS)::value_type; + + // Shift tensor so we handle residue of K in the 0th tile. + gA = domain_offset(make_coord(0, k_residue, 0), gA); + if constexpr (sizeof_bits_v % 8 == 0) { + gB = domain_offset(make_coord(0, k_residue, 0), gB); + } else { + // TODO: Figure out why domain_offset is not returning wrong offset. + gB.data() = recast_ptr(raw_pointer_cast(gB.data()) + gB.layout()(0, k_residue, 0) * cuda::std::min(8, sizeof_bits_v) / 8); + } + gS = domain_offset(make_coord(0, k_residue, 0), gS); + if constexpr (quant_has_bias_v) { + gZ = domain_offset(make_coord(0, k_residue, 0), gZ); + } + + // Define smem layouts. + auto [sA_layout, sB_layout] = make_smem_layouts(cta_tiler); + + // Shared memory buffer. + extern __shared__ char smem_buf[]; + using SharedStorage = SharedStorage; + SharedStorage& smem = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K) + Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K) + + // Define MMA. + auto mma = make_tiled_mma(cta_tiler); + auto num_threads = size(mma); + + // Define copy atoms. + auto [bM, bN, bK] = cta_tiler; + TiledCopy copy_a = make_tiled_copy(num_threads, bM, bK); + TiledCopy copy_b = make_tiled_copy(num_threads, bN, bK); + + // Partition the copying of A/B/C tiles across the threads. + ThrCopy thr_copy_a = copy_a.get_slice(thread_idx); + Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K) + Tensor tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K) + + ThrCopy thr_copy_b = copy_b.get_slice(thread_idx); + Tensor tBgB = thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K) + Tensor tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_M,BCPY_K) + Tensor tBrB_dq = make_fragment_like(tBsB); // (BCPY,BCPY_M,BCPY_K) + Tensor tBgS = thr_copy_b.partition_S(gS); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBrS = make_fragment_like(tBgS(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K) + Tensor tBgZ = thr_copy_b.partition_S(gZ); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBrZ = make_fragment_like(tBgZ(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K) + + // MMA. + ThrMMA thr_mma = mma.get_slice(thread_idx); + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K) + Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) + Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) + + // Predicates for m/n bounds. + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) + Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K) + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) + Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) + Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N) + Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K) + Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K) + Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) + CUTE_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord; + } + CUTE_UNROLL + for (int n = 0; n < size<0>(tBpB); ++n) { + tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord; + } + + // GMEM => RMEM. + auto fetch_gmem = [&](int tile) { + copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA); + copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB); + CUTE_UNROLL + for (int n = 0; n < size<1>(tBrS); ++n) { + if (tBpB(n,0)) { + copy(tBgS(_,n,_,tile), tBrS(_,n,_)); + if constexpr (quant_has_bias_v) { + copy(tBgZ(_,n,_,tile), tBrZ(_,n,_)); + } + } + } + }; + // RMEM => SMEM. + auto store_smem = [&]() { + __syncthreads(); + copy(tArA, tAsA); + CUTE_UNROLL + for (int k = 0; k < size<2>(tBrB); ++k) { + CUTE_UNROLL + for (int n = 0; n < size<1>(tBrB); ++n) { + cute_dequant(tBrB(_,n,k), tBrS(_,n,k), tBrZ(_,n,k), tBrB_dq(_,n,k)); + } + } + copy(tBrB_dq, tBsB); + __syncthreads(); + }; + + // Clear the rmem tiles to account for predicated off loads. + if constexpr (HasKResidue) { + clear(tArA); + clear(tBrB); + clear(tBrS); + clear(tBrZ); + } + + // Prefetch first tile. + if constexpr (HasKResidue) { + Tensor tAgA_k = tAgA(_,_,_,0); + CUTE_UNROLL + for (int k = 0; k < size<2>(tArA); ++k) { + if (get<1>(tAcA(0,0,k)) >= -k_residue) { + copy_if(copy_a, tApA(_,k), tAgA_k(_,_,k), tArA(_,_,k)); + } + } + Tensor tBgB_k = tBgB(_,_,_,0); + Tensor tBgS_k = tBgS(_,_,_,0); + Tensor tBgZ_k = tBgZ(_,_,_,0); + CUTE_UNROLL + for (int k = 0; k < size<2>(tBrB); ++k) { + if (get<1>(tBcB(0,0,k)) >= -k_residue) { + copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k)); + CUTE_UNROLL + for (int n = 0; n < size<1>(tBrS); ++n) { + if (tBpB(n,k)) { + copy(tBgS_k(_,n,k), tBrS(_,n,k)); + if constexpr (quant_has_bias_v) { + copy(tBgZ_k(_,n,k), tBrZ(_,n,k)); + } + } + } + } + } + } else { + fetch_gmem(0); + } + + // Clear accumulators. + clear(tCrC); + + // Loop over CTA tiles. + auto K_TILE_MAX = size<3>(tAgA); + for (int tile = 0; tile < K_TILE_MAX; ++tile) { + store_smem(); + if constexpr (HasKResidue) { + // Avoid fetching full 0th-tile when there is residue. + if (K_TILE_MAX > 1) { + fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); + } + } else { + fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); + } + gemm(mma, tCsA, tCsB, tCrC); + } + + // Epilogue. + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) { + if constexpr ( + cuda::std::is_same_v && + cuda::std::is_same_v) { + // Only nvfp4 supports global scale. + if (global_scale) { + tCgC(i) = Element(tCrC(i) * (*global_scale / (F8E4M3_MAX * F4E2M1_MAX))); + continue; + } + } + tCgC(i) = Element(tCrC(i)); + } + } +} + +template +inline constexpr auto make_scales_layout(int n, int k, int l) { + auto group_size = Int{}; + if constexpr (KMajor) { + return make_layout( + make_shape(n, make_shape(group_size, k / group_size), l), + make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size)); + } else { + return make_layout( + make_shape(make_shape(group_size, n / group_size), k, l), + make_stride(Stride<_0,_1>{}, n / group_size, n * k / group_size)); + } +} + +template +__global__ +__launch_bounds__(decltype(size(make_tiled_mma(CtaTiler{})))::value) +void qmm_naive_kernel( + const Element* A, + const Quant* B, + const Scale* S, + const Element* Z, + const float* global_scale, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, + Element* C, + int m, int n, int k, int l, + bool broadcast_b) { + int thread_idx = int(threadIdx.x); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); + + // Define layouts (mixed). + auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) + auto dB = make_matrix_stride(n, k); // (dN,dK,dL) + auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) + auto S_layout = make_scales_layout(n, k, l); + + // Handle broadcasting. + if (broadcast_b) { + get<2>(dB) = 0; + get<2>(stride(S_layout)) = 0; + } + + // Represent the full tensors. + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, l), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, l), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, l), dC); // (M,N,L) + + Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) + + // For gather, use index lookup for input batch slicing. + uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; + uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord; + + // Get batch slice. + Tensor mA = mA_mkl(_,_,a_batch); // (M,K) + Tensor mB = mB_nkl(_,_,b_batch); // (N,K) + Tensor mC = mC_mnl(_,_,l_coord); // (M,N) + + Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + + // Get the appropriate blocks for this thread block. + auto cta_tiler = CtaTiler{}; + auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + auto gZ = [&]() { + if constexpr (quant_has_bias_v) { + Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) + Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + return local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + } else { + // Dummy tensor; no-bias paths never offset or load gZ. + return gS; + } + }(); + + // Compute tile residues for predication. + int m_max_coord = m - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord + int n_max_coord = n - size<1>(cta_tiler) * n_coord; // N - BLK_N * n_coord + int k_residue = k - size<1>(gA) * size<2>(gA); + + qmm_naive_mainloop( + cta_tiler, + gA, + gB, + gS, + gZ, + gC, + global_scale, + m_max_coord, n_max_coord, k_residue, + thread_idx); +} + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh similarity index 51% rename from Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh rename to Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh index 679e4dade..e6efe0402 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm80.cuh @@ -1,13 +1,12 @@ // Copyright © 2026 Apple Inc. -#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh" -#include "mlx/backend/cuda/quantized/qmm/qmm.h" -#include "mlx/dtype_utils.h" +#include "mlx/backend/cuda/device/cute_dequant.cuh" + +#include // clang-format off -// We can't put kernel code in mlx::core due to name conflicts of "Shape". -namespace cutlass_gemm { +namespace mlx::core::cu { using namespace cute; @@ -26,64 +25,99 @@ union SharedStorage { } epilogue; }; -template -__global__ void qmm_sm80_kernel( - ProblemShape shape_MNKL, CtaTiler cta_tiler, - const Element* A, StrideA dA, SmemLayoutA sA_layout, TiledCopyA g2s_copy_a, S2RAtomA s2r_atom_a, - const Quant* B, StrideB dB, SmemLayoutB sB_layout, TiledCopyB g2s_copy_b, S2RAtomB s2r_atom_b, - Element* C, StrideC dC, SmemLayoutC sC_layout, TiledCopyC s2g_copy_c, R2SAtomC r2s_atom_c, - const Scale* S, const Element* Z, LayoutS S_layout, G2RAtomS g2r_atom_s, TiledMma mma) { - CUTE_STATIC_ASSERT_V(size(g2s_copy_a) == size(mma)); - CUTE_STATIC_ASSERT_V(size(g2s_copy_b) == size(mma)); - CUTE_STATIC_ASSERT_V(size(s2g_copy_c) == size(mma)); - CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA)); - CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB)); - CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC)); - - int thread_idx = int(threadIdx.x); - auto [m_coord, n_coord, l_coord] = static_cast(blockIdx); - - // Represent the full tensors. - Tensor mA_mkl = make_tensor(make_gmem_ptr(A), select<0,2,3>(shape_MNKL), dA); // (M,K,L) - Tensor mB_nkl = make_tensor(make_gmem_ptr(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L) - Tensor mC_mnl = make_tensor(make_gmem_ptr(C), select<0,1,3>(shape_MNKL), dC); // (M,N,L) +template +inline constexpr auto make_smem_layouts(CtaTiler cta_tiler) { + // Note: Kernel launcher assumes cosize being same for all KMajor. + auto [bM, bN, bK] = cta_tiler; - Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) - Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) + // Define the A/B smem layouts (static). + auto swizzle_ab = composition(Swizzle<3,3,3>{}, + Layout>, + Stride<_8,Stride<_1,_64>>>{}); + auto bP = Int<3>{}; // pipeline + auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP)); + auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP)); - // Get batch slice. - Tensor mA = mA_mkl(_,_,l_coord); // (M,K) - Tensor mB = mB_nkl(_,_,l_coord); // (N,K) - Tensor mC = mC_mnl(_,_,l_coord); // (M,N) + // Define the C smem layouts (static). + // TODO: Find a better swizzle. + auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN)); - Tensor mS = mS_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) - Tensor mZ = mZ_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) + return cute::make_tuple(sA_layout, sB_layout, sC_layout); +} - // Get the appropriate blocks for this thread block. - auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) - Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) - Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) +template +inline constexpr auto make_tiled_mma() { + // Note: Kernel launcher assumes num_threads being same for all parameters. + using Atom = cuda::std::conditional_t< + cuda::std::is_same_v, + SM80_16x8x16_F32F16F16F32_TN, + SM80_16x8x16_F32BF16BF16F32_TN>; + if constexpr (TileM >= 32) { + return make_tiled_mma(Atom{}, Layout>{}, Tile<_32,_32,_16>{}); + } else { + return make_tiled_mma(Atom{}, Layout>{}, Tile<_16,_32,_16>{}); + } +} - Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) +template typename Atom, typename NumThreads> +inline constexpr auto make_tiled_copy(NumThreads num_threads) { + return make_tiled_copy( + Copy_Atom>, T>{}, + make_layout(make_shape(Int{}, Int<8>{}), LayoutRight{}), + make_layout(make_shape(Int<1>{}, Int>{}))); +} - // Shared memory buffers. - extern __shared__ char shared_memory[]; +template +CUTE_DEVICE void qmm_sm80_mainloop( + CtaTiler cta_tiler, + TensorA gA, + TensorB gB, + TensorS gS, + TensorZ gZ, + TensorC gC, + int m_max_coord, + int thread_idx) { + // Get the types of operands. + using Element = typename decltype(gA)::value_type; + using Quant = typename decltype(gB)::value_type; + using Scale = typename decltype(gS)::value_type; + + // Define smem layouts. + auto [sA_layout, sB_layout, sC_layout] = make_smem_layouts(cta_tiler); + + // Shared memory buffer. + extern __shared__ char smem_buf[]; using SharedStorage = SharedStorage; - SharedStorage& smem = *reinterpret_cast(shared_memory); + decltype(sA_layout), + decltype(sB_layout), + decltype(sC_layout)>; + SharedStorage& smem = *reinterpret_cast(smem_buf); Tensor sA = make_tensor(make_smem_ptr(smem.mainloop.A.begin()), sA_layout); // (BLK_M,BLK_K) Tensor sB = make_tensor(make_smem_ptr(smem.mainloop.B.begin()), sB_layout); // (BLK_N,BLK_K) Tensor sC = make_tensor(make_smem_ptr(smem.epilogue.C.begin()), sC_layout); // (BLK_M,BLK_N) + // Define MMA. + auto mma = make_tiled_mma(cta_tiler)>(); + auto num_threads = size(mma); + + // Define copy atoms. + constexpr int element_bits = sizeof_bits_v; + constexpr int quant_bits = sizeof_bits_v; + constexpr int qload = 128 / (element_bits / quant_bits); + TiledCopy g2s_copy_a = make_tiled_copy(num_threads); + TiledCopy g2s_copy_b = make_tiled_copy(num_threads); + TiledCopy s2g_copy_c = make_tiled_copy(num_threads); + + Copy_Atom s2r_atom_a; + Copy_Atom>, Quant> s2r_atom_b; + Copy_Atom>, Element> r2s_atom_c; + Copy_Atom, Scale> g2r_atom_s; + // Partition the copying of A/B/C tiles across the threads. ThrCopy g2s_thr_copy_a = g2s_copy_a.get_slice(thread_idx); Tensor tAgA = g2s_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) @@ -136,7 +170,6 @@ __global__ void qmm_sm80_kernel( Tensor g2r_tCrZ = g2r_thr_copy_s.retile_D(tCrZ); // (BCPY,MMA_N,MMA_K) // Predicates for m bound. - auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) Tensor tCpC = make_tensor(make_shape(size<1>(s2g_tCsC), size<2>(s2g_tCsC)), Stride<_1,_0>{}); // (CPY_M,CPY_N) Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) @@ -166,7 +199,7 @@ __global__ void qmm_sm80_kernel( // Copy S/Z: GMEM => RMEM. auto fetch_scales = [&](int tile) { copy(g2r_copy_s, g2r_tCgS(_,_,_,tile), g2r_tCrS); - if constexpr (quant_has_bias_v) { + if constexpr (mlx::core::cu::quant_has_bias_v) { copy(g2r_copy_s, g2r_tCgZ(_,_,_,tile), g2r_tCrZ); } }; @@ -176,7 +209,7 @@ __global__ void qmm_sm80_kernel( copy(s2r_atom_b, s2r_tCsB(_,_,block,smem_pipe_read), s2r_tCrB(_,_,block)); CUTE_UNROLL for (int n = 0; n < size<1>(tCrB); ++n) { - cute_vectorized_dequant( + mlx::core::cu::cute_vectorized_dequant( tCrB(_,n,block), tCrS(_,n,block), tCrZ(_,n,block), @@ -239,57 +272,38 @@ __global__ void qmm_sm80_kernel( copy_if(s2g_copy_c, tCpC, s2g_tCsC, s2g_tCgC); } -template -inline constexpr auto make_mma_atom() { - if constexpr (std::is_same_v) { - return SM80_16x8x16_F32F16F16F32_TN{}; - } - if constexpr (std::is_same_v) { - return SM80_16x8x16_F32BF16BF16F32_TN{}; - } -} - -template -inline constexpr auto make_tiled_mma() { - constexpr auto atom = make_mma_atom(); - if constexpr (TileM >= 32) { - return make_tiled_mma(atom, Layout>{}, Tile<_32,_32,_16>{}); - } else { - return make_tiled_mma(atom, Layout>{}, Tile<_16,_32,_16>{}); - } -} - -template typename Atom, typename NumThreads> -inline auto make_tiled_copy(NumThreads num_threads) { - return make_tiled_copy( - Copy_Atom>, T>{}, - make_layout(make_shape(Int{}, Int<8>{}), LayoutRight{}), - make_layout(make_shape(Int<1>{}, Int>{}))); +template +inline constexpr auto make_scales_layout(int n, int k, int l) { + auto group_size = Int{}; + return make_layout( + make_shape(n, make_shape(group_size, k / group_size), l), + make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size)); } -template -void qmm_sm80( +template +__global__ +__launch_bounds__(decltype(size(make_tiled_mma()))::value) +void qmm_sm80_kernel( const Element* A, - const Quant* B, + const Quant* B, const Scale* S, const Element* Z, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, Element* C, int m, int n, int k, int l, - bool broadcast_b, - GroupSize group_size, - F&& launch_kernel) { - // Define shapes (dynamic). - auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L) + bool broadcast_b) { + int thread_idx = int(threadIdx.x); + int m_coord = int(blockIdx.x); + int n_coord = int(blockIdx.y); + int l_coord = int(blockIdx.z); - // Define TN strides (mixed). + // Define layouts (mixed). auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) auto dB = make_stride(k, Int<1>{}, n * k); // (dN,dK,dL) auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) - - // Define layout of scales/biases (mixed). - auto S_layout = make_layout( - make_shape(n, make_shape(group_size, k / group_size), l), - make_stride(k / group_size, Stride<_0, _1>{}, n * k / group_size)); + auto S_layout = make_scales_layout(n, k, l); // Handle broadcasting. if (broadcast_b) { @@ -297,200 +311,48 @@ void qmm_sm80( get<2>(stride(S_layout)) = 0; } - // Define CTA tile sizes (static). - auto bM = Int{}; - auto bN = Int<128>{}; - auto bK = Int{}; - auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M,BLK_N,BLK_K) - - // Define MMA. - TiledMMA mma = make_tiled_mma(); - auto num_threads = size(mma); - - // Define the A/B smem layouts (static). - auto swizzle_ab = composition(Swizzle<3,3,3>{}, - Layout>, - Stride<_8,Stride<_1,_64>>>{}); - auto bP = Int<3>{}; // pipeline - auto sA_layout = tile_to_shape(swizzle_ab, make_shape(bM, bK, bP)); - auto sB_layout = tile_to_shape(swizzle_ab, make_shape(bN, bK, bP)); - - // Define the C smem layouts (static). - // TODO: Find a better swizzle. - auto sC_layout = tile_to_shape(swizzle_ab, make_shape(bM, bN)); - - // Define the scales/biases smem layouts (static). - auto bS = ceil_div(bK, group_size); - auto sS_layout = make_layout(make_shape(bN, make_shape(group_size, bS)), - make_stride(bS, Stride<_0, _1>{})); - - // Atoms. - constexpr int element_bits = sizeof_bits_v; - constexpr int quant_bits = sizeof_bits_v; - constexpr int qload = 128 / (element_bits / quant_bits); - TiledCopy g2s_copy_a = make_tiled_copy(num_threads); - TiledCopy g2s_copy_b = make_tiled_copy(num_threads); - TiledCopy s2g_copy_c = make_tiled_copy(num_threads); - - Copy_Atom s2r_atom_a; - Copy_Atom>, Quant> s2r_atom_b; - Copy_Atom>, Element> r2s_atom_c; - Copy_Atom, Scale> g2r_atom_s; - - auto* kernel = &qmm_sm80_kernel< - decltype(prob_shape), decltype(cta_tiler), - Element, Quant, Scale, - decltype(dA), decltype(sA_layout), decltype(g2s_copy_a), decltype(s2r_atom_a), - decltype(dB), decltype(sB_layout), decltype(g2s_copy_b), decltype(s2r_atom_b), - decltype(dC), decltype(sC_layout), decltype(s2g_copy_c), decltype(r2s_atom_c), - decltype(S_layout), decltype(g2r_atom_s), decltype(mma)>; - - // Set L1 to be SMEM only. - size_t smem_bytes = sizeof(SharedStorage); - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100); - - dim3 num_blocks(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l); - dim3 block_dims(num_threads); - void* args[] = { - &prob_shape, &cta_tiler, - &A, &dA, &sA_layout, &g2s_copy_a, &s2r_atom_a, - &B, &dB, &sB_layout, &g2s_copy_b, &s2r_atom_b, - &C, &dC, &sC_layout, &s2g_copy_c, &r2s_atom_c, - &S, &Z, &S_layout, &g2r_atom_s, &mma}; - launch_kernel(reinterpret_cast(kernel), num_blocks, block_dims, smem_bytes, args); -} + // Represent the full tensors. + Tensor mA_mkl = make_tensor(make_gmem_ptr(A), make_shape(m, k, l), dA); // (M,K,L) + Tensor mB_nkl = make_tensor(make_gmem_ptr(B), make_shape(n, k, l), dB); // (N,K,L) + Tensor mC_mnl = make_tensor(make_gmem_ptr(C), make_shape(m, n, l), dC); // (M,N,L) -} // namespace cutlass_gemm + Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) + Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) -// clang-format on + // For gather, use index lookup for input batch slicing. + uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord; + uint32_t b_batch = rhs_indices ? rhs_indices[l_coord] : l_coord; -namespace mlx::core { + // Get batch slice. + Tensor mA = mA_mkl(_,_,a_batch); // (M,K) + Tensor mB = mB_nkl(_,_,b_batch); // (N,K) + Tensor mC = mC_mnl(_,_,l_coord); // (M,N) -template -inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { - if (dtype == float16) { - f.template operator()(); - } else if (dtype == bfloat16) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); - } -} + Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) + Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size)) -template -inline void dispatch_groups(int group_size, const char* tag, F&& f) { - if (group_size == 32) { - f.template operator()<32>(); - } else if (group_size == 64) { - f.template operator()<64>(); - } else if (group_size == 128) { - f.template operator()<128>(); - } else { - throw std::invalid_argument( - fmt::format("{} Group size {} is not supported.", tag, group_size)); - } -} + // Get the appropriate blocks for this thread block. + auto cta_tiler = CtaTiler{}; + auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) + Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) -template -inline void dispatch_quant_types( - int bits, - int group_size, - QuantizationMode mode, - const char* tag, - F&& f) { - if (mode == QuantizationMode::Mxfp4) { - f.template operator()(); - } else if (mode == QuantizationMode::Mxfp8) { - f.template operator()(); - } else if (mode == QuantizationMode::Nvfp4) { - f.template operator()(); - } else { - dispatch_groups(group_size, tag, [&]() { - if (bits == 4) { - f.template operator()(); - } else if (bits == 8) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} {}-bit quantization is not supported.", tag, bits)); - } - }); - } -} + Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) -template -void qmm_impl_sm80( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder) { - const char* tag = "[quantized_matmul]"; - int m = out.ndim() > 1 ? out.shape(-2) : 1; - int n = out.shape(-1); - int k = x.shape(-1); - int l = out.size() / (m * n); - bool broadcast_b = w.ndim() == 2; - - dispatch_element_types(out.dtype(), tag, [&]() { - dispatch_quant_types( - bits, - group_size, - mode, - tag, - [&]() { - encoder.set_input_array(x); - encoder.set_input_array(w); - encoder.set_input_array(scales); - if (biases) { - encoder.set_input_array(*biases); - } - encoder.set_output_array(out); - cutlass_gemm::qmm_sm80( - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(scales), - biases ? gpu_ptr(*biases) : nullptr, - gpu_ptr(out), - m, - n, - k, - l, - broadcast_b, - cute::Int{}, - [&](auto* kernel, - dim3 num_blocks, - dim3 block_dims, - uint32_t smem_bytes, - void** args) { - encoder.add_kernel_node_raw( - kernel, num_blocks, block_dims, {}, smem_bytes, args); - }); - }); - }); + // Compute tile residues for predication. + auto m_max_coord = m - size<0>(gA) * m_coord; // M - BLK_M * m_coord + + qmm_sm80_mainloop( + cta_tiler, + gA, + gB, + gS, + gZ, + gC, + m_max_coord, + thread_idx); } -} // namespace mlx::core - -#define QMM_SM80_GPU(TileM) \ - namespace mlx::core { \ - template void qmm_impl_sm80( \ - const array& x, \ - const array& w, \ - const array& scales, \ - const std::optional& biases, \ - array& out, \ - int bits, \ - int group_size, \ - QuantizationMode mode, \ - cu::CommandEncoder& encoder); \ - } +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm90.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm90.cuh new file mode 100644 index 000000000..cd115bcf6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/qmm_sm90.cuh @@ -0,0 +1,91 @@ +// Copyright © 2026 Apple Inc. + +#if defined(__CUDACC_RTC__) + +#include +#include + +// Some CUTLASS headers use following std functions but we can't use STL when +// compiling with NVRTC. +namespace std { +using cuda::std::is_pointer_v; +using cuda::std::max; +using cuda::std::void_t; +} // namespace std + +// The cutlass/floating_point_nvrtc.h file assumes following constants not +// being defined but they are in the CUDA std headers. +#undef FP_NAN +#undef FP_INFINITE +#undef FP_ZERO +#undef FP_SUBNORMAL +#undef FP_NORMAL + +#endif // defined(__CUDACC_RTC__) + +#include +#include +#include +#include +#include +#include + +namespace mlx::core::cu { + +using namespace cute; + +template +CUTLASS_HOST_DEVICE auto make_qmm_sm90_kernel() { + constexpr int AlignmentA = 128 / sizeof_bits::value; + constexpr int AlignmentB = 128 / sizeof_bits::value; + + using Arch = cutlass::arch::Sm90; + using Accumulator = float; + using ClusterShape = Shape<_1, _1, _1>; + + using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + Arch, + cutlass::arch::OpClassTensorOp, + CtaTiler, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + Accumulator, + Accumulator, + // ElementC: + void, + cutlass::layout::ColumnMajor, + AlignmentA, + // ElementD: + Element, + cutlass::layout::ColumnMajor, + AlignmentA, + cutlass::epilogue::TmaWarpSpecializedCooperative>::CollectiveOp; + + // Note that A/B are swapped and transposed to use TMA epilogue. + using Mainloop = typename cutlass::gemm::collective::CollectiveBuilder< + Arch, + cutlass::arch::OpClassTensorOp, + // ElementA: + tuple, + cutlass::layout::RowMajor, + AlignmentB, + // ElementB: + Element, + cutlass::layout::ColumnMajor, + AlignmentA, + Accumulator, + CtaTiler, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename Epilogue::SharedStorage))>, + cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; + + return cutlass::gemm::kernel:: + GemmUniversal, Mainloop, Epilogue>{}; +} + +template +using qmm_sm90_kernel_t = + decltype(make_qmm_sm90_kernel()); + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/scatter.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/scatter.cuh index b2f640350..c5378bfc4 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/scatter.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/scatter.cuh @@ -65,4 +65,91 @@ __global__ void scatter( Op{}(out + out_idx, upd[upd_loc]); } +template +__global__ void masked_scatter( + const T* dst, + const bool* mask, + const int32_t* scatter_offsets, + const T* src, + T* out, + IdxT size, + IdxT src_batch_size, + IdxT mask_batch_size, + const __grid_constant__ Shape dst_shape, + const __grid_constant__ Strides dst_strides, + int32_t dst_ndim, + const __grid_constant__ Shape src_shape, + const __grid_constant__ Strides src_strides, + int32_t src_ndim) { + IdxT index = cg::this_grid().thread_rank(); + if (index >= size) { + return; + } + + T dst_val; + if constexpr (DstContiguous) { + dst_val = dst[index]; + } else { + IdxT dst_loc = + elem_to_loc(index, dst_shape.data(), dst_strides.data(), dst_ndim); + dst_val = dst[dst_loc]; + } + + if (mask[index]) { + IdxT src_index = static_cast(scatter_offsets[index]); + if (src_index < src_batch_size) { + IdxT batch_idx = index / mask_batch_size; + if constexpr (SrcContiguous) { + out[index] = src[batch_idx * src_batch_size + src_index]; + } else { + IdxT src_elem = batch_idx * src_batch_size + src_index; + IdxT src_loc = elem_to_loc( + src_elem, src_shape.data(), src_strides.data(), src_ndim); + out[index] = src[src_loc]; + } + return; + } + } + + out[index] = dst_val; +} + +template +__global__ void masked_scatter_vec_contiguous( + const T* dst, + const bool* mask, + const int32_t* scatter_offsets, + const T* src, + T* out, + IdxT size, + IdxT src_batch_size, + IdxT mask_batch_size) { + IdxT vec_index = cg::this_grid().thread_rank(); + IdxT base = vec_index * N_READS; + if (base >= size) { + return; + } + + auto out_vec = load_vector(dst, vec_index, size, static_cast(0)); + auto mask_vec = load_vector(mask, vec_index, size, false); + auto offset_vec = load_vector(scatter_offsets, vec_index, size, 0); + +#pragma unroll + for (int i = 0; i < N_READS; ++i) { + IdxT index = base + i; + if (index >= size) { + break; + } + if (mask_vec[i]) { + IdxT src_index = static_cast(offset_vec[i]); + if (src_index < src_batch_size) { + IdxT batch_idx = index / mask_batch_size; + out_vec[i] = src[batch_idx * src_batch_size + src_index]; + } + } + } + + store_vector(out, vec_index, out_vec, size); +} + } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/device/utils.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/device/utils.cuh index 2e017ba21..2f9071376 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/device/utils.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/device/utils.cuh @@ -167,6 +167,9 @@ inline __device__ void store_vector( // Type limits utils /////////////////////////////////////////////////////////////////////////////// +constexpr float F8E4M3_MAX = 448.0f; +constexpr float F4E2M1_MAX = 6.0f; + template struct Limits { static constexpr __host__ __device__ T max() { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/dirs.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/dirs.cpp new file mode 100644 index 000000000..bd24853dd --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/dirs.cpp @@ -0,0 +1,49 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/common/utils.h" + +#include +#include + +namespace mlx::core::cu { +namespace { + +namespace fs = std::filesystem; + +std::string resolve_bin_dir(const char* dir) { + fs::path path(dir); + if (path.is_absolute()) { + return path.string(); + } + return fs::absolute(current_binary_dir() / path).string(); +} + +} // namespace + +const char* cccl_dir() { +#if defined(MLX_CCCL_DIR) + return MLX_CCCL_DIR; +#else + return nullptr; +#endif +} + +const char* cuda_bin_dir() { +#if defined(MLX_CUDA_BIN_DIR) + static const std::string dir = resolve_bin_dir(MLX_CUDA_BIN_DIR); + return dir.c_str(); +#else + return nullptr; +#endif +} + +const char* cudnn_bin_dir() { +#if defined(MLX_CUDNN_BIN_DIR) + static const std::string dir = resolve_bin_dir(MLX_CUDNN_BIN_DIR); + return dir.c_str(); +#else + return nullptr; +#endif +} + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp index 1d04c8f37..f8a71d71b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/eval.cpp @@ -2,25 +2,48 @@ #include "mlx/backend/gpu/eval.h" #include "mlx/backend/cuda/allocator.h" -#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/cublas_utils.h" +#include "mlx/backend/cuda/cudnn_utils.h" +#include "mlx/backend/cuda/event.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" +#include "mlx/utils.h" #include namespace mlx::core::gpu { -void new_stream(Stream s) { - // Force initalization of CUDA, so CUDA runtime get destroyed at last. +void init() { + // Force initialization of CUDA, so CUDA runtime get destroyed last. cudaFree(nullptr); - // Make sure CUDA event pool get destroyed after device and stream. - cu::CudaEvent::init_pool(); - // Ensure the static stream objects get created. - cu::get_command_encoder(s); + // Make sure native resources get destroyed after CommandEncoder. + mlx::core::cu::CudaEvent::init_pool(); + init_cublas_handles_cache(); + init_cudnn_handles_cache(); + init_cudnn_conv_cache(); + init_cudnn_sdpa_cache(); +} + +void new_stream(Stream s) { + assert(s.device == Device::gpu); + auto& encoders = cu::get_command_encoders(); + auto& d = cu::device(s.device); + encoders.try_emplace(s.index, d); +} + +void new_thread_unsafe_stream(Stream s) { + assert(s.device == Device::gpu); + auto& encoders = cu::get_global_command_encoders(); + auto& d = cu::device(s.device); + encoders.try_emplace(s.index, d); } void eval(array& arr) { nvtx3::scoped_range r("gpu::eval"); + // Ensure CUDA context is active on this thread. Required when MLX is called + // from threads that have not yet established a CUDA context (e.g. thread + // pools, language runtimes that migrate work across OS threads). + cu::device(arr.primitive().stream().device).make_current(); auto outputs = arr.outputs(); { // If the array is a tracer hold a reference @@ -63,4 +86,11 @@ void synchronize(Stream s) { cu::get_command_encoder(s).synchronize(); } +void clear_streams() { + cu::get_command_encoders().clear(); + if (is_main_thread()) { + cu::get_global_command_encoders().clear(); + } +} + } // namespace mlx::core::gpu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu index 7334c17a8..d3b6f97f5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/event.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/event.cu @@ -113,10 +113,7 @@ void CudaEvent::init_pool() { cuda_event_pool(); } -// Wraps CudaEvent with a few features: -// 1. The class can be copied. -// 2. Make wait/record work with CPU streams. -// 3. Add checks for waiting on un-recorded event. +// Wraps CudaEvent so it can be copied. class CopyableCudaEvent { public: explicit CopyableCudaEvent(Device& d) @@ -126,32 +123,24 @@ class CopyableCudaEvent { cudaEventDisableTiming | cudaEventBlockingSync)) {} void wait() { + check_recorded(); event_->wait(); } void wait(Stream s) { - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this]() mutable { - check_recorded(); - event_->wait(); - }); - } else { - check_recorded(); - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->wait(encoder.stream()); - } + assert(s.device == mlx::core::Device::gpu); + check_recorded(); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->wait(encoder.stream()); } void record(Stream s) { - if (s.device == mlx::core::Device::cpu) { - throw std::runtime_error("CudaEvent can not wait on CPU stream."); - } else { - auto& encoder = cu::get_command_encoder(s); - encoder.commit(); - event_->record(encoder.stream()); - recorded_ = true; - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = cu::get_command_encoder(s); + encoder.commit(); + event_->record(encoder.stream()); + recorded_ = true; } bool is_signaled() const { @@ -213,7 +202,12 @@ auto check_gpu_coherency() { return coherency; } -AtomicEvent::AtomicEvent() { +const CudaStream& signal_stream() { + static CudaStream stream(device(0)); + return stream; +} + +AtomicEvent::AtomicEvent(Device& d) { void* buf; cudaError_t (*cuda_free)(void*); // There are 3 kinds of systems we are implementing for: @@ -223,6 +217,7 @@ AtomicEvent::AtomicEvent() { // => use cuda::atom_ref on pinned host memory // 2. no hardware cpu/gpu coherency // => use cuda::atom_ref on device memory + d.make_current(); auto [concurrent_managed_access, host_native_atomic] = check_gpu_coherency(); if (concurrent_managed_access) { CHECK_CUDA_ERROR(cudaMallocManaged(&buf, sizeof(uint32_t))); @@ -263,14 +258,11 @@ void AtomicEvent::wait(cudaStream_t stream, uint32_t value) { void AtomicEvent::wait(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::wait(s)"); - if (s.device == mlx::core::Device::cpu) { - scheduler::enqueue(s, [*this, value]() mutable { wait(value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - wait(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + wait(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } void AtomicEvent::signal(uint32_t value) { @@ -288,17 +280,11 @@ void AtomicEvent::signal(cudaStream_t stream, uint32_t value) { void AtomicEvent::signal(Stream s, uint32_t value) { nvtx3::scoped_range r("cu::AtomicEvent::signal(s)"); - if (s.device == mlx::core::Device::cpu) { - // Signal through a GPU stream so the atomic is updated in GPU - updating - // the atomic in CPU sometimes does not get GPU notified. - scheduler::enqueue( - s, [*this, value]() mutable { signal(signal_stream(), value); }); - } else { - auto& encoder = get_command_encoder(s); - encoder.commit(); - signal(encoder.stream(), value); - encoder.add_completed_handler([buf = buf_]() {}); - } + assert(s.device == mlx::core::Device::gpu); + auto& encoder = get_command_encoder(s); + encoder.commit(); + signal(encoder.stream(), value); + encoder.add_completed_handler([buf = buf_]() {}); } bool AtomicEvent::is_signaled(uint32_t val) const { @@ -318,9 +304,21 @@ uint32_t AtomicEvent::value() const { } } -const CudaStream& AtomicEvent::signal_stream() { - static CudaStream stream(device(0)); - return stream; +/////////////////////////////////////////////////////////////////////////////// +// EventImpl implementations +/////////////////////////////////////////////////////////////////////////////// + +void EventImpl::ensure_created(Stream s, uint64_t signal_value) { + if (is_created()) { + return; + } + auto& d = cu::device(s.device); + if (s.device == mlx::core::Device::cpu || signal_value > 1) { + nvtx3::mark("Using slow AtomicEvent"); + atomic = std::make_unique(d); + } else { + cuda = std::make_unique(d); + } } } // namespace cu @@ -329,85 +327,85 @@ const CudaStream& AtomicEvent::signal_stream() { // Event implementations /////////////////////////////////////////////////////////////////////////////// -namespace { - -struct EventImpl { - // CudaEvent is preferred when possible because it is fast, however we have - // to fallback to AtomicEvent in following cases: - // 1. the event is used to wait/signal a cpu stream; - // 2. signal value other than 1 has been specified. - std::unique_ptr cuda; - std::unique_ptr atomic; - - bool is_created() const { - return cuda || atomic; - } - - void ensure_created(Stream s, uint64_t signal_value) { - if (is_created()) { - return; - } - if (s.device == mlx::core::Device::cpu || signal_value > 1) { - nvtx3::mark("Using slow AtomicEvent"); - atomic = std::make_unique(); - } else { - cuda = std::make_unique(cu::device(s.device)); - } - } -}; - -} // namespace - Event::Event(Stream s) : stream_(s) { - event_ = std::shared_ptr( - new EventImpl(), [](void* ptr) { delete static_cast(ptr); }); + event_ = std::make_shared(); } void Event::wait() { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + check_error(); + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(); + event.cuda->wait(); } else { - event->atomic->wait(value()); + event.atomic->wait(value()); } CHECK_CUDA_ERROR(cudaPeekAtLastError()); + check_error(); } void Event::wait(Stream s) { - auto* event = static_cast(event_.get()); - assert(event->is_created()); - if (event->cuda) { + auto& event = cast(); + assert(event.is_created()); + if (event.cuda) { assert(value() == 1); - event->cuda->wait(s); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().cuda->wait(); + }); + } else { + event.cuda->wait(s); + } } else { - event->atomic->wait(s, value()); + if (s.device == mlx::core::Device::cpu) { + scheduler::wait_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->wait(value); + }); + } else { + event.atomic->wait(s, value()); + } } } void Event::signal(Stream s) { - auto* event = static_cast(event_.get()); - event->ensure_created(s, value()); - if (event->cuda) { + auto& event = cast(); + event.ensure_created(s, value()); + if (event.cuda) { assert(value() == 1); - event->cuda->record(s); + if (s.device == mlx::core::Device::cpu) { + throw std::runtime_error("CudaEvent can not wait on CPU stream."); + } else { + event.cuda->record(s); + } } else { - event->atomic->signal(s, value()); + if (s.device == mlx::core::Device::cpu) { + // Signal through a GPU stream so the atomic is updated in GPU - updating + // the atomic in CPU sometimes does not get GPU notified. + scheduler::signal_event(s, *this, [value = value()](Event& self) { + self.cast().atomic->signal(cu::signal_stream(), value); + }); + } else { + event.atomic->signal(s, value()); + } } } bool Event::is_signaled() const { - auto* event = static_cast(event_.get()); - if (!event->is_created()) { + auto& event = cast(); + if (!event.is_created()) { return false; } - if (event->cuda) { + if (event.cuda) { assert(value() == 1); - return event->cuda->is_signaled(); + return event.cuda->is_signaled(); } else { - return event->atomic->is_signaled(value()); + return event.atomic->is_signaled(value()); } } +std::atomic& Event::error() { + return cast().error; +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/event.h b/Source/Cmlx/mlx/mlx/backend/cuda/event.h index 68aa4d837..fdeb6a0e7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/event.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/event.h @@ -13,6 +13,7 @@ namespace mlx::core::cu { +class CopyableCudaEvent; class Device; // RAII-managed move-only wrapper of cudaEvent_t. @@ -54,7 +55,7 @@ class CudaEvent { // CudaEvent so the latter should always be preferred when possible. class AtomicEvent { public: - AtomicEvent(); + AtomicEvent(Device& d); void wait(uint32_t value); void wait(cudaStream_t stream, uint32_t value); @@ -66,8 +67,6 @@ class AtomicEvent { uint32_t value() const; private: - const CudaStream& signal_stream(); - uint32_t* ptr() const { return static_cast(buf_.get()); } @@ -76,4 +75,21 @@ class AtomicEvent { std::shared_ptr buf_; }; +struct EventImpl { + std::atomic error; + + // CudaEvent is preferred when possible because it is fast, however we have + // to fallback to AtomicEvent in following cases: + // 1. the event is used to wait/signal a cpu stream; + // 2. signal value other than 1 has been specified. + std::unique_ptr cuda; + std::unique_ptr atomic; + + bool is_created() const { + return cuda || atomic; + } + + void ensure_created(Stream s, uint64_t signal_value); +}; + } // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp index dff5f7441..3a3acdba0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/fence.cpp @@ -9,21 +9,23 @@ namespace mlx::core { struct FenceImpl { uint32_t count; - cu::AtomicEvent event; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; Fence::Fence(Stream s) { - fence_ = std::shared_ptr( - new FenceImpl{0}, [](void* ptr) { delete static_cast(ptr); }); + fence_ = std::make_shared(0, s); + // Ensure that we use AtomicEvent. + cast().event.cast().ensure_created(s, 2); } void Fence::wait(Stream s, const array&) { - auto* fence = static_cast(fence_.get()); - fence->event.wait(fence->count); + cast().event.wait(); } void Fence::update(Stream s, const array& a, bool cross_device) { - auto* fence = static_cast(fence_.get()); + auto& f = cast(); if (cross_device) { // Move to managed memory if there is a device switch auto& cbuf = @@ -34,8 +36,9 @@ void Fence::update(Stream s, const array& a, bool cross_device) { cu::allocator().move_to_unified_memory(cbuf, encoder.stream()); } } - fence->count++; - fence->event.signal(s, fence->count); + f.count++; + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp index 78340457a..e05595b8f 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm.cpp @@ -15,17 +15,14 @@ namespace { cublasComputeType_t dtype_to_compute_type(Dtype dtype) { switch (dtype) { case float16: - return CUBLAS_COMPUTE_32F; case bfloat16: return CUBLAS_COMPUTE_32F; case float32: + case complex64: return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; case float64: return CUBLAS_COMPUTE_64F; - case complex64: - return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32 - : CUBLAS_COMPUTE_32F; default: throw std::runtime_error( fmt::format( @@ -73,6 +70,14 @@ CublasGemm::CublasGemm( batch_count, a_batch_stride, b_batch_stride); + + // alpha and beta are both host pointers + cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST; + CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( + matmul_desc_, + CUBLASLT_MATMUL_DESC_POINTER_MODE, + &pointer_mode, + sizeof(pointer_mode))); } CublasGemm::CublasGemm( @@ -215,8 +220,8 @@ void CublasGemm::execute( const void* a, const void* b, const void* c, - float alpha /* = 1 */, - float beta /* = 0 */) { + const float alpha /* = 1 */, + const float beta /* = 0 */) { const void* alpha_ptr = α const void* beta_ptr = β complex64_t alpha_c, beta_c; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu index e4faba61a..cad152e70 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/cublas_gemm_batched_12_9.cu @@ -182,7 +182,6 @@ void CublasGemm::run_batched( cu::set_mm_device_pointers_nd, num_blocks, block_dims, - 0, gpu_ptr(pointers), gpu_ptr(a), gpu_ptr(b), @@ -199,7 +198,6 @@ void CublasGemm::run_batched( cu::set_mm_device_pointers_g, num_blocks, block_dims, - 0, gpu_ptr(pointers), gpu_ptr(a), gpu_ptr(b), @@ -270,7 +268,6 @@ void CublasGemm::run_batched( cu::set_addmm_device_pointers_nd, num_blocks, block_dims, - 0, gpu_ptr(pointers), gpu_ptr(a), gpu_ptr(b), @@ -289,7 +286,6 @@ void CublasGemm::run_batched( cu::set_addmm_device_pointers_g, num_blocks, block_dims, - 0, gpu_ptr(pointers), gpu_ptr(a), gpu_ptr(b), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.cu b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.cu new file mode 100644 index 000000000..2bb345852 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.cu @@ -0,0 +1,114 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device/gemm_sm70.cuh" + +#include "mlx/backend/cuda/cutlass_utils.cuh" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/dtype_utils.h" + +#include "cuda_jit_sources.h" + +namespace mlx::core { + +namespace { + +inline auto make_cta_tiler(int m) { + int tile_m = std::max(16, std::min(256, next_power_of_2(m))); + int tile_n = 128; + int tile_k = 64; + return cute::make_shape(tile_m, tile_n, tile_k); +} + +} // namespace + +template +inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { + if (dtype == float32) { + f.template operator()(); + } else if (dtype == float16) { + f.template operator()(); + } else if (dtype == bfloat16) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); + } +} + +void gather_mm( + bool a_transposed, + bool b_transposed, + const array& a, + const array& b, + const array& lhs_indices, + const array& rhs_indices, + array& out, + cu::CommandEncoder& encoder) { + int m = out.shape(-2); + int n = out.shape(-1); + int k = a.shape(-1); + int l = out.size() / (m * n); + + bool aligned = (k % 8 == 0); + bool sm80 = encoder.device().compute_capability_major() >= 8; + auto cta_tiler = make_cta_tiler(m); + + std::string module_name = fmt::format( + "gemm_sm70_{}{}_{}_{}_m{}_n{}_k{}", + a_transposed ? "n" : "t", + b_transposed ? "n" : "t", + dtype_to_string(out.dtype()), + aligned ? "aligned" : "unaligned", + cute::size<0>(cta_tiler), + cute::size<1>(cta_tiler), + cute::size<2>(cta_tiler)); + + std::string kernel_name = fmt::format( + "mlx::core::cu::gemm_sm70_kernel<{}, {}, {}, {}, {}, {}>", + !a_transposed, + b_transposed, + aligned, + sm80, + dtype_to_cutlass_type(out.dtype()), + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_gemm_sm70, std::vector{kernel_name}); + }); + + encoder.set_input_array(a); + encoder.set_input_array(b); + encoder.set_input_array(lhs_indices); + encoder.set_input_array(rhs_indices); + encoder.set_output_array(out); + + dim3 num_blocks{ + uint32_t(cute::ceil_div(m, cute::size<0>(cta_tiler))), + uint32_t(cute::ceil_div(n, cute::size<1>(cta_tiler))), + uint32_t(l)}; + dim3 block_dims{uint32_t(cute::size(cu::make_tiled_mma(cta_tiler)))}; + + auto [sA_layout, sB_layout] = cu::make_smem_layouts(cta_tiler); + size_t smem_bytes = + out.itemsize() * (cute::cosize(sA_layout) + cute::cosize(sB_layout)); + + encoder.add_kernel_node_ex( + mod.get_kernel(kernel_name), + num_blocks, + block_dims, + {}, + smem_bytes, + gpu_ptr(a), + gpu_ptr(b), + gpu_ptr(lhs_indices), + gpu_ptr(rhs_indices), + gpu_ptr(out), + m, + n, + k, + l); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.h b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.h new file mode 100644 index 000000000..33819595a --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gather_gemm.h @@ -0,0 +1,23 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +namespace mlx::core { + +namespace cu { +class CommandEncoder; +} + +class array; + +void gather_mm( + bool a_transposed, + bool b_transposed, + const array& a, + const array& b, + const array& lhs_indices, + const array& rhs_indices, + array& out, + cu::CommandEncoder& encoder); + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gemv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gemv.cu index daa275223..e48317daa 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gemv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/gemv.cu @@ -13,6 +13,14 @@ namespace cg = cooperative_groups; static constexpr int rows_per_block = 8; +// Split batches across grid.y/z because both dimensions are capped at 65,535. +inline dim3 get_gemv_grid_dims(uint32_t num_blocks_x, uint32_t batch_size) { + constexpr uint32_t max_grid_yz_dim = 65535; + uint32_t num_blocks_z = cuda::ceil_div(batch_size, max_grid_yz_dim); + uint32_t num_blocks_y = cuda::ceil_div(batch_size, num_blocks_z); + return {num_blocks_x, num_blocks_y, num_blocks_z}; +} + // Accumulator type selection per input element type T. template struct GemvAccType { @@ -88,12 +96,17 @@ __global__ void gemv_batched( T* out, int rows, int cols, + uint32_t batch_size, const __grid_constant__ Shape batch_shape, const __grid_constant__ Strides mat_batch_strides, const __grid_constant__ Strides vec_batch_strides, int batch_ndim) { - auto block = cg::this_thread_block(); - auto batch_idx = block.group_index().y; + auto grid = cg::this_grid(); + auto batch_idx = + grid.block_index().z * grid.dim_blocks().y + grid.block_index().y; + if (batch_idx >= batch_size) { + return; + } auto [vec_offset, mat_offset] = elem_to_loc( batch_idx, batch_shape.data(), @@ -113,6 +126,7 @@ __global__ void gemv_gather( uint32_t* vec_indices, int rows, int cols, + uint32_t batch_size, const __grid_constant__ Shape mat_batch_shape, const __grid_constant__ Strides mat_batch_strides, int mat_batch_ndim, @@ -123,8 +137,12 @@ __global__ void gemv_gather( const __grid_constant__ Strides mat_index_strides, const __grid_constant__ Strides vec_index_strides, int index_batch_ndim) { - auto block = cg::this_thread_block(); - auto indices_idx = block.group_index().y; + auto grid = cg::this_grid(); + auto indices_idx = + grid.block_index().z * grid.dim_blocks().y + grid.block_index().y; + if (indices_idx >= batch_size) { + return; + } uint32_t index_mat, index_vec; if (index_batch_ndim > 1) { auto [mat_idx_offset, vec_idx_offset] = elem_to_loc( @@ -167,7 +185,7 @@ __global__ void gemv_gather( } bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed) { - return K % 32 == 0 && ((M == 1 && b_transposed) || (N == 1 && !a_transposed)); + return (M == 1 && b_transposed) || (N == 1 && !a_transposed); } template @@ -236,7 +254,6 @@ void gemv( kernel, num_blocks_x, block_dims, - 0, mat, vec, gpu_ptr(out), @@ -246,14 +263,14 @@ void gemv( auto kernel = gemv_batched; encoder.add_kernel_node( kernel, - dim3{num_blocks_x, batch_count}, + get_gemv_grid_dims(num_blocks_x, batch_count), block_dims, - 0, mat, vec, gpu_ptr(out), rows, cols, + batch_count, const_param(batch_shape), mat_strides, vec_strides, @@ -300,9 +317,8 @@ void gather_mv( auto kernel = gemv_gather; encoder.add_kernel_node( kernel, - dim3{num_blocks_x, batch_size}, + get_gemv_grid_dims(num_blocks_x, batch_size), block_dims, - 0, mat, vec, gpu_ptr(out), @@ -310,6 +326,7 @@ void gather_mv( gpu_ptr(vec_indices), rows, cols, + batch_size, const_param(mat_.shape()), const_param(mat_.strides()), mat_.ndim() - 2, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm.h b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm.h index 308a1ba9e..844b8f440 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm.h @@ -22,4 +22,18 @@ void cutlass_grouped_gemm_unaligned( array& out, cu::CommandEncoder& encoder); +void cutlass_segmented_mm( + bool a_transposed, + int lda, + bool b_transposed, + int ldb, + int num_segments, + int M, + int N, + const array& a, + const array& b, + const array& segments, + array& out, + cu::CommandEncoder& encoder); + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu index 0c1bb9fef..319c759c2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu @@ -93,6 +93,50 @@ __global__ void prepare_grouped_mm_data( } } +__global__ void prepare_segmented_mm_data( + const uint32_t* segments, + int num_segments, + int M, + int N, + int lda, + int ldb, + int item_size, + bool a_transposed, + bool b_transposed, + int8_t* a_start, + int8_t* b_start, + int8_t* out_start, + ProblemSize* problem_sizes, + int64_t* a_lds, + int64_t* b_lds, + int64_t* out_lds, + void** a_ptrs, + void** b_ptrs, + void** out_ptrs) { + int idx = cg::this_grid().thread_rank(); + if (idx >= num_segments) + return; + + int64_t start = segments[2 * idx]; + int64_t end = segments[2 * idx + 1]; + int K_i = (end > start) ? static_cast(end - start) : 0; + + problem_sizes[idx] = {M, N, K_i}; + a_lds[idx] = lda; + b_lds[idx] = ldb; + out_lds[idx] = N; + + // Offset into K dimension depends on layout: + // A [M,K]: row-major offset = start, col-major offset = start * lda + // B [K,N]: row-major offset = start * ldb, col-major offset = start + int64_t a_offset = a_transposed ? start * lda : start; + int64_t b_offset = b_transposed ? start : start * ldb; + + a_ptrs[idx] = a_start + a_offset * item_size; + b_ptrs[idx] = b_start + b_offset * item_size; + out_ptrs[idx] = out_start + static_cast(idx) * M * N * item_size; +} + } // namespace cu namespace { @@ -158,10 +202,11 @@ class GemmGroupedEncoder : public cutlass::gemm::device::GemmGrouped { public: void encode(cu::CommandEncoder& encoder) { - encoder.add_kernel_node( + encoder.add_kernel_node_ex( cutlass::Kernel, {static_cast(this->params_.threadblock_count), 1, 1}, {GemmKernel::kThreadCount, 1, 1}, + {}, sizeof(typename GemmKernel::SharedStorage), this->params_); } @@ -200,7 +245,7 @@ void grouped_gemm_v2( LayoutB, cutlass::ComplexTransform::kNone, GemmConfiguration::kAlignmentAB, - typename GemmConfiguration::Element, + typename GemmConfiguration::Accumulator, cutlass::layout::RowMajor, typename GemmConfiguration::Accumulator, typename GemmConfiguration::OpClass, @@ -280,6 +325,11 @@ void cutlass_grouped_gemm_unaligned( const array& indices, array& out, cu::CommandEncoder& encoder) { + if (group_count > 1024) { + throw std::runtime_error( + "[gather_mm] Group count can not be larger than 1024."); + } + int K = a.shape(-1); int N = b.shape(-1); @@ -309,10 +359,11 @@ void cutlass_grouped_gemm_unaligned( encoder.set_input_array(indices); encoder.set_output_array(gemm_args); - encoder.add_kernel_node( + encoder.add_kernel_node_ex( cu::prepare_grouped_mm_data, num_blocks, block_dims, + {}, group_count * sizeof(uint32_t), // sizeof(cum_histo) gpu_ptr(indices), indices.size(), @@ -355,4 +406,85 @@ void cutlass_grouped_gemm_unaligned( encoder); } +void cutlass_segmented_mm( + bool a_transposed, + int lda, + bool b_transposed, + int ldb, + int num_segments, + int M, + int N, + const array& a, + const array& b, + const array& segments, + array& out, + cu::CommandEncoder& encoder) { + // Allocate grouped GEMM args on device. + int problem_sizes_nbytes = + num_segments * cuda::ceil_div(sizeof(ProblemSize), 8) * 8; + int nbytes = problem_sizes_nbytes + + num_segments * (3 * sizeof(void*) + 3 * sizeof(int64_t)); + nbytes = cuda::ceil_div(nbytes, 256) * 256; + array gemm_args(cu::malloc_async(nbytes, encoder), {nbytes}, int8); + encoder.add_temporary(gemm_args); + + ProblemSize* problem_sizes = gpu_ptr(gemm_args); + int64_t* a_lds = gpu_ptr(gemm_args) + problem_sizes_nbytes / 8; + int64_t* b_lds = a_lds + num_segments; + int64_t* out_lds = b_lds + num_segments; + void** a_ptrs = reinterpret_cast(out_lds + num_segments); + void** b_ptrs = a_ptrs + num_segments; + void** out_ptrs = b_ptrs + num_segments; + + // Build problem descriptions from segments on the GPU. + int block_size = std::min(num_segments, 256); + int num_blocks = cuda::ceil_div(num_segments, block_size); + + encoder.set_input_array(segments); + encoder.set_output_array(gemm_args); + encoder.add_kernel_node_ex( + cu::prepare_segmented_mm_data, + dim3(num_blocks), + dim3(block_size), + {}, + 0, + gpu_ptr(segments), + num_segments, + M, + N, + static_cast(lda), + static_cast(ldb), + static_cast(out.itemsize()), + a_transposed, + b_transposed, + gpu_ptr(a), + gpu_ptr(b), + gpu_ptr(out), + problem_sizes, + a_lds, + b_lds, + out_lds, + a_ptrs, + b_ptrs, + out_ptrs); + + // Dispatch grouped GEMM. + encoder.set_input_array(a); + encoder.set_input_array(b); + encoder.set_input_array(gemm_args); + encoder.set_output_array(out); + auto* fun = get_grouped_mm_funcion(a.dtype(), N, encoder.device()); + fun(a_transposed, + b_transposed, + num_segments, + problem_sizes, + a_lds, + b_lds, + out_lds, + a_ptrs, + b_ptrs, + out_ptrs, + encoder); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/hadamard.cu b/Source/Cmlx/mlx/mlx/backend/cuda/hadamard.cu index d7def2831..d6778c496 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/hadamard.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/hadamard.cu @@ -122,7 +122,9 @@ void hadamard_mn_contiguous( read_width_n2, read_width_m); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::vector kernel_names = {n2_kernel_name}; if (n1 > 1) { kernel_names.push_back(n1_kernel_name); @@ -142,8 +144,6 @@ void hadamard_mn_contiguous( return std::make_tuple(false, std::move(source), std::move(kernel_names)); }); - auto& encoder = cu::get_command_encoder(s); - if (n1 > 1) { const int64_t num_transforms = x.size() / n1; const uint32_t num_blocks = diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp index f855cd852..68f5ea139 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/indexing.cpp @@ -1,10 +1,13 @@ // Copyright © 2025 Apple Inc. #include "mlx/backend/common/compiled.h" +#include "mlx/backend/common/slicing.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/jit_module.h" #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/gpu/copy.h" +#include "mlx/backend/gpu/scan.h" +#include "mlx/backend/gpu/slicing.h" #include "mlx/dtype_utils.h" #include "mlx/primitives.h" @@ -23,6 +26,8 @@ namespace mlx::core { namespace { constexpr const char* g_scatter_ops[] = {"Max", "Min", "Sum", "Prod", "Assign"}; +constexpr const char* g_slice_ops[] = + {"Maximum", "Minimum", "Add", "Multiply", ""}; void append_indices_arg( cu::KernelArgs& args, @@ -82,7 +87,7 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { dtype_to_string(idx_dtype), nidx); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int large = 0; large <= 1; ++large) { @@ -130,7 +135,8 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { auto kernel = mod.get_kernel(kernel_name); auto [num_blocks, block_dims] = get_launch_args(out, large); - encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args()); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); } void Scatter::eval_gpu(const std::vector& inputs, array& out) { @@ -176,7 +182,9 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { nidx); auto& s = stream(); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int large = 0; large <= 1; ++large) { @@ -225,14 +233,14 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { idx_ndim, large ? "int64_t" : "int32_t"); - auto& encoder = cu::get_command_encoder(s); for (const auto& in : inputs) { encoder.set_input_array(in); } encoder.set_output_array(out); auto kernel = mod.get_kernel(kernel_name); auto [num_blocks, block_dims] = get_launch_args(upd, large); - encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args()); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); } void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { @@ -255,7 +263,7 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { dtype_to_string(out.dtype()), dtype_to_string(idx.dtype())); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int contiguous = 0; contiguous < 4; ++contiguous) { @@ -322,7 +330,8 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { encoder.set_output_array(out); auto kernel = mod.get_kernel(kernel_name); auto [num_blocks, block_dims] = get_launch_args(idx, large); - encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args()); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); } void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { @@ -358,7 +367,9 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { op); auto& s = stream(); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int contiguous = 0; contiguous < 4; ++contiguous) { @@ -421,14 +432,256 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { idx.flags().row_contiguous, large ? "int64_t" : "int32_t"); - auto& encoder = cu::get_command_encoder(s); for (const auto& in : inputs) { encoder.set_input_array(in); } encoder.set_output_array(out); auto kernel = mod.get_kernel(kernel_name); auto [num_blocks, block_dims] = get_launch_args(idx, large); - encoder.add_kernel_node(kernel, num_blocks, block_dims, 0, args.args()); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); +} + +void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("MaskedScatter::eval_gpu"); + assert(inputs.size() == 3); + + const array& dst = inputs[0]; + const array& mask = inputs[1]; + const array& src = inputs[2]; + + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + const size_t total = mask.size(); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + if (total == 0) { + return; + } + + array mask_flat = flatten_in_eval(mask, 1, -1, s); + if (gpu_ptr(mask_flat) != gpu_ptr(mask)) { + encoder.add_temporary(mask_flat); + } + if (!mask_flat.flags().row_contiguous) { + mask_flat = contiguous_copy_gpu(mask_flat, s); + encoder.add_temporary(mask_flat); + } + + array scatter_offsets(mask_flat.shape(), int32, nullptr, {}); + scatter_offsets.set_data(cu::malloc_async(scatter_offsets.nbytes(), encoder)); + encoder.add_temporary(scatter_offsets); + + scan_gpu_inplace( + mask_flat, + scatter_offsets, + Scan::Sum, + /* axis= */ 1, + /* reverse= */ false, + /* inclusive= */ false, + s); + + const size_t batch_count = mask.shape(0); + const size_t mask_batch_size = mask_flat.size() / batch_count; + const size_t src_batch_size = src.size() / src.shape(0); + bool large = total > INT32_MAX || src.size() > INT32_MAX; + bool vectorized = src.flags().row_contiguous && dst.flags().row_contiguous; + constexpr int kMaskedScatterVecSize = 16; + constexpr int kMaskedScatterVecBlockDim = 256; + + std::string module_name = + fmt::format("masked_scatter_{}", dtype_to_string(out.dtype())); + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + std::vector kernel_names; + for (int src_contiguous = 0; src_contiguous <= 1; ++src_contiguous) { + for (int dst_contiguous = 0; dst_contiguous <= 1; ++dst_contiguous) { + for (int use_large = 0; use_large <= 1; ++use_large) { + kernel_names.push_back( + fmt::format( + "mlx::core::cu::masked_scatter<{}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + src_contiguous ? "true" : "false", + dst_contiguous ? "true" : "false", + use_large ? "int64_t" : "int32_t")); + } + } + } + for (int use_large = 0; use_large <= 1; ++use_large) { + kernel_names.push_back( + fmt::format( + "mlx::core::cu::masked_scatter_vec_contiguous<{}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + use_large ? "int64_t" : "int32_t", + kMaskedScatterVecSize)); + } + return std::make_tuple(false, jit_source_scatter, std::move(kernel_names)); + }); + + cu::KernelArgs args; + args.append(dst); + args.append(mask_flat); + args.append(scatter_offsets); + args.append(src); + args.append(out); + if (large) { + args.append(mask_flat.size()); + args.append(src_batch_size); + args.append(mask_batch_size); + } else { + args.append(mask_flat.size()); + args.append(src_batch_size); + args.append(mask_batch_size); + } + if (!vectorized) { + args.append_ndim(dst.shape()); + args.append_ndim(dst.strides()); + args.append(dst.ndim()); + args.append_ndim(src.shape()); + args.append_ndim(src.strides()); + args.append(src.ndim()); + } + + encoder.set_input_array(dst); + encoder.set_input_array(mask_flat); + encoder.set_input_array(scatter_offsets); + encoder.set_input_array(src); + encoder.set_output_array(out); + + std::string kernel_name = vectorized + ? fmt::format( + "mlx::core::cu::masked_scatter_vec_contiguous<{}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + large ? "int64_t" : "int32_t", + kMaskedScatterVecSize) + : fmt::format( + "mlx::core::cu::masked_scatter<{}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + src.flags().row_contiguous ? "true" : "false", + dst.flags().row_contiguous ? "true" : "false", + large ? "int64_t" : "int32_t"); + auto kernel = mod.get_kernel(kernel_name); + auto [num_blocks, block_dims] = vectorized + ? get_launch_args( + mask_flat, large, kMaskedScatterVecSize, kMaskedScatterVecBlockDim) + : get_launch_args(mask_flat, large); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); +} + +void SliceUpdate::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("SliceUpdate::eval_gpu"); + assert(inputs.size() == 2); + if (out.size() == 0) { + return; + } + + auto& in = inputs[0]; + auto& upd = inputs[1]; + + if (upd.size() == 0) { + out.copy_shared_buffer(in); + return; + } + + auto ctype = in.flags().contiguous && in.size() == in.data_size() + ? CopyType::Vector + : CopyType::General; + copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream()); + + // Calculate out strides, initial offset and if copy needs to be made + auto [data_offset, out_strides] = + prepare_slice(out, start_indices_, strides_); + + // Do copy for None reduce type + if (reduce_type_ == SliceUpdate::None) { + copy_gpu_inplace( + /* const array& src = */ upd, + /* array& dst = */ out, + /* const Shape& data_shape = */ upd.shape(), + /* const Strides& i_strides = */ upd.strides(), + /* const Strides& o_strides = */ out_strides, + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ data_offset, + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ stream()); + return; + } + + auto [shape, strides] = + collapse_contiguous_dims(upd.shape(), {upd.strides(), out_strides}); + int nwork = 1; + if (shape.back() % 4 == 0) { + nwork = 4; + } else if (shape.back() % 2 == 0) { + nwork = 2; + } + + const char* op_name = g_slice_ops[reduce_type_]; + auto [ds, rc, cc] = check_contiguity(shape, strides[1]); + bool upd_contiguous = upd.flags().row_contiguous; + bool upd_scalar = upd.data_size() == 1; + bool out_contiguous = rc; + bool large = upd.size() > INT32_MAX; + std::string module_name = + fmt::format("slice_update_{}_{}", op_name, dtype_to_string(out.dtype())); + + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + std::vector kernel_names; + for (int out_c = 0; out_c <= 1; ++out_c) { + for (int upd_c = 0; upd_c <= 1; ++upd_c) { + for (int upd_s = 0; upd_s <= 1; ++upd_s) { + for (int large = 0; large <= 1; ++large) { + for (int nwork = 1; nwork <= 16; nwork *= 2) { + kernel_names.push_back( + fmt::format( + "mlx::core::cu::slice_update_op<{}, {}, mlx::core::cu::{}, {}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + large ? "int64_t" : "int32_t", + op_name, + out_c ? "true" : "false", + upd_c ? "true" : "false", + upd_s ? "true" : "false", + nwork)); + } + } + } + } + } + return std::make_tuple( + false, jit_source_slice_update, std::move(kernel_names)); + }); + + cu::KernelArgs args; + args.append(upd); + args.append(out); + args.append(upd.size()); + args.append_ndim(shape); + args.append_ndim(strides[0]); + args.append(shape.size()); + args.append_ndim(strides[1]); + args.append(data_offset); + + encoder.set_input_array(upd); + encoder.set_output_array(out); + + std::string kernel_name; + kernel_name = fmt::format( + "mlx::core::cu::slice_update_op<{}, {}, mlx::core::cu::{}, {}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + large ? "int64_t" : "int32_t", + op_name, + out_contiguous, + upd_contiguous, + upd_scalar, + nwork); + + auto kernel = mod.get_kernel(kernel_name); + auto [num_blocks, block_dims] = get_launch_args(upd, large, nwork); + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args.args()); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.cpp index b0ebb4019..0d107a8e1 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.cpp @@ -9,12 +9,16 @@ #include #include #include +#include #include #include namespace mlx::core::cu { +// Defined in dirs.cpp to avoid invalidating compile cache. +const char* cccl_dir(); + namespace { #define CHECK_NVRTC_ERROR(cmd) check_nvrtc_error(#cmd, (cmd)) @@ -26,47 +30,79 @@ void check_nvrtc_error(const char* name, nvrtcResult err) { } } +// Return the default path to CUDA toolkit. +const std::filesystem::path& default_cuda_toolkit_path() { +#if defined(_WIN32) + static auto cached_path = []() -> std::filesystem::path { + std::filesystem::path root( + LR"(C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA)"); + for (auto& file : std::filesystem::directory_iterator(root)) { + if (std::filesystem::exists(file.path() / "include" / "cuda.h")) { + return file.path(); + } + } + return {}; + }(); +#else + static std::filesystem::path cached_path = "/usr/local/cuda"; +#endif + return cached_path; +} + +// Get the dirname of nvidia python package that contains CUDA headers. +inline const char* cudart_dirname() { +#if CUDART_VERSION < 13000 + return "cuda_runtime"; +#elif CUDART_VERSION < 14000 + return "cu13"; +#else + static_assert( + false, + "Please find out the newest dirname under site-packages/nvidia " + "and add it in this function."); +#endif +} + // Return the --include-path args used for invoking NVRTC. const std::vector& include_path_args() { static std::vector cached_args = []() { std::vector args; - // Add path to bundled CCCL headers. + // Add path to bundled headers. auto root_dir = current_binary_dir(); #if !defined(_WIN32) root_dir = root_dir.parent_path(); #endif - auto path = root_dir / "include" / "cccl"; -#if defined(MLX_CCCL_DIR) - if (!std::filesystem::exists(path)) { - path = MLX_CCCL_DIR; + auto path = root_dir / "include"; + if (std::filesystem::exists(path)) { + args.push_back(fmt::format("--include-path={}", path.string())); + } + // Add path to CCCL headers. + path = path / "cccl"; + if (!std::filesystem::exists(path) && cccl_dir()) { + path = cccl_dir(); } -#endif if (std::filesystem::exists(path)) { args.push_back(fmt::format("--include-path={}", path.string())); } // Add path to CUDA runtime headers, try local-installed python package // first and then system-installed headers. - path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include"; - if (std::filesystem::exists(path)) { - args.push_back(fmt::format("--include-path={}", path.string())); - } else { + path = root_dir.parent_path() / "nvidia" / cudart_dirname() / "include"; + if (!std::filesystem::exists(path)) { const char* home = std::getenv("CUDA_HOME"); if (!home) { home = std::getenv("CUDA_PATH"); } -#if defined(__linux__) - if (!home) { - home = "/usr/local/cuda"; + path = home ? std::filesystem::path(home) : default_cuda_toolkit_path(); + if (!path.empty()) { + path = path / "include"; } -#endif - if (home && std::filesystem::exists(home)) { - args.push_back(fmt::format("--include-path={}/include", home)); - } else { + if (path.empty() || !std::filesystem::exists(path)) { throw std::runtime_error( "Can not find locations of CUDA headers, please set environment " "variable CUDA_HOME or CUDA_PATH."); } } + args.push_back(fmt::format("--include-path={}", path.string())); return args; }(); return cached_args; @@ -230,8 +266,14 @@ constexpr const char* g_include_names[] = { INCLUDE_PREFIX "cast_op.cuh", INCLUDE_PREFIX "config.h", INCLUDE_PREFIX "complex.cuh", + INCLUDE_PREFIX "cute_dequant.cuh", INCLUDE_PREFIX "fp16_math.cuh", + INCLUDE_PREFIX "gemm_sm70.cuh", + INCLUDE_PREFIX "hadamard.cuh", INCLUDE_PREFIX "indexing.cuh", + INCLUDE_PREFIX "qmm_naive.cuh", + INCLUDE_PREFIX "qmm_sm80.cuh", + INCLUDE_PREFIX "qmm_sm90.cuh", INCLUDE_PREFIX "scatter_ops.cuh", INCLUDE_PREFIX "unary_ops.cuh", INCLUDE_PREFIX "ternary_ops.cuh", @@ -246,8 +288,14 @@ constexpr const char* g_headers[] = { jit_source_cast_op, jit_source_config, jit_source_complex, + jit_source_cute_dequant, jit_source_fp16_math, + jit_source_gemm_sm70, + jit_source_hadamard, jit_source_indexing, + jit_source_qmm_naive, + jit_source_qmm_sm80, + jit_source_qmm_sm90, jit_source_scatter_ops, jit_source_unary_ops, jit_source_ternary_ops, @@ -277,8 +325,11 @@ void compile( CHECK_NVRTC_ERROR(nvrtcAddNameExpression(prog, name.c_str())); } - // Compile program. + // Required for compiling CUTLASS code. std::vector args; + args.push_back("--device-as-default-execution-space"); + + // Target current device. bool use_sass = compiler_supports_device_sass(device); auto cc = device.compute_capability_major(); std::string arch_tag = (cc >= 9) ? "a" : ""; @@ -292,6 +343,8 @@ void compile( for (const auto& include : include_path_args()) { args.push_back(include.c_str()); } + + // Compile program. nvrtcResult compile_result = nvrtcCompileProgram(prog, args.size(), args.data()); if (compile_result != NVRTC_SUCCESS) { @@ -422,20 +475,27 @@ CUfunction JitModule::get_kernel( return get_kernel_and_dims(kernel_name, std::move(configure_kernel)).first; } -std::unordered_map& get_jit_module_cache() { - static std::unordered_map map; - return map; -} - JitModule& get_jit_module( - const mlx::core::Device& device, + Device& device, const std::string& name, const KernelBuilder& builder, - bool cache) { - auto& map = get_jit_module_cache(); - auto it = map.find(name); - if (it == map.end()) { - it = map.try_emplace(name, cu::device(device), name, builder, cache).first; + bool use_disk_cache) { + // The cache are leak intentionally as user code may still be running JIT + // compiled code after main thread teardown. + static auto* cache = new std::unordered_map; + static auto* mtx = new std::shared_mutex; + + { + std::shared_lock rlock(*mtx); + if (auto it = cache->find(name); it != cache->end()) { + return it->second; + } + } + + std::unique_lock wlock(*mtx); + auto it = cache->find(name); + if (it == cache->end()) { + it = cache->try_emplace(name, device, name, builder, use_disk_cache).first; } return it->second; } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.h b/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.h index 3849d8d00..f6796625d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/jit_module.h @@ -91,11 +91,12 @@ class JitModule { Device& device, const std::string& module_name, const KernelBuilder& builder, - bool cache); + bool use_disk_cache); ~JitModule(); JitModule(const JitModule&) = delete; JitModule& operator=(const JitModule&) = delete; + CUfunction get_kernel( const std::string& kernel_name, std::function configure_kernel = nullptr); @@ -109,10 +110,8 @@ class JitModule { kernels_; }; -std::unordered_map& get_jit_module_cache(); - JitModule& get_jit_module( - const mlx::core::Device& device, + Device& device, const std::string& name, const KernelBuilder& builder, bool use_disk_cache = true); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cu b/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cu index b8cef139a..14eae1a48 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cu @@ -49,4 +49,16 @@ std::tuple get_launch_args( return std::make_tuple(num_blocks, block_dim); } +std::pair +get_launch_args_general(int dim0, size_t rest, int work_per_thread /* = 1 */) { + constexpr uint32_t max_grid_yz_dim = 65535; + dim0 = (dim0 + work_per_thread - 1) / work_per_thread; + dim3 block_dims = get_block_dims(dim0, rest, 1); + uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); + uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + uint32_t num_blocks_z = cuda::ceil_div(num_blocks_y, max_grid_yz_dim); + num_blocks_y = cuda::ceil_div(num_blocks_y, num_blocks_z); + return {dim3(num_blocks_x, num_blocks_y, num_blocks_z), block_dims}; +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cuh index 1507a0ae1..344b570f1 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/kernel_utils.cuh @@ -89,7 +89,8 @@ using cuda_type_t = typename CTypeToCudaType::type; template inline constexpr bool is_floating_v = cuda::std::is_same_v || cuda::std::is_same_v || - cuda::std::is_same_v || cuda::std::is_same_v; + cuda::std::is_same_v || cuda::std::is_same_v || + cuda::std::is_same_v || cuda::std::is_same_v; // Type traits for detecting complex numbers. template @@ -145,4 +146,7 @@ inline std::tuple get_launch_args( max_block_dim); } +std::pair +get_launch_args_general(int dim0, size_t rest, int work_per_thread = 1); + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/layer_norm.cu b/Source/Cmlx/mlx/mlx/backend/cuda/layer_norm.cu index da554a1fe..9c8d46fbe 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/layer_norm.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/layer_norm.cu @@ -279,7 +279,6 @@ void LayerNorm::eval_gpu( kernel, n_rows, block_dim(), - 0, gpu_ptr(x), gpu_ptr(w), gpu_ptr(b), @@ -392,7 +391,6 @@ void LayerNormVJP::eval_gpu( kernel, n_rows, block_dim(), - 0, gpu_ptr(x), gpu_ptr(w), gpu_ptr(g), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/logsumexp.cu b/Source/Cmlx/mlx/mlx/backend/cuda/logsumexp.cu index a55a7d01e..25d24faf2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/logsumexp.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/logsumexp.cu @@ -150,7 +150,6 @@ void LogSumExp::eval_gpu(const std::vector& inputs, array& out) { kernel, n_rows, block_dim(), - 0, gpu_ptr(in), gpu_ptr(out), axis_size); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h b/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h index f23f3ff05..d208315a0 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/lru_cache.h @@ -137,9 +137,15 @@ class LRUCache { // Turn a POD struct into a container key by doing bytes compare. // +// IMPORTANT: Do not use aggregate init on the pod field (key.pod = {...}). +// It creates a stack temporary whose padding bytes are uninitialized, and +// trivial copy-assignment copies the entire struct including padding — +// breaking the memcmp-based comparison. Set fields individually instead. +// // Usage: // BytesKey key; -// key.pod = { ... }; +// key.pod.field1 = value1; +// key.pod.field2 = value2; template struct BytesKey { T pod; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/matmul.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/matmul.cpp index 725f1b408..7505ea2cb 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/matmul.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/matmul.cpp @@ -2,7 +2,9 @@ #include "mlx/backend/common/matmul.h" #include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/gemms/block_mask.h" #include "mlx/backend/cuda/gemms/cublas_gemm.h" +#include "mlx/backend/cuda/gemms/gather_gemm.h" #include "mlx/backend/cuda/gemms/gemv.h" #include "mlx/backend/cuda/gemms/grouped_gemm.h" #include "mlx/backend/gpu/copy.h" @@ -136,40 +138,6 @@ void gemm_and_bias( encoder, out, a, b, batch_shape, a_batch_strides, b_batch_strides, alpha); } -void gather_mm_rhs( - const array& a_, - const array& b_, - const array& indices_, - array& out, - cu::CommandEncoder& encoder, - Stream s) { - if (a_.size() / a_.shape(-2) / a_.shape(-1) != indices_.size()) { - throw std::runtime_error("[gather_mm] Broadcasting lhs is not supported."); - } - - int group_count = b_.size() / b_.shape(-1) / b_.shape(-2); - if (group_count > 1024) { - throw std::runtime_error( - "[gather_mm] Group count can not be larger than 1024."); - } - - auto [a_transposed, lda, a] = ensure_batch_contiguous(a_, encoder, s); - auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_, encoder, s); - auto indices = ensure_row_contiguous(indices_, encoder, s); - - cutlass_grouped_gemm_unaligned( - a_transposed, - lda, - b_transposed, - ldb, - group_count, - a, - b, - indices, - out, - encoder); -} - } // namespace void Matmul::eval_gpu(const std::vector& inputs, array& out) { @@ -203,6 +171,82 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b); } +void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("BlockMaskedMM::eval_gpu"); + if (!issubdtype(out.dtype(), floating)) { + throw std::runtime_error( + "[BlockMaskedMM] Does not yet support non-floating point types."); + } + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + auto& a_pre = inputs[0]; + auto& b_pre = inputs[1]; + + // Return 0s if either input is empty. + if (a_pre.size() == 0 || b_pre.size() == 0) { + array zero(0, a_pre.dtype()); + encoder.add_temporary(zero); + fill_gpu(zero, out, s); + return; + } + + int M = a_pre.shape(-2); + int N = b_pre.shape(-1); + int K = a_pre.shape(-1); + + if (M == 0 || N == 0) { + return; + } + if (K == 0) { + array zero(0, a_pre.dtype()); + encoder.add_temporary(zero); + fill_gpu(zero, out, s); + return; + } + + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + bool has_op_mask = inputs.size() > 3; + bool has_out_mask = inputs.size() == 3 || inputs.size() == 5; + + int64_t batch_count = out.size() / (int64_t(M) * N); + + bool a_transposed; + int64_t lda; + array a = a_pre; + bool b_transposed; + int64_t ldb; + array b = b_pre; + + if (has_op_mask) { + // Fused copy + mask in a single pass per matrix. + auto& lhs_mask = inputs[inputs.size() - 2]; + auto& rhs_mask = inputs[inputs.size() - 1]; + a = copy_with_block_mask( + encoder, a_pre, lhs_mask, block_size_, M, K, batch_count); + b = copy_with_block_mask( + encoder, b_pre, rhs_mask, block_size_, K, N, batch_count); + a_transposed = false; + lda = K; + b_transposed = false; + ldb = N; + } else { + std::tie(a_transposed, lda, a) = check_transpose(encoder, s, a_pre); + std::tie(b_transposed, ldb, b) = check_transpose(encoder, s, b_pre); + } + + // Run GEMM. + gemm_and_bias( + encoder, M, N, K, a_transposed, lda, b_transposed, ldb, out, a, b); + + // Apply output mask. + if (has_out_mask) { + auto& out_mask = inputs[2]; + apply_block_mask(encoder, out, out_mask, block_size_, M, N, batch_count); + } +} + void AddMM::eval_gpu(const std::vector& inputs, array& out) { nvtx3::scoped_range r("AddMM::eval_gpu"); auto& s = stream(); @@ -327,14 +371,12 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { auto& encoder = cu::get_command_encoder(s); assert(inputs.size() == 4); - auto& a = inputs[0]; - auto& b = inputs[1]; - auto& lhs_indices = inputs[2]; - auto& rhs_indices = inputs[3]; + auto& a_pre = inputs[0]; + auto& b_pre = inputs[1]; - // Return 0s if either input is empty. - if (a.size() == 0 || b.size() == 0) { - array zero(0, a.dtype()); + // Return 0s if the output or either input is empty. + if (out.size() == 0 || a_pre.size() == 0 || b_pre.size() == 0) { + array zero(0, a_pre.dtype()); encoder.add_temporary(zero); fill_gpu(zero, out, s); return; @@ -343,31 +385,87 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { out.set_data(cu::malloc_async(out.nbytes(), encoder)); // Extract shapes from inputs. - int M = a.shape(-2); - int N = b.shape(-1); - int K = a.shape(-1); + int M = a_pre.shape(-2); + int N = b_pre.shape(-1); + int K = a_pre.shape(-1); + + auto [a_transposed, lda, a] = ensure_batch_contiguous(a_pre, encoder, s); + auto [b_transposed, ldb, b] = ensure_batch_contiguous(b_pre, encoder, s); + auto lhs_indices = ensure_row_contiguous(inputs[2], encoder, s); + auto rhs_indices = ensure_row_contiguous(inputs[3], encoder, s); // We are walking a in order and b is also in order so we can batch up the // matmuls and reuse reading a and b. if (M == 1 && right_sorted_ == true) { - gather_mm_rhs(a, b, rhs_indices, out, encoder, s); + cutlass_grouped_gemm_unaligned( + a_transposed, + lda, + b_transposed, + ldb, + b.size() / b.shape(-1) / b.shape(-2), // group_count + a, + b, + rhs_indices, + out, + encoder); return; } - auto [transposed_a, lda, a_] = check_transpose(encoder, s, a); - auto [transposed_b, ldb, b_] = check_transpose(encoder, s, b); - auto use_gemv = cu::can_use_gemv(M, N, K, transposed_a, transposed_b); + auto use_gemv = cu::can_use_gemv(M, N, K, a_transposed, b_transposed); if (M == 1 && use_gemv) { - gather_mv(b_, a_, rhs_indices, lhs_indices, out, N, K, encoder); + gather_mv(b, a, rhs_indices, lhs_indices, out, N, K, encoder); return; } - if (N == 1 && use_gemv) { - gather_mv(a_, b_, lhs_indices, rhs_indices, out, M, K, encoder); + gather_mv(a, b, lhs_indices, rhs_indices, out, M, K, encoder); + return; + } + + gather_mm( + a_transposed, b_transposed, a, b, lhs_indices, rhs_indices, out, encoder); +} + +void SegmentedMM::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("SegmentedMM::eval_gpu"); + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + assert(inputs.size() == 3); + auto& a_pre = inputs[0]; + auto& b_pre = inputs[1]; + auto& segments_pre = inputs[2]; + + // Return zeros if output is empty or either input is empty. + if (out.size() == 0 || a_pre.size() == 0 || b_pre.size() == 0) { + array zero(0, a_pre.dtype()); + encoder.add_temporary(zero); + fill_gpu(zero, out, s); return; } - throw std::runtime_error("NYI"); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + int M = a_pre.shape(-2); + int N = b_pre.shape(-1); + int num_segments = segments_pre.size() / 2; + + auto [a_transposed, lda, a] = check_transpose(encoder, s, a_pre); + auto [b_transposed, ldb, b] = check_transpose(encoder, s, b_pre); + auto segments = ensure_row_contiguous(segments_pre, encoder, s); + + cutlass_segmented_mm( + a_transposed, + lda, + b_transposed, + ldb, + num_segments, + M, + N, + a, + b, + segments, + out, + encoder); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp index 0caac8de6..94f260767 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/primitives.cpp @@ -24,19 +24,13 @@ namespace mlx::core { throw std::runtime_error(#func " has no CUDA implementation."); \ } -NO_GPU(BlockMaskedMM) -NO_GPU(FFT) -NO_GPU(GatherQMM) -NO_GPU(Hadamard) NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) -NO_GPU(SegmentedMM) NO_GPU_MULTI(SVD) NO_GPU(Inverse) NO_GPU(Cholesky) NO_GPU_MULTI(Eig) NO_GPU_MULTI(Eigh) -NO_GPU(MaskedScatter) namespace distributed { NO_GPU_MULTI(Send) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/ptx.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/ptx.cuh new file mode 100644 index 000000000..b7b87e671 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/ptx.cuh @@ -0,0 +1,125 @@ +#pragma once + +#include + +namespace mlx::core { + +namespace ptx { + +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + +__device__ __forceinline__ void mbarrier_init(uint64_t* mbar, uint32_t count) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.init.shared.b64 [%0], %1;" + : + : "r"(mbar_ptr), "r"(count) + : "memory"); +} + +__device__ __forceinline__ void mbarrier_invalidate(uint64_t* mbar) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.inval.shared.b64 [%0];" : : "r"(mbar_ptr) : "memory"); +} + +// Arrive at barrier (non-master threads) +__device__ __forceinline__ void mbarrier_arrive(uint64_t* mbar) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.arrive.shared.b64 _, [%0];" + : + : "r"(mbar_ptr) + : "memory"); +} + +// Arrive at barrier and set expected transaction count (master thread) +__device__ __forceinline__ void mbarrier_arrive_expect_tx( + uint64_t* mbar, + uint32_t tx_count) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" + : + : "r"(mbar_ptr), "r"(tx_count) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-asynchronous-copy-completion-mechanisms-mbarrier +__device__ __forceinline__ void mbarrier_wait_parity( + uint64_t* mbar, + uint32_t parity) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile( + "{\n\t" + ".reg .pred P;\n\t" + "WAIT_LOOP:\n\t" + "mbarrier.try_wait.parity.shared.b64 P, [%0], %1;\n\t" + "@!P bra WAIT_LOOP;\n\t" + "}\n\t" + : + : "r"(mbar_ptr), "r"(parity) + : "memory"); +} + +// Async bulk tensor copy: global -> shared (2D) +__device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( + void* dst_shmem, + const CUtensorMap* tensor_map, + uint32_t tile_x, + uint32_t tile_y, + uint64_t* mbar) { + uint32_t dst_ptr = __cvta_generic_to_shared(dst_shmem); + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" + : + : "r"(dst_ptr), "l"(tensor_map), "r"(tile_x), "r"(tile_y), "r"(mbar_ptr) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor +// shared::cta -> global +__device__ __forceinline__ void cp_async_bulk_tensor_2d_shared_to_global( + const uint64_t* tensor_map_ptr, + const uint32_t offset_x, + const uint32_t offset_y, + uint64_t* src_shmem) { + uint32_t src_shmem_ptr = __cvta_generic_to_shared(src_shmem); + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%1, %2}], [%3];" :: + "l"(tensor_map_ptr), + "r"(offset_x), + "r"(offset_y), + "r"(src_shmem_ptr) + : "memory"); +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group +template +__device__ __forceinline__ void cp_async_bulk_wait_group_read() { + if constexpr (N == 0) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + } else if constexpr (N == 1) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } else if constexpr (N == 2) { + asm volatile("cp.async.bulk.wait_group.read 2;"); + } else if constexpr (N == 4) { + asm volatile("cp.async.bulk.wait_group.read 4;"); + } +} + +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group +__device__ __forceinline__ void cp_async_bulk_commit_group() { + asm volatile("cp.async.bulk.commit_group;"); +} + +// Ccreates a memory ordering barrier between generic and async proxies +// details: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#proxies +__device__ __forceinline__ void fence_proxy_async_shared_cta() { + asm volatile("fence.proxy.async.shared::cta;"); +} + +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) + +} // namespace ptx +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/affine_quantize.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/affine_quantize.cu index 105547bb3..13bd35dad 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/affine_quantize.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/affine_quantize.cu @@ -1,8 +1,8 @@ // Copyright © 2025 Apple Inc. +#include "mlx/backend/common/quantized.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/quantized_utils.cuh" #include "mlx/dtype_utils.h" #include @@ -27,8 +27,8 @@ affine_quantize(const T* w, uint8_t* out, T* scales, T* biases, size_t size) { constexpr float eps = 1e-7; constexpr int simd_size = WARP_SIZE; constexpr float n_bins = (1 << bits) - 1; - constexpr int pack_factor = get_pack_factor(); - constexpr int bytes_per_pack = get_bytes_per_pack(); + constexpr int pack_factor = get_pack_factor(bits, 8); + constexpr int bytes_per_pack = get_bytes_per_pack(bits); constexpr int values_per_reduce = group_size / simd_size; constexpr int writes_per_reduce = pack_factor / values_per_reduce; constexpr int writes_per_pack = @@ -142,8 +142,8 @@ __global__ void affine_dequantize( auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - constexpr int pack_factor = get_pack_factor(); - constexpr int bytes_per_pack = get_bytes_per_pack(); + constexpr int pack_factor = get_pack_factor(bits, 8); + constexpr int bytes_per_pack = get_bytes_per_pack(bits); size_t offset = tidx + grid_dim_x * size_t(tidy); size_t oindex = offset * pack_factor; @@ -226,6 +226,45 @@ __global__ void affine_dequantize( } // namespace cu +template +void dispatch_groups(int group_size, F&& f) { + switch (group_size) { + case 32: + f(std::integral_constant{}); + break; + case 64: + f(std::integral_constant{}); + break; + case 128: + f(std::integral_constant{}); + break; + } +} + +template +void dispatch_bits(int bits, F&& f) { + switch (bits) { + case 2: + f(std::integral_constant{}); + break; + case 3: + f(std::integral_constant{}); + break; + case 4: + f(std::integral_constant{}); + break; + case 5: + f(std::integral_constant{}); + break; + case 6: + f(std::integral_constant{}); + break; + case 8: + f(std::integral_constant{}); + break; + } +} + void affine_quantize( const array& w, array& wq, @@ -259,7 +298,6 @@ void affine_quantize( kernel, num_blocks, block_dims, - 0, gpu_ptr(w), gpu_ptr(wq), gpu_ptr(scales), @@ -315,7 +353,6 @@ void affine_dequantize( kernel, num_blocks, block_dims, - 0, gpu_ptr(wq), gpu_ptr(scales), gpu_ptr(biases), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp index 959ae9352..f81b6455c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.cpp @@ -13,39 +13,26 @@ namespace mlx::core { namespace { -// Currently cublas supports only mxfp8 and nvfp4 -// quantization modes for block scaled quantization -cudaDataType_t qmode_to_cublas_scale_dtype(std::string mode) { - if (mode == "mxfp8") { - return CUDA_R_8F_UE8M0; - } else if (mode == "nvfp4") { - return CUDA_R_8F_UE4M3; - } else { - throw std::runtime_error( - fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode)); - } -} - -cudaDataType_t qmode_to_cublas_dtype(std::string mode) { - if (mode == "mxfp8") { - return CUDA_R_8F_E4M3; - } else if (mode == "nvfp4") { - return CUDA_R_4F_E2M1; - } else { - throw std::runtime_error( - fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode)); - } -} +struct QuantModeConfig { + cudaDataType_t data_type; + cudaDataType_t scale_dtype; + cublasLtMatmulMatrixScale_t scale_mode; +}; -cublasLtMatmulMatrixScale_t qmode_to_cublas_scale_mode(std::string mode) { +QuantModeConfig get_quant_mode_config(const std::string& mode) { if (mode == "mxfp8") { - return CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + return { + CUDA_R_8F_E4M3, + CUDA_R_8F_UE8M0, + CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0}; } else if (mode == "nvfp4") { - return CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; - } else { - throw std::runtime_error( - fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode)); + return { + CUDA_R_4F_E2M1, + CUDA_R_8F_UE4M3, + CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3}; } + throw std::runtime_error( + fmt::format("Unsupported quantization mode in CublasQQMM: {}.", mode)); } } // namespace @@ -64,21 +51,21 @@ CublasQQMM::CublasQQMM( int64_t a_batch_stride, int64_t b_batch_stride, Dtype out_dtype, - std::string qmode) { + const std::string& qmode) { + auto config = get_quant_mode_config(qmode); + // The compute type must be CUBLAS_COMPUTE_32F. // The scale type must be CUDA_R_32F. cudaDataType_t scale_type = CUDA_R_32F; cublasComputeType_t gemm_compute_type = CUBLAS_COMPUTE_32F; cudaDataType_t output_type = cublas_utils::dtype_to_cublas_type(out_dtype, "CublasQQMM"); - cudaDataType_t data_type = qmode_to_cublas_dtype(qmode); - quantization_mode_ = std::string(qmode); init_base( device, scale_type, gemm_compute_type, - data_type, + config.data_type, output_type, a_transposed, a_rows, @@ -92,8 +79,8 @@ CublasQQMM::CublasQQMM( a_batch_stride, b_batch_stride); - a_scale_mode_ = qmode_to_cublas_scale_mode(qmode); - b_scale_mode_ = qmode_to_cublas_scale_mode(qmode); + a_scale_mode_ = config.scale_mode; + b_scale_mode_ = config.scale_mode; CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( matmul_desc_, @@ -123,7 +110,7 @@ CublasQQMM::CublasQQMM( int64_t b_batch_stride, int64_t c_batch_stride, Dtype out_dtype, - std::string qmode) + const std::string& qmode) : CublasQQMM( device, a_transposed, @@ -158,11 +145,14 @@ void CublasQQMM::run( const array& b, const array& a_scale, const array& b_scale, - float alpha) { + const array& alpha, + const array& beta) { encoder.set_input_array(a); encoder.set_input_array(b); encoder.set_input_array(a_scale); encoder.set_input_array(b_scale); + encoder.set_input_array(alpha); + encoder.set_input_array(beta); encoder.set_output_array(out); execute( @@ -173,19 +163,37 @@ void CublasQQMM::run( gpu_ptr(a_scale), gpu_ptr(b_scale), nullptr, - alpha); + gpu_ptr(alpha), + gpu_ptr(beta)); } -void CublasQQMM::execute( +void CublasQQMM::run( + cu::CommandEncoder& encoder, + array& out, + const array& a, + const array& b, + const array& a_scale, + const array& b_scale) { + encoder.set_input_array(a); + encoder.set_input_array(b); + encoder.set_input_array(a_scale); + encoder.set_input_array(b_scale); + encoder.set_output_array(out); + + execute( + encoder, + gpu_ptr(out), + gpu_ptr(a), + gpu_ptr(b), + gpu_ptr(a_scale), + gpu_ptr(b_scale), + nullptr); +} + +void CublasQQMM::set_scales_ptrs( cu::CommandEncoder& encoder, - void* out, - const void* a, - const void* b, const void* a_scale, - const void* b_scale, - const void* c, - float alpha /* = 1 */, - float beta /* = 0 */) { + const void* b_scale) { CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( matmul_desc_, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, @@ -196,6 +204,49 @@ void CublasQQMM::execute( CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &a_scale, sizeof(a_scale))); +} + +void CublasQQMM::execute( + cu::CommandEncoder& encoder, + void* out, + const void* a, + const void* b, + const void* a_scale, + const void* b_scale, + const void* c, + const void* alpha, + const void* beta) { + set_scales_ptrs(encoder, a_scale, b_scale); + // alpha and beta are both should be device pointers for nvfp4 + // by default cublas uses host pointers + // https://docs.nvidia.com/cuda/cublas/#cublasltpointermode-t + cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE; + CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( + matmul_desc_, + CUBLASLT_MATMUL_DESC_POINTER_MODE, + &pointer_mode, + sizeof(pointer_mode))); + execute_matmul(encoder, out, a, b, c, alpha, beta); +} + +void CublasQQMM::execute( + cu::CommandEncoder& encoder, + void* out, + const void* a, + const void* b, + const void* a_scale, + const void* b_scale, + const void* c, + const float alpha /* = 1 */, + const float beta /* = 0 */) { + set_scales_ptrs(encoder, a_scale, b_scale); + // alpha and beta are both should be host pointers + cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST; + CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute( + matmul_desc_, + CUBLASLT_MATMUL_DESC_POINTER_MODE, + &pointer_mode, + sizeof(pointer_mode))); const void* alpha_ptr = α const void* beta_ptr = β diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h index 0a710f6e1..a9095012c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cublas_qqmm.h @@ -25,7 +25,7 @@ class CublasQQMM : public CublasMatmulBase { int64_t a_batch_stride, int64_t b_batch_stride, Dtype out_dtype, - std::string quantization_mode); + const std::string& quantization_mode); CublasQQMM( cu::Device& device, @@ -43,7 +43,7 @@ class CublasQQMM : public CublasMatmulBase { int64_t b_batch_stride, int64_t c_batch_stride, Dtype out_dtype, - std::string quantization_mode); + const std::string& quantization_mode); void run( cu::CommandEncoder& encoder, @@ -52,20 +52,33 @@ class CublasQQMM : public CublasMatmulBase { const array& b, const array& a_scale, const array& b_scale, - float alpha = 1.0f); + const array& alpha, + const array& beta); - private: - void run_batched( + void run( cu::CommandEncoder& encoder, array& out, const array& a, const array& b, const array& a_scale, - const array& b_scale, - const Shape& batch_shape, - const Strides& a_batch_strides, - const Strides& b_batch_strides, - float alpha); + const array& b_scale); + + private: + void set_scales_ptrs( + cu::CommandEncoder& encoder, + const void* a_scale, + const void* b_scale); + + void execute( + cu::CommandEncoder& encoder, + void* out, + const void* a, + const void* b, + const void* a_scale, + const void* b_scale, + const void* c, + const void* alpha, + const void* beta); void execute( cu::CommandEncoder& encoder, @@ -75,10 +88,9 @@ class CublasQQMM : public CublasMatmulBase { const void* a_scale, const void* b_scale, const void* c, - float alpha = 1, - float beta = 0); + const float alpha = 1.0f, + const float beta = 0.0f); - std::string quantization_mode_; cublasLtMatmulMatrixScale_t a_scale_mode_; cublasLtMatmulMatrixScale_t b_scale_mode_; cublasLtMatmulMatrixScale_t c_scale_mode_; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cuda_fp4.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cuda_fp4.h deleted file mode 100644 index 4500be6b7..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/cuda_fp4.h +++ /dev/null @@ -1,100 +0,0 @@ -#pragma once - -struct __nv_fp8_e8m0 { - __device__ __nv_fp8_e8m0(float x) { - if (!std::isfinite(x)) { - __x = 0xFF; - return; - } - if (x < 0.0f) { - __x = 0x00; - return; - } - float le = std::log2f(x); - int n = static_cast(std::nearbyintf(le)); - - n = n < -127 ? -127 : n; - n = n > 127 ? 127 : n; - __x = static_cast(n + 127); - } - - __device__ operator float() { - if (__x == 0xFF) { - return std::numeric_limits::quiet_NaN(); - } - return std::ldexp(1.0f, static_cast(__x) - 127); - } - - uint8_t __x{0}; -}; - -struct __nv_fp4_e2m1 { - __device__ __nv_fp4_e2m1(float x) { - if (std::isnan(x)) { - __x = 0x7; - return; - } - - const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0; - x = std::abs(x); - - if (x > 5.0f) { - __x = 0x7; - } else if (x >= 3.5f) { - __x = 0x6; - } else if (x > 2.5f) { - __x = 0x5; - } else if (x >= 1.75f) { - __x = 0x4; - } else if (x > 1.25f) { - __x = 0x3; - } else if (x >= 0.75f) { - __x = 0x2; - } else if (x > 0.25f) { - __x = 0x1; - } else { - __x = 0x0; - } - __x |= sign_bit; - } - - __device__ operator float() { - static const float LUT[16] = { - 0.0f, - 0.5f, - 1.0f, - 1.5f, - 2.0f, - 3.0f, - 4.0f, - 6.0f, - -0.0f, - -0.5f, - -1.0f, - -1.5f, - -2.0f, - -3.0f, - -4.0f, - -6.0f}; - - return LUT[__x]; - } - uint8_t __x{0}; -}; - -struct __nv_fp4x4_e2m1 { - __device__ operator float4() { - float4 out; - auto bits = __high & 0xf; - out.x = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__high >> 4) & 0xf; - out.y = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__low) & 0xf; - out.z = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__low >> 4) & 0xf; - out.w = float(*(__nv_fp4_e2m1*)(&bits)); - return out; - } - uint8_t __high{0}; - uint8_t __low{0}; -}; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cu index 3b4e96ef2..b5bfd72f2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cu @@ -1,329 +1,88 @@ // Copyright © 2025 Apple Inc. +#include "mlx/backend/common/quantized.h" #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh" -#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh" +#include "mlx/backend/cuda/quantized/fp_quantize.cuh" #include "mlx/backend/cuda/quantized/quantized.h" -#include "mlx/backend/cuda/quantized/quantized_utils.cuh" #include "mlx/backend/cuda/vector_types.cuh" #include "mlx/dtype_utils.h" #include #include +#include +#include namespace mlx::core { namespace cu { -template -struct Dequantize { - __device__ float operator()(uint8_t x) { - if constexpr (bits == 8) { - return float(*(__nv_fp8_e4m3*)(&x)); - } else { - return float(*(__nv_fp4_e2m1*)(&x)); - } - } -}; - -namespace cg = cooperative_groups; - -template -__global__ void fp_quantize_dequantize(T* w, T* out, size_t size) { - using Tx2 = Vector2_t; - using Tx4 = Vector4_t; - uint32_t rbits = 0; // reserved bits for future use - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - size_t thread_idx = tidx + grid_dim_x * size_t(tidy); - size_t base_idx = thread_idx * group_size; - - if (base_idx >= size) { - return; - } - - auto w_tile = load_vector(w, thread_idx); - float scale = 0.0f; - - Tx2 amax_2x = Tx2{0.0f, 0.0f}; - -#pragma unroll - for (int i = 0; i < group_size; i += 2) { - auto pair = Tx2{w_tile[i], w_tile[i + 1]}; - abs_max_x2(amax_2x, amax_2x, pair); - } - - scale = static_cast( - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y)))); - - scale /= bits == 4 ? 6.0f : 448.0f; - // Convert to mx scale or nv scale - using ScaleType = - std::conditional_t; - auto s = ScaleType(scale); - scale = float(s); - AlignedVector w_hat; - -#pragma unroll - for (int i = 0; i < group_size / 4; i++) { - Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); - float4 dq; - if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); - dq = dequant_fp8(quantized_val); - } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, 1.0f / scale, rbits); - dq = dequant_fp4(quantized_val); - } - w_hat[i * 4] = static_cast(dq.x * scale); - w_hat[i * 4 + 1] = static_cast(dq.y * scale); - w_hat[i * 4 + 2] = static_cast(dq.z * scale); - w_hat[i * 4 + 3] = static_cast(dq.w * scale); - } - store_vector(out, thread_idx, w_hat); +inline void create_2D_tensor_map( + CUtensorMap* tensorMap, + void* input_ptr, + CUtensorMapDataType dtype, + uint64_t rows, + uint64_t cols, + uint32_t tile_y, + uint32_t tile_x, + uint64_t stride_bytes, + CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE) { + constexpr uint32_t rank = 2; // 2D + uint64_t global_dim[rank] = {cols, rows}; + // For row-major layout + uint64_t strides[rank - 1] = {stride_bytes}; + uint32_t tile_dim[rank] = {tile_x, tile_y}; + uint32_t elem_stride[rank] = {1, 1}; + + CHECK_CUDA_ERROR(cuTensorMapEncodeTiled( + tensorMap, + dtype, + rank, + input_ptr, + global_dim, + strides, + tile_dim, + elem_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, + swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE)); } -template -__global__ void -fp_quantize_rowwise(T* w, uint8_t* out, uint8_t* scales, size_t size) { - using Tx2 = Vector2_t; - using Tx4 = Vector4_t; - uint32_t rbits = 0; // reserved bits for future use - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - size_t thread_idx = tidx + grid_dim_x * size_t(tidy); - size_t base_idx = thread_idx * group_size; - - if (base_idx >= size) { - return; - } - - auto w_tile = load_vector(w, thread_idx); - float scale = 0.0f; - - Tx2 amax_2x = Tx2{0.0f, 0.0f}; +inline std::tuple get_columnwise_quantize_mxfp8_launch_args( + size_t grid_dim_x_size, // rows + size_t grid_dim_y_size, // cols + size_t block_size_x, // ROWS_PER_BLOCK + size_t block_size_y, // COL_PER_BLOCK + int threads_per_block) { + dim3 grid; + grid.x = cuda::ceil_div(grid_dim_x_size, block_size_x); + grid.y = cuda::ceil_div(grid_dim_y_size, block_size_y); + grid.z = 1; -#pragma unroll - for (int i = 0; i < group_size; i += 2) { - auto pair = Tx2{w_tile[i], w_tile[i + 1]}; - abs_max_x2(amax_2x, amax_2x, pair); - } + dim3 block(threads_per_block, 1, 1); + return std::make_tuple(grid, block); +} - scale = static_cast( - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y)))); - - scale /= bits == 4 ? 6.0f : 448.0f; - // Convert to mx scale or nv scale - using ScaleType = - std::conditional_t; - auto s = ScaleType(scale); - uint8_t q_scale = s.__x; - scale = float(s); - - scales[thread_idx] = q_scale; - constexpr int elem_per_byte = bits == 8 ? 1 : 2; - AlignedVector quantized; - -#pragma unroll - for (int i = 0; i < group_size / 4; i++) { - Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); - if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); - *reinterpret_cast(&quantized[i * 4]) = quantized_val; - } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, 1.0f / scale, rbits); - *reinterpret_cast(&quantized[i * 2]) = quantized_val; - } +inline CUtensorMapDataType get_tma_dtype(Dtype dtype) { + switch (dtype) { + case float16: + return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + case bfloat16: + return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + case float32: + return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + default: + throw std::runtime_error( + "[fp_quantize_columnwise_tma] Unsupported dtype for TMA"); } - store_vector(out, thread_idx, quantized); } -template -__global__ void fp_quantize_columnwise( - T* w, - uint8_t* out, - uint8_t* scales, +inline std::tuple get_columnwise_quantize_fallback_launch_args( size_t size, + int group_size, int M, int K) { - // Input: [M, K] with strides [1, M] (M-major) - // Quantized output: [M, K/elem_per_byte] row-major (K-major) - // Scales: [M, K/group_size] row-major (K-major) - // Quantize along K (last dimension, groups of group_size elements) - - using Tx2 = Vector2_t; - using Tx4 = Vector4_t; - uint32_t rbits = 0; - - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - - constexpr int BLOCK_X = 32; - constexpr int BLOCK_Y = 32; - constexpr int elem_per_byte = (bits == 8) ? 1 : 2; - constexpr int bytes_per_group = group_size / elem_per_byte; - - constexpr int rows_per_block = BLOCK_X; - constexpr int cols_per_block = BLOCK_Y * group_size; - constexpr int local_cols = cols_per_block / elem_per_byte; - constexpr int bytes_per_block = rows_per_block * local_cols; - - constexpr int SMEM_PAD = 4; - constexpr int padded_local_cols = local_cols + SMEM_PAD; - - auto tidx = idx_in_block.x; - auto tidy = idx_in_block.y; - - int num_col_blocks = (K + cols_per_block - 1) / cols_per_block; - auto bidx = block_idx.x % num_col_blocks; - auto bidy = block_idx.x / num_col_blocks; - - T thread_data[group_size]; - - __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols]; - __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD]; - - int row_base = bidy * rows_per_block + tidx; - int col_base = bidx * cols_per_block + tidy * group_size; - - bool valid = (row_base < M) && (col_base + group_size <= K); - if (valid) { -#pragma unroll - for (int i = 0; i < group_size; i++) { - auto index = row_base + (col_base + i) * M; - thread_data[i] = w[index]; - } - - // Compute scale - Tx2 amax_2x = Tx2{0.0f, 0.0f}; -#pragma unroll - for (int r = 0; r < group_size; r += 2) { - auto pair = Tx2{thread_data[r], thread_data[r + 1]}; - abs_max_x2(amax_2x, amax_2x, pair); - } - float scale = - max(fabsf(static_cast(amax_2x.x)), - fabsf(static_cast(amax_2x.y))); - scale /= (bits == 4) ? 6.0f : 448.0f; - using ScaleType = - std::conditional_t; - auto s = ScaleType(scale); - scale = float(s); - scales_smem[tidx][tidy] = s.__x; - - int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group; - -#pragma unroll - for (int j = 0; j < group_size / 4; j++) { - Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); - if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); - *reinterpret_cast(&quantized_smem[shared_idx + j * 4]) = - quantized_val; - } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, 1.0f / scale, rbits); - *reinterpret_cast(&quantized_smem[shared_idx + j * 2]) = - quantized_val; - } - } - } - __syncthreads(); - - int output_cols = K / elem_per_byte; - int num_groups_per_row = K / group_size; - int linear_tid = tidx + tidy * BLOCK_X; - // Write back quantized values -#pragma unroll - for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) { - int local_row = i / local_cols; - int local_col = i % local_cols; - - int global_row = bidy * rows_per_block + local_row; - int global_col = bidx * local_cols + local_col; - - if (global_row < M && global_col < output_cols) { - int physical_idx = local_row * padded_local_cols + local_col; - out[global_row * output_cols + global_col] = quantized_smem[physical_idx]; - } - } - // Write back scales - constexpr int num_scales = BLOCK_X * BLOCK_Y; -#pragma unroll - for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) { - int local_row = i / BLOCK_Y; - int local_col = i % BLOCK_Y; - - int global_row = bidy * BLOCK_X + local_row; - int global_col = bidx * BLOCK_Y + local_col; - - if (global_row < M && global_col < num_groups_per_row) { - scales[global_row * num_groups_per_row + global_col] = - scales_smem[local_row][local_col]; - } - } -} - -template -__global__ void -fp_dequantize(const uint8_t* w, const uint8_t* scales, T* out, size_t size) { - auto block_size = cg::this_thread_block().dim_threads(); - auto block_idx = cg::this_thread_block().group_index(); - auto idx_in_block = cg::this_thread_block().thread_index(); - - auto tidx = block_idx.x * block_size.x + idx_in_block.x; - auto tidy = block_idx.y * block_size.y + idx_in_block.y; - - auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; - - constexpr int pack_factor = bits == 8 ? 1 : 2; - size_t offset = tidx + grid_dim_x * size_t(tidy); - size_t oindex = offset * pack_factor; - - if (oindex >= size) { - return; - } - - size_t gindex = oindex / group_size; - using ScaleType = - std::conditional_t; - auto scale = float(((ScaleType*)(scales))[gindex]); - - out += oindex; - - uint32_t val = w[offset]; -#pragma clang loop unroll(full) - for (int i = 0; i < pack_factor; i++) { - uint8_t d; - if (bits == 4) { - d = (val >> (bits * i)) & 0x0f; - } else if (bits == 8) { - d = val; - } - out[i] = static_cast(scale * Dequantize{}(d)); - } -} - -inline std::tuple -get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) { - constexpr int BLOCK_X = 32; + constexpr int BLOCK_X = 16; constexpr int BLOCK_Y = 32; int rows_per_block = BLOCK_X; int cols_per_block = BLOCK_Y * group_size; @@ -346,9 +105,13 @@ void fp_quantize_dequantize( array& what, int group_size, int bits, + const std::optional& global_scale /* = std::nullopt */, cu::CommandEncoder& enc, const Stream& s) { enc.set_input_array(w); + if (global_scale.has_value()) { + enc.set_input_array(global_scale.value()); + } enc.set_output_array(what); dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) { using T = cuda_type_t; @@ -367,83 +130,229 @@ void fp_quantize_dequantize( kernel, num_blocks, block_dims, - 0, gpu_ptr(w), gpu_ptr(what), - w.size()); + w.size(), + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); } }); } -void fp_quantize( +void fp_quantize_rowwise( const array& w, array& wq, array& scales, int group_size, int bits, + const std::optional& global_scale /* = std::nullopt */, cu::CommandEncoder& enc, const Stream& s) { enc.set_input_array(w); + if (global_scale.has_value()) { + enc.set_input_array(global_scale.value()); + } enc.set_output_array(wq); enc.set_output_array(scales); - if (w.strides().back() != 1) { - dispatch_float_types(w.dtype(), "fp_quantize_columnwise", [&](auto type_tag) { - using T = cuda_type_t; - if constexpr (!std::is_same_v) { - auto M = w.shape(-2); - auto K = w.shape(-1); - auto kernel = cu::fp_quantize_columnwise; - if (bits == 8) { - kernel = cu::fp_quantize_columnwise; - } else if (group_size == 16) { - kernel = cu::fp_quantize_columnwise; - } - auto [num_blocks, block_dims] = - cu::get_columnwise_quantize_launch_args(w.size(), group_size, M, K); - enc.add_kernel_node( - kernel, - num_blocks, - block_dims, - 0, - gpu_ptr(w), - gpu_ptr(wq), - gpu_ptr(scales), - w.size(), - M, - K); - } else { - throw std::runtime_error( - "[Quantize::eval_gpu] Can not quantize input with type float64."); + dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + auto kernel = cu::fp_quantize_rowwise; + if (bits == 8) { + kernel = cu::fp_quantize_rowwise; + } else if (group_size == 16) { + kernel = cu::fp_quantize_rowwise; } - }); - } else { - dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) { - using T = cuda_type_t; - if constexpr (!std::is_same_v) { - auto kernel = cu::fp_quantize_rowwise; - if (bits == 8) { - kernel = cu::fp_quantize_rowwise; - } else if (group_size == 16) { - kernel = cu::fp_quantize_rowwise; + bool large = w.size() > UINT_MAX; + auto [num_blocks, block_dims] = + get_launch_args(w.size(), w.shape(), w.strides(), large, group_size); + + enc.add_kernel_node( + kernel, + num_blocks, + block_dims, + gpu_ptr(w), + gpu_ptr(wq), + gpu_ptr(scales), + w.size(), + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); + } else { + throw std::runtime_error( + "[Quantize::eval_gpu] Can not quantize input with type float64."); + } + }); +} + +void fp_quantize_columnwise_fallback( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + enc.set_input_array(w); + if (global_scale.has_value()) { + enc.set_input_array(global_scale.value()); + } + enc.set_output_array(wq); + enc.set_output_array(scales); + dispatch_float_types( + w.dtype(), "fp_quantize_columnwise_fallback", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + auto M = w.shape(-2); + auto K = w.shape(-1); + auto kernel = + cu::fp_quantize_columnwise_fallback; + if (bits == 8) { + kernel = cu::fp_quantize_columnwise_fallback; + } else if (group_size == 16) { + kernel = + cu::fp_quantize_columnwise_fallback; + } + auto [num_blocks, block_dims] = + cu::get_columnwise_quantize_fallback_launch_args( + w.size(), group_size, M, K); + enc.add_kernel_node( + kernel, + num_blocks, + block_dims, + gpu_ptr(w), + gpu_ptr(wq), + gpu_ptr(scales), + w.size(), + M, + K, + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); + } else { + throw std::runtime_error( + "[Quantize::eval_gpu] Can not quantize input with type float64."); } - bool large = w.size() > UINT_MAX; - auto [num_blocks, block_dims] = get_launch_args( - w.size(), w.shape(), w.strides(), large, group_size); - - enc.add_kernel_node( - kernel, - num_blocks, - block_dims, - 0, - gpu_ptr(w), - gpu_ptr(wq), - gpu_ptr(scales), - w.size()); - } else { - throw std::runtime_error( - "[Quantize::eval_gpu] Can not quantize input with type float64."); - } - }); + }); +} + +void fp_quantize_columnwise_mxfp8( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + enc.set_input_array(w); + enc.set_output_array(wq); + enc.set_output_array(scales); + + size_t rows = w.shape(-1); + size_t cols = w.size() / rows; + size_t stride_bytes = w.strides(-1) * w.itemsize(); + + dispatch_float_types( + w.dtype(), "fp_quantize_columnwise_mxfp8", [&](auto type_tag) { + using T = cuda_type_t; + if constexpr (!std::is_same_v) { + constexpr int THREADS_PER_BLOCK = 64; + constexpr int ROWS_PER_BLOCK = 64; + constexpr int COLS_PER_BLOCK = 64; + constexpr size_t TILE_M = 32; + constexpr size_t TILE_K = COLS_PER_BLOCK; + constexpr size_t STAGES = ROWS_PER_BLOCK / TILE_M; + + // For columnwise: grid.x = cols, grid.y = rows + // scales_per_stage = TILE_K (one scale per column per stage) + auto [grid, block] = cu::get_columnwise_quantize_mxfp8_launch_args( + cols, rows, COLS_PER_BLOCK, ROWS_PER_BLOCK, THREADS_PER_BLOCK); + + CUtensorMap tensor_map_input; + CUtensorMap tensor_map_output; + + cu::create_2D_tensor_map( + &tensor_map_input, + gpu_ptr(w), + cu::get_tma_dtype(w.dtype()), + rows, + cols, + TILE_M, + TILE_K, + stride_bytes); + + cu::create_2D_tensor_map( + &tensor_map_output, + gpu_ptr(wq), + CU_TENSOR_MAP_DATA_TYPE_UINT8, + cols, + rows, + TILE_K, + TILE_M, + rows); + + auto kernel = cu::fp_quantize_columnwise_mxfp8< + T, + false, + THREADS_PER_BLOCK, + COLS_PER_BLOCK, + ROWS_PER_BLOCK>; + enc.add_kernel_node( + kernel, + grid, + block, + tensor_map_input, + tensor_map_output, + gpu_ptr(scales), + rows, + cols); + } else { + throw std::runtime_error( + "[fp_quantize_columnwise_tma] Cannot quantize input with type float64."); + } + }); +} + +void fp_quantize_columnwise( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + // Use TMA version for SM100+ with MXFP8 (bits=8, group_size=32) + // NVFP4 todo + const size_t rows = w.shape(-1); + const size_t cols = w.size() / rows; + const bool has_full_tma_tiles = ((rows % 128) == 0) && ((cols % 128) == 0); + bool use_tma = + (enc.device().compute_capability_major() >= 10 && bits == 8 && + group_size == 32 && has_full_tma_tiles); + if (use_tma) { + fp_quantize_columnwise_mxfp8( + w, wq, scales, group_size, bits, global_scale, enc, s); + } else { + fp_quantize_columnwise_fallback( + w, wq, scales, group_size, bits, global_scale, enc, s); + } +} + +void fp_quantize( + const array& w, + array& wq, + array& scales, + int group_size, + int bits, + const std::optional& global_scale /* = std::nullopt */, + cu::CommandEncoder& enc, + const Stream& s) { + if (w.strides(-1) == 1) { + fp_quantize_rowwise(w, wq, scales, group_size, bits, global_scale, enc, s); + } else { + fp_quantize_columnwise( + w, wq, scales, group_size, bits, global_scale, enc, s); } } @@ -453,6 +362,7 @@ void fp_dequantize( array& w, int group_size, int bits, + const std::optional& global_scale /* = std::nullopt */, cu::CommandEncoder& enc, const Stream& s) { constexpr int uint8_per_uint32 = 4; @@ -465,6 +375,9 @@ void fp_dequantize( enc.set_input_array(wq); enc.set_input_array(scales); + if (global_scale.has_value()) { + enc.set_input_array(global_scale.value()); + } enc.set_output_array(w); dispatch_float_types(w.dtype(), "fp_dequantize", [&](auto type_tag) { using T = cuda_type_t; @@ -481,11 +394,12 @@ void fp_dequantize( kernel, num_blocks, block_dims, - 0, gpu_ptr(wq), gpu_ptr(scales), gpu_ptr(w), - w.size()); + w.size(), + global_scale.has_value() ? gpu_ptr(global_scale.value()) + : nullptr); } else { throw std::runtime_error( "[Quantize::eval_gpu] Can not dequantize to output with type float64."); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cuh new file mode 100644 index 000000000..bea4e2197 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/fp_quantize.cuh @@ -0,0 +1,601 @@ +// Copyright © 2025 Apple Inc. +#pragma once + +#include "mlx/backend/cuda/ptx.cuh" +#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh" +#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh" +#include "mlx/backend/cuda/vector_types.cuh" +#include "mlx/dtype_utils.h" + +#include +#include +#include +#include + +constexpr float F8E4M3_MAX = 448.0f; +constexpr float F4E2M1_MAX = 6.0f; +constexpr size_t TMA_SHMEM_ALIGNMENT = 128; +constexpr size_t BUFFS_NUM = 2; + +namespace mlx::core { +namespace cu { + +template +struct Dequantize { + __device__ float operator()(uint8_t x) { + if constexpr (bits == 8) { + return float(*(cutlass::float_e4m3_t*)(&x)); + } else { + return float(*(cutlass::float_e2m1_t*)(&x)); + } + } +}; + +template +__device__ __forceinline__ void absmax_x2(T& out, const T& x1, const T& x2) { + if constexpr ( + (std::is_same::value) || + (std::is_same::value)) { + T a = x1; + T b = x2; + out = __hmax2(__habs2(a), __habs2(b)); + } else if constexpr (std::is_same::value) { + float2 a = x1; + float2 b = x2; + out.x = fmaxf(fabsf(a.x), fabsf(b.x)); + out.y = fmaxf(fabsf(a.y), fabsf(b.y)); + } +} + +__device__ __forceinline__ void copy_2d_to_shared( + void* dst, + const CUtensorMap* tensor_map, + uint32_t tile_x, + uint32_t tile_y, + uint32_t num_bytes, + uint64_t* barrier, + const bool is_master_thread) { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + if (is_master_thread) { + // Arrive and tell how many bytes are expected + ptx::mbarrier_arrive_expect_tx(barrier, num_bytes); + // Initiate bulk tensor copy + ptx::cp_async_bulk_tensor_2d_global_to_shared( + dst, tensor_map, tile_x, tile_y, barrier); + } else { + // Other threads just arrive + ptx::mbarrier_arrive(barrier); + } +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) +} + +namespace cg = cooperative_groups; + +template +__global__ void fp_quantize_dequantize( + T* w, + T* out, + size_t size, + float* global_scale = nullptr) { + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f; + + using Tx2 = Vector2_t; + uint32_t rbits = 0; // reserved bits for future use + + size_t thread_idx = cg::this_grid().thread_rank(); + size_t base_idx = thread_idx * group_size; + + if (base_idx >= size) { + return; + } + + auto w_tile = load_vector(w, thread_idx); + float scale_dec_b = 0.0f; + + Tx2 amax_2x = Tx2{0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < group_size; i += 2) { + auto pair = Tx2{w_tile[i], w_tile[i + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + scale_dec_b = static_cast( + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y)))); + + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + float scale_enc_b = scale_enc / float(s); + float scale_dec = float(s) * inv_scale_enc; + AlignedVector w_hat; + +#pragma unroll + for (int i = 0; i < group_size / 8; i++) { + auto& w = *reinterpret_cast*>(&w_tile[i * 8]); + cutlass::NumericArrayConverter fp32_t; + auto scaled = fp32_t(w) * scale_enc_b; + cutlass::Array dq; + if constexpr (bits == 8) { + cutlass::NumericArrayConverter fp8_fp32; + auto quant = fp8_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp8; + dq = fp32_fp8(quant); + } else { + cutlass::NumericArrayConverter fp4_fp32; + auto quant = fp4_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp4; + dq = fp32_fp4(quant); + } + cutlass::NumericArrayConverter t_fp32; + *reinterpret_cast*>(&w_hat[i * 8]) = + t_fp32(dq * scale_dec); + } + store_vector(out, thread_idx, w_hat); +} + +template +__global__ void fp_quantize_rowwise( + T* w, + uint8_t* out, + uint8_t* scales, + size_t size, + float* global_scale = nullptr) { + // NVFP4 conversion: + // Global encode scale: (448 × 6) / *global_scale + // Per-block decode scale: S_dec_b = (block_amax / 6) × S_enc → stored as FP8 + // E4M3 Per-block encode scale: S_enc_b = S_enc / S_dec_b + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + uint32_t rbits = 0; // reserved bits for future use + + size_t thread_idx = cg::this_grid().thread_rank(); + size_t base_idx = thread_idx * group_size; + + if (base_idx >= size) { + return; + } + + auto w_tile = load_vector(w, thread_idx); + float scale_dec_b = 0.0f; + + Tx2 amax_2x = Tx2{0.0f, 0.0f}; + +#pragma unroll + for (int i = 0; i < group_size; i += 2) { + auto pair = Tx2{w_tile[i], w_tile[i + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + scale_dec_b = static_cast( + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y)))); + + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + uint8_t q_scale = s.storage; + float scale_enc_b = scale_enc / float(s); + + scales[thread_idx] = q_scale; + constexpr int elem_per_byte = bits == 8 ? 1 : 2; + AlignedVector quantized; + +#pragma unroll + for (int i = 0; i < group_size / 4; i++) { + Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); + if constexpr (bits == 8) { + uint32_t quantized_val = + scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized[i * 4]) = quantized_val; + } else { + uint16_t quantized_val = + scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized[i * 2]) = quantized_val; + } + } + store_vector(out, thread_idx, quantized); +} + +template < + typename T, + bool USE_SR, + int THREADS_PER_BLOCK, + int COLS_PER_BLOCK, + int ROWS_PER_BLOCK> +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + fp_quantize_columnwise_mxfp8( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_output, + uint8_t* __restrict__ scales, + const size_t rows, + const size_t cols) { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + + constexpr size_t TILE_M = 32; + constexpr size_t TILE_K = COLS_PER_BLOCK; + constexpr size_t STEPS = ROWS_PER_BLOCK / TILE_M; + constexpr int elem_per_byte = 1; + + const auto block_idx = cg::this_thread_block().group_index(); + const auto idx_in_block = cg::this_thread_block().thread_index(); + const int tidx = idx_in_block.x; // Thread handles column tidx + const bool is_master = (tidx == 0); + + const size_t block_offset_col = block_idx.x * COLS_PER_BLOCK; + const size_t block_offset_row = block_idx.y * ROWS_PER_BLOCK; + + constexpr size_t BUFF_ELEMS = TILE_M * TILE_K; + constexpr uint32_t tile_bytes = static_cast(BUFF_ELEMS * sizeof(T)); + + // Note: 128 byte alignment is a requirement for shared memory buffers + // if only we use swizzling while tma copy + // if swizzling is not used, pointers must be 16 byte aligned + __shared__ alignas(16) T in_smem[BUFF_ELEMS * BUFFS_NUM]; + __shared__ alignas(16) uint8_t out_smem[BUFF_ELEMS * BUFFS_NUM]; + + __shared__ alignas(8) uint64_t mbar[STEPS]; + + T thread_data[TILE_M]; + uint32_t rbits = 0; // Reserved for stochastic rounding + const size_t scale_stride = rows / TILE_M; + + // Master thread init memory barriers for all steps + // fence for tma, synchronize threads so all see mbarrier + if (is_master) { +#pragma unroll + for (int iter = 0; iter < STEPS; ++iter) { + ptx::mbarrier_init(&mbar[iter], THREADS_PER_BLOCK); + } + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + // Launch first async copy before entering the loop + copy_2d_to_shared( + &in_smem[0], + &tensor_map_input, + static_cast(block_offset_col), + static_cast(block_offset_row), + tile_bytes, + &mbar[0], + is_master); + +#pragma unroll + for (size_t step = 0; step < STEPS; ++step) { + // buffer memory offset in shared memory (we use double buffering for + // pipelining) + const size_t buff = step % BUFFS_NUM; + const size_t next_step = step + 1; + const size_t step_row_offset = step * TILE_M; + + if (next_step < STEPS) { + // before launching another async copy, check that there is less than 2 + // (to ensure that shared -> global synch is finished and buffer can be + // reused) + ptx::cp_async_bulk_wait_group_read<1>(); + const size_t next_buff = next_step % BUFFS_NUM; + const size_t next_row_offset = block_offset_row + next_step * TILE_M; + const size_t next_buff_elem_offset = next_buff * BUFF_ELEMS; + + copy_2d_to_shared( + &in_smem[next_buff_elem_offset], + &tensor_map_input, + static_cast(block_offset_col), + static_cast(next_row_offset), + tile_bytes, + &mbar[next_step], + is_master); + } + + ptx::fence_proxy_async_shared_cta(); + // Wait until the data is ready, parity is always 0 because for simplicity + // we dont reuse barriers between steps + ptx::mbarrier_wait_parity(&mbar[step], 0); + const size_t buff_offset = buff * BUFF_ELEMS; + // Read the data from shared to registers +#pragma unroll + for (int row = 0; row < TILE_M; ++row) { + thread_data[row] = in_smem[buff_offset + row * TILE_K + tidx]; + } + Tx2 amax_2x = Tx2{T(0.0f), T(0.0f)}; +#pragma unroll + for (int row = 0; row < TILE_M; row += 2) { + auto pair = Tx2{thread_data[row], thread_data[row + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + + float scale = + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y))); + + scale /= F8E4M3_MAX; + + auto s = cutlass::float_ue8m0_t(scale); + scale = float(s); + // Write scale directly to global memory + const size_t global_col = block_offset_col + tidx; + const size_t global_row_group = + (block_offset_row + step_row_offset) / TILE_M; + if (global_col < cols && (block_offset_row + step_row_offset) < rows) { + scales[global_col * scale_stride + global_row_group] = s.storage; + } + const size_t out_buff_offset = buff * BUFF_ELEMS / elem_per_byte; + // Quantize to registers first + constexpr int GROUPS = TILE_M / 4; + uint32_t quantized_regs[GROUPS]; +#pragma unroll + for (int j = 0; j < GROUPS; ++j) { + Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); + quantized_regs[j] = + cu::scale_cvt_Tx4_to_fp8x4(w_Tx4, 1.0f / scale, rbits); + } + // Write output to shared memory with swapped store order to reduce bank + // conflicts. Without swap: stride between threads is TILE_M=32 bytes + // = 8 banks, every 4th thread hits the same bank -> 8-way conflict. + const int lane = tidx % 32; + const int group = (lane / 4) % 2; + const size_t base = out_buff_offset + tidx * TILE_M; + switch (group) { + case 0: + *reinterpret_cast(&out_smem[base + 0]) = { + quantized_regs[0], + quantized_regs[1], + quantized_regs[2], + quantized_regs[3]}; + *reinterpret_cast(&out_smem[base + 16]) = { + quantized_regs[4], + quantized_regs[5], + quantized_regs[6], + quantized_regs[7]}; + break; + case 1: + *reinterpret_cast(&out_smem[base + 16]) = { + quantized_regs[4], + quantized_regs[5], + quantized_regs[6], + quantized_regs[7]}; + *reinterpret_cast(&out_smem[base + 0]) = { + quantized_regs[0], + quantized_regs[1], + quantized_regs[2], + quantized_regs[3]}; + break; + } + __syncthreads(); + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + if (is_master) { + const size_t global_row = block_offset_row + step_row_offset; + const uint32_t out_x = static_cast(global_row); + const uint32_t out_y = static_cast(block_offset_col); + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), + out_x, + out_y, + reinterpret_cast(&out_smem[out_buff_offset])); + ptx::cp_async_bulk_commit_group(); + } + } + // Wait for all TMA stores to complete + ptx::cp_async_bulk_wait_group_read<0>(); + + __syncthreads(); + if (is_master) { +#pragma unroll + for (int iter = 0; iter < STEPS; ++iter) { + ptx::mbarrier_invalidate(&mbar[iter]); + } + } +#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED) +} + +// TODO: add kernel with tma instructions +template +__global__ void fp_quantize_columnwise_fallback( + T* w, + uint8_t* out, + uint8_t* scales, + size_t size, + int M, + int K, + float* global_scale = nullptr) { + // Input: [M, K] with strides [1, M] (M-major) + // Quantized output: [M, K/elem_per_byte] row-major (K-major) + // Scales: [M, K/group_size] row-major (K-major) + // Quantize along K (last dimension, groups of group_size elements) + const bool use_global_scale = global_scale != nullptr; + const float scale_enc = + use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f; + + using Tx2 = Vector2_t; + using Tx4 = Vector4_t; + uint32_t rbits = 0; + + auto block_idx = cg::this_thread_block().group_index(); + auto idx_in_block = cg::this_thread_block().thread_index(); + + constexpr int BLOCK_X = 16; + constexpr int BLOCK_Y = 32; + constexpr int elem_per_byte = (bits == 8) ? 1 : 2; + constexpr int bytes_per_group = group_size / elem_per_byte; + + constexpr int rows_per_block = BLOCK_X; + constexpr int cols_per_block = BLOCK_Y * group_size; + constexpr int local_cols = cols_per_block / elem_per_byte; + constexpr int bytes_per_block = rows_per_block * local_cols; + + constexpr int SMEM_PAD = 4; + constexpr int padded_local_cols = local_cols + SMEM_PAD; + + auto tidx = idx_in_block.x; + auto tidy = idx_in_block.y; + + int num_col_blocks = (K + cols_per_block - 1) / cols_per_block; + auto bidx = block_idx.x % num_col_blocks; + auto bidy = block_idx.x / num_col_blocks; + + T thread_data[group_size]; + + __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols]; + __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD]; + + int row_base = bidy * rows_per_block + tidx; + int col_base = bidx * cols_per_block + tidy * group_size; + + bool valid = (row_base < M) && (col_base + group_size <= K); + if (valid) { +#pragma unroll + for (int i = 0; i < group_size; i++) { + auto index = row_base + (col_base + i) * M; + thread_data[i] = w[index]; + } + + // Compute scale + Tx2 amax_2x = Tx2{0.0f, 0.0f}; +#pragma unroll + for (int r = 0; r < group_size; r += 2) { + auto pair = Tx2{thread_data[r], thread_data[r + 1]}; + absmax_x2(amax_2x, amax_2x, pair); + } + float scale_dec_b = + max(fabsf(static_cast(amax_2x.x)), + fabsf(static_cast(amax_2x.y))); + scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; + scale_dec_b *= scale_enc; + // Convert to mx scale or nv scale + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto s = ScaleType(scale_dec_b); + float scale_enc_b = scale_enc / float(s); + scales_smem[tidx][tidy] = s.storage; + + int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group; + +#pragma unroll + for (int j = 0; j < group_size / 4; j++) { + Tx4 w_Tx4 = *reinterpret_cast(&thread_data[j * 4]); + if constexpr (bits == 8) { + uint32_t quantized_val = + scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized_smem[shared_idx + j * 4]) = + quantized_val; + } else { + uint16_t quantized_val = + scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); + *reinterpret_cast(&quantized_smem[shared_idx + j * 2]) = + quantized_val; + } + } + } + __syncthreads(); + + int output_cols = K / elem_per_byte; + int num_groups_per_row = K / group_size; + int linear_tid = tidx + tidy * BLOCK_X; + // Write back quantized values +#pragma unroll + for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) { + int local_row = i / local_cols; + int local_col = i % local_cols; + + int global_row = bidy * rows_per_block + local_row; + int global_col = bidx * local_cols + local_col; + + if (global_row < M && global_col < output_cols) { + int physical_idx = local_row * padded_local_cols + local_col; + out[global_row * output_cols + global_col] = quantized_smem[physical_idx]; + } + } + // Write back scales + constexpr int num_scales = BLOCK_X * BLOCK_Y; +#pragma unroll + for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) { + int local_row = i / BLOCK_Y; + int local_col = i % BLOCK_Y; + + int global_row = bidy * BLOCK_X + local_row; + int global_col = bidx * BLOCK_Y + local_col; + + if (global_row < M && global_col < num_groups_per_row) { + scales[global_row * num_groups_per_row + global_col] = + scales_smem[local_row][local_col]; + } + } +} + +template +__global__ void fp_dequantize( + const uint8_t* w, + const uint8_t* scales, + T* out, + size_t size, + float* global_scale = nullptr) { + auto block_size = cg::this_thread_block().dim_threads(); + auto block_idx = cg::this_thread_block().group_index(); + auto idx_in_block = cg::this_thread_block().thread_index(); + + auto tidx = block_idx.x * block_size.x + idx_in_block.x; + auto tidy = block_idx.y * block_size.y + idx_in_block.y; + + auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x; + + constexpr int pack_factor = bits == 8 ? 1 : 2; + const bool use_global_scale = global_scale != nullptr; + const float inv_scale_enc = use_mx_scale + ? 1.0f + : (use_global_scale ? (*global_scale) / (F8E4M3_MAX * F4E2M1_MAX) : 1.0f); + size_t offset = tidx + grid_dim_x * size_t(tidy); + size_t oindex = offset * pack_factor; + + if (oindex >= size) { + return; + } + + size_t gindex = oindex / group_size; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; + auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc; + + out += oindex; + + uint32_t val = w[offset]; +#pragma clang loop unroll(full) + for (int i = 0; i < pack_factor; i++) { + uint8_t d; + if (bits == 4) { + d = (val >> (bits * i)) & 0x0f; + } else if (bits == 8) { + d = val; + } + out[i] = static_cast(scale * Dequantize{}(d)); + } +} + +} // namespace cu +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh index 1edf36417..ee8494e24 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/mxfp8_quantize.cuh @@ -1,32 +1,23 @@ #pragma once -#include -#include -#include #include "mlx/backend/cuda/vector_types.cuh" -namespace mlx::core::cu { +#include -// TODO implement fast path -template -__device__ __forceinline__ uint32_t -scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t input, const float scale) { - uint32_t out_fp8x4 = 0; - float4 scaled; - scaled.x = static_cast(input.x) * scale; - scaled.y = static_cast(input.y) * scale; - scaled.z = static_cast(input.z) * scale; - scaled.w = static_cast(input.w) * scale; - out_fp8x4 = __nv_fp8x4_e4m3(scaled).__x; - return out_fp8x4; -} +namespace mlx::core::cu { // Place holder for future fast path implementation template __device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4( - const Vector4_t input, + const Vector4_t& input, const float scale, uint32_t rbits) { - return scale_cvt_Tx4_to_fp8x4_fallback(input, scale); + cutlass::NumericArrayConverter fp32_t; + auto scaled = + fp32_t(*reinterpret_cast*>(&input)) * scale; + cutlass::NumericArrayConverter fp8_fp32; + auto quant = fp8_fp32(scaled); + return *reinterpret_cast(&quant); } -} // namespace mlx::core::cu \ No newline at end of file + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp index 5552e2405..375755bc3 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/no_qqmm_impl.cpp @@ -17,9 +17,8 @@ void qqmm_impl( const array&, const array&, const array&, - Dtype, QuantizationMode, - float) { + const GemmScalars&) { throw std::runtime_error( "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher."); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh index da0df4293..7a5616c76 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/nvfp4_quantize.cuh @@ -1,10 +1,9 @@ #pragma once -#include -#include -#include #include "mlx/backend/cuda/vector_types.cuh" +#include + namespace mlx::core::cu { using bf16x4 = Vector4_t<__nv_bfloat16>; @@ -13,23 +12,15 @@ using f32x4 = Vector4_t; template __device__ __forceinline__ uint16_t -scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t input, const float scale) { +scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t& input, const float scale) { // Fallback implementation for architectures that do not support cvt // instructions or for cuda versions with no fp4 support (< 12.8) -> scalar - uint16_t out_fp4x4 = 0; - fp32x4 scaled; - scaled.x = static_cast(input.x) * scale; - scaled.y = static_cast(input.y) * scale; - scaled.z = static_cast(input.z) * scale; - scaled.w = static_cast(input.w) * scale; - uint8_t q0 = __nv_fp4_e2m1(scaled.x).__x; - uint8_t q1 = __nv_fp4_e2m1(scaled.y).__x; - uint8_t q2 = __nv_fp4_e2m1(scaled.z).__x; - uint8_t q3 = __nv_fp4_e2m1(scaled.w).__x; - out_fp4x4 = (static_cast(q3) << 12) | - (static_cast(q2) << 8) | (static_cast(q1) << 4) | - static_cast(q0); - return out_fp4x4; + cutlass::NumericArrayConverter fp32_t; + auto scaled = + fp32_t(*reinterpret_cast*>(&input)) * scale; + cutlass::NumericArrayConverter fp4_fp32; + auto quant = fp4_fp32(scaled); + return *reinterpret_cast(&quant); } #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ @@ -318,7 +309,7 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast( template __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4( - const Vector4_t input, + const Vector4_t& input, const float scale, uint32_t rbits) { #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ @@ -331,4 +322,5 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4( return scale_cvt_Tx4_to_fp4x4_fallback(input, scale); #endif } -} // namespace mlx::core::cu \ No newline at end of file + +} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/CMakeLists.txt deleted file mode 100644 index 3b88403e8..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/CMakeLists.txt +++ /dev/null @@ -1,19 +0,0 @@ -target_sources( - mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/qmm.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmv.cu - ${CMAKE_CURRENT_SOURCE_DIR}/fp_qmv.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_k.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m16_n.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_k.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m32_n.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_k.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_naive_m64_n.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m16.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m32.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm80_m64.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n16_m1.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n32_m1.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n64_m2.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n128_m2.cu - ${CMAKE_CURRENT_SOURCE_DIR}/qmm_impl_sm90_m128_n256_m2.cu) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.cu index e6f7fc321..5c1d4f76b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.cu @@ -2,25 +2,20 @@ #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/backend/cuda/quantized/qmm/qmm_utils.h" #include namespace mlx::core { -#if defined(MLX_CUDA_SM90A_ENABLED) -// Defined in qmm_impl_sm90_xxx.cu files. -template -void qmm_impl_sm90( - const array& x, - const array& w, - const array& scales, - const array& biases, - array& out, - int bits, - int group_size, - cu::CommandEncoder& encoder, - Stream s); -#endif // defined(MLX_CUDA_SM90A_ENABLED) +namespace { + +inline bool is_last_2_dims_row_contiguous(const array& x) { + return x.flags().contiguous && (x.ndim() >= 2) && (x.strides(-1) == 1) && + (x.strides(-2) == x.shape(-1)); +} + +} // namespace bool supports_qmm_sm90( const array& x, @@ -36,21 +31,25 @@ bool supports_qmm_sm90( if (device.compute_capability_major() != 9) { return false; } - int k = x.shape(-1); + auto [m, n, k, l, broadcast_b] = make_problem_shape(x, w, out); + if ((n * w.itemsize()) % 16 != 0) { // TMA alignment + return false; + } if (k % 64 != 0) { return false; } if (!biases) { return false; } - if (!x.flags().row_contiguous || !w.flags().row_contiguous || - !scales.flags().row_contiguous || !biases->flags().row_contiguous) { + if (!is_last_2_dims_row_contiguous(w) || + !is_last_2_dims_row_contiguous(scales) || + !is_last_2_dims_row_contiguous(*biases)) { return false; } if (!transpose) { return false; } - if (bits % 2 != 0) { + if (bits != 4 && bits != 8) { return false; } if (group_size < k) { @@ -62,55 +61,6 @@ bool supports_qmm_sm90( return true; } -void qmm_sm90( - const array& x, - const array& w, - const array& scales, - const array& biases, - array& out, - int bits, - int group_size, - cu::CommandEncoder& encoder, - Stream s) { -#if defined(MLX_CUDA_SM90A_ENABLED) - auto dispatch = [&]() { - using cute::Int; - using TileShapeMN = cute::Shape, Int>; - using ClusterShape = cute::Shape, Int<1>, Int<1>>; - qmm_impl_sm90( - x, w, scales, biases, out, bits, group_size, encoder, s); - }; - int m = out.ndim() > 1 ? out.shape(-2) : 1; - if (m <= 16) { - dispatch.template operator()<128, 16, 1>(); - } else if (m <= 32) { - dispatch.template operator()<128, 32, 1>(); - } else if (m <= 64) { - dispatch.template operator()<128, 64, 2>(); - } else if (m <= 128) { - dispatch.template operator()<128, 128, 2>(); - } else { - dispatch.template operator()<128, 256, 2>(); - } -#else - throw std::runtime_error( - "[quantized_matmul] Hopper-only kernel is not available."); -#endif // defined(MLX_CUDA_SM90A_ENABLED) -} - -// Defined in qmm_impl_sm80_xxx.cu files. -template -void qmm_impl_sm80( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder); - bool supports_qmm_sm80( const array& x, const array& w, @@ -130,11 +80,11 @@ bool supports_qmm_sm80( if ((n % 128 != 0) || (k % std::max(64, group_size) != 0)) { return false; } - if (!x.flags().row_contiguous || !w.flags().row_contiguous || - !scales.flags().row_contiguous) { + if (!is_last_2_dims_row_contiguous(w) || + !is_last_2_dims_row_contiguous(scales)) { return false; } - if (biases && !biases->flags().row_contiguous) { + if (biases && !is_last_2_dims_row_contiguous(*biases)) { return false; } if (x.dtype() != float16 && x.dtype() != bfloat16) { @@ -149,43 +99,6 @@ bool supports_qmm_sm80( return true; } -void qmm_sm80( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder) { - auto dispatch = [&]() { - qmm_impl_sm80( - x, w, scales, biases, out, bits, group_size, mode, encoder); - }; - int m = out.ndim() > 1 ? out.shape(-2) : 1; - if (m <= 16) { - dispatch.template operator()<16>(); - } else if (m <= 32) { - dispatch.template operator()<32>(); - } else { - dispatch.template operator()<64>(); - } -} - -// Defined in qmm_impl_naive_xxx.cu files. -template -void qmm_impl_naive( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder); - bool supports_qmm_naive( const array& x, const array& w, @@ -198,46 +111,19 @@ bool supports_qmm_naive( QuantizationMode mode, cu::Device& device) { int k = x.shape(-1); - if (k % std::max(64, group_size) != 0) { + if (transpose && (k % std::max(64, group_size) != 0)) { return false; } - if (!x.flags().row_contiguous || !w.flags().row_contiguous || - !scales.flags().row_contiguous) { + if (!is_last_2_dims_row_contiguous(w) || + !is_last_2_dims_row_contiguous(scales)) { return false; } - if (biases && !biases->flags().row_contiguous) { + if (biases && !is_last_2_dims_row_contiguous(*biases)) { return false; } return true; } -void qmm_naive( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - bool transpose, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder) { - auto dispatch = [&]() { - qmm_impl_naive( - x, w, scales, biases, out, bits, group_size, mode, encoder); - }; - dispatch_bool(transpose, [&](auto k_major) { - int m = out.ndim() > 1 ? out.shape(-2) : 1; - if (m <= 16) { - dispatch.template operator()<16, k_major.value>(); - } else if (m <= 32) { - dispatch.template operator()<32, k_major.value>(); - } else { - dispatch.template operator()<64, k_major.value>(); - } - }); -} - bool supports_fp_qmv( const array& x, const array& w, @@ -285,11 +171,11 @@ bool supports_qmv( if (k % 8 != 0) { return false; } - if (!x.flags().row_contiguous || !w.flags().row_contiguous || - !scales.flags().row_contiguous) { + if (!is_last_2_dims_row_contiguous(w) || + !is_last_2_dims_row_contiguous(scales)) { return false; } - if (biases && !biases->flags().row_contiguous) { + if (biases && !is_last_2_dims_row_contiguous(*biases)) { return false; } if (!transpose) { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h index a7464dee6..64cfecfd5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm.h @@ -49,6 +49,8 @@ void qmm_sm80( const array& w, const array& scales, const std::optional& biases, + const std::optional& lhs_indices, + const std::optional& rhs_indices, array& out, int bits, int group_size, @@ -72,6 +74,9 @@ void qmm_naive( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, + const std::optional& lhs_indices, + const std::optional& rhs_indices, array& out, bool transpose, int bits, @@ -118,6 +123,20 @@ void qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder); + +void gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& lhs_indices, + const array& rhs_indices, array& out, int bits, int group_size, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh deleted file mode 100644 index cc387cfc8..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh +++ /dev/null @@ -1,479 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh" -#include "mlx/backend/cuda/quantized/qmm/qmm.h" -#include "mlx/dtype_utils.h" - -// clang-format off - -// We can't put kernel code in mlx::core due to name conflicts of "Shape". -namespace cutlass_gemm { - -using namespace cute; - -template -struct SharedStorage { - ArrayEngine> A; - ArrayEngine> B; -}; - -__device__ __forceinline__ void -cute_naive_dequant(auto w, auto s, auto z, auto out) { - using Element = typename decltype(out)::value_type; - using Quant = typename decltype(w)::value_type; - using Scale = typename decltype(s)::value_type; - transform(w, out, [](Quant q) { return Element(q); } ); - transform(out, s, out, [](Element e, Scale s) { return e * Element(s); }); - if constexpr (quant_has_bias_v) { - transform(out, z, out, plus{}); - } -} - -__device__ __forceinline__ void -cute_dequant(auto w, auto s, auto z, auto out) { - if constexpr (stride(coalesce(w.layout())) == Int<1>{} && - is_static_v) { - cute_vectorized_dequant(w, s, z, out); - } else { - cute_naive_dequant(w, s, z, out); - } -} - -template -__global__ void qmm_naive_kernel( - ProblemShape shape_MNKL, CtaTiler cta_tiler, - const Element* A, StrideA dA, SmemLayoutA sA_layout, TiledCopyA copy_a, - const Quant* B, StrideB dB, SmemLayoutB sB_layout, TiledCopyB copy_b, - Element* C, StrideC dC, - const Scale* S, const Element* Z, LayoutS S_layout, - TiledMma mma) { - CUTE_STATIC_ASSERT_V(size(copy_a) == size(mma)); - CUTE_STATIC_ASSERT_V(size(copy_b) == size(mma)); - CUTE_STATIC_ASSERT_V(congruent(select<0,2,3>(shape_MNKL), dA)); - CUTE_STATIC_ASSERT_V(congruent(select<1,2,3>(shape_MNKL), dB)); - CUTE_STATIC_ASSERT_V(congruent(select<0,1,3>(shape_MNKL), dC)); - - int thread_idx = int(threadIdx.x); - auto [m_coord, n_coord, l_coord] = static_cast(blockIdx); - - // Represent the full tensors. - Tensor mA_mkl = make_tensor(make_gmem_ptr(A), select<0,2,3>(shape_MNKL), dA); // (M,K,L) - Tensor mB_nkl = make_tensor(make_gmem_ptr(B), select<1,2,3>(shape_MNKL), dB); // (N,K,L) - Tensor mC_mnl = make_tensor(make_gmem_ptr(C), select<0,1,3>(shape_MNKL), dC); // (M,N,L) - - Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L) - Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L) - - // Get batch slice. - Tensor mA = mA_mkl(_,_,l_coord); // (M,K) - Tensor mB = mB_nkl(_,_,l_coord); // (N,K) - Tensor mC = mC_mnl(_,_,l_coord); // (M,N) - - Tensor mS = mS_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) - Tensor mZ = mZ_nkl(_,_,l_coord); // (N,(group_size,K/group_size)) - - // Get the appropriate blocks for this thread block. - auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k) - Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) - Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N) - - Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) - - auto m_max_coord = size<0>(shape_MNKL) - size<0>(gA) * m_coord; // M - BLK_M * m_coord - auto n_max_coord = size<1>(shape_MNKL) - size<0>(gB) * n_coord; // N - BLK_N * n_coord - - // Shared memory buffers. - extern __shared__ char shared_memory[]; - using SharedStorage = SharedStorage; - SharedStorage& smem = *reinterpret_cast(shared_memory); - Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout); // (BLK_M,BLK_K) - Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout); // (BLK_N,BLK_K) - - // Partition the copying of A/B/C tiles across the threads. - ThrCopy thr_copy_a = copy_a.get_slice(thread_idx); - Tensor tAgA = thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) - Tensor tAsA = thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K) - Tensor tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K) - - ThrCopy thr_copy_b = copy_b.get_slice(thread_idx); - Tensor tBgB = thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) - Tensor tBsB = thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K) - Tensor tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_M,BCPY_K) - Tensor tBrB_dq = make_fragment_like(tBsB); // (BCPY,BCPY_M,BCPY_K) - Tensor tBgS = thr_copy_b.partition_S(gS); // (BCPY,BCPY_N,BCPY_K,k) - Tensor tBrS = make_fragment_like(tBgS(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K) - Tensor tBgZ = thr_copy_b.partition_S(gZ); // (BCPY,BCPY_N,BCPY_K,k) - Tensor tBrZ = make_fragment_like(tBgZ(_,_,_,0)); // (BCPY,BCPY_N,BCPY_K) - - // MMA. - ThrMMA thr_mma = mma.get_slice(thread_idx); - Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K) - Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K) - Tensor tCgC = thr_mma.partition_C(gC); // (MMA,MMA_M,MMA_N) - Tensor tCrC = thr_mma.make_fragment_C(tCgC); // (MMA,MMA_M,MMA_N) - - // Predicates for m/n bounds. - Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); // (CPY_M,CPY_K) - Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); // (CPY_N,CPY_K) - Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) - Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) - Tensor cC = make_identity_tensor(make_shape(size<0>(gC), size<1>(gC))); // (M,N) - Tensor tAcA = thr_copy_a.partition_S(cA); // (CPY,CPY_M,CPY_K) - Tensor tBcB = thr_copy_b.partition_S(cB); // (CPY,CPY_N,CPY_K) - Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) - CUTE_UNROLL - for (int m = 0; m < size<0>(tApA); ++m) { - tApA(m,0) = get<0>(tAcA(0,m,0)) < m_max_coord; - } - CUTE_UNROLL - for (int n = 0; n < size<0>(tBpB); ++n) { - tBpB(n,0) = get<0>(tBcB(0,n,0)) < n_max_coord; - } - - // GMEM => RMEM. - auto fetch_gmem = [&](int tile) { - copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA); - copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB); - copy(tBgS(_,_,_,tile), tBrS); - copy(tBgZ(_,_,_,tile), tBrZ); - }; - // RMEM => SMEM. - auto store_smem = [&]() { - __syncthreads(); - copy(tArA, tAsA); - CUTE_UNROLL - for (int k = 0; k < size<2>(tBrB); ++k) { - CUTE_UNROLL - for (int n = 0; n < size<1>(tBrB); ++n) { - cute_dequant(tBrB(_,n,k), tBrS(_,n,k), tBrZ(_,n,k), tBrB_dq(_,n,k)); - } - } - copy(tBrB_dq, tBsB); - __syncthreads(); - }; - - // Prefetch first tile. - fetch_gmem(0); - - // Clear accumulators. - clear(tCrC); - - // Loop over CTA tiles. - auto K_TILE_MAX = size<3>(tAgA); - for (int tile = 0; tile < K_TILE_MAX; ++tile) { - store_smem(); - fetch_gmem((tile + 1 < K_TILE_MAX) ? tile + 1 : tile); - gemm(mma, tCsA, tCsB, tCrC); - } - - // Epilogue. - CUTE_UNROLL - for (int i = 0; i < size(tCrC); ++i) { - if ((get<0>(tCcC(i)) < m_max_coord) && (get<1>(tCcC(i)) < n_max_coord)) { - tCgC(i) = Element(tCrC(i)); - } - } -} - -template -inline constexpr auto make_matrix_stride(auto m, auto k) { - if constexpr (KMajor) { - return cute::make_stride(k, cute::Int<1>{}, m * k); - } else { - return cute::make_stride(cute::Int<1>{}, m, m * k); - } -} - -template -inline constexpr auto make_smem_layout(auto bM, auto bK) { - // TODO: Calculate swizzle based on tile shape. - if constexpr (KMajor) { - auto swizzle = composition(Swizzle<3,3,3>{}, - Layout>, - Stride<_8,Stride<_1,_64>>>{}); - return tile_to_shape(swizzle, make_shape(bM, bK)); - } else { - auto swizzle = composition(Swizzle<3,3,3>{}, - Layout, Stride<_1,_64>>{}); - return tile_to_shape(swizzle, make_shape(bM, bK)); - } -} - -template -inline constexpr auto make_tiled_mma() { - using Atom = std::conditional_t< - SM80, - std::conditional_t< - std::is_same_v, - SM80_16x8x16_F32F16F16F32_TN, - std::conditional_t< - std::is_same_v, - SM80_16x8x16_F32BF16BF16F32_TN, - UniversalFMA - > - >, - UniversalFMA>; - if constexpr (!SM80 || std::is_same_v) { - return make_tiled_mma(Atom{}, Layout>{}); - } else { - if constexpr (TileM >= 32) { - return make_tiled_mma(Atom{}, Layout>{}, Tile<_32,_32,_16>{}); - } else { - return make_tiled_mma(Atom{}, Layout>{}, Tile<_16,_32,_16>{}); - } - } -} - -template -inline auto make_tiled_copy(auto num_threads, auto bM, auto bK) { - auto n_read = Int<8>{}; - auto atom = Copy_Atom>>, T>{}; - if constexpr (KMajor) { - auto k_threads = bK / n_read; - return make_tiled_copy( - atom, - make_layout(make_shape(Int{}, k_threads), LayoutRight{}), - make_layout(make_shape(Int<1>{}, n_read))); - } else { - auto m_threads = bM / n_read; - return make_tiled_copy( - atom, - make_layout(make_shape(m_threads, Int{}), LayoutLeft{}), - make_layout(make_shape(n_read, Int<1>{}))); - } -} - -template -inline constexpr auto make_scales_layout(auto n, auto k, auto l, auto group_size) { - if constexpr (KMajor) { - return make_layout( - make_shape(n, make_shape(group_size, k / group_size), l), - make_stride(k / group_size, Stride<_0,_1>{}, n * k / group_size)); - } else { - return make_layout( - make_shape(make_shape(group_size, n / group_size), k, l), - make_stride(Stride<_0,_1>{}, n / group_size, n * k / group_size)); - } -} - -template -void qmm_naive( - const Element* A, - const Quant* B, - const Scale* S, - const Element* Z, - Element* C, - int m, int n, int k, int l, - bool broadcast_b, - auto group_size, - auto&& launch_kernel) { - // Define shapes (dynamic). - auto prob_shape = make_shape(m, n, k, l); // (M,N,K,L) - - // Define TN strides (mixed). - auto dA = make_stride(k, Int<1>{}, m * k); // (dM,dK,dL) - auto dB = make_matrix_stride(n, k); // (dN,dK,dL) - auto dC = make_stride(n, Int<1>{}, m * n); // (dM,dN,dL) - - // Define layout of scales/biases (mixed). - auto S_layout = make_scales_layout(n, k, l, group_size); - - // Handle broadcasting. - if (broadcast_b) { - get<2>(dB) = 0; - get<2>(stride(S_layout)) = 0; - } - - // Define CTA tile sizes (static). - auto bM = Int{}; - auto bN = Int<(!SM80 && group_size > 64) ? 64 : 128>{}; - auto bK = Int{}; - auto cta_tiler = make_shape(bM, bN, bK); // (BLK_M,BLK_N,BLK_K) - - // Define MMA. - TiledMMA mma = make_tiled_mma(); - auto num_threads = size(mma); - - // Define the A/B smem layouts (static). - auto sA_layout = make_smem_layout(bM, bK); - auto sB_layout = make_smem_layout(bN, bK); - - // Atoms. - TiledCopy copy_a = make_tiled_copy(num_threads, bM, bK); - TiledCopy copy_b = make_tiled_copy(num_threads, bN, bK); - - auto* kernel = &qmm_naive_kernel< - decltype(prob_shape), decltype(cta_tiler), - Element, Quant, Scale, - decltype(dA), decltype(sA_layout), decltype(copy_a), - decltype(dB), decltype(sB_layout), decltype(copy_b), - decltype(dC), decltype(S_layout), decltype(mma)>; - - // Set L1 to be SMEM only. - size_t smem_bytes = sizeof(SharedStorage); - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes); - cudaFuncSetAttribute(kernel, cudaFuncAttributePreferredSharedMemoryCarveout, 100); - - dim3 num_blocks(size(ceil_div(m, bM)), size(ceil_div(n, bN)), l); - dim3 block_dims(num_threads); - void* args[] = { - &prob_shape, &cta_tiler, - &A, &dA, &sA_layout, ©_a, - &B, &dB, &sB_layout, ©_b, - &C, &dC, - &S, &Z, &S_layout, - &mma}; - launch_kernel(reinterpret_cast(kernel), num_blocks, block_dims, smem_bytes, args); -} - -} // namespace cutlass_gemm - -// clang-format on - -namespace mlx::core { - -template -inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { - if (dtype == float32) { - f.template operator()(); - } else if (dtype == float16) { - f.template operator()(); - } else if (dtype == bfloat16) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); - } -} - -template -inline void dispatch_groups(int group_size, const char* tag, F&& f) { - if (group_size == 32) { - f.template operator()<32>(); - } else if (group_size == 64) { - f.template operator()<64>(); - } else if (group_size == 128) { - f.template operator()<128>(); - } else { - throw std::invalid_argument( - fmt::format("{} Group size {} is not supported.", tag, group_size)); - } -} - -template -inline void dispatch_quant_types( - int bits, - int group_size, - QuantizationMode mode, - const char* tag, - F&& f) { - if (mode == QuantizationMode::Mxfp4) { - f.template operator()(); - } else if (mode == QuantizationMode::Mxfp8) { - f.template operator()(); - } else if (mode == QuantizationMode::Nvfp4) { - f.template operator()(); - } else { - dispatch_groups(group_size, tag, [&]() { - if (bits == 2) { - f.template operator()(); - } else if (bits == 3) { - f.template operator()(); - } else if (bits == 4) { - f.template operator()(); - } else if (bits == 5) { - f.template operator()(); - } else if (bits == 6) { - f.template operator()(); - } else if (bits == 8) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} {}-bit quantization is not supported.", tag, bits)); - } - }); - } -} - -template -void qmm_impl_naive( - const array& x, - const array& w, - const array& scales, - const std::optional& biases, - array& out, - int bits, - int group_size, - QuantizationMode mode, - cu::CommandEncoder& encoder) { - const char* tag = "[quantized_matmul]"; - int m = out.ndim() > 1 ? out.shape(-2) : 1; - int n = out.shape(-1); - int k = x.shape(-1); - int l = out.size() / (m * n); - bool broadcast_b = w.ndim() == 2; - - bool is_sm80 = encoder.device().compute_capability_major() >= 8; - dispatch_bool(is_sm80, [&](auto sm80) { - dispatch_element_types(out.dtype(), tag, [&]() { - dispatch_quant_types( - bits, - group_size, - mode, - tag, - [&]() { - encoder.set_input_array(x); - encoder.set_input_array(w); - encoder.set_input_array(scales); - if (biases) { - encoder.set_input_array(*biases); - } - encoder.set_output_array(out); - cutlass_gemm::qmm_naive( - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(scales), - biases ? gpu_ptr(*biases) : nullptr, - gpu_ptr(out), - m, - n, - k, - l, - broadcast_b, - cute::Int{}, - [&](auto* kernel, - dim3 num_blocks, - dim3 block_dims, - uint32_t smem_bytes, - void** args) { - encoder.add_kernel_node_raw( - kernel, num_blocks, block_dims, {}, smem_bytes, args); - }); - }); - }); - }); -} - -} // namespace mlx::core - -#define QMM_NAIVE_GPU(TileM, KMajor) \ - namespace mlx::core { \ - template void qmm_impl_naive( \ - const array& x, \ - const array& w, \ - const array& scales, \ - const std::optional& biases, \ - array& out, \ - int bits, \ - int group_size, \ - QuantizationMode mode, \ - cu::CommandEncoder& encoder); \ - } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_k.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_k.cu deleted file mode 100644 index 4bead82a0..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_k.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(16, true) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_n.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_n.cu deleted file mode 100644 index 993243d9e..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m16_n.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(16, false) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_k.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_k.cu deleted file mode 100644 index def1b4e7b..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_k.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(32, true) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_n.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_n.cu deleted file mode 100644 index bf1a500c5..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m32_n.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(32, false) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_k.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_k.cu deleted file mode 100644 index 92f03c788..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_k.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(64, true) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_n.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_n.cu deleted file mode 100644 index 1d1f04004..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_naive_m64_n.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh" - -QMM_NAIVE_GPU(64, false) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu deleted file mode 100644 index cd682be34..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m16.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" - -QMM_SM80_GPU(16) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu deleted file mode 100644 index 1f79364bf..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m32.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" - -QMM_SM80_GPU(32) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu deleted file mode 100644 index 9af44f8e3..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm80_m64.cu +++ /dev/null @@ -1,5 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh" - -QMM_SM80_GPU(64) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh deleted file mode 100644 index ce79dbcee..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh +++ /dev/null @@ -1,261 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/cutlass_utils.cuh" -#include "mlx/backend/cuda/quantized/quantized_utils.h" -#include "mlx/backend/gpu/copy.h" -#include "mlx/dtype_utils.h" - -#include -#include -#include -#include -#include -#include - -#if defined(MLX_CUDA_SM90A_ENABLED) - -// We can't put kernel code in mlx::core due to name conflicts of "Shape". -namespace cutlass_gemm { - -using namespace cute; - -template < - typename TileShapeMN = Shape<_128, _16>, - typename ClusterShape = Shape<_1, _1, _1>, - typename Element, - typename Quant, - typename GroupSize, - typename F> -void qmm_sm90( - const Element* A, - const Quant* B, - const Element* S, - const Element* Z, - Element* D, - int64_t m, - int64_t n, - int64_t k, - int64_t l, - bool broadcast_b, - GroupSize group_size, - F&& launch_kernel) { - constexpr int kAlignmentA = 128 / sizeof_bits::value; - constexpr int kAlignmentB = 128 / sizeof_bits::value; - constexpr int kTileShapeK = - std::max(64, 128 * 8 / sizeof_bits::value); - static_assert(group_size % kTileShapeK == 0); - - using Arch = cutlass::arch::Sm90; - using Accumulator = float; - using TileShape = decltype(append(TileShapeMN{}, Int{})); - - using Epilogue = typename cutlass::epilogue::collective::CollectiveBuilder< - Arch, - cutlass::arch::OpClassTensorOp, - TileShape, - ClusterShape, - cutlass::epilogue::collective::EpilogueTileAuto, - Accumulator, - Accumulator, - // ElementC: - void, - cutlass::layout::ColumnMajor, - kAlignmentA, - // ElementD: - Element, - cutlass::layout::ColumnMajor, - kAlignmentA, - cutlass::epilogue::TmaWarpSpecializedCooperative>::CollectiveOp; - - // Note that A/B are swapped and transposed to use TMA epilogue. - using Mainloop = typename cutlass::gemm::collective::CollectiveBuilder< - Arch, - cutlass::arch::OpClassTensorOp, - // ElementA: - tuple, - cutlass::layout::RowMajor, - kAlignmentB, - // ElementB: - Element, - cutlass::layout::ColumnMajor, - kAlignmentA, - Accumulator, - TileShape, - ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename Epilogue::SharedStorage))>, - cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel:: - GemmUniversal, Mainloop, Epilogue>; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; - - auto dA = make_stride(k, Int<1>{}, m * k); - auto dB = make_stride(k, Int<1>{}, n * k); - auto dS = make_stride(Int<1>{}, n, n * k / group_size); - auto dD = make_stride(Int<1>{}, n, m * n); - if (broadcast_b) { - get<2>(dB) = 0; - get<2>(dS) = 0; - } - - Gemm gemm; - typename Gemm::Arguments args{ - cutlass::gemm::GemmUniversalMode::kGemm, - {int(n), int(m), int(k), int(l)}, - {B, dB, A, dA, S, dS, group_size, Z}, - {{1.f, 0.f}, D, dD, D, dD}}; - - CHECK_CUTLASS_ERROR(gemm.can_implement(args)); - CHECK_CUTLASS_ERROR(gemm.initialize(args, nullptr)); - - auto* kernel = &cutlass::device_kernel; - void* kernel_params[] = {const_cast(&gemm.params())}; - auto cluster = ClusterShape{}; - launch_kernel( - reinterpret_cast(kernel), - gemm.get_grid_shape(gemm.params()), - GemmKernel::get_block_shape(), - {static_cast(get<0>(cluster)), - static_cast(get<1>(cluster)), - static_cast(get<2>(cluster))}, - GemmKernel::SharedStorageSize, - kernel_params); -} - -} // namespace cutlass_gemm - -namespace mlx::core { - -inline array transpose_last_2_dims( - const array& x, - cu::CommandEncoder& encoder, - const Stream& s) { - array transposed = swapaxes_in_eval(x, -1, -2); - array transposed_copy = contiguous_copy_gpu(transposed, s); - encoder.add_temporary(transposed_copy); - return transposed_copy; -} - -template -inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { - if (dtype == float32) { - f.template operator()(); - } else if (dtype == float16) { - f.template operator()(); - } else if (dtype == bfloat16) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); - } -} - -template -inline void dispatch_quant_types(int bits, const char* tag, F&& f) { - if (bits == 2) { - f.template operator()(); - } else if (bits == 4) { - f.template operator()(); - } else if (bits == 8) { - f.template operator()(); - } else { - throw std::invalid_argument( - fmt::format("{} {}-bit quantization is not supported.", tag, bits)); - } -} - -template -inline void dispatch_groups(int group_size, const char* tag, F&& f) { - if (group_size == 64) { - f(cute::Int<64>{}); - } else if (group_size == 128) { - f(cute::Int<128>{}); - } else { - throw std::invalid_argument( - fmt::format("{} Group size {} is not supported.", tag, group_size)); - } -} - -template -void qmm_impl_sm90( - const array& x, - const array& w, - const array& scales_, - const array& biases_, - array& out, - int bits, - int group_size, - cu::CommandEncoder& encoder, - Stream s) { - const char* tag = "[quantized_matmul]"; - int m = out.ndim() > 1 ? out.shape(-2) : 1; - int n = out.shape(-1); - int k = x.shape(-1); - int l = out.size() / (m * n); - bool broadcast_b = w.ndim() == 2; - - // FIXME: Copy happens for every call. - array scales = transpose_last_2_dims(scales_, encoder, s); - array biases = transpose_last_2_dims(biases_, encoder, s); - - dispatch_element_types(out.dtype(), tag, [&]() { - dispatch_quant_types(bits, tag, [&]() { - dispatch_groups(group_size, tag, [&](auto group_size) { - encoder.set_input_array(x); - encoder.set_input_array(w); - encoder.set_input_array(scales); - encoder.set_input_array(biases); - encoder.set_output_array(out); - cutlass_gemm::qmm_sm90( - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(scales), - gpu_ptr(biases), - gpu_ptr(out), - m, - n, - k, - l, - broadcast_b, - group_size, - [&](auto* kernel, - dim3 num_blocks, - dim3 block_dims, - dim3 cluster_shape, - uint32_t smem_bytes, - void** args) { - encoder.add_kernel_node_raw( - kernel, - num_blocks, - block_dims, - cluster_shape, - smem_bytes, - args); - }); - }); - }); - }); -} - -} // namespace mlx::core - -#define QMM_SM90_GPU(TileShapeMN, ClusterShape) \ - namespace mlx::core { \ - template void qmm_impl_sm90( \ - const array& x, \ - const array& w, \ - const array& scales, \ - const array& biases, \ - array& out, \ - int bits, \ - int group_size, \ - cu::CommandEncoder& encoder, \ - Stream s); \ - } - -#else - -#define QMM_SM90_GPU(TileShapeMN, ClusterShape) - -#endif // defined(MLX_CUDA_SM90A_ENABLED) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu deleted file mode 100644 index 8db29f2ad..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n128_m2.cu +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh" - -using namespace cute; - -using TileShapeMN = Shape<_128, _128>; -using ClusterShape = Shape<_2, _1, _1>; - -QMM_SM90_GPU(TileShapeMN, ClusterShape) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu deleted file mode 100644 index ba1a62b1a..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n16_m1.cu +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh" - -using namespace cute; - -using TileShapeMN = Shape<_128, _16>; -using ClusterShape = Shape<_1, _1, _1>; - -QMM_SM90_GPU(TileShapeMN, ClusterShape) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu deleted file mode 100644 index 81f82c72a..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n256_m2.cu +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh" - -using namespace cute; - -using TileShapeMN = Shape<_128, _256>; -using ClusterShape = Shape<_2, _1, _1>; - -QMM_SM90_GPU(TileShapeMN, ClusterShape) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu deleted file mode 100644 index 0955296c7..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n32_m1.cu +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh" - -using namespace cute; - -using TileShapeMN = Shape<_128, _32>; -using ClusterShape = Shape<_1, _1, _1>; - -QMM_SM90_GPU(TileShapeMN, ClusterShape) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu deleted file mode 100644 index 89a2643bb..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_impl_sm90_m128_n64_m2.cu +++ /dev/null @@ -1,10 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/backend/cuda/quantized/qmm/qmm_impl_sm90.cuh" - -using namespace cute; - -using TileShapeMN = Shape<_128, _64>; -using ClusterShape = Shape<_2, _1, _1>; - -QMM_SM90_GPU(TileShapeMN, ClusterShape) diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_naive.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_naive.cu new file mode 100644 index 000000000..c7c5c6049 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_naive.cu @@ -0,0 +1,121 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device/qmm_naive.cuh" + +#include "mlx/backend/cuda/cutlass_utils.cuh" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/backend/cuda/quantized/qmm/qmm_utils.h" + +#include "cuda_jit_sources.h" + +namespace mlx::core { + +namespace { + +inline auto make_cta_tiler(int itemsize, int m, int group_size, bool sm80) { + bool enough_smem = sm80 && itemsize <= 2 && group_size <= 64; + int tile_m = std::max(16, std::min(64, next_power_of_2(m))); + int tile_n = enough_smem ? 128 : 64; + int tile_k = std::max(64, group_size); + return cute::make_shape(tile_m, tile_n, tile_k); +} + +} // namespace + +void qmm_naive( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& global_scale, + const std::optional& lhs_indices, + const std::optional& rhs_indices, + array& out, + bool transpose, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + auto [m, n, k, l, broadcast_b] = make_problem_shape(x, w, out); + bool sm80 = encoder.device().compute_capability_major() >= 8; + auto cta_tiler = make_cta_tiler(x.itemsize(), m, group_size, sm80); + bool has_k_residue = (k % cute::size<2>(cta_tiler)) != 0; + + std::string module_name = fmt::format( + "qmm_naive_t{}_{}_{}_m{}_b{}_g{}_{}", + transpose ? "n" : "t", + has_k_residue ? "residue" : "aligned", + dtype_to_string(x.dtype()), + cute::size<0>(cta_tiler), + bits, + group_size, + quantization_mode_to_string(mode)); + + auto [ctype_x, ctype_q, ctype_s] = get_qmm_cutlass_types(x, bits, mode); + std::string kernel_name = fmt::format( + "mlx::core::cu::qmm_naive_kernel<{}, {}, {}, {}, {}, {}, {}, {}>", + group_size, + transpose, + has_k_residue, + sm80, + ctype_x, + ctype_q, + ctype_s, + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_qmm_naive, std::vector{kernel_name}); + }); + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + if (global_scale) { + encoder.set_input_array(*global_scale); + } + if (lhs_indices) { + encoder.set_input_array(*lhs_indices); + } + if (rhs_indices) { + encoder.set_input_array(*rhs_indices); + } + encoder.set_output_array(out); + + dim3 num_blocks{ + uint32_t(cute::ceil_div(m, cute::size<0>(cta_tiler))), + uint32_t(cute::ceil_div(n, cute::size<1>(cta_tiler))), + uint32_t(l)}; + dim3 block_dims{uint32_t(cute::size(cu::make_tiled_mma(cta_tiler)))}; + + auto [sA_layout, sB_layout] = cu::make_smem_layouts(cta_tiler); + size_t smem_bytes = + x.itemsize() * (cute::cosize(sA_layout) + cute::cosize(sB_layout)); + + encoder.add_kernel_node_ex( + mod.get_kernel(kernel_name), + num_blocks, + block_dims, + {}, + smem_bytes, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, + lhs_indices ? gpu_ptr(*lhs_indices) : nullptr, + rhs_indices ? gpu_ptr(*rhs_indices) : nullptr, + gpu_ptr(out), + m, + n, + k, + l, + broadcast_b); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu new file mode 100644 index 000000000..981823956 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm80.cu @@ -0,0 +1,116 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device/qmm_sm80.cuh" + +#include "mlx/backend/cuda/cutlass_utils.cuh" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/kernel_utils.cuh" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/backend/cuda/quantized/qmm/qmm_utils.h" + +#include "cuda_jit_sources.h" + +namespace mlx::core { + +namespace { + +inline auto make_cta_tiler(int m, int group_size) { + int tile_m = std::max(16, std::min(64, next_power_of_2(m))); + int tile_n = 128; + int tile_k = std::max(64, group_size); + return cute::make_shape(tile_m, tile_n, tile_k); +} + +} // namespace + +void qmm_sm80( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const std::optional& lhs_indices, + const std::optional& rhs_indices, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + auto [m, n, k, l, broadcast_b] = make_problem_shape(x, w, out); + auto cta_tiler = make_cta_tiler(m, group_size); + + std::string module_name = fmt::format( + "qmm_sm80_tn_{}_m{}_b{}_g{}_{}", + dtype_to_string(x.dtype()), + cute::size<0>(cta_tiler), + bits, + group_size, + quantization_mode_to_string(mode)); + + auto [ctype_x, ctype_q, ctype_s] = get_qmm_cutlass_types(x, bits, mode); + std::string kernel_name = fmt::format( + "mlx::core::cu::qmm_sm80_kernel<{}, {}, {}, {}, {}>", + group_size, + ctype_x, + ctype_q, + ctype_s, + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_qmm_sm80, std::vector{kernel_name}); + }); + + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + if (lhs_indices) { + encoder.set_input_array(*lhs_indices); + } + if (rhs_indices) { + encoder.set_input_array(*rhs_indices); + } + encoder.set_output_array(out); + + dim3 num_blocks{ + uint32_t(cute::ceil_div(m, cute::size<0>(cta_tiler))), + uint32_t(cute::ceil_div(n, cute::size<1>(cta_tiler))), + uint32_t(l)}; + dim3 block_dims{uint32_t(cute::size(cu::make_tiled_mma()))}; + + auto [sA_layout, sB_layout, sC_layout] = cu::make_smem_layouts(cta_tiler); + size_t smem_bytes = std::max( + cute::cosize(sA_layout) * x.itemsize() + + cute::cosize(sB_layout) * bits / 8, + cute::cosize(sC_layout) * x.itemsize()); + + auto kernel = mod.get_kernel(kernel_name, [&](CUfunction kernel) { + if (smem_bytes > 48000) { + cuFuncSetAttribute( + kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_bytes); + } + }); + + encoder.add_kernel_node_ex( + kernel, + num_blocks, + block_dims, + {}, + smem_bytes, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + lhs_indices ? gpu_ptr(*lhs_indices) : nullptr, + rhs_indices ? gpu_ptr(*rhs_indices) : nullptr, + gpu_ptr(out), + m, + n, + k, + l, + broadcast_b); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm90.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm90.cu new file mode 100644 index 000000000..7bc5d8343 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_sm90.cu @@ -0,0 +1,207 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cuda/device/qmm_sm90.cuh" + +#include "mlx/backend/cuda/cutlass_utils.cuh" +#include "mlx/backend/cuda/jit_module.h" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" +#include "mlx/backend/cuda/quantized/qmm/qmm_utils.h" +#include "mlx/backend/gpu/copy.h" + +#include "cuda_jit_sources.h" + +namespace mlx::core { + +using namespace cute; + +namespace { + +inline array transpose_last_2_dims( + const array& x, + cu::CommandEncoder& encoder, + const Stream& s) { + array transposed = swapaxes_in_eval(x, -1, -2); + array transposed_copy = contiguous_copy_gpu(transposed, s); + encoder.add_temporary(transposed_copy); + return transposed_copy; +} + +template +inline void dispatch_element_types(Dtype dtype, const char* tag, F&& f) { + if (dtype == float32) { + f.template operator()(); + } else if (dtype == float16) { + f.template operator()(); + } else if (dtype == bfloat16) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} Unsupported dtype: {}.", tag, dtype_to_string(dtype))); + } +} + +template +inline void dispatch_quant_types(int bits, const char* tag, F&& f) { + if (bits == 2) { + f.template operator()(); + } else if (bits == 4) { + f.template operator()(); + } else if (bits == 8) { + f.template operator()(); + } else { + throw std::invalid_argument( + fmt::format("{} {}-bit quantization is not supported.", tag, bits)); + } +} + +template +inline void dispatch_groups(int group_size, const char* tag, F&& f) { + if (group_size == 64) { + f.template operator()<64>(); + } else if (group_size == 128) { + f.template operator()<128>(); + } else { + throw std::invalid_argument( + fmt::format("{} Group size {} is not supported.", tag, group_size)); + } +} + +template +inline void dispatch_tile(int m, F&& f) { + if (m <= 16) { + f.template operator()<16>(); + } else if (m <= 32) { + f.template operator()<32>(); + } else if (m <= 64) { + f.template operator()<64>(); + } else if (m <= 128) { + f.template operator()<128>(); + } else { + f.template operator()<256>(); + } +} + +template +inline void dispatch_gemm( + const array& x, + int n, + int bits, + int group_size, + const char* tag, + F&& f) { + dispatch_element_types(x.dtype(), tag, [&]() { + dispatch_tile(n, [&]() { + dispatch_quant_types(bits, tag, [&]() { + dispatch_groups(group_size, tag, [&]() { + auto cta_tiler = make_shape( + Int<128>{}, + Int{}, + Int)>{}); + auto gemm = cu::make_qmm_sm90_kernel< + GroupSize, + Element, + Quant, + decltype(cta_tiler)>(); + f(cta_tiler, gemm); + }); + }); + }); + }); +} + +} // namespace + +void qmm_sm90( + const array& x, + const array& w, + const array& scales_, + const array& biases_, + array& out, + int bits, + int group_size, + cu::CommandEncoder& encoder, + Stream s) { + const char* tag = "[quantized_matmul]"; + auto [m, n, k, l, broadcast_b] = make_problem_shape(x, w, out); + + auto dA = make_stride(int64_t(k), Int<1>{}, int64_t(m * k)); + auto dB = make_stride(int64_t(k), Int<1>{}, int64_t(n * k)); + auto dS = make_stride(Int<1>{}, int64_t(n), int64_t(n * k / group_size)); + auto dD = make_stride(Int<1>{}, int64_t(n), int64_t(m * n)); + if (broadcast_b) { + get<2>(dB) = 0; + get<2>(dS) = 0; + } + + // FIXME: Copy happens for every call. + array scales = transpose_last_2_dims(scales_, encoder, s); + array biases = transpose_last_2_dims(biases_, encoder, s); + + dispatch_gemm(x, n, bits, group_size, tag, [&](auto cta_tiler, auto gemm) { + // JIT compilation. + std::string module_name = fmt::format( + "qmm_sm90_tn_{}_n{}_b{}_g{}_affine", + dtype_to_string(x.dtype()), + int(size<1>(cta_tiler)), + bits, + group_size); + + auto [ctype_x, ctype_q, ctype_s] = get_qmm_cutlass_types(x, bits); + std::string kernel_name = fmt::format( + "cutlass::device_kernel>", + group_size, + ctype_x, + ctype_q, + cta_tiler_to_string(cta_tiler)); + + cu::JitModule& mod = + cu::get_jit_module(encoder.device(), module_name, [&]() { + return std::make_tuple( + false, jit_source_qmm_sm90, std::vector{kernel_name}); + }); + + // Prepare kernel args. + using Gemm = decltype(gemm); + + const auto* A = gpu_ptr(x); + const auto* B = gpu_ptr(w); + const auto* S = + gpu_ptr(scales); + const auto* Z = + gpu_ptr(biases); + auto* D = gpu_ptr(out); + + auto params = Gemm::to_underlying_arguments( + {cutlass::gemm::GemmUniversalMode::kGemm, + {n, m, k, l}, + {B, dB, A, dA, S, dS, group_size, Z}, + {{1.f, 0.f}, D, dD, D, dD}}, + nullptr); + + size_t smem_bytes = Gemm::SharedStorageSize; + auto kernel = mod.get_kernel(kernel_name, [&](CUfunction kernel) { + if (smem_bytes > 48000) { + cuFuncSetAttribute( + kernel, + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + smem_bytes); + } + }); + + // Append to CUDA graph. + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + encoder.set_input_array(biases); + encoder.set_output_array(out); + encoder.add_kernel_node_ex( + kernel, + Gemm::get_grid_shape(params), + Gemm::get_block_shape(), + {}, + smem_bytes, + params); + }); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_utils.h new file mode 100644 index 000000000..8e54e2e05 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmm_utils.h @@ -0,0 +1,72 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/dtype_utils.h" + +namespace mlx::core { + +inline auto +make_problem_shape(const array& x, const array& w, const array& out) { + int m = out.ndim() > 1 ? out.shape(-2) : 1; + int n = out.shape(-1); + int k = x.shape(-1); + int l = out.size() / (m * n); + bool broadcast_b = (w.ndim() <= 2) || (w.size() != w.data_size()); + return std::make_tuple(m, n, k, l, broadcast_b); +} + +inline const char* get_weight_cutlass_type(const Dtype& dtype) { + switch (dtype) { + case float16: + return "cutlass::half_t"; + case bfloat16: + return "cutlass::bfloat16_t"; + case float32: + return "float"; + default: + throw std::invalid_argument( + fmt::format( + "[quantized_matmul] Unsupported dtype: {}.", + dtype_to_string(dtype))); + } +} + +inline std::tuple +get_quant_cutlass_types(const char* ctype_x, int bits, QuantizationMode mode) { + if (mode == QuantizationMode::Mxfp4) { + return {"cutlass::float_e2m1_t", "cutlass::float_ue8m0_t"}; + } else if (mode == QuantizationMode::Mxfp8) { + return {"cutlass::float_e4m3_t", "cutlass::float_ue8m0_t"}; + } else if (mode == QuantizationMode::Nvfp4) { + return {"cutlass::float_e2m1_t", "cutlass::float_e4m3_t"}; + } else { + if (bits == 2) { + return {"cutlass::uint2b_t", ctype_x}; + } else if (bits == 3) { + return {"cutlass::uint3b_t", ctype_x}; + } else if (bits == 4) { + return {"cutlass::uint4b_t", ctype_x}; + } else if (bits == 5) { + return {"cutlass::uint5b_t", ctype_x}; + } else if (bits == 6) { + return {"cutlass::uint6b_t", ctype_x}; + } else if (bits == 8) { + return {"uint8_t", ctype_x}; + } else { + throw std::invalid_argument( + fmt::format( + "[quantized_matmul] {}-bit quantization is not supported.", + bits)); + } + } +} + +inline std::tuple get_qmm_cutlass_types( + const array& x, + int bits, + QuantizationMode mode = QuantizationMode::Affine) { + auto ctype_x = get_weight_cutlass_type(x.dtype()); + auto [ctype_q, ctype_s] = get_quant_cutlass_types(ctype_x, bits, mode); + return {ctype_x, ctype_q, ctype_s}; +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu index ceda750a7..4ec6b95a7 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmm/qmv.cu @@ -1,7 +1,7 @@ // Copyright © 2026 Apple Inc. +#include "mlx/backend/cuda/device/cute_dequant.cuh" #include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/qmm/cute_dequant.cuh" #include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/dtype_utils.h" @@ -82,7 +82,6 @@ dequant_fma(const T* x, const Q* w, S scale, T bias, float* out) { } template < - int rows_per_block, int elems_per_thread, int group_size, bool has_bias, @@ -90,30 +89,18 @@ template < typename T, typename Q, typename S> -__global__ void qmv_kernel( +__device__ __forceinline__ void qmv_kernel_impl( const T* x, const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, + int row, + int w_batch, int n, - int k, - bool broadcast_w) { - auto grid = cg::this_grid(); - auto block = cg::this_thread_block(); - auto warp = cg::tiled_partition(block); - - // The row that this warp handles. - int row = block.group_index().x * rows_per_block + warp.meta_group_rank(); - if (row >= n) { - return; - } - - // Advance pointers of x/out. - int m = grid.dim_blocks().y; - int l = block.group_index().z; - x += block.group_index().y * k + m * k * l; - out += block.group_index().y * n + m * n * l; + int k) { + auto warp = cg::tiled_partition(cg::this_thread_block()); // For sub-byte Q, pointer moves by 8bits for each advance, e.g. w += 1 would // move past 2 elements for 4-bit Q. @@ -124,7 +111,6 @@ __global__ void qmv_kernel( int groups_per_row = k / group_size; // Advance w/scales/biases to current row. - int w_batch = broadcast_w ? 0 : l; w += (static_cast(row) + n * w_batch) * w_step(k); scales += (static_cast(row) + n * w_batch) * groups_per_row; if constexpr (has_bias) { @@ -170,10 +156,98 @@ __global__ void qmv_kernel( // Write result for current warp, which maps to rows 1-to-1. if (warp.thread_rank() == 0) { + if constexpr ( + cuda::std::is_same_v && + cuda::std::is_same_v) { + // Only nvfp4 supports global scale. + if (global_scale) { + sum *= (*global_scale / (F8E4M3_MAX * F4E2M1_MAX)); + } + } out[row] = static_cast(sum); } } +template < + int rows_per_block, + int elems_per_thread, + int group_size, + bool has_bias, + bool has_residue_k, + typename T, + typename Q, + typename S> +__global__ void qmv_kernel( + const T* x, + const Q* w, + const S* scales, + const T* biases, + const float* global_scale, + T* out, + int n, + int k, + bool broadcast_w) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + // The row that this warp handles. + int row = block.group_index().x * rows_per_block + warp.meta_group_rank(); + if (row >= n) { + return; + } + + // Advance pointers of x/out for M and batch dimensions. + int m = grid.dim_blocks().y; + int l = block.group_index().z; + x += block.group_index().y * k + m * k * l; + out += block.group_index().y * n + m * n * l; + int w_batch = broadcast_w ? 0 : l; + + qmv_kernel_impl( + x, w, scales, biases, global_scale, out, row, w_batch, n, k); +} + +template < + int rows_per_block, + int elems_per_thread, + int group_size, + bool has_bias, + bool has_residue_k, + typename T, + typename Q, + typename S> +__global__ void gather_qmv_kernel( + const T* x, + const Q* w, + const S* scales, + const T* biases, + T* out, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, + int n, + int k) { + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + + int row = block.group_index().x * rows_per_block + warp.meta_group_rank(); + if (row >= n) { + return; + } + + int m = grid.dim_blocks().y; + int l = block.group_index().z; + uint32_t x_idx = lhs_indices[l]; + uint32_t w_idx = rhs_indices[l]; + + x += block.group_index().y * k + m * k * x_idx; + out += block.group_index().y * n + m * n * l; + + qmv_kernel_impl( + x, w, scales, biases, nullptr, out, row, w_idx, n, k); +} + template < int group_size, bool has_bias, @@ -186,6 +260,7 @@ void qmv( const Q* w, const S* scales, const T* biases, + const float* global_scale, T* out, int m, int n, @@ -200,7 +275,8 @@ void qmv( dim3 num_blocks{ uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m), uint32_t(l)}; dim3 block_dims{WARP_SIZE, rows_per_block}; - void* args[] = {&x, &w, &scales, &biases, &out, &n, &k, &broadcast_w}; + void* args[] = { + &x, &w, &scales, &biases, &global_scale, &out, &n, &k, &broadcast_w}; dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) { auto* kernel = &qmv_kernel< @@ -217,6 +293,51 @@ void qmv( }); } +template < + int group_size, + bool has_bias, + typename T, + typename Q, + typename S, + typename F> +void gather_qmv( + const T* x, + const Q* w, + const S* scales, + const T* biases, + T* out, + const uint32_t* lhs_indices, + const uint32_t* rhs_indices, + int m, + int n, + int k, + int l, + F&& launch_kernel) { + constexpr int rows_per_block = 8; + constexpr int elems_per_thread = + (cute::sizeof_bits_v <= 16 && cute::sizeof_bits_v <= 4) ? 16 : 8; + + dim3 num_blocks{ + uint32_t(cuda::ceil_div(n, rows_per_block)), uint32_t(m), uint32_t(l)}; + dim3 block_dims{WARP_SIZE, rows_per_block}; + void* args[] = { + &x, &w, &scales, &biases, &out, &lhs_indices, &rhs_indices, &n, &k}; + + dispatch_bool(k % (WARP_SIZE * elems_per_thread), [&](auto has_residue_k) { + auto* kernel = &gather_qmv_kernel< + rows_per_block, + elems_per_thread, + group_size, + has_bias, + has_residue_k.value, + T, + Q, + S>; + launch_kernel( + reinterpret_cast(kernel), num_blocks, block_dims, args); + }); +} + } // namespace cu template @@ -287,6 +408,7 @@ void qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int bits, int group_size, @@ -297,7 +419,7 @@ void qmv( int n = out.shape(-1); int k = x.shape(-1); int l = out.size() / (m * n); - bool broadcast_w = w.ndim() == 2; + bool broadcast_w = (w.ndim() <= 2) || (w.size() != w.data_size()); dispatch_element_types(out.dtype(), tag, [&]() { dispatch_quant_types( @@ -312,6 +434,9 @@ void qmv( if (biases) { encoder.set_input_array(*biases); } + if (global_scale) { + encoder.set_input_array(*global_scale); + } encoder.set_output_array(out); constexpr bool has_bias = !cutlass::has_negative_zero_v; cu::qmv( @@ -319,6 +444,7 @@ void qmv( gpu_ptr(w), gpu_ptr(scales), biases ? gpu_ptr(*biases) : nullptr, + global_scale ? gpu_ptr(*global_scale) : nullptr, gpu_ptr(out), m, n, @@ -333,4 +459,59 @@ void qmv( }); } +void gather_qmv( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + const array& lhs_indices, + const array& rhs_indices, + array& out, + int bits, + int group_size, + QuantizationMode mode, + cu::CommandEncoder& encoder) { + const char* tag = "[gather_qmm]"; + int m = out.shape(-2); + int n = out.shape(-1); + int k = x.shape(-1); + int l = out.size() / (m * n); + + dispatch_element_types(out.dtype(), tag, [&]() { + dispatch_quant_types( + bits, + group_size, + mode, + tag, + [&]() { + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(scales); + if (biases) { + encoder.set_input_array(*biases); + } + encoder.set_input_array(lhs_indices); + encoder.set_input_array(rhs_indices); + encoder.set_output_array(out); + constexpr bool has_bias = !cutlass::has_negative_zero_v; + cu::gather_qmv( + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(scales), + biases ? gpu_ptr(*biases) : nullptr, + gpu_ptr(out), + gpu_ptr(lhs_indices), + gpu_ptr(rhs_indices), + m, + n, + k, + l, + [&](auto* kernel, dim3 num_blocks, dim3 block_dims, void** args) { + encoder.add_kernel_node_raw( + kernel, num_blocks, block_dims, {}, 0, args); + }); + }); + }); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.cu deleted file mode 100644 index 0cc004647..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.cu +++ /dev/null @@ -1,304 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#include "mlx/backend/cuda/device/utils.cuh" -#include "mlx/backend/cuda/kernel_utils.cuh" -#include "mlx/backend/cuda/quantized/qmv.h" -#include "mlx/backend/cuda/quantized/quantized_utils.cuh" -#include "mlx/dtype_utils.h" - -#include -#include - -namespace mlx::core::cu { - -namespace cg = cooperative_groups; - -static constexpr int rows_per_block = 8; - -template -__device__ void adjust_matrix_offsets( - const T*& x, - const uint32_t*& w, - const uint8_t*& scales, - T*& y, - int output_stride, - const int& x_batch_ndims, - const Shape x_shape, - const Strides x_strides, - const int& w_batch_ndims, - const Shape w_shape, - const Strides w_strides, - const Strides s_strides) { - uint32_t idx = cg::this_grid().block_index().z; - if (x_batch_ndims == 1) { - x += idx * x_strides[0]; - } else { - x += elem_to_loc(idx, x_shape.data(), x_strides.data(), x_batch_ndims); - } - if (w_batch_ndims == 1) { - w += idx * w_strides[0]; - scales += idx * s_strides[0]; - } else { - auto [w_idx, s_idx] = elem_to_loc( - idx, w_shape.data(), w_strides.data(), s_strides.data(), w_batch_ndims); - w += w_idx; - scales += s_idx; - } - y += idx * output_stride; -} - -template < - typename T, - int rows_per_block, - int n_per_thread, - int bits, - int group_size, - bool use_mx_scale> -__device__ void fp_qmv_impl( - const uint32_t* mat, - const uint8_t* scales_, - const T* vec, - T* out, - int rows, - int cols) { - auto block = cg::this_thread_block(); - auto warp = cg::tiled_partition(block); - - constexpr int vals_per_item = bits == 8 ? 4 : 8; - constexpr int nv_per_thread = vals_per_item * n_per_thread; - auto g_idx = block.group_index(); - auto t_idx = block.thread_index(); - int row = g_idx.y * rows_per_block + t_idx.y; - - vec += g_idx.x * cols; - out += g_idx.x * rows; - - using ScaleType = - std::conditional_t; - auto scales = (ScaleType*)(scales_); - auto packed_cols = cols / vals_per_item; - - if (row < rows) { - constexpr int scales_per_step = std::max(nv_per_thread / group_size, 1); - constexpr int scale_step = (WARP_SIZE * nv_per_thread) / group_size; - constexpr int n_per_step = n_per_thread / scales_per_step; - // Offset scales to correct row - scales += row * (cols / group_size) + - (warp.thread_rank() * nv_per_thread) / group_size; - float sum = 0.0f; - for (int col = n_per_thread * warp.thread_rank(); col < packed_cols; - col += (WARP_SIZE * n_per_thread)) { - auto local_vec = - unsafe_load_vector(vec + vals_per_item * col, 0); - auto local_mat = - unsafe_load_vector(mat + row * packed_cols + col, 0); -#pragma unroll - for (int i = 0; i < scales_per_step; ++i) { - float2 local_sum = {0.0f, 0.0f}; -#pragma unroll - for (int j = 0; j < n_per_step; ++j) { - int k = n_per_step * i + j; - if constexpr (bits == 8) { - auto v = dequant_fp8(local_mat[k]); - local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k]); - local_sum.x += - v.y * static_cast(local_vec[vals_per_item * k + 1]); - local_sum.y += - v.z * static_cast(local_vec[vals_per_item * k + 2]); - local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 3]); - } else { - auto v = dequant_fp4(local_mat[k]); - local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k]); - local_sum.y += - v.y * static_cast(local_vec[vals_per_item * k + 1]); - local_sum.x += - v.z * static_cast(local_vec[vals_per_item * k + 2]); - local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 3]); - - v = dequant_fp4(local_mat[k] >> 16); - local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k + 4]); - local_sum.y += - v.y * static_cast(local_vec[vals_per_item * k + 5]); - local_sum.x += - v.z * static_cast(local_vec[vals_per_item * k + 6]); - local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 7]); - } - } - sum += (local_sum.x + local_sum.y) * float(scales[i]); - } - scales += scale_step; - } - - sum = cg::reduce(warp, sum, cg::plus{}); - if (warp.thread_rank() == 0) { - out[row] = static_cast(sum); - } - } -} - -template < - typename T, - int rows_per_block, - int n_per_thread, - int bits, - int group_size, - bool use_mx_scale> -__global__ void fp_qmv_single( - const uint32_t* mat, - const uint8_t* scales, - const T* vec, - T* out, - int rows, - int cols) { - fp_qmv_impl( - mat, scales, vec, out, rows, cols); -} - -template < - typename T, - int rows_per_block, - int n_per_thread, - int bits, - int group_size, - bool use_mx_scale> -__global__ void fp_qmv_batched( - const uint32_t* mat, - const uint8_t* scales, - const T* vec, - T* out, - int rows, - int cols, - int vec_batch_ndims, - const __grid_constant__ Shape vec_shape, - const __grid_constant__ Strides vec_strides, - int mat_batch_ndims, - const __grid_constant__ Shape mat_shape, - const __grid_constant__ Strides mat_strides, - const __grid_constant__ Strides scales_strides) { - adjust_matrix_offsets( - vec, - mat, - scales, - out, - rows * vec_shape[vec_batch_ndims], - vec_batch_ndims, - vec_shape, - vec_strides, - mat_batch_ndims, - mat_shape, - mat_strides, - scales_strides); - fp_qmv_impl( - mat, scales, vec, out, rows, cols); -} - -template -void dispatch_1_2_4(int n, F&& f) { - switch (n) { - case 1: - f(std::integral_constant{}); - break; - case 2: - f(std::integral_constant{}); - break; - case 4: - f(std::integral_constant{}); - break; - } -} - -void fp_qmv( - const array& mat, - const array& scales, - const array& vec, - array& out, - int bits, - int group_size, - int M, - int N, - int K, - CommandEncoder& encoder) { - encoder.set_input_array(mat); - encoder.set_input_array(scales); - encoder.set_input_array(vec); - encoder.set_output_array(out); - dispatch_float_types(out.dtype(), "qmv", [&](auto type_tag) { - using T = cuda_type_t; - if constexpr (!std::is_same_v) { - dim3 block_dims{WARP_SIZE, rows_per_block}; - uint32_t B = out.size() / (M * N); - uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block; - const uint32_t* mat_ptr = gpu_ptr(mat); - const T* vec_ptr = gpu_ptr(vec); - int n = 1; - if (K % 32 == 0 && cu::is_aligned<4>(mat_ptr) && - ((bits == 4 && cu::is_aligned<8>(vec_ptr)) || - cu::is_aligned<4>(vec_ptr))) { - n = 4; - } else if ( - cu::is_aligned<2>(mat_ptr) && - ((bits == 4 && cu::is_aligned<4>(vec_ptr)) || - cu::is_aligned<2>(vec_ptr))) { - n = 2; - } - dispatch_1_2_4(n, [&](auto n) { - dispatch_bool(B > 1, [&](auto batched) { - if (!batched.value) { - auto kernel = - fp_qmv_single; - if (bits == 8) { - kernel = fp_qmv_single; - } else if (group_size == 16) { - kernel = fp_qmv_single; - } - encoder.add_kernel_node( - kernel, - {static_cast(M), blocks_y}, - block_dims, - 0, - mat_ptr, - gpu_ptr(scales), - vec_ptr, - gpu_ptr(out), - N, - K); - } else { - auto kernel = - fp_qmv_batched; - if (bits == 8) { - kernel = fp_qmv_batched; - } else if (group_size == 16) { - kernel = fp_qmv_batched; - } - encoder.add_kernel_node( - kernel, - {static_cast(M), blocks_y, B}, - block_dims, - 0, - mat_ptr, - gpu_ptr(scales), - vec_ptr, - gpu_ptr(out), - N, - K, - vec.ndim() - 2, - const_param(vec.shape()), - const_param(vec.strides()), - mat.ndim() - 2, - const_param(mat.shape()), - const_param(mat.strides()), - const_param(scales.strides())); - } - }); - }); - } - }); -} - -} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.h deleted file mode 100644 index d4f04195e..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qmv.h +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#pragma once - -#include "mlx/backend/cuda/device.h" - -namespace mlx::core::cu { - -void fp_qmv( - const array& w, - const array& scales, - const array& vec, - array& out, - int bits, - int group_size, - int M, - int N, - int K, - CommandEncoder& encoder); - -} // namespace mlx::core::cu diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp index 3b5be7718..196bd9c05 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm.cpp @@ -1,7 +1,7 @@ // Copyright © 2025 Apple Inc. #include "mlx/backend/cuda/device.h" -#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/backend/cuda/quantized/qqmm_impl.h" #include "mlx/backend/cuda/quantized/qqmm_utils.h" #include "mlx/backend/cuda/quantized/quantized.h" @@ -14,129 +14,180 @@ namespace mlx::core { namespace { -array pad_and_swizzle_scales( - const array& scale, +std::tuple quantize_input( + const array& input, cu::CommandEncoder& encoder, - const Stream& s) { - // Compute padded dimensions for full tiles (128 rows × 4 cols) + const Stream& s, + QuantizationMode mode, + int bits, + int group_size, + std::optional global_scale) { + const array x = ensure_contiguous(input, encoder, s); + + // Compute output shapes + auto xq_shape = x.shape(); + xq_shape.back() = x.shape(-1) * bits / 32; + + const int64_t scales_inner = x.shape(-1) / group_size; auto [pad_outer, pad_inner] = - get_padded_scale_dims(scale.shape(-2), scale.shape(-1)); - // cuBLAS requirements for scale factor layout: - // 1. Dimensions must be padded to full tiles (128 rows × 4 cols) - // 2. Out-of-bounds values must be filled with zeros - // 3. Starting addresses must be 16-byte aligned - // - // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout - // Note: cu::malloc_async already provides 256-byte alignment - array scale_tiled( - cu::malloc_async(pad_outer * pad_inner, encoder), - Shape{pad_outer, pad_inner}, - scale.dtype()); - swizzle_scales(scale, scale_tiled, encoder, s); - - encoder.add_temporary(scale_tiled); - return scale_tiled; + get_padded_scale_dims(x.shape(-2), scales_inner); + + auto sshape = x.shape(); + sshape[x.ndim() - 2] = pad_outer; + sshape[x.ndim() - 1] = pad_inner; + sshape.back() = scales_inner; + + // Allocate outputs + const int64_t xq_bytes = x.size() * bits / 8; + const int64_t batch = x.size() / (x.shape(-2) * x.shape(-1)); + const int64_t scales_bytes = batch * (pad_outer * pad_inner); + + array x_q(cu::malloc_async(xq_bytes, encoder), std::move(xq_shape), uint32); + array scales_x( + cu::malloc_async(scales_bytes, encoder), std::move(sshape), uint8); + encoder.add_temporary(x_q); + encoder.add_temporary(scales_x); + // global_scale is not nullopt only for NVFP4 + fp_quantize(x, x_q, scales_x, group_size, bits, global_scale, encoder, s); + return {std::move(x_q), std::move(scales_x)}; +} + +array quantize_dequantize_input( + const array& x_pre, + const std::optional& global_scale, + int bits, + int group_size, + cu::CommandEncoder& encoder, + Stream s) { + bool donate_x = x_pre.is_donatable(); + array x = ensure_row_contiguous(x_pre, encoder, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = donate_x + ? x + : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype()); + if (!donate_x) { + encoder.add_temporary(xhat); + } + fp_quantize_dequantize(x, xhat, group_size, bits, global_scale, encoder, s); + return xhat; +} + +GemmScalars create_nvfp4_scalars( + const array& global_scale_x, + const array& global_scale_w, + cu::CommandEncoder& encoder) { + // NVFP4 requires alpha/beta as device pointers + // alpha = amax_x * amax_w / (448 * 6)^2 + // beta = 0 + array alpha(cu::malloc_async(sizeof(float), encoder), {}, float32); + array beta(cu::malloc_async(sizeof(float), encoder), {}, float32); + compute_qqmm_pointers(alpha, beta, global_scale_x, global_scale_w, encoder); + encoder.add_temporary(alpha); + encoder.add_temporary(beta); + return {alpha, beta}; } } // namespace void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { - assert( - (inputs.size() == 3 && inputs[1].dtype() == uint32) || - (inputs.size() == 2)); nvtx3::scoped_range r("QQMatmul::eval_gpu"); auto& s = stream(); auto& encoder = cu::get_command_encoder(s); auto& device = encoder.device(); - bool w_quantized = (inputs[1].dtype() == uint32); - if (w_quantized && inputs[0].shape(-2) == 1) { - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - - bool donate_x = inputs[0].is_donatable(); - array x = ensure_row_contiguous(inputs[0], encoder, s); - // If x is a copy it should be donatable - donate_x |= x.is_donatable(); - auto xhat = donate_x - ? x - : array(cu::malloc_async(x.nbytes(), encoder), x.shape(), x.dtype()); - if (!donate_x) { - encoder.add_temporary(xhat); - } - fp_quantize_dequantize(x, xhat, group_size_, bits_, encoder, s); + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; - // Make sure the last two dims of w and s are contiguous - array w = ensure_row_contiguous_matrix(inputs[1], encoder, s); - array scales = ensure_row_contiguous_matrix(inputs[2], encoder, s); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); - bool non_batched = w.ndim() == 2; - int K = x.shape(-1); - int M = non_batched ? x.size() / K : x.shape(-2); - int N = out.shape(-1); + // - 2 inputs: x, w (non-quantized w) + // - 3 inputs: x, w, scales_w (quantized w) + bool w_quantized = (w_pre.dtype() == uint32); + int base_size = w_quantized ? 3 : 2; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); - fp_qmv(w, scales, xhat, out, bits_, group_size_, M, N, K, encoder); - return; + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; } - auto cc = device.compute_capability_major() * 100 + - device.compute_capability_minor() * 10; - if (cc < 1000) { - throw std::runtime_error( - "[QQMatmul::eval_gpu] QQMM is only supported on GPUs with compute capability 10.0 or higher."); - } - auto quantize = [&](const array& input, - cu::CommandEncoder& encoder, - const Stream& s) -> std::pair { - auto x = ensure_contiguous(input, encoder, s); - auto xq_shape = x.shape(); - xq_shape.back() = x.shape(-1) * bits_ / 32; - - auto sshape = x.shape(); - const int64_t scales_inner = x.shape(-1) / group_size_; - auto [pad_outer, pad_inner] = - get_padded_scale_dims(x.shape(-2), scales_inner); - sshape[x.ndim() - 2] = pad_outer; - sshape[x.ndim() - 1] = pad_inner; - sshape.back() = scales_inner; - - // Allocate outputs - const int64_t xq_bytes = x.size() * bits_ / 8; - const int64_t batch = x.size() / (x.shape(-2) * x.shape(-1)); - const int64_t scales_bytes = batch * (pad_outer * pad_inner); - - array x_q(cu::malloc_async(xq_bytes, encoder), std::move(xq_shape), uint32); - array scales_x( - cu::malloc_async(scales_bytes, encoder), std::move(sshape), uint8); - - fp_quantize(x, x_q, scales_x, group_size_, bits_, encoder, s); - - encoder.add_temporary(x_q); - encoder.add_temporary(scales_x); - return {x_q, scales_x}; - }; - auto [x_q, scale_x_pre] = quantize(inputs[0], encoder, s); - auto [w_q, scale_w_pre] = !w_quantized ? quantize(inputs[1], encoder, s) - : std::make_pair(inputs[1], inputs[2]); + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input( + w_pre, encoder, s, mode_, bits_, group_size_, global_scale_w) + : std::make_tuple( + ensure_contiguous(w_pre, encoder, s), + ensure_contiguous(inputs[base_size - 1], encoder, s)); - out.set_data(cu::malloc_async(out.nbytes(), encoder)); + // Reroute to qmm when: no support in cuBLAS, or doing GEMV. + bool can_use_cublas = + (mode_ == QuantizationMode::Nvfp4 || mode_ == QuantizationMode::Mxfp8) && + (device.compute_capability_major() >= 10); + int M = x_pre.shape(-2); + bool use_qmm = (!can_use_cublas) || (M == 1); - auto out_dtype = out.dtype(); + if (use_qmm) { + array x = quantize_dequantize_input( + x_pre, global_scale_x, bits_, group_size_, encoder, s); + if (M < 8) { + qmv(x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + bits_, + group_size_, + mode_, + encoder); + } else { + qmm_naive( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + std::nullopt, + std::nullopt, + out, + true, // transpose + bits_, + group_size_, + mode_, + encoder); + } + return; + } - int M = x_q.shape(-2); - int N = w_q.shape(-2); // always transposed - int K_packed = x_q.shape(-1); - int K = K_packed * (32 / bits_); + // Quantize activation. + auto [x_q, scales_x] = quantize_input( + x_pre, encoder, s, mode_, bits_, group_size_, global_scale_x); - // Repack scales from linear to tiled layout for tensor cores - array scale_x = pad_and_swizzle_scales(scale_x_pre, encoder, s); - array scale_w = pad_and_swizzle_scales(scale_w_pre, encoder, s); + int N = w_q.shape(-2); // transposed + int K = x_q.shape(-1) * (32 / bits_); bool x_transposed = false; bool w_transposed = true; // always transposed int64_t lda = K; int64_t ldb = K; + // Repack scales to tiled layout for tensor cores + scales_x = pad_and_swizzle_scales(scales_x, encoder, s); + scales_w = pad_and_swizzle_scales(scales_w, encoder, s); + + GemmScalars scalars; + if (has_global_scales) { + scalars = create_nvfp4_scalars(*global_scale_x, *global_scale_w, encoder); + } + qqmm_impl( encoder, M, @@ -149,10 +200,69 @@ void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { out, x_q, w_q, - scale_x, - scale_w, - out_dtype, - mode_); + scales_x, + scales_w, + mode_, + scalars); +} + +void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("QQMatmul::eval_gpu"); + + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + const array& lhs_indices = ensure_row_contiguous(inputs[2], encoder, s); + const array& rhs_indices = ensure_row_contiguous(inputs[3], encoder, s); + + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + + // - 4 inputs: x, w, lhs_indices, rhs_indices (non-quantized w) + // - 5 inputs: x, w, lhs_indices, rhs_indices, scales_w (quantized w) + bool w_quantized = (w_pre.dtype() == uint32); + int base_size = w_quantized ? 5 : 4; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; + } + + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input( + w_pre, encoder, s, mode_, bits_, group_size_, global_scale_w) + : std::make_tuple( + ensure_contiguous(w_pre, encoder, s), + ensure_contiguous(inputs[base_size - 1], encoder, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, bits_, group_size_, encoder, s); + + // Reroute to qmm. + qmm_naive( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + true, // transpose + bits_, + group_size_, + mode_, + encoder); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp index dd9407dcd..ccf736cf6 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.cpp @@ -19,15 +19,10 @@ void qqmm_impl( const array& b, const array& a_scale, const array& b_scale, - Dtype out_dtype, QuantizationMode mode, - float alpha) { - // Invoke CublasQQMM + const GemmScalars& scalars) { std::string qmode = quantization_mode_to_string(mode); - // Currently only supports non-batched QQMM operations - // that covers all use cases for training, we will just collapse (batch, - // seq_len) into (tokens) CublasQQMM qqmm( encoder.device(), a_transposed, @@ -41,10 +36,22 @@ void qqmm_impl( 1, // batch_count 0, // a_batch_stride 0, // b_batch_stride - out_dtype, + out.dtype(), qmode); - qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha); + if (scalars.has_values()) { + qqmm.run( + encoder, + out, + a, + b, + a_scale, + b_scale, + *scalars.alpha_device, + *scalars.beta_device); + } else { + qqmm.run(encoder, out, a, b, a_scale, b_scale); + } } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.h index 7288e2fd7..ab2b74c19 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_impl.h @@ -1,10 +1,22 @@ -// Copyright © 2026 Apple Inc. +// Copyright © 2025 Apple Inc. #pragma once #include "mlx/backend/cuda/device.h" #include "mlx/primitives.h" +#include + namespace mlx::core { + +struct GemmScalars { + std::optional alpha_device; + std::optional beta_device; + + bool has_values() const { + return alpha_device.has_value(); + } +}; + void qqmm_impl( cu::CommandEncoder& encoder, int M, @@ -19,8 +31,7 @@ void qqmm_impl( const array& b, const array& a_scale, const array& b_scale, - Dtype out_dtype, QuantizationMode mode, - float alpha = 1.0f); + const GemmScalars& scalars = {}); } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu index c8764709b..81d256be2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.cu @@ -8,13 +8,6 @@ namespace mlx::core { -namespace cg = cooperative_groups; - -constexpr int TILE_ROWS = 128; -constexpr int TILE_COLS = 4; -constexpr int TILES_PER_LANE = 1; -constexpr int LANES_PER_BLOCK = 32; - // To pass scales to tensor cores, they need to be repacked into a tiled layout // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout // Tiled layout for scale factors is very well described in CUTLASS @@ -48,6 +41,14 @@ constexpr int LANES_PER_BLOCK = 32; // [252, 253, 254, 255], // [380, 381, 382, 383], // [508, 509, 510, 511]]]]], +namespace cu { + +constexpr int TILE_ROWS = 128; +constexpr int TILE_COLS = 4; +constexpr int TILES_PER_LANE = 1; +constexpr int LANES_PER_BLOCK = 32; + +namespace cg = cooperative_groups; inline std::tuple get_swizzle_launch_args( size_t M_swizzled, @@ -68,7 +69,17 @@ inline std::tuple get_swizzle_launch_args( return std::make_tuple(grid, block); } -namespace cu { +__global__ void compute_qqmm_pointers( + float* alpha_out, + float* beta_out, + const float* tensor_amax_x, + const float* tensor_amax_w) { + // Compute alpha = tensor_amax_x * tensor_amax_w / (448 * 6)^2 + constexpr float inv_scale_sq = + 1.0f / (F8E4M3_MAX * F4E2M1_MAX * F8E4M3_MAX * F4E2M1_MAX); + *alpha_out = (*tensor_amax_x) * (*tensor_amax_w) * inv_scale_sq; + *beta_out = 0.0f; +} __global__ void swizzle_scales( const uint8_t* scales_linear, @@ -210,12 +221,11 @@ void swizzle_scales( size_t output_cols = scales_tiled.shape(-1); auto [num_blocks, block_dims] = - get_swizzle_launch_args(output_rows, output_cols); + cu::get_swizzle_launch_args(output_rows, output_cols); enc.add_kernel_node( cu::swizzle_scales, num_blocks, block_dims, - 0, gpu_ptr(scales), gpu_ptr(scales_tiled), input_rows, @@ -224,4 +234,24 @@ void swizzle_scales( output_cols); } +void compute_qqmm_pointers( + array& alpha_out, + array& beta_out, + const array& tensor_amax_x, + const array& tensor_amax_w, + cu::CommandEncoder& enc) { + enc.set_input_array(tensor_amax_x); + enc.set_input_array(tensor_amax_w); + enc.set_output_array(alpha_out); + enc.set_output_array(beta_out); + enc.add_kernel_node( + cu::compute_qqmm_pointers, + dim3(1), + dim3(1), + gpu_ptr(alpha_out), + gpu_ptr(beta_out), + gpu_ptr(tensor_amax_x), + gpu_ptr(tensor_amax_w)); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.h index 0a9a78f70..fba9ac9d9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/qqmm_utils.h @@ -27,4 +27,36 @@ void swizzle_scales( cu::CommandEncoder& enc, const Stream& s); +inline array pad_and_swizzle_scales( + const array& scale, + cu::CommandEncoder& encoder, + const Stream& s) { + // Compute padded dimensions for full tiles (128 rows × 4 cols) + auto [pad_outer, pad_inner] = + get_padded_scale_dims(scale.shape(-2), scale.shape(-1)); + // cuBLAS requirements for scale factor layout: + // 1. Dimensions must be padded to full tiles (128 rows × 4 cols) + // 2. Out-of-bounds values must be filled with zeros + // 3. Starting addresses must be 16-byte aligned + // https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout + // Note: cu::malloc_async already provides 256-byte alignment + array scale_tiled( + cu::malloc_async(pad_outer * pad_inner, encoder), + Shape{pad_outer, pad_inner}, + scale.dtype()); + swizzle_scales(scale, scale_tiled, encoder, s); + + encoder.add_temporary(scale_tiled); + return scale_tiled; +} + +// Compute alpha = tensor_amax_x * tensor_amax_w / (448 * 6)^2 +// Allocate beta zero on device as well +void compute_qqmm_pointers( + array& alpha_out, + array& beta_out, + const array& tensor_amax_x, + const array& tensor_amax_w, + cu::CommandEncoder& enc); + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp index 30662d55a..4d25f3c3e 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.cpp @@ -2,8 +2,9 @@ #include "mlx/backend/cuda/quantized/quantized.h" #include "mlx/backend/cuda/device.h" -#include "mlx/backend/cuda/quantized/qmv.h" +#include "mlx/backend/cuda/quantized/qmm/qmm.h" #include "mlx/backend/cuda/quantized/quantized_utils.h" +#include "mlx/dtype_utils.h" #include "mlx/fast_primitives.h" #include "mlx/primitives.h" @@ -14,34 +15,266 @@ namespace mlx::core { void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { nvtx3::scoped_range r("QuantizedMatmul::eval_gpu"); auto& s = stream(); - auto& d = cu::device(s.device); - auto& enc = d.get_command_encoder(s); - - out.set_data(cu::malloc_async(out.nbytes(), enc)); - - // Make sure the last two dims of x and w, s, b are contiguous. This should - // be relaxed for x. - array x = ensure_row_contiguous_matrix(inputs[0], enc, s); - array w = ensure_row_contiguous_matrix(inputs[1], enc, s); - array scales = ensure_row_contiguous_matrix(inputs[2], enc, s); - std::optional biases = std::nullopt; - if (inputs.size() == 4) { - biases = ensure_row_contiguous_matrix(inputs[3], enc, s); + auto& encoder = cu::get_command_encoder(s); + + array x = ensure_row_contiguous(inputs[0], encoder, s); + const array& w = inputs[1]; + const array& scales = inputs[2]; + std::optional biases; + if (inputs.size() > 3) { + biases = inputs[3]; } - bool non_batched = w.ndim() == 2 && x.flags().row_contiguous; + auto supports = [&](auto&& f) { + return f( + x, + w, + scales, + biases, + out, + transpose_, + bits_, + group_size_, + mode_, + encoder.device()); + }; + bool can_use_qmm_sm90 = supports(supports_qmm_sm90); + bool can_use_qmm_sm80 = supports(supports_qmm_sm80); + bool can_use_qmm_naive = supports(supports_qmm_naive); + bool can_use_fp_qmv = supports(supports_fp_qmv); + bool can_use_qmv = supports(supports_qmv) || can_use_fp_qmv; + + auto call_qmm_sm90 = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_sm90(x, w, scales, *biases, out, bits_, group_size_, encoder, s); + }; + auto call_qmm_sm80 = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_sm80( + x, + w, + scales, + biases, + std::nullopt, + std::nullopt, + out, + bits_, + group_size_, + mode_, + encoder); + }; + auto call_qmm_naive = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_naive( + x, + w, + scales, + biases, + std::nullopt, + std::nullopt, + std::nullopt, + out, + transpose_, + bits_, + group_size_, + mode_, + encoder); + }; + auto call_qmv = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + if (can_use_fp_qmv) { + fp_qmv(x, w, scales, out, bits_, group_size_, encoder, s); + } else { + qmv(x, + w, + scales, + biases, + std::nullopt, + out, + bits_, + group_size_, + mode_, + encoder); + } + }; + + int M = out.ndim() > 1 ? out.shape(-2) : 1; + int N = out.shape(-1); int K = x.shape(-1); - int M = non_batched ? x.size() / K : x.shape(-2); + int B = out.size() / (M * N); + + if (can_use_qmm_sm90) { + if (can_use_qmv && (M == 1 && B == 1 && N <= 16384 && K <= 16384)) { + call_qmv(); + } else { + call_qmm_sm90(); + } + return; + } + + if (can_use_qmm_sm80) { + if (can_use_qmv && (M * B < 8)) { + call_qmv(); + } else { + call_qmm_sm80(); + } + return; + } + + if (can_use_qmm_naive) { + if (can_use_qmv && (M * B < 8)) { + call_qmv(); + } else { + call_qmm_naive(); + } + return; + } + + if (can_use_qmv) { + call_qmv(); + return; + } + + throw std::runtime_error( + fmt::format( + "[quantized_matmul] No implementation for " + "problem shape: {}x{}x{}x{}, transpose: {}, " + "activation: {}, bits: {}, group size: {}, mode: \"{}\".", + M, + N, + K, + B, + transpose_, + dtype_to_string(x.dtype()), + bits_, + group_size_, + quantization_mode_to_string(mode_))); +} + +void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("GatherQMM::eval_gpu"); + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + array x = ensure_row_contiguous(inputs[0], encoder, s); + const array& w = inputs[1]; + const array& scales = inputs[2]; + std::optional biases; + if (inputs.size() == 6) { + biases = inputs[3]; + } + array lhs_indices = + ensure_row_contiguous(inputs[inputs.size() - 2], encoder, s); + array rhs_indices = + ensure_row_contiguous(inputs[inputs.size() - 1], encoder, s); + + int M = out.ndim() > 1 ? out.shape(-2) : 1; int N = out.shape(-1); + int K = x.shape(-1); + int B = out.size() / (M * N); + + auto supports = [&](auto&& f) { + return f( + x, + w, + scales, + biases, + out, + transpose_, + bits_, + group_size_, + mode_, + encoder.device()); + }; + bool can_use_qmm_sm80 = supports(supports_qmm_sm80); + bool can_use_qmm_naive = supports(supports_qmm_naive); + bool can_use_qmv = supports(supports_qmv); + + auto call_qmm_sm80 = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_sm80( + x, + w, + scales, + biases, + lhs_indices, + rhs_indices, + out, + bits_, + group_size_, + mode_, + encoder); + }; + auto call_qmm_naive = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + qmm_naive( + x, + w, + scales, + biases, + std::nullopt, + lhs_indices, + rhs_indices, + out, + transpose_, + bits_, + group_size_, + mode_, + encoder); + }; + auto call_qmv = [&]() { + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + gather_qmv( + x, + w, + scales, + biases, + lhs_indices, + rhs_indices, + out, + bits_, + group_size_, + mode_, + encoder); + }; - if (M > 8 || !transpose_ || mode_ == QuantizationMode::Affine) { - throw std::runtime_error("QMM NYI"); + if (can_use_qmm_sm80) { + if (can_use_qmv && (M * B < 8)) { + call_qmv(); + } else { + call_qmm_sm80(); + } + return; } - if (transpose_) { - fp_qmv(w, scales, x, out, bits_, group_size_, M, N, K, enc); + if (can_use_qmm_naive) { + if (can_use_qmv && (M * B < 8)) { + call_qmv(); + } else { + call_qmm_naive(); + } return; } + + if (can_use_qmv) { + call_qmv(); + return; + } + + throw std::runtime_error( + fmt::format( + "[gather_qmm] No implementation for " + "problem shape: {}x{}x{}x{}, transpose: {}, " + "activation: {}, bits: {}, group size: {}, mode: \"{}\".", + M, + N, + K, + B, + transpose_, + dtype_to_string(x.dtype()), + bits_, + group_size_, + quantization_mode_to_string(mode_))); } void fast::Quantize::eval_gpu( @@ -49,9 +282,7 @@ void fast::Quantize::eval_gpu( std::vector& outputs) { nvtx3::scoped_range r("Quantize::eval_gpu"); auto& s = stream(); - auto& d = cu::device(s.device); - auto& enc = d.get_command_encoder(s); - + auto& enc = cu::get_command_encoder(s); if (dequantize_) { auto wq = ensure_row_contiguous(inputs[0], enc, s); auto scales = ensure_row_contiguous(inputs[1], enc, s); @@ -63,7 +294,12 @@ void fast::Quantize::eval_gpu( auto biases = ensure_row_contiguous(inputs[2], enc, s); affine_dequantize(wq, scales, biases, w, group_size_, bits_, enc, s); } else { - fp_dequantize(wq, scales, w, group_size_, bits_, enc, s); + // 0 -- xq, 1 -- scales, 2 -- could be global scale for nvfp4 + bool use_global_scale = + mode_ == QuantizationMode::Nvfp4 && inputs.size() > 2; + std::optional global_scale = + use_global_scale ? std::make_optional(inputs[2]) : std::nullopt; + fp_dequantize(wq, scales, w, group_size_, bits_, global_scale, enc, s); } } else { auto w = ensure_contiguous(inputs[0], enc, s); @@ -72,12 +308,17 @@ void fast::Quantize::eval_gpu( wq.set_data(cu::malloc_async(wq.nbytes(), enc)); scales.set_data(cu::malloc_async(scales.nbytes(), enc)); + if (mode_ == QuantizationMode::Affine) { auto& biases = outputs[2]; biases.set_data(cu::malloc_async(biases.nbytes(), enc)); affine_quantize(w, wq, scales, biases, group_size_, bits_, enc, s); } else { - fp_quantize(w, wq, scales, group_size_, bits_, enc, s); + bool use_global_scale = + mode_ == QuantizationMode::Nvfp4 && inputs.size() > 1; + std::optional global_scale = + use_global_scale ? std::make_optional(inputs[1]) : std::nullopt; + fp_quantize(w, wq, scales, group_size_, bits_, global_scale, enc, s); } } } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.h b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.h index 744a12f5c..f15c0f76e 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized.h @@ -1,5 +1,6 @@ // Copyright © 2025 Apple Inc. +#include #include "mlx/backend/cuda/device.h" namespace mlx::core { @@ -30,6 +31,7 @@ void fp_quantize( array& scales, int group_size, int bits, + const std::optional& global_scale, cu::CommandEncoder& enc, const Stream& s); @@ -39,6 +41,7 @@ void fp_dequantize( array& w, int group_size, int bits, + const std::optional& global_scale, cu::CommandEncoder& enc, const Stream& s); @@ -47,6 +50,7 @@ void fp_quantize_dequantize( array& what, int group_size, int bits, + const std::optional& global_scale, cu::CommandEncoder& enc, const Stream& s); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh deleted file mode 100644 index 93d83292d..000000000 --- a/Source/Cmlx/mlx/mlx/backend/cuda/quantized/quantized_utils.cuh +++ /dev/null @@ -1,88 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#include -#include - -namespace mlx::core { - -namespace cu { - -inline __device__ float4 dequant_fp8(uint32_t bits) { - auto out = *(__nv_fp8x4_e4m3*)(&bits); - return out.operator float4(); -} - -inline __device__ float4 dequant_fp4(uint16_t bits) { - auto out = *(__nv_fp4x4_e2m1*)(&bits); - return out.operator float4(); -} - -template -inline constexpr __device__ short get_pack_factor() { - return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits); -} - -template -inline constexpr __device__ short get_bytes_per_pack() { - constexpr int power_of_2_bits = (bits & (bits - 1)) == 0; - return power_of_2_bits ? (wsize / 8) : (bits == 5 ? 5 : 3); -} - -template -__device__ __forceinline__ void abs_max_x2(T& out, const T& x1, const T& x2) { - if constexpr ( - (std::is_same::value) || - (std::is_same::value)) { - T a = x1; - T b = x2; - out = __hmax2(__habs2(a), __habs2(b)); - } else if constexpr (std::is_same::value) { - float2 a = x1; - float2 b = x2; - out.x = fmaxf(fabsf(a.x), fabsf(b.x)); - out.y = fmaxf(fabsf(a.y), fabsf(b.y)); - } -} - -} // namespace cu - -template -void dispatch_groups(int group_size, F&& f) { - switch (group_size) { - case 32: - f(std::integral_constant{}); - break; - case 64: - f(std::integral_constant{}); - break; - case 128: - f(std::integral_constant{}); - break; - } -} - -template -void dispatch_bits(int bits, F&& f) { - switch (bits) { - case 2: - f(std::integral_constant{}); - break; - case 3: - f(std::integral_constant{}); - break; - case 4: - f(std::integral_constant{}); - break; - case 5: - f(std::integral_constant{}); - break; - case 6: - f(std::integral_constant{}); - break; - case 8: - f(std::integral_constant{}); - break; - } -} - -} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/random.cu b/Source/Cmlx/mlx/mlx/backend/cuda/random.cu index e9a310d26..1a485134d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/random.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/random.cu @@ -49,7 +49,7 @@ __global__ void rbitsc( uint8_t* out, dim3 grid_dims, bool odd, - uint32_t bytes_per_key) { + uint64_t bytes_per_key) { auto grid = cg::this_grid(); uint32_t thread_index = grid.thread_rank(); uint32_t index_x = thread_index % grid_dims.x; @@ -89,7 +89,7 @@ __global__ void rbits( uint8_t* out, dim3 grid_dims, bool odd, - uint32_t bytes_per_key, + uint64_t bytes_per_key, int32_t ndim, const __grid_constant__ Shape key_shape, const __grid_constant__ Strides key_strides) { @@ -176,7 +176,6 @@ void RandomBits::eval_gpu(const std::vector& inputs, array& out) { cu::rbitsc, grid, block, - 0, gpu_ptr(keys), gpu_ptr(out), grid_dims, @@ -187,7 +186,6 @@ void RandomBits::eval_gpu(const std::vector& inputs, array& out) { cu::rbits, grid, block, - 0, gpu_ptr(keys), gpu_ptr(out), grid_dims, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce.cu index 269efc034..833fdbf66 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce.cu @@ -15,13 +15,21 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); array in = inputs[0]; - // Make sure no identity reductions trickle down here. assert(!axes_.empty()); - assert(out.size() != in.size()); auto& s = stream(); auto& encoder = cu::get_command_encoder(s); + // When all the reduced axes have size 1 at runtime, which can happen with + // shapeless compilation, the reduction is the identity so just cast-copy + // the input to the output. + if (in.size() > 0 && out.size() == in.size()) { + CopyType ctype = + in.flags().contiguous ? CopyType::Vector : CopyType::General; + copy_gpu(in, out, ctype, s); + return; + } + if (in.size() == 0) { init_reduce(encoder, in, out, reduce_type_); return; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/all_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/all_reduce.cu index 962e80d4f..167a282af 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/all_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/all_reduce.cu @@ -68,12 +68,12 @@ void all_reduce( out.set_data(cu::malloc_async(out.nbytes(), encoder)); - auto get_args = [](int size, int N) { - int threads = std::min(512, (size + N - 1) / N); + auto get_args = [](size_t size, size_t N) { + int threads = + static_cast(std::min(size_t{512}, cuda::ceil_div(size, N))); threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; - int reductions_per_step = threads * N; - size_t steps_needed = - (size + reductions_per_step - 1) / reductions_per_step; + size_t reductions_per_step = threads * N; + size_t steps_needed = cuda::ceil_div(size, reductions_per_step); int blocks; if (steps_needed < 32) { @@ -88,7 +88,7 @@ void all_reduce( blocks = 1024; } - size_t steps_per_block = (steps_needed + blocks - 1) / blocks; + size_t steps_per_block = cuda::ceil_div(steps_needed, blocks); size_t block_step = steps_per_block * reductions_per_step; return std::make_tuple(blocks, threads, block_step); @@ -98,9 +98,7 @@ void all_reduce( size_t block_step; size_t insize = in.size(); Dtype dt = in.dtype(); - - // Cub doesn't like const pointers for load (sigh). - void* indata = const_cast(gpu_ptr(in)); + void* indata = gpu_ptr(in); // Large array so allocate an intermediate and accumulate there std::tie(blocks, threads, block_step) = get_args(insize, N_READS); @@ -120,8 +118,7 @@ void all_reduce( kernel, blocks, threads, - 0, - static_cast(indata), + indata, gpu_ptr(intermediate), block_step, insize); @@ -144,14 +141,7 @@ void all_reduce( using U = typename cu::ReduceResult::type; auto kernel = cu::all_reduce; encoder.add_kernel_node( - kernel, - blocks, - threads, - 0, - static_cast(indata), - gpu_ptr(out), - block_step, - insize); + kernel, blocks, threads, indata, gpu_ptr(out), block_step, insize); }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu index e33551d86..8a6327bf3 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/col_reduce.cu @@ -282,8 +282,6 @@ void col_reduce_looped( using OP = MLX_GET_TYPE(reduce_type_tag); using T = cuda_type_t; using U = typename cu::ReduceResult::type; - // Cub doesn't like const pointers for vectorized loads. (sigh) - T* indata = const_cast(gpu_ptr(in)); constexpr int N_READS = 4; constexpr int BM = 32; @@ -296,8 +294,7 @@ void col_reduce_looped( kernel, grid, blocks, - 0, - indata, + gpu_ptr(in), gpu_ptr(out), static_cast(args), out.size() / args.reduction_stride); @@ -334,7 +331,6 @@ void col_reduce_small( kernel, grid, block, - 0, gpu_ptr(in), gpu_ptr(out), static_cast(args), @@ -391,8 +387,6 @@ void col_reduce_two_pass( using OP = MLX_GET_TYPE(reduce_type_tag); using T = cuda_type_t; using U = typename cu::ReduceResult::type; - // Cub doesn't like const pointers for vectorized loads. (sigh) - T* indata = const_cast(gpu_ptr(in)); constexpr int N_READS = 4; constexpr int BM = 32; @@ -405,8 +399,7 @@ void col_reduce_two_pass( kernel, grid, blocks, - 0, - indata, + gpu_ptr(in), gpu_ptr(intermediate), static_cast(args), out.size() / args.reduction_stride); @@ -444,7 +437,6 @@ void col_reduce_two_pass( kernel, grid, blocks, - 0, gpu_ptr(intermediate), gpu_ptr(out), second_args, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/init_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/init_reduce.cu index df6f90664..463d5d5a9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/init_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/init_reduce.cu @@ -31,6 +31,10 @@ void init_reduce( out.set_data(cu::malloc_async(out.nbytes(), encoder)); } + if (out.size() == 0) { + return; + } + encoder.set_output_array(out); dispatch_all_types(in.dtype(), [&](auto type_tag) { dispatch_reduce_ops(reduce_type, [&](auto reduce_type_tag) { @@ -41,8 +45,7 @@ void init_reduce( dim3 grid = get_2d_grid_dims(out.shape(), out.strides()); dim3 block(grid.x < 1024 ? grid.x : 1024, 1, 1); grid.x = (grid.x + 1023) / 1024; - encoder.add_kernel_node( - kernel, grid, block, 0, gpu_ptr(out), out.size()); + encoder.add_kernel_node(kernel, grid, block, gpu_ptr(out), out.size()); }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/row_reduce.cu b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/row_reduce.cu index ea99e1132..1dd40d14b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/reduce/row_reduce.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/reduce/row_reduce.cu @@ -268,10 +268,15 @@ void row_reduce_simple( kernel = cu::row_reduce_simple; } - T* indata = const_cast(gpu_ptr(in)); int size = plan.shape.back(); encoder.add_kernel_node( - kernel, grid, block, 0, indata, gpu_ptr(out), out.size(), size); + kernel, + grid, + block, + gpu_ptr(in), + gpu_ptr(out), + out.size(), + size); }); }); } @@ -315,7 +320,7 @@ void row_reduce_looped( }); encoder.add_kernel_node( - kernel, grid, block, 0, gpu_ptr(in), gpu_ptr(out), args); + kernel, grid, block, gpu_ptr(in), gpu_ptr(out), args); }); }); } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/rms_norm.cu b/Source/Cmlx/mlx/mlx/backend/cuda/rms_norm.cu index a9fc79f19..e4c0702a9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/rms_norm.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/rms_norm.cu @@ -3,6 +3,7 @@ #include "mlx/backend/cuda/device.h" #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/cuda/reduce/reduce.cuh" +#include "mlx/backend/cuda/steel/utils.cuh" #include "mlx/backend/gpu/copy.h" #include "mlx/dtype_utils.h" #include "mlx/fast_primitives.h" @@ -51,7 +52,18 @@ struct BlockBroadcastReduce { } }; -template +// xs and ws stay in registers +// Each thread does N_CHUNKS_THREAD vectorised interleaved loads of width +// N_READS x, w, out in majority of cases should be 16 bytes aligned -> using +// unsafe_load_vector if not aligned -> fall back to load_vector we do it like +// this because of the register presure: load_vector allocates registers for +// loop load fall back and occupancy drops to 30-40% -> x2 slow down +template < + typename T, + int BLOCK_SIZE, + int N_CHUNKS_THREAD, + bool ALIGNED, + int N_READS = 8> __global__ void rms_norm_small( const T* x, const T* w, @@ -63,40 +75,61 @@ __global__ void rms_norm_small( auto grid = cg::this_grid(); auto block = cg::this_thread_block(); - using BlockReduceT = BlockBroadcastReduce; + using BlockReduceT = BlockBroadcastReduce; __shared__ typename BlockReduceT::TempStorage temp; - auto row = - (grid.block_rank() * block.dim_threads().y) + block.thread_index().y; + auto row = grid.block_rank(); if (row >= n_rows) { return; } x += row * axis_size; out += row * axis_size; - // Normalizer. + AlignedVector xn[N_CHUNKS_THREAD]; + float normalizer = 0; auto index = block.thread_index().x; - auto xn = load_vector(x, index, axis_size, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float t = static_cast(xn[i]); - normalizer += t * t; + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + if constexpr (ALIGNED) { + xn[i] = unsafe_load_vector(x, offset); + } else { + xn[i] = load_vector(x, offset, axis_size, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; ++j) { + float t = static_cast(xn[i][j]); + normalizer += t * t; + } } - normalizer = BlockReduceT{block, temp}.Sum(normalizer); normalizer = rsqrt(normalizer / axis_size + eps); - // Outputs. - auto wn = load_vector(w, index, axis_size, w_stride, T(0)); #pragma unroll - for (int i = 0; i < N_READS; ++i) { - float y = static_cast(xn[i]) * normalizer; - xn[i] = wn[i] * static_cast(y); + for (int i = 0; i < N_CHUNKS_THREAD; i++) { + int offset = BLOCK_SIZE * i + index; + AlignedVector wn; + if constexpr (ALIGNED) { + wn = unsafe_load_vector(w, offset); + } else { + wn = load_vector(w, offset, axis_size, w_stride, T(0)); + } +#pragma unroll + for (int j = 0; j < N_READS; j++) { + float y = static_cast(xn[i][j]) * normalizer; + xn[i][j] = wn[j] * static_cast(y); + } + if constexpr (ALIGNED) { + unsafe_store_vector(out, offset, xn[i]); + } else { + store_vector(out, offset, xn[i], axis_size); + } } - store_vector(out, index, xn, axis_size); } +// TODO: load x to the shared memory and reload from shared memory after the +// reduction template __global__ void rms_norm( const T* x, @@ -148,7 +181,7 @@ template < int BLOCK_DIM, int REDUCE_DIM, int N_READS = 4> -__global__ void rms_norm_vjp_small( +__global__ void rms_norm_vjp_small_fallback( const T* x, const T* w, const T* g, @@ -210,6 +243,154 @@ __global__ void rms_norm_vjp_small( } } +template +__global__ void rms_norm_vjp_small( + const T* x, + const T* w, + const T* g, + T* gx, + float* gw, // accumulate in float always + float eps, + int32_t axis_size, + int32_t n_rows, + int64_t w_stride) { + // persistent kernel, numblocks = number of sms * 2; + // each block is responsible for a row. 128*2 blocks = 256 stride. + // we pipeline loads and computation in shared memory: + // loading the row while doing dw, dx computation for another row + + auto grid = cg::this_grid(); + auto block = cg::this_thread_block(); + + using BlockReduceF2 = BlockBroadcastReduce; + __shared__ typename BlockReduceF2::TempStorage temp; + + // double buffering + constexpr int STAGES = 2; + const int num_blocks = static_cast(grid.num_blocks()); + const int num_tiles = cuda::ceil_div(n_rows, num_blocks); + auto tid = block.thread_index().x; + auto bid = grid.block_rank(); + constexpr int buffer_size = N_READS * BLOCK_SIZE * N_CHUNKS; // axis_size + + // shared memory is dymanic because we need > 48 kb + extern __shared__ char smem_raw[]; + T* smem_x = reinterpret_cast(smem_raw); + T* smem_dy = smem_x + STAGES * buffer_size; + constexpr int BYTES_PER_READ = N_READS * static_cast(sizeof(T)); + constexpr int STAGE_BYTES = buffer_size * static_cast(sizeof(T)); + uint32_t smem_x_addr = __cvta_generic_to_shared(smem_x); + uint32_t smem_dy_addr = __cvta_generic_to_shared(smem_dy); + + AlignedVector wn[N_CHUNKS]; + AlignedVector xn[N_CHUNKS]; + AlignedVector gn[N_CHUNKS]; + // gw_block is summed over every row this block visits + AlignedVector gw_block[N_CHUNKS]; + // zero init + for (int j = 0; j < N_CHUNKS; j++) { + for (int k = 0; k < N_READS; k++) { + gw_block[j][k] = 0.f; + } + } + // shift global pointer for each block + x += axis_size * bid; + g += axis_size * bid; + gx += axis_size * bid; + gw += axis_size * bid; + + // prefetch first row for each block + // and load weights to registers in the same loop + for (int j = 0; j < N_CHUNKS; j++) { + int offset = (j * BLOCK_SIZE + tid) * N_READS; + if (bid < n_rows) { + cp_async( + smem_x_addr + BYTES_PER_READ * (j * BLOCK_SIZE + tid), &x[offset]); + cp_async( + smem_dy_addr + BYTES_PER_READ * (j * BLOCK_SIZE + tid), &g[offset]); + } + wn[j] = load_vector( + w, j * BLOCK_SIZE + tid, axis_size, w_stride, T(0)); + } + // commit all N_CHUNKS 128 byte loads for the first row + cp_async_commit(); + // pipelineing + for (int tile = 0; tile < num_tiles; tile++) { + x += axis_size * num_blocks; + g += axis_size * num_blocks; + + int next = tile + 1; + int index = next % STAGES; + int64_t next_row = + static_cast(bid) + static_cast(next) * num_blocks; + + if (next < num_tiles && next_row < n_rows) { + for (int j = 0; j < N_CHUNKS; j++) { + int offset = (j * BLOCK_SIZE + tid) * N_READS; + cp_async( + smem_x_addr + index * STAGE_BYTES + + BYTES_PER_READ * (j * BLOCK_SIZE + tid), + &x[offset]); + cp_async( + smem_dy_addr + index * STAGE_BYTES + + BYTES_PER_READ * (j * BLOCK_SIZE + tid), + &g[offset]); + } + } + cp_async_commit(); // always commit, empty at the tail + cp_async_wait<1>(); // always wait for 1, the tail is empty + + int64_t cur_row = + static_cast(bid) + static_cast(tile) * num_blocks; + if (cur_row < n_rows) { + // load x and g from shared to registers + // compute the reduction per row + float2 factors = {}; + for (int j = 0; j < N_CHUNKS; j++) { + xn[j] = unsafe_load_vector( + smem_x + (tile % STAGES) * buffer_size, j * BLOCK_SIZE + tid); + gn[j] = unsafe_load_vector( + smem_dy + (tile % STAGES) * buffer_size, j * BLOCK_SIZE + tid); + for (int k = 0; k < N_READS; k++) { + float t = static_cast(xn[j][k]); + float wi = wn[j][k]; + float gi = gn[j][k]; + float wg = wi * gi; + factors = plus_f2(factors, {wg * t, t * t}); + } + } + factors = BlockReduceF2{block, temp}.Reduce(factors, plus_f2, {}); + float meangwx = factors.x / axis_size; + float normalizer = rsqrt(factors.y / axis_size + eps); + float normalizer3 = normalizer * normalizer * normalizer; + + // we store dx after processing each row, accumulate dw + T* gx_row = gx + static_cast(tile) * num_blocks * axis_size; + for (int j = 0; j < N_CHUNKS; j++) { + int offset = j * BLOCK_SIZE + tid; + for (int k = 0; k < N_READS; k++) { + float xi = static_cast(xn[j][k]); + float wi = wn[j][k]; + float gi = gn[j][k]; + if constexpr (HAS_W) { + gw_block[j][k] += gi * xi * normalizer; + } + xn[j][k] = + static_cast(normalizer * wi * gi - xi * meangwx * normalizer3); + } + store_vector(gx_row, offset, xn[j], axis_size); + } + } + } + // store this block's fp32 partial + if constexpr (HAS_W) { + for (int j = 0; j < N_CHUNKS; j++) { + int offset = j * BLOCK_SIZE + tid; + store_vector(gw, offset, gw_block[j], axis_size); + } + } +} + template __global__ void rms_norm_vjp( const T* x, @@ -298,7 +479,7 @@ void dispatch_group_dim(int axis_size, F&& f) { } else if (axis_size <= n_per_thread * 32 * 2) { f(std::integral_constant{}, std::integral_constant(), - std::integral_constant()); + std::integral_constant()); } else if (axis_size <= n_per_thread * 32 * 4) { f(std::integral_constant{}, std::integral_constant(), @@ -318,6 +499,42 @@ void dispatch_group_dim(int axis_size, F&& f) { } } +template +void dispatch_chunks(int n_chunks, F&& f) { + auto block_size = std::integral_constant{}; + if (n_chunks <= 1) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 2) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 3) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 4) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 5) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 6) { + f(block_size, std::integral_constant{}); + } else if (n_chunks <= 7) { + f(block_size, std::integral_constant{}); + } else { + f(block_size, std::integral_constant{}); + } +} + +template +void dispatch_num_chunks(int axis_size, F&& f) { + int nvec = (axis_size + N_READS - 1) / N_READS; + if (axis_size <= N_READS * 64) { + f(std::integral_constant{}, std::integral_constant{}); + } else if (nvec % 128 == 0 && nvec / 128 <= 8) { + dispatch_chunks<128>(nvec / 128, f); + } else if (nvec % 64 == 0 && nvec / 64 <= 8) { + dispatch_chunks<64>(nvec / 64, f); + } else { + dispatch_chunks<128>((nvec + 127) / 128, f); + } +} + // TODO: There are duplicate code with backend/metal/normalization.cpp void RMSNorm::eval_gpu( const std::vector& inputs, @@ -365,34 +582,44 @@ void RMSNorm::eval_gpu( dispatch_float_types(out.dtype(), "rms_norm", [&](auto type_tag) { using DataType = cuda_type_t; constexpr int N_READS = 16 / sizeof(DataType); - if (axis_size <= N_READS * 1024) { - dispatch_group_dim( - axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { - constexpr int block_dim = n_groups() * group_dim(); - auto kernel = - cu::rms_norm_small; - auto n_blocks = - (n_rows + groups_per_block() - 1) / groups_per_block(); - encoder.add_kernel_node( - kernel, - n_blocks, - {block_dim, groups_per_block()}, - 0, - gpu_ptr(x), - gpu_ptr(w), - gpu_ptr(out), - eps_, - axis_size, - n_rows, - w_stride); - }); + if (axis_size <= N_READS * 128 * 8) { + dispatch_num_chunks< + N_READS>(axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + constexpr int N_CHUNKS = n_chunks(); + bool aligned = (axis_size == N_READS * BLOCK_SIZE * N_CHUNKS) && + (w_stride == 1) && + (reinterpret_cast(gpu_ptr(x)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(w)) % 16 == 0) && + (reinterpret_cast(gpu_ptr(out)) % 16 == 0); + // MSVC can't find N_READS two lambda levels down, dispatch_bool would + // add that second level + auto kernel = aligned + ? cu::rms_norm_small + : cu::rms_norm_small< + DataType, + BLOCK_SIZE, + N_CHUNKS, + false, + N_READS>; + encoder.add_kernel_node( + kernel, + n_rows, + BLOCK_SIZE, + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(out), + eps_, + axis_size, + n_rows, + w_stride); + }); } else { auto kernel = cu::rms_norm; encoder.add_kernel_node( kernel, n_rows, 1024, - 0, gpu_ptr(x), gpu_ptr(w), gpu_ptr(out), @@ -403,6 +630,45 @@ void RMSNorm::eval_gpu( }); } +template +inline bool use_rmsnorm_vjp_fast( + const array& x, + const array& w, + const array& g, + const array& gx, + bool has_w, + int32_t axis_size, + int64_t w_stride) { + if (!has_w || w_stride != 1) { + return false; + } + if (axis_size != N_READS * BLOCK_SIZE * N_CHUNKS) { + return false; + } + auto aligned = [](const array& a) { + return reinterpret_cast(gpu_ptr(a)) % 16 == 0; + }; + return aligned(x) && aligned(w) && aligned(g) && aligned(gx); +} + +template +inline int rmsnorm_vjp_num_blocks( + Kernel kernel, + const Stream& s, + int block_size, + size_t smem_bytes, + int32_t n_rows) { + int dev = cu::device(s.device).cuda_device(); + int sm_count = 0; + CHECK_CUDA_ERROR( + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev)); + int blocks_per_sm = 1; + CHECK_CUDA_ERROR(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_sm, kernel, block_size, smem_bytes)); + int64_t want = static_cast(sm_count) * std::max(blocks_per_sm, 1); + return std::max(1, static_cast(std::min(want, n_rows))); +} + void RMSNormVJP::eval_gpu( const std::vector& inputs, std::vector& outputs) { @@ -454,6 +720,91 @@ void RMSNormVJP::eval_gpu( int32_t n_rows = x.data_size() / axis_size; int64_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; + bool handled = false; + + if (has_w && cu::device(s.device).compute_capability_major() >= 8) { + dispatch_float_types(gx.dtype(), "rms_norm_vjp", [&](auto type_tag) { + using DataType = cuda_type_t; + constexpr int N_READS = 16 / sizeof(DataType); + dispatch_num_chunks( + axis_size, [&](auto block_size, auto n_chunks) { + constexpr int BLOCK_SIZE = block_size(); + constexpr int N_CHUNKS = n_chunks(); + if (!use_rmsnorm_vjp_fast( + x, w, g, gx, has_w, axis_size, w_stride)) { + return; + } + constexpr int STAGES = 2; + size_t smem_bytes = + size_t(2) * STAGES * axis_size * sizeof(DataType); + int dev = cu::device(s.device).cuda_device(); + int smem_max = 0; + CHECK_CUDA_ERROR(cudaDeviceGetAttribute( + &smem_max, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); + // we should not be here often, is never true on Ampere and later + if (smem_bytes > static_cast(smem_max)) { + return; + } + auto kernel = cu::rms_norm_vjp_small< + DataType, + /*HAS_W=*/true, + BLOCK_SIZE, + N_CHUNKS, + N_READS>; + if (smem_bytes > 48000) { + CHECK_CUDA_ERROR(cudaFuncSetAttribute( + reinterpret_cast(kernel), + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes)); + } + int num_blocks = rmsnorm_vjp_num_blocks( + kernel, s, BLOCK_SIZE, smem_bytes, n_rows); + array gw_temp({num_blocks, axis_size}, float32, nullptr, {}); + gw_temp.set_data(cu::malloc_async(gw_temp.nbytes(), encoder)); + encoder.add_temporary(gw_temp); + encoder.set_input_array(x); + encoder.set_input_array(w); + encoder.set_input_array(g); + encoder.set_output_array(gx); + encoder.set_output_array(gw_temp); + encoder.add_kernel_node_ex( + kernel, + dim3{static_cast(num_blocks)}, + dim3{static_cast(BLOCK_SIZE)}, + {}, // no cluster + static_cast(smem_bytes), + gpu_ptr(x), + gpu_ptr(w), + gpu_ptr(g), + gpu_ptr(gx), + gpu_ptr(gw_temp), + eps_, + axis_size, + n_rows, + w_stride); + + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, + {num_blocks}, + {axis_size}); + if (gw.dtype() == float32) { + col_reduce( + encoder, gw_temp, gw, Reduce::ReduceType::Sum, {0}, plan); + } else { + array gw_f32({axis_size}, float32, nullptr, {}); + col_reduce( + encoder, gw_temp, gw_f32, Reduce::ReduceType::Sum, {0}, plan); + encoder.add_temporary(gw_f32); + copy_gpu(gw_f32, gw, CopyType::General, s); + } + handled = true; + }); + }); + } + if (handled) { + return; + } + // Allocate a temporary to store the gradients for w and allocate the output // gradient accumulators. array gw_temp = @@ -481,7 +832,8 @@ void RMSNormVJP::eval_gpu( axis_size, [&](auto group_dim, auto n_groups, auto groups_per_block) { constexpr int block_dim = group_dim() * n_groups(); - auto kernel = cu::rms_norm_vjp_small< + static_assert(block_dim <= 32 || groups_per_block() == 1); + auto kernel = cu::rms_norm_vjp_small_fallback< DataType, has_w_constant.value, block_dim, @@ -493,7 +845,6 @@ void RMSNormVJP::eval_gpu( kernel, n_blocks, {block_dim, groups_per_block()}, - 0, gpu_ptr(x), gpu_ptr(w), gpu_ptr(g), @@ -511,7 +862,6 @@ void RMSNormVJP::eval_gpu( kernel, n_rows, 1024, - 0, gpu_ptr(x), gpu_ptr(w), gpu_ptr(g), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/rope.cu b/Source/Cmlx/mlx/mlx/backend/cuda/rope.cu index b0106d47a..cd7e68ceb 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/rope.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/rope.cu @@ -106,8 +106,8 @@ __device__ void rope_impl( const int* offset, float inv_freq, float scale, - const cuda::std::array strides, - const cuda::std::array out_strides, + const cuda::std::array strides, + const cuda::std::array out_strides, int64_t offset_stride, int n_head, uint3 pos, @@ -117,29 +117,31 @@ __device__ void rope_impl( auto batch_idx = (pos.z * N) / n_head_up; auto batch_offset = offset[batch_idx * offset_stride]; float L = scale * static_cast(pos.y + batch_offset); - auto mat_idx = batch_idx * n_head + head_idx; // Compute costheta, sintheta float theta = L * inv_freq; float costheta = cos(theta); float sintheta = sin(theta); + size_t in_batch_head = batch_idx * strides[0] + head_idx * strides[1]; + size_t out_batch_head = + batch_idx * out_strides[0] + head_idx * out_strides[1]; + // Compute the input and output indices size_t in_index_1, in_index_2; size_t out_index_1, out_index_2; if (traditional) { - out_index_1 = 2 * pos.x * out_strides[2] + pos.y * out_strides[1] + - mat_idx * out_strides[0]; - out_index_2 = out_index_1 + 1; - in_index_1 = - 2 * pos.x * strides[2] + pos.y * strides[1] + mat_idx * strides[0]; - in_index_2 = in_index_1 + strides[2]; + out_index_1 = + 2 * pos.x * out_strides[3] + pos.y * out_strides[2] + out_batch_head; + out_index_2 = out_index_1 + out_strides[3]; + in_index_1 = 2 * pos.x * strides[3] + pos.y * strides[2] + in_batch_head; + in_index_2 = in_index_1 + strides[3]; } else { - out_index_1 = pos.x * out_strides[2] + pos.y * out_strides[1] + - mat_idx * out_strides[0]; - out_index_2 = out_index_1 + dims.x * out_strides[2]; - in_index_1 = pos.x * strides[2] + pos.y * strides[1] + mat_idx * strides[0]; - in_index_2 = in_index_1 + dims.x * strides[2]; + out_index_1 = + pos.x * out_strides[3] + pos.y * out_strides[2] + out_batch_head; + out_index_2 = out_index_1 + dims.x * out_strides[3]; + in_index_1 = pos.x * strides[3] + pos.y * strides[2] + in_batch_head; + in_index_2 = in_index_1 + dims.x * strides[3]; } for (int i = 0; i < N && head_idx + i < n_head; ++i) { // Read and write the output @@ -156,10 +158,10 @@ __device__ void rope_impl( } out[out_index_1] = static_cast(rx1); out[out_index_2] = static_cast(rx2); - in_index_1 += strides[0]; - in_index_2 += strides[0]; - out_index_1 += out_strides[0]; - out_index_2 += out_strides[0]; + in_index_1 += strides[1]; + in_index_2 += strides[1]; + out_index_1 += out_strides[1]; + out_index_2 += out_strides[1]; } } @@ -170,8 +172,8 @@ __global__ void rope( const int32_t* offset, float scale, float base, - const __grid_constant__ cuda::std::array strides, - const __grid_constant__ cuda::std::array out_strides, + const __grid_constant__ cuda::std::array strides, + const __grid_constant__ cuda::std::array out_strides, int64_t offset_stride, int n_head, uint3 dims) { @@ -207,8 +209,8 @@ __global__ void rope_freqs( const float* freqs, float scale, float base, - const __grid_constant__ cuda::std::array strides, - const __grid_constant__ cuda::std::array out_strides, + const __grid_constant__ cuda::std::array strides, + const __grid_constant__ cuda::std::array out_strides, int64_t offset_stride, int n_head, uint3 dims, @@ -255,8 +257,8 @@ void RoPE::eval_gpu( auto& offset = inputs[1]; auto& out = outputs[0]; - cuda::std::array strides; - cuda::std::array out_strides; + cuda::std::array strides; + cuda::std::array out_strides; bool donated = false; int ndim = in.ndim(); @@ -264,57 +266,52 @@ void RoPE::eval_gpu( int T = in.shape(-2); int D = in.shape(-1); size_t mat_size = T * D; - int dispatch_ndim = ndim; - while (in.shape(-dispatch_ndim) == 1 && dispatch_ndim > 3) { - dispatch_ndim--; - } int N = 1; for (int i = 1; i < (ndim - 2); ++i) { N *= in.shape(i); } - // We apply rope to less that the whole vector so copy to output and then - // apply in-place. - if (dims_ < D) { - donated = true; - auto ctype = - (in.flags().row_contiguous) ? CopyType::Vector : CopyType::General; - copy_gpu(in, out, ctype, s); - strides[0] = mat_size; - strides[1] = out.strides()[ndim - 2]; - strides[2] = out.strides()[ndim - 1]; - } + // if input has < 4 dims or row_contiguous: if we can donate the input, reuse + // the buffer if in is not donatable, we allocate a new buffer for the output + // with the same strides as input if input has > 4 dims and is not + // row_contigous, we copy the input to a new buffer (we should not be here + // often) in case it is partial rotation: we reuse if input is donatable and + // copy if it is not donatable + bool partial = dims_ < D; + bool layout_ok = ndim <= 4 || in.flags().row_contiguous; - // Either copy or apply in-place - else if (in.flags().row_contiguous) { + if (layout_ok) { if (in.is_donatable()) { donated = true; out.copy_shared_buffer(in); + } else if (!partial) { + out.set_data( + cu::malloc_async(in.data_size() * in.itemsize(), encoder), + in.data_size(), + in.strides(), + in.flags()); } else { - out.set_data(cu::malloc_async(out.nbytes(), encoder)); + donated = true; + auto ctype = + in.flags().row_contiguous ? CopyType::Vector : CopyType::General; + copy_gpu(in, out, ctype, s); } - strides[0] = mat_size; - strides[1] = in.strides()[ndim - 2]; - strides[2] = in.strides()[ndim - 1]; - } else if (dispatch_ndim == 3) { - // Handle non-contiguous 3D inputs - out.set_data(cu::malloc_async(out.nbytes(), encoder)); - strides[0] = in.strides()[ndim - 3]; - strides[1] = in.strides()[ndim - 2]; - strides[2] = in.strides()[ndim - 1]; } else { - // Copy non-contiguous > 3D inputs into the output and treat - // input as donated donated = true; copy_gpu(in, out, CopyType::General, s); - strides[0] = mat_size; - strides[1] = out.strides()[ndim - 2]; - strides[2] = out.strides()[ndim - 1]; } - out_strides[0] = mat_size; - out_strides[1] = out.strides()[ndim - 2]; - out_strides[2] = out.strides()[ndim - 1]; + const auto& src_strides = donated ? out.strides() : in.strides(); + strides[0] = src_strides[0]; + strides[1] = src_strides[ndim - 3]; + strides[2] = src_strides[ndim - 2]; + strides[3] = src_strides[ndim - 1]; + + const auto& dst_strides = out.strides(); + out_strides[0] = dst_strides[0]; + out_strides[1] = dst_strides[ndim - 3]; + out_strides[2] = dst_strides[ndim - 2]; + out_strides[3] = dst_strides[ndim - 1]; // Some flags to help us dispatch below bool single = in.flags().row_contiguous && B == 1 && T == 1; @@ -327,9 +324,9 @@ void RoPE::eval_gpu( } encoder.set_output_array(out); dispatch_float_types(out.dtype(), "rope", [&](auto type_tag) { + using DataType = cuda_type_t; dispatch_bool(traditional_, [&](auto traditional) { dispatch_bool(forward_, [&](auto forward) { - using DataType = cuda_type_t; if (single && !with_freqs) { auto kernel = cu::rope_single; @@ -339,7 +336,6 @@ void RoPE::eval_gpu( kernel, grid, block, - 0, gpu_ptr(donated ? out : in), gpu_ptr(out), gpu_ptr(offset), @@ -356,7 +352,6 @@ void RoPE::eval_gpu( kernel, grid, block, - 0, gpu_ptr(donated ? out : in), gpu_ptr(out), gpu_ptr(offset), @@ -380,7 +375,6 @@ void RoPE::eval_gpu( kernel, grid, block, - 0, gpu_ptr(donated ? out : in), gpu_ptr(out), gpu_ptr(offset), @@ -407,7 +401,6 @@ void RoPE::eval_gpu( kernel, grid, block, - 0, gpu_ptr(donated ? out : in), gpu_ptr(out), gpu_ptr(offset), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp index 54700bdcd..286d500e2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cpp @@ -25,6 +25,18 @@ array prepare_sdpa_input(const array& x, Stream s) { return x; } +array prepare_sdpa_sinks(const array& sinks, Stream s) { + // cuDNN requires sinks to be float32. + if (sinks.dtype() == float32) { + return sinks; + } + array sinks_f32(sinks.shape(), float32, nullptr, {}); + copy_gpu(sinks, sinks_f32, CopyType::Vector, s); + auto& encoder = cu::get_command_encoder(s); + encoder.add_temporary(sinks_f32); + return sinks_f32; +} + void malloc_with_same_layout( cu::CommandEncoder& encoder, array& o, @@ -57,6 +69,58 @@ void malloc_with_same_layout( {true, false, false}); } +bool use_cudnn_for_decoding( + const array& q, + const array& k, + const array& v, + bool has_arr_mask) { + if (q.shape(2) != 1) { + return false; + } + if (has_arr_mask) { + return false; + } + // The cuDNN SDPA is faster than vector kernel but for small sequence the + // overhead would kill the advantage. + constexpr int kv_cache_step = 256; // number is from mlx-lm + if (k.shape(2) < kv_cache_step) { + return false; + } + // When called during graph building the strides is not available, and we + // rely on |supports_sdpa_vector| to decide whether to use fast sdpa since + // we can fallback to |sdpa_vector|. + if ((k.status() != array::evaluated) || (v.status() != array::evaluated)) { + return false; + } + // Check if k/v are slices from fixed-size kv cache. + auto is_slice = [](const array& kv) { + // Get pre-sliced sequence length from strides, and check if the buffer + // belongs to a contiguous kv cache. + int64_t T_kv = kv.strides(1) / kv.strides(2); + if (kv.size() / kv.shape(2) * T_kv != kv.buffer_size() / kv.itemsize()) { + return false; + } + // It is possible to use heuristic to check slices, but for now just make + // mlx-lm work. + return T_kv % kv_cache_step == 0; + }; + return is_slice(k) && is_slice(v); +} + +// Get original kv from slices, i.e. undo keys[..., :offset, :] +array unslice_kv(const array& kv) { + Shape shape = kv.shape(); + shape[2] = /* T_kv */ kv.strides(1) / kv.strides(2); + array copy(shape, kv.dtype(), nullptr, {}); + copy.copy_shared_buffer( + kv, + make_contiguous_strides(shape), + {true, true, false}, + /* data_size */ kv.buffer_size() / kv.itemsize(), + /* offset */ -kv.offset()); + return copy; +} + constexpr int QKV_NDIM = 4; struct SDPACacheKey { @@ -71,6 +135,7 @@ struct SDPACacheKey { bool do_causal; std::array mask_shape; std::array mask_strides; + bool has_sinks; bool output_logsumexp; }; @@ -81,37 +146,43 @@ inline BytesKey build_sdpa_cache_key( const array& v, bool do_causal, const std::optional& mask_arr, - bool output_logsumexp = true) { + const std::optional& sinks, + bool decoding = false, + bool output_logsumexp = false) { BytesKey cache_key; - cache_key.pod = { - encoder.device().cuda_device(), - dtype_to_cudnn_type(q.dtype()), - vector_key(q.shape()), - vector_key(k.shape()), - vector_key(v.shape()), - vector_key(q.strides()), - vector_key(k.strides()), - vector_key(v.strides()), - do_causal, - {}, - {}, - output_logsumexp, - }; + cache_key.pod.device_id = encoder.device().cuda_device(); + cache_key.pod.cudnn_dtype = dtype_to_cudnn_type(q.dtype()); + cache_key.pod.q_shape = vector_key(q.shape()); + cache_key.pod.k_shape = vector_key(k.shape()); + cache_key.pod.v_shape = vector_key(v.shape()); + cache_key.pod.q_strides = vector_key(q.strides()); + cache_key.pod.k_strides = vector_key(k.strides()); + cache_key.pod.v_strides = vector_key(v.strides()); + cache_key.pod.do_causal = do_causal; + cache_key.pod.has_sinks = sinks.has_value(); + cache_key.pod.output_logsumexp = output_logsumexp; if (mask_arr) { cache_key.pod.mask_shape = vector_key(mask_arr->shape()); cache_key.pod.mask_strides = vector_key(mask_arr->strides()); } + if (decoding) { + int64_t T_kv = k.strides(1) / k.strides(2); + cache_key.pod.k_shape[2] = T_kv; + cache_key.pod.v_shape[2] = T_kv; + cache_key.pod.k_strides.fill(0); + cache_key.pod.v_strides.fill(0); + } return cache_key; } auto& sdpa_cache() { - static LRUBytesKeyCache cache( - "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 64); + static thread_local LRUBytesKeyCache cache( + "MLX_CUDA_SDPA_CACHE_SIZE", /* default_capacity */ 256); return cache; } auto& sdpa_backward_cache() { - static LRUBytesKeyCache cache( + static thread_local LRUBytesKeyCache cache( "MLX_CUDA_SDPA_BACKWARD_CACHE_SIZE", /* default_capacity */ 64); return cache; } @@ -122,6 +193,9 @@ enum UIDS { V, SCALE, BIAS, + SINKS, + SEQ_LEN_Q, + SEQ_LEN_KV, O, STATS, // Backward graph: @@ -138,6 +212,9 @@ DnnGraph build_sdpa_graph( const array& v, bool do_causal, const std::optional& mask_arr, + const std::optional& sinks, + const std::optional& seq_len_q, + const std::optional& seq_len_kv, bool output_logsumexp, const array& o, const std::optional& stats) { @@ -157,6 +234,14 @@ DnnGraph build_sdpa_graph( if (mask_arr) { options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr)); } + if (sinks) { + options.set_sink_token(graph.tensor_4d("SINKS", SINKS, *sinks, 1)); + } + if (seq_len_q && seq_len_kv) { + options.set_padding_mask(true); + options.set_seq_len_q(graph.tensor("SEQ_LEN_Q", SEQ_LEN_Q, *seq_len_q)); + options.set_seq_len_kv(graph.tensor("SEQ_LEN_KV", SEQ_LEN_KV, *seq_len_kv)); + } auto [o_, stats_] = graph.sdpa(q_, k_, v_, options); graph.tensor(o_, O, o)->set_output(true); @@ -164,10 +249,10 @@ DnnGraph build_sdpa_graph( graph.tensor(stats_, STATS, *stats)->set_output(true); } - CHECK_CUDNN_FE_ERROR(graph.prepare()); + CHECK_CUDNN_ERROR(graph.prepare()); graph.select_behavior_notes( {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API}); - CHECK_CUDNN_FE_ERROR(graph.build()); + CHECK_CUDNN_ERROR(graph.build()); return graph; } @@ -178,6 +263,7 @@ DnnGraph build_sdpa_backward_graph( const array& v, bool do_causal, const std::optional& mask_arr, + const std::optional& sinks, const array& o, const array& d_o, const array& stats, @@ -202,6 +288,9 @@ DnnGraph build_sdpa_backward_graph( if (mask_arr) { options.set_bias(graph.tensor("BIAS", BIAS, *mask_arr)); } + if (sinks) { + options.set_sink_token(graph.tensor_4d("SINKS", SINKS, *sinks, 1)); + } auto [d_q_, d_k_, d_v_] = graph.sdpa_backward(q_, k_, v_, o_, d_o_, stats_, options); @@ -209,22 +298,28 @@ DnnGraph build_sdpa_backward_graph( graph.tensor(d_k_, D_K, d_k)->set_output(true); graph.tensor(d_v_, D_V, d_v)->set_output(true); - CHECK_CUDNN_FE_ERROR(graph.prepare()); + CHECK_CUDNN_ERROR(graph.prepare()); graph.select_behavior_notes( {fe::BehaviorNote_t::SUPPORTS_CUDA_GRAPH_NATIVE_API}); - CHECK_CUDNN_FE_ERROR(graph.build()); + CHECK_CUDNN_ERROR(graph.build()); return graph; } } // namespace +void init_cudnn_sdpa_cache() { + sdpa_cache(); + sdpa_backward_cache(); +} + bool supports_sdpa_cudnn( const array& q, const array& k, const array& v, + bool has_arr_mask, bool do_causal, Stream s) { - static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SPDA", 1); + static bool enabled = env::get_var("MLX_CUDA_USE_CUDNN_SDPA", 1); if (!enabled) { return false; } @@ -234,8 +329,8 @@ bool supports_sdpa_cudnn( return false; } - // Only use cuDNN for prefilling (T_q > 1) and training (T_q == T_kv). - if ((q.shape(2) == 1) && (q.shape(2) != k.shape(2))) { + // Only use cuDNN for decoding when k/v are slices from fixed-size kv cache. + if ((q.shape(2) == 1) && !use_cudnn_for_decoding(q, k, v, has_arr_mask)) { return false; } @@ -256,20 +351,39 @@ bool supports_sdpa_cudnn( void sdpa_cudnn( const array& q, - const array& k, - const array& v, + array k, + array v, float scale, array& o, std::optional& stats, bool do_causal, const std::optional& mask_arr, + const std::optional& sinks, bool output_logsumexp, Stream s) { auto& encoder = cu::get_command_encoder(s); - auto handle = encoder.device().get_cudnn_handle(); + auto handle = get_cudnn_handle(encoder.device()); malloc_with_same_layout(encoder, o, q); + // For decoding, unslice k/v and apply padding mask. + std::optional seq_len_q; + std::optional seq_len_kv; + bool decoding = use_cudnn_for_decoding(q, k, v, mask_arr.has_value()); + if (decoding) { + int B = q.shape(0); + std::vector seq_len_q_vec(B, q.shape(2)); + std::vector seq_len_kv_vec(B, k.shape(2)); + seq_len_q = array(seq_len_q_vec.begin(), {B, 1, 1, 1}); + seq_len_kv = array(seq_len_kv_vec.begin(), {B, 1, 1, 1}); + encoder.add_temporary(*seq_len_q); + encoder.add_temporary(*seq_len_kv); + k = unslice_kv(k); + v = unslice_kv(v); + encoder.add_temporary(k); + encoder.add_temporary(v); + } + encoder.set_input_array(q); encoder.set_input_array(k); encoder.set_input_array(v); @@ -277,6 +391,13 @@ void sdpa_cudnn( if (mask_arr) { encoder.set_input_array(*mask_arr); } + if (sinks) { + encoder.set_input_array(*sinks); + } + if (seq_len_q && seq_len_kv) { + encoder.set_input_array(*seq_len_q); + encoder.set_input_array(*seq_len_kv); + } if (output_logsumexp) { stats->set_data(cu::malloc_async(stats->nbytes(), encoder)); encoder.set_output_array(*stats); @@ -284,11 +405,22 @@ void sdpa_cudnn( // Search cache. auto cache_key = build_sdpa_cache_key( - encoder, q, k, v, do_causal, mask_arr, output_logsumexp); + encoder, q, k, v, do_causal, mask_arr, sinks, decoding, output_logsumexp); auto it = sdpa_cache().find(cache_key); if (it == sdpa_cache().end()) { auto graph = build_sdpa_graph( - handle, q, k, v, do_causal, mask_arr, output_logsumexp, o, stats); + handle, + q, + k, + v, + do_causal, + mask_arr, + sinks, + seq_len_q, + seq_len_kv, + output_logsumexp, + o, + stats); it = sdpa_cache().emplace(cache_key, std::move(graph)).first; } auto& graph = it->second; @@ -302,11 +434,18 @@ void sdpa_cudnn( if (mask_arr) { variant_pack[BIAS] = gpu_ptr(*mask_arr); } + if (sinks) { + variant_pack[SINKS] = gpu_ptr(*sinks); + } + if (seq_len_q && seq_len_kv) { + variant_pack[SEQ_LEN_Q] = gpu_ptr(*seq_len_q); + variant_pack[SEQ_LEN_KV] = gpu_ptr(*seq_len_kv); + } if (output_logsumexp) { variant_pack[STATS] = gpu_ptr(*stats); } - CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); + CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); } void sdpa_backward_cudnn( @@ -318,13 +457,14 @@ void sdpa_backward_cudnn( const array& stats, bool do_causal, const std::optional& mask_arr, + const std::optional& sinks, const array& d_o, array& d_q, array& d_k, array& d_v, Stream s) { auto& encoder = cu::get_command_encoder(s); - auto handle = encoder.device().get_cudnn_handle(); + auto handle = get_cudnn_handle(encoder.device()); malloc_with_same_layout(encoder, d_q, q); malloc_with_same_layout(encoder, d_k, k); @@ -342,13 +482,29 @@ void sdpa_backward_cudnn( if (mask_arr) { encoder.set_input_array(*mask_arr); } + if (sinks) { + encoder.set_input_array(*sinks); + } // Search cache. - auto cache_key = build_sdpa_cache_key(encoder, q, k, v, do_causal, mask_arr); + auto cache_key = + build_sdpa_cache_key(encoder, q, k, v, do_causal, mask_arr, sinks); auto it = sdpa_backward_cache().find(cache_key); if (it == sdpa_backward_cache().end()) { auto graph = build_sdpa_backward_graph( - handle, q, k, v, do_causal, mask_arr, o, d_o, stats, d_q, d_k, d_v); + handle, + q, + k, + v, + do_causal, + mask_arr, + sinks, + o, + d_o, + stats, + d_q, + d_k, + d_v); it = sdpa_backward_cache().emplace(cache_key, std::move(graph)).first; } auto& graph = it->second; @@ -367,8 +523,11 @@ void sdpa_backward_cudnn( if (mask_arr) { variant_pack[BIAS] = gpu_ptr(*mask_arr); } + if (sinks) { + variant_pack[SINKS] = gpu_ptr(*sinks); + } - CHECK_CUDNN_FE_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); + CHECK_CUDNN_ERROR(graph.encode_graph(encoder, std::move(variant_pack))); } // Defined in scaled_dot_product_attention.cu file. @@ -376,9 +535,7 @@ bool supports_sdpa_vector( const array& q, const array& k, const array& v, - bool has_mask, bool has_arr_mask, - bool do_causal, bool output_logsumexp); void sdpa_vector( const array& q, @@ -392,6 +549,33 @@ void sdpa_vector( namespace fast { +namespace { + +std::tuple has_fused_kernel( + const array& q, + const array& k, + const array& v, + bool has_arr_mask, + bool do_causal, + bool output_logsumexp, + Stream s) { + if (s.device != Device::gpu) { + return {false, "the fused kernels require a GPU stream."}; + } + if (!supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal, s) && + !supports_sdpa_vector(q, k, v, has_arr_mask, output_logsumexp)) { + std::ostringstream msg; + msg << "neither the cuDNN attention nor the vector attention kernel " + << "supports this configuration; got query shape " << q.shape() + << ", key shape " << k.shape() << ", value shape " << v.shape() + << " with dtype " << q.dtype() << "."; + return {false, msg.str()}; + } + return {true, ""}; +} + +} // namespace + bool ScaledDotProductAttention::use_fallback( const array& q, const array& k, @@ -401,14 +585,21 @@ bool ScaledDotProductAttention::use_fallback( bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s) { - if (s.device == Device::cpu) { - return true; + auto [has_fused, reason] = + has_fused_kernel(q, k, v, has_arr_mask, do_causal, output_logsumexp, s); + if (force_fused) { + if (!has_fused) { + std::ostringstream msg; + msg << "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available: " + << reason; + throw std::invalid_argument(msg.str()); + } + return false; } - - return !supports_sdpa_vector( - q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp) && - !supports_sdpa_cudnn(q, k, v, do_causal, s); + return !has_fused; } bool ScaledDotProductAttention::supports_bool_mask() { @@ -433,19 +624,19 @@ void ScaledDotProductAttention::eval_gpu( if (has_arr_mask) { mask_arr = prepare_sdpa_input(inputs[3], s); } + std::optional sinks; + if (has_sinks_) { + sinks = inputs.back(); + } std::optional stats; if (output_logsumexp_) { stats = outputs[1]; } - if (supports_sdpa_vector( - q, k, v, has_mask, has_arr_mask, do_causal_, output_logsumexp_)) { - if (has_sinks_) { - sdpa_vector(q, k, v, scale_, out, do_causal_, inputs.back(), s); - } else { - sdpa_vector(q, k, v, scale_, out, do_causal_, std::nullopt, s); + if (supports_sdpa_cudnn(q, k, v, has_arr_mask, do_causal_, s)) { + if (sinks) { + sinks = prepare_sdpa_sinks(*sinks, s); } - } else { sdpa_cudnn( q, k, @@ -455,8 +646,11 @@ void ScaledDotProductAttention::eval_gpu( stats, do_causal_, mask_arr, + sinks, output_logsumexp_, s); + } else { + sdpa_vector(q, k, v, scale_, out, do_causal_, sinks, s); } } @@ -491,6 +685,10 @@ void ScaledDotProductAttentionVJP::eval_gpu( if (has_arr_mask) { mask_arr = prepare_sdpa_input(inputs[3], s); } + std::optional sinks; + if (has_sinks_) { + sinks = prepare_sdpa_sinks(inputs.back(), s); + } assert(outputs.size() == 3); auto& d_q = outputs[0]; @@ -498,7 +696,20 @@ void ScaledDotProductAttentionVJP::eval_gpu( auto& d_v = outputs[2]; sdpa_backward_cudnn( - q, k, v, scale_, o, stats, do_causal_, mask_arr, d_o, d_q, d_k, d_v, s); + q, + k, + v, + scale_, + o, + stats, + do_causal_, + mask_arr, + sinks, + d_o, + d_q, + d_k, + d_v, + s); } } // namespace fast diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu index 111cfee17..6fafb0ba9 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scaled_dot_product_attention.cu @@ -511,7 +511,6 @@ void sdpa_vector_1pass_fallback( kernel, grid_dim, block_dim, - 0, gpu_ptr(q), gpu_ptr(k), gpu_ptr(v), @@ -599,7 +598,6 @@ void sdpa_vector_2pass_fallback( kernel, grid_dim, block_dim, - 0, gpu_ptr(q), gpu_ptr(k), gpu_ptr(v), @@ -626,7 +624,6 @@ void sdpa_vector_2pass_fallback( kernel, grid_dim, block_dim, - 0, gpu_ptr(intermediate), gpu_ptr(sums), gpu_ptr(maxs), @@ -665,9 +662,7 @@ bool supports_sdpa_vector( const array& q, const array& k, const array& v, - bool has_mask, bool has_arr_mask, - bool do_causal, bool output_logsumexp) { if (output_logsumexp) { return false; diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu index b4d6cf368..a7dd2eea2 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/scan.cu @@ -5,6 +5,7 @@ #include "mlx/backend/cuda/kernel_utils.cuh" #include "mlx/backend/cuda/reduce/reduce_ops.cuh" #include "mlx/backend/gpu/copy.h" +#include "mlx/backend/gpu/scan.h" #include "mlx/dtype_utils.h" #include "mlx/primitives.h" @@ -362,51 +363,38 @@ constexpr bool supports_scan_op() { } } -void Scan::eval_gpu(const std::vector& inputs, array& out) { - nvtx3::scoped_range r("Scan::eval_gpu"); - assert(inputs.size() == 1); - auto in = inputs[0]; - auto& s = stream(); +void scan_gpu_inplace( + const array& in, + array& out, + Scan::ReduceType reduce_type, + int axis, + bool reverse, + bool inclusive, + const Stream& s) { auto& encoder = cu::get_command_encoder(s); - - if (in.flags().contiguous && in.strides()[axis_] != 0) { - if (in.is_donatable() && in.itemsize() == out.itemsize()) { - out.copy_shared_buffer(in); - } else { - out.set_data( - cu::malloc_async(in.data_size() * out.itemsize(), encoder), - in.data_size(), - in.strides(), - in.flags()); - } - } else { - in = contiguous_copy_gpu(in, s); - out.copy_shared_buffer(in); - } - constexpr int N_READS = 4; - int32_t axis_size = in.shape(axis_); - bool contiguous = in.strides()[axis_] == 1; + int32_t axis_size = in.shape(axis); + bool contiguous = in.strides()[axis] == 1; encoder.set_input_array(in); encoder.set_output_array(out); dispatch_all_types(in.dtype(), [&](auto type_tag) { using T = cuda_type_t; - dispatch_scan_ops(reduce_type_, [&](auto scan_op_tag) { + dispatch_scan_ops(reduce_type, [&](auto scan_op_tag) { using Op = MLX_GET_TYPE(scan_op_tag); if constexpr (supports_scan_op()) { using U = typename cu::ScanResult::type; - dispatch_bool(inclusive_, [&](auto inclusive) { - dispatch_bool(reverse_, [&](auto reverse) { + dispatch_bool(inclusive, [&](auto inclusive_tag) { + dispatch_bool(reverse, [&](auto reverse_tag) { if (contiguous) { auto kernel = cu::contiguous_scan< T, U, Op, N_READS, - inclusive.value, - reverse.value>; + inclusive_tag.value, + reverse_tag.value>; int block_dim = cuda::ceil_div(axis_size, N_READS); block_dim = cuda::ceil_div(block_dim, WARP_SIZE) * WARP_SIZE; block_dim = std::min(block_dim, WARP_SIZE * WARP_SIZE); @@ -414,7 +402,6 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { kernel, in.data_size() / axis_size, block_dim, - 0, gpu_ptr(in), gpu_ptr(out), axis_size); @@ -428,9 +415,9 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { N_READS, BM, BN, - inclusive.value, - reverse.value>; - int64_t stride = in.strides()[axis_]; + inclusive_tag.value, + reverse_tag.value>; + int64_t stride = in.strides()[axis]; int64_t stride_blocks = cuda::ceil_div(stride, BN); dim3 num_blocks = get_2d_grid_dims( in.shape(), in.strides(), axis_size * stride); @@ -444,7 +431,6 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { kernel, num_blocks, block_dim, - 0, gpu_ptr(in), gpu_ptr(out), axis_size, @@ -465,4 +451,29 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { }); } +void Scan::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("Scan::eval_gpu"); + assert(inputs.size() == 1); + auto in = inputs[0]; + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + + if (in.flags().contiguous && in.strides()[axis_] != 0) { + if (in.is_donatable() && in.itemsize() == out.itemsize()) { + out.copy_shared_buffer(in); + } else { + out.set_data( + cu::malloc_async(in.data_size() * out.itemsize(), encoder), + in.data_size(), + in.strides(), + in.flags()); + } + } else { + in = contiguous_copy_gpu(in, s); + out.copy_shared_buffer(in); + } + + scan_gpu_inplace(in, out, reduce_type_, axis_, reverse_, inclusive_, s); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp index d45b16df4..525cccc26 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/slicing.cpp @@ -37,7 +37,11 @@ void concatenate_gpu( size_t data_offset = strides[axis] * sizes[i]; out_slice.copy_shared_buffer( out, strides, flags, out_slice.size(), data_offset); - copy_gpu_inplace(inputs[i], out_slice, CopyType::GeneralGeneral, s); + auto ctype = CopyType::GeneralGeneral; + if (axis == 0 && inputs[i].flags().row_contiguous) { + ctype = CopyType::Vector; + } + copy_gpu_inplace(inputs[i], out_slice, ctype, s); } } @@ -56,7 +60,9 @@ array compute_dynamic_offset( dtype_to_cuda_type(dtype), nidx); - cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() { + auto& encoder = cu::get_command_encoder(s); + + cu::JitModule& mod = cu::get_jit_module(encoder.device(), module_name, [&]() { std::string source = R"( #include "mlx/backend/cuda/device/utils.cuh" @@ -81,7 +87,6 @@ array compute_dynamic_offset( return std::make_tuple(false, std::move(source), std::vector{kernel_name}); }); - auto& encoder = cu::get_command_encoder(s); // Prepare output. array offset({1}, int64, nullptr, {}); bool donate = indices.is_donatable() && @@ -103,7 +108,7 @@ array compute_dynamic_offset( args.append(axes); auto kernel = mod.get_kernel(kernel_name); - encoder.add_kernel_node(kernel, 1, 1, 0, args.args()); + encoder.add_kernel_node_raw(kernel, 1, 1, {}, 0, args.args()); return offset; } diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/softmax.cu b/Source/Cmlx/mlx/mlx/backend/cuda/softmax.cu index 4b2a35bb9..d34ac1189 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/softmax.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/softmax.cu @@ -151,7 +151,6 @@ void Softmax::eval_gpu(const std::vector& inputs, array& out) { kernel, n_rows, block_dim(), - 0, gpu_ptr(in), gpu_ptr(out), axis_size); diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu index c258c0138..d9c77387c 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/sort.cu @@ -11,6 +11,7 @@ #include "mlx/dtype_utils.h" #include "mlx/primitives.h" +#include #include #include #include @@ -21,6 +22,8 @@ constexpr int N_PER_THREAD = 8; namespace cu { +namespace cg = cooperative_groups; + template __device__ __forceinline__ T nan_value(); @@ -44,6 +47,12 @@ __device__ __forceinline__ __nv_bfloat16 nan_value<__nv_bfloat16>() { return __float2bfloat16(cuda::std::numeric_limits::quiet_NaN()); } +template <> +__device__ __forceinline__ complex64_t nan_value() { + float qnan = cuda::std::numeric_limits::quiet_NaN(); + return complex64_t{qnan, qnan}; +} + template struct InitValue { __device__ __forceinline__ static T value() { @@ -52,7 +61,9 @@ struct InitValue { }; template -struct InitValue>> { +struct InitValue< + T, + cuda::std::enable_if_t || cu::is_complex_v>> { __device__ __forceinline__ static T value() { return nan_value(); } @@ -65,6 +76,15 @@ __device__ __forceinline__ void thread_swap(T& a, T& b) { b = w; } +template +__device__ __forceinline__ bool check_nan(T a) { + if constexpr (cu::is_complex_v) { + return cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag()); + } else { + return cuda::std::isnan(a); + } +} + template struct LessThan { __device__ __forceinline__ static T init() { @@ -72,9 +92,9 @@ struct LessThan { } __device__ __forceinline__ bool operator()(T a, T b) const { - if constexpr (std::is_floating_point_v) { - bool an = cuda::std::isnan(a); - bool bn = cuda::std::isnan(b); + if constexpr (is_floating_v || cu::is_complex_v) { + bool an = check_nan(a); + bool bn = check_nan(b); if (an | bn) { return (!an) & bn; } @@ -699,6 +719,40 @@ __global__ void mb_block_merge_kernel( } } +template +__device__ __forceinline__ uint32_t +searchsorted_impl(const T* a, T v, uint32_t n, int64_t a_stride) { + LessThan lt; + uint32_t lo = 0; + uint32_t hi = n; + while (lo < hi) { + uint32_t mid = lo + (hi - lo) / 2; + T m = a[static_cast(mid) * a_stride]; + bool below = Right ? !lt(v, m) : lt(m, v); + if (below) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +template +__global__ void searchsorted( + const T* a, + const T* v, + uint32_t* out, + int64_t size, + uint32_t n, + int64_t a_stride) { + int64_t index = cg::this_grid().thread_rank(); + if (index >= size) { + return; + } + out[index] = searchsorted_impl(a, v[index], n, a_stride); +} + } // namespace cu namespace { @@ -745,84 +799,76 @@ void single_block_sort( dispatch_all_types(in.dtype(), [&](auto type_tag) { using CTYPE = MLX_GET_TYPE(type_tag); - if constexpr (!std::is_same_v) { - using ValT = cuda_type_t; - dispatch_block_dim(bn, [&](auto block_dim) { - constexpr int BLOCK_THREADS = block_dim(); - if constexpr (BLOCK_THREADS < 1024) { - dim3 grid(1, n_rows, 1); - dim3 block(BLOCK_THREADS, 1, 1); - - dispatch_bool(argsort, [&](auto arg_tag) { - constexpr bool ARG_SORT = decltype(arg_tag)::value; - using OutT = std::conditional_t; - - if (contiguous) { - auto kernel = cu::block_sort_kernel< - ValT, - OutT, - ARG_SORT, - BLOCK_THREADS, - N_PER_THREAD>; - int64_t in_stride_segment_axis = INT64_MAX; - int64_t out_stride_segment_axis = INT64_MAX; - for (int i = 0; i < nc_shape.size(); i++) { - if (nc_shape[i] == 1) { - continue; - } - if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) { - throw std::runtime_error( - "[Sort::eval_gpu] Stride too large."); - } - in_stride_segment_axis = - std::min(in_stride_segment_axis, in_nc_str[i]); - out_stride_segment_axis = - std::min(out_stride_segment_axis, out_nc_str[i]); + using ValT = cuda_type_t; + dispatch_block_dim(bn, [&](auto block_dim) { + constexpr int BLOCK_THREADS = block_dim(); + if constexpr (BLOCK_THREADS < 1024) { + dim3 grid(1, n_rows, 1); + dim3 block(BLOCK_THREADS, 1, 1); + + dispatch_bool(argsort, [&](auto arg_tag) { + constexpr bool ARG_SORT = decltype(arg_tag)::value; + using OutT = std::conditional_t; + + if (contiguous) { + auto kernel = cu::block_sort_kernel< + ValT, + OutT, + ARG_SORT, + BLOCK_THREADS, + N_PER_THREAD>; + int64_t in_stride_segment_axis = INT64_MAX; + int64_t out_stride_segment_axis = INT64_MAX; + for (int i = 0; i < nc_shape.size(); i++) { + if (nc_shape[i] == 1) { + continue; + } + if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) { + throw std::runtime_error("[Sort::eval_gpu] Stride too large."); } - encoder.add_kernel_node( - kernel, - grid, - block, - 0, - gpu_ptr(in), - gpu_ptr(out), - size_sorted_axis, - in_stride_sorted_axis, - out_stride_sorted_axis, - in_stride_segment_axis, - out_stride_segment_axis); - } else { - auto kernel = cu::block_sort_nc_kernel< - ValT, - OutT, - ARG_SORT, - BLOCK_THREADS, - N_PER_THREAD>; - auto nc_shape_param = const_param(nc_shape); - auto in_nc_strides_param = const_param(in_nc_str); - auto out_nc_strides_param = const_param(out_nc_str); - encoder.add_kernel_node( - kernel, - grid, - block, - 0, - gpu_ptr(in), - gpu_ptr(out), - size_sorted_axis, - in_stride_sorted_axis, - out_stride_sorted_axis, - nc_shape_param, - in_nc_strides_param, - out_nc_strides_param, - nc_dim); + in_stride_segment_axis = + std::min(in_stride_segment_axis, in_nc_str[i]); + out_stride_segment_axis = + std::min(out_stride_segment_axis, out_nc_str[i]); } - }); - } - }); - } else { - throw std::runtime_error( - "CUDA backend does not support sorting complex numbers"); - } + encoder.add_kernel_node( + kernel, + grid, + block, + gpu_ptr(in), + gpu_ptr(out), + size_sorted_axis, + in_stride_sorted_axis, + out_stride_sorted_axis, + in_stride_segment_axis, + out_stride_segment_axis); + } else { + auto kernel = cu::block_sort_nc_kernel< + ValT, + OutT, + ARG_SORT, + BLOCK_THREADS, + N_PER_THREAD>; + auto nc_shape_param = const_param(nc_shape); + auto in_nc_strides_param = const_param(in_nc_str); + auto out_nc_strides_param = const_param(out_nc_str); + encoder.add_kernel_node( + kernel, + grid, + block, + gpu_ptr(in), + gpu_ptr(out), + size_sorted_axis, + in_stride_sorted_axis, + out_stride_sorted_axis, + nc_shape_param, + in_nc_strides_param, + out_nc_strides_param, + nc_dim); + } + }); + } + }); }); } @@ -872,102 +918,94 @@ void multi_block_sort( dispatch_all_types(in.dtype(), [&](auto type_tag) { using CTYPE = MLX_GET_TYPE(type_tag); - if constexpr (!std::is_same_v) { - using ValT = cuda_type_t; - using IdxT = uint32_t; - constexpr int BLOCK_THREADS = sizeof(ValT) == 8 ? 256 : 512; - dim3 grid(n_blocks, n_rows, 1); - dim3 block(BLOCK_THREADS, 1, 1); - - dispatch_bool(argsort, [&](auto arg_tag) { - constexpr bool ARG_SORT = decltype(arg_tag)::value; - auto nc_shape_param = const_param(nc_shape); - auto nc_strides_param = const_param(nc_str); - - auto block_sort_kernel = cu::mb_block_sort_kernel< + using ValT = cuda_type_t; + using IdxT = uint32_t; + constexpr int BLOCK_THREADS = sizeof(ValT) == 8 ? 256 : 512; + dim3 grid(n_blocks, n_rows, 1); + dim3 block(BLOCK_THREADS, 1, 1); + + dispatch_bool(argsort, [&](auto arg_tag) { + constexpr bool ARG_SORT = decltype(arg_tag)::value; + auto nc_shape_param = const_param(nc_shape); + auto nc_strides_param = const_param(nc_str); + + auto block_sort_kernel = cu::mb_block_sort_kernel< + ValT, + IdxT, + ARG_SORT, + BLOCK_THREADS, + N_PER_THREAD>; + encoder.set_input_array(in); + encoder.set_output_array(dev_vals_in); + encoder.set_output_array(dev_idxs_in); + encoder.add_kernel_node( + block_sort_kernel, + grid, + block, + gpu_ptr(in), + gpu_ptr(dev_vals_in), + gpu_ptr(dev_idxs_in), + size_sorted_axis, + stride_sorted_axis, + nc_shape_param, + nc_strides_param, + nc_dim); + + int n_thr_per_group = (n_blocks + 1) < 1024 ? (n_blocks + 1) : 1024; + + for (int merge_tiles = 2; (merge_tiles / 2) < n_blocks; + merge_tiles *= 2) { + auto partition_kernel = cu::mb_block_partition_kernel< ValT, IdxT, ARG_SORT, BLOCK_THREADS, N_PER_THREAD>; - encoder.set_input_array(in); - encoder.set_output_array(dev_vals_in); - encoder.set_output_array(dev_idxs_in); + + encoder.set_input_array(dev_vals_in); + encoder.set_input_array(dev_idxs_in); + encoder.set_output_array(block_partitions); + encoder.add_kernel_node( - block_sort_kernel, - grid, - block, - 0, - gpu_ptr(in), + partition_kernel, + dim3(1, n_rows, 1), + dim3(n_thr_per_group, 1, 1), + gpu_ptr(block_partitions), gpu_ptr(dev_vals_in), gpu_ptr(dev_idxs_in), size_sorted_axis, - stride_sorted_axis, - nc_shape_param, - nc_strides_param, - nc_dim); - - int n_thr_per_group = (n_blocks + 1) < 1024 ? (n_blocks + 1) : 1024; - - for (int merge_tiles = 2; (merge_tiles / 2) < n_blocks; - merge_tiles *= 2) { - auto partition_kernel = cu::mb_block_partition_kernel< - ValT, - IdxT, - ARG_SORT, - BLOCK_THREADS, - N_PER_THREAD>; - - encoder.set_input_array(dev_vals_in); - encoder.set_input_array(dev_idxs_in); - encoder.set_output_array(block_partitions); - - encoder.add_kernel_node( - partition_kernel, - dim3(1, n_rows, 1), - dim3(n_thr_per_group, 1, 1), - 0, - gpu_ptr(block_partitions), - gpu_ptr(dev_vals_in), - gpu_ptr(dev_idxs_in), - size_sorted_axis, - merge_tiles, - n_blocks); - - auto merge_kernel = cu::mb_block_merge_kernel< - ValT, - IdxT, - ARG_SORT, - BLOCK_THREADS, - N_PER_THREAD>; - - encoder.set_input_array(dev_vals_in); - encoder.set_input_array(dev_idxs_in); - encoder.set_input_array(block_partitions); - encoder.set_output_array(dev_vals_out); - encoder.set_output_array(dev_idxs_out); - - encoder.add_kernel_node( - merge_kernel, - dim3(n_blocks, n_rows, 1), - dim3(BLOCK_THREADS, 1, 1), - 0, - gpu_ptr(block_partitions), - gpu_ptr(dev_vals_in), - gpu_ptr(dev_idxs_in), - gpu_ptr(dev_vals_out), - gpu_ptr(dev_idxs_out), - size_sorted_axis, - merge_tiles, - n_blocks); - std::swap(dev_vals_in, dev_vals_out); - std::swap(dev_idxs_in, dev_idxs_out); - } - }); - } else { - throw std::runtime_error( - "CUDA backend does not support sorting complex numbers"); - } + merge_tiles, + n_blocks); + + auto merge_kernel = cu::mb_block_merge_kernel< + ValT, + IdxT, + ARG_SORT, + BLOCK_THREADS, + N_PER_THREAD>; + + encoder.set_input_array(dev_vals_in); + encoder.set_input_array(dev_idxs_in); + encoder.set_input_array(block_partitions); + encoder.set_output_array(dev_vals_out); + encoder.set_output_array(dev_idxs_out); + + encoder.add_kernel_node( + merge_kernel, + dim3(n_blocks, n_rows, 1), + dim3(BLOCK_THREADS, 1, 1), + gpu_ptr(block_partitions), + gpu_ptr(dev_vals_in), + gpu_ptr(dev_idxs_in), + gpu_ptr(dev_vals_out), + gpu_ptr(dev_idxs_out), + size_sorted_axis, + merge_tiles, + n_blocks); + std::swap(dev_vals_in, dev_vals_out); + std::swap(dev_idxs_in, dev_idxs_out); + } + }); }); encoder.add_temporary(dev_vals_out); @@ -1073,4 +1111,48 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { gpu_sort(stream(), inputs[0], out, axis_, false); } -} // namespace mlx::core \ No newline at end of file +void SearchSorted::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("SearchSorted::eval_gpu"); + assert(inputs.size() == 2); + auto& s = stream(); + auto& a = inputs[0]; + auto v = inputs[1]; + + auto& encoder = cu::get_command_encoder(s); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + if (out.size() == 0) { + return; + } + + if (!v.flags().row_contiguous) { + v = contiguous_copy_gpu(v, s); + encoder.add_temporary(v); + } + + encoder.set_input_array(a); + encoder.set_input_array(v); + encoder.set_output_array(out); + + int64_t a_stride = a.strides()[0]; // sequence is 1D + + dispatch_all_types(a.dtype(), [&](auto type_tag) { + using CTYPE = MLX_GET_TYPE(type_tag); + using T = cuda_type_t; + dispatch_bool(right_, [&](auto right) { + auto [num_blocks, block_dims] = + get_launch_args(out, out.size() > INT32_MAX); + encoder.add_kernel_node( + cu::searchsorted, + num_blocks, + block_dims, + gpu_ptr(a), + gpu_ptr(v), + gpu_ptr(out), + static_cast(out.size()), + static_cast(a.size()), + a_stride); + }); + }); +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/ternary.cu b/Source/Cmlx/mlx/mlx/backend/cuda/ternary.cu index e13ef6405..86a9d0d7d 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/ternary.cu +++ b/Source/Cmlx/mlx/mlx/backend/cuda/ternary.cu @@ -53,7 +53,9 @@ __global__ void ternary_g_nd( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -167,7 +169,6 @@ void ternary_op_gpu_inplace( cu::ternary_v, num_blocks, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(c), @@ -193,10 +194,8 @@ void ternary_op_gpu_inplace( if (dim0 >= 4) { work_per_thread = 4; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); if (ndim <= 3) { dispatch_1_2_3(ndim, [&](auto dims_constant) { @@ -208,9 +207,8 @@ void ternary_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(c), @@ -228,9 +226,8 @@ void ternary_op_gpu_inplace( } encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(a), gpu_ptr(b), gpu_ptr(c), diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/unary/unary.cuh b/Source/Cmlx/mlx/mlx/backend/cuda/unary/unary.cuh index cade3f81e..eb305244e 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/unary/unary.cuh +++ b/Source/Cmlx/mlx/mlx/backend/cuda/unary/unary.cuh @@ -48,7 +48,9 @@ __global__ void unary_g( auto block = cg::this_thread_block(); auto grid = cg::this_grid(); IdxT index_rest = - grid.block_index().y * block.dim_threads().y + block.thread_index().y; + (grid.block_index().z * grid.dim_blocks().y + grid.block_index().y) * + block.dim_threads().y + + block.thread_index().y; if (index_rest >= size_rest) { return; } @@ -157,7 +159,6 @@ void unary_op_gpu_inplace( cu::unary_v, num_blocks, block_dims, - 0, gpu_ptr(in), gpu_ptr(out), out.data_size()); @@ -173,15 +174,12 @@ void unary_op_gpu_inplace( kernel = cu::unary_g; work_per_thread = 4; } - dim0 = (dim0 + work_per_thread - 1) / work_per_thread; - auto block_dims = get_block_dims(dim0, rest, 1); - uint32_t num_blocks_x = cuda::ceil_div(dim0, block_dims.x); - uint32_t num_blocks_y = cuda::ceil_div(rest, block_dims.y); + auto [grid_dims, block_dims] = + get_launch_args_general(dim0, rest, work_per_thread); encoder.add_kernel_node( kernel, - {num_blocks_x, num_blocks_y}, + grid_dims, block_dims, - 0, gpu_ptr(in), gpu_ptr(out), rest, diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp index 08f0f3d4a..82272b74b 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/utils.cpp @@ -10,14 +10,6 @@ namespace mlx::core { -void check_cublas_error(const char* name, cublasStatus_t err) { - if (err != CUBLAS_STATUS_SUCCESS) { - // TODO: Use cublasGetStatusString when it is widely available. - throw std::runtime_error( - fmt::format("{} failed with code: {}.", name, static_cast(err))); - } -} - void check_cuda_error(const char* name, cudaError_t err) { if (err != cudaSuccess) { throw std::runtime_error( @@ -33,13 +25,6 @@ void check_cuda_error(const char* name, CUresult err) { } } -void check_cudnn_error(const char* name, cudnnStatus_t err) { - if (err != CUDNN_STATUS_SUCCESS) { - throw std::runtime_error( - fmt::format("{} failed: {}.", name, cudnnGetErrorString(err))); - } -} - const char* dtype_to_cuda_type(const Dtype& dtype) { switch (dtype) { case bool_: diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp b/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp index 495204d3d..de38a06e5 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp +++ b/Source/Cmlx/mlx/mlx/backend/cuda/worker.cpp @@ -7,16 +7,25 @@ namespace mlx::core::cu { Worker::Worker(Device& d) : signal_stream_(d), - signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync), - worker_(&Worker::thread_fn, this) {} + signal_event_(d, cudaEventDisableTiming | cudaEventBlockingSync) {} -Worker::~Worker() { +Worker::~Worker() = default; + +void Worker::start() { + // Note that |shared_from_this| can not be called in constructor. + worker_ = std::thread(&Worker::thread_fn, shared_from_this()); + // Detach the thread and let it free itself after finishing tasks. + // This is to avoid deadlock when joining threads on exit on Windows: + // https://developercommunity.visualstudio.com/t/1654756 + worker_.detach(); +} + +void Worker::stop() { { std::lock_guard lock(mtx_); stop_ = true; } cond_.notify_one(); - worker_.join(); } void Worker::add_task(std::function task) { diff --git a/Source/Cmlx/mlx/mlx/backend/cuda/worker.h b/Source/Cmlx/mlx/mlx/backend/cuda/worker.h index 8f05e7b97..54a4456ac 100644 --- a/Source/Cmlx/mlx/mlx/backend/cuda/worker.h +++ b/Source/Cmlx/mlx/mlx/backend/cuda/worker.h @@ -7,13 +7,14 @@ #include #include #include +#include #include #include namespace mlx::core::cu { // Run tasks in worker thread, synchronized with cuda stream. -class Worker { +class Worker : public std::enable_shared_from_this { public: explicit Worker(Device& d); ~Worker(); @@ -21,6 +22,9 @@ class Worker { Worker(const Worker&) = delete; Worker& operator=(const Worker&) = delete; + void start(); + void stop(); + // Add a pending |task| that will run when consumed or commited. void add_task(std::function task); diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/copy.cpp b/Source/Cmlx/mlx/mlx/backend/gpu/copy.cpp index 403d07d92..699dadd1a 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/copy.cpp +++ b/Source/Cmlx/mlx/mlx/backend/gpu/copy.cpp @@ -4,6 +4,7 @@ #include "mlx/primitives.h" #include +#include namespace mlx::core { @@ -59,20 +60,14 @@ array reshape_in_eval(const array& x, Shape shape, Stream s) { return out; } -array swapaxes_in_eval(const array& x, int axis1, int axis2) { - int ndim = x.ndim(); - if (axis1 < 0) { - axis1 += ndim; - } - if (axis2 < 0) { - axis2 += ndim; +array transpose_in_eval(const array& x, const std::vector& axes) { + Shape shape(axes.size()); + Strides strides(axes.size()); + for (int i = 0; i < axes.size(); ++i) { + shape[i] = x.shape(axes[i]); + strides[i] = x.strides(axes[i]); } - auto shape = x.shape(); - std::swap(shape[axis1], shape[axis2]); - auto strides = x.strides(); - std::swap(strides[axis1], strides[axis2]); - auto [data_size, row_contiguous, col_contiguous] = check_contiguity(shape, strides); bool contiguous = data_size == x.data_size(); @@ -86,4 +81,19 @@ array swapaxes_in_eval(const array& x, int axis1, int axis2) { return out; } +array swapaxes_in_eval(const array& x, int axis1, int axis2) { + int ndim = x.ndim(); + if (axis1 < 0) { + axis1 += ndim; + } + if (axis2 < 0) { + axis2 += ndim; + } + + std::vector axes(ndim); + std::iota(axes.begin(), axes.end(), 0); + std::swap(axes[axis1], axes[axis2]); + return transpose_in_eval(x, axes); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/copy.h b/Source/Cmlx/mlx/mlx/backend/gpu/copy.h index 6e6bc7978..8ebcd8464 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/copy.h +++ b/Source/Cmlx/mlx/mlx/backend/gpu/copy.h @@ -6,6 +6,7 @@ #include "mlx/stream.h" #include +#include namespace mlx::core { @@ -52,6 +53,7 @@ void reshape_gpu(const array& in, array& out, Stream s); // Like the normal ops but safe to call in eval_gpu. array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s); array reshape_in_eval(const array& x, Shape shape, Stream s); +array transpose_in_eval(const array& x, const std::vector& axes); array swapaxes_in_eval(const array& x, int axis1, int axis2); } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/eval.h b/Source/Cmlx/mlx/mlx/backend/gpu/eval.h index f646c2ec9..d75b30abb 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/eval.h +++ b/Source/Cmlx/mlx/mlx/backend/gpu/eval.h @@ -10,9 +10,12 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp index 0138928c0..268d6290b 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/gpu/primitives.cpp @@ -217,41 +217,6 @@ void Slice::eval_gpu(const std::vector& inputs, array& out) { slice_gpu(in, out, start_indices_, strides_, stream()); } -void SliceUpdate::eval_gpu(const std::vector& inputs, array& out) { - assert(inputs.size() == 2); - if (out.size() == 0) { - out.set_data(allocator::malloc(0)); - return; - } - - auto& in = inputs[0]; - auto& upd = inputs[1]; - - if (upd.size() == 0) { - out.copy_shared_buffer(in); - return; - } - - auto ctype = in.flags().contiguous && in.size() == in.data_size() - ? CopyType::Vector - : CopyType::General; - copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream()); - auto [data_offset, out_strides] = - prepare_slice(out, start_indices_, strides_); - - // Do copy - copy_gpu_inplace( - /* const array& src = */ upd, - /* array& dst = */ out, - /* const Shape& data_shape = */ upd.shape(), - /* const Strides& i_strides = */ upd.strides(), - /* const Strides& o_strides = */ out_strides, - /* int64_t i_offset = */ 0, - /* int64_t o_offset = */ data_offset, - /* CopyType ctype = */ CopyType::GeneralGeneral, - /* const Stream& s = */ stream()); -} - void Squeeze::eval_gpu(const std::vector& inputs, array& out) { MLX_PROFILER_RANGE("Squeeze::eval_gpu"); eval(inputs, out); diff --git a/Source/Cmlx/mlx/mlx/backend/gpu/scan.h b/Source/Cmlx/mlx/mlx/backend/gpu/scan.h index dab79c50b..a6dbc6f53 100644 --- a/Source/Cmlx/mlx/mlx/backend/gpu/scan.h +++ b/Source/Cmlx/mlx/mlx/backend/gpu/scan.h @@ -6,7 +6,7 @@ namespace mlx::core { void scan_gpu_inplace( - array in, + const array& in, array& out, Scan::ReduceType reduce_type, int axis, diff --git a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt index 4074e7b1e..ea4a995ad 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/CMakeLists.txt @@ -48,6 +48,7 @@ if(MLX_METAL_JIT) make_jit_source(softmax) make_jit_source(scan) make_jit_source(sort) + make_jit_source(searchsorted kernels/sort.h) make_jit_source( reduce kernels/reduction/reduce_all.h kernels/reduction/reduce_col.h kernels/reduction/reduce_row.h kernels/reduction/reduce_init.h) @@ -71,6 +72,7 @@ if(MLX_METAL_JIT) kernels/steel/conv/loaders/loader_channel_l.h kernels/steel/conv/loaders/loader_channel_n.h) make_jit_source(steel/conv/kernels/steel_conv) + make_jit_source(steel/conv/kernels/steel_conv_3d) make_jit_source(steel/conv/kernels/steel_conv_general kernels/steel/defines.h kernels/steel/conv/loaders/loader_general.h) @@ -78,22 +80,37 @@ if(MLX_METAL_JIT) make_jit_source(quantized kernels/quantized_utils.h) make_jit_source(fp_quantized kernels/quantized_utils.h kernels/fp8.h kernels/fp4.h) + make_jit_source(gemv) make_jit_source(gemv_masked) make_jit_source(steel/attn/kernels/steel_attention) - make_jit_source( - steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h - kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) - make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) - make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) - make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) + if(MLX_METAL_VERSION GREATER_EQUAL 400 + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 + AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) - make_jit_source(quantized_nax kernels/quantized_utils.h) - make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h - kernels/fp4.h) + make_jit_source( + steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h + kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) + make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) + make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) + + make_jit_source(quantized_nax kernels/quantized_utils.h) + make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h + kernels/fp4.h) + + make_jit_source(steel/attn/kernels/steel_attention_nax) - make_jit_source(steel/attn/kernels/steel_attention_nax) + else() + message( + WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and " + "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, " + "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). " + "Building without NAX kernels.") + target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX) + endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/allocator.cpp b/Source/Cmlx/mlx/mlx/backend/metal/allocator.cpp index 06e3d9ae4..903cce218 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/allocator.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/allocator.cpp @@ -7,6 +7,7 @@ #include #include +#include #include namespace mlx::core { @@ -24,26 +25,36 @@ void* Buffer::raw_ptr() { if (!ptr_) { return nullptr; } - return static_cast(ptr_)->contents(); + auto* buf = static_cast(ptr_); + assert(buf->storageMode() != MTL::StorageModePrivate); + return buf->contents(); +} + +bool can_reuse_alien_buffer(void* ptr) { + if (!ptr) { + return true; + } + auto* buf = static_cast(ptr); + return buf->storageMode() != MTL::StorageModePrivate; } } // namespace allocator namespace metal { -MetalAllocator::MetalAllocator() - : device_(device(mlx::core::Device::gpu).mtl_device()), +MetalAllocator::MetalAllocator(Device& d) + : device_(d.mtl_device()), + residency_sets_(d.residency_sets()), buffer_cache_( vm_page_size, [](MTL::Buffer* buf) { return buf->length(); }, [this](MTL::Buffer* buf) { if (!buf->heap()) { - residency_set_.erase(buf); + residency_sets_.erase(buf); } + auto pool = metal::new_scoped_memory_pool(); buf->release(); - }), - residency_set_(device_) { - auto pool = metal::new_scoped_memory_pool(); + }) { const auto& info = gpu::device_info(0); auto memsize = std::get(info.at("memory_size")); auto max_rec_size = @@ -52,28 +63,20 @@ MetalAllocator::MetalAllocator() block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize); gc_limit_ = std::min(static_cast(0.95 * max_rec_size), block_limit_); max_pool_size_ = block_limit_; - device(mlx::core::Device::gpu) - .set_residency_set(residency_set_.mtl_residency_set()); bool is_vm = std::get(info.at("device_name")) == "Apple Paravirtual device"; if (is_vm) { return; } - auto heap_desc = MTL::HeapDescriptor::alloc()->init(); + auto pool = metal::new_scoped_memory_pool(); + auto heap_desc = MTL::HeapDescriptor::alloc()->init()->autorelease(); heap_desc->setResourceOptions(resource_options); heap_desc->setSize(heap_size_); - heap_ = device_->newHeap(heap_desc); - heap_desc->release(); - residency_set_.insert(heap_); + heap_ = NS::TransferPtr(device_->newHeap(heap_desc)); + residency_sets_.insert(heap_.get()); } -MetalAllocator::~MetalAllocator() { - auto pool = metal::new_scoped_memory_pool(); - if (heap_) { - heap_->release(); - } - buffer_cache_.clear(); -} +MetalAllocator::~MetalAllocator() = default; size_t MetalAllocator::set_cache_limit(size_t limit) { std::unique_lock lk(mutex_); @@ -97,7 +100,7 @@ size_t MetalAllocator::get_memory_limit() { size_t MetalAllocator::set_wired_limit(size_t limit) { std::unique_lock lk(mutex_); std::swap(limit, wired_limit_); - residency_set_.resize(wired_limit_); + residency_sets_.resize(wired_limit_); return limit; }; @@ -128,8 +131,6 @@ Buffer MetalAllocator::malloc(size_t size) { if (!buf) { size_t mem_required = get_active_memory() + get_cache_memory() + size; - auto pool = metal::new_scoped_memory_pool(); - // If we have a lot of memory pressure try to reclaim memory from the cache if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) { num_resources_ -= @@ -158,7 +159,7 @@ Buffer MetalAllocator::malloc(size_t size) { lk.lock(); num_resources_++; if (!buf->heap()) { - residency_set_.insert(buf); + residency_sets_.insert(buf); } } @@ -167,7 +168,6 @@ Buffer MetalAllocator::malloc(size_t size) { // Maintain the cache below the requested limit if (get_cache_memory() > max_pool_size_) { - auto pool = metal::new_scoped_memory_pool(); num_resources_ -= buffer_cache_.release_cached_buffers( get_cache_memory() - max_pool_size_); } @@ -177,7 +177,6 @@ Buffer MetalAllocator::malloc(size_t size) { void MetalAllocator::clear_cache() { std::unique_lock lk(mutex_); - auto pool = metal::new_scoped_memory_pool(); num_resources_ -= buffer_cache_.clear(); } @@ -193,7 +192,7 @@ void MetalAllocator::free(Buffer buffer) { } else { num_resources_--; if (!buf->heap()) { - residency_set_.erase(buf); + residency_sets_.erase(buf); } lk.unlock(); auto pool = metal::new_scoped_memory_pool(); @@ -211,7 +210,7 @@ Buffer MetalAllocator::make_buffer(void* ptr, size_t size) { return Buffer{nullptr}; } std::unique_lock lk(mutex_); - residency_set_.insert(buf); + residency_sets_.insert(buf); active_memory_ += buf->length(); peak_memory_ = std::max(peak_memory_, active_memory_); num_resources_++; @@ -226,6 +225,7 @@ void MetalAllocator::release(Buffer buffer) { std::unique_lock lk(mutex_); active_memory_ -= buf->length(); num_resources_--; + residency_sets_.erase(buf); lk.unlock(); auto pool = metal::new_scoped_memory_pool(); buf->release(); @@ -235,7 +235,8 @@ MetalAllocator& allocator() { // By creating the |allocator_| on heap, the destructor of MetalAllocator // will not be called on exit and buffers in the cache will be leaked. This // can save some time at program exit. - static MetalAllocator* allocator_ = new MetalAllocator; + static MetalAllocator* allocator_ = + new MetalAllocator(device(mlx::core::Device::gpu)); return *allocator_; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/allocator.h b/Source/Cmlx/mlx/mlx/backend/metal/allocator.h index 5e177b3d3..885951364 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/allocator.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/allocator.h @@ -9,7 +9,6 @@ #include "mlx/allocator.h" #include "mlx/backend/common/buffer_cache.h" #include "mlx/backend/metal/device.h" -#include "mlx/backend/metal/resident.h" namespace mlx::core::metal { @@ -51,16 +50,18 @@ class MetalAllocator : public allocator::Allocator { // the heap, a heap can have at most heap.size() / 256 buffers. static constexpr int small_size_ = 256; static constexpr int heap_size_ = 1 << 20; - MTL::Heap* heap_; - MetalAllocator(); + + MetalAllocator(Device& d); ~MetalAllocator(); + friend MetalAllocator& allocator(); + NS::SharedPtr heap_; + ResidencySets& residency_sets_; + // Caching allocator BufferCache buffer_cache_; - ResidencySet residency_set_; - // Allocation stats size_t block_limit_; size_t gc_limit_; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp b/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp index 8c0e8c333..7d6b2a713 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/binary.cpp @@ -106,7 +106,7 @@ void binary_op_gpu_inplace( auto kernel = outputs.size() == 2 ? get_binary_two_kernel(d, kernel_name, a.dtype(), out.dtype(), op) : get_binary_kernel(d, kernel_name, a.dtype(), out.dtype(), op); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int arg_idx = 0; @@ -235,23 +235,7 @@ BINARY_GPU(Power) BINARY_GPU(Subtract) void BitwiseBinary::eval_gpu(const std::vector& inputs, array& out) { - switch (op_) { - case BitwiseBinary::And: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::Or: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::Xor: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::LeftShift: - binary_op_gpu(inputs, out, name()); - break; - case BitwiseBinary::RightShift: - binary_op_gpu(inputs, out, name()); - break; - } + binary_op_gpu(inputs, out, name()); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/metal/compiled.cpp index eb51ab750..e95d7b8d5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/compiled.cpp @@ -139,8 +139,8 @@ inline void build_kernel( os += fmt::format(" {0} index_{1} = ", idx_type, xname); if (ndim == 1) { int offset = i * ndim; - os += - fmt::format("elem_to_loc_1(pos.x, in_strides[{0}]);\n", offset); + os += fmt::format( + "elem_to_loc_1<{0}>(pos.x, in_strides[{1}]);\n", idx_type, offset); } else if (ndim == 2) { int offset = i * ndim; os += fmt::format( @@ -208,7 +208,7 @@ inline void build_kernel( " {0} tmp_{1} = ", get_type_string(x.dtype()), namer.get_name(x)); if (is_static_cast(x.primitive())) { os += fmt::format( - "static_cast<{0}>(tmp_{1});\n", + "cast_to<{0}>(tmp_{1});\n", get_type_string(x.dtype()), namer.get_name(x.inputs()[0])); } else { @@ -316,20 +316,20 @@ void Compiled::eval_gpu( /* dynamic_dims = */ false, /* use_big_index = */ false, /* work_per_thread = */ i > 3 ? 2 : 1); - if (i > 1) { - build_kernel( - kernel, - kernel_lib_ + "_strided_" + std::to_string(i) + "_large", - inputs_, - outputs_, - tape_, - is_constant_, - /* contiguous = */ false, - /* ndim = */ i, - /* dynamic_dims = */ false, - /* use_big_index = */ true, - /* work_per_thread = */ i > 3 ? 4 : 1); - } + // Generate int64_t index variant for all ndim, including ndim=1. + // Negative strides force large mode even for small arrays. + build_kernel( + kernel, + kernel_lib_ + "_strided_" + std::to_string(i) + "_large", + inputs_, + outputs_, + tape_, + is_constant_, + /* contiguous = */ false, + /* ndim = */ i, + /* dynamic_dims = */ false, + /* use_big_index = */ true, + /* work_per_thread = */ i > 3 ? 4 : 1); } build_kernel( kernel, @@ -360,11 +360,12 @@ void Compiled::eval_gpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides] = + auto [contiguous, negative_strides, shape, strides] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); - // Whether to use large index. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = + negative_strides || compiled_use_large_index(inputs, outputs, contiguous); // Get the kernel from the lib int ndim = shape.size(); @@ -389,7 +390,7 @@ void Compiled::eval_gpu( kernel_name += "_large"; } auto kernel = d.get_kernel(kernel_name, lib); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); // Put the inputs in diff --git a/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp b/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp index b4a674ff0..509552261 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/conv.cpp @@ -4,6 +4,8 @@ #include #include "mlx/backend/gpu/copy.h" +#include "mlx/backend/gpu/slicing.h" +#include "mlx/backend/metal/binary.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" @@ -19,48 +21,113 @@ namespace mlx::core { namespace { +inline array +ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { + if (x.flags().row_contiguous) { + return x; + } + auto result = contiguous_copy_gpu(x, s); + metal::get_command_encoder(s).add_temporary(result); + return result; +} + +inline int max_unfold_rows(metal::Device& d, size_t row_bytes, int total_rows) { + size_t max_buffer = d.mtl_device()->maxBufferLength(); + size_t max_rows = row_bytes == 0 ? total_rows : max_buffer / row_bytes; + // Force a smaller tile (tests, or to bound the unfold buffer further). + if (int forced = env::get_var("MLX_CONV_UNFOLD_TILE_ROWS", 0); forced > 0) { + max_rows = std::min(max_rows, static_cast(forced)); + } + if (max_rows == 0) { + std::ostringstream msg; + msg << "[conv] A single unfolding row needs " << row_bytes + << " bytes, which exceeds the maximum Metal buffer size of " + << max_buffer << " bytes."; + throw std::runtime_error(msg.str()); + } + return static_cast(std::min(max_rows, static_cast(total_rows))); +} + +inline auto winograd_padded_size(const MLXConvParams<2>& conv_params) { + int64_t pad_h = static_cast(conv_params.iS[0]) + + 2 * static_cast(conv_params.pad[0]); + int64_t pad_w = static_cast(conv_params.iS[1]) + + 2 * static_cast(conv_params.pad[1]); + int padded_h = safe_cast(6 * ceildiv(pad_h - 2, 6) + 2, "conv"); + int padded_w = safe_cast(6 * ceildiv(pad_w - 2, 6) + 2, "conv"); + return std::make_tuple(padded_h, padded_w); +} + +// Return how many rows to compute per each step. +inline int winograd_batch_step( + metal::Device& d, + const array& in, + const MLXConvParams<2>& conv_params) { + int total_n = conv_params.N; + + size_t itemsize = in.itemsize(); + auto [padded_h, padded_w] = winograd_padded_size(conv_params); + int tiles_per_n = + ceildiv(conv_params.oS[0], 6) * ceildiv(conv_params.oS[1], 6); + + // Limit of maximum memory can be used for the step. + size_t working_set = d.mtl_device()->recommendedMaxWorkingSetSize(); + if (int env_ws = env::get_var("MLX_CONV_WINOGRAD_WORKING_SET", 0); + env_ws > 0) { + working_set = env_ws; + } + size_t limit = working_set / 4 * 3; + + // Memory used by inputs. + size_t filt_bytes = + static_cast(8 * 8) * conv_params.C * conv_params.O * itemsize; + size_t io_bytes = itemsize * + (static_cast(total_n) * conv_params.iS[0] * conv_params.iS[1] * + conv_params.C + + static_cast(total_n) * conv_params.oS[0] * conv_params.oS[1] * + conv_params.O); + size_t used = io_bytes + filt_bytes; + size_t budget = limit > used ? limit - used : 0; + + // How many rows to use per step to avoid running over limit. + size_t bytes_per_n = + static_cast(padded_h) * padded_w * conv_params.C * itemsize + + static_cast(8 * 8) * tiles_per_n * + (conv_params.C + conv_params.O) * itemsize; + auto max_n = static_cast(budget / bytes_per_n); + int safe_n = static_cast(std::min(max_n, total_n)); + if (int forced = env::get_var("MLX_CONV_WINOGRAD_TILE_BATCH", 0); + forced > 0) { + return std::min(forced, safe_n); + } + + // When the budget forces tiling, each tile must carry enough gemm rows to + // amortize fixed cost, below that the implicit gemm fallback is much faster. + constexpr int min_rows_per_tile = 32; + if ((safe_n < total_n) && (safe_n * tiles_per_n < min_rows_per_tile)) { + return 0; + } + return safe_n; +} + template void explicit_gemm_conv_ND_gpu( const Stream& s, metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams& conv_params) { // Get gemm shapes - int implicit_M = out.size() / conv_params.O; - int implicit_K = wt.size() / conv_params.O; + int implicit_M = safe_cast(out.size() / conv_params.O, "conv"); + int implicit_K = safe_cast(wt.size() / conv_params.O, "conv"); int implicit_N = conv_params.O; - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate(kname, "naive_unfold_nd_", type_to_name(in_unfolded), "_", N); - auto& compute_encoder = d.get_command_encoder(s.index); + concatenate(kname, "naive_unfold_nd_", type_to_name(in), "_", N); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Reshape weight Shape wt_reshape{implicit_K, implicit_N}; @@ -71,23 +138,72 @@ void explicit_gemm_conv_ND_gpu( wt_flags.col_contiguous = true; wt_reshaped.copy_shared_buffer(wt, wt_restride, wt_flags, wt.data_size()); - // Perform gemm - std::vector copies = {in_unfolded}; - return steel_matmul( - s, - d, - /*a = */ in_unfolded, - /*b = */ wt_reshaped, - /*c = */ out, - /*M = */ implicit_M, - /*N = */ implicit_N, - /*K = */ implicit_K, - /*batch_size_out = */ 1, - /*a_cols = */ implicit_K, - /*b_cols = */ implicit_K, - /*a_transposed = */ false, - /*b_transposed = */ true, - /*copies = */ copies); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{out.strides(-2), out.strides(-1)}; + array out_2d({implicit_M, implicit_N}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = + MTL::Size(conv_params.C, implicit_K / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, implicit_N}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile}; + steel_matmul( + s, + d, + /*a = */ in_tile, + /*b = */ wt_reshaped, + /*c = */ out_tile, + /*M = */ tile_rows, + /*N = */ implicit_N, + /*K = */ implicit_K, + /*batch_size_out = */ 1, + /*a_cols = */ implicit_K, + /*b_cols = */ implicit_K, + /*a_transposed = */ false, + /*b_transposed = */ true, + /*copies = */ copies); + } } template @@ -96,14 +212,14 @@ void explicit_gemm_conv_group_ND_gpu( metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams& conv_params) { const int groups = conv_params.groups; const int C_per_group = conv_params.C / conv_params.groups; const int O_per_group = conv_params.O / conv_params.groups; // Get gemm shapes - const int implicit_M = out.size() / conv_params.O; - const int implicit_K = wt.size() / conv_params.O; + const int implicit_M = safe_cast(out.size() / conv_params.O, "conv"); + const int implicit_K = safe_cast(wt.size() / conv_params.O, "conv"); const int implicit_N = O_per_group; int kernel_size = 1; @@ -111,36 +227,12 @@ void explicit_gemm_conv_group_ND_gpu( kernel_size *= conv_params.wS[i]; } - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K * groups}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate( - kname, "naive_unfold_transpose_nd_", type_to_name(in_unfolded), "_", N); - auto& compute_encoder = d.get_command_encoder(s.index); + concatenate(kname, "naive_unfold_transpose_nd_", type_to_name(in), "_", N); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Transpose kernel weights so that we can slice them by contiguous chunks // of channel groups. @@ -152,29 +244,78 @@ void explicit_gemm_conv_group_ND_gpu( // Materialize array wt_transpose = contiguous_copy_gpu(wt_view, s); - // Perform gemm - std::vector copies = {in_unfolded, wt_transpose}; - return steel_matmul_regular( - /* const Stream& s = */ s, - /* Device& d = */ d, - /* const array& a = */ in_unfolded, - /* const array& b = */ wt_transpose, - /* array& c = */ out, - /* int M = */ implicit_M, - /* int N = */ implicit_N, - /* int K = */ implicit_K, - /* int batch_size_out = */ groups, - /* int lda = */ implicit_K * groups, - /* int ldb = */ implicit_K, - /* int ldd = */ implicit_N * groups, - /* bool transpose_a = */ false, - /* bool transpose_b = */ true, - /* std::vector& copies = */ copies, - /* Shape batch_shape = */ {1}, - /* Strides batch_strides = */ {0}, - /* int64_t A_batch_strides = */ int64_t(implicit_K), - /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, - /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{out.strides(-2), out.strides(-1)}; + array out_2d({implicit_M, conv_params.O}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * groups * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = MTL::Size( + conv_params.C, (implicit_K * groups) / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, conv_params.O}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile, wt_transpose}; + steel_matmul_regular( + /* const Stream& s = */ s, + /* Device& d = */ d, + /* const array& a = */ in_tile, + /* const array& b = */ wt_transpose, + /* array& c = */ out_tile, + /* int M = */ tile_rows, + /* int N = */ implicit_N, + /* int K = */ implicit_K, + /* int batch_size_out = */ groups, + /* int lda = */ implicit_K * groups, + /* int ldb = */ implicit_K, + /* int ldd = */ implicit_N * groups, + /* bool transpose_a = */ false, + /* bool transpose_b = */ true, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ {1}, + /* Strides batch_strides = */ {0}, + /* int64_t A_batch_strides = */ int64_t(implicit_K), + /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, + /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + } } void implicit_gemm_conv_2D_gpu( @@ -182,14 +323,17 @@ void implicit_gemm_conv_2D_gpu( metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams<2>& conv_params) { const int groups = conv_params.groups; const int C_per_group = conv_params.C / conv_params.groups; const int O_per_group = conv_params.O / conv_params.groups; // Deduce implicit gemm size - const int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1]; + const int implicit_M = safe_cast( + static_cast(conv_params.N) * conv_params.oS[0] * + conv_params.oS[1], + "conv"); const int implicit_N = O_per_group; const int implicit_K = conv_params.wS[0] * conv_params.wS[1] * C_per_group; @@ -275,7 +419,7 @@ void implicit_gemm_conv_2D_gpu( small_filter ? 's' : 'l'); // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_conv_kernel( d, kname, @@ -315,10 +459,13 @@ void implicit_gemm_conv_2D_general_gpu( metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams<2>& conv_params) { // Deduce implicit gemm size - int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1]; + int implicit_M = safe_cast( + static_cast(conv_params.N) * conv_params.oS[0] * + conv_params.oS[1], + "conv"); int implicit_N = conv_params.O; int implicit_K = conv_params.wS[0] * conv_params.wS[1] * conv_params.C; @@ -458,7 +605,7 @@ void implicit_gemm_conv_2D_general_gpu( }; // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_conv_general_kernel( d, kname, hash_name, func_consts, out, bm, bn, bk, wm, wn); compute_encoder.set_compute_pipeline_state(kernel); @@ -489,84 +636,458 @@ void implicit_gemm_conv_2D_general_gpu( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -void winograd_conv_2D_gpu( +void implicit_gemm_conv_3D_gpu( const Stream& s, metal::Device& d, const array& in, const array& wt, - array out, - const MLXConvParams<2>& conv_params, - std::vector& copies_w) { - Shape padded_shape = { - conv_params.N, - conv_params.iS[0] + 2 * conv_params.pad[0], - conv_params.iS[1] + 2 * conv_params.pad[1], - conv_params.C}; + array& out, + const MLXConvParams<3>& conv_params) { + const int groups = conv_params.groups; + const int C_per_group = conv_params.C / conv_params.groups; + const int O_per_group = conv_params.O / conv_params.groups; - padded_shape[1] = 6 * ((padded_shape[1] - 2 + 5) / 6) + 2; - padded_shape[2] = 6 * ((padded_shape[2] - 2 + 5) / 6) + 2; + // Deduce implicit gemm size + const int implicit_M = safe_cast( + static_cast(conv_params.N) * conv_params.oS[0] * + conv_params.oS[1] * conv_params.oS[2], + "conv"); + const int implicit_N = O_per_group; + const int implicit_K = + conv_params.wS[0] * conv_params.wS[1] * conv_params.wS[2] * C_per_group; - array in_padded(std::move(padded_shape), in.dtype(), nullptr, {}); + // Determine block and warp tiles + int wm = 2, wn = 2; - // Fill with zeros - array zero_arr = array(0, in.dtype()); - fill_gpu(zero_arr, in_padded, s); - copies_w.push_back(zero_arr); + int bm = implicit_M >= 8192 && C_per_group >= 64 ? 64 : 32; + int bn = (bm == 64 || implicit_N >= 64) ? 64 : 32; + int bk = 16; - // Pick input slice from padded - size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] + - conv_params.pad[1] * in_padded.strides()[2]; - array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {}); - in_padded_slice.copy_shared_buffer( - in_padded, - in_padded.strides(), - in_padded.flags(), - in_padded_slice.size(), - data_offset); - - // Copy input values into the slice - copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s); - - copies_w.push_back(in_padded_slice); - copies_w.push_back(in_padded); + if (implicit_N <= 16) { + bn = 8; + wm = 4; + wn = 1; + } - MLXConvParams<2> conv_params_updated{ - /* const int N = */ static_cast(in_padded.shape(0)), - /* const int C = */ static_cast(in_padded.shape(3)), - /* const int O = */ static_cast(wt.shape(0)), - /* const int iS[NDIM] = */ - {static_cast(in_padded.shape(1)), - static_cast(in_padded.shape(2))}, - /* const int wS[NDIM] = */ - {static_cast(wt.shape(1)), static_cast(wt.shape(2))}, - /* const int oS[NDIM] = */ - {static_cast(out.shape(1)), static_cast(out.shape(2))}, - /* const int str[NDIM] = */ {1, 1}, - /* const int pad[NDIM] = */ {0, 0}, - /* const int kdil[NDIM] = */ {1, 1}, - /* const int idil[NDIM] = */ {1, 1}, - /* const size_t in_strides[NDIM + 2] = */ - {in_padded.strides()[0], - in_padded.strides()[1], - in_padded.strides()[2], - in_padded.strides()[3]}, - /* const size_t wt_strides[NDIM + 2] = */ - {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]}, - /* const size_t out_strides[NDIM + 2] = */ - {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]}, - /* const int groups = */ 1, - /* const bool flip = */ false, + int tn = (implicit_N + bn - 1) / bn; + int tm = (implicit_M + bm - 1) / bm; + int swizzle_log = 0; + + bool small_filter = + (conv_params.wS[0] <= 16 && conv_params.wS[1] <= 16 && + conv_params.wS[2] <= 16); + + int channel_k_iters = ((C_per_group + bk - 1) / bk); + int gemm_k_iters = conv_params.wS[0] * conv_params.wS[1] * conv_params.wS[2] * + channel_k_iters; + + // Fix host side helper params + int sign = (conv_params.flip ? -1 : 1); + int ijw = conv_params.in_strides[3] * conv_params.kdil[2]; + int ijh = conv_params.in_strides[2] * conv_params.kdil[1]; + int ijd = conv_params.in_strides[1] * conv_params.kdil[0]; + + int inp_jump_w = sign * ijw; + int inp_jump_h = sign * (ijh - (conv_params.wS[2] - 1) * ijw); + int inp_jump_d = sign * + (ijd - (conv_params.wS[1] - 1) * ijh - (conv_params.wS[2] - 1) * ijw); + int inp_jump_c = bk - sign * (conv_params.wS[0] - 1) * ijd - + sign * (conv_params.wS[1] - 1) * ijh - + sign * (conv_params.wS[2] - 1) * ijw; + + // Build implicit gemm params + ImplicitGemmConv3DParams gemm_params{ + /* const int M = */ implicit_M, + /* const int N = */ implicit_N, + /* const int K = */ implicit_K, + + /* const int gemm_k_iterations = */ gemm_k_iters, + + /* const int inp_jump_w = */ inp_jump_w, + /* const int inp_jump_h = */ inp_jump_h, + /* const int inp_jump_d = */ inp_jump_d, + /* const int inp_jump_c = */ inp_jump_c, + + /* const int tiles_n = */ tn, + /* const int tiles_m = */ tm, + /* const int swizzle_log = */ swizzle_log}; + + // Determine kernel + std::string kname; + kname.reserve(64); + concatenate( + kname, + "implicit_gemm_conv_3d_", + type_to_name(out), + "_bm", + bm, + "_bn", + bn, + "_bk", + bk, + "_wm", + wm, + "_wn", + wn, + "_filter_", + small_filter ? 's' : 'l'); + + // Encode and dispatch kernel + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = + get_steel_conv_3d_kernel(d, kname, out, bm, bn, bk, wm, wn, small_filter); + compute_encoder.set_compute_pipeline_state(kernel); + + // Deduce grid launch dimensions + int tile = 1 << swizzle_log; + size_t grid_dim_y = (tm + tile - 1) / tile; + size_t grid_dim_x = tn * tile; + + MTL::Size group_dims = MTL::Size(32, wn, wm); + MTL::Size grid_dims = MTL::Size(grid_dim_x, grid_dim_y, groups); + + // Encode arrays + compute_encoder.set_input_array(in, 0); + compute_encoder.set_input_array(wt, 1); + compute_encoder.set_output_array(out, 2); + + // Encode params + compute_encoder.set_bytes(conv_params, 3); + compute_encoder.set_bytes(gemm_params, 4); + + // Launch kernel + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +void pad_and_slice_conv_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const MLXConvParams<3>& conv_params) { + // For now assume conv_params.groups == 1 + int extra_c = ((conv_params.C + 15) / 16) * 16 - conv_params.C; + int extra_o = ((conv_params.O + 15) / 16) * 16 - conv_params.O; + + // Pad function + auto pad_array = [&](const array& x, int pad_ax_first, int pad_ax_last) { + if (pad_ax_first == 0 && pad_ax_last == 0) { + return ensure_row_contiguous(x, d, s); + } + + auto xshape = x.shape(); + xshape.front() += pad_ax_first; + xshape.back() += pad_ax_last; + array x_copy(xshape, x.dtype(), nullptr, {}); + array zero(0, x.dtype()); + pad_gpu(x, zero, x_copy, {0, -1}, {0, 0}, s); + metal::get_command_encoder(s).add_temporary(x_copy); + + return x_copy; }; + // Allocate space for the intermediate output. Don't save it as a temporary + // since it will be sliced to the output so they share the buffer. + auto oshape = out.shape(); + oshape.back() += extra_o; + array intermediate(oshape, out.dtype(), nullptr, {}); + intermediate.set_data(allocator::malloc(intermediate.nbytes())); + + // Actually pad and conv + array in = pad_array(in_pre, 0, extra_c); + array wt = pad_array(wt_pre, extra_o, extra_c); + auto new_params = + MLXConvParams<3>::with_padded_channels(conv_params, extra_o, extra_c); + implicit_gemm_conv_3D_gpu(s, d, in, wt, intermediate, new_params); + + // Slice out + out.copy_shared_buffer( + intermediate, intermediate.strides(), {0}, intermediate.data_size()); +} + +void conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const std::vector& padding, + const std::vector& wt_strides, + const std::vector& wt_dilation, + const std::vector& in_dilation, + const int groups, + bool flip, + std::vector& copies); + +void small_kd_conv_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& conv_params, + std::vector& copies) { + const int H = conv_params.iS[1]; + const int W = conv_params.iS[2]; + const int C = conv_params.C; + const int O = conv_params.O; + const int KD = conv_params.wS[0]; + const int KH = conv_params.wS[1]; + const int KW = conv_params.wS[2]; + const int OD = conv_params.oS[0]; + const int OH = conv_params.oS[1]; + const int OW = conv_params.oS[2]; + + array acc({OD, OH, OW, O}, out.dtype(), nullptr, {}); + for (int kd = 0; kd < KD; ++kd) { + array in_2d({OD, H, W, C}, in.dtype(), nullptr, {}); + in_2d.copy_shared_buffer( + in, + {static_cast(H) * W * C, + static_cast(W) * C, + static_cast(C), + 1}, + {true, true, false}, + static_cast(OD) * H * W * C, + static_cast(kd) * H * W * C); + + // The 2D conv only flips the last two kernel axes, so mirror the depth + // axis here when the convolution is flipped. + const int kd_wt = conv_params.flip ? KD - 1 - kd : kd; + + array wt_2d({O, KH, KW, C}, wt.dtype(), nullptr, {}); + wt_2d.copy_shared_buffer( + wt, + {static_cast(KD) * KH * KW * C, + static_cast(KW) * C, + static_cast(C), + 1}, + {false, false, false}, + static_cast(O - 1) * KD * KH * KW * C + + static_cast(KH) * KW * C, + static_cast(kd_wt) * KH * KW * C); + + array conv_out({OD, OH, OW, O}, out.dtype(), nullptr, {}); + conv_2D_gpu( + s, + d, + in_2d, + wt_2d, + conv_out, + {conv_params.pad[1], conv_params.pad[2]}, + {conv_params.str[1], conv_params.str[2]}, + {conv_params.kdil[1], conv_params.kdil[2]}, + {conv_params.idil[1], conv_params.idil[2]}, + /* groups = */ 1, + conv_params.flip, + copies); + + if (kd == 0) { + acc = conv_out; + } else { + binary_op_gpu_inplace({acc, conv_out}, acc, "Add", s); + copies.push_back(conv_out); + } + } + + // Output shape is [1, OD, OH, OW, O]. + out.copy_shared_buffer( + acc, + {static_cast(OD) * OH * OW * O, + static_cast(OH) * OW * O, + static_cast(OW) * O, + static_cast(O), + 1}, + {true, true, false}, + static_cast(OD) * OH * OW * O, + 0); +} + +// A stride-2, kernel-2 transposed convolution has exactly one valid kernel +// phase for every output coordinate. The explicit unfold path nevertheless +// materializes all eight phases and fills seven of them with zeros. Compute +// each phase as a small regular GEMM instead. This path is intentionally +// narrow: other transposed-convolution configurations retain the general +// implementation below. +bool is_stride_two_conv_transpose_3D(const MLXConvParams<3>& p) { + return p.groups == 1 && p.flip && p.str[0] == 1 && p.str[1] == 1 && + p.str[2] == 1 && p.idil[0] == 2 && p.idil[1] == 2 && p.idil[2] == 2 && + p.kdil[0] == 1 && p.kdil[1] == 1 && p.kdil[2] == 1 && p.wS[0] == 2 && + p.wS[1] == 2 && p.wS[2] == 2 && p.pad[0] == 1 && p.pad[1] == 1 && + p.pad[2] == 1 && static_cast(p.oS[0]) == 2LL * p.iS[0] && + static_cast(p.oS[1]) == 2LL * p.iS[1] && + static_cast(p.oS[2]) == 2LL * p.iS[2]; +} + +void stride_two_conv_transpose_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& p, + std::vector& copies) { + constexpr int kernel_volume = 8; + const int C = p.C; + const int O = p.O; + + // The input and weight are contiguous by the time this helper is called. + // Every phase covers the complete input volume; the phase bit only selects + // the interleaved output coordinates and the corresponding weight slice. + for (int phase = 0; phase < kernel_volume; ++phase) { + const int pd = (phase >> 2) & 1; + const int ph = (phase >> 1) & 1; + const int pw = phase & 1; + const int D = p.iS[0]; + const int H = p.iS[1]; + const int W = p.iS[2]; + + const int M = safe_cast(static_cast(p.N) * D * H * W, "conv"); + + // The weight is [O, 2, 2, 2, C]. Present one spatial phase as a [C, O] + // matrix with the layout expected by a transposed Steel GEMM. + array wt_phase({C, O}, wt.dtype(), nullptr, {}); + array::Flags wt_flags = wt.flags(); + wt_flags.contiguous = false; + wt_flags.row_contiguous = false; + wt_flags.col_contiguous = true; + wt_phase.copy_shared_buffer( + wt, + {1, wt.strides(0)}, + wt_flags, + wt.data_size(), + static_cast(pd * 4 + ph * 2 + pw) * C); + + array in_matrix({M, C}, in.dtype(), nullptr, {}); + in_matrix.copy_shared_buffer(in, {C, 1}, in.flags(), in.data_size()); + + array phase_out({M, O}, out.dtype(), nullptr, {}); + phase_out.set_data(allocator::malloc(phase_out.nbytes())); + + std::vector gemm_copies = {in_matrix, wt_phase}; + steel_matmul( + s, + d, + /* a = */ in_matrix, + /* b = */ wt_phase, + /* out = */ phase_out, + /* M = */ M, + /* N = */ O, + /* K = */ C, + /* batch_size_out = */ 1, + /* lda = */ C, + /* ldb = */ kernel_volume * C, + /* a_transposed = */ false, + /* b_transposed = */ true, + /* copies = */ gemm_copies); + + Shape phase_out_shape{p.N, D, H, W, O}; + array phase_out_nd(phase_out_shape, out.dtype(), nullptr, {}); + phase_out_nd.copy_shared_buffer( + phase_out, + make_contiguous_strides(phase_out_shape), + phase_out.flags(), + phase_out.data_size()); + + Strides out_phase_strides = out.strides(); + out_phase_strides[1] *= 2; + out_phase_strides[2] *= 2; + out_phase_strides[3] *= 2; + array::Flags out_phase_flags = out.flags(); + out_phase_flags.contiguous = false; + out_phase_flags.row_contiguous = false; + out_phase_flags.col_contiguous = false; + array out_phase_view(phase_out_shape, out.dtype(), nullptr, {}); + out_phase_view.copy_shared_buffer( + out, + out_phase_strides, + out_phase_flags, + out.data_size(), + pd * out.strides(1) + ph * out.strides(2) + pw * out.strides(3)); + + copy_gpu_inplace(phase_out_nd, out_phase_view, CopyType::GeneralGeneral, s); + copies.push_back(phase_out); + copies.push_back(phase_out_nd); + copies.push_back(out_phase_view); + } +} + +void dispatch_conv_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const MLXConvParams<3>& conv_params, + std::vector& copies) { + bool is_idil_one = conv_params.idil[0] == 1 && conv_params.idil[1] == 1 && + conv_params.idil[2] == 1; + const int C_per_group = conv_params.C / conv_params.groups; + const int O_per_group = conv_params.O / conv_params.groups; + + bool mod16_channels = + C_per_group % 16 == 0 && (O_per_group <= 16 || O_per_group % 16 == 0); + + // Check if we can do implicit gemm but the channels are not divisible by 16 + // so we can pad and slice. + // + // We check it first because it doesn't need contiguous inputs and it needs + // different output allocation. + if (is_idil_one && !mod16_channels && conv_params.groups == 1) { + return pad_and_slice_conv_3D_gpu(s, d, in_pre, wt_pre, out, conv_params); + } + + // Allocate the output and ensure contiguous inputs + out.set_data(allocator::malloc(out.nbytes())); + auto in = ensure_row_contiguous(in_pre, d, s); + auto wt = ensure_row_contiguous(wt_pre, d, s); + + if (is_stride_two_conv_transpose_3D(conv_params)) { + return stride_two_conv_transpose_3D_gpu( + s, d, in, wt, out, conv_params, copies); + } + + // Decompose 3D conv to per-frame 2D convs + constexpr int kSmallKdLimit3D = 7; + if (is_idil_one && mod16_channels && conv_params.groups == 1 && + conv_params.N == 1 && conv_params.wS[0] <= kSmallKdLimit3D && + conv_params.str[0] == 1 && conv_params.kdil[0] == 1 && + conv_params.pad[0] == 0) { + return small_kd_conv_3D_gpu(s, d, in, wt, out, conv_params, copies); + } + + // Perform the implicit gemm + if (is_idil_one && mod16_channels) { + return implicit_gemm_conv_3D_gpu(s, d, in, wt, out, conv_params); + } + + // Explicit gemms where we unfold and do a matmul + // (separate one for groups > 1) + if (conv_params.groups > 1) { + return explicit_gemm_conv_group_ND_gpu(s, d, in, wt, out, conv_params); + } + return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params); +} + +void winograd_conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<2>& conv_params, + std::vector& copies_w, + int n_step) { int O_c = conv_params.O; int C_c = conv_params.C; + auto [padded_h, padded_w] = winograd_padded_size(conv_params); - int N_tiles_n = conv_params.N; - int N_tiles_h = (conv_params.oS[0] + 5) / 6; - int N_tiles_w = (conv_params.oS[1] + 5) / 6; - int N_tiles = N_tiles_n * N_tiles_h * N_tiles_w; + int N_tiles_h = ceildiv(conv_params.oS[0], 6); + int N_tiles_w = ceildiv(conv_params.oS[1], 6); + int tiles_per_n = N_tiles_h * N_tiles_w; - // Do filter transform + // Do filter transform. Shape filt_wg_shape = {8 * 8, conv_params.C, conv_params.O}; array filt_wg(std::move(filt_wg_shape), wt.dtype(), nullptr, {}); filt_wg.set_data(allocator::malloc(filt_wg.nbytes())); @@ -582,7 +1103,7 @@ void winograd_conv_2D_gpu( type_to_name(out), "_bc", bc); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); compute_encoder.set_compute_pipeline_state(kernel); @@ -598,88 +1119,187 @@ void winograd_conv_2D_gpu( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } - // Do input transform - Shape inp_wg_shape = {8 * 8, N_tiles, conv_params.C}; - array inp_wg(std::move(inp_wg_shape), in.dtype(), nullptr, {}); + // Scratch space reused by every batch tile. + array inp_wg({8 * 8, n_step * tiles_per_n, C_c}, in.dtype(), nullptr, {}); inp_wg.set_data(allocator::malloc(inp_wg.nbytes())); copies_w.push_back(inp_wg); - { - int bc = 32; - int wm = 2; - int wn = 2; - std::string kname; - kname.reserve(32); - concatenate( - kname, - "winograd_conv_2d_input_transform_", - type_to_name(out), - "_bc", - bc); - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in_padded, 0); - compute_encoder.set_output_array(inp_wg, 1); - - compute_encoder.set_bytes(conv_params_updated, 2); - MTL::Size group_dims = MTL::Size(32, wn, wm); - MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n); - - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - } - - // Do batched gemm - Shape out_wg_shape = {8 * 8, N_tiles, conv_params.O}; - array out_wg(std::move(out_wg_shape), in.dtype(), nullptr, {}); + array out_wg({8 * 8, n_step * tiles_per_n, O_c}, in.dtype(), nullptr, {}); out_wg.set_data(allocator::malloc(out_wg.nbytes())); copies_w.push_back(out_wg); - { - std::vector empty_copies; - steel_matmul( - s, - d, - /*a = */ inp_wg, - /*b = */ filt_wg, - /*c = */ out_wg, - /*M = */ N_tiles, - /*N = */ conv_params.O, - /*K = */ conv_params.C, - /*batch_size_out = */ 8 * 8, - /*a_cols = */ conv_params.C, - /*b_cols = */ conv_params.O, - /*a_transposed = */ false, - /*b_transposed = */ false, - /*copies = */ empty_copies); - } - // Do output transform - { - int bc = 32; - int wm = 2; - int wn = 2; - std::string kname; - kname.reserve(32); - concatenate( - kname, - "winograd_conv_2d_output_transform_", - type_to_name(out), - "_bo", - bc); - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); + array in_padded({n_step, padded_h, padded_w, C_c}, in.dtype(), nullptr, {}); + copies_w.push_back(in_padded); - compute_encoder.set_input_array(out_wg, 0); - compute_encoder.set_output_array(out, 1); + // Fill padding with zeros. + array zero_arr = array(0, in.dtype()); + fill_gpu(zero_arr, in_padded, s); + copies_w.push_back(zero_arr); - compute_encoder.set_bytes(conv_params_updated, 2); + int64_t pad_offset = + static_cast(conv_params.pad[0]) * in_padded.strides()[1] + + static_cast(conv_params.pad[1]) * in_padded.strides()[2]; + + // Loop over all rows. + for (int n_offset = 0; n_offset < conv_params.N; n_offset += n_step) { + int tile_n = std::min(n_step, conv_params.N - n_offset); + int N_tiles = tile_n * tiles_per_n; + + // Views for current step. + array in_tile( + {tile_n, conv_params.iS[0], conv_params.iS[1], C_c}, + in.dtype(), + nullptr, + {}); + in_tile.copy_shared_buffer( + in, + in.strides(), + in.flags(), + in_tile.size(), + n_offset * in.strides()[0]); + + array out_tile( + {tile_n, conv_params.oS[0], conv_params.oS[1], O_c}, + out.dtype(), + nullptr, + {}); + out_tile.copy_shared_buffer( + out, + out.strides(), + out.flags(), + out_tile.size(), + n_offset * out.strides()[0]); + + array in_padded_slice(in_tile.shape(), in_padded.dtype(), nullptr, {}); + in_padded_slice.copy_shared_buffer( + in_padded, + in_padded.strides(), + in_padded.flags(), + in_padded_slice.size(), + pad_offset); + + // Copy input values into the slice. + copy_gpu_inplace(in_tile, in_padded_slice, CopyType::GeneralGeneral, s); + copies_w.push_back(in_padded_slice); + + MLXConvParams<2> conv_params_updated{ + /* const int N = */ tile_n, + /* const int C = */ C_c, + /* const int O = */ O_c, + /* const int iS[NDIM] = */ {padded_h, padded_w}, + /* const int wS[NDIM] = */ + {static_cast(wt.shape(1)), static_cast(wt.shape(2))}, + /* const int oS[NDIM] = */ + {static_cast(out.shape(1)), static_cast(out.shape(2))}, + /* const int str[NDIM] = */ {1, 1}, + /* const int pad[NDIM] = */ {0, 0}, + /* const int kdil[NDIM] = */ {1, 1}, + /* const int idil[NDIM] = */ {1, 1}, + /* const size_t in_strides[NDIM + 2] = */ + {in_padded.strides()[0], + in_padded.strides()[1], + in_padded.strides()[2], + in_padded.strides()[3]}, + /* const size_t wt_strides[NDIM + 2] = */ + {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]}, + /* const size_t out_strides[NDIM + 2] = */ + {out.strides()[0], + out.strides()[1], + out.strides()[2], + out.strides()[3]}, + /* const int groups = */ 1, + /* const bool flip = */ false, + }; + + // Do input transform, result layout is (8 x 8 x N_tiles x channels). + { + int bc = 32; + int wm = 2; + int wn = 2; + std::string kname; + kname.reserve(32); + concatenate( + kname, + "winograd_conv_2d_input_transform_", + type_to_name(out), + "_bc", + bc); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel(kname); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(in_padded, 0); + compute_encoder.set_output_array(inp_wg, 1); + + compute_encoder.set_bytes(conv_params_updated, 2); + + MTL::Size group_dims = MTL::Size(32, wn, wm); + MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, tile_n); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } - MTL::Size group_dims = MTL::Size(32, wn, wm); - MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, N_tiles_n); + // Do batched gemm. + { + array inp_wg_tile({8 * 8, N_tiles, C_c}, inp_wg.dtype(), nullptr, {}); + inp_wg_tile.copy_shared_buffer( + inp_wg, + {static_cast(N_tiles) * C_c, C_c, 1}, + inp_wg.flags(), + inp_wg_tile.size()); + + array out_wg_tile({8 * 8, N_tiles, O_c}, out_wg.dtype(), nullptr, {}); + out_wg_tile.copy_shared_buffer( + out_wg, + {static_cast(N_tiles) * O_c, O_c, 1}, + out_wg.flags(), + out_wg_tile.size()); + + std::vector empty_copies; + steel_matmul( + s, + d, + /*a = */ inp_wg_tile, + /*b = */ filt_wg, + /*c = */ out_wg_tile, + /*M = */ N_tiles, + /*N = */ O_c, + /*K = */ C_c, + /*batch_size_out = */ 8 * 8, + /*a_cols = */ C_c, + /*b_cols = */ O_c, + /*a_transposed = */ false, + /*b_transposed = */ false, + /*copies = */ empty_copies); + } - compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + // Do output transform. + { + int bc = 32; + int wm = 2; + int wn = 2; + std::string kname; + kname.reserve(32); + concatenate( + kname, + "winograd_conv_2d_output_transform_", + type_to_name(out), + "_bo", + bc); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel(kname); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(out_wg, 0); + compute_encoder.set_output_array(out_tile, 1); + + compute_encoder.set_bytes(conv_params_updated, 2); + + MTL::Size group_dims = MTL::Size(32, wn, wm); + MTL::Size grid_dims = MTL::Size(N_tiles_w, N_tiles_h, tile_n); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + } } } @@ -688,7 +1308,7 @@ void depthwise_conv_2D_gpu( metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams<2>& conv_params) { std::string base_name; base_name.reserve(32); @@ -728,7 +1348,7 @@ void depthwise_conv_2D_gpu( "_tgp_w_", tw, "_do_flip_", do_flip ? 't' : 'n'); // clang-format on - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(base_name, hash_name, func_consts); compute_encoder.set_compute_pipeline_state(kernel); @@ -740,31 +1360,66 @@ void depthwise_conv_2D_gpu( MTL::Size group_dims = MTL::Size(tc, tw, th); MTL::Size grid_dims = MTL::Size( - conv_params.C / tc, conv_params.oS[1] / tw, (conv_params.oS[0] / th) * N); + conv_params.C / tc, + (conv_params.oS[1] + tw - 1) / tw, + ((conv_params.oS[0] + th - 1) / th) * N); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +void pad_in_channels_conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const MLXConvParams<2>& conv_params) { + // Only the input channels are padded, so the kernel writes straight into out + // and the output stays contiguous. Assumes conv_params.groups == 1. + int extra_c = ((conv_params.C + 15) / 16) * 16 - conv_params.C; + + // Pad function + auto pad_array = [&](const array& x) { + auto xshape = x.shape(); + xshape.back() += extra_c; + array x_copy(xshape, x.dtype(), nullptr, {}); + array zero(0, x.dtype()); + pad_gpu(x, zero, x_copy, {-1}, {0}, s); + metal::get_command_encoder(s).add_temporary(x_copy); + + return x_copy; + }; + + array in = pad_array(in_pre); + array wt = pad_array(wt_pre); + auto new_params = + MLXConvParams<2>::with_padded_channels(conv_params, 0, extra_c); + implicit_gemm_conv_2D_gpu(s, d, in, wt, out, new_params); +} + void dispatch_conv_2D_gpu( const Stream& s, metal::Device& d, const array& in, const array& wt, - array out, + array& out, const MLXConvParams<2>& conv_params, std::vector& copies) { bool is_stride_one = conv_params.str[0] == 1 && conv_params.str[1] == 1; bool is_kdil_one = conv_params.kdil[0] == 1 && conv_params.kdil[1] == 1; bool is_idil_one = conv_params.idil[0] == 1 && conv_params.idil[1] == 1; - if (is_idil_one && conv_params.groups > 1) { + if (conv_params.groups > 1) { const int C_per_group = conv_params.C / conv_params.groups; const int O_per_group = conv_params.O / conv_params.groups; + if (!is_idil_one) { + return explicit_gemm_conv_group_ND_gpu(s, d, in, wt, out, conv_params); + } + if (C_per_group == 1 && O_per_group == 1 && is_kdil_one && conv_params.wS[0] <= 7 && conv_params.wS[1] <= 7 && conv_params.str[0] <= 2 && conv_params.str[1] <= 2 && - conv_params.oS[0] % 8 == 0 && conv_params.oS[1] % 8 == 0 && conv_params.wt_strides[1] == conv_params.wS[1] && conv_params.C % 16 == 0 && conv_params.C == conv_params.O) { return depthwise_conv_2D_gpu(s, d, in, wt, out, conv_params); @@ -779,21 +1434,39 @@ void dispatch_conv_2D_gpu( } // Direct to winograd conv - bool inp_large = - (conv_params.N * conv_params.iS[0] * conv_params.iS[1]) >= 4096; + bool inp_large = (static_cast(conv_params.N) * conv_params.iS[0] * + conv_params.iS[1]) >= 4096; bool channels_large = (conv_params.C + conv_params.O) >= 256; - bool out_large = - (conv_params.N * conv_params.oS[0] * conv_params.oS[1]) >= 256; + bool out_large = (static_cast(conv_params.N) * conv_params.oS[0] * + conv_params.oS[1]) >= 256; if (!conv_params.flip && is_stride_one && is_kdil_one && is_idil_one && conv_params.wS[0] == 3 && conv_params.wS[1] == 3 && conv_params.C % 32 == 0 && conv_params.O % 32 == 0 && inp_large && channels_large) { - return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies); + // Only use winograd conv when having enough memory. + if (int n_step = winograd_batch_step(d, in, conv_params); n_step > 0) { + return winograd_conv_2D_gpu( + s, d, in, wt, out, conv_params, copies, n_step); + } + } + + // Whether the specialized implicit gemm kernel can take the channels as-is. + bool specialized_channels = (conv_params.C <= 4 || conv_params.C % 16 == 0) && + (conv_params.O <= 16 || conv_params.O % 16 == 0); + + // Pad the input channels up to a multiple of 16 to use the faster specialized + // kernel instead of the general one. Only worth the padded work for stride-1 + // convs with a large enough output and kernel, and only when the output + // channels are already aligned so the kernel writes contiguously into out + // (padding those too needs a copy that can cost more than the kernel saves). + bool out_channels_aligned = conv_params.O <= 16 || conv_params.O % 16 == 0; + if (is_idil_one && is_stride_one && out_large && !specialized_channels && + out_channels_aligned && (conv_params.wS[0] * conv_params.wS[1]) >= 9) { + return pad_in_channels_conv_2D_gpu(s, d, in, wt, out, conv_params); } // Direct to implicit gemm conv - if (is_idil_one && (conv_params.C <= 4 || conv_params.C % 16 == 0) && - (conv_params.O <= 16 || conv_params.O % 16 == 0)) { + if (is_idil_one && specialized_channels) { return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params); } @@ -811,22 +1484,18 @@ void depthwise_conv_1D_gpu( const Stream& s, metal::Device& d, const array& in, - array wt, - array out) { + const array& wt, + array& out) { bool large = in.size() > INT32_MAX || in.data_size() > INT32_MAX; std::string base_name; base_name.reserve(32); concatenate( base_name, "depthwise_conv_1d_", - large ? "_large" : "", - type_to_name(out)); + type_to_name(out), + large ? "_large" : ""); - if (!wt.flags().row_contiguous) { - wt = contiguous_copy_gpu(wt, s); - d.add_temporary(wt, s.index); - } - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(base_name); compute_encoder.set_compute_pipeline_state(kernel); @@ -860,9 +1529,9 @@ void depthwise_conv_1D_gpu( void conv_1D_gpu( const Stream& s, metal::Device& d, - const array& in, - const array& wt, - array out, + const array& in_pre, + const array& wt_pre, + array& out, const std::vector& padding, const std::vector& wt_strides, const std::vector& wt_dilation, @@ -870,6 +1539,11 @@ void conv_1D_gpu( int groups, bool flip, std::vector& copies) { + // Allocate space and ensure weights are contiguous + out.set_data(allocator::malloc(out.nbytes())); + auto in = ensure_row_contiguous(in_pre, d, s); + auto wt = ensure_row_contiguous(wt_pre, d, s); + bool is_idil_one = in_dilation[0] == 1; int C = in.shape(2); int O = wt.shape(0); @@ -942,9 +1616,9 @@ void conv_1D_gpu( void conv_2D_gpu( const Stream& s, metal::Device& d, - const array& in, - const array& wt, - array out, + const array& in_pre, + const array& wt_pre, + array& out, const std::vector& padding, const std::vector& wt_strides, const std::vector& wt_dilation, @@ -952,6 +1626,11 @@ void conv_2D_gpu( const int groups, bool flip, std::vector& copies) { + // Allocate space and ensure weights are contiguous + out.set_data(allocator::malloc(out.nbytes())); + auto in = ensure_row_contiguous(in_pre, d, s); + auto wt = ensure_row_contiguous(wt_pre, d, s); + // Make conv params MLXConvParams<2> conv_params{ /* const int N = */ static_cast(in.shape(0)), @@ -989,8 +1668,20 @@ void conv_3D_gpu( const std::vector& wt_strides, const std::vector& wt_dilation, const std::vector& in_dilation, + int groups, bool flip, std::vector& copies) { + // We will use the contiguous strides for the conv params because that is + // what the rest of the code expects. + constexpr int NDIM = 3; + int64_t in_arr_strides[NDIM + 2]; + int64_t wt_arr_strides[NDIM + 2]; + in_arr_strides[NDIM + 1] = wt_arr_strides[NDIM + 1] = 1; + for (int i = NDIM; i >= 0; i--) { + in_arr_strides[i] = in_arr_strides[i + 1] * in.shape(i + 1); + wt_arr_strides[i] = wt_arr_strides[i + 1] * wt.shape(i + 1); + } + // Make conv params MLXConvParams<3> conv_params{ /* const int N = */ static_cast(in.shape(0)), @@ -1015,48 +1706,42 @@ void conv_3D_gpu( /* const int idil[NDIM] = */ {in_dilation[0], in_dilation[1], in_dilation[2]}, /* const size_t in_strides[NDIM + 2] = */ - {in.strides()[0], - in.strides()[1], - in.strides()[2], - in.strides()[3], - in.strides()[4]}, + {in_arr_strides[0], + in_arr_strides[1], + in_arr_strides[2], + in_arr_strides[3], + in_arr_strides[4]}, /* const size_t wt_strides[NDIM + 2] = */ - {wt.strides()[0], - wt.strides()[1], - wt.strides()[2], - wt.strides()[3], - wt.strides()[4]}, + {wt_arr_strides[0], + wt_arr_strides[1], + wt_arr_strides[2], + wt_arr_strides[3], + wt_arr_strides[4]}, /* const size_t out_strides[NDIM + 2] = */ - {out.strides()[0], - out.strides()[1], - out.strides()[2], - out.strides()[3], - out.strides()[4]}, - /* const int groups = */ 1, + {out.strides(0), + out.strides(1), + out.strides(2), + out.strides(3), + out.strides(4)}, + /* const int groups = */ groups, /* const bool flip = */ flip, }; - return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params); + return dispatch_conv_3D_gpu(s, d, in, wt, out, conv_params, copies); } } // namespace void Convolution::eval_gpu(const std::vector& inputs, array& out) { - out.set_data(allocator::malloc(out.nbytes())); auto& s = stream(); auto& d = metal::device(s.device); - // Ensure contiguity + // Intermediates that are put here will be added to the command encoder as + // temporaries. std::vector copies; - auto in = inputs[0]; - auto wt = inputs[1]; - if (!in.flags().row_contiguous) { - in = contiguous_copy_gpu(in, s); - copies.push_back(in); - } - if (!wt.flags().row_contiguous) { - wt = contiguous_copy_gpu(wt, s); - copies.push_back(wt); - } + + // Some shortcuts for brevity + const array& in = inputs[0]; + const array& wt = inputs[1]; // 3D conv if (out.ndim() == 5) { @@ -1070,6 +1755,7 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out) { kernel_strides_, kernel_dilation_, input_dilation_, + groups_, flip_, copies); } @@ -1112,7 +1798,9 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out) { } // Record copies - d.add_temporaries(std::move(copies), s.index); + if (!copies.empty()) { + metal::get_command_encoder(s).add_temporaries(std::move(copies)); + } } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/copy.cpp b/Source/Cmlx/mlx/mlx/backend/metal/copy.cpp index 6b791289c..97c5fccbc 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/copy.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/copy.cpp @@ -107,14 +107,13 @@ void copy_gpu_inplace( auto kernel = dynamic ? get_dynamic_copy_kernel(d, kernel_name, in, out) : get_copy_kernel(d, kernel_name, in, out); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - bool donate_in = in.data_shared_ptr() == nullptr; inp_offset *= size_of(in.dtype()); out_offset *= size_of(out.dtype()); - compute_encoder.set_input_array(donate_in ? out : in, 0, inp_offset); + compute_encoder.set_input_array(in, 0, inp_offset); compute_encoder.set_output_array(out, 1, out_offset); auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); @@ -191,7 +190,7 @@ void fill_gpu(const array& val, array& out, const Stream& s) { std::string kernel_name = large ? "s2" : (work_per_thread > 1 ? "sn" : "s"); concatenate(kernel_name, "_copy", type_to_name(val), type_to_name(out)); auto kernel = get_copy_kernel(d, kernel_name, val, out); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(val, 0); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/custom_kernel.cpp b/Source/Cmlx/mlx/mlx/backend/metal/custom_kernel.cpp index eec6645bd..b73dd72de 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/custom_kernel.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/custom_kernel.cpp @@ -1,328 +1,14 @@ // Copyright © 2024 Apple Inc. -#include -#include - -#include "mlx/backend/common/compiled.h" +#include "mlx/backend/common/metal_kernel.h" #include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/jit/includes.h" #include "mlx/backend/metal/utils.h" -#include "mlx/fast.h" #include "mlx/fast_primitives.h" -#include "mlx/utils.h" - -namespace mlx::core::fast { - -struct CustomKernelCache { - std::unordered_map libraries; -}; - -static CustomKernelCache& cache() { - static CustomKernelCache cache_; - return cache_; -}; - -std::string write_signature( - std::string func_name, - const std::string& header, - const std::string& source, - const std::vector& input_names, - const std::vector& inputs, - const std::vector& output_names, - const std::vector& output_dtypes, - const std::vector>& template_args, - const std::vector& attributes, - const std::vector>& shape_infos, - bool atomic_outputs) { - std::string kernel_source; - kernel_source.reserve(header.size() + source.size() + 16384); - kernel_source += header; - // Auto-generate a function signature based on `template_args` - // and the dtype/shape of the arrays passed as `inputs`. - if (!template_args.empty()) { - kernel_source += "template <"; - int i = 0; - for (const auto& [name, arg] : template_args) { - std::string param_type; - if (std::holds_alternative(arg)) { - param_type = "int"; - } else if (std::holds_alternative(arg)) { - param_type = "bool"; - } else if (std::holds_alternative(arg)) { - param_type = "typename"; - } - if (i > 0) { - kernel_source += ", "; - } - kernel_source += param_type; - kernel_source += " "; - kernel_source += name; - i++; - } - kernel_source += ">\n"; - } - kernel_source += "[[kernel]] void "; - kernel_source += func_name; - kernel_source += "(\n"; - - int index = 0; - constexpr int max_constant_array_size = 8; - // Add inputs - for (int i = 0; i < inputs.size(); ++i) { - const auto& name = input_names[i]; - const auto& arr = inputs[i]; - auto dtype = get_type_string(arr.dtype()); - std::string location = - arr.size() < max_constant_array_size ? "constant" : "device"; - std::string ref = arr.ndim() == 0 ? "&" : "*"; - kernel_source += " const "; - kernel_source += location; - kernel_source += " "; - kernel_source += dtype; - kernel_source += ref; - kernel_source += " "; - kernel_source += name; - kernel_source += " [[buffer("; - kernel_source += std::to_string(index); - kernel_source += ")]],\n"; - index++; - // Add input shape, strides and ndim if present in the source - if (arr.ndim() > 0) { - if (std::get<0>(shape_infos[i])) { - kernel_source += - (" const constant int* " + name + "_shape [[buffer(" + - std::to_string(index) + ")]],\n"); - index++; - } - if (std::get<1>(shape_infos[i])) { - kernel_source += - (" const constant int64_t* " + name + "_strides [[buffer(" + - std::to_string(index) + ")]],\n"); - index++; - } - if (std::get<2>(shape_infos[i])) { - kernel_source += - (" const constant int& " + name + "_ndim [[buffer(" + - std::to_string(index) + ")]],\n"); - index++; - } - } - } - // Add outputs - for (int i = 0; i < output_names.size(); ++i) { - const auto& name = output_names[i]; - const auto& dtype = output_dtypes[i]; - kernel_source += " device "; - auto type_string = get_type_string(dtype); - if (atomic_outputs) { - kernel_source += "atomic<"; - } - kernel_source += type_string; - if (atomic_outputs) { - kernel_source += ">"; - } - kernel_source += "* "; - kernel_source += name; - kernel_source += " [[buffer("; - kernel_source += std::to_string(index); - kernel_source += ")]]"; - if (index < inputs.size() + output_names.size() - 1 || - attributes.size() > 0) { - kernel_source += ",\n"; - } else { - kernel_source += ") {\n"; - } - index++; - } - - index = 0; - for (const auto& attr : attributes) { - kernel_source += attr; - if (index < attributes.size() - 1) { - kernel_source += ",\n"; - } else { - kernel_source += ") {\n"; - } - index++; - } - kernel_source += source; - kernel_source += "\n}\n"; - return kernel_source; -} -std::string write_template( - const std::vector>& template_args) { - std::ostringstream template_def; - template_def << "<"; - int i = 0; - for (const auto& [name, arg] : template_args) { - if (i > 0) { - template_def << ", "; - } - if (std::holds_alternative(arg)) { - template_def << std::get(arg); - } else if (std::holds_alternative(arg)) { - template_def << std::get(arg); - } else if (std::holds_alternative(arg)) { - template_def << get_type_string(std::get(arg)); - } - i++; - } - template_def << ">"; - return template_def.str(); -} - -CustomKernelFunction metal_kernel( - const std::string& name, - const std::vector& input_names, - const std::vector& output_names, - const std::string& source, - const std::string& header /* = "" */, - bool ensure_row_contiguous /* = true */, - bool atomic_outputs /* = false */) { - if (output_names.empty()) { - throw std::invalid_argument( - "[metal_kernel] Must specify at least one output."); - } - std::vector> shape_infos; - for (auto& n : input_names) { - std::tuple shape_info; - std::get<0>(shape_info) = source.find(n + "_shape") != std::string::npos; - std::get<1>(shape_info) = source.find(n + "_strides") != std::string::npos; - std::get<2>(shape_info) = source.find(n + "_ndim") != std::string::npos; - shape_infos.push_back(shape_info); - } - const std::vector> metal_attributes = { - {"dispatch_quadgroups_per_threadgroup", "uint"}, - {"dispatch_simdgroups_per_threadgroup", "uint"}, - {"dispatch_threads_per_threadgroup", "uint3"}, - {"grid_origin", "uint3"}, - {"grid_size", "uint3"}, - {"quadgroup_index_in_threadgroup", "uint"}, - {"quadgroups_per_threadgroup", "uint"}, - {"simdgroup_index_in_threadgroup", "uint"}, - {"simdgroups_per_threadgroup", "uint"}, - {"thread_execution_width", "uint"}, - {"thread_index_in_quadgroup", "uint"}, - {"thread_index_in_simdgroup", "uint"}, - {"thread_index_in_threadgroup", "uint"}, - {"thread_position_in_grid", "uint3"}, - {"thread_position_in_threadgroup", "uint3"}, - {"threadgroup_position_in_grid", "uint3"}, - {"threadgroups_per_grid", "uint3"}, - {"threads_per_grid", "uint3"}, - {"threads_per_simdgroup", "uint"}, - {"threads_per_threadgroup", "uint3"}, - }; +#include - std::vector attributes; - for (const auto& [attr, dtype] : metal_attributes) { - if (source.find(attr) != std::string::npos) { - attributes.push_back(" " + dtype + " " + attr + " [[" + attr + "]]"); - } - } - - return [=, - shape_infos = std::move(shape_infos), - attributes = std::move(attributes)]( - const std::vector& inputs, - const std::vector& output_shapes, - const std::vector& output_dtypes, - std::tuple grid, - std::tuple threadgroup, - const std::vector>& - template_args = {}, - std::optional init_value = std::nullopt, - bool verbose = false, - StreamOrDevice s_ = {}) { - if (inputs.size() != input_names.size()) { - std::ostringstream msg; - msg << "[metal_kernel] Expected `inputs` to have size " - << input_names.size() << " but got size " << inputs.size() << "." - << std::endl; - throw std::invalid_argument(msg.str()); - } - if (output_shapes.size() != output_names.size()) { - std::ostringstream msg; - msg << "[metal_kernel] Expected `output_shapes` to have size " - << output_names.size() << " but got size " << output_shapes.size() - << "." << std::endl; - throw std::invalid_argument(msg.str()); - } - if (output_dtypes.size() != output_names.size()) { - std::ostringstream msg; - msg << "[metal_kernel] Expected `output_dtypes` to have size " - << output_names.size() << " but got size " << output_dtypes.size() - << "." << std::endl; - throw std::invalid_argument(msg.str()); - } - - auto s = to_stream(s_); - if (s.device != Device::gpu) { - throw std::invalid_argument("[metal_kernel] Only supports the GPU."); - } - - std::string kernel_name = "custom_kernel_" + name; - std::string template_def = ""; - if (!template_args.empty()) { - std::regex disallowed_chars("\\<|\\>|(, )"); - template_def = write_template(template_args); - auto template_hash = - std::regex_replace(template_def, disallowed_chars, "_"); - template_hash.pop_back(); - kernel_name += "_"; - kernel_name += template_hash; - } - - std::string kernel_source = write_signature( - kernel_name, - header, - source, - input_names, - inputs, - output_names, - output_dtypes, - template_args, - attributes, - shape_infos, - atomic_outputs); - - if (!template_args.empty()) { - template_def = kernel_name + template_def; - kernel_source += "\ntemplate [[host_name(\""; - kernel_source += kernel_name; - kernel_source += "\")]] [[kernel]] decltype("; - kernel_source += template_def; - kernel_source += ") "; - kernel_source += template_def; - kernel_source += ";\n"; - } - - if (verbose) { - std::cout << "Generated source code for `" << name << "`:" << std::endl - << "```" << std::endl - << kernel_source << std::endl - << "```" << std::endl; - } - - return array::make_arrays( - std::move(output_shapes), - std::move(output_dtypes), - std::make_shared( - s, - std::move(kernel_name), - std::move(kernel_source), - grid, - threadgroup, - shape_infos, - ensure_row_contiguous, - init_value, - std::vector{}, - false, - 0), - std::move(inputs)); - }; -} +namespace mlx::core::fast { void CustomKernel::eval_gpu( const std::vector& inputs, @@ -361,24 +47,12 @@ void CustomKernel::eval_gpu( auto& d = metal::device(s.device); - { - // Clear kernels from the device library cache if needed - auto& kernel_cache = cache(); - if (auto it = kernel_cache.libraries.find(name_); - it != kernel_cache.libraries.end()) { - if (it->second != source_) { - auto& d = metal::device(s.device); - d.clear_library(name_); - it->second = source_; - } - } else { - kernel_cache.libraries.emplace(name_, source_); - } - } - - auto lib = d.get_library(name_, [this] { return metal::utils() + source_; }); + std::string lib_name = fmt::format( + "{}_{:x}_{}", name_, std::hash{}(source_), compile_options_); + auto lib = d.get_library( + lib_name, compile_options_, [this] { return metal::utils() + source_; }); auto kernel = d.get_kernel(name_, lib); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int index = 0; for (int i = 0; i < checked_inputs.size(); i++) { @@ -424,7 +98,7 @@ void CustomKernel::eval_gpu( MTL::Size grid_dims = MTL::Size(gx, gy, gz); compute_encoder.dispatch_threads(grid_dims, group_dims); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } } // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp index 2bc854189..a9a0318b2 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/device.cpp @@ -3,25 +3,71 @@ #include #include +#include + #define NS_PRIVATE_IMPLEMENTATION #define CA_PRIVATE_IMPLEMENTATION #define MTL_PRIVATE_IMPLEMENTATION #include "mlx/backend/common/utils.h" #include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/event.h" #include "mlx/backend/metal/metal.h" #include "mlx/backend/metal/utils.h" #include "mlx/utils.h" +namespace std { + +// Required for putting the pointer in unordered_set. +template +struct hash> { + size_t operator()(const NS::SharedPtr& p) const { + return std::hash{}(p.get()); + } +}; + +} // namespace std + namespace mlx::core::metal { namespace { constexpr const char* default_mtllib_path = METAL_PATH; +void set_compile_options( + MTL::CompileOptions* mtl_options, + const CompileOptions& compile_options) { + if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) { + switch (compile_options.math_mode) { + case MathMode::Safe: + mtl_options->setMathMode(MTL::MathModeSafe); + break; + case MathMode::Relaxed: + mtl_options->setMathMode(MTL::MathModeRelaxed); + break; + case MathMode::Fast: + mtl_options->setMathMode(MTL::MathModeFast); + break; + default: + throw std::invalid_argument("[metal::Device] Invalid math mode."); + } + } else { + if (compile_options.math_mode == MathMode::Relaxed) { + throw std::runtime_error( + "[metal::Device] Metal math mode `relaxed` requires macOS 15, " + "iOS 18, tvOS 18, or visionOS 2."); + } + mtl_options->setFastMathEnabled( + compile_options.math_mode == MathMode::Fast); + } +} + auto get_metal_version() { auto get_metal_version_ = []() { - if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) { + if (__builtin_available(macOS 27, iOS 27, tvOS 27, visionOS 27, *)) { + // TODO: Use MTL::LanguageVersion4_1 after metal-cpp_27 is released. + return static_cast((4 << 16) + 1); + } else if (__builtin_available(macOS 26, iOS 26, tvOS 26, visionOS 26, *)) { return MTL::LanguageVersion4_0; } else if (__builtin_available(macOS 15, iOS 18, tvOS 18, visionOS 2, *)) { return MTL::LanguageVersion3_2; @@ -33,15 +79,28 @@ auto get_metal_version() { return metal_version_; } -auto load_device() { - auto devices = MTL::CopyAllDevices(); - auto device = static_cast(devices->object(0)) - ?: MTL::CreateSystemDefaultDevice(); +NS::SharedPtr load_device() { + auto pool = new_scoped_memory_pool(); + auto devices = NS::TransferPtr(MTL::CopyAllDevices()); + // In headless, sandboxed, or virtualized macOS sessions CopyAllDevices() + // returns an empty NSArray. Indexing object(0) on an empty array raises + // an unrecoverable NSRangeException, so guard the access and fall back to + // CreateSystemDefaultDevice (which can also return null). + MTL::Device* first = nullptr; + if (devices && devices->count() > 0) { + first = static_cast(devices->object(0)); + } + auto device = first ? NS::RetainPtr(first) + : NS::TransferPtr(MTL::CreateSystemDefaultDevice()); if (!device) { - throw std::runtime_error("Failed to load device"); + throw std::runtime_error( + "[metal::load_device] No Metal device available. This typically " + "occurs in headless, sandboxed, or virtualized macOS sessions " + "where the GPU is not accessible."); } return device; } + std::pair load_library_from_path( MTL::Device* device, const char* path) { @@ -134,6 +193,20 @@ std::pair load_swiftpm_library( } MTL::Library* load_default_library(MTL::Device* device) { + // Check override path before automatic lookup + if (!get_metallib_path().empty()) { + auto [lib, error] = + load_library_from_path(device, get_metallib_path().c_str()); + if (!lib) { + throw std::runtime_error( + fmt::format( + "Can not load metallib from specified location \"{}\": {}.", + get_metallib_path(), + error->localizedDescription()->utf8String())); + } + return lib; + } + NS::Error* error[5]; MTL::Library* lib; // First try the colocated mlx.metallib @@ -233,21 +306,40 @@ MTL::Library* load_library( } // namespace -CommandEncoder::CommandEncoder(DeviceStream& stream) : stream_(stream) { - enc_ = stream_.buffer->computeCommandEncoder(MTL::DispatchTypeConcurrent); - enc_->retain(); +CommandEncoder::CommandEncoder( + Device& d, + int index, + ResidencySets& residency_sets) + : device_(d), residency_sets_(residency_sets) { + auto pool = new_scoped_memory_pool(); + queue_ = NS::TransferPtr(device_.mtl_device()->newCommandQueue()); + if (!queue_) { + throw std::runtime_error( + "[metal::CommandEncoder] Failed to make new command queue."); + } + // Sets created later are attached in commit(). + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); + debug_set_stream_queue_label(queue_.get(), index); + buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); } CommandEncoder::~CommandEncoder() { - enc_->endEncoding(); - enc_->release(); + exiting_ = true; + synchronize(); + auto pool = new_scoped_memory_pool(); + buffer_.reset(); + queue_.reset(); } void CommandEncoder::set_buffer( const MTL::Buffer* buf, int idx, int64_t offset /* = 0 */) { - enc_->setBuffer(buf, offset, idx); + // Record as both input and output to ensure synchronization between command + // buffers + all_inputs_.insert((void*)buf); + all_outputs_.insert((void*)buf); + get_command_encoder()->setBuffer(buf, offset, idx); } void CommandEncoder::set_input_array( @@ -255,13 +347,14 @@ void CommandEncoder::set_input_array( int idx, int64_t offset /* = 0 */) { if (all_inputs_.insert(a.buffer().ptr()).second) { - stream_.buffer_sizes += a.data_size(); + buffer_sizes_ += a.data_size(); } auto r_buf = static_cast(const_cast(a.buffer().ptr())); + next_inputs_.insert(r_buf); needs_barrier_ = needs_barrier_ | (prev_outputs_.find(r_buf) != prev_outputs_.end()); auto a_buf = static_cast(a.buffer().ptr()); - enc_->setBuffer(a_buf, a.offset() + offset, idx); + get_command_encoder()->setBuffer(a_buf, a.offset() + offset, idx); } void CommandEncoder::set_output_array( @@ -281,17 +374,34 @@ void CommandEncoder::register_output_array(const array& a) { concurrent_outputs_.insert(buf); } else { next_outputs_.insert(buf); + needs_barrier_ = + needs_barrier_ | (prev_inputs_.find(buf) != prev_inputs_.end()); } } +void CommandEncoder::add_temporary(array arr) { + temporaries_.push_back(std::move(arr)); +} + +void CommandEncoder::add_temporaries(std::vector arrays) { + temporaries_.insert( + temporaries_.end(), + std::make_move_iterator(arrays.begin()), + std::make_move_iterator(arrays.end())); +} + void CommandEncoder::maybeInsertBarrier() { if (needs_barrier_) { - enc_->memoryBarrier(MTL::BarrierScopeBuffers); + get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); needs_barrier_ = false; - prev_outputs_ = std::move(next_outputs_); + // Preserve the hash tables' buckets for reuse across barrier epochs. + prev_inputs_.swap(next_inputs_); + prev_outputs_.swap(next_outputs_); } else { + prev_inputs_.insert(next_inputs_.begin(), next_inputs_.end()); prev_outputs_.insert(next_outputs_.begin(), next_outputs_.end()); } + next_inputs_.clear(); next_outputs_.clear(); } @@ -299,29 +409,195 @@ void CommandEncoder::dispatch_threadgroups( MTL::Size grid_dims, MTL::Size group_dims) { maybeInsertBarrier(); - stream_.buffer_ops++; - enc_->dispatchThreadgroups(grid_dims, group_dims); + buffer_ops_++; + get_command_encoder()->dispatchThreadgroups(grid_dims, group_dims); } void CommandEncoder::dispatch_threads( MTL::Size grid_dims, MTL::Size group_dims) { maybeInsertBarrier(); - stream_.buffer_ops++; - enc_->dispatchThreads(grid_dims, group_dims); + buffer_ops_++; + get_command_encoder()->dispatchThreads(grid_dims, group_dims); } void CommandEncoder::barrier() { - enc_->memoryBarrier(MTL::BarrierScopeBuffers); + get_command_encoder()->memoryBarrier(MTL::BarrierScopeBuffers); +} + +void CommandEncoder::end_encoding() { + // Each command encoder has a unique fence. We also store a map of + // all previous outputs of command encoders to their corresponding fence. + // - The command encoder records its inputs and outputs. + // - Wait on a fence if any inputs in the encoder are outputs of a previous + // encoder. + // - Update the map of outputs to include this command encoder's outputs. + // - Always signal this command encoders fence. + // - Add a completion handler for this command encoder that removes outputs + // from the map to limit the growth of the map and avoid unnecessary waits + // - Temporaries are a special case as they do not cross command encoder + // boundaries. These can be removed early from the encoders inputs and + // outputs since they don't need synchronization. + if (!encoder_) { + return; + } + + // Remove temporaries from inputs and outputs. + for (auto& t : temporaries_) { + all_outputs_.erase(t.buffer().ptr()); + all_inputs_.erase(t.buffer().ptr()); + } + + // Keep references to the fences we waited on and put them in the completion + // handler so they are not prematurely released. + std::unordered_set> waiting_on; + { + std::lock_guard lk(outputs_mtx_); + for (auto& in : all_inputs_) { + if (auto it = prev_ce_outputs_.find(in); it != prev_ce_outputs_.end()) { + // If we've already waited on a fence, don't wait on it again. + if (waiting_on.find(it->second) == waiting_on.end()) { + encoder_->waitForFence(it->second.get()); + waiting_on.insert(it->second); + } + } + } + for (auto& out : all_outputs_) { + prev_ce_outputs_[out] = fence_; + } + } + + encoder_->updateFence(fence_.get()); + buffer_->addCompletedHandler([this, + fence = std::move(fence_), + temporaries = std::move(temporaries_), + all_outputs = std::move(all_outputs_), + waiting_on = std::move(waiting_on)]( + MTL::CommandBuffer*) mutable { + std::lock_guard lk(outputs_mtx_); + for (auto& o : all_outputs) { + if (auto it = prev_ce_outputs_.find(o); it != prev_ce_outputs_.end()) { + if (it->second == fence) { + prev_ce_outputs_.erase(it); + } + } + } + }); + + encoder_->endEncoding(); + encoder_.reset(); + needs_barrier_ = false; + concurrent_ = false; + prev_inputs_.clear(); + next_inputs_.clear(); + prev_outputs_.clear(); + next_outputs_.clear(); + concurrent_outputs_.clear(); + all_inputs_.clear(); +} + +void CommandEncoder::signal_event(Event event, uint64_t value) { + end_encoding(); + buffer_->encodeSignalEvent(event.cast().mtl_event(), value); + signal_events_.push_back({std::move(event), value}); +} + +void CommandEncoder::wait_event(Event event, uint64_t value) { + end_encoding(); + buffer_->encodeWait(event.cast().mtl_event(), value); + wait_events_.push_back(std::move(event)); +} + +bool CommandEncoder::needs_commit() const { + auto [max_ops, max_mb] = device_.get_max_ops_mb_per_buffer(); + return (buffer_ops_ > max_ops) || ((buffer_sizes_ >> 20) > max_mb); +} + +void CommandEncoder::commit(std::function completion) { + // Metal locks a command buffer's residency at commit time, so attach any + // sets created since the last commit first. + residency_sets_.attach_new_sets(queue_.get(), sets_attached_); + buffer_->addCompletedHandler( + [&error_ = error_, + wait_events = std::move(wait_events_), + signal_events = std::move(signal_events_), + completion = std::move(completion)](MTL::CommandBuffer* cbuf) mutable { + if (completion) { + completion(); + } + // If any of the waited event has error in it, poison the encoder. + for (auto& event : wait_events) { + if (error_.store_if_valid(event.load_error())) { + break; + } + } + // Set error only when no error happended before, to preserve the + // earliest error. + bool has_error = error_.valid(); + if (!has_error && cbuf->status() == MTL::CommandBufferStatusError) { + error_.set_message( + std::make_shared(fmt::format( + "[METAL] Command buffer execution failed: {}.", + cbuf->error()->localizedDescription()->utf8String()))); + has_error = true; + } + // Poison all the signaled events when error happened. + if (has_error) { + for (auto& [event, value] : signal_events) { + event.set_error(error_); + } + } + // Metal won't signal the events for us on error, manually signal them + // to avoid infinite waiting. + if (cbuf->status() == MTL::CommandBufferStatusError) { + for (auto& [event, value] : signal_events) { + event.cast().signal(value); + } + } + }); + buffer_->commit(); + buffer_ = NS::RetainPtr(queue_->commandBufferWithUnretainedReferences()); + buffer_ops_ = 0; + buffer_sizes_ = 0; } -Device::Device() { +void CommandEncoder::synchronize() { auto pool = new_scoped_memory_pool(); - device_ = load_device(); - default_library_ = load_default_library(device_); - arch_ = std::string(device_->architecture()->name()->utf8String()); - int ag_tens = arch_[arch_.size() - 3] - '0'; - int ag_ones = arch_[arch_.size() - 2] - '0'; + auto cbuf = buffer_; // retained + end_encoding(); + commit(); + cbuf->waitUntilCompleted(); + + if (!exiting_) { + error_.check(); + } +} + +MTL::ComputeCommandEncoder* CommandEncoder::get_command_encoder() { + if (!encoder_) { + error_.check(); + encoder_ = NS::RetainPtr( + buffer_->computeCommandEncoder(MTL::DispatchTypeConcurrent)); + fence_ = NS::TransferPtr(device_.mtl_device()->newFence()); + } + return encoder_.get(); +} + +Device::Device() : device_(load_device()), residency_sets_(device_.get()) { + auto pool = new_scoped_memory_pool(); + default_library_ = NS::TransferPtr(load_default_library(device_.get())); + arch_ = env::metal_gpu_arch(); + if (arch_.empty()) { + arch_ = std::string(device_->architecture()->name()->utf8String()); + } + int ag_tens = 0; + int ag_ones = 0; + if (arch_.size() >= 3) { + ag_tens = arch_[arch_.size() - 3] - '0'; + ag_ones = arch_[arch_.size() - 2] - '0'; + ag_tens = (ag_tens < 10 && ag_tens >= 0) ? ag_tens : 0; + ag_ones = (ag_ones < 10 && ag_ones >= 0) ? ag_ones : 0; + } arch_gen_ = ag_tens * 10 + ag_ones; auto arch = arch_.back(); switch (arch) { @@ -350,154 +626,7 @@ Device::Device() { max_mb_per_buffer_ = env::max_mb_per_buffer(max_mb_per_buffer_); } -Device::~Device() { - auto pool = new_scoped_memory_pool(); - for (auto& [l, kernel_map] : library_kernels_) { - l->release(); - for (auto& [_, k] : kernel_map) { - k->release(); - } - } - stream_map_.clear(); - device_->release(); -} - -void Device::new_queue(int index) { - auto thread_pool = metal::new_scoped_memory_pool(); - auto q = device_->newCommandQueue(); - debug_set_stream_queue_label(q, index); - if (!q) { - throw std::runtime_error( - "[metal::Device] Failed to make new command queue."); - } - stream_map_.emplace(index, q); - if (residency_set_ != nullptr) { - q->addResidencySet(residency_set_); - } -} - -MTL::CommandQueue* Device::get_queue(Stream stream) { - return get_stream_(stream.index).queue; -} - -bool Device::command_buffer_needs_commit(int index) { - auto& stream = get_stream_(index); - return (stream.buffer_ops > max_ops_per_buffer_) || - ((stream.buffer_sizes >> 20) > max_mb_per_buffer_); -} - -MTL::CommandBuffer* Device::get_command_buffer(int index) { - auto& stream = get_stream_(index); - if (stream.buffer == nullptr) { - stream.buffer = stream.queue->commandBufferWithUnretainedReferences(); - if (!stream.buffer) { - throw std::runtime_error( - "[metal::Device] Unable to create new command buffer"); - } - // Increment ref count so the buffer is not garbage collected - stream.buffer->retain(); - } - return stream.buffer; -} - -void Device::commit_command_buffer(int index) { - auto& stream = get_stream_(index); - stream.buffer->commit(); - stream.buffer->release(); - stream.buffer = nullptr; - stream.buffer_ops = 0; - stream.buffer_sizes = 0; -} - -void Device::add_temporary(array arr, int index) { - get_stream_(index).temporaries.push_back(std::move(arr)); -} - -void Device::add_temporaries(std::vector arrays, int index) { - if (arrays.empty()) { - return; - } - auto& stream = get_stream_(index); - stream.temporaries.insert( - stream.temporaries.end(), - std::make_move_iterator(arrays.begin()), - std::make_move_iterator(arrays.end())); -} - -void Device::end_encoding(int index) { - auto& stream = get_stream_(index); - if (stream.encoder != nullptr) { - // Each command encoder has a unique fence. We also store a map of - // all previous outputs of command encoders to their corresponding fence. - // - The command encoder records its inputs and outputs. - // - Wait on a fence if any inputs in the encoder are outputs of a previous - // encoder. - // - Update the map of outputs to include this command encoder's outputs. - // - Always signal this command encoders fence. - // - Add a completion handler for this command encoder that removes outputs - // from the map to limit the growth of the map and avoid unnecessary waits - // - Temporaries are a special case as they do not cross command encoder - // boundaries. These can be removed early from the encoders inputs and - // outputs since they don't need synchronization. - auto& enc = *stream.encoder; - // Remove temporaries from inputs and outputs - for (auto& t : stream.temporaries) { - enc.outputs().erase(t.buffer().ptr()); - enc.inputs().erase(t.buffer().ptr()); - } - - // Keep references to the fences we waited on and put them - // in the completion handler so they are not prematurely released - std::unordered_set> waiting_on; - { - std::lock_guard lk(stream.fence_mtx); - for (auto in : enc.inputs()) { - if (auto it = stream.outputs.find(in); it != stream.outputs.end()) { - // If we've already waited on a fence, don't wait on it again. - if (waiting_on.find(it->second) == waiting_on.end()) { - enc.wait_for_fence(it->second->fence); - waiting_on.insert(it->second); - } - } - } - for (auto out : enc.outputs()) { - stream.outputs[out] = stream.fence; - } - } - enc.update_fence(stream.fence->fence); - stream.buffer->addCompletedHandler( - [&stream, - waiting_on = std::move(waiting_on), - fence = std::move(stream.fence), - outputs = std::move(enc.outputs()), - temporaries = - std::move(stream.temporaries)](MTL::CommandBuffer*) mutable { - temporaries.clear(); - std::lock_guard lk(stream.fence_mtx); - for (auto o : outputs) { - if (auto it = stream.outputs.find(o); it != stream.outputs.end()) { - if (it->second == fence) { - stream.outputs.erase(it); - } - } - } - }); - } - stream.encoder = nullptr; -} - -CommandEncoder& Device::get_command_encoder(int index) { - auto& stream = get_stream_(index); - if (stream.encoder == nullptr) { - // Ensure there is an active command buffer - if (stream.buffer == nullptr) { - get_command_buffer(index); - } - stream.encoder = std::make_unique(stream); - stream.fence = std::make_shared(device_->newFence()); - } - return *stream.encoder; -} +Device::~Device() = default; MTL::Library* Device::get_library( const std::string& name, @@ -505,32 +634,33 @@ MTL::Library* Device::get_library( { std::shared_lock rlock(library_mtx_); if (auto it = library_map_.find(name); it != library_map_.end()) { - return it->second; + return it->second.get(); } } std::unique_lock wlock(library_mtx_); if (auto it = library_map_.find(name); it != library_map_.end()) { - return it->second; + return it->second.get(); } - auto new_lib = load_library(device_, name, path.c_str()); - library_map_.insert({name, new_lib}); + auto new_lib = load_library(device_.get(), name, path.c_str()); + library_map_.insert({name, NS::TransferPtr(new_lib)}); return new_lib; } -MTL::Library* Device::build_library_(const std::string& source_string) { +NS::SharedPtr Device::build_library_( + const std::string& source_string, + const CompileOptions& compile_options) { auto pool = new_scoped_memory_pool(); auto ns_code = NS::String::string(source_string.c_str(), NS::ASCIIStringEncoding); NS::Error* error = nullptr; - auto options = MTL::CompileOptions::alloc()->init(); - options->setFastMathEnabled(false); + auto options = MTL::CompileOptions::alloc()->init()->autorelease(); + set_compile_options(options, compile_options); options->setLanguageVersion(get_metal_version()); - auto mtl_lib = device_->newLibrary(ns_code, options, &error); - options->release(); + auto mtl_lib = NS::TransferPtr(device_->newLibrary(ns_code, options, &error)); // Throw error if unable to compile library if (!mtl_lib) { @@ -545,17 +675,16 @@ MTL::Library* Device::build_library_(const std::string& source_string) { return mtl_lib; } -MTL::Function* Device::get_function_( +NS::SharedPtr Device::get_function_( const std::string& name, MTL::Library* mtl_lib) { + auto pool = new_scoped_memory_pool(); // Pull kernel from library auto ns_name = NS::String::string(name.c_str(), NS::ASCIIStringEncoding); - auto mtl_function = mtl_lib->newFunction(ns_name); - - return mtl_function; + return NS::TransferPtr(mtl_lib->newFunction(ns_name)); } -MTL::Function* Device::get_function_( +NS::SharedPtr Device::get_function_( const std::string& name, const std::string& specialized_name, const MTLFCList& func_consts, @@ -564,8 +693,11 @@ MTL::Function* Device::get_function_( return get_function_(name, mtl_lib); } + auto pool = new_scoped_memory_pool(); + // Prepare function constants - auto mtl_func_consts = MTL::FunctionConstantValues::alloc()->init(); + auto mtl_func_consts = + MTL::FunctionConstantValues::alloc()->init()->autorelease(); for (auto [value, type, index] : func_consts) { mtl_func_consts->setConstantValue(value, type, index); @@ -580,7 +712,7 @@ MTL::Function* Device::get_function_( // Pull kernel from library NS::Error* error = nullptr; - auto mtl_function = mtl_lib->newFunction(desc, &error); + auto mtl_function = NS::TransferPtr(mtl_lib->newFunction(desc, &error)); // Throw error if unable to build metal function if (!mtl_function) { @@ -592,20 +724,19 @@ MTL::Function* Device::get_function_( throw std::runtime_error(msg.str()); } - mtl_func_consts->release(); - return mtl_function; } -MTL::ComputePipelineState* Device::get_kernel_( +NS::SharedPtr Device::get_kernel_( const std::string& name, const MTL::Function* mtl_function) { // Compile kernel to compute pipeline NS::Error* error = nullptr; - MTL::ComputePipelineState* kernel; + NS::SharedPtr kernel; if (mtl_function) { - kernel = device_->newComputePipelineState(mtl_function, &error); + kernel = + NS::TransferPtr(device_->newComputePipelineState(mtl_function, &error)); } // Throw error if unable to compile metal function @@ -621,7 +752,7 @@ MTL::ComputePipelineState* Device::get_kernel_( return kernel; } -MTL::ComputePipelineState* Device::get_kernel_( +NS::SharedPtr Device::get_kernel_( const std::string& name, const MTL::Function* mtl_function, const MTL::LinkedFunctions* linked_functions) { @@ -636,15 +767,17 @@ MTL::ComputePipelineState* Device::get_kernel_( throw std::runtime_error(msg.str()); } + auto pool = new_scoped_memory_pool(); + // Prepare compute pipeline state descriptor - auto desc = MTL::ComputePipelineDescriptor::alloc()->init(); + auto desc = MTL::ComputePipelineDescriptor::alloc()->init()->autorelease(); desc->setComputeFunction(mtl_function); desc->setLinkedFunctions(linked_functions); // Compile kernel to compute pipeline NS::Error* error = nullptr; - auto kernel = device_->newComputePipelineState( - desc, MTL::PipelineOptionNone, nullptr, &error); + auto kernel = NS::TransferPtr(device_->newComputePipelineState( + desc, MTL::PipelineOptionNone, nullptr, &error)); // Throw error if unable to compile metal function if (!kernel) { @@ -659,62 +792,46 @@ MTL::ComputePipelineState* Device::get_kernel_( return kernel; } -MTL::Library* Device::get_library_(const std::string& name) { - std::shared_lock lock(library_mtx_); - auto it = library_map_.find(name); - return (it != library_map_.end()) ? it->second : nullptr; -} - MTL::Library* Device::get_library( const std::string& name, + const CompileOptions& compile_options, const std::function& builder) { { std::shared_lock rlock(library_mtx_); if (auto it = library_map_.find(name); it != library_map_.end()) { - return it->second; + return it->second.get(); } } std::unique_lock wlock(library_mtx_); if (auto it = library_map_.find(name); it != library_map_.end()) { - return it->second; + return it->second.get(); } - auto mtl_lib = build_library_(builder()); + auto mtl_lib = build_library_(builder(), compile_options); library_map_.insert({name, mtl_lib}); - return mtl_lib; + return mtl_lib.get(); } void Device::clear_library(const std::string& name) { std::unique_lock wlock(library_mtx_); if (auto it = library_map_.find(name); it != library_map_.end()) { - auto kernel_map_it = library_kernels_.find(it->second); - for (auto& [_, kernel] : kernel_map_it->second) { - kernel->release(); - } - library_kernels_.erase(kernel_map_it); - it->second->release(); + library_kernels_.erase(it->second.get()); library_map_.erase(it); } } -MTL::LinkedFunctions* Device::get_linked_functions_( +NS::SharedPtr Device::get_linked_functions_( const std::vector& funcs) { if (funcs.empty()) { return nullptr; } - auto lfuncs = MTL::LinkedFunctions::linkedFunctions(); - - std::vector objs(funcs.size()); - for (int i = 0; i < funcs.size(); i++) { - objs[i] = funcs[i]; - } - - NS::Array* funcs_arr = NS::Array::array(objs.data(), funcs.size()); - + auto pool = new_scoped_memory_pool(); + auto lfuncs = NS::TransferPtr(MTL::LinkedFunctions::linkedFunctions()); + NS::Array* funcs_arr = NS::Array::array( + reinterpret_cast(funcs.data()), funcs.size()); lfuncs->setPrivateFunctions(funcs_arr); - return lfuncs; } @@ -730,7 +847,7 @@ MTL::ComputePipelineState* Device::get_kernel_( // Try loading again to avoid loading twice auto& kernel_map_ = library_kernels_[mtl_lib]; if (auto it = kernel_map_.find(hash_name); it != kernel_map_.end()) { - return it->second; + return it->second.get(); } auto pool = new_scoped_memory_pool(); @@ -740,15 +857,13 @@ MTL::ComputePipelineState* Device::get_kernel_( // Compile kernel to compute pipeline auto mtl_linked_funcs = get_linked_functions_(linked_functions); - auto kernel = get_kernel_(hash_name, mtl_function, mtl_linked_funcs); - - mtl_function->release(); - mtl_linked_funcs->release(); + auto kernel = + get_kernel_(hash_name, mtl_function.get(), mtl_linked_funcs.get()); // Add kernel to cache kernel_map_.insert({hash_name, kernel}); - return kernel; + return kernel.get(); } MTL::ComputePipelineState* Device::get_kernel( @@ -763,9 +878,12 @@ MTL::ComputePipelineState* Device::get_kernel( std::shared_lock lock(kernel_mtx_); // Look for cached kernel - auto& kernel_map_ = library_kernels_[mtl_lib]; - if (auto it = kernel_map_.find(kname); it != kernel_map_.end()) { - return it->second; + auto library_it = library_kernels_.find(mtl_lib); + if (library_it != library_kernels_.end()) { + auto kernel_it = library_it->second.find(kname); + if (kernel_it != library_it->second.end()) { + return kernel_it->second.get(); + } } } return get_kernel_(base_name, mtl_lib, kname, func_consts, linked_functions); @@ -777,35 +895,69 @@ MTL::ComputePipelineState* Device::get_kernel( const MTLFCList& func_consts /* = {} */, const std::vector& linked_functions /* = {} */) { return get_kernel( - base_name, default_library_, hash_name, func_consts, linked_functions); + base_name, + default_library_.get(), + hash_name, + func_consts, + linked_functions); } -void Device::set_residency_set(const MTL::ResidencySet* residency_set) { - if (residency_set_ != nullptr) { - throw std::runtime_error( - "[Device::set_residency_set] Can only be set once."); - } - if (residency_set == nullptr) { - return; - } - residency_set_ = residency_set; - // Attach residency set to existing command queues - for (auto& [_, stream] : stream_map_) { - stream.queue->addResidencySet(residency_set_); +Device& device(mlx::core::Device) { + // Leak singleton device intentionally, to avoid cases where a compute kernel + // returns and tries to access the object after it has been freed by the main + // thread teardown. + static Device* metal_device = new Device; + return *metal_device; +} + +CommandEncoder& get_command_encoder(Stream s) { + auto& encoders = get_command_encoders(); + auto it = encoders.find(s.index); + if (it == encoders.end()) { + auto& global_encoders = get_global_command_encoders(); + it = global_encoders.find(s.index); + if (it == global_encoders.end()) { + throw std::runtime_error( + fmt::format( + "There is no Stream(gpu, {}) in current thread.", s.index)); + } } + return it->second; } -Device& device(mlx::core::Device) { - static Device metal_device; - return metal_device; +std::unordered_map& get_command_encoders() { + static thread_local std::unordered_map encoders; + return encoders; +} + +std::unordered_map& get_global_command_encoders() { + static std::unordered_map encoders; + return encoders; } -std::unique_ptr> new_scoped_memory_pool() { - auto dtor = [](void* ptr) { - static_cast(ptr)->release(); +NS::SharedPtr new_scoped_memory_pool() { + return NS::TransferPtr(NS::AutoreleasePool::alloc()->init()); +} + +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; }; - return std::unique_ptr>( - NS::AutoreleasePool::alloc()->init(), dtor); + static bool is_nax_available_ = _check_nax(); + return is_nax_available_; +#endif } } // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device.h b/Source/Cmlx/mlx/mlx/backend/metal/device.h index e6162d7da..2d2228335 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/device.h @@ -6,11 +6,12 @@ #include #include #include -#include #include #include #include "mlx/array.h" +#include "mlx/backend/common/metal_kernel.h" +#include "mlx/backend/metal/resident.h" #include "mlx/device.h" namespace mlx::core::metal { @@ -18,10 +19,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; @@ -40,29 +44,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) { @@ -71,79 +72,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 { @@ -154,26 +152,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, @@ -181,8 +171,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( @@ -198,40 +195,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); @@ -243,18 +231,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_; @@ -262,28 +249,12 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/backend/metal/device_info.cpp b/Source/Cmlx/mlx/mlx/backend/metal/device_info.cpp index dd18fc6c6..3902e6b5e 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/device_info.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/device_info.cpp @@ -13,7 +13,12 @@ bool is_available() { } int device_count() { - return 1; + try { + metal::device(Device::gpu); + return 1; + } catch (...) { + return 0; + } } const std::unordered_map>& @@ -21,9 +26,10 @@ device_info(int device_index) { auto init_device_info = []() -> std::unordered_map> { auto pool = metal::new_scoped_memory_pool(); - auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device(); + auto& device = metal::device(mlx::core::Device::gpu); + auto raw_device = device.mtl_device(); auto name = std::string(raw_device->name()->utf8String()); - auto arch = std::string(raw_device->architecture()->name()->utf8String()); + auto arch = device.get_architecture(); size_t memsize = 0; size_t length = sizeof(memsize); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/eval.cpp b/Source/Cmlx/mlx/mlx/backend/metal/eval.cpp index bd58a691a..43a6fb993 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/eval.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/eval.cpp @@ -6,29 +6,31 @@ #include "mlx/backend/metal/utils.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" +#include "mlx/utils.h" namespace mlx::core::gpu { -void new_stream(Stream stream) { - if (stream.device == mlx::core::Device::gpu) { - metal::device(stream.device).new_queue(stream.index); - } +void init() {} + +void new_stream(Stream s) { + assert(s.device == Device::gpu); + auto& encoders = metal::get_command_encoders(); + auto& d = metal::device(s.device); + encoders.try_emplace(s.index, d, s.index, d.residency_sets()); } -inline void check_error(MTL::CommandBuffer* cbuf) { - if (cbuf->status() == MTL::CommandBufferStatusError) { - std::ostringstream msg; - msg << "[METAL] Command buffer execution failed: " - << cbuf->error()->localizedDescription()->utf8String(); - throw std::runtime_error(msg.str()); - } +void new_thread_unsafe_stream(Stream s) { + assert(s.device == Device::gpu); + auto& encoders = metal::get_global_command_encoders(); + auto& d = metal::device(s.device); + encoders.try_emplace(s.index, d, s.index, d.residency_sets()); } void eval(array& arr) { auto pool = metal::new_scoped_memory_pool(); auto s = arr.primitive().stream(); - auto& d = metal::device(s.device); - auto command_buffer = d.get_command_buffer(s.index); + auto& encoder = metal::get_command_encoder(s); + auto* command_buffer = encoder.get_command_buffer(); auto outputs = arr.outputs(); { @@ -54,44 +56,35 @@ void eval(array& arr) { buffers.erase(it); } - if (d.command_buffer_needs_commit(s.index)) { - d.end_encoding(s.index); + if (encoder.needs_commit()) { + encoder.end_encoding(); scheduler::notify_new_task(s); - command_buffer->addCompletedHandler( - [s, buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) { - scheduler::notify_task_completion(s); - check_error(cbuf); - }); - d.commit_command_buffer(s.index); - d.get_command_buffer(s.index); + encoder.commit([s, buffers = std::move(buffers)]() { + scheduler::notify_task_completion(s); + }); } else { command_buffer->addCompletedHandler( - [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) { - check_error(cbuf); - }); + [buffers = std::move(buffers)](MTL::CommandBuffer* cbuf) {}); } } void finalize(Stream s) { auto pool = metal::new_scoped_memory_pool(); - auto& d = metal::device(s.device); - auto cb = d.get_command_buffer(s.index); - d.end_encoding(s.index); - cb->addCompletedHandler([](MTL::CommandBuffer* cbuf) { check_error(cbuf); }); - d.commit_command_buffer(s.index); - d.get_command_buffer(s.index); + auto& encoder = metal::get_command_encoder(s); + auto* cb = encoder.get_command_buffer(); + encoder.end_encoding(); + encoder.commit(); } void synchronize(Stream s) { - auto pool = metal::new_scoped_memory_pool(); - auto& d = metal::device(s.device); - auto cb = d.get_command_buffer(s.index); - cb->retain(); - d.end_encoding(s.index); - d.commit_command_buffer(s.index); - cb->waitUntilCompleted(); - check_error(cb); - cb->release(); + metal::get_command_encoder(s).synchronize(); +} + +void clear_streams() { + metal::get_command_encoders().clear(); + if (is_main_thread()) { + metal::get_global_command_encoders().clear(); + } } } // namespace mlx::core::gpu diff --git a/Source/Cmlx/mlx/mlx/backend/metal/event.cpp b/Source/Cmlx/mlx/mlx/backend/metal/event.cpp index 39a49230d..38a387c9c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/event.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/event.cpp @@ -1,62 +1,83 @@ // Copyright © 2024 Apple Inc. -#include "mlx/event.h" -#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/event.h" #include "mlx/scheduler.h" namespace mlx::core { -Event::Event(Stream stream) : stream_(stream) { - auto dtor = [](void* ptr) { - auto p = metal::new_scoped_memory_pool(); - static_cast(ptr)->release(); - }; - auto p = metal::new_scoped_memory_pool(); - event_ = std::shared_ptr( - metal::device(Device::gpu).mtl_device()->newSharedEvent(), dtor); - if (event_ == nullptr) { +/////////////////////////////////////////////////////////////////////////////// +// EventImpl implementations +/////////////////////////////////////////////////////////////////////////////// + +namespace metal { + +EventImpl::EventImpl(Device& d) { + auto p = new_scoped_memory_pool(); + mtl_event_ = NS::TransferPtr(d.mtl_device()->newSharedEvent()); + if (!mtl_event_) { throw std::runtime_error( "[Event::Event] Failed to create Metal shared event."); } } +EventImpl::~EventImpl() { + auto p = new_scoped_memory_pool(); + mtl_event_.reset(); +} + +void EventImpl::wait(uint64_t value) { + mtl_event_->waitUntilSignaledValue(value, -1); // never times out +} + +void EventImpl::signal(uint64_t value) { + mtl_event_->setSignaledValue(value); +} + +} // namespace metal + +/////////////////////////////////////////////////////////////////////////////// +// Event implementations +/////////////////////////////////////////////////////////////////////////////// + +Event::Event(Stream stream) : stream_(stream) { + event_ = std::make_shared(metal::device(stream.device)); +} + void Event::wait() { - if (!static_cast(event_.get()) - ->waitUntilSignaledValue(value(), -1)) { - throw std::runtime_error("[Event::wait] Timed out"); - } + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [*this]() mutable { wait(); }); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); + }); } else { - auto& d = metal::device(stream.device); - d.end_encoding(stream.index); - auto command_buffer = d.get_command_buffer(stream.index); - command_buffer->encodeWait(static_cast(event_.get()), value()); - command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {}); + auto& encoder = metal::get_command_encoder(stream); + encoder.wait_event(*this, value()); } } void Event::signal(Stream stream) { if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [*this]() mutable { - static_cast(event_.get())->setSignaledValue(value()); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + self.cast().signal(value); }); } else { - auto& d = metal::device(stream.device); - d.end_encoding(stream.index); - auto command_buffer = d.get_command_buffer(stream.index); - command_buffer->encodeSignalEvent( - static_cast(event_.get()), value()); - command_buffer->addCompletedHandler([*this](MTL::CommandBuffer*) {}); + auto& encoder = metal::get_command_encoder(stream); + encoder.signal_event(*this, value()); } } bool Event::is_signaled() const { - return static_cast(event_.get())->signaledValue() >= - value(); + auto* mtl_event = cast().mtl_event(); + return mtl_event->signaledValue() >= value(); +} + +std::atomic& Event::error() { + return cast().error(); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/event.h b/Source/Cmlx/mlx/mlx/backend/metal/event.h new file mode 100644 index 000000000..d1e43fa02 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/event.h @@ -0,0 +1,31 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/metal/device.h" +#include "mlx/event.h" + +namespace mlx::core::metal { + +class EventImpl { + public: + EventImpl(Device& d); + ~EventImpl(); + + void wait(uint64_t value); + void signal(uint64_t value); + + auto& error() { + return error_; + } + + auto* mtl_event() const { + return mtl_event_.get(); + } + + private: + // All streams outlive events so pointers would be always valid. + std::atomic error_; + + NS::SharedPtr mtl_event_; +}; + +} // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp index da73d5d91..70dd0e33b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/fence.cpp @@ -7,8 +7,8 @@ namespace mlx::core { struct FenceImpl { - FenceImpl() { - auto d = metal::device(Device::gpu).mtl_device(); + FenceImpl(Stream stream) { + auto d = metal::device(stream.device).mtl_device(); if (!d->supportsFamily(MTL::GPUFamilyMetal3)) { use_fast = false; } else if (__builtin_available(macOS 15, iOS 18, *)) { @@ -16,8 +16,7 @@ struct FenceImpl { } if (!use_fast) { - auto p = metal::new_scoped_memory_pool(); - fence = static_cast(d->newSharedEvent()); + event = std::make_unique(stream); } else { auto buf = allocator::malloc(sizeof(uint32_t)).ptr(); fence = static_cast(buf); @@ -26,17 +25,14 @@ struct FenceImpl { } ~FenceImpl() { - if (!use_fast) { - // Wraps Metal SharedEvent - auto p = metal::new_scoped_memory_pool(); - static_cast(fence)->release(); - } else { + if (use_fast) { allocator::free(allocator::Buffer{static_cast(fence)}); } } bool use_fast{false}; uint32_t count{0}; void* fence; + std::unique_ptr event; std::atomic_uint* cpu_value() { return static_cast( @@ -44,24 +40,21 @@ struct FenceImpl { } }; -Fence::Fence(Stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl{}, dtor); +Fence::Fence(Stream stream) { + fence_ = std::make_shared(stream); } void Fence::wait(Stream stream, const array& x) { auto& f = *static_cast(fence_.get()); + if (!f.use_fast) { + f.event->wait(stream); + return; + } + if (stream.device == Device::cpu) { scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable { auto& f = *static_cast(fence_.get()); - if (!f.use_fast) { - if (!static_cast(f.fence)->waitUntilSignaledValue( - count, -1)) { - throw std::runtime_error("[Fence::wait] Timed out"); - } - return; - } while (f.cpu_value()[0] < count) { } }); @@ -69,18 +62,7 @@ void Fence::wait(Stream stream, const array& x) { } auto& d = metal::device(stream.device); - auto idx = stream.index; - - if (!f.use_fast) { - d.end_encoding(idx); - auto command_buffer = d.get_command_buffer(idx); - command_buffer->encodeWait(static_cast(f.fence), f.count); - command_buffer->addCompletedHandler( - [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); - return; - } - - auto& compute_encoder = d.get_command_encoder(idx); + auto& compute_encoder = metal::get_command_encoder(stream); // Register outputs to ensure that no kernels which depends on the // output starts before this one is done @@ -95,7 +77,7 @@ void Fence::wait(Stream stream, const array& x) { compute_encoder.set_bytes(f.count, 1); compute_encoder.dispatch_threads(kernel_dims, kernel_dims); - d.get_command_buffer(idx)->addCompletedHandler( + compute_encoder.get_command_buffer()->addCompletedHandler( [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); } @@ -103,33 +85,24 @@ void Fence::update(Stream stream, const array& x, bool cross_device) { auto& f = *static_cast(fence_.get()); f.count++; + if (!f.use_fast) { + f.event->set_value(f.count); + f.event->signal(stream); + return; + } + if (stream.device == Device::cpu) { scheduler::enqueue(stream, [fence_ = fence_, count = f.count]() mutable { auto& f = *static_cast(fence_.get()); - if (!f.use_fast) { - static_cast(f.fence)->setSignaledValue(count); - return; - } - f.cpu_value()[0] = count; }); return; } auto& d = metal::device(stream.device); - auto idx = stream.index; - if (!f.use_fast) { - d.end_encoding(idx); - auto command_buffer = d.get_command_buffer(idx); - command_buffer->encodeSignalEvent( - static_cast(f.fence), f.count); - command_buffer->addCompletedHandler( - [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); - return; - } + auto& compute_encoder = metal::get_command_encoder(stream); // Launch input visibility kernels - auto& compute_encoder = d.get_command_encoder(idx); if (cross_device) { auto kernel = d.get_kernel("input_coherent"); uint32_t nthreads = (x.data_size() * x.itemsize() + sizeof(uint32_t) - 1) / @@ -155,7 +128,7 @@ void Fence::update(Stream stream, const array& x, bool cross_device) { compute_encoder.set_bytes(f.count, 1); compute_encoder.dispatch_threads(kernel_dims, kernel_dims); - d.get_command_buffer(idx)->addCompletedHandler( + compute_encoder.get_command_buffer()->addCompletedHandler( [fence_ = fence_](MTL::CommandBuffer* cbuf) {}); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp b/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp index d99e1badb..3c9195f8e 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/fft.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "mlx/3rdparty/pocketfft.h" #include "mlx/backend/common/utils.h" @@ -26,6 +27,10 @@ using MTLFC = std::tuple; #define MIN_THREADGROUP_MEM_SIZE 256 // For strided reads/writes, coalesce at least this many complex64s #define MIN_COALESCE_WIDTH 4 +// Precomputed radix twiddles pay off for the largest batched fused plan. +constexpr int BLUESTEIN_TWIDDLE_TABLE_FFT_SIZE = 4096; +constexpr int BLUESTEIN_TWIDDLE_TABLE_SIZE = 584; +constexpr int MIN_BLUESTEIN_TWIDDLE_TABLE_BATCH = 1024; inline const std::vector supported_radices() { // Ordered by preference in decomposition. @@ -127,6 +132,21 @@ FFTPlan plan_fft(int n) { // Rough heuristic for choosing faster powers of two when we can plan.n2 = n > 65536 ? 1024 : 64; plan.n1 = n / plan.n2; + // Each step is a single threadgroup FFT, so neither factor can be larger + // than the largest Stockham size. The no transpose kernels cannot + // decompose a step any further, so grow n2 instead of nesting four step. + if (plan.n1 > MAX_STOCKHAM_FFT_SIZE) { + plan.n1 = MAX_STOCKHAM_FFT_SIZE; + plan.n2 = n / plan.n1; + } + if (plan.n2 > MAX_STOCKHAM_FFT_SIZE) { + std::ostringstream msg; + msg << "[FFT] GPU FFT is not supported for size " << n + << ", the largest supported size is " + << MAX_STOCKHAM_FFT_SIZE * MAX_STOCKHAM_FFT_SIZE + << ". Run the FFT on the CPU stream instead."; + throw std::runtime_error(msg.str()); + } return plan; } else if (n > MAX_STOCKHAM_FFT_SIZE) { // Otherwise we use a multi-upload Bluestein's @@ -300,7 +320,10 @@ std::tuple compute_raders_constants( } // Bluestein -std::pair compute_bluestein_constants(int n, int bluestein_n) { +std::pair compute_bluestein_constants( + int n, + int bluestein_n, + int radix_twiddle_size = 0) { // We need to calculate the Bluestein twiddle factors // in double precision for the overall numerical stability // of Bluestein's FFT algorithm to be acceptable. @@ -327,7 +350,7 @@ std::pair compute_bluestein_constants(int n, int bluestein_n) { w_k.set_data(allocator::malloc(w_k.nbytes())); std::copy(w_k_vec.begin(), w_k_vec.end(), w_k.data()); - array w_q({bluestein_n}, complex64, nullptr, {}); + array w_q({bluestein_n + radix_twiddle_size}, complex64, nullptr, {}); w_q.set_data(allocator::malloc(w_q.nbytes())); auto w_q_ptr = reinterpret_cast*>(w_q.data()); @@ -515,6 +538,7 @@ void fft_op( bool inplace, const Stream& s) { auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); size_t n = out.dtype() == float32 ? out.shape(axis) : in.shape(axis); if (n == 1) { @@ -581,7 +605,7 @@ void fft_op( auto plan = plan_fft(n); if (plan.four_step) { four_step_fft(in, out, axis, inverse, real, plan, copies, s, inplace); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); return; } @@ -632,18 +656,28 @@ void fft_op( size = out.size(); } int total_batch_size = size / n; + static const std::vector bluestein_twiddle_table_plan = { + 0, 0, 4, 0, 0, 0, 0, 0, 0}; + bool use_bluestein_twiddle_table = !real && + plan.bluestein_n == BLUESTEIN_TWIDDLE_TABLE_FFT_SIZE && + plan.stockham == bluestein_twiddle_table_plan && + total_batch_size >= MIN_BLUESTEIN_TWIDDLE_TABLE_BATCH; + if (plan.bluestein_n > 0 && !real) { + func_consts.push_back(make_bool(&use_bluestein_twiddle_table, 22)); + } int threads_per_fft = (fft_size + elems_per_thread - 1) / elems_per_thread; // We batch among threadgroups for improved efficiency when n is small int threadgroup_batch_size = std::max(MIN_THREADGROUP_MEM_SIZE / fft_size, 1); if (four_step_params.required) { - // Require a threadgroup batch size of at least 4 for four step FFT - // so we can coalesce the memory accesses. - threadgroup_batch_size = - std::max(threadgroup_batch_size, MIN_COALESCE_WIDTH); + // Batch the four step FFT so we can coalesce the memory accesses, but + // never past what fits in threadgroup memory. + threadgroup_batch_size = std::max( + threadgroup_batch_size, + std::min(MIN_COALESCE_WIDTH, MAX_STOCKHAM_FFT_SIZE / fft_size)); } int threadgroup_mem_size = next_power_of_2(threadgroup_batch_size * fft_size); - // FFTs up to 2^20 are currently supported + // The plan keeps every step within the threadgroup memory limit assert(threadgroup_mem_size <= MAX_STOCKHAM_FFT_SIZE); // ceil divide @@ -654,7 +688,6 @@ void fft_op( // We can perform 2 RFFTs at once so the batch size is halved. batch_size = (batch_size + 2 - 1) / 2; } - auto& compute_encoder = d.get_command_encoder(s.index); auto in_type_str = in.dtype() == float32 ? "float" : "float2"; auto out_type_str = out.dtype() == float32 ? "float" : "float2"; // Only required by four step @@ -685,6 +718,9 @@ void fft_op( std::string base_name = kname.str(); // We use a specialized kernel for each FFT size kname << "_n" << fft_size << "_inv_" << inverse; + if (use_bluestein_twiddle_table) { + kname << "_precomputed"; + } std::string hash_name = kname.str(); auto template_def = func_name == "four_step_fft" ? get_template_definition( base_name, @@ -703,22 +739,45 @@ void fft_op( auto kernel = get_fft_kernel(d, base_name, hash_name, func_consts, template_def); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(in_contiguous, 0); - compute_encoder.set_output_array(out, 1); - if (plan.bluestein_n > 0) { // Precomputed twiddle factors for Bluestein's - auto [w_k, w_q] = compute_bluestein_constants(n, plan.bluestein_n); + auto [w_k, w_q] = use_bluestein_twiddle_table + ? compute_bluestein_constants( + n, plan.bluestein_n, BLUESTEIN_TWIDDLE_TABLE_SIZE) + : compute_bluestein_constants(n, plan.bluestein_n); copies.push_back(w_q); copies.push_back(w_k); + if (use_bluestein_twiddle_table) { + auto twiddle_kernel = + get_fft_twiddle_kernel(d, base_name, template_def); + compute_encoder.set_compute_pipeline_state(twiddle_kernel); + compute_encoder.set_output_array( + w_q, + 0, + plan.bluestein_n * static_cast(sizeof(complex64_t))); + auto twiddle_group_size = std::min( + static_cast(BLUESTEIN_TWIDDLE_TABLE_SIZE), + twiddle_kernel->maxTotalThreadsPerThreadgroup()); + compute_encoder.dispatch_threads( + MTL::Size(BLUESTEIN_TWIDDLE_TABLE_SIZE, 1, 1), + MTL::Size(twiddle_group_size, 1, 1)); + } + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_input_array(w_q, 2); // w_q compute_encoder.set_input_array(w_k, 3); // w_k compute_encoder.set_bytes(n, 4); compute_encoder.set_bytes(plan.bluestein_n, 5); compute_encoder.set_bytes(total_batch_size, 6); } else if (plan.rader_n > 1) { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + auto [b_q, g_q, g_minus_q] = compute_raders_constants(plan.rader_n, s); copies.push_back(b_q); copies.push_back(g_q); @@ -731,10 +790,18 @@ void fft_op( compute_encoder.set_bytes(total_batch_size, 6); compute_encoder.set_bytes(plan.rader_n, 7); } else if (four_step_params.required) { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(four_step_params.n1, 2); compute_encoder.set_bytes(four_step_params.n2, 3); compute_encoder.set_bytes(total_batch_size, 4); } else { + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in_contiguous, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_bytes(n, 2); compute_encoder.set_bytes(total_batch_size, 3); } @@ -745,7 +812,7 @@ void fft_op( compute_encoder.dispatch_threads(grid_dims, group_dims); } - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } void fft_op( @@ -789,8 +856,7 @@ void nd_fft_op( fft_op(in_arr, out_arr, axis, inverse, step_real, inplace, s); } - auto& d = metal::device(s.device); - d.add_temporaries(std::move(temp_arrs), s.index); + metal::get_command_encoder(s).add_temporaries(std::move(temp_arrs)); } void FFT::eval_gpu(const std::vector& inputs, array& out) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/hadamard.cpp b/Source/Cmlx/mlx/mlx/backend/metal/hadamard.cpp index 65a877151..c696e4d0f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/hadamard.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/hadamard.cpp @@ -69,7 +69,8 @@ void hadamard_mn_contiguous( int n = n1 * n2; int read_width_n1 = n1 == 2 ? 2 : 4; int read_width_n2 = n2 == 2 ? 2 : 4; - int read_width_m = (n == 2 || m == 28) ? 2 : 4; + // The m stage strides by n / read_width_m, so cap the read width at n. + int read_width_m = (n == 1) ? 1 : ((n == 2 || m == 28) ? 2 : 4); int max_radix_1 = std::min(n1, 16); int max_radix_2 = std::min(n2, 16); float scale_n1 = 1.0; @@ -97,17 +98,16 @@ void hadamard_mn_contiguous( auto lib = d.get_library(kname, [&]() { std::string kernel; concatenate( - kernel, - metal::utils(), - gen_hadamard_codelet(m), - metal::hadamard(), - get_template_definition( - "n2" + kname, - "hadamard_n", - get_type_string(x.dtype()), - n2, - max_radix_2, - read_width_n2)); + kernel, metal::utils(), gen_hadamard_codelet(m), metal::hadamard()); + if (n2 > 1) { + kernel += get_template_definition( + "n2" + kname, + "hadamard_n", + get_type_string(x.dtype()), + n2, + max_radix_2, + read_width_n2); + } if (n1 > 1) { kernel += get_template_definition( "n1" + kname, @@ -132,7 +132,7 @@ void hadamard_mn_contiguous( // Launch the strided transform for n1 if (n1 > 1) { - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel("n1" + kname, lib); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(x, 0); @@ -142,19 +142,22 @@ void hadamard_mn_contiguous( } // Launch the transform for n2 - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = d.get_kernel("n2" + kname, lib); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(n1 > 1 ? y : x, 0); - compute_encoder.set_output_array(y, 1); - compute_encoder.set_bytes(scale_n2, 2); - compute_encoder.dispatch_threads(grid_dims_n2, group_dims_n2); + auto& compute_encoder = metal::get_command_encoder(s); + if (n2 > 1) { + auto kernel = d.get_kernel("n2" + kname, lib); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(n1 > 1 ? y : x, 0); + compute_encoder.set_output_array(y, 1); + compute_encoder.set_bytes(scale_n2, 2); + compute_encoder.dispatch_threads(grid_dims_n2, group_dims_n2); + } // Launch the strided transform for m if (m > 1) { auto kernel = d.get_kernel("m" + kname, lib); compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(y, 0); + // With n == 1 no earlier stage copied x into y, so read from x. + compute_encoder.set_input_array(n > 1 ? y : x, 0); compute_encoder.set_output_array(y, 1); compute_encoder.set_bytes(scale_m, 2); compute_encoder.dispatch_threads(grid_dims_m, group_dims_m); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp b/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp index 48311dc6a..117e402f7 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/indexing.cpp @@ -3,13 +3,15 @@ #include #include "mlx/backend/common/compiled.h" +#include "mlx/backend/common/slicing.h" #include "mlx/backend/common/utils.h" #include "mlx/backend/gpu/copy.h" +#include "mlx/backend/gpu/scan.h" +#include "mlx/backend/gpu/slicing.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/jit/includes.h" #include "mlx/backend/metal/jit/indexing.h" #include "mlx/backend/metal/kernels.h" -#include "mlx/backend/metal/scan.h" #include "mlx/backend/metal/utils.h" #include "mlx/dtype.h" #include "mlx/primitives.h" @@ -36,6 +38,22 @@ std::pair make_index_args( return {idx_args.str(), idx_arr.str()}; } +template +inline std::string make_op(typename T::ReduceType r, const std::string& dt) { + switch (r) { + case T::None: + return "None"; + case T::Sum: + return fmt::format("Sum<{0}>", dt); + case T::Prod: + return fmt::format("Prod<{0}>", dt); + case T::Max: + return fmt::format("Max<{0}>", dt); + case T::Min: + return fmt::format("Min<{0}>", dt); + } +} + void Gather::eval_gpu(const std::vector& inputs, array& out) { auto& src = inputs[0]; int nidx = inputs.size() - 1; @@ -93,7 +111,7 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { return kernel_source; }); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kernel_name, lib); compute_encoder.set_compute_pipeline_state(kernel); @@ -146,7 +164,7 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { return kernel_source; }); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kernel_name, lib); compute_encoder.set_compute_pipeline_state(kernel); @@ -212,7 +230,8 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { } void Scatter::eval_gpu(const std::vector& inputs, array& out) { - if (size_of(out.dtype()) == 8) { + if (size_of(out.dtype()) == 8 && + !(out.dtype() == complex64 && reduce_type_ == Scatter::Sum)) { std::ostringstream msg; msg << "[Scatter::eval_gpu] Does not support " << out.dtype(); throw std::invalid_argument(msg.str()); @@ -307,27 +326,7 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { std::string out_type_str = get_type_string(out.dtype()); std::string idx_type_str = nidx ? get_type_string(inputs[1].dtype()) : "bool"; - std::string op_type; - switch (reduce_type_) { - case Scatter::None: - op_type = "None"; - break; - case Scatter::Sum: - op_type = "Sum<{0}>"; - break; - case Scatter::Prod: - op_type = "Prod<{0}>"; - break; - case Scatter::Max: - op_type = "Max<{0}>"; - break; - case Scatter::Min: - op_type = "Min<{0}>"; - break; - } - if (reduce_type_ != Scatter::None) { - op_type = fmt::format(fmt::runtime(op_type), out_type_str); - } + std::string op_type = make_op(reduce_type_, out_type_str); auto [idx_args, idx_arr] = make_index_args(idx_type_str, nidx); kernel_source += fmt::format( @@ -345,7 +344,7 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { return kernel_source; }); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kernel_name, lib); size_t nthreads = upd.size(); @@ -484,7 +483,7 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { return kernel_source; }); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kernel_name, lib); compute_encoder.set_compute_pipeline_state(kernel); @@ -521,6 +520,13 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { } void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { + if (size_of(out.dtype()) == 8 && + !(out.dtype() == complex64 && reduce_type_ == ScatterAxis::Sum)) { + std::ostringstream msg; + msg << "[ScatterAxis::eval_gpu] Does not support " << out.dtype(); + throw std::invalid_argument(msg.str()); + } + auto& src = inputs[0]; auto& idx = inputs[1]; auto& upd = inputs[2]; @@ -600,7 +606,7 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { return kernel_source; }); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kernel_name, lib); compute_encoder.set_compute_pipeline_state(kernel); @@ -651,6 +657,7 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); const size_t total = mask.size(); const CopyType ct = (total == 1) @@ -663,18 +670,18 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { array mask_flat = flatten_in_eval(mask, 1, -1, s); if (mask_flat.data() != mask.data()) { - d.add_temporary(mask_flat, s.index); + compute_encoder.add_temporary(mask_flat); } if (!mask_flat.flags().row_contiguous) { mask_flat = contiguous_copy_gpu(mask_flat, s); - d.add_temporary(mask_flat, s.index); + compute_encoder.add_temporary(mask_flat); } // Prefix (exclusive) of mask → scatter_offsets array scatter_offsets(mask_flat.shape(), uint32, nullptr, {}); scatter_offsets.set_data(allocator::malloc(scatter_offsets.nbytes())); - d.add_temporary(scatter_offsets, s.index); + compute_encoder.add_temporary(scatter_offsets); scan_gpu_inplace( mask_flat, @@ -706,7 +713,6 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { // Binding int bind_idx = 0; const int ndim = static_cast(src.ndim()); - auto& compute_encoder = d.get_command_encoder(s.index); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(mask_flat, bind_idx++); compute_encoder.set_input_array(scatter_offsets, bind_idx++); @@ -724,4 +730,151 @@ void MaskedScatter::eval_gpu(const std::vector& inputs, array& out) { compute_encoder.dispatch_threads(grid_dims, group_dims); } +void SliceUpdate::eval_gpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + if (out.size() == 0) { + out.set_data(allocator::malloc(0)); + return; + } + + auto& in = inputs[0]; + auto& upd = inputs[1]; + + if (upd.size() == 0) { + out.copy_shared_buffer(in); + return; + } + + auto ctype = in.flags().contiguous && in.size() == in.data_size() + ? CopyType::Vector + : CopyType::General; + copy_gpu(in, out, in.data_size() == 1 ? CopyType::Scalar : ctype, stream()); + auto [data_offset, out_strides] = + prepare_slice(out, start_indices_, strides_); + + // Do copy + if (reduce_type_ == SliceUpdate::None) { + copy_gpu_inplace( + /* const array& src = */ upd, + /* array& dst = */ out, + /* const Shape& data_shape = */ upd.shape(), + /* const Strides& i_strides = */ upd.strides(), + /* const Strides& o_strides = */ out_strides, + /* int64_t i_offset = */ 0, + /* int64_t o_offset = */ data_offset, + /* CopyType ctype = */ CopyType::GeneralGeneral, + /* const Stream& s = */ stream()); + return; + } + + std::string op_name; + switch (reduce_type_) { + case SliceUpdate::None: + op_name = "none"; + break; + case SliceUpdate::Sum: + op_name = "sum"; + break; + case SliceUpdate::Prod: + op_name = "prod"; + break; + case SliceUpdate::Max: + op_name = "max"; + break; + case SliceUpdate::Min: + op_name = "min"; + break; + } + + bool upd_contiguous = upd.flags().row_contiguous; + bool upd_scalar = upd.data_size() == 1; + + Shape shape; + std::vector strides; + if (upd_scalar) { + std::tie(shape, strides) = + collapse_contiguous_dims(upd.shape(), {out_strides, out_strides}); + } else { + std::tie(shape, strides) = + collapse_contiguous_dims(upd.shape(), {upd.strides(), out_strides}); + } + + int ndim_constant = shape.size(); + if (ndim_constant > 3) { + ndim_constant = 0; + } + + int nwork = 1; + if (shape.back() % 4 == 0) { + nwork = 4; + } else if (shape.back() % 2 == 0) { + nwork = 2; + } + + auto [ds, rc, cc] = check_contiguity(shape, strides[1]); + bool out_contiguous = rc; + bool large = upd.size() > INT32_MAX; + std::string kernel_name = fmt::format( + "slice_update_{0}_{1}{2}_{3}_{4}_{5}_nw{6}_nd{7}", + op_name, + type_to_name(out), + large ? "int64_t" : "int", + out_contiguous ? "oc_true" : "oc_false", + upd_contiguous ? "updc_true" : "updc_false", + upd_scalar ? "upds_true" : "upds_false", + nwork, + ndim_constant); + + auto& s = stream(); + auto& d = metal::device(s.device); + + auto lib = d.get_library(kernel_name, [&]() { + std::string kernel_source = metal::utils(); + concatenate(kernel_source, metal::reduce_utils(), metal::scatter()); + + std::string out_type = get_type_string(out.dtype()); + std::string op_type = make_op(reduce_type_, out_type); + + kernel_source += fmt::format( + slice_update_op_kernel, + kernel_name, + out_type, + large ? "int64_t" : "int", + op_type, + out_contiguous, + upd_contiguous, + upd_scalar, + nwork, + ndim_constant); + + return kernel_source; + }); + + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = d.get_kernel(kernel_name, lib); + compute_encoder.set_compute_pipeline_state(kernel); + + // Set all the buffers + int ndim = shape.size(); + int64_t size = upd.size(); + compute_encoder.set_input_array(upd, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_vector_bytes(shape, 2); + compute_encoder.set_vector_bytes(strides[0], 3); + compute_encoder.set_bytes(ndim, 4); + compute_encoder.set_bytes(size, 5); + compute_encoder.set_vector_bytes(strides[1], 6); + compute_encoder.set_bytes(data_offset, 7); + + // Launch grid + int64_t dim0 = ndim > 0 ? shape[ndim - 1] : 1; + int64_t dim1 = ndim > 1 ? shape[ndim - 2] : 1; + int64_t rest = size / (dim0 * dim1); + dim0 /= nwork; + + auto group_dims = get_block_dims(dim0, dim1, rest); + MTL::Size grid_dims(dim0, dim1, rest); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h index a6ef0f14a..4fb1be111 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit/includes.h @@ -31,6 +31,7 @@ const char* scan(); const char* scatter_axis(); const char* softmax(); const char* sort(); +const char* searchsorted(); const char* reduce(); const char* gemm(); @@ -41,7 +42,9 @@ const char* steel_gemm_gather(); const char* steel_gemm_segmented(); const char* conv(); const char* steel_conv(); +const char* steel_conv_3d(); const char* steel_conv_general(); +const char* gemv(); const char* gemv_masked(); const char* steel_attention(); @@ -49,6 +52,7 @@ const char* gemm_nax(); const char* steel_gemm_fused_nax(); const char* steel_gemm_gather_nax(); const char* steel_gemm_splitk_nax(); +const char* steel_gemm_segmented_nax(); const char* quantized_nax(); const char* fp_quantized_nax(); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit/indexing.h b/Source/Cmlx/mlx/mlx/backend/metal/jit/indexing.h index fa141fccf..b36093b85 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit/indexing.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit/indexing.h @@ -74,3 +74,9 @@ constexpr std::string_view scatter_kernels = R"( constexpr std::string_view masked_assign_kernel = R"( template [[host_name("{0}")]] [[kernel]] decltype(masked_assign_impl<{1}, {2}>) masked_assign_impl<{1}, {2}>; )"; + +constexpr std::string_view slice_update_op_kernel = R"( +template [[host_name("{0}")]] +[[kernel]] decltype(slice_update_op_impl<{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}>) +slice_update_op_impl<{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}>; +)"; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp b/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp index b657457e9..8dfe30a15 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/jit_kernels.cpp @@ -8,6 +8,40 @@ using namespace fmt::literals; namespace mlx::core { +#ifdef MLX_METAL_NO_NAX +// NAX JIT preambles are only generated (via make_jit_source) when the SDK +// requirement is met. On older SDKs they are skipped and MLX_METAL_NO_NAX is +// defined, so is_nax_available() returns false and the get_*_nax_kernel entry +// points below are never reached. These empty definitions only exist to satisfy +// the linker for this translation unit. +namespace metal { +const char* gemm_nax() { + return ""; +} +const char* steel_gemm_fused_nax() { + return ""; +} +const char* steel_gemm_gather_nax() { + return ""; +} +const char* steel_gemm_splitk_nax() { + return ""; +} +const char* steel_gemm_segmented_nax() { + return ""; +} +const char* quantized_nax() { + return ""; +} +const char* fp_quantized_nax() { + return ""; +} +const char* steel_attention_nax() { + return ""; +} +} // namespace metal +#endif // MLX_METAL_NO_NAX + MTL::ComputePipelineState* get_arange_kernel( metal::Device& d, const std::string& kernel_name, @@ -399,6 +433,25 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + bool right) { + auto lib = d.get_library(kernel_name, [&]() { + std::ostringstream kernel_source; + // The kernel compares through LessThan, which lives in sort.h. + kernel_source << metal::utils() << metal::sort() << metal::searchsorted(); + kernel_source << get_template_definition( + kernel_name, + "searchsorted", + get_type_string(in.dtype()), + right ? "true" : "false"); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_mb_sort_kernel( metal::Device& d, const std::string& kernel_name, @@ -698,6 +751,70 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_gemv_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + bool transpose_mat, + int bm, + int bn, + int sm, + int sn, + int tm, + int tn, + bool nc, + bool axpby) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + transpose_mat ? "gemv_t" : "gemv", + get_type_string(out.dtype()), + bm, + bn, + sm, + sn, + tm, + tn, + nc ? 1 : 0, + axpby ? 1 : 0); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + +MTL::ComputePipelineState* get_gemv_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + bool transpose_mat, + int bm, + int bn, + int sm, + int sn, + int tm, + int tn) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + transpose_mat ? "gemv_t_gather" : "gemv_gather", + get_type_string(out.dtype()), + bm, + bn, + sm, + sn, + tm, + tn); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_gemv_masked_kernel( metal::Device& d, const std::string& kernel_name, @@ -739,6 +856,52 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + "gemv_wide", + get_type_string(out.dtype()), + vecs_per_tg, + k_lanes); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::gemv() + << get_template_definition( + lib_name, + "gemv_wide_gather", + get_type_string(out.dtype()), + vecs_per_tg, + k_lanes); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + MTL::ComputePipelineState* get_steel_conv_kernel( metal::Device& d, const std::string& kernel_name, @@ -770,6 +933,35 @@ MTL::ComputePipelineState* get_steel_conv_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_steel_conv_3d_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + int bm, + int bn, + int bk, + int wm, + int wn, + bool small_filter) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::utils() << metal::conv() << metal::steel_conv_3d() + << get_template_definition( + lib_name, + "implicit_gemm_conv_3d", + get_type_string(out.dtype()), + bm, + bn, + bk, + wm, + wn, + small_filter); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_steel_conv_general_kernel( metal::Device& d, const std::string& kernel_name, @@ -816,6 +1008,18 @@ MTL::ComputePipelineState* get_fft_kernel( return d.get_kernel(kernel_name, lib, hash_name, func_consts); } +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string& library_name, + const std::string& template_def) { + auto lib = d.get_library(library_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::fft() << template_def; + return kernel_source.str(); + }); + return d.get_kernel("generate_bluestein_twiddles", lib); +} + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, @@ -954,7 +1158,7 @@ MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel( const std::string& kernel_name, const std::string& hash_name, const metal::MTLFCList& func_consts, - const array& out, + const array& in, bool transpose_a, bool transpose_b, int bm, @@ -970,6 +1174,40 @@ MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel( << get_template_definition( lib_name, "gemm_splitk_nax", + get_type_string(in.dtype()), + bm, + bn, + bk, + wm, + wn, + transpose_a, + transpose_b); + return kernel_source.str(); + }); + return d.get_kernel(kernel_name, lib, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + bool transpose_a, + bool transpose_b, + int bm, + int bn, + int bk, + int wm, + int wn) { + const auto& lib_name = kernel_name; + auto lib = d.get_library(lib_name, [&]() { + std::ostringstream kernel_source; + kernel_source << metal::utils() << metal::gemm_nax() + << metal::steel_gemm_segmented_nax() + << get_template_definition( + lib_name, + "segmented_mm_nax", get_type_string(out.dtype()), bm, bn, @@ -1092,7 +1330,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bd, int wm, int wn, - const array& m) { + const array& m, + bool split_d) { const auto& lib_name = kernel_name; auto lib = d.get_library(lib_name, [&]() { std::string kernel_source; @@ -1102,7 +1341,7 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( metal::steel_attention_nax(), get_template_definition( lib_name, - "attention_nax", + split_d ? "attention_nax_dsplit" : "attention_nax", get_type_string(q.dtype()), bq, bk, diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h index 82aa4f976..18a56ebf1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels.h @@ -81,6 +81,12 @@ MTL::ComputePipelineState* get_sort_kernel( int bn, int tn); +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& in, + bool right); + MTL::ComputePipelineState* get_mb_sort_kernel( metal::Device& d, const std::string& kernel_name, @@ -201,6 +207,43 @@ MTL::ComputePipelineState* get_steel_conv_kernel( int n_channel_specialization, bool small_filter); +MTL::ComputePipelineState* get_steel_conv_3d_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + int bm, + int bn, + int bk, + int wm, + int wn, + bool small_filter); + +MTL::ComputePipelineState* get_gemv_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + bool transpose_mat, + int bm, + int bn, + int sm, + int sn, + int tm, + int tn, + bool nc, + bool axpby); + +MTL::ComputePipelineState* get_gemv_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out, + bool transpose_mat, + int bm, + int bn, + int sm, + int sn, + int tm, + int tn); + MTL::ComputePipelineState* get_gemv_masked_kernel( metal::Device& d, const std::string& kernel_name, @@ -216,6 +259,24 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( int tn, bool contiguous); +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes); + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& out, + int vecs_per_tg, + int k_lanes); + MTL::ComputePipelineState* get_steel_conv_general_kernel( metal::Device& d, const std::string& kernel_name, @@ -235,6 +296,11 @@ MTL::ComputePipelineState* get_fft_kernel( const metal::MTLFCList& func_consts, const std::string& template_def); +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string& library_name, + const std::string& template_def); + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, @@ -287,6 +353,20 @@ MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel( bool rhs); MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array& in, + bool transpose_a, + bool transpose_b, + int bm, + int bn, + int bk, + int wm, + int wn); + +MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( metal::Device& d, const std::string& kernel_name, const std::string& hash_name, @@ -346,7 +426,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int bd, int wm, int wn, - const array& m); + const array& m, + bool split_d); // Create a GPU kernel template definition for JIT compilation template diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt b/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt index d2ddf38d9..ecabbda55 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/CMakeLists.txt @@ -49,7 +49,7 @@ endfunction(build_kernel) build_kernel(arg_reduce) build_kernel(conv steel/conv/params.h) -build_kernel(gemv steel/utils.h) +build_kernel(dot) build_kernel(layer_norm) build_kernel(random) build_kernel(rms_norm) @@ -68,6 +68,7 @@ set(STEEL_HEADERS steel/conv/loaders/loader_channel_n.h steel/conv/loaders/loader_general.h steel/conv/kernels/steel_conv.h + steel/conv/kernels/steel_conv_3d.h steel/conv/kernels/steel_conv_general.h steel/gemm/gemm.h steel/gemm/mma.h @@ -109,7 +110,8 @@ set(STEEL_NAX_HEADERS steel/utils/integral_constant.h steel/gemm/kernels/steel_gemm_fused_nax.h steel/gemm/kernels/steel_gemm_gather_nax.h - steel/gemm/kernels/steel_gemm_splitk_nax.h) + steel/gemm/kernels/steel_gemm_splitk_nax.h + steel/gemm/kernels/steel_gemm_segmented_nax.h) set(STEEL_NAX_ATTN_HEADERS steel/defines.h @@ -140,25 +142,31 @@ if(NOT MLX_METAL_JIT) build_kernel(scan scan.h) build_kernel(softmax softmax.h) build_kernel(logsumexp logsumexp.h) + build_kernel(searchsorted searchsorted.h sort.h) build_kernel(sort sort.h) build_kernel(ternary ternary.h ternary_ops.h) build_kernel(unary unary.h unary_ops.h) build_kernel(steel/conv/kernels/steel_conv ${STEEL_HEADERS}) + build_kernel(steel/conv/kernels/steel_conv_3d ${STEEL_HEADERS}) build_kernel(steel/conv/kernels/steel_conv_general ${STEEL_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_fused ${STEEL_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_gather ${STEEL_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_masked ${STEEL_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_splitk ${STEEL_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_segmented ${STEEL_HEADERS}) + build_kernel(gemv steel/utils.h) build_kernel(gemv_masked steel/utils.h) build_kernel(steel/attn/kernels/steel_attention ${STEEL_ATTN_HEADERS}) - if((MLX_METAL_VERSION GREATER_EQUAL 400) AND (MACOS_SDK_VERSION GREATER_EQUAL - 26.2)) + if(MLX_METAL_VERSION GREATER_EQUAL 400 + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 + AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS}) build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS}) + build_kernel(steel/gemm/kernels/steel_gemm_segmented_nax + ${STEEL_NAX_HEADERS}) build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS}) build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h @@ -168,14 +176,24 @@ if(NOT MLX_METAL_JIT) ${STEEL_NAX_ATTN_HEADERS}) else() + message( + WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and " + "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, " + "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). " + "Building without NAX kernels.") target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX) endif() endif() +set(METAL_LINK_FLAGS) +if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") + set(METAL_LINK_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}") +endif() + add_custom_command( OUTPUT ${MLX_METAL_PATH}/mlx.metallib - COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o + COMMAND xcrun -sdk macosx metal ${METAL_LINK_FLAGS} ${KERNEL_AIR} -o ${MLX_METAL_PATH}/mlx.metallib DEPENDS ${KERNEL_AIR} COMMENT "Building mlx.metallib" diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/atomic.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/atomic.h index dacfa939a..3835fff02 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/atomic.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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; + } + } } } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h index 863d6369e..37650c7d9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/binary_ops.h @@ -16,20 +16,27 @@ struct Add { struct FloorDivide { template - T operator()(T x, T y) thread { + 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) thread { - 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) thread { - 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) thread { - 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; } }; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h index 5f654332b..fe78e0050 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/complex.h @@ -4,164 +4,286 @@ #include +#include "mlx/backend/metal/kernels/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) thread : real(real), - imag(imag) {}; - constexpr complex64_t() thread : 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)) { @@ -172,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/mlx/backend/metal/kernels/conv.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal index 50a14ad3e..6f3ce0cd4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/conv.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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++) @@ -30,7 +31,7 @@ template out_pixels *= params->oS[i]; // Set out - out += gid.z * filter_size + gid.y * (params->C); + out += (size_t)gid.z * filter_size + (size_t)gid.y * (params->C); // Coordinates in input int is[N] = {0}; @@ -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++) @@ -93,7 +96,8 @@ template out_pixels *= params->oS[i]; // Set out - out += gid.z * filter_size + gid.x * (filter_size / params->C); + out += + (size_t)gid.z * filter_size + (size_t)gid.x * (filter_size / params->C); // Coordinates in input int is[N] = {0}; @@ -102,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; @@ -149,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 \ @@ -156,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) \ @@ -205,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; @@ -276,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/mlx/backend/metal/kernels/copy.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/copy.h index cf22347ee..95ed69b76 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/copy.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/dot.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/dot.h new file mode 100644 index 000000000..6917b99f0 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/dot.h @@ -0,0 +1,68 @@ +#pragma once + +#include + +#include "mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/dot.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/dot.metal new file mode 100644 index 000000000..5c9ea3be6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/dot.metal @@ -0,0 +1,18 @@ +#include + +#include "mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/fence.air b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fence.air deleted file mode 100644 index eb751cf8c..000000000 Binary files a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fence.air and /dev/null differ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.h index e478a85b6..1d847331d 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/fft.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.metal index 590b558ef..4d457928a 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft.metal @@ -38,18 +38,22 @@ real) // clang-format off -#define instantiate_ffts(tg_mem_size) \ - instantiate_fft(tg_mem_size, float2, float2) \ +#define instantiate_c2c_ffts(tg_mem_size, complex_T) \ + instantiate_fft(tg_mem_size, complex_T, complex_T) \ + instantiate_rader(tg_mem_size, complex_T, complex_T) \ + instantiate_bluestein(tg_mem_size, complex_T, complex_T) \ + instantiate_four_step(tg_mem_size, complex_T, complex_T, 0, /*real=*/false) \ + instantiate_four_step(tg_mem_size, complex_T, complex_T, 1, /*real=*/false) + +#define instantiate_ffts(tg_mem_size) \ + instantiate_c2c_ffts(tg_mem_size, float2) \ + instantiate_c2c_ffts(tg_mem_size, complex32_t) \ instantiate_fft(tg_mem_size, float, float2) \ instantiate_fft(tg_mem_size, float2, float) \ - instantiate_rader(tg_mem_size, float2, float2) \ instantiate_rader(tg_mem_size, float, float2) \ instantiate_rader(tg_mem_size, float2, float) \ - instantiate_bluestein(tg_mem_size, float2, float2) \ instantiate_bluestein(tg_mem_size, float, float2) \ instantiate_bluestein(tg_mem_size, float2, float) \ - instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/false) \ - instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/false) \ instantiate_four_step(tg_mem_size, float, float2, 0, /*real=*/true) \ instantiate_four_step(tg_mem_size, float2, float2, 1, /*real=*/true) \ instantiate_four_step(tg_mem_size, float2, float2, 0, /*real=*/true) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/radix.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/radix.h index bd61eef6d..9d742eb33 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/radix.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/fft/readwrite.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/readwrite.h index 3d1b23f4e..86ba7df1f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/readwrite.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fft/readwrite.h @@ -2,6 +2,7 @@ #include +#include "mlx/backend/metal/kernels/complex.h" #include "mlx/backend/metal/kernels/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,7 +85,7 @@ 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_, @@ -73,21 +108,35 @@ struct ReadWriter { } // ifft(x) = 1/n * conj(fft(conj(x))) - METAL_FUNC float2 post_in(float2 elem) const thread { - 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 thread { - 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 thread { - 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 thread { - 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 thread { @@ -96,83 +145,96 @@ struct ReadWriter { return grid_index >= batch_size; } - METAL_FUNC void load() const thread { + 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 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) + 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) + 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); } } } @@ -199,53 +261,72 @@ struct ReadWriter { tg_idx / coalesce_width * elems_per_thread; } - // Four Step FFT First Step + // Four-step FFT I/O METAL_FUNC void load_strided(int stride, int overall_n) thread { - 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]); + 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) thread { - 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); + 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) thread { - // 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) thread { - 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 @@ -310,7 +391,7 @@ METAL_FUNC void ReadWriter::write() const thread { 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); } } } @@ -335,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; } @@ -368,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); } } } @@ -426,10 +507,10 @@ METAL_FUNC void ReadWriter::load() const thread { 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; } } @@ -487,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); @@ -520,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; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp4.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp4.h index 47bf4dda6..af4cbb02c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp4.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp4.h @@ -1,5 +1,8 @@ #pragma once +constant constexpr float F8E4M3_MAX = 448.0f; +constant constexpr float F4E2M1_MAX = 6.0f; + struct fp4_e2m1 { fp4_e2m1(float x) thread { if (metal::isnan(x)) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h index 796dd2163..42c5ec128 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp8.h @@ -78,3 +78,14 @@ 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); +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h index 279258b46..7061771d5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.h @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. #include #include @@ -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)); } } @@ -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/mlx/backend/metal/kernels/fp_quantized.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal index 01a2df021..d8d462288 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized.metal @@ -1,4 +1,4 @@ -// Copyright © 2025 Apple Inc. +// Copyright © 2025-2026 Apple Inc. // clang-format off #include "mlx/backend/metal/kernels/utils.h" @@ -9,19 +9,34 @@ #define instantiate_quantized(mode, name, type, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits, \ - fp_ ## name, \ - type, \ - group_size, \ - bits) + fp_ ## name, \ + type, \ + group_size, \ + bits) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + true) #define instantiate_quantized_batched(mode, name, type, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched, \ fp_ ## name, \ - type, \ - group_size, \ - bits, \ - batched) + type, \ + group_size, \ + bits, \ + batched) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_batch_" #batched "_hgs", \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + batched, \ + true) #define instantiate_quantized_aligned(mode, name, type, aligned, group_size, bits) \ instantiate_kernel( \ @@ -42,6 +57,30 @@ aligned, \ batched) +#define instantiate_quantized_qmv_fast(mode, type, results, batched, group_size, bits) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched, \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + false, \ + results) \ + instantiate_kernel( \ + #mode "_qmv_fast_" #type "_gs_" #group_size "_b_" #bits "_r_" #results "_batch_" #batched "_hgs", \ + fp_qmv_fast, \ + type, \ + group_size, \ + bits, \ + batched, \ + true, \ + results) + +#define instantiate_quantized_qmv_fast_r2(mode, type, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 1, group_size, bits) \ + instantiate_quantized_qmv_fast(mode, type, 2, 0, group_size, bits) + #define instantiate_quantized_quad(mode, name, type, D, batched, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_d_" #D "_batch_" #batched, \ @@ -52,6 +91,17 @@ D, \ batched) +#define instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, batched) \ + instantiate_kernel( \ + #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched, \ + fp_ ## name, \ + type, \ + group_size, \ + bits, \ + vecs_per_tg, \ + k_lanes, \ + batched) + #define instantiate_quantized_split_k(mode, name, type, split_k, group_size, bits) \ instantiate_kernel( \ #mode "_" #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \ @@ -105,47 +155,68 @@ instantiate_quantized_quad(mode, qmv_quad, type, 128, 1, group_size, bits) \ instantiate_quantized_quad(mode, qmv_quad, type, 128, 0, group_size, bits) +// vecs_per_tg (input-vector tile) 2..5; the fp path uses k_lanes=16. +#define instantiate_quantized_wide_wrap(mode, name, type, vecs_per_tg, k_lanes, group_size, bits) \ + instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 0) \ + instantiate_quantized_wide(mode, name, type, vecs_per_tg, k_lanes, group_size, bits, 1) + +#define instantiate_quantized_all_wide(type, mode, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 2, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 3, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 4, 16, group_size, bits) \ + instantiate_quantized_wide_wrap(mode, qmv_wide, type, 5, 16, group_size, bits) + #define instantiate_quantized_all_splitk(type, mode, group_size, bits) \ instantiate_quantized_split_k(mode, qvm_split_k, type, 8, group_size, bits) \ - instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits) + instantiate_quantized_split_k(mode, qvm_split_k, type, 32, group_size, bits) \ + instantiate_quantized_aligned(mode, qmm_t_splitk, type, true, group_size, bits) \ + instantiate_quantized_aligned(mode, qmm_t_splitk, type, false, group_size, bits) #define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nt, type, 16, 32, 32, 1, 2, true, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs, gather_qmm_rhs_nn, type, 16, 32, 32, 1, 2, false, mode, group_size, bits) -#define instantiate_quantize_dequantize(type, mode, group_size, bits) \ - instantiate_kernel( \ - #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits, \ +#define instantiate_quantize_dequantize(type, mode, group_size, bits, has_global_scale) \ + instantiate_kernel( \ + #mode "_quantize_dequantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ fp_quantize_dequantize, \ - type, \ - group_size, \ - bits) \ - instantiate_kernel( \ - #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits, \ - fp_quantize, \ - type, \ - group_size, \ - bits) \ - instantiate_kernel( \ - #mode "_dequantize_" #type "_gs_" #group_size "_b_" #bits, \ - fp_dequantize, \ - type, \ - group_size, \ - bits) + type, \ + group_size, \ + bits, \ + has_global_scale) \ + instantiate_kernel( \ + #mode "_quantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ + fp_quantize, \ + type, \ + group_size, \ + bits, \ + has_global_scale) \ + instantiate_kernel( \ + #mode "_dequantize_" #type "_gs_" #group_size "_b_" #bits "_hgs_" #has_global_scale, \ + fp_dequantize, \ + type, \ + group_size, \ + bits, \ + has_global_scale) #define instantiate_quantized_modes(type, mode, group_size, bits) \ instantiate_quantized_all_batched(type, mode, group_size, bits) \ instantiate_quantized_all_single(type, mode, group_size, bits) \ instantiate_quantized_all_quad(type, mode, group_size, bits) \ + instantiate_quantized_all_wide(type, mode, group_size, bits) \ instantiate_quantized_all_splitk(type, mode, group_size, bits) \ instantiate_quantized_all_aligned(type, mode, group_size, bits) \ - instantiate_quantized_all_rhs(type, mode, group_size, bits) \ - instantiate_quantize_dequantize(type, mode, group_size, bits) + instantiate_quantized_all_rhs(type, mode, group_size, bits) #define instantiate_quantized_types(type) \ - instantiate_quantized_modes(type, nvfp4, 16, 4) \ - instantiate_quantized_modes(type, mxfp8, 32, 8) \ - instantiate_quantized_modes(type, mxfp4, 32, 4) + instantiate_quantized_modes(type, nvfp4, 16, 4) \ + instantiate_quantized_modes(type, mxfp8, 32, 8) \ + instantiate_quantized_modes(type, mxfp4, 32, 4) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, false) \ + instantiate_quantize_dequantize(type, nvfp4, 16, 4, true) \ + instantiate_quantize_dequantize(type, mxfp8, 32, 8, false) \ + instantiate_quantize_dequantize(type, mxfp4, 32, 4, false) \ + instantiate_quantized_qmv_fast_r2(nvfp4, type, 16, 4) instantiate_quantized_types(float) instantiate_quantized_types(bfloat16_t) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h index 16676f0b5..946bce786 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.h @@ -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)); } } @@ -178,7 +179,7 @@ struct QuantizedBlockLoader { 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/mlx/backend/metal/kernels/fp_quantized_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal index 4d65a384d..771b2a963 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/fp_quantized_nax.metal @@ -24,7 +24,7 @@ type, \ group_size, \ bits, \ - aligned) + aligned, bm, bk, bn, wm, wn) #define instantiate_quantized_aligned_batched(mode, name, type, bm, bn, bk, wm, wn, aligned, batched, group_size, bits) \ instantiate_kernel( \ @@ -34,7 +34,7 @@ group_size, \ bits, \ aligned, \ - batched) + batched, bm, bk, bn, wm, wn) #define instantiate_gather_qmm_rhs(func, name, type, bm, bn, bk, wm, wn, transpose, mode, group_size, bits) \ instantiate_kernel( \ @@ -57,12 +57,18 @@ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 1, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, true, 0, group_size, bits) \ instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 1, group_size, bits) \ - instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 64, 64, 64, 2, 2, false, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, true, 0, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 1, group_size, bits) \ + instantiate_quantized_aligned_batched(mode, qmm_t_nax, type, 32, 64, 64, 2, 2, false, 0, group_size, bits) #define instantiate_quantized_all_rhs(type, mode, group_size, bits) \ instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 64, 64, 64, 2, 2, true, mode, group_size, bits) \ - instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits) + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 64, 64, 64, 2, 2, false, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nt, type, 32, 64, 64, 2, 2, true, mode, group_size, bits) \ + instantiate_gather_qmm_rhs(fp_gather_qmm_rhs_nax, gather_qmm_rhs_nax_nn, type, 32, 64, 64, 2, 2, false, mode, group_size, bits) #define instantiate_quantized_modes(type, mode, group_size, bits) \ instantiate_quantized_all_aligned(type, mode, group_size, bits) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.h new file mode 100644 index 000000000..e5535bf01 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.h @@ -0,0 +1,1044 @@ +// Copyright © 2023-2024 Apple Inc. + +#pragma once + +#include +#include + +#include "mlx/backend/metal/kernels/utils.h" + +#include "mlx/backend/metal/kernels/steel/utils.h" + +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); +} diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.metal index 277b68aa8..5637ecf73 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/gemv.metal @@ -3,510 +3,10 @@ #include #include -#include "mlx/backend/metal/kernels/utils.h" - -#include "mlx/backend/metal/kernels/steel/utils.h" +#include "mlx/backend/metal/kernels/gemv.h" 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 += 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[(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[(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 -/////////////////////////////////////////////////////////////////////////////// - -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); -} - #define instantiate_gemv_helper( \ name, itype, bm, bn, sm, sn, tm, tn, nc, axpby) \ instantiate_kernel( \ @@ -545,96 +45,6 @@ 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) */ - 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); -} - // clang-format off #define instantiate_gemv_bs_helper(nm, itype, bm, bn, sm, sn, tm, tn) \ instantiate_kernel( \ @@ -652,82 +62,6 @@ instantiate_gemv_bs_blocks(float16, half); instantiate_gemv_bs_blocks(bfloat16, bfloat16_t); instantiate_gemv_bs_blocks(complex64, complex64_t); -/////////////////////////////////////////////////////////////////////////////// -/// 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 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); -} - // clang-format off #define instantiate_gemv_t_helper( \ name, itype, bm, bn, sm, sn, tm, tn, nc, axpby) \ @@ -756,96 +90,6 @@ 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) */ - 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); -} - // clang-format off #define instantiate_gemv_t_bs_helper( \ nm, itype, bm, bn, sm, sn, tm, tn) \ @@ -866,3 +110,39 @@ 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 + +// Batch addressing and axpby are function constants. M == 1 stays with +// gemv, so tiles start at two vectors. +// clang-format off +#define instantiate_gemv_wide_helper(name, itype, nv, kl) \ + instantiate_kernel( \ + "gemv_wide_" #name "_nv" #nv "_kl" #kl, \ + gemv_wide, itype, nv, kl) + +#define instantiate_gemv_wide(name, itype, kl) \ + instantiate_gemv_wide_helper(name, itype, 2, kl) \ + instantiate_gemv_wide_helper(name, itype, 3, kl) \ + instantiate_gemv_wide_helper(name, itype, 4, kl) \ + instantiate_gemv_wide_helper(name, itype, 5, kl) // clang-format on + +instantiate_gemv_wide(float16, half, 16); +instantiate_gemv_wide(bfloat16, bfloat16_t, 16); +instantiate_gemv_wide(float16, half, 32); +instantiate_gemv_wide(bfloat16, bfloat16_t, 32); + +// clang-format off +#define instantiate_gemv_wide_gather_helper(name, itype, nv, kl) \ + instantiate_kernel( \ + "gemv_wide_gather_" #name "_nv" #nv "_kl" #kl, \ + gemv_wide_gather, itype, nv, kl) + +#define instantiate_gemv_wide_gather(name, itype, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 2, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 3, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 4, kl) \ + instantiate_gemv_wide_gather_helper(name, itype, 5, kl) // clang-format on + +instantiate_gemv_wide_gather(float16, half, 16); +instantiate_gemv_wide_gather(bfloat16, bfloat16_t, 16); +instantiate_gemv_wide_gather(float16, half, 32); +instantiate_gemv_wide_gather(bfloat16, bfloat16_t, 32); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h index f0217b336..3604f20f6 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/indexing/scatter.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/moe_stream.air b/Source/Cmlx/mlx/mlx/backend/metal/kernels/moe_stream.air deleted file mode 100644 index 2a8386a80..000000000 Binary files a/Source/Cmlx/mlx/mlx/backend/metal/kernels/moe_stream.air and /dev/null differ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h index d60bc882e..42855e14d 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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); } } } @@ -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/mlx/backend/metal/kernels/quantized.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal index f734b9bce..069482cba 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized.metal @@ -52,6 +52,17 @@ D, \ batched) +#define instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, batched) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_nv_" #vecs_per_tg "_kl_" #k_lanes "_batch_" #batched, \ + name, \ + type, \ + group_size, \ + bits, \ + vecs_per_tg, \ + k_lanes, \ + batched) + #define instantiate_quantized_split_k(name, type, group_size, bits, split_k) \ instantiate_kernel( \ #name "_" #type "_gs_" #group_size "_b_" #bits "_spk_" #split_k, \ @@ -107,9 +118,34 @@ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 1) \ instantiate_quantized_quad(affine_qmv_quad, type, group_size, bits, 128, 0) +// vecs_per_tg (input-vector tile) 2..5; affine uses k_lanes=8 (more rows per +// simdgroup) where the fp path uses 16. +#define instantiate_quantized_wide_wrap(name, type, group_size, bits, vecs_per_tg, k_lanes) \ + instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 0) \ + instantiate_quantized_wide(name, type, group_size, bits, vecs_per_tg, k_lanes, 1) + +#define instantiate_quantized_all_wide(type, group_size, bits) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 2, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 3, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 4, 8) \ + instantiate_quantized_wide_wrap(affine_qmv_wide, type, group_size, bits, 5, 8) + #define instantiate_quantized_all_splitk(type, group_size, bits) \ instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 8) \ - instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 32) + instantiate_quantized_split_k(affine_qvm_split_k, type, group_size, bits, 32) \ + +#define instantiate_quantized_splitk_qmm(name, type, group_size, bits, aligned) \ + instantiate_kernel( \ + #name "_" #type "_gs_" #group_size "_b_" #bits "_alN_" #aligned, \ + name, \ + type, \ + group_size, \ + bits, \ + aligned) + +#define instantiate_quantized_all_splitk_qmm(type, group_size, bits) \ + instantiate_quantized_splitk_qmm(affine_qmm_t_splitk, type, group_size, bits, true) \ + instantiate_quantized_splitk_qmm(affine_qmm_t_splitk, type, group_size, bits, false) #define instantiate_quantized_all_rhs(type, group_size, bits) \ instantiate_gather_qmm_rhs(affine_gather_qmm_rhs, affine_gather_qmm_rhs_nt, type, group_size, bits, 16, 32, 32, 1, 2, true) \ @@ -120,7 +156,9 @@ instantiate_quantized_all_batched(type, group_size, bits) \ instantiate_quantized_all_aligned(type, group_size, bits) \ instantiate_quantized_all_quad(type, group_size, bits) \ + instantiate_quantized_all_wide(type, group_size, bits) \ instantiate_quantized_all_splitk(type, group_size, bits) \ + instantiate_quantized_all_splitk_qmm(type, group_size, bits) \ instantiate_quantized_all_rhs(type, group_size, bits) #define instantiate_quantized_types(group_size, bits) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h index 22238112d..ed32eb59a 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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); } } } @@ -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/mlx/backend/metal/kernels/quantized_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal index 5a9c9fb87..9557fd838 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/quantized_nax.metal @@ -74,11 +74,17 @@ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 1) \ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, true, 0) \ instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 1) \ - instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 64, 64, 64, 2, 2, false, 0) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 1) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, true, 0) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 1) \ + instantiate_quantized_aligned_batched(affine_qmm_t_nax, type, group_size, bits, 32, 64, 64, 2, 2, false, 0) #define instantiate_quantized_all_rhs(type, group_size, bits) \ instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 64, 64, 64, 2, 2, true) \ - instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false) + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 64, 64, 64, 2, 2, false) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nt, type, group_size, bits, 32, 64, 64, 2, 2, true) \ + instantiate_gather_qmm_rhs(affine_gather_qmm_rhs_nax, affine_gather_qmm_rhs_nax_nn, type, group_size, bits, 32, 64, 64, 2, 2, false) #define instantiate_quantized_funcs(type, group_size, bits) \ instantiate_quantized_all_batched(type, group_size, bits) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/random.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/random.metal index ccbd464d3..13d7fd931 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/random.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/reduce.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal index de5dfbad7..3f5f91be4 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduce.metal @@ -124,11 +124,14 @@ instantiate_init_min_max(max, Max) instantiate_row_reduce_general(name##tname, itype, otype, op) \ instantiate_col_reduce_general(name##tname, itype, otype, op) -#define instantiate_and_or(name, op) \ - instantiate_reduce_functions(name, bool_, bool, bool, op) \ - instantiate_reduce_functions(name, int16, int16_t, bool, op) \ - instantiate_reduce_functions(name, int32, int32_t, bool, op) \ - instantiate_reduce_functions(name, int64, int64_t, bool, op) +#define instantiate_and_or(name, op) \ + instantiate_reduce_functions(name, bool_, bool, bool, op) \ + instantiate_reduce_functions(name, int16, int16_t, bool, op) \ + instantiate_reduce_functions(name, int32, int32_t, bool, op) \ + instantiate_reduce_functions(name, int64, int64_t, bool, op) \ + instantiate_reduce_functions(name, float16, float16_t, bool, op) \ + instantiate_reduce_functions(name, bfloat16, bfloat16_t, bool, op) \ + instantiate_reduce_functions(name, float32, float, bool, op) instantiate_and_or(and, And) instantiate_and_or(or, Or) diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/ops.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/ops.h index b7a9cacb3..f875e51c9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/ops.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/ops.h @@ -133,6 +133,15 @@ struct Sum { 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) thread { return a + b; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h index e0d08392c..47ad63fbb 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_all.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/reduction/reduce_col.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h index c109faf0b..b1546adb5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_col.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/reduction/reduce_row.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_row.h index 936d75bb5..09d1d8834 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/reduction/reduce_row.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/rms_norm.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/rms_norm.metal index 62f2457b7..eb9b5c0af 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/rms_norm.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/scaled_dot_product_attention.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal index c668d9d8c..187e6fffc 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/scaled_dot_product_attention.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/scan.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scan.h index a6bfde001..2e5493961 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/scan.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/scan.h @@ -99,15 +99,24 @@ template struct CumMax { static constexpr constant U init = Limits::min; + 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; + } + template U operator()(U a, T b) thread { - return (a >= b) ? a : b; + 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; } @@ -122,15 +131,24 @@ template struct CumMin { static constexpr constant U init = Limits::max; + 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; + } + template U operator()(U a, T b) thread { - return (a <= b) ? a : b; + 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; } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h index c7652c108..2f466a567 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sdpa_vector.h @@ -472,6 +472,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/mlx/backend/metal/kernels/searchsorted.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/searchsorted.h new file mode 100644 index 000000000..87b7c3f5b --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/searchsorted.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/searchsorted.metal new file mode 100644 index 000000000..a48ee2c32 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/searchsorted.metal @@ -0,0 +1,28 @@ +// Copyright © 2026 Apple Inc. + +#include + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/sort.h" +#include "mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/sort.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sort.h index 0d357333a..ea2640bac 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sort.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/sort.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sort.metal index 7f198d55d..d9bf20476 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/sort.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/sort.metal @@ -50,6 +50,7 @@ instantiate_block_sort_bn(bfloat16, bfloat16_t) instantiate_block_sort_long(uint64, uint64_t) instantiate_block_sort_long(int64, int64_t) +instantiate_block_sort_long(complex64, complex64_t) #define instantiate_multi_block_sort( \ vtname, vtype, itname, itype, arg_sort, bn, tn) \ @@ -77,4 +78,5 @@ instantiate_multi_block_sort_base(bfloat16, bfloat16_t) instantiate_multi_block_sort(vtname, vtype, uint32, uint32_t, true, 256, 4) instantiate_multi_block_sort_long(uint64, uint64_t) -instantiate_multi_block_sort_long(int64, int64_t) // clang-format on +instantiate_multi_block_sort_long(int64, int64_t) +instantiate_multi_block_sort_long(complex64, complex64_t) // clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h index 4de11b081..29fa7ba39 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h @@ -15,16 +15,6 @@ constant bool has_mask [[function_constant(300)]]; constant bool do_causal [[function_constant(301)]]; constant bool has_sinks [[function_constant(302)]]; -template -struct TransformScale { - T scale; - METAL_FUNC TransformScale(T scale_) : scale(scale_) {} - - METAL_FUNC T apply(T x) const { - return scale * x; - } -}; - struct MaxOp { template METAL_FUNC static constexpr T apply(T x, T y) { @@ -173,7 +163,7 @@ template < VBlockLoader loader_v( V, params->V_strides[2], Vs, simd_group_id, simd_lane_id); - TransformScale ts(static_cast(params->scale * M_LOG2E_F)); + const AccumType scale = params->scale * M_LOG2E_F; // Prepare MMA tiles constexpr short kFragSize = 8; // MMAFrag size @@ -216,13 +206,12 @@ template < threadgroup_barrier(mem_flags::mem_threadgroup); - // Load Q blocks apply scale + // Load Q blocks if (!align_Q && int(tid.x) == (params->NQ_aligned)) { loader_q.load_safe(short2(BD, params->qL_rem)); } else { loader_q.load_unsafe(); } - loader_q.apply_inplace_op(ts); // Init row reduction variables constexpr short kRowsPT = decltype(Stile)::kRowsPerThread; @@ -245,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 @@ -281,6 +275,12 @@ template < tile_matmad(Stile, Qtile, Ktile, Stile); } + // Apply scale in float32 + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < decltype(Stile)::kElemsPerTile; ii++) { + Stile.elems()[ii] *= scale; + } + // Mask out length sequence if (!align_K && kb == (params->NK_aligned)) { using stile_t = decltype(Stile); @@ -303,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; @@ -428,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); } @@ -438,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/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal index 0ff9d91b0..fbd84004f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h index bc7e5afee..4a5a9716f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h @@ -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/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal index d40df50a4..66d55539a 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.metal @@ -11,10 +11,18 @@ "_wm" #wm "_wn" #wn "_mask" #mname, \ attention_nax, dtype, bq, bk, bd, wm, wn, mtype, float) -#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ - instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ +#define instantiate_attn_dsplit(tname, dtype, bq, bk, bd, wm, wn, mname, mtype) \ + instantiate_kernel( \ + "steel_attention_dsplit_" #tname "_bq" #bq "_bk" #bk "_bd" #bd \ + "_wm" #wm "_wn" #wn "_mask" #mname, \ + attention_nax_dsplit, dtype, bq, bk, bd, wm, wn, mtype, float) + +#define instantiate_attn_shapes_helper(iname, itype, mname, mtype) \ + instantiate_attn_dsplit(iname, itype, 64, 32, 256, 4, 2, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 128, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 96, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 32, 64, 4, 1, mname, mtype) \ + instantiate_attn(iname, itype, 64, 64, 128, 4, 1, mname, mtype) \ instantiate_attn(iname, itype, 64, 64, 64, 4, 1, mname, mtype) #define instantiate_attn_mask_helper(iname, itype) \ diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/nax.h index 331cf18e1..6d978ebfb 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/attn/nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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]; + 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>(); - METAL_FUNC constexpr void clear() thread { + // 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) thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; - } - METAL_FUNC thread T* elems() thread { - return reinterpret_cast(val_frags); - } - - METAL_FUNC const thread T* elems() const thread { - return reinterpret_cast(val_frags); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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 { - 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 = {}) thread { - 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 thread { + // 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 = {}) thread { - 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)); - } - } - } + 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>(); - 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 = {}) thread { + // 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_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_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[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 thread { + // 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::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_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 thread { + // 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 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 thread { - 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{}); - }); - }); - } -}; + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); -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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - CType>(); - - STEEL_PRAGMA_UNROLL - for (short mm = 0; mm < TM; mm++) { + // Copy out results 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]; - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + 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 +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 kNumFrags = kTileRows * kTileCols; + STEEL_CONST short kElemsPerTile = kNumFrags * kElemsPerFrag; - STEEL_CONST short kSubTiles = kTileRows * kTileCols; - STEEL_CONST short kElemsPerTile = kSubTiles * kElemsPerSubTile; + STEEL_CONST short kFragThrRows = NAXFrag_t::kElemRows; + STEEL_CONST short kFragThrCols = NAXFrag_t::kElemCols; + STEEL_CONST short kFragRowsJump = NAXFrag_t::kElemRowsJump; - STEEL_CONST short kRowsPerThread = kTileRows * NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kColsPerThread = kTileCols * NAXSubTile_t::kColsPerThread; + STEEL_CONST short kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - NAXSubTile_t val_subtiles[kSubTiles]; + 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 < 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) thread { - 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 thread { - return val_subtiles[i * kTileCols + j]; + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + 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); + } + } + + 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_subtiles[0].elems()); + return reinterpret_cast(val_frags); } METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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( + 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 thread { - 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( + 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) thread { - 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>{}); - } - } + 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 { - 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>{}); - } - } + 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) thread { - 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) thread { - 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 thread { - 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 thread { - 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 @@ -993,7 +811,8 @@ struct NAXTile { 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 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 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{}); + } } } } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h index 850ec15be..f559596b7 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_3d.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_3d.h index d2fbac0fc..f2ccc1c03 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_3d.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h index 1241f7735..38250f9b8 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h index 7385158f3..15dc13ccc 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h @@ -447,5 +447,509 @@ struct Conv2DWeightBlockLoader { } }; +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DInputBlockLoaderLargeFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<3>* params; + const constant ImplicitGemmConv3DParams* gemm_params; + + short weight_d; + short weight_h; + short weight_w; + + short kdil_d; + short kdil_h; + short kdil_w; + + const device T* src[n_rows]; + + int read_n[n_rows]; + int read_id[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + /* Constructor */ + METAL_FUNC Conv3DInputBlockLoaderLargeFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + 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]]) thread + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size*(thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_d(0), + weight_h(0), + weight_w(0), + kdil_d(params_->flip ? -params_->kdil[0] : params_->kdil[0]), + kdil_h(params_->flip ? -params_->kdil[1] : params_->kdil[1]), + kdil_w(params_->flip ? -params_->kdil[2] : params_->kdil[2]) { + int out_n_pixels = params->oS[0] * params->oS[1] * params->oS[2]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_ndhw = offsets.y + bi + i * TROWS; + int n = offset_ndhw / out_n_pixels; + int dhw = offset_ndhw % out_n_pixels; + int od = dhw / (params->oS[1] * params->oS[2]); + int hw = dhw % (params->oS[1] * params->oS[2]); + int oh = hw / params->oS[2]; + int ow = hw % params->oS[2]; + + int id = od * params->str[0] - params->pad[0]; + int ih = oh * params->str[1] - params->pad[1]; + int iw = ow * params->str[2] - params->pad[2]; + + read_n[i] = n; + + if (params->flip) { + read_id[i] = id + (params->wS[0] - 1) * params->kdil[0]; + read_ih[i] = ih + (params->wS[1] - 1) * params->kdil[1]; + read_iw[i] = iw + (params->wS[2] - 1) * params->kdil[2]; + } else { + read_id[i] = id; + read_ih[i] = ih; + read_iw[i] = iw; + } + + // Adjust for flip + if (params->flip) { + id += (params->wS[0] - 1) * params->kdil[0]; + ih += (params->wS[1] - 1) * params->kdil[1]; + iw += (params->wS[2] - 1) * params->kdil[2]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + id * params->in_strides[1] + + ih * params->in_strides[2] + iw * params->in_strides[3] + bj; + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const thread { + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + int id = read_id[i] + weight_d * kdil_d; + int ih = read_ih[i] + weight_h * kdil_h; + int iw = read_iw[i] + weight_w * kdil_w; + + // Read from input if in bounds + if ((n < params->N) && (id >= 0 && id < params->iS[0]) && + (ih >= 0 && ih < params->iS[1]) && (iw >= 0 && iw < params->iS[2])) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() thread { + if (++weight_w < params->wS[2]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + if (++weight_d < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_d; + } + + return; + } + + weight_d = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DInputBlockLoaderSmallFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + using mask_t = short; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<3>* params; + const constant ImplicitGemmConv3DParams* gemm_params; + + short weight_d; + short weight_h; + short weight_w; + + const device T* src[n_rows]; + + mask_t mask_d[n_rows]; + mask_t mask_h[n_rows]; + mask_t mask_w[n_rows]; + + /* Constructor */ + METAL_FUNC Conv3DInputBlockLoaderSmallFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + 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]]) thread + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size*(thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_d(0), + weight_h(0), + weight_w(0) { + int out_n_pixels = params->oS[0] * params->oS[1] * params->oS[2]; + + int read_n[n_rows]; + int read_id[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_ndhw = offsets.y + bi + i * TROWS; + int n = offset_ndhw / out_n_pixels; + int dhw = offset_ndhw % out_n_pixels; + int od = dhw / (params->oS[1] * params->oS[2]); + int hw = dhw % (params->oS[1] * params->oS[2]); + int oh = hw / params->oS[2]; + int ow = hw % params->oS[2]; + + int id = od * params->str[0] - params->pad[0]; + int ih = oh * params->str[1] - params->pad[1]; + int iw = ow * params->str[2] - params->pad[2]; + + read_n[i] = n; + read_id[i] = id; + read_ih[i] = ih; + read_iw[i] = iw; + + // Adjust for flip + if (params->flip) { + id += (params->wS[0] - 1) * params->kdil[0]; + ih += (params->wS[1] - 1) * params->kdil[1]; + iw += (params->wS[2] - 1) * params->kdil[2]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + id * params->in_strides[1] + + ih * params->in_strides[2] + iw * params->in_strides[3] + bj; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + mask_d[i] = 0; + mask_h[i] = 0; + mask_w[i] = 0; + } + + for (short kd = 0; kd < params->wS[0]; kd++) { + short flip_d = params->flip ? params->wS[0] - kd - 1 : kd; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int n = read_n[i]; + int id = read_id[i] + flip_d * params->kdil[0]; + + bool in_bounds = n < params->N && id >= 0 && id < params->iS[0]; + + mask_d[i] |= (in_bounds << kd); + } + } + + for (short kh = 0; kh < params->wS[1]; kh++) { + short flip_h = params->flip ? params->wS[1] - kh - 1 : kh; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int ih = read_ih[i] + flip_h * params->kdil[1]; + + bool in_bounds = ih >= 0 && ih < params->iS[1]; + + mask_h[i] |= (in_bounds << kh); + } + } + + for (short kw = 0; kw < params->wS[2]; kw++) { + short flip_w = params->flip ? params->wS[2] - kw - 1 : kw; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int iw = read_iw[i] + flip_w * params->kdil[2]; + + bool in_bounds = iw >= 0 && iw < params->iS[2]; + + mask_w[i] |= (in_bounds << kw); + } + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + 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; + + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Read from input if in bounds + if ((mask_d[i] & d_mask) && (mask_h[i] & h_mask) && + (mask_w[i] & w_mask)) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() thread { + if (++weight_w < params->wS[2]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + if (++weight_d < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_d; + } + + return; + } + + weight_d = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DWeightBlockLoader { + // Destination dimensions + STEEL_CONST short BROWS = BN; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = + (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4); + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Leading dimension for src + const int src_ld; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + const device T* src; + + const constant MLXConvParams<3>* params; + + int weight_dhw; + int weight_step; + + const int read_n; + const bool do_read; + + /* Constructor */ + METAL_FUNC Conv3DWeightBlockLoader( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + 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]]) 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)), + dst(dst_ + bi * dst_ld + bj), + src(src_ + bi * src_ld + bj), + params(params_), + weight_dhw(0), + weight_step(params->C / params->groups), + read_n(offsets.y + bi), + 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 thread { + if (BN != 8 || do_read) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BN; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } + } else { + for (short i = 0; i < BN; i += TROWS) { + if ((read_n + i) < params->O) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() thread { + if (++weight_dhw < (params->wS[0] * params->wS[1] * params->wS[2])) { + src += weight_step; + return; + } + + weight_dhw = 0; + + src += + BK - (params->wS[0] * params->wS[1] * params->wS[2] - 1) * weight_step; + } +}; + } // namespace steel } // namespace mlx diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/params.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/params.h index 61b8474f8..67d382744 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/params.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/conv/params.h @@ -4,21 +4,45 @@ template struct MLXConvParams { - const int N; // Batch size - const int C; // In channels - const int O; // Out channels - const int iS[NDIM]; // Input spatial dim - const int wS[NDIM]; // Weight spatial dim - const int oS[NDIM]; // Output spatial dim - const int str[NDIM]; // Kernel strides - const int pad[NDIM]; // Input padding - const int kdil[NDIM]; // Kernel dilation - const int idil[NDIM]; // Input dilation - const int64_t in_strides[NDIM + 2]; // In strides - const int64_t wt_strides[NDIM + 2]; // Wt strides - const int64_t out_strides[NDIM + 2]; // Out strides - const int groups; // Input channel groups - const bool flip; + int N; // Batch size + int C; // In channels + int O; // Out channels + int iS[NDIM]; // Input spatial dim + int wS[NDIM]; // Weight spatial dim + int oS[NDIM]; // Output spatial dim + int str[NDIM]; // Kernel strides + int pad[NDIM]; // Input padding + int kdil[NDIM]; // Kernel dilation + int idil[NDIM]; // Input dilation + int64_t in_strides[NDIM + 2]; // In strides + int64_t wt_strides[NDIM + 2]; // Wt strides + int64_t out_strides[NDIM + 2]; // Out strides + int groups; // Input channel groups + bool flip; + + static MLXConvParams + with_padded_channels(MLXConvParams other, int pad_out, int pad_in) { + MLXConvParams params = other; + + // Update strides + for (int i = 0; i < NDIM + 1; i++) { + params.in_strides[i] = + (params.in_strides[i] / params.C) * (params.C + pad_in); + params.wt_strides[i] = + (params.wt_strides[i] / params.C) * (params.C + pad_in); + params.out_strides[i] = + (params.out_strides[i] / params.O) * (params.O + pad_out); + } + params.in_strides[NDIM + 1] = 1; + params.wt_strides[NDIM + 1] = 1; + params.out_strides[NDIM + 1] = 1; + + // Update channels + params.C += pad_in; + params.O += pad_out; + + return params; + }; }; namespace mlx { @@ -40,6 +64,23 @@ struct ImplicitGemmConv2DParams { const int swizzle_log; }; +struct ImplicitGemmConv3DParams { + const int M; + const int N; + const int K; + + const int gemm_k_iterations; + + const int inp_jump_w; + const int inp_jump_h; + const int inp_jump_d; + const int inp_jump_c; + + const int tiles_n; + const int tiles_m; + const int swizzle_log; +}; + struct Conv2DGeneralJumpParams { const int f_wgt_jump_h; const int f_wgt_jump_w; diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h index 26645cfb5..c42aa1306 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/gemm_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h index 4ff926065..8267707ad 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal index 366f1258f..97cbe4ea1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_fused_nax.metal @@ -22,6 +22,10 @@ #define instantiate_gemm_shapes_helper(iname, itype, oname, otype) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 64, 256, 2, 2) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 128, 64, 2, 4) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 64, 128, 256, 2, 4) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 64, 4, 4) \ + instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 256, 4, 4) \ instantiate_gemm_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4) instantiate_gemm_shapes_helper(float16, half, float16, half); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h index 67cd73784..7fec27060 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal index 690c7a305..6141297ba 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.metal @@ -35,3 +35,4 @@ instantiate_gather_mm_shapes_helper(float16, half, float16, half); instantiate_gather_mm_shapes_helper(bfloat16, bfloat, bfloat16, bfloat); +instantiate_gather_mm_shapes_helper(float32, float, float32, float); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h new file mode 100644 index 000000000..9cfece5c5 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.metal b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.metal new file mode 100644 index 000000000..78085cede --- /dev/null +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.metal @@ -0,0 +1,31 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "mlx/backend/metal/kernels/utils.h" + +#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h" +#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented_nax.h" + +// clang-format off +#define instantiate_segmented_mm(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \ + instantiate_kernel( \ + "steel_segmented_mm_nax_" #tname "_" #iname "_" #oname \ + "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn, \ + segmented_mm_nax, itype, bm, bn, bk, wm, wn, trans_a, trans_b, float) + +#define instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \ + instantiate_segmented_mm(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \ + instantiate_segmented_mm(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \ + instantiate_segmented_mm(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \ + instantiate_segmented_mm(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn) + +#define instantiate_segmented_mm_shapes_helper(iname, itype, oname, otype) \ + instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 256, 2, 2) \ + instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 128, 2, 2) \ + instantiate_segmented_mm_transpose_helper(iname, itype, oname, otype, 64, 64, 64, 2, 2) + +instantiate_segmented_mm_shapes_helper(float16, half, float16, half); +instantiate_segmented_mm_shapes_helper(bfloat16, bfloat, bfloat16, bfloat); +instantiate_segmented_mm_shapes_helper(float32, float, float32, float); +// clang-format on diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h index 1b6b82808..982c770c6 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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/mlx/backend/metal/kernels/steel/gemm/mma.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/mma.h index a78e57465..1d9e0a782 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/mma.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h index 39d4f6ae6..6d978ebfb 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/gemm/nax.h @@ -7,7 +7,6 @@ #include #include "mlx/backend/metal/kernels/steel/defines.h" -#include "mlx/backend/metal/kernels/steel/gemm/transforms.h" #include "mlx/backend/metal/kernels/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]; + 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>(); - METAL_FUNC constexpr void clear() thread { + // 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) thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC constexpr const thread frag_type& frag_at( - const short i, - const short j) const thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr thread frag_type& frag_at() thread { - return val_frags[i * kSubTileCols + j]; - } - - template - METAL_FUNC constexpr const thread frag_type& frag_at() const thread { - return val_frags[i * kSubTileCols + j]; - } - - METAL_FUNC thread T* elems() thread { - return reinterpret_cast(val_frags); - } - - METAL_FUNC const thread T* elems() const thread { - return reinterpret_cast(val_frags); - } - - template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { + // 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) thread { + // 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 = {}) thread { - 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 thread { + // 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 = {}) thread { - 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)); - } - } - } + 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>(); - 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 = {}) thread { + // 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_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_a[i] = Am0[i]; + ct_a[kElemsPerFrag + i] = Am1[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 thread { + // 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::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_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 thread { + // 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 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 thread { - 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< - metal::remove_addrspace_t, - metal::remove_addrspace_t, - 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]; - } - } - } + // Do matmul + gemm_op.run(ct_a, ct_b, ct_c); - // Load B into right operand registers - STEEL_PRAGMA_UNROLL - for (short nn = 0; nn < TN; nn++) { + // Copy out results 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]; - } + for (short i = 0; i < kElemsPerFrag; i++) { + Cm0[i] = ct_c[i]; + Cm1[i] = ct_c[kElemsPerFrag + 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 +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 kRowsPerThread = kTileRows * NAXFrag_t::kElemRows; + STEEL_CONST short kColsPerThread = kTileCols * NAXFrag_t::kElemCols; - STEEL_CONST short kSubTileThrRows = NAXSubTile_t::kRowsPerThread; - STEEL_CONST short kSubTileThrCols = NAXSubTile_t::kColsPerThread; + typedef typename NAXFrag_t::template dtype_frag_t frag_type; - NAXSubTile_t val_subtiles[kSubTiles]; + 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 < 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) thread { - 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 thread { - return val_subtiles[i * kTileCols + j]; + return val_frags[i * kTileCols + j]; + } + + template + METAL_FUNC constexpr thread frag_type& frag_at() thread { + return val_frags[i * kTileCols + j]; } template - METAL_FUNC constexpr const thread NAXSubTile_t& subtile_at() const thread { - return val_subtiles[i * kTileCols + 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); + } + } + + 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_subtiles[0].elems()); + return reinterpret_cast(val_frags); } METAL_FUNC const thread elem_type* elems() const thread { - return reinterpret_cast(val_subtiles[0].elems()); + return reinterpret_cast(val_frags); } template - METAL_FUNC void row_reduce(thread metal::vec& vals) const thread { - 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) thread { - 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) thread { - 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( + 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 thread { - 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( + 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) thread { - 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>{}); - } - } + 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 { - 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>{}); - } - } + 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 { - 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)], + 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) thread { - 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 thread { - 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_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 thread { - 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 @@ -1001,7 +811,8 @@ struct NAXTile { 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/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h index 2c9a1b5ba..2f153f482 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/steel/utils/integral_constant.h @@ -23,10 +23,6 @@ struct integral_constant { METAL_FUNC constexpr operator value_type() const thread noexcept { return value; } - - // METAL_FUNC constexpr value_type operator()() const noexcept { - // return value; - // } }; template @@ -132,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/mlx/backend/metal/kernels/utils.h b/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h index f2a8362ae..f15928282 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/kernels/utils.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/kernels/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()); }; /////////////////////////////////////////////////////////////////////////////// @@ -446,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/mlx/backend/metal/logsumexp.cpp b/Source/Cmlx/mlx/mlx/backend/metal/logsumexp.cpp index 2cfdcdc8a..8f7cbe3af 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/logsumexp.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/logsumexp.cpp @@ -19,14 +19,15 @@ void LogSumExp::eval_gpu(const std::vector& inputs, array& out) { } auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); // Make sure that the last dimension is contiguous - auto ensure_contiguous = [&s, &d](const array& x) { + auto ensure_contiguous = [&](const array& x) { if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) { return x; } else { array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + compute_encoder.add_temporary(x_copy); return x_copy; } }; @@ -66,7 +67,6 @@ void LogSumExp::eval_gpu(const std::vector& inputs, array& out) { kernel_name += type_to_name(out); auto kernel = get_logsumexp_kernel(d, kernel_name, out); - auto& compute_encoder = d.get_command_encoder(s.index); { MTL::Size grid_dims, group_dims; if (axis_size <= looped_limit) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/make_compiled_preamble.sh b/Source/Cmlx/mlx/mlx/backend/metal/make_compiled_preamble.sh index bb55ed3a3..1b1d9cf45 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/make_compiled_preamble.sh +++ b/Source/Cmlx/mlx/mlx/backend/metal/make_compiled_preamble.sh @@ -34,11 +34,32 @@ mkdir -p "$OUTPUT_DIR" CCC="xcrun -sdk macosx metal -x metal" HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null ) -# Remove any included system frameworks (for MetalPerformancePrimitive headers) -HDRS=$(echo "$HDRS" | grep -v "Xcode") +# Fail early if the Metal compiler returned errors instead of header paths. +# Valid lines start with '.' chars (depth indicators); anything else means +# xcrun/metal failed (e.g. missing SDK, misconfigured toolchain). +if [ -n "$HDRS" ]; then + invalid_lines=$(echo "$HDRS" | grep -v '^\.*\.' || true) + if [ -n "$invalid_lines" ]; then + echo "Error: Metal compiler header resolution failed for ${INPUT_FILE}" >&2 + echo "Expected lines starting with '.' but got:" >&2 + echo "$invalid_lines" >&2 + echo "" >&2 + echo "This usually means xcrun or the Metal toolchain is not configured correctly." >&2 + echo "Try running: xcrun -sdk macosx metal --version" >&2 + exit 1 + fi +fi + +# Remove any included system frameworks (for MetalPerformancePrimitive headers). +# Match "Xcode.app" rather than bare "Xcode" so project checkouts that live under +# a directory named Xcode/ are not filtered out along with the SDK headers. +HDRS=$(echo "$HDRS" | grep -v "Xcode.app") # Use the header depth to sort the files in order of inclusion -declare -a HDRS_LIST=($HDRS) +declare -a HDRS_LIST=() +while read -r dots path; do + [ -n "$dots" ] && HDRS_LIST+=("$dots" "$path") +done <<< "$HDRS" declare -a HDRS_STACK=() declare -a HDRS_SORTED=() diff --git a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp index 98d7b668e..8c0b46d6d 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/matmul.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "mlx/backend/common/broadcasting.h" @@ -14,6 +15,7 @@ #include "mlx/backend/metal/kernels/defines.h" #include "mlx/backend/metal/kernels/steel/gemm/params.h" #include "mlx/backend/metal/matmul.h" +#include "mlx/backend/metal/reduce.h" #include "mlx/backend/metal/utils.h" #include "mlx/primitives.h" #include "mlx/utils.h" @@ -44,7 +46,7 @@ inline array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { if (!x.flags().row_contiguous) { array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + metal::get_command_encoder(s).add_temporary(x_copy); return x_copy; } else { return x; @@ -75,7 +77,7 @@ ensure_batch_contiguous(const array& x, metal::Device& d, const Stream& s) { } array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + metal::get_command_encoder(s).add_temporary(x_copy); return std::make_tuple(false, x_copy.strides()[x_copy.ndim() - 2], x_copy); } @@ -204,6 +206,15 @@ void steel_matmul_regular_axpby_nax( int bm = 128, bn = 128, bk = 512; int wm = 4, wn = 4; + // Temp routing for larger devices + char devc = d.get_architecture().back(); + if (devc == 's' || devc == 'c' || devc == 'd') { + bk = (K >= 8192 && K > (M + N)) ? 64 : 256; + + bm = 64; + wm = 2; + } + // Prepare kernel name std::ostringstream kname; @@ -245,7 +256,7 @@ void steel_matmul_regular_axpby_nax( std::string hash_name = kname.str(); // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_fused_nax_kernel( /* metal::Device& d = */ d, /* const std::string& kernel_name = */ base_name, @@ -268,6 +279,9 @@ void steel_matmul_regular_axpby_nax( // TODO: Explore device-based tuning for swizzle int swizzle_log = tm <= 3 ? 0 : 1; + if (devc == 's' || devc == 'c' || devc == 'd') { + swizzle_log = 2; + } // Prepare steel matmul params GEMMParams params{/* const int M = */ M, @@ -322,7 +336,7 @@ void steel_matmul_regular_axpby_nax( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); // Record copies - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } template @@ -351,35 +365,6 @@ void steel_matmul_regular_axpby( int64_t C_batch_stride /* = 0*/, float alpha /* = 1.0f */, float beta /* = 0.0f */) { - if (metal::is_nax_available() && !issubdtype(a.dtype(), complexfloating) && - (env::enable_tf32() || a.dtype() != float32)) { - return steel_matmul_regular_axpby_nax( - /* const Stream& s = */ s, - /* metal::Device& d = */ d, - /* const array& a = */ a, - /* const array& b = */ b, - /* const array& c = */ c, - /* array& out = */ out, - /* int M = */ M, - /* int N = */ N, - /* int K = */ K, - /* int batch_size_out = */ batch_size_out, - /* int lda = */ lda, - /* int ldb = */ ldb, - /* int ldd = */ ldd, - /* bool transpose_a = */ transpose_a, - /* bool transpose_b = */ transpose_b, - /* std::vector& copies = */ copies, - /* Shape batch_shape = */ batch_shape, - /* Strides batch_strides = */ batch_strides, - /* int64_t A_batch_stride = */ A_batch_stride, - /* int64_t B_batch_stride = */ B_batch_stride, - /* int64_t matrix_stride_out = */ matrix_stride_out, - /* int64_t C_batch_stride = */ C_batch_stride, - /* float alpha = */ alpha, - /* float beta = */ beta); - } - using namespace mlx::steel; // Determine dispatch kernel @@ -430,7 +415,7 @@ void steel_matmul_regular_axpby( std::string hash_name = kname.str(); // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_fused_kernel( /* metal::Device& d = */ d, /* const std::string& kernel_name = */ base_name, @@ -507,7 +492,7 @@ void steel_matmul_regular_axpby( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); // Record copies - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } /////////////////////////////////////////////////////////////////////////////// @@ -575,7 +560,7 @@ void steel_gemm_splitk_axpby( << "_K_" << (k_aligned ? "t" : "n") << "aligned"; // clang-format on // Encode and dispatch gemm kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_splitk_kernel( /* metal::Device& d = */ d, /* const std::string& kernel_name = */ kname.str(), @@ -664,7 +649,7 @@ void steel_gemm_splitk_axpby( compute_encoder.dispatch_threads(grid_dims, group_dims); } - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } /////////////////////////////////////////////////////////////////////////////// @@ -692,14 +677,26 @@ void steel_gemm_splitk_axpby_nax( float beta = 0.0f) { using namespace mlx::steel; - constexpr int bm = 128, bn = 128, bk = 512; - constexpr int wm = 4, wn = 4; + int bm = 128, bn = 128, bk = 512; + int wm = 4, wn = 4; + int split_k_partition_size = 4096; + + if ((M + N) / 2 < 512 || K <= 4096) { + bm = bn = 64; + bk = 256; + wm = wn = 2; + } + if (K <= 1024) { + split_k_partition_size = K / 2; + } else if (K <= 2048) { + split_k_partition_size = 1024; + } else if (K <= 4096) { + split_k_partition_size = 2048; + } // Determine how many partitions to split K into - constexpr int split_k_partition_size = 3072; int split_k_partitions = (K + split_k_partition_size - 1) / split_k_partition_size; - const int bk_iters_per_partition = split_k_partition_size / bk; const int split_k_partition_stride = M * N; @@ -736,13 +733,13 @@ void steel_gemm_splitk_axpby_nax( std::string hash_name = kname.str(); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_splitk_nax_kernel( /* metal::Device& d = */ d, /* const std::string& kernel_name = */ base_name, /* const std::string& hash_name = */ hash_name, /* const metal::MTLFCList& func_consts = */ func_consts, - /* const array& out = */ C_split, + /* const array& in = */ a, /* bool transpose_a = */ transpose_a, /* bool transpose_b = */ transpose_b, /* int bm = */ bm, @@ -836,7 +833,7 @@ void steel_gemm_splitk_axpby_nax( compute_encoder.dispatch_threads(grid_dims, group_dims); } - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } /////////////////////////////////////////////////////////////////////////////// @@ -916,13 +913,17 @@ void steel_matmul_axpby( int _tm = (M + 16 - 1) / 16; int _tn = (N + 16 - 1) / 16; int _tk = K / 16; + int64_t matrix_size = static_cast(M) * N; + bool use_nax = metal::is_nax_available() && + !issubdtype(a.dtype(), complexfloating) && + (env::enable_tf32() || a.dtype() != float32); + char devc = d.get_architecture().back(); + int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024; // Case 1: Small M×N with large K, use SIMD split-K - char devc = d.get_architecture().back(); // Max and Ultra dispatch larger sizes to splitk - int min_tmn_threshold = (devc == 's' || devc == 'd') ? 2048 : 1024; - if (batch_size_out == 1 && (_tm * _tn) <= min_tmn_threshold && _tk >= 8 && - K >= std::max(M, N)) { + if (!use_nax && batch_size_out == 1 && (_tm * _tn) <= min_tmn_threshold && + _tk >= 8 && K >= std::max(M, N)) { return steel_gemm_splitk_axpby( /* const Stream& s = */ s, /* metal::Device& d = */ d, @@ -944,14 +945,9 @@ void steel_matmul_axpby( } // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K - // TODO: Add device-specific tuning for more NAX GPUs in the future - constexpr int min_mn_threshold = 2048 * 2048; - constexpr int min_k_threshold = 10240; - if (batch_size_out == 1 && metal::is_nax_available() && - !issubdtype(a.dtype(), complexfloating) && - (env::enable_tf32() || a.dtype() != float32) && - int64_t(M) * N >= min_mn_threshold && K >= min_k_threshold && - K >= (3 * std::max(M, N))) { + if (use_nax && batch_size_out == 1 && + (K >= 3 * std::max(M, N) || + (std::max(M, N) <= 1024 && K > 2 * std::max(M, N)))) { return steel_gemm_splitk_axpby_nax( /* const Stream& s = */ s, /* metal::Device& d = */ d, @@ -986,6 +982,34 @@ void steel_matmul_axpby( int64_t B_batch_stride_ = B_batch_stride.empty() ? 0 : B_batch_stride.back(); int64_t C_batch_stride_ = C_batch_stride.empty() ? 0 : C_batch_stride.back(); + if (use_nax) { + return steel_matmul_regular_axpby_nax( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ a, + /* const array& b = */ b, + /* const array& c = */ c, + /* array& out = */ out, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* int batch_size_out = */ batch_size_out, + /* int lda = */ lda, + /* int ldb = */ ldb, + /* int ldd = */ N, + /* bool transpose_a = */ transpose_a, + /* bool transpose_b = */ transpose_b, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ std::move(batch_shape), + /* Strides batch_strides = */ std::move(batch_strides), + /* int64_t A_batch_stride = */ A_batch_stride_, + /* int64_t B_batch_stride = */ B_batch_stride_, + /* int64_t matrix_stride_out = */ int64_t(M) * N, + /* int64_t C_batch_stride = */ C_batch_stride_, + /* float alpha = */ alpha, + /* float beta = */ beta); + } + return steel_matmul_regular_axpby( /* const Stream& s = */ s, /* metal::Device& d = */ d, @@ -1017,6 +1041,54 @@ void steel_matmul_axpby( // GEMV dispatch /////////////////////////////////////////////////////////////////////////////// +void dot_product( + const Stream& s, + metal::Device& d, + const array& a, + const array& b, + array& out, + int K, + std::vector& copies) { + constexpr int thread_group_size = 512; + constexpr int items_per_thread = 32; + constexpr int simd_groups = thread_group_size / 32; + auto& compute_encoder = metal::get_command_encoder(s); + std::string kname = "dot_product_" + type_to_name(a); + concatenate( + kname, + "_it", + items_per_thread, + "_tg", + thread_group_size, + "_sg", + simd_groups); + auto kernel = d.get_kernel(kname); + + int n = K; + int threads = (n + items_per_thread - 1) / items_per_thread; + int blocks = (threads + thread_group_size - 1) / thread_group_size; + + array partials({blocks}, float32, nullptr, {}); + partials.set_data(allocator::malloc(partials.nbytes())); + copies.push_back(partials); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(b, 1); + compute_encoder.set_output_array(partials, 2); + compute_encoder.set_bytes(n, 3); + compute_encoder.dispatch_threads( + MTL::Size(size_t(blocks) * thread_group_size, 1, 1), + MTL::Size(thread_group_size, 1, 1)); + + array tempResult(out.shape(), float32, nullptr, {}); + tempResult.set_data(allocator::malloc(tempResult.nbytes())); + copies.push_back(tempResult); + all_reduce_dispatch(partials, tempResult, "sum", compute_encoder, d, s); + copy_gpu(tempResult, out, CopyType::Scalar, s); + + compute_encoder.add_temporaries(std::move(copies)); +} + template void gemv_axbpy( const Stream& s, @@ -1119,8 +1191,20 @@ void gemv_axbpy( << "_axpby" << do_axpby; // clang-format on // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = d.get_kernel(kname.str()); + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gemv_kernel( + d, + kname.str(), + out, + transpose_mat, + bm, + bn, + sm, + sn, + tm, + tn, + !contiguous_kernel, + do_axpby); compute_encoder.set_compute_pipeline_state(kernel); int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp; @@ -1154,7 +1238,7 @@ void gemv_axbpy( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } inline void gemv( @@ -1196,6 +1280,170 @@ inline void gemv( /* Strides B_batch_stride = */ B_batch_stride); } +struct GemvWideConfig { + int vecs_per_tg; + int k_lanes; + int grid_x; +}; + +std::string gemv_wide_kernel_name( + const char* prefix, + const array& out, + const GemvWideConfig& config) { + std::ostringstream kname; + kname << prefix << type_to_name(out) << "_nv" << config.vecs_per_tg << "_kl" + << config.k_lanes; + return kname.str(); +} + +inline std::optional gemv_wide_config( + metal::Device& d, + const array& out, + int M, + int N, + int K, + int mat_ld, + int vec_ld, + int64_t mat_offset, + int64_t vec_offset, + const Strides& mat_batch_stride, + const Strides& vec_batch_stride) { + // Pre-M3 generations are limited by load issue rate rather than + // bandwidth and do not profit from the amortized stream; they keep the + // existing kernels. + if (d.get_architecture_gen() < 15) { + return std::nullopt; + } + + auto all_aligned = [](const Strides& strides, int alignment) { + return std::all_of( + strides.begin(), strides.end(), [alignment](int64_t stride) { + return stride % alignment == 0; + }); + }; + + bool vec4_aligned = mat_offset % 8 == 0 && vec_offset % 8 == 0 && + mat_ld % 4 == 0 && vec_ld % 4 == 0 && all_aligned(mat_batch_stride, 4) && + all_aligned(vec_batch_stride, 4); + + // Wide amortizes one matrix stream over several vectors; M == 1 has + // nothing to amortize and stays with gemv. + if (M <= 1 || N <= 1 || K % 4 != 0 || + (out.dtype() != float16 && out.dtype() != bfloat16) || !vec4_aligned) { + return std::nullopt; + } + + // A register tile holds at most five vectors (wider collapses occupancy); + // M balances across passes that each re-stream the matrix, and a fourth + // pass no longer pays. + int passes = (M + 4) / 5; + if (passes > 3) { + return std::nullopt; + } + + // Wide enough that the row grid alone saturates the GPU: N/4 threadgroups, + // orders of magnitude beyond concurrent execution capacity. + bool vocab_wide = N >= 65536; + + // Full rows double the threads per group and halve each lane's K share + // — more in flight where the grid is thinnest: one-pass tiles have a + // single grid column, and tiny outputs starve every column. + bool full_simd = passes == 1 || N <= 64; + int grid_x = vocab_wide ? 1 : passes; + return GemvWideConfig{(M + passes - 1) / passes, full_simd ? 32 : 16, grid_x}; +} + +bool gemv_wide( + const Stream& s, + metal::Device& d, + const array& mat, + const array& in_vec, + const array* bias, + array& out, + int M, + int N, + int K, + int mat_ld, + int vec_ld, + int batch_size_out, + const Shape& batch_shape, + const Strides& mat_batch_stride, + const Strides& vec_batch_stride, + std::vector& copies, + const Strides* bias_batch_stride, + float alpha, + float beta) { + auto plan = gemv_wide_config( + d, + out, + M, + N, + K, + mat_ld, + vec_ld, + mat.offset(), + in_vec.offset(), + mat_batch_stride, + vec_batch_stride); + if (!plan) { + return false; + } + + const bool has_batch = batch_shape.size() != 1; + const bool do_axpby = bias != nullptr; + int batch_ndim = batch_shape.size(); + + MTL::Size group_dims(32, plan->k_lanes / 8, 1); + MTL::Size grid_dims(plan->grid_x, (N + 3) / 4, batch_size_out); + + std::string base_name = gemv_wide_kernel_name("gemv_wide_", out, *plan); + std::ostringstream hash_name; + hash_name << base_name << "_nc" << has_batch << "_axpby" << do_axpby; + metal::MTLFCList func_consts = { + {&has_batch, MTL::DataType::DataTypeBool, 0}, + {&do_axpby, MTL::DataType::DataTypeBool, 1}, + }; + + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gemv_wide_kernel( + d, + base_name, + hash_name.str(), + func_consts, + out, + plan->vecs_per_tg, + plan->k_lanes); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(mat, 0); + compute_encoder.set_input_array(in_vec, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(N, 5); + compute_encoder.set_bytes(M, 6); + compute_encoder.set_bytes(mat_ld, 7); + compute_encoder.set_bytes(vec_ld, 8); + compute_encoder.set_bytes(batch_ndim, 11); + compute_encoder.set_vector_bytes(batch_shape, 12); + compute_encoder.set_vector_bytes(vec_batch_stride, 13); + compute_encoder.set_vector_bytes(mat_batch_stride, 14); + + if (do_axpby) { + compute_encoder.set_input_array(*bias, 2); + compute_encoder.set_bytes(alpha, 9); + compute_encoder.set_bytes(beta, 10); + compute_encoder.set_vector_bytes(*bias_batch_stride, 15); + int bias_ld = bias->strides()[bias->ndim() - 2]; + int bias_fd = bias->strides()[bias->ndim() - 1]; + compute_encoder.set_bytes(bias_ld, 16); + compute_encoder.set_bytes(bias_fd, 17); + } + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + compute_encoder.add_temporaries(std::move(copies)); + return true; +} + /////////////////////////////////////////////////////////////////////////////// // Matmul implementation /////////////////////////////////////////////////////////////////////////////// @@ -1207,6 +1455,7 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { } auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); auto& a_pre = inputs[0]; auto& b_pre = inputs[1]; @@ -1214,7 +1463,7 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { if (a_pre.size() == 0 || b_pre.size() == 0) { array zero = array(0, a_pre.dtype()); fill_gpu(zero, out, s); - d.add_temporary(std::move(zero), s.index); + compute_encoder.add_temporary(std::move(zero)); return; } @@ -1255,6 +1504,43 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { ///////////////////////////////////////////////////////////////////////////// // Gemv specialization + if (M == 1 && N == 1 && batch_size_out == 1 && a.flags().row_contiguous && + b.flags().row_contiguous && a.dtype() != complex64) { + return dot_product( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& a = */ a, + /* const array& b = */ b, + /* array& out = */ out, + /* int K = */ K, + /* std::vector& copies = */ copies); + } + // The wide gemv route streams the weight matrix once per <= 5 input + // vectors instead of running a row-padded GEMM tile. + if (!a_transposed && b_transposed && + gemv_wide( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& mat = */ b, + /* const array& in_vec = */ a, + /* const array* bias = */ nullptr, + /* array& out = */ out, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* int mat_ld = */ b_cols, + /* int vec_ld = */ a_cols, + /* int batch_size_out = */ static_cast(batch_size_out), + /* const Shape& batch_shape = */ batch_shape, + /* const Strides& mat_batch_stride = */ B_batch_stride, + /* const Strides& vec_batch_stride = */ A_batch_stride, + /* std::vector& copies = */ copies, + /* const Strides* bias_batch_stride = */ nullptr, + /* float alpha = */ 1.0f, + /* float beta = */ 0.0f)) { + return; + } + // Route to gemv if needed if (std::min(M, N) == 1) { return gemv( @@ -1319,6 +1605,7 @@ void AddMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); // Handle empty matrix case (K=0) if (inputs[0].shape(-1) == 0) { @@ -1332,7 +1619,7 @@ void AddMM::eval_gpu(const std::vector& inputs, array& out) { } else { array beta_scalar = array(beta_, c.dtype()); binary_op_gpu({c, beta_scalar}, out, "Multiply", s); - d.add_temporary(std::move(beta_scalar), s.index); + compute_encoder.add_temporary(std::move(beta_scalar)); } return; } @@ -1386,6 +1673,32 @@ void AddMM::eval_gpu(const std::vector& inputs, array& out) { ///////////////////////////////////////////////////////////////////////////// // Gemv specialization + // The wide gemv route streams the weight matrix once per <= 5 input + // vectors instead of running a row-padded GEMM tile. + if (!transpose_a && transpose_b && + gemv_wide( + /* const Stream& s = */ s, + /* metal::Device& d = */ d, + /* const array& mat = */ b, + /* const array& in_vec = */ a, + /* const array* bias = */ &c, + /* array& out = */ out, + /* int M = */ M, + /* int N = */ N, + /* int K = */ K, + /* int mat_ld = */ ldb, + /* int vec_ld = */ lda, + /* int batch_size_out = */ static_cast(batch_size_out), + /* const Shape& batch_shape = */ batch_shape, + /* const Strides& mat_batch_stride = */ B_batch_stride, + /* const Strides& vec_batch_stride = */ A_batch_stride, + /* std::vector& copies = */ copies, + /* const Strides* bias_batch_stride = */ &C_batch_stride, + /* float alpha = */ alpha_, + /* float beta = */ beta_)) { + return; + } + // Route to gemv if needed if (std::min(M, N) == 1) { return gemv_axbpy( @@ -1452,6 +1765,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { } auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); auto& a_pre = inputs[0]; auto& b_pre = inputs[1]; @@ -1459,7 +1773,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { if (a_pre.size() == 0 || b_pre.size() == 0) { array zero = array(0, a_pre.dtype()); fill_gpu(zero, out, s); - d.add_temporary(std::move(zero), s.index); + compute_encoder.add_temporary(std::move(zero)); return; } @@ -1630,7 +1944,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { tn, contiguous_kernel); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp; @@ -1712,7 +2026,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); return; } @@ -1735,7 +2049,6 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { << "_K_" << (k_aligned ? "t" : "n") << "aligned"; // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); auto kernel = get_steel_gemm_masked_kernel( d, kname.str(), @@ -1822,7 +2135,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { compute_encoder.dispatch_threadgroups(grid_dims, group_dims); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } /////////////////////////////////////////////////////////////////////////////// @@ -1861,7 +2174,7 @@ void gather_mm_rhs( int K = a.shape(-1); int M = a.size() / K; int N = b.shape(-1); - int lda = a.strides()[a.ndim() - 2]; // should be K + int lda = K; // Define the dispatch blocks int bm = 16, bn = 64, bk = 16; @@ -1913,7 +2226,7 @@ void gather_mm_rhs( align_K ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_gather_kernel( d, base_name, @@ -1994,7 +2307,7 @@ void gather_mm_rhs_nax( int K = a.shape(-1); int M = a.size() / K; int N = b.shape(-1); - int lda = a.strides()[a.ndim() - 2]; // should be K + int lda = K; int E = b.shape(0); // Define the dispatch blocks @@ -2056,7 +2369,7 @@ void gather_mm_rhs_nax( align_K ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_gemm_gather_nax_kernel( d, base_name, @@ -2116,12 +2429,14 @@ void gather_mv( bool is_mv, metal::Device& d, const Stream& s) { + auto& compute_encoder = metal::get_command_encoder(s); + // Copy if needed std::vector copies; auto [transpose_mat, mat_cols, mat] = check_transpose(copies, s, mat_, N == 1); auto [transpose_vec, vec_cols, vec] = check_transpose(copies, s, vec_, true); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); // If we are doing vector matrix instead of matrix vector we need to flip the // matrix transposition. Basically m @ v = v @ m.T assuming that v is treated @@ -2188,8 +2503,8 @@ void gather_mv( << tm << "_tn" << tn; // Encode and dispatch kernel - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = d.get_kernel(kname.str()); + auto kernel = get_gemv_gather_kernel( + d, kname.str(), out, transpose_mat, bm, bn, sm, sn, tm, tn); compute_encoder.set_compute_pipeline_state(kernel); int n_tgp = (out_vector_len + n_out_per_tgp - 1) / n_out_per_tgp; @@ -2222,6 +2537,138 @@ void gather_mv( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +// Dispatches gathers whose b is a transposed view (x @ W.T layouts) onto +// gemv_wide_gather; returns false when the layout or shape wants the existing +// routes instead. The route is decided from raw strides so no copy is encoded +// unless it will be dispatched. +bool gather_mm_wide( + const array& a_, + const array& b_, + const array& vec_indices, + const array& mat_indices, + array& out, + int M, + int N, + int K, + metal::Device& d, + const Stream& s) { + // b must already be a transposed view (second-to-last dim contiguous, + // inner stride != 1). check_transpose can't decide this: it copies any + // other layout to row-contiguous and reports it as non-transposed. + if (b_.strides()[b_.ndim() - 1] == 1 || b_.strides()[b_.ndim() - 2] != 1) { + return false; + } + int ldb = b_.strides()[b_.ndim() - 1]; + // a is served directly when row-contiguous, by a contiguous copy when + // neither layout matches, and declines when transposed. + bool a_direct = a_.strides()[a_.ndim() - 1] == 1; + if (!a_direct && a_.strides()[a_.ndim() - 2] == 1) { + return false; + } + int lda; + int64_t vec_offset; + Strides vec_batch_stride; + if (a_direct) { + lda = a_.strides()[a_.ndim() - 2]; + vec_offset = a_.offset(); + vec_batch_stride = Strides(a_.strides().begin(), a_.strides().end() - 2); + } else { + lda = a_.shape(-1); + vec_offset = 0; + vec_batch_stride = Strides(a_.ndim() - 2); + int64_t stride = int64_t(a_.shape(-2)) * a_.shape(-1); + for (int i = a_.ndim() - 3; i >= 0; --i) { + vec_batch_stride[i] = stride; + stride *= a_.shape(i); + } + } + Strides mat_batch_stride(b_.strides().begin(), b_.strides().end() - 2); + + auto cfg = gemv_wide_config( + d, + out, + M, + N, + K, + ldb, + lda, + b_.offset(), + vec_offset, + mat_batch_stride, + vec_batch_stride); + if (!cfg) { + return false; + } + + std::vector copies; + array a = a_; + if (!a_direct) { + a = contiguous_copy_gpu(a_, s); + copies.push_back(a); + } + const array& b = b_; + + int batch_size_out = out.size() / M / N; + int batch_ndim = out.ndim() - 2; + int batch_ndim_vec = a.ndim() - 2; + int batch_ndim_mat = b.ndim() - 2; + Strides index_strides = vec_indices.strides(); + index_strides.insert( + index_strides.end(), + mat_indices.strides().begin(), + mat_indices.strides().end()); + if (index_strides.empty()) { + // Scalar indices: the kernel still loads a stride (scaled by tid.z == 0). + index_strides = {0, 0}; + } + + MTL::Size group_dims(32, cfg->k_lanes / 8, 1); + MTL::Size grid_dims(cfg->grid_x, (N + 3) / 4, batch_size_out); + + const bool do_axpby = false; + std::string base_name = gemv_wide_kernel_name("gemv_wide_gather_", out, *cfg); + std::ostringstream hash_name; + hash_name << base_name << "_axpby0"; + metal::MTLFCList func_consts = { + {&do_axpby, MTL::DataType::DataTypeBool, 1}, + }; + + auto& compute_encoder = metal::get_command_encoder(s); + auto kernel = get_gemv_wide_gather_kernel( + d, + base_name, + hash_name.str(), + func_consts, + out, + cfg->vecs_per_tg, + cfg->k_lanes); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(b, 0); + compute_encoder.set_input_array(a, 1); + compute_encoder.set_output_array(out, 3); + compute_encoder.set_bytes(K, 4); + compute_encoder.set_bytes(N, 5); + compute_encoder.set_bytes(M, 6); + compute_encoder.set_bytes(ldb, 7); + compute_encoder.set_bytes(lda, 8); + compute_encoder.set_bytes(batch_ndim, 11); + compute_encoder.set_vector_bytes(out.shape(), 12); + compute_encoder.set_vector_bytes(index_strides, 13); + compute_encoder.set_bytes(batch_ndim_vec, 14); + compute_encoder.set_vector_bytes(a.shape(), 15); + compute_encoder.set_vector_bytes(a.strides(), 16); + compute_encoder.set_bytes(batch_ndim_mat, 17); + compute_encoder.set_vector_bytes(b.shape(), 18); + compute_encoder.set_vector_bytes(b.strides(), 19); + compute_encoder.set_input_array(vec_indices, 20); + compute_encoder.set_input_array(mat_indices, 21); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + compute_encoder.add_temporaries(std::move(copies)); + return true; +} + void gather_mm( const array& a_, const array& b_, @@ -2233,11 +2680,13 @@ void gather_mm( int K, metal::Device& d, const Stream& s) { + auto& compute_encoder = metal::get_command_encoder(s); + // Copy if needed std::vector copies; auto [transpose_a, lda, a] = check_transpose(copies, s, a_, false); auto [transpose_b, ldb, b] = check_transpose(copies, s, b_, false); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); // Determine dispatch kernel int bm = 64, bn = 64, bk = 16; @@ -2301,7 +2750,6 @@ void gather_mm( align_K ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); auto kernel = get_steel_gemm_gather_kernel( d, base_name, @@ -2331,7 +2779,7 @@ void gather_mm( (batch_ndim > 0) ? lhs_indices.strides()[0] : 0, /* const int64_t batch_stride_b = */ (batch_ndim > 0) ? rhs_indices.strides()[0] : 0, - /* const int64_t batch_stride_d = */ M * N, + /* const int64_t batch_stride_d = */ int64_t(M) * N, /* const int swizzle_log = */ 0, /* const int gemm_k_iterations_aligned = */ (K / bk), /* const int batch_ndim = */ batch_ndim}; @@ -2363,6 +2811,7 @@ void gather_mm( void GatherMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); auto& a = inputs[0]; auto& b = inputs[1]; @@ -2373,7 +2822,7 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { if (a.size() == 0 || b.size() == 0) { array zero = array(0, a.dtype()); fill_gpu(zero, out, s); - d.add_temporary(std::move(zero), s.index); + compute_encoder.add_temporary(std::move(zero)); return; } @@ -2405,6 +2854,11 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { return; } + // The wide gather route streams each gathered matrix once per <= 5 rows. + if (gather_mm_wide(a, b, lhs_indices, rhs_indices, out, M, N, K, d, s)) { + return; + } + // Route to non specialized gather mm gather_mm(a, b, lhs_indices, rhs_indices, out, M, N, K, d, s); } @@ -2419,7 +2873,9 @@ void segmented_mm( int K, metal::Device& d, const Stream& s) { - auto check_segments_layout = [&d, &s](const array& x) { + auto& compute_encoder = metal::get_command_encoder(s); + + auto check_segments_layout = [&](const array& x) { // Contiguous so return early if (x.flags().row_contiguous) { return std::make_tuple(true, x); @@ -2440,7 +2896,7 @@ void segmented_mm( } array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + compute_encoder.add_temporary(x_copy); return std::make_tuple(true, x_copy); }; @@ -2449,7 +2905,7 @@ void segmented_mm( auto [transpose_a, lda, a] = check_transpose(copies, s, a_, false); auto [transpose_b, ldb, b] = check_transpose(copies, s, b_, false); auto [segments_contiguous, segments] = check_segments_layout(segments_); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); // Determine dispatch kernel int bm = 64, bn = 64, bk = 16; @@ -2459,15 +2915,36 @@ void segmented_mm( char devc = d.get_architecture().back(); GEMM_TPARAM_MACRO(devc) + bool use_nax = metal::is_nax_available() && + (env::enable_tf32() || out.dtype() != float32); + const bool align_M = (M % bm) == 0; const bool align_N = (N % bn) == 0; - // Define the kernel name + metal::MTLFCList func_consts = { + {&segments_contiguous, MTL::DataType::DataTypeBool, 199}, + {&align_M, MTL::DataType::DataTypeBool, 200}, + {&align_N, MTL::DataType::DataTypeBool, 201}, + }; + std::string base_name; base_name.reserve(128); + base_name += "steel_segmented_mm_"; + + // Use NAX kernel if available + if (use_nax) { + int average_k = K / batch_size_out; + bm = 64; + bn = 64; + bk = (average_k >= 256) ? 256 : (average_k >= 128) ? 128 : 64; + wm = 2; + wn = 2; + + base_name += "nax_"; + } + concatenate( base_name, - "steel_segmented_mm_", transpose_a ? 't' : 'n', transpose_b ? 't' : 'n', "_", @@ -2485,13 +2962,6 @@ void segmented_mm( "_wn", wn); - metal::MTLFCList func_consts = { - {&segments_contiguous, MTL::DataType::DataTypeBool, 199}, - {&align_M, MTL::DataType::DataTypeBool, 200}, - {&align_N, MTL::DataType::DataTypeBool, 201}, - }; - - // And the kernel hash that includes the function constants std::string hash_name; hash_name.reserve(128); concatenate( @@ -2505,20 +2975,32 @@ void segmented_mm( align_N ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); - auto kernel = get_steel_gemm_segmented_kernel( - d, - base_name, - hash_name, - func_consts, - out, - transpose_a, - transpose_b, - bm, - bn, - bk, - wm, - wn); + auto kernel = (use_nax) ? get_steel_gemm_segmented_nax_kernel( + d, + base_name, + hash_name, + func_consts, + out, + transpose_a, + transpose_b, + bm, + bn, + bk, + wm, + wn) + : get_steel_gemm_segmented_kernel( + d, + base_name, + hash_name, + func_consts, + out, + transpose_a, + transpose_b, + bm, + bn, + bk, + wm, + wn); compute_encoder.set_compute_pipeline_state(kernel); // Prepare the matmul params @@ -2532,15 +3014,16 @@ void segmented_mm( /* const int tiles_m = */ (M + bm - 1) / bm, /* const int64_t batch_stride_a = */ 0, /* const int64_t batch_stride_b = */ 0, - /* const int64_t batch_stride_d = */ M * N, + /* const int64_t batch_stride_d = */ int64_t(M) * N, /* const int swizzle_log = */ 0, /* const int gemm_k_iterations_aligned = */ 0, /* const int batch_ndim = */ 0}; // Prepare the grid MTL::Size group_dims = MTL::Size(32, wn, wm); - MTL::Size grid_dims = - MTL::Size(params.tiles_n, params.tiles_m, batch_size_out); + MTL::Size grid_dims = (use_nax && bk == 64) + ? MTL::Size(batch_size_out, params.tiles_n, params.tiles_m) + : MTL::Size(params.tiles_n, params.tiles_m, batch_size_out); // Launch kernel compute_encoder.set_input_array(a, 0); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/metal.cpp b/Source/Cmlx/mlx/mlx/backend/metal/metal.cpp index 51bd2e62f..e95b6f0b9 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/metal.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/metal.cpp @@ -7,6 +7,12 @@ namespace mlx::core::metal { +namespace { + +std::string g_metallib_path; + +} // namespace + bool is_available() { return true; } @@ -14,7 +20,7 @@ bool is_available() { void start_capture(std::string path, NS::Object* object) { auto pool = new_scoped_memory_pool(); - auto descriptor = MTL::CaptureDescriptor::alloc()->init(); + auto descriptor = MTL::CaptureDescriptor::alloc()->init()->autorelease(); descriptor->setCaptureObject(object); if (!path.empty()) { @@ -27,7 +33,6 @@ void start_capture(std::string path, NS::Object* object) { auto manager = MTL::CaptureManager::sharedCaptureManager(); NS::Error* error; bool started = manager->startCapture(descriptor, &error); - descriptor->release(); if (!started) { std::ostringstream msg; msg << "[metal::start_capture] Failed to start: " @@ -47,4 +52,12 @@ void stop_capture() { manager->stopCapture(); } +void set_metallib_path(const std::string& path) { + g_metallib_path = path; +} + +const std::string& get_metallib_path() { + return g_metallib_path; +} + } // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/metal.h b/Source/Cmlx/mlx/mlx/backend/metal/metal.h index 6662e21eb..d686c2971 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/metal.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/metal.h @@ -22,4 +22,9 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/backend/metal/no_metal.cpp b/Source/Cmlx/mlx/mlx/backend/metal/no_metal.cpp index 9d6a73461..3b72ad005 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/no_metal.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/no_metal.cpp @@ -3,7 +3,6 @@ #include #include "mlx/backend/metal/metal.h" -#include "mlx/fast.h" namespace mlx::core { @@ -22,21 +21,13 @@ device_info() { "[metal::device_info] Cannot get device info without metal backend"); }; -} // namespace metal +void set_metallib_path(const std::string& path) {} -namespace fast { - -CustomKernelFunction metal_kernel( - const std::string&, - const std::vector&, - const std::vector&, - const std::string&, - const std::string&, - bool, - bool) { - throw std::runtime_error("[metal_kernel] No Metal back-end."); +const std::string& get_metallib_path() { + throw std::runtime_error( + "[metal::get_metallib_path] Cannot get metallib path without metal backend"); } -} // namespace fast +} // namespace metal } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp index 533b1927c..3795a6fb2 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/nojit_kernels.cpp @@ -100,6 +100,14 @@ MTL::ComputePipelineState* get_sort_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_searchsorted_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + bool) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_mb_sort_kernel( metal::Device& d, const std::string& kernel_name, @@ -226,6 +234,36 @@ MTL::ComputePipelineState* get_steel_gemm_segmented_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_gemv_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + bool, + int, + int, + int, + int, + int, + int, + bool, + bool) { + return d.get_kernel(kernel_name); +} + +MTL::ComputePipelineState* get_gemv_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + bool, + int, + int, + int, + int, + int, + int) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_gemv_masked_kernel( metal::Device& d, const std::string& kernel_name, @@ -243,6 +281,28 @@ MTL::ComputePipelineState* get_gemv_masked_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_gemv_wide_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + +MTL::ComputePipelineState* get_gemv_wide_gather_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_steel_conv_kernel( metal::Device& d, const std::string& kernel_name, @@ -257,6 +317,19 @@ MTL::ComputePipelineState* get_steel_conv_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_steel_conv_3d_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&, + int, + int, + int, + int, + int, + bool) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_steel_conv_general_kernel( metal::Device& d, const std::string& kernel_name, @@ -280,6 +353,13 @@ MTL::ComputePipelineState* get_fft_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_fft_twiddle_kernel( + metal::Device& d, + const std::string&, + const std::string&) { + return d.get_kernel("generate_bluestein_twiddles"); +} + MTL::ComputePipelineState* get_quantized_kernel( metal::Device& d, const std::string& kernel_name, @@ -355,6 +435,22 @@ MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel( return d.get_kernel(kernel_name, hash_name, func_consts); } +MTL::ComputePipelineState* get_steel_gemm_segmented_nax_kernel( + metal::Device& d, + const std::string& kernel_name, + const std::string& hash_name, + const metal::MTLFCList& func_consts, + const array&, + bool, + bool, + int, + int, + int, + int, + int) { + return d.get_kernel(kernel_name, hash_name, func_consts); +} + MTL::ComputePipelineState* get_qmm_nax_kernel( metal::Device& d, const std::string& kernel_name, @@ -407,7 +503,8 @@ MTL::ComputePipelineState* get_steel_attention_nax_kernel( int, int, int, - const array&) { + const array&, + bool) { return d.get_kernel(kernel_name, hash_name, func_consts); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/normalization.cpp b/Source/Cmlx/mlx/mlx/backend/metal/normalization.cpp index da0160b24..f3370f087 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/normalization.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/normalization.cpp @@ -60,7 +60,7 @@ void RMSNorm::eval_gpu( op_name += "_looped"; } op_name += type_to_name(out); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); { auto kernel = d.get_kernel(op_name); @@ -82,8 +82,7 @@ void RMSNorm::eval_gpu( uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - x.data_shared_ptr() == nullptr ? out : x, 0); + compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); compute_encoder.set_output_array(out, 2); compute_encoder.set_bytes(eps_, 3); @@ -98,6 +97,7 @@ void RMSNormVJP::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); // Ensure row contiguity. We could relax this step by checking that the array // is contiguous (no broadcasts or holes) and that the input strides are the @@ -109,11 +109,9 @@ void RMSNormVJP::eval_gpu( array x_copy = contiguous_copy_gpu(x, s); return {x_copy, true}; }; - bool donate_g = inputs[2].is_donatable(); auto [x, copied] = check_input(inputs[0]); const array& w = inputs[1]; auto [g, g_copied] = check_input(inputs[2]); - donate_g |= g_copied; array& gx = outputs[0]; array& gw = outputs[1]; @@ -131,23 +129,28 @@ void RMSNormVJP::eval_gpu( gx.set_data(allocator::malloc(gx.nbytes())); } if (g_copied && !g_in_gx) { - d.add_temporary(g, s.index); + compute_encoder.add_temporary(g); } auto axis_size = static_cast(x.shape().back()); int n_rows = x.data_size() / axis_size; + const int target_groups = 512; + uint32_t rows_per_group = 1; + int n_groups = n_rows; + if (axis_size <= RMS_LOOPED_LIMIT) { + rows_per_group = (n_rows + target_groups - 1) / target_groups; + n_groups = (n_rows + rows_per_group - 1) / rows_per_group; + } + // Allocate the gradient accumulator gw and a temporary to store the // gradients before they are accumulated. - array gw_temp = - (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w; + array gw_temp = (has_w) + ? array({n_groups, x.shape().back()}, gw.dtype(), nullptr, {}) + : w; if (has_w) { - if (!g_in_gx && donate_g) { - gw_temp.copy_shared_buffer(g); - } else { - gw_temp.set_data(allocator::malloc(gw_temp.nbytes())); - d.add_temporary(gw_temp, s.index); - } + gw_temp.set_data(allocator::malloc(gw_temp.nbytes())); + compute_encoder.add_temporary(gw_temp); } gw.set_data(allocator::malloc(gw.nbytes())); @@ -165,7 +168,6 @@ void RMSNormVJP::eval_gpu( {&has_w, MTL::DataType::DataTypeBool, 20}, }; - auto& compute_encoder = d.get_command_encoder(s.index); { auto kernel = d.get_kernel(op_name, hash_name, func_consts); @@ -175,7 +177,7 @@ void RMSNormVJP::eval_gpu( size_t simds_needed = (threadgroup_needed + simd_size - 1) / simd_size; size_t threadgroup_size = simd_size * simds_needed; assert(threadgroup_size <= kernel->maxTotalThreadsPerThreadgroup()); - size_t n_threads = n_rows * threadgroup_size; + size_t n_threads = n_groups * threadgroup_size; grid_dims = MTL::Size(n_threads, 1, 1); group_dims = MTL::Size(threadgroup_size, 1, 1); } else { @@ -195,12 +197,16 @@ void RMSNormVJP::eval_gpu( compute_encoder.set_bytes(eps_, 5); compute_encoder.set_bytes(axis_size, 6); compute_encoder.set_bytes(w_stride, 7); + if (axis_size <= looped_limit) { + compute_encoder.set_bytes(static_cast(n_rows), 8); + compute_encoder.set_bytes(rows_per_group, 9); + } compute_encoder.dispatch_threads(grid_dims, group_dims); } if (has_w) { ReductionPlan plan( - ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size}); + ReductionOpType::ContiguousStridedReduce, {n_groups}, {axis_size}); strided_reduce_general_dispatch( gw_temp, gw, "sum", plan, {0}, compute_encoder, d, s); } @@ -258,7 +264,7 @@ void LayerNorm::eval_gpu( n_reads = 4; } op_name += type_to_name(out); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); { auto kernel = d.get_kernel(op_name); @@ -287,8 +293,7 @@ void LayerNorm::eval_gpu( uint32_t w_stride = (w.ndim() == 1) ? w.strides()[0] : 0; uint32_t b_stride = (b.ndim() == 1) ? b.strides()[0] : 0; compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - x.data_shared_ptr() == nullptr ? out : x, 0); + compute_encoder.set_input_array(x, 0); compute_encoder.set_input_array(w, 1); compute_encoder.set_input_array(b, 2); compute_encoder.set_output_array(out, 3); @@ -305,6 +310,7 @@ void LayerNormVJP::eval_gpu( std::vector& outputs) { auto& s = stream(); auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); // Ensure row contiguity. We could relax this step by checking that the array // is contiguous (no broadcasts or holes) and that the input strides are the @@ -341,7 +347,7 @@ void LayerNormVJP::eval_gpu( gx.set_data(allocator::malloc(gx.nbytes())); } if (g_copied && !g_in_gx) { - d.add_temporary(g, s.index); + compute_encoder.add_temporary(g); } auto axis_size = static_cast(x.shape().back()); @@ -356,14 +362,13 @@ void LayerNormVJP::eval_gpu( gw_temp.copy_shared_buffer(g); } else { gw_temp.set_data(allocator::malloc(gw_temp.nbytes())); - d.add_temporary(gw_temp, s.index); + compute_encoder.add_temporary(gw_temp); } } gw.set_data(allocator::malloc(gw.nbytes())); gb.set_data(allocator::malloc(gb.nbytes())); // Finish with the gradient for b in case we had a b - auto& compute_encoder = d.get_command_encoder(s.index); if (gb.ndim() == 1 && gb.size() == axis_size) { ReductionPlan plan( ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size}); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp index 5f6376c5e..d1d0e781c 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/primitives.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include "mlx/backend/common/slicing.h" #include "mlx/backend/common/utils.h" @@ -11,6 +12,8 @@ #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/utils.h" +#include "mlx/dtype_utils.h" +#include "mlx/fast_primitives.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" #include "mlx/utils.h" @@ -37,48 +40,21 @@ void Arange::eval_gpu(const std::vector& inputs, array& out) { MTL::Size grid_dims = MTL::Size(nthreads, 1, 1); MTL::Size group_dims = MTL::Size( std::min(nthreads, kernel->maxTotalThreadsPerThreadgroup()), 1, 1); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - switch (out.dtype()) { - case bool_: // unsupported + dispatch_all_types(out.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + if constexpr (std::is_same_v) { throw std::runtime_error("[Arange::eval_gpu] Does not support bool"); - case uint8: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case uint64: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int8: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case int64: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case float16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case float32: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - case bfloat16: - arange_set_scalars(start_, start_ + step_, compute_encoder); - break; - default: + } else if constexpr ( + std::is_integral_v || std::is_same_v || + std::is_same_v || std::is_same_v) { + arange_set_scalars(start_, start_ + step_, compute_encoder); + } else { throw std::runtime_error("[Arange::eval_gpu] Does not support type."); - } + } + }); compute_encoder.set_output_array(out, 2); compute_encoder.dispatch_threads(grid_dims, group_dims); @@ -116,7 +92,7 @@ void ArgReduce::eval_gpu(const std::vector& inputs, array& out) { // ArgReduce int simd_size = 32; int n_reads = 4; - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); { auto kernel = d.get_kernel(op_name + type_to_name(in)); NS::UInteger thread_group_size = std::min( @@ -183,7 +159,7 @@ void RandomBits::eval_gpu(const std::vector& inputs, array& out) { // organize into grid nkeys x elem_per_key MTL::Size grid_dims = MTL::Size(num_keys, half_size + odd, 1); auto group_dims = get_block_dims(num_keys, half_size + odd, 1); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(keys, 0); compute_encoder.set_output_array(out, 1); @@ -239,4 +215,26 @@ void LUF::eval_gpu( throw std::runtime_error("[LUF::eval_gpu] Metal LU factorization NYI."); } +namespace fast { + +// There is no fused Metal cross entropy kernel yet +bool CrossEntropy::use_fallback(Stream s) { + return true; +} + +void CrossEntropy::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + throw std::runtime_error("[CrossEntropy::eval_gpu] Metal cross entropy NYI."); +} + +void CrossEntropyVJP::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + throw std::runtime_error( + "[CrossEntropyVJP::eval_gpu] Metal cross entropy NYI."); +} + +} // namespace fast + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp index cb67c74f7..1ae99a8e5 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/quantized.cpp @@ -1,5 +1,6 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. +#include "mlx/backend/common/quantized.h" #include "mlx/backend/common/broadcasting.h" #include "mlx/backend/common/compiled.h" #include "mlx/backend/gpu/copy.h" @@ -54,7 +55,7 @@ inline array ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { if (!x.flags().row_contiguous) { array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + metal::get_command_encoder(s).add_temporary(x_copy); return x_copy; } else { return x; @@ -77,15 +78,30 @@ inline array ensure_row_contiguous_matrix( } } array x_copy = contiguous_copy_gpu(x, s); - d.add_temporary(x_copy, s.index); + metal::get_command_encoder(s).add_temporary(x_copy); return x_copy; } inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { - auto arch = d.get_architecture(); - auto arch_size = arch.back(); - auto arch_gen = arch.substr(arch.size() - 3, 2); - if (arch_gen == "13" || arch_gen == "14") { + auto arch_size = d.get_architecture().back(); + auto arch_gen = d.get_architecture_gen(); + if (arch_gen >= 17 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 33; + } else if (D <= 4096 && O <= 4096) { + return 25; + } else { + return 13; + } + } else if (arch_gen >= 15 && arch_size != 'd') { + if (D <= 2048 && O <= 2048) { + return 13; + } else if (D <= 4096 && O <= 4096) { + return 15; + } else { + return 13; + } + } else if (arch_gen >= 13) { switch (arch_size) { case 'd': if (D <= 2048 && O <= 2048) { @@ -126,6 +142,12 @@ inline int get_qmv_batch_limit(int D, int O, metal::Device& d) { } } +// Must match the K step in qmv_fast_impl (kernels/quantized.h): +// pack_factor() * (bits == 2 ? 1 : 2) * SIMD_SIZE +inline int qmv_fast_k_alignment(int bits) { + return get_pack_factor(bits, 32) * (bits == 2 ? 1 : 2) * 32; +} + inline int add_strides_and_shapes( CommandEncoder& compute_encoder, bool skip, @@ -135,7 +157,7 @@ inline int add_strides_and_shapes( const std::optional& biases, int offset) { if (skip) { - return 0; + return offset; } // TODO: Collapse batch dimensions @@ -173,6 +195,209 @@ inline int add_gather_strides_and_shapes( return offset; } +auto get_quantize_kernel_dims( + MTL::ComputePipelineState* kernel, + const array& w, + const array& out, + int group_size, + int bits, + bool dequantize = false) { + // Treat uint32 as uint8 in kernel + constexpr int uint8_per_uint32 = 4; + constexpr int simd_size = 32; + int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; + int per_thread = + dequantize ? packs_per_int : std::max(group_size / simd_size, 1); + size_t nthreads = + dequantize ? out.size() / packs_per_int : w.size() / per_thread; + + NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + if (thread_group_size > nthreads) { + thread_group_size = nthreads; + } + auto group_dims = MTL::Size(thread_group_size, 1, 1); + bool use_2d = nthreads > UINT_MAX; + auto grid_shape = w.shape(); + if (dequantize) { + grid_shape.back() *= uint8_per_uint32; + } else { + grid_shape.back() /= per_thread; + } + MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) + : MTL::Size(nthreads, 1, 1); + return std::make_tuple(grid_dims, group_dims); +} + +void quantize_impl( + const std::vector& inputs, + std::vector& outputs, + QuantizationMode mode, + int group_size, + int bits, + bool dequantize, + Stream s) { + auto& w_pre = inputs[0]; + auto& out = outputs[0]; + out.set_data(allocator::malloc(out.nbytes())); + + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + bool has_biases = (mode == QuantizationMode::Affine); + bool has_global_scale = !has_biases && (inputs.size() > (1 + dequantize)); + + auto w = ensure_row_contiguous(w_pre, d, s); + if (dequantize) { + auto scales = ensure_row_contiguous(inputs[1], d, s); + if (has_biases) { + auto biases = ensure_row_contiguous(inputs[2], d, s); + compute_encoder.set_input_array(biases, 2); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[2], 2); + } + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + compute_encoder.set_output_array(out, 3); + } else { + auto& scales = outputs[1]; + scales.set_data(allocator::malloc(scales.nbytes())); + if (has_biases) { + auto& biases = outputs[2]; + biases.set_data(allocator::malloc(biases.nbytes())); + compute_encoder.set_output_array(biases, 3); + } else if (has_global_scale) { + compute_encoder.set_input_array(inputs[1], 3); + } + compute_encoder.set_input_array(w, 0); + compute_encoder.set_output_array(out, 1); + compute_encoder.set_output_array(scales, 2); + } + + auto type_string = dequantize ? get_type_string(out.dtype()) + : get_type_string(w_pre.dtype()); + auto mode_string = quantization_mode_to_string(mode); + std::string kname; + concatenate( + kname, + mode_string + (dequantize ? "_dequantize" : "_quantize"), + "_", + type_string, + "_gs_", + group_size, + "_b_", + bits); + if (!has_biases) { + concatenate(kname, "_hgs_", has_global_scale ? "true" : "false"); + } + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + dequantize ? "dequantize" : "quantize", + mode_string, + type_string, + group_size, + bits, + has_global_scale); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits, dequantize); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +auto quantize_input( + const array& w, + const std::optional& global_scale, + QuantizationMode mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + auto wq_shape = w.shape(); + wq_shape.back() = w.shape(-1) * bits / 32; + auto scales_shape = w.shape(); + scales_shape.back() = w.shape(-1) / group_size; + + std::vector inputs{w}; + if (global_scale) { + inputs.push_back(*global_scale); + } + std::vector outputs{ + array(wq_shape, uint32, nullptr, {}), + array(scales_shape, uint8, nullptr, {})}; + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.add_temporary(outputs[0]); + compute_encoder.add_temporary(outputs[1]); + quantize_impl(inputs, outputs, mode, group_size, bits, false, s); + return std::make_tuple(outputs[0], outputs[1]); +} + +void fp_quantize_dequantize( + const array& in, + const std::optional& global_scale, + array& out, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + const Stream& s) { + auto& compute_encoder = metal::get_command_encoder(s); + + auto w = ensure_row_contiguous(in, d, s); + compute_encoder.set_input_array(w, 0); + if (global_scale) { + compute_encoder.set_input_array(*global_scale, 1); + } + compute_encoder.set_output_array(out, 2); + auto type_string = get_type_string(in.dtype()); + std::string kname; + concatenate( + kname, + mode + "_quantize_dequantize_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_hgs_", + global_scale ? "true" : "false"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "quantize_dequantize", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + + auto [grid_dims, group_dims] = + get_quantize_kernel_dims(kernel, w, out, group_size, bits); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + +array quantize_dequantize_input( + const array& x_pre, + const std::optional& global_scale, + const std::string& mode, + int group_size, + int bits, + metal::Device& d, + Stream s) { + bool donate_x = x_pre.is_donatable(); + array x = ensure_row_contiguous(x_pre, d, s); + // If x is a copy it should be donatable + donate_x |= x.is_donatable(); + auto xhat = + donate_x ? x : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); + if (!donate_x) { + metal::get_command_encoder(s).add_temporary(xhat); + } + fp_quantize_dequantize(x, global_scale, xhat, mode, group_size, bits, d, s); + return xhat; +} + } // namespace void qmv_quad( @@ -215,7 +440,7 @@ void qmv_quad( B > 1 ? "_batch_1" : "_batch_0"); auto kernel = get_quantized_kernel_wrapped( d, kname, "qmv_quad", mode, type_string, group_size, bits, K, B > 1); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -238,6 +463,7 @@ void qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -251,13 +477,19 @@ void qmv( int bn = 8; int bk = 32; - MTL::Size group_dims(bk, 2, 1); - MTL::Size grid_dims(M, (N + bn - 1) / bn, B); std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % 512 == 0; + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; + // A narrower output tile reduces register pressure for large + // floating-point quantized matrix-vector products on M5 Max GPUs. + bool use_narrow_qmv = fast && N >= 4096 && d.get_architecture_gen() == 17 && + d.get_architecture().back() == 's' && mode == "nvfp4"; + int results_per_simdgroup = use_narrow_qmv ? 2 : 4; + bn = 2 * results_per_simdgroup; + MTL::Size group_dims(bk, 2, 1); + MTL::Size grid_dims(M, (N + bn - 1) / bn, B); concatenate( kname, @@ -267,7 +499,9 @@ void qmv( group_size, "_b_", bits, - B > 1 ? "_batch_1" : "_batch_0"); + use_narrow_qmv ? "_r_2" : "", + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( d, kname, @@ -276,9 +510,102 @@ void qmv( type_string, group_size, bits, - B > 1); + B > 1, + global_scale.has_value(), + results_per_simdgroup); + + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); + if (biases) { + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); + } + int c = 3; + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(out, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); +} + +// affine qmv_wide only beats qmv on gen-15+; fp benefits on every gen. +inline bool use_qmv_wide(const std::string& mode, metal::Device& d) { + return mode != "affine" || d.get_architecture_gen() >= 15; +} + +// Dispatches qmv_wide (fp modes -> fp_qmv_wide, affine -> affine_qmv_wide): +// vecs_per_tg input vectors streamed and reused per weight group. +void qmv_wide( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // vecs_per_tg is the per-threadgroup input-vector tile. Each tile re-reads + // the weights, so use the fewest tiles, then the smallest tile that fills + // them. + int n_tiles = (M + 4) / 5; // ceil(M / 5); tile size caps at 5 + int vecs_per_tg = (M + n_tiles - 1) / n_tiles; + + // k_lanes: lanes reducing K per output row (32/k_lanes rows per simdgroup). + // The affine subchunk decode has enough ALU per weight load to favor more + // rows per simdgroup (kl8); the fp modes' vectorized dot is balanced at 16. + int k_lanes = mode == "affine" ? 8 : 16; + constexpr int num_simdgroups = 2; + int B = out.size() / M / N; + bool batched = B > 1; + // Output rows per threadgroup: (32 / k_lanes) per simdgroup x num_simdgroups. + int rows_per_tg = (32 / k_lanes) * num_simdgroups; + + MTL::Size group_dims(32, num_simdgroups, 1); + MTL::Size grid_dims( + (M + vecs_per_tg - 1) / vecs_per_tg, + (N + rows_per_tg - 1) / rows_per_tg, + B); + + std::string kname; + kname.reserve(64); + std::string type_string = get_type_string(x.dtype()); + concatenate( + kname, + mode + "_qmv_wide_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + "_nv_", + vecs_per_tg, + "_kl_", + k_lanes, + batched ? "_batch_1" : "_batch_0"); + auto kernel = get_quantized_kernel_wrapped( + d, + kname, + "qmv_wide", + mode, + type_string, + group_size, + bits, + vecs_per_tg, + k_lanes, + batched); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -291,7 +618,8 @@ void qmv( compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(K, c++); compute_encoder.set_bytes(N, c++); - add_strides_and_shapes(compute_encoder, B <= 1, x, w, scales, biases, c); + compute_encoder.set_bytes(M, c++); + add_strides_and_shapes(compute_encoder, !batched, x, w, scales, biases, c); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -310,6 +638,8 @@ void qvm_split_k( metal::Device& d, const Stream& s, const std::string& mode) { + auto& compute_encoder = metal::get_command_encoder(s); + int split_k = K > 8192 ? 32 : 8; int split_D = (K + split_k - 1) / split_k; int B = out.size() / M / N; @@ -319,7 +649,7 @@ void qvm_split_k( constexpr int bk = 32; int bn = std::min(group_size, 32) * num_simdgroups; MTL::Size group_dims = MTL::Size(bk, num_simdgroups, 1); - MTL::Size grid_dims = MTL::Size(M, N / bn, B); + MTL::Size grid_dims = MTL::Size(M, (N + bn - 1) / bn, B); auto x_shape = x.shape(); auto x_strides = x.strides(); @@ -357,7 +687,7 @@ void qvm_split_k( temp_shape.insert(temp_shape.end() - 2, split_k); array intermediate(temp_shape, x.dtype(), nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); - d.add_temporary(intermediate, s.index); + compute_encoder.add_temporary(intermediate); std::string type_string = get_type_string(x.dtype()); std::string kname; @@ -377,7 +707,6 @@ void qvm_split_k( auto kernel = get_quantized_kernel_wrapped( d, kname, "qvm_split_k", mode, type_string, group_size, bits, split_k); - auto& compute_encoder = d.get_command_encoder(s.index); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -421,6 +750,7 @@ void qvm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -449,18 +779,29 @@ void qvm( group_size, "_b_", bits, - B > 1 ? "_batch_1" : "_batch_0"); + B > 1 ? "_batch_1" : "_batch_0", + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( - d, kname, "qvm", mode, type_string, group_size, bits, B > 1); - auto& compute_encoder = d.get_command_encoder(s.index); + d, + kname, + "qvm", + mode, + type_string, + group_size, + bits, + B > 1, + global_scale.has_value()); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_output_array(out, c++); compute_encoder.set_bytes(K, c++); @@ -489,7 +830,8 @@ void qmm_nax( int wm = 2; int wn = 2; - int bm = 64; + // Use smaller bm when one block covers all of M. + int bm = (M <= 32) ? 32 : 64; int bn = 64; int bk = 64; MTL::Size group_dims(32, wn, wm); @@ -554,7 +896,7 @@ void qmm_nax( wm, wn); } - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -596,7 +938,9 @@ void gather_qmm_nax( int wn = 2; int bm = 64; int bn = 64; - int bk = 32; + // The gather qmm NAX kernels are instantiated with BK = 64 only; any + // other value here makes the kernel name lookup fail. + int bk = 64; MTL::Size group_dims(32, wn, wm); MTL::Size grid_dims((N + bn - 1) / bn, (M + bm - 1) / bm, B); @@ -628,7 +972,7 @@ void gather_qmm_nax( kernel = get_qmm_nax_kernel_wrapped( d, kname, - "gather_qmm_t_nax_", + "gather_qmm_t_nax", mode, type_string, group_size, @@ -643,7 +987,7 @@ void gather_qmm_nax( kernel = get_qmm_nax_kernel_wrapped( d, kname, - "gather_qmm_n_nax_", + "gather_qmm_n_nax", mode, type_string, group_size, @@ -655,7 +999,7 @@ void gather_qmm_nax( wn); } - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -692,7 +1036,9 @@ void qmm( metal::Device& d, const Stream& s, const std::string& mode) { - if (metal::is_nax_available() && transpose && (K % 64 == 0) && + bool has_nax_kernel = + metal::is_nax_available() && (transpose || mode == "affine"); + if (has_nax_kernel && transpose && (K % 64 == 0) && (env::enable_tf32() || x.dtype() != float32)) { return qmm_nax( /* const array& x = */ x, @@ -752,7 +1098,7 @@ void qmm( kernel = get_quantized_kernel_wrapped( d, kname, "qmm_n", mode, type_string, group_size, bits, batched); } - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -771,6 +1117,105 @@ void qmm( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } +void qmm_splitk( + const array& x, + const array& w, + const array& scales, + const std::optional& biases, + array& out, + int group_size, + int bits, + int M, + int N, + int K, + metal::Device& d, + const Stream& s, + const std::string& mode) { + // Choose split_k to target ~512 threadgroups + int bm = 32, bn = 32; + int n_tiles = (N + bn - 1) / bn; + int m_tiles = (M + bm - 1) / bm; + int current_tgs = n_tiles * m_tiles; + int split_k = std::max(1, 512 / current_tgs); + + // Each K partition must be a whole number of BK-wide (32) K-tiles as well as + // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do + // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's + // group_size=16) would over-read into the next group's weights/scales. + int k_align = group_size > 32 ? group_size : 32; + split_k = std::min(split_k, K / k_align); + + // Ensure K divides evenly by split_k * k_align + while (split_k > 1 && (K % (split_k * k_align) != 0)) { + split_k--; + } + if (split_k <= 1) { + return qmm( + x, w, scales, biases, out, true, group_size, bits, M, N, K, d, s, mode); + } + + int k_partition_size = K / split_k; + int split_k_partition_stride = M * N; + + // Allocate intermediate buffer: insert split_k at the front so that + // partition_stride = M * N matches the leading stride of the buffer. + auto& compute_encoder = metal::get_command_encoder(s); + auto temp_shape = out.shape(); + if (temp_shape.size() == 1) { + temp_shape.insert(temp_shape.begin(), 1); + } + temp_shape.insert(temp_shape.begin(), split_k); + array intermediate(temp_shape, x.dtype(), nullptr, {}); + intermediate.set_data(allocator::malloc(intermediate.nbytes())); + compute_encoder.add_temporary(intermediate); + + // Grid: (N_tiles, M_tiles, split_k) + MTL::Size group_dims(32, 2, 2); + MTL::Size grid_dims(n_tiles, m_tiles, split_k); + + bool aligned = N % 32 == 0; + std::string type_string = get_type_string(x.dtype()); + std::string kname; + kname.reserve(64); + concatenate( + kname, + mode + "_qmm_t_splitk_", + type_string, + "_gs_", + group_size, + "_b_", + bits, + aligned ? "_alN_true" : "_alN_false"); + auto kernel = get_quantized_kernel_wrapped( + d, kname, "qmm_t_splitk", mode, type_string, group_size, bits, aligned); + + compute_encoder.set_compute_pipeline_state(kernel); + + int c = 0; + compute_encoder.set_input_array(w, c++); + compute_encoder.set_input_array(scales, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); + } + compute_encoder.set_input_array(x, c++); + compute_encoder.set_output_array(intermediate, c++); + compute_encoder.set_bytes(K, c++); + compute_encoder.set_bytes(N, c++); + compute_encoder.set_bytes(M, c++); + compute_encoder.set_bytes(k_partition_size, c++); + compute_encoder.set_bytes(split_k_partition_stride, c++); + + compute_encoder.dispatch_threadgroups(grid_dims, group_dims); + + // Sum across split_k dimension (axis 0) + ReductionPlan plan( + ReductionOpType::ContiguousStridedReduce, + {intermediate.shape(0)}, + {intermediate.strides(0)}); + strided_reduce_general_dispatch( + intermediate, out, "sum", plan, {0}, compute_encoder, d, s); +} + void gather_qmm( const array& x, const array& w, @@ -840,7 +1285,7 @@ void gather_qmm( d, kname, "gather_qmm_n", mode, type_string, group_size, bits); } - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); int c = 0; @@ -867,6 +1312,7 @@ void gather_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -888,7 +1334,7 @@ void gather_qmv( std::string kname; kname.reserve(64); std::string type_string = get_type_string(x.dtype()); - bool fast = N % bn == 0 && K % 512 == 0; + bool fast = N % bn == 0 && K % qmv_fast_k_alignment(bits) == 0; concatenate( kname, mode + (fast ? "_gather_qmv_fast_" : "_gather_qmv_"), @@ -896,7 +1342,8 @@ void gather_qmv( "_gs_", group_size, "_b_", - bits); + bits, + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( d, @@ -905,17 +1352,20 @@ void gather_qmv( mode, type_string, group_size, - bits); + bits, + global_scale.has_value()); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -933,6 +1383,7 @@ void gather_qvm( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, const array& lhs_indices, const array& rhs_indices, array& out, @@ -962,18 +1413,28 @@ void gather_qvm( "_gs_", group_size, "_b_", - bits); + bits, + global_scale ? "_hgs" : ""); auto kernel = get_quantized_kernel_wrapped( - d, kname, "gather_qvm", mode, type_string, group_size, bits); - auto& compute_encoder = d.get_command_encoder(s.index); + d, + kname, + "gather_qvm", + mode, + type_string, + group_size, + bits, + global_scale.has_value()); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); - int c = 0; - compute_encoder.set_input_array(w, c++); - compute_encoder.set_input_array(scales, c++); + compute_encoder.set_input_array(w, 0); + compute_encoder.set_input_array(scales, 1); if (biases) { - compute_encoder.set_input_array(*biases, c++); + compute_encoder.set_input_array(*biases, 2); + } else if (global_scale) { + compute_encoder.set_input_array(*global_scale, 2); } + int c = 3; compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(lhs_indices, c++); compute_encoder.set_input_array(rhs_indices, c++); @@ -1026,9 +1487,15 @@ void gather_qmm_rhs_nax( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } - // TODO: Tune the block sizes - int bm = 64, bn = 64, bk = 64; + // Use smaller bm for many experts and few tokens. + int E = w.size() / w.shape(-1) / w.shape(-2); + int bm = (M / E < 64) ? 32 : 64; + int bn = 64, bk = 64; int wm = 2, wn = 2; const bool align_M = (M % bm) == 0; @@ -1079,7 +1546,7 @@ void gather_qmm_rhs_nax( align_K ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_nax_kernel( d, kname, @@ -1104,9 +1571,8 @@ void gather_qmm_rhs_nax( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); @@ -1177,6 +1643,10 @@ void gather_qmm_rhs( array x = broadcast_with_indices(x_); array w = ensure_row_contiguous(w_, d, s); array scales = ensure_row_contiguous(scales_, d, s); + std::optional biases; + if (biases_) { + biases = ensure_row_contiguous(*biases_, d, s); + } // TODO: Tune the block sizes int bm = 16, bn = 32, bk = 32; @@ -1229,7 +1699,7 @@ void gather_qmm_rhs( align_K ? 't' : 'n'); // Get and set the kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_gather_qmm_kernel( d, kname, @@ -1254,9 +1724,8 @@ void gather_qmm_rhs( compute_encoder.set_input_array(x, c++); compute_encoder.set_input_array(w, c++); compute_encoder.set_input_array(scales, c++); - if (biases_) { - array biases = ensure_row_contiguous(*biases_, d, s); - compute_encoder.set_input_array(biases, c++); + if (biases) { + compute_encoder.set_input_array(*biases, c++); } compute_encoder.set_input_array(indices, c++); compute_encoder.set_output_array(out, c++); @@ -1272,6 +1741,7 @@ void dispatch_qmv( const array& w, const array& scales, const std::optional& biases, + const std::optional& global_scale, array& out, int group_size, int bits, @@ -1282,13 +1752,31 @@ void dispatch_qmv( const Stream& s, const std::string& mode) { // It is a qmv with a small inner dimension so route to qmv_quad kernel - if ((K == 128 || K == 64) && is_power_of_2(bits)) { + if ((K == 128 || K == 64) && is_power_of_2(bits) && !global_scale) { qmv_quad(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); return; } - // Run of the mill qmv - qmv(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + // Small batch so route to qmv_wide, which reuses each weight group across the + // M vectors. + if (M >= 2 && use_qmv_wide(mode, d) && !global_scale) { + qmv_wide(x, w, scales, biases, out, group_size, bits, M, N, K, d, s, mode); + return; + } + qmv(x, + w, + scales, + biases, + global_scale, + out, + group_size, + bits, + M, + N, + K, + d, + s, + mode); } void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { @@ -1317,6 +1805,13 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { auto mode = quantization_mode_to_string(mode_); // It is a matrix matrix product. if (M >= vector_limit) { + // Use split-K qmm for small M with transposed weights (non-batched only) + int B = out.size() / M / N; + if (transpose_ && B == 1) { + qmm_splitk( + x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + return; + } qmm(x, w, scales, @@ -1337,13 +1832,39 @@ void QuantizedMatmul::eval_gpu(const std::vector& inputs, array& out) { // Run of the mill qmv if (transpose_) { dispatch_qmv( - x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); return; } // Run of the mill qvm if (K < 1024) { - qvm(x, w, scales, biases, out, group_size_, bits_, M, N, K, d, s, mode); + qvm(x, + w, + scales, + biases, + std::nullopt, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); return; } @@ -1429,6 +1950,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1448,6 +1970,7 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { w, scales, biases, + std::nullopt, lhs_indices, rhs_indices, out, @@ -1461,183 +1984,133 @@ void GatherQMM::eval_gpu(const std::vector& inputs, array& out) { mode); } -void quantize_dequantize( - const array& in, - array& out, - std::string mode, - int group_size, - int bits, - metal::Device& d, - const Stream& s) { - auto& compute_encoder = d.get_command_encoder(s.index); - - auto w = ensure_row_contiguous(in, d, s); - compute_encoder.set_input_array(w, 0); - compute_encoder.set_output_array(out, 1); - auto type_string = get_type_string(in.dtype()); - std::string kname; - concatenate( - kname, - mode + "_quantize_dequantize_", - type_string, - "_gs_", - group_size, - "_b_", - bits); - auto kernel = get_quantized_kernel_wrapped( - d, kname, "quantize_dequantize", mode, type_string, group_size, bits); +void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { + auto& s = stream(); + auto& d = metal::device(s.device); - compute_encoder.set_compute_pipeline_state(kernel); + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + auto mode = quantization_mode_to_string(mode_); - constexpr int uint8_per_uint32 = 4; - constexpr int simd_size = 32; - int packs_per_int = (bits == 3 || bits == 5) ? 8 : bits == 6 ? 4 : 8 / bits; - int per_thread = std::max(group_size / simd_size, 1); - size_t nthreads = w.size() / per_thread; + out.set_data(allocator::malloc(out.nbytes())); - NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - if (thread_group_size > nthreads) { - thread_group_size = nthreads; + // - 2 inputs: x, w (non-quantized w) + // - 3 inputs: x, w, scales_w (quantized w) + bool w_quantized = (inputs[1].dtype() == uint32); + int base_size = w_quantized ? 3 : 2; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; } - auto group_dims = MTL::Size(thread_group_size, 1, 1); - bool use_2d = nthreads > UINT_MAX; - auto grid_shape = w.shape(); - grid_shape.back() /= per_thread; - MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) - : MTL::Size(nthreads, 1, 1); - compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); + + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); + + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + dispatch_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); } -void QQMatmul::eval_gpu(const std::vector& inputs, array& out) { +void GatherQQMM::eval_gpu(const std::vector& inputs, array& out) { auto& s = stream(); auto& d = metal::device(s.device); + const array& x_pre = inputs[0]; + const array& w_pre = inputs[1]; + const array& lhs_indices = ensure_row_contiguous(inputs[2], d, s); + const array& rhs_indices = ensure_row_contiguous(inputs[3], d, s); auto mode = quantization_mode_to_string(mode_); - bool w_quantized = (inputs[1].dtype() == uint32); - if (w_quantized && inputs[0].shape(-2) == 1) { - out.set_data(allocator::malloc(out.nbytes())); - - bool donate_x = inputs[0].is_donatable(); - array x = ensure_row_contiguous(inputs[0], d, s); - // If x is a copy it should be donatable - donate_x |= x.is_donatable(); - auto xhat = donate_x - ? x - : array(allocator::malloc(x.nbytes()), x.shape(), x.dtype()); - quantize_dequantize(x, xhat, mode, group_size_, bits_, d, s); - - // Make sure the last two dims of w and s are contiguous - array w = ensure_row_contiguous_matrix(inputs[1], d, s); - array scales = ensure_row_contiguous_matrix(inputs[2], d, s); - - bool non_batched = w.ndim() == 2; - int K = x.shape(-1); - int M = non_batched ? x.size() / K : x.shape(-2); - int N = out.shape(-1); - dispatch_qmv( - xhat, - w, - scales, - std::nullopt, - out, - group_size_, - bits_, - M, - N, - K, - d, - s, - mode); - return; - } else { - throw std::runtime_error("[QQMatmul] NYI for the general case"); - } -} -void fast::Quantize::eval_gpu( - const std::vector& inputs, - std::vector& outputs) { - auto& w_pre = inputs[0]; - auto& out = outputs[0]; out.set_data(allocator::malloc(out.nbytes())); - auto& s = stream(); - auto& d = metal::device(s.device); - auto& compute_encoder = d.get_command_encoder(s.index); - - auto w = ensure_row_contiguous(w_pre, d, s); - if (dequantize_) { - auto scales = ensure_row_contiguous(inputs[1], d, s); - if (mode_ == QuantizationMode::Affine) { - auto biases = ensure_row_contiguous(inputs[2], d, s); - compute_encoder.set_input_array(biases, 2); - } - compute_encoder.set_input_array(w, 0); - compute_encoder.set_input_array(scales, 1); - compute_encoder.set_output_array(out, 3); - } else { - auto& scales = outputs[1]; - scales.set_data(allocator::malloc(scales.nbytes())); - if (mode_ == QuantizationMode::Affine) { - auto& biases = outputs[2]; - biases.set_data(allocator::malloc(biases.nbytes())); - compute_encoder.set_output_array(biases, 3); - } - compute_encoder.set_input_array(w, 0); - compute_encoder.set_output_array(out, 1); - compute_encoder.set_output_array(scales, 2); + // - 4 inputs: x, w (non-quantized w) + // - 5 inputs: x, w, scales_w (quantized w) + bool w_quantized = (inputs[1].dtype() == uint32); + int base_size = w_quantized ? 5 : 4; + // For nvfp4, global scales are optional but must be both present or both + // absent If present, they add 2 more inputs (global_scale_x, global_scale_w) + bool has_global_scales = + mode_ == QuantizationMode::Nvfp4 && inputs.size() == base_size + 2; + assert(inputs.size() == base_size || has_global_scales); + + std::optional global_scale_x; + std::optional global_scale_w; + if (has_global_scales) { + global_scale_x = inputs[inputs.size() - 2]; + global_scale_w = inputs[inputs.size() - 1]; } - auto type_string = dequantize_ ? get_type_string(out.dtype()) - : get_type_string(w_pre.dtype()); - auto mode = quantization_mode_to_string(mode_); - std::string kname; - concatenate( - kname, - mode + (dequantize_ ? "_dequantize" : "_quantize"), - "_", - type_string, - "_gs_", - group_size_, - "_b_", - bits_); - auto kernel = get_quantized_kernel_wrapped( - d, - kname, - dequantize_ ? "dequantize" : "quantize", - mode, - type_string, - group_size_, - bits_); + // Quantize weights. + auto [w_q, scales_w] = !w_quantized + ? quantize_input(w_pre, global_scale_w, mode_, group_size_, bits_, d, s) + : std::make_tuple( + ensure_row_contiguous_matrix(w_pre, d, s), + ensure_row_contiguous_matrix(inputs[base_size - 1], d, s)); - compute_encoder.set_compute_pipeline_state(kernel); + // Quantize activation. + array x = quantize_dequantize_input( + x_pre, global_scale_x, mode, group_size_, bits_, d, s); - // Treat uint32 as uint8 in kernel - constexpr int uint8_per_uint32 = 4; - constexpr int simd_size = 32; - int packs_per_int = (bits_ == 3 || bits_ == 5) ? 8 - : bits_ == 6 ? 4 - : 8 / bits_; - int per_thread = - dequantize_ ? packs_per_int : std::max(group_size_ / simd_size, 1); - size_t nthreads = - dequantize_ ? out.size() / packs_per_int : w.size() / per_thread; + bool non_batched = w_q.ndim() == 2; + int K = x.shape(-1); + int M = non_batched ? x.size() / K : x.shape(-2); + int N = out.shape(-1); + gather_qmv( + x, + w_q, + scales_w, + std::nullopt, + global_scale_w, + lhs_indices, + rhs_indices, + out, + group_size_, + bits_, + M, + N, + K, + d, + s, + mode); +} - NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - if (thread_group_size > nthreads) { - thread_group_size = nthreads; - } - auto group_dims = MTL::Size(thread_group_size, 1, 1); - bool use_2d = nthreads > UINT_MAX; - auto grid_shape = w.shape(); - if (dequantize_) { - grid_shape.back() *= uint8_per_uint32; - } else { - grid_shape.back() /= per_thread; - } - MTL::Size grid_dims = use_2d ? get_2d_grid_dims(grid_shape, w.strides()) - : MTL::Size(nthreads, 1, 1); - compute_encoder.dispatch_threads(grid_dims, group_dims); +void fast::Quantize::eval_gpu( + const std::vector& inputs, + std::vector& outputs) { + quantize_impl( + inputs, outputs, mode_, group_size_, bits_, dequantize_, stream()); } void fast::ConvertFP8::eval_gpu( diff --git a/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp b/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp index 504943d82..ac11ac635 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/reduce.cpp @@ -273,6 +273,21 @@ std::pair remap_reduce_types( } return {in.dtype(), in.dtype()}; } else if (op_name == "and" || op_name == "or") { + // Integers can be tested as whatever type has the same width, since only + // their bits matter. Floats cannot: -0.0 compares equal to zero but has a + // bit set, so it has to be tested as a float. complex64 stays on the + // integer path, since testing it as a complex would only look at the real + // part and miss 1j. + switch (in.dtype()) { + case float16: + return {float16, bool_}; + case bfloat16: + return {bfloat16, bool_}; + case float32: + return {float32, bool_}; + default: + break; + } if (in.dtype().size() == 1) { return {bool_, bool_}; } else if (in.dtype().size() == 2) { @@ -292,6 +307,9 @@ void init_reduce( CommandEncoder& compute_encoder, metal::Device& d, const Stream& s) { + if (out.size() == 0) { + return; + } auto [_, out_type] = remap_reduce_types(out, op_name); const std::string func_name = "init_reduce"; std::string kname = func_name; @@ -358,7 +376,7 @@ void all_reduce_dispatch( // Allocate an intermediate tensor to hold results if needed array intermediate({n_rows}, out_type, nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); - d.add_temporary(intermediate, s.index); + compute_encoder.add_temporary(intermediate); // 1st pass size_t row_size = (in_size + n_rows - 1) / n_rows; @@ -652,7 +670,7 @@ void strided_reduce_longcolumn( intermediate_shape.end(), out.shape().begin(), out.shape().end()); array intermediate(std::move(intermediate_shape), out_type, nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); - d.add_temporary(intermediate, s.index); + compute_encoder.add_temporary(intermediate); // Prepare the arguments for the kernel args.reduce_shape.push_back(args.reduction_size); @@ -823,7 +841,7 @@ void strided_reduce_2pass( intermediate_shape.end(), out.shape().begin(), out.shape().end()); array intermediate(std::move(intermediate_shape), out_type, nullptr, {}); intermediate.set_data(allocator::malloc(intermediate.nbytes())); - d.add_temporary(intermediate, s.index); + compute_encoder.add_temporary(intermediate); // Prepare the arguments for the kernel args.reduce_shape.push_back(args.reduction_size); @@ -952,9 +970,17 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { assert(inputs.size() == 1); array in = inputs[0]; - // Make sure no identity reductions trickle down here assert(!axes_.empty()); - assert(out.size() != in.size()); + + // When all the reduced axes have size 1 at runtime, which can happen with + // shapeless compilation, the reduction is the identity so just cast-copy + // the input to the output. + if (in.size() > 0 && out.size() == in.size()) { + CopyType ctype = + in.flags().contiguous ? CopyType::Vector : CopyType::General; + copy_gpu(in, out, ctype, stream()); + return; + } // Continue with reduction operation // Minimum of 4 bytes since we use size 4 structs for all reduce @@ -986,7 +1012,7 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { // Initialize output auto& s = stream(); auto& d = metal::device(s.device); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); // Reduce if (in.size() > 0) { @@ -1000,7 +1026,7 @@ void Reduce::eval_gpu(const std::vector& inputs, array& out) { // stride. if (plan.type == GeneralReduce) { array in_copy = contiguous_copy_gpu(in, s); - d.add_temporary(in_copy, s.index); + compute_encoder.add_temporary(in_copy); in = in_copy; plan = get_reduction_plan(in, axes_); } diff --git a/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp b/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp index 798824c2f..243699363 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/resident.cpp @@ -1,19 +1,40 @@ // Copyright © 2024 Apple Inc. +#include +#include +#include +#include +#include + +#include "mlx/backend/metal/device.h" #include "mlx/backend/metal/resident.h" +#include "mlx/utils.h" namespace mlx::core::metal { -ResidencySet::ResidencySet(MTL::Device* d) { +ResidencySets::ResidencySets(MTL::Device* d) { if (!d->supportsFamily(MTL::GPUFamilyMetal3)) { - return; - } else if (__builtin_available(macOS 15, iOS 18, *)) { - auto pool = new_scoped_memory_pool(); - auto desc = MTL::ResidencySetDescriptor::alloc()->init(); - NS::Error* error; - wired_set_ = d->newResidencySet(desc, &error); - desc->release(); - if (!wired_set_) { + return; // enabled_ stays false, so everything below is a no-op + } + if (__builtin_available(macOS 15, iOS 18, *)) { + device_ = d; + enabled_ = true; + int pct = env::residency_set_max_pct(); + if (pct <= 0 || pct >= 100) { + max_bytes_per_set_ = 0; // a single set holds everything + } else { + size_t ws = static_cast(d->recommendedMaxWorkingSetSize()); + // Floor the cap so a small working set or a small percentage cannot ask + // for an absurd number of sets. 64 MiB comfortably exceeds the + // allocator's heap. + max_bytes_per_set_ = std::max( + (ws / 100) * static_cast(pct), size_t(64) << 20); + } + std::lock_guard lk(mtx_); + // Set 0 always exists and is the fallback when a later set cannot be + // made, so failing to create it is fatal. + NS::Error* error = nullptr; + if (!add_set_locked(&error)) { std::ostringstream msg; msg << "[metal::Device] Unable to construct residency set.\n"; if (error) { @@ -21,80 +42,198 @@ ResidencySet::ResidencySet(MTL::Device* d) { } throw std::runtime_error(msg.str()); } - wired_set_->requestResidency(); } } -void ResidencySet::insert(MTL::Allocation* buf) { - if (!wired_set_) { - return; +ResidencySets::~ResidencySets() = default; + +bool ResidencySets::add_set_locked(NS::Error** error_out) { + NS::SharedPtr set; + if (__builtin_available(macOS 15, iOS 18, *)) { + auto pool = new_scoped_memory_pool(); + auto desc = MTL::ResidencySetDescriptor::alloc()->init()->autorelease(); + NS::Error* error = nullptr; + set = NS::TransferPtr(device_->newResidencySet(desc, &error)); + if (set) { + // A standing request, so allocations added to this set later are + // covered without requesting residency again on every insert. + set->requestResidency(); + } else if (error_out) { + *error_out = error; + } } - if (wired_set_->allocatedSize() + buf->allocatedSize() <= capacity_) { - wired_set_->addAllocation(buf); - wired_set_->commit(); - } else { - unwired_set_.insert(buf); + if (!set) { + return false; + } + sets_.push_back(Set{std::move(set), 0}); + num_sets_.store(sets_.size(), std::memory_order_release); + if (env::residency_debug()) { + fprintf( + stderr, + "[residency] created residency set %zu (max_bytes_per_set=%zu MB)\n", + sets_.size() - 1, + max_bytes_per_set_ >> 20); + } + return true; +} + +uint32_t ResidencySets::choose_set_locked(size_t bytes) { + if (max_bytes_per_set_ == 0) { + return 0; + } + const uint32_t none = static_cast(sets_.size()); + uint32_t empty = none; + uint32_t emptiest = 0; + for (uint32_t i = 0; i < sets_.size(); ++i) { + const size_t size = sets_[i].size; + if (size + bytes <= max_bytes_per_set_) { + return i; // first fit + } + // Only reached by an allocation larger than the cap, which needs a set of + // its own. Reusing an emptied one keeps the set count from growing. + if (i != 0 && size == 0 && empty == none) { + empty = i; + } + if (size < sets_[emptiest].size) { + emptiest = i; + } } + if (empty != none) { + return empty; + } + if (sets_.size() < kMaxSets && add_set_locked()) { + return static_cast(sets_.size() - 1); + } + // At the cap, or the driver would not make another set. Filling the emptiest + // set keeps them evenly sized, so each one still holds a bounded fraction of + // the wired bytes. + return emptiest; +} + +uint32_t ResidencySets::add_to_set_locked( + const MTL::Allocation* buf, + Placement& at) { + uint32_t idx = choose_set_locked(at.bytes); + auto& s = sets_[idx]; + s.set->addAllocation(buf); + s.size += at.bytes; + total_wired_ += at.bytes; + at.set_id = idx; + return idx; +} + +void ResidencySets::remove_from_set_locked( + const MTL::Allocation* buf, + Placement& at) { + auto& s = sets_[at.set_id]; + s.set->removeAllocation(buf); + // Subtract the size recorded at insert; allocatedSize() is never read back. + s.size -= at.bytes; + total_wired_ -= at.bytes; + at.set_id = kNoSet; } -void ResidencySet::erase(MTL::Allocation* buf) { - if (!wired_set_) { +void ResidencySets::insert(MTL::Allocation* buf) { + if (!enabled_) { return; } - if (auto it = unwired_set_.find(buf); it != unwired_set_.end()) { - unwired_set_.erase(it); - } else { - wired_set_->removeAllocation(buf); - wired_set_->commit(); + const size_t bytes = buf->allocatedSize(); + std::lock_guard lk(mtx_); + + auto [it, inserted] = buf_to_set_.try_emplace(buf, Placement{kNoSet, bytes}); + if (!inserted) { + assert(false && "allocation is already tracked"); + return; } + // Stay within the wired limit. The excess is tracked but left out of any + // set, and is added to one by resize() if the limit is raised later. + if (total_wired_ + bytes > capacity_) { + return; + } + uint32_t idx = add_to_set_locked(buf, it->second); + sets_[idx].set->commit(); } -void ResidencySet::resize(size_t size) { - if (!wired_set_) { +void ResidencySets::erase(MTL::Allocation* buf) { + if (!enabled_) { + return; + } + std::lock_guard lk(mtx_); + auto it = buf_to_set_.find(buf); + if (it == buf_to_set_.end()) { + assert(false && "erasing an allocation that was never inserted"); return; } + if (it->second.set_id != kNoSet) { + const uint32_t idx = it->second.set_id; + remove_from_set_locked(buf, it->second); + sets_[idx].set->commit(); + } + buf_to_set_.erase(it); +} +void ResidencySets::resize(size_t size) { + if (!enabled_) { + return; + } + std::lock_guard lk(mtx_); if (capacity_ == size) { return; } capacity_ = size; - size_t current_size = wired_set_->allocatedSize(); - - if (current_size < size) { - auto pool = new_scoped_memory_pool(); - // Add unwired allocations to the set - for (auto it = unwired_set_.begin(); it != unwired_set_.end();) { - auto buf_size = (*it)->allocatedSize(); - if (current_size + buf_size > size) { - it++; - } else { - current_size += buf_size; - wired_set_->addAllocation(*it); - unwired_set_.erase(it++); + auto pool = new_scoped_memory_pool(); + std::vector touched(sets_.size(), false); + // The loops below only mutate map values, never insert or erase, so the + // iterators stay valid across the whole walk. + if (total_wired_ < capacity_) { + // The budget grew: add allocations that now fit. + for (auto& [buf, at] : buf_to_set_) { + if (at.set_id != kNoSet || total_wired_ + at.bytes > capacity_) { + continue; + } + uint32_t idx = add_to_set_locked(buf, at); + if (idx >= touched.size()) { + touched.resize(sets_.size(), false); // a new set was made } + touched[idx] = true; } - wired_set_->commit(); - } else if (current_size > size) { - auto pool = new_scoped_memory_pool(); - // Remove wired allocations until under capacity - auto allocations = wired_set_->allAllocations(); - auto num_allocations = wired_set_->allocationCount(); - for (int i = 0; i < num_allocations && current_size > size; ++i) { - auto buf = static_cast(allocations->object(i)); - wired_set_->removeAllocation(buf); - current_size -= buf->allocatedSize(); - unwired_set_.insert(buf); + } else { + // The budget shrank: remove allocations until we are back under it. + for (auto& [buf, at] : buf_to_set_) { + if (total_wired_ <= capacity_) { + break; + } + if (at.set_id == kNoSet) { + continue; + } + touched[at.set_id] = true; + remove_from_set_locked(buf, at); + } + } + for (size_t i = 0; i < touched.size(); ++i) { + if (touched[i]) { + sets_[i].set->commit(); } - wired_set_->commit(); } } -ResidencySet::~ResidencySet() { - if (wired_set_) { - auto pool = new_scoped_memory_pool(); - wired_set_->release(); +void ResidencySets::attach_new_sets(MTL::CommandQueue* q, uint64_t& attached) { + // Lock-free when there is nothing new, which is the common case. It also + // covers the disabled case, where the count stays 0. + if (num_sets_.load(std::memory_order_acquire) == attached) { + return; + } + // Attach the new sets and record how far we got under a single lock hold, + // so a set created in between cannot be missed. + std::lock_guard lk(mtx_); + std::vector sets; + sets.reserve(sets_.size() - attached); + for (uint64_t id = attached; id < sets_.size(); ++id) { + sets.push_back(sets_[id].set.get()); } + q->addResidencySets(sets.data(), sets.size()); + attached = sets_.size(); } } // namespace mlx::core::metal diff --git a/Source/Cmlx/mlx/mlx/backend/metal/resident.h b/Source/Cmlx/mlx/mlx/backend/metal/resident.h index 5db558286..23960ad13 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/resident.h +++ b/Source/Cmlx/mlx/mlx/backend/metal/resident.h @@ -2,31 +2,122 @@ #pragma once -#include "mlx/backend/metal/device.h" +#include +#include +#include +#include +#include + +#include namespace mlx::core::metal { -class ResidencySet { +// 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: - ResidencySet(MTL::Device* d); - ~ResidencySet(); - - ResidencySet(const ResidencySet&) = delete; - ResidencySet& operator=(const ResidencySet&) = delete; + ResidencySets(MTL::Device* d); + ~ResidencySets(); - const MTL::ResidencySet* mtl_residency_set() { - return wired_set_; - } + 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: - MTL::ResidencySet* wired_set_{nullptr}; - std::unordered_set unwired_set_; + // 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 diff --git a/Source/Cmlx/mlx/mlx/backend/metal/rope.cpp b/Source/Cmlx/mlx/mlx/backend/metal/rope.cpp index ca0a66221..1147361ba 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/rope.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/rope.cpp @@ -118,7 +118,7 @@ void RoPE::eval_gpu( {&head_seq_transpose, MTL::DataType::DataTypeBool, 3}}; auto kernel = d.get_kernel(kname, hash_name, func_consts); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); float base = std::log2(base_); compute_encoder.set_compute_pipeline_state(kernel); @@ -133,8 +133,9 @@ void RoPE::eval_gpu( if (single) { compute_encoder.set_bytes(out_strides, 1, 4); uint32_t dim0 = dims_ / 2; - group_dims = get_block_dims(dim0, N, 1); - grid_dims = MTL::Size(dim0, N, 1); + uint32_t dim1 = B * N; + group_dims = get_block_dims(dim0, dim1, 1); + grid_dims = MTL::Size(dim0, dim1, 1); } else { compute_encoder.set_bytes(strides, 3, 4); compute_encoder.set_bytes(out_strides, 3, 5); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp b/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp index c990cda37..981d2b181 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/scaled_dot_product_attention.cpp @@ -28,13 +28,14 @@ void sdpa_full_self_attention_nax( const std::optional& sinks) { using namespace mlx::steel; - int wm = 4; - int wn = 1; - int bd = q.shape(-1); int bq = 64; int bk = 32; + bool split_d = bd == 256; + int wm = 4; + int wn = split_d ? 2 : 1; + int B = q.shape(0); int H = q.shape(1); int D = q.shape(3); @@ -43,6 +44,36 @@ void sdpa_full_self_attention_nax( int qL = q.shape(2); int kL = k.shape(2); + // The causal offset describes the true diagonal even if kL is widened + // below. + int qL_off = kL - qL; + + // Check if K/V are from chunked KV cache, and assume aligned K/V if so. + auto has_backing_rows = [](const array& kv, int rows) { + auto& st = kv.strides(); + if ((st[0] < 0) || (st[1] <= 0) || (st[2] <= 0) || (st[1] % st[2] != 0)) { + return false; + } + int64_t itemsize = kv.itemsize(); + // The rows must stay inside the head's row pitch (so they belong to the + // cache the slice was taken from) ... + int64_t pitch = st[1] / st[2]; + int64_t row0 = ((kv.offset() / itemsize) % st[1]) / st[2]; + if (row0 + rows > pitch) { + return false; + } + // ... and inside the buffer. + int64_t end = (kv.shape(0) - 1) * st[0] + (kv.shape(1) - 1) * st[1] + + (rows - 1) * st[2] + kv.shape(3); + return kv.offset() + end * itemsize <= int64_t(kv.buffer_size()); + }; + if (split_d && do_causal_ && !mask.has_value() && (kL % bk)) { + int kLp = bk * ((kL + bk - 1) / bk); + if (has_backing_rows(k, kLp) && has_backing_rows(v, kLp)) { + kL = kLp; + } + } + const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; const bool has_mask = mask.has_value(); @@ -59,7 +90,7 @@ void sdpa_full_self_attention_nax( std::string base_name; concatenate( base_name, - "steel_attention_", + split_d ? "steel_attention_dsplit_" : "steel_attention_", type_to_name(q), "_bq", bq, @@ -89,7 +120,7 @@ void sdpa_full_self_attention_nax( "_has_sinks_", (has_sinks ? 't' : 'n')); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_attention_nax_kernel( d, @@ -102,7 +133,8 @@ void sdpa_full_self_attention_nax( bd, wm, wn, - (has_mask ? *mask : q)); + (has_mask ? *mask : q), + split_d); compute_encoder.set_compute_pipeline_state(kernel); @@ -131,7 +163,7 @@ void sdpa_full_self_attention_nax( /* int qL_rem = */ (qL - NQ_aligned * bq), /* int kL_rem = */ (kL - NK_aligned * bk), - /* int qL_off = */ (kL - qL), + /* int qL_off = */ qL_off, /* int64_t Q_strides[3] = */ {q.strides(0), q.strides(1), q.strides(2)}, /* int64_t K_strides[3] = */ {k.strides(0), k.strides(1), k.strides(2)}, @@ -160,6 +192,7 @@ void sdpa_full_self_attention_nax( MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -174,7 +207,16 @@ void sdpa_full_self_attention_metal( bool do_causal_, const std::optional& mask, const std::optional& sinks) { - if (metal::is_nax_available() && q.shape(3) != 80 && + int B = q.shape(0); + int H = q.shape(1); + int D = q.shape(3); + int gqa_factor = q.shape(1) / k.shape(1); + + int qL = q.shape(2); + int kL = k.shape(2); + + if (metal::is_nax_available() && + (D == 64 || D == 96 || D == 128 || D == 256) && (env::enable_tf32() || q.dtype() != float32)) { return sdpa_full_self_attention_nax( /* const Stream& s = */ s, @@ -198,14 +240,6 @@ void sdpa_full_self_attention_metal( int bq = 32; int bk = bd < 128 ? 32 : 16; - int B = q.shape(0); - int H = q.shape(1); - int D = q.shape(3); - int gqa_factor = q.shape(1) / k.shape(1); - - int qL = q.shape(2); - int kL = k.shape(2); - const bool align_Q = (qL % bq) == 0; const bool align_K = (kL % bk) == 0; const bool has_mask = mask.has_value(); @@ -252,7 +286,7 @@ void sdpa_full_self_attention_metal( "_has_sinks_", (has_sinks ? 't' : 'n')); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = get_steel_attention_kernel( d, @@ -323,6 +357,7 @@ void sdpa_full_self_attention_metal( MTL::Size grid_dims = MTL::Size(NQ, H, B); MTL::Size group_dims = MTL::Size(32, wm, wn); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } @@ -378,8 +413,9 @@ void sdpa_vector( hash_name += has_sinks ? "_sinks" : "_nosinks"; // Get the kernel - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname, hash_name, func_consts); + check_kernel_threadgroup_size(kernel, group_dims, hash_name); compute_encoder.set_compute_pipeline_state(kernel); // Set its arguments @@ -429,7 +465,13 @@ void sdpa_vector_2pass( // Set the kernel name std::string kname; kname.reserve(64); - kname += "sdpa_vector_2pass_1_"; + kname += "sdpa_vector_2pass_1"; + if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) && + q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) && + k.shape(2) >= 8192) { + kname += "_gqa"; + } + kname += "_"; kname += get_type_string(q.dtype()); kname += "_"; kname += std::to_string(q.shape(-1)); @@ -474,6 +516,11 @@ void sdpa_vector_2pass( blocks = 32; } } + if (int blocks_env = env::get_var("MLX_SDPA_BLOCKS", 0); blocks_env > 0) { + // The 2-pass reduction consumes the partials in simd-width (32) chunks + // and silently drops the tail otherwise, so round up to a multiple of 32. + blocks = ((blocks_env + 31) / 32) * 32; + } size_t k_head_stride = k.shape(1) == 1 ? k.strides(0) : k.strides(1); size_t k_seq_stride = k.strides()[2]; size_t v_head_stride = v.shape(1) == 1 ? v.strides(0) : v.strides(1); @@ -495,9 +542,10 @@ void sdpa_vector_2pass( intermediate.set_data(allocator::malloc(intermediate.nbytes())); sums.set_data(allocator::malloc(sums.nbytes())); maxs.set_data(allocator::malloc(maxs.nbytes())); - d.add_temporary(intermediate, s.index); - d.add_temporary(sums, s.index); - d.add_temporary(maxs, s.index); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.add_temporary(intermediate); + compute_encoder.add_temporary(sums); + compute_encoder.add_temporary(maxs); bool has_mask = mask.has_value(); bool bool_mask = has_mask && (*mask).dtype() == bool_; @@ -521,7 +569,6 @@ void sdpa_vector_2pass( hash_name += std::to_string(blocks); // Get the kernel - auto& compute_encoder = d.get_command_encoder(s.index); auto kernel = d.get_kernel(kname, hash_name, func_consts); check_kernel_threadgroup_size(kernel, group_dims, hash_name); @@ -583,28 +630,23 @@ void sdpa_vector_2pass( compute_encoder.dispatch_threadgroups(grid_dims, group_dims); } -} // namespace - -bool ScaledDotProductAttention::use_fallback( +std::tuple has_fused_kernel( const array& q, const array& k, const array& v, bool has_mask, bool has_arr_mask, bool do_causal, - bool is_training, bool output_logsumexp, Stream s) { - if (is_training) { - // It's faster for training on Metal to use the unfused SDPA for both - // forward and backward. - return true; + if (s.device != Device::gpu) { + return {false, "the fused kernels require a GPU (Metal) stream."}; } if (output_logsumexp) { - return true; - } - if (s.device == Device::cpu) { - return true; + return { + false, + "the fused forward does not produce the logsumexp required for " + "the fused VJP; use default routing when training."}; } const int value_head_dim = v.shape(-1); @@ -615,25 +657,113 @@ bool ScaledDotProductAttention::use_fallback( const int num_kv_heads = k.shape(1); const int gqa_factor = num_query_heads / num_kv_heads; - const bool sdpa_vector_supported_head_dim = - query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 96 || query_head_dim == 128 || - query_head_dim == 256); - const bool sdpa_full_supported_head_dim = query_head_dim == value_head_dim && - (query_head_dim == 64 || query_head_dim == 80 || query_head_dim == 128); + std::ostringstream msg; + if (query_sequence_length > 8) { + const bool supported_head_dim = query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 72 || query_head_dim == 80 || + query_head_dim == 96 || query_head_dim == 128 || + query_head_dim == 192 || query_head_dim == 256); + if (!supported_head_dim) { + msg << "the full attention kernel supports head dims " + << "{64, 72, 80, 96, 128, 192, 256} with matching query/value head " + << "dims; got query head dim " << query_head_dim + << " and value head dim " << value_head_dim << "."; + return {false, msg.str()}; + } + if (has_mask && !has_arr_mask && + !(query_sequence_length <= key_sequence_length && do_causal)) { + msg << "the full attention kernel with a causal mask requires the " + << "query sequence to be no longer than the key sequence; got " + << "query length " << query_sequence_length << " and key length " + << key_sequence_length << "."; + return {false, msg.str()}; + } + } else { + const bool supported_head_dim = + (query_head_dim == value_head_dim && + (query_head_dim == 64 || query_head_dim == 96 || + query_head_dim == 128 || query_head_dim == 192 || + query_head_dim == 256)) || + (query_head_dim == 192 && value_head_dim == 128); + if (!supported_head_dim) { + msg << "the vector attention kernel supports head dims " + << "{64, 96, 128, 192, 256} with matching query/value head dims, " + << "or query head dim 192 with value head dim 128; got query head " + << "dim " << query_head_dim << " and value head dim " + << value_head_dim << "."; + return {false, msg.str()}; + } + if (query_sequence_length > key_sequence_length) { + msg << "the vector attention kernel requires the query sequence to be " + << "no longer than the key sequence; got query length " + << query_sequence_length << " and key length " << key_sequence_length + << "."; + return {false, msg.str()}; + } + if (query_sequence_length * gqa_factor > 32) { + msg << "the vector attention kernel requires the query length times " + << "the GQA factor to be at most 32; got query length " + << query_sequence_length << " and GQA factor " << gqa_factor << "."; + return {false, msg.str()}; + } + } + return {true, ""}; +} - const bool sdpa_full_supported_mask = !has_mask || has_arr_mask || - (query_sequence_length <= key_sequence_length && do_causal); +} // namespace - const bool supports_sdpa_full = query_sequence_length > 8 && - sdpa_full_supported_mask && sdpa_full_supported_head_dim; +bool ScaledDotProductAttention::use_fallback( + const array& q, + const array& k, + const array& v, + bool has_mask, + bool has_arr_mask, + bool do_causal, + bool is_training, + bool output_logsumexp, + bool force_fused, + Stream s) { + auto [has_fused, reason] = has_fused_kernel( + q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, s); + if (force_fused) { + if (!has_fused) { + std::ostringstream msg; + msg << "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available: " + << reason; + throw std::invalid_argument(msg.str()); + } + return false; + } - const bool supports_sdpa_vector = (query_sequence_length <= 8) && - (query_sequence_length <= key_sequence_length) && - sdpa_vector_supported_head_dim && - (query_sequence_length * gqa_factor) <= 32; + if (is_training) { + // It's faster for training on Metal to use the unfused SDPA for both + // forward and backward. + return true; + } + if (!has_fused) { + return true; + } + + const int query_sequence_length = q.shape(2); + const int query_head_dim = q.shape(-1); + const int value_head_dim = v.shape(-1); - return !(supports_sdpa_full || supports_sdpa_vector); + // Use headdim-split kernel when NAX is enabled and there are enough query + // blocks to fill the machine. + if (metal::is_nax_available() && + (env::enable_tf32() || q.dtype() != float32) && + query_sequence_length >= 1024 && query_head_dim == 256 && do_causal && + !has_arr_mask) { + return false; + } + + // Unfused path is faster for following shapes. + if (query_sequence_length > 8) { + return query_head_dim == 192 || query_head_dim == 256; + } else { + return query_head_dim == value_head_dim && query_head_dim == 192; + } } bool ScaledDotProductAttention::supports_bool_mask() { @@ -711,7 +841,8 @@ void ScaledDotProductAttention::eval_gpu( return (strides[0] == strides[1] * shape[1]); }; - const auto& q = copy_unless(q_copy_unless, q_pre); + bool q_copied = !q_copy_unless(q_pre); + array q = (q_copied) ? contiguous_copy_gpu(q_pre, s) : q_pre; const auto& k = copy_unless(kv_copy_unless, k_pre); const auto& v = copy_unless(kv_copy_unless, v_pre); @@ -719,6 +850,9 @@ void ScaledDotProductAttention::eval_gpu( if (q.is_donatable() && q.flags().row_contiguous && q.size() == o.size()) { o.copy_shared_buffer(q); } else { + if (q_copied) { + copies.push_back(q); + } o.set_data(allocator::malloc(o.nbytes())); } @@ -778,7 +912,7 @@ void ScaledDotProductAttention::eval_gpu( s, d, q, k, v, scale_, o, do_causal_, mask, sinks); } - d.add_temporaries(std::move(copies), s.index); + metal::get_command_encoder(s).add_temporaries(std::move(copies)); } bool ScaledDotProductAttentionVJP::use_fallback(const array& q, Stream s) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp index b35de64ce..18935191f 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/scan.cpp @@ -4,16 +4,16 @@ #include #include "mlx/backend/gpu/copy.h" +#include "mlx/backend/gpu/scan.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" -#include "mlx/backend/metal/scan.h" #include "mlx/backend/metal/utils.h" #include "mlx/primitives.h" namespace mlx::core { void scan_gpu_inplace( - array in, + const array& in, array& out, Scan::ReduceType reduce_type, int axis, @@ -59,14 +59,14 @@ void scan_gpu_inplace( auto kernel = get_scan_kernel(d, kname, reverse, inclusive, reduce_type_str, in, out); - if (contiguous) { - auto& compute_encoder = d.get_command_encoder(s.index); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(out, 1); - size_t size = in.shape(axis); - compute_encoder.set_bytes(size, 2); + auto& compute_encoder = metal::get_command_encoder(s); + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(out, 1); + size_t size = in.shape(axis); + compute_encoder.set_bytes(size, 2); + if (contiguous) { // Compute the thread grid int n_reads = (in.itemsize() <= 4) ? 4 : 2; constexpr int simd_size = 32; @@ -89,16 +89,9 @@ void scan_gpu_inplace( MTL::Size group_dims(thread_group_size, 1, 1); compute_encoder.dispatch_threads(grid_dims, group_dims); } else { - auto& compute_encoder = d.get_command_encoder(s.index); - compute_encoder.set_compute_pipeline_state(kernel); - compute_encoder.set_input_array( - in.data_shared_ptr() == nullptr ? out : in, 0); - compute_encoder.set_output_array(out, 1); - size_t size = in.shape(axis); size_t stride = in.strides()[axis]; int bn = 32; size_t stride_blocks = (stride + bn - 1) / bn; - compute_encoder.set_bytes(size, 2); compute_encoder.set_bytes(stride, 3); compute_encoder.set_bytes(stride_blocks, 4); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/scan.h b/Source/Cmlx/mlx/mlx/backend/metal/scan.h deleted file mode 100644 index dab79c50b..000000000 --- a/Source/Cmlx/mlx/mlx/backend/metal/scan.h +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "mlx/array.h" -#include "mlx/primitives.h" - -namespace mlx::core { - -void scan_gpu_inplace( - array in, - array& out, - Scan::ReduceType reduce_type, - int axis, - bool reverse, - bool inclusive, - const Stream& s); - -} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp b/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp index 1e14c35c8..f6ff41abc 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/slicing.cpp @@ -31,7 +31,7 @@ void concatenate_gpu( flags.col_contiguous = false; flags.contiguous = false; auto& d = metal::device(s.device); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); auto concurrent_ctx = compute_encoder.start_concurrent(); for (int i = 0; i < inputs.size(); i++) { array out_slice(inputs[i].shape(), out.dtype(), nullptr, {}); @@ -48,6 +48,7 @@ array compute_dynamic_offset( const std::vector& axes, const Stream& s) { auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); // Kernel to compute offset here. array offset({1}, int64, nullptr, {}); @@ -57,8 +58,8 @@ array compute_dynamic_offset( offset.copy_shared_buffer(indices); } else { offset.set_data(allocator::malloc(offset.itemsize())); + compute_encoder.add_temporary(offset); } - d.add_temporary(offset, s.index); auto dtype = indices.dtype(); std::string lib_name = "compute_dynamic_offset_" + type_to_name(dtype); @@ -83,7 +84,6 @@ array compute_dynamic_offset( }); auto kernel = d.get_kernel(lib_name, lib); - auto& compute_encoder = d.get_command_encoder(s.index); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(indices, 0); compute_encoder.set_output_array(offset, 1); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/softmax.cpp b/Source/Cmlx/mlx/mlx/backend/metal/softmax.cpp index 0b1a1848d..f455cf2d1 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/softmax.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/softmax.cpp @@ -58,7 +58,7 @@ void Softmax::eval_gpu(const std::vector& inputs, array& out) { kernel_name += type_to_name(out); auto kernel = get_softmax_kernel(d, kernel_name, precise_, out); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); { MTL::Size grid_dims, group_dims; if (axis_size <= looped_limit) { diff --git a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp index 3c84022f2..65f144c02 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/sort.cpp @@ -63,7 +63,7 @@ void single_block_sort( auto kernel = get_sort_kernel(d, kname.str(), in, out, bn, tn); // Prepare command encoder - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); // Set inputs @@ -160,7 +160,7 @@ void multi_block_sort( dev_vals_0, dev_vals_1, dev_idxs_0, dev_idxs_1, block_partitions}; // Prepare command encoder - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); // Do blockwise sort { @@ -268,7 +268,7 @@ void multi_block_sort( (axis == in.ndim() - 1) ? CopyType::Vector : CopyType::General, s); - d.add_temporaries(std::move(copies), s.index); + compute_encoder.add_temporaries(std::move(copies)); } void gpu_merge_sort( @@ -365,4 +365,51 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { gpu_merge_sort(s, d, in, out, axis_, false); } +void SearchSorted::eval_gpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 2); + auto& a = inputs[0]; + auto v = inputs[1]; + + out.set_data(allocator::malloc(out.nbytes())); + if (out.size() == 0) { + return; + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + if (a.size() == 0) { + array zero = array(0, out.dtype()); + fill_gpu(zero, out, s); + compute_encoder.add_temporary(std::move(zero)); + return; + } + + if (!v.flags().row_contiguous) { + v = contiguous_copy_gpu(v, s); + compute_encoder.add_temporary(v); + } + + int64_t a_stride = a.strides()[0]; // sequence is 1D + auto n = static_cast(a.size()); + + std::string kernel_name = "searchsorted_"; + concatenate(kernel_name, type_to_name(a), right_ ? "_right" : "_left"); + auto kernel = get_searchsorted_kernel(d, kernel_name, a, right_); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(a, 0); + compute_encoder.set_input_array(v, 1); + compute_encoder.set_output_array(out, 2); + compute_encoder.set_bytes(n, 3); + compute_encoder.set_bytes(a_stride, 4); + + size_t thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + thread_group_size = std::min(thread_group_size, out.size()); + MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1); + MTL::Size grid_dims = get_2d_grid_dims(out.shape(), out.strides()); + compute_encoder.dispatch_threads(grid_dims, group_dims); +} + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/metal/ternary.cpp b/Source/Cmlx/mlx/mlx/backend/metal/ternary.cpp index 252815aae..9bf07f31b 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/ternary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/ternary.cpp @@ -78,7 +78,7 @@ void ternary_op_gpu_inplace( auto kernel = get_ternary_kernel(d, kernel_name, out.dtype(), op); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(a, 0); compute_encoder.set_input_array(b, 1); diff --git a/Source/Cmlx/mlx/mlx/backend/metal/unary.cpp b/Source/Cmlx/mlx/mlx/backend/metal/unary.cpp index 833b23f63..354f199c6 100644 --- a/Source/Cmlx/mlx/mlx/backend/metal/unary.cpp +++ b/Source/Cmlx/mlx/mlx/backend/metal/unary.cpp @@ -57,7 +57,7 @@ void unary_op_gpu_inplace( auto kernel = get_unary_kernel(d, kernel_name, in.dtype(), out.dtype(), op); auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); - auto& compute_encoder = d.get_command_encoder(s.index); + auto& compute_encoder = metal::get_command_encoder(s); compute_encoder.set_compute_pipeline_state(kernel); compute_encoder.set_input_array(in, 0); compute_encoder.set_output_array(out, 1); diff --git a/Source/Cmlx/mlx/mlx/backend/no_cpu/compiled.cpp b/Source/Cmlx/mlx/mlx/backend/no_cpu/compiled.cpp index 2eeddab47..09987e7d5 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_cpu/compiled.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_cpu/compiled.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include "mlx/compile_impl.h" #include "mlx/primitives.h" @@ -12,6 +12,7 @@ namespace detail { bool compile_available_for_device(const Device& device) { return device == Device::gpu; } + } // namespace detail void Compiled::eval_cpu( diff --git a/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp index ae51dd9b2..e52230715 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_cpu/primitives.cpp @@ -71,6 +71,7 @@ NO_CPU(Gather) NO_CPU(GatherAxis) NO_CPU(GatherMM) NO_CPU(GatherQMM) +NO_CPU(GatherQQMM) NO_CPU(Greater) NO_CPU(GreaterEqual) NO_CPU(Hadamard) @@ -106,6 +107,7 @@ NO_CPU(Round) NO_CPU(Scan) NO_CPU(Scatter) NO_CPU(ScatterAxis) +NO_CPU(SearchSorted) NO_CPU(Select) NO_CPU(SegmentedMM) NO_CPU(Sigmoid) diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/allocator.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/allocator.cpp index abb83e50e..a800e381a 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/allocator.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/allocator.cpp @@ -1,10 +1,12 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include #include #include "mlx/allocator.h" +#include "mlx/backend/common/buffer_cache.h" #include "mlx/memory.h" +#include "mlx/stream.h" #ifdef __APPLE__ #include "mlx/backend/no_gpu/apple_memory.h" @@ -18,10 +20,22 @@ size_t get_memory_size() { namespace mlx::core { +namespace { + +void synchronize_cpu_streams() { + for (const auto& s : get_streams()) { + if (s.device == Device::cpu) { + synchronize(s); + } + } +} + +} // namespace + namespace allocator { class CommonAllocator : public Allocator { - /** A general CPU allocator. */ + /** A general CPU allocator with buffer caching. */ public: virtual Buffer malloc(size_t size) override; virtual void free(Buffer buffer) override; @@ -46,42 +60,50 @@ class CommonAllocator : public Allocator { return limit; } + size_t get_cache_memory() const; + size_t set_cache_limit(size_t limit); + void clear_cache(); + private: + friend CommonAllocator& common_allocator(); + CommonAllocator(); + + static size_t get_buffer_size(void* buf); + size_t memory_limit_; + size_t cache_limit_; size_t active_memory_{0}; size_t peak_memory_{0}; - std::mutex mutex_; - CommonAllocator() : memory_limit_(0.8 * get_memory_size()) { - if (memory_limit_ == 0) { - memory_limit_ = 1UL << 33; - } - }; - - friend CommonAllocator& common_allocator(); + mutable std::mutex mutex_; + mutable BufferCache buffer_cache_; }; -CommonAllocator& common_allocator() { - static CommonAllocator allocator_; - return allocator_; -} - -Allocator& allocator() { - return common_allocator(); -} - -void* Buffer::raw_ptr() { - if (!ptr_) { - return nullptr; +CommonAllocator::CommonAllocator() + : memory_limit_(0.8 * get_memory_size()), + cache_limit_(32UL << 20), // 32 MB default cache limit + buffer_cache_(/* page_size */ 4096, get_buffer_size, std::free) { + if (memory_limit_ == 0) { + memory_limit_ = 1ULL << 33; } - return static_cast(ptr_) + 1; } Buffer CommonAllocator::malloc(size_t size) { + std::unique_lock lk(mutex_); + // Try cache first + void* cached = buffer_cache_.reuse_from_cache(size); + if (cached) { + active_memory_ += get_buffer_size(cached); + peak_memory_ = std::max(active_memory_, peak_memory_); + return Buffer{cached}; + } + lk.unlock(); + + // Cache miss: allocate from OS void* ptr = std::malloc(size + sizeof(size_t)); if (ptr != nullptr) { *static_cast(ptr) = size; } - std::unique_lock lk(mutex_); + lk.lock(); active_memory_ += size; peak_memory_ = std::max(active_memory_, peak_memory_); return Buffer{ptr}; @@ -89,16 +111,75 @@ Buffer CommonAllocator::malloc(size_t size) { void CommonAllocator::free(Buffer buffer) { auto sz = size(buffer); - std::free(buffer.ptr()); std::unique_lock lk(mutex_); active_memory_ -= sz; + + if (sz > 0 && buffer_cache_.cache_size() + sz <= cache_limit_) { + buffer_cache_.recycle_to_cache(buffer.ptr()); + } else { + lk.unlock(); + std::free(buffer.ptr()); + } } size_t CommonAllocator::size(Buffer buffer) const { - if (buffer.ptr() == nullptr) { + return get_buffer_size(buffer.ptr()); +} + +size_t CommonAllocator::get_cache_memory() const { + std::unique_lock lk(mutex_); + return buffer_cache_.cache_size(); +} + +size_t CommonAllocator::set_cache_limit(size_t limit) { + synchronize_cpu_streams(); + std::unique_lock lk(mutex_); + std::swap(cache_limit_, limit); + if (buffer_cache_.cache_size() > cache_limit_) { + buffer_cache_.release_cached_buffers( + buffer_cache_.cache_size() - cache_limit_); + } + return limit; +} + +void CommonAllocator::clear_cache() { + { + std::unique_lock lk(mutex_); + if (buffer_cache_.cache_size() == 0) { + return; + } + } + synchronize_cpu_streams(); + std::unique_lock lk(mutex_); + buffer_cache_.clear(); +} + +size_t CommonAllocator::get_buffer_size(void* buf) { + if (buf) { + return *static_cast(buf); + } else { return 0; } - return *static_cast(buffer.ptr()); +} + +CommonAllocator& common_allocator() { + static CommonAllocator allocator_; + return allocator_; +} + +Allocator& allocator() { + return common_allocator(); +} + +void* Buffer::raw_ptr() { + if (!ptr_) { + return nullptr; + } + return static_cast(ptr_) + 1; +} + +bool can_reuse_alien_buffer(void*) { + return true; } } // namespace allocator @@ -119,16 +200,17 @@ size_t get_memory_limit() { return allocator::common_allocator().get_memory_limit(); } -// No-ops for common allocator size_t get_cache_memory() { - return 0; + return allocator::common_allocator().get_cache_memory(); } -size_t set_cache_limit(size_t) { - return 0; +size_t set_cache_limit(size_t limit) { + return allocator::common_allocator().set_cache_limit(limit); } size_t set_wired_limit(size_t) { return 0; } -void clear_cache() {} +void clear_cache() { + allocator::common_allocator().clear_cache(); +} } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/eval.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/eval.cpp index 5d88970a0..b08f6ec6f 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/eval.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/eval.cpp @@ -7,7 +7,17 @@ namespace mlx::core::gpu { -void new_stream(Stream) {} +void init() {} + +void new_stream(Stream) { + throw std::invalid_argument( + "[new_stream] Cannot make gpu stream without gpu backend."); +} + +void new_thread_unsafe_stream(Stream) { + throw std::invalid_argument( + "[new_thread_unsafe_stream] Cannot make gpu stream without gpu backend."); +} void eval(array&) { throw std::runtime_error("[gpu::eval] GPU backend is not available"); @@ -18,7 +28,9 @@ void finalize(Stream) { } void synchronize(Stream) { - throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); + throw std::runtime_error("[gpu::synchronize] GPU backend is not available"); } +void clear_streams() {} + } // namespace mlx::core::gpu diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/event.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/event.cpp index 6dde047ab..8966b7761 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/event.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/event.cpp @@ -12,42 +12,52 @@ struct EventCounter { uint64_t value{0}; std::mutex mtx; std::condition_variable cv; + std::atomic error; + + void wait(uint64_t val) { + std::unique_lock lk(mtx); + if (value >= val) { + return; + } + cv.wait(lk, [this, val] { return value >= val; }); + } }; Event::Event(Stream stream) : stream_(stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - event_ = std::shared_ptr(new EventCounter{}, dtor); + event_ = std::make_shared(); } void Event::wait() { - auto ec = static_cast(event_.get()); - std::unique_lock lk(ec->mtx); - if (ec->value >= value()) { - return; - } - ec->cv.wait(lk, [value = value(), ec] { return ec->value >= value; }); + check_error(); + cast().wait(value()); + check_error(); } void Event::wait(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { wait(); }); + scheduler::wait_event(stream, *this, [value = value()](Event& self) { + self.cast().wait(value); + }); } void Event::signal(Stream stream) { - scheduler::enqueue(stream, [*this]() mutable { - auto ec = static_cast(event_.get()); + scheduler::signal_event(stream, *this, [value = value()](Event& self) { + auto& ec = self.cast(); { - std::lock_guard lk(ec->mtx); - ec->value = value(); + std::lock_guard lk(ec.mtx); + ec.value = value; } - ec->cv.notify_all(); + ec.cv.notify_all(); }); } bool Event::is_signaled() const { - auto ec = static_cast(event_.get()); - { - std::lock_guard lk(ec->mtx); - return (ec->value >= value()); - } + auto& ec = cast(); + std::lock_guard lk(ec.mtx); + return ec.value >= value(); +} + +std::atomic& Event::error() { + return cast().error; } + } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp index cd66d23cf..05852c860 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/fence.cpp @@ -1,54 +1,30 @@ // Copyright © 2024 Apple Inc. -#include -#include - #include "mlx/fence.h" -#include "mlx/scheduler.h" +#include "mlx/event.h" namespace mlx::core { struct FenceImpl { - uint32_t count{0}; - uint32_t value{0}; - std::mutex mtx; - std::condition_variable cv; + uint32_t count; + Event event; + + FenceImpl(uint32_t count, Stream s) : count(count), event(s) {} }; -Fence::Fence(Stream) { - auto dtor = [](void* ptr) { delete static_cast(ptr); }; - fence_ = std::shared_ptr(new FenceImpl{}, dtor); +Fence::Fence(Stream s) { + fence_ = std::make_shared(0, s); } -void Fence::wait(Stream stream, const array&) { - auto& f = *static_cast(fence_.get()); - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - if (f.value >= count) { - return; - } - f.cv.wait(lk, [&f, count] { return f.value >= count; }); - }); - } else { - throw std::runtime_error("[Fence::wait] Invalid stream."); - } +void Fence::wait(Stream s, const array&) { + cast().event.wait(s); } -void Fence::update(Stream stream, const array&, bool) { - auto& f = *static_cast(fence_.get()); +void Fence::update(Stream s, const array&, bool) { + auto& f = cast(); f.count++; - if (stream.device == Device::cpu) { - scheduler::enqueue(stream, [count = f.count, fence_ = fence_]() mutable { - auto& f = *static_cast(fence_.get()); - std::unique_lock lk(f.mtx); - f.value = count; - f.cv.notify_all(); - }); - } else { - throw std::runtime_error("[Fence::update] Invalid stream."); - } + f.event.set_value(f.count); + f.event.signal(s); } } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp index 4819ed272..f17f12cfa 100644 --- a/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/backend/no_gpu/primitives.cpp @@ -32,20 +32,26 @@ bool fast::ScaledDotProductAttention::use_fallback( bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s) { + if (force_fused) { + throw std::invalid_argument( + "[scaled_dot_product_attention] force_fused=True but no fused " + "kernel is available in CPU backend."); + } return true; } -bool fast::ScaledDotProductAttention::supports_bool_mask() { - return false; -} - bool fast::ScaledDotProductAttentionVJP::use_fallback( const array& q, Stream s) { return true; } +bool fast::ScaledDotProductAttention::supports_bool_mask() { + return false; +} + NO_GPU(Abs) NO_GPU(Add) NO_GPU(AddMM) @@ -98,6 +104,7 @@ NO_GPU(Gather) NO_GPU(GatherAxis) NO_GPU(GatherMM) NO_GPU(GatherQMM) +NO_GPU(GatherQQMM) NO_GPU(Greater) NO_GPU(GreaterEqual) NO_GPU(Hadamard) @@ -133,6 +140,7 @@ NO_GPU(Round) NO_GPU(Scan) NO_GPU(Scatter) NO_GPU(ScatterAxis) +NO_GPU(SearchSorted) NO_GPU(Select) NO_GPU(SegmentedMM) NO_GPU(Sigmoid) @@ -162,6 +170,8 @@ NO_GPU(View) NO_GPU(MaskedScatter) namespace fast { +NO_GPU_USE_FALLBACK(CrossEntropy) +NO_GPU_MULTI(CrossEntropyVJP) NO_GPU_USE_FALLBACK(LayerNorm) NO_GPU_MULTI(LayerNormVJP) NO_GPU_USE_FALLBACK(RMSNorm) diff --git a/Source/Cmlx/mlx/mlx/compile.cpp b/Source/Cmlx/mlx/mlx/compile.cpp index 48708cd60..bb17c4496 100644 --- a/Source/Cmlx/mlx/mlx/compile.cpp +++ b/Source/Cmlx/mlx/mlx/compile.cpp @@ -1,6 +1,9 @@ // Copyright © 2023-2024 Apple Inc. + +#include #include #include +#include #include #include #include @@ -212,7 +215,7 @@ std::vector Compiled::output_shapes(const std::vector& inputs) { namespace detail { -CompileMode& compile_mode() { +std::atomic& compile_mode() { auto get_val = []() { if (std::getenv("MLX_DISABLE_COMPILE")) { return CompileMode::disabled; @@ -220,7 +223,7 @@ CompileMode& compile_mode() { return CompileMode::enabled; } }; - static CompileMode compile_mode_ = get_val(); + static std::atomic compile_mode_ = get_val(); return compile_mode_; } @@ -296,7 +299,7 @@ std::uintptr_t get_function_address(const std::function& fun) { return reinterpret_cast(*fun_ptr); } -class CompilerCache { +class CompileCache { public: struct CacheEntry { CacheEntry(Stream stream, bool shapeless) @@ -311,15 +314,38 @@ class CompilerCache { std::shared_ptr extra; }; - // Returns a reference to a CacheEntry which can be updated - // by the caller to avoid copying large tapes / inputs / outputs - CacheEntry& find( + CompileCache() { + // Make sure the allocator is fully initialized before the compiler cache. + allocator::allocator(); + } + + // Returns a reference to a CacheEntry which can be updated by the caller to + // avoid copying large tapes / inputs / outputs, with the shared_ptr of + // entries to avoid getting erased during compilation. + std::tuple>> find( std::uintptr_t fun_id, const std::vector& inputs, bool shapeless, const std::vector& constants) { - // Find the cache entries for |fun_id|. - std::vector& entries = cache_[fun_id]; + // Find the cache entries for |fun_id| in a thread-safe way. + auto entries_ptr = [&]() { + // Lookup with shared lock. + { + std::shared_lock lock(mutex_); + auto it = cache_.find(fun_id); + if (it != cache_.end()) { + return it->second; + } + } + // Insertion with exclusive lock. + std::unique_lock lock(mutex_); + auto& ptr = cache_[fun_id]; + if (!ptr) { + ptr = std::make_shared>(); + } + return ptr; + }(); + auto& entries = *entries_ptr; // Compare if 2 arrays have same shape and dtype. auto has_same_shape_and_dtype = [shapeless]( @@ -357,36 +383,37 @@ class CompilerCache { // Check the inputs match and return if so if (has_same_shape_and_dtype(inputs, entry.inputs) && constants == entry.constants) { - return entry; + return {entry, std::move(entries_ptr)}; } } // Otherwise append a new cache entry entries.push_back(CacheEntry{stream, shapeless}); - return entries.back(); + return {entries.back(), std::move(entries_ptr)}; } void erase(std::uintptr_t fun_id) { + std::unique_lock lock(mutex_); cache_.erase(fun_id); } void clear() { + std::unique_lock lock(mutex_); cache_.clear(); } private: - CompilerCache() { - // Make sure the allocator is fully - // initialized before the compiler cache - allocator::allocator(); - } - - friend CompilerCache& compiler_cache(); - std::unordered_map> cache_; + // The cache may get its key erased from a separate thread, but its value is + // only added and modified in the thread of creation. + // Put value in a shared_ptr to avoid race condition when erasing happened + // during compilation for the same function. + std::unordered_map>> + cache_; + std::shared_mutex mutex_; }; -CompilerCache& compiler_cache() { - static CompilerCache compiler_cache_; - return compiler_cache_; +std::shared_ptr& compile_cache_unsafe() { + static thread_local auto cache = std::make_shared(); + return cache; } std::tuple, std::vector, std::shared_ptr> @@ -857,7 +884,8 @@ void compile_fuse( // are not fusable except for broadcast which we can split to avoid // stopping fusion if (!all_parents_in) { - if (a.has_primitive() && is_broadcast(a.primitive())) { + if (a.has_primitive() && is_broadcast(a.primitive()) && + input_set.size() < max_compile_arrays) { array b = split_one(a, parents_map, cache); recurse(b, depth, s, shape); } else { @@ -1113,12 +1141,12 @@ ArrayFnWithExtra compile( } // Find a cache entry with the correct inputs - auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants); + auto [entry, entries_ptr] = + compile_cache_unsafe()->find(fun_id, inputs, shapeless, constants); + static_assert(std::is_reference_v); // No matching cache entry existed, so compile if (entry.empty) { - // Mark the entry as not empty since we are about to fill it - entry.empty = false; // Set the constants entry.constants = std::move(constants); // Trace to build the graph @@ -1133,16 +1161,22 @@ ArrayFnWithExtra compile( compile_dfs(entry.inputs, entry.outputs, inputs); // Simplify the tape - if (compile_mode() != CompileMode::no_simplify) { + auto mode = compile_mode().load(); + if (mode != CompileMode::no_simplify) { compile_simplify( entry.tape, parents_map, entry.outputs, /* passes */ 3); } // Kernel fusion to generate Compiled primitives. The tape and // new outputs must be updated accordingly - if (compile_mode() != CompileMode::no_fuse) { + if (mode != CompileMode::no_fuse) { compile_fuse(entry.tape, parents_map, entry.inputs, entry.outputs); } + + // Mark the entry as filled only after every step above completed, so + // a throwing first trace leaves the entry empty and a later call + // re-traces cleanly instead of hitting a half-filled entry + entry.empty = false; } // At this point we must have a tape, now replace the placeholders @@ -1181,12 +1215,20 @@ std::function(const std::vector&)> compile( }; } -void compile_erase(std::uintptr_t fun_id) { - detail::compiler_cache().erase(fun_id); +CompileCacheWeakPtr compile_cache() { + return compile_cache_unsafe(); } -void compile_clear_cache() { - detail::compiler_cache().clear(); +void compile_erase(const CompileCacheWeakPtr& cache, std::uintptr_t fun_id) { + if (auto p = cache.lock()) { + p->erase(fun_id); + } +} + +void compile_clear_cache(const CompileCacheWeakPtr& cache) { + if (auto p = cache.lock()) { + p->clear(); + } } } // namespace detail @@ -1206,8 +1248,8 @@ std::function(const std::vector&)> compile( auto pfun = std::shared_ptr< std::function(const std::vector&)>>( new std::function(const std::vector&)>{fun}, - [](auto* p) { - detail::compile_erase(reinterpret_cast(p)); + [cache = detail::compile_cache()](auto* p) { + detail::compile_erase(cache, reinterpret_cast(p)); delete p; }); fun_id = reinterpret_cast(pfun.get()); diff --git a/Source/Cmlx/mlx/mlx/compile_impl.h b/Source/Cmlx/mlx/mlx/compile_impl.h index 238a8b940..1afe07b93 100644 --- a/Source/Cmlx/mlx/mlx/compile_impl.h +++ b/Source/Cmlx/mlx/mlx/compile_impl.h @@ -27,12 +27,19 @@ MLX_API ArrayFnWithExtra compile( bool shapeless, std::vector constants); -// Erase cached compile functions -MLX_API void compile_erase(std::uintptr_t fun_id); +// Get the compiler cache of current thread. +class CompileCache; +using CompileCacheWeakPtr = std::weak_ptr; +MLX_API CompileCacheWeakPtr compile_cache(); + +// Erase cached compile function. +MLX_API void compile_erase( + const CompileCacheWeakPtr& cache, + std::uintptr_t fun_id); // Clear the compiler cache causing a recompilation of all compiled functions // when called again. -MLX_API void compile_clear_cache(); +MLX_API void compile_clear_cache(const CompileCacheWeakPtr& cache); bool compile_available_for_device(const Device& device); diff --git a/Source/Cmlx/mlx/mlx/core/moe_stream_op.cpp b/Source/Cmlx/mlx/mlx/core/moe_stream_op.cpp index 6316097d4..c287667a3 100644 --- a/Source/Cmlx/mlx/mlx/core/moe_stream_op.cpp +++ b/Source/Cmlx/mlx/mlx/core/moe_stream_op.cpp @@ -53,7 +53,7 @@ class LoadSSDExpert : public Primitive { eval_impl(inputs, outputs); auto& d = metal::device(mlx::core::Device::gpu); - d.add_temporary(outputs[0], stream().index); + metal::get_command_encoder(stream()).add_temporary(outputs[0]); } void eval_cpu(const std::vector& inputs, std::vector& outputs) override { diff --git a/Source/Cmlx/mlx/mlx/device.cpp b/Source/Cmlx/mlx/mlx/device.cpp index 4a62036e8..f0c868f21 100644 --- a/Source/Cmlx/mlx/mlx/device.cpp +++ b/Source/Cmlx/mlx/mlx/device.cpp @@ -25,20 +25,12 @@ void set_default_device(const Device& d) { mutable_default_device() = d; } -bool operator==(const Device& lhs, const Device& rhs) { - return lhs.type == rhs.type && lhs.index == rhs.index; -} - -bool operator!=(const Device& lhs, const Device& rhs) { - return !(lhs == rhs); -} - bool is_available(const Device& d) { switch (d.type) { case Device::cpu: - return cpu::is_available(); + return cpu::is_available() && (d.index < cpu::device_count()); case Device::gpu: - return gpu::is_available(); + return gpu::is_available() && (d.index < gpu::device_count()); } // appease compiler return false; diff --git a/Source/Cmlx/mlx/mlx/device.h b/Source/Cmlx/mlx/mlx/device.h index f89ad189f..ff75fe0cc 100644 --- a/Source/Cmlx/mlx/mlx/device.h +++ b/Source/Cmlx/mlx/mlx/device.h @@ -5,6 +5,7 @@ #include "mlx/api.h" #include +#include #include #include @@ -23,15 +24,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/mlx/mlx/distributed/distributed.cpp b/Source/Cmlx/mlx/mlx/distributed/distributed.cpp index 3cde6a263..34be2ec60 100644 --- a/Source/Cmlx/mlx/mlx/distributed/distributed.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/distributed.cpp @@ -138,9 +138,30 @@ Group Group::split(int color, int key /* = -1 */) const { return Group(group_->split(color, key)); } -Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { +namespace { + +std::unordered_map>& +get_backends() { static std::unordered_map> backends; + return backends; +} + +Group register_group(std::shared_ptr group, std::string bk) { + auto& backends = get_backends(); + if (group == nullptr) { + group = std::make_shared(); + } else { + backends.insert({"any", group}); + } + backends.insert({std::move(bk), group}); + return Group(group); +} + +} // namespace + +Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { + auto& backends = get_backends(); // Already initialized so return the group. if (auto g = backends.find(bk); g != backends.end()) { @@ -185,13 +206,31 @@ Group init(bool strict /* = false */, const std::string& bk /* = "any" */) { throw std::invalid_argument(msg.str()); } - if (group == nullptr) { - group = std::make_shared(); - } else { - backends.insert({"any", group}); + return register_group(std::move(group), std::move(bk_)); +} + +Group init(bool strict, const std::string& bk, AllGatherFactory factory) { + if (bk != "jaccl") { + std::ostringstream msg; + msg << "[distributed] An all gather factory is only supported with the " + << "'jaccl' backend but '" << bk << "' was provided."; + throw std::invalid_argument(msg.str()); } - backends.insert({std::move(bk_), group}); - return Group(group); + + auto& backends = get_backends(); + + // Already initialized so return the cached group. The factory is ignored in + // this case since the backend is already up. + if (auto g = backends.find(bk); g != backends.end()) { + return Group(g->second); + } + + auto group = jaccl::init(strict, std::move(factory)); + return register_group(std::move(group), bk); +} + +void clear_backends() { + get_backends().clear(); } } // namespace mlx::core::distributed diff --git a/Source/Cmlx/mlx/mlx/distributed/distributed.h b/Source/Cmlx/mlx/mlx/distributed/distributed.h index 00c7a80e8..0d6e4d41e 100644 --- a/Source/Cmlx/mlx/mlx/distributed/distributed.h +++ b/Source/Cmlx/mlx/mlx/distributed/distributed.h @@ -2,6 +2,8 @@ #pragma once +#include +#include #include #include "mlx/api.h" @@ -13,7 +15,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(); @@ -58,4 +79,15 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/.gitignore b/Source/Cmlx/mlx/mlx/distributed/jaccl/.gitignore new file mode 100644 index 000000000..ebf4281dc --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/.gitignore @@ -0,0 +1 @@ +!lib diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/CMakeLists.txt b/Source/Cmlx/mlx/mlx/distributed/jaccl/CMakeLists.txt index 567ed672a..359af8b2b 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/CMakeLists.txt @@ -1,12 +1,9 @@ if(MLX_BUILD_CPU AND ${CMAKE_SYSTEM_NAME} MATCHES "Darwin" - AND MACOS_SDK_VERSION GREATER_EQUAL 26.2) - target_sources( - mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/mesh.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/ring.cpp) + AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 + AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/jaccl.cpp) + target_link_libraries(mlx PRIVATE jaccl) else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_jaccl.cpp) endif() diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.cpp index 5935a98f1..99a4ac3be 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.cpp @@ -1,178 +1,193 @@ // Copyright © 2025 Apple Inc. -#include -#include +#include "mlx/distributed/jaccl/jaccl.h" +#include "mlx/backend/cpu/encoder.h" +#include "mlx/distributed/distributed_impl.h" +#include "mlx/dtype_utils.h" -#include +#include +#include -#include "mlx/distributed/distributed_impl.h" -#include "mlx/distributed/jaccl/mesh.h" -#include "mlx/distributed/jaccl/ring.h" -#include "mlx/distributed/jaccl/utils.h" +#include +#include using GroupImpl = mlx::core::distributed::detail::GroupImpl; -using json = nlohmann::json; + +namespace mlx::core::distributed::jaccl { namespace { -struct DeviceFile { - DeviceFile(const char* dev_file) { - std::ifstream f(dev_file); - json devices = json::parse(f); - if (!devices.is_array()) { - throw std::runtime_error( - "[jaccl] The device file should start with an array"); - } - - devices_.resize(devices.size()); - for (int rank = 0; rank < devices.size(); rank++) { - auto conn = devices[rank]; - if (!conn.is_array()) { - throw std::runtime_error( - "[jaccl] The device file should have an array of arrays"); - } - if (conn.size() != devices_.size()) { - std::ostringstream msg; - msg << "[jaccl] The device file should contain the connectivity of each rank to " - << "all other ranks but rank " << rank << " contains only " - << conn.size() << " entries."; - throw std::runtime_error(msg.str()); - } - - devices_[rank].resize(conn.size()); - for (int dst = 0; dst < conn.size(); dst++) { - auto names = conn[dst]; - if (names.is_string()) { - devices_[rank][dst].push_back(names); - } else if (names.is_array()) { - for (auto name_it = names.begin(); name_it != names.end(); - name_it++) { - devices_[rank][dst].push_back(*name_it); - } - } else if (!names.is_null()) { - throw std::runtime_error( - "[jaccl] Device names should be null, a string or array of strings."); - } - } - } +/** + * Map MLX Dtype to JACCL Dtype for dispatch. + */ +int dtype_to_jaccl_dtype(Dtype dt) { + switch (dt) { + case bool_: + return ::jaccl::Dtype::Bool; + case int8: + return ::jaccl::Dtype::Int8; + case int16: + return ::jaccl::Dtype::Int16; + case int32: + return ::jaccl::Dtype::Int32; + case int64: + return ::jaccl::Dtype::Int64; + case uint8: + return ::jaccl::Dtype::UInt8; + case uint16: + return ::jaccl::Dtype::UInt16; + case uint32: + return ::jaccl::Dtype::UInt32; + case uint64: + return ::jaccl::Dtype::UInt64; + case float16: + return ::jaccl::Dtype::Float16; + case bfloat16: + return ::jaccl::Dtype::BFloat16; + case float32: + return ::jaccl::Dtype::Float32; + case float64: + return ::jaccl::Dtype::Float64; + case complex64: + return ::jaccl::Dtype::Complex64; + default: + throw std::runtime_error("[jaccl] Unsupported dtype for JACCL operation"); + } +} + +/** + * Adapter that wraps a standalone jaccl::Group to implement + * MLX's distributed::detail::GroupImpl interface. + * + * This bridges mlx::core::array to raw pointers for JACCL operations. + */ +class JACCLGroup : public GroupImpl { + public: + JACCLGroup(std::shared_ptr<::jaccl::Group> group) + : group_(std::move(group)) {} + + Stream communication_stream(StreamOrDevice s) override { + return to_stream(s, Device::cpu); + } + + int rank() override { + return group_->rank(); + } + + int size() override { + return group_->size(); + } + + void all_sum(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->all_sum(in_ptr, out_ptr, n_bytes, dtype); + }); } - int size() { - return devices_.size(); + void all_max(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->all_max(in_ptr, out_ptr, n_bytes, dtype); + }); } - bool is_valid_mesh() { - for (int src = 0; src < size(); src++) { - for (int dst = 0; dst < size(); dst++) { - if (devices_[src][dst].size() != static_cast(src != dst)) { - return false; - } - } - } + void all_min(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->all_min(in_ptr, out_ptr, n_bytes, dtype); + }); + } - return true; + void all_gather(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = input.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() { + group_->all_gather(in_ptr, out_ptr, n_bytes); + }); } - bool is_valid_ring() { - int num_connections = devices_[0][1].size(); - if (num_connections == 0) { - return false; - } - - for (int src = 0; src < size(); src++) { - int left = (src + size() - 1) % size(); - int right = (src + 1) % size(); - for (int dst = 0; dst < size(); dst++) { - if (dst != left && dst != right) { - if (devices_[src][dst].size() != 0) { - return false; - } - } else { - if (devices_[src][dst].size() != num_connections) { - return false; - } - } - } - } - - return true; + void send(const array& input, int dst, Stream stream) override { + auto data = input.data(); + size_t n_bytes = input.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.dispatch( + [data, n_bytes, dst, this]() { group_->send(data, n_bytes, dst); }); } - std::vector extract_mesh_connectivity(int rank) { - std::vector devices(size()); - for (int dst = 0; dst < size(); dst++) { - if (dst != rank) { - devices[dst] = devices_[rank][dst][0]; - } - } - return devices; + void recv(array& out, int src, Stream stream) override { + auto data = out.data(); + size_t n_bytes = out.nbytes(); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_output_array(out); + encoder.dispatch( + [data, n_bytes, src, this]() { group_->recv(data, n_bytes, src); }); } - std::pair, std::vector> - extract_ring_connectivity(int rank) { - int left = (rank + size() - 1) % size(); - int right = (rank + 1) % size(); + void sum_scatter(const array& input, array& output, Stream stream) override { + auto in_ptr = input.data(); + auto out_ptr = output.data(); + size_t n_bytes = output.nbytes(); + int dtype = dtype_to_jaccl_dtype(output.dtype()); + auto& encoder = cpu::get_command_encoder(stream); + encoder.set_input_array(input); + encoder.set_output_array(output); + encoder.dispatch([in_ptr, out_ptr, n_bytes, dtype, this]() { + group_->sum_scatter(in_ptr, out_ptr, n_bytes, dtype); + }); + } - return std::make_pair(devices_[rank][left], devices_[rank][right]); + std::shared_ptr split(int color, int key = -1) override { + throw std::runtime_error("[jaccl] Group split not supported."); } - std::vector>> devices_; + private: + std::shared_ptr<::jaccl::Group> group_; }; } // namespace -namespace mlx::core::distributed::jaccl { - bool is_available() { - return ibv().is_available(); + return ::jaccl::is_available(); } std::shared_ptr init(bool strict /* = false */) { - const char* dev_file = std::getenv("MLX_IBV_DEVICES"); - const char* coordinator = std::getenv("MLX_JACCL_COORDINATOR"); - const char* rank_str = std::getenv("MLX_RANK"); - const char* ring = std::getenv("MLX_JACCL_RING"); - - if (!is_available() || !dev_file || !coordinator || !rank_str) { - if (strict) { - std::ostringstream msg; - msg << "[jaccl] You need to provide via environment variables a rank (MLX_RANK), " - << "a device file (MLX_IBV_DEVICES) and a coordinator ip/port (MLX_JACCL_COORDINATOR) " - << "but provided MLX_RANK=\"" << ((rank_str) ? rank_str : "") - << "\", MLX_IBV_DEVICES=\"" << ((dev_file) ? dev_file : "") - << "\" and MLX_JACCL_COORDINATOR=\"" - << ((coordinator) ? coordinator : ""); - throw std::runtime_error(msg.str()); - } + auto group = ::jaccl::init(strict); + if (group == nullptr) { return nullptr; } + return std::make_shared(std::move(group)); +} - auto rank = std::atoi(rank_str); - bool prefer_ring = ring != nullptr; - DeviceFile devices(dev_file); - - if (rank >= devices.size() || rank < 0) { - std::ostringstream msg; - msg << "[jaccl] Invalid rank " << rank << ". It should be between 0 and " - << devices.size(); - throw std::runtime_error(msg.str()); - } - - if (prefer_ring && devices.is_valid_ring()) { - auto [left, right] = devices.extract_ring_connectivity(rank); - return std::make_shared( - rank, devices.size(), left, right, coordinator); - } else if (devices.is_valid_mesh()) { - auto device_names = devices.extract_mesh_connectivity(rank); - return std::make_shared(rank, device_names, coordinator); - } else if (devices.is_valid_ring()) { - auto [left, right] = devices.extract_ring_connectivity(rank); - return std::make_shared( - rank, devices.size(), left, right, coordinator); - } else { - throw std::runtime_error( - "[jaccl] The device file should define a valid mesh or a valid ring."); +std::shared_ptr init(bool strict, AllGatherFactory factory) { + auto group = ::jaccl::init(strict, factory); + if (group == nullptr) { + return nullptr; } + return std::make_shared(std::move(group)); } } // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.h index d07f9cccc..0e2210ac7 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/jaccl.h @@ -1,5 +1,9 @@ // Copyright © 2025 Apple Inc. +#pragma once + +#include + #include "mlx/distributed/distributed.h" namespace mlx::core::distributed::jaccl { @@ -8,5 +12,6 @@ using GroupImpl = mlx::core::distributed::detail::GroupImpl; bool is_available(); std::shared_ptr init(bool strict = false); +std::shared_ptr init(bool strict, AllGatherFactory factory); } // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/CMakeLists.txt b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/CMakeLists.txt new file mode 100644 index 000000000..0cfee400d --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/CMakeLists.txt @@ -0,0 +1,83 @@ +cmake_minimum_required(VERSION 3.24) +project(jaccl LANGUAGES CXX) + +# Default to Release when built standalone. Without this CMake uses an empty +# build type (-O0) which severely impacts the reduction / memcpy hot paths. +if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE + Release + CACHE STRING "Build type" FORCE) +endif() + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +include(FetchContent) + +# nlohmann/json for device file parsing +message(STATUS "Downloading json for JACCL") +FetchContent_Declare( + json + URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz) +FetchContent_MakeAvailable(json) + +# Check platform and SDK version requirements +if(NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin") + message(STATUS "JACCL requires macOS (Darwin). Skipping JACCL build.") + return() +endif() + +# Try to determine MACOS_SDK_VERSION if not set +if(NOT DEFINED MACOS_SDK_VERSION) + execute_process( + COMMAND xcrun --sdk macosx --show-sdk-version + OUTPUT_VARIABLE MACOS_SDK_VERSION + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(MACOS_SDK_VERSION) + message(STATUS "Detected macOS SDK version: ${MACOS_SDK_VERSION}") + endif() +endif() + +if(DEFINED MACOS_SDK_VERSION AND MACOS_SDK_VERSION VERSION_LESS "26.2") + message(STATUS "JACCL requires macOS SDK >= 26.2. Skipping JACCL build.") + return() +endif() + +add_library( + jaccl + ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/tcp.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/rdma.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/mesh.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/ring.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/jaccl.cpp) + +target_include_directories( + jaccl PUBLIC $ + $) + +target_include_directories( + jaccl PRIVATE $) + +target_compile_features(jaccl PUBLIC cxx_std_20) + +# Install targets +install( + TARGETS jaccl + EXPORT jacclTargets + LIBRARY DESTINATION lib + ARCHIVE DESTINATION lib + RUNTIME DESTINATION bin + INCLUDES + DESTINATION include) + +install( + DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/jaccl/ + DESTINATION include/jaccl + FILES_MATCHING + PATTERN "*.h") + +install( + EXPORT jacclTargets + FILE jacclTargets.cmake + NAMESPACE jaccl:: + DESTINATION lib/cmake/jaccl) diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/README.md b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/README.md new file mode 100644 index 000000000..d682a06bb --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/README.md @@ -0,0 +1,333 @@ +# JACCL + +**JACCL** is a low-latency distributed communication library designed for macOS +systems with Thunderbolt 5 connectivity. + +## Overview + +JACCL leverages RDMA (Remote Direct Memory Access) over Thunderbolt to achieve +communication latency an order of magnitude lower than traditional TCP-based +approaches. This makes it ideal for: + +- Tensor parallelism in large model inference +- High-performance distributed training +- Low-latency collective operations between Macs + +JACCL was made possible by Apple's RDMA over Thunderbolt technology introduced +in macOS 26.2. + +## Features + +- **Mesh Topology**: Fully connected communication where each node can directly + communicate with any other node +- **Ring Topology**: High-bandwidth ring all-reduce for large messages +- **Collective Operations**: + - `all_sum`: Sum values across all nodes + - `all_max`: Element-wise maximum across all nodes + - `all_min`: Element-wise minimum across all nodes + - `all_gather`: Gather data from all nodes +- **Point-to-Point Operations**: + - `send`: Send data to a specific node + - `recv`: Receive data from a specific node +- **Synchronization**: + - `barrier`: Block until all nodes in the group reach this point +- **Type Support**: Bool, Int8-64, UInt8-64, Float16, BFloat16, Float32, + Float64, Complex64 + +## Requirements + +- macOS SDK >= 26.2 +- Thunderbolt 5 connectivity between nodes +- RDMA over Thunderbolt enabled (requires macOS recovery mode setup) + +## Enabling RDMA over Thunderbolt + +RDMA over Thunderbolt must be enabled in macOS recovery mode: + +1. Start your Mac in [recovery mode](https://support.apple.com/en-us/102518) +2. Open Terminal from Utilities -> Terminal +3. Run: `rdma_ctl enable` +4. Reboot + +To verify RDMA is enabled, run: + +```bash +ibv_devices +``` + +You should see output like: + +``` +device node GUID +------ ---------------- +rdma_en2 8096a9d9edbaac05 +rdma_en3 8196a9d9edbaac05 +rdma_en5 8396a9d9edbaac05 +``` + +## Building + +JACCL can be built as a standalone library: + +```bash +cd mlx/distributed/jaccl/lib +mkdir build && cd build +cmake .. +make +``` + +You can also include it in your own project via CMake: + +``` +FetchContent_Declare( + jaccl + GIT_REPOSITORY https://github.com/ml-explore/mlx.git + GIT_TAG main + SOURCE_SUBDIR mlx/distributed/jaccl/lib +) +FetchContent_MakeAvailable(jaccl) +``` + +## Usage + +### Environment Variables + +The easiest way to intiialize JACCL is by using the following environment +variables: + +- **JACCL_RANK** / **MLX_RANK**: The rank of this process (0-based integer) +- **JACCL_IBV_DEVICES** / **MLX_IBV_DEVICES**: Path to a JSON file describing + device connectivity +- **JACCL_COORDINATOR** / **MLX_JACCL_COORDINATOR**: IP:port of the coordinator + (rank 0 listener) +- **JACCL_RING** / **MLX_JACCL_RING**: (Optional) Prefer ring topology over + mesh + +### Device File Format + +The device file is a JSON array where each entry describes the RDMA devices +connecting that rank to all other ranks: + +```json +[ + [null, "rdma_en5", "rdma_en4", "rdma_en3"], + ["rdma_en5", null, "rdma_en3", "rdma_en4"], + ["rdma_en4", "rdma_en3", null, "rdma_en5"], + ["rdma_en3", "rdma_en4", "rdma_en5", null] +] +``` + +For a valid mesh, `devices[i][j]` should contain the device name connecting +rank `i` to rank `j`, or `null` if `i == j`. + +For a valid ring, only adjacent nodes should have device names (all others +should be null). + +### Basic Example + +```cpp +#include +#include + +int main() { + // Initialize JACCL group + auto group = jaccl::init(); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl; + + // Perform all-reduce sum + float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + float output[10]; + + group->all_sum(input, output, sizeof(input), jaccl::Float32); + + std::cout << "Result: " << output[0] << std::endl; + + return 0; +} +``` + +You can also manually define the configuration instead of reading it from +environment variables. + +```cpp +#include +#include + +int main() { + auto cfg = jaccl::Config() + .set_rank(0) + .set_coordinator("192.168.1.1:32132") + .set_devices({ + {{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}}, + {{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}}, + {{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}}, + {{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}} + }); + auto group = jaccl::init(cfg); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl; + + // Perform all-reduce sum + float input[10] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + float output[10]; + + group->all_sum(input, output, sizeof(input), jaccl::Float32); + + std::cout << "Result: " << output[0] << std::endl; + + return 0; +} +``` + +### Using with MLX + +JACCL integrates seamlessly with MLX's distributed communication: + +```python +import mlx.core as mx + +# Initialize with JACCL backend +world = mx.distributed.init(backend="jaccl") + +# Perform distributed operations +x = mx.ones((10,)) +result = mx.distributed.all_sum(x, group=world) +``` + +Launch with `mlx.launch`: + +```bash +mlx.launch --backend jaccl --hostfile hosts.json my_script.py +``` + +## Hostfile Example + +For use with `mlx.launch`, create a hostfile JSON: + +```json +{ + "backend": "jaccl", + "hosts": [ + { + "ssh": "m3-ultra-1", + "ips": ["192.168.1.1"], + "rdma": [null, "rdma_en5", "rdma_en4", "rdma_en3"] + }, + { + "ssh": "m3-ultra-2", + "ips": [], + "rdma": ["rdma_en5", null, "rdma_en3", "rdma_en4"] + }, + { + "ssh": "m3-ultra-3", + "ips": [], + "rdma": ["rdma_en4", "rdma_en3", null, "rdma_en5"] + }, + { + "ssh": "m3-ultra-4", + "ips": [], + "rdma": ["rdma_en3", "rdma_en4", "rdma_en5", null] + } + ] +} +``` + +## Automatic Configuration + +MLX provides `mlx.distributed_config` to automatically discover and configure +Thunderbolt connectivity: + +```bash +# Visualize connections +mlx.distributed_config --verbose \ + --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ + --over thunderbolt --dot | dot -Tpng | open -f -a Preview + +# Auto-configure and generate hostfile +mlx.distributed_config --verbose \ + --hosts m3-ultra-1,m3-ultra-2,m3-ultra-3,m3-ultra-4 \ + --over thunderbolt --backend jaccl \ + --auto-setup --output m3-ultra-jaccl.json +``` + +## API + +The main API of JACCL is the communication group. It provides efficient +high-level collectives. + +**Note: JACCL does no memory allocation. All output pointers should point to a +location with sufficient memory allocated to hold the result.** + +```cpp +class Group { + public: + virtual ~Group() {} + + // Helper functions to know which process we are in the group + virtual int rank() = 0; + virtual int size() = 0; + + // All reduce implementations. Input and output of the same size the + // reduction happens according to dtype and across the group. + virtual void all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0; + virtual void all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0; + virtual void all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0; + + // All gather implementation. The output is group->size() * n_bytes. + virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0; + + // Simple send/recv primitives. + virtual void send(const void* input, size_t n_bytes, int dst) = 0; + virtual void recv(void* output, size_t n_bytes, int src) = 0; + + // Block until every rank reaches this point. + virtual void barrier() = 0; +}; +``` + +All that is left to use JACCL (except the communication group) is + +```cpp +std::shared_ptr init(bool strict = false); +std::shared_ptr init(const Config& cfg, bool strict = false); +``` + +that create the communication group from environment variables or from the +configuration object. The latter allows one to configure JACCL using means +other than environment variables. + +```cpp +class Config { + public: + + Config(); + + Config& set_rank(int rank); + Config& set_coordinator(std::string coordinator); + Config& set_devices(std::vector>> devices); + Config& prefer_ring(bool prefer = true); + + bool is_valid_mesh() const; + bool is_valid_ring() const; +} +``` + +## License + +JACCL is part of MLX and is released under the same license. + +## Acknowledgments + +The name JACCL (pronounced Jackal) stands for Jack and Angelos’ Collective +Communication Library and it is an obvious pun to Nvidia’s NCCL but also +tribute to Jack Beasley who led the development of RDMA over Thunderbolt at +Apple. diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/CMakeLists.txt b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/CMakeLists.txt new file mode 100644 index 000000000..869f5dc2b --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/CMakeLists.txt @@ -0,0 +1,41 @@ +cmake_minimum_required(VERSION 3.24) +project(jaccl_examples LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +include(FetchContent) + +# cmake-format: off +# +# Treating the nearby directory like a fetched dependency to simulate what a +# user would do in their own C++ project. Ideally you would write something like +# the following: +# +# FetchContent_Declare( +# jaccl +# GIT_REPOSITORY https://github.com/ml-explore/mlx.git +# GIT_TAG main +# SOURCE_SUBDIR mlx/distributed/jaccl/lib +# ) +# +# cmake-format: on +FetchContent_Declare(jaccl SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/../) +FetchContent_MakeAvailable(jaccl) + +# Helper function to build examples +function(build_example SRCFILE) + get_filename_component(example_name ${SRCFILE} NAME_WE) + set(target "jaccl_${example_name}") + add_executable(${target} ${SRCFILE}) + target_link_libraries(${target} PRIVATE jaccl) + message(STATUS "Building JACCL example: ${target}") +endfunction() + +# Examples +build_example(minimal_env.cpp) +build_example(minimal_cfg.cpp) +build_example(minimal_barrier.cpp) + +# Benchmarks +build_example(allreduce_bench.cpp) diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp new file mode 100644 index 000000000..ae1fa5a74 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/allreduce_bench.cpp @@ -0,0 +1,286 @@ +// Copyright © 2025 Apple Inc. +// +// JACCL All-Reduce Benchmark +// +// Measures bandwidth and latency of all_sum across a sweep of message sizes, +// similar in spirit to the NCCL all-reduce benchmark (nccl-tests). +// +// Usage: +// Set the environment variables described in jaccl.h, then run: +// +// ./jaccl_allreduce_bench [-w ] [-n ] +// [-b ] [-e ] +// [-f ] [-d ] +// [-c] [-h] +// +// Or use the MLX launcher: +// +// mlx.launch --hostfile hosts.json ./jaccl_allreduce_bench +// +// The arguments are: +// +// -w Warmup iterations per message size (default: 5) +// -n Timed iterations per message size (default: 20) +// -b Minimum message size in bytes (default: 1K) +// -e Maximum message size in bytes (default: 256M) +// -f Multiplicative step factor (default: 2) +// -d Datatype: float32, float16, bfloat16 (default: float32) +// -c Check correctness (default: off) +// -h Print this help message + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static void usage(const char* prog) { + std::cerr + << "Usage: " << prog << " [options]\n" + << " -w Warmup iterations (default: 5)\n" + << " -n Timed iterations (default: 20)\n" + << " -b Minimum message size (default: 1K)\n" + << " -e Maximum message size (default: 256M)\n" + << " -f Multiplicative step factor (default: 2)\n" + << " -d float32|float16|bfloat16 (default: float32)\n" + << " -c Check correctness\n" + << " -h Show this help\n"; +} + +static size_t parse_size(const char* s) { + char* end = nullptr; + double val = std::strtod(s, &end); + if (end && (*end == 'K' || *end == 'k')) + val *= 1024; + else if (end && (*end == 'M' || *end == 'm')) + val *= 1024 * 1024; + else if (end && (*end == 'G' || *end == 'g')) + val *= 1024 * 1024 * 1024; + return static_cast(val); +} + +static std::string fmt_bytes(size_t bytes) { + const char* units[] = {"B", "KB", "MB", "GB", "TB"}; + int idx = 0; + double val = static_cast(bytes); + while (val >= 1024.0 && idx < 4) { + val /= 1024.0; + idx++; + } + char buf[32]; + if (val == static_cast(val)) + std::snprintf(buf, sizeof(buf), "%d %s", static_cast(val), units[idx]); + else + std::snprintf(buf, sizeof(buf), "%.2f %s", val, units[idx]); + return buf; +} + +// Conversion from algorithm bandwidth to bus bandwidth for a ring reduce. +static double bus_factor(int nranks) { + return 2.0 * (nranks - 1) / static_cast(nranks); +} + +int main(int argc, char** argv) { + int warmup_iters = 5; + int timed_iters = 20; + size_t min_bytes = 1024; + size_t max_bytes = 256 * 1024 * 1024; + int step_factor = 2; + std::string dtype_str = "float32"; + bool check = false; + + for (int i = 1; i < argc; i++) { + std::string arg = argv[i]; + if (arg == "-h" || arg == "--help") { + usage(argv[0]); + return 0; + } else if (arg == "-w" && i + 1 < argc) { + warmup_iters = std::atoi(argv[++i]); + } else if (arg == "-n" && i + 1 < argc) { + timed_iters = std::atoi(argv[++i]); + } else if (arg == "-b" && i + 1 < argc) { + min_bytes = parse_size(argv[++i]); + } else if (arg == "-e" && i + 1 < argc) { + max_bytes = parse_size(argv[++i]); + } else if (arg == "-f" && i + 1 < argc) { + step_factor = std::atoi(argv[++i]); + } else if (arg == "-d" && i + 1 < argc) { + dtype_str = argv[++i]; + } else if (arg == "-c") { + check = true; + } else { + std::cerr << "Unknown option: " << arg << "\n"; + usage(argv[0]); + return 1; + } + } + + jaccl::Dtype dtype; + size_t elem_size; + if (dtype_str == "float32") { + dtype = jaccl::Float32; + elem_size = 4; + } else if (dtype_str == "float16") { + dtype = jaccl::Float16; + elem_size = 2; + } else if (dtype_str == "bfloat16") { + dtype = jaccl::BFloat16; + elem_size = 2; + } else { + std::cerr << "Unsupported dtype: " << dtype_str << "\n"; + return 1; + } + + auto group = jaccl::init(/* strict= */ true); + + int rank = group->rank(); + int nranks = group->size(); + + if (rank == 0) { + std::cout << "# JACCL All-Reduce Benchmark\n" + << "# Ranks: " << nranks << "\n" + << "# Dtype: " << dtype_str << "\n" + << "# Warmup: " << warmup_iters << " iters\n" + << "# Timed: " << timed_iters << " iters\n" + << "# Sizes: " << fmt_bytes(min_bytes) << " .. " + << fmt_bytes(max_bytes) << " (x" << step_factor << ")\n" + << "#\n"; + + // Table header (NCCL-style) + std::cout << std::left << std::setw(14) << "# size" << std::right + << std::setw(12) << "count" << std::setw(12) << "type" + << std::setw(14) << "time (us)" << std::setw(14) << "algo BW" + << std::setw(14) << "bus BW"; + if (check) + std::cout << std::setw(10) << "check"; + std::cout << "\n"; + std::cout << std::left << std::setw(14) << "# (bytes)" << std::right + << std::setw(12) << "(elems)" << std::setw(12) << "" + << std::setw(14) << "" << std::setw(14) << "(GB/s)" + << std::setw(14) << "(GB/s)"; + if (check) + std::cout << std::setw(10) << ""; + std::cout << "\n"; + } + + size_t max_elems = max_bytes / elem_size; + max_bytes = max_elems * elem_size; + + std::vector sendbuf(max_bytes); + + // Fill send buffer with a simple deterministic pattern (rank + 1) casted to + // the target type, so correctness checks are straightforward: after all_sum + // every element should equal sum_{r=0}^{nranks-1} (r + 1) = + // nranks*(nranks+1)/2. + auto fill_buffer = [&](char* buf, size_t n_bytes) { + float val = static_cast(rank + 1); + if (dtype == jaccl::Float32) { + auto* p = reinterpret_cast(buf); + size_t n = n_bytes / sizeof(float); + for (size_t i = 0; i < n; i++) { + p[i] = val; + } + } else if (dtype == jaccl::Float16) { + // Write via the library's float16_t + auto* p = reinterpret_cast(buf); + size_t n = n_bytes / sizeof(jaccl::float16_t); + for (size_t i = 0; i < n; i++) { + p[i] = jaccl::float16_t(val); + } + } else if (dtype == jaccl::BFloat16) { + auto* p = reinterpret_cast(buf); + size_t n = n_bytes / sizeof(jaccl::bfloat16_t); + for (size_t i = 0; i < n; i++) { + p[i] = jaccl::bfloat16_t(val); + } + } + }; + + auto check_buffer = [&](const char* buf, size_t n_bytes) -> bool { + float expected = static_cast(nranks) * (nranks + 1) / 2.0f; + float tol = (dtype == jaccl::Float32) ? 1e-5f : 1e-1f; + size_t n = n_bytes / elem_size; + + for (size_t i = 0; i < n; i++) { + float val; + if (dtype == jaccl::Float32) { + val = reinterpret_cast(buf)[i]; + } else if (dtype == jaccl::Float16) { + val = static_cast( + reinterpret_cast(buf)[i]); + } else { + val = static_cast( + reinterpret_cast(buf)[i]); + } + if (std::abs(val - expected) > tol) { + return false; + } + } + return true; + }; + + double bf = bus_factor(nranks); + + for (size_t nbytes = min_bytes; nbytes <= max_bytes; + nbytes *= static_cast(step_factor)) { + // Round down to element boundary + size_t n = std::max((nbytes / elem_size) * elem_size, elem_size); + size_t count = n / elem_size; + + fill_buffer(sendbuf.data(), n); + + // Warmup + for (int i = 0; i < warmup_iters; i++) { + group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype); + } + + // Timed iterations + auto t0 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < timed_iters; i++) { + group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype); + } + auto t1 = std::chrono::high_resolution_clock::now(); + + double elapsed_us = + std::chrono::duration(t1 - t0).count(); + double avg_us = elapsed_us / timed_iters; + + // Bandwidth in GB/s + double algo_bw = (static_cast(n) / avg_us) / 1e3; + double bus_bw = algo_bw * bf; + + // Correctness check + std::string check_result; + if (check) { + fill_buffer(sendbuf.data(), n); + group->all_sum(sendbuf.data(), sendbuf.data(), n, dtype); + check_result = check_buffer(sendbuf.data(), n) ? "OK" : "FAIL"; + } + + if (rank == 0) { + std::cout << std::left << std::setw(14) << n << std::right + << std::setw(12) << count << std::setw(12) << dtype_str + << std::setw(14) << std::fixed << std::setprecision(1) << avg_us + << std::setw(14) << std::fixed << std::setprecision(2) + << algo_bw << std::setw(14) << std::fixed + << std::setprecision(2) << bus_bw; + if (check) + std::cout << std::setw(10) << check_result; + std::cout << "\n"; + } + } + + if (rank == 0) { + std::cout << "# Done.\n"; + } + + return 0; +} diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp new file mode 100644 index 000000000..0d94ffb6e --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_barrier.cpp @@ -0,0 +1,42 @@ +// Copyright © 2026 Apple Inc. +// +// Exercises Group::barrier(). Ranks arrive at the barrier at staggered times; +// after the barrier returns we do a small all_sum to confirm the group is +// healthy and that barrier() carried the correct fence semantics. + +#include +#include +#include + +#include + +int main() { + auto group = jaccl::init(); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + int rank = group->rank(); + int size = group->size(); + + std::this_thread::sleep_for(std::chrono::milliseconds(100 * rank)); + std::cout << "rank " << rank << " entering barrier" << std::endl; + + group->barrier(); + + std::cout << "rank " << rank << " exited barrier" << std::endl; + + int in = rank + 1; + int out = 0; + group->all_sum(&in, &out, sizeof(in), jaccl::Int32); + int expected = size * (size + 1) / 2; + if (out != expected) { + std::cerr << "rank " << rank << ": post-barrier all_sum mismatch (got " + << out << ", expected " << expected << ")" << std::endl; + return 1; + } + std::cout << "rank " << rank << ": post-barrier all_sum OK (" << out << ")" + << std::endl; + return 0; +} diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp new file mode 100644 index 000000000..4aec4b2e2 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_cfg.cpp @@ -0,0 +1,36 @@ +#include +#include + +int main() { + auto cfg = + jaccl::Config() + .set_rank(0) // should be different per node + .set_coordinator("192.168.1.1:32132") // rank 0 will listen here + .set_devices( + {{{}, {"rdma_en5"}, {"rdma_en4"}, {"rdma_en3"}}, + {{"rdma_en5"}, {}, {"rdma_en3"}, {"rdma_en4"}}, + {{"rdma_en4"}, {"rdma_en3"}, {}, {"rdma_en5"}}, + {{"rdma_en3"}, {"rdma_en4"}, {"rdma_en5"}, {}}}); + auto group = jaccl::init(cfg); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl; + + // Perform all-reduce sum + float input[10] = { + 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + float output[10]; + + group->all_sum(input, output, sizeof(input), jaccl::Float32); + + std::cout << "Result: "; + for (auto o : output) { + std::cout << o << " "; + } + std::cout << std::endl; + + return 0; +} diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_env.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_env.cpp new file mode 100644 index 000000000..2bc04da56 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/examples/minimal_env.cpp @@ -0,0 +1,28 @@ +#include +#include + +int main() { + // Initialize JACCL group + auto group = jaccl::init(); + if (!group) { + std::cerr << "Failed to initialize JACCL" << std::endl; + return 1; + } + + std::cout << "Rank " << group->rank() << " of " << group->size() << std::endl; + + // Perform all-reduce sum + float input[10] = { + 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + float output[10]; + + group->all_sum(input, output, sizeof(input), jaccl::Float32); + + std::cout << "Result: "; + for (auto o : output) { + std::cout << o << " "; + } + std::cout << std::endl; + + return 0; +} diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/group.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/group.h new file mode 100644 index 000000000..603765a51 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/group.h @@ -0,0 +1,69 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include +#include + +namespace jaccl { + +/** + * Abstract base class for a JACCL communication group. + */ +class Group { + public: + virtual ~Group() {} + + virtual int rank() = 0; + virtual int size() = 0; + + virtual void + all_sum(const void* input, void* output, size_t n_bytes, int dtype) = 0; + + virtual void + all_max(const void* input, void* output, size_t n_bytes, int dtype) = 0; + + virtual void + all_min(const void* input, void* output, size_t n_bytes, int dtype) = 0; + + virtual void all_gather(const void* input, void* output, size_t n_bytes) = 0; + + /** + * Reduce scatter with a sum reduction. + * + * The input holds size() contiguous chunks of n_bytes each (total + * size() * n_bytes bytes). After the call, output (n_bytes bytes) on rank r + * contains the elementwise sum over all ranks of the r-th input chunk. + */ + virtual void + sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) = 0; + + virtual void send(const void* input, size_t n_bytes, int dst) = 0; + virtual void recv(void* output, size_t n_bytes, int src) = 0; + virtual void barrier() = 0; +}; + +/** + * Type IDs for dispatch in the standalone JACCL library. + * + * Users pass one of these to all_sum/all_max/all_min so JACCL knows how to + * interpret the data for typed reduction operations. + */ +enum Dtype { + Bool = 0, + Int8, + Int16, + Int32, + Int64, + UInt8, + UInt16, + UInt32, + UInt64, + Float16, + BFloat16, + Float32, + Float64, + Complex64, +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp new file mode 100644 index 000000000..7715022fc --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.cpp @@ -0,0 +1,302 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include + +#include + +#include "jaccl/jaccl.h" +#include "jaccl/mesh.h" +#include "jaccl/ring.h" + +using json = nlohmann::json; + +namespace { + +std::vector>> parse_devices_json( + const char* dev_file) { + std::ifstream f(dev_file); + json devices = json::parse(f); + if (!devices.is_array()) { + throw std::runtime_error( + "[jaccl] The device file should start with an array"); + } + + std::vector>> result(devices.size()); + for (int rank = 0; rank < devices.size(); rank++) { + auto conn = devices[rank]; + if (!conn.is_array()) { + throw std::runtime_error( + "[jaccl] The device file should have an array of arrays"); + } + if (conn.size() != devices.size()) { + std::ostringstream msg; + msg << "[jaccl] The device file should contain the connectivity of each rank to " + << "all other ranks but rank " << rank << " contains only " + << conn.size() << " entries."; + throw std::runtime_error(msg.str()); + } + + result[rank].resize(conn.size()); + for (int dst = 0; dst < conn.size(); dst++) { + auto names = conn[dst]; + if (names.is_string()) { + result[rank][dst].push_back(names); + } else if (names.is_array()) { + for (auto name_it = names.begin(); name_it != names.end(); name_it++) { + result[rank][dst].push_back(*name_it); + } + } else if (!names.is_null()) { + throw std::runtime_error( + "[jaccl] Device names should be null, a string or array of strings."); + } + } + } + + return result; +} + +template +const char* getenv(First first, Rest... rest) { + const char* rs = std::getenv(first); + if (rs != nullptr) { + return rs; + } + if constexpr (sizeof...(rest) > 0) { + return getenv(rest...); + } + return rs; +} + +} // namespace + +namespace jaccl { + +Config::Config() : rank_(0), size_(0) {} + +Config& Config::set_rank(const char* rank_str) { + if (rank_str) { + rank_ = std::atoi(rank_str); + } + return *this; +} + +Config& Config::set_rank(int rank) { + rank_ = rank; + return *this; +} + +Config& Config::set_coordinator(const char* coordinator) { + if (coordinator != nullptr) { + coordinator_ = coordinator; + } + return *this; +} + +Config& Config::set_coordinator(std::string coordinator) { + coordinator_ = std::move(coordinator); + return *this; +} + +Config& Config::set_devices_from_file(const char* dev_file) { + if (dev_file) { + set_devices(parse_devices_json(dev_file)); + } + return *this; +} + +Config& Config::set_devices( + std::vector>> devices) { + devices_ = std::move(devices); + size_ = devices_.size(); + for (int r = 0; r < size_; r++) { + if (size_ != devices_[r].size()) { + std::ostringstream msg; + msg << "[jaccl] The full connectivity matrix should be provided but we have " + << size_ << " rows and row " << r << " has " << devices_[r].size() + << " columns."; + throw std::invalid_argument(msg.str()); + } + } + return *this; +} + +Config& Config::prefer_ring(bool prefer /* = true */) { + prefer_ring_ = prefer; + return *this; +} + +Config& Config::set_all_gather(AllGatherFn agf) { + all_gather_fn_ = std::move(agf); + return *this; +} + +Config& Config::set_all_gather_factory( + std::function factory) { + all_gather_factory_ = std::move(factory); + return *this; +} + +bool Config::is_valid_mesh() const { + if (size_ < 2) { + return false; + } + + for (int src = 0; src < size_; src++) { + for (int dst = 0; dst < size_; dst++) { + if ((src == dst && devices_[src][dst].size() != 0) || + (src != dst && devices_[src][dst].size() == 0)) { + return false; + } + } + } + return true; +} + +bool Config::is_valid_ring() const { + if (size_ < 2) { + return false; + } + + int num_connections = devices_[0][1].size(); + for (int src = 0; src < size_; src++) { + int left = (src + size_ - 1) % size_; + int right = (src + 1) % size_; + for (int dst = 0; dst < size_; dst++) { + if (dst == left || dst == right) { + if (devices_[src][dst].size() != num_connections) { + return false; + } + } + } + } + return true; +} + +bool Config::is_valid() const { + return size_ >= 1 && rank_ < size_ && rank_ >= 0 && + (!coordinator_.empty() || all_gather_factory_ || all_gather_fn_) && + (is_valid_mesh() || is_valid_ring()); +} + +std::vector Config::get_mesh_connectivity() const { + if (!is_valid_mesh()) { + throw std::runtime_error("[jaccl] The devices do not form a valid mesh."); + } + std::vector devices(size_); + for (int dst = 0; dst < size_; dst++) { + if (dst != rank_) { + devices[dst] = devices_[rank_][dst][0]; + } + } + return devices; +} + +std::pair, std::vector> +Config::get_ring_connectivity() const { + if (!is_valid_ring()) { + throw std::runtime_error("[jaccl] The devices do not form a valid ring."); + } + int left = (rank_ + size_ - 1) % size_; + int right = (rank_ + 1) % size_; + + return std::make_pair(devices_[rank_][left], devices_[rank_][right]); +} + +SideChannel Config::get_side_channel() const { + if (all_gather_factory_) { + return SideChannel(rank_, size_, all_gather_factory_(rank_, size_)); + } + + if (all_gather_fn_) { + return SideChannel(rank_, size_, all_gather_fn_); + } + + auto tcp = + std::make_shared(rank_, size_, get_coordinator().c_str()); + return SideChannel( + rank_, + size_, + [tcp = std::move(tcp)](const char* src, char* dst, size_t n_bytes) { + (*tcp)(src, dst, n_bytes); + }); +} + +Config Config::from_env() { + const char* dev_file = getenv("JACCL_IBV_DEVICES", "MLX_IBV_DEVICES"); + const char* coordinator = + getenv("JACCL_COORDINATOR", "MLX_JACCL_COORDINATOR"); + const char* rank_str = getenv("JACCL_RANK", "MLX_RANK"); + const char* ring = getenv("JACCL_RING", "MLX_JACCL_RING"); + + return Config() + .set_rank(rank_str) + .set_coordinator(coordinator) + .set_devices_from_file(dev_file) + .prefer_ring(ring != nullptr); +} + +bool is_available() { + return ibv().is_available(); +} + +std::shared_ptr init(bool strict /* = false */) { + auto cfg = Config::from_env(); + if (!cfg.is_valid()) { + if (strict) { + std::ostringstream msg; + msg << "[jaccl] You need to provide via environment variables a rank " + << "(JACCL_RANK/MLX_RANK), a device file (JACCL_IBV_DEVICES/" + << "MLX_IBV_DEVICES) and a coordinator ip/port (JACCL_COORDINATOR/" + << "MLX_JACCL_COORDINATOR)."; + throw std::runtime_error(msg.str()); + } + return nullptr; + } + + return init(cfg, strict); +} + +std::shared_ptr init( + bool strict, + std::function factory) { + auto cfg = Config::from_env().set_all_gather_factory(factory); + if (!cfg.is_valid()) { + if (strict) { + std::ostringstream msg; + msg << "[jaccl] You need to provide via environment variables a rank " + << "(JACCL_RANK/MLX_RANK) and a device file " + << "(JACCL_IBV_DEVICES/MLX_IBV_DEVICES)"; + throw std::runtime_error(msg.str()); + } + return nullptr; + } + + return init(cfg, strict); +} + +std::shared_ptr init(const Config& cfg, bool strict /* = false */) { + if (cfg.get_prefer_ring() && cfg.is_valid_ring()) { + auto [left, right] = cfg.get_ring_connectivity(); + return std::make_shared( + cfg.get_rank(), cfg.get_size(), left, right, cfg.get_side_channel()); + } else if (cfg.is_valid_mesh()) { + auto mesh = cfg.get_mesh_connectivity(); + return std::make_shared( + cfg.get_rank(), mesh, cfg.get_side_channel()); + } else if (cfg.is_valid_ring()) { + auto [left, right] = cfg.get_ring_connectivity(); + return std::make_shared( + cfg.get_rank(), cfg.get_size(), left, right, cfg.get_side_channel()); + } else { + if (!strict) { + return nullptr; + } + + throw std::runtime_error( + "[jaccl] The configuration should define a valid mesh or a valid ring."); + } +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h new file mode 100644 index 000000000..39feb34b4 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/jaccl.h @@ -0,0 +1,108 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include +#include +#include + +#include "jaccl/group.h" +#include "jaccl/rdma.h" + +namespace jaccl { + +class Config { + public: + Config(); + + Config& set_rank(const char* rank_str); + Config& set_rank(int rank); + Config& set_coordinator(const char* coordinator); + Config& set_coordinator(std::string coordinator); + Config& set_devices_from_file(const char* dev_file); + Config& set_devices( + std::vector>> devices); + Config& prefer_ring(bool prefer = true); + Config& set_all_gather(AllGatherFn agf); + Config& set_all_gather_factory(std::function factory); + + bool is_valid_mesh() const; + bool is_valid_ring() const; + bool is_valid() const; + + int get_rank() const { + return rank_; + } + + int get_size() const { + return size_; + } + + std::string get_coordinator() const { + return coordinator_; + } + + bool get_prefer_ring() const { + return prefer_ring_; + } + + static Config from_env(); + + friend std::shared_ptr init(const Config& cfg, bool strict); + + private: + std::vector get_mesh_connectivity() const; + std::pair, std::vector> + get_ring_connectivity() const; + SideChannel get_side_channel() const; + + int rank_; + int size_; + std::string coordinator_; + std::vector>> devices_; + bool prefer_ring_; + AllGatherFn all_gather_fn_; + std::function all_gather_factory_; +}; + +/** + * Check if JACCL (RDMA over Thunderbolt) is available on this system. + */ +bool is_available(); + +/** + * Initialize a JACCL communication group from environment variables. + * + * Reads configuration from environment variables: + * - JACCL_RANK / MLX_RANK: The rank of this process + * - JACCL_IBV_DEVICES / MLX_IBV_DEVICES: Path to the device connectivity + * JSON file + * - JACCL_COORDINATOR / MLX_JACCL_COORDINATOR: IP:port of the coordinator + * - JACCL_RING / MLX_JACCL_RING: If set, prefer ring topology + * + * Args: + * strict: If true, throw on failure. If false, return nullptr. + * + * Returns: + * A shared_ptr to the Group, or nullptr on failure. + */ +std::shared_ptr init(bool strict = false); + +/** + * Initialize a JACCL communication group from environment variables, using a + * custom all-gather factory for the side channel. + * + * The factory is called once per rank with the rank and group size, and must + * return an all-gather function that will be used to exchange RDMA connection + * metadata during setup. + */ +std::shared_ptr init( + bool strict, + std::function factory); + +/** + * Initialize a JACCL communication group from an explicit Config object. + */ +std::shared_ptr init(const Config& cfg, bool strict = false); + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp new file mode 100644 index 000000000..5b749f2bd --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.cpp @@ -0,0 +1,229 @@ +// Copyright © 2026 Apple Inc. + +#include "jaccl/mesh.h" +#include "jaccl/reduction_ops.h" +#include "jaccl/types.h" + +namespace jaccl { + +MeshGroup::MeshGroup( + int rank, + const std::vector& device_names, + SideChannel sc) + : rank_(rank), + size_(device_names.size()), + side_channel_(std::move(sc)), + connections_(create_connections(device_names)) { + if (size_ > MESH_MAX_PEERS) { + std::ostringstream msg; + msg << "[jaccl] The JACCL mesh supports up to " << MESH_MAX_PEERS + << " peers but " << size_ << " were provided."; + throw std::runtime_error(msg.str()); + } + + // Initialize all the connections and allocate buffers + initialize(); + + // Make sure every node has reached here before continuing + side_channel_.barrier(); + + // Create the mesh implementation object + mesh_ = MeshImpl(rank_, size_, connections_, buffers_, scatter_buffers_); +} + +void MeshGroup::initialize() { + // Create the queue pairs + for (auto& conn : connections_) { + if (conn.ctx == nullptr) { + continue; + } + conn.allocate_protection_domain(); + conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_queue_pair(); + } + + allocate_buffers(); + + // First init all connections + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + connections_[peer].queue_pair_init(); + } + + // Gather the information to be exchanged, this also serves as a barrier + // so that all peers have initialized their connections before attempting + // to transition to RTS. + std::vector info; + for (auto& conn : connections_) { + info.emplace_back(conn.info()); + } + auto all_infos = side_channel_.all_gather(info); + + // Transition queue pairs to RTS + for (int peer = 0; peer < size_; peer++) { + if (peer == rank_) { + continue; + } + auto peer_info = all_infos[peer][rank_]; + connections_[peer].queue_pair_rtr(peer_info); + connections_[peer].queue_pair_rts(); + } +} + +void MeshGroup::allocate_buffers() { + // Deregister any buffers and free the memory + buffers_.clear(); + scatter_buffers_.clear(); + + // Allocate the memory + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + // Mesh buffers + for (int j = 0; j < size_; j++) { + buffers_.emplace_back(FRAME_SIZE * (1 << k)); + } + // Scatter buffers (size_ send slots followed by size_ recv slots) + for (int j = 0; j < 2 * size_; j++) { + scatter_buffers_.emplace_back(FRAME_SIZE * (1 << k)); + } + } + } + + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + // Mesh buffers + for (int j = 0; j < size_; j++) { + if (j == rank_) { + // This is our send buffer so register it with all pds so we can + // send it to all connected devices. + for (auto& conn : connections_) { + if (conn.ctx != nullptr) { + buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] + .register_to_protection_domain(conn.protection_domain); + } + } + } else { + // This is the recv buffer from rank j so register it to rank j's + // protection domain. + buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] + .register_to_protection_domain(connections_[j].protection_domain); + } + } + + // Scatter buffers. Slot p (send to peer p) and slot size_ + p (recv from + // peer p) are both registered to peer p's protection domain. The slots + // for our own rank are unused but kept for uniform indexing. + int scatter_base = k * NUM_BUFFERS * 2 * size_ + i * 2 * size_; + for (int j = 0; j < size_; j++) { + if (j == rank_) { + continue; + } + scatter_buffers_[scatter_base + j].register_to_protection_domain( + connections_[j].protection_domain); + scatter_buffers_[scatter_base + size_ + j] + .register_to_protection_domain(connections_[j].protection_domain); + } + } + } +} + +void MeshGroup::all_sum( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, SumOp{}); + }); +} + +void MeshGroup::all_max( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, MaxOp{}); + }); +} + +void MeshGroup::all_min( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, MinOp{}); + }); +} + +void MeshGroup::all_gather(const void* input, void* output, size_t n_bytes) { + mesh_.all_gather( + static_cast(input), static_cast(output), n_bytes); +} + +void MeshGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + +void MeshGroup::send(const void* input, size_t n_bytes, int dst) { + mesh_.send(static_cast(input), n_bytes, dst); +} + +void MeshGroup::recv(void* output, size_t n_bytes, int src) { + mesh_.recv(static_cast(output), n_bytes, src); +} + +void MeshGroup::barrier() { + uint8_t b = 0; + all_sum(&b, &b, sizeof(b), Dtype::UInt8); +} + +template +void MeshGroup::all_reduce( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + if (size_ > 2 && n_bytes > 32 * 1024) { + // Large messages are bandwidth bound so use the reduce scatter + all gather + // path which moves size_x less data per link than the fully connected + // all_reduce. + mesh_.all_reduce_scatter_gather(in_ptr, out_ptr, count, reduce_op); + } else { + // Small messages are latency bound so use the single phase fully + // connected all_reduce cause it is a bit better. + mesh_.all_reduce(in_ptr, out_ptr, count, reduce_op); + } +} + +template +void MeshGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + // n_bytes is the size of the output (one chunk). The input holds size_ such + // chunks laid out contiguously. + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + mesh_.sum_scatter(in_ptr, out_ptr, count, reduce_op); +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h new file mode 100644 index 000000000..8e5164aee --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh.h @@ -0,0 +1,91 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include "jaccl/group.h" +#include "jaccl/mesh_impl.h" +#include "jaccl/rdma.h" + +namespace jaccl { + +/** + * The JACCL communication group for a fully connected mesh. We expect one + * connection per peer and it should be the lowest latency communication group + * for small to medium size messages. + * + * Like all JACCL groups it uses a side channel to exchange the necessary + * information and then configure the connections to be ready for RDMA + * operations. + */ +class MeshGroup : public Group { + public: + MeshGroup( + int rank, + const std::vector& device_names, + SideChannel sc); + + int rank() override { + return rank_; + } + + int size() override { + return size_; + } + + void all_sum(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_max(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_min(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_gather(const void* input, void* output, size_t n_bytes) override; + + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void send(const void* input, size_t n_bytes, int dst) override; + void recv(void* output, size_t n_bytes, int src) override; + + void barrier() override; + + private: + template + void all_reduce( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + + /** + * Performs the connection initialization. Namely, after this call all + * Connection objects should have a queue pair in RTS state and all buffers + * should have been allocated. + */ + void initialize(); + + /** + * Allocate all the buffers that we will use in the communication group. + */ + void allocate_buffers(); + + int rank_; + int size_; + SideChannel side_channel_; + std::vector connections_; + std::vector buffers_; + std::vector scatter_buffers_; + + MeshImpl mesh_; +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h new file mode 100644 index 000000000..75c540e91 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/mesh_impl.h @@ -0,0 +1,632 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include + +#include "jaccl/rdma.h" + +constexpr int MESH_MAX_PEERS = 8; +constexpr int MESH_PIPELINE = 2; +constexpr int64_t MAX_BUFFER_SIZE = FRAME_SIZE * (1 << (BUFFER_SIZES - 1)); + +namespace jaccl { + +class MeshImpl { + public: + MeshImpl( + int rank, + int size, + std::vector& conns, + std::vector& buffers, + std::vector& scatter_buffers) + : rank_(rank), + size_(size), + connections_(conns), + buffers_(buffers), + scatter_buffers_(scatter_buffers), + staging_mem_( + std::make_unique(MESH_PIPELINE * MAX_BUFFER_SIZE)) {} + + MeshImpl() : rank_(0), size_(1) {} + + template + void all_reduce(const T* in, T* out, int64_t size, ReduceOp reduce_op) { + // Fully connected all reduce with deterministic reduction order. + // + // We copy rank 0's data to the output buffer and then we reduce every + // subsequent rank in-place in the output. + // + // Our own data is copied to a staging buffer to ensure we can reduce it in + // the output when needed. + + auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T)); + int64_t N = buffer_size / sizeof(T); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; + int64_t total = static_cast(size); + int num_peers = size_ - 1; + + // A helper for convenient access to the staging buffer. + auto local_staging = [&](int buff) -> T* { + return reinterpret_cast(staging_mem_.get() + buff * MAX_BUFFER_SIZE); + }; + + // Counters to maintain the state of transfers + int in_flight = 0; + int64_t read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int recv_end[MESH_MAX_PEERS] = {0}; + int reduce_chunk = 0; + int reduce_rank = 0; + + // Total number of chunks + int64_t total_chunks = (total + N - 1) / N; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + post_recv_all(sz, buff); + + // Copy the local data to send buffer and staging buffer + int64_t elems = std::min(N, total - read_offset); + std::copy( + in + read_offset, in + read_offset + elems, local_staging(buff)); + std::copy( + in + read_offset, + in + read_offset + elems, + send_buffer(sz, buff).begin()); + recv_end[rank_]++; + post_send_all(sz, buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + while (reduce_chunk < total_chunks) { + // Poll the hardware for completions. + // + // If a send was completed mark how many completions we have received + // for that buffer. If we have sent the buffer to all peers we can + // reuse the send buffer so copy the next chunk of data into it and + // post the next send. local_staging is refilled later in the reduce + // loop, since the chunk in it may still be needed for the own-rank + // reduction step. + // + // If a receive is completed then advance the pointer of completed + // receives. + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + int64_t elems = std::min(N, total - read_offset); + std::copy( + in + read_offset, + in + read_offset + elems, + send_buffer(sz, buff).begin()); + post_send_all(sz, buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + else if (work_type == RECV_WR) { + recv_end[rank]++; + } + } + + // Process the received chunks in order. + // + // Rank 0 is always copied as is. Our rank is always read from the + // staging area. + while (reduce_chunk < total_chunks) { + // w is our write location so break if it is ahead of the read location. + int64_t w = static_cast(reduce_chunk) * N; + if (w >= read_offset) { + break; + } + // We want to reduce the 'reduce_chunk' chunk but it hasn't arrived + // yet. + if (recv_end[reduce_rank] <= reduce_chunk) { + break; + } + int b = reduce_chunk % PIPELINE; + int64_t elems = std::min(N, total - w); + + // Data is read from the staging area + if (reduce_rank == rank_) { + if (reduce_rank == 0) { + std::copy_n(local_staging(b), elems, out + w); + } else { + reduce_op(local_staging(b), out + w, elems); + } + + // Refill local_staging(b) for the next own-rank chunk and bump + // recv_end[rank_] to unblock its reduction. + int64_t next_local_chunk = + static_cast(reduce_chunk) + PIPELINE; + if (next_local_chunk < total_chunks) { + int64_t off = next_local_chunk * N; + int64_t next_elems = std::min(N, total - off); + std::copy(in + off, in + off + next_elems, local_staging(b)); + recv_end[rank_]++; + } + } + + // Data is read from the recv buffers + else { + if (reduce_rank == 0) { + std::copy_n( + recv_buffer(sz, b, reduce_rank).begin(), elems, out + w); + } else { + reduce_op( + recv_buffer(sz, b, reduce_rank).begin(), out + w, elems); + } + + // Check if we need to post another receive + int64_t next_chunk = static_cast(reduce_chunk) + PIPELINE; + if (next_chunk < total_chunks) { + recv_from(sz, reduce_rank, b); + in_flight++; + } + } + + // Means we processed that chunk so move to the next one + reduce_rank++; + if (reduce_rank >= size_) { + reduce_rank = 0; + reduce_chunk++; + } + } + } + + // Drain remaining in-flight completions (outstanding sends). + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + in_flight -= n; + } + } + + void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) { + // Copy our data to the appropriate place. Skip when in place (the scatter + // gather all reduce passes our own reduced shard which already lives at + // out_ptr + rank_ * n_bytes). + if (in_ptr != out_ptr + rank_ * n_bytes) { + std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); + } + + // Fully connected all gather + char* data = out_ptr; + char* our_data = out_ptr + rank_ * n_bytes; + auto [sz, N] = buffer_size_from_message(n_bytes); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; + int64_t total = static_cast(n_bytes); + int num_peers = size_ - 1; + + // Counters to maintain the state of transfers + int in_flight = 0; + int read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int write_offset[MESH_MAX_PEERS] = {0}; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + post_recv_all(sz, buff); + std::copy( + our_data + read_offset, + our_data + std::min(read_offset + N, total), + send_buffer(sz, buff).begin()); + post_send_all(sz, buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + // Send completed. If all sends completed then send the next chunk. + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + std::copy( + our_data + read_offset, + our_data + std::min(read_offset + N, total), + send_buffer(sz, buff).begin()); + post_send_all(sz, buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + // Recv completed. If we have more chunks then post another recv. + else if (work_type == RECV_WR) { + std::copy( + recv_buffer(sz, buff, rank).begin(), + recv_buffer(sz, buff, rank).begin() + + std::min(N, total - write_offset[rank]), + data + rank * n_bytes + write_offset[rank]); + write_offset[rank] += N; + if (write_offset[rank] + N * (PIPELINE - 1) < total) { + recv_from(sz, rank, buff); + in_flight++; + } + } + } + } + } + + template + void sum_scatter(const T* in, T* out, int64_t count, ReduceOp reduce_op) { + // Fully connected reduce scatter. + // + // The input holds size_ contiguous chunks of `count` elements each. Every + // rank keeps chunk rank_ and needs the elementwise reduction of that chunk + // across all ranks. To that end each rank p sends its chunk j to rank j and + // receives every peer's chunk rank_. The output is seeded with this rank's + // own chunk rank_ and every received chunk is reduced into it. + // + // Unlike all_reduce/all_gather each peer receives a *different* chunk, so + // we use the dedicated scatter buffers: per (sz, buff) tile there are size_ + // send slots (slot p -> peer p) and size_ recv slots (slot p <- peer p). + + const T* our_chunk = in + static_cast(rank_) * count; + + auto [sz, buffer_size] = buffer_size_from_message(count * sizeof(T)); + int64_t N = buffer_size / sizeof(T); + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; + int64_t total = static_cast(count); + int num_peers = size_ - 1; + + // Seed the output with our own chunk. + std::copy_n(our_chunk, total, out); + + // Counters to maintain the state of transfers + int in_flight = 0; + int64_t read_offset = 0; + int completed_send_count[PIPELINE] = {0}; + int64_t write_offset[MESH_MAX_PEERS] = {0}; + + // Prefill the pipeline + int buff = 0; + while (read_offset < total && buff < PIPELINE) { + int64_t elems = std::min(N, total - read_offset); + scatter_post_recv_all(sz, buff); + // Stage the chunk destined for each peer p (our input chunk p) into that + // peer's send buffer. + for (int p = 0; p < size_; p++) { + if (p == rank_) { + continue; + } + const T* src = in + static_cast(p) * count + read_offset; + std::copy( + src, src + elems, scatter_send_buffer(sz, buff, p).begin()); + } + scatter_post_send_all(sz, buff); + + buff++; + in_flight += 2 * num_peers; + read_offset += N; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll(connections_, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + // Send completed. Once every peer received this buffer, refill it with + // the next slice of each peer's chunk and post the next sends. + if (work_type == SEND_WR && read_offset < total) { + completed_send_count[buff]++; + if (completed_send_count[buff] == num_peers) { + int64_t elems = std::min(N, total - read_offset); + for (int p = 0; p < size_; p++) { + if (p == rank_) { + continue; + } + const T* src = in + static_cast(p) * count + read_offset; + std::copy( + src, + src + elems, + scatter_send_buffer(sz, buff, p).begin()); + } + scatter_post_send_all(sz, buff); + + completed_send_count[buff] = 0; + in_flight += num_peers; + read_offset += N; + } + } + + // Recv completed. Reduce the peer's contribution into our chunk and, if + // there is more data to fetch from that peer, post another recv. + else if (work_type == RECV_WR) { + int64_t elems = std::min(N, total - write_offset[rank]); + reduce_op( + scatter_recv_buffer(sz, buff, rank).begin(), + out + write_offset[rank], + elems); + write_offset[rank] += N; + if (write_offset[rank] + N * (PIPELINE - 1) < total) { + scatter_recv_from(sz, rank, buff); + in_flight++; + } + } + } + } + } + + template + void all_reduce_scatter_gather( + const T* in, + T* out, + int64_t count, + ReduceOp reduce_op) { + // Bandwidth optimal all reduce for large messages: a reduce scatter + // followed by an all gather. Compared to the fully connected all_reduce + // (which sends every rank's whole input to every peer) this moves size_x + // less data per link at the cost of an extra communication phase, so it is + // preferred for large messages where bandwidth dominates latency. + // + // The input is split into size_ equal chunks of `chunk` elements. The + // reduce scatter reduces chunk rank_ across all ranks and leaves it at + // out + rank_ * chunk. The all gather then reads that reduced shard in + // place and distributes every rank's reduced chunk into the output. No + // intermediate buffer is needed: the shard already lives where the all + // gather expects our contribution. + // + // Any trailing elements that do not divide evenly (fewer than size_) are + // handled with the fully connected all_reduce on the tail. + + int64_t chunk = count / size_; + int64_t base = chunk * size_; + + if (chunk > 0) { + // Reduce scatter our chunk directly into its final location in the + // output. sum_scatter only reads `in` and writes out + rank_ * chunk so + // this is safe even when in aliases out. + T* shard = out + static_cast(rank_) * chunk; + sum_scatter(in, shard, chunk, reduce_op); + + // All gather every rank's reduced chunk into the output. Our own shard is + // already in place so all_gather skips the self copy and only fills the + // other ranks' slices, which held the (already sent) input chunks. + all_gather( + reinterpret_cast(shard), + reinterpret_cast(out), + chunk * sizeof(T)); + } + + // Reduce the trailing elements (fewer than size_) with the fully connected + // all reduce so every rank ends up with the same tail. + if (base < count) { + all_reduce(in + base, out + base, count - base, reduce_op); + } + } + + void send(const char* in_ptr, int64_t n_bytes, int dst) { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + auto [sz, N] = buffer_size_from_message(n_bytes); + + int in_flight = 0; + int64_t read_offset = 0; + + // Prefill the pipeline + int buff = 0; + while (read_offset < n_bytes && buff < PIPELINE) { + std::copy( + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, n_bytes), + send_buffer(sz, buff).begin()); + send_to(sz, dst, buff); + + buff++; + read_offset += N; + in_flight++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a send was completed and we have more data to send then go ahead + // and send them. + ibv_wc wc[WC_NUM]; + int n = connections_[dst].poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + if (read_offset < n_bytes) { + std::copy( + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, n_bytes), + send_buffer(sz, buff).begin()); + send_to(sz, dst, buff); + + read_offset += N; + in_flight++; + } + } + } + } + + void recv(char* out_ptr, int64_t n_bytes, int src) { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + auto [sz, N] = buffer_size_from_message(n_bytes); + + int in_flight = 0; + int64_t write_offset = 0; + + // Prefill the pipeline + int buff = 0; + while (N * buff < n_bytes && buff < PIPELINE) { + recv_from(sz, src, buff); + + in_flight++; + buff++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a recv was completed copy it to the output and if we have more + // data to fetch post another recv. + ibv_wc wc[WC_NUM]; + int n = connections_[src].poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + int rank = wc[i].wr_id & 0xff; + + in_flight--; + + std::copy( + recv_buffer(sz, buff, src).begin(), + recv_buffer(sz, buff, src).begin() + + std::min(n_bytes - write_offset, static_cast(N)), + out_ptr + write_offset); + write_offset += N; + + if (write_offset + (PIPELINE - 1) * N < n_bytes) { + recv_from(sz, src, buff); + + in_flight++; + } + } + } + } + + private: + void send_to(int sz, int rank, int buff) { + connections_[rank].post_send( + send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank); + } + + void recv_from(int sz, int rank, int buff) { + connections_[rank].post_recv( + recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank); + } + + SharedBuffer& send_buffer(int sz, int buff) { + return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank_]; + } + + SharedBuffer& recv_buffer(int sz, int buff, int rank) { + return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank]; + } + + void post_send_all(int sz, int buff) { + auto& b = send_buffer(sz, buff); + int wr_id = SEND_WR << 16 | buff << 8; + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + connections_[i].post_send(b, wr_id | i); + } + } + + void post_recv_all(int sz, int buff) { + int b = sz * NUM_BUFFERS * size_ + buff * size_; + int wr_id = RECV_WR << 16 | buff << 8; + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + connections_[i].post_recv(buffers_[b + i], wr_id | i); + } + } + + // Scatter buffer helpers. Per (sz, buff) tile there are 2 * size_ slots: the + // first size_ are send buffers (slot p is sent to peer p) and the next size_ + // are recv buffers (slot p receives from peer p). + SharedBuffer& scatter_send_buffer(int sz, int buff, int peer) { + return scatter_buffers_ + [sz * NUM_BUFFERS * 2 * size_ + buff * 2 * size_ + peer]; + } + + SharedBuffer& scatter_recv_buffer(int sz, int buff, int peer) { + return scatter_buffers_ + [sz * NUM_BUFFERS * 2 * size_ + buff * 2 * size_ + size_ + peer]; + } + + void scatter_send_to(int sz, int rank, int buff) { + connections_[rank].post_send( + scatter_send_buffer(sz, buff, rank), SEND_WR << 16 | buff << 8 | rank); + } + + void scatter_recv_from(int sz, int rank, int buff) { + connections_[rank].post_recv( + scatter_recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank); + } + + void scatter_post_send_all(int sz, int buff) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + scatter_send_to(sz, i, buff); + } + } + + void scatter_post_recv_all(int sz, int buff) { + for (int i = 0; i < size_; i++) { + if (i == rank_) { + continue; + } + scatter_recv_from(sz, i, buff); + } + } + + int rank_; + int size_; + std::span connections_; + std::span buffers_; + std::span scatter_buffers_; + std::unique_ptr staging_mem_; +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/utils.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.cpp similarity index 78% rename from Source/Cmlx/mlx/mlx/distributed/jaccl/utils.cpp rename to Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.cpp index dc87d428c..f449d8c81 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/utils.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.cpp @@ -2,10 +2,12 @@ #include #include +#include #include #include +#include -#include "mlx/distributed/jaccl/utils.h" +#include "jaccl/rdma.h" #define LOAD_SYMBOL(symbol, variable) \ { \ @@ -31,7 +33,7 @@ void* page_aligned_alloc(size_t num_bytes) { } // namespace -namespace mlx::core::distributed::jaccl { +namespace jaccl { IBVWrapper::IBVWrapper() { librdma_handle_ = dlopen("librdma.dylib", RTLD_NOW | RTLD_GLOBAL); @@ -170,7 +172,12 @@ void Connection::create_queue_pair() { queue_pair = ibv().create_qp(protection_domain, &init_attr); if (queue_pair == nullptr) { - throw std::runtime_error("[jaccl] Couldn't create queue pair"); + int err = errno; + std::string error_message = std::generic_category().message(err); + std::ostringstream msg; + msg << "[jaccl] Creating the queue pair failed with '" << error_message + << " (" << errno << ")'."; + throw std::runtime_error(msg.str()); } } @@ -181,12 +188,33 @@ const Destination& Connection::info() { ibv_port_attr port_attr; ibv().query_port(ctx, 1, &port_attr); - ibv_gid gid; - ibv().query_gid(ctx, 1, 1, &gid); + ibv_gid gid = {}; + bool found_gid = false; + for (int i = 0; i < port_attr.gid_tbl_len; i++) { + ibv_gid tmp; + if (ibv().query_gid(ctx, 1, i, &tmp) == 0) { + if (*(uint64_t*)&tmp.raw[0] == 0 && *(uint16_t*)&tmp.raw[8] == 0 && + *(uint16_t*)&tmp.raw[10] == 0xffff) { + gid = tmp; + found_gid = true; + break; + } + } + } + + // Fail here rather than hand an unset GID to the queue pair. + if (!found_gid) { + std::ostringstream msg; + msg << "[jaccl] No IPv4-mapped GID for this device. Thunderbolt RDMA ports " + << "only publish one once the interface has an IPv4 address; assign a " + << "link-local address to it, for example: ifconfig inet " + << "169.254.0.1 netmask 255.255.0.0 alias"; + throw std::runtime_error(msg.str()); + } src.local_id = port_attr.lid; src.queue_pair_number = queue_pair->qp_num; - src.packet_sequence_number = 7; // TODO: Change to sth random + src.packet_sequence_number = 7; src.global_identifier = gid; return src; @@ -267,6 +295,7 @@ std::vector create_connections( } // Search for the name and try to open the device + bool found = false; for (int i = 0; i < num_devices; i++) { if (name == ibv().get_device_name(devices[i])) { auto ctx = ibv().open_device(devices[i]); @@ -276,21 +305,31 @@ std::vector create_connections( throw std::runtime_error(msg.str()); } connections.emplace_back(ctx); + found = true; break; } } + + // Returning a shorter vector than device_names would leave the callers, + // which size themselves from device_names, indexing past its end. + if (!found) { + std::ostringstream msg; + msg << "[jaccl] Could not find device " << name << " (" << num_devices + << " available)"; + throw std::runtime_error(msg.str()); + } } ibv().free_device_list(devices); return connections; } -SideChannel::SideChannel(int rank, int size, const char* addr) +TCPAllGather::TCPAllGather(int rank, int size, const char* addr) : rank_(rank), size_(size) { - auto address = detail::parse_address(addr); + auto address = parse_address(addr); if (rank_ == 0) { - detail::TCPSocket server(IBV_TAG); + TCPSocket server(IBV_TAG); server.listen(IBV_TAG, address); for (int i = 0; i < size - 1; i++) { @@ -311,7 +350,7 @@ SideChannel::SideChannel(int rank, int size, const char* addr) } } else { sockets_.push_back( - detail::TCPSocket::connect( + TCPSocket::connect( IBV_TAG, address, 4, 1000, [](int attempt, int wait) { std::cerr << IBV_TAG << " Connection attempt " << attempt << " waiting " << wait << " ms" << std::endl; @@ -320,10 +359,31 @@ SideChannel::SideChannel(int rank, int size, const char* addr) } } +void TCPAllGather::operator()(const char* src, char* dst, size_t n_bytes) { + std::lock_guard lock(mutex_); + if (rank_ == 0) { + std::copy(src, src + n_bytes, dst); + for (int i = 1; i < size_; i++) { + sockets_[i - 1].recv(IBV_TAG, dst + i * n_bytes, n_bytes); + } + for (int i = 1; i < size_; i++) { + sockets_[i - 1].send(IBV_TAG, dst, size_ * n_bytes); + } + } else { + sockets_[0].send(IBV_TAG, src, n_bytes); + sockets_[0].recv(IBV_TAG, dst, size_ * n_bytes); + } +} + +SideChannel::SideChannel(int rank, int size, AllGatherFn agf) + : rank_(rank), size_(size), all_gather_fn_(std::move(agf)) {} + SideChannel::SideChannel(SideChannel&& sc) - : rank_(sc.rank_), size_(sc.size_), sockets_(std::move(sc.sockets_)) { + : rank_(sc.rank_), + size_(sc.size_), + all_gather_fn_(std::move(sc.all_gather_fn_)) { sc.rank_ = -1; sc.size_ = -1; } -} // namespace mlx::core::distributed::jaccl +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/utils.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.h similarity index 76% rename from Source/Cmlx/mlx/mlx/distributed/jaccl/utils.h rename to Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.h index 5b700f49b..01d9ed37d 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/utils.h +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/rdma.h @@ -4,10 +4,14 @@ #include +#include +#include +#include +#include #include #include -#include "mlx/distributed/utils.h" +#include "jaccl/tcp.h" constexpr const char* IBV_TAG = "[jaccl]"; constexpr int SEND_WR = 1; @@ -18,8 +22,6 @@ constexpr int BUFFER_SIZES = 8; constexpr int NUM_BUFFERS = 2; constexpr int FRAME_SIZE = 4096; -namespace detail = mlx::core::distributed::detail; - namespace { template @@ -44,7 +46,7 @@ inline std::pair buffer_size_from_message(int64_t msg) { } // namespace -namespace mlx::core::distributed::jaccl { +namespace jaccl { /** * Wrapper for the ibverbs API. @@ -221,7 +223,7 @@ std::vector create_connections( const std::vector& device_names); inline int poll( - const std::vector& connections, + std::span connections, int num_completions, ibv_wc* work_completions) { int completions = 0; @@ -244,8 +246,8 @@ inline int poll( } inline int poll( - const std::vector& connections_1, - const std::vector& connections_2, + std::span connections_1, + std::span connections_2, int num_completions, ibv_wc* work_completions) { int completions = 0; @@ -257,6 +259,36 @@ inline int poll( return completions; } +/** + * A function that performs an all-gather across ranks. + * + * 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[r * n_bytes, (r+1) * n_bytes] contains the data from rank r. + * n_bytes: The number of bytes contributed by each rank. + */ +using AllGatherFn = + std::function; + +class TCPAllGather { + public: + TCPAllGather(int rank, int size, const char* addr); + + TCPAllGather(const TCPAllGather&) = delete; + TCPAllGather(TCPAllGather&&) = delete; + TCPAllGather& operator=(const TCPAllGather&) = delete; + TCPAllGather& operator=(TCPAllGather&&) = delete; + + void operator()(const char* src, char* dst, size_t n_bytes); + + private: + int rank_; + int size_; + std::vector sockets_; + std::mutex mutex_; +}; + /** * Implement a TCP side channel to exchange information about the RDMA * connections. @@ -266,7 +298,7 @@ inline int poll( */ class SideChannel { public: - SideChannel(int rank, int size, const char* addr); + SideChannel(int rank, int size, AllGatherFn agf); SideChannel(SideChannel&& sc); SideChannel(const SideChannel&) = delete; @@ -280,63 +312,52 @@ class SideChannel { if constexpr (is_container::value) { using U = typename T::value_type; - // Share the lengths first and set the communication size to be the + // Share the lengths first to set the communication size to be the // maximum length of the containers. auto lengths = all_gather(v.size()); auto max_len = *std::max_element(lengths.begin(), lengths.end()); - for (auto& s : result) { - s.resize(max_len); - } - // All gather of length max_len - if (rank_ == 0) { - std::copy(v.begin(), v.end(), result[rank_].begin()); - for (int i = 1; i < size_; i++) { - sockets_[i - 1].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); - } - for (int i = 1; i < size_; i++) { - for (int j = 0; j < size_; j++) { - sockets_[i - 1].send( - IBV_TAG, result[j].data(), sizeof(U) * max_len); - } - } - } else { - std::copy(v.begin(), v.end(), result[rank_].begin()); - sockets_[0].send(IBV_TAG, result[rank_].data(), sizeof(U) * max_len); - for (int i = 0; i < size_; i++) { - sockets_[0].recv(IBV_TAG, result[i].data(), sizeof(U) * max_len); - } - } + // Allocate flat memory for the all gather + std::vector buffer; + buffer.resize(max_len * (size_ + 1)); + + // Copy our value and share it + std::copy(v.begin(), v.end(), buffer.begin()); + all_gather_fn_( + reinterpret_cast(&buffer[0]), + reinterpret_cast(&buffer[max_len]), + max_len * sizeof(U)); - // Resize the outputs back to the original length + // Put the values into the individual containers for (int i = 0; i < size_; i++) { - result[i].resize(lengths[i]); + std::copy( + buffer.begin() + (i + 1) * max_len, + buffer.begin() + (i + 1) * max_len + lengths[i], + std::inserter(result[i], result[i].end())); } } // T is a scalar else { - if (rank_ == 0) { - result[rank_] = v; - for (int i = 1; i < size_; i++) { - sockets_[i - 1].recv(IBV_TAG, &result[i], sizeof(T)); - } - for (int i = 1; i < size_; i++) { - sockets_[i - 1].send(IBV_TAG, result.data(), size_ * sizeof(T)); - } - } else { - sockets_[0].send(IBV_TAG, &v, sizeof(T)); - sockets_[0].recv(IBV_TAG, result.data(), size_ * sizeof(T)); - } + all_gather_fn_( + reinterpret_cast(&v), + reinterpret_cast(result.data()), + sizeof(T)); } return result; } + void barrier() { + // Twice has proven to be more robust to initialization issues. + all_gather(0); + all_gather(0); + } + private: int rank_; int size_; - std::vector sockets_; + AllGatherFn all_gather_fn_; }; -} // namespace mlx::core::distributed::jaccl +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h new file mode 100644 index 000000000..4dd395539 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/reduction_ops.h @@ -0,0 +1,217 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include + +#include "jaccl/types.h" + +namespace jaccl { + +// Each reduction op has an in place form out[i] OP= in[i] and an out of place +// form out[i] = a[i] OP b[i]. The out of place pointers are __restrict, so +// callers must only use it when a, b and output are distinct buffers. + +template +struct SumOp { + void operator()(const T* input, T* output, size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = output[i] + input[i]; + } + } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = a[i] + b[i]; + } + } +}; + +template +struct MaxOp { + void operator()(const T* input, T* output, size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (output[i] > input[i]) ? output[i] : input[i]; + } + } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] > b[i]) ? a[i] : b[i]; + } + } +}; + +template +struct MinOp { + void operator()(const T* input, T* output, size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (output[i] < input[i]) ? output[i] : input[i]; + } + } + void operator()( + const T* __restrict a, + const T* __restrict b, + T* __restrict output, + size_t N) const { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] < b[i]) ? a[i] : b[i]; + } + } +}; + +// +// The last piece of the puzzle to use the native bf16 while compiling a single +// binary for all Macs is to compile these functions with +// target("arch=armv8.6-a"). +// +// Now we can simply check in runtime and call them only when they are +// supported. +// + +#if defined(__aarch64__) + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_sum(const void* input, void* output, size_t N) { + auto in = reinterpret_cast(input); + auto out = reinterpret_cast<__bf16*>(output); + for (size_t i = 0; i < N; i++) { + out[i] = out[i] + in[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_max(const void* input, void* output, size_t N) { + auto in = reinterpret_cast(input); + auto out = reinterpret_cast<__bf16*>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (out[i] > in[i]) ? out[i] : in[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_min(const void* input, void* output, size_t N) { + auto in = reinterpret_cast(input); + auto out = reinterpret_cast<__bf16*>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (out[i] < in[i]) ? out[i] : in[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_sum(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = pa[i] + pb[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_max(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (pa[i] > pb[i]) ? pa[i] : pb[i]; + } +} + +__attribute__((target("arch=armv8.6-a"))) inline void +native_bf16_min(const void* a, const void* b, void* output, size_t N) { + auto pa = reinterpret_cast(a); + auto pb = reinterpret_cast(b); + auto out = reinterpret_cast<__bf16* __restrict>(output); + for (size_t i = 0; i < N; i++) { + out[i] = (pa[i] < pb[i]) ? pa[i] : pb[i]; + } +} + +template <> +struct SumOp { + void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const { + if (has_native_bf16_support()) { + native_bf16_sum(input, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = output[i] + input[i]; + } + } + } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_sum(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = a[i] + b[i]; + } + } + } +}; + +template <> +struct MaxOp { + void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const { + if (has_native_bf16_support()) { + native_bf16_max(input, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (output[i] > input[i]) ? output[i] : input[i]; + } + } + } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_max(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] > b[i]) ? a[i] : b[i]; + } + } + } +}; + +template <> +struct MinOp { + void operator()(const bfloat16_t* input, bfloat16_t* output, size_t N) const { + if (has_native_bf16_support()) { + native_bf16_min(input, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (output[i] < input[i]) ? output[i] : input[i]; + } + } + } + void operator()( + const bfloat16_t* __restrict a, + const bfloat16_t* __restrict b, + bfloat16_t* __restrict output, + size_t N) const { + if (has_native_bf16_support()) { + native_bf16_min(a, b, output, N); + } else { + for (size_t i = 0; i < N; i++) { + output[i] = (a[i] < b[i]) ? a[i] : b[i]; + } + } + } +}; + +#endif // defined(__aarch64__) + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.cpp new file mode 100644 index 000000000..a08035f3f --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.cpp @@ -0,0 +1,258 @@ +// Copyright © 2026 Apple Inc. + +#include "jaccl/ring.h" +#include "jaccl/reduction_ops.h" +#include "jaccl/types.h" + +namespace jaccl { + +RingGroup::RingGroup( + int rank, + int size, + const std::vector& left_devices, + const std::vector& right_devices, + SideChannel sc) + : rank_(rank), + size_(size), + n_conns_(left_devices.size()), + side_channel_(std::move(sc)), + left_(create_connections(left_devices)), + right_(create_connections(right_devices)), + pool_(n_conns_ > 0 ? n_conns_ - 1 : 0) { + if (left_.size() > RING_MAX_CONNS || right_.size() > RING_MAX_CONNS) { + std::ostringstream msg; + msg << "[jaccl] Up to " << RING_MAX_CONNS << " per direction supported but " + << left_.size() << " were provided."; + throw std::runtime_error(msg.str()); + } + + // Initialize all the connections and allocate buffers + initialize(); + + // Make sure every node has reached here before continuing. + side_channel_.barrier(); + + // Create the ring implementation object + ring_ = RingImpl( + rank_, size_, left_, right_, send_buffers_, recv_buffers_, &pool_); +} + +void RingGroup::initialize() { + // Create the queue pairs + for (auto& conn : left_) { + conn.allocate_protection_domain(); + conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_queue_pair(); + } + for (auto& conn : right_) { + conn.allocate_protection_domain(); + conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); + conn.create_queue_pair(); + } + + // Allocate the buffers + allocate_buffers(); + + // Initialize the conections + for (auto& conn : left_) { + conn.queue_pair_init(); + } + for (auto& conn : right_) { + conn.queue_pair_init(); + } + + // Gather the information to be exchanged, this also serves as a barrier + // so that all peers have initialized their connections before attempting + // to transition to RTS. + std::vector left_info; + for (auto& conn : left_) { + left_info.emplace_back(conn.info()); + } + std::vector right_info; + for (auto& conn : right_) { + right_info.emplace_back(conn.info()); + } + auto all_left_infos = side_channel_.all_gather(left_info); + auto all_right_infos = side_channel_.all_gather(right_info); + + // Transition queue pairs to RTS + int left_peer = (rank_ + size_ - 1) % size_; + for (int i = 0; i < left_.size(); i++) { + auto peer_info = all_right_infos[left_peer][i]; + left_[i].queue_pair_rtr(peer_info); + left_[i].queue_pair_rts(); + } + int right_peer = (rank_ + 1) % size_; + for (int i = 0; i < right_.size(); i++) { + auto peer_info = all_left_infos[right_peer][i]; + right_[i].queue_pair_rtr(peer_info); + right_[i].queue_pair_rts(); + } +} + +void RingGroup::allocate_buffers() { + // Deregister any buffers and free the memory + send_buffers_.clear(); + recv_buffers_.clear(); + + // Allocate the memory + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + for (int j = 0; j < n_conns_ * 2; j++) { + send_buffers_.emplace_back(FRAME_SIZE * (1 << k)); + recv_buffers_.emplace_back(FRAME_SIZE * (1 << k)); + } + } + } + + // Register the buffers with the corresponding connections + for (int k = 0; k < BUFFER_SIZES; k++) { + for (int i = 0; i < NUM_BUFFERS; i++) { + for (int j = 0; j < n_conns_ * 2; j++) { + int wire = j % n_conns_; + int lr = j / n_conns_; + if (lr) { + send_buffers_[k * NUM_BUFFERS * n_conns_ * 2 + i * n_conns_ * 2 + j] + .register_to_protection_domain(left_[wire].protection_domain); + recv_buffers_[k * NUM_BUFFERS * n_conns_ * 2 + i * n_conns_ * 2 + j] + .register_to_protection_domain(right_[wire].protection_domain); + } else { + send_buffers_[k * NUM_BUFFERS * n_conns_ * 2 + i * n_conns_ * 2 + j] + .register_to_protection_domain(right_[wire].protection_domain); + recv_buffers_[k * NUM_BUFFERS * n_conns_ * 2 + i * n_conns_ * 2 + j] + .register_to_protection_domain(left_[wire].protection_domain); + } + } + } + } +} + +void RingGroup::all_sum( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, SumOp{}); + }); +} + +void RingGroup::all_max( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, MaxOp{}); + }); +} + +void RingGroup::all_min( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + all_reduce(input, output, n_bytes, MinOp{}); + }); +} + +void RingGroup::all_gather(const void* input, void* output, size_t n_bytes) { + ring_.all_gather( + static_cast(input), + static_cast(output), + n_bytes, + n_conns_); +} + +void RingGroup::sum_scatter( + const void* input, + void* output, + size_t n_bytes, + int dtype) { + dispatch_all_types(dtype, [&](auto type_tag) { + using T = JACCL_GET_TYPE(type_tag); + reduce_scatter(input, output, n_bytes, SumOp{}); + }); +} + +void RingGroup::send(const void* input, size_t n_bytes, int dst) { + int right = (rank_ + 1) % size_; + int left = (rank_ + size_ - 1) % size_; + if (dst != right && dst != left) { + std::ostringstream msg; + msg << "[jaccl] In ring mode send is only supported to direct neighbors " + << "but tried to send to " << dst << " from " << rank_ << std::endl; + throw std::runtime_error(msg.str()); + } + ring_.send(static_cast(input), n_bytes, dst, n_conns_); +} + +void RingGroup::recv(void* output, size_t n_bytes, int src) { + int right = (rank_ + 1) % size_; + int left = (rank_ + size_ - 1) % size_; + if (src != right && src != left) { + std::ostringstream msg; + msg << "[jaccl] In ring mode recv is only supported to direct neighbors " + << "but tried to recv from " << src << " to " << rank_ << std::endl; + throw std::runtime_error(msg.str()); + } + ring_.recv(static_cast(output), n_bytes, src, n_conns_); +} + +void RingGroup::barrier() { + uint8_t b = 0; + all_sum(&b, &b, sizeof(b), Dtype::UInt8); +} + +template +void RingGroup::all_reduce( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + + if (n_bytes <= 32768 || count < size_ * 2 * n_conns_) { + ring_.all_reduce<1, T, ReduceOp>(in_ptr, out_ptr, count, 1, reduce_op); + return; + } + + ring_.all_reduce<2, T, ReduceOp>(in_ptr, out_ptr, count, n_conns_, reduce_op); +} + +template +void RingGroup::reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op) { + // n_bytes is the size of the output (one chunk). The input holds size_ such + // chunks laid out contiguously so the full element count is size_ * count. + auto in_ptr = static_cast(input); + auto out_ptr = static_cast(output); + int64_t count = n_bytes / sizeof(T); + if (size_ == 1) { + std::copy_n(in_ptr, count, out_ptr); + return; + } + + int64_t total = static_cast(size_) * count; + + // Mirror the all_reduce heuristics: use a single wire for small messages and + // scale up to n_conns_ wires for large ones where bandwidth dominates. + if (total < size_ * 2 * n_conns_ || n_bytes <= 65536) { + ring_.reduce_scatter(in_ptr, out_ptr, total, 1, reduce_op); + return; + } + + ring_.reduce_scatter( + in_ptr, out_ptr, total, n_conns_, reduce_op); +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.h new file mode 100644 index 000000000..c901c15a8 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring.h @@ -0,0 +1,96 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include "jaccl/group.h" +#include "jaccl/rdma.h" +#include "jaccl/ring_impl.h" + +namespace jaccl { + +/** + * The JACCL communication group for a ring where each node is connected to its + * two neighboring nodes. It should be the highest bandwidth communication + * group for large messages when many connections per peer are used. + * + * Like all JACCL groups it uses a side channel to exchange the necessary + * information and then configure the connections to be ready for RDMA + * operations. + */ +class RingGroup : public Group { + public: + RingGroup( + int rank, + int size, + const std::vector& left_devices, + const std::vector& right_devices, + SideChannel sc); + + int rank() override { + return rank_; + } + + int size() override { + return size_; + } + + void all_sum(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_max(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_min(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void all_gather(const void* input, void* output, size_t n_bytes) override; + + void sum_scatter(const void* input, void* output, size_t n_bytes, int dtype) + override; + + void send(const void* input, size_t n_bytes, int dst) override; + void recv(void* output, size_t n_bytes, int src) override; + + void barrier() override; + + private: + template + void all_reduce( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + + template + void reduce_scatter( + const void* input, + void* output, + size_t n_bytes, + ReduceOp reduce_op); + + /** + * Performs the connection initialization. Namely, after this call all + * Connection objects should have a queue pair in RTS state and all buffers + * should have been allocated. + */ + void initialize(); + + /** + * Allocate all the buffers that we will use in the communication group. + */ + void allocate_buffers(); + + int rank_; + int size_; + int n_conns_; + SideChannel side_channel_; + std::vector left_; + std::vector right_; + std::vector send_buffers_; + std::vector recv_buffers_; + // Declared before ring_ so it outlives the RingImpl that points at it. + ThreadPool pool_; + RingImpl ring_; +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h new file mode 100644 index 000000000..15b81a217 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -0,0 +1,814 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include + +#include "jaccl/rdma.h" +#include "jaccl/threadpool.h" + +constexpr int RING_MAX_CONNS = 4; + +namespace jaccl { + +class RingImpl { + public: + RingImpl( + int rank, + int size, + std::vector& left, + std::vector& right, + std::vector& send_buffers, + std::vector& recv_buffers, + ThreadPool* pool = nullptr) + : rank_(rank), + size_(size), + n_conns_(left.size()), + left_(left), + right_(right), + send_buffers_(send_buffers), + recv_buffers_(recv_buffers), + pool_(pool) {} + + RingImpl( + int rank, + int size, + Connection* left_begin, + Connection* right_begin, + size_t n_conns, + std::vector& send_buffers, + std::vector& recv_buffers, + ThreadPool* pool = nullptr) + : rank_(rank), + size_(size), + n_conns_(n_conns), + left_(left_begin, n_conns), + right_(right_begin, n_conns), + send_buffers_(send_buffers), + recv_buffers_(recv_buffers), + pool_(pool) {} + + RingImpl() : rank_(0), size_(1), n_conns_(0), pool_(nullptr) {} + + // Copy the received chunk into the output. Used by the all gather passes. + struct CopyOp { + template + inline void operator()(const T* recv, const T*, T* out, int64_t n) const { + std::copy(recv, recv + std::max(0, n), out); + } + }; + + // Reduce the received chunk with this rank's own input for the chunk. When + // in place (input aliases output) we use the fast two argument kernel; + // otherwise the fused out of place kernel seeds the output from the input. + template + struct ReduceRecvOp { + ReduceOp reduce_op; + template + inline void operator()(const T* recv, const T* base, T* out, int64_t n) + const { + n = std::max(0, n); + if constexpr (INPLACE) { + reduce_op(recv, out, n); + } else { + reduce_op(base, recv, out, n); + } + } + }; + + template + void all_reduce( + const T* in_ptr, + T* out_ptr, + int64_t size, + int n_wires, + ReduceOp reduce_op) { + int64_t chunk_size = (size + size_ - 1) / size_; + int64_t size_per_wire = + (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires); + + // Split the reduce scatter + all gather across the available wires. Each + // wire handles a contiguous slice of each chunk in every direction. + dispatch_wires(n_wires, [&](int lw) { + all_reduce_wire( + in_ptr, + out_ptr, + size, + chunk_size, + size_per_wire, + n_wires, + lw, + reduce_op); + }); + } + + // Perform the ring all reduce (reduce scatter followed by all gather) for a + // single wire lw. + // + // Every chunk of chunk_size elements is divided into MAX_DIR directional + // regions and each region is further split into n_wires contiguous slices of + // size_per_wire elements. This function is responsible for slice lw in every + // direction and only touches left_[lw] / right_[lw] and the buffers that + // belong to wire lw, so several wires can run concurrently. + template + void all_reduce_wire( + const T* in_ptr, + T* out_ptr, + int64_t size, + int64_t chunk_size, + int64_t size_per_wire, + int n_wires, + int lw, + ReduceOp reduce_op) { + // This wire's slice within the chunk: [wire_offset, wire_end). Clamp to the + // per wire end (not the whole region) so the last frame can't spill into + // the next wire's slice when size_per_wire is not a multiple of N. + int64_t wire_offset[MAX_DIR]; + int64_t wire_end[MAX_DIR]; + int64_t send_offset[MAX_DIR]; + int64_t recv_offset[MAX_DIR]; + for (int lr = 0; lr < MAX_DIR; lr++) { + wire_offset[lr] = lr * n_wires * size_per_wire + + static_cast(lw) * size_per_wire; + int64_t region_end = + std::min(chunk_size, (lr + 1) * n_wires * size_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + size_per_wire); + send_offset[lr] = rank_ * chunk_size; + } + recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; + if constexpr (MAX_DIR == 2) { + recv_offset[1] = ((rank_ + 1) % size_) * chunk_size; + } + + // First reduce scatter, then all gather. The offsets carry over from the + // reduce scatter into the all gather. Pick the in place vs out of place + // reduction kernel once so the hot loop stays branch free. + if (in_ptr == out_ptr) { + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + wire_end, + send_offset, + recv_offset, + in_ptr, + out_ptr, + ReduceRecvOp{reduce_op}); + } else { + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + wire_end, + send_offset, + recv_offset, + in_ptr, + out_ptr, + ReduceRecvOp{reduce_op}); + } + ring_pass( + lw, + size, + chunk_size, + size_ * chunk_size, + size_per_wire, + wire_offset, + wire_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); + } + + void + all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes, int n_wires) { + // Copy our data to the appropriate place + std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); + + // Split the all gather across the available wires. Each wire handles a + // contiguous slice of every rank's data in both directions. + size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); + dispatch_wires(n_wires, [&](int lw) { + all_gather_wire( + out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + }); + } + + // Perform the ring all gather for a single wire lw. + // + // The wire is responsible for the contiguous slice + // [lw * n_bytes_per_wire, (lw + 1) * n_bytes_per_wire) + // of every rank's n_bytes region, sent in the left direction (lr == 0), and + // the mirrored slice sent in the right direction (lr == 1). + // + // This function only ever touches left_[lw] / right_[lw] and the buffers that + // belong to wire lw so several wires can run concurrently. + void all_gather_wire( + char* out_ptr, + int64_t n_bytes, + int64_t n_bytes_per_wire, + int lw) { + // Both directions send the same contiguous slice of each rank's region. + int64_t slice = static_cast(lw) * n_bytes_per_wire; + int64_t wire_offset[2] = {slice, slice}; + // Clamp to this wire's slice end so the last frame can't spill into the + // next wire's slice. + int64_t wire_end_bytes = std::min(n_bytes, slice + n_bytes_per_wire); + int64_t wire_end[2] = {wire_end_bytes, wire_end_bytes}; + int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; + int64_t recv_offset[2] = { + ((rank_ + size_ - 1) % size_) * n_bytes, + ((rank_ + 1) % size_) * n_bytes}; + + ring_pass<2, char>( + lw, + n_bytes, + n_bytes, + n_bytes * size_, + n_bytes_per_wire, + wire_offset, + wire_end, + send_offset, + recv_offset, + out_ptr, + out_ptr, + CopyOp{}); + } + + // Standalone ring reduce scatter. + // + // The input holds size_ contiguous chunks of `chunk = size / size_` elements + // (size is guaranteed divisible by size_ for a sum_scatter). Rank r produces + // the full reduction of chunk r in out_ptr (chunk elements). Compared to the + // reduce scatter phase embedded in the ring all_reduce this converges every + // rank on its *own* chunk index so it can be exposed directly, and it writes + // straight into the chunk sized output using it as the rolling accumulator. + template + void reduce_scatter( + const T* in_ptr, + T* out_ptr, + int64_t size, + int n_wires, + ReduceOp reduce_op) { + int64_t chunk = size / size_; + // Two directional regions, each split across the wires. + int64_t size_per_wire = (chunk + (2 * n_wires) - 1) / (2 * n_wires); + + dispatch_wires(n_wires, [&](int lw) { + reduce_scatter_wire( + in_ptr, out_ptr, chunk, size_per_wire, n_wires, lw, reduce_op); + }); + } + + // Perform the dual direction ring reduce scatter for a single wire lw. + // + // The output chunk is divided into 2 directional regions and each region is + // split into n_wires contiguous slices of size_per_wire elements. Direction 0 + // flows left (send right / recv left), direction 1 flows right. Both + // directions converge on chunk rank_, so together they fill the whole output + // chunk. This wire owns slice lw of every region and only touches + // left_[lw] / right_[lw]. + // + // The rolling partial for each direction lives in the chunk sized output + // itself: every step sends the partial produced by the previous step and the + // recv writes the next partial into the same location. To avoid clobbering a + // partial that is still being staged for sending, a recv slice is only + // reduced into the output once the matching send slice has been staged (i.e. + // recv_count[lr] < send_count[lr]). + template + void reduce_scatter_wire( + const T* in_ptr, + T* out_ptr, + int64_t chunk, + int64_t size_per_wire, + int n_wires, + int lw, + ReduceOp reduce_op) { + constexpr int MAX_DIR = 2; + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * 2 * MAX_DIR; + + auto [sz, buffer_bytes] = + buffer_size_from_message(size_per_wire * sizeof(T)); + int64_t N = buffer_bytes / sizeof(T); + int64_t total = static_cast(size_) * chunk; + + // This wire's slice within the chunk: [wire_offset, wire_end). Clamp to the + // per wire end (not the whole region) so the last frame can't spill into + // the next wire's slice when size_per_wire is not a multiple of N. + int64_t wire_offset[MAX_DIR]; + int64_t wire_end[MAX_DIR]; + for (int lr = 0; lr < MAX_DIR; lr++) { + wire_offset[lr] = lr * n_wires * size_per_wire + + static_cast(lw) * size_per_wire; + int64_t region_end = std::min(chunk, (lr + 1) * n_wires * size_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + size_per_wire); + } + + // Input windows (count space, chunk aligned). Both directions converge on + // chunk rank_ (verified: rank r finishes chunk r). Direction 0 flows left + // and rotates backward; direction 1 flows right and rotates forward. + int64_t in_send_offset[MAX_DIR]; + int64_t in_recv_offset[MAX_DIR]; + in_send_offset[0] = ((rank_ - 1 + size_) % size_) * chunk; + in_recv_offset[0] = ((rank_ - 2 + 2 * size_) % size_) * chunk; + in_send_offset[1] = ((rank_ + 1) % size_) * chunk; + in_recv_offset[1] = ((rank_ + 2) % size_) * chunk; + + int64_t n_steps = (size_per_wire + N - 1) / N; + + int in_flight = 0; + int send_count[MAX_DIR] = {0}; + int recv_count[MAX_DIR] = {0}; + // Recv completions that arrived before their matching send slice was staged + // are deferred until the send catches up so we never overwrite a partial + // that is still being read for sending. + int deferred_recv[MAX_DIR] = {0}; + + // Apply as many deferred recv slices for direction lr as are now unblocked + // by staged sends. Recv completions arrive in FIFO order per queue pair, so + // recv slice recv_count[lr] lives in pipeline buffer recv_count[lr] % + // PIPELINE. + auto drain_deferred = [&](int lr) { + while (deferred_recv[lr] > 0 && recv_count[lr] < send_count[lr]) { + int slice = recv_count[lr]; + int b = slice % PIPELINE; + int64_t offset = wire_offset[lr] + static_cast(slice) * N; + int64_t n = std::min(N, wire_end[lr] - offset); + reduce_op( + recv_buffer(sz, b, lr, lw).template begin(), + in_ptr + in_recv_offset[lr] + offset, + out_ptr + offset, + std::max(0, n)); + recv_count[lr]++; + deferred_recv[lr]--; + if (recv_count[lr] + (PIPELINE - 1) < n_steps) { + recv_from(sz, b, lr, lw); + in_flight++; + } + } + }; + + for (int k = 0; k < size_ - 1; k++) { + // Step 0 forwards this rank's own input; later steps forward the partial + // accumulated in the output by the previous step. + const T* send_base = (k == 0) ? in_ptr : out_ptr; + // For step 0 the send window is in the input (count space); for later + // steps the partial lives at wire_offset in the chunk sized output. + int64_t send_base_offset[MAX_DIR]; + for (int lr = 0; lr < MAX_DIR; lr++) { + send_base_offset[lr] = (k == 0) ? in_send_offset[lr] : 0; + } + + // Prefill the pipeline + int buff = 0; + while (buff < n_steps && buff < PIPELINE) { + for (int lr = 0; lr < MAX_DIR; lr++) { + recv_from(sz, buff, lr, lw); + } + for (int lr = 0; lr < MAX_DIR; lr++) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_base_offset[lr] + offset, + send_base + send_base_offset[lr] + + std::max(offset, std::min(offset + N, wire_end[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); + } + + buff++; + in_flight += 2 * MAX_DIR; + } + + // Main loop + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll_wire(lw, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int lr = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR) { + if (send_count[lr] < n_steps) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_base_offset[lr] + offset, + send_base + send_base_offset[lr] + + std::max(offset, std::min(offset + N, wire_end[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); + in_flight++; + } + // A newly staged send may unblock deferred recvs. + drain_deferred(lr); + } + + else if (work_type == RECV_WR) { + // Only reduce into the output once the matching send slice has been + // staged; otherwise defer until send_count catches up. + deferred_recv[lr]++; + drain_deferred(lr); + } + } + } + + // Advance the input windows around the ring for the next step and reset + // the per step counters. + in_send_offset[0] = (in_send_offset[0] + total - chunk) % total; + in_recv_offset[0] = (in_recv_offset[0] + total - chunk) % total; + in_send_offset[1] = (in_send_offset[1] + chunk) % total; + in_recv_offset[1] = (in_recv_offset[1] + chunk) % total; + for (int lr = 0; lr < MAX_DIR; lr++) { + send_count[lr] = recv_count[lr] = 0; + deferred_recv[lr] = 0; + } + } + } + + // Run size_ - 1 pipelined ring steps for a single wire in every direction. + // + // At every step each direction sends its current slice to a neighbor and + // receives the neighbor's slice, which recv_op either reduces or copies into + // out_ptr. After each step the send and recv windows rotate around the ring + // (direction 0 backward, direction 1 forward). This is the shared engine + // behind both all reduce passes and all gather. + // + // The first step sends from in_ptr, later steps from out_ptr; this lets the + // all reduce seed its output without an up front copy. Callers that stage + // into out_ptr (both all gather paths) pass out_ptr for in_ptr. wire_offset + // is this wire's element offset within a chunk, wire_end the end of its + // slice. send_offset and recv_offset are the starting windows, updated in + // place so callers can chain passes. stride is the elements a window moves + // per step and total is the wrap around modulus. + template + inline void ring_pass( + int lw, + int64_t size, + int64_t stride, + int64_t total, + int64_t size_per_wire, + const int64_t (&wire_offset)[MAX_DIR], + const int64_t (&wire_end)[MAX_DIR], + int64_t (&send_offset)[MAX_DIR], + int64_t (&recv_offset)[MAX_DIR], + const T* in_ptr, + T* out_ptr, + RecvOp recv_op) { + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE * 2 * MAX_DIR; + auto [sz, buffer_bytes] = + buffer_size_from_message(size_per_wire * sizeof(T)); + int64_t N = buffer_bytes / sizeof(T); + int64_t n_steps = (size_per_wire + N - 1) / N; + + int in_flight = 0; + int64_t send_limits[MAX_DIR]; + int64_t recv_limits[MAX_DIR]; + int send_count[MAX_DIR] = {0}; + int recv_count[MAX_DIR] = {0}; + auto set_limits = [&]() { + for (int lr = 0; lr < MAX_DIR; lr++) { + send_limits[lr] = std::min( + wire_end[lr], std::max(0, size - send_offset[lr])); + recv_limits[lr] = std::min( + wire_end[lr], std::max(0, size - recv_offset[lr])); + } + }; + set_limits(); + + for (int k = 0; k < size_ - 1; k++) { + // Step 0 sends this rank's own input; later steps send the accumulated + // partial from out_ptr. + const T* send_base = (k == 0) ? in_ptr : out_ptr; + + // Prefill the pipeline + int buff = 0; + while (buff < n_steps && buff < PIPELINE) { + for (int lr = 0; lr < MAX_DIR; lr++) { + recv_from(sz, buff, lr, lw); + } + for (int lr = 0; lr < MAX_DIR; lr++) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_offset[lr] + offset, + send_base + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_count[lr]++; + send_to(sz, buff, lr, lw); + } + + buff++; + in_flight += 2 * MAX_DIR; + } + + // Main loop + // + // Keep going until we have no longer data in flight. + while (in_flight > 0) { + ibv_wc wc[WC_NUM]; + int n = poll_wire(lw, WC_NUM, wc); + for (int i = 0; i < n; i++) { + int work_type = wc[i].wr_id >> 16; + int buff = (wc[i].wr_id >> 8) & 0xff; + int lr = wc[i].wr_id & 0xff; + + in_flight--; + + if (work_type == SEND_WR && send_count[lr] < n_steps) { + int64_t offset = wire_offset[lr] + send_count[lr] * N; + std::copy( + send_base + send_offset[lr] + offset, + send_base + send_offset[lr] + + std::max(offset, std::min(offset + N, send_limits[lr])), + send_buffer(sz, buff, lr, lw).template begin()); + send_to(sz, buff, lr, lw); + in_flight++; + send_count[lr]++; + } + + else if (work_type == RECV_WR) { + int64_t offset = wire_offset[lr] + recv_count[lr] * N; + // The base operand is this rank's own input for the chunk; the all + // gather passes copy and ignore it. + recv_op( + recv_buffer(sz, buff, lr, lw).template begin(), + in_ptr + recv_offset[lr] + offset, + out_ptr + recv_offset[lr] + offset, + std::min(N, recv_limits[lr] - offset)); + recv_count[lr]++; + if (recv_count[lr] + (PIPELINE - 1) < n_steps) { + recv_from(sz, buff, lr, lw); + in_flight++; + } + } + } + } + + // Rotate the windows around the ring for the next step. + send_offset[0] = (send_offset[0] + total - stride) % total; + recv_offset[0] = (recv_offset[0] + total - stride) % total; + if constexpr (MAX_DIR == 2) { + send_offset[1] = (send_offset[1] + stride) % total; + recv_offset[1] = (recv_offset[1] + stride) % total; + } + set_limits(); + for (int lr = 0; lr < MAX_DIR; lr++) { + send_count[lr] = recv_count[lr] = 0; + } + } + } + + void send(const char* in_ptr, int64_t n_bytes, int dst, int n_wires) { + int left = (rank_ + size_ - 1) % size_; + + // In the case that size_ == 2 then left == right so we bias send towards + // left and recv towards right so that the selections will be correct for + // the 2 node case. + int dir = dst == left; + + int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; + + // Split the send across the available wires. Each wire handles the + // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). + dispatch_wires(n_wires, [&](int lw) { + send_wire(in_ptr, n_bytes, dir, bytes_per_wire, lw); + }); + } + + // Perform a point-to-point send for a single wire lw. + // + // Only touches the connection and buffers of wire lw so several wires can run + // concurrently. + void send_wire( + const char* in_ptr, + int64_t n_bytes, + int dir, + int64_t bytes_per_wire, + int lw) { + auto& conns = dir ? left_ : right_; + + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + + auto [sz, N] = buffer_size_from_message(bytes_per_wire); + + int in_flight = 0; + int64_t read_offset = std::min(lw * bytes_per_wire, n_bytes); + int64_t limit = std::min((lw + 1) * bytes_per_wire, n_bytes); + + // Prefill the pipeline + int buff = 0; + while (read_offset < limit && buff < PIPELINE) { + std::copy( + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, limit), + send_buffer(sz, buff, dir, lw).begin()); + send_to(sz, buff, dir, lw); + + buff++; + read_offset += N; + in_flight++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a send was completed and we have more data to send then go ahead + // and send them. + ibv_wc wc[WC_NUM]; + int n = conns[lw].poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + + in_flight--; + + if (read_offset < limit) { + std::copy( + in_ptr + read_offset, + in_ptr + std::min(read_offset + N, limit), + send_buffer(sz, buff, dir, lw).begin()); + send_to(sz, buff, dir, lw); + + read_offset += N; + in_flight++; + } + } + } + } + + void recv(char* out_ptr, int64_t n_bytes, int src, int n_wires) { + int right = (rank_ + 1) % size_; + + // In the case that size_ == 2 then left == right so we bias send towards + // left and recv towards right so that the selections will be correct for + // the 2 node case. + int dir = src == right; + + int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; + + // Split the recv across the available wires. Each wire handles the + // contiguous slice [lw * bytes_per_wire, (lw + 1) * bytes_per_wire). + dispatch_wires(n_wires, [&](int lw) { + recv_wire(out_ptr, n_bytes, dir, bytes_per_wire, lw); + }); + } + + // Perform a point-to-point recv for a single wire lw. + // + // Only touches the connection and buffers of wire lw so several wires can run + // concurrently. + void recv_wire( + char* out_ptr, + int64_t n_bytes, + int dir, + int64_t bytes_per_wire, + int lw) { + auto& conns = dir ? right_ : left_; + + constexpr int PIPELINE = 2; + constexpr int WC_NUM = PIPELINE; + + auto [sz, N] = buffer_size_from_message(bytes_per_wire); + + int in_flight = 0; + int64_t write_offset = std::min(lw * bytes_per_wire, n_bytes); + int64_t limit = std::min((lw + 1) * bytes_per_wire, n_bytes); + + // Prefill the pipeline + int buff = 0; + while (write_offset + N * buff < limit && buff < PIPELINE) { + recv_from(sz, buff, dir, lw); + + buff++; + in_flight++; + } + + // Main loop + while (in_flight > 0) { + // Poll the hardware for completions. + // + // If a recv was completed copy it to the output and if we have more + // data to fetch post another recv. + ibv_wc wc[WC_NUM]; + int n = conns[lw].poll(WC_NUM, wc); + for (int i = 0; i < n; i++) { + int buff = (wc[i].wr_id >> 8) & 0xff; + + in_flight--; + + std::copy( + recv_buffer(sz, buff, dir, lw).begin(), + recv_buffer(sz, buff, dir, lw).begin() + + std::max( + 0, std::min(limit - write_offset, N)), + out_ptr + write_offset); + write_offset += N; + + if (write_offset + (PIPELINE - 1) * N < limit) { + recv_from(sz, buff, dir, lw); + + in_flight++; + } + } + } + } + + private: + void send_to(int sz, int buff, int left_right, int wire) { + auto& conns = left_right ? left_ : right_; + conns[wire].post_send( + send_buffer(sz, buff, left_right, wire), + SEND_WR << 16 | buff << 8 | left_right); + } + + void recv_from(int sz, int buff, int left_right, int wire) { + auto& conns = left_right ? right_ : left_; + conns[wire].post_recv( + recv_buffer(sz, buff, left_right, wire), + RECV_WR << 16 | buff << 8 | left_right); + } + + SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) { + return send_buffers_ + [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + + left_right * n_conns_ + wire]; + } + + SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) { + return recv_buffers_ + [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + + left_right * n_conns_ + wire]; + } + + // Poll the completion queues that belong to a single wire. + // + // A wire always uses both its left and right connections (direction 0 sends + // right / receives left, direction 1 sends left / receives right) so both are + // polled here. Restricting polling to a single wire's connections is what + // allows several wires to run concurrently. + int poll_wire(int wire, int num_completions, ibv_wc* work_completions) { + return poll( + std::span(&left_[wire], 1), + std::span(&right_[wire], 1), + num_completions, + work_completions); + } + + // Run fn(lw) for each wire, the first n_wires - 1 on the pool and the last + // inline, then wait for the pool calls before returning. + template + void dispatch_wires(int n_wires, Fn&& fn) { + if (n_wires <= 1 || pool_ == nullptr) { + for (int lw = 0; lw < n_wires; lw++) { + fn(lw); + } + return; + } + + std::vector> futures; + futures.reserve(n_wires - 1); + for (int lw = 0; lw < n_wires - 1; lw++) { + futures.emplace_back(pool_->enqueue(fn, lw)); + } + + // Wait for the pool calls even if the inline one throws, so they never + // outlive this frame. + try { + fn(n_wires - 1); + } catch (...) { + for (auto& f : futures) { + f.wait(); + } + throw; + } + for (auto& f : futures) { + f.wait(); + } + } + + int rank_; + int size_; + int n_conns_; + std::span left_; + std::span right_; + std::span send_buffers_; + std::span recv_buffers_; + ThreadPool* pool_; +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.cpp new file mode 100644 index 000000000..08f2338c6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.cpp @@ -0,0 +1,206 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include +#include + +#include "jaccl/tcp.h" + +namespace jaccl { + +/** + * Parse a sockaddr from an ip and port provided as strings. + */ +address_t parse_address(const std::string& ip, const std::string& port) { + struct addrinfo hints, *res; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + + int status = getaddrinfo(ip.c_str(), port.c_str(), &hints, &res); + if (status != 0) { + std::ostringstream msg; + msg << "Can't parse address " << ip << ":" << port; + throw std::runtime_error(msg.str()); + } + + address_t result; + memcpy(&result.addr, res->ai_addr, res->ai_addrlen); + result.len = res->ai_addrlen; + freeaddrinfo(res); + + return result; +} + +/** + * Parse a sockaddr provided as an : string. + */ +address_t parse_address(const std::string& ip_port) { + auto colon = ip_port.find(":"); + if (colon == std::string::npos) { + std::ostringstream msg; + msg << "Can't parse address " << ip_port; + throw std::runtime_error(msg.str()); + } + std::string ip(ip_port.begin(), ip_port.begin() + colon); + std::string port(ip_port.begin() + colon + 1, ip_port.end()); + + return parse_address(ip, port); +} + +TCPSocket::TCPSocket(const char* tag) { + sock_ = socket(AF_INET, SOCK_STREAM, 0); + if (sock_ < 0) { + std::ostringstream msg; + msg << tag << " Couldn't create socket (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } +} + +TCPSocket::TCPSocket(TCPSocket&& s) { + sock_ = s.sock_; + s.sock_ = -1; +} + +TCPSocket& TCPSocket::operator=(TCPSocket&& s) { + if (this != &s) { + sock_ = s.sock_; + s.sock_ = -1; + } + return *this; +} + +TCPSocket::TCPSocket(int s) : sock_(s) {} + +TCPSocket::~TCPSocket() { + if (sock_ > 0) { + shutdown(sock_, 2); + close(sock_); + } +} + +int TCPSocket::detach() { + int s = sock_; + sock_ = -1; + return s; +} + +void TCPSocket::listen(const char* tag, const address_t& addr) { + int success; + + // Make sure we can launch immediately after shutdown by setting the + // reuseaddr option so that we don't get address already in use errors + int enable = 1; + success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int)); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't enable reuseport (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + // Bind the socket to the address and port + success = bind(sock_, addr.get(), addr.len); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't bind socket (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + // Prepare waiting for connections + success = ::listen(sock_, 0); + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't listen (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } +} + +TCPSocket TCPSocket::accept(const char* tag) { + int peer = ::accept(sock_, nullptr, nullptr); + if (peer < 0) { + std::ostringstream msg; + msg << tag << " Accept failed (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + return TCPSocket(peer); +} + +void TCPSocket::send(const char* tag, const void* data, size_t len) { + while (len > 0) { + auto n = ::send(sock_, data, len, 0); + if (n <= 0) { + std::ostringstream msg; + msg << tag << " Send failed with errno=" << errno; + throw std::runtime_error(msg.str()); + } + len -= n; + data = static_cast(data) + n; + } +} + +void TCPSocket::recv(const char* tag, void* data, size_t len) { + while (len > 0) { + auto n = ::recv(sock_, data, len, 0); + if (n <= 0) { + std::ostringstream msg; + msg << tag << " Recv failed with errno=" << errno; + throw std::runtime_error(msg.str()); + } + len -= n; + data = static_cast(data) + n; + } +} + +TCPSocket TCPSocket::connect( + const char* tag, + const address_t& addr, + int num_retries, + int wait, + std::function cb) { + int sock, success; + + // Attempt to connect `num_retries` times with exponential backoff. + for (int attempt = 0; attempt < num_retries; attempt++) { + // Create the socket + sock = socket(AF_INET, SOCK_STREAM, 0); + if (sock < 0) { + std::ostringstream msg; + msg << tag << " Couldn't create socket to connect (error: " << errno + << ")"; + throw std::runtime_error(msg.str()); + } + + success = ::connect(sock, addr.get(), addr.len); + if (success == 0) { + break; + } + + if (cb != nullptr) { + cb(attempt, wait); + } + if (wait > 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(wait)); + } + + wait <<= 1; + } + + if (success < 0) { + std::ostringstream msg; + msg << tag << " Couldn't connect (error: " << errno << ")"; + throw std::runtime_error(msg.str()); + } + + return TCPSocket(sock); +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.h new file mode 100644 index 000000000..8c8791353 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/tcp.h @@ -0,0 +1,67 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace jaccl { + +struct address_t { + sockaddr_storage addr; + socklen_t len; + + const sockaddr* get() const { + return (struct sockaddr*)&addr; + } +}; + +/** + * Parse a sockaddr from an ip and port provided as strings. + */ +address_t parse_address(const std::string& ip, const std::string& port); + +/** + * Parse a sockaddr provided as an : string. + */ +address_t parse_address(const std::string& ip_port); + +/** + * Small wrapper over a TCP socket to simplify initiating connections. + */ +class TCPSocket { + public: + TCPSocket(const char* tag); + TCPSocket(const TCPSocket&) = delete; + TCPSocket& operator=(const TCPSocket&) = delete; + TCPSocket(TCPSocket&& s); + TCPSocket& operator=(TCPSocket&&); + ~TCPSocket(); + + void listen(const char* tag, const address_t& addr); + TCPSocket accept(const char* tag); + + void send(const char* tag, const void* data, size_t len); + void recv(const char* tag, void* data, size_t len); + + int detach(); + + operator int() const { + return sock_; + } + + static TCPSocket connect( + const char* tag, + const address_t& addr, + int num_retries = 1, + int wait = 0, + std::function cb = nullptr); + + private: + TCPSocket(int sock); + + int sock_; +}; + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h new file mode 100644 index 000000000..b5d05153e --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/threadpool.h @@ -0,0 +1,137 @@ +// This code was modified from https://github.com/progschj/ThreadPool +// The original License is copied below: +// +// Copyright (c) 2012 Jakob Progsch, Václav Zeman +// This software is provided 'as-is', without any express or implied +// warranty. In no event will the authors be held liable for any damages +// arising from the use of this software. +// +// Permission is granted to anyone to use this software for any purpose, +// including commercial applications, and to alter it and redistribute it +// freely, subject to the following restrictions: +// +// 1. The origin of this software must not be misrepresented; you must not +// claim that you wrote the original software. If you use this software +// in a product, an acknowledgment in the product documentation would be +// appreciated but is not required. +// +// 2. Altered source versions must be plainly marked as such, and must not be +// misrepresented as being the original software. +// +// 3. This notice may not be removed or altered from any source +// distribution. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace jaccl { + +class ThreadPool { + public: + ThreadPool(size_t); + template + auto enqueue(F&& f, Args&&... args) + -> std::future>; + void resize(size_t); + ~ThreadPool(); + + private: + void stop_and_wait(); + void start_threads(size_t); + + std::vector workers; + std::queue> tasks; + std::mutex queue_mutex; + std::condition_variable condition; + bool stop; +}; + +inline ThreadPool::ThreadPool(size_t threads) : stop(false) { + start_threads(threads); +} + +template +auto ThreadPool::enqueue(F&& f, Args&&... args) + -> std::future> { + using return_type = typename std::invoke_result_t; + + auto task = std::make_shared>( + std::bind(std::forward(f), std::forward(args)...)); + + std::future res = task->get_future(); + { + std::unique_lock lock(queue_mutex); + + if (stop) { + throw std::runtime_error( + "[ThreadPool::enqueue] Not allowed on stopped ThreadPool"); + } + + tasks.emplace([task]() { (*task)(); }); + } + condition.notify_one(); + return res; +} + +inline void ThreadPool::resize(size_t threads) { + if (workers.size() == threads) { + return; + } + + if (workers.size() > threads) { + stop_and_wait(); + } + start_threads(threads - workers.size()); +} + +inline ThreadPool::~ThreadPool() { + stop_and_wait(); +} + +inline void ThreadPool::stop_and_wait() { + // Stop the current threads and wait until they finish + { + std::unique_lock lock(queue_mutex); + stop = true; + } + condition.notify_all(); + for (std::thread& worker : workers) { + worker.join(); + } + + // Reset the member variables so that the threadpool is reusable + stop = false; + workers.clear(); +} + +inline void ThreadPool::start_threads(size_t threads) { + for (size_t i = 0; i < threads; ++i) { + workers.emplace_back([this] { + for (;;) { + std::function task; + + { + std::unique_lock lock(this->queue_mutex); + this->condition.wait( + lock, [this] { return this->stop || !this->tasks.empty(); }); + if (this->stop && this->tasks.empty()) + return; + task = std::move(this->tasks.front()); + this->tasks.pop(); + } + + task(); + } + }); + } +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/types.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/types.h new file mode 100644 index 000000000..defae7319 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/lib/jaccl/types.h @@ -0,0 +1,263 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#include +#include +#include + +#if defined(__aarch64__) && defined(__APPLE__) +#include +#endif + +namespace jaccl { + +namespace { + +union float_bits { + float f; + uint32_t u; +}; + +} // namespace + +#ifdef __ARM_FEATURE_FP16_SCALAR_ARITHMETIC + +#include + +#else + +#define __JACCL_HALF_NAN__ 0x7D00 + +// +// The MLX float16 compatibility fallback. Redefined here to keep JACCL +// standalone. Also it doesn't really need all the ops MXL defines. +// +struct float16_t { + uint16_t bits_; + + float16_t(const float& x) : bits_(0) { + float_bits in; + in.f = x; + + // Extract sign + uint32_t x_sign_32 = in.u & uint32_t(0x80000000); + uint16_t x_sign_16 = (x_sign_32 >> 16); + + if (std::isnan(x)) { + bits_ = x_sign_16 | uint16_t(__JACCL_HALF_NAN__); + } else { + // Union + float_bits inf_scale, zero_scale, magic_bits; + + // Find exponent bits and take the max supported by half + uint32_t x_expo_32 = in.u & uint32_t(0x7f800000); + uint32_t max_expo_32 = uint32_t(0x38800000); + x_expo_32 = x_expo_32 < max_expo_32 ? max_expo_32 : x_expo_32; + x_expo_32 += uint32_t(15) << 23; + + // Handle scaling to inf as needed + inf_scale.u = uint32_t(0x77800000); + zero_scale.u = uint32_t(0x08800000); + + // Combine with magic and let addition do rounding + magic_bits.u = x_expo_32; + magic_bits.f += (std::abs(x) * inf_scale.f) * zero_scale.f; + + // Take the lower 5 bits of the exponent + uint32_t x_expo_16 = ((magic_bits.u >> 13) & uint32_t(0x7c00)); + + // Collect the lower 12 bits which have the mantissa + uint32_t x_mant_16 = magic_bits.u & uint32_t(0x0fff); + + // Combine sign, exp and mantissa + bits_ = (x_sign_16 | uint16_t(x_expo_16 + x_mant_16)); + } + } + + operator float() const { + float_bits out; + + uint32_t x_sign_32 = (bits_ << 16) & uint32_t(0x80000000); + uint32_t base = (bits_ << 16); + uint32_t two_base = base + base; + + uint32_t denorm_max = 1u << 27; + if (two_base < denorm_max) { + out.u = uint32_t(126) << 23; // magic mask + out.u |= (two_base >> 17); // Bits from fp16 + out.f -= 0.5f; // magic bias + } else { + out.u = uint32_t(0xE0) << 23; // exponent offset + out.u += (two_base >> 4); // Bits from fp16 + float out_unscaled = out.f; // Store value + out.u = uint32_t(0x7800000); // exponent scale + out.f *= out_unscaled; + } + + // Add sign + out.u |= x_sign_32; + + return out.f; + } + + bool operator<(float16_t x) { + return static_cast(*this) < static_cast(x); + } + + bool operator>(float16_t x) { + return static_cast(*this) > static_cast(x); + } + + float16_t operator+(float16_t x) { + return static_cast(*this) + static_cast(x); + } + + float16_t& operator+=(float16_t x) { + *this = *this + x; + return *this; + } +}; + +#endif + +// +// Check at runtime if the CPU supports native bfloat16 (FEAT_BF16). +// +// This allows us to compile once for all Macs but still enable the feature if +// it is supported. +// +inline bool has_native_bf16_support() { +#if defined(__aarch64__) && defined(__APPLE__) + static bool has_support = []() { + int value = 0; + size_t value_size = sizeof(value); + int success = sysctlbyname( + "hw.optional.arm.FEAT_BF16", &value, &value_size, nullptr, 0); + return success == 0 & value != 0; + }(); + return has_support; +#else + return false; +#endif +} + +// +// The MLX bfloat16 compatibility fallback. +// +// Contrary to float16 we always define it and we 'll use +// has_native_bf16_support to decide if we are going to use __bf16 instead +// during runtime. +// + +#define __JACCL_BFLOAT_NAN__ 0x7FC0 + +struct bfloat16_t { + uint16_t bits_; + + bfloat16_t(const float& x) { + if (std::isnan(x)) { + bits_ = __JACCL_BFLOAT_NAN__; + } else { + float_bits in; + in.f = x; + in.u += (in.u >> 16 & 1) + uint32_t(0x7FFF); + bits_ = in.u >> 16; + } + } + + operator float() const { + float_bits out; + out.u = ((uint32_t)bits_) << 16; + return out.f; + } + + bool operator<(bfloat16_t x) { + return static_cast(*this) < static_cast(x); + } + + bool operator>(bfloat16_t x) { + return static_cast(*this) > static_cast(x); + } + + bfloat16_t operator+(bfloat16_t x) { + return static_cast(*this) + static_cast(x); + } + + bfloat16_t& operator+=(bfloat16_t x) { + *this = *this + x; + return *this; + } +}; + +using complex64_t = std::complex; + +inline bool operator<(complex64_t lhs, complex64_t rhs) { + return lhs.real() < rhs.real() || + (lhs.real() == rhs.real() && lhs.imag() < rhs.imag()); +} + +inline bool operator>(complex64_t lhs, complex64_t rhs) { + return lhs.real() > rhs.real() || + (lhs.real() == rhs.real() && lhs.imag() > rhs.imag()); +} + +template +struct type_identity { + using type = T; +}; + +#define JACCL_GET_TYPE(x) typename decltype(x)::type + +/** + * Dispatch a function for all supported types based on a Dtype. + */ +template +void dispatch_all_types(int dtype, F&& f) { + switch (dtype) { + case Dtype::Bool: + f(type_identity{}); + break; + case Dtype::Int8: + f(type_identity{}); + break; + case Dtype::Int16: + f(type_identity{}); + break; + case Dtype::Int32: + f(type_identity{}); + break; + case Dtype::Int64: + f(type_identity{}); + break; + case Dtype::UInt8: + f(type_identity{}); + break; + case Dtype::UInt16: + f(type_identity{}); + break; + case Dtype::UInt32: + f(type_identity{}); + break; + case Dtype::UInt64: + f(type_identity{}); + break; + case Dtype::Float16: + f(type_identity{}); + break; + case Dtype::BFloat16: + f(type_identity{}); + break; + case Dtype::Float32: + f(type_identity{}); + break; + case Dtype::Float64: + f(type_identity{}); + break; + case Dtype::Complex64: + f(type_identity{}); + break; + } +} + +} // namespace jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.cpp deleted file mode 100644 index c8df4e674..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.cpp +++ /dev/null @@ -1,451 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/distributed/jaccl/mesh.h" -#include "mlx/backend/cpu/encoder.h" -#include "mlx/distributed/reduction_ops.h" -#include "mlx/dtype_utils.h" - -constexpr int MAX_PEERS = 8; - -namespace mlx::core::distributed::jaccl { - -MeshGroup::MeshGroup( - int rank, - const std::vector& device_names, - const char* coordinator_addr) - : rank_(rank), - size_(device_names.size()), - side_channel_(rank_, size_, coordinator_addr), - connections_(create_connections(device_names)) { - if (size_ > MAX_PEERS) { - std::ostringstream msg; - msg << "[jaccl] The JACCL mesh supports up to " << MAX_PEERS - << " peers but " << size_ << " were provided."; - throw std::runtime_error(msg.str()); - } - - // Initialize all the connections and allocate buffers - initialize(); - - // Make sure every node has reached here before continuing - side_channel_.all_gather(0); -} - -void MeshGroup::initialize() { - // Create the queue pairs - for (auto& conn : connections_) { - if (conn.ctx == nullptr) { - continue; - } - conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); - conn.create_queue_pair(); - } - - allocate_buffers(); - - // First init all connections - for (int peer = 0; peer < size_; peer++) { - if (peer == rank_) { - continue; - } - connections_[peer].queue_pair_init(); - } - - // Gather the information to be exchanged, this also serves as a barrier so - // that all peers have initialized their connections before attempting to - // transition to RTS. - std::vector info; - for (auto& conn : connections_) { - info.emplace_back(conn.info()); - } - auto all_infos = side_channel_.all_gather(info); - - // Transition queue pairs to RTS - for (int peer = 0; peer < size_; peer++) { - if (peer == rank_) { - continue; - } - auto peer_info = all_infos[peer][rank_]; - connections_[peer].queue_pair_rtr(peer_info); - connections_[peer].queue_pair_rts(); - } -} - -void MeshGroup::allocate_buffers() { - // Deregister any buffers and free the memory - buffers_.clear(); - - // Allocate the memory - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - for (int j = 0; j < size_; j++) { - buffers_.emplace_back(FRAME_SIZE * (1 << k)); - } - } - } - - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - for (int j = 0; j < size_; j++) { - // This is our send buffer so register it with all pds so we can send - // it to all connected devices. - if (j == rank_) { - for (auto& conn : connections_) { - if (conn.ctx != nullptr) { - buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] - .register_to_protection_domain(conn.protection_domain); - } - } - } - - // This is the recv buffer from rank j so register it to rank j's - // protection domain. - else { - buffers_[k * NUM_BUFFERS * size_ + i * size_ + j] - .register_to_protection_domain(connections_[j].protection_domain); - } - } - } - } -} - -void MeshGroup::all_sum(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::SumOp{}); - }); -} - -void MeshGroup::all_max(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MaxOp{}); - }); -} - -void MeshGroup::all_min(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MinOp{}); - }); -} - -void MeshGroup::all_gather(const array& input, array& output, Stream stream) { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - size_t n_bytes = input.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.set_output_array(output); - encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - - // Fully connected all gather - char* data = out_ptr; - char* our_data = out_ptr + rank_ * n_bytes; - auto [sz, N] = buffer_size_from_message(n_bytes); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; - int64_t total = static_cast(n_bytes); - int num_peers = size_ - 1; - - // Counters to maintain the state of transfers - int in_flight = 0; - int read_offset = 0; - int completed_send_count[PIPELINE] = {0}; - int write_offset[MAX_PEERS] = {0}; - - // Prefill the pipeline - int buff = 0; - while (read_offset < total && buff < PIPELINE) { - post_recv_all(sz, buff); - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - buff++; - in_flight += 2 * num_peers; - read_offset += N; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(connections_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - // Send completed. If all sends completed then send the next chunk. - if (work_type == SEND_WR && read_offset < total) { - completed_send_count[buff]++; - if (completed_send_count[buff] == num_peers) { - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - completed_send_count[buff] = 0; - in_flight += num_peers; - read_offset += N; - } - } - - // Recv completed. If we have more chunks then post another recv. - else if (work_type == RECV_WR) { - std::copy( - recv_buffer(sz, buff, rank).begin(), - recv_buffer(sz, buff, rank).begin() + - std::min(N, total - write_offset[rank]), - data + rank * n_bytes + write_offset[rank]); - write_offset[rank] += N; - if (write_offset[rank] + N * (PIPELINE - 1) < total) { - recv_from(sz, rank, buff); - in_flight++; - } - } - } - } - }); -} - -void MeshGroup::send(const array& input, int dst, Stream stream) { - auto data = input.data(); - int64_t n_bytes = input.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.dispatch([data, n_bytes, dst, this]() { - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE; - auto [sz, N] = buffer_size_from_message(n_bytes); - - int in_flight = 0; - int64_t read_offset = 0; - - // Prefill the pipeline - int buff = 0; - while (read_offset < n_bytes && buff < PIPELINE) { - std::copy( - data + read_offset, - data + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); - send_to(sz, dst, buff); - - buff++; - read_offset += N; - in_flight++; - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed and we have more data to send then go ahead - // and send them. - ibv_wc wc[WC_NUM]; - int n = connections_[dst].poll(WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - if (read_offset < n_bytes) { - std::copy( - data + read_offset, - data + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); - send_to(sz, dst, buff); - - read_offset += N; - in_flight++; - } - } - } - }); -} - -void MeshGroup::recv(array& out, int src, Stream stream) { - auto data = out.data(); - int64_t n_bytes = out.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_output_array(out); - encoder.dispatch([data, n_bytes, src, this]() { - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE; - auto [sz, N] = buffer_size_from_message(n_bytes); - - int in_flight = 0; - int64_t write_offset = 0; - - // Prefill the pipeline - int buff = 0; - while (N * buff < n_bytes && buff < PIPELINE) { - recv_from(sz, src, buff); - - in_flight++; - buff++; - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a recv was completed copy it to the output and if we have more - // data to fetch post another recv. - ibv_wc wc[WC_NUM]; - int n = connections_[src].poll(WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - std::copy( - recv_buffer(sz, buff, src).begin(), - recv_buffer(sz, buff, src).begin() + - std::min(n_bytes - write_offset, static_cast(N)), - data + write_offset); - write_offset += N; - - if (write_offset + (PIPELINE - 1) * N < n_bytes) { - recv_from(sz, src, buff); - - in_flight++; - } - } - } - }); -} - -template -void MeshGroup::all_reduce( - const array& input, - array& output, - Stream stream, - ReduceOp reduce_op) { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.set_output_array(output); - encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() { - // If not inplace all reduce then copy the input to the output first - if (in_ptr != out_ptr) { - std::memcpy(out_ptr, in_ptr, size * sizeof(T)); - } - - // Fully connected all reduce - T* data = out_ptr; - auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T)); - int64_t N = buffer_size / sizeof(T); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_PEERS * 2; - int64_t total = static_cast(size); - int num_peers = size_ - 1; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t read_offset = 0; - int completed_send_count[PIPELINE] = {0}; - int completed_recv_begin[MAX_PEERS] = {0}; - int completed_recv_end[MAX_PEERS] = {0}; - - // Prefill the pipeline - int buff = 0; - while (read_offset < total && buff < PIPELINE) { - post_recv_all(sz, buff); - std::copy( - data + read_offset, - data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - buff++; - in_flight += 2 * num_peers; - read_offset += N; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed mark how many completions we have received - // for that buffer. If we have sent the buffer to all peers we can - // reuse the buffer so copy the next chunk of data and send it to all. - // - // If a receive is completed then advance the pointer of completed - // receives. - ibv_wc wc[WC_NUM]; - int n = poll(connections_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - if (work_type == SEND_WR && read_offset < total) { - completed_send_count[buff]++; - if (completed_send_count[buff] == num_peers) { - std::copy( - data + read_offset, - data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - completed_send_count[buff] = 0; - in_flight += num_peers; - read_offset += N; - } - } - - else if (work_type == RECV_WR) { - completed_recv_end[rank]++; - } - } - - // Process the completed recv - // - // For each rank we have a range of completed recv defined by a begin - // and end inclusive and exlusive in standard C++ fashion. - // - // When there is an unprocessed receive we first check if we have - // finished sending the write location. If so then we reduce in-place - // and then check if there is more to be received and post a recv. - for (int r = 0; r < size_; r++) { - int s = completed_recv_begin[r]; - int e = completed_recv_end[r]; - int w = s * N; - while (w < read_offset && e - s > 0) { - int buff = s % PIPELINE; - reduce_op( - recv_buffer(sz, buff, r).begin(), - data + w, - std::min(N, total - w)); - w += N; - s++; - if (w + (PIPELINE - 1) * N < total) { - recv_from(sz, r, buff); - in_flight++; - } - } - completed_recv_begin[r] = s; - } - } - }); -} - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.h deleted file mode 100644 index 6f779e9cc..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh.h +++ /dev/null @@ -1,122 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#pragma once - -#include "mlx/distributed/distributed_impl.h" -#include "mlx/distributed/jaccl/utils.h" - -using GroupImpl = mlx::core::distributed::detail::GroupImpl; - -namespace mlx::core::distributed::jaccl { - -/** - * The JACCL communication group for a fully connected mesh. We expect one - * connection per peer and it should be the lowest latency communication group - * for small to medium size messages. - * - * Like all JACCL groups it uses a side channel to exchange the necessary - * information and then configure the connections to be ready for RDMA - * operations. - */ -class MeshGroup : public GroupImpl { - public: - MeshGroup( - int rank, - const std::vector& device_names, - const char* coordinator_addr); - - Stream communication_stream(StreamOrDevice s) override { - return to_stream(s, Device::cpu); - } - - int rank() override { - return rank_; - } - - int size() override { - return size_; - } - - void all_sum(const array& input, array& output, Stream stream) override; - void all_max(const array& input, array& output, Stream stream) override; - void all_min(const array& input, array& output, Stream stream) override; - void all_gather(const array& input, array& output, Stream stream) override; - void send(const array& input, int dst, Stream stream) override; - void recv(array& out, int src, Stream stream) override; - - void sum_scatter(const array& input, array& output, Stream stream) override { - throw std::runtime_error("[jaccl] sum_scatter not supported."); - } - - std::shared_ptr split(int color, int key = -1) override { - throw std::runtime_error("[jaccl] Group split not supported."); - } - - private: - template - void all_reduce( - const array& input, - array& output, - Stream stream, - ReduceOp reduce_op); - - /** - * Performs the connection initialization. Namely, after this call all - * Connection objects should have a queue pair in RTS state and all buffers - * should have been allocated. - */ - void initialize(); - - /** - * Allocate all the buffers that we will use in the communication group. - */ - void allocate_buffers(); - - void send_to(int sz, int rank, int buff) { - connections_[rank].post_send( - send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank); - } - - void recv_from(int sz, int rank, int buff) { - connections_[rank].post_recv( - recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank); - } - - SharedBuffer& send_buffer(int sz, int buff) { - return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank_]; - } - - SharedBuffer& recv_buffer(int sz, int buff, int rank) { - return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank]; - } - - void post_send_all(int sz, int buff) { - auto& b = send_buffer(sz, buff); - int wr_id = SEND_WR << 16 | buff << 8; - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - connections_[i].post_send(b, wr_id | i); - } - } - - void post_recv_all(int sz, int buff) { - int b = sz * NUM_BUFFERS * size_ + buff * size_; - int wr_id = RECV_WR << 16 | buff << 8; - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - connections_[i].post_recv(buffers_[b + i], wr_id | i); - } - } - - int rank_; - int size_; - SideChannel side_channel_; - std::vector connections_; - std::vector buffers_; -}; - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh_impl.h deleted file mode 100644 index fc486a396..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/mesh_impl.h +++ /dev/null @@ -1,358 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#pragma once - -#include - -#include "mlx/distributed/jaccl/utils.h" - -constexpr int MESH_MAX_PEERS = 8; - -namespace mlx::core::distributed::jaccl { - -class MeshImpl { - public: - MeshImpl( - int rank, - int size, - std::vector& conns, - std::vector& buffers) - : rank_(rank), size_(size), connections_(conns), buffers_(buffers) {} - - MeshImpl() : rank_(0), size_(1) {} - - template - void - all_reduce(const T* in_ptr, T* out_ptr, int64_t size, ReduceOp reduce_op) { - // If not inplace all reduce then copy the input to the output first - if (in_ptr != out_ptr) { - std::memcpy(out_ptr, in_ptr, size * sizeof(T)); - } - - // Fully connected all reduce - T* data = out_ptr; - auto [sz, buffer_size] = buffer_size_from_message(size * sizeof(T)); - int64_t N = buffer_size / sizeof(T); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; - int64_t total = static_cast(size); - int num_peers = size_ - 1; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t read_offset = 0; - int completed_send_count[PIPELINE] = {0}; - int completed_recv_begin[MESH_MAX_PEERS] = {0}; - int completed_recv_end[MESH_MAX_PEERS] = {0}; - - // Prefill the pipeline - int buff = 0; - while (read_offset < total && buff < PIPELINE) { - post_recv_all(sz, buff); - std::copy( - data + read_offset, - data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - buff++; - in_flight += 2 * num_peers; - read_offset += N; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed mark how many completions we have received - // for that buffer. If we have sent the buffer to all peers we can - // reuse the buffer so copy the next chunk of data and send it to all. - // - // If a receive is completed then advance the pointer of completed - // receives. - ibv_wc wc[WC_NUM]; - int n = poll(connections_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - if (work_type == SEND_WR && read_offset < total) { - completed_send_count[buff]++; - if (completed_send_count[buff] == num_peers) { - std::copy( - data + read_offset, - data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - completed_send_count[buff] = 0; - in_flight += num_peers; - read_offset += N; - } - } - - else if (work_type == RECV_WR) { - completed_recv_end[rank]++; - } - } - - // Process the completed recv - // - // For each rank we have a range of completed recv defined by a begin - // and end inclusive and exlusive in standard C++ fashion. - // - // When there is an unprocessed receive we first check if we have - // finished sending the write location. If so then we reduce in-place - // and then check if there is more to be received and post a recv. - for (int r = 0; r < size_; r++) { - int s = completed_recv_begin[r]; - int e = completed_recv_end[r]; - int w = s * N; - while (w < read_offset && e - s > 0) { - int buff = s % PIPELINE; - reduce_op( - recv_buffer(sz, buff, r).begin(), - data + w, - std::min(N, total - w)); - w += N; - s++; - if (w + (PIPELINE - 1) * N < total) { - recv_from(sz, r, buff); - in_flight++; - } - } - completed_recv_begin[r] = s; - } - } - } - - void all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes) { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - - // Fully connected all gather - char* data = out_ptr; - char* our_data = out_ptr + rank_ * n_bytes; - auto [sz, N] = buffer_size_from_message(n_bytes); - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MESH_MAX_PEERS * 2; - int64_t total = static_cast(n_bytes); - int num_peers = size_ - 1; - - // Counters to maintain the state of transfers - int in_flight = 0; - int read_offset = 0; - int completed_send_count[PIPELINE] = {0}; - int write_offset[MESH_MAX_PEERS] = {0}; - - // Prefill the pipeline - int buff = 0; - while (read_offset < total && buff < PIPELINE) { - post_recv_all(sz, buff); - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - buff++; - in_flight += 2 * num_peers; - read_offset += N; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(connections_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - // Send completed. If all sends completed then send the next chunk. - if (work_type == SEND_WR && read_offset < total) { - completed_send_count[buff]++; - if (completed_send_count[buff] == num_peers) { - std::copy( - our_data + read_offset, - our_data + std::min(read_offset + N, total), - send_buffer(sz, buff).begin()); - post_send_all(sz, buff); - - completed_send_count[buff] = 0; - in_flight += num_peers; - read_offset += N; - } - } - - // Recv completed. If we have more chunks then post another recv. - else if (work_type == RECV_WR) { - std::copy( - recv_buffer(sz, buff, rank).begin(), - recv_buffer(sz, buff, rank).begin() + - std::min(N, total - write_offset[rank]), - data + rank * n_bytes + write_offset[rank]); - write_offset[rank] += N; - if (write_offset[rank] + N * (PIPELINE - 1) < total) { - recv_from(sz, rank, buff); - in_flight++; - } - } - } - } - } - - void send(const char* in_ptr, int64_t n_bytes, int dst) { - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE; - auto [sz, N] = buffer_size_from_message(n_bytes); - - int in_flight = 0; - int64_t read_offset = 0; - - // Prefill the pipeline - int buff = 0; - while (read_offset < n_bytes && buff < PIPELINE) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); - send_to(sz, dst, buff); - - buff++; - read_offset += N; - in_flight++; - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed and we have more data to send then go ahead - // and send them. - ibv_wc wc[WC_NUM]; - int n = connections_[dst].poll(WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - if (read_offset < n_bytes) { - std::copy( - in_ptr + read_offset, - in_ptr + std::min(read_offset + N, n_bytes), - send_buffer(sz, buff).begin()); - send_to(sz, dst, buff); - - read_offset += N; - in_flight++; - } - } - } - } - - void recv(char* out_ptr, int64_t n_bytes, int src) { - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE; - auto [sz, N] = buffer_size_from_message(n_bytes); - - int in_flight = 0; - int64_t write_offset = 0; - - // Prefill the pipeline - int buff = 0; - while (N * buff < n_bytes && buff < PIPELINE) { - recv_from(sz, src, buff); - - in_flight++; - buff++; - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a recv was completed copy it to the output and if we have more - // data to fetch post another recv. - ibv_wc wc[WC_NUM]; - int n = connections_[src].poll(WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int rank = wc[i].wr_id & 0xff; - - in_flight--; - - std::copy( - recv_buffer(sz, buff, src).begin(), - recv_buffer(sz, buff, src).begin() + - std::min(n_bytes - write_offset, static_cast(N)), - out_ptr + write_offset); - write_offset += N; - - if (write_offset + (PIPELINE - 1) * N < n_bytes) { - recv_from(sz, src, buff); - - in_flight++; - } - } - } - } - - private: - void send_to(int sz, int rank, int buff) { - connections_[rank].post_send( - send_buffer(sz, buff), SEND_WR << 16 | buff << 8 | rank); - } - - void recv_from(int sz, int rank, int buff) { - connections_[rank].post_recv( - recv_buffer(sz, buff, rank), RECV_WR << 16 | buff << 8 | rank); - } - - SharedBuffer& send_buffer(int sz, int buff) { - return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank_]; - } - - SharedBuffer& recv_buffer(int sz, int buff, int rank) { - return buffers_[sz * NUM_BUFFERS * size_ + buff * size_ + rank]; - } - - void post_send_all(int sz, int buff) { - auto& b = send_buffer(sz, buff); - int wr_id = SEND_WR << 16 | buff << 8; - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - connections_[i].post_send(b, wr_id | i); - } - } - - void post_recv_all(int sz, int buff) { - int b = sz * NUM_BUFFERS * size_ + buff * size_; - int wr_id = RECV_WR << 16 | buff << 8; - for (int i = 0; i < size_; i++) { - if (i == rank_) { - continue; - } - connections_[i].post_recv(buffers_[b + i], wr_id | i); - } - } - - int rank_; - int size_; - std::span connections_; - std::span buffers_; -}; - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/no_jaccl.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/no_jaccl.cpp index 12fd6ab2f..cab81eae3 100644 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/no_jaccl.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/jaccl/no_jaccl.cpp @@ -17,4 +17,11 @@ std::shared_ptr init(bool strict /* = false */) { return nullptr; } +std::shared_ptr init(bool strict, AllGatherFactory /* factory */) { + if (strict) { + throw std::runtime_error("Cannot initialize jaccl distributed backend."); + } + return nullptr; +} + } // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.cpp b/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.cpp deleted file mode 100644 index fe27a2e9c..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.cpp +++ /dev/null @@ -1,692 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#include "mlx/distributed/jaccl/ring.h" -#include "mlx/backend/cpu/encoder.h" -#include "mlx/distributed/reduction_ops.h" -#include "mlx/dtype_utils.h" - -namespace mlx::core::distributed::jaccl { - -RingGroup::RingGroup( - int rank, - int size, - const std::vector& left_devices, - const std::vector& right_devices, - const char* coordinator_addr) - : rank_(rank), - size_(size), - side_channel_(rank_, size_, coordinator_addr), - left_(create_connections(left_devices)), - right_(create_connections(right_devices)) { - if (left_.size() > MAX_CONNS || right_.size() > MAX_CONNS) { - std::ostringstream msg; - msg << "[jaccl] Up to " << MAX_CONNS << " per direction supported but " - << left_.size() << " were provided."; - throw std::runtime_error(msg.str()); - } - - // Initialize all the connections and allocate buffers - initialize(); - - // Make sure every node has reached here before continuing - side_channel_.all_gather(0); -} - -void RingGroup::initialize() { - // Create the queue pairs - for (auto& conn : left_) { - conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); - conn.create_queue_pair(); - } - for (auto& conn : right_) { - conn.allocate_protection_domain(); - conn.create_completion_queue(MAX_SEND_WR + MAX_RECV_WR); - conn.create_queue_pair(); - } - - // Allocate the buffers - allocate_buffers(); - - // Initialize the conections - for (auto& conn : left_) { - conn.queue_pair_init(); - } - for (auto& conn : right_) { - conn.queue_pair_init(); - } - - // Gather the information to be exchanged, this also serves as a barrier so - // that all peers have initialized their connections before attempting to - // transition to RTS. - std::vector left_info; - for (auto& conn : left_) { - left_info.emplace_back(conn.info()); - } - std::vector right_info; - for (auto& conn : right_) { - right_info.emplace_back(conn.info()); - } - auto all_left_infos = side_channel_.all_gather(left_info); - auto all_right_infos = side_channel_.all_gather(right_info); - - // Transition queue pairs to RTS - int left_peer = (rank_ + size_ - 1) % size_; - for (int i = 0; i < left_.size(); i++) { - auto peer_info = all_right_infos[left_peer][i]; - left_[i].queue_pair_rtr(peer_info); - left_[i].queue_pair_rts(); - } - int right_peer = (rank_ + 1) % size_; - for (int i = 0; i < right_.size(); i++) { - auto peer_info = all_left_infos[right_peer][i]; - right_[i].queue_pair_rtr(peer_info); - right_[i].queue_pair_rts(); - } -} - -void RingGroup::allocate_buffers() { - // Deregister any buffers and free the memory - send_buffers_.clear(); - recv_buffers_.clear(); - - // Allocate the memory - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - for (int j = 0; j < MAX_CONNS * 2; j++) { - send_buffers_.emplace_back(FRAME_SIZE * (1 << k)); - recv_buffers_.emplace_back(FRAME_SIZE * (1 << k)); - } - } - } - - // Register the buffers with the corresponding connections - for (int k = 0; k < BUFFER_SIZES; k++) { - for (int i = 0; i < NUM_BUFFERS; i++) { - for (int j = 0; j < MAX_CONNS * 2; j++) { - int wire = j % MAX_CONNS; - int lr = j / MAX_CONNS; - if (wire >= left_.size()) { - continue; - } - if (lr) { - send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j] - .register_to_protection_domain(left_[wire].protection_domain); - recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j] - .register_to_protection_domain(right_[wire].protection_domain); - } else { - send_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j] - .register_to_protection_domain(right_[wire].protection_domain); - recv_buffers_[k * NUM_BUFFERS * MAX_CONNS * 2 + i * MAX_CONNS * 2 + j] - .register_to_protection_domain(left_[wire].protection_domain); - } - } - } - } -} - -void RingGroup::all_sum(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::SumOp{}); - }); -} - -void RingGroup::all_max(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MaxOp{}); - }); -} - -void RingGroup::all_min(const array& input, array& output, Stream stream) { - dispatch_all_types(output.dtype(), [&](auto type_tag) { - using T = MLX_GET_TYPE(type_tag); - all_reduce(input, output, stream, detail::MinOp{}); - }); -} - -void RingGroup::all_gather(const array& input, array& output, Stream stream) { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - size_t n_bytes = input.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.set_output_array(output); - encoder.dispatch([in_ptr, out_ptr, n_bytes, this]() { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * 2; - int n_wires = left_.size(); - size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); - size_t out_bytes = n_bytes * size_; - auto [sz, N] = buffer_size_from_message(n_bytes_per_wire); - int n_steps = (n_bytes_per_wire + N - 1) / N; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t send_offset[2]; - int64_t recv_offset[2]; - int64_t limits[2]; - int send_count[2 * MAX_CONNS] = {0}; - int recv_count[2 * MAX_CONNS] = {0}; - send_offset[0] = send_offset[1] = rank_ * n_bytes; - recv_offset[0] = ((rank_ + size_ - 1) % size_) * n_bytes; - recv_offset[1] = ((rank_ + 1) % size_) * n_bytes; - limits[0] = n_wires * n_bytes_per_wire; - limits[1] = n_bytes; - - // Possible perf improvement by not syncing at every step but running ahead - // as needed. - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff); - for (int lr = 0; lr < 2; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * MAX_CONNS + lw] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff); - - buff++; - in_flight += 2 * 2 * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / MAX_CONNS; - int lw = wire % MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + out_bytes - n_bytes) % out_bytes; - recv_offset[0] = (recv_offset[0] + out_bytes - n_bytes) % out_bytes; - send_offset[1] = (send_offset[1] + n_bytes) % out_bytes; - recv_offset[1] = (recv_offset[1] + n_bytes) % out_bytes; - for (int i = 0; i < 2 * MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } - }); -} - -void RingGroup::send(const array& input, int dst, Stream stream) { - int right = (rank_ + 1) % size_; - int left = (rank_ + size_ - 1) % size_; - if (dst != right && dst != left) { - std::ostringstream msg; - msg << "[jaccl] In ring mode send is only supported to direct neighbors " - << "but tried to send to " << dst << " from " << rank_ << std::endl; - throw std::runtime_error(msg.str()); - } - auto data = input.data(); - int64_t n_bytes = input.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.dispatch([data, n_bytes, dst, left, this]() { - // In the case that size_ == 2 then left == right so we bias send towards - // left and recv towards right so that the selections will be correct for - // the 2 node case. - auto& conns = (dst == left) ? left_ : right_; - int dir = dst == left; - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_CONNS; - - int n_wires = conns.size(); - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; - auto [sz, N] = buffer_size_from_message(bytes_per_wire); - - int in_flight = 0; - int64_t read_offset[MAX_CONNS]; - int64_t limits[MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - read_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } - - // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (read_offset[lw] < limits[lw] && buff < PIPELINE) { - std::copy( - data + read_offset[lw], - data + std::min(read_offset[lw] + N, limits[lw]), - send_buffer(sz, buff, dir, lw).begin()); - send_to(sz, buff, dir, lw); - - buff++; - read_offset[lw] += N; - in_flight++; - } - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed and we have more data to send then go ahead - // and send them. - ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % MAX_CONNS; - - in_flight--; - - if (read_offset[lw] < limits[lw]) { - std::copy( - data + read_offset[lw], - data + std::min(read_offset[lw] + N, limits[lw]), - send_buffer(sz, buff, dir, lw).begin()); - send_to(sz, buff, dir, lw); - - read_offset[lw] += N; - in_flight++; - } - } - } - }); -} - -void RingGroup::recv(array& out, int src, Stream stream) { - int right = (rank_ + 1) % size_; - int left = (rank_ + size_ - 1) % size_; - if (src != right && src != left) { - std::ostringstream msg; - msg << "[jaccl] In ring mode recv is only supported to direct neighbors " - << "but tried to recv from " << src << " to " << rank_ << std::endl; - throw std::runtime_error(msg.str()); - } - auto data = out.data(); - int64_t n_bytes = out.nbytes(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_output_array(out); - encoder.dispatch([data, n_bytes, src, right, this]() { - // In the case that size_ == 2 then left == right so we bias send towards - // left and recv towards right so that the selections will be correct for - // the 2 node case. - auto& conns = (src == right) ? right_ : left_; - int dir = src == right; - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_CONNS; - - int n_wires = conns.size(); - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; - auto [sz, N] = buffer_size_from_message(bytes_per_wire); - - int in_flight = 0; - int64_t write_offset[MAX_CONNS]; - int64_t limits[MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - write_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } - - // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (N * buff < limits[lw] && buff < PIPELINE) { - recv_from(sz, buff, dir, lw); - - buff++; - in_flight++; - } - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a recv was completed copy it to the output and if we have more - // data to fetch post another recv. - ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % MAX_CONNS; - - in_flight--; - - std::copy( - recv_buffer(sz, buff, dir, lw).begin(), - recv_buffer(sz, buff, dir, lw).begin() + - std::max( - 0, std::min(limits[lw] - write_offset[lw], N)), - data + write_offset[lw]); - write_offset[lw] += N; - - if (write_offset[lw] + (PIPELINE - 1) * N < limits[lw]) { - recv_from(sz, buff, dir, lw); - - in_flight++; - } - } - } - }); -} - -template -void RingGroup::all_reduce( - const array& input, - array& output, - Stream stream, - ReduceOp reduce_op) { - auto in_ptr = input.data(); - auto out_ptr = output.data(); - auto& encoder = cpu::get_command_encoder(stream); - encoder.set_input_array(input); - encoder.set_output_array(output); - encoder.dispatch([in_ptr, out_ptr, size = input.size(), this, reduce_op]() { - if (size < size_ * 2 * left_.size()) { - all_reduce_impl<1, T, ReduceOp>(in_ptr, out_ptr, size, 1, reduce_op); - return; - } - - all_reduce_impl<2, T, ReduceOp>( - in_ptr, out_ptr, size, left_.size(), reduce_op); - }); -} - -template -void RingGroup::all_reduce_impl( - const T* in_ptr, - T* out_ptr, - int64_t size, - int n_wires, - ReduceOp reduce_op) { - // If not inplace all reduce then copy the input to the output first - if (in_ptr != out_ptr) { - std::memcpy(out_ptr, in_ptr, size * sizeof(T)); - } - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * MAX_CONNS * 2 * MAX_DIR; - int64_t chunk_size = (size + size_ - 1) / size_; - int64_t size_per_wire = - (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires); - auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T)); - N /= sizeof(T); - int64_t n_steps = (size_per_wire + N - 1) / N; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t chunk_multiple_size = size_ * chunk_size; - int64_t send_offset[MAX_DIR]; - int64_t recv_offset[MAX_DIR]; - int64_t send_limits[MAX_DIR]; - int64_t recv_limits[MAX_DIR]; - int send_count[MAX_DIR * MAX_CONNS] = {0}; - int recv_count[MAX_DIR * MAX_CONNS] = {0}; - send_offset[0] = rank_ * chunk_size; - recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = rank_ * chunk_size; - recv_offset[1] = ((rank_ + 1) % size_) * chunk_size; - send_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - - // First reduce scatter - // - // Possible perf improvement by not syncing at every step but running ahead - // as needed. - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / MAX_CONNS; - int lw = wire % MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - reduce_op( - recv_buffer(sz, buff, lr, lw).begin(), - out_ptr + recv_offset[lr] + offset, - std::max(0, std::min(N, recv_limits[lr] - offset))); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } - - // Secondly all gather - // - // The offsets are correct from the scatter reduce - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / MAX_CONNS; - int lw = wire % MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, recv_limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } -} - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.h deleted file mode 100644 index a59ceb3dd..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring.h +++ /dev/null @@ -1,178 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#pragma once - -#include "mlx/distributed/distributed_impl.h" -#include "mlx/distributed/jaccl/utils.h" - -using GroupImpl = mlx::core::distributed::detail::GroupImpl; - -constexpr int MAX_CONNS = 4; - -namespace mlx::core::distributed::jaccl { - -/** - * The JACCL communication group for a ring where each node is connected to its - * two neighboring nodes. It should be the highest bandwidth communication - * group for large messages when many connections per peer are used. - * - * Like all JACCL groups it uses a side channel to exchange the necessary - * information and then configure the connections to be ready for RDMA - * operations. - */ -class RingGroup : public GroupImpl { - public: - RingGroup( - int rank, - int size, - const std::vector& left_devices, - const std::vector& right_devices, - const char* coordinator_addr); - - Stream communication_stream(StreamOrDevice s) override { - return to_stream(s, Device::cpu); - } - - int rank() override { - return rank_; - } - - int size() override { - return size_; - } - - void all_sum(const array& input, array& output, Stream stream) override; - void all_max(const array& input, array& output, Stream stream) override; - void all_min(const array& input, array& output, Stream stream) override; - void all_gather(const array& input, array& output, Stream stream) override; - void send(const array& input, int dst, Stream stream) override; - void recv(array& out, int src, Stream stream) override; - - void sum_scatter(const array& input, array& output, Stream stream) override { - throw std::runtime_error("[jaccl] sum_scatter not supported."); - } - - std::shared_ptr split(int color, int key = -1) override { - throw std::runtime_error("[jaccl] Group split not supported."); - } - - private: - template - void all_reduce( - const array& input, - array& output, - Stream stream, - ReduceOp reduce_op); - - template - void all_reduce_impl( - const T* in_ptr, - T* out_ptr, - int64_t size, - int n_wires, - ReduceOp reduce_op); - - /** - * Performs the connection initialization. Namely, after this call all - * Connection objects should have a queue pair in RTS state and all buffers - * should have been allocated. - */ - void initialize(); - - /** - * Allocate all the buffers that we will use in the communication group. - */ - void allocate_buffers(); - - void send_to(int sz, int buff, int left_right, int wire) { - if (left_right) { - left_[wire].post_send( - send_buffer_left(sz, buff, wire), - SEND_WR << 16 | buff << 8 | (MAX_CONNS + wire)); - } else { - right_[wire].post_send( - send_buffer_right(sz, buff, wire), SEND_WR << 16 | buff << 8 | wire); - } - } - - void recv_from(int sz, int buff, int left_right, int wire) { - if (left_right) { - right_[wire].post_recv( - recv_buffer_right(sz, buff, wire), - RECV_WR << 16 | buff << 8 | (MAX_CONNS + wire)); - } else { - left_[wire].post_recv( - recv_buffer_left(sz, buff, wire), RECV_WR << 16 | buff << 8 | wire); - } - } - - SharedBuffer& send_buffer_right(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire]; - } - - SharedBuffer& send_buffer_left(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS + - wire]; - } - - SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + - left_right * MAX_CONNS + wire]; - } - - SharedBuffer& recv_buffer_left(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + wire]; - } - - SharedBuffer& recv_buffer_right(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + MAX_CONNS + - wire]; - } - - SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * MAX_CONNS * 2 + buff * MAX_CONNS * 2 + - left_right * MAX_CONNS + wire]; - } - - template - void post_recv_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - recv_from(sz, buff, lr, lw); - } - } - } - - void post_recv_all(int sz, int buff) { - post_recv_all<2>(sz, buff, left_.size()); - } - - template - void post_send_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - send_to(sz, buff, lr, lw); - } - } - } - - void post_send_all(int sz, int buff) { - post_send_all<2>(sz, buff, left_.size()); - } - - int rank_; - int size_; - SideChannel side_channel_; - std::vector left_; - std::vector right_; - std::vector send_buffers_; - std::vector recv_buffers_; -}; - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring_impl.h b/Source/Cmlx/mlx/mlx/distributed/jaccl/ring_impl.h deleted file mode 100644 index ce883d1fc..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/jaccl/ring_impl.h +++ /dev/null @@ -1,631 +0,0 @@ -// Copyright © 2026 Apple Inc. - -#pragma once - -#include - -#include "mlx/distributed/jaccl/utils.h" - -constexpr int RING_MAX_CONNS = 4; - -namespace mlx::core::distributed::jaccl { - -class RingImpl { - public: - RingImpl( - int rank, - int size, - std::vector& left, - std::vector& right, - std::vector& send_buffers, - std::vector& recv_buffers) - : rank_(rank), - size_(size), - n_conns_(left.size()), - left_(left), - right_(right), - send_buffers_(send_buffers), - recv_buffers_(recv_buffers) {} - - RingImpl( - int rank, - int size, - Connection* left_begin, - Connection* right_begin, - size_t n_conns, - std::vector& send_buffers, - std::vector& recv_buffers) - : rank_(rank), - size_(size), - n_conns_(n_conns), - left_(left_begin, n_conns), - right_(right_begin, n_conns), - send_buffers_(send_buffers), - recv_buffers_(recv_buffers) {} - - RingImpl() : rank_(0), size_(1), n_conns_(0) {} - - template - void all_reduce( - const T* in_ptr, - T* out_ptr, - int64_t size, - int n_wires, - ReduceOp reduce_op) { - // If not inplace all reduce then copy the input to the output first - if (in_ptr != out_ptr) { - std::memcpy(out_ptr, in_ptr, size * sizeof(T)); - } - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * MAX_DIR; - int64_t chunk_size = (size + size_ - 1) / size_; - int64_t size_per_wire = - (chunk_size + (MAX_DIR * n_wires) - 1) / (MAX_DIR * n_wires); - auto [sz, N] = buffer_size_from_message(size_per_wire * sizeof(T)); - N /= sizeof(T); - int64_t n_steps = (size_per_wire + N - 1) / N; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t chunk_multiple_size = size_ * chunk_size; - int64_t send_offset[MAX_DIR]; - int64_t recv_offset[MAX_DIR]; - int64_t send_limits[MAX_DIR]; - int64_t recv_limits[MAX_DIR]; - int send_count[MAX_DIR * RING_MAX_CONNS] = {0}; - int recv_count[MAX_DIR * RING_MAX_CONNS] = {0}; - send_offset[0] = rank_ * chunk_size; - recv_offset[0] = ((rank_ + size_ - 1) % size_) * chunk_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = rank_ * chunk_size; - recv_offset[1] = ((rank_ + 1) % size_) * chunk_size; - send_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - - // First reduce scatter - // - // Possible perf improvement by not syncing at every step but running ahead - // as needed. - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - reduce_op( - recv_buffer(sz, buff, lr, lw).begin(), - out_ptr + recv_offset[lr] + offset, - std::max(0, std::min(N, recv_limits[lr] - offset))); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } - - // Secondly all gather - // - // The offsets are correct from the scatter reduce - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff, n_wires); - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff, n_wires); - - buff++; - in_flight += 2 * MAX_DIR * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, send_limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * size_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, recv_limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - recv_offset[0] = (recv_offset[0] + chunk_multiple_size - chunk_size) % - chunk_multiple_size; - if constexpr (MAX_DIR == 2) { - send_offset[1] = (send_offset[1] + chunk_size) % chunk_multiple_size; - recv_offset[1] = (recv_offset[1] + chunk_size) % chunk_multiple_size; - send_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - send_offset[0])); - send_limits[1] = - std::min(chunk_size, std::max(0, size - send_offset[1])); - recv_limits[0] = std::min( - n_wires * size_per_wire, - std::max(0, size - recv_offset[0])); - recv_limits[1] = - std::min(chunk_size, std::max(0, size - recv_offset[1])); - } else { - send_limits[0] = - std::min(chunk_size, std::max(0, size - send_offset[0])); - recv_limits[0] = - std::min(chunk_size, std::max(0, size - recv_offset[0])); - } - for (int i = 0; i < MAX_DIR * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } - } - - void - all_gather(const char* in_ptr, char* out_ptr, int64_t n_bytes, int n_wires) { - // Copy our data to the appropriate place - std::memcpy(out_ptr + rank_ * n_bytes, in_ptr, n_bytes); - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS * 2 * 2; - size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); - size_t out_bytes = n_bytes * size_; - auto [sz, N] = buffer_size_from_message(n_bytes_per_wire); - int n_steps = (n_bytes_per_wire + N - 1) / N; - - // Counters to maintain the state of transfers - int in_flight = 0; - int64_t send_offset[2]; - int64_t recv_offset[2]; - int64_t limits[2]; - int send_count[2 * RING_MAX_CONNS] = {0}; - int recv_count[2 * RING_MAX_CONNS] = {0}; - send_offset[0] = send_offset[1] = rank_ * n_bytes; - recv_offset[0] = ((rank_ + size_ - 1) % size_) * n_bytes; - recv_offset[1] = ((rank_ + 1) % size_) * n_bytes; - limits[0] = n_wires * n_bytes_per_wire; - limits[1] = n_bytes; - - // Possible perf improvement by not syncing at every step but running ahead - // as needed. - for (int k = 0; k < size_ - 1; k++) { - // Prefill the pipeline - int buff = 0; - while (buff < n_steps && buff < PIPELINE) { - post_recv_all(sz, buff); - for (int lr = 0; lr < 2; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - int64_t offset = lw * N + - send_count[lr * RING_MAX_CONNS + lw] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_count[lr * RING_MAX_CONNS + lw]++; - } - } - post_send_all(sz, buff); - - buff++; - in_flight += 2 * 2 * n_wires; - } - - // Main loop - // - // Keep going until we have no longer data in flight. - while (in_flight > 0) { - ibv_wc wc[WC_NUM]; - int n = poll(left_, right_, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int work_type = wc[i].wr_id >> 16; - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lr = wire / RING_MAX_CONNS; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (work_type == SEND_WR && send_count[wire] < n_steps) { - int64_t offset = lw * N + send_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - out_ptr + send_offset[lr] + offset, - out_ptr + send_offset[lr] + - std::max(offset, std::min(offset + N, limits[lr])), - send_buffer(sz, buff, lr, lw).begin()); - send_to(sz, buff, lr, lw); - in_flight++; - send_count[wire]++; - } - - else if (work_type == RECV_WR) { - int64_t offset = lw * N + recv_count[wire] * n_wires * N + - lr * n_wires * n_bytes_per_wire; - std::copy( - recv_buffer(sz, buff, lr, lw).begin(), - recv_buffer(sz, buff, lr, lw).begin() + - std::max(0, std::min(N, limits[lr] - offset)), - out_ptr + recv_offset[lr] + offset); - recv_count[wire]++; - if (recv_count[wire] + (PIPELINE - 1) < n_steps) { - recv_from(sz, buff, lr, lw); - in_flight++; - } - } - } - } - - send_offset[0] = (send_offset[0] + out_bytes - n_bytes) % out_bytes; - recv_offset[0] = (recv_offset[0] + out_bytes - n_bytes) % out_bytes; - send_offset[1] = (send_offset[1] + n_bytes) % out_bytes; - recv_offset[1] = (recv_offset[1] + n_bytes) % out_bytes; - for (int i = 0; i < 2 * RING_MAX_CONNS; i++) { - send_count[i] = recv_count[i] = 0; - } - } - } - - void send(const char* in_ptr, int64_t n_bytes, int dst, int n_wires) { - int left = (rank_ + size_ - 1) % size_; - - // In the case that size_ == 2 then left == right so we bias send towards - // left and recv towards right so that the selections will be correct for - // the 2 node case. - auto& conns = (dst == left) ? left_ : right_; - int dir = dst == left; - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS; - - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; - auto [sz, N] = buffer_size_from_message(bytes_per_wire); - - int in_flight = 0; - int64_t read_offset[RING_MAX_CONNS]; - int64_t limits[RING_MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - read_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } - - // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (read_offset[lw] < limits[lw] && buff < PIPELINE) { - std::copy( - in_ptr + read_offset[lw], - in_ptr + std::min(read_offset[lw] + N, limits[lw]), - send_buffer(sz, buff, dir, lw).begin()); - send_to(sz, buff, dir, lw); - - buff++; - read_offset[lw] += N; - in_flight++; - } - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a send was completed and we have more data to send then go ahead - // and send them. - ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - if (read_offset[lw] < limits[lw]) { - std::copy( - in_ptr + read_offset[lw], - in_ptr + std::min(read_offset[lw] + N, limits[lw]), - send_buffer(sz, buff, dir, lw).begin()); - send_to(sz, buff, dir, lw); - - read_offset[lw] += N; - in_flight++; - } - } - } - } - - void recv(char* out_ptr, int64_t n_bytes, int src, int n_wires) { - int right = (rank_ + 1) % size_; - - // In the case that size_ == 2 then left == right so we bias send towards - // left and recv towards right so that the selections will be correct for - // the 2 node case. - auto& conns = (src == right) ? right_ : left_; - int dir = src == right; - - constexpr int PIPELINE = 2; - constexpr int WC_NUM = PIPELINE * RING_MAX_CONNS; - - int64_t bytes_per_wire = (n_bytes + n_wires - 1) / n_wires; - auto [sz, N] = buffer_size_from_message(bytes_per_wire); - - int in_flight = 0; - int64_t write_offset[RING_MAX_CONNS]; - int64_t limits[RING_MAX_CONNS]; - for (int lw = 0; lw < n_wires; lw++) { - write_offset[lw] = std::min(lw * bytes_per_wire, n_bytes); - limits[lw] = std::min((lw + 1) * bytes_per_wire, n_bytes); - } - - // Prefill the pipeline - for (int lw = 0; lw < n_wires; lw++) { - int buff = 0; - while (N * buff < limits[lw] && buff < PIPELINE) { - recv_from(sz, buff, dir, lw); - - buff++; - in_flight++; - } - } - - // Main loop - while (in_flight > 0) { - // Poll the hardware for completions. - // - // If a recv was completed copy it to the output and if we have more - // data to fetch post another recv. - ibv_wc wc[WC_NUM]; - int n = poll(conns, WC_NUM, wc); - for (int i = 0; i < n; i++) { - int buff = (wc[i].wr_id >> 8) & 0xff; - int wire = wc[i].wr_id & 0xff; - int lw = wire % RING_MAX_CONNS; - - in_flight--; - - std::copy( - recv_buffer(sz, buff, dir, lw).begin(), - recv_buffer(sz, buff, dir, lw).begin() + - std::max( - 0, std::min(limits[lw] - write_offset[lw], N)), - out_ptr + write_offset[lw]); - write_offset[lw] += N; - - if (write_offset[lw] + (PIPELINE - 1) * N < limits[lw]) { - recv_from(sz, buff, dir, lw); - - in_flight++; - } - } - } - } - - private: - void send_to(int sz, int buff, int left_right, int wire) { - if (left_right) { - left_[wire].post_send( - send_buffer_left(sz, buff, wire), - SEND_WR << 16 | buff << 8 | (RING_MAX_CONNS + wire)); - } else { - right_[wire].post_send( - send_buffer_right(sz, buff, wire), SEND_WR << 16 | buff << 8 | wire); - } - } - - void recv_from(int sz, int buff, int left_right, int wire) { - if (left_right) { - right_[wire].post_recv( - recv_buffer_right(sz, buff, wire), - RECV_WR << 16 | buff << 8 | (RING_MAX_CONNS + wire)); - } else { - left_[wire].post_recv( - recv_buffer_left(sz, buff, wire), RECV_WR << 16 | buff << 8 | wire); - } - } - - SharedBuffer& send_buffer_right(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + wire]; - } - - SharedBuffer& send_buffer_left(int sz, int buff, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + n_conns_ + - wire]; - } - - SharedBuffer& send_buffer(int sz, int buff, int left_right, int wire) { - return send_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + - left_right * n_conns_ + wire]; - } - - SharedBuffer& recv_buffer_left(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + wire]; - } - - SharedBuffer& recv_buffer_right(int sz, int buff, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + n_conns_ + - wire]; - } - - SharedBuffer& recv_buffer(int sz, int buff, int left_right, int wire) { - return recv_buffers_ - [sz * NUM_BUFFERS * n_conns_ * 2 + buff * n_conns_ * 2 + - left_right * n_conns_ + wire]; - } - - template - void post_recv_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - recv_from(sz, buff, lr, lw); - } - } - } - - void post_recv_all(int sz, int buff) { - post_recv_all<2>(sz, buff, n_conns_); - } - - template - void post_send_all(int sz, int buff, int n_wires) { - for (int lr = 0; lr < MAX_DIR; lr++) { - for (int lw = 0; lw < n_wires; lw++) { - send_to(sz, buff, lr, lw); - } - } - } - - void post_send_all(int sz, int buff) { - post_send_all<2>(sz, buff, n_conns_); - } - - int rank_; - int size_; - int n_conns_; - std::span left_; - std::span right_; - std::span send_buffers_; - std::span recv_buffers_; -}; - -} // namespace mlx::core::distributed::jaccl diff --git a/Source/Cmlx/mlx/mlx/distributed/mpi/mpi.cpp b/Source/Cmlx/mlx/mlx/distributed/mpi/mpi.cpp index 3b176e6e6..ea3960edc 100644 --- a/Source/Cmlx/mlx/mlx/distributed/mpi/mpi.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/mpi/mpi.cpp @@ -166,6 +166,11 @@ struct MPIWrapper { if (!is_available()) { return false; } + // MPI_Init is an error to call twice, and init() can run more than once + // when it returns without a group. + if (initialized_) { + return true; + } bool success = init(nullptr, nullptr) == MPI_SUCCESS; // Initialize custom types and ops @@ -495,6 +500,19 @@ std::shared_ptr init(bool strict /* = false */) { return nullptr; } + // Open MPI initializes a world of size 1 for a program that was not started + // with mpirun, which is not a distributed group. + int size = 1; + mpi().size(mpi().world(), &size); + if (size <= 1) { + if (strict) { + throw std::runtime_error( + "[mpi] The world has a single process. Launch with mpirun to " + "initialize the mpi backend."); + } + return nullptr; + } + return std::make_shared(mpi().world(), true); } diff --git a/Source/Cmlx/mlx/mlx/distributed/nccl/CMakeLists.txt b/Source/Cmlx/mlx/mlx/distributed/nccl/CMakeLists.txt index 644d060d7..d67f18e48 100644 --- a/Source/Cmlx/mlx/mlx/distributed/nccl/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/distributed/nccl/CMakeLists.txt @@ -1,25 +1,11 @@ -if(MLX_BUILD_CUDA) - target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp) +if(MLX_BUILD_CUDA AND NOT WIN32) find_package(NCCL) if(NCCL_FOUND) target_link_libraries(mlx PRIVATE ${NCCL_LIBRARIES}) target_include_directories(mlx PRIVATE ${NCCL_INCLUDE_DIRS}) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nccl.cpp) else() - message( - STATUS - "NCCL not found, using stubs. To run distributed with NCCL backend, install NCCL." - ) - file( - DOWNLOAD - "https://raw.githubusercontent.com/NVIDIA/nccl/refs/tags/v2.27.5-1/src/nccl.h.in" - "${CMAKE_CURRENT_BINARY_DIR}/nccl.h") - add_library(nccl_stub OBJECT - ${CMAKE_CURRENT_SOURCE_DIR}/nccl_stub/nccl_stubs.cpp) - target_include_directories( - nccl_stub - PRIVATE ${CUDAToolkit_INCLUDE_DIRS} - PUBLIC ${CMAKE_CURRENT_BINARY_DIR}) - target_link_libraries(mlx PRIVATE $) + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) endif() else() target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/no_nccl.cpp) diff --git a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl.cpp b/Source/Cmlx/mlx/mlx/distributed/nccl/nccl.cpp index d8244bf94..faaee3183 100644 --- a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/nccl/nccl.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include "mlx/backend/cuda/device.h" @@ -267,28 +268,54 @@ inline void bootstrap_unique_id( } // namespace detail +// helper struct to manage communicator +struct NCCLComm { + ncclComm_t comm; + int rank_; + int size_; + + NCCLComm(ncclComm_t c, int rank, int size) + : comm(c), rank_(rank), size_(size) {} + + static std::shared_ptr + create(int numRanks, int rank, ncclUniqueId commId) { + ncclComm_t raw; + CHECK_NCCL(ncclCommInitRank(&raw, numRanks, commId, rank)); + return std::make_shared(raw, rank, numRanks); + } + + static std::shared_ptr split(NCCLComm* source, int color, int key) { + ncclComm_t raw; + // default config, blocking comm creation + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + CHECK_NCCL(ncclCommSplit(source->comm, color, key, &raw, &config)); + int new_rank, new_size; + CHECK_NCCL(ncclCommUserRank(raw, &new_rank)); + CHECK_NCCL(ncclCommCount(raw, &new_size)); + return std::make_shared(raw, new_rank, new_size); + } + + NCCLComm(const NCCLComm&) = delete; + NCCLComm& operator=(const NCCLComm&) = delete; +}; + using GroupImpl = mlx::core::distributed::detail::GroupImpl; class NCCLGroup : public GroupImpl { public: NCCLGroup(int worldRank, int worldSize, const std::string initMethod) - : rank_(worldRank), - size_(worldSize), - comm_(nullptr), - initMethod_(initMethod) { + : rank_(worldRank), size_(worldSize), initMethod_(initMethod) { if (initialized_) return; int ndev; CHECK_CUDA(cudaGetDeviceCount(&ndev)); CHECK_CUDA(cudaSetDevice(rank_ % ndev)); detail::bootstrap_unique_id(uniqueId_, rank_, size_, initMethod_); - CHECK_NCCL(ncclCommInitRank(&comm_, size_, uniqueId_, rank_)); + comm_ = NCCLComm::create(size_, rank_, uniqueId_); initialized_ = true; } - - ~NCCLGroup() { - ncclCommDestroy(comm_); - initialized_ = false; - } + // Used by split() to wrap an already-created communicator + NCCLGroup(std::shared_ptr comm, int rank, int size) + : rank_(rank), size_(size), comm_(std::move(comm)) {} Stream communication_stream(StreamOrDevice s) override { return to_stream(s, Device::gpu); @@ -309,8 +336,11 @@ class NCCLGroup : public GroupImpl { }); } - virtual std::shared_ptr split(int color, int key = -1) override { - throw std::runtime_error("[nccl] Group split not supported."); + std::shared_ptr split(int color, int key = -1) override { + key = (key < 0) ? rank() : key; + auto new_comm = NCCLComm::split(comm_.get(), color, key); + return std::make_shared( + new_comm, new_comm->rank_, new_comm->size_); } void all_gather(const array& input, array& output, Stream stream) override { @@ -322,7 +352,7 @@ class NCCLGroup : public GroupImpl { gpu_ptr(output), input.size(), dt, - comm_, + comm_->comm, encoder.stream())); }); } @@ -371,7 +401,7 @@ class NCCLGroup : public GroupImpl { input.size(), dt, op, - comm_, + comm_->comm, encoder.stream())); } @@ -390,14 +420,15 @@ class NCCLGroup : public GroupImpl { output.size(), dt, op, - comm_, + comm_->comm, encoder.stream())); } - int rank_, size_; + int rank_; + int size_; std::string initMethod_; ncclUniqueId uniqueId_; - ncclComm_t comm_; + std::shared_ptr comm_; bool initialized_ = false; }; diff --git a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt b/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt deleted file mode 100644 index 8b1378917..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp b/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp deleted file mode 100644 index cd1b52229..000000000 --- a/Source/Cmlx/mlx/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#if !defined(_WIN32) -#include -#endif -#include -#include - -ncclResult_t ncclGetUniqueId(ncclUniqueId*) { - return ncclSuccess; -} - -const char* ncclGetErrorString(ncclResult_t result) { - return nullptr; -} - -ncclResult_t -ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank) { - return ncclSuccess; -} - -ncclResult_t ncclCommDestroy(ncclComm_t comm) { - return ncclSuccess; -} - -ncclResult_t ncclAllGather( - const void* sendbuff, - void* recvbuff, - size_t sendcount, - ncclDataType_t datatype, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} - -ncclResult_t ncclAllReduce( - const void* sendbuff, - void* recvbuff, - size_t count, - ncclDataType_t datatype, - ncclRedOp_t op, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} - -ncclResult_t ncclReduceScatter( - const void* sendbuff, - void* recvbuff, - size_t recvcount, - ncclDataType_t datatype, - ncclRedOp_t op, - ncclComm_t comm, - cudaStream_t stream) { - return ncclSuccess; -} diff --git a/Source/Cmlx/mlx/mlx/distributed/ops.cpp b/Source/Cmlx/mlx/mlx/distributed/ops.cpp index 1762f0e6b..3b338dda6 100644 --- a/Source/Cmlx/mlx/mlx/distributed/ops.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/ops.cpp @@ -20,60 +20,44 @@ Group to_group(std::optional group) { } } -} // namespace - -array all_sum( +array all_reduce( const array& x, - std::optional group_ /* = std::nullopt */, - StreamOrDevice s /* = {} */) { + const std::optional& group_, + StreamOrDevice s, + AllReduce::ReduceType reduce_type) { auto group = to_group(group_); - if (group.size() == 1) { return x; } auto stream = detail::communication_stream(group, s); - return array( x.shape(), x.dtype(), - std::make_shared(stream, group, AllReduce::Sum), + std::make_shared(stream, group, reduce_type), {x}); } -array all_max( +} // namespace + +array all_sum( const array& x, std::optional group_ /* = std::nullopt */, StreamOrDevice s /* = {} */) { - auto group = to_group(group_); - - if (group.size() == 1) { - return x; - } - auto stream = detail::communication_stream(group, s); + return all_reduce(x, group_, s, AllReduce::Sum); +} - return array( - x.shape(), - x.dtype(), - std::make_shared(stream, group, AllReduce::Max), - {x}); +array all_max( + const array& x, + std::optional group_ /* = std::nullopt */, + StreamOrDevice s /* = {} */) { + return all_reduce(x, group_, s, AllReduce::Max); } array all_min( const array& x, std::optional group_ /* = std::nullopt */, StreamOrDevice s /* = {} */) { - auto group = to_group(group_); - - if (group.size() == 1) { - return x; - } - auto stream = detail::communication_stream(group, s); - - return array( - x.shape(), - x.dtype(), - std::make_shared(stream, group, AllReduce::Min), - {x}); + return all_reduce(x, group_, s, AllReduce::Min); } array all_gather( @@ -165,11 +149,11 @@ array sum_scatter( if (group.size() == 1) { return x; } - if (x.shape()[0] % group.size() != 0) { + if (x.ndim() == 0 || x.shape()[0] % group.size() != 0) { std::ostringstream msg; msg << "[sum_scatter] Invalid shape=" << x.shape() << " for a group of size " << group.size() - << ". The first dimension (axis 0) must be divisible by the group size."; + << ". The first dimension (axis 0) must be divisible by the group size and input mst have at least one."; throw std::invalid_argument(msg.str()); } diff --git a/Source/Cmlx/mlx/mlx/distributed/ring/ring.cpp b/Source/Cmlx/mlx/mlx/distributed/ring/ring.cpp index ea4004284..9a81010e3 100644 --- a/Source/Cmlx/mlx/mlx/distributed/ring/ring.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/ring/ring.cpp @@ -5,12 +5,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include @@ -21,72 +23,13 @@ #include "mlx/distributed/distributed_impl.h" #include "mlx/distributed/reduction_ops.h" #include "mlx/distributed/utils.h" +#include "mlx/dtype_utils.h" #include "mlx/threadpool.h" #ifndef SOL_TCP #define SOL_TCP IPPROTO_TCP #endif -#define SWITCH_TYPE(x, ...) \ - switch ((x).dtype()) { \ - case bool_: { \ - using T = bool; \ - __VA_ARGS__; \ - } break; \ - case int8: { \ - using T = int8_t; \ - __VA_ARGS__; \ - } break; \ - case int16: { \ - using T = int16_t; \ - __VA_ARGS__; \ - } break; \ - case int32: { \ - using T = int32_t; \ - __VA_ARGS__; \ - } break; \ - case int64: { \ - using T = int64_t; \ - __VA_ARGS__; \ - } break; \ - case uint8: { \ - using T = uint8_t; \ - __VA_ARGS__; \ - } break; \ - case uint16: { \ - using T = uint16_t; \ - __VA_ARGS__; \ - } break; \ - case uint32: { \ - using T = uint32_t; \ - __VA_ARGS__; \ - } break; \ - case uint64: { \ - using T = uint64_t; \ - __VA_ARGS__; \ - } break; \ - case bfloat16: { \ - using T = bfloat16_t; \ - __VA_ARGS__; \ - } break; \ - case float16: { \ - using T = float16_t; \ - __VA_ARGS__; \ - } break; \ - case float32: { \ - using T = float; \ - __VA_ARGS__; \ - } break; \ - case float64: { \ - using T = double; \ - __VA_ARGS__; \ - } break; \ - case complex64: { \ - using T = complex64_t; \ - __VA_ARGS__; \ - } break; \ - } - namespace mlx::core::distributed::ring { constexpr const size_t ALL_SUM_SIZE = 8 * 1024 * 1024; @@ -94,6 +37,10 @@ constexpr const size_t ALL_SUM_BUFFERS = 2; constexpr const int CONN_ATTEMPTS = 5; constexpr const int CONN_WAIT = 1000; constexpr const char* RING_TAG = "[ring]"; +// send(2) and recv(2) reject a length above INT_MAX with EINVAL, so a single +// transfer of 2 GiB or more fails outright rather than being carried in +// pieces. +constexpr const size_t MAX_IO_BYTES = 1024 * 1024 * 1024; using GroupImpl = mlx::core::distributed::detail::GroupImpl; using json = nlohmann::json; @@ -232,12 +179,16 @@ class SocketThread { if (!recvs_.empty()) { auto& task = recvs_.front(); - ssize_t r = ::recv(fd_, task.buffer, task.size, 0); + ssize_t r = + ::recv(fd_, task.buffer, std::min(task.size, MAX_IO_BYTES), 0); if (r > 0) { task.buffer = static_cast(task.buffer) + r; task.size -= r; delete_recv = task.size == 0; error_count = 0; + } else if (r == 0) { + error_count++; + log_info(true, "Socket", fd_, "was closed by the peer"); } else if (errno != EAGAIN) { error_count++; log_info( @@ -246,12 +197,16 @@ class SocketThread { } if (!sends_.empty()) { auto& task = sends_.front(); - ssize_t r = ::send(fd_, task.buffer, task.size, 0); + ssize_t r = + ::send(fd_, task.buffer, std::min(task.size, MAX_IO_BYTES), 0); if (r > 0) { task.buffer = static_cast(task.buffer) + r; task.size -= r; delete_send = task.size == 0; error_count = 0; + } else if (r == 0) { + error_count++; + log_info(true, "Sending to socket", fd_, "made no progress"); } else if (errno != EAGAIN) { error_count++; log_info(true, "Sending to socket", fd_, "failed with errno", errno); @@ -259,7 +214,18 @@ class SocketThread { } if (error_count >= 10) { - log_info(true, "Too many send/recv errors. Aborting..."); + log_info(true, "Too many send/recv errors. Failing pending tasks..."); + // Throw exception in invoker's thread. + auto error = std::make_exception_ptr( + std::runtime_error("[ring] connection to a peer was lost")); + for (auto& task : recvs_) { + task.promise.set_exception(error); + } + for (auto& task : sends_) { + task.promise.set_exception(error); + } + recvs_.clear(); + sends_.clear(); return; } } @@ -475,18 +441,24 @@ class RingGroup : public GroupImpl { } void all_sum(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::SumOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::SumOp()); + }); } void all_max(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::MaxOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MaxOp()); + }); } void all_min(const array& input, array& output, Stream stream) override { - SWITCH_TYPE( - output, all_reduce(input, output, stream, detail::MinOp())); + dispatch_all_types(output.dtype(), [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + all_reduce(input, output, stream, detail::MinOp()); + }); } std::shared_ptr split(int color, int key = -1) override { @@ -525,7 +497,7 @@ class RingGroup : public GroupImpl { (i % 2) ? -1 : 1))); } for (auto& f : all_gathers) { - f.wait(); + f.get(); } }); } @@ -650,7 +622,7 @@ class RingGroup : public GroupImpl { reduce_op))); } for (auto& f : all_sums) { - f.wait(); + f.get(); } }); } @@ -737,8 +709,8 @@ class RingGroup : public GroupImpl { } if (j >= 0) { - sends[b].wait(); - recvs[b].wait(); + sends[b].get(); + recvs[b].get(); if (2 * j < send_plan.size()) { reduce_op( recv_buffers[j % ALL_SUM_BUFFERS], @@ -749,8 +721,13 @@ class RingGroup : public GroupImpl { std::swap(a, b); } - sends[b].wait(); - recvs[b].wait(); + // Check valid() to avoid consuming same future twice for single packet. + if (sends[b].valid()) { + sends[b].get(); + } + if (recvs[b].valid()) { + recvs[b].get(); + } } void all_gather_impl( @@ -784,8 +761,8 @@ class RingGroup : public GroupImpl { send_segment = (send_segment + size_ + direction) % size_; recv_segment = (recv_segment + size_ + direction) % size_; - sent.wait(); - recvd.wait(); + sent.get(); + recvd.get(); } } @@ -804,7 +781,7 @@ class RingGroup : public GroupImpl { std::min(data_size, (i + 1) * segment_size) - i * segment_size)); } for (auto& f : sends) { - f.wait(); + f.get(); } } @@ -822,7 +799,7 @@ class RingGroup : public GroupImpl { std::min(data_size, (i + 1) * segment_size) - i * segment_size)); } for (auto& f : recvs) { - f.wait(); + f.get(); } } diff --git a/Source/Cmlx/mlx/mlx/distributed/utils.cpp b/Source/Cmlx/mlx/mlx/distributed/utils.cpp index 2de994cf4..11f8df85c 100644 --- a/Source/Cmlx/mlx/mlx/distributed/utils.cpp +++ b/Source/Cmlx/mlx/mlx/distributed/utils.cpp @@ -10,6 +10,16 @@ namespace mlx::core::distributed::detail { +namespace { + +[[noreturn]] void throw_socket_error(const char* tag, const char* action) { + std::ostringstream msg; + msg << tag << " " << action << " (error: " << errno << ")"; + throw std::runtime_error(msg.str()); +} + +} // namespace + /** * Parse a sockaddr from an ip and port provided as strings. */ @@ -53,9 +63,7 @@ address_t parse_address(const std::string& ip_port) { TCPSocket::TCPSocket(const char* tag) { sock_ = socket(AF_INET, SOCK_STREAM, 0); if (sock_ < 0) { - std::ostringstream msg; - msg << tag << " Couldn't create socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't create socket"); } } @@ -95,40 +103,30 @@ void TCPSocket::listen(const char* tag, const address_t& addr) { int enable = 1; success = setsockopt(sock_, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't enable reuseaddr (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't enable reuseaddr"); } success = setsockopt(sock_, SOL_SOCKET, SO_REUSEPORT, &enable, sizeof(int)); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't enable reuseport (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't enable reuseport"); } // Bind the socket to the address and port success = bind(sock_, addr.get(), addr.len); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't bind socket (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't bind socket"); } // Prepare waiting for connections success = ::listen(sock_, 0); if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't listen (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't listen"); } } TCPSocket TCPSocket::accept(const char* tag) { int peer = ::accept(sock_, nullptr, nullptr); if (peer < 0) { - std::ostringstream msg; - msg << tag << " Accept failed (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Accept failed"); } return TCPSocket(peer); @@ -173,10 +171,7 @@ TCPSocket TCPSocket::connect( // Create the socket sock = socket(AF_INET, SOCK_STREAM, 0); if (sock < 0) { - std::ostringstream msg; - msg << tag << " Couldn't create socket to connect (error: " << errno - << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't create socket to connect"); } success = ::connect(sock, addr.get(), addr.len); @@ -195,9 +190,7 @@ TCPSocket TCPSocket::connect( } if (success < 0) { - std::ostringstream msg; - msg << tag << " Couldn't connect (error: " << errno << ")"; - throw std::runtime_error(msg.str()); + throw_socket_error(tag, "Couldn't connect"); } return TCPSocket(sock); diff --git a/Source/Cmlx/mlx/mlx/einsum.cpp b/Source/Cmlx/mlx/mlx/einsum.cpp index 629088772..dcaa8c51e 100644 --- a/Source/Cmlx/mlx/mlx/einsum.cpp +++ b/Source/Cmlx/mlx/mlx/einsum.cpp @@ -95,10 +95,14 @@ std::pair, std::string> parse(std::string subscripts) { std::sort(rhs.begin(), rhs.end()); } std::vector input_list; - std::stringstream ss(lhs); - std::string token; - while (getline(ss, token, ',')) { - input_list.push_back(token); + for (size_t start = 0;;) { + auto pos = lhs.find(',', start); + if (pos == std::string::npos) { + input_list.push_back(lhs.substr(start)); + break; + } + input_list.push_back(lhs.substr(start, pos - start)); + start = pos + 1; } return {input_list, rhs}; } @@ -356,7 +360,7 @@ array batch_tensordot( std::vector b_batch, std::vector b_concat, StreamOrDevice s) { - // Broadcast contracting dimensions + // Broadcast contracting and batch dimensions. { auto a_shape = a.shape(); auto b_shape = b.shape(); @@ -365,6 +369,11 @@ array batch_tensordot( a_shape[a_contract[i]] = d; b_shape[b_contract[i]] = d; } + for (int i = 0; i < a_batch.size(); ++i) { + auto d = std::max(a.shape(a_batch[i]), b.shape(b_batch[i])); + a_shape[a_batch[i]] = d; + b_shape[b_batch[i]] = d; + } a = broadcast_to(a, a_shape, s); b = broadcast_to(b, b_shape, s); } @@ -800,7 +809,7 @@ std::pair, PathInfo> einsum_path_helper( max_size = std::max(max_size, term_size(in, dim_map)); } - PathInfo path_info; + PathInfo path_info{}; // Get the full naive cost std::tie(path_info.naive_cost, path_info.naive_scaling) = @@ -813,6 +822,8 @@ std::pair, PathInfo> einsum_path_helper( std::iota(positions.begin(), positions.end(), 0); path.emplace_back( std::move(inputs), std::move(output), std::move(positions)); + path_info.optimized_cost = path_info.naive_cost; + path_info.optimized_scaling = path_info.naive_scaling; } else { std::tie(path, path_info.optimized_cost, path_info.optimized_scaling) = greedy_path(inputs, output, dim_map, path_info.naive_cost, max_size); diff --git a/Source/Cmlx/mlx/mlx/error.h b/Source/Cmlx/mlx/mlx/error.h new file mode 100644 index 000000000..ba1164f19 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/error.h @@ -0,0 +1,49 @@ +// 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 diff --git a/Source/Cmlx/mlx/mlx/event.h b/Source/Cmlx/mlx/mlx/event.h index 66a6a75df..cf2d5cc7d 100644 --- a/Source/Cmlx/mlx/mlx/event.h +++ b/Source/Cmlx/mlx/mlx/event.h @@ -5,6 +5,7 @@ #include #include +#include "mlx/error.h" #include "mlx/stream.h" namespace mlx::core { @@ -26,6 +27,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; @@ -47,7 +68,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/mlx/mlx/export.cpp b/Source/Cmlx/mlx/mlx/export.cpp index 9a773cd7a..bcc0ad161 100644 --- a/Source/Cmlx/mlx/mlx/export.cpp +++ b/Source/Cmlx/mlx/mlx/export.cpp @@ -5,6 +5,7 @@ #include "mlx/fast_primitives.h" #include "mlx/graph_utils.h" #include "mlx/primitives.h" +#include "mlx/transforms_impl.h" #include "mlx/utils.h" #include "mlx/version.h" @@ -120,6 +121,9 @@ void serialize(Writer& os, T v) { } else if constexpr (std::is_enum_v) { serialize(os, static_cast(v)); } else if constexpr (std::is_same_v) { + } else if constexpr (std::is_same_v) { + serialize(os, static_cast(v.size())); + os.write(v.data(), v.size()); } else if constexpr (is_iterable) { serialize(os, static_cast(v.size())); for (const auto& t : v) { @@ -155,6 +159,19 @@ T deserialize(Reader& is) { return static_cast(deserialize(is)); } else if constexpr (std::is_same_v) { return nullptr; + } else if constexpr (std::is_same_v) { + auto size = deserialize(is); + // Bound the allocation by what the file can still contain + auto pos = is.tell(); + is.seek(0, std::ios_base::end); + auto remaining = static_cast(is.tell() - pos); + is.seek(pos); + if (size > remaining) { + throw std::runtime_error("[import_function] Invalid string size."); + } + T v(size, '\0'); + is.read(v.data(), size); + return v; } else if constexpr (is_iterable) { T v; auto size = deserialize(is); @@ -169,7 +186,7 @@ T deserialize(Reader& is) { } else if constexpr (is_optional) { auto has_value = deserialize(is); if (has_value) { - return deserialize(is); + return T{deserialize(is)}; } else { return std::nullopt; } @@ -279,6 +296,8 @@ void extract_state(const T state, std::vector& unpacked_state) { unpacked_state.push_back(state); } else if constexpr (std::is_enum_v) { unpacked_state.push_back(static_cast(state)); + } else if constexpr (std::is_same_v) { + unpacked_state.push_back(state); } else if constexpr (is_iterable) { unpacked_state.push_back(state); } else if constexpr (is_pair || is_tuple) { @@ -411,6 +430,7 @@ struct PrimitiveFactory { "CumLogaddexp"), SERIALIZE_PRIMITIVE(Scatter), SERIALIZE_PRIMITIVE(ScatterAxis), + SERIALIZE_PRIMITIVE(SearchSorted), SERIALIZE_PRIMITIVE(Select), SERIALIZE_PRIMITIVE(Sigmoid), SERIALIZE_PRIMITIVE(Sign), @@ -446,6 +466,7 @@ struct PrimitiveFactory { SERIALIZE_PRIMITIVE(ScaledDotProductAttention), SERIALIZE_PRIMITIVE(CustomKernel)}; std::unordered_map name_remap; + std::unordered_map stream_map; PrimitiveFactory() { for (auto& [n, f] : factory) { @@ -471,13 +492,25 @@ struct PrimitiveFactory { } }; - std::shared_ptr load(Reader& is) { - auto stream = deserialize(is); - if (get_stream(stream.index) != stream) { - std::ostringstream msg; - msg << "[import_function] Invalid stream encountered " << stream << "."; - throw std::invalid_argument(msg.str()); + Stream resolve_stream(const Stream& stream) { + if (auto it = stream_map.find(stream.index); it != stream_map.end()) { + return it->second; } + // Try to find an existing stream on the same device + for (auto& s : get_streams()) { + if (s.device == stream.device) { + stream_map.emplace(stream.index, s); + return s; + } + } + // No stream on that device, make a new one + Stream s = new_stream(stream.device); + stream_map.emplace(stream.index, s); + return s; + } + + std::shared_ptr load(Reader& is) { + auto stream = resolve_stream(deserialize(is)); auto name = deserialize(is); if (auto it = factory.find(name); it != factory.end()) { return it->second.deserialize(is, stream); @@ -505,10 +538,15 @@ struct PrimitiveFactory { }; }; -void write_header(Writer& os, int count, bool shapeless) { +void write_header( + Writer& os, + int count, + bool shapeless, + const std::string& metadata) { serialize(os, std::string(version())); serialize(os, count); serialize(os, shapeless); + serialize(os, metadata); } // A struct to hold and retrieve the graphs that are exported / imported @@ -658,14 +696,16 @@ FunctionTable::Function* FunctionTable::find( FunctionExporter::FunctionExporter( const std::string& file, std::function(const Args&, const Kwargs&)> fun, - bool shapeless) + bool shapeless, + std::string metadata) : os(file), fun(std::move(fun)), - ftable(std::make_shared(shapeless)) { + ftable(std::make_shared(shapeless)), + metadata_(std::move(metadata)) { if (!os.is_open()) { throw std::runtime_error("[export_function] Failed to open " + file); } - write_header(os, count, shapeless); + write_header(os, count, shapeless, metadata_); } FunctionExporter::FunctionExporter( @@ -774,8 +814,14 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) { }; // Trace to build the graph - auto [trace_inputs, trace_outputs, extra] = - detail::compile_trace(flat_fun, inputs, ftable->shapeless); + std::vector trace_inputs; + std::vector trace_outputs; + std::shared_ptr extra; + { + detail::InExportTracing export_tracing; + std::tie(trace_inputs, trace_outputs, extra) = + detail::compile_trace(flat_fun, inputs, ftable->shapeless); + } // DFS the graph and get the tape auto [tape, parents_map] = @@ -797,7 +843,7 @@ void FunctionExporter::export_function(const Args& args, const Kwargs& kwargs) { // Update the header auto pos = os.tell(); os.seek(0); - write_header(os, count, ftable->shapeless); + write_header(os, count, ftable->shapeless, metadata_); os.seek(pos); serialize(os, kwarg_keys); @@ -876,44 +922,51 @@ void FunctionExporter::operator()(const Args& args, const Kwargs& kwargs) { FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { return FunctionExporter{ file, [fun](const Args& args, const Kwargs&) { return fun(args); }, - shapeless}; + shapeless, + metadata}; } FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, - bool shapeless /* = false */) { + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { return exporter( file, [fun](const Args&, const Kwargs kwargs) { return fun(kwargs); }, - shapeless); + shapeless, + metadata); } FunctionExporter exporter( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, - bool shapeless /* = false */) { - return FunctionExporter{file, fun, shapeless}; + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + return FunctionExporter{file, fun, shapeless, metadata}; } void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(args); } void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(kwargs); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(kwargs); } void export_function( @@ -921,8 +974,9 @@ void export_function( const std::function(const Args&, const Kwargs&)>& fun, const Args& args, const Kwargs& kwargs, - bool shapeless /* = false */) { - exporter(file, fun, shapeless)(args, kwargs); + bool shapeless /* = false */, + const std::string& metadata /* = "" */) { + exporter(file, fun, shapeless, metadata)(args, kwargs); } FunctionExporter exporter( @@ -1032,6 +1086,7 @@ ImportedFunction::ImportedFunction(const std::string& file) auto mlx_version = deserialize(is); auto function_count = deserialize(is); ftable->shapeless = deserialize(is); + metadata_ = deserialize(is); std::unordered_map constants; auto import_one = [&]() { diff --git a/Source/Cmlx/mlx/mlx/export.h b/Source/Cmlx/mlx/mlx/export.h index 5532f7c81..8135919b2 100644 --- a/Source/Cmlx/mlx/mlx/export.h +++ b/Source/Cmlx/mlx/mlx/export.h @@ -50,17 +50,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. @@ -69,20 +72,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/mlx/mlx/export_impl.h b/Source/Cmlx/mlx/mlx/export_impl.h index 467a5f0d6..7fd053a90 100644 --- a/Source/Cmlx/mlx/mlx/export_impl.h +++ b/Source/Cmlx/mlx/mlx/export_impl.h @@ -27,17 +27,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&, @@ -57,7 +60,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, @@ -77,6 +81,7 @@ struct MLX_API FunctionExporter { int count{0}; bool closed{false}; std::shared_ptr ftable; + std::string metadata_; }; struct MLX_API ImportedFunction { @@ -88,12 +93,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/mlx/mlx/fast.cpp b/Source/Cmlx/mlx/mlx/fast.cpp index d6065e536..d6321ed36 100644 --- a/Source/Cmlx/mlx/mlx/fast.cpp +++ b/Source/Cmlx/mlx/mlx/fast.cpp @@ -187,6 +187,107 @@ bool RMSNormVJP::is_equivalent(const Primitive& other) const { return eps_ == a_other.eps_; } +array cross_entropy( + const array& logits, + const array& targets, + StreamOrDevice s_ /* = {} */) { + if (logits.ndim() < 1) { + throw std::invalid_argument( + "[cross_entropy] logits must have at least 1 dimension but got input " + "with 0 dimensions."); + } + auto expected = logits.shape(); + expected.pop_back(); + if (targets.shape() != expected) { + std::ostringstream msg; + msg << "[cross_entropy] targets shape " << targets.shape() + << " does not match logits shape " << logits.shape() + << " with the last axis removed."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(logits.dtype(), floating)) { + std::ostringstream msg; + msg << "[cross_entropy] Received unsupported logits type " << logits.dtype() + << "."; + throw std::invalid_argument(msg.str()); + } + if (!issubdtype(targets.dtype(), integer)) { + std::ostringstream msg; + msg << "[cross_entropy] targets must be integer class indices but got " + << targets.dtype() << "."; + throw std::invalid_argument(msg.str()); + } + + auto s = to_stream(s_); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto loss = subtract(logsumexp(x, -1, /* keepdims= */ false, s), score, s); + return std::vector{astype(loss, float32, s)}; + }; + + auto passed_targets = astype(targets, int32, s); + + if (!CrossEntropy::use_fallback(s)) { + return array( + expected, + float32, + std::make_shared(s, fallback), + {logits, passed_targets}); + } + return fallback({logits, passed_targets})[0]; +} + +std::vector CrossEntropy::vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) { + assert(primals.size() == 2); + assert(outputs.size() == 1); + assert(cotangents.size() == 1); + + for (auto arg : argnums) { + if (arg != 0) { + throw std::invalid_argument( + "[cross_entropy] Cannot differentiate with respect to the targets."); + } + } + + auto s = stream(); + auto fallback = [s](const std::vector& inputs) { + auto& x = inputs[0]; + auto& y = inputs[1]; + auto& loss = inputs[2]; + auto& g = inputs[3]; + + auto score = + squeeze(take_along_axis(x, expand_dims(y, -1, s), -1, s), -1, s); + auto lse = add(loss, astype(score, float32, s), s); + auto p = + exp(subtract(astype(x, float32, s), expand_dims(lse, -1, s), s), s); + Shape class_shape(x.ndim(), 1); + class_shape.back() = x.shape(-1); + auto onehot = astype( + equal( + expand_dims(y, -1, s), + reshape(arange(x.shape(-1), y.dtype(), s), class_shape, s), + s), + float32, + s); + auto gx = multiply(expand_dims(g, -1, s), subtract(p, onehot, s), s); + return std::vector{astype(gx, x.dtype(), s)}; + }; + + return {array( + primals[0].shape(), + primals[0].dtype(), + std::make_shared(s, fallback), + {primals[0], primals[1], outputs[0], cotangents[0]})}; +} + array layer_norm( const array& x, const std::optional& weight, @@ -332,7 +433,7 @@ std::vector LayerNorm::vjp( // df/db if (b.ndim() == 0) { - vjps.push_back(zeros_like(w, s)); + vjps.push_back(zeros_like(b, s)); } else { vjps.push_back(sum(g, axes, /* keepdims= */ false, s)); } @@ -406,6 +507,23 @@ array rope( if (offset.dtype().size() != 4) { inputs[1] = astype(offset, int32, s); } + if (dims <= 0) { + std::ostringstream msg; + msg << "[rope] dims must be positive but got " << dims << "."; + throw std::invalid_argument(msg.str()); + } + if (dims % 2 != 0) { + std::ostringstream msg; + msg << "[rope] dims must be even but got " << dims << "."; + throw std::invalid_argument(msg.str()); + } + if (dims > x.shape(-1)) { + std::ostringstream msg; + msg << "[rope] dims must not exceed the input's last dimension (" + << x.shape(-1) << ") but got " << dims << "."; + throw std::invalid_argument(msg.str()); + } + if (inputs.size() == 3 && (inputs[2].ndim() != 1 || inputs[2].shape(0) != dims / 2)) { std::ostringstream msg; @@ -601,7 +719,8 @@ array scaled_dot_product_attention( const std::string& mask_mode /* = "" */, std::optional mask_arr /* = {} */, const std::optional& sinks /* = {} */, - StreamOrDevice s /* = {}*/) { + bool force_fused /* = false */, + StreamOrDevice s /* = {} */) { for (const auto& tensor : {queries, keys, values}) { if (tensor.ndim() != 4) { std::ostringstream msg; @@ -817,6 +936,7 @@ array scaled_dot_product_attention( do_causal, is_training, output_logsumexp, + force_fused, stream)) { if (has_bool_mask && !ScaledDotProductAttention::supports_bool_mask()) { // Convert bool mask to additive mask. @@ -829,7 +949,13 @@ array scaled_dot_product_attention( } Shape out_shape{q.shape(0), q.shape(1), q.shape(2), v.shape(-1)}; auto primitive = std::make_shared( - stream, fallback, scale, do_causal, has_sinks, output_logsumexp); + stream, + fallback, + scale, + do_causal, + has_sinks, + output_logsumexp, + force_fused); if (output_logsumexp) { return array::make_arrays( {std::move(out_shape), Shape{q.shape(0), q.shape(1), q.shape(2), 1}}, @@ -895,7 +1021,8 @@ bool ScaledDotProductAttention::is_equivalent(const Primitive& other) const { static_cast(other); return scale_ == a_other.scale_ && do_causal_ == a_other.do_causal_ && has_sinks_ == a_other.has_sinks_ && - output_logsumexp_ == a_other.output_logsumexp_; + output_logsumexp_ == a_other.output_logsumexp_ && + force_fused_ == a_other.force_fused_; } bool ScaledDotProductAttentionVJP::is_equivalent(const Primitive& other) const { diff --git a/Source/Cmlx/mlx/mlx/fast.h b/Source/Cmlx/mlx/mlx/fast.h index 801d3893c..e79f300ce 100644 --- a/Source/Cmlx/mlx/mlx/fast.h +++ b/Source/Cmlx/mlx/mlx/fast.h @@ -6,6 +6,7 @@ #include #include "mlx/api.h" +#include "mlx/backend/common/metal_kernel.h" #include "mlx/utils.h" namespace mlx::core::fast { @@ -23,6 +24,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, @@ -52,6 +57,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; @@ -75,7 +81,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/mlx/mlx/fast_primitives.h b/Source/Cmlx/mlx/mlx/fast_primitives.h index 443483087..63021cd0c 100644 --- a/Source/Cmlx/mlx/mlx/fast_primitives.h +++ b/Source/Cmlx/mlx/mlx/fast_primitives.h @@ -3,6 +3,7 @@ #include #include +#include "mlx/backend/common/metal_kernel.h" #include "mlx/primitives.h" namespace mlx::core::fast { @@ -96,6 +97,67 @@ class RMSNormVJP : public Custom { float eps_; }; +// loss is always fp32 and the logits never have to be upcast in the graph. +class CrossEntropy : public Custom { + public: + CrossEntropy( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + static bool use_fallback(Stream stream); + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + std::vector vjp( + const std::vector& primals, + const std::vector& cotangents, + const std::vector& argnums, + const std::vector& outputs) override; + + DEFINE_NAME(CrossEntropy) + bool is_equivalent(const Primitive& other) const override { + return true; + } + std::vector output_shapes(const std::vector& inputs) override { + return {inputs[1].shape()}; + } + + auto state() const { + return std::monostate{}; + } +}; + +class CrossEntropyVJP : public Custom { + public: + CrossEntropyVJP( + Stream stream, + std::function(std::vector)> fallback) + : Custom(stream, std::move(fallback)) {} + + void eval_cpu(const std::vector& inputs, std::vector& outputs) + override { + throw std::runtime_error("NYI"); + } + void eval_gpu(const std::vector& inputs, std::vector& outputs) + override; + + DEFINE_NAME(CrossEntropyVJP) + bool is_equivalent(const Primitive& other) const override { + return true; + } + DEFINE_INPUT_OUTPUT_SHAPE() + + auto state() const { + return std::monostate{}; + } +}; + class LayerNorm : public Custom { public: LayerNorm( @@ -211,12 +273,14 @@ class ScaledDotProductAttention : public Custom { float scale, bool do_causal, bool has_sinks, - bool output_logsumexp) + bool output_logsumexp, + bool force_fused) : Custom(stream, std::move(fallback)), scale_(scale), do_causal_(do_causal), has_sinks_(has_sinks), - output_logsumexp_(output_logsumexp) {} + output_logsumexp_(output_logsumexp), + force_fused_(force_fused) {} static bool use_fallback( const array& q, @@ -227,6 +291,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal, bool is_training, bool output_logsumexp, + bool force_fused, Stream s); static bool supports_bool_mask(); @@ -250,7 +315,12 @@ class ScaledDotProductAttention : public Custom { DEFINE_INPUT_OUTPUT_SHAPE() auto state() const { return std::make_tuple( - nullptr, scale_, do_causal_, has_sinks_, output_logsumexp_); + nullptr, + scale_, + do_causal_, + has_sinks_, + output_logsumexp_, + force_fused_); } private: @@ -258,6 +328,7 @@ class ScaledDotProductAttention : public Custom { bool do_causal_; bool has_sinks_; bool output_logsumexp_; + bool force_fused_; }; class ScaledDotProductAttentionVJP : public Custom { @@ -375,7 +446,8 @@ class CustomKernel : public Primitive { std::optional init_value, std::vector scalar_arguments, bool is_precompiled, - int shared_memory) + int shared_memory, + CompileOptions::Data compile_options = {}) : Primitive(stream), name_(std::move(name)), source_(std::move(source)), @@ -386,7 +458,8 @@ class CustomKernel : public Primitive { init_value_(init_value), scalar_arguments_(std::move(scalar_arguments)), is_precompiled_(is_precompiled), - shared_memory_(shared_memory) {} + shared_memory_(shared_memory), + compile_options_(compile_options) {} void eval_cpu(const std::vector& inputs, std::vector& outputs) override { @@ -408,7 +481,8 @@ class CustomKernel : public Primitive { init_value_, scalar_arguments_, is_precompiled_, - shared_memory_); + shared_memory_, + compile_options_); } private: @@ -422,6 +496,7 @@ class CustomKernel : public Primitive { std::vector scalar_arguments_; bool is_precompiled_; int shared_memory_; + CompileOptions::Data compile_options_; }; } // namespace mlx::core::fast diff --git a/Source/Cmlx/mlx/mlx/fence.h b/Source/Cmlx/mlx/mlx/fence.h index 0ececdb6d..3fd5da333 100644 --- a/Source/Cmlx/mlx/mlx/fence.h +++ b/Source/Cmlx/mlx/mlx/fence.h @@ -32,8 +32,13 @@ class Fence { void update(Stream stream, const array& x, bool cross_device); void wait(Stream stream, const array& x); + template + auto& cast() const { + return *static_cast(fence_.get()); + } + private: - std::shared_ptr fence_{nullptr}; + std::shared_ptr fence_; }; } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/fft.cpp b/Source/Cmlx/mlx/mlx/fft.cpp index 6510faec1..06860a0e3 100644 --- a/Source/Cmlx/mlx/mlx/fft.cpp +++ b/Source/Cmlx/mlx/mlx/fft.cpp @@ -1,4 +1,6 @@ // Copyright © 2023 Apple Inc. +#include +#include #include #include @@ -9,12 +11,30 @@ namespace mlx::core::fft { +double fft_scale_factor(const Shape& n, FFTNorm norm, bool inverse) { + if (n.empty()) { + return 1.0; + } + double n_elements = + std::accumulate(n.begin(), n.end(), 1.0, std::multiplies()); + switch (norm) { + case FFTNorm::Backward: + return 1.0; + case FFTNorm::Ortho: + return inverse ? std::sqrt(n_elements) : 1.0 / std::sqrt(n_elements); + case FFTNorm::Forward: + return inverse ? n_elements : 1.0 / n_elements; + } + throw std::invalid_argument("[fftn] Invalid FFT normalization mode."); +} + array fft_impl( const array& a, Shape n, const std::vector& axes, bool real, bool inverse, + FFTNorm norm, StreamOrDevice s) { if (a.ndim() < 1) { throw std::invalid_argument( @@ -91,11 +111,16 @@ array fft_impl( auto in_type = real && !inverse ? float32 : complex64; auto out_type = real && inverse ? float32 : complex64; - return array( + auto out = array( out_shape, out_type, std::make_shared(to_stream(s), valid_axes, inverse, real), {astype(in, in_type, s)}); + auto scale = fft_scale_factor(n, norm, inverse); + if (scale != 1.0) { + return multiply(out, array(scale, out.dtype()), s); + } + return out; } array fft_impl( @@ -103,6 +128,7 @@ array fft_impl( const std::vector& axes, bool real, bool inverse, + FFTNorm norm, StreamOrDevice s) { Shape n; for (auto ax : axes) { @@ -111,88 +137,119 @@ array fft_impl( if (real && inverse && a.ndim() > 0) { n.back() = (n.back() - 1) * 2; } - return fft_impl(a, n, axes, real, inverse, s); + return fft_impl(a, n, axes, real, inverse, norm, s); } -array fft_impl(const array& a, bool real, bool inverse, StreamOrDevice s) { +array fft_impl( + const array& a, + bool real, + bool inverse, + FFTNorm norm, + StreamOrDevice s) { std::vector axes(a.ndim()); std::iota(axes.begin(), axes.end(), 0); - return fft_impl(a, axes, real, inverse, s); + return fft_impl(a, axes, real, inverse, norm, s); } array fftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, n, axes, false, false, s); + return fft_impl(a, n, axes, false, false, norm, s); } array fftn( const array& a, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, axes, false, false, s); + return fft_impl(a, axes, false, false, norm, s); } -array fftn(const array& a, StreamOrDevice s /* = {} */) { - return fft_impl(a, false, false, s); +array fftn( + const array& a, + FFTNorm norm /* = FFTNorm::Backward */, + StreamOrDevice s /* = {} */) { + return fft_impl(a, false, false, norm, s); } array ifftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, n, axes, false, true, s); + return fft_impl(a, n, axes, false, true, norm, s); } array ifftn( const array& a, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, axes, false, true, s); + return fft_impl(a, axes, false, true, norm, s); } -array ifftn(const array& a, StreamOrDevice s /* = {} */) { - return fft_impl(a, false, true, s); +array ifftn( + const array& a, + FFTNorm norm /* = FFTNorm::Backward */, + StreamOrDevice s /* = {} */) { + return fft_impl(a, false, true, norm, s); } array rfftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, n, axes, true, false, s); + return fft_impl(a, n, axes, true, false, norm, s); } array rfftn( const array& a, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, axes, true, false, s); + return fft_impl(a, axes, true, false, norm, s); } -array rfftn(const array& a, StreamOrDevice s /* = {} */) { - return fft_impl(a, true, false, s); +array rfftn( + const array& a, + FFTNorm norm /* = FFTNorm::Backward */, + StreamOrDevice s /* = {} */) { + return fft_impl(a, true, false, norm, s); } array irfftn( const array& a, const Shape& n, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, n, axes, true, true, s); + return fft_impl(a, n, axes, true, true, norm, s); } array irfftn( const array& a, const std::vector& axes, + FFTNorm norm /* = FFTNorm::Backward */, StreamOrDevice s /* = {} */) { - return fft_impl(a, axes, true, true, s); + return fft_impl(a, axes, true, true, norm, s); } -array irfftn(const array& a, StreamOrDevice s /* = {} */) { - return fft_impl(a, true, true, s); +array irfftn( + const array& a, + FFTNorm norm /* = FFTNorm::Backward */, + StreamOrDevice s /* = {} */) { + return fft_impl(a, true, true, norm, s); } -array fftshift( +namespace { + +// Shared implementation for fftshift/ifftshift: validates axes and computes +// the per-axis roll amount, differing only in shift sign and error prefix. +array fftshift_impl( + const char* name, const array& a, const std::vector& axes, - StreamOrDevice s /* = {} */) { + bool inverse, + StreamOrDevice s) { if (axes.empty()) { return a; } @@ -203,41 +260,32 @@ array fftshift( int axis = ax < 0 ? ax + a.ndim() : ax; if (axis < 0 || axis >= a.ndim()) { std::ostringstream msg; - msg << "[fftshift] Invalid axis " << ax << " for array with " << a.ndim() - << " dimensions."; + msg << "[" << name << "] Invalid axis " << ax << " for array with " + << a.ndim() << " dimensions."; throw std::invalid_argument(msg.str()); } // Match NumPy's implementation - shifts.push_back(a.shape(axis) / 2); + int shift = a.shape(axis) / 2; + shifts.push_back(inverse ? -shift : shift); } return roll(a, shifts, axes, s); } -array ifftshift( +} // namespace + +array fftshift( const array& a, const std::vector& axes, StreamOrDevice s /* = {} */) { - if (axes.empty()) { - return a; - } - - Shape shifts; - for (int ax : axes) { - // Convert negative axes to positive - int axis = ax < 0 ? ax + a.ndim() : ax; - if (axis < 0 || axis >= a.ndim()) { - std::ostringstream msg; - msg << "[ifftshift] Invalid axis " << ax << " for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } - // Match NumPy's implementation - int size = a.shape(axis); - shifts.push_back(-(size / 2)); - } + return fftshift_impl("fftshift", a, axes, false, s); +} - return roll(a, shifts, axes, s); +array ifftshift( + const array& a, + const std::vector& axes, + StreamOrDevice s /* = {} */) { + return fftshift_impl("ifftshift", a, axes, true, s); } // Default versions that operate on all axes @@ -259,4 +307,32 @@ array ifftshift(const array& a, StreamOrDevice s /* = {} */) { return ifftshift(a, axes, s); } +array fftfreq(int n, double d /* = 1.0 */, StreamOrDevice s /* = {} */) { + if (n <= 0) { + throw std::invalid_argument("[fftfreq] `n` must be greater than 0."); + } + if (d == 0.0) { + throw std::invalid_argument("[fftfreq] `d` must be non-zero."); + } + auto pos = arange(0, (n + 1) / 2, float32, s); + auto neg = arange(-(n / 2), 0, float32, s); + auto freqs = concatenate({pos, neg}, s); + auto scale = + array(static_cast(1.0 / (static_cast(n) * d)), float32); + return multiply(freqs, scale, s); +} + +array rfftfreq(int n, double d /* = 1.0 */, StreamOrDevice s /* = {} */) { + if (n <= 0) { + throw std::invalid_argument("[rfftfreq] `n` must be greater than 0."); + } + if (d == 0.0) { + throw std::invalid_argument("[rfftfreq] `d` must be non-zero."); + } + auto freqs = arange(0, n / 2 + 1, float32, s); + auto scale = + array(static_cast(1.0 / (static_cast(n) * d)), float32); + return multiply(freqs, scale, s); +} + } // namespace mlx::core::fft diff --git a/Source/Cmlx/mlx/mlx/fft.h b/Source/Cmlx/mlx/mlx/fft.h index 9abf2b18c..f69b8aa2c 100644 --- a/Source/Cmlx/mlx/mlx/fft.h +++ b/Source/Cmlx/mlx/mlx/fft.h @@ -2,6 +2,7 @@ #pragma once +#include #include #include "array.h" @@ -11,40 +12,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. */ @@ -52,14 +87,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. */ @@ -67,14 +104,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. */ @@ -82,34 +121,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. */ @@ -117,14 +184,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`. */ @@ -132,15 +201,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/mlx/mlx/io/CMakeLists.txt b/Source/Cmlx/mlx/mlx/io/CMakeLists.txt index eebd7618b..1fedcd46b 100644 --- a/Source/Cmlx/mlx/mlx/io/CMakeLists.txt +++ b/Source/Cmlx/mlx/mlx/io/CMakeLists.txt @@ -17,6 +17,11 @@ if(MLX_BUILD_GGUF) PRIVATE $) add_library(gguflib STATIC ${gguflib_SOURCE_DIR}/fp16.c ${gguflib_SOURCE_DIR}/gguflib.c) + # gguflib uses assert() to reject malformed tensor headers (e.g. ndim > 8). + # Those checks are otherwise compiled out by -DNDEBUG in release builds, which + # leaves out-of-bounds reads/writes unguarded when loading untrusted GGUF + # files. Force NDEBUG off for this target so the asserts stay live. + target_compile_options(gguflib PRIVATE -UNDEBUG) target_link_libraries(mlx PRIVATE $) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gguf.cpp ${CMAKE_CURRENT_SOURCE_DIR}/gguf_quants.cpp) diff --git a/Source/Cmlx/mlx/mlx/io/gguf.cpp b/Source/Cmlx/mlx/mlx/io/gguf.cpp index 206f6fb31..6c27c4698 100644 --- a/Source/Cmlx/mlx/mlx/io/gguf.cpp +++ b/Source/Cmlx/mlx/mlx/io/gguf.cpp @@ -124,8 +124,7 @@ void set_mx_value_from_gguf( value = array(val->boolval, bool_); break; case GGUF_VALUE_TYPE_STRING: - value = - std::string(val->string.string, static_cast(val->string.len)); + value = std::string(val->string.string, val->string.len); break; case GGUF_VALUE_TYPE_FLOAT64: value = array(val->float64, float32); @@ -161,7 +160,7 @@ void set_mx_value_from_gguf( value = array(reinterpret_cast(data), {size}, uint64); break; case GGUF_VALUE_TYPE_INT64: - value = array(reinterpret_cast(data), {size}, int64); + value = array(reinterpret_cast(data), {size}, int64); break; case GGUF_VALUE_TYPE_FLOAT32: value = array(reinterpret_cast(data), {size}, float32); @@ -174,7 +173,7 @@ void set_mx_value_from_gguf( for (auto& str : strs) { auto str_val = reinterpret_cast(data); data += (str_val->len + sizeof(gguf_string)); - str = std::string(str_val->string, static_cast(str_val->len)); + str = std::string(str_val->string, str_val->len); ctx->off += (str_val->len + sizeof(gguf_string)); } value = std::move(strs); @@ -200,10 +199,102 @@ void set_mx_value_from_gguf( } } +inline size_t gguf_value_type_size(uint32_t type) { + switch (type) { + case GGUF_VALUE_TYPE_BOOL: + case GGUF_VALUE_TYPE_UINT8: + case GGUF_VALUE_TYPE_INT8: + return 1; + case GGUF_VALUE_TYPE_UINT16: + case GGUF_VALUE_TYPE_INT16: + return 2; + case GGUF_VALUE_TYPE_UINT32: + case GGUF_VALUE_TYPE_INT32: + case GGUF_VALUE_TYPE_FLOAT32: + return 4; + case GGUF_VALUE_TYPE_UINT64: + case GGUF_VALUE_TYPE_INT64: + case GGUF_VALUE_TYPE_FLOAT64: + return 8; + default: + return 0; + } +} + +void check_metadata_value_in_file( + const gguf_ctx* ctx, + uint32_t type, + const gguf_value* val) { + auto end = ctx->data + ctx->size; + // Bytes available from a pointer up to the end of the mapping; 0 if the + // pointer lies outside [ctx->data, end]. + auto avail = [&](const uint8_t* p) -> size_t { + return (p < ctx->data || p > end) ? 0 : static_cast(end - p); + }; + auto base = reinterpret_cast(val); + auto fail = [](const char* what) { + std::ostringstream msg; + msg << "[load_gguf] " << what + << " Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + }; + + size_t fixed = gguf_value_type_size(type); + if (fixed) { + if (fixed > avail(base)) { + fail("Metadata value extends past the end of the file."); + } + return; + } + + auto check_string = [&](const uint8_t* p) -> const uint8_t* { + uint64_t len = reinterpret_cast(p)->len; + if (sizeof(uint64_t) + len > avail(p)) { + fail("String metadata value extends past the end of the file."); + } + return p + sizeof(uint64_t) + len; + }; + + if (type == GGUF_VALUE_TYPE_STRING) { + if (sizeof(uint64_t) > avail(base)) { + fail("String metadata value extends past the end of the file."); + } + check_string(base); + return; + } + + if (type == GGUF_VALUE_TYPE_ARRAY) { + if (gguf_array_header_size > avail(base)) { + fail("Metadata value extends past the end of the file."); + } + const uint8_t* elt = base + gguf_array_header_size; + size_t elt_size = gguf_value_type_size(val->array.type); + if (elt_size) { + if (val->array.len > avail(elt) / elt_size) { + fail("Array metadata value extends past the end of the file."); + } + return; + } + if (val->array.type == GGUF_VALUE_TYPE_STRING) { + const uint8_t* p = elt; + for (uint64_t i = 0; i < val->array.len; i++) { + if (sizeof(uint64_t) > avail(p)) { + fail("Array metadata value extends past the end of the file."); + } + p = check_string(p); + } + } + return; + } + + throw std::runtime_error("[load_gguf] Received unexpected type."); +} + std::unordered_map load_metadata(gguf_ctx* ctx) { std::unordered_map metadata; gguf_key key; while (gguf_get_key(ctx, &key)) { + check_metadata_value_in_file(ctx, key.type, key.val); std::string key_name = std::string(key.name, key.namelen); auto& val = metadata.insert({key_name, GGUFMetaData{}}).first->second; set_mx_value_from_gguf(ctx, key.type, key.val, val); @@ -211,6 +302,24 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { return metadata; } +void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) { + auto fail = [&tensor](const std::string& what) { + std::ostringstream msg; + msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) + << "' " << what << ". Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + }; + if (tensor.offset < ctx->data_off) { + fail("has a data offset that overflows the data section"); + } + if (tensor.offset > ctx->size) { + fail("has a data offset past the end of the file"); + } + if (tensor.bsize > ctx->size - tensor.offset) { + fail("extends past the end of the file"); + } +} + std::unordered_map load_arrays(gguf_ctx* ctx) { std::unordered_map array_map; gguf_tensor tensor; @@ -225,6 +334,7 @@ std::unordered_map load_arrays(gguf_ctx* ctx) { }; while (gguf_get_tensor(ctx, &tensor)) { + check_tensor_in_file(ctx, tensor); if (tensor.type == GGUF_TYPE_Q4_0 || tensor.type == GGUF_TYPE_Q4_1 || tensor.type == GGUF_TYPE_Q8_0) { gguf_load_quantized(array_map, tensor); diff --git a/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp b/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp index 148ed6c47..83dfe1f82 100644 --- a/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp +++ b/Source/Cmlx/mlx/mlx/io/gguf_quants.cpp @@ -8,10 +8,12 @@ namespace mlx::core { -void unpack_32_4(uint8_t* data, int8_t* dst) { +// Unpacks a block of 32 4-bit weights. `data` must point at the first +// quant byte, past the block header. +void unpack_32_4(const uint8_t* data, int8_t* dst) { std::fill_n(dst, 16, 0); for (int j = 0; j < 16; ++j) { - uint8_t x = (data[j + 2] & 0x0F); // j+2 to skip scale bytes. + uint8_t x = (data[j] & 0x0F); if (j % 2 != 0) { x <<= 4; } @@ -19,7 +21,7 @@ void unpack_32_4(uint8_t* data, int8_t* dst) { } // Last 16 weights are in the higher bits for (int j = 0; j < 16; ++j) { - uint8_t x = (data[j + 2] >> 4); + uint8_t x = (data[j] >> 4); if (j % 2 != 0) { x <<= 4; } @@ -42,7 +44,7 @@ void extract_q4_0_data( for (int64_t i = 0; i < scales_arr.size(); i++) { scales[i] = *((float16_t*)data); biases[i] = -8 * scales[i]; - unpack_32_4(data, weights); + unpack_32_4(data + 2, weights); // skip the scale bytes weights += 16; data += bytes_per_block; } @@ -64,7 +66,7 @@ void extract_q4_1_data( for (int64_t i = 0; i < scales_arr.size(); i++) { scales[i] = *((float16_t*)data); biases[i] = *((float16_t*)(data) + 1); - unpack_32_4(data, weights); + unpack_32_4(data + 4, weights); // skip the scale and bias bytes weights += 16; data += bytes_per_block; } diff --git a/Source/Cmlx/mlx/mlx/io/load.cpp b/Source/Cmlx/mlx/mlx/io/load.cpp index bd572ce1d..10a9f72a9 100644 --- a/Source/Cmlx/mlx/mlx/io/load.cpp +++ b/Source/Cmlx/mlx/mlx/io/load.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include #include #include @@ -162,10 +162,6 @@ void save(std::shared_ptr out_stream, array a) { a = contiguous(a, true); a.eval(); - if (a.nbytes() == 0) { - throw std::invalid_argument("[save] cannot serialize an empty array"); - } - //////////////////////////////////////////////////////// // Check file if (!out_stream->good() || !out_stream->is_open()) { @@ -226,7 +222,11 @@ void save(std::shared_ptr out_stream, array a) { out_stream->write(magic_ver_len.str().c_str(), magic_ver_len.str().length()); out_stream->write(header.str().c_str(), header.str().length()); - out_stream->write(a.data(), a.nbytes()); + // An empty array has no data to write, and asking for its pointer is not + // meaningful, so stop after the header. + if (a.nbytes() > 0) { + out_stream->write(a.data(), a.nbytes()); + } } /** Save array to file in .npy format */ @@ -349,13 +349,15 @@ array load(std::string file, StreamOrDevice s) { namespace io { ThreadPool& thread_pool() { - static ThreadPool pool_{MLX_IO_THREAD_COUNT}; - return pool_; + // Leak - see Scheduler singleton comment in scheduler.cpp. + static ThreadPool* pool_ = new ThreadPool{MLX_IO_THREAD_COUNT}; + return *pool_; } ThreadPool& ParallelFileReader::thread_pool() { - static ThreadPool thread_pool{MLX_IO_THREAD_COUNT}; - return thread_pool; + // Leak - see Scheduler singleton comment in scheduler.cpp. + static ThreadPool* thread_pool = new ThreadPool{MLX_IO_THREAD_COUNT}; + return *thread_pool; } void ParallelFileReader::read(char* data, size_t n) { diff --git a/Source/Cmlx/mlx/mlx/io/load.h b/Source/Cmlx/mlx/mlx/io/load.h index e9157eda7..00c3ff1d4 100644 --- a/Source/Cmlx/mlx/mlx/io/load.h +++ b/Source/Cmlx/mlx/mlx/io/load.h @@ -96,7 +96,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); } @@ -164,7 +166,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/mlx/mlx/io/safetensors.cpp b/Source/Cmlx/mlx/mlx/io/safetensors.cpp index b8f91ba9b..75c8a6ddc 100644 --- a/Source/Cmlx/mlx/mlx/io/safetensors.cpp +++ b/Source/Cmlx/mlx/mlx/io/safetensors.cpp @@ -1,7 +1,8 @@ // Copyright © 2023 Apple Inc. -// + #include #include +#include #include #include "mlx/backend/cuda/cuda.h" @@ -27,6 +28,7 @@ using json = nlohmann::json; #define ST_U32 "U32" #define ST_U64 "U64" #define ST_F8_E4M3 "F8_E4M3" +#define ST_F8_E8M0 "F8_E8M0" // Note: Complex numbers aren't in the spec yet so this could change - // https://github.com/huggingface/safetensors/issues/389 @@ -96,9 +98,12 @@ Dtype dtype_from_safetensor_str(std::string_view str) { return complex64; } else if (str == ST_F8_E4M3) { return uint8; + } else if (str == ST_F8_E8M0) { + return uint8; } else { - throw std::runtime_error( - "[safetensor] unsupported dtype " + std::string(str)); + std::ostringstream msg; + msg << "[safetensor] unsupported dtype " << str; + throw std::runtime_error(msg.str()); } } @@ -109,8 +114,9 @@ SafetensorsLoad load_safetensors( //////////////////////////////////////////////////////// // Open and check file if (!in_stream->good() || !in_stream->is_open()) { - throw std::runtime_error( - "[load_safetensors] Failed to open " + in_stream->label()); + std::ostringstream msg; + msg << "[load_safetensors] Failed to open " << in_stream->label(); + throw std::runtime_error(msg.str()); } auto stream = cu::is_available() ? to_stream(s) : to_stream(s, Device::cpu); @@ -120,19 +126,37 @@ SafetensorsLoad load_safetensors( constexpr uint64_t kMaxJsonHeaderLength = 100000000; in_stream->read(reinterpret_cast(&jsonHeaderLength), 8); if (jsonHeaderLength <= 0 || jsonHeaderLength >= kMaxJsonHeaderLength) { - throw std::runtime_error( - "[load_safetensors] Invalid json header length " + in_stream->label()); + std::ostringstream msg; + msg << "[load_safetensors] Invalid json header length " + << in_stream->label(); + throw std::runtime_error(msg.str()); } + + // Determine file size to be able to validate that reads are within the + // bounds of the file (at least at creation time) + in_stream->seek(0, std::ios_base::end); + size_t file_size = in_stream->tell(); + in_stream->seek(8, std::ios_base::beg); + // Load the json metadata + if (file_size < jsonHeaderLength + 8) { + std::ostringstream msg; + msg << "[load_safetensors] The JSON header is " << jsonHeaderLength + << " bytes long but the file is only " << file_size << " bytes. " + << "Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + } auto rawJson = std::make_unique(jsonHeaderLength); in_stream->read(rawJson.get(), jsonHeaderLength); auto metadata = json::parse(rawJson.get(), rawJson.get() + jsonHeaderLength); // Should always be an object on the top-level if (!metadata.is_object()) { - throw std::runtime_error( - "[load_safetensors] Invalid json metadata " + in_stream->label()); + std::ostringstream msg; + msg << "[load_safetensors] Invalid json metadata " << in_stream->label(); + throw std::runtime_error(msg.str()); } size_t offset = jsonHeaderLength + 8; + // Load the arrays using metadata std::unordered_map res; std::unordered_map metadata_map; @@ -147,6 +171,36 @@ SafetensorsLoad load_safetensors( const Shape& shape = item.value().at("shape"); const std::vector& data_offsets = item.value().at("data_offsets"); Dtype type = dtype_from_safetensor_str(dtype); + if (data_offsets.size() != 2) { + std::ostringstream msg; + msg << "[load_safetensors] Tensor '" << item.key() + << "' data_offsets must have exactly 2 entries but has " + << data_offsets.size(); + throw std::runtime_error(msg.str()); + } + { + size_t expected_nbytes = type.size(); + for (auto dim : shape) { + expected_nbytes *= static_cast(dim); + } + if (data_offsets[1] < data_offsets[0] || + data_offsets[1] - data_offsets[0] != expected_nbytes) { + std::ostringstream msg; + msg << "[load_safetensors] Tensor '" << item.key() + << "' invalid data offsets (" << data_offsets[0] << ", " + << data_offsets[1] << "). Expecting " << expected_nbytes + << " bytes."; + throw std::runtime_error(msg.str()); + } + } + if (offset + data_offsets[1] > file_size) { + std::ostringstream msg; + msg << "[load_safetensors] Tensor '" << item.key() + << "' invalid data offsets (" << data_offsets[0] << ", " + << data_offsets[1] << ") exceeding the size of the file. " + << "Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); + } res.insert( {item.key(), array( @@ -170,8 +224,9 @@ void save_safetensors( //////////////////////////////////////////////////////// // Check file if (!out_stream->good() || !out_stream->is_open()) { - throw std::runtime_error( - "[save_safetensors] Failed to open " + out_stream->label()); + std::ostringstream msg; + msg << "[save_safetensors] Failed to open " << out_stream->label(); + throw std::runtime_error(msg.str()); } //////////////////////////////////////////////////////// @@ -195,11 +250,6 @@ void save_safetensors( size_t offset = 0; for (auto& [key, arr] : a) { - if (arr.nbytes() == 0) { - throw std::invalid_argument( - "[save_safetensors] cannot serialize an empty array key: " + key); - } - json child; child["dtype"] = dtype_to_safetensor_str(arr.dtype()); child["shape"] = arr.shape(); @@ -213,7 +263,11 @@ void save_safetensors( out_stream->write(reinterpret_cast(&header_len), 8); out_stream->write(header.c_str(), header_len); for (auto& [key, arr] : a) { - out_stream->write(arr.data(), arr.nbytes()); + // An empty tensor contributes a zero length span and has no data pointer + // worth asking for. + if (arr.nbytes() > 0) { + out_stream->write(arr.data(), arr.nbytes()); + } } } diff --git a/Source/Cmlx/mlx/mlx/linalg.cpp b/Source/Cmlx/mlx/mlx/linalg.cpp index 7ac080dab..0d06b5646 100644 --- a/Source/Cmlx/mlx/mlx/linalg.cpp +++ b/Source/Cmlx/mlx/mlx/linalg.cpp @@ -84,8 +84,8 @@ inline array matrix_norm( bool keepdims, StreamOrDevice s) { auto dtype = at_least_float(a.dtype()); - auto row_axis = axis[0]; - auto col_axis = axis[1]; + int row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; + int col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; if (ord == -1.0) { col_axis -= (!keepdims && col_axis > row_axis && col_axis > 0); return astype( @@ -111,8 +111,6 @@ inline array matrix_norm( dtype, s); } else if (ord == 2.0 || ord == -2.0) { - row_axis = (axis[0] < 0) ? axis[0] + a.ndim() : axis[0]; - col_axis = (axis[1] < 0) ? axis[1] + a.ndim() : axis[1]; auto a_matrix = (row_axis > col_axis) ? moveaxis(moveaxis(a, row_axis, -1, s), col_axis, -1, s) : moveaxis(moveaxis(a, col_axis, -1, s), row_axis, -2, s); @@ -168,7 +166,14 @@ array norm( bool keepdims /* = false */, StreamOrDevice s /* = {} */) { if (!axis) { - return norm(flatten(a, s), std::vector{0}, keepdims, s); + auto out = norm(flatten(a, s), std::vector{0}, keepdims, s); + if (keepdims) { + // The flatten above collapses the input to one dimension, so keepdims + // has to restore the rank of the original array rather than the + // flattened one. + out = reshape(out, Shape(a.ndim(), 1), s); + } + return out; } if (axis.value().size() > 2) { @@ -364,6 +369,15 @@ array pinv(const array& a, StreamOrDevice s /* = {} */) { throw std::invalid_argument(msg.str()); } + // The cutoff below reduces over the singular values, which cannot run on an + // empty array. Nothing needs computing anyway, and the result is the shape + // of the transposed input, like numpy. + if (a.size() == 0) { + auto out_shape = a.shape(); + std::swap(out_shape[a.ndim() - 1], out_shape[a.ndim() - 2]); + return zeros(std::move(out_shape), a.dtype(), s); + } + int m = a.shape(-2); int n = a.shape(-1); int k = std::min(m, n); @@ -433,12 +447,8 @@ array cross( int axis /* = -1 */, StreamOrDevice s /* = {} */) { auto check_ax = [axis](const array& arr) { - if (axis >= static_cast(arr.ndim()) || axis + arr.ndim() < 0) { - std::ostringstream msg; - msg << "[linalg::cross] axis " << axis << " invalid for array with " - << arr.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + normalize_axis_index(axis, arr.ndim(), "[linalg::cross] "); if (arr.shape(axis) < 2 || arr.shape(axis) > 3) { throw std::invalid_argument( "[linalg::cross] The specified axis must have size 2 or 3."); @@ -705,4 +715,170 @@ array solve_triangular( return matmul(a_inv, b, s); } -} // namespace mlx::core::linalg \ No newline at end of file +void validate_det( + const array& a, + const StreamOrDevice& stream, + const std::string& fname) { + check_cpu_stream(stream, fname); + if (issubdtype(a.dtype(), complexfloating)) { + throw std::invalid_argument(fname + " Complex inputs are not supported."); + } + if (a.ndim() < 2) { + std::ostringstream msg; + msg << fname + << " Arrays must have >= 2 dimensions. Received array " + "with " + << a.ndim() << " dimensions."; + throw std::invalid_argument(msg.str()); + } + + if (a.shape(-1) != a.shape(-2)) { + throw std::invalid_argument(fname + " Only defined for square matrices."); + } +} + +array det_raw_small(const array& a, StreamOrDevice s) { + int n = a.shape(-1); + + // Empty 0x0 matrix: determinant is the empty product = 1 + if (n == 0) { + Shape out_shape(a.shape().begin(), a.shape().end() - 2); + return broadcast_to(array(1.0f, a.dtype()), std::move(out_shape), s); + } + + // Helper to extract a[..., i, j] from the last two dims + auto elem = [&](int i, int j) { + auto starts = Shape(a.ndim(), 0); + auto stops = a.shape(); + starts[a.ndim() - 2] = i; + stops[a.ndim() - 2] = i + 1; + starts[a.ndim() - 1] = j; + stops[a.ndim() - 1] = j + 1; + return squeeze(squeeze(slice(a, starts, stops, s), -1, s), -1, s); + }; + + if (n == 1) { + return elem(0, 0); + } else if (n == 2) { + return subtract( + multiply(elem(0, 0), elem(1, 1), s), + multiply(elem(0, 1), elem(1, 0), s), + s); + } else { + // 3x3: a00*(a11*a22 - a12*a21) - a01*(a10*a22 - a12*a20) + a02*(a10*a21 - + // a11*a20) + auto a00 = elem(0, 0), a01 = elem(0, 1), a02 = elem(0, 2); + auto a10 = elem(1, 0), a11 = elem(1, 1), a12 = elem(1, 2); + auto a20 = elem(2, 0), a21 = elem(2, 1), a22 = elem(2, 2); + return add( + subtract( + multiply( + a00, + subtract(multiply(a11, a22, s), multiply(a12, a21, s), s), + s), + multiply( + a01, + subtract(multiply(a10, a22, s), multiply(a12, a20, s), s), + s), + s), + multiply( + a02, subtract(multiply(a10, a21, s), multiply(a11, a20, s), s), s), + s); + } +} + +std::pair slogdet_impl(const array& input, StreamOrDevice s) { + int n = input.shape(-1); + auto dtype = input.dtype(); + + // Small-matrix fast path + if (n <= 3) { + auto raw = det_raw_small(input, s); + auto abs_raw = abs(raw, s); + auto sgn = sign(raw, s); + auto logabs = log(abs_raw, s); + return std::make_pair(sgn, logabs); + } + + // General LU-based path + auto [LU, pivots] = lu_factor(input, s); + + // Extract diagonal of U + auto diag = diagonal(LU, 0, -2, -1, s); + + // Permutation parity: count positions where pivot[i] != i + int k = std::min(input.shape(-2), input.shape(-1)); + auto iota = arange(0, k, uint32, s); + auto parity = astype( + sum(not_equal(pivots, iota, s), + /* axis = */ -1, + /* keepdims = */ false, + s), + int32, + s); + + // Count negative diagonal elements + auto num_neg = astype( + sum(less(diag, array(0.0f, dtype), s), + /* axis = */ -1, + /* keepdims = */ false, + s), + int32, + s); + + // sign = (-1)^(parity + num_neg) + auto total = add(parity, num_neg, s); + auto sign_val = astype( + subtract( + array(1, int32), + multiply(array(2, int32), remainder(total, array(2, int32), s), s), + s), + dtype, + s); + + // logabsdet = sum(log(abs(diag))) + auto logabsdet = + sum(log(abs(diag, s), s), /* axis = */ -1, /* keepdims = */ false, s); + + // Handle singular matrices: any zero on diagonal + auto is_zero = + any(equal(diag, array(0.0f, dtype), s), + /* axis = */ -1, + /* keepdims = */ false, + s); + sign_val = where(is_zero, array(0.0f, dtype), sign_val, s); + logabsdet = where( + is_zero, + array(-std::numeric_limits::infinity(), dtype), + logabsdet, + s); + + return std::make_pair(sign_val, logabsdet); +} + +std::pair slogdet(const array& a, StreamOrDevice s /* = {} */) { + validate_det(a, s, "[linalg::slogdet]"); + + auto dtype = at_least_float(a.dtype()); + auto input = astype(a, dtype, s); + return slogdet_impl(input, s); +} + +array det(const array& a, StreamOrDevice s /* = {} */) { + validate_det(a, s, "[linalg::det]"); + + auto dtype = at_least_float(a.dtype()); + auto input = astype(a, dtype, s); + int n = input.shape(-1); + + // Small-matrix fast path: compute directly, skip log/exp round-trip + if (n <= 3) { + return det_raw_small(input, s); + } + + // General case: det = sign * exp(logabsdet) + auto [sign_val, logabsdet] = slogdet_impl(input, s); + return multiply(sign_val, exp(logabsdet, s), s); +} + +} // namespace mlx::core::linalg diff --git a/Source/Cmlx/mlx/mlx/linalg.h b/Source/Cmlx/mlx/mlx/linalg.h index fe3f83c22..08e0bdd7a 100644 --- a/Source/Cmlx/mlx/mlx/linalg.h +++ b/Source/Cmlx/mlx/mlx/linalg.h @@ -112,4 +112,8 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/ops.cpp b/Source/Cmlx/mlx/mlx/ops.cpp index 076de699b..5fe3d9be5 100644 --- a/Source/Cmlx/mlx/mlx/ops.cpp +++ b/Source/Cmlx/mlx/mlx/ops.cpp @@ -1,4 +1,4 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2026 Apple Inc. // Required for using M_PI in MSVC. #define _USE_MATH_DEFINES @@ -10,20 +10,27 @@ #include #include "mlx/backend/cuda/cuda.h" +#include "mlx/backend/metal/metal.h" #include "mlx/fast_primitives.h" #include "mlx/ops.h" #include "mlx/primitives.h" #include "mlx/transforms.h" #include "mlx/transforms_impl.h" +#include "mlx/types/limits.h" #include "mlx/utils.h" namespace mlx::core { namespace { +// Pass the reduction's name in `op_without_identity` when it has no identity +// element, as max and min do not. Those are undefined over an empty axis, so +// naming them here has that checked. Reductions with an identity, such as sum +// and any, reduce an empty axis fine and leave it unset. std::tuple, bool> compute_reduce_shape( const std::vector& axes, - const Shape& shape) { + const Shape& shape, + std::string_view op_without_identity = {}) { bool is_noop = true; std::set axes_set; auto ndim = shape.size(); @@ -45,11 +52,23 @@ std::tuple, bool> compute_reduce_shape( if (axes_set.count(i) == 0) { out_shape.push_back(shape[i]); } else { + if (!op_without_identity.empty() && shape[i] == 0) { + std::ostringstream msg; + msg << "[" << op_without_identity << "] Cannot " << op_without_identity + << " reduce over axis " << i << " with size 0."; + throw std::invalid_argument(msg.str()); + } out_shape.push_back(1); } is_noop &= (out_shape.back() == shape[i]); } std::vector sorted_axes(axes_set.begin(), axes_set.end()); + // Dimensions that are size 1 at trace time can be dynamic when compiling + // shapeless, so the reduction cannot be elided from the graph. Reductions + // over no axes stay no-ops since they are shape independent. + if (is_noop && !sorted_axes.empty() && detail::in_dynamic_tracing()) { + is_noop = false; + } return {out_shape, sorted_axes, is_noop}; } @@ -58,16 +77,33 @@ Dtype at_least_float(const Dtype& d) { } array indices_or_default( - std::optional indices, + std::string_view tag, + const std::optional& indices, const array& x, StreamOrDevice s) { + if (x.ndim() < 2) { + std::ostringstream msg; + msg << tag + << " Input must have at least two dimensions but got input with shape " + << x.shape() << "."; + throw std::invalid_argument(msg.str()); + } + if (indices.has_value()) { - return indices.value(); + if (!issubdtype(indices->dtype(), integer)) { + std::ostringstream msg; + msg << tag + << " Got indices with invalid dtype. Indices must be integral."; + throw std::invalid_argument(msg.str()); + } + return astype(indices.value(), uint32); } Shape shape(x.shape().begin(), x.shape().end() - 2); - int total = - std::reduce(shape.begin(), shape.end(), 1, std::multiplies()); + int total = safe_cast( + std::reduce( + shape.begin(), shape.end(), int64_t{1}, std::multiplies{}), + "gather"); return reshape(arange(total, uint32, s), std::move(shape), s); } @@ -165,6 +201,10 @@ array arange( throw std::invalid_argument("[arange] Cannot compute length."); } + if (step == 0) { + throw std::invalid_argument("[arange] step must be nonzero."); + } + // Check if start and stop specify a valid range because if not, we have to // return an empty array if (std::isinf(step) && @@ -232,6 +272,7 @@ array linspace( double start, double stop, int num /* = 50 */, + bool endpoint /* = true */, Dtype dtype /* = float32 */, StreamOrDevice s /* = {} */) { if (num < 0) { @@ -243,8 +284,11 @@ array linspace( return astype(array({start}), dtype, s); } auto inner_type = dtype == float64 ? float64 : float32; + // Without the endpoint the samples are spaced so that `stop` would be the + // next one after the last, i.e. the step is (stop - start) / num. + auto denominator = endpoint ? num - 1 : num; array t = - divide(arange(0, num, inner_type, s), array(num - 1, inner_type), s); + divide(arange(0, num, inner_type, s), array(denominator, inner_type), s); array t_bar = subtract(array(1, inner_type), t, s); return astype( add(multiply(t_bar, array(start, inner_type), s), @@ -254,8 +298,10 @@ array linspace( s); } -array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) { - if (dtype == a.dtype()) { +// Private API used by python bindings. +MLX_API array +astype(array a, Dtype dtype, bool force_copy, StreamOrDevice s = {}) { + if (dtype == a.dtype() && !force_copy) { return a; } auto copied_shape = a.shape(); // |a| will be moved @@ -266,12 +312,20 @@ array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) { {std::move(a)}); } +array astype(array a, Dtype dtype, StreamOrDevice s /* = {} */) { + return astype(std::move(a), dtype, false, s); +} + array as_strided( array a, Shape shape, Strides strides, size_t offset, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument( + "[as_strided] Negative dimensions not allowed."); + } auto copied_shape = shape; // |shape| will be moved auto dtype = a.dtype(); // |a| will be moved return array( @@ -331,7 +385,11 @@ array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { } array zeros_like(const array& a, StreamOrDevice s /* = {} */) { - return full_like(a, 0, a.dtype(), to_stream(s)); + return zeros_like(a, a.dtype(), s); +} + +array zeros_like(const array& a, Dtype dtype, StreamOrDevice s /* = {} */) { + return full_like(a, 0, dtype, to_stream(s)); } array ones(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { @@ -339,14 +397,23 @@ array ones(const Shape& shape, Dtype dtype, StreamOrDevice s /* = {} */) { } array ones_like(const array& a, StreamOrDevice s /* = {} */) { - return full_like(a, 1, a.dtype(), to_stream(s)); + return ones_like(a, a.dtype(), s); +} + +array ones_like(const array& a, Dtype dtype, StreamOrDevice s /* = {} */) { + return full_like(a, 1, dtype, to_stream(s)); } array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s /* = {} */) { - if (n <= 0 || m <= 0) { + if (n < 0 || m < 0) { throw std::invalid_argument("[eye] N and M must be positive integers."); } array result = zeros({n, m}, dtype, s); + + if (n == 0 || m == 0) { + return result; + } + if (k >= m || -k >= n) { return result; } @@ -432,7 +499,7 @@ array unflatten( } } if (infer_idx >= 0) { - shape[infer_idx] = a.shape(ax) / size; + shape[infer_idx] = safe_cast(a.shape(ax) / size, "unflatten"); size *= shape[infer_idx]; } if (size != a.shape(ax)) { @@ -525,26 +592,50 @@ array hadamard_transform( {astype(a, dtype, s)}); } -array squeeze_impl( +namespace { + +std::vector normalize_squeeze_axes( const array& a, - std::vector axes, - StreamOrDevice s /* = {} */) { - for (auto& ax : axes) { - auto new_ax = ax < 0 ? ax + a.ndim() : ax; - if (new_ax < 0 || new_ax >= a.ndim()) { - std::ostringstream msg; - msg << "[squeeze] Invalid axes " << ax << " for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + const std::vector& axes) { + int ndim = a.ndim(); + std::set unique_axes; + for (auto ax : axes) { + int new_ax = normalize_axis_index(ax, ndim, "[squeeze] "); if (a.shape(new_ax) != 1) { std::ostringstream msg; msg << "[squeeze] Cannot squeeze axis " << ax << " with size " - << a.shape(ax) << " which is not equal to 1."; + << a.shape(new_ax) << " which is not equal to 1."; throw std::invalid_argument(msg.str()); } - ax = new_ax; + unique_axes.insert(new_ax); + } + if (unique_axes.size() != axes.size()) { + throw std::invalid_argument("[squeeze] Received duplicate axes."); + } + return std::vector(unique_axes.begin(), unique_axes.end()); +} + +std::vector normalize_expand_dims_axes( + const array& a, + const std::vector& axes) { + int out_ndim = a.ndim() + axes.size(); + std::set unique_axes; + for (auto ax : axes) { + unique_axes.insert(normalize_axis_index(ax, out_ndim, "[expand_dims] ")); } + if (unique_axes.size() != axes.size()) { + throw std::invalid_argument("[expand_dims] Received duplicate axes."); + } + return std::vector(unique_axes.begin(), unique_axes.end()); +} + +} // namespace + +// Assumes the axes are non-negative, sorted, unique, and valid for a. +array squeeze_impl( + const array& a, + std::vector axes, + StreamOrDevice s /* = {} */) { auto shape = Squeeze::output_shape(a, axes); return array( std::move(shape), @@ -560,19 +651,11 @@ array squeeze( if (axes.empty()) { return a; } - std::set unique_axes; - for (auto ax : axes) { - unique_axes.insert(ax < 0 ? ax + a.ndim() : ax); - } - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[squeeze] Received duplicate axes."); - } - std::vector sorted_axes(unique_axes.begin(), unique_axes.end()); - return squeeze_impl(a, std::move(sorted_axes), s); + return squeeze_impl(a, normalize_squeeze_axes(a, axes), s); } array squeeze(const array& a, int axis, StreamOrDevice s /* = {} */) { - return squeeze_impl(a, {axis}, s); + return squeeze_impl(a, normalize_squeeze_axes(a, {axis}), s); } array squeeze(const array& a, StreamOrDevice s /* = {} */) { @@ -585,21 +668,11 @@ array squeeze(const array& a, StreamOrDevice s /* = {} */) { return squeeze_impl(a, std::move(axes), s); } +// Assumes the axes are non-negative, sorted, unique and valid for the output. array expand_dims_impl( const array& a, std::vector axes, StreamOrDevice s /* = {} */) { - auto out_ndim = a.ndim() + axes.size(); - for (auto& ax : axes) { - auto new_ax = ax < 0 ? ax + out_ndim : ax; - if (new_ax < 0 || new_ax >= out_ndim) { - std::ostringstream msg; - msg << "[expand_dims] Invalid axis " << ax << " for output array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - ax = new_ax; - } auto shape = ExpandDims::output_shape(a, axes); return array( std::move(shape), @@ -609,7 +682,7 @@ array expand_dims_impl( } array expand_dims(const array& a, int axis, StreamOrDevice s /* = {} */) { - return expand_dims_impl(a, {axis}, s); + return expand_dims_impl(a, normalize_expand_dims_axes(a, {axis}), s); } array expand_dims( @@ -619,23 +692,34 @@ array expand_dims( if (axes.empty()) { return a; } - { // Check for repeats - std::set unique_axes(axes.begin(), axes.end()); - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[expand_dims] Received duplicate axes."); - } - } - // Check for repeats again - auto out_ndim = a.ndim() + axes.size(); - std::set unique_axes; + return expand_dims_impl(a, normalize_expand_dims_axes(a, axes), s); +} + +array flip( + const array& a, + const std::vector& axes, + StreamOrDevice s /* = {} */) { + auto ndim = static_cast(a.ndim()); + Shape start(ndim, 0); + Shape stop = a.shape(); + Shape strides(ndim, 1); for (auto ax : axes) { - unique_axes.insert(ax < 0 ? ax + out_ndim : ax); - } - if (unique_axes.size() != axes.size()) { - throw std::invalid_argument("[expand_dims] Received duplicate axes."); + int axis = normalize_axis_index(ax, ndim, "[flip] "); + start[axis] = a.shape(axis) - 1; + stop[axis] = -a.shape(axis) - 1; + strides[axis] = -1; } - std::vector sorted_axes(unique_axes.begin(), unique_axes.end()); - return expand_dims_impl(a, std::move(sorted_axes), s); + return slice(a, std::move(start), std::move(stop), std::move(strides), s); +} + +array flip(const array& a, int axis, StreamOrDevice s /* = {} */) { + return flip(a, std::vector{axis}, s); +} + +array flip(const array& a, StreamOrDevice s /* = {} */) { + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + return flip(a, axes, s); } // Slice helper @@ -849,7 +933,11 @@ array slice_update( src.shape(), src.dtype(), std::make_shared( - to_stream(s), std::move(start), std::move(stop), std::move(strides)), + to_stream(s), + SliceUpdate::None, + std::move(start), + std::move(stop), + std::move(strides)), {src, upd}); } @@ -894,18 +982,168 @@ array slice_update( {src, upd, start}); } +array slice_update( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + SliceUpdate::ReduceType mode, + StreamOrDevice s) { + if (start.size() != src.ndim() || stop.size() != src.ndim() || + strides.size() != src.ndim()) { + std::ostringstream msg; + msg << "[slice_update] Invalid number of indices or strides for " + << "array with dimension " << src.ndim() << "."; + throw std::invalid_argument(msg.str()); + } + + auto [has_neg_strides, upd_shape] = + normalize_slice(src.shape(), start, stop, strides); + + auto upd = broadcast_to(astype(update, src.dtype(), s), upd_shape, s); + + if (!has_neg_strides && upd_shape == src.shape()) { + switch (mode) { + case SliceUpdate::None: + return upd; + case SliceUpdate::Sum: + return add(src, upd, s); + case SliceUpdate::Prod: + return multiply(src, upd, s); + case SliceUpdate::Max: + return maximum(src, upd, s); + case SliceUpdate::Min: + return minimum(src, upd, s); + } + } + + return array( + src.shape(), + src.dtype(), + std::make_shared( + to_stream(s), + mode, + std::move(start), + std::move(stop), + std::move(strides)), + {src, upd}); +} + +array slice_update_add( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s /*= {}*/) { + return slice_update( + src, + update, + std::move(start), + std::move(stop), + std::move(strides), + SliceUpdate::Sum, + s); +} + +array slice_update_add( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s /*= {}*/) { + return slice_update_add( + src, update, std::move(start), std::move(stop), Shape(src.ndim(), 1), s); +} + +array slice_update_prod( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s /*= {}*/) { + return slice_update( + src, + update, + std::move(start), + std::move(stop), + std::move(strides), + SliceUpdate::Prod, + s); +} + +array slice_update_prod( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s /*= {}*/) { + return slice_update_prod( + src, update, std::move(start), std::move(stop), Shape(src.ndim(), 1), s); +} + +array slice_update_max( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s /*= {}*/) { + return slice_update( + src, + update, + std::move(start), + std::move(stop), + std::move(strides), + SliceUpdate::Max, + s); +} + +array slice_update_max( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s /*= {}*/) { + return slice_update_max( + src, update, std::move(start), std::move(stop), Shape(src.ndim(), 1), s); +} + +array slice_update_min( + const array& src, + const array& update, + Shape start, + Shape stop, + Shape strides, + StreamOrDevice s /*= {}*/) { + return slice_update( + src, + update, + std::move(start), + std::move(stop), + std::move(strides), + SliceUpdate::Min, + s); +} + +array slice_update_min( + const array& src, + const array& update, + Shape start, + Shape stop, + StreamOrDevice s /*= {}*/) { + return slice_update_min( + src, update, std::move(start), std::move(stop), Shape(src.ndim(), 1), s); +} + std::vector split( const array& a, const Shape& indices, int axis, StreamOrDevice s /* = {} */) { - auto ax = axis < 0 ? axis + a.ndim() : axis; - if (ax < 0 || ax >= a.ndim()) { - std::ostringstream msg; - msg << "Invalid axis (" << axis << ") passed to split" - << " for array with shape " << a.shape() << "."; - throw std::invalid_argument(msg.str()); - } + auto ax = normalize_axis_index(axis, a.ndim(), "[split] "); if (indices.empty()) { return {a}; @@ -947,14 +1185,14 @@ split(const array& a, const Shape& indices, StreamOrDevice s /* = {} */) { std::vector split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) { - auto ax = axis < 0 ? axis + a.ndim() : axis; - if (ax < 0 || ax >= a.ndim()) { + auto ax = normalize_axis_index(axis, a.ndim(), "[split] "); + if (num_splits <= 0) { std::ostringstream msg; - msg << "Invalid axis " << axis << " passed to split" - << " for array with shape " << a.shape() << "."; + msg << "[split] num_splits must be positive and non-zero but got " + << num_splits << "."; throw std::invalid_argument(msg.str()); } - auto q_and_r = std::ldiv(a.shape(axis), num_splits); + auto q_and_r = std::ldiv(a.shape(ax), num_splits); if (q_and_r.rem) { std::ostringstream msg; msg << "Array split does not result in sub arrays with equal size:" @@ -967,7 +1205,7 @@ split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) { for (int i = 0; i < indices.size(); ++i) { indices[i] = (i + 1) * split_size; } - return split(a, indices, axis, s); + return split(a, indices, ax, s); } std::vector @@ -975,6 +1213,26 @@ split(const array& a, int num_splits, StreamOrDevice s /* = {} */) { return split(a, num_splits, 0, to_stream(s)); } +std::vector +unstack(const array& a, int axis, StreamOrDevice s /* = {} */) { + auto ndim = static_cast(a.ndim()); + auto ax = normalize_axis_index(axis, ndim, "[unstack] "); + auto n = a.shape(ax); + std::vector res; + res.reserve(n); + if (n == 0) { + return res; + } + for (auto& part : split(a, n, ax, s)) { + res.push_back(squeeze(part, ax, s)); + } + return res; +} + +std::vector unstack(const array& a, StreamOrDevice s /* = {} */) { + return unstack(a, 0, s); +} + std::vector meshgrid( const std::vector& arrays, bool sparse /* = false */, @@ -1054,7 +1312,9 @@ array concatenate( }; auto shape = arrays[0].shape(); - shape[ax] = 0; + // Accumulate the concatenation axis in 64 bits so a total that does not fit + // in a shape dimension is reported rather than silently wrapping. + int64_t concat_size = 0; // Make the output shape and validate that all arrays have the same shape // except for the concatenation axis. for (auto& a : arrays) { @@ -1073,8 +1333,9 @@ array concatenate( throw_invalid_shapes(); } } - shape[ax] += a.shape(ax); + concat_size += a.shape(ax); } + shape[ax] = safe_cast(concat_size, "concatenate"); // Promote all the arrays to the same type auto dtype = result_type(arrays); @@ -1133,7 +1394,9 @@ array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) { } if (repeats == 0) { - return array({}, arr.dtype()); + auto shape = arr.shape(); + shape[axis] = 0; + return array(std::initializer_list{}, shape, arr.dtype()); } if (repeats == 1) { @@ -1148,7 +1411,8 @@ array repeat(const array& arr, int repeats, int axis, StreamOrDevice s) { // Reshape back into a contiguous array where S_axis is now S_axis * repeats shape.erase(shape.begin() + axis + 1); - shape[axis] *= repeats; + shape[axis] = + safe_cast(static_cast(shape[axis]) * repeats, "repeat"); out = reshape(out, shape, s); return out; @@ -1180,7 +1444,8 @@ array tile( } expand_shape.push_back(shape[i]); broad_shape.push_back(shape[i]); - final_shape.push_back(reps[i] * shape[i]); + final_shape.push_back( + safe_cast(static_cast(reps[i]) * shape[i], "tile")); } auto x = reshape(arr, std::move(expand_shape), s); @@ -1188,6 +1453,91 @@ array tile( return reshape(x, std::move(final_shape), s); } +array reflect_pad( + const array& a, + const std::vector& axes, + const Shape& low_pad_size, + const Shape& high_pad_size, + const Shape& out_shape, + bool include_edge, + StreamOrDevice s /* = {} */) { + array out = zeros(out_shape, a.dtype(), s); + Shape starts(a.ndim(), 0); + auto stops = a.shape(); + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + starts[ax] = low_pad_size[i]; + stops[ax] += low_pad_size[i]; + } + // Copy over values from the unpadded array + array padded = slice_update(out, a, starts, stops, s); + + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + int n = a.shape(ax); + int L = low_pad_size[i]; + int H = high_pad_size[i]; + if (L == 0 && H == 0) { + continue; + } + if (n == 0) { + std::ostringstream msg; + msg << "[pad] Cannot pad empty axis " << ax << " using mode '" + << (include_edge ? "symmetric" : "reflect") << "'."; + throw std::invalid_argument(msg.str()); + } + // reflect skips the edge value (period 2(n-1)); symmetric repeats it + // (period 2n). + int offset = (!include_edge && n > 1) ? 1 : 0; + int tile = n - offset; + + if (L > 0) { + int filled_start = low_pad_size[i]; + int remaining = L; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_start + offset; + src_stops[ax] = filled_start + offset + chunk; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_start - chunk; + dst_stops[ax] = filled_start; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_start -= chunk; + remaining -= chunk; + } + } + + if (H > 0) { + int filled_end = low_pad_size[i] + n; + int remaining = H; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_end - offset - chunk; + src_stops[ax] = filled_end - offset; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_end; + dst_stops[ax] = filled_end + chunk; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_end += chunk; + remaining -= chunk; + } + } + } + return padded; +} + array edge_pad( const array& a, const std::vector& axes, @@ -1269,7 +1619,10 @@ array pad( } auto ax = axes[i] < 0 ? a.ndim() + axes[i] : axes[i]; - out_shape[ax] += low_pad_size[i] + high_pad_size[i]; + out_shape[ax] = safe_cast( + static_cast(out_shape[ax]) + low_pad_size[i] + + high_pad_size[i], + "pad"); } if (mode == "constant") { @@ -1280,6 +1633,12 @@ array pad( {a, astype(pad_value, a.dtype(), s)}); } else if (mode == "edge") { return edge_pad(a, axes, low_pad_size, high_pad_size, out_shape, s); + } else if (mode == "reflect") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, false, s); + } else if (mode == "symmetric") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, true, s); } else { std::ostringstream msg; msg << "Invalid padding mode (" << mode << ") passed to pad"; @@ -1435,6 +1794,11 @@ array broadcast_to( const array& a, const Shape& shape, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument( + "[broadcast_to] Negative dimensions not allowed."); + } + if (a.shape() == shape) { return a; } @@ -1492,7 +1856,7 @@ std::vector broadcast_arrays( outputs.push_back(in); } else { outputs.push_back(array( - std::move(out_shape), + out_shape, in.dtype(), std::make_shared(to_stream(s), out_shape), {in})); @@ -1509,23 +1873,19 @@ std::vector broadcast_arrays( for (int i = 0; i < inputs.size(); ++i) { auto& in = inputs[i]; auto out_shape = check_and_get_shape(in); - if (in.shape() == out_shape) { - outputs.push_back(in); - } else { - // broadcasted array goes first followed by other stopgrad inputs - std::vector p_inputs = {in}; - for (int j = 0; j < inputs.size(); ++j) { - if (j == i) { - continue; - } - p_inputs.push_back(stop_grad_inputs[j]); + // broadcasted array goes first followed by other stopgrad inputs + std::vector p_inputs = {in}; + for (int j = 0; j < inputs.size(); ++j) { + if (j == i) { + continue; } - outputs.push_back(array( - std::move(out_shape), - in.dtype(), - std::make_shared(to_stream(s), ignore_axes), - std::move(p_inputs))); + p_inputs.push_back(stop_grad_inputs[j]); } + outputs.push_back(array( + out_shape, + in.dtype(), + std::make_shared(to_stream(s), ignore_axes), + std::move(p_inputs))); } return outputs; } @@ -1560,23 +1920,19 @@ std::vector broadcast_arrays( } for (int i = 0; i < inputs.size(); ++i) { auto& in = inputs[i]; - if (in.shape() == shape) { - outputs.push_back(in); - } else { - // broadcasted array goes first followed by other stopgrad inputs - std::vector p_inputs = {in}; - for (int j = 0; j < inputs.size(); ++j) { - if (j == i) { - continue; - } - p_inputs.push_back(stop_grad_inputs[j]); + // broadcasted array goes first followed by other stopgrad inputs + std::vector p_inputs = {in}; + for (int j = 0; j < inputs.size(); ++j) { + if (j == i) { + continue; } - outputs.push_back(array( - shape, - in.dtype(), - std::make_shared(to_stream(s), shape), - std::move(p_inputs))); + p_inputs.push_back(stop_grad_inputs[j]); } + outputs.push_back(array( + shape, + in.dtype(), + std::make_shared(to_stream(s), shape), + std::move(p_inputs))); } return outputs; } @@ -1599,7 +1955,7 @@ std::pair broadcast_arrays( array equal(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -1607,7 +1963,7 @@ array equal(const array& a, const array& b, StreamOrDevice s /* = {} */) { array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, @@ -1618,7 +1974,7 @@ array not_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) { array greater(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -1629,7 +1985,7 @@ array greater_equal( StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, @@ -1640,7 +1996,7 @@ array greater_equal( array less(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -1648,7 +2004,7 @@ array less(const array& a, const array& b, StreamOrDevice s /* = {} */) { array less_equal(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, @@ -1688,6 +2044,9 @@ array isinf(const array& a, StreamOrDevice s /* = {} */) { if (issubdtype(a.dtype(), integer) || a.dtype() == bool_) { return full(a.shape(), false, bool_, s); } + if (issubdtype(a.dtype(), complexfloating)) { + return logical_or(isinf(real(a, s), s), isinf(imag(a, s), s), s); + } return logical_or(isposinf(a, s), isneginf(a, s), s); } @@ -1742,11 +2101,11 @@ array nan_to_num( auto type_to_max = [](const auto& dtype) -> float { if (dtype == float32) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == bfloat16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else if (dtype == float16) { - return std::numeric_limits::max(); + return numeric_limits::max(); } else { std::ostringstream msg; msg << "[nan_to_num] Does not yet support given type: " << dtype << "."; @@ -1924,7 +2283,33 @@ array sum( int axis, bool keepdims /* = false */, StreamOrDevice s /* = {} */) { - return sum(a, std::vector{axis}, keepdims, s); + return sum(a, std::vector{axis}, keepdims, s); +} + +array count_nonzero( + const array& a, + bool keepdims /* = false */, + StreamOrDevice s /* = {} */) { + std::vector axes(a.ndim()); + std::iota(axes.begin(), axes.end(), 0); + return count_nonzero(a, axes, keepdims, s); +} + +array count_nonzero( + const array& a, + int axis, + bool keepdims /* = false */, + StreamOrDevice s /* = {} */) { + return count_nonzero(a, std::vector{axis}, keepdims, s); +} + +array count_nonzero( + const array& a, + const std::vector& axes, + bool keepdims /* = false */, + StreamOrDevice s /* = {} */) { + auto nz = astype(not_equal(a, array(0, a.dtype()), s), int32, s); + return sum(nz, axes, keepdims, s); } array mean(const array& a, bool keepdims, StreamOrDevice s /* = {}*/) { @@ -2026,6 +2411,16 @@ array median( array(0.5, dtype), s); } + // Sorting moves NaN to the end, so the midpoint slice never selects it. + // Propagate it explicitly to stay consistent with max, min and mean. + if (issubdtype(a.dtype(), inexact)) { + median_a = where( + any(isnan(flat_a, s), -1, /* keepdims = */ true, s), + array(std::numeric_limits::quiet_NaN(), dtype), + median_a, + s); + } + median_a = squeeze(median_a, -1, s); if (keepdims) { median_a = expand_dims(median_a, sorted_axes, s); @@ -2059,7 +2454,17 @@ array var( StreamOrDevice s /* = {}*/) { auto dtype = at_least_float(a.dtype()); auto mu = mean(a, axes, /* keepdims= */ true, s); - auto v = sum(square(subtract(a, mu, s), s), axes, keepdims, s); + auto d = subtract(a, mu, s); + // The variance of complex values is the mean squared magnitude. Squaring the + // deviations directly gives a complex result which can even be negative, so + // multiply by the conjugate instead. + auto sq = issubdtype(dtype, complexfloating) + ? real(multiply(d, conjugate(d, s), s), s) + : square(d, s); + if (issubdtype(dtype, complexfloating)) { + dtype = float32; + } + auto v = sum(sq, axes, keepdims, s); if (ddof != 0) { auto normalizer = maximum( @@ -2140,7 +2545,7 @@ array prod( out_type = int32; } auto out = (is_noop) - ? a + ? astype(a, out_type, s) : array( std::move(out_shape), out_type, @@ -2171,11 +2576,8 @@ array max( const std::vector& axes, bool keepdims /* = false */, StreamOrDevice s /* = {}*/) { - if (a.size() == 0) { - throw std::invalid_argument("[max] Cannot max reduce zero size array."); - } auto [out_shape, sorted_axes, is_noop] = - compute_reduce_shape(axes, a.shape()); + compute_reduce_shape(axes, a.shape(), "max"); auto out = (is_noop) ? a : array( @@ -2208,14 +2610,11 @@ array min( const std::vector& axes, bool keepdims /* = false */, StreamOrDevice s /* = {}*/) { - if (a.size() == 0) { - throw std::invalid_argument("[min] Cannot min reduce zero size array."); - } if (axes.empty()) { return a; } auto [out_shape, sorted_axes, is_noop] = - compute_reduce_shape(axes, a.shape()); + compute_reduce_shape(axes, a.shape(), "min"); auto out = (is_noop) ? a : array( @@ -2311,6 +2710,82 @@ array argmax( return out; } +array bartlett(int M, StreamOrDevice s /* = {} */) { + if (M < 1) { + return array({}); + } + if (M == 1) { + return ones({1}, float32, s); + } + + auto n = arange(0, M, float32, s); + float factor_val = 2.0f / (M - 1); + auto factor = array(factor_val, float32); + auto term = subtract(multiply(factor, n, s), array(1.0f, float32), s); + return subtract(array(1.0f, float32), abs(term, s), s); +} + +array hanning(int M, StreamOrDevice s /* = {} */) { + if (M < 1) { + return array({}); + } + if (M == 1) { + return ones({1}, float32, s); + } + + auto n = arange(0, M, float32, s); + array factor(M_PI / (M - 1), float32); + return square(sin(multiply(factor, n, s), s), s); +} + +array hamming(int M, StreamOrDevice s /* = {} */) { + if (M < 1) { + return array({}); + } + if (M == 1) { + return ones({1}, float32, s); + } + + auto n = arange(0, M, float32, s); + float factor_val = (2.0 * M_PI) / (M - 1); + auto factor = array(factor_val, float32); + + auto arg = multiply(factor, n, s); + auto cos_vals = cos(arg, s); + + auto left_coef = array(0.54f, float32); + auto right_coef = array(0.46f, float32); + + return subtract(left_coef, multiply(right_coef, cos_vals, s), s); +} + +array blackman(int M, StreamOrDevice s /* = {} */) { + if (M < 1) { + return array({}); + } + if (M == 1) { + return ones({1}, float32, s); + } + + auto n = arange(0, M, float32, s); + + float arg_val = (2.0 * M_PI) / (M - 1); + auto x = multiply(array(arg_val, float32), n, s); + + auto cos_x = cos(x, s); + + auto alpha = array(0.34f, float32); + auto beta = array(0.5f, float32); + auto gamma = array(0.16f, float32); + + auto term1 = multiply(beta, cos_x, s); + + auto cos_sq = square(cos_x, s); + auto term2 = multiply(gamma, cos_sq, s); + + return add(subtract(alpha, term1, s), term2, s); +} + /** Returns a sorted copy of the flattened array. */ array sort(const array& a, StreamOrDevice s /* = {} */) { int size = a.size(); @@ -2319,17 +2794,9 @@ array sort(const array& a, StreamOrDevice s /* = {} */) { /** Returns a sorted copy of the array along a given axis. */ array sort(const array& a, int axis, StreamOrDevice s /* = {} */) { - // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[sort] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - + auto ax = normalize_axis_index(axis, a.ndim(), "[sort] "); return array( - a.shape(), a.dtype(), std::make_shared(to_stream(s), axis), {a}); + a.shape(), a.dtype(), std::make_shared(to_stream(s), ax), {a}); } /** Returns indices that sort the flattened array. */ @@ -2340,17 +2807,9 @@ array argsort(const array& a, StreamOrDevice s /* = {} */) { /** Returns indices that sort the array along a given axis. */ array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) { - // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[argsort] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - + auto ax = normalize_axis_index(axis, a.ndim(), "[argsort] "); return array( - a.shape(), uint32, std::make_shared(to_stream(s), axis), {a}); + a.shape(), uint32, std::make_shared(to_stream(s), ax), {a}); } /** @@ -2372,14 +2831,7 @@ array partition( int axis, StreamOrDevice s /* = {} */) { // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[partition] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - int axis_ = axis < 0 ? axis + a.ndim() : axis; + int axis_ = normalize_axis_index(axis, a.ndim(), "[partition] "); int kth_ = kth < 0 ? kth + a.shape(axis) : kth; if (kth_ < 0 || kth_ >= a.shape(axis_)) { std::ostringstream msg; @@ -2413,14 +2865,7 @@ array argpartition( int axis, StreamOrDevice s /* = {} */) { // Check for valid axis - if (axis + static_cast(a.ndim()) < 0 || - axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[argpartition] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } - int axis_ = axis < 0 ? axis + a.ndim() : axis; + int axis_ = normalize_axis_index(axis, a.ndim(), "[argpartition] "); int kth_ = kth < 0 ? kth + a.shape(axis) : kth; if (kth_ < 0 || kth_ >= a.shape(axis_)) { std::ostringstream msg; @@ -2435,6 +2880,37 @@ array argpartition( {a}); } +array searchsorted( + const array& sorted_sequence, + const array& values, + const std::string& side /* = "left" */, + StreamOrDevice s /* = {} */) { + if (side != "left" && side != "right") { + std::ostringstream msg; + msg << "[searchsorted] Invalid side '" << side + << "'. Expected 'left' or 'right'."; + throw std::invalid_argument(msg.str()); + } + if (sorted_sequence.ndim() != 1) { + std::ostringstream msg; + msg << "[searchsorted] The sorted sequence must be 1-D but has " + << sorted_sequence.ndim() << " dimensions."; + throw std::invalid_argument(msg.str()); + } + if (sorted_sequence.size() > UINT32_MAX) { + std::ostringstream msg; + msg << "[searchsorted] The sorted sequence has " << sorted_sequence.size() + << " elements, more than the uint32 output can index."; + throw std::invalid_argument(msg.str()); + } + auto dtype = promote_types(sorted_sequence.dtype(), values.dtype()); + return array( + values.shape(), + uint32, + std::make_shared(to_stream(s), side == "right"), + {astype(sorted_sequence, dtype, s), astype(values, dtype, s)}); +} + /** Returns topk elements of the flattened array. */ array topk(const array& a, int k, StreamOrDevice s /* = {}*/) { int size = a.size(); @@ -2444,13 +2920,7 @@ array topk(const array& a, int k, StreamOrDevice s /* = {}*/) { /** Returns topk elements of the array along a given axis. */ array topk(const array& a, int k, int axis, StreamOrDevice s /* = {}*/) { // Check for valid axis - int axis_ = axis < 0 ? axis + a.ndim() : axis; - if (axis_ < 0 || axis_ >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[topk] Received invalid axis " << axis << " for array with " - << a.ndim() << " dimensions."; - throw std::invalid_argument(msg.str()); - } + int axis_ = normalize_axis_index(axis, a.ndim(), "[topk] "); if (k < 0 || k > a.shape(axis_)) { std::ostringstream msg; msg << "[topk] Received invalid k=" << k << " along axis " << axis @@ -2511,7 +2981,7 @@ array logsumexp( std::make_shared(to_stream(s)), {astype(a, dtype, s)}); if (!keepdims) { - out = squeeze(out, -1, s); + out = squeeze(out, axes, s); } return out; } @@ -2541,6 +3011,10 @@ array abs(const array& a, StreamOrDevice s /* = {} */) { return out; } +array positive(const array& a, StreamOrDevice s /* = {} */) { + return array(a); +} + array negative(const array& a, StreamOrDevice s /* = {} */) { if (a.dtype() == bool_) { auto msg = "[negative] Not supported for bool, use logical_not instead."; @@ -2568,7 +3042,7 @@ array logical_not(const array& a, StreamOrDevice s /* = {} */) { array logical_and(const array& a, const array& b, StreamOrDevice s /* = {} */) { // Broadcast arrays to a common shape auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, @@ -2582,7 +3056,7 @@ array operator&&(const array& a, const array& b) { array logical_or(const array& a, const array& b, StreamOrDevice s /* = {} */) { // Broadcast arrays to a common shape auto inputs = broadcast_arrays({astype(a, bool_, s), astype(b, bool_, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, bool_, @@ -2593,6 +3067,10 @@ array operator||(const array& a, const array& b) { return logical_or(a, b); } +array logical_xor(const array& a, const array& b, StreamOrDevice s /* = {} */) { + return not_equal(astype(a, bool_, s), astype(b, bool_, s), s); +} + array reciprocal(const array& a, StreamOrDevice s /* = {} */) { auto dtype = at_least_float(a.dtype()); return divide(array(1.0f, dtype), a, to_stream(s)); @@ -2602,7 +3080,7 @@ array add(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -2615,7 +3093,7 @@ array subtract(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, @@ -2631,7 +3109,7 @@ array multiply(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, @@ -2647,7 +3125,7 @@ array divide(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = at_least_float(promote_types(a.dtype(), b.dtype())); auto inputs = broadcast_arrays( {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -2671,16 +3149,19 @@ array floor_divide( } auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); } array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto dtype = promote_types(a.dtype(), b.dtype()); + if (issubdtype(dtype, complexfloating)) { + throw std::invalid_argument("[remainder] Complex type not supported."); + } auto inputs = broadcast_arrays( {astype(a, dtype, s), astype(b, dtype, to_stream(s))}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, dtype, @@ -2710,7 +3191,7 @@ array maximum(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, @@ -2722,7 +3203,7 @@ array minimum(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = promote_types(a.dtype(), b.dtype()); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, @@ -2740,11 +3221,22 @@ array floor(const array& a, StreamOrDevice s /* = {} */) { array ceil(const array& a, StreamOrDevice s /* = {} */) { if (a.dtype() == complex64) { - throw std::invalid_argument("[floor] Not supported for complex64."); + throw std::invalid_argument("[ceil] Not supported for complex64."); } return array(a.shape(), a.dtype(), std::make_shared(to_stream(s)), {a}); } +array trunc(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[trunc] Not supported for complex64."); + } + if (issubdtype(a.dtype(), integer)) { + return array(a); + } + auto zero = array(0, a.dtype()); + return where(less(a, zero, s), ceil(a, s), floor(a, s), s); +} + array square(const array& a, StreamOrDevice s /* = {} */) { return array( a.shape(), a.dtype(), std::make_shared(to_stream(s)), {a}); @@ -2757,6 +3249,9 @@ array exp(const array& a, StreamOrDevice s /* = {} */) { } array expm1(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[expm1] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); auto input = astype(a, dtype, s); return array( @@ -2803,9 +3298,12 @@ array arctan(const array& a, StreamOrDevice s /* = {} */) { } array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64 || b.dtype() == complex64) { + throw std::invalid_argument("[arctan2] Not supported for complex64."); + } auto dtype = at_least_float(promote_types(a.dtype(), b.dtype())); auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, dtype, std::make_shared(to_stream(s)), std::move(inputs)); } @@ -2901,7 +3399,7 @@ array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) { auto out_type = at_least_float(promote_types(a.dtype(), b.dtype())); auto inputs = broadcast_arrays({astype(a, out_type, s), astype(b, out_type, s)}, s); - auto& shape = inputs[0].shape(); + auto shape = inputs[0].shape(); return array( shape, out_type, @@ -2910,6 +3408,9 @@ array logaddexp(const array& a, const array& b, StreamOrDevice s /* = {} */) { } array sigmoid(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[sigmoid] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); auto input = astype(a, dtype, s); return array( @@ -2917,6 +3418,9 @@ array sigmoid(const array& a, StreamOrDevice s /* = {} */) { } array erf(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[erf] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); return array( a.shape(), @@ -2926,6 +3430,9 @@ array erf(const array& a, StreamOrDevice s /* = {} */) { } array erfinv(const array& a, StreamOrDevice s /* = {} */) { + if (a.dtype() == complex64) { + throw std::invalid_argument("[erfinv] Not supported for complex64."); + } auto dtype = at_least_float(a.dtype()); return array( a.shape(), @@ -3000,9 +3507,14 @@ array matmul( } // We can batch the multiplication by reshaping a - if (in_a.ndim() > 2 && in_b.ndim() <= 2) { + bool flattened_a = false; + + // Avoid flatten/unflatten during dynamic tracing because unflatten stores the + // trace-time shape, which breaks shapeless replay with dynamic batch sizes. + if (in_a.ndim() > 2 && in_b.ndim() <= 2 && !detail::in_dynamic_tracing()) { a = flatten(a, 0, -2, s); - } else if (in_b.ndim() > 2) { + flattened_a = true; + } else if (in_a.ndim() > 2 || in_b.ndim() > 2) { std::tie(a, b) = broadcast_arrays(a, b, {-2, -1}, s); } @@ -3014,7 +3526,8 @@ array matmul( out_type, std::make_shared(to_stream(s)), {a, b}); - if (in_a.ndim() > 2 && in_b.ndim() <= 2) { + + if (flattened_a) { auto orig_shape = in_a.shape(); orig_shape.pop_back(); out = unflatten(out, 0, std::move(orig_shape), s); @@ -3131,11 +3644,13 @@ array kron(const array& a, const array& b, StreamOrDevice s /* = {} */) { for (int i = ndim - 1, j = a.ndim() - 1; j >= 0; j--, i--) { a_shape[2 * i] = a.shape(j); - out_shape[i] *= a.shape(j); + out_shape[i] = + safe_cast(static_cast(out_shape[i]) * a.shape(j), "kron"); } for (int i = ndim - 1, j = b.ndim() - 1; j >= 0; j--, i--) { b_shape[2 * i + 1] = b.shape(j); - out_shape[i] *= b.shape(j); + out_shape[i] = + safe_cast(static_cast(out_shape[i]) * b.shape(j), "kron"); } return reshape( @@ -3231,12 +3746,8 @@ array take_along_axis( const array& indices, int axis, StreamOrDevice s /* = {} */) { - if (axis + a.ndim() < 0 || axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << "[take_along_axis] Received invalid axis for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + axis = normalize_axis_index(axis, a.ndim(), "[take_along_axis] "); if (indices.ndim() != a.ndim()) { std::ostringstream msg; @@ -3245,9 +3756,6 @@ array take_along_axis( throw std::invalid_argument(msg.str()); } - // Allow negative axis - axis = axis < 0 ? a.ndim() + axis : axis; - // Broadcast indices and input ignoring the take axis auto inputs = broadcast_arrays({a, indices}, std::vector{axis - int(a.ndim())}, s); @@ -3269,12 +3777,8 @@ array scatter_axis( StreamOrDevice s) { std::string prefix = (mode == ScatterAxis::None) ? "[put_along_axis]" : "[scatter_add_axis]"; - if (axis + a.ndim() < 0 || axis >= static_cast(a.ndim())) { - std::ostringstream msg; - msg << prefix << " Received invalid axis for array with " << a.ndim() - << " dimensions."; - throw std::invalid_argument(msg.str()); - } + // Normalizes and validates the axis, including negative out of bounds ones + axis = normalize_axis_index(axis, a.ndim(), prefix + " "); if (indices.ndim() != a.ndim()) { std::ostringstream msg; @@ -3299,9 +3803,6 @@ array scatter_axis( auto inputs = broadcast_arrays({indices, upd}, s); inputs.insert(inputs.begin(), a); - // Allow negative axis - axis = axis < 0 ? a.ndim() + axis : axis; - // Broadcast src, indices, values while ignoring the take axis inputs = broadcast_arrays(inputs, {axis - int(a.ndim())}, s); @@ -3409,7 +3910,8 @@ array scatter( } // TODO, remove when scatter supports 64-bit outputs - if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8) { + if (to_stream(s).device == Device::gpu && size_of(a.dtype()) == 8 && + !(a.dtype() == complex64 && mode == Scatter::Sum)) { std::ostringstream msg; msg << "[scatter] GPU scatter does not yet support " << a.dtype() << " for the input or updates."; @@ -3640,6 +4142,7 @@ array cumsum( int axis, bool reverse /* = false*/, bool inclusive /* = true*/, + std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { int ndim = a.ndim(); if (axis >= ndim || axis < -ndim) { @@ -3649,21 +4152,24 @@ array cumsum( throw std::invalid_argument(msg.str()); } axis = (axis + a.ndim()) % a.ndim(); - auto out_type = a.dtype() == bool_ ? int32 : a.dtype(); + auto x = dtype ? astype(a, *dtype, s) : a; + auto out_type = x.dtype() == bool_ ? int32 : x.dtype(); return array( - a.shape(), + x.shape(), out_type, std::make_shared( to_stream(s), Scan::ReduceType::Sum, axis, reverse, inclusive), - {a}); + {x}); } array cumsum( const array& a, bool reverse /* = false*/, bool inclusive /* = true*/, + std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { - return cumsum(flatten(a, to_stream(s)), 0, reverse, inclusive, to_stream(s)); + return cumsum( + flatten(a, to_stream(s)), 0, reverse, inclusive, dtype, to_stream(s)); } array cumprod( @@ -3671,6 +4177,7 @@ array cumprod( int axis, bool reverse /* = false*/, bool inclusive /* = true*/, + std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { int ndim = a.ndim(); if (axis >= ndim || axis < -ndim) { @@ -3680,20 +4187,22 @@ array cumprod( throw std::invalid_argument(msg.str()); } axis = (axis + a.ndim()) % a.ndim(); + auto x = dtype ? astype(a, *dtype, s) : a; return array( - a.shape(), - a.dtype(), + x.shape(), + x.dtype(), std::make_shared( to_stream(s), Scan::ReduceType::Prod, axis, reverse, inclusive), - {a}); + {x}); } array cumprod( const array& a, bool reverse /* = false*/, bool inclusive /* = true*/, + std::optional dtype /* = std::nullopt*/, StreamOrDevice s /* = {}*/) { - return cumprod(flatten(a, s), 0, reverse, inclusive, s); + return cumprod(flatten(a, s), 0, reverse, inclusive, dtype, s); } array cummax( @@ -3756,6 +4265,33 @@ array cummin( return cummin(flatten(a, s), 0, reverse, inclusive, s); } +array diff( + const array& a, + int n /* = 1 */, + int axis /* = -1 */, + StreamOrDevice s /* = {} */) { + int ndim = static_cast(a.ndim()); + int ax = axis < 0 ? axis + ndim : axis; + if (ax < 0 || ax >= ndim) { + throw std::invalid_argument("[diff] Axis is out of bounds for the array."); + } + if (n < 0) { + throw std::invalid_argument("[diff] Order `n` must be non-negative."); + } + array x = a; + for (int i = 0; i < n; ++i) { + Shape upper_start(x.ndim(), 0); + Shape lower_stop = x.shape(); + Shape strides(x.ndim(), 1); + upper_start[ax] = 1; + lower_stop[ax] = x.shape(ax) - 1; + auto upper = slice(x, upper_start, x.shape(), strides, s); + auto lower = slice(x, Shape(x.ndim(), 0), lower_stop, strides, s); + x = subtract(upper, lower, s); + } + return x; +} + array logcumsumexp( const array& a, int axis, @@ -3932,22 +4468,29 @@ array conv_transpose_general( std::vector padding_lo(padding.size()); std::vector padding_hi(padding.size()); for (int i = 0; i < padding.size(); ++i) { - int wt_size = 1 + dilation[i] * (weight.shape(1 + i) - 1); - padding_lo[i] = wt_size - padding[i] - 1; + int64_t wt_size = + 1 + static_cast(dilation[i]) * (weight.shape(1 + i) - 1); + padding_lo[i] = safe_cast(wt_size - padding[i] - 1, "conv"); - int conv_output_shape = (input.shape(i + 1) - 1) * stride[i] - - 2 * padding[i] + dilation[i] * (weight.shape(i + 1) - 1) + 1; + int64_t conv_output_shape = + static_cast(input.shape(i + 1) - 1) * stride[i] - + 2 * static_cast(padding[i]) + + static_cast(dilation[i]) * (weight.shape(i + 1) - 1) + 1; - int in_size = 1 + (conv_output_shape - 1); - int out_size = 1 + stride[i] * (input.shape(1 + i) - 1); - padding_hi[i] = in_size - out_size + padding[i] + - output_padding[i]; // Adjust with output_padding + int64_t in_size = 1 + (conv_output_shape - 1); + int64_t out_size = + 1 + static_cast(stride[i]) * (input.shape(1 + i) - 1); + // Adjust with output_padding + padding_hi[i] = + safe_cast(in_size - out_size + padding[i] + output_padding[i], "conv"); } + auto ndim = stride.size(); + return conv_general( /* const array& input = */ input, /* const array& weight = */ weight, - /* std::vector stride = */ std::vector(stride.size(), 1), + /* std::vector stride = */ std::vector(ndim, 1), /* std::vector padding_lo = */ std::move(padding_lo), /* std::vector padding_hi = */ std::move(padding_hi), /* std::vector kernel_dilation = */ std::move(dilation), @@ -4090,7 +4633,8 @@ array conv_general( for (int i = 0; i < spatial_dims; i++) { if (padding_lo[i] < 0) { - starts[i + 1] -= padding_lo[i]; + starts[i + 1] = safe_cast( + starts[i + 1] - static_cast(padding_lo[i]), "conv"); padding_lo[i] = 0; } @@ -4209,11 +4753,39 @@ std::pair validate_mode_with_type( } } +void validate_global_scale( + std::string_view tag, + QuantizationMode qmode, + const std::optional& global_scale) { + if (global_scale.has_value()) { + if (qmode != QuantizationMode::Nvfp4) { + std::ostringstream msg; + msg << "[" << tag << "] Global scale is only supported for 'nvfp4' " + << "quantization mode."; + throw std::invalid_argument(msg.str()); + } else { + if (global_scale->size() != 1) { + std::ostringstream msg; + msg << "[" << tag << "] Global scale must be a scalar but got shape " + << global_scale->shape() << "."; + throw std::invalid_argument(msg.str()); + } + // TODO: not sure if type should be restricted to float32 + if (global_scale->dtype() != float32) { + std::ostringstream msg; + msg << "[" << tag << "] Global scale must have dtype float32 but got " + << global_scale->dtype() << "."; + throw std::invalid_argument(msg.str()); + } + } + } +} + 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 */, @@ -4251,7 +4823,6 @@ array quantized_matmul( if (x.ndim() > 2 && w.ndim() > 2) { inputs = broadcast_arrays(inputs, {-2, -1}, s); } - auto out_shape = inputs[0].shape(); out_shape.back() = w_outer_dims; return array( @@ -4263,11 +4834,14 @@ array quantized_matmul( } void validate_qqmm_inputs( - array x, - array w, - std::optional scales_w, + const array& x, + const array& w, + const std::optional& scales_w, int group_size, - int bits) { + int bits, + const std::optional& global_scale_x, + const std::optional& global_scale_w, + QuantizationMode qmode) { // check 2D (for now) if (x.ndim() > 2 || w.ndim() > 2) { std::ostringstream msg; @@ -4304,12 +4878,25 @@ void validate_qqmm_inputs( << "first argument dtype == " << x.dtype() << "."; throw std::invalid_argument(msg.str()); } + // validate global scales + validate_global_scale("qqmm", qmode, global_scale_x); + validate_global_scale("qqmm", qmode, global_scale_w); + // For nvfp4 mode, both global scales must be provided together or neither + if (qmode == QuantizationMode::Nvfp4) { + bool has_x = global_scale_x.has_value(); + bool has_w = global_scale_w.has_value(); + if (has_x != has_w) { + throw std::invalid_argument( + "[qqmm] For nvfp4 mode, either both global_scale_x and " + "global_scale_w must be provided, or neither."); + } + } } std::pair extract_qqmm_dims( - array x, - array w, - std::optional scales_w, + const array& x, + const array& w, + const std::optional& scales_w, int group_size, int bits) { if (w.dtype() != uint32) { @@ -4337,22 +4924,17 @@ std::pair extract_qqmm_dims( } array qqmm( - array in_x, - array w, - std::optional scales_w, + const array& in_x, + const array& w, + const std::optional& scales_w, 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 */, StreamOrDevice s /* = {} */) { auto stream = to_stream(s); auto qmode = string_to_quantization_mode(mode, "qqmm"); - // cuBLAS block scaled matmul only supports nvfp4 and mxfp8 - if (qmode != QuantizationMode::Nvfp4 && qmode != QuantizationMode::Mxfp8) { - std::ostringstream msg; - msg << "[qqmm] Only 'nvfp4' and 'mxfp8' quantization modes are supported but '" - << mode << "' was provided."; - throw std::invalid_argument(msg.str()); - } // we need to check 2 cases: // 1. w is quantized, scales is provided // 2. w is not quantized, scales is not provided @@ -4369,7 +4951,8 @@ array qqmm( } // validate inputs - validate_qqmm_inputs(x, w, scales_w, group_size, bits); + validate_qqmm_inputs( + x, w, scales_w, group_size, bits, global_scale_x, global_scale_w, qmode); // validate and extract shapes auto [w_inner_dims, w_outer_dims] = extract_qqmm_dims(x, w, scales_w, group_size, bits); @@ -4380,6 +4963,11 @@ array qqmm( if (scales_w.has_value()) { inputs.push_back(*scales_w); } + if (global_scale_x.has_value() && global_scale_w.has_value()) { + inputs.push_back(*global_scale_x); + inputs.push_back(*global_scale_w); + } + auto out_shape = inputs[0].shape(); out_shape.back() = w_outer_dims; auto out = array( @@ -4515,6 +5103,7 @@ std::vector fp_quantize( int group_size, int bits, QuantizationMode mode, + const std::optional& global_scale /* = std::nullopt */, Stream s) { int expected_gs = mode == QuantizationMode::Nvfp4 ? 16 : 32; int expected_bits = mode == QuantizationMode::Mxfp8 ? 8 : 4; @@ -4532,6 +5121,12 @@ std::vector fp_quantize( << bits << "."; throw std::invalid_argument(msg.str()); } + + auto inputs = std::vector{w}; + if (global_scale.has_value()) { + inputs.push_back(global_scale.value()); + } + auto fallback = [bits = bits, group_size = group_size, s]( const std::vector& inputs) -> std::vector { auto& w = inputs[0]; @@ -4543,16 +5138,27 @@ std::vector fp_quantize( divide(max(abs(wq, s), -1, true, s), array(maxval, w.dtype()), s); if (group_size == 16) { // convert to e4m3 + auto scale_encode = inputs.size() > 1 + ? divide(array(448.0f * 6.0f, float32), inputs[1], s) + : array(1.0f, float32); + scales = multiply(scales, scale_encode, s); scales = to_fp8(scales, s); - wq = divide(wq, from_fp8(scales, w.dtype(), s), s); + wq = multiply( + divide(wq, from_fp8(scales, w.dtype(), s), s), scale_encode, s); } else { // convert to e8m0 auto z = array(0, scales.dtype()); - scales = where( - equal(scales, z, s), - z, - astype(round(log2(scales, s), s), int32, s), + // Round the scale up so the block maximum stays representable, + // matching the CUDA backend. + auto exponent = astype(round(log2(scales, s), s), int32, s); + auto decoded = + power(array(2.0f, float32), astype(exponent, float32, s), s); + exponent = where( + less(decoded, astype(scales, float32, s), s), + add(exponent, array(1, int32), s), + exponent, s); + scales = where(equal(scales, z, s), z, exponent, s); wq = divide(wq, power(array(2.0f, w.dtype()), scales, s), s); scales = astype(add(scales, array(127, int32), s), uint8, s); @@ -4600,9 +5206,9 @@ std::vector fp_quantize( {uint32, uint8}, std::make_shared( s, fallback, group_size, bits, mode, false), - {w}); + inputs); } - return fallback({w}); + return fallback(inputs); } std::vector quantize( @@ -4610,6 +5216,7 @@ std::vector quantize( std::optional group_size_ /* = std::nullopt */, std::optional bits_ /* = std::nullopt */, const std::string& mode /* = "affine" */, + const std::optional& global_scale /* = std::nullopt */, StreamOrDevice s /* = {} */) { auto qmode = string_to_quantization_mode(mode, "quantize"); auto [group_size, bits] = @@ -4636,11 +5243,11 @@ std::vector quantize( << " matrix has shape " << w.shape(); throw std::invalid_argument(msg.str()); } - + validate_global_scale("quantize", qmode, global_scale); if (qmode == QuantizationMode::Affine) { return affine_quantize(w, group_size, bits, s); } else { - return fp_quantize(w, group_size, bits, qmode, to_stream(s)); + return fp_quantize(w, group_size, bits, qmode, global_scale, to_stream(s)); } } @@ -4745,6 +5352,7 @@ array fp_dequantize( int bits, Dtype out_type, QuantizationMode mode, + const std::optional& global_scale /* = std::nullopt */, Stream s) { int expected_gs = mode == QuantizationMode::Nvfp4 ? 16 : 32; int expected_bits = mode == QuantizationMode::Mxfp8 ? 8 : 4; @@ -4789,6 +5397,11 @@ array fp_dequantize( throw std::invalid_argument(msg.str()); } + auto inputs = std::vector{w, scales}; + if (global_scale.has_value()) { + inputs.push_back(global_scale.value()); + } + auto fallback = [wshape = std::move(wshape), sshape = std::move(sshape), @@ -4831,13 +5444,17 @@ array fp_dequantize( out = reshape(out, {-1, group_size}, s); scales = reshape(scales, {-1, 1}, s); if (group_size == 16) { - scales = from_fp8(scales, out_type, s); + array inv_scale_enc = inputs.size() > 2 + ? divide(inputs[2], array(448.0f * 6.0f, out_type), s) + : array(1.0f, out_type); + scales = multiply(from_fp8(scales, out_type, s), inv_scale_enc, s); } else { scales = subtract(astype(scales, out_type, s), array(127, out_type), s); scales = power(array(2.0f, out_type), scales, s); } return {reshape(multiply(out, scales, s), wshape, s)}; }; + if (s.device == Device::gpu) { auto out_shape = w.shape(); out_shape.back() = out_size; @@ -4846,9 +5463,9 @@ array fp_dequantize( out_type, std::make_shared( s, fallback, group_size, bits, mode, true), - {w, scales}); + inputs); } - return fallback({w, scales})[0]; + return fallback(inputs)[0]; } array dequantize( @@ -4858,6 +5475,7 @@ array dequantize( std::optional group_size_ /* = std::nullopt */, std::optional bits_ /* = std::nullopt */, const std::string& mode /* = "affine" */, + const std::optional& global_scale /* = std::nullopt */, std::optional dtype /* = std::nullopt */, StreamOrDevice s /* = {} */) { auto [out_type, qmode] = @@ -4884,6 +5502,7 @@ array dequantize( << "but it has only " << w.ndim() << "."; throw std::invalid_argument(msg.str()); } + validate_global_scale("dequantize", qmode, global_scale); if (qmode == QuantizationMode::Affine) { return astype( @@ -4892,7 +5511,14 @@ array dequantize( s); } else { return fp_dequantize( - w, scales, group_size, bits, out_type, qmode, to_stream(s)); + w, + scales, + group_size, + bits, + out_type, + qmode, + global_scale, + to_stream(s)); } } @@ -4968,30 +5594,11 @@ array gather_qmm( } // Extract indices and broadcast them - array lhs_indices = indices_or_default(lhs_indices_, x, s); - array rhs_indices = indices_or_default(rhs_indices_, w, s); + array lhs_indices = indices_or_default("[gather_qmm]", lhs_indices_, x, s); + array rhs_indices = indices_or_default("[gather_qmm]", rhs_indices_, w, s); std::tie(lhs_indices, rhs_indices) = broadcast_arrays(lhs_indices, rhs_indices, s); - if (!issubdtype(lhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_qmm] Got lhs_indices with invalid dtype. Indices must be integral."); - } - - if (!issubdtype(rhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_qmm] Got rhs_indices with invalid dtype. Indices must be integral."); - } - if (x.ndim() < 2) { - std::ostringstream msg; - msg << "[gather_qmm] Non-quantized input must have at least two" - << " dimensions but got input with shape " << x.shape() << "."; - throw std::invalid_argument(msg.str()); - } - - lhs_indices = astype(lhs_indices, uint32, s); - rhs_indices = astype(rhs_indices, uint32, s); - // Compute the full output shape auto out_shape = lhs_indices.shape(); out_shape.push_back(x.shape(-2)); @@ -5027,6 +5634,56 @@ array gather_qmm( std::move(inputs)); } +array gather_qqmm( + const array& x, + const array& w, + const std::optional& scales_w, + const std::optional& lhs_indices_, + const std::optional& rhs_indices_, + std::optional group_size_, + std::optional bits_, + const std::string& mode, + const std::optional& global_scale_x, + const std::optional& global_scale_w, + bool sorted_indices, + StreamOrDevice s) { + auto stream = to_stream(s); + auto qmode = string_to_quantization_mode(mode, "gather_qqmm"); + auto [group_size, bits] = + quantization_params_from_mode(qmode, group_size_, bits_); + + // Extract indices and broadcast them + array lhs_indices = indices_or_default("[gather_qqmm]", lhs_indices_, x, s); + array rhs_indices = indices_or_default("[gather_qqmm]", rhs_indices_, w, s); + std::tie(lhs_indices, rhs_indices) = + broadcast_arrays(lhs_indices, rhs_indices, s); + + std::vector inputs = { + x, + w, + lhs_indices, + rhs_indices, + }; + if (scales_w.has_value()) { + inputs.push_back(*scales_w); + } + if (global_scale_x.has_value() && global_scale_w.has_value()) { + inputs.push_back(*global_scale_x); + inputs.push_back(*global_scale_w); + } + + auto [w_inner_dims, w_outer_dims] = + extract_qqmm_dims(x, w, scales_w, group_size, bits); + auto out_shape = lhs_indices.shape(); + out_shape.push_back(x.shape(-2)); + out_shape.push_back(w_outer_dims); + return array( + std::move(out_shape), + x.dtype(), + std::make_shared(stream, group_size, bits, qmode), + std::move(inputs)); +} + array tensordot( const array& a, const array& b, @@ -5129,6 +5786,28 @@ array inner(const array& a, const array& b, StreamOrDevice s /* = {} */) { return tensordot(a, b, {-1}, {-1}, s); } +array vecdot( + const array& a, + const array& b, + int axis /* = -1 */, + StreamOrDevice s /* = {} */) { + if (a.ndim() == 0 || b.ndim() == 0) { + throw std::invalid_argument("[vecdot] inputs must be at least 1D."); + } + int ax = axis < 0 ? axis + a.ndim() : axis; + if (ax < 0 || ax >= a.ndim()) { + throw std::invalid_argument("[vecdot] axis is out of bounds."); + } + if (axis < 0 ? axis + b.ndim() != ax : axis >= b.ndim()) { + throw std::invalid_argument("[vecdot] axis is out of bounds."); + } + if (a.shape(ax) != b.shape(ax)) { + throw std::invalid_argument( + "[vecdot] a and b must have the same size along axis."); + } + return sum(multiply(conjugate(a, s), b, s), ax, false, s); +} + /** Compute D = beta * C + alpha * (A @ B) */ array addmm( array c, @@ -5503,28 +6182,14 @@ array gather_mm( b = astype(b, out_type, s); // Handle broadcasting - array lhs_indices = indices_or_default(lhs_indices_, a, s); - array rhs_indices = indices_or_default(rhs_indices_, b, s); - - if (!issubdtype(lhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_mm] Got lhs_indices with invalid dtype. Indices must be integral."); - } - - if (!issubdtype(rhs_indices.dtype(), integer)) { - throw std::invalid_argument( - "[gather_mm] Got rhs_indices with invalid dtype. Indices must be integral."); - } - - lhs_indices = astype(lhs_indices, uint32, s); - rhs_indices = astype(rhs_indices, uint32, s); + array lhs_indices = indices_or_default("[gather_mm]", lhs_indices_, a, s); + array rhs_indices = indices_or_default("[gather_mm]", rhs_indices_, b, s); + std::tie(lhs_indices, rhs_indices) = + broadcast_arrays(lhs_indices, rhs_indices, s); int M = a.shape(-2); int N = b.shape(-1); - std::tie(lhs_indices, rhs_indices) = - broadcast_arrays(lhs_indices, rhs_indices, s); - auto out_shape = lhs_indices.shape(); out_shape.push_back(M); out_shape.push_back(N); @@ -5659,6 +6324,11 @@ array diag(const array& a, int k /* = 0 */, StreamOrDevice s /* = {} */) { int a_size = a.size(); int n = a_size + std::abs(k); auto res = zeros({n, n}, a.dtype(), s); + if (a_size == 0) { + // Nothing to place on the diagonal, and scattering into the 0x0 output + // produced when k is zero is an error. + return res; + } std::vector indices; auto s1 = std::max(0, -k); @@ -6015,7 +6685,10 @@ array roll( if (size == 0) { continue; // skip rolling this axis if it has size 0 } - auto split_index = (sh < 0) ? (-sh) % size : size - sh % size; + // Promote to 64-bit so negating a shift of INT_MIN does not overflow. + int64_t sh64 = sh; + auto split_index = static_cast( + (sh64 < 0) ? (-sh64) % size : size - sh64 % size); auto parts = split(result, Shape{split_index}, ax, s); std::swap(parts[0], parts[1]); @@ -6034,11 +6707,11 @@ array roll(const array& a, int shift, StreamOrDevice s /* = {} */) { } array roll(const array& a, const Shape& shift, StreamOrDevice s /* = {} */) { - int total_shift = 0; - for (auto& s : shift) { - total_shift += s; + int64_t total_shift = 0; + for (auto& sh : shift) { + total_shift += sh; } - return roll(a, total_shift, s); + return roll(a, safe_cast(total_shift, "roll"), s); } array roll(const array& a, int shift, int axis, StreamOrDevice s /* = {} */) { @@ -6059,11 +6732,12 @@ array roll( const Shape& shift, int axis, StreamOrDevice s /* = {} */) { - int total_shift = 0; - for (auto& s : shift) { - total_shift += s; + int64_t total_shift = 0; + for (auto& sh : shift) { + total_shift += sh; } - return roll(a, Shape{total_shift}, std::vector{axis}, s); + return roll( + a, Shape{safe_cast(total_shift, "roll")}, std::vector{axis}, s); } array real(const array& a, StreamOrDevice s /* = {} */) { diff --git a/Source/Cmlx/mlx/mlx/ops.h b/Source/Cmlx/mlx/mlx/ops.h index 1ff3bbfaa..f597753b1 100644 --- a/Source/Cmlx/mlx/mlx/ops.h +++ b/Source/Cmlx/mlx/mlx/ops.h @@ -38,13 +38,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 = {}); @@ -89,14 +101,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. */ @@ -171,6 +187,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, @@ -224,6 +250,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 = {}); @@ -234,6 +332,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, @@ -531,6 +634,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 = {}) { @@ -666,6 +783,18 @@ min(const array& a, MLX_API array min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); +/** Returns the Hanning window of size M. */ +MLX_API array hanning(int M, StreamOrDevice s = {}); + +/** Returns the Hamming window of size M. */ +MLX_API array hamming(int M, StreamOrDevice s = {}); + +/** Returns the bartlett window of size M. */ +MLX_API array bartlett(int M, StreamOrDevice s = {}); + +/** Returns the Blackmann window of size M. */ +MLX_API array blackman(int M, StreamOrDevice s = {}); + /** Returns the index of the minimum value in the array. */ MLX_API array argmin(const array& a, bool keepdims, StreamOrDevice s = {}); inline array argmin(const array& a, StreamOrDevice s = {}) { @@ -686,16 +815,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 = {}); /** @@ -724,6 +865,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 = {}); @@ -768,6 +916,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); @@ -787,6 +938,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 = {}); @@ -864,6 +1019,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 = {}); @@ -1218,7 +1376,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( @@ -1226,14 +1389,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( @@ -1241,7 +1418,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( @@ -1273,6 +1459,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, @@ -1375,10 +1565,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, @@ -1391,6 +1581,7 @@ MLX_API std::vector quantize( std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "affine", + const std::optional& global_scale = std::nullopt, StreamOrDevice s = {}); /** Dequantize a matrix produced by quantize() */ @@ -1401,17 +1592,20 @@ MLX_API array dequantize( std::optional group_size = std::nullopt, std::optional bits = std::nullopt, const std::string& mode = "affine", + const std::optional& global_scale = std::nullopt, std::optional dtype = std::nullopt, StreamOrDevice s = {}); MLX_API array qqmm( - array x, // input activations - array w, // maybe quantized weights - 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, StreamOrDevice s = {}); /** Convert an E4M3 float8 to the given floating point dtype. */ @@ -1435,6 +1629,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, @@ -1455,6 +1663,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/mlx/mlx/primitives.cpp b/Source/Cmlx/mlx/mlx/primitives.cpp index b07862514..9a1771394 100644 --- a/Source/Cmlx/mlx/mlx/primitives.cpp +++ b/Source/Cmlx/mlx/mlx/primitives.cpp @@ -221,7 +221,9 @@ std::vector Abs::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + assert(primals.size() == 1); + assert(argnums.size() == 1); + return {multiply(cotangents[0], sign(primals[0], stream()), stream())}; } std::vector Abs::jvp( @@ -230,7 +232,14 @@ std::vector Abs::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(argnums.size() == 1); - return {multiply(tangents[0], sign(primals[0], stream()), stream())}; + auto s = sign(primals[0], stream()); + if (issubdtype(primals[0].dtype(), complexfloating)) { + // |z| is real-valued, so its directional derivative is real: + // d|z| = Re(conj(z) * t) / |z| = Re(conj(sign(z)) * t). + return {real( + multiply(tangents[0], conjugate(s, stream()), stream()), stream())}; + } + return {multiply(tangents[0], s, stream())}; } std::pair, std::vector> Abs::vmap( @@ -345,6 +354,13 @@ bool AddMM::is_equivalent(const Primitive& other) const { return (alpha_ == a_other.alpha_ && beta_ == a_other.beta_); } +std::vector AddMM::output_shapes(const std::vector& inputs) { + // inputs are {a, b, c}, output shape is a's shape with last dim from b + auto out_shape = inputs[0].shape(); + out_shape.back() = inputs[1].shape(-1); + return {std::move(out_shape)}; +} + std::pair, std::vector> AddMM::vmap( const std::vector& inputs, const std::vector& axes) { @@ -511,7 +527,7 @@ std::vector ArcTan2::vjp( const std::vector& argnums, const std::vector&) { assert(primals.size() == 2); - assert(argnums.size() == 2); + assert(cotangents.size() == 1); const auto& s = stream(); const array& x1 = primals[0]; @@ -538,18 +554,23 @@ std::vector ArcTan2::jvp( const std::vector& tangents, const std::vector& argnums) { assert(primals.size() == 2); - assert(argnums.size() == 2); + assert(tangents.size() == argnums.size()); const auto& s = stream(); const array& x1 = primals[0]; const array& x2 = primals[1]; - const array& dx1 = tangents[0]; - const array& dx2 = tangents[1]; - return {divide( - subtract(multiply(x2, dx1, s), multiply(x1, dx2, s), s), - add(square(x1, s), square(x2, s), s), - s)}; + auto numerator = [&]() { + if (argnums.size() == 2) { + return subtract( + multiply(x2, tangents[0], s), multiply(x1, tangents[1], s), s); + } else if (argnums[0] == 0) { + return multiply(x2, tangents[0], s); + } else { + return negative(multiply(x1, tangents[0], s), s); + } + }(); + return {divide(numerator, add(square(x1, s), square(x2, s), s), s)}; } std::pair, std::vector> ArcTan2::vmap( @@ -595,7 +616,7 @@ std::pair, std::vector> ArgPartition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{argpartition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{argpartition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } std::vector ArgPartition::vjp( @@ -787,11 +808,8 @@ std::vector BitwiseBinary::jvp( const std::vector& tangents, const std::vector& argnums) { assert(primals.size() == 2); - std::vector vjps = {zeros_like(tangents[0], stream())}; - if (argnums.size() > 1) { - vjps.push_back(vjps.back()); - } - return vjps; + auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); + return {zeros(shape, tangents[0].dtype(), stream())}; } std::vector BitwiseBinary::vjp( @@ -799,7 +817,11 @@ std::vector BitwiseBinary::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + std::vector vjps; + for (auto arg : argnums) { + vjps.push_back(zeros_like(primals[arg], stream())); + } + return vjps; } std::vector @@ -897,7 +919,40 @@ std::vector BroadcastAxes::jvp( std::pair, std::vector> BroadcastAxes::vmap( const std::vector& inputs, const std::vector& axes) { - throw std::invalid_argument("[BroadcastAxes] VMAP NYI"); + std::vector new_inputs = inputs; + std::vector new_axes = axes; + size_t ndim = 0; + bool have_batch = false; + for (int i = 0; i < inputs.size(); i++) { + have_batch |= axes[i] >= 0; + ndim = std::max(inputs[i].ndim(), ndim); + } + + std::vector expand; + expand.reserve(ndim); + for (int i = 0; i < inputs.size(); i++) { + int extra = ndim - inputs[i].ndim(); + if (axes[i] >= 0 && extra > 0) { + new_axes[i] += extra; + expand.resize(extra); + std::iota(expand.begin(), expand.end(), 0); + new_inputs[i] = expand_dims(new_inputs[i], expand, stream()); + } + + if (new_axes[i] > 0) { + new_inputs[i] = moveaxis(new_inputs[i], new_axes[i], 0, stream()); + } + } + + auto shape = output_shape(new_inputs, ignore_axes_); + auto dtype = new_inputs[0].dtype(); + return { + {array( + shape, + dtype, + std::make_shared(stream(), ignore_axes_), + std::move(new_inputs))}, + {have_batch ? 0 : -1}}; } bool BroadcastAxes::is_equivalent(const Primitive& other) const { @@ -1132,12 +1187,29 @@ bool Concatenate::is_equivalent(const Primitive& other) const { std::vector Concatenate::output_shapes( const std::vector& inputs) { auto shape = inputs[0].shape(); + int64_t concat_size = shape[axis_]; for (int i = 1; i < inputs.size(); ++i) { - shape[axis_] += inputs[i].shape(axis_); + concat_size += inputs[i].shape(axis_); } + shape[axis_] = safe_cast(concat_size, "concatenate"); return {std::move(shape)}; } +std::vector Conjugate::vjp( + const std::vector&, + const std::vector& cotangents, + const std::vector&, + const std::vector&) { + return {conjugate(cotangents[0], stream())}; +} + +std::vector Conjugate::jvp( + const std::vector&, + const std::vector& tangents, + const std::vector&) { + return {conjugate(tangents[0], stream())}; +} + std::pair, std::vector> Conjugate::vmap( const std::vector& inputs, const std::vector& axes) { @@ -1187,9 +1259,11 @@ array conv_weight_backward_patches( // padded shape for (int i = 1; i < in.ndim() - 1; i++) { - in_padded_shape[i] += padding_lo[i - 1] + padding_hi[i - 1]; - padding_ends[i] += padding_lo[i - 1]; - padding_starts[i] += padding_lo[i - 1]; + int64_t lo = padding_lo[i - 1]; + int64_t hi = padding_hi[i - 1]; + in_padded_shape[i] = safe_cast(in_padded_shape[i] + lo + hi, "conv"); + padding_ends[i] = safe_cast(padding_ends[i] + lo, "conv"); + padding_starts[i] = safe_cast(padding_starts[i] + lo, "conv"); } // padded strides (contiguous) @@ -1245,12 +1319,18 @@ array conv_weight_backward_patches( namespace { // Conv helpers -inline int conv_out_axis_size(int in_dim, int wt_dim, int stride, int padding) { +// Computed in 64 bits so extreme but in-range int32 parameters do not overflow. +inline int64_t conv_out_axis_size( + int64_t in_dim, + int64_t wt_dim, + int64_t stride, + int64_t padding) { return ((in_dim + padding - wt_dim) / stride) + 1; } // Conv helpers -inline int dilate_size(int dim, int dil) { +// Computed in 64 bits so extreme but in-range int32 parameters do not overflow. +inline int64_t dilate_size(int64_t dim, int64_t dil) { return 1 + dil * (dim - 1); } @@ -1329,13 +1409,16 @@ Shape Convolution::conv_out_shape( throw std::invalid_argument(msg.str()); } - int kd = dilate_size(wt_shape[i], kernel_dilation[i - 1]); - int id = dilate_size(in_shape[i], input_dilation[i - 1]); + int64_t kd = dilate_size(wt_shape[i], kernel_dilation[i - 1]); + int64_t id = dilate_size(in_shape[i], input_dilation[i - 1]); - out_shape[i] = conv_out_axis_size( - id, kd, strides[i - 1], pads_lo[i - 1] + pads_hi[i - 1]); + int64_t out_size = conv_out_axis_size( + id, + kd, + strides[i - 1], + static_cast(pads_lo[i - 1]) + pads_hi[i - 1]); - if (out_shape[i] <= 0) { + if (out_size <= 0) { std::ostringstream msg; msg << "[conv] Spatial dimensions of input after padding" << " cannot be smaller than weight spatial dimensions." @@ -1344,6 +1427,8 @@ Shape Convolution::conv_out_shape( << ", and weight of shape " << wt_shape << "."; throw std::invalid_argument(msg.str()); } + + out_shape[i] = safe_cast(out_size, "conv"); } out_shape[i] = O; @@ -1387,12 +1472,12 @@ std::vector Convolution::vjp( std::vector padding_hi = padding_hi_; for (int i = 0; i < padding_lo.size(); ++i) { - int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1); - padding_lo[i] = wt_size - padding_lo_[i] - 1; + int64_t wt_size = dilate_size(wt.shape(1 + i), kernel_dilation_[i]); + padding_lo[i] = safe_cast(wt_size - padding_lo_[i] - 1, "conv"); - int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1); - int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1); - padding_hi[i] = in_size - out_size + padding_hi_[i]; + int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); + int64_t out_size = dilate_size(cotan.shape(1 + i), kernel_strides_[i]); + padding_hi[i] = safe_cast(in_size - out_size + padding_hi_[i], "conv"); } // Check for negative padding @@ -1424,7 +1509,8 @@ std::vector Convolution::vjp( for (int i = 0; i < grad.ndim() - 2; i++) { if (padding_lo[i] < 0) { - starts[i + 1] -= padding_lo[i]; + starts[i + 1] = safe_cast( + starts[i + 1] - static_cast(padding_lo[i]), "conv"); } if (padding_hi[i] < 0) { stops[i + 1] += padding_hi[i]; @@ -1452,10 +1538,12 @@ std::vector Convolution::vjp( auto padding_hi = padding_lo_; for (int i = 0; i < padding_hi.size(); ++i) { - int in_size = 1 + input_dilation_[i] * (in.shape(1 + i) - 1); - int out_size = 1 + kernel_strides_[i] * (cotan.shape(1 + i) - 1); - int wt_size = 1 + kernel_dilation_[i] * (wt.shape(1 + i) - 1); - padding_hi[i] = out_size - in_size + wt_size - padding_hi[i] - 1; + int64_t in_size = dilate_size(in.shape(1 + i), input_dilation_[i]); + int64_t out_size = + dilate_size(cotan.shape(1 + i), kernel_strides_[i]); + int64_t wt_size = dilate_size(wt.shape(1 + i), kernel_dilation_[i]); + padding_hi[i] = safe_cast( + out_size - in_size + wt_size - padding_hi[i] - 1, "conv"); } auto cotan_trans = swapaxes(cotan, 0, -1, stream()); @@ -1631,7 +1719,10 @@ std::vector Cosh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Cosh::jvp( @@ -1751,7 +1842,9 @@ std::vector DivMod::jvp( const std::vector& primals, const std::vector& tangents, const std::vector& argnums) { - return {zeros_like(primals[0], stream())}; + // One tangent per output; inputs are pre-broadcast and pre-promoted + auto zero = zeros_like(primals[0], stream()); + return {zero, zero}; } std::pair, std::vector> DivMod::vmap( @@ -1789,7 +1882,9 @@ std::pair, std::vector> Divide::vmap( const std::vector& inputs, const std::vector& axes) { auto [a, b, to_ax] = vmap_binary_op(inputs, axes, stream()); - return {{divide(a, b, stream())}, {to_ax}}; + auto out = issubdtype(a.dtype(), integer) ? floor_divide(a, b, stream()) + : divide(a, b, stream()); + return {{out}, {to_ax}}; } std::vector Remainder::vjp( @@ -1846,6 +1941,11 @@ std::pair, std::vector> Equal::vmap( return {{equal(a, b, stream())}, {to_ax}}; } +bool Equal::is_equivalent(const Primitive& other) const { + const Equal& e_other = static_cast(other); + return equal_nan_ == e_other.equal_nan_; +} + std::vector Equal::vjp( const std::vector& primals, const std::vector& cotangents, @@ -1863,7 +1963,7 @@ std::vector Equal::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::vector Erf::vjp( @@ -1935,7 +2035,7 @@ std::vector Exp::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector& outputs) { - return {multiply(cotangents[0], outputs[0], stream())}; + return {multiply(cotangents[0], conjugate(outputs[0], stream()), stream())}; } std::vector Exp::jvp( @@ -2079,12 +2179,12 @@ bool Flatten::is_equivalent(const Primitive& other) const { Shape Flatten::output_shape(const array& input, int start_axis, int end_axis) { Shape shape = input.shape(); - auto flat_size = input.shape(start_axis); + int64_t flat_size = input.shape(start_axis); for (int ax = start_axis + 1; ax <= end_axis; ++ax) { flat_size *= input.shape(ax); } shape.erase(shape.begin() + start_axis + 1, shape.begin() + end_axis + 1); - shape[start_axis] = flat_size; + shape[start_axis] = safe_cast(flat_size, "flatten"); return shape; } @@ -2158,12 +2258,13 @@ std::pair, std::vector> FFT::vmap( if (fft_ax >= ax) { fft_ax++; } - if (real_) { - auto n = out_shape[fft_ax]; - out_shape[fft_ax] = inverse_ ? 2 * (n - 1) : n / 2 + 1; - } } } + // Only the last transformed axis changes size in a real transform + if (real_) { + auto n = out_shape[fft_axes.back()]; + out_shape[fft_axes.back()] = inverse_ ? 2 * (n - 1) : n / 2 + 1; + } return { {array( out_shape, @@ -2210,8 +2311,10 @@ std::vector FFT::vjp( two, one, stream()); - return { - multiply(fft::rfftn(cotangents[0], axes, stream()), mask, stream())}; + return {multiply( + fft::rfftn(cotangents[0], axes, fft::FFTNorm::Backward, stream()), + mask, + stream())}; } else if (real_) { Shape n; for (auto ax : axes_) { @@ -2237,17 +2340,22 @@ std::vector FFT::vjp( one, stream()); return {multiply( - fft::irfftn(multiply(cotangents[0], mask, stream()), n, axes, stream()), + fft::irfftn( + multiply(cotangents[0], mask, stream()), + n, + axes, + fft::FFTNorm::Backward, + stream()), array(n_elements, in.dtype()), stream())}; } else if (inverse_) { return {multiply( - fft::fftn(cotangents[0], axes, stream()), + fft::fftn(cotangents[0], axes, fft::FFTNorm::Backward, stream()), array(1 / n_elements, complex64), stream())}; } else { return {multiply( - fft::ifftn(cotangents[0], axes, stream()), + fft::ifftn(cotangents[0], axes, fft::FFTNorm::Backward, stream()), array(n_elements, complex64), stream())}; } @@ -2260,14 +2368,15 @@ std::vector FFT::jvp( assert(primals.size() == 1); assert(argnums.size() == 1); auto& tan = tangents[0]; + std::vector axes(axes_.begin(), axes_.end()); if (real_ & inverse_) { - return {fft::irfftn(tan, stream())}; + return {fft::irfftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else if (real_) { - return {fft::rfftn(tan, stream())}; + return {fft::rfftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else if (inverse_) { - return {fft::ifftn(tan, stream())}; + return {fft::ifftn(tan, axes, fft::FFTNorm::Backward, stream())}; } else { - return {fft::fftn(tan, stream())}; + return {fft::fftn(tan, axes, fft::FFTNorm::Backward, stream())}; } } @@ -2397,9 +2506,9 @@ std::vector Gather::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); std::vector inds(primals.begin() + 1, primals.end()); @@ -2415,7 +2524,8 @@ std::vector Gather::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather] Cannot calculate JVP with respect to indices."); + "[gather] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } std::vector inds(primals.begin() + 1, primals.end()); return {gather(tangents[0], inds, axes_, slice_sizes_, stream())}; @@ -2462,9 +2572,9 @@ std::vector GatherAxis::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather_axis] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); vjps.push_back(array( @@ -2483,7 +2593,8 @@ std::vector GatherAxis::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather_axis] Cannot calculate JVP with respect to indices."); + "[gather_axis] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } return {take_along_axis(tangents[0], primals[1], axis_, stream())}; } @@ -2530,7 +2641,7 @@ std::vector Greater::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::pair, std::vector> GreaterEqual::vmap( @@ -2557,7 +2668,7 @@ std::vector GreaterEqual::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::vector Imag::vjp( @@ -2614,7 +2725,7 @@ std::vector Less::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::pair, std::vector> LessEqual::vmap( @@ -2641,7 +2752,7 @@ std::vector LessEqual::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::vector Log::vjp( @@ -2649,7 +2760,15 @@ std::vector Log::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // return jvp(primals, cotangents, argnums); + assert(primals.size() == 1); + assert(argnums.size() == 1); + auto out = divide(cotangents[0], conjugate(primals[0], stream()), stream()); + if (base_ != Base::e) { + auto scale = 1 / std::log(base_ == Base::ten ? 10.0f : 2.0f); + out = multiply(array(scale, out.dtype()), out, stream()); + } + return {out}; } std::vector Log::jvp( @@ -2681,12 +2800,20 @@ std::pair, std::vector> Log::vmap( axes}; } +bool Log::is_equivalent(const Primitive& other) const { + const Log& l_other = static_cast(other); + return base_ == l_other.base_; +} + std::vector Log1p::vjp( const std::vector& primals, const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Log1p::jvp( @@ -2875,10 +3002,16 @@ std::vector LogSumExp::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(tangents.size() == 1); - return {multiply( - tangents[0], - softmax(primals[0], std::vector{-1}, true, stream()), - stream())}; + // d/dt logsumexp(x) = sum(softmax(x) * t) over the reduced axis. The result + // must be summed back to the output shape, not left as a per-element product. + return { + sum(multiply( + tangents[0], + softmax(primals[0], std::vector{-1}, true, stream()), + stream()), + -1, + /* keepdims = */ true, + stream())}; } std::vector LogSumExp::output_shapes(const std::vector& inputs) { @@ -3074,30 +3207,30 @@ std::vector Select::jvp( const std::vector& tangents, const std::vector& argnums) { assert(primals.size() == 3); - assert(tangents.size() == 3); + assert(tangents.size() == argnums.size()); auto jvp_fun = [&](int i) { int arg = argnums[i]; if (arg == 0) { - return zeros_like(primals[0], stream()); + return zeros_like(primals[1], stream()); } else if (arg == 1) { return multiply( - astype(primals[0], tangents[1].dtype(), stream()), - tangents[1], + astype(primals[0], tangents[i].dtype(), stream()), + tangents[i], stream()); } else { return multiply( astype( - logical_not(primals[0], stream()), tangents[2].dtype(), stream()), - tangents[2], + logical_not(primals[0], stream()), tangents[i].dtype(), stream()), + tangents[i], stream()); } }; - array jvp = jvp_fun(argnums[0]); + array jvp = jvp_fun(0); for (int i = 1; i < argnums.size(); i++) { - jvp = add(jvp, jvp_fun(argnums[i])); + jvp = add(jvp, jvp_fun(i)); } return {jvp}; } @@ -3188,7 +3321,7 @@ std::vector NotEqual::jvp( const std::vector& tangents, const std::vector& argnums) { auto shape = broadcast_shapes(primals[0].shape(), primals[1].shape()); - return {zeros(shape, bool_, stream())}; + return {zeros(shape, tangents[0].dtype(), stream())}; } std::vector Pad::vjp( @@ -3231,7 +3364,32 @@ std::vector Pad::jvp( std::pair, std::vector> Pad::vmap( const std::vector& inputs, const std::vector& axes) { - throw std::runtime_error("Pad vmap is NYI."); + assert(inputs.size() == 2); + assert(axes.size() == 2); + + if (axes[1] >= 0) { + throw std::invalid_argument( + "[Pad::vmap] Vmap over padding value is not supported."); + } + + auto ax = axes[0]; + auto pad_axes = axes_; + if (ax >= 0) { + for (auto& pad_ax : pad_axes) { + pad_ax = (pad_ax >= ax) ? pad_ax + 1 : pad_ax; + } + } + + return { + {pad( + inputs[0], + pad_axes, + low_pad_size_, + high_pad_size_, + inputs[1], + "constant", + stream())}, + {ax}}; } bool Pad::is_equivalent(const Primitive& other) const { @@ -3246,7 +3404,10 @@ std::vector Partition::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - auto sort_idx = argpartition(primals[0], kth_, axis_, stream()); + // The permutation is locally constant in the input, so cut the gradient + // there to keep higher order derivatives working. + auto sort_idx = + stop_gradient(argpartition(primals[0], kth_, axis_, stream()), stream()); return {put_along_axis( zeros_like(primals[0], stream()), sort_idx, @@ -3261,7 +3422,8 @@ std::vector Partition::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(tangents.size() == 1); - auto sort_idx = argpartition(primals[0], kth_, axis_, stream()); + auto sort_idx = + stop_gradient(argpartition(primals[0], kth_, axis_, stream()), stream()); auto out = take_along_axis(tangents[0], sort_idx, axis_, stream()); return {out}; } @@ -3273,7 +3435,7 @@ std::pair, std::vector> Partition::vmap( assert(axes.size() == 1); int axis_left = axes[0] >= 0 && axes[0] <= axis_; - return {{partition(inputs[0], axis_ + axis_left, stream())}, axes}; + return {{partition(inputs[0], kth_, axis_ + axis_left, stream())}, axes}; } bool Partition::is_equivalent(const Primitive& other) const { @@ -3312,12 +3474,12 @@ std::vector Power::jvp( const std::vector& tangents, const std::vector& argnums) { auto output = power(primals[0], primals[1], stream()); - auto grads = vjp(primals, tangents, argnums, {output}); - if (argnums.size() > 1) { - return {add(grads[0], grads[1], stream())}; - } else { - return grads; + auto jvp = vjp(primals, {tangents[0]}, {argnums[0]}, {output})[0]; + for (int i = 1; i < argnums.size(); ++i) { + jvp = add( + jvp, vjp(primals, {tangents[i]}, {argnums[i]}, {output})[0], stream()); } + return {jvp}; } std::pair, std::vector> Power::vmap( @@ -3424,6 +3586,7 @@ std::vector QuantizedMatmul::vjp( group_size_, bits_, quantization_mode_to_string(mode_), + {}, // placeholder for amax std::nullopt, stream()); wq = unflatten(wq, -1, {-1, group_size_}, stream()); @@ -3484,14 +3647,14 @@ std::vector QQMatmul::output_shapes(const std::vector& inputs) { } std::vector QQMatmul::vjp( - const std::vector& primals, // non quantized x, non quantized w + const std::vector& primals, // non quantized x, non quantized w, if + // nvfp4 global_scale_x, global_scale_w const std::vector& cotangents, // non quantized upstream grads const std::vector& argnums, const std::vector&) { - if (primals.size() != 2) { - throw std::runtime_error( - "[QQMatmul::vjp] Expected exactly 2 non-quantized primal inputs (x, w)."); - } + bool is_nvfp4 = mode_ == QuantizationMode::Nvfp4; + assert(primals.size() == 2 || (is_nvfp4 && primals.size() == 4)); + std::vector vjps; auto& cotan = cotangents[0]; auto& s = stream(); @@ -3499,6 +3662,15 @@ std::vector QQMatmul::vjp( // primal[0] -- non quantized activations (M, K) // cotan -- non quantized grads (M, N) auto qmode = quantization_mode_to_string(mode_); + std::optional cotan_amax = (primals.size() == 4) + ? std::make_optional(astype(max(abs(cotan, s), s), float32, s)) + : std::nullopt; + + auto get_primal_scale = [&](int idx) { + return (primals.size() == 4) ? std::make_optional(primals[idx]) + : std::nullopt; + }; + for (auto arg : argnums) { if (arg == 0) { // gradient wrt to x // We transpose weights -> quantize along N @@ -3509,6 +3681,8 @@ std::vector QQMatmul::vjp( group_size_, bits_, qmode, + cotan_amax, + get_primal_scale(3), // global_scale_w (for w.T) s)); } else if (arg == 1) { // gradient wrt to weights vjps.push_back(qqmm( @@ -3518,7 +3692,11 @@ std::vector QQMatmul::vjp( group_size_, bits_, qmode, + cotan_amax, + get_primal_scale(2), // global_scale_x (for x.T) s)); + } else { + vjps.push_back(zeros_like(primals[arg], s)); } } return vjps; @@ -3643,6 +3821,7 @@ std::vector GatherQMM::vjp( bits_, quantization_mode_to_string(mode_), std::nullopt, + std::nullopt, // amax placeholder stream()), -1, {-1, group_size_}, @@ -3670,6 +3849,34 @@ bool GatherQMM::is_equivalent(const Primitive& other) const { mode_ == qm_other.mode_ && transpose_ == qm_other.transpose_; } +std::vector GatherQMM::output_shapes(const std::vector& inputs) { + const auto& x = inputs[0]; + const auto& w = inputs[1]; + const auto& lhs_indices = + (mode_ == QuantizationMode::Affine) ? inputs[4] : inputs[3]; + int w_outer = transpose_ ? w.shape(-2) : w.shape(-1) * 32 / bits_; + auto out_shape = lhs_indices.shape(); + out_shape.push_back(x.shape(-2)); + out_shape.push_back(w_outer); + return {out_shape}; +} + +bool GatherQQMM::is_equivalent(const Primitive& other) const { + const GatherQQMM& qm_other = static_cast(other); + return group_size_ == qm_other.group_size_ && bits_ == qm_other.bits_ && + mode_ == qm_other.mode_; +} + +std::vector GatherQQMM::output_shapes(const std::vector& inputs) { + const auto& x = inputs[0]; + const auto& w = inputs[1]; + const auto& lhs_indices = inputs[2]; + auto out_shape = lhs_indices.shape(); + out_shape.push_back(x.shape(-2)); + out_shape.push_back(w.shape(-2)); + return {out_shape}; +} + std::pair, std::vector> RandomBits::vmap( const std::vector& inputs, const std::vector& axes) { @@ -3807,7 +4014,7 @@ Shape Reshape::output_shape(const array& input, Shape shape) { // Infer the shape if (size > 0 && infer_idx >= 0) { - shape[infer_idx] = input.size() / size; + shape[infer_idx] = safe_cast(input.size() / size, "reshape"); size *= shape[infer_idx]; } else if (infer_idx >= 0) { throw std::invalid_argument( @@ -3845,7 +4052,7 @@ std::vector Reduce::vjp( auto p1 = cumprod(x, axis, /*reverse=*/false, /*inclusive=*/false, s); auto p2 = cumprod(x, axis, /*reverse=*/true, /*inclusive=*/false, s); auto exclusive_prod = multiply(p1, p2, s); - return multiply(exclusive_prod, cotan, s); + return multiply(conjugate(exclusive_prod, s), cotan, s); }; // To compute a numerically stable gradient for prod we need an exclusive @@ -4156,8 +4363,66 @@ std::vector Scan::vjp( where(eq_zero, z, grad, stream()), stream())}; } else { - // Can probably be implemented by equals and then cummax to make the mask - throw std::runtime_error("VJP is not implemented for cumulative min/max"); + // Cumulative max/min: route each position's cotangent to the input element + // that holds the running extreme there, breaking ties toward the latest + // element in scan order. The owning index is reconstructed by marking + // where the input equals the inclusive running extreme and taking a running + // extreme over those indices (cummax picks the largest tied index for a + // forward scan and cummin the smallest for a reverse scan). The cotangents + // are then scatter-added. For exclusive scans, outputs[0] is already the + // inclusive running extreme shifted by one, so we slice the input, + // outputs[0], and the cotangents to length n-1; scatter_add_axis broadcasts + // the shorter updates into the full-sized gradient. + auto in = primals[0]; + auto cotan = cotangents[0]; + int n = in.shape(axis_); + auto s = stream(); + + auto in_window = in; + auto extreme_window = outputs[0]; + Shape start; + Shape stop; + + if (!inclusive_) { + start = Shape(in.ndim(), 0); + stop = in.shape(); + + start[axis_] = (reverse_) ? 1 : 0; + stop[axis_] = (reverse_) ? n : n - 1; + in_window = slice(in, start, stop, s); + + start[axis_] = (reverse_) ? 0 : 1; + stop[axis_] = (reverse_) ? n - 1 : n; + extreme_window = slice(outputs[0], start, stop, s); + } + + Shape index_shape(in.ndim(), 1); + index_shape[axis_] = (inclusive_) ? n : n - 1; + auto iota = reshape( + (inclusive_) ? arange(0, n, s) + : ((reverse_) ? arange(1, n, s) : arange(0, n - 1, s)), + index_shape, + s); + auto masked = where( + equal(in_window, extreme_window, s), + iota, + array(reverse_ ? n : -1, int32), + s); + // The owner indices are locally constant in the input, so cut the + // gradient there to keep higher order derivatives working. + auto owner = stop_gradient( + astype( + reverse_ ? cummin(masked, axis_, /* reverse = */ true, true, s) + : cummax(masked, axis_, /* reverse = */ false, true, s), + uint32, + s), + s); + + if (!inclusive_) { + cotan = slice(cotan, std::move(start), std::move(stop), s); + } + + return {scatter_add_axis(zeros_like(in, s), owner, cotan, axis_, s)}; } } @@ -4274,7 +4539,8 @@ std::vector Scatter::vjp( } } else { throw std::invalid_argument( - "[scatter] Cannot calculate VJP with respect to indices."); + "[scatter] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4386,7 +4652,8 @@ std::vector ScatterAxis::vjp( vjps.push_back(take_along_axis(cotangents[0], indices, axis_, stream())); } else { throw std::invalid_argument( - "[scatter_axis] Cannot calculate VJP with respect to indices."); + "[scatter_axis] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4399,7 +4666,8 @@ std::vector ScatterAxis::jvp( for (auto arg : argnums) { if (arg == 1) { throw std::invalid_argument( - "[scatter_axis] Cannot calculate JVP with respect to indices."); + "[scatter_axis] Cannot calculate JVP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } if (argnums.size() == 2) { @@ -4507,7 +4775,8 @@ std::vector MaskedScatter::vjp( vjps.push_back(reshape(gsrc_flat, src.shape(), s)); } else { throw std::invalid_argument( - "[masked_scatter] Cannot calculate VJP with respect to mask."); + "[masked_scatter] Cannot calculate VJP with respect to mask. " + "Use stop_gradient on mask to stop gradients from being computed."); } } return vjps; @@ -4528,15 +4797,16 @@ std::vector MaskedScatter::jvp( } array out = zeros_like(dst, s); - for (int arg : argnums) { + for (int i = 0; i < argnums.size(); ++i) { + int arg = argnums[i]; if (arg == 0) { - out = where(mask_b, out, tangents[0], s); + out = where(mask_b, out, tangents[i], s); } else if (arg == 2) { out = array( out.shape(), out.dtype(), std::make_shared(s), - {out, mask, tangents[1]}); + {out, mask, tangents[i]}); } else { throw std::invalid_argument("[masked_scatter] invalid arg index in JVP"); } @@ -4669,7 +4939,10 @@ std::vector Sin::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Sin::jvp( @@ -4694,7 +4967,10 @@ std::vector Sinh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Sinh::jvp( @@ -4776,10 +5052,17 @@ std::pair, std::vector> SliceUpdate::vmap( // No vmapping needed if (src_ax == -1 && upd_ax == -1) { - return {{slice_update(src, upd, start, stop, strides, stream())}, {-1}}; + return { + {array( + src.shape(), + src.dtype(), + std::make_shared( + stream(), reduce_type_, start, stop, strides), + {src, upd})}, + {-1}}; } - // Broadcast src + // Broadcast Src if (src_ax == -1) { src = expand_dims(src, upd_ax, stream()); auto shape = src.shape(); @@ -4802,37 +5085,99 @@ std::pair, std::vector> SliceUpdate::vmap( stop.insert(stop.begin() + src_ax, src.shape(src_ax)); strides.insert(strides.begin() + src_ax, 1); - return {{slice_update(src, upd, start, stop, strides, stream())}, {src_ax}}; + return { + {array( + src.shape(), + src.dtype(), + std::make_shared( + stream(), reduce_type_, start, stop, strides), + {src, upd})}, + {src_ax}}; } std::vector SliceUpdate::vjp( const std::vector& primals, const std::vector& cotangents, const std::vector& argnums, - const std::vector&) { + const std::vector& outputs) { // Check inputs assert(primals.size() == 2); - auto& cotan = cotangents[0]; - auto& upd = primals[1]; + const array& result = outputs[0]; + const array& values = primals[0]; + const array& updates = primals.back(); + const array& cotan = cotangents[0]; std::vector vjps; for (int num : argnums) { // Vjp for source if (num == 0) { - vjps.push_back(slice_update( - cotan, - zeros_like(upd, stream()), - start_indices_, - end_indices_, - strides_, - stream())); + switch (reduce_type_) { + case SliceUpdate::None: + vjps.push_back(array( + cotan.shape(), + cotan.dtype(), + std::make_shared( + stream(), + reduce_type_, + start_indices_, + end_indices_, + strides_), + {cotan, zeros_like(updates, stream())})); + break; + case SliceUpdate::Sum: + vjps.push_back(cotan); + break; + case SliceUpdate::Max: + case SliceUpdate::Min: + vjps.push_back(where( + equal(result, values, stream()), + cotan, + array(0, cotan.dtype()), + stream())); + break; + case SliceUpdate::Prod: + vjps.push_back(array( + cotan.shape(), + cotan.dtype(), + std::make_shared( + stream(), + reduce_type_, + start_indices_, + end_indices_, + strides_), + {cotan, updates})); + break; + } } // Vjp fpr updates else { - vjps.push_back( - slice(cotan, start_indices_, end_indices_, strides_, stream())); + auto sliced_cotan = + slice(cotan, start_indices_, end_indices_, strides_, stream()); + switch (reduce_type_) { + case SliceUpdate::None: + case SliceUpdate::Sum: + vjps.emplace_back(std::move(sliced_cotan)); + break; + case SliceUpdate::Max: + case SliceUpdate::Min: { + auto sliced_result = + slice(result, start_indices_, end_indices_, strides_, stream()); + vjps.push_back(where( + equal(sliced_result, updates, stream()), + sliced_cotan, + array(0, cotan.dtype()), + stream())); + break; + } + case SliceUpdate::Prod: { + auto sliced_values = + slice(values, start_indices_, end_indices_, strides_, stream()); + vjps.push_back(multiply(sliced_cotan, sliced_values, stream())); + break; + } + } } } @@ -4845,18 +5190,48 @@ std::vector SliceUpdate::jvp( const std::vector& argnums) { // Check inputs assert(primals.size() == 2); - return {slice_update( - tangents[0], - tangents[1], - start_indices_, - end_indices_, - strides_, - stream())}; + + array result_tan = zeros_like(primals[0], stream()); + array update_tan = zeros_like(primals[1], stream()); + for (int i = 0; i < argnums.size(); ++i) { + if (argnums[i] == 0) { + result_tan = tangents[i]; + } else if (argnums[i] == 1) { + update_tan = tangents[i]; + } + } + + switch (reduce_type_) { + case SliceUpdate::None: + return {array( + result_tan.shape(), + result_tan.dtype(), + std::make_shared( + stream(), reduce_type_, start_indices_, end_indices_, strides_), + {result_tan, update_tan})}; + case SliceUpdate::Sum: + return {array( + result_tan.shape(), + result_tan.dtype(), + std::make_shared( + stream(), reduce_type_, start_indices_, end_indices_, strides_), + {result_tan, update_tan})}; + case SliceUpdate::Prod: + case SliceUpdate::Max: + case SliceUpdate::Min: { + throw std::runtime_error( + "[SliceUpdate] JVP for product, minimum and maximum not implemented."); + } + } + + // Appease gcc (although no path reaches here). + return {}; } bool SliceUpdate::is_equivalent(const Primitive& other) const { const auto& s_other = static_cast(other); return ( + reduce_type_ == s_other.reduce_type_ && start_indices_ == s_other.start_indices_ && end_indices_ == s_other.end_indices_ && strides_ == s_other.strides_); } @@ -4993,8 +5368,20 @@ std::vector DynamicSliceUpdate::vjp( std::vector DynamicSliceUpdate::jvp( const std::vector& primals, const std::vector& tangents, - const std::vector&) { - return {slice_update(tangents[0], tangents[1], primals[2], axes_, stream())}; + const std::vector& argnums) { + array tan_src = zeros_like(primals[0], stream()); + array tan_upd = zeros_like(primals[1], stream()); + for (int i = 0; i < argnums.size(); ++i) { + if (argnums[i] == 0) { + tan_src = tangents[i]; + } else if (argnums[i] == 1) { + tan_upd = tangents[i]; + } else { + throw std::invalid_argument( + "[DynamicSliceUpdate::jvp] Not supported for start indices"); + } + } + return {slice_update(tan_src, tan_upd, primals[2], axes_, stream())}; } bool DynamicSliceUpdate::is_equivalent(const Primitive& other) const { @@ -5052,6 +5439,47 @@ bool Softmax::is_equivalent(const Primitive& other) const { return precise_ == s_other.precise_; } +std::pair, std::vector> SearchSorted::vmap( + const std::vector& inputs, + const std::vector& axes) { + if (axes[0] != -1) { + throw std::invalid_argument( + "[searchsorted] Cannot vmap over the sorted sequence, only over the " + "values being searched for."); + } + auto side = right_ ? "right" : "left"; + return {{searchsorted(inputs[0], inputs[1], side, stream())}, {axes[1]}}; +} + +std::vector SearchSorted::vjp( + const std::vector& primals, + const std::vector&, + const std::vector& argnums, + const std::vector&) { + std::vector vjps; + for (auto arg : argnums) { + vjps.push_back(zeros_like(primals[arg], stream())); + } + return vjps; +} + +std::vector SearchSorted::jvp( + const std::vector& primals, + const std::vector&, + const std::vector&) { + return {zeros(primals[1].shape(), uint32, stream())}; +} + +bool SearchSorted::is_equivalent(const Primitive& other) const { + const SearchSorted& r_other = static_cast(other); + return right_ == r_other.right_; +} + +std::vector SearchSorted::output_shapes( + const std::vector& inputs) { + return {inputs[1].shape()}; +} + std::pair, std::vector> Sort::vmap( const std::vector& inputs, const std::vector& axes) { @@ -5067,7 +5495,18 @@ std::vector Sort::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // Sort applies a permutation to the input, so the cotangents must be + // scattered back to the original positions (the transpose of the + // permutation), not gathered forward as in the jvp. + // The permutation is locally constant in the input, so cut the gradient + // there to keep higher order derivatives working. + auto sort_idx = stop_gradient(argsort(primals[0], axis_, stream()), stream()); + return {put_along_axis( + zeros_like(primals[0], stream()), + sort_idx, + cotangents[0], + axis_, + stream())}; } std::vector Sort::jvp( @@ -5076,7 +5515,7 @@ std::vector Sort::jvp( const std::vector& argnums) { assert(primals.size() == 1); assert(tangents.size() == 1); - auto sort_idx = argsort(primals[0], axis_, stream()); + auto sort_idx = stop_gradient(argsort(primals[0], axis_, stream()), stream()); auto out = take_along_axis(tangents[0], sort_idx, axis_, stream()); return {out}; } @@ -5120,7 +5559,10 @@ std::vector Square::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Square::jvp( @@ -5292,7 +5734,10 @@ std::vector Tan::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Tan::jvp( @@ -5318,7 +5763,10 @@ std::vector Tanh::vjp( const std::vector& cotangents, const std::vector& argnums, const std::vector&) { - return jvp(primals, cotangents, argnums); + // The vjp conjugates the jvp's multiplier (a no-op for real inputs). + return {conjugate( + jvp(primals, {conjugate(cotangents[0], stream())}, argnums)[0], + stream())}; } std::vector Tanh::jvp( @@ -5406,7 +5854,8 @@ std::vector BlockMaskedMM::vjp( if ((needs_lhs_mask_vjp && primals[op_mask_idx].dtype() == bool_) || (needs_rhs_mask_vjp && primals[op_mask_idx + 1].dtype() == bool_)) { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } auto expand_mask = [&](array mask, int Y, int X) { @@ -5603,7 +6052,8 @@ std::vector BlockMaskedMM::vjp( } else { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } } return vjps; @@ -5668,7 +6118,8 @@ std::vector GatherMM::vjp( stream())); } else { throw std::invalid_argument( - "[GatherMM] Cannot calculate VJP with respect to indices."); + "[GatherMM] Cannot calculate VJP with respect to indices. " + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -5680,6 +6131,16 @@ bool GatherMM::is_equivalent(const Primitive& other) const { right_sorted_ == g_other.right_sorted_; } +std::vector GatherMM::output_shapes(const std::vector& inputs) { + const auto& a = inputs[0]; + const auto& b = inputs[1]; + const auto& lhs_indices = inputs[2]; + auto out_shape = lhs_indices.shape(); + out_shape.push_back(a.shape(-2)); + out_shape.push_back(b.shape(-1)); + return {out_shape}; +} + bool BlockMaskedMM::is_equivalent(const Primitive& other) const { const BlockMaskedMM& a_other = static_cast(other); return (block_size_ == a_other.block_size_); diff --git a/Source/Cmlx/mlx/mlx/primitives.h b/Source/Cmlx/mlx/mlx/primitives.h index 4091aafcf..0cfc71bf0 100644 --- a/Source/Cmlx/mlx/mlx/primitives.h +++ b/Source/Cmlx/mlx/mlx/primitives.h @@ -200,6 +200,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_}; }; @@ -542,6 +543,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_); } @@ -697,6 +699,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() @@ -972,9 +975,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"; @@ -1322,9 +1325,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_; }; @@ -1698,6 +1701,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_); @@ -1712,6 +1716,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) @@ -2055,12 +2094,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) {} @@ -2070,14 +2113,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_; @@ -2150,6 +2211,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/mlx/mlx/random.cpp b/Source/Cmlx/mlx/mlx/random.cpp index def3169cb..2932b3eaf 100644 --- a/Source/Cmlx/mlx/mlx/random.cpp +++ b/Source/Cmlx/mlx/mlx/random.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include #include #include @@ -38,6 +39,10 @@ array bits( int width /* 4 */, const std::optional& key_ /*= nullopt */, StreamOrDevice s /* = {} */) { + if (std::any_of(shape.begin(), shape.end(), [](auto i) { return i < 0; })) { + throw std::invalid_argument("[bits] Negative dimensions not allowed."); + } + auto key = key_ ? *key_ : KeySequence::default_().next(); if (key.dtype() != uint32) { std::ostringstream msg; @@ -182,8 +187,7 @@ array normal( return complex_normal(shape, loc, scale, key, s); } else if (!issubdtype(dtype, floating)) { throw std::invalid_argument( - "[normal] Can only generate uniform numbers with " - "floating point type."); + "[normal] Only floating point and complex64 types are supported."); } auto stream = to_stream(s); @@ -280,7 +284,12 @@ array randint( "[randint] randint only accepts integer dtypes and bool."); } auto u = uniform(low, high, shape, float32, key, s); - return astype(maximum(u, low, s), dtype, s); + // astype truncates decimal parts so -1.7 becomes -1, use floor. + auto out = astype(floor(u, s), dtype, s); + // low/high may not be representable in float32 so actual range of uniform + // may be larger than [low, high) and we have to clamp to [low, high - 1]. + auto hi = astype(subtract(high, array(1, high.dtype()), s), dtype, s); + return maximum(minimum(out, hi, s), astype(low, dtype, s), s); } array bernoulli( @@ -379,12 +388,42 @@ int get_valid_axis(int axis, int ndim) { return ax; } +// O(N + M) in memory, where the Gumbel-max trick needs O(N * M) +array categorical_inverse_cdf( + const array& logits, + const Shape& shape, + const std::optional& key, + StreamOrDevice s) { + auto dtype = promote_types(logits.dtype(), float32); + auto x = astype(logits, dtype, s); + auto m = max(x, s); + + // exp(x - m) is NaN when m is infinite, where the distribution is uniform + // over the maximal entries + auto w = where( + isinf(m, s), + astype(equal(x, m, s), dtype, s), + exp(subtract(x, m, s), s), + s); + + auto cdf = cumsum(w, 0, /* reverse = */ false, /* inclusive = */ false, s); + auto u = multiply(uniform(shape, float32, key, s), sum(w, s), s); + + return subtract(searchsorted(cdf, u, "right", s), array(1u, uint32), s); +} + array categorical_impl( const array& logits, int axis, const Shape& shape, const std::optional& key /*= nullopt */, StreamOrDevice s) { + // searchsorted only takes 1D sequence. + auto n = logits.shape(axis); + if (n > 0 && logits.size() == static_cast(n) && !shape.empty()) { + return categorical_inverse_cdf(reshape(logits, {n}, s), shape, key, s); + } + auto gumbel_shape = shape; auto offset = axis + shape.size() - logits.ndim() + 1; gumbel_shape.insert(gumbel_shape.begin() + offset, logits.shape(axis)); @@ -449,8 +488,7 @@ array laplace( StreamOrDevice s /* = {} */) { if (!issubdtype(dtype, floating)) { throw std::invalid_argument( - "[laplace] Can only generate uniform numbers with real" - "floating point type."); + "[laplace] Only real floating point types are supported."); } auto stream = to_stream(s); diff --git a/Source/Cmlx/mlx/mlx/random.h b/Source/Cmlx/mlx/mlx/random.h index a23c25572..0c11de637 100644 --- a/Source/Cmlx/mlx/mlx/random.h +++ b/Source/Cmlx/mlx/mlx/random.h @@ -19,20 +19,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/mlx/mlx/scheduler.cpp b/Source/Cmlx/mlx/mlx/scheduler.cpp index 87ff98e22..572e236bb 100644 --- a/Source/Cmlx/mlx/mlx/scheduler.cpp +++ b/Source/Cmlx/mlx/mlx/scheduler.cpp @@ -1,72 +1,174 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2026 Apple Inc. -#include "mlx/scheduler.h" -#include "mlx/backend/gpu/device_info.h" +#include +#include + +#include "mlx/backend/cpu/eval.h" #include "mlx/backend/gpu/eval.h" +#include "mlx/compile_impl.h" +#include "mlx/scheduler.h" +#include "mlx/utils.h" namespace mlx::core { -Stream default_stream(Device d) { - if (!gpu::is_available() && d == Device::gpu) { - throw std::invalid_argument( - "[default_stream] Cannot get gpu stream without gpu backend."); - } - return scheduler::scheduler().get_default_stream(d); -} - -void set_default_stream(Stream s) { - if (!gpu::is_available() && s.device == Device::gpu) { - throw std::invalid_argument( - "[set_default_stream] Cannot set gpu stream without gpu backend."); - } - return scheduler::scheduler().set_default_stream(s); -} - -Stream get_stream(int index) { - return scheduler::scheduler().get_stream(index); -} - -Stream new_stream(Device d) { - if (!gpu::is_available() && d == Device::gpu) { - throw std::invalid_argument( - "[new_stream] Cannot make gpu stream without gpu backend."); - } - return scheduler::scheduler().new_stream(d); -} - -Stream new_stream() { - return scheduler::scheduler().new_stream(default_device()); -} - void synchronize(Stream s) { if (s.device == mlx::core::Device::cpu) { auto p = std::make_shared>(); std::future f = p->get_future(); scheduler::enqueue(s, [p = std::move(p)]() { p->set_value(); }); f.wait(); + scheduler::check_error(s); } else { gpu::synchronize(s); } } +void synchronize(ThreadLocalStream s) { + synchronize(stream_from_thread_local_stream(s)); +} + void synchronize() { synchronize(default_stream(default_device())); } +void clear_streams() { + detail::compile_clear_cache(detail::compile_cache()); + cpu::clear_streams(); + gpu::clear_streams(); +} + namespace scheduler { -/** A singleton scheduler to manage devices, streams, and task execution. */ +struct StreamThread { + std::mutex mtx; + std::queue> q; + std::condition_variable cond; + bool stop; + std::thread thread; + Error error; + + 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(); + } + } + + void enqueue(std::function f) { + if (is_main_thread()) { + error.check(); + } + { + std::lock_guard lk(mtx); + if (stop) { + throw std::runtime_error( + "Cannot enqueue work after stream is stopped."); + } + q.emplace(std::move(f)); + } + cond.notify_one(); + } +}; + +Scheduler::Scheduler() { + is_main_thread(); + gpu::init(); +} + +Scheduler::~Scheduler() = default; + +void Scheduler::enqueue(Stream s, std::function task) { + auto& st = get_thread(s); + st.enqueue([&st, task = std::move(task)]() mutable { + try { + task(); + } catch (const std::exception& error) { + // Set error to stream only when no error happended before, to preserve + // the earliest error. + if (!st.error.valid()) { + st.error.set_message(std::make_shared(error.what())); + } + } + }); +} + +void Scheduler::wait_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + auto& st = get_thread(s); + st.enqueue([&st, event = std::move(event), task = std::move(task)]() mutable { + task(event); + // Poison current stream if the waited event has error. + st.error.store_if_valid(event.load_error()); + }); +} + +void Scheduler::signal_event( + Stream s, + Event event, + std::function task) { + assert(s.device == Device::cpu); + auto& st = get_thread(s); + st.enqueue([&st, event = std::move(event), task = std::move(task)]() mutable { + // Poison the signal event if current stream has error. + if (st.error.valid()) { + event.set_error(st.error); + } + task(event); + }); +} + +void Scheduler::check_error(Stream s) { + get_thread(s).error.check(); +} + +StreamThread& Scheduler::get_thread(Stream s) { + { + std::shared_lock lock(threads_mtx_); + auto it = threads_.find(s.index); + if (it != threads_.end()) { + return *it->second.get(); + } + } + std::unique_lock lock(threads_mtx_); + auto it = threads_.find(s.index); + if (it == threads_.end()) { + it = threads_.emplace(s.index, std::make_unique()).first; + } + return *it->second.get(); +} + +// Leak the scheduler singleton on all platforms. During static destruction, +// worker threads may still be executing JIT-compiled code that has been +// unmapped, causing SIGSEGV (macOS/Linux) or join() deadlocks (Windows/MSVC +// CRT). +// The OS reclaims all resources at process exit anyway. Scheduler& scheduler() { - // Leak the scheduler on Windows to avoid joining threads on exit, can be - // removed after Visual Studio fixes bug: - // https://developercommunity.visualstudio.com/t/1654756 -#ifdef _WIN32 static Scheduler* scheduler = new Scheduler; return *scheduler; -#else - static Scheduler scheduler; - return scheduler; -#endif } } // namespace scheduler diff --git a/Source/Cmlx/mlx/mlx/scheduler.h b/Source/Cmlx/mlx/mlx/scheduler.h index 4b949459f..7bce5a9fe 100644 --- a/Source/Cmlx/mlx/mlx/scheduler.h +++ b/Source/Cmlx/mlx/mlx/scheduler.h @@ -3,75 +3,24 @@ #pragma once #include -#include #include -#include +#include #include #include "mlx/api.h" #include "mlx/backend/gpu/eval.h" #include "mlx/device.h" #include "mlx/stream.h" +#include "mlx/utils.h" 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; @@ -79,30 +28,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); - } - - 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) { { @@ -134,40 +63,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/mlx/mlx/stream.cpp b/Source/Cmlx/mlx/mlx/stream.cpp new file mode 100644 index 000000000..b78ee67d6 --- /dev/null +++ b/Source/Cmlx/mlx/mlx/stream.cpp @@ -0,0 +1,109 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/stream.h" +#include "mlx/backend/cpu/device_info.h" +#include "mlx/backend/cpu/eval.h" +#include "mlx/backend/gpu/device_info.h" +#include "mlx/backend/gpu/eval.h" + +#include +#include +#include +#include + +namespace mlx::core { + +namespace { + +auto& default_stream_storage(Device d) { + // Each device has its own default stream in each thread. + static thread_local auto default_streams = []() { + std::array>, 2> streams; + streams[static_cast(Device::cpu)].resize(cpu::device_count()); + streams[static_cast(Device::gpu)].resize(gpu::device_count()); + return streams; + }(); + return default_streams[static_cast(d.type)].at(d.index); +} + +auto& all_streams() { + static std::tuple, std::shared_mutex> streams_and_mtx; + return streams_and_mtx; +} + +auto& thread_local_streams() { + static std::tuple, std::mutex> streams_and_mtx; + return streams_and_mtx; +} + +} // namespace + +Stream default_stream(Device d) { + if (!gpu::is_available() && d.type == Device::gpu) { + throw std::invalid_argument( + "[default_stream] Cannot get gpu stream without gpu backend."); + } + auto& s = default_stream_storage(d); + if (!s.has_value()) { + s = new_stream(d.type); + } + return s.value(); +} + +void set_default_stream(Stream s) { + if (!gpu::is_available() && s.device == Device::gpu) { + throw std::invalid_argument( + "[set_default_stream] Cannot set gpu stream without gpu backend."); + } + default_stream_storage(s.device) = s; +} + +std::vector get_streams() { + auto& [streams, mtx] = all_streams(); + std::shared_lock lock(mtx); + return streams; +} + +Stream new_stream(Device d) { + auto& [streams, mtx] = all_streams(); + std::unique_lock lock(mtx); + int index = streams.size(); + auto& s = streams.emplace_back(index, d); + if (d == Device::gpu) { + gpu::new_stream(s); + } else { + cpu::new_stream(s); + } + return s; +} + +Stream new_thread_unsafe_stream(Device d) { + auto& [streams, mtx] = all_streams(); + std::unique_lock lock(mtx); + int index = streams.size(); + auto& s = streams.emplace_back(index, d); + if (d == Device::gpu) { + gpu::new_thread_unsafe_stream(s); + } else { + cpu::new_thread_unsafe_stream(s); + } + return s; +} + +ThreadLocalStream new_thread_local_stream(Device d) { + auto& [streams, mtx] = thread_local_streams(); + std::lock_guard lock(mtx); + int index = streams.size(); + return streams.emplace_back(index, d); +} + +Stream stream_from_thread_local_stream(ThreadLocalStream tls) { + static thread_local std::map streams; + auto it = streams.find(tls); + if (it == streams.end()) { + it = streams.emplace(tls, new_stream(tls.device)).first; + } + return it->second; +} + +} // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/stream.h b/Source/Cmlx/mlx/mlx/stream.h index 9bd14db2e..3099ee268 100644 --- a/Source/Cmlx/mlx/mlx/stream.h +++ b/Source/Cmlx/mlx/mlx/stream.h @@ -2,6 +2,9 @@ #pragma once +#include +#include + #include "mlx/api.h" #include "mlx/device.h" @@ -11,27 +14,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); -inline bool operator==(const Stream& lhs, const Stream& rhs) { - return lhs.index == rhs.index; -} +/** 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 == rhs); -} +/** Get the stream for current thread from ThreadLocalStream. */ +MLX_API Stream stream_from_thread_local_stream(ThreadLocalStream tls); + +/** Get all available streams. */ +MLX_API std::vector get_streams(); /* Synchronize with the default stream. */ MLX_API void synchronize(); @@ -39,4 +53,10 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/transforms.cpp b/Source/Cmlx/mlx/mlx/transforms.cpp index 4967c50a8..694bca53d 100644 --- a/Source/Cmlx/mlx/mlx/transforms.cpp +++ b/Source/Cmlx/mlx/mlx/transforms.cpp @@ -24,6 +24,28 @@ namespace mlx::core { static constexpr int MAX_ACTIVE_TASKS = 10; +namespace { + +// Create a tracer copy of a primal for use in vjp/jvp. If the primal is a +// stale Copy from a previous transform call (not an active tracer), peel it +// off to prevent copy-chain accumulation when containers feed tracers back. +array make_tracer(const array& p) { + auto s = p.has_primitive() ? p.primitive().stream() + : default_stream(default_device()); + auto source = p; + if (!p.is_tracer() && p.has_primitive() && !p.inputs().empty()) { + auto& prim = p.primitive(); + if (typeid(prim) == typeid(Copy)) { + source = p.inputs()[0]; + } + } + auto out = copy(source, s); + out.set_tracer(true); + return out; +} + +} // namespace + /* This class is only meant to be used in eval * for synchronizing with the main thread. */ class Synchronizer : public Primitive { @@ -41,13 +63,19 @@ class Synchronizer : public Primitive { // These are used to implement the in_tracing() function the returns true if we // are currently under a function transformation and the retain_graph() // function which returns true if we are forced to retain the graph during -// evaluation. +// evaluation. They are thread_local since a trace belongs to the thread that +// runs it, so that concurrent transforms on different threads do not race on +// shared tracing state. std::vector>& detail::InTracing::trace_stack() { - static std::vector> trace_stack_; + static thread_local std::vector> trace_stack_; return trace_stack_; } -int detail::InTracing::grad_counter{0}; -int detail::RetainGraph::tracing_counter{0}; +thread_local int detail::InTracing::grad_counter{0}; +int& detail::InExportTracing::counter() { + static thread_local int counter_; + return counter_; +} +thread_local int detail::RetainGraph::tracing_counter{0}; array eval_impl(std::vector outputs, bool async) { std::deque tape; @@ -196,93 +224,119 @@ array eval_impl(std::vector outputs, bool async) { } } - std::unordered_set open_streams; - while (!tape.empty()) { - auto arr = std::move(tape.back()); - tape.pop_back(); + std::set open_streams; + try { + while (!tape.empty()) { + auto arr = std::move(tape.back()); + tape.pop_back(); - auto stream = arr.primitive().stream(); - open_streams.insert(stream.index); + auto stream = arr.primitive().stream(); + open_streams.insert(stream); - if (async) { - // Lookup corresponding event - auto e = events.find(stream.index); - if (e == events.end()) { - e = events.emplace(stream.index, Event{stream}).first; - } - e->second.set_value(1); - arr.attach_event(e->second); - for (auto& s : arr.siblings()) { - s.attach_event(e->second); + if (async) { + // Lookup corresponding event + auto e = events.find(stream.index); + if (e == events.end()) { + e = events.emplace(stream.index, Event{stream}).first; + } + e->second.set_value(1); + arr.attach_event(e->second); + for (auto& s : arr.siblings()) { + s.attach_event(e->second); + } } - } - for (auto& in : arr.inputs()) { - if (auto it = needs_fence.find(in.id()); it != needs_fence.end()) { - // Use fence to wait within a single eval - // Get the input array's stream fence and wait on the - // output arrays stream - fences[it->second.first].wait(stream, in); - } else if (in.event().valid()) { - if (in.event().is_signaled()) { - in.detach_event(); - } else if (in.event().stream() != stream) { - // Use event to wait across async eval - in.event().wait(stream); + for (auto& in : arr.inputs()) { + if (auto it = needs_fence.find(in.id()); it != needs_fence.end()) { + // Use fence to wait within a single eval + // Get the input array's stream fence and wait on the + // output arrays stream + fences[it->second.first].wait(stream, in); + } else if (in.event().valid()) { + if (in.event().is_signaled()) { + in.detach_event(); + } else if (in.event().stream() != stream) { + // Use event to wait across async eval + in.event().wait(stream); + } } } - } - if (arr.primitive().device() == Device::gpu) { - gpu::eval(arr); - } else { - cpu::eval(arr); - } + if (arr.primitive().device() == Device::gpu) { + gpu::eval(arr); + } else { + cpu::eval(arr); + } - if (scheduler::n_active_tasks() > MAX_ACTIVE_TASKS || - (get_active_memory() > get_memory_limit() && - scheduler::n_active_tasks() > 0)) { - // Commit any open streams - for (auto i : open_streams) { - auto s = get_stream(i); - if (s.device == Device::gpu) { - gpu::finalize(s); + if (scheduler::n_active_tasks() > MAX_ACTIVE_TASKS || + (get_active_memory() > get_memory_limit() && + scheduler::n_active_tasks() > 0)) { + // Commit any open streams + for (auto& s : open_streams) { + if (s.device == Device::gpu) { + gpu::finalize(s); + } } - } - scheduler::wait_for_one(); - while (get_active_memory() > get_memory_limit() && - scheduler::n_active_tasks() > 0) { scheduler::wait_for_one(); - } - } - - auto maybe_update_fence = [&fences, &needs_fence, stream](const array& a) { - if (auto nf = needs_fence.find(a.id()); nf != needs_fence.end()) { - auto it = fences.find(stream.index); - if (it == fences.end()) { - it = fences.emplace(stream.index, Fence{stream}).first; + while (get_active_memory() > get_memory_limit() && + scheduler::n_active_tasks() > 0) { + scheduler::wait_for_one(); } - it->second.update(stream, a, nf->second.second); } - }; - arr.set_status(array::Status::evaluated); - // TODO Maybe always want the fence coherent kernel in the same cbuf - // as the other kernels? - maybe_update_fence(arr); - for (auto& sib : arr.siblings()) { - sib.set_status(array::Status::evaluated); - maybe_update_fence(sib); + auto maybe_update_fence = + [&fences, &needs_fence, stream](const array& a) { + if (auto nf = needs_fence.find(a.id()); nf != needs_fence.end()) { + auto it = fences.find(stream.index); + if (it == fences.end()) { + it = fences.emplace(stream.index, Fence{stream}).first; + } + it->second.update(stream, a, nf->second.second); + } + }; + + arr.set_status(array::Status::evaluated); + // TODO Maybe always want the fence coherent kernel in the same cbuf + // as the other kernels? + maybe_update_fence(arr); + for (auto& sib : arr.siblings()) { + sib.set_status(array::Status::evaluated); + maybe_update_fence(sib); + } + if (!arr.is_tracer()) { + arr.detach(); + } } - if (!arr.is_tracer()) { - arr.detach(); + } catch (...) { + // A primitive threw from inside its eval (e.g. argument validation in + // eval_gpu, or a JIT compile failure). Arrays evaluated earlier in this + // tape are already marked evaluated, but their kernels sit in pending + // command buffers that only the epilogue below would commit, and events + // attached during this eval would never be signaled. Left that way, a + // later read of an affected array returns an unwritten buffer or blocks + // forever. Signal the events and flush the touched streams, then let the + // exception propagate. + for (auto& [idx, e] : events) { + try { + auto es = e.stream(); + e.signal(es); + open_streams.insert(es); + } catch (...) { + } } + for (auto& s : open_streams) { + try { + synchronize(s); + } catch (...) { + // Preserve the original exception. + } + } + throw; } // Signal the event in its stream - for (auto i : open_streams) { - auto s = get_stream(i); - if (auto e = events.find(i); e != events.end()) { + for (auto& s : open_streams) { + if (auto e = events.find(s.index); e != events.end()) { e->second.signal(s); } if (s.device == Device::gpu) { @@ -335,10 +389,7 @@ std::pair, std::vector> vjp( // Make tracers from given primals std::vector primals_; for (auto& p : primals) { - auto s = p.has_primitive() ? p.primitive().stream() - : default_stream(default_device()); - primals_.push_back(copy(p, s)); // Does not do a deep copy - primals_.back().set_tracer(true); + primals_.push_back(make_tracer(p)); } // Pass tracer primals through the function @@ -543,10 +594,7 @@ std::pair, std::vector> jvp( std::vector primals_; for (auto& p : primals) { - auto s = p.has_primitive() ? p.primitive().stream() - : default_stream(default_device()); - primals_.push_back(copy(p, s)); // Does not do a deep copy - primals_.back().set_tracer(true); + primals_.push_back(make_tracer(p)); } auto outputs = fun(primals_); @@ -620,6 +668,8 @@ std::pair, std::vector> jvp( auto jvps = a.primitive().jvp(a.inputs(), tangents, argnums); auto outputs = a.outputs(); + // A primitive's jvp returns one tangent per output + assert(jvps.size() <= outputs.size()); for (int i = 0; i < jvps.size(); ++i) { tan_map.insert({outputs[i].id(), jvps[i]}); } diff --git a/Source/Cmlx/mlx/mlx/transforms_impl.h b/Source/Cmlx/mlx/mlx/transforms_impl.h index eff458c40..4fc78a1ac 100644 --- a/Source/Cmlx/mlx/mlx/transforms_impl.h +++ b/Source/Cmlx/mlx/mlx/transforms_impl.h @@ -44,10 +44,26 @@ struct InTracing { } private: - static int grad_counter; + static thread_local int grad_counter; static std::vector>& trace_stack(); }; +struct InExportTracing { + InExportTracing() { + counter()++; + } + ~InExportTracing() { + counter()--; + } + + static bool in_export_tracing() { + return counter() > 0; + } + + private: + static int& counter(); +}; + struct RetainGraph { RetainGraph() { tracing_counter++; @@ -61,7 +77,7 @@ struct RetainGraph { } private: - static int tracing_counter; + static thread_local int tracing_counter; }; /** Return true if we are currently performing a function transformation in @@ -70,6 +86,11 @@ inline bool in_tracing() { return detail::InTracing::in_tracing(); } +/** Return true if the current trace is being recorded for export. */ +inline bool in_export_tracing() { + return detail::InExportTracing::in_export_tracing(); +} + /** Return true if we are in a dynamic (shapeless) trace used for compiling or * exporting graphs with dynamic shapes. */ inline bool in_dynamic_tracing() { diff --git a/Source/Cmlx/mlx/mlx/types/bf16.h b/Source/Cmlx/mlx/mlx/types/bf16.h index 595194174..7feaaa664 100644 --- a/Source/Cmlx/mlx/mlx/types/bf16.h +++ b/Source/Cmlx/mlx/mlx/types/bf16.h @@ -8,6 +8,7 @@ #include #define __MLX_BFLOAT_NAN__ 0x7FC0 +#define __MLX_BFLOAT_ONE__ 0x3F80 namespace mlx::core { @@ -29,8 +30,8 @@ struct _MLX_BFloat16 { // Appease std::vector for being special _MLX_BFloat16& operator=(std::vector::reference x) { - bits_ = x; - return *this; + bits_ = (x) ? __MLX_BFLOAT_ONE__ : 0; + return (*this); } _MLX_BFloat16& operator=(const float& x) { diff --git a/Source/Cmlx/mlx/mlx/types/fp16.h b/Source/Cmlx/mlx/mlx/types/fp16.h index c174afebc..31b0a78d1 100644 --- a/Source/Cmlx/mlx/mlx/types/fp16.h +++ b/Source/Cmlx/mlx/mlx/types/fp16.h @@ -8,6 +8,7 @@ #include #define __MLX_HALF_NAN__ 0x7D00 +#define __MLX_HALF_ONE__ 0x3C00 namespace mlx::core { @@ -29,8 +30,8 @@ struct _MLX_Float16 { // Appease std::vector for being special _MLX_Float16& operator=(std::vector::reference x) { - bits_ = x; - return *this; + bits_ = (x) ? __MLX_HALF_ONE__ : 0; + return (*this); } _MLX_Float16& operator=(const float& x) { diff --git a/Source/Cmlx/mlx/mlx/types/half_types.h b/Source/Cmlx/mlx/mlx/types/half_types.h index d9d6b9bf5..560f73735 100644 --- a/Source/Cmlx/mlx/mlx/types/half_types.h +++ b/Source/Cmlx/mlx/mlx/types/half_types.h @@ -36,10 +36,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); \ @@ -52,7 +53,23 @@ 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 diff --git a/Source/Cmlx/mlx/mlx/types/limits.h b/Source/Cmlx/mlx/mlx/types/limits.h index 5f2b1e9e0..0caf67228 100644 --- a/Source/Cmlx/mlx/mlx/types/limits.h +++ b/Source/Cmlx/mlx/mlx/types/limits.h @@ -30,6 +30,9 @@ struct numeric_limits { constexpr static float16_t lowest() { return bits_to_half(0xFBFF); } + static constexpr float16_t min() { + return bits_to_half(0x0400); + } static constexpr float16_t max() { return bits_to_half(0x7BFF); } @@ -56,6 +59,9 @@ struct numeric_limits { constexpr static bfloat16_t lowest() { return bits_to_bfloat(0xFF7F); } + static constexpr bfloat16_t min() { + return bits_to_bfloat(0x0080); + } static constexpr bfloat16_t max() { return bits_to_bfloat(0x7F7F); } diff --git a/Source/Cmlx/mlx/mlx/utils.cpp b/Source/Cmlx/mlx/mlx/utils.cpp index 1ae259b1e..4e29e8847 100644 --- a/Source/Cmlx/mlx/mlx/utils.cpp +++ b/Source/Cmlx/mlx/mlx/utils.cpp @@ -1,8 +1,10 @@ // Copyright © 2023 Apple Inc. #include +#include #include #include +#include #include #include "mlx/dtype_utils.h" @@ -16,6 +18,10 @@ Stream to_stream(StreamOrDevice s) { return default_stream(default_device()); } else if (std::holds_alternative(s)) { return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return stream_from_thread_local_stream(std::get(s)); } else { return std::get(s); } @@ -26,6 +32,10 @@ Stream to_stream(StreamOrDevice s, Device default_) { return default_stream(default_); } else if (std::holds_alternative(s)) { return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return default_stream(std::get(s)); + } else if (std::holds_alternative(s)) { + return stream_from_thread_local_stream(std::get(s)); } else { return std::get(s); } @@ -57,23 +67,51 @@ inline void PrintFormatter::print(std::ostream& os, uint64_t val) { os << val; } inline void PrintFormatter::print(std::ostream& os, float16_t val) { - os << val; + if (format_options.precision == -1) { + os << val; + } else { + os << std::fixed << std::setprecision(format_options.precision) << val; + } } inline void PrintFormatter::print(std::ostream& os, bfloat16_t val) { - os << val; + if (format_options.precision == -1) { + os << val; + } else { + os << std::fixed << std::setprecision(format_options.precision) << val; + } } inline void PrintFormatter::print(std::ostream& os, float val) { - os << val; + if (format_options.precision == -1) { + os << val; + } else { + os << std::fixed << std::setprecision(format_options.precision) << val; + } } inline void PrintFormatter::print(std::ostream& os, double val) { - os << val; + if (format_options.precision == -1) { + os << val; + } else { + os << std::fixed << std::setprecision(format_options.precision) << val; + } } inline void PrintFormatter::print(std::ostream& os, complex64_t val) { - os << val.real(); - if (val.imag() >= 0 || std::isnan(val.imag())) { - os << "+" << val.imag() << "j"; + if (format_options.precision == -1) { + os << val.real(); + if (val.imag() >= 0 || std::isnan(val.imag())) { + os << "+" << val.imag() << "j"; + } else { + os << "-" << -val.imag() << "j"; + } } else { - os << "-" << -val.imag() << "j"; + os << std::fixed << std::setprecision(format_options.precision) + << val.real(); + if (val.imag() >= 0 || std::isnan(val.imag())) { + os << "+" << std::fixed << std::setprecision(format_options.precision) + << val.imag() << "j"; + } else { + os << "-" << std::fixed << std::setprecision(format_options.precision) + << -val.imag() << "j"; + } } } @@ -82,11 +120,14 @@ PrintFormatter& get_global_formatter() { return formatter; } -void abort_with_exception(const std::exception& error) { - std::ostringstream msg; - msg << "Terminating due to uncaught exception: " << error.what(); - std::cerr << msg.str() << std::endl; - std::abort(); +void set_printoptions(PrintOptions options) { + auto& formatter = get_global_formatter(); + formatter.format_options = options; +} + +bool is_main_thread() { + static auto main_thread_id = std::this_thread::get_id(); + return main_thread_id == std::this_thread::get_id(); } Dtype result_type(const std::vector& arrays) { @@ -258,13 +299,28 @@ int get_var(const char* name, int default_value) { } } +std::string get_var(const char* name, const char* default_value) { + if (const char* buff_str = std::getenv(name)) { + return buff_str; + } else { + return default_value; + } +} + } // namespace env template -void set_finfo_limits(double& min, double& max, double& eps) { +void set_finfo_limits( + int& bits, + double& min, + double& max, + double& eps, + double& smallest_normal) { + bits = 8 * sizeof(T); min = numeric_limits::lowest(); max = numeric_limits::max(); eps = numeric_limits::epsilon(); + smallest_normal = numeric_limits::min(); } finfo::finfo(Dtype dtype) : dtype(dtype) { @@ -274,16 +330,16 @@ finfo::finfo(Dtype dtype) : dtype(dtype) { throw std::invalid_argument(msg.str()); } if (dtype == float32) { - set_finfo_limits(min, max, eps); + set_finfo_limits(bits, min, max, eps, smallest_normal); } else if (dtype == float16) { - set_finfo_limits(min, max, eps); + set_finfo_limits(bits, min, max, eps, smallest_normal); } else if (dtype == bfloat16) { - set_finfo_limits(min, max, eps); + set_finfo_limits(bits, min, max, eps, smallest_normal); } else if (dtype == float64) { - set_finfo_limits(min, max, eps); + set_finfo_limits(bits, min, max, eps, smallest_normal); } else if (dtype == complex64) { this->dtype = float32; - set_finfo_limits(min, max, eps); + set_finfo_limits(bits, min, max, eps, smallest_normal); } } diff --git a/Source/Cmlx/mlx/mlx/utils.h b/Source/Cmlx/mlx/mlx/utils.h index bb2de466b..0c62ba30d 100644 --- a/Source/Cmlx/mlx/mlx/utils.h +++ b/Source/Cmlx/mlx/mlx/utils.h @@ -3,6 +3,10 @@ #pragma once #include +#include +#include +#include +#include #include #include "mlx/api.h" @@ -13,7 +17,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_); @@ -38,6 +47,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); @@ -53,20 +66,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. */ @@ -88,6 +106,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). */ @@ -136,6 +172,7 @@ inline int next_power_of_2(int n) { namespace env { int get_var(const char* name, int default_value); +std::string get_var(const char* name, const char* default_value); inline int bfs_max_width() { static int bfs_max_width_ = get_var("MLX_BFS_MAX_WIDTH", 20); @@ -154,6 +191,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; @@ -169,6 +221,11 @@ inline int nccl_timeout(int default_value) { return nccl_timeout; } +inline const std::string& metal_gpu_arch() { + static std::string gpu_arch_ = get_var("MLX_METAL_GPU_ARCH", ""); + return gpu_arch_; +} + } // namespace env } // namespace mlx::core diff --git a/Source/Cmlx/mlx/mlx/version.h b/Source/Cmlx/mlx/mlx/version.h index 0667a9804..9de9f2e55 100644 --- a/Source/Cmlx/mlx/mlx/version.h +++ b/Source/Cmlx/mlx/mlx/version.h @@ -5,8 +5,8 @@ #include "mlx/api.h" #define MLX_VERSION_MAJOR 0 -#define MLX_VERSION_MINOR 30 -#define MLX_VERSION_PATCH 6 +#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/mlx/python/mlx/__array_api_info.py b/Source/Cmlx/mlx/python/mlx/__array_api_info.py new file mode 100644 index 000000000..847a0bcbf --- /dev/null +++ b/Source/Cmlx/mlx/python/mlx/__array_api_info.py @@ -0,0 +1,82 @@ +class ArrayNamespaceInfo: + def capabilities(self): + return { + "boolean indexing": False, + "data-dependent shapes": False, + "max dimensions": 10, + } + + def default_device(self): + import mlx.core as mx + + return mx.default_device() + + def default_dtypes(self, *, device=None): + import mlx.core as mx + + if device is not None and not isinstance(device, mx.Device): + raise TypeError("Expected a mlx Device") + return { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + } + + def devices(self): + import mlx.core as mx + + devices = [ + mx.Device(dev_type, i) + for dev_type in (mx.cpu, mx.gpu) + for i in range(mx.device_count(dev_type)) + ] + return tuple(devices) + + def dtypes(self, *, device=None, kind=None): + import mlx.core as mx + + if device is not None and not isinstance(device, mx.Device): + raise TypeError("Expected a mlx Device") + device = device if device is not None else self.default_device() + + dtypes = { + "bool": mx.bool_, + "int8": mx.int8, + "int16": mx.int16, + "int32": mx.int32, + "int64": mx.int64, + "uint8": mx.uint8, + "uint16": mx.uint16, + "uint32": mx.uint32, + "uint64": mx.uint64, + "float32": mx.float32, + "complex64": mx.complex64, + } + if device.type == mx.cpu: + dtypes["float64"] = mx.float64 + if kind is None: + return dtypes + + signed = {"int8", "int16", "int32", "int64"} + unsigned = {"uint8", "uint16", "uint32", "uint64"} + real = {"float32", "float64"} + complex_ = {"complex64"} + kinds = { + "bool": {"bool"}, + "signed integer": signed, + "unsigned integer": unsigned, + "integral": signed | unsigned, + "real floating": real, + "complex floating": complex_, + "numeric": signed | unsigned | real | complex_, + } + kind = (kind,) if isinstance(kind, str) else kind + if not isinstance(kind, tuple) or any(k not in kinds for k in kind): + raise ValueError(f"Unsupported dtype kind: {kind!r}") + names = {name for k in kind for name in kinds[k]} + return {name: dtype for name, dtype in dtypes.items() if name in names} + + +def __array_namespace_info__(): + return ArrayNamespaceInfo() diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py index d57747a32..6289da0c4 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/common.py @@ -69,8 +69,8 @@ def from_file(cls, hostfile): envs = [] hosts = [] if isinstance(data, dict): - backend = data["backend"] - envs = data["envs"] + backend = data.get("backend", backend) + envs = data.get("envs", envs) hosts = data["hosts"] elif isinstance(data, list): hosts = data @@ -90,7 +90,7 @@ def from_file(cls, hostfile): @classmethod def from_list(cls, hostlist, repeats=1): hosts = [] - for i, h in enumerate(hostlist.split(",")): + for h in hostlist.split(","): if h == "": raise ValueError("Hostname cannot be empty") try: @@ -98,8 +98,8 @@ def from_list(cls, hostlist, repeats=1): ips = [h] except ValueError: ips = [] - for i in range(repeats): - hosts.append(Host(i, h, ips, [])) + for _ in range(repeats): + hosts.append(Host(len(hosts), h, ips, [])) return cls(hosts) diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py index 1c0bbfe6c..a4dabed62 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/config.py @@ -295,6 +295,8 @@ def tb_connectivity_to_dot(hosts, tb_hosts, uuid_reverse_index): for p in h.ports: if not p.connected_to: continue + if p.connected_to not in uuid_reverse_index: + continue dst = uuid_reverse_index[p.connected_to] if dst[0] < i: continue @@ -359,13 +361,13 @@ def check_valid_ring(hosts, rings, strict=True): log_error("Try passing --dot to visualize the connectivity") if len(rings) > 0: log_error("Rings found:") - for r in rings: + for r, _ in rings: log_error(f" - {','.join(hosts[i].ssh_hostname for i in r)}") sys.exit(1) return has_ring -def check_ssh_connections(hosts): +def check_ssh_connections(hosts, ignore_unreachable=False): results = [None] * len(hosts) def _check(hostname, i): @@ -417,7 +419,7 @@ def _check(hostname, i): for t in threads: t.join() - if not all(results): + if not all(results) and not ignore_unreachable: log_error("Could not ssh to the following hosts:") for i, h in enumerate(hosts): if not results[i]: @@ -493,13 +495,13 @@ def configure_jaccl_ring(args, hosts, ips, ring, sshinfo): peer_left = ring[i - 1] peer_right = ring[(i + 1) % num_nodes] rdmas = [] - for j in range(len(hosts)): - if j not in (peer_left, peer_right): + for other in ring: + if other not in (peer_left, peer_right): rdmas.append(None) else: rdma = [] for c in range(count): - rdma.append(f"rdma_{ips.ips[i, j][c][0]}") + rdma.append(f"rdma_{ips.ips[node, other][c][0]}") rdmas.append(rdma[0] if count == 1 else rdma) jaccl_hosts.append(Host(i, h.ssh_hostname, h.ips, rdmas)) hostfile = Hostfile(jaccl_hosts, "jaccl-ring", args.env) @@ -573,13 +575,17 @@ def main(): parser.add_argument( "--hosts", default="127.0.0.1", help="A comma separated list of hosts" ) + parser.add_argument( + "--ignore-unreachable", + action="store_true", + help="Ignore hosts that are not reachable via ssh", + ) parser.add_argument("--hostfile", help="The file containing the hosts") parser.add_argument( "--over", choices=["thunderbolt", "ethernet"], default="thunderbolt", help="What type of connectivity to configure", - required=True, ) parser.add_argument( "--output-hostfile", help="If provided, save the hostfile to this path" @@ -609,17 +615,22 @@ def main(): ) args = parser.parse_args() - if args.hostfile is not None: - hosts = Hostfile.from_file(args.hostfile).hosts - else: - hosts = Hostfile.from_list(args.hosts).hosts + try: + if args.hostfile is not None: + hosts = Hostfile.from_file(args.hostfile).hosts + else: + hosts = Hostfile.from_list(args.hosts).hosts + except ValueError as e: + parser.error(str(e)) # Check that we can ssh log( args.verbose, f"Checking for ssh access for {', '.join(h.ssh_hostname for h in hosts)}", ) - sshinfo = check_ssh_connections(hosts) + sshinfo = check_ssh_connections(hosts, args.ignore_unreachable) + hosts = [h for r, h in zip(sshinfo, hosts) if r] + sshinfo = [r for r in sshinfo if r] # Prepare a hostfile for communication over ethernet using the ips of the # provided hostnames diff --git a/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py b/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py index d866eacaf..0e661e358 100644 --- a/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py +++ b/Source/Cmlx/mlx/python/mlx/_distributed_utils/launch.py @@ -47,9 +47,11 @@ def terminate(self): class RemoteProcess(CommandProcess): def __init__(self, rank, host, python, cwd, files, env, command): is_local = host == "127.0.0.1" - cmd = RemoteProcess.make_launch_script(rank, cwd, files, env, command, is_local) + cmd = RemoteProcess.make_launch_script( + rank, python, cwd, files, env, command, is_local + ) if not is_local: - cmd = f"ssh -tt -o LogLevel=QUIET {host} {shlex.quote(cmd)}" + cmd = f"ssh -tt -o LogLevel=QUIET {shlex.quote(host)} {shlex.quote(cmd)}" self._host = host self._pidfile = None @@ -91,7 +93,7 @@ def terminate(self): # Kill the remote program if possible cmd = RemoteProcess.make_kill_script(self._pidfile) if not self._is_local: - cmd = f"ssh {self._host} {shlex.quote(cmd)}" + cmd = f"ssh {shlex.quote(self._host)} {shlex.quote(cmd)}" c = run( cmd, check=True, @@ -104,7 +106,7 @@ def terminate(self): self._killed = c.stdout.strip() == "1" @staticmethod - def make_launch_script(rank, cwd, files, env, command, is_local): + def make_launch_script(rank, python, cwd, files, env, command, is_local): script = "" # Disable echo @@ -119,11 +121,12 @@ def make_launch_script(rank, cwd, files, env, command, is_local): # Change the working directory if one was requested. Otherwise attempt to # change to the current one but don't fail if it wasn't possible. d = cwd or os.getcwd() - script += f"if [[ -d {repr(d)} ]]; then " - script += f" cd {repr(d)}; " + qd = shlex.quote(d) + script += f"if [[ -d {qd} ]]; then " + script += f" cd {qd}; " if cwd is not None: script += "else " - script += f" echo 'Failed to change directory to' {repr(d)} >2; " + script += f" echo 'Failed to change directory to' {qd} >&2; " script += "fi; " # Add the environment variables that were requested @@ -146,6 +149,10 @@ def make_launch_script(rank, cwd, files, env, command, is_local): # Finally add the rank script += f"export MLX_RANK={rank}; " + # Run the command with the explicitly requested python interpreter + if python is not None: + command = [python, *command] + # Replace the process with the script script += f"cmd=({' '.join(map(shlex.quote, command))}); " script += 'exec "${cmd[@]}"' @@ -326,7 +333,7 @@ def launch_ring(parser, hosts, args, command): def launch_nccl(parser, hosts, args, command): if not hosts[0].ips: - raise ValueError("Rank 0 should have an IP reachable from all other ranks") + parser.error("Rank 0 should have an IP reachable from all other ranks") master_host = hosts[0].ips[0] master_port = args.nccl_port @@ -365,13 +372,35 @@ def launch_nccl(parser, hosts, args, command): def launch_jaccl(parser, hosts, args, command): if not hosts[0].ips: - raise ValueError("Rank 0 should have an IP reachable from all other ranks") + parser.error("Rank 0 should have an IP reachable from all other ranks") jaccl_ring = args.backend == "jaccl-ring" have_rdmas = all(len(h.rdma) == len(hosts) for h in hosts) + if not have_rdmas: + parser.error( + "The hostfile is malformed: number of RDMA devices does not match hosts" + ) have_nulls = all(h.rdma[i] is None for i, h in enumerate(hosts)) - if not have_rdmas or not have_nulls: - raise ValueError("Malformed hostfile for jaccl backend") + if not have_nulls: + parser.error("The hostfile is malformed: RDMA device of self should be null") + + # Find pairs that miss rmda in hostfile. + n = len(hosts) + missing_rdma = [ + (i, j) + for i, h in enumerate(hosts) + for j in (((i - 1) % n, (i + 1) % n) if jaccl_ring else range(n)) + if i != j and h.rdma[j] is None + ] + + if missing_rdma: + pairs = ", ".join( + f"{hosts[i].ssh_hostname} to {hosts[j].ssh_hostname}" + for i, j in missing_rdma[:3] + ) + if len(missing_rdma) > 3: + pairs += f" and {len(missing_rdma) - 3} more" + parser.error(f"The hostfile is malformed: no RDMA device is listed for {pairs}") coordinator = hosts[0].ips[0] env = args.env @@ -514,19 +543,15 @@ def main(): help="The port to use for the NCCL communication (only for nccl backend)", ) parser.add_argument( - "--no-verify-script", - action="store_false", - dest="verify_script", - help="Do not verify that the script exists", - ) - parser.add_argument( - "--python", default=sys.executable, help="Use this python on the remote hosts" + "--python", + default=None, + help="Launch the command with this python interpreter on the remote hosts", ) args, rest = parser.parse_known_args() if args.print_python: - print(args.python) + print(args.python or sys.executable) return if len(rest) == 0: @@ -535,10 +560,13 @@ def main(): rest.pop(0) # Try to extract a list of hosts and corresponding ips - if args.hostfile is not None: - hostfile = Hostfile.from_file(args.hostfile) - else: - hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts) + try: + if args.hostfile is not None: + hostfile = Hostfile.from_file(args.hostfile) + else: + hostfile = Hostfile.from_list(args.hosts, args.repeat_hosts) + except ValueError as e: + parser.error(str(e)) # Extract extra arguments from the hostfile if hostfile.backend != "" and args.backend is None: @@ -547,13 +575,11 @@ def main(): args.backend = "nccl" if mx.cuda.is_available() else "ring" args.env = hostfile.envs + args.env - # Check if the script is a file and convert it to a full path - if (script := Path(rest[0])).exists() and script.is_file(): - rest[0:1] = [args.python, str(script.resolve())] - elif (command := shutil.which(rest[0])) is not None: - rest[0] = command - elif args.verify_script: - raise ValueError(f"Invalid script or command {rest[0]}") + # Add terminal size + cols, lines = shutil.get_terminal_size((0, 0)) + if cols > 0 and not any(e.startswith("COLUMNS=") for e in args.env): + args.env.append(f"COLUMNS={cols}") + args.env.append(f"LINES={lines}") # Launch if args.backend == "ring": diff --git a/Source/Cmlx/mlx/python/mlx/_stub_patterns.txt b/Source/Cmlx/mlx/python/mlx/_stub_patterns.txt index 7ca436980..90afb5598 100644 --- a/Source/Cmlx/mlx/python/mlx/_stub_patterns.txt +++ b/Source/Cmlx/mlx/python/mlx/_stub_patterns.txt @@ -1,36 +1,44 @@ mlx.core.__prefix__: - from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, ParamSpec, TypeVar - import sys - if sys.version_info >= (3, 10): - from typing import TypeAlias - else: - from typing_extensions import TypeAlias + from typing import Any, BinaryIO as file, Literal, ParamSpec, Protocol, TypeAlias, TypeVar P = ParamSpec("P") R = TypeVar("R") + class DLPackCompatible(Protocol): + def __dlpack__(self, *args: Any, **kwargs: Any) -> Any: ... + def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any: ... mlx.core.__suffix__: - from typing import Union - scalar: TypeAlias = Union[int, float, bool] - list_or_scalar: TypeAlias = Union[scalar, list["list_or_scalar"]] + scalar: TypeAlias = int | float | bool | complex + list_or_scalar: TypeAlias = scalar | list["list_or_scalar"] + StreamOrDevice: TypeAlias = Stream | ThreadLocalStream | Device | DeviceType | None bool_: Dtype = ... +mlx.core.matrix_norm: + matrix_norm = linalg.norm + +mlx.core.array.__(eq|ne)__: + @overload + def __\1__(self, other: bool | int | float | array | Annotated[NDArray, dict(writable=False)] | complex) -> array: ... + @overload + def __\1__(self, other: ArrayLike) -> array | bool: ... + @overload + def __\1__(self, other: object) -> Any: ... + +mlx.core._PrintOptionsContext: + class _PrintOptionsContext: + def __init__(self, arg: PrintOptions, /) -> None: ... + def __enter__(self) -> _PrintOptionsContext: ... + def __exit__(self, *args) -> None: ... + mlx.core.distributed.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from mlx.core.distributed import Group - from typing import Sequence, Optional, Union + from mlx.core import array, Dtype, StreamOrDevice + from collections.abc import Callable, Sequence mlx.core.fast.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Union + from mlx.core import array, StreamOrDevice mlx.core.linalg.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Tuple, Union - -mlx.core.metal.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar - from typing import Sequence, Optional, Union + from mlx.core import array, StreamOrDevice mlx.core.random.__prefix__: - from mlx.core import array, Dtype, Device, Stream, scalar, float32, int32 - from typing import Sequence, Optional, Union + from mlx.core import array, Dtype, StreamOrDevice, scalar, float32, int32 + from collections.abc import Sequence diff --git a/Source/Cmlx/mlx/python/mlx/extension.py b/Source/Cmlx/mlx/python/mlx/extension.py index c426d5953..d794fbe28 100644 --- a/Source/Cmlx/mlx/python/mlx/extension.py +++ b/Source/Cmlx/mlx/python/mlx/extension.py @@ -30,13 +30,14 @@ def build_extension(self, ext: CMakeExtension) -> None: debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug cfg = "Debug" if debug else "Release" - # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON - # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code - # from Python. + # Point CMake at the interpreter running the build. Otherwise + # find_package(Python) picks whichever interpreter it finds first, + # which is not the one nanobind and mlx are installed into. cmake_args = [ f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}", f"-DCMAKE_BUILD_TYPE={cfg}", "-DBUILD_SHARED_LIBS=ON", + f"-DPython_EXECUTABLE={sys.executable}", ] build_args = [] # Adding CMake arguments set as environment variable diff --git a/Source/Cmlx/mlx/python/mlx/nn/__init__.py b/Source/Cmlx/mlx/python/mlx/nn/__init__.py index df72e2a94..5b51f8a01 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/__init__.py +++ b/Source/Cmlx/mlx/python/mlx/nn/__init__.py @@ -2,4 +2,7 @@ from mlx.nn import init, losses from mlx.nn.layers import * -from mlx.nn.utils import average_gradients, value_and_grad +from mlx.nn.utils import ( + average_gradients, + value_and_grad, +) diff --git a/Source/Cmlx/mlx/python/mlx/nn/init.py b/Source/Cmlx/mlx/python/mlx/nn/init.py index 8818cb9fd..f0bda14a6 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/init.py +++ b/Source/Cmlx/mlx/python/mlx/nn/init.py @@ -194,12 +194,13 @@ def glorot_uniform( ) -> Callable[[mx.array, float], mx.array]: r"""A Glorot uniform initializer. - This initializer samples from a uniform distribution with a range - computed from the number of input (``fan_in``) and output (``fan_out``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is computed from the number of input (``fan_in``) and output (``fan_out``) units according to: .. math:: - \sigma = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}} + \text{limit} = \gamma \sqrt{\frac{6.0}{\text{fan\_in} + \text{fan\_out}}} For more details see the original reference: `Understanding the difficulty of training deep feedforward neural networks @@ -295,13 +296,14 @@ def he_uniform( ) -> Callable[[mx.array, Literal["fan_in", "fan_out"], float], mx.array]: r"""A He uniform (Kaiming uniform) initializer. - This initializer samples from a uniform distribution with a range - computed from the number of input (``fan_in``) or output (``fan_out``) + This initializer samples from a uniform distribution on the interval + :math:`[-\text{limit}, \text{limit}]`, where the bound :math:`\text{limit}` + is computed from the number of input (``fan_in``) or output (``fan_out``) units according to: .. math:: - \sigma = \gamma \sqrt{\frac{3.0}{\text{fan}}} + \text{limit} = \gamma \sqrt{\frac{3.0}{\text{fan}}} where :math:`\text{fan}` is either the number of input units when the ``mode`` is ``"fan_in"`` or output units when the ``mode`` is @@ -358,9 +360,15 @@ def sparse( ) -> Callable[[mx.array], mx.array]: r"""An initializer that returns a sparse matrix. + Sparsity is applied along each row: every row has the same fraction of its + entries set to zero. With a weight matrix applied as ``x @ w.T``, this + limits each output feature to at most a ``1 - sparsity`` fraction of the + input features, following the sparse initialization of Martens, J. (2010), + "Deep learning via Hessian-free optimization". + Args: - sparsity (float): The fraction of elements in each column to be set to - zero. + sparsity (float): The fraction of elements in each row to be set to + zero. mean (float, optional): Mean of the normal distribution. Default: ``0.0``. std (float, optional): Standard deviation of the normal distribution. @@ -376,8 +384,8 @@ def sparse( >>> init_fn = nn.init.sparse(sparsity=0.5) >>> init_fn(mx.zeros((2, 2))) - array([[-1.91187, -0.117483], - [0, 0]], dtype=float32) + array([[-1.91187, 0], + [0, -0.117483]], dtype=float32) """ def initializer(a: mx.array) -> mx.array: @@ -387,6 +395,10 @@ def initializer(a: mx.array) -> mx.array: rows, cols = a.shape num_zeros = int(math.ceil(sparsity * cols)) + # Zero out `num_zeros` random entries in each row, so every row (output + # feature) drops the same fraction of its input connections. Sorting a + # per-row random key (axis=1) picks the zeroed columns independently for + # every row. order = mx.argsort(mx.random.uniform(shape=a.shape), axis=1) a = mx.random.normal(shape=a.shape, scale=std, loc=mean, dtype=dtype) @@ -415,8 +427,8 @@ def orthogonal( def initializer(a: mx.array) -> mx.array: if a.ndim != 2: raise ValueError( - f"Orthogonal initialization requires a 2D array but got" - " a {a.ndim}D array." + "Orthogonal initialization requires a 2D array but got" + f" a {a.ndim}D array." ) rows, cols = a.shape diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/__init__.py b/Source/Cmlx/mlx/python/mlx/nn/layers/__init__.py index c2fba5834..25b96c0a7 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/__init__.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/__init__.py @@ -64,9 +64,11 @@ ) from mlx.nn.layers.distributed import ( AllToShardedLinear, + FullyShardedModule, QuantizedAllToShardedLinear, QuantizedShardedToAllLinear, ShardedToAllLinear, + fully_shard, ) from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d from mlx.nn.layers.embedding import Embedding diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py b/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py index 7b7bb138e..305b5ec7c 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/activations.py @@ -66,8 +66,12 @@ def leaky_relu(x, negative_slope=0.01): def log_softmax(x, axis=-1): r"""Applies the Log Softmax function. - Applies :math:`x + \log \sum_i e^{x_i}` element wise. + Applies :math:`x - \log \sum_i e^{x_i}` element wise. """ + # Shift by the max first. Subtracting the logsumexp of x directly loses the + # normalizer, since adding it to a large max rounds away before the + # subtraction happens. + x = x - mx.stop_gradient(mx.max(x, axis=axis, keepdims=True)) return x - mx.logsumexp(x, axis=axis, keepdims=True) @@ -176,7 +180,7 @@ def gelu_approx(x): .. math:: - x = 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) + x = 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) """ return 0.5 * x * (1 + mx.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * x**3))) @@ -225,7 +229,8 @@ def step(x: mx.array, threshold: float = 0.0): r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set - to 1 if the input is greater than a specified threshold, and 0 otherwise. + to 1 if the input is greater than or equal to a specified threshold, and 0 + otherwise. .. math:: \text{step}(x) = \begin{cases} @@ -234,10 +239,10 @@ def step(x: mx.array, threshold: float = 0.0): \end{cases} Args: - threshold: The value to threshold at. + threshold: The value to threshold at. Default: ``0.0``. """ - return mx.where(x > threshold, 1, 0) + return mx.where(x >= threshold, 1, 0) @partial(mx.compile, shapeless=True) @@ -552,7 +557,7 @@ class GELU(Module): However, if ``approx`` is set to 'precise' or 'fast' it applies .. math:: - \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left((\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) \\ + \textrm{GELUApprox}(x) &= 0.5 * x * \left(1 + \text{Tanh}\left(\sqrt{2 / \pi} * \left(x + 0.044715 * x^3\right)\right)\right) \\ \textrm{GELUFast}(x) &= x * \sigma\left(1.702 * x\right) respectively. @@ -606,7 +611,8 @@ class Step(Module): r"""Applies the Step Activation Function. This function implements a binary step activation, where the output is set - to 1 if the input is greater than a specified threshold, and 0 otherwise. + to 1 if the input is greater than or equal to a specified threshold, and 0 + otherwise. .. math:: \text{step}(x) = \begin{cases} @@ -615,7 +621,7 @@ class Step(Module): \end{cases} Args: - threshold: The value to threshold at. + threshold: The value to threshold at. Default: ``0.0``. """ def __init__(self, threshold: float = 0.0): @@ -642,7 +648,6 @@ class HardTanh(Module): """ -@_make_activation_module(hard_shrink) class HardShrink(Module): r"""Applies the HardShrink function. @@ -652,6 +657,13 @@ class HardShrink(Module): lambd: the :math:`\lambda` value for Hardshrink. Default: ``0.5`` """ + def __init__(self, lambd=0.5): + super().__init__() + self.lambd = lambd + + def __call__(self, x): + return hard_shrink(x, self.lambd) + @_make_activation_module(softmin) class Softmin(Module): diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/base.py b/Source/Cmlx/mlx/python/mlx/nn/layers/base.py index fca65d787..e980f8343 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/base.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/base.py @@ -129,7 +129,7 @@ def load_weights( Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list. Args: - file_or_weights (str or list(tuple(str, mx.array))): The path to + file_or_weights (str or list(tuple(str, array))): The path to the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list of pairs of parameter names and arrays. strict (bool, optional): If ``True`` then checks that the provided @@ -255,13 +255,13 @@ def filter_and_map( but it can also be used to extract any subset of the module's parameters. Args: - filter_fn (Callable): Given a value, the key in which it is found - and the containing module, decide whether to keep the value or + filter_fn (Callable): Given the containing module, the key in which + it is found and the value, decide whether to keep the value or drop it. map_fn (Callable, optional): Optionally transform the value before returning it. - is_leaf_fn (Callable, optional): Given a value, the key in which it - is found and the containing module decide if it is a leaf. + is_leaf_fn (Callable, optional): Given the containing module, the + key in which it is found and the value decide if it is a leaf. Returns: A dictionary containing the contents of the module recursively filtered @@ -340,6 +340,13 @@ def apply(dst, parameters): raise ValueError(f'Module does not have parameter named "{k}".') elif isinstance(parameters, list): for i in range(len(parameters)): + if i >= len(dst): + if strict: + raise ValueError( + f"List index {i} is out of bounds for " + f"destination of length {len(dst)}." + ) + continue current_value = dst[i] new_value = parameters[i] if isinstance(current_value, mx.array): @@ -454,6 +461,16 @@ def _validate_keys(self, keys, strict): raise KeyError(f"Module doesn't contain member {k}.") return keys + def _validate_keys_recursive(self, keys): + # A key such as "bias" is expected to be present somewhere in the model + # but not necessarily in every single submodule, so validate against the + # whole model rather than each module in isolation. + keys = keys if isinstance(keys, list) else [keys] + modules = self.modules() + for k in keys: + if not any(k in m for m in modules): + raise KeyError(f"Module doesn't contain member {k}.") + def freeze( self, *, @@ -504,6 +521,9 @@ def _freeze_impl(_, m): m._no_grad.update(local_keys) if recurse: + if strict and keys is not None: + self._validate_keys_recursive(keys) + strict = False self.apply_to_modules(_freeze_impl) else: _freeze_impl("", self) @@ -554,6 +574,9 @@ def _unfreeze_impl(_, m): m._no_grad.difference_update(local_keys) if recurse: + if strict and keys is not None: + self._validate_keys_recursive(keys) + strict = False self.apply_to_modules(_unfreeze_impl) else: _unfreeze_impl("", self) @@ -627,6 +650,13 @@ def _update_modules(dst, modules, strict): raise ValueError(f'Module does not have sub-module named "{k}".') elif isinstance(modules, list): for i in range(len(modules)): + if i >= len(dst): + if strict: + raise ValueError( + f"List index {i} is out of bounds for " + f"destination of length {len(dst)}." + ) + continue current_value = dst[i] new_value = modules[i] if Module.is_module(current_value) and Module.is_module(new_value): diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py index 88b97add0..40f34ff56 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution.py @@ -25,6 +25,7 @@ class Conv1d(Module): padding (int, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int, optional): The dilation of the convolution. + Default: ``1``. groups (int, optional): The number of groups for the convolution. Default: ``1``. bias (bool, optional): If ``True`` add a learnable bias to the output. @@ -66,7 +67,7 @@ def __init__( def _extra_repr(self): return ( - f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " + f"{self.weight.shape[-1] * self.groups}, {self.weight.shape[0]}, " f"kernel_size={self.weight.shape[1]}, stride={self.stride}, " f"padding={self.padding}, dilation={self.dilation}, " f"groups={self.groups}, " @@ -101,6 +102,7 @@ class Conv2d(Module): padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. + Default: ``1``. groups (int, optional): The number of groups for the convolution. Default: ``1``. bias (bool, optional): If ``True`` add a learnable bias to the @@ -146,7 +148,7 @@ def __init__( def _extra_repr(self): return ( - f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " + f"{self.weight.shape[-1] * self.groups}, {self.weight.shape[0]}, " f"kernel_size={self.weight.shape[1:3]}, stride={self.stride}, " f"padding={self.padding}, dilation={self.dilation}, " f"groups={self.groups}, " @@ -179,9 +181,10 @@ class Conv3d(Module): kernel_size (int or tuple): The size of the convolution filters. stride (int or tuple, optional): The size of the stride when applying the filter. Default: ``1``. - dilation (int or tuple, optional): The dilation of the convolution. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. + dilation (int or tuple, optional): The dilation of the convolution. + Default: ``1``. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py index a11c4cb40..1fe8935fc 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/convolution_transpose.py @@ -25,6 +25,7 @@ class ConvTranspose1d(Module): padding (int, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int, optional): The dilation of the convolution. + Default: ``1``. output_padding(int, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the output. @@ -100,6 +101,7 @@ class ConvTranspose2d(Module): padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. + Default: ``1``. output_padding(int or tuple, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the @@ -140,7 +142,7 @@ def __init__( def _extra_repr(self): return ( f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " - f"kernel_size={self.weight.shape[1:2]}, stride={self.stride}, " + f"kernel_size={self.weight.shape[1:3]}, stride={self.stride}, " f"padding={self.padding}, dilation={self.dilation}, " f"output_padding={self.output_padding}, " f"bias={'bias' in self}" @@ -180,6 +182,7 @@ class ConvTranspose3d(Module): padding (int or tuple, optional): How many positions to 0-pad the input with. Default: ``0``. dilation (int or tuple, optional): The dilation of the convolution. + Default: ``1``. output_padding(int or tuple, optional): Additional size added to one side of the output shape. Default: ``0``. bias (bool, optional): If ``True`` add a learnable bias to the @@ -222,7 +225,7 @@ def __init__( def _extra_repr(self): return ( f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " - f"kernel_size={self.weight.shape[1:3]}, stride={self.stride}, " + f"kernel_size={self.weight.shape[1:4]}, stride={self.stride}, " f"padding={self.padding}, dilation={self.dilation}, " f"output_padding={self.output_padding}, " f"bias={'bias' in self}" diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py b/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py index c7d79e55c..16979b509 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/distributed.py @@ -1,14 +1,14 @@ # Copyright © 2024 Apple Inc. import math -from functools import lru_cache +from functools import lru_cache, reduce from typing import Callable, Optional, Union import mlx.core as mx from mlx.nn.layers.base import Module from mlx.nn.layers.linear import Linear from mlx.nn.layers.quantized import QuantizedLinear -from mlx.utils import tree_map_with_path +from mlx.utils import tree_flatten, tree_map_with_path, tree_unflatten @lru_cache @@ -199,7 +199,7 @@ class AllToShardedLinear(Module): Args: input_dims (int): The dimensionality of the input features output_dims (int): The dimensionality of the output features - bias (bool, optional): If set to ``False`` the the layer will not use a + bias (bool, optional): If set to ``False`` the layer will not use a bias. Default is ``True``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is @@ -283,7 +283,7 @@ class ShardedToAllLinear(Module): Args: input_dims (int): The dimensionality of the input features output_dims (int): The dimensionality of the output features - bias (bool, optional): If set to ``False`` the the layer will not use a + bias (bool, optional): If set to ``False`` the layer will not use a bias. Default is ``True``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is @@ -371,6 +371,8 @@ class QuantizedAllToShardedLinear(Module): weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. + mode (str, optional): The quantization method to use (see + :func:`~mlx.core.quantize`). Default: ``"affine"``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. @@ -383,6 +385,7 @@ def __init__( bias: bool = True, group_size: int = 64, bits: int = 4, + mode: str = "affine", group: Optional[mx.distributed.Group] = None, ): super().__init__() @@ -390,6 +393,7 @@ def __init__( # Quantization config self.group_size = group_size self.bits = bits + self.mode = mode # Initialize the quantized weight scale = math.sqrt(1.0 / input_dims) @@ -406,7 +410,10 @@ def __init__( high=scale, shape=(output_dims // N, input_dims), ) - self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits) + self.weight, self.scales, *biases = mx.quantize( + weight, group_size, bits, mode=mode + ) + self.biases = biases[0] if biases else None # And bias if needed if bias: @@ -427,7 +434,7 @@ def _extra_repr(self) -> str: out_dims *= self.group.size() return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " - f"group_size={self.group_size}, bits={self.bits}" + f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" ) def __call__(self, x: mx.array) -> mx.array: @@ -438,10 +445,11 @@ def __call__(self, x: mx.array) -> mx.array: x, self["weight"], scales=self["scales"], - biases=self["biases"], + biases=self.get("biases"), transpose=True, group_size=self.group_size, bits=self.bits, + mode=self.mode, ) if "bias" in self: x = x + self["bias"] @@ -465,6 +473,7 @@ def from_quantized_linear( hasattr(quantized_linear_layer, "bias"), group_size=quantized_linear_layer.group_size, bits=quantized_linear_layer.bits, + mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, ) sl.update( @@ -497,6 +506,8 @@ class QuantizedShardedToAllLinear(Module): weight. See :func:`~mlx.core.quantize`. Default: ``64``. bits (int, optional): The bit width to use for the quantized weight. See :func:`~mlx.core.quantize`. Default: ``4``. + mode (str, optional): The quantization method to use (see + :func:`~mlx.core.quantize`). Default: ``"affine"``. group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. @@ -509,6 +520,7 @@ def __init__( bias: bool = True, group_size: int = 64, bits: int = 4, + mode: str = "affine", group: Optional[mx.distributed.Group] = None, ): super().__init__() @@ -516,6 +528,7 @@ def __init__( # Quantization config self.group_size = group_size self.bits = bits + self.mode = mode # Initialize the quantized weight scale = math.sqrt(1.0 / input_dims) @@ -532,7 +545,10 @@ def __init__( high=scale, shape=(output_dims, input_dims // N), ) - self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits) + self.weight, self.scales, *biases = mx.quantize( + weight, group_size, bits, mode=mode + ) + self.biases = biases[0] if biases else None # And bias if needed if bias: @@ -552,7 +568,7 @@ def _extra_repr(self) -> str: in_dims = (in_dims * 32) // self.bits * self.group.size() return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " - f"group_size={self.group_size}, bits={self.bits}" + f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" ) def __call__(self, x: mx.array) -> mx.array: @@ -560,10 +576,11 @@ def __call__(self, x: mx.array) -> mx.array: x, self["weight"], scales=self["scales"], - biases=self["biases"], + biases=self.get("biases"), transpose=True, group_size=self.group_size, bits=self.bits, + mode=self.mode, ) x = mx.distributed.all_sum(x, group=self.group) if "bias" in self: @@ -588,6 +605,7 @@ def from_quantized_linear( hasattr(quantized_linear_layer, "bias"), group_size=quantized_linear_layer.group_size, bits=quantized_linear_layer.bits, + mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, ) sl.update( @@ -599,3 +617,147 @@ def from_quantized_linear( ) return sl + + +def _make_gather_fn(group, full_shapes, shard_sizes, compute_dtype): + N = group.size() + indices = reduce(lambda acc, w: acc + [acc[-1] + w], shard_sizes, [0]) + split_indices = indices[1:-1] + shard_shapes = [(shape[0] // N,) + tuple(shape[1:]) for shape in full_shapes] + + def _maybe_cast(x, dtype): + if dtype is None or x.dtype == dtype: + return x + return x.astype(dtype) + + @mx.custom_function + def gather(shards): + shard = mx.concatenate( + [_maybe_cast(s.reshape(1, -1), compute_dtype) for s in shards], axis=1 + ) + full = mx.distributed.all_gather(shard, group=group) + parts = mx.split(full, split_indices, axis=1) + return [p.reshape(shape) for p, shape in zip(parts, full_shapes)] + + @gather.vjp + def gather_vjp(shards, cotangents, _): + local_full = mx.concatenate([c.reshape(N, -1) for c in cotangents], axis=1) + local_shard = mx.distributed.sum_scatter(local_full, group=group) / N + parts = mx.split(local_shard, split_indices, axis=1) + return [ + _maybe_cast(p.reshape(shape), s.dtype) + for p, shape, s in zip(parts, shard_shapes, shards) + ] + + return gather + + +def _maybe_shard(m, k, v): + if isinstance(v, FullyShardedModule): + return False + return Module.valid_parameter_filter(m, k, v) + + +class FullyShardedModule(Module): + """Wrap a module so each member of the group holds only a shard of its + parameters. + + The full parameters are gathered for the forward pass and the gradients + are reduce-scattered in the backward pass, so during training + each member of the group stores and updates only its own shard. + + Every parameter is sharded along axis 0, so each parameter's size along + that axis must be divisible by the size of ``group``. + + Use :func:`~mlx.nn.layers.distributed.fully_shard` to wrap a module. + + Args: + module (mlx.nn.Module): The module whose parameters will be sharded. + group (mlx.core.distributed.Group, optional): The group to shard + across. If not set, the global group is used. Default: ``None``. + compute_dtype (mlx.core.Dtype, optional): If set, the gathered + parameters are cast to this dtype for the forward pass. + Default: ``None``. + """ + + def __init__( + self, + module: Module, + group: Optional[mx.distributed.Group] = None, + compute_dtype: Optional[mx.Dtype] = None, + ): + super().__init__() + group = group or mx.distributed.init() + N = group.size() + + shard_params = module.filter_and_map(_maybe_shard) + flat = tree_flatten(shard_params) + for path, a in flat: + if a.ndim == 0: + raise ValueError( + f"Cannot shard parameter '{path}' because it is a scalar." + ) + if a.shape[0] % N != 0: + raise ValueError( + f"Cannot shard parameter '{path}' with shape {a.shape} " + f"across {N} devices: axis 0 must be divisible by {N}." + ) + + super(Module, self).__setattr__("_paths", [k for k, _ in flat]) + full_shapes = [a.shape for _, a in flat] + shard_sizes = [a.size // N for _, a in flat] + + module.update(_shard(shard_params, lambda p, w: 0, group)) + + self.module = module + self._gather_fn = _make_gather_fn( + group, full_shapes, shard_sizes, compute_dtype + ) + + def _extra_repr(self) -> str: + return f"num_sharded_params={len(self._paths)}" + + def _gathered_call(self, fn, *args, **kwargs): + shard_tree = self.module.filter_and_map(_maybe_shard) + shards = [a for _, a in tree_flatten(shard_tree)] + fulls = self._gather_fn(shards) + self.module.update(tree_unflatten(list(zip(self._paths, fulls)))) + try: + return fn(*args, **kwargs) + finally: + self.module.update(shard_tree) + + def __call__(self, *args, **kwargs): + return self._gathered_call(self.module, *args, **kwargs) + + def as_linear(self, *args, **kwargs): + return self._gathered_call(self.module.as_linear, *args, **kwargs) + + +def fully_shard( + module: Module, + *, + group: Optional[mx.distributed.Group] = None, + compute_dtype: Optional[mx.Dtype] = None, +) -> Module: + """Wrap ``module`` in a :class:`FullyShardedModule`. + + Args: + module (mlx.nn.Module): The module to wrap. + group (mlx.core.distributed.Group, optional): The group to shard + across. If not set, the global group is used. Default: ``None``. + compute_dtype (mlx.core.Dtype, optional): If set, the gathered + parameters are cast to this dtype for the forward pass. + Default: ``None``. + + Returns: + The wrapped :class:`FullyShardedModule`, or ``module`` unchanged. + """ + group = group or mx.distributed.init() + if group.size() == 1: + return module + if isinstance(module, FullyShardedModule): + return module + + wrapped = FullyShardedModule(module, group=group, compute_dtype=compute_dtype) + return wrapped if wrapped._paths else module diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py b/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py index a6a690996..ef74ba9b6 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/dropout.py @@ -12,7 +12,7 @@ class Dropout(Module): expected value of a given element will remain the same. Args: - p (float): The probability to zero an element + p (float): The probability to zero an element. Default: ``0.5``. """ def __init__(self, p: float = 0.5): @@ -40,7 +40,7 @@ class Dropout2d(Module): Randomly zero out entire channels independently with probability :math:`p`. This layer expects the channels to be last, i.e. the input shape should be - ``NWHC`` or ``WHC`` where:``N`` is the batch dimension,``H`` is the input + ``NHWC`` or ``HWC`` where:``N`` is the batch dimension,``H`` is the input image height,``W`` is the input image width, and``C`` is the number of input channels @@ -56,6 +56,7 @@ class Dropout2d(Module): Args: p (float): Probability of zeroing a channel during training. + Default: ``0.5``. """ def __init__(self, p: float = 0.5): @@ -105,6 +106,7 @@ class Dropout3d(Module): Args: p (float): Probability of zeroing a channel during training. + Default: ``0.5``. """ def __init__(self, p: float = 0.5): diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py b/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py index 891627be3..e05cfb5f0 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/embedding.py @@ -45,6 +45,9 @@ def to_quantized( group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = "affine", + quantize_input: bool = False, ): """Return a :obj:`QuantizedEmbedding` layer that approximates this embedding layer.""" + if quantize_input: + raise ValueError("Quantized input is not supported.") return QuantizedEmbedding.from_embedding(self, group_size, bits, mode) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py b/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py index 8f00ff180..a1da07c46 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/linear.py @@ -5,7 +5,7 @@ import mlx.core as mx from mlx.nn.layers.base import Module -from mlx.nn.layers.quantized import QuantizedLinear +from mlx.nn.layers.quantized import QQLinear, QuantizedLinear class Identity(Module): @@ -32,7 +32,6 @@ class Linear(Module): y = x W^\top + b - where: where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``. The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, @@ -75,8 +74,36 @@ def to_quantized( group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = "affine", + quantize_input: bool = False, ): - """Return a :obj:`QuantizedLinear` layer that approximates this layer.""" + """Return a quantized approximation of this layer. + + If ``quantize_input`` is ``False``, returns a :obj:`QuantizedLinear` + (weights are quantized). If ``quantize_input`` is ``True``, returns + a :obj:`QQLinear` (weights and activations are quantized). + + Args: + group_size (Optional[int]): The quantization group size (see + :func:`mlx.core.quantize`). Default: ``None``. + bits (Optional[int]): The number of bits per parameter (see + :func:`mlx.core.quantize`). Default: ``None``. + mode (str): The quantization method to use (see + :func:`mlx.core.quantize`). Default: ``"affine"``. + quantize_input (bool): Whether to quantize input. Default: ``False``. + + Returns: + QuantizedLinear or QQLinear: A quantized version of this layer. + + Notes: + Quantized input is only supported for ``"nvfp4"`` and ``"mxfp8"`` + modes. + """ + if quantize_input: + if mode not in ["nvfp4", "mxfp8"]: + raise ValueError( + f"Quantized activations are only supported for 'nvfp4' and 'mxfp8' modes, got {mode}." + ) + return QQLinear.from_linear(self, group_size, bits, mode) return QuantizedLinear.from_linear(self, group_size, bits, mode) @@ -87,10 +114,10 @@ class Bilinear(Module): .. math:: - y_i = x_1^\top W_i x_2 + b_i + y_i = x_2^\top W_i x_1 + b_i where: - :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``, + :math:`W` has shape ``[output_dims, input2_dims, input1_dims]``, :math:`b` has shape ``[output_dims]``, and :math:`i` indexes the output dimension. The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`, diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py b/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py index bdddb6ccf..97f6942f0 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/normalization.py @@ -25,7 +25,8 @@ class InstanceNorm(Module): affine (bool): Default: ``False``. Shape: - - Input: :math:`(..., C)` where :math:`C` is equal to :attr:`dims`. + - Input: :math:`(N, ..., C)` where :math:`C` is equal to :attr:`dims`. + The input must have at least 3 dimensions. - Output: Same shape as the input. Examples: @@ -46,6 +47,8 @@ def __init__( affine: bool = False, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[InstanceNorm] 'eps' must be positive but got {eps}.") if affine: self.weight = mx.ones((dims,)) self.bias = mx.zeros((dims,)) @@ -56,12 +59,24 @@ def _extra_repr(self): return f"{self.dims}, eps={self.eps}, affine={'weight' in self}" def __call__(self, x: mx.array) -> mx.array: - reduction_axes = tuple(range(1, x.ndim - 1)) - # Compute stats - mean = mx.mean(x, axis=reduction_axes, keepdims=True) - var = mx.var(x, axis=reduction_axes, keepdims=True) - # Normalize - x = (x - mean) * mx.rsqrt(var + self.eps) + if x.ndim < 3: + raise ValueError( + f"InstanceNorm expects inputs with at least 3 dimensions" + f" (N, ..., C) but the input has {x.ndim} dimensions." + ) + batch_size, features = x.shape[0], x.shape[-1] + spatial_shape = x.shape[1:-1] + channels_first = mx.transpose(x, (0, x.ndim - 1, *range(1, x.ndim - 1))) + x = mx.fast.layer_norm( + channels_first.reshape(batch_size, features, -1), + None, + None, + self.eps, + ) + x = mx.transpose( + x.reshape(batch_size, features, *spatial_shape), + (0, *range(2, len(spatial_shape) + 2), 1), + ) # Scale and shift if necessary return (self.weight * x + self.bias) if "weight" in self else x @@ -73,7 +88,7 @@ class LayerNorm(Module): .. math:: - y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, + y = \frac{x - E[x]}{\sqrt{Var[x] + \epsilon}} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. @@ -82,18 +97,21 @@ class LayerNorm(Module): Args: dims (int): The feature dimension of the input to normalize over - eps (float): A small additive constant for numerical stability + eps (float): A small additive constant for numerical stability. + Default: ``1e-5``. affine (bool): If True learn an affine transform to apply after the - normalization + normalization. Default: ``True``. bias (bool): If True include a translation to the affine transformation. If set to False the transformation is not really affine - just scaling. + just scaling. Default: ``True``. """ def __init__( self, dims: int, eps: float = 1e-5, affine: bool = True, bias: bool = True ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[LayerNorm] 'eps' must be positive but got {eps}.") if affine: self.weight = mx.ones((dims,)) if bias: @@ -128,11 +146,14 @@ class RMSNorm(Module): Args: dims (int): The feature dimension of the input to normalize over - eps (float): A small additive constant for numerical stability + eps (float): A small additive constant for numerical stability. + Default: ``1e-5``. """ def __init__(self, dims: int, eps: float = 1e-5): super().__init__() + if eps <= 0.0: + raise ValueError(f"[RMSNorm] 'eps' must be positive but got {eps}.") self.weight = mx.ones((dims,)) self.eps = eps @@ -150,7 +171,7 @@ class GroupNorm(Module): .. math:: - y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, + y = \frac{x - E[x]}{\sqrt{Var[x] + \epsilon}} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. However, the mean and @@ -166,11 +187,12 @@ class GroupNorm(Module): Args: num_groups (int): Number of groups to separate the features into dims (int): The feature dimensions of the input to normalize over - eps (float): A small additive constant for numerical stability + eps (float): A small additive constant for numerical stability. + Default: ``1e-5``. affine (bool): If True learn an affine transform to apply after the - normalization. + normalization. Default: ``True``. pytorch_compatible (bool): If True perform the group normalization in - the same order/grouping as PyTorch. + the same order/grouping as PyTorch. Default: ``False``. """ def __init__( @@ -182,6 +204,21 @@ def __init__( pytorch_compatible: bool = False, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[GroupNorm] 'eps' must be positive but got {eps}.") + if num_groups <= 0: + raise ValueError( + f"The number of groups ({num_groups}) must be a positive integer." + ) + if dims <= 0: + raise ValueError( + f"The number of features ({dims}) must be a positive integer." + ) + if dims % num_groups != 0: + raise ValueError( + f"The number of features ({dims}) must be evenly divisible" + f" by the number of groups ({num_groups})." + ) if affine: self.bias = mx.zeros((dims,)) self.weight = mx.ones((dims,)) @@ -238,13 +275,13 @@ def __call__(self, x): class BatchNorm(Module): - r"""Applies Batch Normalization over a 2D or 3D input. + r"""Applies Batch Normalization over a 2D, 3D or 4D input. Computes .. math:: - y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta, + y = \frac{x - E[x]}{\sqrt{Var[x] + \epsilon}} \gamma + \beta, where :math:`\gamma` and :math:`\beta` are learned per feature dimension parameters initialized at 1 and 0 respectively. @@ -287,6 +324,8 @@ def __init__( track_running_stats: bool = True, ): super().__init__() + if eps <= 0.0: + raise ValueError(f"[BatchNorm] 'eps' must be positive but got {eps}.") self.num_features = num_features self.eps = eps @@ -315,13 +354,14 @@ def _extra_repr(self): f"track_running_stats={self.track_running_stats}" ) - def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: + def _calc_stats(self, x: mx.array, ddof: int = 0) -> Tuple[mx.array, mx.array]: """ Calculate the mean and variance of the input tensor across the batch and spatial dimensions. Args: x (array): Input tensor. + ddof (int): Delta degrees of freedom for variance. Returns: tuple: Tuple containing mean and variance. @@ -329,7 +369,7 @@ def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]: reduction_axes = tuple(range(0, x.ndim - 1)) mean = mx.mean(x, axis=reduction_axes) - var = mx.var(x, axis=reduction_axes) + var = mx.var(x, axis=reduction_axes, ddof=ddof) return mean, var @@ -348,13 +388,23 @@ def __call__(self, x: mx.array) -> mx.array: f"Expected input tensor to have 2, 3 or 4 dimensions, but got {x.ndim}" ) + if self.training: + stats_size = 1 + for size in x.shape[:-1]: + stats_size *= size + if stats_size == 1: + raise ValueError( + "BatchNorm training requires more than one value per channel." + ) + # Calculate the mean and variance used to normalize the input x. If we # are in training mode update the running stats if needed. mean, var = self._calc_stats(x) if self.training and self.track_running_stats: mu = self.momentum + _, running_var = self._calc_stats(x, ddof=1) self.running_mean = (1 - mu) * self.running_mean + mu * mean - self.running_var = (1 - mu) * self.running_var + mu * var + self.running_var = (1 - mu) * self.running_var + mu * running_var elif self.track_running_stats: mean = self.running_mean var = self.running_var diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py b/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py index 203117634..eaf50228e 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/pooling.py @@ -85,6 +85,19 @@ class _Pool(Module): def __init__(self, pooling_function, kernel_size, stride, padding, padding_value): super().__init__() + class_name = type(self).__name__ + for name, values in (("kernel_size", kernel_size), ("stride", stride)): + if any(v <= 0 for v in values): + raise ValueError( + f"[{class_name}] '{name}' must be positive but got " + f"{tuple(values)}." + ) + if any(p[0] < 0 for p in padding): + raise ValueError( + f"[{class_name}] 'padding' must be non-negative but got " + f"{tuple(p[0] for p in padding)}." + ) + self._pooling_function = pooling_function self._kernel_size = kernel_size self._stride = stride diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py b/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py index 0d7632d80..49f300d4b 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/positional_encoding.py @@ -85,7 +85,19 @@ def __init__( ): super().__init__() - one_zero = 1 - mx.arange(0, dims // 2) / (dims // 2 - 1) + # Sinusoidal embeddings are constructed from sin/cos pairs, so the + # embedding dimension must be positive and even. + if dims <= 0 or dims % 2 != 0: + raise ValueError( + f"[SinusoidalPositionalEncoding] dims must be positive and even but got {dims}." + ) + + # Avoid division by zero when dims == 2 (one frequency pair). + if dims == 2: + one_zero = mx.array([1.0]) + else: + one_zero = 1 - mx.arange(0, dims // 2) / (dims // 2 - 1) + min_freq = math.log(min_freq) max_freq = math.log(max_freq) @@ -95,7 +107,7 @@ def __init__( self._sigmas = self._sigmas * (2 * math.pi) # Save some constants that define the implementation - self.scale = scale or (2 / dims) ** 0.5 + self.scale = scale if scale is not None else (2 / dims) ** 0.5 self.cos_first = cos_first def __call__(self, x): diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py b/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py index 15bbfa76d..7cebb5b2d 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/quantized.py @@ -15,6 +15,11 @@ def _defaults_for_mode(mode, group_size, bits): "nvfp4": (16, 4), "mxfp8": (32, 8), } + if mode not in mode_defaults: + raise ValueError( + f"Invalid quantization mode '{mode}'. " + f"Valid modes are: {', '.join(mode_defaults)}." + ) default_group_size, default_bits = mode_defaults[mode] return group_size or default_group_size, bits or default_bits @@ -25,13 +30,18 @@ def quantize( bits: int = None, *, mode: str = "affine", + quantize_input: bool = False, class_predicate: Optional[Callable[[str, Module], Union[bool, dict]]] = None, ): """Quantize the sub-modules of a module according to a predicate. - By default all layers that define a ``to_quantized(group_size, bits)`` - method will be quantized. Both :obj:`Linear` and :obj:`Embedding` layers - will be quantized. Note also, the module is updated in-place. + By default all layers that define a ``to_quantized()`` method will be + quantized. Both :obj:`Linear` and :obj:`Embedding` layers will be + quantized. The module is updated in-place. + + Note: + ``quantize_input=True`` is only supported for ``"nvfp4"`` and ``"mxfp8"`` + modes and :obj:`Linear` layers. Args: model (mlx.nn.Module): The model whose leaf modules may be quantized. @@ -41,12 +51,23 @@ def quantize( :func:`mlx.core.quantize`). Default: ``None``. mode (str): The quantization method to use (see :func:`mlx.core.quantize`). Default: ``"affine"``. + quantize_input (bool): Whether to quantize activations. Default: ``False``. class_predicate (Optional[Callable]): A callable which receives the - :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a - dict of params for `to_quantized` if it should be quantized and - ``False`` otherwise. If ``None``, then all layers that define a - ``to_quantized(group_size, bits)`` method are quantized. - Default: ``None``. + :obj:`Module` path and :obj:`Module` itself and returns ``True`` or a + dict of params for ``to_quantized`` if it should be quantized and + ``False`` otherwise. If ``None``, then all layers that define a + ``to_quantized()`` method are quantized. Default: ``None``. + + Example: + Weight only quantization for all layers that define a ``to_quantized()`` method: + + >>> import mlx.nn as nn + >>> nn.quantize(model, group_size=64, bits=4, mode="affine") + + Weight and input quantization for all linear layers: + + >>> predicate = lambda p, m: isinstance(m, nn.Linear) + >>> nn.quantize(model, mode="nvfp4", quantize_input=True, class_predicate=predicate) """ class_predicate = class_predicate or (lambda _, m: hasattr(m, "to_quantized")) @@ -54,8 +75,15 @@ def _maybe_quantize(path, m): if bool_or_params := class_predicate(path, m): if hasattr(m, "to_quantized"): if isinstance(bool_or_params, bool): - return m.to_quantized(group_size=group_size, bits=bits, mode=mode) + kwargs = {"group_size": group_size, "bits": bits, "mode": mode} + if quantize_input: + kwargs["quantize_input"] = quantize_input + return m.to_quantized(**kwargs) elif isinstance(bool_or_params, dict): + if ("quantize_input" in bool_or_params) and not bool_or_params[ + "quantize_input" + ]: + bool_or_params.pop("quantize_input") return m.to_quantized(**bool_or_params) else: raise ValueError( diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/recurrent.py b/Source/Cmlx/mlx/python/mlx/nn/layers/recurrent.py index 3ffa7654c..92a75e8d9 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/recurrent.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/recurrent.py @@ -33,7 +33,7 @@ class RNN(Module): hidden_size (int): Dimension of the hidden state, ``H``. bias (bool, optional): Whether to use a bias. Default: ``True``. nonlinearity (callable, optional): Non-linearity to use. If ``None``, - then func:`tanh` is used. Default: ``None``. + then :func:`tanh` is used. Default: ``None``. """ def __init__( @@ -184,6 +184,8 @@ def __call__(self, x, hidden=None): if hidden is not None: n = n + r * h_proj_n + elif self.bhn is not None: + n = n + r * self.bhn n = mx.tanh(n) if hidden is not None: diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py b/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py index d856f2554..31ec916d1 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/transformer.py @@ -228,7 +228,7 @@ def __call__(self, x, memory, x_mask, memory_mask): y = self.dropout1(y) x = self.ln1(x + y) - y = self.cross_attention(y, memory, memory, memory_mask) + y = self.cross_attention(x, memory, memory, memory_mask) y = self.dropout2(y) x = self.ln2(x + y) @@ -327,27 +327,33 @@ def __init__( ): super().__init__() - self.encoder = custom_encoder or TransformerEncoder( - num_encoder_layers, - dims, - num_heads, - mlp_dims, - dropout, - activation, - norm_first, - checkpoint, - ) + if custom_encoder is not None: + self.encoder = custom_encoder + else: + self.encoder = TransformerEncoder( + num_encoder_layers, + dims, + num_heads, + mlp_dims, + dropout, + activation, + norm_first, + checkpoint, + ) - self.decoder = custom_decoder or TransformerDecoder( - num_decoder_layers, - dims, - num_heads, - mlp_dims, - dropout, - activation, - norm_first, - checkpoint, - ) + if custom_decoder is not None: + self.decoder = custom_decoder + else: + self.decoder = TransformerDecoder( + num_decoder_layers, + dims, + num_heads, + mlp_dims, + dropout, + activation, + norm_first, + checkpoint, + ) def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask): memory = self.encoder(src, src_mask) diff --git a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py index e6bd282af..0be9f87bf 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py +++ b/Source/Cmlx/mlx/python/mlx/nn/layers/upsample.py @@ -3,6 +3,7 @@ import operator from functools import partial, reduce from itertools import product +from math import ceil from typing import Callable, Literal, Tuple, Union import mlx.core as mx @@ -12,7 +13,7 @@ def _scaled_indices(N, scale, align_corners, dim, ndims): M = int(scale * N) if align_corners: - indices = mx.arange(M, dtype=mx.float32) * ((N - 1) / (M - 1)) + indices = mx.arange(M, dtype=mx.float32) * ((N - 1) / max(M - 1, 1)) else: step = 1 / scale start = ((M - 1) * step - N + 1) / 2 @@ -51,6 +52,118 @@ def _linear_indices(N, scale, align_corners, dim, ndims): ) +def _aa_indices(N, scale, align_corners, dim, ndims, kernel_fn, kernel_radius): + """Compute antialiased interpolation indices for a given kernel. + + When downscaling (scale < 1), the kernel support widens by 1/scale to + act as a low-pass filter, preventing aliasing. Out-of-bounds taps are + zeroed and weights are renormalized per output pixel. + + For upscaling (scale >= 1), the kernel is applied at its native width + without widening. This matches PyTorch's F.interpolate(antialias=True) + behavior where the kernel coefficient (e.g. a=-0.5 for cubic) is used + for both up and downsampling when antialias=True. + + Args: + N: input size for this dimension + scale: scale factor for this dimension + align_corners: align_corners flag + dim: which spatial dimension + ndims: number of spatial dimensions + kernel_fn: callable(distance) -> weight, operating in normalized + filter coordinates where |distance| < kernel_radius has support + kernel_radius: support radius of the kernel in normalized coords + (1.0 for triangle/linear, 2.0 for cubic) + """ + indices = _scaled_indices(N, scale, align_corners, dim, ndims) + + # For downscale, widen the filter by 1/scale + if scale < 1: + inv_scale = 1.0 / scale + else: + inv_scale = 1.0 + + support = kernel_radius * inv_scale + num_taps = ceil(support) + 1 + + # Compute per-tap weights, zero out-of-bounds, then normalize + all_idx = [] + all_w = [] + for k in range(-num_taps + 1, num_taps): + idx = mx.floor(indices) + k + # Map distance to normalized filter coordinates + dist = mx.abs(indices - idx) / inv_scale + w = kernel_fn(dist) + # Zero out-of-bounds taps + w = mx.where((idx >= 0) & (idx < N), w, 0.0) + all_idx.append(idx) + all_w.append(w) + + # Normalize so weights sum to 1 per output pixel + w_sum = sum(all_w) + w_sum = mx.where(w_sum > 0, w_sum, 1.0) + + result = [] + for idx, w in zip(all_idx, all_w): + w = mx.expand_dims(w / w_sum, -1) + idx = mx.clip(idx, a_min=0, a_max=N - 1).astype(mx.uint32) + result.append((idx, w)) + + return tuple(result) + + +def _triangle_kernel(x): + """Triangle (linear) filter kernel. Support radius = 1.""" + return mx.maximum(1.0 - x, 0.0) + + +def _cubic_kernel(x): + """Keys cubic kernel with a=-0.5 (PIL/Pillow convention). + + This coefficient is used by PyTorch when antialias=True for both + bilinear and bicubic modes. The non-antialiased cubic path uses + a=-0.75 (OpenCV convention) -- see ``_cubic_indices``. + + Support radius = 2. + """ + a = -0.5 + w_inner = ((a + 2.0) * x - (a + 3.0)) * x * x + 1 + w_outer = (((x - 5) * x + 8) * x - 4) * a + return mx.where(x <= 1.0, w_inner, mx.where(x <= 2.0, w_outer, 0.0)) + + +def _linear_aa_indices(N, scale, align_corners, dim, ndims): + """Linear interpolation with antialiasing (triangle kernel).""" + return _aa_indices( + N, + scale, + align_corners, + dim, + ndims, + kernel_fn=_triangle_kernel, + kernel_radius=1.0, + ) + + +def _cubic_aa_indices(N, scale, align_corners, dim, ndims): + """Cubic interpolation with antialiasing (Keys cubic, a=-0.5). + + Note: the non-antialiased cubic path (``_cubic_indices``) uses a=-0.75 + (OpenCV convention). When ``antialias=True``, PyTorch switches to a=-0.5 + (PIL convention). This coefficient change affects the interpolant shape, + not just the filter width. See ``_cubic_kernel`` for details. + """ + return _aa_indices( + N, + scale, + align_corners, + dim, + ndims, + kernel_fn=_cubic_kernel, + kernel_radius=2.0, + ) + + def _cubic_indices(N, scale, align_corners, dim, ndims): indices = _scaled_indices(N, scale, align_corners, dim, ndims) indices_l1 = mx.floor(indices) @@ -60,8 +173,9 @@ def _cubic_indices(N, scale, align_corners, dim, ndims): @partial(mx.compile, shapeless=True) def _get_weight(ind, grid, dist): - # PyTorch uses -0.5 for antialiasing=true (compatibility with PIL) - # and uses -0.75 for antialiasing=false (compatibility with OpenCV) + # a=-0.75 (OpenCV convention) for non-antialiased cubic. + # When antialias=True, _cubic_aa_indices uses a=-0.5 (PIL convention) + # via _cubic_kernel instead. a = -0.75 x = mx.abs(ind - grid) if dist == 1: @@ -89,6 +203,14 @@ def _get_weight(ind, grid, dist): ) +def _validate_antialias_align_corners(align_corners, antialias): + if antialias and align_corners: + raise ValueError( + "[Upsample] antialias=True with align_corners=True is not " + "supported. Use align_corners=False for antialiased interpolation." + ) + + def upsample_nearest(x: mx.array, scale_factor: Tuple): dims = x.ndim - 2 if dims != len(scale_factor): @@ -145,7 +267,41 @@ def _interpolate( return sum(wi * xi for wi, xi in zip(weights, samples)) -def upsample_linear(x: mx.array, scale_factor: Tuple, align_corners: bool = False): +def _interpolate_separable( + x: mx.array, scale_factor: Tuple, indices_fn: Callable, align_corners: bool = False +): + dims = x.ndim - 2 + if dims != len(scale_factor): + raise ValueError("A scale needs to be provided for each spatial dimension") + + _, *N, _ = x.shape + out = x + + for i, (n, s) in enumerate(zip(N, scale_factor)): + axis = i + 1 + samples = [] + for idx, weight in indices_fn(n, s, align_corners, i, dims): + sample = mx.take(out, idx.reshape(-1), axis=axis) + samples.append(sample * weight) + out = sum(samples) + + return out + + +def upsample_linear( + x: mx.array, + scale_factor: Tuple, + align_corners: bool = False, + antialias: bool = False, +): + _validate_antialias_align_corners(align_corners, antialias) + if antialias: + return _interpolate_separable( + x=x, + scale_factor=scale_factor, + indices_fn=_linear_aa_indices, + align_corners=align_corners, + ) return _interpolate( x=x, scale_factor=scale_factor, @@ -154,7 +310,20 @@ def upsample_linear(x: mx.array, scale_factor: Tuple, align_corners: bool = Fals ) -def upsample_cubic(x: mx.array, scale_factor: Tuple, align_corners: bool = False): +def upsample_cubic( + x: mx.array, + scale_factor: Tuple, + align_corners: bool = False, + antialias: bool = False, +): + _validate_antialias_align_corners(align_corners, antialias) + if antialias: + return _interpolate_separable( + x=x, + scale_factor=scale_factor, + indices_fn=_cubic_aa_indices, + align_corners=align_corners, + ) return _interpolate( x=x, scale_factor=scale_factor, @@ -185,6 +354,23 @@ class Upsample(Module): ``align_corners=True`` then the top and left edge of the input and output will be matching as will the bottom right edge. + .. note:: + When ``antialias=True`` is used with ``"linear"`` or ``"cubic"`` mode, + an antialiased filter is applied during downsampling (scale factor < 1), + producing smoother results by avoiding aliasing artifacts. For 2D + integer-ratio downscales with ``align_corners=False``, this matches the + behavior of PyTorch's ``F.interpolate(antialias=True)``. Non-integer + scale factors are supported but may differ from PyTorch because of + existing index-selection differences. + + For ``"cubic"`` mode, enabling ``antialias`` also changes the cubic + kernel coefficient from ``a=-0.75`` (OpenCV convention) to ``a=-0.5`` + (PIL/Pillow convention), matching PyTorch's behavior. This affects the + interpolant shape, not just the filter width. + + ``antialias=True`` with ``align_corners=True`` is not supported and + will raise a ``ValueError``. + Parameters: scale_factor (float or tuple): The multiplier for the spatial size. If a ``float`` is provided, it is the multiplier for all spatial dimensions. @@ -195,6 +381,11 @@ class Upsample(Module): align_corners (bool, optional): Changes the way the corners are treated during ``"linear"`` and ``"cubic"`` upsampling. See the note above and the examples below for more details. Default: ``False``. + antialias (bool, optional): If ``True``, apply an antialiasing filter + when downsampling with ``"linear"`` or ``"cubic"`` mode. For + ``"cubic"`` mode this also switches the kernel coefficient to + ``a=-0.5``. Not supported with ``"nearest"`` mode or with + ``align_corners=True``. Default: ``False``. Examples: >>> import mlx.core as mx @@ -230,22 +421,35 @@ def __init__( scale_factor: Union[float, Tuple], mode: Literal["nearest", "linear", "cubic"] = "nearest", align_corners: bool = False, + antialias: bool = False, ): super().__init__() if mode not in ["nearest", "linear", "cubic"]: raise ValueError(f"[Upsample] Got unsupported upsampling algorithm: {mode}") + if antialias and mode == "nearest": + raise ValueError( + "[Upsample] Antialiasing is not supported for nearest neighbor upsampling" + ) if isinstance(scale_factor, (list, tuple)): - self.scale_factor = tuple(map(float, scale_factor)) + scale_factor = tuple(map(float, scale_factor)) else: - self.scale_factor = float(scale_factor) + scale_factor = float(scale_factor) + + _validate_antialias_align_corners(align_corners, antialias) + + self.scale_factor = scale_factor self.mode = mode self.align_corners = align_corners + self.antialias = antialias def _extra_repr(self) -> str: - return ( + repr_str = ( f"scale_factor={self.scale_factor}, mode={self.mode!r}, " f"align_corners={self.align_corners}" ) + if self.antialias: + repr_str += ", antialias=True" + return repr_str def __call__(self, x: mx.array) -> mx.array: dims = x.ndim - 2 @@ -270,8 +474,8 @@ def __call__(self, x: mx.array) -> mx.array: if self.mode == "nearest": return upsample_nearest(x, scale_factor) elif self.mode == "linear": - return upsample_linear(x, scale_factor, self.align_corners) + return upsample_linear(x, scale_factor, self.align_corners, self.antialias) elif self.mode == "cubic": - return upsample_cubic(x, scale_factor, self.align_corners) + return upsample_cubic(x, scale_factor, self.align_corners, self.antialias) else: raise Exception(f"Unknown interpolation mode: {self.mode}") diff --git a/Source/Cmlx/mlx/python/mlx/nn/losses.py b/Source/Cmlx/mlx/python/mlx/nn/losses.py index b765d25de..b98d2765d 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/losses.py +++ b/Source/Cmlx/mlx/python/mlx/nn/losses.py @@ -63,9 +63,21 @@ def cross_entropy( >>> targets = mx.array([[0.9, 0.1], [0.1, 0.9]]) >>> nn.losses.cross_entropy(logits, targets) array([0.348587, 0.348587], dtype=float32) + >>> + >>> # Half precision logits with class indices as targets. On CUDA a + >>> # fused kernel accumulates the reduction in float32: + >>> logits = mx.array([[2.0, -1.0], [-1.0, 2.0]], mx.bfloat16) + >>> targets = mx.array([0, 1]) + >>> nn.losses.cross_entropy(logits, targets) + array([0.0485873, 0.0485873], dtype=float32) + >>> + >>> # Metal and the CPU reduce in the dtype of the logits, so upcast + >>> # them to get the same accuracy: + >>> nn.losses.cross_entropy(logits.astype(mx.float32), targets) + array([0.0485873, 0.0485873], dtype=float32) """ if label_smoothing < 0 or label_smoothing >= 1: - raise ValueError(f"Label smoothing must in [0, 1), got {label_smoothing}.") + raise ValueError(f"Label smoothing must be in [0, 1), got {label_smoothing}.") # Whether targets are class indices or probabilities targets_as_probs = targets.ndim == logits.ndim @@ -83,26 +95,38 @@ def _drop_dim(shape, axis): f"Targets shape {targets.shape} does not match logits shape {logits.shape}." ) - if targets_as_probs: - score = mx.sum(logits * targets, axis=axis) + use_fast = ( + mx.cuda.is_available() + and mx.default_device() == mx.gpu + and not targets_as_probs + and label_smoothing == 0 + and axis in (-1, logits.ndim - 1) + and mx.issubdtype(logits.dtype, mx.floating) + and mx.issubdtype(targets.dtype, mx.integer) + ) + + if use_fast: + loss = mx.fast.cross_entropy(logits, targets).astype(logits.dtype) else: - score = mx.take_along_axis(logits, mx.expand_dims(targets, axis), axis).squeeze( - axis - ) + logits = logits - mx.stop_gradient(mx.max(logits, axis=axis, keepdims=True)) - logsumexp_logits = mx.logsumexp(logits, axis=axis) - if label_smoothing > 0: - # Adjust the true class score with label smoothing - adjusted_score = (1 - label_smoothing) * score + if targets_as_probs: + score = mx.sum(logits * targets, axis=axis) + else: + score = mx.take_along_axis( + logits, mx.expand_dims(targets, axis), axis + ).squeeze(axis) - # Calculate the mean logit across the classes for smoothed loss - mean_logits = logits.mean(axis=axis) - smoothed_loss = -mean_logits * label_smoothing + logsumexp_logits = mx.logsumexp(logits, axis=axis) + if label_smoothing > 0: + adjusted_score = (1 - label_smoothing) * score - # Combine the adjusted score and smoothed loss with the logsumexp logits - loss = logsumexp_logits - adjusted_score + smoothed_loss - else: - loss = logsumexp_logits - score + mean_logits = logits.mean(axis=axis) + smoothed_loss = -mean_logits * label_smoothing + + loss = logsumexp_logits - adjusted_score + smoothed_loss + else: + loss = logsumexp_logits - score # Apply weights if provided if weights is not None: @@ -136,8 +160,8 @@ def binary_cross_entropy( inputs (array): The predicted values. If ``with_logits`` is ``True``, then ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities. targets (array): The binary target values in {0, 1}. - with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``. weights (array, optional): Optional weights for each target. Default: ``None``. + with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. @@ -249,7 +273,9 @@ def nll_loss( Returns: array: The computed NLL loss. """ - loss = -mx.take_along_axis(inputs, targets[..., None], axis).squeeze(-1) + loss = -mx.take_along_axis(inputs, mx.expand_dims(targets, axis), axis).squeeze( + axis + ) return _reduce(loss, reduction) @@ -272,7 +298,7 @@ def gaussian_nll_loss( \ \epsilon\right)\right) + \frac{\left(\text{inputs} - \text{targets} \right)^2} {\max\left(\text{vars}, \ \epsilon \right)}\right) + \text{const.} - where ``inputs`` are the predicted means and ``vars`` are the the + where ``inputs`` are the predicted means and ``vars`` are the predicted variances. Args: @@ -284,7 +310,7 @@ def gaussian_nll_loss( eps (float, optional): Small positive constant for numerical stability. Default: ``1e-6``. reduction (str, optional): Specifies the reduction to apply to the output: - ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. + ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. Returns: array: The Gaussian NLL loss. @@ -319,7 +345,7 @@ def kl_div_loss( .. code-block:: python - mx.exp(targets) * (targets - inputs).sum(axis) + (mx.exp(targets) * (targets - inputs)).sum(axis) Args: inputs (array): Log probabilities for the predicted distribution. @@ -408,20 +434,23 @@ def triplet_loss( axis (int, optional): The distribution axis. Default: ``-1``. p (int, optional): The norm degree for pairwise distance. Default: ``2``. margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``. - eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``. + eps (float, optional): Small positive constant added to the p-norm sum + before taking the ``1 / p`` power. Defaults to ``1e-6``. reduction (str, optional): Specifies the reduction to apply to the output: ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: - array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input; - if reduction is "mean" or "sum", returns a scalar tensor. + array: Computed triplet loss. If reduction is ``"none"``, returns a tensor with the + same shape as the inputs but with the ``axis`` dimension removed; if reduction + is ``"mean"`` or ``"sum"``, returns a scalar tensor. """ - loss = mx.maximum( - mx.sqrt(mx.power(anchors - positives, p).sum(axis) + eps) - - mx.sqrt(mx.power(anchors - negatives, p).sum(axis) + eps) - + margin, - 0, + pos_dist = mx.power( + mx.power(mx.abs(anchors - positives), p).sum(axis) + eps, 1.0 / p + ) + neg_dist = mx.power( + mx.power(mx.abs(anchors - negatives), p).sum(axis) + eps, 1.0 / p ) + loss = mx.maximum(pos_dist - neg_dist + margin, 0) return _reduce(loss, reduction) @@ -500,8 +529,7 @@ def log_cosh_loss( .. math:: \text{logcosh}(y_{\text{true}}, y_{\text{pred}}) = - \frac{1}{n} \sum_{i=1}^{n} - \log(\cosh(y_{\text{pred}}^{(i)} - y_{\text{true}}^{(i)})) + \log(\cosh(y_{\text{pred}} - y_{\text{true}})) Args: @@ -536,8 +564,8 @@ def cosine_similarity_loss( \frac{x_1 \cdot x_2}{\max(\|x_1\| \cdot \|x_2\|, \epsilon)} Args: - x1 (mx.array): The first set of inputs. - x2 (mx.array): The second set of inputs. + x1 (array): The first set of inputs. + x2 (array): The second set of inputs. axis (int, optional): The embedding axis. Default: ``1``. eps (float, optional): The minimum value of the denominator used for numerical stability. Default: ``1e-8``. @@ -545,7 +573,7 @@ def cosine_similarity_loss( ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``. Returns: - mx.array: The computed cosine similarity loss. + array: The computed cosine similarity loss. """ x1_norm = mx.linalg.norm(x1, axis=axis) x2_norm = mx.linalg.norm(x2, axis=axis) @@ -595,7 +623,7 @@ def margin_ranking_loss( >>> inputs2 = mx.array([0.75596, 0.225763, 0.256995]) >>> loss = nn.losses.margin_ranking_loss(inputs1, inputs2, targets) >>> loss - array(0.773433, dtype=float32) + array([1.32937, 0.990929, 0], dtype=float32) """ if not (inputs1.shape == inputs2.shape == targets.shape): raise ValueError( diff --git a/Source/Cmlx/mlx/python/mlx/nn/utils.py b/Source/Cmlx/mlx/python/mlx/nn/utils.py index 97354d112..a60468db6 100644 --- a/Source/Cmlx/mlx/python/mlx/nn/utils.py +++ b/Source/Cmlx/mlx/python/mlx/nn/utils.py @@ -5,7 +5,7 @@ import mlx.core as mx -from ..utils import tree_flatten, tree_map, tree_unflatten +from ..utils import tree_flatten, tree_map, tree_reduce, tree_unflatten from .layers.base import Module @@ -71,11 +71,35 @@ def wrapped_checkpointed_fn(*args, **kwargs): return wrapped_checkpointed_fn +def _extract_info(flat): + keys = [k for k, _ in flat] + shapes = [g.shape for _, g in flat] + sizes = [g.size for _, g in flat] + dtypes = [g.dtype for _, g in flat] + return keys, shapes, sizes, dtypes + + +def _group_by_size(keys, sizes, itemsize, communication_size): + grad_groups = [] + grad_group = [] + grad_group_size = 0 + for i in range(len(keys)): + grad_group.append(i) + grad_group_size += sizes[i] * itemsize + if grad_group_size >= communication_size: + grad_groups.append(grad_group) + grad_group = [] + grad_group_size = 0 + if grad_group: + grad_groups.append(grad_group) + grad_group = [] + return grad_groups + + def average_gradients( gradients: Any, group: Optional[mx.distributed.Group] = None, all_reduce_size: int = 32 * 1024**2, - communication_type: Optional[mx.Dtype] = None, communication_stream: Optional[mx.Stream] = None, ): """Average the gradients across the distributed processes in the passed group. @@ -92,10 +116,7 @@ def average_gradients( all_reduce_size (int): Group arrays until their size in bytes exceeds this number. Perform one communication step per group of arrays. If less or equal to 0 array grouping is disabled. Default: ``32MiB``. - communication_type (Optional[mlx.core.Dtype]): If provided cast to this - type before performing the communication. Typically cast to a - smaller float to reduce the communication size. Default: ``None``. - communication_stream (Optional[mlx.core.Stream]): The stream to usse + communication_stream (Optional[mlx.core.Stream]): The stream to use for the communication. If unspecified the default communication stream is used which can vary by back-end. Default: ``None``. """ @@ -105,13 +126,16 @@ def average_gradients( if N == 1: return gradients - def _average(x): - dt = x.dtype - x = x.astype(communication_type) if communication_type is not None else x - return mx.distributed.all_sum(x, stream=communication_stream).astype(dt) / N - if all_reduce_size <= 0: - return tree_map(_average, gradients) + return tree_map( + lambda x: mx.distributed.all_sum( + x, + group=group, + stream=communication_stream, + ) + / N, + gradients, + ) else: flat_grads = tree_flatten(gradients) @@ -119,34 +143,13 @@ def _average(x): return gradients # Extract some info for the gradient - keys = [k for k, _ in flat_grads] - shapes = [v.shape for _, v in flat_grads] - sizes = [v.size for _, v in flat_grads] - dtypes = [v.dtype for _, v in flat_grads] + keys, shapes, sizes, dtypes = _extract_info(flat_grads) # We can't group them if they have mixed types if not all(dt == dtypes[0] for dt in dtypes): - return average_gradients(gradients, group, 0, communication_type) - itemsize = ( - communication_type.size - if communication_type is not None - else dtypes[0].size - ) - + return average_gradients(gradients, group, 0) # Gather the gradients in groups that are just above or equal to all_reduce_size - grad_groups = [] - grad_group = [] - grad_group_size = 0 - for i in range(len(keys)): - grad_group.append(i) - grad_group_size += sizes[i] * itemsize - if grad_group_size >= all_reduce_size: - grad_groups.append(grad_group) - grad_group = [] - grad_group_size = 0 - if grad_group: - grad_groups.append(grad_group) - grad_group = [] + grad_groups = _group_by_size(keys, sizes, dtypes[0].size, all_reduce_size) # Concatenate-reduce-split new_flat_grads = [] @@ -155,7 +158,12 @@ def _average(x): big_grad = mx.concatenate( [flat_grads[i][1].reshape(-1) for i in grad_group] ) - big_grad = _average(big_grad) + big_grad = ( + mx.distributed.all_sum( + big_grad, stream=communication_stream, group=group + ) + / N + ) big_grad = mx.split(big_grad, indices[1:-1]) new_flat_grads.extend( (keys[j], big_grad[i].reshape(shapes[j])) @@ -163,3 +171,38 @@ def _average(x): ) return tree_unflatten(new_flat_grads) + + +def clip_grad_norm_sharded( + gradients: Any, + max_norm: float, + group: Optional[mx.distributed.Group] = None, +): + """Clip the global norm of gradients that are sharded across a group. + + This is the sharded equivalent of + :func:`mlx.optimizers.clip_grad_norm`. Each member of the group holds only + a shard of the gradients, so the global norm is computed by summing the + local squared norms across the group before rescaling. It is useful for + clipping the gradients of a module wrapped with :func:`mlx.nn.fully_shard`. + + Args: + gradients (Any): A Python tree containing the local shard of the + gradient arrays. + max_norm (float): The maximum allowed global norm of the gradients. + group (Optional[mlx.core.distributed.Group]): The group across which + the gradients are sharded. If set to ``None`` the global group is + used. Default: ``None``. + + Returns: + (Any, mlx.core.array): The possibly rescaled local shard of the + gradients and the global gradient norm. + """ + local_norm_squared = tree_reduce( + lambda acc, g: acc + g.square().sum(), gradients, 0.0 + ) + global_norm_squared = mx.distributed.all_sum(local_norm_squared, group=group) + grad_norm = mx.sqrt(global_norm_squared) + normalizer = mx.minimum(max_norm / (grad_norm + 1e-6), 1.0) + clipped_gradients = tree_map(lambda g: g * normalizer, gradients) + return clipped_gradients, grad_norm diff --git a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py index 2cc9e26b1..8be344247 100644 --- a/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py +++ b/Source/Cmlx/mlx/python/mlx/optimizers/optimizers.py @@ -39,7 +39,7 @@ def init(self, parameters: dict): :meth:`Optimizer.update`. Args: - model (dict): A Python tree of parameters. + parameters (dict): A Python tree of parameters. Example: >>> optimizer = optim.SGD(learning_rate=1e-1, momentum=0.9) @@ -77,7 +77,7 @@ def init_single(self, parameter: mx.array, state: dict): state initialization. Args: - parameter (mx.array): A single parameter that will be optimized. + parameter (array): A single parameter that will be optimized. state (dict): The optimizer's state. """ raise NotImplementedError() @@ -112,8 +112,8 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): """To be extended by derived classes to implement the optimizer's update. Args: - gradient (mx.array): The ``parameter`` gradient. - parameter (mx.array): The ``parameter`` to update. + gradient (array): The ``parameter`` gradient. + parameter (array): The ``parameter`` to update. state (dict): The optimizer's state. """ raise NotImplementedError() @@ -165,7 +165,7 @@ class MultiOptimizer(Optimizer): Args: optimizers (list[Optimizer]): A list of optimizers to delegate to - filters (list[Callable[[str, array], bool]): A list of predicates that + filters (list[Callable[[str, array], bool]]): A list of predicates that should be one less than the provided optimizers. """ @@ -328,7 +328,7 @@ def __init__( raise ValueError( f"RMSprop alpha should be >=0, {self.alpha} was provided instead" ) - if self.eps < 0.0: + if self.eps <= 0.0: raise ValueError( f"RMSprop epsilon should be >0, {self.eps} was provided instead" ) @@ -379,7 +379,7 @@ def __init__( self._maybe_schedule("learning_rate", learning_rate) self.eps = eps - if self.eps < 0.0: + if self.eps <= 0.0: raise ValueError( f"Adagrad epsilon should be >0, {self.eps} was provided instead" ) @@ -419,7 +419,7 @@ class AdaDelta(Optimizer): rho (float, optional): The coefficient :math:`\rho` used for computing a running average of squared gradients. Default: ``0.9`` eps (float, optional): The term :math:`\epsilon` added to the denominator to improve - numerical stability. Default: `1e-8` + numerical stability. Default: ``1e-6`` """ def __init__( @@ -437,7 +437,7 @@ def __init__( raise ValueError( f"AdaDelta rho should be >=0, {self.rho} was provided instead" ) - if self.eps < 0.0: + if self.eps <= 0.0: raise ValueError( f"AdaDelta epsilon should be >0, {self.eps} was provided instead" ) @@ -499,6 +499,16 @@ def __init__( ): super().__init__() + for i, beta in enumerate(betas): + if not 0.0 <= beta < 1.0: + raise ValueError( + f"Adam beta{i + 1} should be in [0, 1), {beta} was provided " + "instead" + ) + + if not 0.0 <= eps: + raise ValueError(f"Adam epsilon should be >=0, {eps} was provided instead") + self._maybe_schedule("learning_rate", learning_rate) self.betas = betas self.eps = eps @@ -606,8 +616,9 @@ class Adamax(Adam): Args: learning_rate (float or callable): The learning rate :math:`\lambda`. betas (Tuple[float, float], optional): The coefficients - :math:`(\beta_1, \beta_2)` used for computing running averages of the - gradient and its square. Default: ``(0.9, 0.999)`` + :math:`(\beta_1, \beta_2)` used for computing the running average of + the gradient and the exponentially weighted infinity norm. + Default: ``(0.9, 0.999)`` eps (float, optional): The term :math:`\epsilon` added to the denominator to improve numerical stability. Default: ``1e-8`` """ @@ -619,10 +630,6 @@ def __init__( eps: float = 1e-8, ): super().__init__(learning_rate, betas, eps) - if not 0.0 <= eps: - raise ValueError( - f"Epsilon value should be >=0, {self.eps} was provided instead" - ) def init_single(self, parameter: mx.array, state: dict): """Initialize optimizer state""" @@ -682,6 +689,13 @@ def __init__( ): super().__init__() + for i, beta in enumerate(betas): + if not 0.0 <= beta < 1.0: + raise ValueError( + f"Lion beta{i + 1} should be in [0, 1), {beta} was provided " + "instead" + ) + self._maybe_schedule("learning_rate", learning_rate) self.betas = betas self.weight_decay = weight_decay @@ -797,9 +811,7 @@ def _approximate_exp_moving_avg(self, exp_avg_sq_row, exp_avg_sq_col): exp_avg_sq_row / mx.mean(exp_avg_sq_row, axis=-1, keepdims=True) ) c_factor = mx.rsqrt(exp_avg_sq_col) - return mx.matmul( - mx.expand_dims(r_factor, axis=-1), mx.expand_dims(c_factor, axis=0) - ) + return mx.expand_dims(r_factor, axis=-1) * mx.expand_dims(c_factor, axis=-2) def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict): """Performs the Adafactor parameter and state update.""" @@ -969,6 +981,9 @@ def clip_grad_norm(grads, max_norm): (dict, float): The possibly rescaled gradients and the original gradient norm. """ + if max_norm < 0: + raise ValueError(f"max_norm should be >=0, {max_norm} was provided instead") + norm_squared = tree_reduce(lambda acc, g: acc + g.square().sum(), grads, 0.0) total_norm = mx.sqrt(norm_squared) normalizer = mx.minimum(max_norm / (total_norm + 1e-6), 1.0) diff --git a/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py b/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py index 67e4e29cd..99712097c 100644 --- a/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py +++ b/Source/Cmlx/mlx/python/mlx/optimizers/schedulers.py @@ -51,6 +51,8 @@ def step_decay(init: float, decay_rate: float, step_size: int) -> Callable: >>> optimizer.learning_rate array(0.081, dtype=float32) """ + if step_size < 1: + raise ValueError(f"step_size must be greater than 0, but got {step_size}.") def schedule(step): return init * (decay_rate ** (step // step_size)) @@ -79,6 +81,8 @@ def cosine_decay(init: float, decay_steps: int, end: float = 0.0) -> Callable: >>> optimizer.learning_rate array(0.0999961, dtype=float32) """ + if decay_steps < 1: + raise ValueError(f"decay_steps must be greater than 0, but got {decay_steps}.") def schedule(step): s = mx.minimum(step, decay_steps) diff --git a/Source/Cmlx/mlx/python/mlx/utils.py b/Source/Cmlx/mlx/python/mlx/utils.py index f4aafe1e3..f8595ef9a 100644 --- a/Source/Cmlx/mlx/python/mlx/utils.py +++ b/Source/Cmlx/mlx/python/mlx/utils.py @@ -1,4 +1,5 @@ # Copyright © 2023 Apple Inc. + from collections import defaultdict from itertools import zip_longest from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -96,12 +97,13 @@ def tree_map_with_path( elif isinstance(tree, (list, tuple)): prefix = f"{path}." if path else "" TreeType = type(tree) - return TreeType( + subtrees = ( tree_map_with_path( fn, child, *(r[i] for r in rest), is_leaf=is_leaf, path=f"{prefix}{i}" ) for i, child in enumerate(tree) ) + return TreeType(*subtrees) if hasattr(tree, "_fields") else TreeType(subtrees) elif isinstance(tree, dict): prefix = f"{path}." if path else "" return { @@ -227,16 +229,24 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any: next_idx = "" if not next_idx else next_idx[0] children[current_idx].append((next_idx, value)) - # Assume they are a list and fail to dict if the keys are not all integers + # Assume list when all keys are integers. try: - keys = sorted((int(idx), idx) for idx in children.keys()) + keys = {} + for idx in children: + keys[int(idx)] = idx + # Guard against "01" and "1" treated as one key. + is_list = len(keys) == len(children) + except ValueError: + is_list = False + + if is_list: l = [] - for i, k in keys: + for i, k in sorted(keys.items()): # if i <= len(l), no {} will be appended. l.extend([{} for _ in range(i - len(l))]) l.append(tree_unflatten(children[k])) return l - except ValueError: + else: return {k: tree_unflatten(v) for k, v in children.items()} @@ -284,7 +294,8 @@ def tree_reduce(fn, tree, initializer=None, is_leaf=None): def tree_merge(tree_a, tree_b, merge_fn=None): """Merge two Python trees in one containing the values of both. It can be - thought of as a deep dict.update method. + thought of as a deep dict.update method. Empty containers are treated as + empty subtrees. Args: tree_a (Any): The first Python tree. @@ -295,9 +306,19 @@ def tree_merge(tree_a, tree_b, merge_fn=None): The Python tree containing the values of both ``tree_a`` and ``tree_b``. """ - if isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0: + empty_a = isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0 + empty_b = isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0 + + if empty_a and empty_b: + if type(tree_a) is not type(tree_b): + raise ValueError( + f"Cannot merge {type(tree_a).__name__} with {type(tree_b).__name__}" + ) + return type(tree_a)() + + if empty_a: tree_a = None - if isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0: + if empty_b: tree_b = None if tree_a is None and tree_b is not None: return tree_b @@ -306,9 +327,8 @@ def tree_merge(tree_a, tree_b, merge_fn=None): if isinstance(tree_a, (list, tuple)) and isinstance(tree_b, (list, tuple)): TreeType = type(tree_a) - return TreeType( - tree_merge(a, b, merge_fn) for a, b in zip_longest(tree_a, tree_b) - ) + subtrees = (tree_merge(a, b, merge_fn) for a, b in zip_longest(tree_a, tree_b)) + return TreeType(*subtrees) if hasattr(tree_a, "_fields") else TreeType(subtrees) elif isinstance(tree_a, dict) and isinstance(tree_b, dict): return { k: tree_merge(tree_a.get(k, None), tree_b.get(k, None), merge_fn) diff --git a/Source/Cmlx/mlx/python/src/CMakeLists.txt b/Source/Cmlx/mlx/python/src/CMakeLists.txt index 69152f502..0798add41 100644 --- a/Source/Cmlx/mlx/python/src/CMakeLists.txt +++ b/Source/Cmlx/mlx/python/src/CMakeLists.txt @@ -2,6 +2,7 @@ nanobind_add_module( core NB_STATIC STABLE_ABI + FREE_THREADED LTO NOMINSIZE NB_DOMAIN @@ -27,7 +28,8 @@ nanobind_add_module( ${CMAKE_CURRENT_SOURCE_DIR}/linalg.cpp ${CMAKE_CURRENT_SOURCE_DIR}/constants.cpp ${CMAKE_CURRENT_SOURCE_DIR}/trees.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp) + ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/print.cpp) if(MLX_BUILD_PYTHON_STUBS) nanobind_add_stub( diff --git a/Source/Cmlx/mlx/python/src/array.cpp b/Source/Cmlx/mlx/python/src/array.cpp index 231474c2d..fbab84618 100644 --- a/Source/Cmlx/mlx/python/src/array.cpp +++ b/Source/Cmlx/mlx/python/src/array.cpp @@ -2,16 +2,19 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include #include "mlx/backend/metal/metal.h" +#include "mlx/utils.h" #include "python/src/buffer.h" #include "python/src/convert.h" #include "python/src/indexing.h" @@ -97,9 +100,6 @@ class ArrayPythonIterator { }; void init_array(nb::module_& m) { - // Set Python print formatting options - mx::get_global_formatter().capitalize_bool = true; - // Types nb::class_( m, @@ -202,6 +202,10 @@ void init_array(nb::module_& m) { Get information on floating-point types. )pbdoc") .def(nb::init()) + .def_ro( + "bits", + &mx::finfo::bits, + R"pbdoc(The number of bits occupied by the type.)pbdoc") .def_ro( "min", &mx::finfo::min, @@ -217,6 +221,10 @@ void init_array(nb::module_& m) { The difference between 1.0 and the next smallest representable number larger than 1.0. )pbdoc") + .def_ro( + "smallest_normal", + &mx::finfo::smallest_normal, + R"pbdoc(The smallest positive normal number.)pbdoc") .def_ro("dtype", &mx::finfo::dtype, R"pbdoc(The :obj:`Dtype`.)pbdoc") .def("__repr__", [](const mx::finfo& f) { std::ostringstream os; @@ -255,7 +263,8 @@ void init_array(nb::module_& m) { "ArrayAt", R"pbdoc( A helper object to apply updates at specific indices. - )pbdoc") + )pbdoc", + nb::pooled(/* capacity = */ 128)) .def("__getitem__", &ArrayAt::set_indices, "indices"_a.none()) .def("add", &ArrayAt::add, "value"_a) .def("subtract", &ArrayAt::subtract, "value"_a) @@ -293,16 +302,17 @@ void init_array(nb::module_& m) { "array", R"pbdoc(An N-dimensional array object.)pbdoc", nb::type_slots(array_slots), - nb::is_weak_referenceable()) + nb::is_weak_referenceable(), + nb::pooled(/* capacity = */ 128)) .def( "__init__", - [](mx::array* aptr, ArrayInitType v, std::optional t) { - new (aptr) mx::array(create_array(v, t)); + [](mx::array* aptr, nb::object v, std::optional t) { + new (aptr) mx::array(create_array(v, t, true)); }, "val"_a, "dtype"_a = nb::none(), nb::sig( - "def __init__(self: array, val: Union[scalar, list, tuple, numpy.ndarray, array], dtype: Optional[Dtype] = None)")) + "def __init__(self: array, val: scalar | list | tuple | DLPackCompatible | array, dtype: Dtype | None = None)")) .def_prop_ro( "size", &mx::array::size, @@ -375,7 +385,9 @@ void init_array(nb::module_& m) { )pbdoc") .def( "astype", - &mx::astype, + [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); + }, "dtype"_a, "stream"_a = nb::none(), R"pbdoc( @@ -481,7 +493,7 @@ void init_array(nb::module_& m) { throw std::invalid_argument( "Invalid pickle state: expected (ndarray, Dtype::Val)"); } - using ND = nb::ndarray; + using ND = nb::ndarray; ND nd = nb::cast(state[0]); auto val = static_cast(nb::cast(state[1])); if (val == mx::Dtype::Val::bfloat16) { @@ -492,13 +504,26 @@ void init_array(nb::module_& m) { reinterpret_cast(nd.shape_ptr()), owner, nullptr, - nb::bfloat16), + nb::dtype()), mx::bfloat16)); } else { new (&arr) mx::array(nd_array_to_mlx(nd, std::nullopt)); } }) - .def("__dlpack__", [](const mx::array& a) { return mlx_to_dlpack(a); }) + .def( + "__dlpack__", + [](const mx::array& a, + nb::object, + nb::object, + std::optional> dl_device, + std::optional copy) { + return mlx_to_dlpack(a, copy.value_or(false), dl_device); + }, + nb::kw_only(), + "stream"_a = nb::none(), + "max_version"_a = nb::none(), + "dl_device"_a = nb::none(), + "copy"_a = nb::none()) .def( "__dlpack_device__", [](const mx::array& a) { @@ -506,13 +531,18 @@ void init_array(nb::module_& m) { // https://github.com/dmlc/dlpack/blob/5c210da409e7f1e51ddf445134a4376fdbd70d7d/include/dlpack/dlpack.h#L74 if (mx::metal::is_available()) { return nb::make_tuple(8, 0); - } else if (mx::cu::is_available()) { - return nb::make_tuple(13, 0); } else { // CPU device return nb::make_tuple(1, 0); } }) + .def( + "__array__", + [](const mx::array& self, nb::object dtype, nb::object copy) { + return mlx_to_np_array(self); + }, + "dtype"_a = nb::none(), + "copy"_a = nb::none()) .def("__copy__", [](const mx::array& self) { return mx::array(self); }) .def( "__deepcopy__", @@ -999,6 +1029,11 @@ void init_array(nb::module_& m) { nb::rv_policy::none) .def("__int__", [](mx::array& a) { return nb::int_(to_scalar(a)); }) .def("__float__", [](mx::array& a) { return nb::float_(to_scalar(a)); }) + .def( + "__complex__", + [](mx::array& a) { + return nb::cast>(to_scalar(a)); + }) .def( "__format__", [](mx::array& a, nb::object format_spec) { diff --git a/Source/Cmlx/mlx/python/src/buffer.h b/Source/Cmlx/mlx/python/src/buffer.h index 272a91888..4b194b3d2 100644 --- a/Source/Cmlx/mlx/python/src/buffer.h +++ b/Source/Cmlx/mlx/python/src/buffer.h @@ -43,7 +43,7 @@ std::string buffer_format(const mx::array& a) { case mx::float32: return "f"; case mx::bfloat16: - return "B"; + return "bfloat16"; case mx::float64: return "d"; case mx::complex64: diff --git a/Source/Cmlx/mlx/python/src/convert.cpp b/Source/Cmlx/mlx/python/src/convert.cpp index e004ac9cf..a3da76ef5 100644 --- a/Source/Cmlx/mlx/python/src/convert.cpp +++ b/Source/Cmlx/mlx/python/src/convert.cpp @@ -1,12 +1,29 @@ // Copyright © 2024 Apple Inc. +#include +#include +#include +#include + #include +#include #include "python/src/convert.h" #include "python/src/utils.h" +#include "mlx/allocator.h" +#include "mlx/backend/common/utils.h" +#include "mlx/backend/cuda/cuda.h" +#include "mlx/backend/metal/metal.h" +#include "mlx/dtype_utils.h" +#include "mlx/ops.h" #include "mlx/utils.h" +// Defined in ops.cpp. +namespace mlx::core { +array astype(array a, Dtype dtype, bool force_copy, StreamOrDevice s = {}); +} + enum PyScalarT { pybool = 0, pyint = 1, @@ -14,158 +31,360 @@ enum PyScalarT { pycomplex = 3, }; -namespace nanobind { -template <> -struct ndarray_traits { - static constexpr bool is_complex = false; - static constexpr bool is_float = true; - static constexpr bool is_bool = false; - static constexpr bool is_int = false; - static constexpr bool is_signed = true; -}; -}; // namespace nanobind - int check_shape_dim(int64_t dim) { - if (dim > std::numeric_limits::max()) { - throw std::invalid_argument( - "Shape dimension falls outside supported `int` range."); + if (dim > std::numeric_limits::max() || + dim < std::numeric_limits::min()) { + std::ostringstream msg; + msg << "Shape dimension " << dim << " is outside the supported range [" + << std::numeric_limits::min() << ", " + << std::numeric_limits::max() + << "]. MLX currently uses 32-bit integers for shape dimensions."; + PyErr_SetString(PyExc_OverflowError, msg.str().c_str()); + nb::detail::raise_python_error(); } return static_cast(dim); } -template -mx::array nd_array_to_mlx_contiguous( - nb::ndarray nd_array, - const mx::Shape& shape, - mx::Dtype dtype) { - // Make a copy of the numpy buffer - // Get buffer ptr pass to array constructor - auto data_ptr = nd_array.data(); - return mx::array(static_cast(data_ptr), shape, dtype); -} - -mx::array nd_array_to_mlx( - nb::ndarray nd_array, - std::optional dtype) { - // Compute the shape and size +mx::Shape get_shape(const nb::ndarray& nd_array) { mx::Shape shape; shape.reserve(nd_array.ndim()); for (int i = 0; i < nd_array.ndim(); i++) { shape.push_back(check_shape_dim(nd_array.shape(i))); } - auto type = nd_array.dtype(); + return shape; +} + +mx::Strides get_strides(const nb::ndarray& nd_array) { + mx::Strides strides; + strides.reserve(nd_array.ndim()); + for (int i = 0; i < nd_array.ndim(); i++) { + strides.push_back(nd_array.stride(i)); + } + return strides; +} - // Copy data and make array +size_t strided_storage_size( + const mx::Shape& shape, + const mx::Strides& strides) { + size_t storage_size = 1; + for (int i = 0; i < shape.size(); i++) { + if (shape[i] == 0) { + return 0; + } + if (strides[i] < 0) { + throw std::invalid_argument( + "Cannot convert DLPack arrays with negative strides to mlx array."); + } + storage_size += (shape[i] - 1) * strides[i]; + } + return storage_size; +} + +auto get_strided_layout( + const nb::ndarray& nd_array, + const mx::Shape& shape) { + auto strides = get_strides(nd_array); + auto storage_size = strided_storage_size(shape, strides); + auto [no_bsx_size, is_row_contiguous, is_col_contiguous] = shape.empty() + ? std::make_tuple(storage_size, true, true) + : mx::check_contiguity(shape, strides); + mx::array::Flags flags{ + no_bsx_size == storage_size, + is_row_contiguous, + is_col_contiguous, + }; + return std::make_tuple(storage_size, std::move(strides), flags); +} + +template +auto dispatch_dlpack_dtype( + nb::dlpack::dtype type, + F&& f, + const char* error_message) { if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::bool_)); + return f.template operator()(mx::bool_); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::uint8)); + return f.template operator()(mx::uint8); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::uint16)); + return f.template operator()(mx::uint16); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::uint32)); + return f.template operator()(mx::uint32); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::uint64)); + return f.template operator()(mx::uint64); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::int8)); + return f.template operator()(mx::int8); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::int16)); + return f.template operator()(mx::int16); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::int32)); + return f.template operator()(mx::int32); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::int64)); + return f.template operator()(mx::int64); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::float16)); - } else if (type == nb::bfloat16) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::bfloat16)); + return f.template operator()(mx::float16); + } else if (type == nb::dtype()) { + return f.template operator()(mx::bfloat16); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::float32)); + return f.template operator()(mx::float32); } else if (type == nb::dtype()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::float32)); + return f.template operator()(mx::float32); } else if (type == nb::dtype>()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::complex64)); + return f.template operator()(mx::complex64); } else if (type == nb::dtype>()) { - return nd_array_to_mlx_contiguous( - nd_array, shape, dtype.value_or(mx::complex64)); + return f.template operator()(mx::complex64); } else { - throw std::invalid_argument("Cannot convert numpy array to mlx array."); + throw std::invalid_argument(error_message); } } -template -nb::ndarray mlx_to_nd_array_impl( - mx::array a, - std::optional t = {}) { - { - nb::gil_scoped_release nogil; - a.eval(); +mx::Dtype mlx_dtype_from_dlpack( + nb::dlpack::dtype type, + const char* error_message) { + return dispatch_dlpack_dtype( + type, [](mx::Dtype dtype) { return dtype; }, error_message); +} + +nb::dlpack::dtype mlx_dtype_to_dl_dtype(mx::Dtype dtype) { + nb::dlpack::dtype result; + dispatch_all_types(dtype, [&](auto type_tag) { + using T = MLX_GET_TYPE(type_tag); + result = nb::dtype(); + }); + return result; +} + +template +mx::array cpu_nd_array_to_mlx( + nb::ndarray nd_array, + const mx::Shape& shape, + mx::Dtype dst_dtype) { + auto out = mx::array(shape, dst_dtype, nullptr, {}); + auto [storage_size, strides, flags] = get_strided_layout(nd_array, shape); + out.set_data( + mx::allocator::malloc(storage_size * mx::size_of(dst_dtype)), + storage_size, + std::move(strides), + flags); + if (storage_size > 0) { + dispatch_all_types( + dst_dtype, [&, storage_size = storage_size](auto type_tag) { + using DstT = MLX_GET_TYPE(type_tag); + auto src = static_cast(nd_array.data()); + auto dst = out.data(); + std::copy(src, src + storage_size, dst); + }); } - std::vector shape(a.shape().begin(), a.shape().end()); - return nb::ndarray( - a.data(), - a.ndim(), - shape.data(), - /* owner= */ nb::none(), - a.strides().data(), - t.value_or(nb::dtype())); + out.set_status(mx::array::Status::available); + return out; } -template -nb::ndarray mlx_to_nd_array(const mx::array& a) { - switch (a.dtype()) { - case mx::bool_: - return mlx_to_nd_array_impl(a); - case mx::uint8: - return mlx_to_nd_array_impl(a); - case mx::uint16: - return mlx_to_nd_array_impl(a); - case mx::uint32: - return mlx_to_nd_array_impl(a); - case mx::uint64: - return mlx_to_nd_array_impl(a); - case mx::int8: - return mlx_to_nd_array_impl(a); - case mx::int16: - return mlx_to_nd_array_impl(a); - case mx::int32: - return mlx_to_nd_array_impl(a); - case mx::int64: - return mlx_to_nd_array_impl(a); - case mx::float16: - return mlx_to_nd_array_impl(a); - case mx::bfloat16: - throw nb::type_error("bfloat16 arrays cannot be converted to NumPy."); - case mx::float32: - return mlx_to_nd_array_impl(a); - case mx::float64: - return mlx_to_nd_array_impl(a); - case mx::complex64: - return mlx_to_nd_array_impl, NDParams...>(a); +// Try to adopt a CPU host buffer as an mlx array without copying. On unified +// memory the host pointer is GPU-addressable, so we wrap it directly via the +// allocator instead of copying. The bytes are reinterpreted rather than +// converted, so the source element width must already match the destination +// dtype. The source ndarray is kept alive for the lifetime of the returned +// array. +// +// Returns std::nullopt when the buffer cannot be adopted (no Metal backend, +// dtype width mismatch, or a pointer the platform will not wrap), so the caller +// can fall back to a copy or raise. +std::optional cpu_nd_array_to_mlx_no_copy( + nb::ndarray nd_array, + const mx::Shape& shape, + mx::Dtype dst_dtype) { + if (!mx::metal::is_available() || + nd_array.itemsize() != mx::size_of(dst_dtype)) { + return std::nullopt; + } + + auto [storage_size, strides, flags] = get_strided_layout(nd_array, shape); + auto buf = mx::allocator::make_buffer( + const_cast(nd_array.data()), + storage_size * mx::size_of(dst_dtype)); + // make_buffer returns a null buffer when the pointer cannot be wrapped, e.g. + // when its alignment is not accepted by the platform. + if (buf.ptr() == nullptr) { + return std::nullopt; + } + + mx::array out(shape, dst_dtype, nullptr, {}); + out.set_data( + buf, + storage_size, + std::move(strides), + flags, + nd_array.byte_offset(), + // The buffer wraps caller-owned memory, so release the wrapper rather + // than returning it to the allocator's reuse pool, which must only + // recycle buffers it allocated itself. + [owner = std::move(nd_array)](mx::allocator::Buffer b) { + mx::allocator::release(b); + }); + out.set_status(mx::array::Status::available); + return out; +} + +mx::array metal_nd_array_to_mlx( + nb::ndarray nd_array, + mx::Dtype src_dtype, + mx::Dtype dst_dtype, + bool copy) { + if (!mx::metal::is_available()) { + throw std::invalid_argument("Metal DLPack import is not available."); + } + auto shape = get_shape(nd_array); + if (nd_array.itemsize() != mx::size_of(src_dtype)) { + throw std::invalid_argument( + "Cannot convert Metal DLPack dtype to mlx dtype."); + } + auto [storage_size, strides, flags] = get_strided_layout(nd_array, shape); + auto data_handle = nd_array.data_handle(); + mx::array out(shape, src_dtype, nullptr, {}); + out.set_data( + mx::allocator::Buffer(data_handle), + storage_size, + std::move(strides), + flags, + nd_array.byte_offset(), + [owner = std::move(nd_array)](mx::allocator::Buffer) {}); + out.set_status(mx::array::Status::available); + + if (copy) { + auto result = mx::astype(out, dst_dtype, true, mx::Device::gpu); + result.eval(); + return result; + } + return out; +} + +mx::array nd_array_to_mlx( + nb::ndarray nd_array, + std::optional requested_dtype, + std::optional src_dlpack_dtype_override, + std::optional copy) { + auto src_dlpack_dtype = src_dlpack_dtype_override.value_or(nd_array.dtype()); + auto src_mlx_dtype = mlx_dtype_from_dlpack( + src_dlpack_dtype, "[convert] Cannot convert array to mlx."); + auto dst_dtype = requested_dtype.value_or(src_mlx_dtype); + auto device_type = nd_array.device_type(); + + // A dtype change requires converting the elements, which cannot be done + // without a copy. + bool no_copy = copy.has_value() && !copy.value(); + if (no_copy && dst_dtype != src_mlx_dtype) { + throw std::invalid_argument( + "[convert] Cannot convert array to the requested dtype without a " + "copy."); + } + + switch (device_type) { + case nb::device::cpu::value: { + auto shape = get_shape(nd_array); + // For copy=None (try to share) and copy=False (must share), attempt a + // zero-copy adoption of the host buffer first. A copy is passed by value + // so the source is preserved for the fallback below. + if (!copy.value_or(false)) { + if (auto out = + cpu_nd_array_to_mlx_no_copy(nd_array, shape, dst_dtype)) { + return *out; + } + if (no_copy) { + throw std::invalid_argument( + "[convert] Cannot import a CPU array without a copy."); + } + } + // copy=True, or copy=None where adoption was not possible: copy. + return dispatch_dlpack_dtype( + src_dlpack_dtype, + [&](mx::Dtype) { + return cpu_nd_array_to_mlx(nd_array, shape, dst_dtype); + }, + "[convert] Cannot convert array to mlx."); + } + case nb::device::metal::value: { + // A Metal buffer can be adopted without a copy only if the active + // allocator recognizes it. + bool can_reuse_buffer = + mx::allocator::can_reuse_alien_buffer(nd_array.data_handle()); + if (no_copy && !can_reuse_buffer) { + throw std::invalid_argument( + "[convert] Cannot import a private Metal buffer without a copy."); + } + bool should_copy = copy.value_or(false) || dst_dtype != src_mlx_dtype || + !can_reuse_buffer; + return metal_nd_array_to_mlx( + nd_array, src_mlx_dtype, dst_dtype, should_copy); + } + case nb::device::cuda::value: + case nb::device::cuda_managed::value: + throw std::invalid_argument("[convert] CUDA import is not supported."); default: - throw nb::type_error("type cannot be converted to NumPy."); + throw std::invalid_argument("[convert] Unsupported device."); } } +template +nb::ndarray mlx_to_nd_array( + const mx::array& a, + std::optional> dl_device) { + auto default_device = mx::metal::is_available() + ? std::tuple{nb::device::metal::value, 0} + : std::tuple{nb::device::cpu::value, 0}; + auto [device_type, device_id] = dl_device.value_or(default_device); + + if (device_type == nb::device::cuda::value || + device_type == nb::device::cuda_managed::value) { + throw nb::buffer_error("CUDA DLPack export is not supported."); + } + if (device_type != nb::device::cpu::value && + device_type != nb::device::metal::value) { + throw nb::buffer_error( + "Cannot export mlx array to requested DLPack device."); + } + if (device_type == nb::device::metal::value && !mx::metal::is_available()) { + throw nb::buffer_error("Metal DLPack export is not available."); + } + + auto arr = a; + void* data = nullptr; + uint64_t byte_offset = 0; + { + nb::gil_scoped_release nogil; + arr.eval(); + } + data = device_type == nb::device::cpu::value ? arr.buffer().raw_ptr() + : arr.buffer().ptr(); + byte_offset = arr.offset(); + + std::vector shape(arr.shape().begin(), arr.shape().end()); + auto owner = nb::cast(arr); + return nb::ndarray( + data, + arr.ndim(), + shape.data(), + /* owner= */ owner, + arr.strides().data(), + mlx_dtype_to_dl_dtype(arr.dtype()), + device_type, + device_id, + '\0', + byte_offset); +} + nb::ndarray mlx_to_np_array(const mx::array& a) { - return mlx_to_nd_array(a); + if (a.dtype() == mx::bfloat16) { + throw nb::type_error("bfloat16 arrays cannot be converted to NumPy."); + } + return mlx_to_nd_array(a, std::tuple{nb::device::cpu::value, 0}); } -nb::ndarray<> mlx_to_dlpack(const mx::array& a) { - return mlx_to_nd_array<>(a); +nb::ndarray<> mlx_to_dlpack( + const mx::array& a, + bool force_copy, + std::optional> dl_device) { + if (force_copy) { + return mlx_to_nd_array<>(mx::astype(a, a.dtype(), true), dl_device); + } + return mlx_to_nd_array<>(a, dl_device); } nb::object to_scalar(mx::array& a) { @@ -286,7 +505,8 @@ PyScalarT validate_shape( T list, const mx::Shape& shape, int idx, - bool& all_python_primitive_elements) { + bool& all_python_primitive_elements, + bool& has_wide_int) { if (idx >= shape.size()) { throw std::invalid_argument("Initialization encountered extra dimension."); } @@ -305,13 +525,18 @@ PyScalarT validate_shape( PyScalarT t; if (nb::isinstance(l)) { t = validate_shape( - nb::cast(l), shape, idx + 1, all_python_primitive_elements); + nb::cast(l), + shape, + idx + 1, + all_python_primitive_elements, + has_wide_int); } else if (nb::isinstance(*list.begin())) { t = validate_shape( nb::cast(l), shape, idx + 1, - all_python_primitive_elements); + all_python_primitive_elements, + has_wide_int); } else if (nb::isinstance(l)) { all_python_primitive_elements = false; auto arr = nb::cast(l); @@ -330,6 +555,13 @@ PyScalarT validate_shape( t = pybool; } else if (nb::isinstance(l)) { t = pyint; + // Match the scalar path, which widens to int64 rather than failing + // when a python int does not fit in int32. + auto val = nb::cast(l); + if (val > std::numeric_limits::max() || + val < std::numeric_limits::min()) { + has_wide_int = true; + } } else if (nb::isinstance(l)) { t = pyfloat; } else if (PyComplex_Check(l.ptr())) { @@ -375,7 +607,8 @@ mx::array array_from_list_impl( T pl, const PyScalarT& inferred_type, std::optional specified_type, - const mx::Shape& shape) { + const mx::Shape& shape, + bool has_wide_int) { // Make the array switch (inferred_type) { case pybool: { @@ -384,7 +617,8 @@ mx::array array_from_list_impl( return mx::array(vals.begin(), shape, specified_type.value_or(mx::bool_)); } case pyint: { - auto dtype = specified_type.value_or(mx::int32); + auto dtype = + specified_type.value_or(has_wide_int ? mx::int64 : mx::int32); if (dtype == mx::int64) { std::vector vals; fill_vector(pl, vals); @@ -444,17 +678,19 @@ mx::array array_from_list_impl(T pl, std::optional dtype) { // Validate the shape and type bool all_python_primitive_elements = true; - auto type = validate_shape(pl, shape, 0, all_python_primitive_elements); + bool has_wide_int = false; + auto type = + validate_shape(pl, shape, 0, all_python_primitive_elements, has_wide_int); if (all_python_primitive_elements) { // `pl` does not contain mlx arrays - return array_from_list_impl(pl, type, dtype, shape); + return array_from_list_impl(pl, type, dtype, shape, has_wide_int); } // `pl` contains mlx arrays std::vector arrays; for (auto l : pl) { - arrays.push_back(create_array(nb::cast(l), dtype)); + arrays.push_back(create_array(nb::cast(l), dtype)); } return mx::stack(arrays); } @@ -467,38 +703,60 @@ mx::array array_from_list(nb::tuple pl, std::optional dtype) { return array_from_list_impl(pl, dtype); } -mx::array create_array(ArrayInitType v, std::optional t) { - if (auto pv = std::get_if(&v); pv) { - return mx::array(nb::cast(*pv), t.value_or(mx::bool_)); - } else if (auto pv = std::get_if(&v); pv) { - auto val = nb::cast(*pv); +mx::array create_array( + nb::object v, + std::optional t, + std::optional copy) { + if (!nb::isinstance(v) && nb::ndarray_check(v)) { + using ContigArray = nb::ndarray; + ContigArray nd; + std::optional nb_dtype; + // Nanobind does not recognize bfloat16 numpy array: + // https://github.com/wjakob/nanobind/discussions/560 + if (nb::hasattr(v, "dtype") && v.attr("dtype").equal(nb::str("bfloat16"))) { + nd = nb::cast(v.attr("view")("uint16")); + nb_dtype = nb::dtype(); + } else { + nd = nb::cast(v); + } + return nd_array_to_mlx(nd, t, nb_dtype, copy); + } + + if (copy.has_value() && copy.value() == false) { + throw std::invalid_argument( + "Unable to avoid copy while creating an array as requested."); + } + + if (nb::isinstance(v)) { + return mx::array(nb::cast(v), t.value_or(mx::bool_)); + } else if (nb::isinstance(v)) { + auto val = nb::cast(v); auto default_type = (val > std::numeric_limits::max() || val < std::numeric_limits::min()) ? mx::int64 : mx::int32; return mx::array(val, t.value_or(default_type)); - } else if (auto pv = std::get_if(&v); pv) { + } else if (nb::isinstance(v)) { auto out_type = t.value_or(mx::float32); if (out_type == mx::float64) { - return mx::array(nb::cast(*pv), out_type); + return mx::array(nb::cast(v), out_type); } else { - return mx::array(nb::cast(*pv), out_type); + return mx::array(nb::cast(v), out_type); } - } else if (auto pv = std::get_if>(&v); pv) { + } else if (PyComplex_Check(v.ptr())) { return mx::array( - static_cast(*pv), t.value_or(mx::complex64)); - } else if (auto pv = std::get_if(&v); pv) { - return array_from_list(*pv, t); - } else if (auto pv = std::get_if(&v); pv) { - return array_from_list(*pv, t); - } else if (auto pv = std::get_if< - nb::ndarray>(&v); - pv) { - return nd_array_to_mlx(*pv, t); - } else if (auto pv = std::get_if(&v); pv) { - return mx::astype(*pv, t.value_or((*pv).dtype())); + static_cast(nb::cast>(v)), + t.value_or(mx::complex64)); + } else if (nb::isinstance(v)) { + return array_from_list(nb::cast(v), t); + } else if (nb::isinstance(v)) { + return array_from_list(nb::cast(v), t); + } else if (nb::isinstance(v)) { + auto arr = nb::cast(v); + auto dtype = t.value_or(arr.dtype()); + return mx::astype(arr, dtype, copy.value_or(false)); } else { - auto arr = to_array_with_accessor(std::get(v).obj); - return mx::astype(arr, t.value_or(arr.dtype())); + auto arr = to_array_with_accessor(v); + return mx::astype(arr, t.value_or(arr.dtype()), copy.value_or(false)); } } diff --git a/Source/Cmlx/mlx/python/src/convert.h b/Source/Cmlx/mlx/python/src/convert.h index 4f8a77abe..133f443ed 100644 --- a/Source/Cmlx/mlx/python/src/convert.h +++ b/Source/Cmlx/mlx/python/src/convert.h @@ -2,49 +2,88 @@ #pragma once #include +#include #include #include #include "mlx/array.h" -#include "mlx/ops.h" namespace mx = mlx::core; namespace nb = nanobind; namespace nanobind { -static constexpr dlpack::dtype bfloat16{4, 16, 1}; -}; // namespace nanobind + +template <> +struct ndarray_traits { + static constexpr bool is_complex = false; + static constexpr bool is_float = true; + static constexpr bool is_bool = false; + static constexpr bool is_int = false; + static constexpr bool is_signed = true; +}; + +template <> +struct ndarray_traits { + static constexpr bool is_complex = false; + static constexpr bool is_float = true; + static constexpr bool is_bool = false; + static constexpr bool is_int = false; + static constexpr bool is_signed = true; +}; + +namespace detail { + +template <> +struct dtype_traits { + static constexpr dlpack::dtype value{ + /* code */ uint8_t(nb::dlpack::dtype_code::Float), + /* bits */ 16, + /* lanes */ 1}; + static constexpr const char* name = "float16"; +}; + +template <> +struct dtype_traits { + static constexpr dlpack::dtype value{ + /* code */ uint8_t(nb::dlpack::dtype_code::Bfloat), + /* bits */ 16, + /* lanes */ 1}; + static constexpr const char* name = "bfloat16"; +}; + +} // namespace detail + +} // namespace nanobind struct ArrayLike { ArrayLike(nb::object obj) : obj(obj) {}; nb::object obj; }; -using ArrayInitType = std::variant< - nb::bool_, - nb::int_, - nb::float_, - // Must be above ndarray - mx::array, - // Must be above complex - nb::ndarray, - std::complex, - nb::list, - nb::tuple, - ArrayLike>; - mx::array nd_array_to_mlx( - nb::ndarray nd_array, - std::optional dtype); + nb::ndarray nd_array, + std::optional mx_dtype, + std::optional src_dlpack_dtype_override = std::nullopt, + std::optional copy = std::nullopt); nb::ndarray mlx_to_np_array(const mx::array& a); -nb::ndarray<> mlx_to_dlpack(const mx::array& a); +nb::ndarray<> mlx_to_dlpack( + const mx::array& a, + bool force_copy, + std::optional> dl_device); nb::object to_scalar(mx::array& a); nb::object tolist(mx::array& a); -mx::array create_array(ArrayInitType v, std::optional t); +mx::array create_array( + nb::object v, + std::optional t, + std::optional copy = true); mx::array array_from_list(nb::list pl, std::optional dtype); mx::array array_from_list(nb::tuple pl, std::optional dtype); + +// Narrow a Python-side shape dimension (int64) to a C++ mx::ShapeElem (int32), +// raising a clear error if the value would overflow. +int check_shape_dim(int64_t dim); diff --git a/Source/Cmlx/mlx/python/src/device.cpp b/Source/Cmlx/mlx/python/src/device.cpp index f15f7f92d..83d32d1cb 100644 --- a/Source/Cmlx/mlx/python/src/device.cpp +++ b/Source/Cmlx/mlx/python/src/device.cpp @@ -1,8 +1,10 @@ // Copyright © 2023-2025 Apple Inc. +#include #include #include +#include #include #include #include @@ -59,11 +61,13 @@ void init_device(nb::module_& m) { "set_default_device", &mx::set_default_device, "device"_a, + nb::sig("def set_default_device(device: Device | DeviceType) -> None"), R"pbdoc(Set the default device.)pbdoc"); m.def( "is_available", &mx::is_available, "device"_a, + nb::sig("def is_available(device: Device | DeviceType) -> bool"), R"pbdoc(Check if a back-end is available for the given device.)pbdoc"); m.def( "device_count", @@ -80,8 +84,12 @@ void init_device(nb::module_& m) { )pbdoc"); m.def( "device_info", - &mx::device_info, - nb::arg("d") = mx::default_device(), + [](std::optional d) { + return mx::device_info(d.value_or(mx::default_device())); + }, + "d"_a = nb::none(), + nb::sig( + "def device_info(d: None | Device | DeviceType = None) -> dict[str, str | int]"), R"pbdoc( Get information about a device. diff --git a/Source/Cmlx/mlx/python/src/distributed.cpp b/Source/Cmlx/mlx/python/src/distributed.cpp index 9f4a7cb59..0ad91c6e4 100644 --- a/Source/Cmlx/mlx/python/src/distributed.cpp +++ b/Source/Cmlx/mlx/python/src/distributed.cpp @@ -12,6 +12,9 @@ #include "python/src/small_vector.h" #include "python/src/utils.h" +#include +#include + namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; @@ -75,10 +78,54 @@ void init_distributed(nb::module_& parent_module) { m.def( "init", - &mx::distributed::init, + [](bool strict, + const std::string& backend, + std::optional all_gather_factory) + -> mx::distributed::Group { + if (!all_gather_factory.has_value()) { + return mx::distributed::init(strict, backend); + } + + if (backend != "jaccl") { + throw std::invalid_argument( + "all_gather_factory is only supported with backend='jaccl'."); + } + + auto py_factory = std::move(*all_gather_factory); + auto cpp_factory = [py_factory = std::move(py_factory)]( + int rank, + int size) -> mx::distributed::AllGatherFn { + nb::gil_scoped_acquire gil; + nb::object py_inner = py_factory(rank, size); + if (!PyCallable_Check(py_inner.ptr())) { + throw std::invalid_argument( + "all_gather_factory must return a callable"); + } + return [py_inner = std::move(py_inner), size]( + const void* src, void* dst, size_t n_bytes) { + nb::gil_scoped_acquire gil; + nb::bytes src_bytes(src, n_bytes); + nb::object result_obj = py_inner(src_bytes, n_bytes); + nb::bytes result = nb::cast(result_obj); + size_t expected = static_cast(size) * n_bytes; + if (result.size() != expected) { + std::ostringstream msg; + msg << "Custom all-gather returned " << result.size() + << " bytes but expected " << expected; + throw std::runtime_error(msg.str()); + } + std::memcpy(dst, result.data(), expected); + }; + }; + + return mx::distributed::init(strict, backend, std::move(cpp_factory)); + }, "strict"_a = false, "backend"_a = "any", - nb::sig("def init(strict: bool = False, backend: str = 'any') -> Group"), + nb::kw_only(), + "all_gather_factory"_a = nb::none(), + nb::sig( + "def init(strict: bool = False, backend: str = 'any', *, all_gather_factory: Callable[[int, int], Callable[[bytes, int], bytes]] | None = None) -> Group"), R"pbdoc( Initialize the communication backend and create the global communication group. @@ -99,6 +146,13 @@ void init_distributed(nb::module_& parent_module) { set to ``any`` all available backends are tried and the first one that succeeds becomes the global group which will be returned in subsequent calls. Default: ``any`` + all_gather_factory (Callable, optional): A factory used only with the + ``jaccl`` backend. It is called once per rank with ``(rank, size)`` + and must return a callable with signature + ``f(src: bytes, n_bytes: int) -> bytes``. The returned callable + performs a byte-level all-gather used as the JACCL side channel + when exchanging RDMA connection metadata. The returned bytes must + have length ``size * n_bytes``. Returns: Group: The group representing all the launched processes. @@ -116,7 +170,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_sum(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_sum(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce sum. @@ -145,7 +199,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_max(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_max(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce max. @@ -174,7 +228,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_min(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_min(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( All reduce min. @@ -203,7 +257,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def all_gather(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def all_gather(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Gather arrays from all processes. @@ -236,7 +290,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def send(x: array, dst: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def send(x: array, dst: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Send an array from the current process to the process that has rank ``dst`` in the group. @@ -264,7 +318,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def recv(shape: Sequence[int], dtype: Dtype, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Recv an array with shape ``shape`` and dtype ``dtype`` from process with rank ``src``. @@ -297,7 +351,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def recv_like(x: array, src: int, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def recv_like(x: array, src: int, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Recv an array with shape and type like ``x`` from process with rank ``src``. @@ -330,7 +384,7 @@ void init_distributed(nb::module_& parent_module) { "group"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def sum_scatter(x: array, *, group: Optional[Group] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def sum_scatter(x: array, *, group: Group | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sum ``x`` across all processes in the group and shard the result along the first axis across ranks. ``x.shape[0]`` must be divisible by the group size. @@ -349,4 +403,10 @@ void init_distributed(nb::module_& parent_module) { Returns: array: The output array with shape ``[x.shape[0] // group.size(), *x.shape[1:]]``. )pbdoc"); + + // Ensure the distributed backend cache is cleared before the interpreter + // goes away, so that any Python objects held by cached groups are released + // while Python is still alive. + auto atexit = nb::module_::import_("atexit"); + atexit.attr("register")(nb::cpp_function(&mx::distributed::clear_backends)); } diff --git a/Source/Cmlx/mlx/python/src/export.cpp b/Source/Cmlx/mlx/python/src/export.cpp index 3d3ffdf6b..7748f87b1 100644 --- a/Source/Cmlx/mlx/python/src/export.cpp +++ b/Source/Cmlx/mlx/python/src/export.cpp @@ -124,8 +124,8 @@ auto wrap_export_function(nb::callable fun) { outputs_.push_back(nb::cast(outputs)); } else if (!nb::try_cast(outputs, outputs_)) { throw std::invalid_argument( - "[export_function] Outputs can be either a single array " - "a tuple or list of arrays."); + "[export_function] Outputs can be either a single array, " + "a tuple, or a list of arrays."); } return outputs_; }; @@ -138,6 +138,7 @@ void init_export(nb::module_& m) { const nb::callable& fun, const nb::args& args, bool shapeless, + const std::optional& metadata, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs(args, kwargs, "[export_function]"); @@ -147,8 +148,14 @@ void init_export(nb::module_& m) { wrap_export_function(fun), args_, kwargs_, - shapeless); + shapeless, + metadata.value_or("")); } else { + if (metadata && !metadata->empty()) { + throw std::invalid_argument( + "[export_function] The metadata argument is only supported " + "when exporting to a file, not when using a callback."); + } auto callback = nb::cast(file_or_callback); auto wrapped_callback = [callback](const mx::ExportCallbackInput& input) { @@ -163,7 +170,10 @@ void init_export(nb::module_& m) { "args"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), "kwargs"_a, + nb::sig( + "def export_function(file_or_callback: str | Callable, fun: Callable, *args, shapeless: bool = False, metadata: str | None = None, **kwargs) -> None"), R"pbdoc( Export an MLX function. @@ -178,16 +188,23 @@ void init_export(nb::module_& m) { versions of MLX may not be compatible with future versions. Args: - file (str or Callable): Either a file path to export the function - to or a callback. + file_or_callback (str or Callable): Either a file path to export + the function to or a callback. fun (Callable): A function which takes as input zero or more :class:`array` and returns one or more :class:`array`. *args (array): Example array inputs to the function. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (str, optional): A string to save alongside the + function, for example a JSON encoded model configuration. Only + supported when exporting to a file. Read it back with + :func:`import_function`. Default: ``None``. **kwargs (array): Additional example keyword array inputs to the function. + Raises: + ValueError: If ``metadata`` is given when exporting with a callback. + Example: .. code-block:: python @@ -201,17 +218,25 @@ void init_export(nb::module_& m) { )pbdoc"); m.def( "import_function", - [](const std::string& file) { - return nb::cpp_function( - [fn = mx::import_function(file)]( + [](const std::string& file, bool return_metadata) -> nb::object { + auto imported = mx::import_function(file); + auto metadata = imported.metadata(); + auto fn = nb::cpp_function( + [imported = std::move(imported)]( const nb::args& args, const nb::kwargs& kwargs) { auto [args_, kwargs_] = validate_and_extract_inputs( args, kwargs, "[import_function::call]"); - return nb::tuple(nb::cast(fn(args_, kwargs_))); + return nb::tuple(nb::cast(imported(args_, kwargs_))); }); + if (return_metadata) { + return nb::make_tuple(fn, nb::cast(metadata)); + } + return fn; }, "file"_a, - nb::sig("def import_function(file: str) -> Callable"), + "return_metadata"_a = false, + nb::sig( + "def import_function(file: str, return_metadata: bool = False) -> Callable | tuple[Callable, str]"), R"pbdoc( Import a function from a file. @@ -228,15 +253,20 @@ void init_export(nb::module_& m) { Args: file (str): The file path to import the function from. + return_metadata (bool, optional): If ``True`` also return the + metadata string saved with the function. Default: ``False``. Returns: - Callable: The imported function. + Callable or tuple: + The imported function. If ``return_metadata`` is ``True`` a + tuple of the imported function and the metadata string is + returned instead. Example: >>> fn = mx.import_function("function.mlxfn") >>> out = fn(a, b, x=x, y=y)[0] >>> - >>> out = fn((a, b), {"x": x, "y": y}[0] + >>> out = fn((a, b), {"x": x, "y": y})[0] )pbdoc"); nb::class_( @@ -247,7 +277,7 @@ void init_export(nb::module_& m) { A context managing class for exporting multiple traces of the same function to a file. - Make an instance of this class by calling fun:`mx.exporter`. + Make an instance of this class by calling :func:`mx.exporter`. )pbdoc") .def("close", &PyFunctionExporter::close) .def("__enter__", [](PyFunctionExporter& exporter) { return &exporter; }) @@ -272,14 +302,23 @@ void init_export(nb::module_& m) { m.def( "exporter", - [](const std::string& file, nb::callable fun, bool shapeless) { + [](const std::string& file, + nb::callable fun, + bool shapeless, + const std::optional& metadata) { return PyFunctionExporter{ - mx::exporter(file, wrap_export_function(fun), shapeless), fun}; + mx::exporter( + file, + wrap_export_function(fun), + shapeless, + metadata.value_or("")), + fun}; }, "file"_a, "fun"_a, nb::kw_only(), "shapeless"_a = false, + "metadata"_a = nb::none(), R"pbdoc( Make a callable object to export multiple traces of a function to a file. @@ -293,6 +332,9 @@ void init_export(nb::module_& m) { file (str): File path to export the function to. shapeless (bool, optional): Whether or not the function allows inputs with variable shapes. Default: ``False``. + metadata (str, optional): A string to save alongside the + function, for example a JSON encoded model configuration. Read + it back with :func:`import_function`. Default: ``None``. Example: diff --git a/Source/Cmlx/mlx/python/src/fast.cpp b/Source/Cmlx/mlx/python/src/fast.cpp index 1a43d89d9..67c3442cf 100644 --- a/Source/Cmlx/mlx/python/src/fast.cpp +++ b/Source/Cmlx/mlx/python/src/fast.cpp @@ -75,6 +75,43 @@ struct PyCustomKernelFunction { const char* tag_; }; +mx::MathMode parse_metal_math_mode(const std::string& math_mode) { + if (math_mode == "safe") { + return mx::MathMode::Safe; + } else if (math_mode == "relaxed") { + return mx::MathMode::Relaxed; + } else if (math_mode == "fast") { + return mx::MathMode::Fast; + } + throw std::invalid_argument( + "[metal_kernel] Expected math_mode to be 'safe', 'relaxed', or 'fast'."); +} + +mx::CompileOptions parse_compile_options(const nb::object& obj) { + mx::CompileOptions result; + if (obj.is_none()) { + return result; + } + + if (!nb::isinstance(obj)) { + throw std::invalid_argument( + "[metal_kernel] Expected `compile_options` to be a dict."); + } + + nb::dict dict = nb::cast(obj); + for (auto [key, value] : dict) { + auto key_str = nb::cast(key); + if (key_str == "math_mode") { + result.math_mode = parse_metal_math_mode(nb::cast(value)); + } else { + std::ostringstream msg; + msg << "[metal_kernel] Unknown compile option `" << key_str << "`."; + throw std::invalid_argument(msg.str()); + } + } + return result; +} + } // namespace void init_fast(nb::module_& parent_module) { @@ -90,7 +127,7 @@ void init_fast(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def rms_norm(x: array, weight: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), + "def rms_norm(x: array, weight: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Root Mean Square normalization (RMS norm). @@ -117,7 +154,7 @@ void init_fast(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def layer_norm(x: array, weight: Optional[array], bias: Optional[array], eps: float, *, stream: Union[None, Stream, Device] = None) -> array"), + "def layer_norm(x: array, weight: array | None, bias: array | None, eps: float, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Layer normalization. @@ -137,6 +174,36 @@ void init_fast(nb::module_& parent_module) { array: The output array. )pbdoc"); + m.def( + "cross_entropy", + &mx::fast::cross_entropy, + "logits"_a, + "targets"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def cross_entropy(logits: array, targets: array, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Cross entropy loss with class indices as targets. + + Computes ``logsumexp(logits, axis=-1) - logits[..., target]`` in a + fused kernel with accumulation in float32. + + Note: Currently is implemented only on CUDA, fallback to unfused version with + manual casting on Metal and CPU. + + Args: + logits (array): The unnormalized logits. The loss is computed over + the last axis. + targets (array): Class indices. The shape should match the shape of + ``logits`` with the last axis removed. The indices must be in + ``[0, logits.shape[-1])``. + + Returns: + array: The per-element loss in float32, with the shape of + ``targets``. + )pbdoc"); + m.def( "rope", [](const mx::array& a, @@ -160,7 +227,7 @@ void init_fast(nb::module_& parent_module) { "freqs"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def rope(a: array, dims: int, *, traditional: bool, base: Optional[float], scale: float, offset: Union[int, array], freqs: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def rope(a: array, dims: int, *, traditional: bool, base: float | None, scale: float, offset: int | array, freqs: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Apply rotary positional encoding to the input. @@ -197,6 +264,7 @@ void init_fast(nb::module_& parent_module) { const float scale, const std::variant& mask, const std::optional& sinks, + bool force_fused, mx::StreamOrDevice s) { bool has_mask = !std::holds_alternative(mask); bool has_str_mask = @@ -213,16 +281,32 @@ void init_fast(nb::module_& parent_module) { throw std::invalid_argument(msg.str()); } return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, mask_str, std::nullopt, sinks, s); + queries, + keys, + values, + scale, + mask_str, + std::nullopt, + sinks, + force_fused, + s); } else { auto mask_arr = std::get(mask); return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, "", mask_arr, sinks, s); + queries, + keys, + values, + scale, + "", + mask_arr, + sinks, + force_fused, + s); } } else { return mx::fast::scaled_dot_product_attention( - queries, keys, values, scale, "", {}, sinks, s); + queries, keys, values, scale, "", {}, sinks, force_fused, s); } }, "q"_a, @@ -232,9 +316,10 @@ void init_fast(nb::module_& parent_module) { "scale"_a, "mask"_a = nb::none(), "sinks"_a = nb::none(), + "force_fused"_a = false, "stream"_a = nb::none(), nb::sig( - "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: Union[None, str, array] = None, sinks: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float, mask: None | str | array = None, sinks: array | None = None, force_fused: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( A fast implementation of multi-head attention: ``O = softmax(Q @ K.T, dim=-1) @ V``. @@ -276,6 +361,11 @@ void init_fast(nb::module_& parent_module) { last query aligns with the last key. sinks (array, optional): An optional array of attention sinks. Default: ``None``. + force_fused (bool, optional): If ``True``, use a fused kernel + regardless of the builtin heuristics and raise error when no + fused kernel is available. For certain configurations this would + result in slower kernel getting used but can reduce memory + consumption. Default: ``False``. Returns: array: The output array. @@ -304,7 +394,8 @@ void init_fast(nb::module_& parent_module) { const std::string& source, const std::string& header, bool ensure_row_contiguous, - bool atomic_outputs) { + bool atomic_outputs, + const nb::object& compile_options) { auto kernel = mx::fast::metal_kernel( name, input_names, @@ -312,7 +403,8 @@ void init_fast(nb::module_& parent_module) { source, header, ensure_row_contiguous, - atomic_outputs); + atomic_outputs, + parse_compile_options(compile_options)); return nb::cpp_function( PyCustomKernelFunction(std::move(kernel), "[metal_kernel]"), nb::kw_only(), @@ -326,7 +418,7 @@ void init_fast(nb::module_& parent_module) { "verbose"_a = false, "stream"_a = nb::none(), nb::sig( - "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"), + "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"), R"pbdoc( Run the kernel. @@ -356,6 +448,7 @@ void init_fast(nb::module_& parent_module) { "header"_a = "", "ensure_row_contiguous"_a = true, "atomic_outputs"_a = false, + "compile_options"_a = nb::none(), R"pbdoc( A jit-compiled custom Metal kernel defined from a source string. @@ -376,6 +469,12 @@ void init_fast(nb::module_& parent_module) { before the kernel runs. Default: ``True``. atomic_outputs (bool): Whether to use atomic outputs in the function signature e.g. ``device atomic``. Default: ``False``. + compile_options (dict, optional): Options to compile the Metal kernel + with. Supported options: + + * ``"math_mode"``: The Metal math mode: ``"safe"``, ``"relaxed"``, + or ``"fast"``. ``"safe"`` preserves IEEE behavior for special + values such as ``exp(-inf) == 0``. Default: ``"safe"``. Returns: Callable ``metal_kernel``. @@ -443,7 +542,7 @@ void init_fast(nb::module_& parent_module) { "verbose"_a = false, "stream"_a = nb::none(), nb::sig( - "def __call__(self, *, inputs: List[Union[scalar, array]], output_shapes: List[Sequence[int]], output_dtypes: List[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: Optional[List[Tuple[str, Union[bool, int, Dtype]]]] = None, init_value: Optional[float] = None, verbose: bool = false, stream: Union[None, Stream, Device] = None)"), + "def __call__(self, *, inputs: list[scalar | array], output_shapes: list[Sequence[int]], output_dtypes: list[Dtype], grid: tuple[int, int, int], threadgroup: tuple[int, int, int], template: list[tuple[str, bool | int | Dtype]] | None = None, init_value: float | None = None, verbose: bool = false, stream: StreamOrDevice = None)"), R"pbdoc( Run the kernel. diff --git a/Source/Cmlx/mlx/python/src/fft.cpp b/Source/Cmlx/mlx/python/src/fft.cpp index 522e1064c..99f8478bb 100644 --- a/Source/Cmlx/mlx/python/src/fft.cpp +++ b/Source/Cmlx/mlx/python/src/fft.cpp @@ -2,18 +2,40 @@ #include #include +#include #include #include #include +#include #include "mlx/fft.h" #include "mlx/ops.h" #include "python/src/small_vector.h" +#include "python/src/utils.h" namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; +namespace { + +mx::fft::FFTNorm parse_norm(std::string_view norm, std::string_view op) { + if (norm == "backward") { + return mx::fft::FFTNorm::Backward; + } + if (norm == "ortho") { + return mx::fft::FFTNorm::Ortho; + } + if (norm == "forward") { + return mx::fft::FFTNorm::Forward; + } + throw std::invalid_argument( + std::string("[") + std::string(op) + + "] Invalid norm. Expected one of {'backward', 'ortho', 'forward'}."); +} + +} // namespace + void init_fft(nb::module_& parent_module) { auto m = parent_module.def_submodule( "fft", "mlx.core.fft: Fast Fourier Transforms."); @@ -22,16 +44,19 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, int axis, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "fft"); if (n.has_value()) { - return mx::fft::fft(a, n.value(), axis, s); + return mx::fft::fft(a, n.value(), axis, fft_norm, s); } else { - return mx::fft::fft(a, axis, s); + return mx::fft::fft(a, axis, fft_norm, s); } }, "a"_a, "n"_a = nb::none(), "axis"_a = -1, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( One dimensional discrete Fourier Transform. @@ -43,6 +68,8 @@ void init_fft(nb::module_& parent_module) { zeros to match ``n``. The default value is ``a.shape[axis]``. axis (int, optional): Axis along which to perform the FFT. The default is ``-1``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The DFT of the input along the given axis. @@ -52,16 +79,19 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, int axis, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "ifft"); if (n.has_value()) { - return mx::fft::ifft(a, n.value(), axis, s); + return mx::fft::ifft(a, n.value(), axis, fft_norm, s); } else { - return mx::fft::ifft(a, axis, s); + return mx::fft::ifft(a, axis, fft_norm, s); } }, "a"_a, "n"_a = nb::none(), "axis"_a = -1, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( One dimensional inverse discrete Fourier Transform. @@ -73,6 +103,8 @@ void init_fft(nb::module_& parent_module) { zeros to match ``n``. The default value is ``a.shape[axis]``. axis (int, optional): Axis along which to perform the FFT. The default is ``-1``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The inverse DFT of the input along the given axis. @@ -82,21 +114,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "fft2"); if (axes.has_value() && n.has_value()) { - return mx::fft::fftn(a, n.value(), axes.value(), s); + return mx::fft::fftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::fftn(a, axes.value(), s); + return mx::fft::fftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[fft2] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::fftn(a, s); + return mx::fft::fftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a.none() = std::vector{-2, -1}, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( Two dimensional discrete Fourier Transform. @@ -109,6 +144,8 @@ void init_fft(nb::module_& parent_module) { sizes of ``a`` along ``axes``. axes (list(int), optional): Axes along which to perform the FFT. The default is ``[-2, -1]``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The DFT of the input along the given axes. @@ -118,21 +155,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "ifft2"); if (axes.has_value() && n.has_value()) { - return mx::fft::ifftn(a, n.value(), axes.value(), s); + return mx::fft::ifftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::ifftn(a, axes.value(), s); + return mx::fft::ifftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[ifft2] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::ifftn(a, s); + return mx::fft::ifftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a.none() = std::vector{-2, -1}, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( Two dimensional inverse discrete Fourier Transform. @@ -145,6 +185,8 @@ void init_fft(nb::module_& parent_module) { sizes of ``a`` along ``axes``. axes (list(int), optional): Axes along which to perform the FFT. The default is ``[-2, -1]``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The inverse DFT of the input along the given axes. @@ -154,21 +196,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "fftn"); if (axes.has_value() && n.has_value()) { - return mx::fft::fftn(a, n.value(), axes.value(), s); + return mx::fft::fftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::fftn(a, axes.value(), s); + return mx::fft::fftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[fftn] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::fftn(a, s); + return mx::fft::fftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a = nb::none(), + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( n-dimensional discrete Fourier Transform. @@ -182,6 +227,8 @@ void init_fft(nb::module_& parent_module) { axes (list(int), optional): Axes along which to perform the FFT. The default is ``None`` in which case the FFT is over the last ``len(s)`` axes are or all axes if ``s`` is also ``None``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The DFT of the input along the given axes. @@ -191,21 +238,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "ifftn"); if (axes.has_value() && n.has_value()) { - return mx::fft::ifftn(a, n.value(), axes.value(), s); + return mx::fft::ifftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::ifftn(a, axes.value(), s); + return mx::fft::ifftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[ifftn] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::ifftn(a, s); + return mx::fft::ifftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a = nb::none(), + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( n-dimensional inverse discrete Fourier Transform. @@ -219,6 +269,8 @@ void init_fft(nb::module_& parent_module) { axes (list(int), optional): Axes along which to perform the FFT. The default is ``None`` in which case the FFT is over the last ``len(s)`` axes or all axes if ``s`` is also ``None``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The inverse DFT of the input along the given axes. @@ -228,16 +280,19 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, int axis, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "rfft"); if (n.has_value()) { - return mx::fft::rfft(a, n.value(), axis, s); + return mx::fft::rfft(a, n.value(), axis, fft_norm, s); } else { - return mx::fft::rfft(a, axis, s); + return mx::fft::rfft(a, axis, fft_norm, s); } }, "a"_a, "n"_a = nb::none(), "axis"_a = -1, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( One dimensional discrete Fourier Transform on a real input. @@ -253,6 +308,8 @@ void init_fft(nb::module_& parent_module) { zeros to match ``n``. The default value is ``a.shape[axis]``. axis (int, optional): Axis along which to perform the FFT. The default is ``-1``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The DFT of the input along the given axis. The output @@ -263,16 +320,19 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, int axis, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "irfft"); if (n.has_value()) { - return mx::fft::irfft(a, n.value(), axis, s); + return mx::fft::irfft(a, n.value(), axis, fft_norm, s); } else { - return mx::fft::irfft(a, axis, s); + return mx::fft::irfft(a, axis, fft_norm, s); } }, "a"_a, "n"_a = nb::none(), "axis"_a = -1, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( The inverse of :func:`rfft`. @@ -288,6 +348,8 @@ void init_fft(nb::module_& parent_module) { ``a.shape[axis] // 2 + 1``. axis (int, optional): Axis along which to perform the FFT. The default is ``-1``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The real array containing the inverse of :func:`rfft`. @@ -297,21 +359,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "rfft2"); if (axes.has_value() && n.has_value()) { - return mx::fft::rfftn(a, n.value(), axes.value(), s); + return mx::fft::rfftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::rfftn(a, axes.value(), s); + return mx::fft::rfftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[rfft2] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::rfftn(a, s); + return mx::fft::rfftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a.none() = std::vector{-2, -1}, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( Two dimensional real discrete Fourier Transform. @@ -329,6 +394,8 @@ void init_fft(nb::module_& parent_module) { sizes of ``a`` along ``axes``. axes (list(int), optional): Axes along which to perform the FFT. The default is ``[-2, -1]``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The real DFT of the input along the given axes. The output @@ -339,21 +406,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "irfft2"); if (axes.has_value() && n.has_value()) { - return mx::fft::irfftn(a, n.value(), axes.value(), s); + return mx::fft::irfftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::irfftn(a, axes.value(), s); + return mx::fft::irfftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[irfft2] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::irfftn(a, s); + return mx::fft::irfftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a.none() = std::vector{-2, -1}, + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( The inverse of :func:`rfft2`. @@ -372,6 +442,8 @@ void init_fft(nb::module_& parent_module) { sizes of ``a`` along ``axes``. axes (list(int), optional): Axes along which to perform the FFT. The default is ``[-2, -1]``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The real array containing the inverse of :func:`rfft2`. @@ -381,21 +453,24 @@ void init_fft(nb::module_& parent_module) { [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "rfftn"); if (axes.has_value() && n.has_value()) { - return mx::fft::rfftn(a, n.value(), axes.value(), s); + return mx::fft::rfftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::rfftn(a, axes.value(), s); + return mx::fft::rfftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[rfftn] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::rfftn(a, s); + return mx::fft::rfftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a = nb::none(), + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( n-dimensional real discrete Fourier Transform. @@ -414,30 +489,36 @@ void init_fft(nb::module_& parent_module) { axes (list(int), optional): Axes along which to perform the FFT. The default is ``None`` in which case the FFT is over the last ``len(s)`` axes or all axes if ``s`` is also ``None``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The real DFT of the input along the given axes. The output + data type will be complex. )pbdoc"); m.def( "irfftn", [](const mx::array& a, const std::optional& n, const std::optional>& axes, + const std::string& norm, mx::StreamOrDevice s) { + auto fft_norm = parse_norm(norm, "irfftn"); if (axes.has_value() && n.has_value()) { - return mx::fft::irfftn(a, n.value(), axes.value(), s); + return mx::fft::irfftn(a, n.value(), axes.value(), fft_norm, s); } else if (axes.has_value()) { - return mx::fft::irfftn(a, axes.value(), s); + return mx::fft::irfftn(a, axes.value(), fft_norm, s); } else if (n.has_value()) { throw std::invalid_argument( "[irfftn] `axes` should not be `None` if `s` is not `None`."); } else { - return mx::fft::irfftn(a, s); + return mx::fft::irfftn(a, fft_norm, s); } }, "a"_a, "s"_a = nb::none(), "axes"_a = nb::none(), + "norm"_a = "backward", "stream"_a = nb::none(), R"pbdoc( The inverse of :func:`rfftn`. @@ -456,19 +537,61 @@ void init_fft(nb::module_& parent_module) { axes (list(int), optional): Axes along which to perform the FFT. The default is ``None`` in which case the FFT is over the last ``len(s)`` axes or all axes if ``s`` is also ``None``. + norm (str, optional): One of ``"backward"``, ``"ortho"``, or + ``"forward"``. Default is ``"backward"``. Returns: array: The real array containing the inverse of :func:`rfftn`. )pbdoc"); + m.def( + "fftfreq", + [](int n, double d, mx::StreamOrDevice s) { + return mx::fft::fftfreq(n, d, s); + }, + "n"_a, + "d"_a = 1.0, + "stream"_a = nb::none(), + R"pbdoc( + Return the discrete Fourier Transform sample frequencies. + + Args: + n (int): Window length. + d (float, optional): Sample spacing. The default is ``1.0``. + + Returns: + array: The sample frequencies as a one-dimensional array of type ``float32``. + )pbdoc"); + m.def( + "rfftfreq", + [](int n, double d, mx::StreamOrDevice s) { + return mx::fft::rfftfreq(n, d, s); + }, + "n"_a, + "d"_a = 1.0, + "stream"_a = nb::none(), + R"pbdoc( + Return the discrete Fourier Transform sample frequencies + for use with :func:`rfft` and :func:`irfft`. + + The returned array contains the non-negative frequency terms + in the range ``[0, floor(n/2)]``. + + Args: + n (int): Window length. + d (float, optional): Sample spacing. The default is ``1.0``. + + Returns: + array: The sample frequencies as a one-dimensional array of type ``float32``. + )pbdoc"); m.def( "fftshift", - [](const mx::array& a, - const std::optional>& axes, - mx::StreamOrDevice s) { - if (axes.has_value()) { - return mx::fft::fftshift(a, axes.value(), s); - } else { + [](const mx::array& a, const IntOrVec& axes, mx::StreamOrDevice s) { + if (std::holds_alternative(axes)) { return mx::fft::fftshift(a, s); + } else if (auto pv = std::get_if(&axes); pv) { + return mx::fft::fftshift(a, {*pv}, s); + } else { + return mx::fft::fftshift(a, std::get>(axes), s); } }, "a"_a, @@ -479,7 +602,7 @@ void init_fft(nb::module_& parent_module) { Args: a (array): The input array. - axes (list(int), optional): Axes over which to perform the shift. + axes (int or list(int), optional): Axis or axes over which to perform the shift. If ``None``, shift all axes. Returns: @@ -487,13 +610,13 @@ void init_fft(nb::module_& parent_module) { )pbdoc"); m.def( "ifftshift", - [](const mx::array& a, - const std::optional>& axes, - mx::StreamOrDevice s) { - if (axes.has_value()) { - return mx::fft::ifftshift(a, axes.value(), s); - } else { + [](const mx::array& a, const IntOrVec& axes, mx::StreamOrDevice s) { + if (std::holds_alternative(axes)) { return mx::fft::ifftshift(a, s); + } else if (auto pv = std::get_if(&axes); pv) { + return mx::fft::ifftshift(a, {*pv}, s); + } else { + return mx::fft::ifftshift(a, std::get>(axes), s); } }, "a"_a, @@ -505,7 +628,7 @@ void init_fft(nb::module_& parent_module) { Args: a (array): The input array. - axes (list(int), optional): Axes over which to perform the inverse shift. + axes (int or list(int), optional): Axis or axes over which to perform the inverse shift. If ``None``, shift all axes. Returns: diff --git a/Source/Cmlx/mlx/python/src/indexing.cpp b/Source/Cmlx/mlx/python/src/indexing.cpp index 564c4cb45..1dce53788 100644 --- a/Source/Cmlx/mlx/python/src/indexing.cpp +++ b/Source/Cmlx/mlx/python/src/indexing.cpp @@ -3,10 +3,11 @@ #include #include -#include "python/src/convert.h" -#include "python/src/indexing.h" +#include #include "mlx/ops.h" +#include "python/src/convert.h" +#include "python/src/indexing.h" bool is_none_slice(const nb::slice& in_slice) { return ( @@ -15,8 +16,26 @@ bool is_none_slice(const nb::slice& in_slice) { nb::getattr(in_slice, "step").is_none()); } +bool is_index_scalar(const nb::object& obj) { + if (nb::isinstance(obj)) { + return false; + } + if (!PyIndex_Check(obj.ptr())) { + return false; + } + if (nb::ndarray_check(obj) && nb::cast>(obj).ndim() > 0) { + return false; + } + return true; +} + int safe_to_int32(nb::object obj) { - auto val = nb::cast(nb::cast(obj)); + auto idx = nb::steal(PyNumber_Index(obj.ptr())); + if (!idx.is_valid()) { + throw nb::python_error(); + } + + auto val = nb::cast(nb::cast(idx)); if (val > INT32_MAX || val < INT32_MIN) { throw std::invalid_argument("Slice indices must be 32-bit integers."); } @@ -25,7 +44,7 @@ int safe_to_int32(nb::object obj) { int get_slice_int(nb::object obj, int default_val) { if (!obj.is_none()) { - if (!nb::isinstance(obj)) { + if (!is_index_scalar(obj)) { throw std::invalid_argument("Slice indices must be integers or None."); } return safe_to_int32(obj); @@ -60,7 +79,7 @@ mx::array get_int_index(nb::object idx, int axis_size) { } bool is_valid_index_type(const nb::object& obj) { - return nb::isinstance(obj) || nb::isinstance(obj) || + return nb::isinstance(obj) || is_index_scalar(obj) || nb::isinstance(obj) || obj.is_none() || nb::ellipsis().is(obj) || nb::isinstance(obj); } @@ -102,7 +121,7 @@ mx::array mlx_get_item_array(const mx::array& src, const mx::array& indices) { return take(src, indices, 0); } -mx::array mlx_get_item_int(const mx::array& src, const nb::int_& idx) { +mx::array mlx_get_item_int(const mx::array& src, const nb::object& idx) { // Check input and raise error if 0 dim for parity with np if (src.ndim() == 0) { throw std::invalid_argument( @@ -139,7 +158,7 @@ mx::array mlx_gather_nd( gather_indices.push_back(arange(start, end, stride, mx::uint32)); num_slices++; is_slice[i] = true; - } else if (nb::isinstance(idx)) { + } else if (is_index_scalar(idx)) { gather_indices.push_back(get_int_index(idx, src.shape(i))); } else if (nb::isinstance(idx)) { auto arr = nb::cast(idx); @@ -289,7 +308,7 @@ mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) { bool have_non_array = false; bool gather_first = false; for (auto& idx : indices) { - if (nb::isinstance(idx) || (nb::isinstance(idx))) { + if (nb::isinstance(idx) || is_index_scalar(idx)) { if (have_array && have_non_array) { gather_first = true; break; @@ -312,7 +331,7 @@ mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) { // Then find the last array for (last_array = indices.size() - 1; last_array >= 0; last_array--) { auto& idx = indices[last_array]; - if (nb::isinstance(idx) || nb::isinstance(idx)) { + if (nb::isinstance(idx) || is_index_scalar(idx)) { break; } } @@ -348,7 +367,7 @@ mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) { } else { for (int i = 0; i < indices.size(); i++) { auto& idx = indices[i]; - if (nb::isinstance(idx) || nb::isinstance(idx)) { + if (nb::isinstance(idx) || is_index_scalar(idx)) { break; } else if (idx.is_none()) { remaining_indices.push_back(idx); @@ -385,8 +404,8 @@ mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) { int axis = 0; for (auto& idx : remaining_indices) { if (!idx.is_none()) { - if (!have_array && nb::isinstance(idx)) { - int st = nb::cast(idx); + if (!have_array && is_index_scalar(idx)) { + int st = safe_to_int32(idx); st = (st < 0) ? st + src.shape(axis) : st; starts[axis] = st; @@ -419,7 +438,7 @@ mx::array mlx_get_item_nd(mx::array src, const nb::tuple& entries) { auto& idx = remaining_indices[axis]; if (unsqueeze_needed && idx.is_none()) { unsqueeze_axes.push_back(axis - squeeze_axes.size()); - } else if (squeeze_needed && nb::isinstance(idx)) { + } else if (squeeze_needed && is_index_scalar(idx)) { squeeze_axes.push_back(axis - unsqueeze_axes.size()); } } @@ -439,8 +458,8 @@ mx::array mlx_get_item(const mx::array& src, const nb::object& obj) { return mlx_get_item_slice(src, nb::cast(obj)); } else if (nb::isinstance(obj)) { return mlx_get_item_array(src, nb::cast(obj)); - } else if (nb::isinstance(obj)) { - return mlx_get_item_int(src, nb::cast(obj)); + } else if (is_index_scalar(obj)) { + return mlx_get_item_int(src, obj); } else if (nb::isinstance(obj)) { return mlx_get_item_nd(src, nb::cast(obj)); } else if (nb::isinstance(obj)) { @@ -457,7 +476,7 @@ mx::array mlx_get_item(const mx::array& src, const nb::object& obj) { std::tuple, mx::array, std::vector> mlx_scatter_args_int( const mx::array& src, - const nb::int_& idx, + const nb::object& idx, const mx::array& update) { if (src.ndim() == 0) { throw std::invalid_argument( @@ -685,7 +704,7 @@ mlx_scatter_args_nd( // Add the shape to the update update_shape[ax - 1] = 1; } - } else if (nb::isinstance(pyidx)) { + } else if (is_index_scalar(pyidx)) { // Add index to arrays arr_indices.push_back(get_int_index(pyidx, src.shape(ax++))); // Add the shape to the update @@ -755,10 +774,12 @@ mlx_compute_scatter_args( return mlx_scatter_args_slice(src, nb::cast(obj), vals); } else if (nb::isinstance(obj)) { return mlx_scatter_args_array(src, nb::cast(obj), vals); - } else if (nb::isinstance(obj)) { - return mlx_scatter_args_int(src, nb::cast(obj), vals); + } else if (is_index_scalar(obj)) { + return mlx_scatter_args_int(src, obj, vals); } else if (nb::isinstance(obj)) { return mlx_scatter_args_nd(src, nb::cast(obj), vals); + } else if (nb::isinstance(obj)) { + return {{}, broadcast_to(vals, src.shape()), {}}; } else if (obj.is_none()) { return {{}, broadcast_to(vals, src.shape()), {}}; } else if (nb::isinstance(obj)) { @@ -769,59 +790,68 @@ mlx_compute_scatter_args( throw std::invalid_argument("Cannot index mlx array using the given type."); } -auto mlx_slice_update( +std::tuple, mx::Shape, mx::Shape, mx::Shape> +mlx_compute_slice_update_args( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + // Build the slice params + mx::Shape starts(src.ndim(), 0); + mx::Shape stops = src.shape(); + mx::Shape strides(src.ndim(), 1); + // Can't route to slice update if not slice, tuple, or int - if (src.ndim() == 0 || nb::isinstance(obj) || + if (src.ndim() == 0 || (!nb::isinstance(obj) && !nb::isinstance(obj) && - !nb::isinstance(obj))) { - return std::make_pair(false, src); + !is_index_scalar(obj))) { + return std::make_tuple( + std::nullopt, std::move(starts), std::move(stops), std::move(strides)); } if (nb::isinstance(obj)) { // Can't route to slice update if any arrays are present for (auto idx : nb::cast(obj)) { if (nb::isinstance(idx) || nb::isinstance(idx)) { - return std::make_pair(false, src); + return std::make_tuple( + std::nullopt, + std::move(starts), + std::move(stops), + std::move(strides)); } } } - // Should be able to route to slice update - // Pre process tuple - auto upd = to_array(v, src.dtype()); + // Should be able to route to slice update just extract the update value and + // and the slice arguments. + + // Cast v to an array and ensure it is the right type + auto update = to_array(v, src.dtype()); // Remove extra leading singletons dimensions from the update int s = 0; - for (; s < static_cast(upd.ndim()) - 1 && upd.shape(s) == 1 && - (upd.ndim() - s) > src.ndim(); + for (; s < static_cast(update.ndim()) - 1 && update.shape(s) == 1 && + (update.ndim() - s) > src.ndim(); s++) { }; auto squeeze_axes = std::vector(s); std::iota(squeeze_axes.begin(), squeeze_axes.end(), 0); - auto up = mx::squeeze(upd, squeeze_axes); + update = mx::squeeze(update, squeeze_axes); - // Build slice update params - mx::Shape starts(src.ndim(), 0); - mx::Shape stops = src.shape(); - mx::Shape strides(src.ndim(), 1); - if (nb::isinstance(obj)) { + // Single int then make it a slice of size 1 + if (is_index_scalar(obj)) { if (src.ndim() < 1) { std::ostringstream msg; msg << "Too many indices for array with " << src.ndim() << " dimensions."; throw std::invalid_argument(msg.str()); } - auto idx = nb::cast(obj); + auto idx = safe_to_int32(obj); idx = idx < 0 ? idx + stops[0] : idx; starts[0] = idx; stops[0] = idx + 1; - auto out = slice_update( - src, up, std::move(starts), std::move(stops), std::move(strides)); - return std::make_pair(true, out); + return std::make_tuple( + update, std::move(starts), std::move(stops), std::move(strides)); } - // If it's just a simple slice, just do a slice update and return + // Simple slice, just extract it into the first dim if (nb::isinstance(obj)) { // Read slice arguments get_slice_params( @@ -830,16 +860,14 @@ auto mlx_slice_update( strides[0], nb::cast(obj), src.shape(0)); - - // Do slice update - auto out = slice_update(src, up, starts, stops, strides); - return std::make_pair(true, out); + return std::make_tuple( + update, std::move(starts), std::move(stops), std::move(strides)); } // It must be a tuple auto entries = nb::cast(obj); - // Expand ellipses into a series of ':' slices + // Expand ellipsis into a series of ':' slices auto [non_none_indices, indices] = mlx_expand_ellipsis(src.shape(), entries); // Dimension check @@ -851,15 +879,20 @@ auto mlx_slice_update( // If no non-None indices return the broadcasted update if (non_none_indices == 0) { - return std::make_pair(true, broadcast_to(up, src.shape())); + return std::make_tuple( + broadcast_to(update, src.shape()), + std::move(starts), + std::move(stops), + std::move(strides)); } + // Parse the update slice int unspecified = src.ndim() - non_none_indices; std::vector squeeze_dims; std::vector expand_dims; for (int i = indices.size() - 1, ax = non_none_indices - 1, - upd_ax = upd.ndim() - unspecified - 1; + upd_ax = update.ndim() - unspecified - 1; i >= 0; --i) { auto& pyidx = indices[i]; @@ -872,8 +905,8 @@ auto mlx_slice_update( src.shape(ax)); ax--; upd_ax--; - } else if (nb::isinstance(pyidx)) { - int st = nb::cast(pyidx); + } else if (is_index_scalar(pyidx)) { + int st = safe_to_int32(pyidx); st = (st < 0) ? st + src.shape(i) : st; starts[ax] = st; stops[ax] = st + 1; @@ -887,15 +920,15 @@ auto mlx_slice_update( } } } + update = mx::squeeze( + mx::expand_dims(update, std::move(expand_dims)), std::move(squeeze_dims)); - up = mx::squeeze( - mx::expand_dims(up, std::move(expand_dims)), std::move(squeeze_dims)); - auto out = slice_update(src, up, starts, stops, strides); - return std::make_pair(true, out); + return std::make_tuple( + update, std::move(starts), std::move(stops), std::move(strides)); } std::optional extract_boolean_mask(const nb::object& obj) { - using NDArray = nb::ndarray; + using NDArray = nb::ndarray; if (nb::isinstance(obj)) { return mx::array(nb::cast(obj), mx::bool_); } else if (nb::isinstance(obj)) { @@ -921,9 +954,11 @@ void mlx_set_item( mx::array& src, const nb::object& obj, const ScalarOrArray& v) { - auto [success, out] = mlx_slice_update(src, obj, v); - if (success) { - src.overwrite_descriptor(out); + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + src.overwrite_descriptor( + slice_update(src, *update, starts, stops, strides)); return; } @@ -947,6 +982,12 @@ mx::array mlx_add_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_add(src, *update, starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_add(src, indices, updates, axes); @@ -959,6 +1000,12 @@ mx::array mlx_subtract_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_add(src, -(*update), starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_add(src, indices, -updates, axes); @@ -971,6 +1018,12 @@ mx::array mlx_multiply_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_prod(src, *update, starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_prod(src, indices, updates, axes); @@ -983,6 +1036,12 @@ mx::array mlx_divide_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_prod(src, reciprocal(*update), starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_prod(src, indices, reciprocal(updates), axes); @@ -995,6 +1054,12 @@ mx::array mlx_maximum_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_max(src, *update, starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_max(src, indices, updates, axes); @@ -1007,6 +1072,12 @@ mx::array mlx_minimum_item( const mx::array& src, const nb::object& obj, const ScalarOrArray& v) { + auto [update, starts, stops, strides] = + mlx_compute_slice_update_args(src, obj, v); + if (update) { + return slice_update_min(src, *update, starts, stops, strides); + } + auto [indices, updates, axes] = mlx_compute_scatter_args(src, obj, v); if (indices.size() > 0) { return scatter_min(src, indices, updates, axes); diff --git a/Source/Cmlx/mlx/python/src/linalg.cpp b/Source/Cmlx/mlx/python/src/linalg.cpp index 58f9be776..a5e82ae85 100644 --- a/Source/Cmlx/mlx/python/src/linalg.cpp +++ b/Source/Cmlx/mlx/python/src/linalg.cpp @@ -55,7 +55,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def norm(a: array, /, ord: Union[None, int, float, str] = None, axis: Union[None, int, list[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def norm(a: array, /, ord: None | int | float | str = None, axis: None | int | list[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix or vector norm. @@ -177,7 +177,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def qr(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def qr(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( The QR factorization of the input matrix. @@ -220,7 +220,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def svd(a: array, compute_uv: bool = True, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"), + "def svd(a: array, compute_uv: bool = True, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( The Singular Value Decomposition (SVD) of the input matrix. @@ -246,8 +246,7 @@ void init_linalg(nb::module_& parent_module) { "a"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def inv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def inv(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a square matrix. @@ -271,7 +270,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tri_inv(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tri_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a triangular square matrix. @@ -296,7 +295,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cholesky(a: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cholesky(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the Cholesky decomposition of a real symmetric positive semi-definite matrix. @@ -326,7 +325,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cholesky_inv(L: array, upper: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cholesky_inv(a: array, upper: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the inverse of a real symmetric positive semi-definite matrix using it's Cholesky decomposition. @@ -347,7 +346,8 @@ void init_linalg(nb::module_& parent_module) { If the input matrix is not a triangular matrix behaviour is undefined. Args: - L (array): Input array. + a (array): Input array. This is the Cholesky factor + :math:`\mathbf{L}`, not :math:`\mathbf{A}` itself. upper (bool, optional): If ``True``, return the upper triangular Cholesky factor. If ``False``, return the lower triangular Cholesky factor. Default: ``False``. stream (Stream, optional): Stream or device. Defaults to ``None`` @@ -362,8 +362,7 @@ void init_linalg(nb::module_& parent_module) { "a"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def pinv(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def pinv(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the (Moore-Penrose) pseudo-inverse of a matrix. @@ -389,7 +388,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cross(a: array, b: array, axis: int = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cross(a: array, b: array, axis: int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the cross product of two arrays along a specified axis. @@ -448,7 +447,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eig(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def eig(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Compute the eigenvalues and eigenvectors of a square matrix. @@ -526,7 +525,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eigh(a: array, UPLO: str = 'L', *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def eigh(a: array, UPLO: str = 'L', *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Compute the eigenvalues and eigenvectors of a complex Hermitian or real symmetric matrix. @@ -572,7 +571,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def lu(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array, array]"), + "def lu(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( Compute the LU factorization of the given matrix ``A``. @@ -603,7 +602,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def lu_factor(a: array, *, stream: Union[None, Stream, Device] = None) -> Tuple[array, array]"), + "def lu_factor(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), R"pbdoc( Computes a compact representation of the LU factorization. @@ -623,7 +622,7 @@ void init_linalg(nb::module_& parent_module) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def solve(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def solve(a: array, b: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the solution to a system of linear equations ``AX = B``. @@ -645,7 +644,7 @@ void init_linalg(nb::module_& parent_module) { "upper"_a = false, "stream"_a = nb::none(), nb::sig( - "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def solve_triangular(a: array, b: array, *, upper: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Computes the solution of a triangular system of linear equations ``AX = B``. @@ -660,4 +659,78 @@ void init_linalg(nb::module_& parent_module) { Returns: array: The unique solution to the system ``AX = B``. )pbdoc"); + + m.def( + "det", + &mx::linalg::det, + "a"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig("def det(a: array, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + 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. + + Args: + a (array): Input array. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + array: The determinant(s) of the input matrix (matrices). + + Example: + >>> A = mx.array([[1., 2.], [3., 4.]]) + >>> mx.linalg.det(A, stream=mx.cpu) + array(-2, dtype=float32) + )pbdoc"); + + m.def( + "slogdet", + [](const mx::array& a, mx::StreamOrDevice s) { + auto result = mx::linalg::slogdet(a, s); + return nb::make_tuple(result.first, result.second); + }, + "a"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def slogdet(a: array, *, stream: StreamOrDevice = None) -> tuple[array, array]"), + R"pbdoc( + 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. + + Args: + a (array): Input array. + stream (Stream, optional): Stream or device. Defaults to ``None`` + in which case the default stream of the default device is used. + + Returns: + tuple(array, array): The ``sign`` and ``logabsdet`` of the + determinant. ``sign`` is -1, 0, or +1. ``logabsdet`` is the + natural log of the absolute value of the determinant. + + Example: + >>> A = mx.array([[1., 2.], [3., 4.]]) + >>> sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + >>> sign + array(-1, dtype=float32) + >>> logabsdet + array(0.693147, dtype=float32) + )pbdoc"); + // Array API standard aliases (https://data-apis.org/array-api/latest/). + parent_module.attr("matrix_norm") = m.attr("norm"); } diff --git a/Source/Cmlx/mlx/python/src/load.cpp b/Source/Cmlx/mlx/python/src/load.cpp index f98307f19..605049864 100644 --- a/Source/Cmlx/mlx/python/src/load.cpp +++ b/Source/Cmlx/mlx/python/src/load.cpp @@ -162,8 +162,12 @@ class PyFileReader : public mx::io::Reader { private: void _read(char* data, size_t n) { - auto memview = PyMemoryView_FromMemory(data, n, PyBUF_WRITE); - nb::object bytes_read = readinto_func_(nb::handle(memview)); + nb::object memview = + nb::steal(PyMemoryView_FromMemory(data, n, PyBUF_WRITE)); + if (!memview.is_valid()) { + throw std::runtime_error("[load] Failed to create memoryview for read"); + } + nb::object bytes_read = readinto_func_(memview); if (bytes_read.is_none() || nb::cast(bytes_read) < n) { throw std::runtime_error("[load] Failed to read from python stream"); @@ -374,9 +378,12 @@ class PyFileWriter : public mx::io::Writer { void write(const char* data, size_t n) override { nb::gil_scoped_acquire gil; - auto memview = - PyMemoryView_FromMemory(const_cast(data), n, PyBUF_READ); - nb::object bytes_written = write_func_(nb::handle(memview)); + nb::object memview = nb::steal( + PyMemoryView_FromMemory(const_cast(data), n, PyBUF_READ)); + if (!memview.is_valid()) { + throw std::runtime_error("[load] Failed to create memoryview for write"); + } + nb::object bytes_written = write_func_(memview); if (bytes_written.is_none() || nb::cast(bytes_written) < n) { throw std::runtime_error("[load] Failed to write to python stream"); diff --git a/Source/Cmlx/mlx/python/src/mlx.cpp b/Source/Cmlx/mlx/python/src/mlx.cpp index 2829b3219..243449b38 100644 --- a/Source/Cmlx/mlx/python/src/mlx.cpp +++ b/Source/Cmlx/mlx/python/src/mlx.cpp @@ -23,6 +23,7 @@ void init_constants(nb::module_&); void init_fast(nb::module_&); void init_distributed(nb::module_&); void init_export(nb::module_&); +void init_print(nb::module_&); NB_MODULE(core, m) { m.doc() = "mlx: A framework for machine learning on Apple silicon."; @@ -30,6 +31,10 @@ NB_MODULE(core, m) { auto reprlib_fix = nb::module_::import_("mlx._reprlib_fix"); nb::set_leak_warnings(false); + auto array_namespace_info = nb::module_::import_("mlx.__array_api_info"); + m.attr("__array_namespace_info__") = + array_namespace_info.attr("__array_namespace_info__"); + init_mlx_func(m); init_device(m); init_stream(m); @@ -46,6 +51,7 @@ NB_MODULE(core, m) { init_fast(m); init_distributed(m); init_export(m); + init_print(m); m.attr("__version__") = mx::version(); } diff --git a/Source/Cmlx/mlx/python/src/mlx_func.cpp b/Source/Cmlx/mlx/python/src/mlx_func.cpp index 09aceabe9..9955e134f 100644 --- a/Source/Cmlx/mlx/python/src/mlx_func.cpp +++ b/Source/Cmlx/mlx/python/src/mlx_func.cpp @@ -2,6 +2,8 @@ #include "python/src/mlx_func.h" +#include + // A garbage collected function which wraps nb::cpp_function // See https://github.com/wjakob/nanobind/discussions/919 diff --git a/Source/Cmlx/mlx/python/src/ops.cpp b/Source/Cmlx/mlx/python/src/ops.cpp index de51ad502..677d998af 100644 --- a/Source/Cmlx/mlx/python/src/ops.cpp +++ b/Source/Cmlx/mlx/python/src/ops.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. +#include #include #include #include @@ -15,6 +16,7 @@ #include "mlx/einsum.h" #include "mlx/ops.h" #include "mlx/utils.h" +#include "python/src/convert.h" #include "python/src/load.h" #include "python/src/small_vector.h" #include "python/src/utils.h" @@ -23,11 +25,14 @@ namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; -using Scalar = std::variant; +using Scalar = std::variant; mx::Dtype scalar_to_dtype(Scalar s) { - if (std::holds_alternative(s)) { - return mx::int32; + if (auto pv = std::get_if(&s); pv) { + return (*pv > std::numeric_limits::max() || + *pv < std::numeric_limits::min()) + ? mx::int64 + : mx::int32; } else if (std::holds_alternative(s)) { return mx::float32; } else { @@ -36,7 +41,7 @@ mx::Dtype scalar_to_dtype(Scalar s) { } double scalar_to_double(Scalar s) { - if (auto pv = std::get_if(&s); pv) { + if (auto pv = std::get_if(&s); pv) { return static_cast(*pv); } else if (auto pv = std::get_if(&s); pv) { return *pv; @@ -45,6 +50,13 @@ double scalar_to_double(Scalar s) { } } +mx::Shape to_shape(const nb::object& shape) { + if (nb::isinstance(shape)) { + return {check_shape_dim(nb::cast(shape))}; + } + return nb::cast(shape); +} + void init_ops(nb::module_& m) { m.def( "reshape", @@ -55,7 +67,7 @@ void init_ops(nb::module_& m) { "stream"_a = nb::none(), nb::sig( "def reshape(a: array, /, shape: Sequence[int], *, stream: " - "Union[None, Stream, Device] = None) -> array"), + "StreamOrDevice = None) -> array"), R"pbdoc( Reshape an array while preserving the size. @@ -83,7 +95,7 @@ void init_ops(nb::module_& m) { "stream"_a = nb::none(), nb::sig( "def flatten(a: array, /, start_axis: int = 0, end_axis: int = " - "-1, *, stream: Union[None, Stream, Device] = None) -> array"), + "-1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Flatten an array. @@ -119,7 +131,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def unflatten(a: array, /, axis: int, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Unflatten an axis of an array to a shape. @@ -156,8 +168,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def squeeze(a: array, /, axis: Union[None, int, Sequence[int]] = " - "None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def squeeze(a: array, /, axis: None | int | Sequence[int] = " + "None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Remove length one axes from an array. @@ -169,6 +181,59 @@ void init_ops(nb::module_& m) { Returns: array: The output array with size one axes removed. )pbdoc"); + m.def( + "flip", + [](const mx::array& a, const IntOrVec& v, const mx::StreamOrDevice& s) { + if (std::holds_alternative(v)) { + return mx::flip(a, s); + } else if (auto pv = std::get_if(&v); pv) { + return mx::flip(a, *pv, s); + } else { + return mx::flip(a, std::get>(v), s); + } + }, + nb::arg(), + "axis"_a = nb::none(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def flip(a: array, /, axis: None | int | Sequence[int] = None, " + "*, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Reverse the order of elements along the given axis. + + Args: + a (array): Input array. + axis (int or tuple(int), optional): Axis or axes to flip over. + Defaults to ``None`` in which case all axes are flipped. + + Returns: + array: The flipped array. + )pbdoc"); + m.def( + "unstack", + [](const mx::array& a, int axis, mx::StreamOrDevice s) { + return mx::unstack(a, axis, s); + }, + nb::arg(), + nb::kw_only(), + "axis"_a = 0, + "stream"_a = nb::none(), + nb::sig( + "def unstack(x: array, /, *, axis: int = 0, stream: StreamOrDevice = None) -> list[array]"), + R"pbdoc( + Split an array into a sequence of arrays along the given axis. + + The inverse of :func:`stack`. The given axis is removed from each of + the returned arrays. + + Args: + x (array): Input array. + axis (int, optional): Axis along which to unstack. Default: ``0``. + + Returns: + list(array): A list of arrays, one for each index along ``axis``. + )pbdoc"); m.def( "expand_dims", [](const mx::array& a, @@ -185,8 +250,8 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def expand_dims(a: array, /, axis: Union[int, Sequence[int]], " - "*, stream: Union[None, Stream, Device] = None) -> array"), + "def expand_dims(a: array, /, axis: int | Sequence[int], " + "*, stream: StreamOrDevice = None) -> array"), R"pbdoc( Add a size one dimension at the given axis. @@ -206,7 +271,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def abs(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def abs(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise absolute value. @@ -225,7 +290,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sign(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sign(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise sign. @@ -235,6 +300,23 @@ void init_ops(nb::module_& m) { Returns: array: The sign of ``a``. )pbdoc"); + m.def( + "positive", + &mx::positive, + nb::arg(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def positive(a: array, /, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Element-wise unary plus. Returns a copy of the input. + + Args: + a (array): Input array. + + Returns: + array: A copy of ``a``. + )pbdoc"); m.def( "negative", [](const ScalarOrArray& a, mx::StreamOrDevice s) { @@ -244,7 +326,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def negative(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def negative(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise negation. @@ -267,7 +349,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def add(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def add(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise addition. @@ -294,7 +376,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def subtract(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def subtract(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise subtraction. @@ -321,7 +403,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def multiply(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def multiply(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise multiplication. @@ -348,7 +430,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise division. @@ -375,7 +457,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def divmod(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def divmod(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise quotient and remainder. @@ -403,7 +485,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def floor_divide(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def floor_divide(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise integer division. @@ -430,7 +512,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def remainder(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def remainder(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise remainder of division. @@ -458,7 +540,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise equality. @@ -485,7 +567,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def not_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def not_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise not equal. @@ -512,7 +594,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def less(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def less(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise less than. @@ -539,7 +621,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def less_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def less_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise less than or equal. @@ -566,7 +648,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def greater(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def greater(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise greater than. @@ -593,7 +675,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def greater_equal(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def greater_equal(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise greater or equal. @@ -622,7 +704,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def array_equal(a: Union[scalar, array], b: Union[scalar, array], equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def array_equal(a: scalar | array, b: scalar | array, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Array equality check. @@ -647,7 +729,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def matmul(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def matmul(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication. @@ -671,6 +753,23 @@ void init_ops(nb::module_& m) { Returns: array: The matrix product of ``a`` and ``b``. )pbdoc"); + m.def( + "trunc", + &mx::trunc, + nb::arg(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def trunc(a: array, /, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Element-wise truncation towards zero. + + Args: + a (array): Input array. + + Returns: + array: The truncated array. + )pbdoc"); m.def( "square", [](const ScalarOrArray& a, mx::StreamOrDevice s) { @@ -680,7 +779,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def square(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def square(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise square. @@ -699,7 +798,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise square root. @@ -718,7 +817,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def rsqrt(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def rsqrt(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise reciprocal and square root. @@ -737,7 +836,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def reciprocal(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def reciprocal(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise reciprocal. @@ -756,7 +855,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_not(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_not(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical not. @@ -776,7 +875,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_and(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_and(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical and. @@ -798,7 +897,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logical_or(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logical_or(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logical or. @@ -809,6 +908,27 @@ void init_ops(nb::module_& m) { Returns: array: The boolean array containing the logical or of ``a`` and ``b``. )pbdoc"); + m.def( + "logical_xor", + [](const ScalarOrArray& a, const ScalarOrArray& b, mx::StreamOrDevice s) { + return mx::logical_xor(to_array(a), to_array(b), s); + }, + nb::arg(), + nb::arg(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def logical_xor(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Element-wise logical exclusive or. + + Args: + a (array): First input array or scalar. + b (array): Second input array or scalar. + + Returns: + array: The boolean array containing the logical xor of ``a`` and ``b``. + )pbdoc"); m.def( "logaddexp", [](const ScalarOrArray& a_, @@ -822,14 +942,14 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logaddexp(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logaddexp(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise log-add-exp. This is a numerically stable log-add-exp of two arrays with numpy-style broadcasting semantics. Either or both input arrays can also be scalars. - The computation is is a numerically stable version of ``log(exp(a) + exp(b))``. + The computation is a numerically stable version of ``log(exp(a) + exp(b))``. Args: a (array): Input array or scalar. @@ -847,7 +967,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def exp(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def exp(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise exponential. @@ -866,7 +986,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def expm1(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def expm1(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise exponential minus 1. @@ -887,7 +1007,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def erf(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def erf(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise error function. @@ -909,7 +1029,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def erfinv(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def erfinv(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse of :func:`erf`. @@ -928,7 +1048,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sin(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise sine. @@ -947,7 +1067,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cos(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise cosine. @@ -966,7 +1086,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tan(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise tangent. @@ -985,7 +1105,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arcsin(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arcsin(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse sine. @@ -1004,7 +1124,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arccos(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arccos(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse cosine. @@ -1023,7 +1143,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctan(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctan(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse tangent. @@ -1041,7 +1161,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctan2(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctan2(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse tangent of the ratio of two arrays. @@ -1061,7 +1181,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sinh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic sine. @@ -1080,7 +1200,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def cosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def cosh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic cosine. @@ -1099,7 +1219,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tanh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise hyperbolic tangent. @@ -1118,7 +1238,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arcsinh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arcsinh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic sine. @@ -1137,7 +1257,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arccosh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arccosh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic cosine. @@ -1156,7 +1276,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arctanh(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arctanh(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise inverse hyperbolic tangent. @@ -1175,7 +1295,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def degrees(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def degrees(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from radians to degrees. @@ -1194,7 +1314,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def radians(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def radians(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert angles from degrees to radians. @@ -1213,7 +1333,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise natural logarithm. @@ -1232,7 +1352,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log2(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log2(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise base-2 logarithm. @@ -1251,7 +1371,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log10(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log10(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise base-10 logarithm. @@ -1270,7 +1390,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def log1p(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def log1p(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise natural log of one plus the array. @@ -1287,7 +1407,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def stop_gradient(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def stop_gradient(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Stop gradients from being computed. @@ -1311,7 +1431,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sigmoid(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sigmoid(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise logistic sigmoid. @@ -1339,7 +1459,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def power(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def power(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise power operation. @@ -1356,21 +1476,25 @@ void init_ops(nb::module_& m) { m.def( "arange", [](Scalar start, - Scalar stop, + std::optional stop, const std::optional& step, const std::optional& dtype_, mx::StreamOrDevice s) { + if (!stop) { + stop = start; + start = 0; + } // Determine the final dtype based on input types mx::Dtype dtype = dtype_ ? *dtype_ : mx::promote_types( scalar_to_dtype(start), step ? mx::promote_types( - scalar_to_dtype(stop), scalar_to_dtype(*step)) - : scalar_to_dtype(stop)); + scalar_to_dtype(*stop), scalar_to_dtype(*step)) + : scalar_to_dtype(*stop)); return mx::arange( scalar_to_double(start), - scalar_to_double(stop), + scalar_to_double(*stop), step ? scalar_to_double(*step) : 1.0, dtype, s); @@ -1382,7 +1506,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arange(start : Union[int, float], stop : Union[int, float], step : Union[None, int, float], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def arange(start : int | float, stop : None | int | float, step : None | int | float, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generates ranges of numbers. @@ -1391,9 +1515,9 @@ void init_ops(nb::module_& m) { Args: start (float or int, optional): Starting value which defaults to ``0``. - stop (float or int): Stopping value. + stop (float or int, optional): Stopping value. step (float or int, optional): Increment which defaults to ``1``. - dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``. + dtype (Dtype, optional): Specifies the data type of the output. If unspecified will default to ``float32`` if any of ``start``, ``stop``, or ``step`` are ``float``. Otherwise will default to ``int32``, or ``int64`` if any of ``start``, ``stop``, or ``step`` does not fit in ``int32``. Returns: array: The range of values. @@ -1427,28 +1551,122 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def arange(stop : Union[int, float], step : Union[None, int, float] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array")); + "def arange(stop : int | float, step : None | int | float = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array")); + m.def( + "bartlett", + &mlx::core::bartlett, + "M"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"pbdoc( + Return the Bartlett window. + + The Bartlett window is a taper formed by using a weighted cosine. + + .. math:: + w(n) = 1 - \frac{2|n - (M-1)/2|}{M-1} + \qquad 0 \le n \le M-1 + + Args: + M (int): Number of points in the output window. + + Returns: + array: The window, with the maximum value normalized to one (the value one + appears only if the number of samples is odd). + )pbdoc"); + m.def( + "hanning", + &mlx::core::hanning, + "M"_a, + nb::kw_only(), + "stream"_a = nb::none(), + R"pbdoc( + Return the Hanning window. + + The Hanning window is a taper formed by using a weighted cosine. + + .. math:: + w(n) = 0.5 - 0.5 \cos\left(\frac{2\pi n}{M-1}\right) + \qquad 0 \le n \le M-1 + + Args: + M (int): Number of points in the output window. + + Returns: + array: The window, with the maximum value normalized to one (the value one + appears only if the number of samples is odd). + )pbdoc"); + m.def( + "hamming", + &mlx::core::hamming, + "M"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig("def hamming(M: int, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Return the Hamming window. + + The Hamming window is a taper formed by using a weighted cosine. + + .. math:: + w(n) = 0.54 - 0.46 \cos\left(\frac{2\pi n}{M-1}\right) + \qquad 0 \le n \le M-1 + + Args: + M (int): Number of points in the output window. + + Returns: + array: The window, with the maximum value normalized to one (the value one + appears only if the number of samples is odd). + )pbdoc"); + m.def( + "blackman", + &mlx::core::blackman, + "M"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def blackman(M: int, *, stream: StreamOrDevice = None) -> array"), // <--- J'ai rajouté ça + R"pbdoc( + Return the Blackman window. + + The Blackman window is a taper formed by using the first three terms of a summation of cosines. + + .. math:: + w(n) = 0.42 - 0.5 \cos\left(\frac{2\pi n}{M-1}\right) + 0.08 \cos\left(\frac{4\pi n}{M-1}\right) + \qquad 0 \le n \le M-1 + + Args: + M (int): Number of points in the output window. + + Returns: + array: The window, with the maximum value normalized to one (the value one + appears only if the number of samples is odd). + )pbdoc"); m.def( "linspace", [](Scalar start, Scalar stop, int num, + bool endpoint, std::optional dtype, mx::StreamOrDevice s) { return mx::linspace( scalar_to_double(start), scalar_to_double(stop), num, + endpoint, dtype.value_or(mx::float32), s); }, "start"_a, "stop"_a, "num"_a = 50, + "endpoint"_a = true, "dtype"_a.none() = mx::float32, "stream"_a = nb::none(), nb::sig( - "def linspace(start: scalar, stop: scalar, num: Optional[int] = 50, dtype: Optional[Dtype] = float32, stream: Union[None, Stream, Device] = None) -> array"), + "def linspace(start: scalar, stop: scalar, num: int | None = 50, endpoint: bool = True, dtype: Dtype | None = float32, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate ``num`` evenly spaced numbers over interval ``[start, stop]``. @@ -1456,6 +1674,9 @@ void init_ops(nb::module_& m) { start (scalar): Starting value. stop (scalar): Stopping value. num (int, optional): Number of samples, defaults to ``50``. + endpoint (bool, optional): If ``True``, ``stop`` is the last + sample. Otherwise it is not included and the samples are spaced + over the half-open interval ``[start, stop)``. Default: ``True``. dtype (Dtype, optional): Specifies the data type of the output, default to ``float32``. @@ -1470,14 +1691,14 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def kron(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def kron(a: array, b: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the Kronecker product of two arrays ``a`` and ``b``. Args: a (array): The first input array. b (array): The second input array. - stream (Union[None, Stream, Device], optional): Optional stream or + stream (StreamOrDevice, optional): Optional stream or device for execution. Default: ``None``. Returns: @@ -1514,7 +1735,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def take(a: array, /, indices: Union[int, array], axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def take(a: array, /, indices: int | array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Take elements along an axis. @@ -1547,11 +1768,11 @@ void init_ops(nb::module_& m) { }, nb::arg(), "indices"_a, - "axis"_a.none(), + "axis"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def take_along_axis(a: array, /, indices: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def take_along_axis(a: array, /, indices: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Take values along an axis at the specified indices. @@ -1590,7 +1811,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def put_along_axis(a: array, /, indices: array, values: array, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def put_along_axis(a: array, /, indices: array, values: array, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Put values along an axis at the specified indices. @@ -1610,15 +1831,11 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "full", - [](const std::variant& shape, + [](const nb::object& shape, const ScalarOrArray& vals, std::optional dtype, mx::StreamOrDevice s) { - if (auto pv = std::get_if(&shape); pv) { - return mx::full({*pv}, to_array(vals, dtype), s); - } else { - return mx::full(std::get(shape), to_array(vals, dtype), s); - } + return mx::full(to_shape(shape), to_array(vals, dtype), s); }, "shape"_a, "vals"_a, @@ -1626,7 +1843,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def full(shape: Union[int, Sequence[int]], vals: Union[scalar, array], dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def full(shape: int | Sequence[int], vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array with the given value. @@ -1642,24 +1859,48 @@ void init_ops(nb::module_& m) { Returns: array: The output array with the specified shape and values. )pbdoc"); + m.def( + "full_like", + [](const mx::array& a, + const ScalarOrArray& vals, + std::optional dtype, + mx::StreamOrDevice s) { + auto t = dtype.value_or(a.dtype()); + return mx::full_like(a, to_array(vals, t), t, s); + }, + nb::arg(), + "vals"_a, + "dtype"_a = nb::none(), + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def full_like(a: array, vals: scalar | array, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + An array filled with ``vals`` with the same shape as the input. + + Args: + a (array): The input to take the shape from. + vals (float or int or array): Values to fill the array with. + dtype (Dtype, optional): Data type of the output array. If + unspecified the type of the input is used. + + Returns: + array: The output array. + )pbdoc"); m.def( "zeros", - [](const std::variant& shape, + [](const nb::object& shape, std::optional dtype, mx::StreamOrDevice s) { auto t = dtype.value_or(mx::float32); - if (auto pv = std::get_if(&shape); pv) { - return mx::zeros({*pv}, t, s); - } else { - return mx::zeros(std::get(shape), t, s); - } + return mx::zeros(to_shape(shape), t, s); }, "shape"_a, "dtype"_a.none() = mx::float32, nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def zeros(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def zeros(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array of zeros. @@ -1673,59 +1914,95 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "asarray", - [](const ArrayInitType& a, std::optional dtype) { - return create_array(a, dtype); - }, + [](const nb::object& a, + std::optional dtype, + std::optional copy) { return create_array(a, dtype, copy); }, nb::arg(), "dtype"_a = nb::none(), + nb::kw_only(), + "copy"_a = nb::none(), nb::sig( - "def asarray(a: Union[scalar, array, Sequence], dtype: " - "Optional[Dtype] = None) -> array"), + "def asarray(a: scalar | array | Sequence | DLPackCompatible, dtype: " + "Dtype | None = None, *, copy: bool | None = None) -> array"), R"pbdoc( Convert the input to an array. Args: a: Input data. dtype (Dtype, optional): The desired data-type for the array. + copy (bool, optional): Whether to copy the input. If ``True``, + always copy. If ``False``, never copy. If ``None``, share memory + when possible and copy otherwise. Zero-copy DLPack imports + preserve the DLPack strides. Returns: array: An array interpretation of the input. + + Raises: + ValueError: If ``copy`` is ``False`` and a copy is required. + )pbdoc"); + m.def( + "from_dlpack", + [](nb::ndarray x, std::optional copy) { + return nd_array_to_mlx(x, std::nullopt, std::nullopt, copy); + }, + nb::arg(), + nb::kw_only(), + "copy"_a = nb::none(), + nb::sig( + "def from_dlpack(x: DLPackCompatible, /, *, copy: bool | None = None) -> array"), + R"pbdoc( + Create an array from an object that supports DLPack. + + Args: + x: Input object implementing ``__dlpack__`` and + ``__dlpack_device__``. + copy (bool, optional): Whether to copy the input. If ``True``, + always copy. If ``False``, never copy. If ``None``, share memory + when possible and copy otherwise. Zero-copy imports preserve the + DLPack strides. + + Returns: + array: An array containing the input data. )pbdoc"); m.def( "zeros_like", - &mx::zeros_like, + [](const mx::array& a, + std::optional dtype, + mx::StreamOrDevice s) { + return mx::zeros_like(a, dtype.value_or(a.dtype()), s); + }, nb::arg(), + "dtype"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def zeros_like(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def zeros_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array of zeros like the input. Args: - a (array): The input to take the shape and type from. + a (array): The input to take the shape from. + dtype (Dtype, optional): Output data type. If ``None``, the output + type defaults to the input array's data type. Returns: array: The output array filled with zeros. )pbdoc"); m.def( "ones", - [](const std::variant& shape, + [](const nb::object& shape, std::optional dtype, mx::StreamOrDevice s) { auto t = dtype.value_or(mx::float32); - if (auto pv = std::get_if(&shape); pv) { - return mx::ones({*pv}, t, s); - } else { - return mx::ones(std::get(shape), t, s); - } + return mx::ones(to_shape(shape), t, s); }, "shape"_a, "dtype"_a.none() = mx::float32, nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ones(shape: Union[int, Sequence[int]], dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ones(shape: int | Sequence[int], dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array of ones. @@ -1739,17 +2016,24 @@ void init_ops(nb::module_& m) { )pbdoc"); m.def( "ones_like", - &mx::ones_like, + [](const mx::array& a, + std::optional dtype, + mx::StreamOrDevice s) { + return mx::ones_like(a, dtype.value_or(a.dtype()), s); + }, nb::arg(), + "dtype"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ones_like(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ones_like(a: array, /, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array of ones like the input. Args: - a (array): The input to take the shape and type from. + a (array): The input to take the shape from. + dtype (Dtype, optional): Output data type. If ``None``, the output + type defaults to the input array's data type. Returns: array: The output array filled with ones. @@ -1770,7 +2054,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def eye(n: int, m: Optional[int] = None, k: int = 0, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def eye(n: int, m: int | None = None, k: int = 0, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create an identity matrix or a general diagonal matrix. @@ -1794,7 +2078,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def identity(n: int, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) -> array"), + "def identity(n: int, dtype: Dtype | None = float32, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create a square identity matrix. @@ -1822,7 +2106,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tri(n: int, m: int, k: int, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tri(n: int, m: int, k: int, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An array with ones at and below the given diagonal and zeros elsewhere. @@ -1844,7 +2128,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tril(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tril(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array above the given diagonal. @@ -1864,7 +2148,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def triu(x: array, k: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def triu(x: array, k: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Zeros the array below the given diagonal. @@ -1887,7 +2171,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def allclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Approximate comparison of two arrays. @@ -1924,7 +2208,7 @@ void init_ops(nb::module_& m) { "equal_nan"_a = false, "stream"_a = nb::none(), nb::sig( - "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def isclose(a: array, b: array, /, rtol: float = 1e-05, atol: float = 1e-08, *, equal_nan: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a boolean array where two arrays are element-wise equal within a tolerance. @@ -1965,7 +2249,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def all(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def all(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An `and` reduction over the given axes. @@ -1994,7 +2278,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def any(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def any(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An `or` reduction over the given axes. @@ -2022,7 +2306,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def minimum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def minimum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise minimum. @@ -2049,7 +2333,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def maximum(a: Union[scalar, array], b: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def maximum(a: scalar | array, b: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise maximum. @@ -2072,7 +2356,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def floor(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def floor(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise floor. @@ -2091,7 +2375,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def ceil(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def ceil(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise ceil. @@ -2109,8 +2393,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isnan(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isnan(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are NaN. @@ -2128,8 +2411,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isinf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isinf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are +/- inifnity. @@ -2147,8 +2429,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isfinite(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isfinite(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are finite. @@ -2168,14 +2449,13 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isposinf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isposinf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are positive infinity. Args: a (array): Input array. - stream (Union[None, Stream, Device]): Optional stream or device. + stream (StreamOrDevice): Optional stream or device. Returns: array: The boolean array indicating which elements are positive infinity. @@ -2188,14 +2468,13 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def isneginf(a: array, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def isneginf(a: array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return a boolean array indicating which elements are negative infinity. Args: a (array): Input array. - stream (Union[None, Stream, Device]): Optional stream or device. + stream (StreamOrDevice): Optional stream or device. Returns: array: The boolean array indicating which elements are negative infinity. @@ -2209,7 +2488,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def moveaxis(a: array, /, source: int, destination: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def moveaxis(a: array, /, source: int, destination: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Move an axis to a new position. @@ -2230,7 +2509,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: Union[None, Stream, Device] = None) -> array"), + "def swapaxes(a: array, /, axis1 : int, axis2: int, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Swap two axes of an array. @@ -2258,7 +2537,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def transpose(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def transpose(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Transpose the dimensions of the array. @@ -2286,7 +2565,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def permute_dims(a: array, /, axes: Optional[Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def permute_dims(a: array, /, axes: Sequence[int] | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( See :func:`transpose`. )pbdoc"); @@ -2304,7 +2583,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sum(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sum(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sum reduce the array over the given axes. @@ -2319,6 +2598,41 @@ void init_ops(nb::module_& m) { Returns: array: The output array with the corresponding axes reduced. )pbdoc"); + m.def( + "count_nonzero", + [](const mx::array& a, + const IntOrVec& axis, + bool keepdims, + mx::StreamOrDevice s) { + if (std::holds_alternative(axis)) { + return mx::count_nonzero(a, keepdims, s); + } else if (auto pv = std::get_if(&axis); pv) { + return mx::count_nonzero(a, *pv, keepdims, s); + } else { + return mx::count_nonzero( + a, std::get>(axis), keepdims, s); + } + }, + nb::arg(), + "axis"_a = nb::none(), + nb::kw_only(), + "keepdims"_a = false, + "stream"_a = nb::none(), + nb::sig( + "def count_nonzero(a: array, /, *, axis: None | int | Sequence[int] = None, keepdims: bool = False, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Count the number of non-zero elements along the given axis. + + Args: + a (array): Input array. + axis (int or tuple(int), optional): Axis or axes to count over. + Defaults to ``None`` in which case the whole array is counted. + keepdims (bool, optional): Keep the reduced axes as size one. + Default: ``False``. + + Returns: + array: The counts as an ``int32`` array. + )pbdoc"); m.def( "prod", [](const mx::array& a, @@ -2333,7 +2647,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def prod(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def prod(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( An product reduction over the given axes. @@ -2362,7 +2676,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def min(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def min(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `min` reduction over the given axes. @@ -2391,7 +2705,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def max(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def max(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `max` reduction over the given axes. @@ -2427,7 +2741,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def logcumsumexp(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def logcumsumexp(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative logsumexp of the elements along the given axis. @@ -2457,7 +2771,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def logsumexp(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def logsumexp(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( A `log-sum-exp` reduction over the given axes. @@ -2492,7 +2806,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def mean(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def mean(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the mean(s) over the given axes. @@ -2521,7 +2835,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def median(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def median(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the median(s) over the given axes. @@ -2552,7 +2866,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def var(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def var(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the variance(s) over the given axes. @@ -2585,7 +2899,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def std(a: array, /, axis: Union[None, int, Sequence[int]] = None, keepdims: bool = False, ddof: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def std(a: array, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the standard deviation(s) over the given axes. @@ -2621,7 +2935,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def split(a: array, /, indices_or_sections: Union[int, Sequence[int]], axis: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def split(a: array, /, indices_or_sections: int | Sequence[int], axis: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Split an array along a given axis. @@ -2665,7 +2979,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argmin(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argmin(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Indices of the minimum values along the axis. @@ -2697,7 +3011,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argmax(a: array, /, axis: Union[None, int] = None, keepdims: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argmax(a: array, /, axis: None | int = None, keepdims: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Indices of the maximum values along the axis. @@ -2725,10 +3039,13 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def sort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def sort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a sorted copy of the array. + The sort is stable, meaning equal elements preserve their relative + order. ``NaN`` values are placed at the end. + Args: a (array): Input array. axis (int or None, optional): Optional axis to sort over. @@ -2752,10 +3069,13 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argsort(a: array, /, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argsort(a: array, /, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the indices that sort the array. + The sort is stable, meaning equal elements preserve their relative + order. ``NaN`` values are placed at the end. + Args: a (array): Input array. axis (int or None, optional): Optional axis to sort over. @@ -2783,7 +3103,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def partition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def partition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns a partitioned copy of the array such that the smaller ``kth`` elements are first. @@ -2821,7 +3141,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def argpartition(a: array, /, kth: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def argpartition(a: array, /, kth: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the indices that partition the array. @@ -2842,6 +3162,40 @@ void init_ops(nb::module_& m) { Returns: array: The ``uint32`` array containing indices that partition the input. )pbdoc"); + m.def( + "searchsorted", + &mx::searchsorted, + nb::arg(), + nb::arg(), + "side"_a = "left", + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def searchsorted(sorted_sequence: array, values: array, /, side: str = 'left', *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Find the indices that keep ``sorted_sequence`` sorted when inserting ``values``. + + Args: + sorted_sequence (array): A 1-D array sorted in ascending order. + values (array): The values to insert. May have any shape. + side (str, optional): Either ``'left'`` or ``'right'``. With + ``'left'`` the first suitable index is returned, so the result is + the number of elements strictly less than the value. With + ``'right'`` the last is returned, so the result is the number of + elements less than or equal to it. The two differ only where a + value is already present. Default: ``'left'``. + + Returns: + array: A ``uint32`` array with the same shape as ``values``, holding + indices in ``[0, sorted_sequence.size]``. + + Example: + >>> a = mx.array([1, 2, 2, 4]) + >>> mx.searchsorted(a, mx.array([0, 2, 3, 5])) + array([0, 1, 3, 4], dtype=uint32) + >>> mx.searchsorted(a, mx.array([0, 2, 3, 5]), side="right") + array([0, 3, 3, 4], dtype=uint32) + )pbdoc"); m.def( "topk", [](const mx::array& a, @@ -2860,7 +3214,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def topk(a: array, /, k: int, axis: Union[None, int] = -1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def topk(a: array, /, k: int, axis: None | int = -1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the ``k`` largest elements from the input along a given axis. @@ -2886,7 +3240,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def broadcast_to(a: Union[scalar, array], /, shape: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def broadcast_to(a: scalar | array, /, shape: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Broadcast an array to the given shape. @@ -2908,7 +3262,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def broadcast_arrays(*arrays: array, stream: Union[None, Stream, Device] = None) -> Tuple[array, ...]"), + "def broadcast_arrays(*arrays: array, stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Broadcast arrays against one another. @@ -2934,7 +3288,7 @@ void init_ops(nb::module_& m) { "precise"_a = false, "stream"_a = nb::none(), nb::sig( - "def softmax(a: array, /, axis: Union[None, int, Sequence[int]] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def softmax(a: array, /, axis: None | int | Sequence[int] = None, *, precise: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the softmax along the given axis. @@ -2949,6 +3303,10 @@ void init_ops(nb::module_& m) { axis (int or list(int), optional): Optional axis or axes to compute the softmax over. If unspecified this performs the softmax over the full array. + precise (bool, optional): Accumulate in ``float32`` for inputs of + lower precision. Otherwise the accumulation type matches the + input, which can lose precision over long reduction axes. + Default: ``False``. Returns: array: The output of the softmax. @@ -2969,7 +3327,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def concatenate(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def concatenate(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Concatenate the arrays along the given axis. @@ -2997,7 +3355,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def concat(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def concat(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( See :func:`concatenate`. )pbdoc"); @@ -3017,7 +3375,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def stack(arrays: list[array], axis: Optional[int] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def stack(arrays: list[array], axis: int | None = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Stacks the arrays along a new axis. @@ -3038,14 +3396,14 @@ void init_ops(nb::module_& m) { mx::StreamOrDevice s) { std::vector arrays = nb::cast>(arrays_); - return mx::meshgrid(arrays, sparse, indexing, s); + return nb::tuple(nb::cast(mx::meshgrid(arrays, sparse, indexing, s))); }, "arrays"_a, "sparse"_a = false, "indexing"_a = "xy", "stream"_a = nb::none(), nb::sig( - "def meshgrid(*arrays: array, sparse: Optional[bool] = False, indexing: Optional[str] = 'xy', stream: Union[None, Stream, Device] = None) -> array"), + "def meshgrid(*arrays: array, sparse: bool | None = False, indexing: str | None = 'xy', stream: StreamOrDevice = None) -> tuple[array, ...]"), R"pbdoc( Generate multidimensional coordinate grids from 1-D coordinate arrays @@ -3058,7 +3416,7 @@ void init_ops(nb::module_& m) { Defaults to ``'xy'``. Returns: - list(array): The output arrays. + tuple(array): The output arrays. )pbdoc"); m.def( "repeat", @@ -3078,7 +3436,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def repeat(array: array, repeats: int, axis: int | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Repeat an array along a specified axis. @@ -3115,7 +3473,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def clip(a: array, /, a_min: Union[scalar, array, None], a_max: Union[scalar, array, None], *, stream: Union[None, Stream, Device] = None) -> array"), + "def clip(a: array, /, a_min: scalar | array | None, a_max: scalar | array | None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Clip the values of the array between the given minimum and maximum. @@ -3143,18 +3501,23 @@ void init_ops(nb::module_& m) { const ScalarOrArray& constant_value, mx::StreamOrDevice s) { if (auto pv = std::get_if(&pad_width); pv) { - return mx::pad(a, *pv, to_array(constant_value), mode, s); + return mx::pad(a, *pv, to_array(constant_value, a.dtype()), mode, s); } else if (auto pv = std::get_if>(&pad_width); pv) { return mx::pad( - a, std::get<0>(*pv), to_array(constant_value), mode, s); + a, + std::get<0>(*pv), + to_array(constant_value, a.dtype()), + mode, + s); } else if (auto pv = std::get_if>(&pad_width); pv) { - return mx::pad(a, *pv, to_array(constant_value), mode, s); + return mx::pad(a, *pv, to_array(constant_value, a.dtype()), mode, s); } else { auto v = std::get>>(pad_width); if (v.size() == 1) { - return mx::pad(a, v[0], to_array(constant_value), mode, s); + return mx::pad( + a, v[0], to_array(constant_value, a.dtype()), mode, s); } else { - return mx::pad(a, v, to_array(constant_value), mode, s); + return mx::pad(a, v, to_array(constant_value, a.dtype()), mode, s); } } }, @@ -3165,7 +3528,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def pad(a: array, pad_width: Union[int, tuple[int], tuple[int, int], list[tuple[int, int]]], mode: Literal['constant', 'edge'] = 'constant', constant_values: Union[scalar, array] = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge', 'reflect', 'symmetric'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Pad an array with a constant value @@ -3180,7 +3543,9 @@ void init_ops(nb::module_& m) { mode: Padding mode. One of the following strings: "constant" (default): Pads with a constant value. "edge": Pads with the edge values of array. - constant_value (array or scalar, optional): Optional constant value + "reflect": Pads with the reflection of the array, without repeating the edge values. + "symmetric": Pads with the reflection of the array, repeating the edge values. + constant_values (array or scalar, optional): Optional constant value to pad the edges of the array with. Returns: @@ -3212,7 +3577,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def as_strided(a: array, /, shape: Optional[Sequence[int]] = None, strides: Optional[Sequence[int]] = None, offset: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def as_strided(a: array, /, shape: Sequence[int] | None = None, strides: Sequence[int] | None = None, offset: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Create a view into the array with the given shape and strides. @@ -3239,38 +3604,61 @@ void init_ops(nb::module_& m) { Returns: array: The output array which is the strided view of the input. )pbdoc"); + m.def( + "astype", + [](const mx::array& a, mx::Dtype dtype, mx::StreamOrDevice s) { + return mx::astype(a, dtype, s); + }, + nb::arg(), + "dtype"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def astype(a: array, dtype: Dtype, /, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Cast the array to a specified type. + + Args: + a (array): Input array. + dtype (Dtype): Type to which the array is cast. + + Returns: + array: The array with type ``dtype``. + )pbdoc"); m.def( "cumsum", [](const mx::array& a, std::optional axis, bool reverse, bool inclusive, + std::optional dtype, mx::StreamOrDevice s) { if (axis) { - return mx::cumsum(a, *axis, reverse, inclusive, s); - } else { - return mx::cumsum(mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); + return mx::cumsum(a, *axis, reverse, inclusive, dtype, s); } + return mx::cumsum(a, reverse, inclusive, dtype, s); }, nb::arg(), "axis"_a = nb::none(), nb::kw_only(), "reverse"_a = false, "inclusive"_a = true, + "dtype"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def cumsum(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cumsum(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative sum of the elements along the given axis. Args: - a (array): Input array + a (array): Input array. axis (int, optional): Optional axis to compute the cumulative sum over. If unspecified the cumulative sum of the flattened array is returned. reverse (bool): Perform the cumulative sum in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. + dtype (Dtype, optional): Cast the input to this type before summing. Returns: array: The output array. @@ -3281,32 +3669,34 @@ void init_ops(nb::module_& m) { std::optional axis, bool reverse, bool inclusive, + std::optional dtype, mx::StreamOrDevice s) { if (axis) { - return mx::cumprod(a, *axis, reverse, inclusive, s); - } else { - return mx::cumprod(mx::reshape(a, {-1}, s), 0, reverse, inclusive, s); + return mx::cumprod(a, *axis, reverse, inclusive, dtype, s); } + return mx::cumprod(a, reverse, inclusive, dtype, s); }, nb::arg(), "axis"_a = nb::none(), nb::kw_only(), "reverse"_a = false, "inclusive"_a = true, + "dtype"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def cumprod(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cumprod(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative product of the elements along the given axis. Args: - a (array): Input array + a (array): Input array. axis (int, optional): Optional axis to compute the cumulative product - over. If unspecified the cumulative product of the flattened array is - returned. + over. If unspecified the cumulative product of the flattened array + is returned. reverse (bool): Perform the cumulative product in reverse. inclusive (bool): The i-th element of the output includes the i-th element of the input. + dtype (Dtype, optional): Cast the input to this type before multiplying. Returns: array: The output array. @@ -3331,7 +3721,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def cummax(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cummax(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative maximum of the elements along the given axis. @@ -3367,7 +3757,7 @@ void init_ops(nb::module_& m) { "inclusive"_a = true, "stream"_a = nb::none(), nb::sig( - "def cummin(a: array, /, axis: Optional[int] = None, *, reverse: bool = False, inclusive: bool = True, stream: Union[None, Stream, Device] = None) -> array"), + "def cummin(a: array, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the cumulative minimum of the elements along the given axis. @@ -3383,6 +3773,28 @@ void init_ops(nb::module_& m) { Returns: array: The output array. )pbdoc"); + m.def( + "diff", + &mx::diff, + nb::arg(), + "n"_a = 1, + "axis"_a = -1, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def diff(a: array, /, n: int = 1, axis: int = -1, *, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + The n-th discrete difference along the given axis. + + Args: + a (array): Input array. + n (int, optional): The number of times to difference. Default: ``1``. + axis (int, optional): The axis along which to difference. + Default: ``-1``. + + Returns: + array: The n-th differences. + )pbdoc"); m.def( "conj", [](const ScalarOrArray& a, mx::StreamOrDevice s) { @@ -3391,8 +3803,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::kw_only(), "stream"_a = nb::none(), - nb::sig( - "def conj(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig("def conj(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the elementwise complex conjugate of the input. Alias for `mx.conjugate`. @@ -3412,7 +3823,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conjugate(a: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conjugate(a: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the elementwise complex conjugate of the input. Alias for `mx.conj`. @@ -3486,7 +3897,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: Union[None, Stream, Device] = None) -> array)"), + R"(def convolve(a: array, v: array, /, mode: str = "full", *, stream: StreamOrDevice = None) -> array)"), R"pbdoc( The discrete convolution of 1D arrays. @@ -3513,7 +3924,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 1D convolution over an input with several channels @@ -3571,7 +3982,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv2d(input: array, weight: array, /, stride: Union[int, tuple[int, int]] = 1, padding: Union[int, tuple[int, int]] = 0, dilation: Union[int, tuple[int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 2D convolution over an input with several channels @@ -3641,7 +4052,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv3d(input: array, weight: array, /, stride: Union[int, tuple[int, int, int]] = 1, padding: Union[int, tuple[int, int, int]] = 0, dilation: Union[int, tuple[int, int, int]] = 1, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 3D convolution over an input with several channels @@ -3677,7 +4088,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose1d(input: array, weight: array, /, stride: int = 1, padding: int = 0, dilation: int = 1, output_padding: int = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 1D transposed convolution over an input with several channels @@ -3752,7 +4163,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose2d(input: array, weight: array, /, stride: Union[int, Tuple[int, int]] = 1, padding: Union[int, Tuple[int, int]] = 0, dilation: Union[int, Tuple[int, int]] = 1, output_padding: Union[int, Tuple[int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose2d(input: array, weight: array, /, stride: int | tuple[int, int] = 1, padding: int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, output_padding: int | tuple[int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 2D transposed convolution over an input with several channels @@ -3838,7 +4249,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_transpose3d(input: array, weight: array, /, stride: Union[int, Tuple[int, int, int]] = 1, padding: Union[int, Tuple[int, int, int]] = 0, dilation: Union[int, Tuple[int, int, int]] = 1, output_padding: Union[int, Tuple[int, int, int]] = 0, groups: int = 1, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_transpose3d(input: array, weight: array, /, stride: int | tuple[int, int, int] = 1, padding: int | tuple[int, int, int] = 0, dilation: int | tuple[int, int, int] = 1, output_padding: int | tuple[int, int, int] = 0, groups: int = 1, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( 3D transposed convolution over an input with several channels @@ -3940,7 +4351,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def conv_general(input: array, weight: array, /, stride: Union[int, Sequence[int]] = 1, padding: Union[int, Sequence[int], tuple[Sequence[int], Sequence[int]]] = 0, kernel_dilation: Union[int, Sequence[int]] = 1, input_dilation: Union[int, Sequence[int]] = 1, groups: int = 1, flip: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def conv_general(input: array, weight: array, /, stride: int | Sequence[int] = 1, padding: int | Sequence[int] | tuple[Sequence[int], Sequence[int]] = 0, kernel_dilation: int | Sequence[int] = 1, input_dilation: int | Sequence[int] = 1, groups: int = 1, flip: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( General convolution over an input with several channels @@ -3973,8 +4384,7 @@ void init_ops(nb::module_& m) { &mlx_save_helper, "file"_a, "arr"_a, - nb::sig( - "def save(file: Union[file, str, pathlib.Path], arr: array) -> None"), + nb::sig("def save(file: file | str | pathlib.Path, arr: array) -> None"), R"pbdoc( Save the array to a binary file in ``.npy`` format. @@ -3990,8 +4400,7 @@ void init_ops(nb::module_& m) { "file"_a, "args"_a, "kwargs"_a, - nb::sig( - "def savez(file: Union[file, str, pathlib.Path], *args, **kwargs)"), + nb::sig("def savez(file: file | str | pathlib.Path, *args, **kwargs)"), R"pbdoc( Save several arrays to a binary file in uncompressed ``.npz`` format. @@ -4025,7 +4434,7 @@ void init_ops(nb::module_& m) { "args"_a, "kwargs"_a, nb::sig( - "def savez_compressed(file: Union[file, str, pathlib.Path], *args, **kwargs)"), + "def savez_compressed(file: file | str | pathlib.Path, *args, **kwargs)"), R"pbdoc( Save several arrays to a binary file in compressed ``.npz`` format. @@ -4044,7 +4453,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def load(file: Union[file, str, pathlib.Path], /, format: Optional[str] = None, return_metadata: bool = False, *, stream: Union[None, Stream, Device] = None) -> Union[array, dict[str, array], Tuple[dict[str, array], dict[str, Any]]]"), + "def load(file: file | str | pathlib.Path, /, format: str | None = None, return_metadata: bool = False, *, stream: StreamOrDevice = None) -> array | dict[str, array] | tuple[dict[str, array], dict[str, Any]]"), R"pbdoc( Load array(s) from a binary file. @@ -4079,7 +4488,7 @@ void init_ops(nb::module_& m) { "arrays"_a, "metadata"_a = nb::none(), nb::sig( - "def save_safetensors(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: Optional[dict[str, str]] = None)"), + "def save_safetensors(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, str] | None = None)"), R"pbdoc( Save array(s) to a binary file in ``.safetensors`` format. @@ -4101,7 +4510,7 @@ void init_ops(nb::module_& m) { "arrays"_a, "metadata"_a = nb::none(), nb::sig( - "def save_gguf(file: Union[file, str, pathlib.Path], arrays: dict[str, array], metadata: dict[str, Union[array, str, list[str]]])"), + "def save_gguf(file: file | str | pathlib.Path, arrays: dict[str, array], metadata: dict[str, array | str | list[str]])"), R"pbdoc( Save array(s) to a binary file in ``.gguf`` format. @@ -4132,7 +4541,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def where(condition: Union[scalar, array], x: Union[scalar, array], y: Union[scalar, array], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def where(condition: scalar | array, x: scalar | array, y: scalar | array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Select from ``x`` or ``y`` according to ``condition``. @@ -4164,7 +4573,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def nan_to_num(a: Union[scalar, array], nan: float = 0, posinf: Optional[float] = None, neginf: Optional[float] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def nan_to_num(a: scalar | array, nan: float = 0, posinf: float | None = None, neginf: float | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Replace NaN and Inf values with finite numbers. @@ -4191,7 +4600,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def round(a: array, /, decimals: int = 0, stream: Union[None, Stream, Device] = None) -> array"), + "def round(a: array, /, decimals: int = 0, stream: StreamOrDevice = None) -> array"), R"pbdoc( Round to the given number of decimals. @@ -4224,7 +4633,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def quantized_matmul(x: array, w: array, /, scales: array, biases: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, stream: Union[None, Stream, Device] = None) -> array"), + "def quantized_matmul(x: array, w: array, /, scales: array, biases: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the matrix multiplication with the quantized matrix ``w``. The quantization uses one floating point scale and bias per ``group_size`` of @@ -4258,10 +4667,11 @@ void init_ops(nb::module_& m) { "group_size"_a = nb::none(), "bits"_a = nb::none(), "mode"_a = "affine", + "global_scale"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def quantize(w: array, /, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, stream: Union[None, Stream, Device] = None) -> tuple[array, array, array]"), + "def quantize(w: array, /, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, global_scale: array | None = None, stream: StreamOrDevice = None) -> tuple[array, array, array]"), R"pbdoc( Quantize the array ``w``. @@ -4286,6 +4696,8 @@ void init_ops(nb::module_& m) { ``w`` in the quantized array. See supported values and defaults in the :ref:`table of quantization modes `. Default: ``None``. mode (str, optional): The quantization mode. Default: ``"affine"``. + global_scale (array, optional): The per-input float32 scale used for + ``"nvfp4"`` quantization if provided. Default: ``None``. Returns: tuple: A tuple with either two or three elements containing: @@ -4355,11 +4767,12 @@ void init_ops(nb::module_& m) { "group_size"_a = nb::none(), "bits"_a = nb::none(), "mode"_a = "affine", + "global_scale"_a = nb::none(), "dtype"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def dequantize(w: array, /, scales: array, biases: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def dequantize(w: array, /, scales: array, biases: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', global_scale: array | None = None, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Dequantize the matrix ``w`` using quantization parameters. @@ -4374,11 +4787,13 @@ void init_ops(nb::module_& m) { bits (int, optional): The number of bits occupied by each element of ``w`` in the quantized array. See supported values and defaults in the :ref:`table of quantization modes `. Default: ``None``. + mode (str, optional): The quantization mode. Default: ``"affine"``. + global_scale (array, optional): The per-input float32 scale used for + ``"nvfp4"`` quantization if provided. Default: ``None``. dtype (Dtype, optional): The data type of the dequantized output. If ``None`` the return type is inferred from the scales and biases when possible and otherwise defaults to ``bfloat16``. Default: ``None``. - mode (str, optional): The quantization mode. Default: ``"affine"``. Returns: array: The dequantized version of ``w`` @@ -4412,7 +4827,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_qmm(x: array, w: array, /, scales: array, biases: Optional[array] = None, lhs_indices: Optional[array] = None, rhs_indices: Optional[array] = None, transpose: bool = True, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def gather_qmm(x: array, w: array, /, scales: array, biases: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, transpose: bool = True, group_size: int | None = None, bits: int | None = None, mode: str = 'affine', *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform quantized matrix multiplication with matrix-level gather. @@ -4449,6 +4864,58 @@ void init_ops(nb::module_& m) { array: The result of the multiplication of ``x`` with ``w`` after gathering using ``lhs_indices`` and ``rhs_indices``. )pbdoc"); + m.def( + "gather_qqmm", + &mx::gather_qqmm, + nb::arg(), + nb::arg(), + "scales"_a = nb::none(), + "lhs_indices"_a = nb::none(), + "rhs_indices"_a = nb::none(), + "group_size"_a = nb::none(), + "bits"_a = nb::none(), + "mode"_a = "nvfp4", + "global_scale_x"_a = nb::none(), + "global_scale_w"_a = nb::none(), + nb::kw_only(), + "sorted_indices"_a = false, + "stream"_a = nb::none(), + nb::sig( + "def gather_qqmm(x: array, w: array, /, scales: array | None = None, lhs_indices: array | None = None, rhs_indices: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Fused :func:`qqmm` with matrix-level gather. + + Similar to :func:`gather_mm`, the indices ``lhs_indices`` and + ``rhs_indices`` contain flat indices along the batch dimensions (i.e. + all but the last two dimensions) of ``x`` and ``w`` respectively. + + Args: + x (array): Input array. + w (array): Weight matrix. If quantized, it is packed in unsigned integers. + scales (array, optional): The scales to use per ``group_size`` elements of + ``w`` if ``w`` is quantized. Default: ``None``. + lhs_indices (array, optional): Integer indices for ``x``. Default: ``None``. + rhs_indices (array, optional): Integer indices for ``w``. Default: ``None``. + group_size (int, optional): Number of elements in ``x`` and ``w`` that + share a scale. See supported values and defaults in the + :ref:`table of quantization modes `. Default: ``None``. + bits (int, optional): Number of bits used to represent each element of + ``x`` and ``w``. See supported values and defaults in the + :ref:`table of quantization modes `. Default: ``None``. + mode (str, optional): The quantization mode. Default: ``"nvfp4"``. + Supported modes are ``nvfp4`` and ``mxfp8``. See the + :ref:`table of quantization modes ` for details. + global_scale_x (array, optional): The per-input float32 scale used for x + with ``"nvfp4"`` quantization. Default: ``None``. + global_scale_w (array, optional): The per-input float32 scale used for w + with ``"nvfp4"`` quantization. Default: ``None``. + sorted_indices (bool, optional): May allow a faster implementation + if the passed indices are sorted. Default: ``False``. + + Returns: + array: The result of the multiplication of quantized ``x`` with quantized ``w``. + needed). + )pbdoc"); m.def( "segmented_mm", &mx::segmented_mm, @@ -4458,7 +4925,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def segmented_mm(a: array, b: array, /, segments: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def segmented_mm(a: array, b: array, /, segments: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform a matrix multiplication but segment the inner dimension and save the result for each segment separately. @@ -4494,7 +4961,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tensordot(a: array, b: array, /, axes: Union[int, list[Sequence[int]]] = 2, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tensordot(a: array, b: array, /, axes: int | list[Sequence[int]] = 2, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the tensor dot product along the specified axes. @@ -4518,7 +4985,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def inner(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def inner(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Ordinary inner product of vectors for 1-D arrays, in higher dimensions a sum product over the last axes. @@ -4529,6 +4996,27 @@ void init_ops(nb::module_& m) { Returns: array: The inner product. )pbdoc"); + m.def( + "vecdot", + &mx::vecdot, + nb::arg(), + nb::arg(), + "axis"_a = -1, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def vecdot(a: array, b: array, /, *, axis: int = -1, stream: StreamOrDevice = None) -> array"), + R"pbdoc( + Compute the vector dot product of two arrays along an axis. + + Args: + a (array): Input array + b (array): Input array + axis (int, optional): Axis over which to compute the dot product. Default: ``-1``. + + Returns: + array: The vector dot product. + )pbdoc"); m.def( "outer", &mx::outer, @@ -4537,7 +5025,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def outer(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def outer(a: array, b: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Compute the outer product of two 1-D arrays, if the array's passed are not 1-D a flatten op will be run beforehand. @@ -4564,7 +5052,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def tile(a: array, reps: Union[int, Sequence[int]], /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def tile(a: array, reps: int | Sequence[int], /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Construct an array by repeating ``a`` the number of times given by ``reps``. @@ -4586,7 +5074,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def addmm(c: array, a: array, b: array, /, alpha: float = 1.0, beta: float = 1.0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with addition and optional scaling. @@ -4616,7 +5104,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: Optional[array] = None, mask_lhs: Optional[array] = None, mask_rhs: Optional[array] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def block_masked_mm(a: array, b: array, /, block_size: int = 64, mask_out: array | None = None, mask_lhs: array | None = None, mask_rhs: array | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with block masking. @@ -4655,7 +5143,7 @@ void init_ops(nb::module_& m) { "sorted_indices"_a = false, "stream"_a = nb::none(), nb::sig( - "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: Union[None, Stream, Device] = None) -> array"), + "def gather_mm(a: array, b: array, /, lhs_indices: array, rhs_indices: array, *, sorted_indices: bool = False, stream: StreamOrDevice = None) -> array"), R"pbdoc( Matrix multiplication with matrix-level gather. @@ -4697,7 +5185,7 @@ void init_ops(nb::module_& m) { "axis2"_a = 1, "stream"_a = nb::none(), nb::sig( - "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: Union[None, Stream, Device] = None) -> array"), + "def diagonal(a: array, offset: int = 0, axis1: int = 0, axis2: int = 1, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return specified diagonals. @@ -4729,7 +5217,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def diag(a: array, /, k: int = 0, *, stream: Union[None, Stream, Device] = None) -> array"), + "def diag(a: array, /, k: int = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Extract a diagonal or construct a diagonal matrix. If ``a`` is 1-D then a diagonal matrix is constructed with ``a`` on the @@ -4765,7 +5253,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"), + "def trace(a: array, /, offset: int = 0, axis1: int = 0, axis2: int = 1, dtype: Dtype | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Return the sum along a specified diagonal in the given array. @@ -4795,13 +5283,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_1d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_1d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least one dimension. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least one dimension. @@ -4818,13 +5306,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_2d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_2d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least two dimensions. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least two dimensions. @@ -4841,13 +5329,13 @@ void init_ops(nb::module_& m) { "arys"_a, "stream"_a = nb::none(), nb::sig( - "def atleast_3d(*arys: array, stream: Union[None, Stream, Device] = None) -> Union[array, list[array]]"), + "def atleast_3d(*arys: array, stream: StreamOrDevice = None) -> array | list[array]"), R"pbdoc( Convert all arrays to have at least three dimensions. Args: *arys: Input arrays. - stream (Union[None, Stream, Device], optional): The stream to execute the operation on. + stream (StreamOrDevice, optional): The stream to execute the operation on. Returns: array or list(array): An array or list of arrays with at least three dimensions. @@ -4877,7 +5365,7 @@ void init_ops(nb::module_& m) { ""_a, ""_a, nb::sig( - "def issubdtype(arg1: Union[Dtype, DtypeCategory], arg2: Union[Dtype, DtypeCategory]) -> bool"), + "def issubdtype(arg1: Dtype | DtypeCategory, arg2: Dtype | DtypeCategory) -> bool"), R"pbdoc( Check if a :obj:`Dtype` or :obj:`DtypeCategory` is a subtype of another. @@ -4926,6 +5414,132 @@ void init_ops(nb::module_& m) { >>> mx.issubdtype(mx.signedinteger, mx.floating) False )pbdoc"); + m.def( + "result_type", + [](const nb::args& arrays_and_dtypes) { + auto to_dtype = [](const nb::handle& v) -> mx::Dtype { + if (nb::isinstance(v)) { + return nb::cast(v).dtype(); + } else if (nb::isinstance(v)) { + return nb::cast(v); + } else { + throw std::invalid_argument( + "[result_type] Inputs must be arrays or dtypes."); + } + }; + if (arrays_and_dtypes.size() == 0) { + throw std::invalid_argument( + "[result_type] At least one array or dtype is required."); + } + mx::Dtype t = to_dtype(arrays_and_dtypes[0]); + for (size_t i = 1; i < arrays_and_dtypes.size(); ++i) { + t = mx::promote_types(t, to_dtype(arrays_and_dtypes[i])); + } + return t; + }, + nb::sig("def result_type(*arrays_and_dtypes: array | Dtype) -> Dtype"), + R"pbdoc( + The type that results from applying type promotion to the inputs. + + Args: + *arrays_and_dtypes (array or Dtype): A variable number of arrays + or dtypes. + + Returns: + Dtype: The result type. + )pbdoc"); + m.def( + "can_cast", + [](const nb::object& from_, const mx::Dtype& to) { + mx::Dtype from_dtype = mx::bool_; + if (nb::isinstance(from_)) { + from_dtype = nb::cast(from_).dtype(); + } else if (nb::isinstance(from_)) { + from_dtype = nb::cast(from_); + } else { + throw std::invalid_argument( + "[can_cast] `from_` must be an array or a dtype."); + } + return mx::promote_types(from_dtype, to) == to; + }, + "from_"_a, + "to"_a, + nb::sig("def can_cast(from_: array | Dtype, to: Dtype) -> bool"), + R"pbdoc( + Determine if one data type can be cast to another according to type + promotion rules. + + ``from_`` can be cast to ``to`` if promoting the two together gives + back ``to``. + + Args: + from_ (array or Dtype): The source array or dtype. + to (Dtype): The destination dtype. + + Returns: + bool: Whether the cast can be performed. + )pbdoc"); + m.def( + "isdtype", + [](const mx::Dtype& dtype, const nb::object& kind) { + auto check_one = [&dtype](const nb::handle& k) -> bool { + if (nb::isinstance(k)) { + return dtype == nb::cast(k); + } else if (nb::isinstance(k)) { + auto s = nb::cast(k); + if (s == "bool") { + return dtype == mx::bool_; + } else if (s == "signed integer") { + return mx::issubdtype(dtype, mx::signedinteger); + } else if (s == "unsigned integer") { + return mx::issubdtype(dtype, mx::unsignedinteger); + } else if (s == "integral") { + return mx::issubdtype(dtype, mx::integer); + } else if (s == "real floating") { + return mx::issubdtype(dtype, mx::floating); + } else if (s == "complex floating") { + return mx::issubdtype(dtype, mx::complexfloating); + } else if (s == "numeric") { + return mx::issubdtype(dtype, mx::number); + } else { + std::ostringstream msg; + msg << "[isdtype] Unknown data type kind: '" << s << "'."; + throw std::invalid_argument(msg.str()); + } + } else { + throw std::invalid_argument( + "[isdtype] `kind` must be a dtype, a string, or a tuple of " + "dtypes and strings."); + } + }; + if (nb::isinstance(kind)) { + for (auto k : nb::cast(kind)) { + if (check_one(k)) { + return true; + } + } + return false; + } + return check_one(kind); + }, + "dtype"_a, + "kind"_a, + nb::sig( + "def isdtype(dtype: Dtype, kind: Dtype | str | tuple[Dtype | str, ...]) -> bool"), + R"pbdoc( + Test whether a dtype belongs to one or more data type kinds. + + Args: + dtype (Dtype): The dtype to test. + kind (Dtype, str, or tuple): A dtype, a kind string, or a tuple + of dtypes and kind strings. Supported kind strings are + ``"bool"``, ``"signed integer"``, ``"unsigned integer"``, + ``"integral"``, ``"real floating"``, ``"complex floating"``, + and ``"numeric"``. + + Returns: + bool: ``True`` if ``dtype`` matches any of the given kinds. + )pbdoc"); m.def( "bitwise_and", [](const ScalarOrArray& a_, @@ -4939,7 +5553,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_and(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_and(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise and. @@ -4966,7 +5580,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_or(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_or(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise or. @@ -4993,7 +5607,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_xor(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_xor(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise xor. @@ -5021,7 +5635,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def left_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def left_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise left shift. @@ -5049,7 +5663,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def right_shift(a: Union[scalar, array], b: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def right_shift(a: scalar | array, b: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise right shift. @@ -5074,7 +5688,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def bitwise_invert(a: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"), + "def bitwise_invert(a: scalar | array, stream: StreamOrDevice = None) -> array"), R"pbdoc( Element-wise bitwise inverse. @@ -5096,7 +5710,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def view(a: Union[scalar, array], dtype: Dtype, stream: Union[None, Stream, Device] = None) -> array"), + "def view(a: scalar | array, dtype: Dtype, stream: StreamOrDevice = None) -> array"), R"pbdoc( View the array as a different type. @@ -5122,7 +5736,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def hadamard_transform(a: array, scale: Optional[float] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def hadamard_transform(a: array, scale: float | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the Walsh-Hadamard transform along the final axis. @@ -5186,7 +5800,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def einsum(subscripts: str, *operands, stream: Union[None, Stream, Device] = None) -> array"), + "def einsum(subscripts: str, *operands, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform the Einstein summation convention on the operands. @@ -5221,7 +5835,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def roll(a: array, shift: Union[int, Tuple[int]], axis: Union[None, int, Tuple[int]] = None, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def roll(a: array, shift: int | tuple[int], axis: None | int | tuple[int] = None, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Roll array elements along a given axis. @@ -5249,7 +5863,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def real(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def real(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the real part of a complex array. @@ -5268,7 +5882,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def imag(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"), + "def imag(a: array, /, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Returns the imaginary part of a complex array. @@ -5295,7 +5909,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def slice(a: array, start_indices: array, axes: Sequence[int], slice_size: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Extract a sub-array from the input array. @@ -5334,7 +5948,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: Union[None, Stream, Device] = None) -> array"), + "def slice_update(a: array, update: array, start_indices: array, axes: Sequence[int], *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Update a sub-array of the input array. @@ -5363,7 +5977,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: Union[None, Stream, Device] = None) -> array"), + "def contiguous(a: array, /, allow_col_major: bool = False, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Force an array to be row contiguous. Copy if necessary. @@ -5392,7 +6006,7 @@ void init_ops(nb::module_& m) { return nb::tuple(nb::cast(result)); }, - nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> Tuple[int]"), + nb::sig("def broadcast_shapes(*shapes: Sequence[int]) -> tuple[int]"), R"pbdoc( Broadcast shapes. @@ -5443,7 +6057,7 @@ void init_ops(nb::module_& m) { nb::arg(), nb::arg(), nb::sig( - "def depends(inputs: Union[array, Sequence[array]], dependencies: Union[array, Sequence[array]])"), + "def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array])"), R"pbdoc( Insert dependencies between arrays in the graph. The outputs are identical to ``inputs`` but with dependencies on ``dependencies``. @@ -5465,10 +6079,12 @@ void init_ops(nb::module_& m) { "group_size"_a = nb::none(), "bits"_a = nb::none(), "mode"_a = "nvfp4", + "global_scale_x"_a = nb::none(), + "global_scale_w"_a = nb::none(), nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def qqmm(x: array, w: array, scales: Optional[array] = None, group_size: Optional[int] = None, bits: Optional[int] = None, mode: str = 'nvfp4', *, stream: Union[None, Stream, Device] = None) -> array"), + "def qqmm(x: array, w: array, scales: array | None = None, group_size: int | None = None, bits: int | None = None, mode: str = 'nvfp4', global_scale_x: array | None = None, global_scale_w: array | None = None, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Perform a matrix multiplication using a possibly quantized weight matrix ``w`` and a non-quantized input ``x``. The input ``x`` is quantized on the @@ -5486,6 +6102,7 @@ void init_ops(nb::module_& m) { If ``x`` and `w`` are not quantized, their data types must be ``float32``, ``float16``, or ``bfloat16``. If ``w`` is quantized, it must be packed in unsigned integers. + ``global_scale_x`` and ``global_scale_w`` are only used for ``nvfp4`` quantization. Args: x (array): Input array. @@ -5501,7 +6118,10 @@ void init_ops(nb::module_& m) { mode (str, optional): The quantization mode. Default: ``"nvfp4"``. Supported modes are ``nvfp4`` and ``mxfp8``. See the :ref:`table of quantization modes ` for details. - + global_scale (array, optional): The per-input float32 scale used for x + with ``"nvfp4"`` quantization. Default: ``None``. + global_scale_w (array, optional): The per-input float32 scale used for w + with ``"nvfp4"`` quantization. Default: ``None``. Returns: array: The result of the multiplication of quantized ``x`` with quantized ``w``. needed). @@ -5514,7 +6134,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: Union[None, Stream, Device] = None) -> array"), + "def from_fp8(x: array, dtype: Dtype = bfloat16, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert the array from fp8 (e4m3) to another floating-point type. @@ -5532,7 +6152,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def to_fp8(x: array, *, stream: Union[None, Stream, Device] = None) -> array"), + "def to_fp8(x: array, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Convert the array to fp8 (e4m3) from another floating-point type. @@ -5542,4 +6162,20 @@ void init_ops(nb::module_& m) { Returns: array: The array converted to fp8 with type ``uint8``. )pbdoc"); + // Array API standard aliases (https://data-apis.org/array-api/latest/). + m.attr("acos") = m.attr("arccos"); + m.attr("acosh") = m.attr("arccosh"); + m.attr("asin") = m.attr("arcsin"); + m.attr("asinh") = m.attr("arcsinh"); + m.attr("atan") = m.attr("arctan"); + m.attr("atanh") = m.attr("arctanh"); + m.attr("atan2") = m.attr("arctan2"); + m.attr("bitwise_left_shift") = m.attr("left_shift"); + m.attr("bitwise_right_shift") = m.attr("right_shift"); + m.attr("cumulative_prod") = m.attr("cumprod"); + m.attr("cumulative_sum") = m.attr("cumsum"); + m.attr("empty") = m.attr("zeros"); + m.attr("empty_like") = m.attr("zeros_like"); + m.attr("matrix_transpose") = m.attr("transpose"); + m.attr("pow") = m.attr("power"); } diff --git a/Source/Cmlx/mlx/python/src/print.cpp b/Source/Cmlx/mlx/python/src/print.cpp index 03c3fdaa7..6c08d88d3 100644 --- a/Source/Cmlx/mlx/python/src/print.cpp +++ b/Source/Cmlx/mlx/python/src/print.cpp @@ -44,7 +44,7 @@ void init_print(nb::module_& m) { Example: >>> print(x) # Uses default precision - >>> mx.set_printoptions(precision=3): + >>> mx.set_printoptions(precision=3) >>> print(x) # Uses precision of 3 >>> print(x) # Uses precision of 3 (again) diff --git a/Source/Cmlx/mlx/python/src/random.cpp b/Source/Cmlx/mlx/python/src/random.cpp index c832c5a9e..10b82b892 100644 --- a/Source/Cmlx/mlx/python/src/random.cpp +++ b/Source/Cmlx/mlx/python/src/random.cpp @@ -9,6 +9,7 @@ #include "mlx/ops.h" #include "mlx/random.h" +#include "python/src/random.h" #include "python/src/small_vector.h" #include "python/src/utils.h" @@ -18,50 +19,109 @@ using namespace nb::literals; class PyKeySequence { public: - explicit PyKeySequence(uint64_t seed) { - state_.append(mx::random::key(seed)); + ~PyKeySequence() { + if (state_.has_value()) { + nb::gil_scoped_acquire gil; + state_.reset(); + } + } + + void reset() { + state_.reset(); } void seed(uint64_t seed) { - state_[0] = mx::random::key(seed); + state()[0] = mx::random::key(seed); } mx::array next() { - auto out = mx::random::split(nb::cast(state_[0])); - state_[0] = out.first; + auto out = mx::random::split(nb::cast(state()[0])); + state()[0] = out.first; return out.second; } - nb::list state() { - return state_; - } - - void release() { - nb::gil_scoped_acquire gil; - state_.release().dec_ref(); + nb::list& state() { + if (!state_) { + static auto time_seed = []() { + auto now = std::chrono::system_clock::now(); + return std::chrono::duration_cast( + now.time_since_epoch()) + .count(); + }(); + state_ = nb::list(); + state_->append(mx::random::key(time_seed)); + } + return *state_; } private: - nb::list state_; + std::optional state_; }; PyKeySequence& default_key() { - auto get_current_time_seed = []() { - auto now = std::chrono::system_clock::now(); - return std::chrono::duration_cast( - now.time_since_epoch()) - .count(); - }; - static PyKeySequence ks(get_current_time_seed()); + // Each thread has its own random key to avoid race condition. + static thread_local PyKeySequence ks; return ks; } +void reset_random_state() { + default_key().reset(); +} + +// A process-global sentinel for `mx.random.state`. Since it is the same object +// on every thread, capturing it (e.g. with `mx.compile`) is thread-independent; +// the pytree traversal in trees.cpp resolves it to the calling thread's key. +class RandomState {}; + +nb::object random_state_sentinel() { + static nb::object sentinel = []() { + auto sentinel = nb::cast(RandomState{}); + sentinel.inc_ref(); + return sentinel; + }(); + + return sentinel; +} + +mx::array random_state_key() { + return nb::cast(default_key().state()[0]); +} + +void set_random_state_key(const mx::array& key) { + default_key().state()[0] = nb::cast(key); +} + void init_random(nb::module_& parent_module) { auto m = parent_module.def_submodule( "random", "mlx.core.random: functionality related to random number generation"); - m.attr("state") = default_key().state(); + nb::class_(m, "_RandomState") + .def("__len__", [](const RandomState&) { return 1; }) + .def( + "__getitem__", + [](const RandomState&, int i) -> nb::object { + if (i != 0 && i != -1) { + throw nb::index_error("random state index out of range"); + } + return default_key().state()[0]; + }, + "index"_a) + .def("__iter__", [](const RandomState&) { + return nb::iter(default_key().state()); + }); + + m.def("__getattr__", [&](nb::handle key) -> nb::object { + // Create random.state lazily to avoid initializing device during import. + if (nb::isinstance(key) && nb::cast(key) == "state") { + default_key().state(); + return random_state_sentinel(); + } + return nb::steal(PyErr_Format( + PyExc_AttributeError, + "Module 'random' has no attribute %R", + key.ptr())); + }); m.def( "seed", [](uint64_t seed) { default_key().seed(seed); }, @@ -93,7 +153,7 @@ void init_random(nb::module_& parent_module) { "num"_a = 2, "stream"_a = nb::none(), nb::sig( - "def split(key: array, num: int = 2, stream: Union[None, Stream, Device] = None) -> array"), + "def split(key: array, num: int = 2, stream: StreamOrDevice = None) -> array"), R"pbdoc( Split a PRNG key into sub keys. @@ -128,7 +188,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def uniform(low: Union[scalar, array] = 0, high: Union[scalar, array] = 1, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def uniform(low: scalar | array = 0, high: scalar | array = 1, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate uniformly distributed random numbers. @@ -171,7 +231,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Union[scalar, array, None] = None, scale: Union[scalar, array, None] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def normal(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: scalar | array | None = None, scale: scalar | array | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate normally distributed random numbers. @@ -211,7 +271,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate jointly-normal random samples given a mean and covariance. @@ -258,7 +318,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def randint(low: Union[scalar, array], high: Union[scalar, array], shape: Sequence[int] = [], dtype: Optional[Dtype] = int32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def randint(low: scalar | array, high: scalar | array, shape: Sequence[int] = [], dtype: Dtype | None = int32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate random integers from the given interval. @@ -266,6 +326,12 @@ void init_random(nb::module_& parent_module) { half-open interval ``[low, high)``. The lower and upper bound can be scalars or arrays and must be broadcastable to ``shape``. + .. note:: + The samples are drawn from a ``float32`` uniform and clamped to + ``[low, high - 1]``, so not every integer in the range is reachable + once the bounds or the width of the interval go beyond the + ``2**24`` integer resolution of ``float32``. + Args: low (scalar or array): Lower bound of the interval. high (scalar or array): Upper bound of the interval. @@ -295,7 +361,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def bernoulli(p: Union[scalar, array] = 0.5, shape: Optional[Sequence[int]] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def bernoulli(p: scalar | array = 0.5, shape: Sequence[int] | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate Bernoulli random values. @@ -339,7 +405,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def truncated_normal(lower: Union[scalar, array], upper: Union[scalar, array], shape: Optional[Sequence[int]] = None, dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def truncated_normal(lower: scalar | array, upper: scalar | array, shape: Sequence[int] | None = None, dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate values from a truncated normal distribution. @@ -373,7 +439,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def gumbel(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def gumbel(shape: Sequence[int] = [], dtype: Dtype | None = float32, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample from the standard Gumbel distribution. @@ -419,7 +485,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def categorical(logits: array, axis: int = -1, shape: Optional[Sequence[int]] = None, num_samples: Optional[int] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def categorical(logits: array, axis: int = -1, shape: Sequence[int] | None = None, num_samples: int | None = None, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample from a categorical distribution. @@ -462,7 +528,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def laplace(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: float = 0.0, scale: float = 1.0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def laplace(shape: Sequence[int] = [], dtype: Dtype | None = float32, loc: float = 0.0, scale: float = 1.0, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Sample numbers from a Laplace distribution. @@ -495,7 +561,7 @@ void init_random(nb::module_& parent_module) { "key"_a = nb::none(), "stream"_a = nb::none(), nb::sig( - "def permutation(x: Union[int, array], axis: int = 0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), + "def permutation(x: int | array, axis: int = 0, key: array | None = None, stream: StreamOrDevice = None) -> array"), R"pbdoc( Generate a random permutation or permute the entries of an array. @@ -510,7 +576,10 @@ void init_random(nb::module_& parent_module) { array: The generated random permutation or randomly permuted input array. )pbdoc"); - // Register static Python object cleanup before the interpreter exits + + // Ensure the main thread cleanup will happen before the interpreter goes + // away. As a result if the other threads join the main thread we should have + // a clean tear-down. auto atexit = nb::module_::import_("atexit"); - atexit.attr("register")(nb::cpp_function([]() { default_key().release(); })); + atexit.attr("register")(nb::cpp_function([]() { default_key().reset(); })); } diff --git a/Source/Cmlx/mlx/python/src/random.h b/Source/Cmlx/mlx/python/src/random.h new file mode 100644 index 000000000..02d81d4c7 --- /dev/null +++ b/Source/Cmlx/mlx/python/src/random.h @@ -0,0 +1,20 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include + +#include "mlx/array.h" + +namespace mx = mlx::core; +namespace nb = nanobind; + +// Clear the `mx.random.state` python object in current thread. +void reset_random_state(); + +// The process-global `mx.random.state` sentinel. +nb::object random_state_sentinel(); + +// Read/write the calling thread's current PRNG key. +mx::array random_state_key(); +void set_random_state_key(const mx::array& key); diff --git a/Source/Cmlx/mlx/python/src/small_vector.h b/Source/Cmlx/mlx/python/src/small_vector.h index 2be8dcec1..0b2e6bb2e 100644 --- a/Source/Cmlx/mlx/python/src/small_vector.h +++ b/Source/Cmlx/mlx/python/src/small_vector.h @@ -2,6 +2,11 @@ #pragma once +#include +#include +#include +#include + #include "mlx/small_vector.h" #include @@ -10,15 +15,23 @@ NAMESPACE_BEGIN(NB_NAMESPACE) NAMESPACE_BEGIN(detail) template -struct type_caster> { - using List = mlx::core::SmallVector; +struct type_caster<::mlx::core::SmallVector> { + using List = ::mlx::core::SmallVector; using Caster = make_caster; + // For narrow integer element types we fetch each element through a wider + // integer caster so we can emit a clean OverflowError on overflow instead of + // nanobind's generic "incompatible function arguments" TypeError. + static constexpr bool kNarrowInt = std::is_integral_v && + !std::is_same_v && (sizeof(Type) < sizeof(int64_t)); + NB_TYPE_CASTER( List, const_name("tuple[") + make_caster::Name + const_name(", ...]")) - bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept { + // Not noexcept: on overflow of a narrow integer element we raise + // OverflowError so nanobind surfaces a clean error to the user. + bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) { size_t size; PyObject* temp; @@ -29,19 +42,39 @@ struct type_caster> { value.clear(); value.reserve(size); - Caster caster; bool success = o != nullptr; flags = flags_for_local_caster(flags); for (size_t i = 0; i < size; ++i) { - if (!caster.from_python(o[i], flags, cleanup) || - !caster.template can_cast()) { - success = false; - break; + if constexpr (kNarrowInt) { + make_caster wide; + if (!wide.from_python(o[i], flags, cleanup) || + !wide.template can_cast()) { + success = false; + break; + } + int64_t v = wide.operator cast_t(); + if (v > std::numeric_limits::max() || + v < std::numeric_limits::min()) { + std::ostringstream msg; + msg << "Integer value " << v << " is outside the supported range [" + << static_cast(std::numeric_limits::min()) << ", " + << static_cast(std::numeric_limits::max()) << "]."; + Py_XDECREF(temp); + PyErr_SetString(PyExc_OverflowError, msg.str().c_str()); + raise_python_error(); + } + value.push_back(static_cast(v)); + } else { + Caster caster; + if (!caster.from_python(o[i], flags, cleanup) || + !caster.template can_cast()) { + success = false; + break; + } + value.push_back(caster.operator cast_t()); } - - value.push_back(caster.operator cast_t()); } Py_XDECREF(temp); diff --git a/Source/Cmlx/mlx/python/src/stream.cpp b/Source/Cmlx/mlx/python/src/stream.cpp index e10f4751c..467518e99 100644 --- a/Source/Cmlx/mlx/python/src/stream.cpp +++ b/Source/Cmlx/mlx/python/src/stream.cpp @@ -9,6 +9,7 @@ #include "mlx/stream.h" #include "mlx/utils.h" +#include "python/src/random.h" namespace mx = mlx::core; namespace nb = nanobind; @@ -61,12 +62,34 @@ void init_stream(nb::module_& m) { s == nb::cast(other); }); + nb::class_( + m, + "ThreadLocalStream", + R"pbdoc( + A stream that will be unique per thread and can be used to run operations on a given device. + )pbdoc") + .def_ro("device", &mx::ThreadLocalStream::device) + .def( + "__repr__", + [](const mx::ThreadLocalStream& s) { + std::ostringstream os; + os << "ThreadLocalStream(" << s.device << ", " << s.index << ")"; + return os.str(); + }) + .def( + "__eq__", + [](const mx::ThreadLocalStream& s, const nb::object& other) { + return nb::isinstance(other) && + s == nb::cast(other); + }); + nb::implicitly_convertible(); m.def( "default_stream", &mx::default_stream, "device"_a, + nb::sig("def default_stream(device: Device | DeviceType) -> Stream"), R"pbdoc(Get the device's default stream.)pbdoc"); m.def( "set_default_stream", @@ -85,7 +108,41 @@ void init_stream(nb::module_& m) { "new_stream", &mx::new_stream, "device"_a, - R"pbdoc(Make a new stream on the given device.)pbdoc"); + nb::sig("def new_stream(device: Device | DeviceType) -> Stream"), + R"pbdoc( + Make a new stream on the given device. + + The stream can only be used on the thread where it was created on, using + it in any other thread would result in errors. + )pbdoc"); + m.def( + "new_thread_unsafe_stream", + &mx::new_thread_unsafe_stream, + "device"_a, + nb::sig( + "def new_thread_unsafe_stream(device: Device | DeviceType) -> Stream"), + R"pbdoc( + Make a new stream that can be used in any thread. + + Unlike :func:`new_stream` which can only work on the thread of creation, + streams created by this API can be passed to and evaluated anywhere, but + note that currently all nodes in a graph must be evaluated in sequence + and it is user's responsibilty to ensure there is no race condition. + )pbdoc"); + m.def( + "new_thread_local_stream", + &mx::new_thread_local_stream, + "device"_a, + nb::sig( + "def new_thread_local_stream(device: Device | DeviceType) -> ThreadLocalStream"), + R"pbdoc(Make a new stream that will be unique per thread.)pbdoc"); + m.def( + "clear_streams", + []() { + reset_random_state(); + mx::clear_streams(); + }, + R"pbdoc(Destroy all streams created in current thread.)pbdoc"); nb::class_(m, "StreamContext", R"pbdoc( A context manager for setting the current device and stream. @@ -132,15 +189,20 @@ void init_stream(nb::module_& m) { )pbdoc"); m.def( "synchronize", - [](const std::optional& s) { - s ? mx::synchronize(s.value()) : mx::synchronize(); + [](mx::StreamOrDevice s) { + if (std::holds_alternative(s)) { + mx::synchronize(); + } else { + mx::synchronize(mx::to_stream(s)); + } }, "stream"_a = nb::none(), R"pbdoc( Synchronize with the given stream. Args: - stream (Stream, optional): The stream to synchronize with. If ``None`` + stream (Stream, optional): Stream to synchronize. If device is + provided the default stream for that device is used. If ``None`` then the default stream of the default device is used. Default: ``None``. )pbdoc"); diff --git a/Source/Cmlx/mlx/python/src/transforms.cpp b/Source/Cmlx/mlx/python/src/transforms.cpp index c615b12ec..e6a778eca 100644 --- a/Source/Cmlx/mlx/python/src/transforms.cpp +++ b/Source/Cmlx/mlx/python/src/transforms.cpp @@ -412,6 +412,7 @@ struct PyCompiledFun { nb::object captured_inputs; nb::object captured_outputs; bool shapeless; + mx::detail::CompileCacheWeakPtr cache; // Data to attach to the compiled function that contains the python output // structure and the number of arrays in said structure. @@ -421,6 +422,11 @@ struct PyCompiledFun { AttachedData(nb::object output_structure_, int num_outputs_) : output_structure(output_structure_), num_outputs(num_outputs_) {} + + ~AttachedData() { + nb::gil_scoped_acquire gil; + output_structure.reset(); + } }; PyCompiledFun( @@ -439,14 +445,17 @@ struct PyCompiledFun { PyCompiledFun& operator=(PyCompiledFun&& other) = delete; PyCompiledFun(PyCompiledFun&& other) : fun(std::move(other.fun)), - fun_id(reinterpret_cast(fun.ptr())) { + fun_id(reinterpret_cast(fun.ptr())), + captured_inputs(std::move(other.captured_inputs)), + captured_outputs(std::move(other.captured_outputs)), + shapeless(other.shapeless), + cache(other.cache) { other.fun_id = 0; - captured_inputs = std::move(other.captured_inputs); - captured_outputs = std::move(other.captured_outputs); - shapeless = other.shapeless; }; nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) { + cache = mx::detail::compile_cache(); + // Flat array inputs std::vector inputs; @@ -580,7 +589,7 @@ struct PyCompiledFun { ~PyCompiledFun() { nb::gil_scoped_acquire gil; - mx::detail::compile_erase(fun_id); + mx::detail::compile_erase(cache, fun_id); fun.reset(); captured_inputs.reset(); captured_outputs.reset(); @@ -1333,7 +1342,7 @@ void init_transforms(nb::module_& m) { "argnums"_a = nb::none(), "argnames"_a = std::vector{}, nb::sig( - "def value_and_grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Tuple[R, Any]]"), + "def value_and_grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, tuple[R, Any]]"), R"pbdoc( Returns a function which computes the value and gradient of ``fun``. @@ -1402,7 +1411,7 @@ void init_transforms(nb::module_& m) { "argnums"_a = nb::none(), "argnames"_a = std::vector{}, nb::sig( - "def grad(fun: Callable[P, R], argnums: Optional[Union[int, Sequence[int]]] = None, argnames: Union[str, Sequence[str]] = []) -> Callable[P, Any]"), + "def grad(fun: Callable[P, R], argnums: int | Sequence[int] | None = None, argnames: str | Sequence[str] = []) -> Callable[P, Any]"), R"pbdoc( Returns a function which computes the gradient of ``fun``. @@ -1472,7 +1481,7 @@ void init_transforms(nb::module_& m) { "outputs"_a = nb::none(), "shapeless"_a = false, nb::sig( - "def compile(fun: Callable[P, R], inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) -> Callable[P, R]"), + "def compile(fun: Callable[P, R], inputs: object | None = None, outputs: object | None = None, shapeless: bool = False) -> Callable[P, R]"), R"pbdoc( Returns a compiled function which produces the same output as ``fun``. @@ -1499,7 +1508,7 @@ void init_transforms(nb::module_& m) { Returns: Callable: A compiled function which has the same input arguments - as ``fun`` and returns the the same output(s). + as ``fun`` and returns the same output(s). )pbdoc"); m.def( "disable_compile", @@ -1535,8 +1544,10 @@ void init_transforms(nb::module_& m) { computation. )pbdoc"); - // Register static Python object cleanup before the interpreter exits + // Clean up main thread compile cache before python interpreter shuts down. auto atexit = nb::module_::import_("atexit"); atexit.attr("register")( - nb::cpp_function([]() { mx::detail::compile_clear_cache(); })); + nb::cpp_function([cache = mx::detail::compile_cache()]() { + mx::detail::compile_clear_cache(cache); + })); } diff --git a/Source/Cmlx/mlx/python/src/trees.cpp b/Source/Cmlx/mlx/python/src/trees.cpp index 4b9ca9e12..cf05d4e25 100644 --- a/Source/Cmlx/mlx/python/src/trees.cpp +++ b/Source/Cmlx/mlx/python/src/trees.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include "python/src/trees.h" +#include "python/src/random.h" template void validate_subtrees(const std::vector& subtrees) { @@ -152,8 +153,12 @@ void tree_visit( void tree_visit(nb::handle tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (nb::isinstance(subtree) || + if (subtree.is(random_state)) { + visitor(nb::cast(random_state_key())); + } else if ( + nb::isinstance(subtree) || nb::isinstance(subtree)) { for (auto item : subtree) { recurse(item); @@ -174,8 +179,14 @@ void tree_visit_update( nb::object tree, std::function visitor) { std::function recurse; + auto random_state = random_state_sentinel(); recurse = [&](nb::handle subtree) { - if (nb::isinstance(subtree)) { + if (subtree.is(random_state)) { + // Read/write the calling thread's key; keep the sentinel in the tree. + set_random_state_key( + nb::cast(visitor(nb::cast(random_state_key())))); + return nb::cast(subtree); + } else if (nb::isinstance(subtree)) { auto l = nb::cast(subtree); for (int i = 0; i < l.size(); ++i) { l[i] = recurse(l[i]); @@ -262,14 +273,12 @@ nb::object tree_unflatten( } nb::object structure_sentinel() { - static nb::object sentinel; - - if (sentinel.ptr() == nullptr) { - sentinel = nb::capsule(&sentinel); - // probably not needed but this should make certain that we won't ever - // delete the sentinel + static nb::object sentinel = []() { + PyObject* raw_obj = PyObject_New(PyObject, &PyBaseObject_Type); + nb::object sentinel = nb::steal(raw_obj); sentinel.inc_ref(); - } + return sentinel; + }(); return sentinel; } diff --git a/Source/Cmlx/mlx/python/src/utils.cpp b/Source/Cmlx/mlx/python/src/utils.cpp index 32d8a0164..84214cd47 100644 --- a/Source/Cmlx/mlx/python/src/utils.cpp +++ b/Source/Cmlx/mlx/python/src/utils.cpp @@ -31,16 +31,20 @@ mx::array to_array( return mx::array(val, (out_t == mx::bool_) ? mx::int32 : out_t); } else if (auto pv = std::get_if(&v); pv) { auto out_t = dtype.value_or(mx::float32); - return mx::array( - nb::cast(*pv), - mx::issubdtype(out_t, mx::floating) ? out_t : mx::float32); + if (!mx::issubdtype(out_t, mx::floating)) { + out_t = mx::float32; + } + if (out_t == mx::float64) { + // Cast straight to double: going through float would round the value + // to float32 precision while the result still advertises float64. + return mx::array(nb::cast(*pv), out_t); + } + return mx::array(nb::cast(*pv), out_t); } else if (auto pv = std::get_if>(&v); pv) { return mx::array(static_cast(*pv), mx::complex64); } else if (auto pv = std::get_if(&v); pv) { return *pv; - } else if (auto pv = std::get_if< - nb::ndarray>(&v); - pv) { + } else if (auto pv = std::get_if>(&v); pv) { return nd_array_to_mlx(*pv, dtype); } else { return to_array_with_accessor(std::get(v).obj); diff --git a/Source/Cmlx/mlx/python/src/utils.h b/Source/Cmlx/mlx/python/src/utils.h index 583d72a08..d2fcbd037 100644 --- a/Source/Cmlx/mlx/python/src/utils.h +++ b/Source/Cmlx/mlx/python/src/utils.h @@ -24,7 +24,7 @@ using ScalarOrArray = std::variant< // Must be above ndarray mx::array, // Must be above complex - nb::ndarray, + nb::ndarray, std::complex, ArrayLike>; diff --git a/Source/Cmlx/mlx/python/tests/__main__.py b/Source/Cmlx/mlx/python/tests/__main__.py deleted file mode 100644 index 5230bd535..000000000 --- a/Source/Cmlx/mlx/python/tests/__main__.py +++ /dev/null @@ -1,5 +0,0 @@ -from . import mlx_tests - -__unittest = True - -mlx_tests.MLXTestRunner(module=None) diff --git a/Source/Cmlx/mlx/python/tests/cuda_skip.py b/Source/Cmlx/mlx/python/tests/cuda_skip.py deleted file mode 100644 index 20793d5c9..000000000 --- a/Source/Cmlx/mlx/python/tests/cuda_skip.py +++ /dev/null @@ -1,62 +0,0 @@ -cuda_skip = { - "TestLayers.test_quantized_embedding", - # Block masked matmul NYI - "TestBlas.test_block_masked_matmul", - # Gather matmul NYI - "TestBlas.test_gather_matmul", - "TestBlas.test_gather_matmul_grad", - "TestBlas.test_gather_mm_sorted_vjp", - # Segmented matmul NYI - "TestBlas.test_segmented_mm", - # Hadamard NYI - "TestOps.test_hadamard", - "TestOps.test_hadamard_grad_vmap", - # FFTs NYI - "TestFFT.test_fft", - "TestFFT.test_fft_big_powers_of_two", - "TestFFT.test_fft_contiguity", - "TestFFT.test_fft_exhaustive", - "TestFFT.test_fft_grads", - "TestFFT.test_fft_into_ifft", - "TestFFT.test_fft_large_numbers", - "TestFFT.test_fft_shared_mem", - "TestFFT.test_fftn", - # Lapack ops NYI - "TestLinalg.test_cholesky", - "TestLinalg.test_cholesky_inv", - "TestLinalg.test_eig", - "TestLinalg.test_eigh", - "TestLinalg.test_inverse", - "TestVmap.test_vmap_inverse", - "TestLinalg.test_lu", - "TestLinalg.test_lu_factor", - "TestLinalg.test_pseudo_inverse", - "TestLinalg.test_qr_factorization", - "TestInit.test_orthogonal", - "TestLinalg.test_svd_decomposition", - "TestVmap.test_vmap_svd", - "TestLinalg.test_tri_inverse", - # Quantization NYI - "TestQuantized.test_gather_matmul_grad", - "TestQuantized.test_gather_qmm", - "TestQuantized.test_gather_qmm_sorted", - "TestQuantized.test_gather_qmm_grad", - "TestQuantized.test_non_multiples", - "TestQuantized.test_qmm", - "TestQuantized.test_qmm_jvp", - "TestQuantized.test_qmm_shapes", - "TestQuantized.test_qmm_vjp", - "TestQuantized.test_qmv", - "TestQuantized.test_fp_qvm", - "TestQuantized.test_qvm", - "TestQuantized.test_qvm_splitk", - "TestQuantized.test_qmv_small_non_multiples", - "TestQuantized.test_small_matrix", - "TestQuantized.test_throw", - "TestQuantized.test_vjp_scales_biases", - "TestExportImport.test_export_quantized_model", - # Masked scatter - "TestOps.test_masked_scatter", - "TestVmap.test_vmap_masked_scatter", - "TestArray.test_setitem_with_boolean_mask", -} diff --git a/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py b/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py index 8db456662..cbb966304 100644 --- a/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py +++ b/Source/Cmlx/mlx/python/tests/mlx_distributed_tests.py @@ -1,10 +1,12 @@ # Copyright © 2025 Apple Inc. +import math + import mlx.core as mx import mlx.nn as nn import mlx_tests from mlx.nn.layers.distributed import shard_inplace, shard_linear -from mlx.nn.utils import average_gradients +from mlx.nn.utils import average_gradients, clip_grad_norm_sharded class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase): @@ -47,17 +49,6 @@ def new_all_sum(x, **kwargs): self.assertTrue(all(mx.all(g == 1) for g in new_grads)) self.assertEqual(n_calls, 10) - n_calls = 0 - xtype = mx.float16 - new_grads = average_gradients( - grads, all_reduce_size=2 * 50, communication_type=mx.float16 - ) - mx.eval(new_grads) - self.assertEqual(len(new_grads), 10) - self.assertTrue(all(g.dtype == mx.float32 for g in new_grads)) - self.assertTrue(all(mx.all(g == 1) for g in new_grads)) - self.assertEqual(n_calls, 2) - finally: mx.distributed.all_sum = original_all_sum @@ -146,7 +137,7 @@ def test_shard_linear(self): self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol)) self.assertTrue(mx.allclose(y[part], y1, atol=self.atol, rtol=self.rtol)) - # And their quant versions (QuintizedMatmul is not supported on CUDA) + # And their quant versions (QuantizedMatmul is not supported on CUDA) if not mx.cuda.is_available(): qlin = lin.to_quantized() slin1 = shard_linear(qlin, "all-to-sharded") @@ -157,6 +148,27 @@ def test_shard_linear(self): self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol)) self.assertTrue(mx.allclose(y[part], y1)) + # Test non-affine quantization modes (mxfp8) + qlin_mxfp8 = lin.to_quantized(group_size=32, bits=8, mode="mxfp8") + self.assertEqual(qlin_mxfp8.mode, "mxfp8") + + slin1_mxfp8 = shard_linear(qlin_mxfp8, "all-to-sharded") + slin2_mxfp8 = shard_linear(qlin_mxfp8, "sharded-to-all") + + # Verify mode is propagated + self.assertEqual(slin1_mxfp8.mode, "mxfp8") + self.assertEqual(slin2_mxfp8.mode, "mxfp8") + + # Verify biases parameter is not set for mxfp8 + self.assertIsNone(slin1_mxfp8.get("biases")) + self.assertIsNone(slin2_mxfp8.get("biases")) + + y = qlin_mxfp8(x) + y1 = slin1_mxfp8(x) + y2 = slin2_mxfp8(x[part]) + self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol)) + self.assertTrue(mx.allclose(y[part], y1)) + # Check the backward works as expected def dummy_loss(model, x, y): return (model(x) * y).sum() @@ -312,3 +324,51 @@ def test_all_gather(self): y = mx.distributed.all_gather(x) self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + + def test_clip_grad_norm_sharded(self): + world = mx.distributed.init() + N = world.size() + + value = 3.0 + grads_slice = {"a": mx.ones((4, 3)) * value, "b": mx.ones((5,)) * value} + local_numel = 4 * 3 + 5 + expected_norm = math.sqrt(N * local_numel) * value + + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=1e9, group=world + ) + mx.eval(clipped, grad_norm) + self.assertTrue( + mx.allclose( + grad_norm, mx.array(expected_norm), atol=self.atol, rtol=self.rtol + ) + ) + for k in grads_slice: + self.assertTrue( + mx.allclose(clipped[k], grads_slice[k], atol=self.atol, rtol=self.rtol) + ) + + max_norm = 1.0 + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=max_norm, group=world + ) + mx.eval(clipped, grad_norm) + scale = max_norm / (expected_norm + 1e-6) + for k in grads_slice: + self.assertTrue( + mx.allclose( + clipped[k], grads_slice[k] * scale, atol=self.atol, rtol=self.rtol + ) + ) + + def test_jaccl_all_gather_factory_validation(self): + # A custom side-channel factory is only valid with the jaccl backend. + with self.assertRaises(ValueError): + mx.distributed.init( + backend="ring", + all_gather_factory=lambda rank, size: lambda src, n_bytes: b"", + ) + + # The factory must be callable. + with self.assertRaises(TypeError): + mx.distributed.init(backend="jaccl", all_gather_factory="not_callable") diff --git a/Source/Cmlx/mlx/python/tests/mlx_tests.py b/Source/Cmlx/mlx/python/tests/mlx_tests.py index c344e7c86..2b60f4615 100644 --- a/Source/Cmlx/mlx/python/tests/mlx_tests.py +++ b/Source/Cmlx/mlx/python/tests/mlx_tests.py @@ -1,14 +1,8 @@ # Copyright © 2023 Apple Inc. import os - -# Use regular fp32 precision for tests -os.environ["MLX_ENABLE_TF32"] = "0" - -# Do not abort on cache thrashing -os.environ["MLX_ENABLE_CACHE_THRASHING_CHECK"] = "0" - import platform +import sys import unittest from typing import Any, Callable, List, Tuple, Union @@ -18,38 +12,14 @@ class MLXTestRunner(unittest.TestProgram): def __init__(self, *args, **kwargs): + # Do not exit in runTests + kwargs["exit"] = False super().__init__(*args, **kwargs) - def createTests(self, *args, **kwargs): - super().createTests(*args, **kwargs) - - # Asume CUDA backend in this case - device = os.getenv("DEVICE", None) - if device is not None: - device = getattr(mx, device) - else: - device = mx.default_device() - - if not (device == mx.gpu and not mx.metal.is_available()): - return - - from cuda_skip import cuda_skip - - filtered_suite = unittest.TestSuite() - - def filter_and_add(t): - if isinstance(t, unittest.TestSuite): - for sub_t in t: - filter_and_add(sub_t) - else: - t_id = ".".join(t.id().split(".")[-2:]) - if t_id in cuda_skip: - print(f"Skipping {t_id}") - else: - filtered_suite.addTest(t) - - filter_and_add(self.test) - self.test = filtered_suite + def runTests(self): + super().runTests() + mx.clear_streams() + sys.exit(0 if self.result.wasSuccessful() else 1) class MLXTestCase(unittest.TestCase): diff --git a/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py b/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py index e6eed8a99..db9bbb2e6 100644 --- a/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py +++ b/Source/Cmlx/mlx/python/tests/nccl_test_distributed.py @@ -1,8 +1,11 @@ # Copyright © 2024 Apple Inc. import mlx.core as mx +import mlx.nn as nn import mlx_distributed_tests import mlx_tests +from mlx.nn.utils import average_gradients +from mlx.utils import tree_flatten, tree_map class TestNCCLDistributed(mlx_distributed_tests.MLXDistributedCommonTestCase): @@ -47,6 +50,134 @@ def test_sum_scatter(self): maxrelerror = maxrelerror.max() self.assertLessEqual(maxrelerror, rtol) + def test_groups(self): + world = mx.distributed.init() + self.assertEqual(world.size(), 8) + self.assertTrue(0 <= world.rank() < 8) + + world2 = mx.distributed.init() + self.assertEqual(world.size(), world2.size()) + self.assertEqual(world.rank(), world2.rank()) + + sub = world.split(world.rank() % 2) + self.assertEqual(sub.size(), 4) + self.assertEqual(sub.rank(), world.rank() // 2) + + sub = world.split(world.rank() // 2) + self.assertEqual(sub.size(), 2) + + def test_all_reduce_split(self): + world = mx.distributed.init() + dtypes = [ + (mx.float32, 1e-6), + (mx.float16, 5e-3), + (mx.bfloat16, 1e-1), + ] + sizes = [ + (7,), + (10,), + (1024,), + (1024, 1024), + ] + key = mx.random.key(0) + group = world.split(world.rank() % 2) + + for dt, rtol in dtypes: + for sh in sizes: + x = ( + mx.random.uniform(shape=(group.size(),) + sh, key=key) * 10 + ).astype(dt) + + # All sum + y = mx.distributed.all_sum(x[group.rank()], group=group) + z = x.sum(0) + maxrelerror = (y - z).abs() + if rtol > 0: + maxrelerror /= z.abs() + maxrelerror = maxrelerror.max() + self.assertLessEqual(maxrelerror, rtol) + + # All max + y = mx.distributed.all_max(x[group.rank()], group=group) + z = x.max(0) + self.assertTrue(mx.all(y == z)) + + # All min + y = mx.distributed.all_min(x[group.rank()], group=group) + z = x.min(0) + self.assertTrue(mx.all(y == z)) + + def test_all_gather_split(self): + world = mx.distributed.init() + dtypes = [mx.float32, mx.float16, mx.bfloat16] + sub = world.split(world.rank() % 2) + for dt in dtypes: + x = mx.ones((2, 2, 4), dtype=dt) + y = mx.distributed.all_gather(x, group=sub) + self.assertEqual(y.shape, (sub.size() * 2, 2, 4)) + self.assertTrue(mx.all(y == 1)) + + def test_fully_shard_grads(self): + dtypes = [ + (mx.float32, 1e-6, 1e-6), + (mx.bfloat16, 1e-3, 1e-3), + ] + + world = mx.distributed.init() + N = world.size() + rank = world.rank() + dims = 8 * N + part = slice(rank * dims // N, (rank + 1) * dims // N) + + class MLP(nn.Module): + def __init__(self, dims): + super().__init__() + self.l1 = nn.Linear(dims, dims) + self.l2 = nn.Linear(dims, dims) + + def __call__(self, x): + return self.l2(nn.relu(self.l1(x))) + + def loss_fn(model, x, y): + logits = model(x).astype(mx.float32) + return ((logits - y) ** 2).mean() + + for dtype, atol, rtol in dtypes: + + mx.random.seed(0xF0F0F0F0) + + kx, ky = mx.random.split(mx.random.key(rank)) + x = mx.random.normal((4, dims), dtype=dtype, key=kx) + y = mx.random.normal((4, dims), key=ky) + + # DDP reference: replicated params, gradients averaged across ranks + model = MLP(dims) + params = model.trainable_parameters() + + # Cast parameters to dtype for forward + model.update(tree_map(lambda p: p.astype(dtype), params)) + loss, grads = mx.value_and_grad(loss_fn)(model, x, y) + grads = average_gradients(grads, group=world) + loss = mx.distributed.all_sum(loss, group=world) / N + mx.eval(loss, grads) + + # Shard the model + model_sharded = MLP(dims) + model_sharded.update(params) + model_sharded = nn.fully_shard(model_sharded, compute_dtype=dtype) + loss_sharded, grads_sharded = mx.value_and_grad(loss_fn)( + model_sharded, x, y + ) + + loss_sharded = mx.distributed.all_sum(loss_sharded, group=world) / N + mx.eval(loss_sharded, grads_sharded) + grads_ref = dict(tree_flatten(grads)) + self.assertTrue(mx.allclose(loss, loss_sharded, atol=1e-4, rtol=1e-4)) + for k, gs in tree_flatten(grads_sharded["module"]): + self.assertTrue( + mx.allclose(gs, grads_ref[k][part], atol=atol, rtol=rtol) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/ring_test_distributed.py b/Source/Cmlx/mlx/python/tests/ring_test_distributed.py index dab40e48d..31db2ed5f 100644 --- a/Source/Cmlx/mlx/python/tests/ring_test_distributed.py +++ b/Source/Cmlx/mlx/python/tests/ring_test_distributed.py @@ -14,8 +14,7 @@ def setUpClass(cls): def test_groups(self): world = mx.distributed.init() - self.assertEqual(world.size(), 8) - self.assertTrue(0 <= world.rank() < 8) + self.assertTrue(0 <= world.rank() < world.size()) world2 = mx.distributed.init() self.assertEqual(world.size(), world2.size()) diff --git a/Source/Cmlx/mlx/python/tests/run.py b/Source/Cmlx/mlx/python/tests/run.py new file mode 100644 index 000000000..df96d4dd3 --- /dev/null +++ b/Source/Cmlx/mlx/python/tests/run.py @@ -0,0 +1,18 @@ +import os +import sys + +# Use regular fp32 precision for tests +os.environ["MLX_ENABLE_TF32"] = "0" + +# Do not abort on cache thrashing +os.environ["MLX_ENABLE_CACHE_THRASHING_CHECK"] = "0" + +__unittest = True + +import mlx_tests + +if __name__ == "__main__": + # Run all tests by default. + dirname = os.path.dirname(os.path.realpath(__file__)) + argv = [sys.argv[0], "discover", dirname, *sys.argv[1:]] + mlx_tests.MLXTestRunner(argv=argv, module=None) diff --git a/Source/Cmlx/mlx/python/tests/test_array.py b/Source/Cmlx/mlx/python/tests/test_array.py index 4efed9dac..80459fabe 100644 --- a/Source/Cmlx/mlx/python/tests/test_array.py +++ b/Source/Cmlx/mlx/python/tests/test_array.py @@ -1,21 +1,16 @@ # Copyright © 2023-2024 Apple Inc. -import gc import operator import os import pickle import platform import sys import unittest +import warnings import weakref from copy import copy, deepcopy from itertools import permutations -if platform.system() == "Windows": - import psutil -else: - import resource - import mlx.core as mx import mlx_tests import numpy as np @@ -24,10 +19,27 @@ import tensorflow as tf has_tf = True -except ImportError as e: +except ImportError: has_tf = False +try: + import torch + + torch_version = [int(v) for v in torch.__version__.split("+")[0].split(".")] + is_torch_212 = torch_version[0] > 2 or ( + torch_version[0] == 2 and torch_version[1] >= 12 + ) + has_torch_mps = ( + is_torch_212 + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ) +except ImportError: + torch = None + has_torch_mps = False + + class TestVersion(mlx_tests.MLXTestCase): def test_version(self): v = mx.__version__ @@ -37,6 +49,37 @@ def test_version(self): self.assertEqual(v, mx.__version__[: len(v)]) +class TestArrayNamespsceInfo(mlx_tests.MLXTestCase): + def test(self): + namespace = mx.__array_namespace_info__() + + self.assertEqual(namespace.default_device(), mx.default_device()) + self.assertEqual( + namespace.default_dtypes(), + { + "real floating": mx.float32, + "complex floating": mx.complex64, + "integral": mx.int32, + "indexing": mx.int32, + }, + ) + self.assertEqual( + namespace.dtypes(device=mx.Device(mx.cpu), kind="real floating"), + {"float32": mx.float32, "float64": mx.float64}, + ) + if mx.is_available(mx.gpu): + self.assertEqual( + namespace.dtypes(device=mx.Device(mx.gpu), kind="real floating"), + {"float32": mx.float32}, + ) + self.assertEqual( + namespace.dtypes(kind=("bool", "complex floating")), + {"bool": mx.bool_, "complex64": mx.complex64}, + ) + with self.assertRaises(ValueError): + namespace.dtypes(kind="invalid") + + class TestDtypes(mlx_tests.MLXTestCase): def test_dtypes(self): self.assertEqual(mx.bool_.size, 1) @@ -104,13 +147,33 @@ def test_finfo(self): self.assertEqual(mx.finfo(mx.float32).min, np.finfo(np.float32).min) self.assertEqual(mx.finfo(mx.float32).max, np.finfo(np.float32).max) self.assertEqual(mx.finfo(mx.float32).eps, np.finfo(np.float32).eps) + self.assertEqual(mx.finfo(mx.float32).bits, np.finfo(np.float32).bits) + self.assertEqual( + mx.finfo(mx.float32).smallest_normal, + float(np.finfo(np.float32).smallest_normal), + ) self.assertEqual(mx.finfo(mx.float32).dtype, mx.float32) self.assertEqual(mx.finfo(mx.float16).min, np.finfo(np.float16).min) self.assertEqual(mx.finfo(mx.float16).max, np.finfo(np.float16).max) self.assertEqual(mx.finfo(mx.float16).eps, np.finfo(np.float16).eps) + self.assertEqual(mx.finfo(mx.float16).bits, np.finfo(np.float16).bits) + self.assertEqual( + mx.finfo(mx.float16).smallest_normal, + float(np.finfo(np.float16).smallest_normal), + ) self.assertEqual(mx.finfo(mx.float16).dtype, mx.float16) + # bfloat16 has no numpy equivalent; check against known IEEE values. + self.assertEqual(mx.finfo(mx.bfloat16).bits, 16) + self.assertAlmostEqual( + mx.finfo(mx.bfloat16).smallest_normal, 2.0**-126, places=40 + ) + + # finfo of a complex type reports its real component (array API). + self.assertEqual(mx.finfo(mx.complex64).dtype, mx.float32) + self.assertEqual(mx.finfo(mx.complex64).bits, 32) + def test_iinfo(self): with self.assertRaises(ValueError): mx.iinfo(mx.float32) @@ -123,6 +186,55 @@ def test_iinfo(self): self.assertEqual(mx.iinfo(mx.uint32).max, np.iinfo(np.uint32).max) self.assertEqual(mx.iinfo(mx.int8).dtype, mx.int8) + def test_result_type(self): + self.assertEqual(mx.result_type(mx.int8, mx.int16), mx.int16) + self.assertEqual(mx.result_type(mx.float32, mx.float64), mx.float64) + # Accepts arrays as well as dtypes, and more than two inputs. + self.assertEqual( + mx.result_type(mx.array([1], dtype=mx.int8), mx.int16, mx.int32), + mx.int32, + ) + self.assertEqual( + mx.result_type(mx.array(1.0), mx.array(1, dtype=mx.int32)), + mx.float32, + ) + with self.assertRaises(ValueError): + mx.result_type() + + def test_can_cast(self): + self.assertTrue(mx.can_cast(mx.int8, mx.int16)) + self.assertFalse(mx.can_cast(mx.int16, mx.int8)) + self.assertTrue(mx.can_cast(mx.float32, mx.float64)) + self.assertFalse(mx.can_cast(mx.float64, mx.float32)) + self.assertTrue(mx.can_cast(mx.uint8, mx.int16)) + self.assertFalse(mx.can_cast(mx.uint16, mx.int16)) + # Accepts an array for the source. + self.assertTrue(mx.can_cast(mx.array([1, 2, 3], dtype=mx.int8), mx.int32)) + + def test_isdtype(self): + self.assertTrue(mx.isdtype(mx.int32, mx.int32)) + self.assertFalse(mx.isdtype(mx.int32, mx.int16)) + self.assertTrue(mx.isdtype(mx.int32, "signed integer")) + self.assertTrue(mx.isdtype(mx.uint8, "unsigned integer")) + self.assertFalse(mx.isdtype(mx.uint8, "signed integer")) + self.assertTrue(mx.isdtype(mx.int16, "integral")) + self.assertTrue(mx.isdtype(mx.float32, "real floating")) + self.assertTrue(mx.isdtype(mx.complex64, "complex floating")) + self.assertTrue(mx.isdtype(mx.bool_, "bool")) + self.assertFalse(mx.isdtype(mx.bool_, "numeric")) + self.assertTrue(mx.isdtype(mx.float32, "numeric")) + # Tuple of kinds (any match). + self.assertTrue(mx.isdtype(mx.float32, ("integral", "real floating"))) + self.assertTrue(mx.isdtype(mx.int8, (mx.int8, mx.int16))) + self.assertFalse(mx.isdtype(mx.int32, ("bool", "real floating"))) + with self.assertRaises(ValueError): + mx.isdtype(mx.int32, "not a kind") + + # Reachable through the array API namespace. + xp = mx.array(1.0).__array_namespace__() + for name in ("result_type", "can_cast", "isdtype", "vecdot"): + self.assertTrue(hasattr(xp, name), msg=name) + class TestEquality(mlx_tests.MLXTestCase): def test_array_eq_array(self): @@ -198,7 +310,7 @@ def test_list_not_equals_array(self): def test_dlx_device_type(self): a = mx.array([1, 2, 3]) device_type, device_id = a.__dlpack_device__() - self.assertIn(device_type, [1, 8, 13]) + self.assertIn(device_type, [1, 8]) self.assertEqual(device_id, 0) if device_type == 8: @@ -442,6 +554,33 @@ def test_double_keeps_precision(self): out = mx.array([x], dtype=mx.float64).item() self.assertEqual(out, x) + def test_construction_from_lists_wide_ints(self): + # A python int that does not fit in int32 widens to int64, the same + # rule the scalar path already uses. It used to raise std::bad_cast. + for value in (2**31, 2**40, -(2**31) - 1, -(2**40)): + for make in ( + lambda v: [v], + lambda v: (v,), + lambda v: [[v]], + lambda v: [v, 1], + ): + x = mx.array(make(value)) + self.assertEqual(x.dtype, mx.int64, msg=f"{value} {make(value)}") + self.assertEqual(x.flatten()[0].item(), value) + self.assertEqual(mx.array(value).dtype, mx.int64) + + # Values that still fit keep int32, including both boundaries. + for value in (0, 1, 2**31 - 1, -(2**31)): + x = mx.array([value]) + self.assertEqual(x.dtype, mx.int32, msg=str(value)) + self.assertEqual(x[0].item(), value) + + # An explicit dtype still wins. + self.assertEqual(mx.array([2**40], mx.int64).dtype, mx.int64) + self.assertEqual(mx.array([1, 2], mx.int64).dtype, mx.int64) + # A float in the list still makes it float, not int64. + self.assertEqual(mx.array([2**40, 1.5]).dtype, mx.float32) + def test_construction_from_lists_of_mlx_arrays(self): dtypes = [ mx.bool_, @@ -601,6 +740,26 @@ def test_array_repr(self): x = mx.array([1 - 1j], dtype=mx.complex64) expected = "array([1-1j], dtype=complex64)" + def test_array_repr_precision(self): + x = mx.array([1.123456789], dtype=mx.float32) + expected = "array([1.12346], dtype=float32)" + self.assertEqual(str(x), expected) + + with mx.printoptions(precision=4): + expected = "array([1.1235], dtype=float32)" + self.assertEqual(str(x), expected) + mx.set_printoptions(precision=2) + expected = "array([1.12], dtype=float32)" + self.assertEqual(str(x), expected) + + x = mx.sin(x) + expected = "array([0.90], dtype=float32)" + self.assertEqual(str(x), expected) + + with mx.printoptions(precision=4): + expected = "array([0.9016], dtype=float32)" + self.assertEqual(str(x), expected) + def test_array_to_list(self): types = [mx.bool_, mx.uint32, mx.int32, mx.int64, mx.float32] for t in types: @@ -751,10 +910,13 @@ def test_array_from_noncontiguous_np(self): def test_array_np_shape_dim_check(self): a_npy = np.empty(2**31, dtype=np.bool_) - with self.assertRaises(ValueError) as e: + with self.assertRaises(OverflowError) as e: mx.array(a_npy) self.assertEqual( - str(e.exception), "Shape dimension falls outside supported `int` range." + str(e.exception), + "Shape dimension 2147483648 is outside the supported range " + "[-2147483648, 2147483647]. MLX currently uses 32-bit integers " + "for shape dimensions.", ) def test_dtype_promotion(self): @@ -964,6 +1126,11 @@ def test_indexing(self): a_sliced_mlx = a_mlx[-1] self.assertTrue(np.array_equal(a_sliced_mlx, a_npy[-1])) + # NumPy integer scalar indexing + a_sliced_mlx = a_mlx[np.int64(5)] + a_sliced_npy = np.asarray(a_sliced_mlx) + self.assertTrue(np.array_equal(a_sliced_npy, a_npy[np.int64(5)])) + # Basic content check, empty index a_sliced_mlx = a_mlx[()] a_sliced_npy = np.asarray(a_sliced_mlx) @@ -1105,7 +1272,7 @@ def test_indexing_grad(self): ind = mx.array([0, 1, 0]).astype(mx.float32) def index_fn(x, ind): - return x[ind.astype(mx.int32)].sum() + return x[mx.stop_gradient(ind.astype(mx.int32))].sum() grad_x, grad_ind = mx.grad(index_fn, argnums=(0, 1))(x, ind) expected = mx.array([[2, 2], [1, 1]]) @@ -1125,8 +1292,11 @@ def test_setitem(self): a[-1] = 2 self.assertEqual(a.tolist(), [2, 2, 2]) + a[np.int64(1)] = 9 + self.assertEqual(a.tolist(), [2, 9, 2]) + a[0] = mx.array([[[1]]]) - self.assertEqual(a.tolist(), [1, 2, 2]) + self.assertEqual(a.tolist(), [1, 9, 2]) a[:] = 0 self.assertEqual(a.tolist(), [0, 0, 0]) @@ -1140,6 +1310,28 @@ def test_setitem(self): a[0:2] = 3 self.assertEqual(a.tolist(), [3, 3, 1]) + # Assigning through a bare Ellipsis, like a[:] and a[None] + e = mx.zeros((2, 3), mx.int32) + e[...] = 5 + self.assertEqual(e.tolist(), [[5, 5, 5], [5, 5, 5]]) + + # Broadcasting an array update through Ellipsis + e[...] = mx.array([1, 2, 3]) + self.assertEqual(e.tolist(), [[1, 2, 3], [1, 2, 3]]) + + e[...] = mx.zeros((2, 3), mx.int32) + self.assertEqual(e.tolist(), [[0, 0, 0], [0, 0, 0]]) + + # Scalar array + e = mx.array(0) + e[...] = 7 + self.assertEqual(e.item(), 7) + + # Shapes that cannot broadcast are still rejected + e = mx.zeros((2, 3), mx.int32) + with self.assertRaises(ValueError): + e[...] = mx.array([1, 2]) + a[0:3] = 4 self.assertEqual(a.tolist(), [4, 4, 4]) @@ -1423,6 +1615,161 @@ def test_array_at(self): src = src.at[0:1].add(update) self.assertTrue(mx.array_equal(src, mx.array([[2.0, 4.0]]))) + # Test all array.at ops with slice-only indices + a = mx.random.uniform(shape=(10, 5, 2)) + update = mx.ones((2, 5)) + a[1:3, :, 0] = 0 + a = a.at[1:3, :, 0].add(update) + self.assertEqualArray(a[1:3, :, 0], update) + a = a.at[1:3, :, 0].subtract(update) + self.assertEqualArray(a[1:3, :, 0], mx.zeros_like(update)) + a = a.at[1:3, :, 0].add(2 * update) + self.assertEqualArray(a[1:3, :, 0], 2 * update) + a = a.at[1:3, :, 0].multiply(2 * update) + self.assertEqualArray(a[1:3, :, 0], 4 * update) + a = a.at[1:3, :, 0].divide(3 * update) + self.assertEqualArray(a[1:3, :, 0], (4 / 3) * update) + a[1:3, :, 0] = 5 + update = mx.arange(10).reshape(2, 5) + a = a.at[1:3, :, 0].maximum(update) + self.assertEqualArray(a[1:3, :, 0], mx.maximum(a[1:3, :, 0], update)) + a[1:3, :, 0] = 5 + a = a.at[1:3, :, 0].minimum(update) + self.assertEqualArray(a[1:3, :, 0], mx.minimum(a[1:3, :, 0], update)) + + @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") + def test_array_at_complex_add_gpu(self): + n = 4096 + base = [1 + 10j, 2 + 20j, 3 + 30j, 4 + 40j] + + with mx.stream(mx.gpu): + a = mx.array(base, dtype=mx.complex64) + update_indices = mx.full((n,), 3, dtype=mx.int32) + updates = mx.full((n,), 1 + 3j, dtype=mx.complex64) + out = a.at[update_indices].add(updates) + mx.eval(out) + + indices = mx.array([1, 1, 3]) + x = mx.array([1 + 0j, 3 + 4j, 6 + 8j, 5 + 12j], dtype=mx.complex64) + + def loss(z): + return mx.square(mx.abs(z[indices])).sum() + + _, gradient = mx.value_and_grad(loss)(x) + mx.eval(gradient) + + expected = base.copy() + expected[-1] += n * (1 + 3j) + self.assertEqual(out.tolist(), expected) + np.testing.assert_allclose( + np.array(gradient), + np.array([0, 12 + 16j, 0, 10 + 24j], dtype=np.complex64), + rtol=0, + atol=1e-5, + ) + + def test_array_at_slice_update_extensive(self): + # Test with transposed inputs + a = mx.zeros((4, 5)) + update = mx.ones((5, 2)).T # Shape (2, 5) + a = a.at[1:3, :].add(update) + self.assertEqualArray(a[1:3, :], update) + + # Test with transposed updates on transposed slice + a = mx.zeros((5, 4)) + update = mx.ones((2, 5)) + a = a.at[:, 1:3].add(update.T) + self.assertEqualArray(a[:, 1:3], update.T) + + # Test with slice of another array as update + source = mx.arange(20, dtype=mx.float32).reshape(4, 5) + a = mx.zeros((4, 5)) + update = source[1:3, :] # Shape (2, 5) + a = a.at[0:2, :].add(update) + self.assertEqualArray(a[0:2, :], source[1:3, :]) + + # Test with both input and update being slices + source = mx.arange(30, dtype=mx.float32).reshape(5, 6) + a = mx.zeros((5, 6)) + a = a.at[1:4, 1:5].add(source[0:3, 0:4]) + self.assertEqualArray(a[1:4, 1:5], source[0:3, 0:4]) + + # Test with transposed slice of another array + source = mx.arange(20, dtype=mx.float32).reshape(4, 5) + a = mx.zeros((5, 4)) + update = source[1:3, :].T # Shape (5, 2) + a = a.at[:, 1:3].add(update) + self.assertEqualArray(a[:, 1:3], update) + + # Test with negative indexing in slices + a = mx.zeros((5, 5)) + update = mx.ones((2, 5)) + a = a.at[-3:-1, :].add(update) + self.assertEqualArray(a[-3:-1, :], update) + + # Test with strided slices + a = mx.zeros((6, 6)) + update = mx.ones((2, 3)) + a = a.at[1:5:2, 0:6:2].add(update) + self.assertEqualArray(a[1:5:2, 0:6:2], update) + + # Test with slice of transposed array + source = mx.arange(20, dtype=mx.float32).reshape(4, 5) + a = mx.zeros((5, 4)) + update = source.T[:, 1:3] # Shape (5, 2) + a = a.at[:, 1:3].add(update) + self.assertEqualArray(a[:, 1:3], update) + + # Test with 3D arrays and transposed updates + a = mx.zeros((3, 4, 5)) + update = mx.ones((4, 3, 5)).transpose(1, 0, 2) # Shape (3, 4, 5) + a = a.at[:, :, :].add(update) + self.assertEqualArray(a, update) + + # Test with slice of 3D array + source = mx.arange(60, dtype=mx.float32).reshape(3, 4, 5) + a = mx.zeros((3, 4, 5)) + update = source[0:2, :, :] + a = a.at[1:3, :, :].add(update) + self.assertEqualArray(a[1:3, :, :], source[0:2, :, :]) + + # Test with mixed slice and index + a = mx.zeros((4, 5, 6)) + update = mx.ones((2, 6)) + a = a.at[1:3, 2, :].add(update) + self.assertEqualArray(a[1:3, 2, :], update) + + # Test with update from strided slice + source = mx.arange(60, dtype=mx.float32).reshape(3, 4, 5) + a = mx.zeros((3, 2, 5)) + update = source[:, ::2, :] # Shape (3, 2, 5) + a = a.at[:, :, :].add(update) + self.assertEqualArray(a, update) + + def test_slice_update_contiguous_2d(self): + for shape in [(32, 32), (64, 64), (17, 33), (128, 128)]: + for upd_shape in [(16, 16), (8, 8), (4, 4)]: + if upd_shape[0] > shape[0] or upd_shape[1] > shape[1]: + continue + y = mx.zeros(shape) + x = mx.random.normal(upd_shape) + z = y.at[: upd_shape[0], : upd_shape[1]].add(x) + diff = z[: upd_shape[0], : upd_shape[1]] - x + self.assertTrue(mx.allclose(diff, mx.zeros_like(diff))) + + # Test non-zero offset slice update + y = mx.zeros((32, 32)) + x = mx.random.normal((16, 16)) + z = y.at[16:, 16:].add(x) + self.assertTrue(mx.allclose(z[16:, 16:] - x, mx.zeros_like(x))) + + # Test with size divisible by 4, 2, and odd + for cols in [32, 18, 15]: + y = mx.zeros((32, cols)) + x = mx.random.normal((16, cols)) + z = y.at[:16, :].add(x) + self.assertTrue(mx.allclose(z[:16, :] - x, mx.zeros_like(x))) + def test_slice_negative_step(self): a_np = np.arange(20) a_mx = mx.array(a_np) @@ -1641,11 +1988,10 @@ def test_buffer_protocol(self): mv_mx = memoryview(a_mx) self.assertEqual(mv_mx.strides, (8, 2)) self.assertEqual(mv_mx.shape, (3, 4)) - self.assertEqual(mv_mx.format, "B") - with self.assertRaises(RuntimeError) as cm: + self.assertIn(mv_mx.format, "bfloat16") + with self.assertRaises(ValueError) as cm: np.array(a_mx) - e = cm.exception - self.assertTrue("Item size 2 for PEP 3118 buffer format string" in str(e)) + self.assertIn("bfloat16", str(cm.exception)) # Test buffer protocol with non-arrays ie bytes a = ord("a") * 257 + mx.arange(10).astype(mx.int16) @@ -1678,6 +2024,15 @@ def test_array_view_ref_counting(self): a_np = None self.assertIsNone(wr()) + def test_create_from_buffer(self): + x = mx.array(b"Hello") + self.assertEqual(x.dtype, mx.uint8) + self.assertEqual(x.tolist(), [72, 101, 108, 108, 111]) + + x = mx.array(bytearray([1, 2, 3])) + self.assertEqual(x.dtype, mx.uint8) + self.assertEqual(x.tolist(), [1, 2, 3]) + @unittest.skipIf(not has_tf, "requires TensorFlow") def test_buffer_protocol_tf(self): dtypes_list = [ @@ -1866,19 +2221,395 @@ def test_add_numpy(self): self.assertEqual(z.item(), 3) def test_dlpack(self): + class CpuDLPack: + def __init__(self, array): + self.array = array + + def __dlpack_device__(self): + return (1, 0) + + def __dlpack__(self, *args, **kwargs): + kwargs["dl_device"] = (1, 0) + return self.array.__dlpack__(*args, **kwargs) + x = mx.array(1, dtype=mx.int32) - y = np.from_dlpack(x) + y = np.from_dlpack(CpuDLPack(x)) self.assertTrue(mx.array_equal(y, x)) x = mx.array([[1.0, 2.0], [3.0, 4.0]]) - y = np.from_dlpack(x) + y = np.from_dlpack(CpuDLPack(x)) self.assertTrue(mx.array_equal(y, x)) x = mx.arange(16).reshape(4, 4) x = x[::2, ::2] - y = np.from_dlpack(x) + y = np.from_dlpack(CpuDLPack(x)) self.assertTrue(mx.array_equal(y, x)) + def test_from_dlpack_cpu(self): + x = np.arange(3, dtype=np.float32) + + # copy=None may adopt the buffer or copy; either way the values match + # the source at import time. + y = mx.from_dlpack(x) + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + + # copy=True always copies, so later mutations of the source are not seen. + y = mx.from_dlpack(x, copy=True) + x += 10 + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + + # copy=False adopts the buffer when possible and raises otherwise; it + # must never silently copy. + x = np.arange(3, dtype=np.float32) + try: + y = mx.from_dlpack(x, copy=False) + x += 10 + except ValueError: + pass + else: + self.assertEqual(y.tolist(), [10.0, 11.0, 12.0]) + + def test_dlpack_cpu_dtype_mapping(self): + class CpuDLPack: + def __init__(self, array): + self.array = array + + def __dlpack_device__(self): + return (1, 0) + + def __dlpack__(self, *args, **kwargs): + kwargs["dl_device"] = (1, 0) + return self.array.__dlpack__(*args, **kwargs) + + dlpack_to_mlx = [ + (np.bool_, mx.bool_), + (np.uint8, mx.uint8), + (np.uint16, mx.uint16), + (np.uint32, mx.uint32), + (np.uint64, mx.uint64), + (np.int8, mx.int8), + (np.int16, mx.int16), + (np.int32, mx.int32), + (np.int64, mx.int64), + (np.float16, mx.float16), + (np.float32, mx.float32), + (np.float64, mx.float32), + (np.complex64, mx.complex64), + (np.complex128, mx.complex64), + ] + for np_dtype, mlx_dtype in dlpack_to_mlx: + with self.subTest(direction="import", dtype=np_dtype): + x = np.ones(3, dtype=np_dtype) + y = mx.from_dlpack(x) + self.assertEqual(y.dtype, mlx_dtype) + + if torch is not None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + x = torch.ones(3, dtype=torch.complex32) + with self.assertRaises(ValueError): + mx.from_dlpack(x) + + mlx_to_dlpack = [ + (mx.bool_, np.bool_), + (mx.uint8, np.uint8), + (mx.uint16, np.uint16), + (mx.uint32, np.uint32), + (mx.uint64, np.uint64), + (mx.int8, np.int8), + (mx.int16, np.int16), + (mx.int32, np.int32), + (mx.int64, np.int64), + (mx.float16, np.float16), + (mx.float32, np.float32), + (mx.complex64, np.complex64), + ] + for mlx_dtype, np_dtype in mlx_to_dlpack: + with self.subTest(direction="export", dtype=mlx_dtype): + x = mx.ones((3,), dtype=mlx_dtype) + y = np.from_dlpack(CpuDLPack(x)) + self.assertEqual(y.dtype, np_dtype) + + if torch is not None and has_torch_mps: + x = mx.ones((3,), dtype=mx.bfloat16) + y = torch.from_dlpack(x) + self.assertEqual(y.dtype, torch.bfloat16) + + def test_from_dlpack_cpu_strided(self): + x = np.arange(12, dtype=np.float32).reshape(3, 4) + view = x.T + y = mx.from_dlpack(view) + + self.assertEqual(y.tolist(), view.tolist()) + self.assertFalse(memoryview(y).c_contiguous) + self.assertEqual(memoryview(y).strides, view.strides) + + stepped = np.arange(20, dtype=np.int32)[2:10:2] + y = mx.from_dlpack(stepped) + self.assertEqual(y.tolist(), [2, 4, 6, 8]) + self.assertFalse(memoryview(y).c_contiguous) + self.assertEqual(memoryview(y).strides, stepped.strides) + + broadcast = np.broadcast_to(np.array([7], dtype=np.int32), (3,)) + y = mx.from_dlpack(broadcast) + self.assertEqual(y.tolist(), [7, 7, 7]) + self.assertFalse(memoryview(y).c_contiguous) + self.assertEqual(memoryview(y).strides, broadcast.strides) + + negative_stride = np.arange(5, dtype=np.float32)[::-1] + with self.assertRaises(ValueError): + mx.from_dlpack(negative_stride) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_import(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + self.assertEqual(x.__dlpack_device__()[0], 8) + + y = mx.asarray(x) + self.assertEqual(y.dtype, mx.float32) + torch.mps.synchronize() + self.assertEqual(y.tolist(), x.cpu().numpy().tolist()) + self.assertIn("array(", repr(y)) + mv = memoryview(y) + self.assertEqual(mv.tolist(), x.cpu().numpy().tolist()) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_array_copies_dlpack_input(self): + assert torch is not None + x = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + y = mx.array(x) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_asarray_copy_true_copies_dlpack_input(self): + assert torch is not None + x = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + y = mx.asarray(x, copy=True) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_zero_copy_shares_updates(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + torch.mps.synchronize() + y = mx.asarray(x) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), x.cpu().numpy().tolist()) + + y += 10 + mx.eval(y) + self.assertEqual(x.cpu().numpy().tolist(), y.tolist()) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_matching_dtype_argument_shares_updates(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + torch.mps.synchronize() + y = mx.asarray(x, dtype=mx.float32, copy=False) + self.assertEqual(y.dtype, mx.float32) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), x.cpu().numpy().tolist()) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_different_dtype_argument_copies(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + torch.mps.synchronize() + z = mx.asarray(x, dtype=mx.float16) + expected = x.to(torch.float16).cpu().numpy().tolist() + + self.assertEqual(z.dtype, mx.float16) + self.assertEqual(z.tolist(), expected) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(z.tolist(), expected) + + with self.assertRaises(ValueError): + mx.asarray(x, dtype=mx.float16, copy=False) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_data_offset(self): + assert torch is not None + view = torch.arange(12, device="mps", dtype=torch.float32)[3:9] + view_mx = mx.asarray(view) + torch.mps.synchronize() + self.assertEqual(view_mx.tolist(), view.cpu().numpy().tolist()) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_strided_view(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + view = x.T + torch.mps.synchronize() + y = mx.asarray(view, copy=False) + self.assertEqual(y.tolist(), view.cpu().numpy().tolist()) + + x[0, 1] = 99 + torch.mps.synchronize() + self.assertEqual(y.tolist(), view.cpu().numpy().tolist()) + + y_copy = mx.asarray(view, copy=True) + expected = view.cpu().numpy().tolist() + self.assertFalse(memoryview(y_copy).c_contiguous) + self.assertEqual( + memoryview(y_copy).strides, + tuple(s * view.element_size() for s in view.stride()), + ) + x[0, 2] = 77 + torch.mps.synchronize() + self.assertEqual(y_copy.tolist(), expected) + + z = mx.asarray(view, dtype=mx.float16) + self.assertEqual(z.dtype, mx.float16) + self.assertFalse(memoryview(z).c_contiguous) + self.assertEqual(memoryview(z).strides, tuple(s * 2 for s in view.stride())) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_stepped_view(self): + x = torch.arange(20, device="mps", dtype=torch.int32) + view = x[2:10:2] + torch.mps.synchronize() + y = mx.asarray(view, copy=False) + self.assertEqual(y.tolist(), [2, 4, 6, 8]) + + x[4] = 99 + torch.mps.synchronize() + self.assertEqual(y.tolist(), [2, 99, 6, 8]) + + y_copy = mx.asarray(view, copy=True) + expected = y.tolist() + x[6] = 77 + torch.mps.synchronize() + self.assertEqual(y_copy.tolist(), expected) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_broadcast_stride(self): + assert torch is not None + x = torch.tensor([7], device="mps", dtype=torch.int32) + view = x.expand(3) + torch.mps.synchronize() + y = mx.asarray(view, copy=False) + self.assertEqual(y.tolist(), [7, 7, 7]) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0, 0, 0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_dlpack_bfloat16(self): + assert torch is not None + x = torch.arange(12, device="mps", dtype=torch.float32).reshape(3, 4) + bf = x.to(torch.bfloat16) + bf_mx = mx.asarray(bf) + + self.assertEqual(bf_mx.dtype, mx.bfloat16) + torch.mps.synchronize() + self.assertEqual( + bf_mx.astype(mx.float32).tolist(), + bf.to(torch.float32).cpu().numpy().tolist(), + ) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_torch_mps_array_operand(self): + assert torch is not None + a = mx.array([1]) + b = torch.tensor([2]) + self.assertTrue(mx.array_equal(a + b, mx.array([3]))) + + b_mps = b.to("mps") + torch.mps.synchronize() + self.assertTrue(mx.array_equal(a + b_mps, mx.array([3]))) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_mlx_dlpack_exports_mps_tensor_to_torch(self): + assert torch is not None + x = mx.array([1]).astype(mx.float16) + mx.eval(x) + y = torch.utils.dlpack.from_dlpack(x) + torch.mps.synchronize() + + self.assertEqual(y.device.type, "mps") + self.assertEqual(y.dtype, torch.float16) + self.assertEqual(y.cpu().numpy().tolist(), [1.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_mlx_dlpack_exports_mps_tensor_to_torch_tensor(self): + assert torch is not None + x = mx.array([1]).astype(mx.float16) + mx.eval(x) + y = torch.tensor(x) + torch.mps.synchronize() + + self.assertEqual(y.device.type, "mps") + self.assertEqual(y.dtype, torch.float16) + self.assertEqual(y.cpu().numpy().tolist(), [1.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_mlx_dlpack_export_torch_update_writes_mlx_buffer(self): + x = mx.arange(8, dtype=mx.float32) + y = x[2:6] + mx.eval(y) + t = torch.utils.dlpack.from_dlpack(y) + + self.assertEqual(t.device.type, "mps") + self.assertEqual(t.cpu().numpy().tolist(), [2.0, 3.0, 4.0, 5.0]) + + t.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 0.0, 0.0, 0.0]) + self.assertEqual(x.tolist(), [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 6.0, 7.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_from_dlpack_torch_mps_copy_none_shares_updates(self): + assert torch is not None + x = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + y = mx.from_dlpack(x) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 0.0, 0.0]) + + y += 10 + mx.eval(y) + self.assertEqual(x.cpu().numpy().tolist(), [10.0, 10.0, 10.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_from_dlpack_torch_mps_copy_false_shares_updates(self): + assert torch is not None + x = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + y = mx.from_dlpack(x, copy=False) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 0.0, 0.0]) + + @unittest.skipUnless(has_torch_mps, "PyTorch MPS is required") + def test_from_dlpack_torch_mps_copy_true_copies(self): + assert torch is not None + x = torch.arange(3, device="mps", dtype=torch.float32) + torch.mps.synchronize() + y = mx.from_dlpack(x, copy=True) + + x.zero_() + torch.mps.synchronize() + self.assertEqual(y.tolist(), [0.0, 1.0, 2.0]) + def test_getitem_with_list(self): a = mx.array([1, 2, 3, 4, 5]) idx = [0, 2, 4] @@ -1992,6 +2723,18 @@ def test_array_namespace_asarray(self): arr_pass = xp.asarray(existing) self.assertEqual(arr_pass.tolist(), [4, 5, 6]) + def test_asarray_copy(self): + existing = mx.array([1, 2, 3]) + + self.assertEqual(mx.asarray(existing, copy=True).tolist(), [1, 2, 3]) + self.assertEqual( + mx.asarray(existing, dtype=mx.float32, copy=True).dtype, mx.float32 + ) + with self.assertRaises(ValueError): + mx.asarray(existing, copy=False) + with self.assertRaises(ValueError): + mx.asarray(existing, dtype=mx.float32, copy=False) + def test_asarray(self): # List inputs self.assertEqual(mx.asarray([1, 2, 3]).tolist(), [1, 2, 3]) @@ -2014,17 +2757,33 @@ def test_asarray(self): # MLX array inputs arr = mx.array([1, 2, 3]) self.assertEqual(mx.asarray(arr).tolist(), [1, 2, 3]) + self.assertEqual(mx.asarray(arr, copy=True).tolist(), [1, 2, 3]) + with self.assertRaises(ValueError): + mx.asarray(arr, copy=False) arr_int = mx.array([1, 2, 3], dtype=mx.int32) arr_float = mx.asarray(arr_int, dtype=mx.float32) self.assertEqual(arr_float.dtype, mx.float32) self.assertEqual(arr_float.tolist(), [1.0, 2.0, 3.0]) + with self.assertRaises(ValueError): + mx.asarray(arr_int, dtype=mx.float32, copy=False) # NumPy array inputs np_arr = np.array([1.0, 2.0, 3.0], dtype=np.float32) mx_arr = mx.asarray(np_arr) self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0]) self.assertEqual(mx_arr.dtype, mx.float32) + # copy=False adopts the buffer when possible and raises otherwise; it + # must never silently copy. + try: + mx_arr = mx.asarray(np_arr, copy=False) + except ValueError: + pass + else: + self.assertEqual(mx_arr.tolist(), [1.0, 2.0, 3.0]) + + with self.assertRaises(ValueError): + mx.asarray([1, 2, 3], copy=False) # dtype parameter self.assertEqual(mx.asarray([1, 2, 3], dtype=mx.float32).dtype, mx.float32) @@ -2034,16 +2793,23 @@ def test_to_scalar(self): a = mx.array(1) self.assertEqual(int(a), 1) self.assertEqual(float(a), 1) + self.assertEqual(complex(a), 1 + 0j) a = mx.array(1.5) self.assertEqual(float(a), 1.5) self.assertEqual(int(a), 1) + self.assertEqual(complex(a), 1.5 + 0j) + + a = mx.array(1 + 2j, dtype=mx.complex64) # type: ignore + self.assertEqual(complex(a), 1 + 2j) a = mx.zeros((2, 1)) with self.assertRaises(ValueError): float(a) with self.assertRaises(ValueError): int(a) + with self.assertRaises(ValueError): + complex(a) def test_format(self): a = mx.arange(3) @@ -2085,31 +2851,6 @@ def test_deep_graphs(self): x = mx.sin(x) mx.eval(x) - def test_siblings_without_eval(self): - def get_mem(): - if platform.system() == "Windows": - process = psutil.Process(os.getpid()) - return process.memory_info().peak_wset - else: - return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - - key = mx.array([1, 2]) - - def t(): - a, b = mx.split(key, 2) - a = mx.reshape(a, []) - b = mx.reshape(b, []) - return b - - mx.synchronize() - t() - gc.collect() - expected = get_mem() - for _ in range(100): - t() - used = get_mem() - self.assertEqual(expected, used) - def test_scalar_integer_conversion_overflow(self): y = mx.array(2000000000, dtype=mx.int32) x = 3000000000 diff --git a/Source/Cmlx/mlx/python/tests/test_autograd.py b/Source/Cmlx/mlx/python/tests/test_autograd.py index c37161a4d..c85fcf09a 100644 --- a/Source/Cmlx/mlx/python/tests/test_autograd.py +++ b/Source/Cmlx/mlx/python/tests/test_autograd.py @@ -1,10 +1,20 @@ # Copyright © 2023 Apple Inc. import gc +import itertools +import math import unittest import mlx.core as mx import mlx_tests +import numpy as np + +try: + import torch + + has_torch = True +except ImportError: + has_torch = False class TestAutograd(mlx_tests.MLXTestCase): @@ -30,6 +40,100 @@ def test_jvp(self): self.assertEqual(out[0].item(), 4.0 * 1.0 + 2.0 * 3.0) self.assertEqual(out[1].item(), 4.0 * 1.0 + 6.0 * 3.0) + def test_jvp_comparison_tangent_dtype(self): + # Comparison op JVP tangents should preserve the input tangent's + # dtype (e.g. float32), not return bool. Using bool tangents causes + # downstream ops like negative to crash. (issue #3081) + x = mx.array([1.0, -2.0, 3.0]) + t = mx.ones_like(x) + + for op in [ + mx.greater, + mx.less, + mx.equal, + mx.greater_equal, + mx.less_equal, + mx.not_equal, + ]: + _, tangents = mx.jvp(lambda x, _op=op: _op(x, 0.0), [x], [t]) + self.assertEqual(tangents[0].dtype, mx.float32) + + def test_jvp_with_constant_inputs(self): + # JVPs of primitives with only a subset of inputs traced used to + # index tangents out of bounds and silently return wrong tangents + # (issue #3627) + + # d/dt (t + 1)^2 = 2 at t = 0 + cases = [ + lambda t: mx.where((t + 1) ** 2 > -1, (t + 1) ** 2, 999.0), + lambda t: mx.where((t + 1) ** 2 < -1, 999.0, (t + 1) ** 2), + lambda t: mx.where(mx.array(True), (t + 1) ** 2, 999.0), + lambda t: mx.where(mx.array(False), 999.0, (t + 1) ** 2), + ] + for fun in cases: + _, (dout,) = mx.jvp(fun, [mx.array(0.0)], [mx.array(1.0)]) + self.assertEqual(dout.item(), 2.0) + + # Constant condition with both branches traced + _, (dout,) = mx.jvp( + lambda a, b: mx.where(mx.array([True, False]), a, b), + [mx.zeros(2), mx.zeros(2)], + [mx.array([1.0, 2.0]), mx.array([3.0, 4.0])], + ) + self.assertTrue(mx.array_equal(dout, mx.array([1.0, 4.0]))) + + # The tangent of a where with only the condition traced is zero + # with the output's dtype + _, (dout,) = mx.jvp( + lambda c: mx.where(c > 0, 2.0, 3.0), [mx.array(1.0)], [mx.array(1.0)] + ) + self.assertEqual(dout.item(), 0.0) + self.assertEqual(dout.dtype, mx.float32) + + # d/dy atan2(y, x) = x / (x^2 + y^2) + _, (dout,) = mx.jvp( + lambda y: mx.arctan2(y, mx.array(2.0)), [mx.array(1.0)], [mx.array(1.0)] + ) + self.assertAlmostEqual(dout.item(), 0.4, places=6) + + # d/dx atan2(y, x) = -y / (x^2 + y^2) + _, (dout,) = mx.jvp( + lambda x: mx.arctan2(mx.array(2.0), x), [mx.array(1.0)], [mx.array(1.0)] + ) + self.assertAlmostEqual(dout.item(), -0.4, places=6) + + # masked_scatter with a constant destination + mask = mx.array([True, False, True, False]) + + def masked_set(src): + dst = mx.zeros(4) + dst[mask] = src + return dst + + _, (dout,) = mx.jvp(masked_set, [mx.ones(2)], [mx.ones(2)]) + self.assertTrue(mx.array_equal(dout, mx.array([1.0, 0.0, 1.0, 0.0]))) + + def test_jvp_through_bitwise_ops(self): + # JVPs of bitwise ops returned one tangent per traced input instead + # of one per output which corrupted the tangents of downstream + # outputs (issue #3629). The corruption is out-of-bounds UB that can + # go unnoticed in release builds; the assert in the jvp transform + # catches it deterministically in debug builds. + def fun(x): + b = (x == 0) & (x > -1) + return x + b.astype(mx.float32) + + x = mx.array([0.0, 1.0, 2.0]) + _, (dout,) = mx.jvp(fun, [x], [mx.ones_like(x)]) + self.assertTrue(mx.array_equal(dout, mx.ones_like(x))) + + def fun(x): + b = (x == 0) | (x > 1) + return x * b.astype(mx.float32) + + _, (dout,) = mx.jvp(fun, [x], [mx.ones_like(x)]) + self.assertTrue(mx.array_equal(dout, mx.array([1.0, 0.0, 1.0]))) + def test_vjp(self): fun = lambda x: 2 * x out, dout = mx.vjp(fun, [mx.array(1.0)], [mx.array(2.0)]) @@ -52,6 +156,50 @@ def test_vjp(self): self.assertEqual(out[1].item(), 2.0 * 1.0 + 6.0 * 3.0) self.assertEqual(out[2].item(), 4.0 * 3.0) + def test_jvp_with_partly_traced_inputs(self): + # power: each traced input must use its own tangent (issue #3634) + fun = lambda a, b: a**b + primals = [mx.array(2.0), mx.array(3.0)] + dyda = 12.0 # d/da a^b = b * a^(b - 1) + dydb = math.log(2.0) * 8.0 # d/db a^b = ln(a) * a^b + _, (j,) = mx.jvp(fun, primals, [mx.array(1.0), mx.array(0.0)]) + self.assertAlmostEqual(j.item(), dyda, places=4) + _, (j,) = mx.jvp(fun, primals, [mx.array(0.0), mx.array(1.0)]) + self.assertAlmostEqual(j.item(), dydb, places=4) + _, (j,) = mx.jvp(fun, primals, [mx.array(1.0), mx.array(1.0)]) + self.assertAlmostEqual(j.item(), dyda + dydb, places=4) + fun = lambda a: a ** mx.array(3.0) + _, (j,) = mx.jvp(fun, [mx.array(2.0)], [mx.array(1.0)]) + self.assertAlmostEqual(j.item(), dyda, places=4) + fun = lambda b: mx.array(2.0) ** b + _, (j,) = mx.jvp(fun, [mx.array(3.0)], [mx.array(1.0)]) + self.assertAlmostEqual(j.item(), dydb, places=4) + + # divmod: one tangent per output, and consuming the second output + # used to crash (issue #3634) + def fun(x): + q, r = mx.divmod(x, mx.array(3.0)) + return q, -r + + _, (dq, dr) = mx.jvp(fun, [mx.array(7.0)], [mx.array(1.0)]) + self.assertEqual(dq.item(), 0.0) + self.assertEqual(dr.item(), 0.0) + + # slice_update with dynamic start indices and a subset of traced + # inputs (issue #3634) + src = mx.zeros(4) + upd = mx.ones(2) + start = mx.array([1]) + fun = lambda u: mx.slice_update(src, u, start, axes=[0]) + _, (j,) = mx.jvp(fun, [upd], [mx.ones(2)]) + self.assertEqual(j.tolist(), [0.0, 1.0, 1.0, 0.0]) + fun = lambda s: mx.slice_update(s, upd, start, axes=[0]) + _, (j,) = mx.jvp(fun, [src], [mx.ones(4)]) + self.assertEqual(j.tolist(), [1.0, 0.0, 0.0, 1.0]) + fun = lambda s, u: mx.slice_update(s, u, start, axes=[0]) + _, (j,) = mx.jvp(fun, [src, upd], [mx.ones(4), mx.full(2, 2.0)]) + self.assertEqual(j.tolist(), [1.0, 2.0, 2.0, 1.0]) + def test_grad(self): fun = lambda x: x * x @@ -282,65 +430,230 @@ def fun(x, idx): x[idx] = 2.0 return x.sum() - dfdx = mx.grad(fun)(mx.array([1.0, 2.0, 3.0]), mx.array([1])) - self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 0.0, 1.0]))) + dfdx = mx.grad(fun)(mx.array([1.0, 2.0, 3.0, 4.0]), mx.array([1, 3])) + self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 0.0, 1.0, 0.0]))) self.assertEqual(dfdx.dtype, mx.float32) - y = mx.array([0.0, 1.0, 2.0]) + y = mx.array([0.0, 1.0, 2.0, 3.0]) def fun(x, idx): y[idx] = x return y.sum() - dfdx = mx.grad(fun)(mx.array([2.0]), mx.array([1])) - self.assertTrue(mx.array_equal(dfdx, mx.array([1.0]))) + dfdx = mx.grad(fun)(mx.array([2.0, 3.0]), mx.array([1, 3])) + self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 1.0]))) self.assertEqual(dfdx.dtype, mx.float32) + def test_index_vjp_requires_stop_gradient(self): + msg = "stop_gradient" + x = mx.array([1.0, 2.0, 3.0, 4.0]) + idx = mx.array([1, 3]) + updates = mx.array([5.0, 6.0]) + x_axis = x[:, None] + idx_axis = idx[:, None] + updates_axis = updates[:, None] + + def gather_fun(x, idx): + return mx.take(x, idx) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_fun, [x, idx], [mx.ones((2,))]) + + def gather_axis_fun(x, idx): + return mx.take_along_axis(x, idx, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_axis_fun, [x_axis, idx_axis], [mx.ones((2, 1))]) + + def scatter_fun(x, idx, updates): + return x.at[idx].add(updates) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(scatter_fun, [x, idx, updates], [mx.ones((4,))]) + + def scatter_axis_fun(x, idx, updates): + return mx.put_along_axis(x, idx, updates, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp( + scatter_axis_fun, + [x_axis, idx_axis, updates_axis], + [mx.ones((4, 1))], + ) + + def test_stop_gradient_computed_indices(self): + def gather_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + return mx.take(w, idx).sum() + + grad = mx.grad(gather_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([0.0, 0.0, 1.0, 1.0]))) + + def gather_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + return mx.take_along_axis(w, idx, axis=1).sum() + + grad = mx.grad(gather_axis_fun)(mx.array([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0]])) + self.assertTrue( + mx.array_equal( + grad, + mx.array([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]), + ) + ) + + def scatter_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + out = mx.zeros((4,)) + out[idx] = w[:2] + return out.sum() + + grad = mx.grad(scatter_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([1.0, 1.0, 0.0, 0.0]))) + + def scatter_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + updates = w.sum(axis=1, keepdims=True) + out = mx.put_along_axis(mx.zeros((3, 3)), idx, updates, axis=1) + return out.sum() + + grad = mx.grad(scatter_axis_fun)(mx.ones((3, 3))) + self.assertTrue(mx.array_equal(grad, mx.ones((3, 3)))) + + def test_take_along_axis_complex_vjp(self): + x = mx.zeros((4,), dtype=mx.complex64) + indices = mx.array([3, 3], dtype=mx.int32) + cotangent = mx.array([1 + 2j, 3 + 4j], dtype=mx.complex64) + _, (gradient,) = mx.vjp( + lambda z: mx.take_along_axis(z, indices, axis=0), + [x], + [cotangent], + ) + mx.eval(gradient) + self.assertEqualArray( + gradient, + mx.array([0j, 0j, 0j, 4 + 6j], dtype=mx.complex64), + atol=0, + rtol=0, + ) + + def test_scatter_add_vjp(self): + def fun(src, updates): + x = src.at[mx.array([1, 3])].add(updates) + return x + + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([1.0, 2.0]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0]))) + def test_scatter_max_vjp(self): def fun(src, updates): - x = src.at[1].maximum(updates) + x = src.at[mx.array([1, 3])].maximum(updates) return x - cotan = mx.array([4.0, 5.0, 6.0]) - _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan]) + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([1.0, 2.0]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) mx.eval(vjps) - # Update larger than value - self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0]))) - self.assertTrue(mx.allclose(vjps[1], mx.array([5.0]))) + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([0.0, 0.0]))) - cotan = mx.array([[4.0], [5.0], [6.0]]) - _, vjps = mx.vjp( - fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan] - ) + updates = mx.array([5.0, 6.0]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) mx.eval(vjps) - # Update and value are equal - self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]]))) - self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]]))) + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0, 0.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0]))) def test_scatter_min_vjp(self): def fun(src, updates): - x = src.at[1].minimum(updates) + x = src.at[mx.array([1, 3])].minimum(updates) return x - cotan = mx.array([4.0, 5.0, 6.0]) - _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0]), mx.array([[3.0]])], [cotan]) + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([5.0, 6.0]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) mx.eval(vjps) - # Update larger than value - self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0]))) - self.assertTrue(mx.allclose(vjps[1], mx.array([0.0]))) + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([0.0, 0.0]))) - cotan = mx.array([[4.0], [5.0], [6.0]]) - _, vjps = mx.vjp( - fun, [mx.array([[1.0], [2.0], [3.0]]), mx.array([[[2.0]]])], [cotan] - ) + updates = mx.array([1.0, 1.0]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 6.0, 0.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([5.0, 7.0]))) + + def test_slice_update_max_vjp(self): + def fun(src, updates): + x = src.at[1:3].maximum(updates) + return x + + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([[1.0, 2.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[0.0, 0.0]]))) + + updates = mx.array([[5.0, 6.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 0.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]]))) + + def test_slice_update_min_vjp(self): + def fun(src, updates): + x = src.at[1:3].minimum(updates) + return x + + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([[5.0, 6.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[0.0, 0.0]]))) + + updates = mx.array([[1.0, 1.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) mx.eval(vjps) - # Update and value are equal - self.assertTrue(mx.allclose(vjps[0], mx.array([[4.0], [5.0], [6.0]]))) - self.assertTrue(mx.allclose(vjps[1], mx.array([[[5.0]]]))) + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 0.0, 0.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]]))) + + def test_slice_update_add_vjp(self): + def fun(src, updates): + x = src.at[1:3].add(updates) + return x + + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([[1.0, 2.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 5.0, 6.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[5.0, 6.0]]))) + + def test_slice_update_multiply_vjp(self): + def fun(src, updates): + x = src.at[1:3].multiply(updates) + return x + + cotan = mx.array([4.0, 5.0, 6.0, 7.0]) + updates = mx.array([[2.0, 3.0]]) + _, vjps = mx.vjp(fun, [mx.array([1.0, 2.0, 3.0, 4.0]), updates], [cotan]) + mx.eval(vjps) + + self.assertTrue(mx.allclose(vjps[0], mx.array([4.0, 10.0, 18.0, 7.0]))) + self.assertTrue(mx.allclose(vjps[1], mx.array([[10.0, 18.0]]))) def test_split_against_slice(self): def f_split(x): @@ -499,6 +812,168 @@ def fun(y): expected = mx.array([0.0, 0.0, 0.0, 9.0, 1.0]) self.assertTrue(mx.allclose(out, expected)) + def test_cummax_grad(self): + # Ties route to the latest occurrence, matching the cummax indices. + a = mx.array([3.0, 3.0, 1.0, 5.0, 5.0]) + + def fun(y): + return mx.cummax(y).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([1.0, 2.0, 0.0, 1.0, 1.0])) + ) + + def fun(y): + return mx.cummax(y, inclusive=False).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([1.0, 2.0, 0.0, 1.0, 0.0])) + ) + + def fun(y): + return mx.cummax(y, reverse=True).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([0.0, 0.0, 0.0, 4.0, 1.0])) + ) + + def fun(y): + return mx.cummax(y, reverse=True, inclusive=False).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([0.0, 0.0, 0.0, 3.0, 1.0])) + ) + + # Non-uniform cotangents are routed to the owning index. + cot = mx.array([10.0, 1.0, 1.0, 100.0, 1000.0]) + _, vjps = mx.vjp(lambda y: mx.cummax(y), (a,), (cot,)) + self.assertTrue(mx.allclose(vjps[0], mx.array([10.0, 2.0, 0.0, 100.0, 1000.0]))) + + # 2D along an inner axis. + m = mx.array([[1.0, 3.0, 3.0], [4.0, 2.0, 4.0]]) + + def fun(y): + return mx.cummax(y, axis=1).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(m), mx.array([[1.0, 1.0, 1.0], [2.0, 0.0, 1.0]])) + ) + + def fun(y): + return mx.cummax(y, axis=1, inclusive=False).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(m), mx.array([[1.0, 1.0, 0.0], [2.0, 0.0, 0.0]])) + ) + + def test_cummin_grad(self): + a = mx.array([3.0, 3.0, 1.0, 5.0, 5.0]) + + def fun(y): + return mx.cummin(y).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([1.0, 1.0, 3.0, 0.0, 0.0])) + ) + + def fun(y): + return mx.cummin(y, inclusive=False).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([1.0, 1.0, 2.0, 0.0, 0.0])) + ) + + def fun(y): + return mx.cummin(y, reverse=True).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([0.0, 0.0, 3.0, 1.0, 1.0])) + ) + + def fun(y): + return mx.cummin(y, reverse=True, inclusive=False).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(a), mx.array([0.0, 0.0, 2.0, 1.0, 1.0])) + ) + + cot = mx.array([10.0, 1.0, 1.0, 100.0, 1000.0]) + _, vjps = mx.vjp(lambda y: mx.cummin(y), (a,), (cot,)) + self.assertTrue(mx.allclose(vjps[0], mx.array([10.0, 1.0, 1101.0, 0.0, 0.0]))) + + # 2D along the outer axis. + m = mx.array([[1.0, 3.0, 3.0], [4.0, 2.0, 4.0]]) + + def fun(y): + return mx.cummin(y, axis=0).sum() + + self.assertTrue( + mx.allclose(mx.grad(fun)(m), mx.array([[2.0, 1.0, 2.0], [0.0, 1.0, 0.0]])) + ) + + @unittest.skipIf(not has_torch, "requires Torch") + def test_cummax_cummin_grad_vs_torch(self): + # Cross-check the cumulative max/min VJP against PyTorch autograd over + # axes, scan direction, inclusive/exclusive modes, ties, and weighted + # cotangents. Torch has no reverse or exclusive scan, so reverse is + # emulated by flipping along the axis and exclusive by shifting the + # inclusive scan one step (its leading element carries no gradient). + def torch_scan(x, axis, reverse, inclusive, op): + xf = torch.flip(x, [axis]) if reverse else x + scan = torch.cummax if op == "max" else torch.cummin + c = scan(xf, axis).values + if not inclusive: + n = c.size(axis) + head = torch.zeros_like(c.narrow(axis, 0, 1)) # constant, no grad + c = torch.cat([head, c.narrow(axis, 0, n - 1)], dim=axis) + return torch.flip(c, [axis]) if reverse else c + + def mx_scan(z, axis, reverse, inclusive, op): + scan = mx.cummax if op == "max" else mx.cummin + return scan(z, axis=axis, reverse=reverse, inclusive=inclusive) + + inputs = [ + np.array([3.0, 3.0, 1.0, 5.0, 5.0, 1.0, 5.0, 2.0], dtype=np.float32), + np.array( + [[1.0, 3.0, 3.0, 2.0], [4.0, 2.0, 4.0, 4.0], [4.0, 3.0, 1.0, 4.0]], + dtype=np.float32, + ), + ] + rng = np.random.default_rng(0) + for x_np in inputs: + cotangents = { + "ones": np.ones_like(x_np), + "weighted": rng.uniform(1.0, 9.0, x_np.shape).astype(np.float32), + } + for axis in range(x_np.ndim): + for op, reverse, inclusive in itertools.product( + ("max", "min"), (False, True), (True, False) + ): + for cot_name, cot_np in cotangents.items(): + with self.subTest( + shape=x_np.shape, + axis=axis, + op=op, + reverse=reverse, + inclusive=inclusive, + cotangent=cot_name, + ): + _, (mx_grad,) = mx.vjp( + lambda z: mx_scan(z, axis, reverse, inclusive, op), + (mx.array(x_np),), + (mx.array(cot_np),), + ) + + xt = torch.tensor(x_np, requires_grad=True) + out = torch_scan(xt, axis, reverse, inclusive, op) + (out * torch.tensor(cot_np)).sum().backward() + + self.assertTrue( + np.allclose( + np.array(mx_grad), xt.grad.numpy(), atol=1e-5 + ) + ) + def test_topk_grad(self): a = mx.array([[1, 2, 6, 4, 5], [9, 5, 6, 7, 8]], mx.float32) @@ -509,6 +984,54 @@ def fun(x): expected = mx.array([[0, 0, 1, 0, 1], [1, 0, 0, 0, 1]], mx.float32) self.assertTrue(mx.array_equal(out, expected)) + def test_sort_grad(self): + # Sort permutes the input, so its vjp must scatter the cotangents back + # to the original positions (the transpose of the permutation), not + # gather them forward. A non-involutive permutation exposes the bug. + x = mx.array([3.0, 1.0, 2.0, 5.0, 4.0]) + cotan = mx.array([10.0, 20.0, 30.0, 40.0, 50.0]) + grad = mx.vjp(lambda a: mx.sort(a), (x,), (cotan,))[1][0] + self.assertTrue(mx.array_equal(grad, mx.array([30.0, 10.0, 20.0, 50.0, 40.0]))) + + # vjp must be the transpose of the jvp (adjoint test) along each axis. + mx.random.seed(0) + for axis in (0, 1, -1): + a = mx.random.normal((4, 6)) + v = mx.random.normal(a.shape) + w = mx.random.normal(a.shape) + jv = mx.jvp(lambda z: mx.sort(z, axis=axis), (a,), (v,))[1][0] + jtw = mx.vjp(lambda z: mx.sort(z, axis=axis), (a,), (w,))[1][0] + self.assertAlmostEqual( + mx.sum(w * jv).item(), mx.sum(v * jtw).item(), places=4 + ) + + def test_logsumexp_grad(self): + # The jvp of logsumexp reduces along the axis (sum of softmax * tangent), + # so the tangent it returns must have the reduced output shape, not the + # input shape. + x = mx.array([[1.0, 2.0, 3.0], [4.0, 1.0, 0.0]]) + v = mx.array([[1.0, 0.0, -1.0], [2.0, 1.0, 0.0]]) + jv = mx.jvp(lambda z: mx.logsumexp(z, axis=-1, keepdims=True), (x,), (v,))[1][0] + self.assertEqual(jv.shape, (2, 1)) + expected = mx.sum(mx.softmax(x, axis=-1) * v, axis=-1, keepdims=True) + self.assertTrue(mx.allclose(jv, expected)) + + # vjp must be the transpose of the jvp (adjoint test). + mx.random.seed(0) + for keepdims in (True, False): + a = mx.random.normal((4, 6)) + v = mx.random.normal(a.shape) + + def fun(z): + return mx.logsumexp(z, axis=-1, keepdims=keepdims) + + w = mx.random.normal(fun(a).shape) + jv = mx.jvp(fun, (a,), (v,))[1][0] + jtw = mx.vjp(fun, (a,), (w,))[1][0] + self.assertAlmostEqual( + mx.sum(w * jv).item(), mx.sum(v * jtw).item(), places=4 + ) + def test_custom_function(self): # Make a custom function my_exp = mx.custom_function(mx.exp) @@ -875,6 +1398,160 @@ def test_reduce_jvp(self): out, jout = mx.jvp(mx.max, primals=(a,), tangents=(b,)) self.assertEqual(jout[0].item(), 0) + def test_complex_prod_vjp(self): + def prod(x): + return x.prod(axis=0) + + primal = mx.random.normal((2, 20), dtype=mx.complex64) + cotangent = mx.random.normal((20,), dtype=mx.complex64) + + _, vjps = mx.vjp(prod, [primal], [cotangent]) + + expected = mx.stack( + [mx.conj(primal[1]) * cotangent, mx.conj(primal[0]) * cotangent] + ) + + # Check against hand-computed vjps + self.assertTrue(mx.array_equal(vjps[0], expected)) + + # Ensure that prod agrees with multiply for complex values + _, vjps_multiply = mx.vjp(mx.multiply, [primal[0], primal[1]], [cotangent]) + + self.assertTrue(mx.array_equal(mx.stack(vjps_multiply), vjps[0])) + + def test_complex_exp_vjp(self): + primal = mx.random.normal((3, 4, 5), dtype=mx.complex64) + cotangent = mx.random.normal( + ( + 3, + 4, + 5, + ), + dtype=mx.complex64, + ) + + _, vjps = mx.vjp(mx.exp, [primal], [cotangent]) + + expected = cotangent * mx.conj(mx.exp(primal)) + + # Check against hand-computed vjps + self.assertTrue(mx.allclose(vjps[0], expected)) + + def test_complex_log_vjp(self): + primal = mx.random.normal((3, 4, 5), dtype=mx.complex64) + cotangent = mx.random.normal( + ( + 3, + 4, + 5, + ), + dtype=mx.complex64, + ) + + # guard against values too close to the origin + primal = mx.where(abs(primal) < 1e-3, 1e-3, primal) + + _, vjps = mx.vjp(mx.log, [primal], [cotangent]) + + expected = cotangent * mx.conj(1 / primal) + + # Check against hand-computed vjps + self.assertTrue(mx.allclose(vjps[0], expected)) + + def test_complex_unary_vjps(self): + # For a holomorphic f the vjp is cotangent * conj(f'(z)); these ops used + # to delegate to their jvp and drop the conjugate for complex inputs. + mx.random.seed(0) + z = mx.random.normal((3, 4, 5), dtype=mx.complex64) + cotangent = mx.random.normal((3, 4, 5), dtype=mx.complex64) + z = mx.where(abs(z) < 1e-3, 1e-3 + 0j, z) + + ops = { + mx.square: lambda x: 2 * x, + mx.sin: mx.cos, + mx.sinh: mx.cosh, + mx.cosh: mx.sinh, + mx.tan: lambda x: 1 / mx.cos(x) ** 2, + mx.tanh: lambda x: 1 - mx.tanh(x) ** 2, + mx.log1p: lambda x: 1 / (1 + x), + } + for fn, deriv in ops.items(): + _, (vjp,) = mx.vjp(fn, [z], [cotangent]) + expected = cotangent * mx.conj(deriv(z)) + self.assertTrue(mx.allclose(vjp, expected, atol=1e-5), msg=str(fn)) + + def test_complex_abs_grad(self): + mx.random.seed(0) + primal = mx.random.normal((3, 4, 5), dtype=mx.complex64) + # guard against values too close to the origin where |z| is not smooth + primal = mx.where(abs(primal) < 1e-3, 1e-3 + 0j, primal) + + # |z| is real-valued, so its jvp is real: + # d|z| = Re(conj(z) * t) / |z| + tangent = mx.random.normal(primal.shape, dtype=mx.complex64) + _, (jvp,) = mx.jvp(mx.abs, [primal], [tangent]) + expected = mx.real(mx.conj(primal) * tangent) / mx.abs(primal) + self.assertEqual(jvp.dtype, mx.float32) + self.assertTrue(mx.allclose(jvp, expected, atol=1e-5)) + + # The vjp's real and imaginary parts are the gradients w.r.t. Re(z) and + # Im(z); for a real cotangent this is cotangent * sign(z). + cotangent = mx.random.normal(primal.shape) + _, (vjp,) = mx.vjp(mx.abs, [primal], [cotangent]) + self.assertTrue( + mx.allclose(vjp, cotangent * (primal / mx.abs(primal)), atol=1e-5) + ) + + # Real inputs are unaffected. + x = mx.random.normal((10,)) + t = mx.random.normal((10,)) + _, (jvp,) = mx.jvp(mx.abs, [x], [t]) + self.assertTrue(mx.allclose(jvp, mx.sign(x) * t)) + + def test_second_order_permutation_ops(self): + # The permutation these ops apply is locally constant in the input, so + # the indices must not carry a gradient. Otherwise differentiating the + # vjp a second time fails with "Cannot calculate VJP with respect to + # indices". + def hvp(f, x, v): + return mx.grad(lambda a: mx.sum(mx.grad(f)(a) * v))(x) + + def numerical_hvp(f, x, v, eps=1e-3): + # The values below are well separated, so the permutation does not + # change over this step and the difference is exact enough. + return (mx.grad(f)(x + eps * v) - mx.grad(f)(x - eps * v)) / (2 * eps) + + x = mx.array([3.0, 1.0, 2.0, 5.0]) + v = mx.array([1.0, -2.0, 0.5, 1.5]) + + for fn in ( + lambda a: mx.sort(a), + lambda a: mx.partition(a, 2), + lambda a: mx.topk(a, 2), + lambda a: mx.cummax(a, axis=0), + lambda a: mx.cummin(a, axis=0), + lambda a: mx.cummax(a, axis=0, reverse=True), + lambda a: mx.cummax(a, axis=0, inclusive=False), + lambda a: mx.cummin(a, axis=0, reverse=True, inclusive=False), + ): + f = lambda a: mx.sum(fn(a) ** 2) + self.assertTrue( + mx.allclose(hvp(f, x, v), numerical_hvp(f, x, v), atol=1e-3) + ) + + # A non-trailing axis + y = mx.array([[3.0, 1.0], [2.0, 5.0]]) + w = mx.array([[1.0, -2.0], [0.5, 1.5]]) + for fn in ( + lambda a: mx.sort(a, axis=0), + lambda a: mx.partition(a, 1, axis=0), + lambda a: mx.cummax(a, axis=0), + ): + f = lambda a: mx.sum(fn(a) ** 2) + self.assertTrue( + mx.allclose(hvp(f, y, w), numerical_hvp(f, y, w), atol=1e-3) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_bf16.py b/Source/Cmlx/mlx/python/tests/test_bf16.py index 2e4e2e0c3..d25364c14 100644 --- a/Source/Cmlx/mlx/python/tests/test_bf16.py +++ b/Source/Cmlx/mlx/python/tests/test_bf16.py @@ -15,6 +15,13 @@ except ImportError as e: has_torch = False +try: + import ml_dtypes + + has_ml_dtypes = True +except ImportError: + has_ml_dtypes = False + class TestBF16(mlx_tests.MLXTestCase): def __test_ops( @@ -191,6 +198,31 @@ def test_conversion(self): self.assertEqual(a_mx.dtype, mx.bfloat16) self.assertTrue(mx.array_equal(a_mx, expected)) + @unittest.skipIf(not has_ml_dtypes, "requires ml_dtypes") + def test_conversion_ml_dtypes(self): + x_scalar = np.array(1.5, dtype=ml_dtypes.bfloat16) + a_scalar = mx.array(x_scalar) + self.assertEqual(a_scalar.dtype, mx.bfloat16) + self.assertEqual(a_scalar.shape, ()) + self.assertEqual(a_scalar.item(), 1.5) + + data = [1.5, 2.5, 3.5] + x_vector = np.array(data, dtype=ml_dtypes.bfloat16) + a_vector = mx.array(x_vector) + expected = mx.array(data, dtype=mx.bfloat16) + self.assertEqual(a_vector.dtype, mx.bfloat16) + self.assertEqual(a_vector.shape, (3,)) + self.assertTrue(mx.array_equal(a_vector, expected)) + + a_cast = mx.array(x_scalar, dtype=mx.float32) + self.assertEqual(a_cast.dtype, mx.float32) + self.assertEqual(a_cast.item(), 1.5) + + a_asarray = mx.asarray(x_vector) + self.assertEqual(a_asarray.dtype, mx.bfloat16) + self.assertEqual(a_asarray.shape, (3,)) + self.assertTrue(mx.array_equal(a_asarray, expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_blas.py b/Source/Cmlx/mlx/python/tests/test_blas.py index dedfa5d4f..7dbbaea04 100644 --- a/Source/Cmlx/mlx/python/tests/test_blas.py +++ b/Source/Cmlx/mlx/python/tests/test_blas.py @@ -8,6 +8,9 @@ import mlx_tests import numpy as np +# Ignore matmul warnings. +np.seterr(divide="ignore", over="ignore", invalid="ignore") + class TestBlas(mlx_tests.MLXTestCase): @property @@ -535,6 +538,125 @@ def test_matrix_vector_edgecases(self): ) self.assertTrue(np.array_equal(c_mlx, c_npy)) + def test_dot_product(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + def run_test(dtype, size, offset, atol): + with self.subTest(dtype=str(dtype), size=size, offset=offset): + np.random.seed(42) + scale = size**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, size + offset).astype(np.float32) + ).astype(dtype)[offset:] + b_mx = mx.array( + np.random.normal(0.0, scale, size + offset).astype(np.float32) + ).astype(dtype)[offset:] + + expected = np.inner( + np.array(a_mx.astype(mx.float32)), + np.array(b_mx.astype(mx.float32)), + ) + actual = np.array(mx.inner(a_mx, b_mx).astype(mx.float32)) + self.assertTrue(np.allclose(actual, expected, atol=atol)) + + for dtype, atol in ( + (mx.float32, 1e-5), + (mx.float16, 2e-3), + (mx.bfloat16, 2e-3), + ): + for size in (1023, 1024, 1025, 16385, 131072, 1000000): + for offset in (0, 1): + run_test(dtype, size, offset, atol) + + def test_wide_matmul(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible a @ b.T products of a few rows take the wide gemv route on metal; + # cover numerical correctness across the routing boundaries: row tails, + # K not divisible by 4, offset and sliced views, and batching. + # Inputs scale as K**-0.5 so outputs stay O(K**-0.5); atol scales + # with them, above output rounding for every dtype but below a + # dropped reduction block. + def run_test(dtype, shape_a, shape_b, f_np_b, f_mx_b): + with self.subTest(dtype=str(dtype), shape_a=shape_a, shape_b=shape_b): + np.random.seed(7) + scale = shape_a[-1] ** -0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, shape_a).astype(np.float32) + ).astype(dtype) + b_mx = mx.array( + np.random.normal(0.0, scale, shape_b).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + b_np = np.array(b_mx.astype(mx.float32)) + + out_np = a_np @ f_np_b(b_np) + out_mx = (a_mx @ f_mx_b(b_mx)).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + + nt_np = lambda b: b.swapaxes(-1, -2) + nt_mx = lambda b: mx.swapaxes(b, -1, -2) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (1, 2, 3, 5, 8, 11, 16): + for K, N in ( + (64, 128), + (512, 128), + (2048, 256), + (2048, 32), + (2052, 1000), + ): + run_test(dtype, (M, K), (N, K), nt_np, nt_mx) + + # K % 4 != 0 falls back to the general kernels + run_test(dtype, (5, 514), (333, 514), nt_np, nt_mx) + + # sliced weights: leading dimension != K, plus offset views at 16- + # and 8-byte alignment, and a slice with an odd vec4 tail + run_test( + dtype, + (5, 512), + (333, 576), + lambda b: b[:, :512].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, :512], -1, -2), + ) + run_test( + dtype, + (5, 512), + (333, 512), + lambda b: b[7:, :].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[7:, :], -1, -2), + ) + run_test( + dtype, + (3, 512), + (333, 520), + lambda b: b[:, 4:516].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, 4:516], -1, -2), + ) + run_test( + dtype, + (3, 2052), + (129, 2056), + lambda b: b[:, :2052].swapaxes(-1, -2), + lambda b: mx.swapaxes(b[:, :2052], -1, -2), + ) + + # batched: regular, broadcast weights, and multi-dim batch + run_test(dtype, (4, 3, 512), (4, 257, 512), nt_np, nt_mx) + run_test( + dtype, + (4, 3, 512), + (1, 257, 512), + lambda b: np.broadcast_to(b, (4, 257, 512)).swapaxes(-1, -2), + lambda b: mx.swapaxes(mx.broadcast_to(b, (4, 257, 512)), -1, -2), + ) + run_test(dtype, (2, 3, 5, 512), (2, 3, 129, 512), nt_np, nt_mx) + def test_mismatch_stride_mm(self): np.random.seed(0) a_npy = np.random.normal(0.0, 1.0 / 128, (4, 16, 16)).astype(np.float32) @@ -730,6 +852,61 @@ def test_addmm(self): expected = 0.5 * (a @ b) + 2.0 * c self.assertTrue(mx.allclose(out, expected, rtol=tol, atol=tol)) + def test_wide_addmm(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible few-row addmm shapes take the wide gemv route on metal; cover + # the axpby epilogue against bias shapes, scales, and batching. + def run_test(dtype, B, M, K, N, c_shape, alpha, beta): + with self.subTest(dtype=str(dtype), c_shape=c_shape, alpha=alpha): + np.random.seed(3) + shape_a = (M, K) if B is None else (B, M, K) + shape_b = (N, K) if B is None else (B, N, K) + scale = K**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, shape_a).astype(np.float32) + ).astype(dtype) + b_mx = mx.array( + np.random.normal(0.0, scale, shape_b).astype(np.float32) + ).astype(dtype) + c_mx = mx.array( + np.random.normal(0.0, scale, c_shape).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + b_np = np.array(b_mx.astype(mx.float32)) + c_np = np.array(c_mx.astype(mx.float32)) + + out_np = alpha * (a_np @ b_np.swapaxes(-1, -2)) + beta * c_np + out_mx = mx.addmm( + c_mx, + a_mx, + mx.swapaxes(b_mx, -1, -2), + alpha, + beta, + ).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + atol = 0.05 * (abs(alpha) + abs(beta)) * scale + self.assertTrue(np.allclose(out_mx, out_np, atol=atol)) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (2, 5, 12): + for c_shape in ((250,), (1, 250), (M, 250)): + for alpha, beta in ((1.0, 1.0), (2.5, 0.5), (1.0, 0.0)): + run_test(dtype, None, M, 512, 250, c_shape, alpha, beta) + run_test(dtype, 3, 4, 512, 250, (3, 4, 250), 1.0, 1.0) + + # the epilogue must scale the accumulator before narrowing: a + # product past the fp16 max rescued by alpha stays finite (M = 4 + # keeps the 2-byte shape routed on every supported generation) + with self.subTest(dtype=str(dtype), case="alpha rescue"): + a = mx.full((4, 512), 2.0, dtype=dtype) + b = mx.full((250, 512), 100.0, dtype=dtype) + c = mx.ones((4, 250), dtype=dtype) + out = mx.addmm(c, a, mx.swapaxes(b, -1, -2), 0.125, 0.0) + self.assertTrue(np.allclose(out.astype(mx.float32), 12800.0)) + def test_addmm_grad(self): def make_ref_addmm(alpha, beta): return lambda c, a, b: alpha * (a @ b) + beta * c @@ -1190,6 +1367,67 @@ def test_shape( self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5)) + def test_gather_mm_blocks(self): + if mx.default_device() == mx.cpu: + self.skipTest("requires GPU") + + # Eligible gathered products with few-row blocks route to the wide + # gather on metal; check block indexing against a per-entry reference. + def run_test(dtype, G, M, K, N, E, idx_shape): + with self.subTest(dtype=str(dtype), G=G, M=M, idx=idx_shape): + np.random.seed(11) + scale = K**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, (G, M, K)).astype(np.float32) + ).astype(dtype) + w_mx = mx.array( + np.random.normal(0.0, scale, (E, N, K)).astype(np.float32) + ).astype(dtype) + a_np = np.array(a_mx.astype(mx.float32)) + w_np = np.array(w_mx.astype(mx.float32)) + rhs = np.random.randint(0, E, size=idx_shape).astype(np.uint32) + + out_np = np.stack( + [a_np[i % G] @ w_np[r].T for i, r in enumerate(rhs.reshape(-1))] + ).reshape(*idx_shape, M, N) + out_mx = mx.gather_mm( + a_mx.reshape(*idx_shape, M, K), + mx.swapaxes(w_mx, -1, -2), + None, + mx.array(rhs), + ).astype(mx.float32) + + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + for M in (2, 4, 5, 11): + run_test(dtype, 6, M, 2048, 1024, 8, (6,)) + run_test(dtype, 6, 3, 2048, 1024, 8, (2, 3)) + + # scalar indices: batch_ndim == 0 must still bind index strides + with self.subTest(dtype=str(dtype), idx="scalar"): + np.random.seed(11) + scale = 512**-0.5 + a_mx = mx.array( + np.random.normal(0.0, scale, (4, 512)).astype(np.float32) + ).astype(dtype) + w_mx = mx.array( + np.random.normal(0.0, scale, (8, 129, 512)).astype(np.float32) + ).astype(dtype) + out_np = ( + np.array(a_mx.astype(mx.float32)) + @ np.array(w_mx[3].astype(mx.float32)).T + ) + out_mx = mx.gather_mm( + a_mx, + mx.swapaxes(w_mx, -1, -2), + None, + mx.array(3, dtype=mx.uint32), + ).astype(mx.float32) + self.assertListEqual(list(out_np.shape), list(out_mx.shape)) + self.assertTrue(np.allclose(out_mx, out_np, atol=0.05 * scale)) + def test_gather_matmul_grad(self): lhs_indices = mx.array([[7, 6], [4, 1], [0, 2]], dtype=mx.uint32) rhs_indices = mx.array([[2], [0], [1]], dtype=mx.uint32) @@ -1236,6 +1474,24 @@ def f_test(a, b): self.assertEqual(r.shape, t.shape) self.assertTrue(mx.allclose(r, t, atol=1e-4).item()) + def test_gather_matmul_index_vjp_requires_stop_gradient(self): + a = mx.ones((4, 1, 2, 2)) + b = mx.ones((4, 1, 2, 2)) + + def fun(w): + indices = mx.reshape(mx.argsort(w)[:2], (1, 2)) + return mx.gather_mm(a, b, indices, indices).sum() + + with self.assertRaisesRegex(ValueError, "stop_gradient"): + mx.grad(fun)(mx.array([3.0, 1.0, 2.0, 0.0])) + + def fun_stopped(w): + indices = mx.stop_gradient(mx.reshape(mx.argsort(w)[:2], (1, 2))) + return mx.gather_mm(a, b, indices, indices).sum() + + grad = mx.grad(fun_stopped)(mx.array([3.0, 1.0, 2.0, 0.0])) + self.assertTrue(mx.array_equal(grad, mx.zeros((4,)))) + def test_gather_mm_sorted(self): def gather_mm_ref(a, b, rhs): b = b[rhs] diff --git a/Source/Cmlx/mlx/python/tests/test_compile.py b/Source/Cmlx/mlx/python/tests/test_compile.py index d64c057fd..663ec295b 100644 --- a/Source/Cmlx/mlx/python/tests/test_compile.py +++ b/Source/Cmlx/mlx/python/tests/test_compile.py @@ -4,6 +4,7 @@ import inspect import io import math +import threading from functools import partial, wraps from io import StringIO @@ -44,6 +45,104 @@ def fun(x, y): self.assertEqual(out.dtype, mx.int32) self.assertTrue(mx.array_equal(out, mx.array([2, 4]))) + def test_compile_nonfinite_constants(self): + # Regression test: a non-finite scalar constant (NaN / infinity) baked + # into a fused compiled kernel used to stream a bare token (e.g. `nan`) + # into the generated kernel source, which is not a valid identifier and + # broke compilation (notably on the Metal backend). + x = mx.array([1.0, -1.0]) + + for dtype in (mx.float32, mx.float16, mx.bfloat16): + xd = x.astype(dtype) + + nan_fn = mx.compile( + lambda a: mx.where(a > 0, a, mx.array(float("nan"), dtype=a.dtype)) + ) + out = nan_fn(xd) + mx.eval(out) + self.assertEqual(out[0].item(), 1.0) + self.assertTrue(math.isnan(out[1].item())) + + neg_inf_fn = mx.compile( + lambda a: mx.where(a > 0, a, mx.array(float("-inf"), dtype=a.dtype)) + ) + out = neg_inf_fn(xd) + mx.eval(out) + self.assertEqual(out[0].item(), 1.0) + self.assertEqual(out[1].item(), float("-inf")) + + def test_compile_tuple_output_in_thread(self): + @mx.compile + def fun(x): + return x + 1, x * 2 + + results = [] + errors = [] + + def worker(): + try: + x = mx.array([1.0]) + y, z = fun(x) + mx.eval(y, z) + results.append((y.item(), z.item())) + except Exception as e: + errors.append(e) + mx.clear_streams() + + for _ in range(3): + thread = threading.Thread(target=worker) + thread.start() + thread.join() + gc.collect() + + if errors: + raise errors[0] + self.assertEqual(results, [(2.0, 2.0)] * 3) + + def test_compile_release_on_another_thread(self): + # A function traced on one thread but released on another must still + # drop its cache entry, otherwise a later compile of the same id gets + # handed the dead function's tape instead of being traced again. + traces = [] + + def fun(x): + traces.append(1) + return x + 1 + + holder = {} + traced = threading.Event() + released = threading.Event() + errors = [] + + def worker(): + try: + holder["fn"] = mx.compile(fun) + mx.eval(holder["fn"](mx.array([1.0]))) + traced.set() + self.assertTrue(released.wait(10)) + # The same callable, so the same id. + fn = mx.compile(fun) + mx.eval(fn(mx.array([1.0]))) + except Exception as e: + errors.append(e) + finally: + traced.set() + mx.clear_streams() + + # The tracing thread has to outlive the release, on exit it would tear + # down its cache anyway. + thread = threading.Thread(target=worker) + thread.start() + self.assertTrue(traced.wait(10)) + holder.clear() + gc.collect() + released.set() + thread.join() + + if errors: + raise errors[0] + self.assertEqual(len(traces), 2) + def test_compile_grad(self): def loss_fn(x): return mx.exp(x).sum() @@ -385,6 +484,76 @@ def fun(): self.assertFalse(mx.allclose(fun(), fun(), 1e-2, 1e-2)) + def test_compile_rng_across_threads(self): + # A function compiled with inputs/outputs=mx.random.state on one thread + # must still use (and advance/seed) the calling thread's RNG state when + # invoked from another thread, whether captured directly or nested. + + # The state sentinel is a single global object shared across threads. + state_from_thread = {} + + def grab(): + state_from_thread["s"] = mx.random.state + mx.clear_streams() + + t = threading.Thread(target=grab) + t.start() + t.join() + self.assertIs(mx.random.state, state_from_thread["s"]) + + direct = partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)( + lambda: mx.random.uniform(shape=(10, 10)) + ) + + nested_state = [{"unused": mx.array(0.0)}, mx.random.state] + nested = partial(mx.compile, inputs=nested_state, outputs=nested_state)( + lambda: mx.random.uniform(shape=(10, 10)) + ) + + for fun in (direct, nested): + results = {} + + def worker(): + with mx.stream(mx.cpu): + a = fun() + b = fun() + results["advances"] = not bool(mx.allclose(a, b, 1e-2, 1e-2).item()) + mx.random.seed(42) + c = fun() + mx.random.seed(42) + d = fun() + results["seed_reproducible"] = bool(mx.allclose(c, d).item()) + mx.random.seed(1234) + e = fun() + results["seed_changes"] = not bool( + mx.allclose(c, e, 1e-2, 1e-2).item() + ) + mx.clear_streams() + + t = threading.Thread(target=worker) + t.start() + t.join() + + self.assertTrue(results["advances"]) + self.assertTrue(results["seed_reproducible"]) + self.assertTrue(results["seed_changes"]) + + def test_compile_state_capture_with_rng_updates_in_place(self): + # Capturing mx.random.state alongside other state via outputs= must not + # break in-place updates of the other captured containers. + counter = {"v": mx.array(0.0)} + state = [counter, mx.random.state] + + @partial(mx.compile, inputs=state, outputs=state) + def step(): + counter["v"] = counter["v"] + 1.0 + return mx.random.uniform(shape=(2,)) + + for _ in range(3): + step() + mx.eval(counter["v"]) + self.assertEqual(counter["v"].item(), 3.0) + def test_compile_kwargs(self): @mx.compile def fun(x, y, z): @@ -504,6 +673,45 @@ def ones_fun(x): self.assertEqual(compiled_zero_like(y).shape, y_shape) self.assertEqual(compiled_ones_like(y).shape, y_shape) + def test_shapeless_compile_gather_qmm(self): + K, N, num_experts = 64, 32, 4 + + w = mx.random.normal((num_experts, N, K)) + qw, s, b = mx.quantize(w) + mx.eval(qw, s, b) + + idx = mx.array([0, 1, 2, 3]) + x4 = mx.ones((num_experts, 4, K)) + x8 = mx.ones((num_experts, 8, K)) + + def fn(x): + return mx.gather_qmm( + x, qw, s, b, lhs_indices=idx, rhs_indices=idx, transpose=True + ) + + cfn = mx.compile(fn, shapeless=True) + + self.assertEqual(cfn(x4).shape, fn(x4).shape) + self.assertEqual(cfn(x8).shape, fn(x8).shape) + + def test_shapeless_compile_gather_mm(self): + K, N, num_experts = 64, 32, 4 + + idx = mx.array([0, 1, 2, 3]) + b = mx.random.normal((num_experts, K, N)) + mx.eval(b) + + x4 = mx.ones((num_experts, 4, K)) + x8 = mx.ones((num_experts, 8, K)) + + def fn(x): + return mx.gather_mm(x, b, lhs_indices=idx, rhs_indices=idx) + + cfn = mx.compile(fn, shapeless=True) + + self.assertEqual(cfn(x4).shape, fn(x4).shape) + self.assertEqual(cfn(x8).shape, fn(x8).shape) + def test_compile_with_constant(self): # Test float @partial(mx.compile) @@ -884,6 +1092,34 @@ def test_shapeless_compile_matmul(self): fun = mx.compile(lambda a, b: a @ b, shapeless=True) self.assertTrue(mx.allclose(fun(a, b), a @ b)) + def test_shapeless_compile_addmm(self): + def fun(c, a, b): + return mx.addmm(c, a, b) + + cfun = mx.compile(fun, shapeless=True) + + # First shape + c = mx.ones((2, 4)) + a = mx.ones((2, 3)) + b = mx.ones((3, 4)) + self.assertTrue(mx.allclose(cfun(c, a, b), fun(c, a, b))) + + # Different shape, same ranks — should not recompile + c = mx.ones((3, 5)) + a = mx.ones((3, 6)) + b = mx.ones((6, 5)) + self.assertTrue(mx.allclose(cfun(c, a, b), fun(c, a, b))) + + # With alpha and beta + fun2 = mx.compile( + lambda c, a, b: mx.addmm(c, a, b, alpha=2.0, beta=3.0), shapeless=True + ) + c = mx.ones((2, 4)) + a = mx.ones((2, 3)) + b = mx.ones((3, 4)) + expected = 3.0 * c + 2.0 * (a @ b) + self.assertTrue(mx.allclose(fun2(c, a, b), expected)) + def test_shapeless_compile_slice_update(self): def fun(x): x[2] = mx.array([3.0]) @@ -989,6 +1225,45 @@ def fun(args): self.assertEqual(out[0].shape, (3, 1, 4, 2)) self.assertEqual(out[1].shape, (2, 2, 5)) + def test_shapeless_compile_reduce_after_gather(self): + # Reductions over dimensions that happen to have size 1 at trace time + # used to be elided from the graph, so replays with larger dynamic + # shapes returned stale values (issue #3201) + buf = mx.array([10.0, 20.0, 30.0, 40.0, 50.0]) + reductions = [ + mx.sum, + mx.mean, + mx.prod, + mx.min, + mx.max, + mx.all, + mx.any, + mx.argmin, + mx.argmax, + ] + for reduction in reductions: + + def fun(buf, idx): + return reduction(mx.take(buf, idx, axis=0)) + + # Trace with a size-1 reduction and replay with larger sizes + cfun = mx.compile(fun, shapeless=True) + for n in [1, 2, 3, 4]: + idx = mx.arange(n) + self.assertTrue( + mx.array_equal(cfun(buf, idx), fun(buf, idx)), + f"{reduction.__name__} failed for n={n}", + ) + + # Replay with size 1 so the reduction is an identity at runtime + cfun = mx.compile(fun, shapeless=True) + for n in [2, 1]: + idx = mx.arange(n) + self.assertTrue( + mx.array_equal(cfun(buf, idx), fun(buf, idx)), + f"{reduction.__name__} failed for replay with n={n}", + ) + def test_leaks(self): gc.collect() if mx.metal.is_available(): @@ -1049,6 +1324,28 @@ def fun(x, y, z): self.assertTrue(mx.allclose(d[0], d_hat[0])) self.assertTrue(mx.allclose(d[1], d_hat[1])) + def test_compile_large_graph_with_broadcasts(self): + N = 20 + _as = [mx.array(2 * i, dtype=mx.float32) for i in range(N)] + _bs = [mx.array(i, dtype=mx.float32) for i in range(N)] + _c = mx.array(0.0) + x = mx.random.normal((2, 2)) + + def f(x): + y = 0 + for i in range(N): + y = y + _as[i] * x * _bs[i] * _c + return y + + ref = f(x) + mx.eval(ref) + f = mx.compile(f) + for i in range(2): + y = f(x) + mx.eval(y) + + self.assertTrue(mx.allclose(y, ref)) + def test_wrap_compiled(self): @mx.compile def inner(): @@ -1272,6 +1569,93 @@ def fun(x): np.asarray(out, copy=False).__array_interface__["data"][0], in_ptr ) + def test_compile_negative_strides(self): + # 1D negative stride with elementwise expression + @mx.compile + def f(x): + return 2.0 * x[::-1] + + x = mx.arange(8, dtype=mx.float32) + expected = 2.0 * x[::-1] + self.assertTrue(mx.array_equal(f(x), expected)) + + # 1D negative stride with slice update + def g_eager(x): + base = mx.zeros_like(x) + base[::-1] += 2.0 * x[::-1] + return base + + g_compiled = mx.compile(g_eager) + expected = g_eager(x) + self.assertTrue(mx.array_equal(g_compiled(x), expected)) + + # 2D negative stride + @mx.compile + def h(x): + return x[::-1] + 1.0 + + y = mx.arange(12, dtype=mx.float32).reshape(3, 4) + expected = y[::-1] + 1.0 + self.assertTrue(mx.array_equal(h(y), expected)) + + # Mixed positive and negative strides + @mx.compile + def m(x): + return x[::-1, ::2] * 3.0 + + z = mx.arange(24, dtype=mx.float32).reshape(4, 6) + expected = z[::-1, ::2] * 3.0 + self.assertTrue(mx.array_equal(m(z), expected)) + + # 4D negative stride (exercises work_per_thread > 1 path) + @mx.compile + def p(x): + return x + 1.0 + + w = mx.arange(120, dtype=mx.float32).reshape(2, 3, 4, 5) + expected = w[::-1, :, ::-1, :] + 1.0 + self.assertTrue(mx.array_equal(p(w[::-1, :, ::-1, :]), expected)) + + def test_compile_abs_unsigned(self): + # abs has to compile for the wider unsigned types too + fun = lambda x: mx.abs(x) + 1 + for dtype in [mx.uint8, mx.uint16, mx.uint32, mx.uint64]: + x = mx.array([1, 2, 3], dtype) + self.assertTrue(mx.array_equal(mx.compile(fun)(x), fun(x))) + + def test_compiled_subnormal_bool_cast(self): + f32_sub = mx.array(np.array([0x00000001] * 4, dtype=np.uint32)).view(mx.float32) + f16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.float16) + bf16_sub = mx.array(np.array([0x0001] * 4, dtype=np.uint16)).view(mx.bfloat16) + + # A single-op compile does not fuse; the fused path needs >= 2 ops. + fn = mx.compile(lambda x: mx.broadcast_to(x, (2, 4)).astype(mx.bool_)) + for sub in (f32_sub, f16_sub, bf16_sub): + self.assertTrue(mx.all(fn(sub)).item()) + + def test_compile_different_log_bases(self): + # The logs are intermediates, since outputs are not simplified. + def entropies(p): + nats = -mx.sum(p * mx.log(p)) + bits = -mx.sum(p * mx.log2(p)) + return mx.stack([nats, bits]) + + p = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32) + expected = np.array( + [-(p * np.log(p)).sum(), -(p * np.log2(p)).sum()], dtype=np.float32 + ) + out = mx.compile(entropies)(mx.array(p)) + self.assertTrue(np.allclose(out, expected, atol=1e-5)) + + def test_compile_equal_nan(self): + def fun(x): + return mx.stack( + [mx.array_equal(x, x), mx.array_equal(x, x, equal_nan=True)] + ) + + x = mx.array([1.0, float("nan"), 3.0]) + self.assertTrue(mx.array_equal(mx.compile(fun)(x), mx.array([False, True]))) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_conv.py b/Source/Cmlx/mlx/python/tests/test_conv.py index cef912aee..062841c33 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv.py +++ b/Source/Cmlx/mlx/python/tests/test_conv.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest from itertools import permutations @@ -13,7 +14,7 @@ import torch.nn.functional as F has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -308,9 +309,11 @@ def run_conv2D( lambda x: mx.array(x).astype(mx_dtype), (in_np, wt_np) ) in_pt, wt_pt = map( - lambda x: torch.from_numpy(x.transpose(0, 3, 1, 2)) - .to("cpu") - .to(torch_dtype), + lambda x: ( + torch.from_numpy(x.transpose(0, 3, 1, 2)) + .to("cpu") + .to(torch_dtype) + ), (in_np, wt_np), ) @@ -410,7 +413,9 @@ def run_conv2D_grad( ) in_np = np.random.normal(0.0, scale, (N, iH, iW, C)).astype(np_dtype) - wt_np = np.random.normal(0.0, scale, (O, kH, kW, C)).astype(np_dtype) + wt_np = np.random.normal(0.0, scale, (O, kH, kW, C // groups)).astype( + np_dtype + ) ct_np = np.random.normal(0.0, scale, (N, oH, oW, O)).astype(np_dtype) in_mx, wt_mx, ct_mx = map(mx.array, (in_np, wt_np, ct_np)) @@ -479,6 +484,19 @@ def f(a, b): N, C, O, idim, kdim, stride, padding, dilation, dtype=dtype ) + for groups in (2, 4, 32): + run_conv2D_grad( + 4, + 32, + 32, + (32, 32), + (3, 3), + (2, 2), + (1, 1), + groups=groups, + dtype=dtype, + ) + @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_3D(self): def run_conv3D( @@ -550,7 +568,6 @@ def run_conv3D( (1, 1, 6), (4, 16, 32), ): - continue for idim, kdim, stride, padding in ( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), @@ -1054,7 +1071,6 @@ def test_repeated_conv(self): @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_depthwise(self): - # fmt: off shapes = ( # N, H, W, C kH, kW, O, strides, padding, groups @@ -1085,7 +1101,7 @@ def test_torch_conv_depthwise(self): groups=groups, flip=flip, np_dtype=dtype, - atol=2e-5 if dtype == np.float32 else 5e-4, + atol=2e-5 if dtype == np.float32 else 5e-3, ) @unittest.skipIf(not has_torch, "requires Torch") @@ -1152,7 +1168,8 @@ def loss_fn(kernel, inputs, strides, groups): ) self.assertEqual(grads.shape, k_shape) - def test_1d_conv_with_2d(self): + @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "flaky in CI") + def test_conv_1d_with_2d(self): x = mx.random.uniform(shape=(2, 10, 16)) y = mx.random.normal(shape=(16, 3, 16)) @@ -1186,6 +1203,110 @@ def test_conv2d_unaligned_channels(self): y_hat = mx.conv2d(x, w) self.assertTrue(mx.allclose(y, y_hat)) + x = mx.random.uniform(shape=(2, 16, 16, 24)) + w = mx.random.uniform(shape=(32, 5, 5, 24)) + y = mx.conv2d(x, w, padding=2, stream=mx.cpu) + y_hat = mx.conv2d(x, w, padding=2) + self.assertTrue(mx.allclose(y, y_hat)) + + x = mx.random.uniform(shape=(2, 16, 16, 24)) + w = mx.random.uniform(shape=(32, 3, 3, 24)) + y = mx.conv_transpose2d(x, w, stream=mx.cpu) + y_hat = mx.conv_transpose2d(x, w) + self.assertTrue(mx.allclose(y, y_hat)) + + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv2d_winograd_batch_tiling(self): + # Use envs to test tiling without allocating large buffers. + tile_key = "MLX_CONV_WINOGRAD_TILE_BATCH" + ws_key = "MLX_CONV_WINOGRAD_WORKING_SET" + prev = {k: os.environ.get(k) for k in (tile_key, ws_key)} + + # Winograd needs 3x3 stride-1, channels in multiples of 32, + # C + O >= 256 and N * iH * iW >= 4096. + cases = ( + ((8, 48, 48, 64), (192, 3, 3, 64)), + ((5, 52, 44, 128), (128, 3, 3, 128)), + ((4, 48, 48, 192), (96, 3, 3, 192)), + ) + + def run(x, w, env={}): + for k in (tile_key, ws_key): + os.environ.pop(k, None) + os.environ.update(env) + y = mx.conv2d(x, w, padding=1) + mx.eval(y) + return np.array(y) + + try: + for in_shape, wt_shape in cases: + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + # Small weights keep the output near unit scale. + w = mx.array( + (np.random.normal(size=wt_shape) * 0.05).astype(np.float32) + ) + b = mx.zeros((wt_shape[0],)) + mx.eval(x, w, b) + cpu_ref = np.array(mx.conv2d(x, w, padding=1, stream=mx.cpu)) + + untiled = run(x, w) + self.assertGreater(np.abs(untiled).max(), 0) + self.assertTrue(np.allclose(untiled, cpu_ref, atol=1e-3)) + + # Tiled winograd keeps the same per-element reduction order, + # so it is bit-identical to untiled; the implicit gemm + # fallback never is. Exact equality pins each run to its path. + # 3 divides none of the batches, so it also covers a short + # final tile. + for tile in (1, 3): + with self.subTest(in_shape=in_shape, tile=tile): + tiled = run(x, w, {tile_key: str(tile)}) + self.assertTrue(np.array_equal(untiled, tiled)) + + # A consumer op checks the output is fenced across + # command encoders. + os.environ[tile_key] = str(tile) + fused = mx.conv2d(x, w, padding=1) + b + mx.eval(fused) + self.assertTrue(np.allclose(untiled, fused, atol=1e-4)) + os.environ.pop(tile_key, None) + + # Budget for ~2 batch elements so the selector itself must + # tile; mirrors the winograd_batch_step arithmetic. + n, iH, iW, C = in_shape + O = wt_shape[0] + pH = 6 * ((iH + 2 - 2 + 5) // 6) + 2 + pW = 6 * ((iW + 2 - 2 + 5) // 6) + 2 + per_n = ( + pH * pW * C * 4 + + 64 * ((iH + 5) // 6) * ((iW + 5) // 6) * (C + O) * 4 + ) + used = (n * iH * iW * (C + O) + 64 * C * O) * 4 + with self.subTest(in_shape=in_shape, budget="tiled"): + budget = str(int((used + 5 * per_n // 2) / 0.75)) + tiled = run(x, w, {ws_key: budget}) + self.assertTrue(np.array_equal(untiled, tiled)) + + # Too small for even one batch element: must fall back. + with self.subTest(in_shape=in_shape, budget="infeasible"): + fallback = run(x, w, {ws_key: "1"}) + self.assertFalse(np.array_equal(untiled, fallback)) + self.assertTrue(np.allclose(fallback, cpu_ref, atol=1e-3)) + + # A forced tile is capped by the budget, so this must still + # fall back. + with self.subTest(in_shape=in_shape, budget="forced+infeasible"): + capped = run(x, w, {ws_key: "1", tile_key: "1"}) + self.assertFalse(np.array_equal(untiled, capped)) + self.assertTrue(np.allclose(capped, cpu_ref, atol=1e-3)) + finally: + for k, v in prev.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + def test_conv2d_large_filter_small_channels(self): x = mx.random.normal(shape=(1, 181, 181, 1)) w = mx.random.normal(shape=(1, 182, 182, 1)) @@ -1193,6 +1314,32 @@ def test_conv2d_large_filter_small_channels(self): y_hat = mx.conv2d(x, w, (1, 1), (1, 1)) self.assertTrue(mx.allclose(y, y_hat, rtol=1e-3, atol=1e-3)) + def test_conv_3D_small_kd_decomposition(self): + # Exercises the small kernel-depth 3D -> KD x 2D decomposition (#3625): + # N=1, small KD, depth stride/dilation 1, no depth padding, mod16 channels. + # Validated against the CPU reference, which uses a different code path. + for T, H, W, Cin, Cout, kd, kh, kw in [ + (5, 16, 16, 32, 32, 3, 3, 3), # canonical 3x3x3 (2D hits Winograd) + (4, 12, 10, 16, 48, 3, 3, 3), # Cout != Cin + (6, 14, 14, 32, 32, 1, 3, 3), # KD = 1 + (5, 12, 12, 16, 16, 5, 1, 1), # larger KD, 1x1 spatial + (4, 10, 10, 32, 16, 2, 3, 3), # KD = 2 + ]: + x = mx.random.normal((1, T, H, W, Cin)) + w = mx.random.normal((Cout, kd, kh, kw, Cin)) + # flip mirrors every kernel axis, including the decomposed depth + for flip in (False, True): + y_gpu = mx.conv_general(x, w, stride=(1, 1, 1), flip=flip) + y_cpu = mx.conv_general( + x, w, stride=(1, 1, 1), flip=flip, stream=mx.cpu + ) + mx.eval(y_gpu, y_cpu) + self.assertTrue( + mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4), + f"3D small-kd mismatch T{T} H{H} W{W} " + f"C{Cin}->{Cout} k{kd}{kh}{kw} flip={flip}", + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py index 7289955ed..7d12cb16f 100644 --- a/Source/Cmlx/mlx/python/tests/test_conv_transpose.py +++ b/Source/Cmlx/mlx/python/tests/test_conv_transpose.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest from itertools import permutations @@ -485,6 +486,8 @@ def run_conv_transpose3D( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), ((15, 15, 15), (3, 3, 3), (3, 3, 3), (2, 2, 2)), + # Exercises the Metal phase-aware stride-2/kernel-2 path. + ((3, 4, 5), (2, 2, 2), (2, 2, 2), (0, 0, 0)), ): run_conv_transpose3D( N, C, O, idim, kdim, stride, padding, dtype=dtype @@ -805,6 +808,43 @@ def run_conv_transpose_3d_output_padding( dtype=dtype, ) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv_transpose_unfold_tiling(self): + # The explicit-GEMM conv path unfolds into one buffer that can exceed + # maxBufferLength for large outputs; it unfolds and runs the gemm in row + # tiles instead (issue #3082). + key = "MLX_CONV_UNFOLD_TILE_ROWS" + prev = os.environ.get(key) + cases = ( + (mx.conv_transpose1d, (2, 9, 4), (5, 3, 4), {"stride": 2}), + (mx.conv_transpose2d, (2, 5, 5, 4), (5, 3, 3, 4), {"stride": 2}), + (mx.conv_transpose3d, (1, 4, 4, 4, 2), (3, 2, 2, 2, 2), {"stride": 2}), + (mx.conv_transpose1d, (2, 9, 4), (6, 3, 2), {"stride": 2, "groups": 2}), + ) + try: + # tile_rows=1 forces uniform tiles; 7 does not divide any output, so + # it exercises a partial final tile too. + for conv, in_shape, wt_shape, kwargs in cases: + for tile_rows in (1, 7): + with self.subTest( + conv=conv.__name__, tile_rows=tile_rows, **kwargs + ): + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + w = mx.array(np.random.normal(size=wt_shape).astype(np.float32)) + os.environ.pop(key, None) + untiled = conv(x, w, **kwargs) + mx.eval(untiled) + os.environ[key] = str(tile_rows) + tiled = conv(x, w, **kwargs) + mx.eval(tiled) + self.assertTrue(np.allclose(untiled, tiled, atol=1e-4)) + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_double.py b/Source/Cmlx/mlx/python/tests/test_double.py index f957e8fc9..3186e7e57 100644 --- a/Source/Cmlx/mlx/python/tests/test_double.py +++ b/Source/Cmlx/mlx/python/tests/test_double.py @@ -296,11 +296,55 @@ def test_conversion(self): b = a.tolist() self.assertEqual(b, [1.0, 2.0]) + def test_python_float_keeps_double_precision(self): + with mx.stream(mx.cpu): + # https://github.com/ml-explore/mlx/issues/4160 + for v in (0.37, 0.1, math.pi): + self.assertEqual(mx.full((3,), v, dtype=mx.float64).tolist(), [v] * 3) + + # https://github.com/ml-explore/mlx/issues/4159 + a = mx.array([1.0], dtype=mx.float64) + for v in (0.1, 1e-4, math.pi): + out = a * v + self.assertEqual(out.dtype, mx.float64) + self.assertEqual(out.tolist(), [v]) + + # values outside the float32 range used to saturate to inf or zero + self.assertEqual((a * 1e300).tolist(), [1e300]) + self.assertEqual((a * 1e-300).tolist(), [1e-300]) + + # every op that pairs a scalar with an array goes through the same + # conversion + zero = mx.array([0.0], dtype=mx.float64) + self.assertEqual(mx.maximum(zero, 0.1).tolist(), [0.1]) + self.assertEqual(mx.minimum(a, 0.1).tolist(), [0.1]) + self.assertEqual(mx.clip(a, None, 0.1).tolist(), [0.1]) + self.assertEqual(mx.where(mx.array([False]), zero, 0.1).tolist(), [0.1]) + self.assertEqual( + mx.pad(a, 1, constant_values=0.1).tolist(), [0.1, 1.0, 0.1] + ) + + # the python float is still weak: it does not widen the array + for dtype in (mx.float16, mx.bfloat16, mx.float32): + self.assertEqual((mx.array([1.0], dtype=dtype) * 0.1).dtype, dtype) + self.assertEqual((mx.array([1], dtype=mx.int32) * 0.1).dtype, mx.float32) + + # pad() keeps returning the input's dtype for every input type + for dtype in (mx.int32, mx.float16, mx.bfloat16, mx.float32): + padded = mx.pad(mx.array([1.0], dtype=dtype), 1, constant_values=0.5) + self.assertEqual(padded.dtype, dtype) + def test_linspace(self): with mx.stream(mx.cpu): - vals = mx.linspace(0, math.pi, 2, mx.float64) + vals = mx.linspace(0, math.pi, 2, dtype=mx.float64) self.assertEqual(vals.tolist()[1], math.pi) + vals = mx.linspace(0, math.pi, 4, endpoint=False, dtype=mx.float64) + self.assertEqual(vals.dtype, mx.float64) + self.assertTrue( + np.allclose(vals.tolist(), np.linspace(0, math.pi, 4, endpoint=False)) + ) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_einsum.py b/Source/Cmlx/mlx/python/tests/test_einsum.py index a73ea3818..08884ae98 100644 --- a/Source/Cmlx/mlx/python/tests/test_einsum.py +++ b/Source/Cmlx/mlx/python/tests/test_einsum.py @@ -65,6 +65,39 @@ def test_longer_paths(self): mx_path = mx.einsum_path(case, *inputs) self.assertEqual(np_path[0][1:], mx_path[0]) + def test_scalar_operands(self): + # An empty subscript is a scalar operand. A trailing one used to be + # dropped by the parser, so "i,->i" looked like a single input. + s1 = mx.array(2.0) + s2 = mx.array(3.0) + v = mx.random.uniform(shape=(3,)) + m = mx.random.uniform(shape=(2, 3)) + + cases = [ + ("->", (s1,)), + (",->", (s1, s2)), + (",,->", (s1, s2, s1)), + ("i,->i", (v, s1)), + (",i->i", (s1, v)), + ("ij,->ij", (m, s1)), + (",ij->ij", (s1, m)), + ("i,,->i", (v, s1, s2)), + ] + for spec, operands in cases: + mx_out = mx.einsum(spec, *operands) + np_out = np.einsum(spec, *[np.array(o) for o in operands]) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + # Operand count still has to match the number of subscripts + with self.assertRaises(ValueError): + mx.einsum(",->", s1) + with self.assertRaises(ValueError): + mx.einsum("i,->i", v) + # An empty subscript requires a 0-d operand + with self.assertRaises(ValueError): + mx.einsum(",->", v, s1) + def test_simple_einsum(self): a = mx.arange(4 * 4).reshape(4, 4) a_mx = mx.einsum("ii->i", a) @@ -188,7 +221,6 @@ def test_broadcasting(self): a = mx.full((5, 1), 1.0) b = mx.full((8, 2), 1.0) a_mx = mx.einsum("ab,bc->c", a, b) - return a_np = np.einsum("ab,bc->c", a, b) self.assertTrue(np.array_equal(a_mx, a_np)) @@ -358,6 +390,40 @@ def inputs_for_case(test_case): with self.assertRaises(ValueError): mx.einsum(test_case[1], *inputs) + def test_ellipses_broadcast(self): + # Size 1 batch dimensions covered by an ellipsis have to broadcast + # against the other operands, including when the smaller operand + # comes first. + shape_pairs = [ + ((1, 3, 4), (2, 4, 5)), + ((2, 3, 4), (1, 4, 5)), + ((1, 1, 3, 4), (5, 2, 4, 5)), + ((5, 1, 3, 4), (1, 2, 4, 5)), + ] + for sa, sb in shape_pairs: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...ij,...jk->...ik", a, b) + np_out = np.einsum("...ij,...jk->...ik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + for sa, sb in [((1, 4), (5, 4)), ((5, 4), (1, 4))]: + a = mx.random.uniform(shape=sa) + b = mx.random.uniform(shape=sb) + mx_out = mx.einsum("...i,...i->...", a, b) + np_out = np.einsum("...i,...i->...", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + + # Same thing with explicit labels rather than an ellipsis + a = mx.random.uniform(shape=(1, 3, 4)) + b = mx.random.uniform(shape=(2, 4, 5)) + mx_out = mx.einsum("bij,bjk->bik", a, b) + np_out = np.einsum("bij,bjk->bik", np.array(a), np.array(b)) + self.assertEqual(mx_out.shape, np_out.shape) + self.assertTrue(np.allclose(mx_out, np_out, rtol=1e-4, atol=1e-4)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_eval.py b/Source/Cmlx/mlx/python/tests/test_eval.py index 5d6daaec2..265a090aa 100644 --- a/Source/Cmlx/mlx/python/tests/test_eval.py +++ b/Source/Cmlx/mlx/python/tests/test_eval.py @@ -195,6 +195,47 @@ def test_multistream_deadlock(self): mx.eval(z) mx.set_memory_limit(old_limit) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_eval_exception_does_not_corrupt_state(self): + # An exception thrown from inside a primitive's eval (here a Metal + # compile error raised lazily at eval time) must not corrupt arrays + # evaluated earlier in the same batch: they are already marked + # evaluated, so their pending command buffers must still be + # committed before the exception propagates. + a = mx.full((1024,), 3.0) + b = a * 2.0 # encoded in the same eval batch as the failing kernel + + kernel = mx.fast.metal_kernel( + name="test_eval_exception_bad_kernel", + input_names=["inp"], + output_names=["out"], + source="this is not metal code {", + ) + with self.assertRaises(Exception): + (y,) = kernel( + inputs=[b], + output_shapes=[b.shape], + output_dtypes=[b.dtype], + grid=(1, 1, 1), + threadgroup=(1, 1, 1), + ) + mx.eval(y) + + self.assertTrue(mx.all(b == 6.0).item()) + + # Fresh computations after the failure stay correct. + x = mx.full((512,), 2.0) + self.assertEqual((x + 1.0).sum().item(), 512.0 * 3.0) + + @unittest.skipIf( + mx.cuda.is_available(), "CUDA backend waits cpu stream synchronously" + ) + def test_async_eval_error_in_synchronize(self): + a = mx.linalg.inv(mx.array([[1.0, 2.0], [2.0, 4.0]]), stream=mx.cpu) + mx.async_eval(a) + with self.assertRaises(RuntimeError): + mx.synchronize(mx.cpu) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_export_import.py b/Source/Cmlx/mlx/python/tests/test_export_import.py index c8f99bbb6..23f0645b3 100644 --- a/Source/Cmlx/mlx/python/tests/test_export_import.py +++ b/Source/Cmlx/mlx/python/tests/test_export_import.py @@ -1,6 +1,7 @@ # Copyright © 2024 Apple Inc. import gc +import json import os import tempfile import unittest @@ -313,6 +314,24 @@ def fun(a, b, c): out = imported_fun(x, y, z)[0] self.assertTrue(mx.array_equal(expected, out)) + def test_export_searchsorted(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + # both sides, since the side is the primitive's only state and a lost + # state would still round trip for the default + for side in ("left", "right"): + + def fun(a, v): + return mx.searchsorted(a, v, side=side) + + x = mx.sort(mx.random.uniform(shape=(32,))) + y = mx.random.uniform(shape=(3, 5)) + mx.export_function(path, fun, (x, y)) + imported_fun = mx.import_function(path) + expected = fun(x, y) + out = imported_fun(x, y)[0] + self.assertTrue(mx.array_equal(expected, out)) + def test_export_conv(self): path = os.path.join(self.test_dir, "fn.mlxfn") @@ -501,13 +520,14 @@ def fn(x, y): def test_export_with_callback(self): def fn(x, y): - return mx.log(mx.abs(x - y)) + return mx.log(mx.abs(x - y)).astype(mx.int32) n_in = None n_out = None n_const = None keywords = None primitives = [] + primitive_args = [] def callback(args): nonlocal n_in, n_out, n_const, keywords, primitives @@ -522,6 +542,7 @@ def callback(args): n_const = len(args["constants"]) elif t == "primitive": primitives.append(args["name"]) + primitive_args.append(args["arguments"]) mx.export_function(callback, fn, mx.array(1.0), y=mx.array(1.0)) self.assertEqual(n_in, 2) @@ -529,7 +550,11 @@ def callback(args): self.assertEqual(n_const, 0) self.assertEqual(len(keywords), 1) self.assertEqual(keywords[0][0], "y") - self.assertEqual(primitives, ["Subtract", "Abs", "Log"]) + self.assertEqual(primitives, ["Subtract", "Abs", "Log", "AsType"]) + self.assertEqual(primitive_args[0], []) + self.assertEqual(primitive_args[1], []) + self.assertEqual(primitive_args[2], [2]) + self.assertEqual(primitive_args[3], [mx.int32]) @unittest.skipIf(not mx.is_available(mx.gpu), "No GPU available") def test_export_import_custom_kernel(self): @@ -575,6 +600,95 @@ def call(a): out = imported(a)[0] self.assertTrue(mx.allclose(expected, out)) + def test_export_custom_metal_kernel_without_evaluation(self): + source = """ + uint elem = thread_position_in_grid.x; + out[elem] = a[elem]; + """ + kernel = mx.fast.metal_kernel( + name="export_only", + input_names=["a"], + output_names=["out"], + source=source, + ) + + def call(a): + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(min(a.size, 256), 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.zeros((2, 2)) + for shapeless in (False, True): + path = os.path.join( + self.test_dir, + f"metal_kernel_export_only_{shapeless}.mlxfn", + ) + mx.export_function(path, call, a, shapeless=shapeless) + self.assertTrue(os.path.exists(path)) + + # A shapeless import can't be called since CustomKernel does + # not support shape inference + if mx.metal.is_available() and not shapeless: + imported = mx.import_function(path) + self.assertTrue(mx.array_equal(imported(a)[0], call(a))) + + def call_cpu(a): + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(min(a.size, 256), 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.cpu, + )[0] + + path = os.path.join(self.test_dir, "metal_kernel_export_cpu.mlxfn") + with self.assertRaisesRegex(ValueError, "Only supports the GPU"): + mx.export_function(path, call_cpu, a) + + if not mx.metal.is_available(): + with self.assertRaisesRegex(RuntimeError, "No Metal back-end"): + call(a) + with self.assertRaisesRegex(RuntimeError, "No Metal back-end"): + mx.eval(mx.compile(call)(a)) + + def test_export_custom_metal_kernel_with_math_mode(self): + source = """ + uint elem = thread_position_in_grid.x; + out[elem] = metal::exp(a[elem]); + """ + kernel = mx.fast.metal_kernel( + name="math_mode_export", + input_names=["a"], + output_names=["out"], + source=source, + compile_options={"math_mode": "safe"}, + ) + + def call(a): + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(min(a.size, 256), 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.array([-float("inf"), 0.0]) + path = os.path.join(self.test_dir, "metal_kernel_math_mode.mlxfn") + mx.export_function(path, call, a) + self.assertTrue(os.path.exists(path)) + + if mx.metal.is_available(): + imported = mx.import_function(path) + self.assertTrue(mx.array_equal(imported(a)[0], call(a))) + def test_export_import_multi_with_constants(self): path = os.path.join(self.test_dir, "fn.mlxfn") @@ -609,6 +723,86 @@ def fun(x, y, z): imported = mx.import_function(path) self.assertTrue(mx.array_equal(imported(x, y, z)[0], fun(x, y, z))) + def test_export_matmul_shapeless_mid_dim(self): + path = os.path.join(self.test_dir, "matmul_shapeless.mlxfn") + + E, H = 64, 17 + arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H)) + + def fn(x): + return mx.matmul(x, arr) + + sample = mx.zeros((1, 40, E), dtype=mx.float32) + mx.export_function(path, fn, sample, shapeless=True) + imported = mx.import_function(path) + + for seq_len in (40, 248, 623): + with self.subTest(seq_len=seq_len): + x = mx.arange(seq_len * E, dtype=mx.float32).reshape((1, seq_len, E)) + expected = fn(x) + (y,) = imported(x) + self.assertEqual(y.shape, (1, seq_len, H)) + self.assertTrue(mx.allclose(y, expected)) + + def test_export_matmul_shapeless_batch_and_mid_dim(self): + path = os.path.join(self.test_dir, "matmul_shapeless_batch.mlxfn") + + B, E, H = 2, 32, 8 + arr = mx.arange(E * H, dtype=mx.float32).reshape((E, H)) * (1.0 / (E * H)) + + def fn(x): + return mx.matmul(x, arr) + + sample = mx.zeros((B, 10, E), dtype=mx.float32) + mx.export_function(path, fn, sample, shapeless=True) + imported = mx.import_function(path) + + for seq_len in (10, 50, 100): + with self.subTest(seq_len=seq_len): + x = mx.arange(B * seq_len * E, dtype=mx.float32).reshape( + (B, seq_len, E) + ) + expected = fn(x) + (y,) = imported(x) + self.assertEqual(y.shape, (B, seq_len, H)) + self.assertTrue(mx.allclose(y, expected)) + + def test_export_import_metadata(self): + path = os.path.join(self.test_dir, "fn.mlxfn") + + def fun(x): + return mx.abs(x) + + x = mx.array([1.0, -2.0, 3.0]) + metadata = json.dumps({"name": "model", "params": 7_000_000_000, "lr": 0.1}) + + mx.export_function(path, fun, x, metadata=metadata) + + imported = mx.import_function(path) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + imported, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + self.assertEqual(json.loads(imported_metadata)["params"], 7_000_000_000) + self.assertTrue(mx.array_equal(imported(x)[0], fun(x))) + + mx.export_function(path, fun, x) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, "") + + # Metadata survives the per-trace header rewrite of a multi-trace export + with mx.exporter(path, fun, metadata=metadata) as exporter: + exporter(mx.array([1.0])) + exporter(mx.array([1.0, 2.0])) + _, imported_metadata = mx.import_function(path, return_metadata=True) + self.assertEqual(imported_metadata, metadata) + + with self.assertRaises(TypeError): + mx.export_function(path, fun, x, metadata={"name": "model"}) + + with self.assertRaises(ValueError): + mx.export_function(lambda x: None, fun, x, metadata=metadata) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_fast.py b/Source/Cmlx/mlx/python/tests/test_fast.py index 63c5b06ac..200781d37 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast.py +++ b/Source/Cmlx/mlx/python/tests/test_fast.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest import mlx.core as mx @@ -169,6 +170,55 @@ def test_rope(self): x, dims, traditional=traditional, base=base, scale=scale, offset=offset ) + @unittest.skipIf("CI" in os.environ, "Allocates too much memory for CI") + def test_rope_large_input(self): + dims, seq_len, batch_size, n_heads = 32, 8192, 8, 32 + base, scale, offset, traditional = 10000.0, 1.0, 0, False + x = mx.random.normal(shape=[batch_size, seq_len, n_heads, dims]).astype( + mx.float32 + ) + x = x.swapaxes(1, 2) + rx_fast = mx.fast.rope( + x, dims, traditional=traditional, base=base, scale=scale, offset=offset + ) + ref = rope_orig(x, dims, traditional, base, scale, offset) + self.assertLess(mx.abs(ref - rx_fast).max(), 5e-3) + + def test_rope_dims_validation(self): + T = 4 + feature_dim = 64 + x = mx.random.uniform(shape=(1, T, feature_dim)) + + # dims = 0 should raise + with self.assertRaises(ValueError): + mx.fast.rope( + x, dims=0, traditional=False, base=10000.0, scale=1.0, offset=0 + ) + + # negative dims should raise + with self.assertRaises(ValueError): + mx.fast.rope( + x, dims=-2, traditional=False, base=10000.0, scale=1.0, offset=0 + ) + + # odd dims should raise + with self.assertRaises(ValueError): + mx.fast.rope( + x, dims=7, traditional=False, base=10000.0, scale=1.0, offset=0 + ) + + # dims > feature_dim should raise + with self.assertRaises(ValueError): + mx.fast.rope( + x, dims=128, traditional=False, base=10000.0, scale=1.0, offset=0 + ) + + # valid dims should not raise + mx.fast.rope(x, dims=32, traditional=False, base=10000.0, scale=1.0, offset=0) + mx.fast.rope( + x, dims=feature_dim, traditional=False, base=10000.0, scale=1.0, offset=0 + ) + def test_rope_with_freqs(self): mx.random.seed(0) @@ -330,6 +380,28 @@ def test_rope_batch(self): rx = rope_orig(x, dims, traditional, base, scale, offset) self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5) + def test_rope_single_batch(self): + base = 10000.0 + scale = 1.0 + offset = 5 + + for traditional in [True, False]: + for B in [2, 4, 8]: + for n_head in [1, 4, 7]: + for dims in [64, 128]: + x = mx.random.uniform(shape=(B, n_head, 1, dims)) + mx.eval(x) + rx_fast = mx.fast.rope( + x, + dims, + traditional=traditional, + base=base, + scale=scale, + offset=offset, + ) + rx = rope_orig(x, dims, traditional, base, scale, offset) + self.assertLess(mx.abs(rx - rx_fast).max(), 1e-5) + def test_rope_with_large_offset(self): x = mx.random.normal(shape=(1, 1, 1024, 32)) rx_fp32 = mx.fast.rope( @@ -356,7 +428,7 @@ def test_rms_norm(self): dtypes = [mx.float32, mx.float16, mx.bfloat16] epss = [1e-3, 1e-5] - dimss = [31, 32, 33] + dimss = [31, 32, 33, 256, 512] defaults = (mx.float32, 1e-5, 32) for dtype in dtypes: @@ -411,23 +483,23 @@ def test_rms_norm(self): mx.fast.rms_norm(x, mx.ones((4,)), 1e-5) def test_rms_norm_grad(self): - D = 32 eps = 1e-5 f1 = lambda x, w, y: (rms_norm(x, w, eps) * y).sum() f2 = lambda x, w, y: (mx.fast.rms_norm(x, w, eps) * y).sum() f3 = lambda x, y: (rms_norm(x, mx.ones((x.shape[-1],)), eps) * y).sum() f4 = lambda x, y: (mx.fast.rms_norm(x, None, eps) * y).sum() - x = mx.random.uniform(shape=(8, 100, D)) - w = mx.random.uniform(shape=(D,)) - y = mx.random.uniform(shape=(8, 100, D)) - gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) - gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) - self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) - self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) - gx1 = mx.grad(f3, argnums=(0,))(x, y) - gx2 = mx.grad(f4, argnums=(0,))(x, y) - self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) + for D in [32, 256]: + x = mx.random.uniform(shape=(8, 100, D)) + w = mx.random.uniform(shape=(D,)) + y = mx.random.uniform(shape=(8, 100, D)) + gx1, gw1 = mx.grad(f1, argnums=(0, 1))(x, w, y) + gx2, gw2 = mx.grad(f2, argnums=(0, 1))(x, w, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) + self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + gx1 = mx.grad(f3, argnums=(0,))(x, y) + gx2 = mx.grad(f4, argnums=(0,))(x, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) D = 8192 x = mx.random.uniform(shape=(2, 2, D)) @@ -453,6 +525,67 @@ def inner(x, w, y): self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + def test_cross_entropy(self): + def cross_entropy_ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits.astype(mx.float32), axis=-1) - score.astype( + mx.float32 + ) + + tolerances = {mx.float32: 1e-5, mx.float16: 3e-2, mx.bfloat16: 3e-1} + + for V in [7, 32, 128, 255, 256, 1000, 4096, 8192]: + for dtype in [mx.float32, mx.float16, mx.bfloat16]: + logits = (mx.random.normal(shape=(4, 7, V), scale=3.0) * 2).astype( + dtype + ) + targets = mx.random.randint(0, V, shape=(4, 7)) + expected = cross_entropy_ref(logits, targets) + out = mx.fast.cross_entropy(logits, targets) + self.assertEqual(out.dtype, mx.float32) + self.assertEqual(out.shape, targets.shape) + self.assertLess(mx.abs(out - expected).max().item(), tolerances[dtype]) + + def test_cross_entropy_shape_checks(self): + logits = mx.random.normal(shape=(4, 16)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((5,), mx.int32)) + with self.assertRaises(ValueError): + # Probability targets are not supported by the fused op. + mx.fast.cross_entropy(logits, mx.zeros((4, 16), mx.int32)) + with self.assertRaises(ValueError): + mx.fast.cross_entropy(logits, mx.zeros((4,), mx.float32)) + + def test_cross_entropy_grad(self): + def ref(logits, targets): + score = mx.take_along_axis(logits, mx.expand_dims(targets, -1), -1).squeeze( + -1 + ) + return mx.logsumexp(logits, axis=-1) - score + + f1 = lambda x, y: ref(x, y).mean() + f2 = lambda x, y: mx.fast.cross_entropy(x, y).mean() + + for V in [7, 128, 1000, 4096]: + logits = mx.random.normal(shape=(4, 7, V), scale=2.0) + targets = mx.random.randint(0, V, shape=(4, 7)) + g1 = mx.grad(f1, argnums=0)(logits, targets) + g2 = mx.grad(f2, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + + w = mx.random.uniform(shape=(4, 7)) + f3 = lambda x, y: (ref(x, y) * w).sum() + f4 = lambda x, y: (mx.fast.cross_entropy(x, y) * w).sum() + logits = mx.random.normal(shape=(4, 7, 512), scale=2.0) + targets = mx.random.randint(0, 512, shape=(4, 7)) + g1 = mx.grad(f3, argnums=0)(logits, targets) + g2 = mx.grad(f4, argnums=0)(logits, targets) + self.assertEqual(g2.shape, logits.shape) + self.assertLess(mx.abs(g1 - g2).max().item(), 1e-6) + def test_layer_norm_dim_check(self): with self.assertRaises(ValueError): weight = mx.ones((129,)) @@ -601,6 +734,42 @@ def inner(x, w, b, y): self.assertLess(mx.abs(gb1).max(), 1e-9) self.assertLess(mx.abs(gb2).max(), 1e-9) + def test_layer_norm_grad_no_bias(self): + # Second-order gradient through layer_norm with weight but no bias. + # Regression test: the VJP fallback had zeros_like(w) instead of + # zeros_like(b) for the bias placeholder gradient, causing a shape + # mismatch that crashes on higher-order differentiation. + D = 8 + eps = 1e-5 + x = mx.random.uniform(shape=(2, 4, D)) + w = mx.random.uniform(shape=(D,)) + y = mx.random.uniform(shape=(2, 4, D)) + mx.eval(x, w, y) + + f_ref = lambda x, w, y: (layer_norm(x, w, None, eps) * y).sum() + f_fast = lambda x, w, y: (mx.fast.layer_norm(x, w, None, eps) * y).sum() + + # First order should match reference + gx1, gw1 = mx.grad(f_ref, argnums=(0, 1))(x, w, y) + gx2, gw2 = mx.grad(f_fast, argnums=(0, 1))(x, w, y) + self.assertLess(mx.abs(gx1 - gx2).max(), 1e-5) + self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 1e-5) + + # Second order — this crashes without the fix due to shape mismatch + # in the bias placeholder gradient: zeros_like(w) shape (D,) vs + # expected zeros_like(b) shape () + def gf(f): + def inner(x, w, y): + gx, gw = mx.grad(f, argnums=(0, 1))(x, w, y) + return ((gx + gw) * y).sum() + + return inner + + gx1, gw1 = mx.grad(gf(f_ref), argnums=(0, 1))(x, w, y) + gx2, gw2 = mx.grad(gf(f_fast), argnums=(0, 1))(x, w, y) + self.assertLess(mx.abs(gx1 - gx2).max() / mx.abs(gx1).mean(), 5e-5) + self.assertLess(mx.abs(gw1 - gw2).max() / mx.abs(gw1).mean(), 5e-5) + def test_layer_norm_grad_no_params(self): eps = 1e-5 f1 = lambda x: layer_norm(x, None, None, eps).sum() @@ -918,6 +1087,158 @@ def call_kernel(a: mx.array, source): out = call_kernel(a, source) self.assertTrue(mx.array_equal(out, mx.ones_like(out))) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_kernel_same_name_different_source_one_eval(self): + # Regression test for #3832: two kernels sharing a name but with + # different sources, dispatched in a SINGLE eval batch, must each run + # their own compiled code instead of silently reusing the first's. + def call_kernel(a, source): + kernel = mx.fast.metal_kernel( + name="dup_name", + input_names=["inp"], + output_names=["out"], + source=source, + ) + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.arange(32, dtype=mx.float32) + out_a = call_kernel( + a, "uint e = thread_position_in_grid.x; out[e] = inp[e] * 2.0f;" + ) + out_b = call_kernel( + a, "uint e = thread_position_in_grid.x; out[e] = inp[e] + 100.0f;" + ) + mx.eval(out_a, out_b) # one batch — the reported failure case + self.assertTrue(mx.array_equal(out_a, a * 2.0)) + self.assertTrue(mx.array_equal(out_b, a + 100.0)) + + @unittest.skipIf(not mx.cuda.is_available(), "CUDA is not available") + def test_cuda_kernel_same_name_different_source(self): + # The CUDA module cache was keyed on the kernel name alone, so the + # second kernel here silently ran the first one's code. Metal had the + # same bug, fixed in #3833. + def call_kernel(a, source): + kernel = mx.fast.cuda_kernel( + name="dup_name", + input_names=["inp"], + output_names=["out"], + source=source, + ) + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.arange(32, dtype=mx.float32) + elem = "auto e = cooperative_groups::this_grid().thread_rank();" + out_a = call_kernel(a, f"{elem} out[e] = inp[e] * 2.0f;") + out_b = call_kernel(a, f"{elem} out[e] = inp[e] + 100.0f;") + mx.eval(out_a, out_b) + self.assertTrue(mx.array_equal(out_a, a * 2.0)) + self.assertTrue(mx.array_equal(out_b, a + 100.0)) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_metal_kernel_math_mode(self): + with self.assertRaises(ValueError): + mx.fast.metal_kernel( + name="invalid_math_mode", + input_names=["inp"], + output_names=["out"], + source="out[0] = inp[0];", + compile_options={"math_mode": "precise"}, + ) + + with self.assertRaises(ValueError): + mx.fast.metal_kernel( + name="invalid_compile_options", + input_names=["inp"], + output_names=["out"], + source="out[0] = inp[0];", + compile_options={"unknown": "value"}, + ) + + # Numerical special cases such as exp(-inf) can agree between math + # modes, so they don't reliably detect whether the mode was applied. + # Branch on the compiler's __FAST_MATH__ macro instead: it is defined + # only when fast math is enabled, so the test fails if the selected + # math mode is not forwarded to the Metal compiler. + source = """ + uint elem = thread_position_in_grid.x; + #if defined(__FAST_MATH__) && __FAST_MATH__ + out[elem] = 1.0f; + #else + out[elem] = 0.0f; + #endif + """ + + a = mx.zeros((4,), dtype=mx.float32) + expected = { + "safe": mx.zeros_like(a), + "fast": mx.ones_like(a), + } + + # Reuse the same kernel name across modes so the library cache is forced + # to rebuild when the math mode changes, guarding against a stale build + # being returned for a different mode. + for mode, expected_out in expected.items(): + kernel = mx.fast.metal_kernel( + name="math_mode", + input_names=["inp"], + output_names=["out"], + source=source, + compile_options={"math_mode": mode}, + ) + out = kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + self.assertTrue(mx.array_equal(out, expected_out)) + + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_custom_kernel_mixed_dtypes(self): + # Calling the same kernel with different input dtypes in a single + # graph should not invalidate pipeline states that are still in use + # by an uncommitted command buffer + kernel = mx.fast.metal_kernel( + name="mixed_dtypes", + input_names=["inp"], + output_names=["out"], + source=""" + uint elem = thread_position_in_grid.x; + out[elem] = inp[elem] + inp[elem]; + """, + ) + + def call_kernel(a: mx.array): + return kernel( + inputs=[a], + grid=(a.size, 1, 1), + threadgroup=(a.size, 1, 1), + output_shapes=[a.shape], + output_dtypes=[a.dtype], + stream=mx.gpu, + )[0] + + a = mx.full((32,), 1.5, dtype=mx.float16) + b = mx.full((32,), 2.5, dtype=mx.float32) + out = call_kernel(a).astype(mx.float32) + call_kernel(b) + self.assertTrue(mx.allclose(out, mx.full((32,), 8.0))) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py b/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py index fa6d03985..b11ce884e 100644 --- a/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py +++ b/Source/Cmlx/mlx/python/tests/test_fast_sdpa.py @@ -1,5 +1,7 @@ import math +import os import unittest +from itertools import product import mlx.core as mx import mlx_tests @@ -74,31 +76,23 @@ def do_attention(f, q, k, v, scale, mask=None, transpose=False): def prepare_inputs(B, qL, kL, D, qH, kH, mask, transpose, dtype): - np.random.seed(0) - np_dtype = getattr(np, dtype) + mx.random.seed(0) + scale = 1.0 / math.sqrt(D) shape_q = (B, qL, qH, D) if transpose else (B, qH, qL, D) shape_kv = (B, kL, kH, D) if transpose else (B, kH, kL, D) - scale = 1.0 / math.sqrt(D) - - q_np = np.random.normal(0.0, 0.5, shape_q).astype(np_dtype) - k_np = np.random.normal(0.0, 0.5, shape_kv).astype(np_dtype) - v_np = np.random.normal(0.0, scale, shape_kv).astype(np_dtype) - - q_mx = mx.array(q_np) - k_mx = mx.array(k_np) - v_mx = mx.array(v_np) + q = mx.random.uniform(0.0, 0.5, shape_q, dtype) + k = mx.random.uniform(0.0, 0.5, shape_kv, dtype) + v = mx.random.uniform(0.0, scale, shape_kv, dtype) if mask is not None: if mask == "additive": - mask_np = np.random.normal(0.0, 0.5, (B, qH, qL, kL)).astype(np_dtype) - mask = mx.array(mask_np) + mask = mx.random.uniform(0.0, 0.5, (B, qH, qL, kL), dtype) elif mask == "bool": - mask_np = np.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5 - mask = mx.array(mask_np) + mask = mx.random.uniform(0.0, 1.0, (B, qH, qL, kL)) < 0.5 - return q_mx, k_mx, v_mx, scale, mask + return q, k, v, scale, mask # SDPA for MHA (n_heads == n_kv_heads) @@ -122,192 +116,129 @@ def mlx_primitives_sdpa(q, k, v, scale, mask=None): return scores @ v -# SDPA for GQA (n_heads > n_kv_heads, n_kv_heads > 1, n_heads % n_kv_heads == 0) -def mlx_primitives_sdpa_with_gqa(q, k, v, scale, mask=None): - n_repeats = q.shape[1] // k.shape[1] - - # borrowing kv cache tiling from mlx-examples/llms/mistral/mistral.py - n_heads = q.shape[1] - B = q.shape[0] - L = k.shape[2] - - def repeat(a): - a = mx.concatenate([mx.expand_dims(a, 2)] * n_repeats, axis=2) - return a.reshape([B, n_heads, L, -1]) - - k, v = map(repeat, (k, v)) - - return mlx_primitives_sdpa(q, k, v, scale, mask=mask) - - -class TestFastSelfAttentionSDPA(mlx_tests.MLXTestCase): - def test_fast_sdpa(self): - # Not yet supported: - # * K pre-transposed in kernel, V pre-transposed in kernel - np.random.seed(0) - R = 20 - L = R - Dk = 64 - H = 3 - scale = float(1.0 / np.sqrt(Dk)) - q_npy = np.random.normal(0.0, 1.0, (1, H, R, Dk)).astype(np.float32) - k_npy = np.random.normal(0.0, 1.0, (1, H, L, Dk)).astype(np.float32) - v_npy = np.random.normal(0.0, 1.0, (1, H, L, Dk)).astype(np.float32) - - q_mlx = mx.array(q_npy) - k_mlx = mx.array(k_npy) - v_mlx = mx.array(v_npy) - - reference = mlx_primitives_sdpa(q_mlx, k_mlx, v_mlx, scale) - - o_mlx = mx.fast.scaled_dot_product_attention( - q_mlx, k_mlx, v_mlx, scale=scale, mask=None - ) - - self.assertListEqual(list(reference.shape), list(o_mlx.shape)) - self.assertTrue(mx.allclose(o_mlx, reference, atol=1e-4)) - - dtypes = [np.float32] - - Dk = 64 - - if mx.is_available(mx.gpu): - dtypes.append(np.half) - - for SEQUENCE_LENGTH in [63, 129, 400]: - for DTYPE in dtypes: - B = 2 - H = 24 - n_kv_heads = H - q_npy = np.random.normal(0.0, 1.0, (B, H, SEQUENCE_LENGTH, Dk)).astype( - DTYPE +class TestFastSDPA(mlx_tests.MLXTestCase): + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_72(self): + B, D, qH, kH = (1, 72, 8, 2) + for qL, kL, dtype, mask_str in product( + (64, 65), + (128, 127), + (mx.float16, mx.bfloat16, mx.float32), + (None, "additive", "bool", "causal"), + ): + with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, False, dtype ) - k_npy = np.random.normal( - 0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk) - ).astype(DTYPE) - v_npy = np.random.normal( - 0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk) - ).astype(DTYPE) - - q_mlx = mx.array(q_npy) - k_mlx = mx.array(k_npy) - v_mlx = mx.array(v_npy) - - reference = mlx_primitives_sdpa_with_gqa(q_mlx, k_mlx, v_mlx, scale) - o_mlx = mx.fast.scaled_dot_product_attention( - q_mlx, - k_mlx, - v_mlx, - scale=scale, + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask ) - self.assertListEqual(list(reference.shape), list(o_mlx.shape)) - rtol = 1e-3 - atol = 1e-2 - - if SEQUENCE_LENGTH > 500: - rtol = 1e-2 - - if DTYPE == np.half: - rtol = 1e-2 - - self.assertTrue(mx.allclose(o_mlx, reference, rtol=rtol, atol=atol)) - + if dtype == mx.float32: + atol = 1e-5 + elif dtype == mx.bfloat16: + atol = 5e-3 + else: + atol = 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_head_dim_96(self): + B, D, qH, kH = (1, 96, 8, 2) + for qL, kL, dtype, mask_str in product( + (64, 65), + (128, 127), + (mx.float16, mx.bfloat16, mx.float32), + (None, "additive", "bool", "causal"), + ): + with self.subTest(qL=qL, kL=kL, dtype=dtype, mask=mask_str): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, False, dtype + ) + ref = mlx_ref_attn(q, k, v, scale, mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) -class TestFastSDPA(mlx_tests.MLXTestCase): - def test_fast_sdpa(self): - # Not yet supported: - # * K pre-transposed in kernel, V pre-transposed in kernel - np.random.seed(0) - L = 43 - R = 1 - Dk = 128 - scale = float(1.0 / np.sqrt(128.0)) - q_npy = np.random.normal(0.0, 1.0, (1, 32, R, Dk)).astype(np.float32) - k_npy = np.random.normal(0.0, 1.0, (1, 32, L, Dk)).astype(np.float32) - v_npy = np.random.normal(0.0, 1.0, (1, 32, L, Dk)).astype(np.float32) - - q_mlx = mx.array(q_npy) - k_mlx = mx.array(k_npy) - v_mlx = mx.array(v_npy) - - reference = mlx_primitives_sdpa(q_mlx, k_mlx, v_mlx, scale) - - o_mlx = mx.fast.scaled_dot_product_attention( - q_mlx, k_mlx, v_mlx, scale=scale, mask=None - ) - - self.assertListEqual(list(reference.shape), list(o_mlx.shape)) - self.assertTrue(mx.allclose(o_mlx, reference, atol=1e-4)) - - B = 1 - H = 32 - dtypes = [np.float32] - if mx.is_available(mx.gpu): - dtypes.append(np.half) - - for SEQUENCE_LENGTH in [1, 7, 9, 32, 63, 67, 129, 400, 2000]: - for DO_GQA in [0, 1]: - for DTYPE in dtypes: - n_kv_heads = 8 if DO_GQA else 32 - q_npy = np.random.normal(0.0, 1.0, (B, H, R, Dk)).astype(DTYPE) - k_npy = np.random.normal( - 0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk) - ).astype(DTYPE) - v_npy = np.random.normal( - 0.0, 1.0, (B, n_kv_heads, SEQUENCE_LENGTH, Dk) - ).astype(DTYPE) - - q_mlx = mx.array(q_npy) - k_mlx = mx.array(k_npy) - v_mlx = mx.array(v_npy) - - reference = mlx_primitives_sdpa_with_gqa(q_mlx, k_mlx, v_mlx, scale) - o_mlx = mx.fast.scaled_dot_product_attention( - q_mlx, k_mlx, v_mlx, scale=scale + if dtype == mx.float32: + atol = 1e-5 + elif dtype == mx.bfloat16: + atol = 5e-3 + else: + atol = 3e-4 + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) + + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_full_head_dim_256(self): + # On NAX devices, large nearly-square causal blocks take the fused + # path; everything else takes the unfused fallback. Ragged lengths + # exercise the kernel's unaligned pipelines, and K/V sliced out of a + # longer preallocated cache (the way mlx-lm hands them over) exercise + # the dispatch reading the slice past its end. All of it must be + # correct. + D = 256 + Nq, Nkv = 8, 2 + scale = D**-0.5 + mx.random.seed(0) + cases = [ + # fused on NAX: aligned square, ragged square (unaligned Q and + # K/V), aligned rectangle at the routing boundary, ragged + # rectangle + (2048, 2048, "causal", None), + (2049, 2049, "causal", None), + (2048, 2560, "causal", None), + (2049, 2560, "causal", None), + # fused on NAX with ragged K/V sliced out of a longer cache: the + # rows behind the slice must not leak into the output + (2049, 2049, "causal", 2304), + (2048, 2500, "causal", 2560), + (1031, 2049, "causal", 2304), + ] + for dtype in (mx.float32, mx.bfloat16): + for qL, kL, mask, cache_len in cases: + with self.subTest( + dtype=dtype, qL=qL, kL=kL, mask=mask, cache_len=cache_len + ): + q = (5e-1 * mx.random.normal(shape=(1, Nq, qL, D))).astype(dtype) + if cache_len is None: + k = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype( + dtype + ) + v = (5e-1 * mx.random.normal(shape=(1, Nkv, kL, D))).astype( + dtype + ) + else: + # Large, finite stale rows behind the slice: any of + # them reaching the output is loud. + k_cache = 1e2 * mx.random.normal(shape=(1, Nkv, cache_len, D)) + v_cache = 1e3 * mx.random.normal(shape=(1, Nkv, cache_len, D)) + k_cache[..., :kL, :] = 5e-1 * mx.random.normal( + shape=(1, Nkv, kL, D) + ) + v_cache[..., :kL, :] = 5e-1 * mx.random.normal( + shape=(1, Nkv, kL, D) + ) + k = k_cache.astype(dtype)[..., :kL, :] + v = v_cache.astype(dtype)[..., :kL, :] + k_rep = mx.repeat(k, Nq // Nkv, axis=1) + v_rep = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, k_rep, v_rep, scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask ) - - self.assertListEqual(list(reference.shape), list(o_mlx.shape)) - rtol = 1e-5 - atol = 1e-1 - - if SEQUENCE_LENGTH > 500: - rtol = 1e-2 - - if DTYPE == np.half: - rtol = 1e-2 - - self.assertTrue(mx.allclose(o_mlx, reference, rtol=rtol, atol=atol)) - q = mx.random.normal(shape=(1, 32, 1, Dk)) - k = mx.random.normal(shape=(1, 32, 32, Dk)) - v = mx.random.normal(shape=(1, 32, 128, Dk)) - - atol = 1e-6 - y = mlx_primitives_sdpa(q, k, v[:, :, :32], scale) - y_hat = mx.fast.scaled_dot_product_attention(q, k, v[:, :, :32], scale=scale) - self.assertTrue(mx.allclose(y, y_hat, atol=atol)) - - # Test with per-example mask - q = mx.random.normal(shape=(2, 8, 4, 32)) - k = mx.random.normal(shape=(2, 2, 8, 32)) - v = mx.random.normal(shape=(2, 2, 8, 32)) - mask = 10 * mx.random.normal(shape=(2, 1, 4, 8)) - y = mlx_primitives_sdpa_with_gqa(q, k, v, scale, mask=mask) - y_hat = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask) - self.assertTrue(mx.allclose(y, y_hat, atol=atol)) - - # Test with boolean causal mask - indices = mx.arange(8) - bool_mask = indices[:, None] >= indices[None] - additive_mask = (~bool_mask).astype(mx.float32) * mx.finfo(mx.float32).min - x = mx.random.normal(shape=(1, 2, 8, 32)) - y = mlx_primitives_sdpa_with_gqa(x, x, x, scale, mask=additive_mask) - y_hat = mx.fast.scaled_dot_product_attention( - x, x, x, scale=scale, mask=bool_mask - ) - self.assertTrue(mx.allclose(y, y_hat, atol=atol)) - - def test_fast_sdpa_vector_kv_transposed_head_seq(self): + self.assertEqual(out.shape, ref.shape) + if dtype == mx.float32: + # The fused shapes run through tf32 tensor ops when + # MLX_ENABLE_TF32 is on (the default). + tol = 1e-3 if qL >= 2048 else 1e-4 + else: + tol = 5e-3 + self.assertTrue(mx.allclose(ref, out, atol=tol, rtol=tol)) + + def test_sdpa_vector_kv_transposed_head_seq(self): D = 64 Nq = 4 Nkv = 1 @@ -339,7 +270,7 @@ def test_fast_sdpa_vector_kv_transposed_head_seq(self): ) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) - def test_fast_sdpa_vector(self): + def test_sdpa_vector(self): D = 64 L = 43 Nq = 4 @@ -411,7 +342,21 @@ def test_fast_sdpa_vector(self): ) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) - def test_fully_masked(self): + def test_sdpa_vector_gqa_long(self): + scale = 1.0 + mx.random.seed(0) + for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]: + for L in [8192, 8201]: + q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D)) + k = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + v = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L] + kr = mx.repeat(k, Nq // Nkv, axis=1) + vr = mx.repeat(v, Nq // Nkv, axis=1) + ref = mlx_primitives_sdpa(q, kr, vr, scale) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + + def test_sdpa_fully_masked(self): Lkv = 8 mask = mx.array(False) for D in [128]: @@ -423,7 +368,7 @@ def test_fully_masked(self): out = mx.fast.scaled_dot_product_attention(q, k, v, mask=mask, scale=1) self.assertFalse(mx.any(mx.isnan(out))) - def test_inf_score(self): + def test_sdpa_inf_score(self): Lkv = 8 for D in [4, 128]: for Lq in [1, 8]: @@ -435,7 +380,7 @@ def test_inf_score(self): out = mx.fast.scaled_dot_product_attention(q, k, v, mask=None, scale=1) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) - def test_fast_sdpa_few_query(self): + def test_sdpa_few_query(self): D = 64 L = 43 Lq = 8 @@ -494,7 +439,7 @@ def test_fast_sdpa_few_query(self): self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) @unittest.skip("Different head and value dims is not enabled") - def test_fast_sdpa_vector_value_dims(self): + def test_sdpa_vector_value_dims(self): D = 192 V = 128 Nq = 4 @@ -550,91 +495,145 @@ def test_sdpa_vector_batched(self): ref = mlx_ref_attn(q, k, v, mask=mask) self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(not mx.is_available(mx.gpu), "GPU kernel path only") + def test_sdpa_blocks_env_override(self): + # MLX_SDPA_BLOCKS used to be applied as-is, and values that are not + # a multiple of 32 silently corrupted the 2-pass vector output. The + # override is now rounded up to a multiple of 32. + D = 128 + q = mx.random.normal(shape=(1, 32, 1, D), dtype=mx.float16) + k = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) + v = mx.random.normal(shape=(1, 8, 8192, D), dtype=mx.float16) + ref = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) + try: + for blocks in (16, 33, 48, 100): + os.environ["MLX_SDPA_BLOCKS"] = str(blocks) + out = mx.fast.scaled_dot_product_attention(q, k, v, scale=D**-0.5) + self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4)) + finally: + del os.environ["MLX_SDPA_BLOCKS"] -class TestSDPA(mlx_tests.MLXTestCase): + @unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU") def test_sdpa(self): - if not mx.is_available(mx.gpu): - return - # fmt: off - shapes_64 = ( + shapes_64 = [ # ( B, qsl, ksl, head_dim, n_qh, n_kvh) + ( 1, 20, 20, 64, 3, 3), + ( 1, 63, 63, 64, 24, 24), + ( 1, 129, 129, 64, 24, 24), + ( 1, 400, 400, 64, 24, 24), ( 1, 128, 128, 64, 32, 32), ( 1, 64, 128, 64, 32, 32), ( 1, 65, 128, 64, 32, 8), ( 1, 64, 127, 64, 32, 8), ( 1, 65, 127, 64, 32, 8), ( 1, 127, 65, 64, 32, 8), - ) - - shapes_128 = ( + ] + shapes_128 = [ # ( B, qsl, ksl, head_dim, n_qh, n_kvh) ( 1, 128, 128, 128, 32, 8), ( 1, 64, 128, 128, 32, 8), ( 1, 65, 127, 128, 32, 8), ( 1, 127, 65, 128, 32, 8), - ) + ] + for ksl in [7, 9, 32, 63, 67, 129, 400, 2000]: + shapes_128.append((1, 1, ksl, 128, 32, 32)) + shapes_128.append((1, 1, ksl, 128, 32, 8)) # fmt: on shapes = shapes_64 + shapes_128 - dtypes = ["float32", "float16"] + dtypes = [mx.float16] + if mx.metal.is_available(): + dtypes.append(mx.float32) masks = [None, "additive", "bool", "causal"] transposes = (False, True) - for dtype in dtypes: - for t in transposes: - for mask_str in masks: - for B, qL, kL, D, qH, kH in shapes: - with self.subTest( - B=B, - qsl=qL, - ksl=kL, - head_dim=D, - n_q_heads=qH, - n_kv_heads=kH, - mask=mask_str, - transpose=t, - dtype=dtype, - ): - - np.random.seed(0) - q_mx, k_mx, v_mx, scale, mask = prepare_inputs( - B, qL, kL, D, qH, kH, mask_str, t, dtype - ) - - out_ref = do_attention( - mlx_ref_attn, q_mx, k_mx, v_mx, scale, mask, t - ) - - out_fst = do_attention( - mx.fast.scaled_dot_product_attention, - q_mx, - k_mx, - v_mx, - scale, - mask, - t, - ) - - # For causal mask when qL > kL, first qL-kL rows are undefined - # Compare only the valid portion - if mask_str == "causal" and qL > kL: - offset = qL - kL - if t: # transpose=True: shape is (B, qL, qH, D) - out_ref = out_ref[:, offset:, :, :] - out_fst = out_fst[:, offset:, :, :] - else: # transpose=False: shape is (B, qH, qL, D) - out_ref = out_ref[:, :, offset:, :] - out_fst = out_fst[:, :, offset:, :] - - atol = 2e-5 if dtype == "float32" else 3e-4 - - self.assertListEqual( - list(out_ref.shape), list(out_fst.shape) - ) - - diff = mx.abs(out_fst - out_ref) - atol * mx.abs(out_ref) - self.assertLessEqual(mx.max(diff).item(), atol) + for dtype, t, mask_str, (B, qL, kL, D, qH, kH) in product( + dtypes, transposes, masks, shapes + ): + with self.subTest( + B=B, + qsl=qL, + ksl=kL, + head_dim=D, + n_q_heads=qH, + n_kv_heads=kH, + mask=mask_str, + transpose=t, + dtype=dtype, + ): + q, k, v, scale, mask = prepare_inputs( + B, qL, kL, D, qH, kH, mask_str, t, dtype + ) + + out_ref = do_attention(mlx_ref_attn, q, k, v, scale, mask, t) + + out_fst = do_attention( + mx.fast.scaled_dot_product_attention, + q, + k, + v, + scale, + mask, + t, + ) + + # For causal mask when qL > kL, first qL-kL rows are undefined + # Compare only the valid portion + if mask_str == "causal" and qL > kL: + offset = qL - kL + if t: # transpose=True: shape is (B, qL, qH, D) + out_ref = out_ref[:, offset:, :, :] + out_fst = out_fst[:, offset:, :, :] + else: # transpose=False: shape is (B, qH, qL, D) + out_ref = out_ref[:, :, offset:, :] + out_fst = out_fst[:, :, offset:, :] + + atol = 2e-5 if dtype == mx.float32 else 3e-4 + + self.assertListEqual(list(out_ref.shape), list(out_fst.shape)) + + diff = mx.abs(out_fst - out_ref) - atol * mx.abs(out_ref) + self.assertLessEqual(mx.max(diff).item(), atol) + + @unittest.skipIf(not mx.is_available(mx.gpu), "too slow on CPU") + @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "not enough memory") + def test_sdpa_long_masked_sequence(self): + # Test for int16 overflow in steel_attention_nax.h mask + # indexing (col_pos declared as short, overflows when kL > 32767). + D = 64 + dtype = mx.float16 + atol = 1e-3 # Slightly looser than test_sdpa due to long masked sequences + + for kL, active in [ + (8192, 1024), + (36864, 1024), + (49152, 1024), + (66048, 1024), + ]: + with self.subTest(kL=kL, active=active): + mx.random.seed(0) + qH, kH, qL = 32, 16, 512 + scale = 1.0 / math.sqrt(D) + + q = mx.random.normal(shape=(1, qH, qL, D)).astype(dtype) + k = mx.random.normal(shape=(1, kH, kL, D)).astype(dtype) + v = mx.random.normal(shape=(1, kH, kL, D)).astype(dtype) + + # Additive mask: -1e4 for inactive, 0 for last `active` positions + mask = mx.full((1, 1, 1, kL), -1e4, dtype=dtype) + mask[..., kL - active :] = 0.0 + + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + + self.assertFalse(mx.isnan(out).any().item()) + self.assertListEqual(list(out.shape), list(ref.shape)) + + diff = mx.abs(out - ref) - atol * mx.abs(ref) + self.assertLessEqual(mx.max(diff).item(), atol) def test_sdpa_broadcast_mask(self): mask = mx.array(True) @@ -738,19 +737,24 @@ def test_sdpa_attention_sinks(self): with self.assertRaises(ValueError): mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, sinks=sinks) - for T_kv in [128, 4096]: - for T_q in [1, 128]: - for N_kv in [2, 8]: - q = mx.random.normal(shape=(B, N_q, T_q, D)) - k = mx.random.normal(shape=(B, N_kv, T_kv, D)) - v = mx.random.normal(shape=(B, N_kv, T_kv, D)) - sinks = 10 * mx.random.normal(shape=(N_q,)) - - expected = mlx_ref_attn(q, k, v, scale, sinks=sinks) - out = mx.fast.scaled_dot_product_attention( - q, k, v, scale=scale, sinks=sinks - ) - self.assertTrue(mx.allclose(out, expected, atol=1e-5)) + for T_q, T_kv, N_kv, dtype in product( + (1, 128), + (128, 4096), + (2, 8), + (mx.float16, mx.float32), + ): + with self.subTest(T_q=T_q, T_kv=T_kv, N_kv=N_kv, dtype=dtype): + q = mx.random.normal(shape=(B, N_q, T_q, D), dtype=dtype) + k = mx.random.normal(shape=(B, N_kv, T_kv, D), dtype=dtype) + v = mx.random.normal(shape=(B, N_kv, T_kv, D), dtype=dtype) + sinks = 10 * mx.random.normal(shape=(N_q,), dtype=dtype) + + expected = mlx_ref_attn(q, k, v, scale, sinks=sinks) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, sinks=sinks + ) + atol = 1e-5 if dtype == mx.float32 else 1e-2 + self.assertTrue(mx.allclose(out, expected, atol=atol)) def test_sdpa_grad(self): # High tolerance due to cuDNN SDPA kernel requiring tf32. @@ -799,6 +803,132 @@ def test_grad(slow, fast, args): ).sum() test_grad(loss_slow, loss_fast, [q, k, v]) + @unittest.skipIf(not mx.metal.is_available(), "Metal kernel path only") + def test_sdpa_force_fused_metal(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(qL, kL, D, qH=8, kH=8): + q = mx.random.normal((1, qH, qL, D), mx.float16) + k = mx.random.normal((1, kH, kL, D), mx.float16) + v = mx.random.normal((1, kH, kL, D), mx.float16) + return q, k, v + + # Full attention kernel. + for D, qL, mask in product((192, 256), (9, 16), (None, "causal")): + with self.subTest(head_dim=D, qL=qL, mask=mask): + q, k, v = make_qkv(qL, 512, D, 8, 4) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + # Vector attention kernel. + for D in (192, 256): + with self.subTest(head_dim=D): + q, k, v = make_qkv(4, 16385, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + # No full attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(16, 512, 512) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=512**-0.5, force_fused=True + ) + with self.assertRaisesRegex( + ValueError, "query sequence to be no longer than the key sequence" + ): + q, k, v = make_qkv(32, 16, 64) + mx.fast.scaled_dot_product_attention( + q, + k, + v, + scale=64**-0.5, + mask="causal", + force_fused=True, + ) + + # No vector attention fused kernels. + with self.assertRaisesRegex(ValueError, "supports head dims"): + q, k, v = make_qkv(1, 128, 72) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=72**-0.5, force_fused=True + ) + with self.assertRaisesRegex(ValueError, "GQA factor to be at most 32"): + q, k, v = make_qkv(8, 128, 64, qH=8, kH=1) + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + # No CPU fused kernel. + with mx.stream(mx.cpu): + q, k, v = make_qkv(8, 128, 8) + with self.assertRaisesRegex(ValueError, "require a GPU"): + mx.fast.scaled_dot_product_attention( + q, k, v, scale=64**-0.5, force_fused=True + ) + + @unittest.skipIf(not mx.cuda.is_available(), "CUDA kernel path only") + def test_sdpa_force_fused_cuda(self): + if mx.default_device() != mx.gpu: + self.skipTest("requires GPU") + + def make_qkv(qL, kL, D, qH=8, kH=8): + q = mx.random.normal((1, qH, qL, D), mx.float16) + k = mx.random.normal((1, kH, kL, D), mx.float16) + v = mx.random.normal((1, kH, kL, D), mx.float16) + return q, k, v + + # Vector attention kernel. + for D in (64, 96, 128): + with self.subTest(head_dim=D): + q, k, v = make_qkv(3, 128, D, 4, 2) + scale = D**-0.5 + ref = mlx_ref_attn(q, k, v, scale=scale) + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, force_fused=True + ) + self.assertTrue(mx.allclose(ref, out, atol=1e-3, rtol=1e-3)) + + def test_sdpa_sliced(self): + N = 8 + D = 64 + scale = D**-0.5 + + for B, T_q, T_kv, offset, mask in product( + (1, 2, 4), + (1, 8), + (256, 512), + (8, 9, 64, 79), + (None, "causal"), + ): + with self.subTest(B=B, T_q=T_q, T_kv=T_kv, offset=offset, mask=mask): + q = mx.random.normal((B, N, T_q, D), mx.float16) + k = mx.random.normal((B, N, T_kv, D), mx.float16) + v = mx.random.normal((B, N, T_kv, D), mx.float16) + + k = k[..., :offset, :] + v = v[..., :offset, :] + + ref = mlx_ref_attn(q, k, v, scale=scale, mask=mask) + + for i in range(2): + out = mx.fast.scaled_dot_product_attention( + q, k, v, scale=scale, mask=mask + ) + if B == 1: + tolerance = {"rtol": 1e-3, "atol": 1e-3} + else: + tolerance = {"rtol": 1e-2, "atol": 1e-2} + self.assertTrue(mx.allclose(ref, out, **tolerance)) + if __name__ == "__main__": mlx_tests.MLXTestRunner(failfast=True) diff --git a/Source/Cmlx/mlx/python/tests/test_fft.py b/Source/Cmlx/mlx/python/tests/test_fft.py index 07ab62672..1f96aad56 100644 --- a/Source/Cmlx/mlx/python/tests/test_fft.py +++ b/Source/Cmlx/mlx/python/tests/test_fft.py @@ -91,6 +91,44 @@ def test_fftn(self): np_op = getattr(np.fft, op) self.check_mx_np(mx_op, np_op, x, axes=ax, s=s) + # Explicitly exercise transposed layouts and axes that are not + # physically last in memory order. + xt = np.transpose(a, (1, 2, 0)) + self.check_mx_np(mx.fft.fftn, np.fft.fftn, xt, axes=(2, 0)) + self.check_mx_np(mx.fft.ifftn, np.fft.ifftn, xt, axes=(2, 0)) + + rt = np.transpose(r, (1, 2, 0)) + self.check_mx_np(mx.fft.rfftn, np.fft.rfftn, rt, axes=(2, 0)) + irfft_in = np.ascontiguousarray(np.fft.rfftn(rt, axes=(2, 0))) + self.check_mx_np(mx.fft.irfftn, np.fft.irfftn, irfft_in, axes=(2, 0)) + + def test_fft_norm(self): + norms = ["backward", "ortho", "forward"] + + r = np.random.randn(8, 6).astype(np.float32) + i = np.random.randn(8, 6).astype(np.float32) + c = r + 1j * i + + for norm in norms: + self.check_mx_np(mx.fft.fft, np.fft.fft, c, axis=1, norm=norm) + self.check_mx_np(mx.fft.ifft, np.fft.ifft, c, axis=1, norm=norm) + self.check_mx_np(mx.fft.rfft, np.fft.rfft, r, axis=1, norm=norm) + + cr = np.fft.rfft(r, axis=1) + self.check_mx_np(mx.fft.irfft, np.fft.irfft, cr, axis=1, norm=norm) + + self.check_mx_np(mx.fft.fft2, np.fft.fft2, c, axes=(0, 1), norm=norm) + self.check_mx_np(mx.fft.ifft2, np.fft.ifft2, c, axes=(0, 1), norm=norm) + self.check_mx_np(mx.fft.fftn, np.fft.fftn, c, axes=(0, 1), norm=norm) + self.check_mx_np(mx.fft.ifftn, np.fft.ifftn, c, axes=(0, 1), norm=norm) + + self.check_mx_np(mx.fft.rfft2, np.fft.rfft2, r, axes=(0, 1), norm=norm) + self.check_mx_np(mx.fft.rfftn, np.fft.rfftn, r, axes=(0, 1), norm=norm) + + cr2 = np.fft.rfft2(r, axes=(0, 1)) + self.check_mx_np(mx.fft.irfft2, np.fft.irfft2, cr2, axes=(0, 1), norm=norm) + self.check_mx_np(mx.fft.irfftn, np.fft.irfftn, cr2, axes=(0, 1), norm=norm) + def _run_ffts(self, shape, atol=1e-4, rtol=1e-4): np.random.seed(9) @@ -130,6 +168,29 @@ def test_fft_shared_mem(self): atol = 1e-4 if num < 1025 else 1e-3 self._run_ffts((batch_size, num), atol=atol) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_batched_bluestein_twiddle_table(self): + for batch, n in ((1023, 1031), (1024, 1031), (1024, 1531)): + with self.subTest(batch=batch, n=n), mx.stream(mx.gpu): + index = mx.arange(n, dtype=mx.float32) + base = mx.cos(index * 0.017) + 0.25 * mx.sin(index * 0.031) + base = base + 1j * ( + mx.sin(index * 0.023) - 0.125 * mx.cos(index * 0.047) + ) + scale_index = mx.arange(batch, dtype=mx.float32) + scales = (0.75 + 0.25 * mx.cos(scale_index * 0.013)) * ( + mx.cos(scale_index * 0.019) + 1j * mx.sin(scale_index * 0.019) + ) + signal = scales[:, None] * base[None, :] + + for transform, atol in ((mx.fft.fft, 1e-3), (mx.fft.ifft, 1e-4)): + expected = scales[:, None] * transform(base)[None, :] + output = transform(signal) + mx.eval(expected, output) + self.assertTrue( + mx.allclose(output, expected, atol=atol, rtol=1e-4).item() + ) + @unittest.skip("Too slow for CI but useful for local testing.") def test_fft_exhaustive(self): nums = range(2, 4097) @@ -147,6 +208,21 @@ def test_fft_big_powers_of_two(self): for k in range(17, 20): self._run_ffts((3, 2**k), atol=1e-2) + # Past 2**20 the four step plan has to grow n2 to keep both factors + # inside threadgroup memory + for k in range(20, 25): + self._run_ffts((1, 2**k), atol=1e-2, rtol=1e-3) + + @unittest.skipIf( + not mx.metal.is_available(), "the size limit is specific to the Metal FFT plan" + ) + def test_fft_too_large(self): + # Larger than the four step plan can decompose, so it has to throw + # rather than run a kernel that silently returns the wrong answer. + # CUDA hands this to cuFFT instead and has no such limit. + with self.assertRaises(RuntimeError): + mx.eval(mx.fft.fft(mx.zeros(2**25))) + def test_fft_large_numbers(self): numbers = [ 1037, # prime > 2048 @@ -205,6 +281,50 @@ def test_fft_throws(self): x = mx.array(3.0) with self.assertRaises(ValueError): mx.fft.irfftn(x) + with self.assertRaises(ValueError): + mx.fft.fft(mx.array([1.0]), norm="invalid") + + def test_fftfreq(self): + for n, d in [(1, 1.0), (4, 0.5), (5, 0.25), (8, -0.5), (6, 1.0)]: + out = mx.fft.fftfreq(n, d=d) + expected = np.fft.fftfreq(n, d=d).astype(np.float32) + self.assertEqual(out.dtype, mx.float32) + np.testing.assert_allclose(out, expected, atol=0.0, rtol=0.0) + + with self.assertRaises(ValueError): + mx.fft.fftfreq(0) + + with self.assertRaises(ValueError): + mx.fft.fftfreq(-1) + + # Test default d=1.0 + out = mx.fft.fftfreq(8) + expected = np.fft.fftfreq(8).astype(np.float32) + np.testing.assert_allclose(out, expected, atol=0.0, rtol=0.0) + + with self.assertRaises(ValueError): + mx.fft.fftfreq(4, d=0.0) + + def test_rfftfreq(self): + for n, d in [(1, 1.0), (4, 0.5), (5, 0.25), (8, -0.5), (6, 1.0)]: + out = mx.fft.rfftfreq(n, d=d) + expected = np.fft.rfftfreq(n, d=d).astype(np.float32) + self.assertEqual(out.dtype, mx.float32) + np.testing.assert_allclose(out, expected, atol=0.0, rtol=0.0) + + # Test default d=1.0 + out = mx.fft.rfftfreq(8) + expected = np.fft.rfftfreq(8).astype(np.float32) + np.testing.assert_allclose(out, expected, atol=0.0, rtol=0.0) + + with self.assertRaises(ValueError): + mx.fft.rfftfreq(0) + + with self.assertRaises(ValueError): + mx.fft.rfftfreq(-1) + + with self.assertRaises(ValueError): + mx.fft.rfftfreq(4, d=0.0) def test_fftshift(self): # Test 1D arrays @@ -213,11 +333,14 @@ def test_fftshift(self): # Test with specific axis r = np.random.rand(4, 6).astype(np.float32) + self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=0) + self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=1) self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[0]) self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[1]) self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[0, 1]) # Test with negative axes + self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=-1) self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[-1]) # Test with odd lengths @@ -238,11 +361,14 @@ def test_ifftshift(self): # Test with specific axis r = np.random.rand(4, 6).astype(np.float32) + self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=0) + self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=1) self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[0]) self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[1]) self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[0, 1]) # Test with negative axes + self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=-1) self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[-1]) # Test with odd lengths @@ -310,14 +436,51 @@ def g(x): x = mx.random.normal(sh) else: x = mx.random.normal((*sh, 2)).view(mx.complex64).squeeze() + mx.eval(x) + x_torch = torch.tensor(x, device="cpu") fx = f(x) - gx = g(torch.tensor(x)) + gx = g(x_torch) self.assertLess((fx - gx).abs().max() / gx.abs().mean(), 1e-4) dfdx = mx.grad(f)(x) - dgdx = torch.func.grad(g)(torch.tensor(x)) + dgdx = torch.func.grad(g)(x_torch) self.assertLess((dfdx - dgdx).abs().max() / dgdx.abs().mean(), 1e-4) + def make_ffts(self): + mxffts = { + (True, True): mx.fft.irfftn, + (True, False): mx.fft.rfftn, + (False, True): mx.fft.ifftn, + (False, False): mx.fft.fftn, + } + shape = (3, 8, 6) + r = np.random.rand(*shape).astype(np.float32) + i = np.random.rand(*shape).astype(np.float32) + for (real, inverse), fftn in mxffts.items(): + a_np = r if real and not inverse else r + 1j * i + for axes in [(-1,), (0,), (-2, -1), (-1, -2), (0, 1)]: + yield fftn, a_np, axes + + def test_fft_vmap(self): + for fftn, a_np, axes in self.make_ffts(): + a = mx.array(a_np) + f = lambda x: fftn(x, axes=axes) + expected = mx.stack([f(a[i]) for i in range(a.shape[0])]) + out = mx.vmap(f)(a) + self.assertEqual(tuple(out.shape), tuple(expected.shape)) + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + + def test_fft_jvp(self): + # The fft is linear so the jvp is the fft of the tangent + for fftn, a_np, axes in self.make_ffts(): + a = mx.array(a_np) + t = mx.array(np.random.rand(*a_np.shape).astype(a_np.dtype)) + f = lambda x: fftn(x, axes=axes) + expected = f(t) + out = mx.jvp(f, [a], [t])[1][0] + self.assertEqual(tuple(out.shape), tuple(expected.shape)) + np.testing.assert_allclose(out, expected, atol=1e-5, rtol=1e-5) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_init.py b/Source/Cmlx/mlx/python/tests/test_init.py index 046a6e836..bd8d9fa4f 100644 --- a/Source/Cmlx/mlx/python/tests/test_init.py +++ b/Source/Cmlx/mlx/python/tests/test_init.py @@ -1,4 +1,5 @@ # Copyright © 2023 Apple Inc. +import math import unittest import mlx.core as mx @@ -106,6 +107,19 @@ def test_sparse(self): with self.assertRaises(ValueError): result = initializer(mx.zeros((1,))) + def test_sparse_zeros_per_row(self): + # Sparsity is applied along each row: every row drops exactly + # ceil(sparsity * cols) of its entries, independent of matrix shape, so + # each output feature (a row of `w` used as `x @ w.T`) keeps the same + # number of input connections. + for sparsity, shape in [(0.5, (4, 10)), (0.3, (5, 10)), (0.5, (2, 2))]: + _, cols = shape + expected = int(math.ceil(sparsity * cols)) + result = init.sparse(sparsity)(mx.zeros(shape)) + zeros_per_row = mx.sum(result == 0, axis=1) + with self.subTest(shape=shape, sparsity=sparsity): + self.assertTrue(mx.all(zeros_per_row == expected).item()) + def test_orthogonal(self): initializer = init.orthogonal(gain=1.0, dtype=mx.float32) @@ -134,6 +148,11 @@ def test_orthogonal(self): "Orthogonal init failed on a rectangular matrix.", ) + # A non-2D input reports its actual number of dimensions + with self.assertRaises(ValueError) as cm: + initializer(mx.zeros((2, 3, 4), dtype=mx.float32)) + self.assertIn("3D array", str(cm.exception)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_linalg.py b/Source/Cmlx/mlx/python/tests/test_linalg.py index 8e2444f20..52e02e4d9 100644 --- a/Source/Cmlx/mlx/python/tests/test_linalg.py +++ b/Source/Cmlx/mlx/python/tests/test_linalg.py @@ -38,6 +38,7 @@ def test_norm(self): with self.subTest( shape=shape, ord=o, axis=axis, keepdims=keepdims ): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -51,6 +52,7 @@ def test_norm(self): out_np = np.linalg.norm(x_np, ord=o, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, ord=o, keepdims=keepdims) with self.subTest(shape=shape, ord=o, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue( np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) ) @@ -63,8 +65,29 @@ def test_norm(self): out_np = np.linalg.norm(x_np, keepdims=keepdims) out_mx = mx.linalg.norm(x_mx, keepdims=keepdims) with self.subTest(shape=shape, keepdims=keepdims): + self.assertEqual(out_mx.shape, out_np.shape) self.assertTrue(np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6)) + # tests for negative indexing: -1/1/inf/-inf/ + norms = [-1, 1, -float("inf"), float("inf")] + for shape in [(3, 3), (2, 3, 3), (2, 3, 3, 3)]: + x_mx = mx.arange(1, math.prod(shape) + 1, dtype=mx.float32).reshape(shape) + x_np = np.arange(1, math.prod(shape) + 1, dtype=np.float32).reshape(shape) + neg_indices = [-i for i in range(1, x_np.ndim + 1)] + neg_axes = [list(p) for p in itertools.permutations(neg_indices, 2)] + for ord in norms: + for axes in neg_axes: + out_np = np.linalg.norm( + x_np, + ord=ord, + axis=tuple(axes), + ) + out_mx = mx.linalg.norm(x_mx, ord=ord, axis=axes) + with self.subTest(ord=ord, axes=axes): + self.assertTrue( + np.allclose(out_np, out_mx, atol=1e-5, rtol=1e-6) + ) + def test_complex_norm(self): for shape in [(3,), (2, 3), (2, 3, 3)]: x_np = np.random.uniform(size=shape).astype( @@ -134,6 +157,16 @@ def test_qr_factorization(self): mx.allclose(out, mx.eye(min(A.shape)), rtol=1e-4, atol=1e-6) ) + # Zero-size inputs. Both factors carry min(M, N) as a dimension, so + # both are empty whichever dimension is zero. + for shape in [(0, 0), (3, 0, 0), (0, 4, 4), (5, 0), (0, 5)]: + A_np = np.zeros(shape, dtype=np.float32) + Q, R = mx.linalg.qr(mx.array(A_np), stream=mx.cpu) + mx.eval(Q, R) + Q_np, R_np = np.linalg.qr(A_np) + self.assertEqual(Q.shape, Q_np.shape) + self.assertEqual(R.shape, R_np.shape) + def test_svd_decomposition(self): A = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float32) U, S, Vt = mx.linalg.svd(A, compute_uv=True, stream=mx.cpu) @@ -168,6 +201,23 @@ def test_svd_decomposition(self): ) ) + # Zero-size inputs. When only one of the dimensions is zero the + # factors are not empty and hold the identity, like numpy. + for shape in [(0, 4, 4), (3, 0, 0), (2, 5, 0), (2, 0, 5), (5, 0), (0, 5)]: + a_np = np.zeros(shape, dtype=np.float32) + U, S, Vt = mx.linalg.svd(mx.array(a_np), stream=mx.cpu) + mx.eval(U, S, Vt) + U_np, S_np, Vt_np = np.linalg.svd(a_np) + self.assertEqual(U.shape, U_np.shape) + self.assertEqual(S.shape, S_np.shape) + self.assertEqual(Vt.shape, Vt_np.shape) + self.assertTrue(np.array_equal(np.array(U), U_np)) + self.assertTrue(np.array_equal(np.array(Vt), Vt_np)) + + S_only = mx.linalg.svd(mx.array(a_np), compute_uv=False, stream=mx.cpu) + mx.eval(S_only) + self.assertEqual(S_only.shape, S_np.shape) + # Test float64 - use CPU stream since float64 is not supported on GPU with mx.stream(mx.cpu): A_f64 = mx.array( @@ -273,6 +323,13 @@ def test_pseudo_inverse(self): A_plus = mx.linalg.pinv(A, stream=mx.cpu) self.assertTrue(mx.allclose(A @ A_plus @ A, A)) + # Zero-size inputs. The result takes the shape of the transposed input. + for shape in [(0, 0), (0, 3), (3, 0), (0, 2, 2), (2, 0, 0), (0, 4, 3)]: + A_np = np.zeros(shape, dtype=np.float32) + A_plus = mx.linalg.pinv(mx.array(A_np), stream=mx.cpu) + mx.eval(A_plus) + self.assertEqual(A_plus.shape, np.linalg.pinv(A_np).shape) + def test_cholesky_inv(self): mx.random.seed(7) @@ -415,6 +472,19 @@ def check_eigs_and_vecs(A_np, kwargs={}): self.assertEqual(eig_vals_c64.dtype, mx.complex64) self.assertEqual(eig_vecs_c64.dtype, mx.complex64) + # Zero-size inputs. The input is square, so both outputs are empty. + for shape in [(0, 0), (3, 0, 0), (0, 4, 4)]: + A_np = np.zeros(shape, dtype=np.float32) + eig_vals, eig_vecs = mx.linalg.eig(mx.array(A_np), stream=mx.cpu) + mx.eval(eig_vals, eig_vecs) + vals_np, vecs_np = np.linalg.eig(A_np) + self.assertEqual(eig_vals.shape, vals_np.shape) + self.assertEqual(eig_vecs.shape, vecs_np.shape) + + vals_only = mx.linalg.eigvals(mx.array(A_np), stream=mx.cpu) + mx.eval(vals_only) + self.assertEqual(vals_only.shape, vals_np.shape) + # Test error cases with self.assertRaises(ValueError): mx.linalg.eig(mx.array([1.0, 2.0])) # 1D array @@ -473,6 +543,29 @@ def check_eigs_and_vecs(A_np, kwargs={}): A_np = A_np + A_np.T.conj() check_eigs_and_vecs(A_np) + # UPLO picks the triangle like numpy; only observable when the two + # triangles disagree + A_np = np.array([[1.0, 999.0], [2.0, 3.0]], dtype=np.float32) + for uplo in ("L", "U"): + w = mx.linalg.eigvalsh(mx.array(A_np), UPLO=uplo, stream=mx.cpu) + w_np = np.linalg.eigvalsh(A_np, UPLO=uplo) + self.assertTrue(np.allclose(w, w_np, atol=1e-5)) + + # Zero-size inputs + for shape in [(0, 4, 4), (3, 0, 0), (0, 0)]: + a_np = np.zeros(shape, dtype=np.float32) + w, v = mx.linalg.eigh(mx.array(a_np), stream=mx.cpu) + mx.eval(w, v) + w_np, v_np = np.linalg.eigh(a_np) + self.assertEqual(w.shape, w_np.shape) + self.assertEqual(v.shape, v_np.shape) + self.assertTrue(np.array_equal(np.array(w), w_np)) + self.assertTrue(np.array_equal(np.array(v), v_np)) + + w_only = mx.linalg.eigvalsh(mx.array(a_np), stream=mx.cpu) + mx.eval(w_only) + self.assertEqual(w_only.shape, w_np.shape) + # Test error cases with self.assertRaises(ValueError): mx.linalg.eigh(mx.array([1.0, 2.0])) # 1D array @@ -520,6 +613,19 @@ def test_lu(self): P, L, U = mx.linalg.lu(a, stream=mx.cpu) self.assertTrue(mx.allclose(L[P, :] @ U, a)) + # Test singular matrix (should not throw) + a = mx.array( + [ + [1.0, 2.0, 3.0, 4.0], + [2.0, 4.0, 6.0, 8.0], + [0.0, 1.0, 1.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ] + ) + P, L, U = mx.linalg.lu(a, stream=mx.cpu) + L_permuted = mx.take_along_axis(L, P[..., None], axis=-2) + self.assertTrue(mx.allclose(L_permuted @ U, a)) + def test_lu_factor(self): mx.random.seed(7) @@ -616,6 +722,248 @@ def test_solve_triangular(self): expected = np.linalg.solve(a, b) self.assertTrue(np.allclose(result, expected)) + def test_det(self): + # 1x1 fast path + A = mx.array([[5.0]]) + self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 5.0)) + + # 2x2 fast path + A = mx.array([[1.0, 2.0], [3.0, 4.0]]) + d = mx.linalg.det(A, stream=mx.cpu) + self.assertTrue(np.allclose(d, -2.0)) + + # 3x3 fast path + A = mx.array([[1.0, 2.0, 3.0], [0.0, 1.0, 4.0], [5.0, 6.0, 0.0]]) + d = mx.linalg.det(A, stream=mx.cpu) + expected = np.linalg.det(np.array(A)) + self.assertTrue(np.allclose(d, expected, atol=1e-5)) + + # 4x4 LU path: compare with numpy + np.random.seed(42) + A_np = np.random.randn(4, 4).astype(np.float32) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4)) + + # 5x5 LU path + A_np = np.random.randn(5, 5).astype(np.float32) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4)) + + # Identity matrix + A = mx.eye(5) + self.assertTrue(np.allclose(mx.linalg.det(A, stream=mx.cpu), 1.0)) + + # Batched: (3, 4, 4) + A_np = np.random.randn(3, 4, 4).astype(np.float32) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4)) + + # Multi-batch: (2, 3, 3, 3) + A_np = np.random.randn(2, 3, 3, 3).astype(np.float32) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-4)) + + # Integer input auto-promotes to float + A = mx.array([[1, 2], [3, 4]]) + d = mx.linalg.det(A, stream=mx.cpu) + self.assertTrue(np.allclose(d, -2.0)) + + # float64 + A_np = np.random.randn(4, 4).astype(np.float64) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-10)) + + # Singular 4x4 matrix (LU path): det should be 0 + A = mx.array( + [ + [1.0, 2.0, 3.0, 4.0], + [2.0, 4.0, 6.0, 8.0], + [0.0, 1.0, 1.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ] + ) + d = mx.linalg.det(A, stream=mx.cpu) + self.assertTrue(np.allclose(d, 0.0, atol=1e-5)) + + # Singular 5x5 matrix (LU path) + A_np = np.ones((5, 5), dtype=np.float32) + A_mx = mx.array(A_np) + d = mx.linalg.det(A_mx, stream=mx.cpu) + self.assertTrue(np.allclose(d, 0.0, atol=1e-5)) + + # Batched singular matrices (LU path) + A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)]) + A_mx = mx.array(A_np) + d_mx = mx.linalg.det(A_mx, stream=mx.cpu) + d_np = np.linalg.det(A_np) + self.assertTrue(np.allclose(d_mx, d_np, atol=1e-5)) + + # Empty 0x0 matrix: det is the empty product = 1 + d = mx.linalg.det(mx.zeros((0, 0)), stream=mx.cpu) + self.assertEqual(d.shape, ()) + self.assertEqual(float(d), 1.0) + + # Batched empty matrices: shape preserves batch dims + d = mx.linalg.det(mx.zeros((3, 0, 0)), stream=mx.cpu) + self.assertTrue(np.allclose(d, np.linalg.det(np.zeros((3, 0, 0))))) + + # Error: non-square + with self.assertRaises(ValueError): + mx.linalg.det(mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu) + + # Error: 1D + with self.assertRaises(ValueError): + mx.linalg.det(mx.array([1.0, 2.0]), stream=mx.cpu) + + # Error: complex unsupported (small-matrix path) + with self.assertRaises(ValueError): + mx.linalg.det(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu) + + # Error: complex unsupported (LU path) + with self.assertRaises(ValueError): + mx.linalg.det(mx.eye(4).astype(mx.complex64), stream=mx.cpu) + + def test_slogdet(self): + # 2x2: det = -2 => sign = -1, logabsdet = log(2) + A = mx.array([[1.0, 2.0], [3.0, 4.0]]) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertTrue(np.allclose(sign, -1.0)) + self.assertTrue(np.allclose(logabsdet, np.log(2.0), atol=1e-5)) + + # Identity: sign = 1, logabsdet = 0 + A = mx.eye(4) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertTrue(np.allclose(sign, 1.0)) + self.assertTrue(np.allclose(logabsdet, 0.0, atol=1e-6)) + + # Compare with numpy for random matrices + np.random.seed(42) + for n in [1, 2, 3, 4, 5]: + A_np = np.random.randn(n, n).astype(np.float32) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(A_np) + with self.subTest(n=n): + self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5)) + self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4)) + + # Singular matrix 2x2 (fast path): sign = 0, logabsdet = -inf + A = mx.array([[1.0, 2.0], [2.0, 4.0]]) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertEqual(float(sign), 0.0) + self.assertEqual(float(logabsdet), float("-inf")) + + # Singular 4x4 matrix (LU path): sign = 0, logabsdet = -inf + A = mx.array( + [ + [1.0, 2.0, 3.0, 4.0], + [2.0, 4.0, 6.0, 8.0], + [0.0, 1.0, 1.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ] + ) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertEqual(float(sign), 0.0) + self.assertEqual(float(logabsdet), float("-inf")) + + # Singular 5x5 matrix (LU path): all-ones matrix + A = mx.array(np.ones((5, 5), dtype=np.float32)) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertEqual(float(sign), 0.0) + self.assertEqual(float(logabsdet), float("-inf")) + + # Batched with mix of singular and non-singular (LU path) + A_np = np.array([np.diag([1.0, 2.0, 0.0, 3.0]), np.eye(4, dtype=np.float32)]) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(A_np) + self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5)) + # Check -inf for singular, 0.0 for identity + self.assertEqual(float(logabs_mx[0]), float("-inf")) + self.assertTrue(np.allclose(logabs_mx[1], 0.0, atol=1e-6)) + + # Batched + A_np = np.random.randn(3, 4, 4).astype(np.float32) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(A_np) + self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5)) + self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4)) + + # Multi-batch + A_np = np.random.randn(2, 3, 3, 3).astype(np.float32) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(A_np) + self.assertTrue(np.allclose(sign_mx, sign_np, atol=1e-5)) + self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-4)) + + # Numerical stability: large matrix where det overflows + # 0.1 * I_100 has det = 0.1^100 which underflows in float32 + # but slogdet should give sign=1, logabsdet = 100*log(0.1) + n = 100 + A = mx.array(0.1) * mx.eye(n) + sign, logabsdet = mx.linalg.slogdet(A, stream=mx.cpu) + self.assertTrue(np.allclose(sign, 1.0)) + self.assertTrue(np.allclose(logabsdet, n * np.log(0.1), atol=1e-3)) + + # Verify det = sign * exp(logabsdet) for non-singular cases + A_np = np.random.randn(5, 5).astype(np.float32) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + det_mx = mx.linalg.det(A_mx, stream=mx.cpu) + reconstructed = float(sign_mx) * np.exp(float(logabs_mx)) + self.assertTrue(np.allclose(float(det_mx), reconstructed, rtol=1e-4)) + + # float64 + A_np = np.random.randn(4, 4).astype(np.float64) + A_mx = mx.array(A_np) + sign_mx, logabs_mx = mx.linalg.slogdet(A_mx, stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(A_np) + self.assertTrue(np.allclose(sign_mx, sign_np)) + self.assertTrue(np.allclose(logabs_mx, logabs_np, atol=1e-10)) + + # Empty 0x0 matrix: sign = 1, logabsdet = 0 (empty product) + sign, logabsdet = mx.linalg.slogdet(mx.zeros((0, 0)), stream=mx.cpu) + self.assertEqual(sign.shape, ()) + self.assertEqual(logabsdet.shape, ()) + self.assertEqual(float(sign), 1.0) + self.assertEqual(float(logabsdet), 0.0) + + # Batched empty matrices + sign, logabsdet = mx.linalg.slogdet(mx.zeros((3, 0, 0)), stream=mx.cpu) + sign_np, logabs_np = np.linalg.slogdet(np.zeros((3, 0, 0))) + self.assertTrue(np.allclose(sign, sign_np)) + self.assertTrue(np.allclose(logabsdet, logabs_np)) + + # Error: non-square + with self.assertRaises(ValueError): + mx.linalg.slogdet( + mx.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stream=mx.cpu + ) + + # Error: 1D + with self.assertRaises(ValueError): + mx.linalg.slogdet(mx.array([1.0, 2.0]), stream=mx.cpu) + + # Error: complex unsupported (small-matrix path) + with self.assertRaises(ValueError): + mx.linalg.slogdet(mx.array([[1.0 + 1j, 2.0], [3.0, 4.0]]), stream=mx.cpu) + + # Error: complex unsupported (LU path) + with self.assertRaises(ValueError): + mx.linalg.slogdet(mx.eye(4).astype(mx.complex64), stream=mx.cpu) + if __name__ == "__main__": mlx_tests.MLXTestRunner() diff --git a/Source/Cmlx/mlx/python/tests/test_load.py b/Source/Cmlx/mlx/python/tests/test_load.py index 97fbf6c4c..f5947de47 100644 --- a/Source/Cmlx/mlx/python/tests/test_load.py +++ b/Source/Cmlx/mlx/python/tests/test_load.py @@ -2,6 +2,7 @@ import os import platform +import struct import tempfile import unittest from pathlib import Path @@ -87,6 +88,40 @@ def test_load_npy_dtype(self): with self.assertRaises(Exception): out = mx.load(save_file, stream=mx.cpu) + def test_load_npy_read_error(self): + save_file = os.path.join(self.test_dir, "truncated.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + out = mx.load(save_file, stream=mx.cpu) + with self.assertRaises(RuntimeError): + mx.eval(out) + + def test_async_load_npy_read_error_across_streams(self): + save_file = os.path.join(self.test_dir, "truncated_async.npy") + expected = np.arange(16, dtype=np.float32) + np.save(save_file, expected) + with open(save_file, "r+b") as f: + f.truncate(os.path.getsize(save_file) - expected.nbytes) + + producer_stream = mx.new_stream(mx.cpu) + consumer_stream = mx.new_stream(mx.cpu) + out = mx.add( + mx.load(save_file, stream=producer_stream), + 1.0, + stream=consumer_stream, + ) + with self.assertRaises(RuntimeError): + mx.eval(out) + # Depending on backend the error might be caught early before poisoning + # the producer_stream, but still sync to clear the errors. + try: + mx.synchronize(producer_stream) + except Exception: + pass + def test_save_and_load_safetensors(self): test_file = os.path.join(self.test_dir, "test.safetensors") with self.assertRaises(Exception): @@ -164,6 +199,87 @@ def test_save_and_load_gguf(self): self.assertTrue("test" in load_dict) self.assertTrue(mx.array_equal(load_dict["test"], save_dict["test"])) + @unittest.skipIf(platform.system() == "Windows", "GGUF is disabled on Windows") + def test_load_gguf_quantized(self): + # Write a minimal GGUF v3 file with one quantized tensor and check + # that the loaded (weight, scales, biases) triplet dequantizes to the + # values prescribed by the GGML block formats. + rows, cols = 2, 64 + n_blocks = rows * cols // 32 + np.random.seed(7) + + def write_gguf(path, ggml_type_id, blocks): + with open(path, "wb") as f: + f.write(b"GGUF") + # version, tensor count, metadata kv count + f.write(struct.pack("