diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 708479f79..fa9fa4daf 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -97,6 +97,56 @@ jobs: name: katago-macos-opencl path: cpp/katago + build-macos-coreml: + runs-on: macos-latest + permissions: + contents: read + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install dependencies + run: | + brew install ninja zlib libzip + brew tap chinchangyang/katagocoreml-cpp + brew install katagocoreml + + - name: Cache CMake build + uses: actions/cache@v4 + with: + path: | + cpp/CMakeCache.txt + cpp/CMakeFiles + cpp/build.ninja + cpp/.ninja_deps + cpp/.ninja_log + key: ${{ runner.os }}-cmake-coreml-${{ hashFiles('**/CMakeLists.txt') }} + restore-keys: | + ${{ runner.os }}-cmake-coreml- + + - name: Configure CMake + working-directory: cpp + run: | + cmake . -G Ninja -DUSE_BACKEND=COREML -DCMAKE_BUILD_TYPE=Release + + - name: Build + working-directory: cpp + run: | + ninja + + - name: Run tests + working-directory: cpp + run: | + ./katago runtests + + - name: Upload artifact + if: github.event_name == 'push' && github.ref == 'refs/heads/master' + uses: actions/upload-artifact@v4 + with: + name: katago-macos-coreml + path: cpp/katago + build-windows: runs-on: windows-latest permissions: diff --git a/.gitignore b/.gitignore index 2e933d553..b509ec1ff 100644 --- a/.gitignore +++ b/.gitignore @@ -90,3 +90,7 @@ cpp/.ninja_log cpp/build.ninja cpp/KataGoSwift.* cpp/include/KataGoSwift/KataGoSwift-swift.h + +# For CoreML Backend +cpp/KataGoCoreML.* +cpp/include/KataGoCoreML/KataGoCoreML-swift.h diff --git a/Compiling.md b/Compiling.md index 648fea548..80cd7df36 100644 --- a/Compiling.md +++ b/Compiling.md @@ -118,7 +118,7 @@ As also mentioned in the instructions below but repeated here for visibility, if * If using OpenCL, you will want to verify that KataGo is picking up the correct device (e.g. some systems may have both an Intel CPU OpenCL and GPU OpenCL, if KataGo appears to pick the wrong one, you can correct this by specifying `openclGpuToUse` in `configs/gtp_example.cfg`). ## MacOS - * TLDR: + * TLDR (Metal backend - recommended for most users): ``` git clone https://github.com/lightvector/KataGo.git cd KataGo/cpp @@ -127,11 +127,23 @@ As also mentioned in the instructions below but repeated here for visibility, if cmake -G Ninja -DUSE_BACKEND=METAL -DBUILD_DISTRIBUTED=1 ninja ``` + * TLDR (CoreML backend - hybrid CPU+GPU+Neural Engine for maximum throughput): + ``` + # First, install the katagocoreml library via Homebrew + brew tap chinchangyang/katagocoreml-cpp + brew install katagocoreml + + git clone https://github.com/lightvector/KataGo.git + cd KataGo/cpp + cmake -G Ninja -DUSE_BACKEND=COREML -DBUILD_DISTRIBUTED=1 + ninja + ``` * Requirements * [Homebrew](https://brew.sh): `/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"` * CMake with a minimum version of 3.18.2: `brew install cmake`. * AppleClang and Swift compilers: `xcode-select --install`. - * If using the Metal backend, [Ninja](https://ninja-build.org): `brew install ninja` + * If using the Metal or CoreML backend, [Ninja](https://ninja-build.org): `brew install ninja` + * If using the CoreML backend, katagocoreml library: `brew tap chinchangyang/katagocoreml-cpp && brew install katagocoreml` * libzip: `brew install libzip`. * If you want to do self-play training and research, probably Google perftools `brew install gperftools` for TCMalloc or some other better malloc implementation. For unknown reasons, the allocation pattern in self-play with large numbers of threads and parallel games causes a lot of memory fragmentation under glibc malloc that will eventually run your machine out of memory, but better mallocs handle it fine. * If compiling to contribute to public distributed training runs, OpenSSL is required (`brew install openssl`). @@ -139,14 +151,14 @@ As also mentioned in the instructions below but repeated here for visibility, if * `git clone https://github.com/lightvector/KataGo.git` * Compile using CMake and make in the cpp directory: * `cd KataGo/cpp` - * `cmake . -G Ninja -DUSE_BACKEND=METAL` or `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=EIGEN` depending on which backend you want. + * `cmake . -G Ninja -DUSE_BACKEND=METAL` or `cmake . -G Ninja -DUSE_BACKEND=COREML` or `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=EIGEN` depending on which backend you want. * Specify also `-DUSE_TCMALLOC=1` if using TCMalloc. * Compiling will also call git commands to embed the git hash into the compiled executable, specify also `-DNO_GIT_REVISION=1` to disable it if this is causing issues for you. * Specify `-DUSE_AVX2=1` to also compile Eigen with AVX2 and FMA support, which will make it incompatible with old CPUs but much faster. Intel-based Macs with new processors support AVX2, but Apple Silicon Macs do not support AVX2 natively. (If you want to go further, you can also add `-DCMAKE_CXX_FLAGS='-march=native'` which will specialize to precisely your machine's CPU, but the exe might not run on other machines at all). * Specify `-DBUILD_DISTRIBUTED=1` to compile with support for contributing data to public distributed training runs. * If building distributed, you will also need to build with Git revision support, including building within a clone of the repo, as opposed to merely an unzipped copy of its source. * Only builds from specific tagged versions or branches can contribute, in particular, instead of the `master` branch, use either the latest [release](https://github.com/lightvector/KataGo/releases) tag or the tip of the `stable` branch. To minimize the chance of any data incompatibilities or bugs, please do NOT attempt to contribute with custom changes or circumvent these limitations. - * `ninja` for Metal backend, or `make` for other backends. + * `ninja` for Metal or CoreML backend, or `make` for other backends. * Done! You should now have a compiled `katago` executable in your working directory. * Pre-trained neural nets are available at [the main training website](https://katagotraining.org/). * You will probably want to edit `configs/gtp_example.cfg` (see "Tuning for Performance" above). diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 254d23233..ab2b3dc69 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1,5 +1,5 @@ cmake_minimum_required(VERSION 3.18.2) -if(USE_BACKEND STREQUAL "METAL") +if(USE_BACKEND STREQUAL "METAL" OR USE_BACKEND STREQUAL "COREML") project(katago LANGUAGES CXX Swift) else() project(katago) @@ -32,7 +32,7 @@ endif() set(BUILD_DISTRIBUTED 0 CACHE BOOL "Build with http support for contributing to distributed training") set(USE_BACKEND CACHE STRING "Neural net backend") string(TOUPPER "${USE_BACKEND}" USE_BACKEND) -set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN) +set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL COREML) set(USE_TCMALLOC 0 CACHE BOOL "Use TCMalloc") set(NO_GIT_REVISION 0 CACHE BOOL "Disable embedding the git revision into the compiled exe") @@ -118,12 +118,42 @@ elseif(USE_BACKEND STREQUAL "METAL") neuralnet/metalbackend.swift) _swift_generate_cxx_header( KataGoSwift - "${CMAKE_CURRENT_BINARY_DIR}/include/KataGoSwift/KataGoSwift-swift.h" - SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/neuralnet/metalbackend.swift") + "${CMAKE_CURRENT_BINARY_DIR}/include/KataGoSwift/KataGoSwift-swift.h") target_include_directories(KataGoSwift PUBLIC "${CMAKE_CURRENT_BINARY_DIR}/include") set_target_properties(KataGoSwift PROPERTIES Swift_MODULE_NAME "KataGoSwift") target_compile_options(KataGoSwift PUBLIC "$<$:-cxx-interoperability-mode=default>") +elseif(USE_BACKEND STREQUAL "COREML") + message(STATUS "-DUSE_BACKEND=COREML, using Core ML backend.") + if(NOT "${CMAKE_GENERATOR}" STREQUAL "Ninja") + message(FATAL_ERROR "Bidirectional C++ Interop requires Ninja generator. Have ${CMAKE_GENERATOR}") + endif() + if("${CMAKE_Swift_COMPILER_VERSION}" VERSION_LESS 5.9) + message(FATAL_ERROR "Bidirectional C++ Interop requires Swift 5.9 or greater. Have ${CMAKE_Swift_COMPILER_VERSION}") + endif() + if(NOT "${CMAKE_CXX_COMPILER_ID}" STREQUAL "AppleClang") + message(FATAL_ERROR "Project requires building with AppleClang. Have ${CMAKE_CXX_COMPILER_ID}") + endif() + # Find katagocoreml library for native Core ML model conversion using pkg-config + find_package(PkgConfig REQUIRED) + pkg_check_modules(KATAGOCOREML REQUIRED katagocoreml) + list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/external/macos/cmake/modules") + include(InitializeSwift) + include(AddSwift) + set(CMAKE_OSX_DEPLOYMENT_TARGET 13.0) + set(NEURALNET_BACKEND_SOURCES + neuralnet/coremlbackend.cpp + ) + add_library(KataGoCoreML STATIC + neuralnet/coremlbackend.swift + neuralnet/mpsgraphlayers.swift) + _swift_generate_cxx_header( + KataGoCoreML + "${CMAKE_CURRENT_BINARY_DIR}/include/KataGoCoreML/KataGoCoreML-swift.h") + target_include_directories(KataGoCoreML PUBLIC "${CMAKE_CURRENT_BINARY_DIR}/include") + set_target_properties(KataGoCoreML PROPERTIES Swift_MODULE_NAME "KataGoCoreML") + target_compile_options(KataGoCoreML PUBLIC + "$<$:-cxx-interoperability-mode=default>") elseif(USE_BACKEND STREQUAL "OPENCL") message(STATUS "-DUSE_BACKEND=OPENCL, using OpenCL backend.") set(NEURALNET_BACKEND_SOURCES @@ -403,6 +433,19 @@ elseif(USE_BACKEND STREQUAL "METAL") if("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64") message(WARNING "You are currently running cmake on an Intel-based processor. It is known that running KataGo in this configuration may encounter performance issues. It is recommended to switch to a cmake version designed for ARM64 architecture for optimal performance.") endif() +elseif(USE_BACKEND STREQUAL "COREML") + target_compile_definitions(katago PRIVATE USE_COREML_BACKEND) + target_include_directories(katago PRIVATE ${KATAGOCOREML_INCLUDE_DIRS}) + # Link katagocoreml with full path and add all required library directories + find_library(KATAGOCOREML_LIB katagocoreml HINTS /usr/local/lib REQUIRED) + target_link_directories(katago PRIVATE ${KATAGOCOREML_LIBRARY_DIRS}) + # Link MetalPerformanceShadersGraph for hybrid MPSGraph + CoreML backend + target_link_libraries(katago KataGoCoreML ${KATAGOCOREML_LIB} ${KATAGOCOREML_LDFLAGS} + "-framework MetalPerformanceShaders" + "-framework MetalPerformanceShadersGraph") + if("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64") + message(WARNING "You are currently running cmake on an Intel-based processor. Core ML backend may not work optimally. ARM64 architecture is recommended.") + endif() elseif(USE_BACKEND STREQUAL "OPENCL") target_compile_definitions(katago PRIVATE USE_OPENCL_BACKEND) find_package(OpenCL) diff --git a/cpp/command/benchmark.cpp b/cpp/command/benchmark.cpp index 3100fb1b1..b29b9325b 100644 --- a/cpp/command/benchmark.cpp +++ b/cpp/command/benchmark.cpp @@ -260,6 +260,9 @@ int MainCmds::benchmark(const vector& args) { #ifdef USE_METAL_BACKEND cout << "You are currently using the Metal version of KataGo." << endl; #endif +#ifdef USE_COREML_BACKEND + cout << "You are currently using the Core ML version of KataGo." << endl; +#endif #ifdef USE_OPENCL_BACKEND cout << "You are currently using the OpenCL version of KataGo." << endl; cout << "If you have a strong GPU capable of FP16 tensor cores (e.g. RTX2080), " diff --git a/cpp/main.cpp b/cpp/main.cpp index 0fcc36dea..c15703cd8 100644 --- a/cpp/main.cpp +++ b/cpp/main.cpp @@ -244,6 +244,8 @@ string Version::getKataGoVersionFullInfo() { out << "Using TensorRT backend" << endl; #elif defined(USE_METAL_BACKEND) out << "Using Metal backend" << endl; +#elif defined(USE_COREML_BACKEND) + out << "Using Core ML backend" << endl; #elif defined(USE_OPENCL_BACKEND) out << "Using OpenCL backend" << endl; #elif defined(USE_EIGEN_BACKEND) @@ -280,6 +282,8 @@ string Version::getGitRevisionWithBackend() { s += "-trt"; #elif defined(USE_METAL_BACKEND) s += "-metal"; +#elif defined(USE_COREML_BACKEND) + s += "-coreml"; #elif defined(USE_OPENCL_BACKEND) s += "-opencl"; #elif defined(USE_EIGEN_BACKEND) diff --git a/cpp/neuralnet/coremlbackend.cpp b/cpp/neuralnet/coremlbackend.cpp new file mode 100644 index 000000000..faf38363b --- /dev/null +++ b/cpp/neuralnet/coremlbackend.cpp @@ -0,0 +1,1059 @@ +#ifdef USE_COREML_BACKEND + +#include "../neuralnet/modelversion.h" +#include "../neuralnet/nneval.h" +#include "../neuralnet/nninputs.h" +#include "../neuralnet/nninterface.h" +#include "../neuralnet/coremlbackend.h" + +#include +#include +#include +#include +#include +#include + +using namespace std; + +//------------------------------------------------------------------------------ +// CoreML Model Conversion - Native C++ using katagocoreml library +//------------------------------------------------------------------------------ + +namespace CoreMLConversion { + +// Get cache directory: ~/Documents/KataGo/CoreMLModels/ +static string getCacheDirectory() { + const char* homeDir = getenv("HOME"); + if(homeDir == nullptr) { + struct passwd* pw = getpwuid(getuid()); + homeDir = pw ? pw->pw_dir : "/tmp"; + } + return string(homeDir) + "/Documents/KataGo/CoreMLModels"; +} + +// Create directory if it doesn't exist +static bool createDirectoryIfNeeded(const string& path) { + struct stat st; + if(stat(path.c_str(), &st) == 0) { + return S_ISDIR(st.st_mode); + } + // Create parent directories recursively + size_t pos = path.find_last_of('/'); + if(pos != string::npos && pos > 0) { + createDirectoryIfNeeded(path.substr(0, pos)); + } + return mkdir(path.c_str(), 0755) == 0; +} + +// Check if file/directory exists +static bool fileExists(const string& path) { + struct stat st; + return stat(path.c_str(), &st) == 0; +} + +// Generate cache key with _native suffix to distinguish from Python-converted models +static string generateCacheKey( + const string& modelSHA256, + int boardX, + int boardY, + bool useFP16, + bool optimizeMask +) { + string precisionSuffix = useFP16 ? "fp16" : "fp32"; + string maskSuffix = optimizeMask ? "nomask" : "mask"; + return modelSHA256 + "_" + to_string(boardX) + "x" + to_string(boardY) + + "_" + precisionSuffix + "_" + maskSuffix + "_native"; +} + +// Convert KataGo model to CoreML using native katagocoreml library +static void convertModelToCoreML( + const string& modelPath, + const string& outputPath, + int boardXSize, + int boardYSize, + bool useFP16, + bool optimizeIdentityMask +) { + katagocoreml::ConversionOptions opts; + opts.board_x_size = boardXSize; + opts.board_y_size = boardYSize; + opts.compute_precision = useFP16 ? "FLOAT16" : "FLOAT32"; + opts.optimize_identity_mask = optimizeIdentityMask; + opts.specification_version = 8; // iOS 17+ / macOS 14+ + + katagocoreml::KataGoConverter::convert(modelPath, outputPath, opts); +} + +// Ensure model is converted and cached, returns path to .mlpackage +static string ensureModelConverted( + const string& modelPath, + const string& modelSHA256, + int boardX, + int boardY, + bool useFP16, + bool optimizeMask, + int serverThreadIdx +) { + string cacheDir = getCacheDirectory(); + string cacheKey = generateCacheKey(modelSHA256, boardX, boardY, useFP16, optimizeMask); + string mlpackagePath = cacheDir + "/" + cacheKey + ".mlpackage"; + + // Check if already cached + if(fileExists(mlpackagePath)) { + cerr << "Core ML backend " << serverThreadIdx << ": Using cached model at " << mlpackagePath << endl; + return mlpackagePath; + } + + // Create cache directory if needed + if(!createDirectoryIfNeeded(cacheDir)) { + throw runtime_error("Failed to create cache directory: " + cacheDir); + } + + // Convert model + cerr << "Core ML backend " << serverThreadIdx << ": Converting model to " << mlpackagePath << endl; + try { + convertModelToCoreML(modelPath, mlpackagePath, boardX, boardY, useFP16, optimizeMask); + cerr << "Core ML backend " << serverThreadIdx << ": Conversion completed successfully" << endl; + } catch(const exception& e) { + throw runtime_error(string("Core ML model conversion failed: ") + e.what()); + } + + return mlpackagePath; +} + +} // namespace CoreMLConversion + +//------------------------------------------------------------------------------ +// Model Descriptor Conversion - C++ to Swift types for MPSGraph +//------------------------------------------------------------------------------ + +namespace CoreMLProcess { + +/// Converts a ConvLayerDesc instance from C++ to Swift +SWConvLayerDesc convLayerDescToSwift(const ConvLayerDesc* desc) { + return createSWConvLayerDesc( + desc->convYSize, + desc->convXSize, + desc->inChannels, + desc->outChannels, + desc->dilationY, + desc->dilationX, + (float*)desc->weights.data()); +} + +/// Converts a BatchNormLayerDesc instance from C++ to Swift +SWBatchNormLayerDesc batchNormLayerDescToSwift(const BatchNormLayerDesc* desc) { + return createSWBatchNormLayerDesc( + desc->numChannels, + (float*)desc->mergedScale.data(), + (float*)desc->mergedBias.data()); +} + +/// Convert an activation layer description from C++ to Swift +ActivationKind activationLayerDescToSwift(const ActivationLayerDesc* desc) { + switch(desc->activation) { + case ACTIVATION_RELU: + return ActivationKind::relu(); + case ACTIVATION_MISH: + return ActivationKind::mish(); + case ACTIVATION_MISH_SCALE8: + return ActivationKind::identity(); // Metal/CoreML does not use scaled mish + case ACTIVATION_IDENTITY: + return ActivationKind::identity(); + default: + return ActivationKind::identity(); + } +} + +/// Convert a matrix multiplication layer description from C++ to Swift +SWMatMulLayerDesc matMulLayerDescToSwift(const MatMulLayerDesc* desc) { + return createSWMatMulLayerDesc( + desc->inChannels, + desc->outChannels, + (float*)desc->weights.data()); +} + +/// Convert a matrix bias layer description from C++ to Swift +SWMatBiasLayerDesc matBiasLayerDescToSwift(const MatBiasLayerDesc* desc) { + return createSWMatBiasLayerDesc(desc->numChannels, (float*)desc->weights.data()); +} + +/// Convert a residual block description from C++ to Swift +SWResidualBlockDesc residualBlockDescToSwift(const ResidualBlockDesc* desc) { + SWBatchNormLayerDesc preBN = batchNormLayerDescToSwift(&desc->preBN); + ActivationKind preActivationKind = activationLayerDescToSwift(&desc->preActivation); + SWConvLayerDesc regularConv = convLayerDescToSwift(&desc->regularConv); + SWBatchNormLayerDesc midBN = batchNormLayerDescToSwift(&desc->midBN); + ActivationKind midActivationKind = activationLayerDescToSwift(&desc->midActivation); + SWConvLayerDesc finalConv = convLayerDescToSwift(&desc->finalConv); + + return createSWResidualBlockDesc( + preBN, + preActivationKind, + regularConv, + midBN, + midActivationKind, + finalConv); +} + +/// Convert a global pooling residual block description from C++ to Swift +SWGlobalPoolingResidualBlockDesc globalPoolingResidualBlockDescToSwift(const GlobalPoolingResidualBlockDesc* desc) { + SWBatchNormLayerDesc preBN = batchNormLayerDescToSwift(&desc->preBN); + ActivationKind preActivationKind = activationLayerDescToSwift(&desc->preActivation); + SWConvLayerDesc regularConv = convLayerDescToSwift(&desc->regularConv); + SWConvLayerDesc gpoolConv = convLayerDescToSwift(&desc->gpoolConv); + SWBatchNormLayerDesc gpoolBN = batchNormLayerDescToSwift(&desc->gpoolBN); + ActivationKind gpoolActivationKind = activationLayerDescToSwift(&desc->gpoolActivation); + SWMatMulLayerDesc gpoolToBiasMul = matMulLayerDescToSwift(&desc->gpoolToBiasMul); + SWBatchNormLayerDesc midBN = batchNormLayerDescToSwift(&desc->midBN); + ActivationKind midActivationKind = activationLayerDescToSwift(&desc->midActivation); + SWConvLayerDesc finalConv = convLayerDescToSwift(&desc->finalConv); + + return createSWGlobalPoolingResidualBlockDesc( + preBN, + preActivationKind, + regularConv, + gpoolConv, + gpoolBN, + gpoolActivationKind, + gpoolToBiasMul, + midBN, + midActivationKind, + finalConv); +} + +// Forward declaration for mutual recursion +swift::Array residualBlocksToSwift(const vector>& blocks); + +/// Convert a nested bottleneck residual block description from C++ to Swift +SWNestedBottleneckResidualBlockDesc nestedBottleneckResidualBlockDescToSwift(const NestedBottleneckResidualBlockDesc* desc) { + SWBatchNormLayerDesc preBN = batchNormLayerDescToSwift(&desc->preBN); + ActivationKind preActivationKind = activationLayerDescToSwift(&desc->preActivation); + SWConvLayerDesc preConv = convLayerDescToSwift(&desc->preConv); + auto swBlocks = residualBlocksToSwift(desc->blocks); + SWBatchNormLayerDesc postBN = batchNormLayerDescToSwift(&desc->postBN); + ActivationKind postActivationKind = activationLayerDescToSwift(&desc->postActivation); + SWConvLayerDesc postConv = convLayerDescToSwift(&desc->postConv); + + return createSWNestedBottleneckResidualBlockDesc( + preBN, + preActivationKind, + preConv, + swBlocks, + postBN, + postActivationKind, + postConv); +} + +/// Convert residual blocks from C++ to Swift +swift::Array residualBlocksToSwift(const vector>& blocks) { + auto builder = createBlockDescriptorBuilder(); + + for(size_t i = 0; i < blocks.size(); i++) { + void* blockDesc = blocks[i].second.get(); + + if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) { + BlockDescriptor descriptor = globalPoolingResidualBlockDescToSwift((GlobalPoolingResidualBlockDesc*)blockDesc); + builder.enque(descriptor); + } else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) { + BlockDescriptor descriptor = nestedBottleneckResidualBlockDescToSwift((NestedBottleneckResidualBlockDesc*)blockDesc); + builder.enque(descriptor); + } else { + BlockDescriptor descriptor = residualBlockDescToSwift((ResidualBlockDesc*)blockDesc); + builder.enque(descriptor); + } + } + + return builder.getBlockDescriptors(); +} + +/// Convert a SGF metadata encoder description from C++ to Swift +swift::Optional sGFMetadataEncoderDescToSwift(const SGFMetadataEncoderDesc* desc) { + SWMatMulLayerDesc mul1 = matMulLayerDescToSwift(&desc->mul1); + SWMatBiasLayerDesc bias1 = matBiasLayerDescToSwift(&desc->bias1); + ActivationKind act1 = activationLayerDescToSwift(&desc->act1); + SWMatMulLayerDesc mul2 = matMulLayerDescToSwift(&desc->mul2); + SWMatBiasLayerDesc bias2 = matBiasLayerDescToSwift(&desc->bias2); + ActivationKind act2 = activationLayerDescToSwift(&desc->act2); + SWMatMulLayerDesc mul3 = matMulLayerDescToSwift(&desc->mul3); + + return createSWSGFMetadataEncoderDesc( + desc->metaEncoderVersion, + desc->numInputMetaChannels, + mul1, + bias1, + act1, + mul2, + bias2, + act2, + mul3); +} + +/// Convert a trunk description from C++ to Swift +SWTrunkDesc trunkDescToSwift(const TrunkDesc* trunk) { + SWConvLayerDesc initialConv = convLayerDescToSwift(&trunk->initialConv); + SWMatMulLayerDesc initialMatMul = matMulLayerDescToSwift(&trunk->initialMatMul); + auto sgfMetadataEncoder = sGFMetadataEncoderDescToSwift(&trunk->sgfMetadataEncoder); + auto swBlocks = residualBlocksToSwift(trunk->blocks); + SWBatchNormLayerDesc trunkTipBN = batchNormLayerDescToSwift(&trunk->trunkTipBN); + ActivationKind trunkTipActivation = activationLayerDescToSwift(&trunk->trunkTipActivation); + + return createSWTrunkDesc( + trunk->modelVersion, + trunk->trunkNumChannels, + trunk->midNumChannels, + trunk->regularNumChannels, + trunk->gpoolNumChannels, + initialConv, + initialMatMul, + sgfMetadataEncoder, + swBlocks, + trunkTipBN, + trunkTipActivation); +} + +/// Convert a policy head description from C++ to Swift +SWPolicyHeadDesc policyHeadDescToSwift(const PolicyHeadDesc* policyHead) { + SWConvLayerDesc p1Conv = convLayerDescToSwift(&policyHead->p1Conv); + SWConvLayerDesc g1Conv = convLayerDescToSwift(&policyHead->g1Conv); + SWBatchNormLayerDesc g1BN = batchNormLayerDescToSwift(&policyHead->g1BN); + ActivationKind g1Activation = activationLayerDescToSwift(&policyHead->g1Activation); + SWMatMulLayerDesc gpoolToBiasMul = matMulLayerDescToSwift(&policyHead->gpoolToBiasMul); + SWBatchNormLayerDesc p1BN = batchNormLayerDescToSwift(&policyHead->p1BN); + ActivationKind p1Activation = activationLayerDescToSwift(&policyHead->p1Activation); + SWConvLayerDesc p2Conv = convLayerDescToSwift(&policyHead->p2Conv); + SWMatMulLayerDesc gpoolToPassMul = matMulLayerDescToSwift(&policyHead->gpoolToPassMul); + SWMatBiasLayerDesc gpoolToPassBias = matBiasLayerDescToSwift(&policyHead->gpoolToPassBias); + ActivationKind passActivation = activationLayerDescToSwift(&policyHead->passActivation); + SWMatMulLayerDesc gpoolToPassMul2 = matMulLayerDescToSwift(&policyHead->gpoolToPassMul2); + + return createSWPolicyHeadDesc( + policyHead->modelVersion, + p1Conv, + g1Conv, + g1BN, + g1Activation, + gpoolToBiasMul, + p1BN, + p1Activation, + p2Conv, + gpoolToPassMul, + gpoolToPassBias, + passActivation, + gpoolToPassMul2); +} + +/// Convert a value head description from C++ to Swift +SWValueHeadDesc valueHeadDescToSwift(const ValueHeadDesc* valueHead) { + SWConvLayerDesc v1Conv = convLayerDescToSwift(&valueHead->v1Conv); + SWBatchNormLayerDesc v1BN = batchNormLayerDescToSwift(&valueHead->v1BN); + ActivationKind v1Activation = activationLayerDescToSwift(&valueHead->v1Activation); + SWMatMulLayerDesc v2Mul = matMulLayerDescToSwift(&valueHead->v2Mul); + SWMatBiasLayerDesc v2Bias = matBiasLayerDescToSwift(&valueHead->v2Bias); + ActivationKind v2Activation = activationLayerDescToSwift(&valueHead->v2Activation); + SWMatMulLayerDesc v3Mul = matMulLayerDescToSwift(&valueHead->v3Mul); + SWMatBiasLayerDesc v3Bias = matBiasLayerDescToSwift(&valueHead->v3Bias); + SWMatMulLayerDesc sv3Mul = matMulLayerDescToSwift(&valueHead->sv3Mul); + SWMatBiasLayerDesc sv3Bias = matBiasLayerDescToSwift(&valueHead->sv3Bias); + SWConvLayerDesc vOwnershipConv = convLayerDescToSwift(&valueHead->vOwnershipConv); + + return createSWValueHeadDesc( + valueHead->modelVersion, + v1Conv, + v1BN, + v1Activation, + v2Mul, + v2Bias, + v2Activation, + v3Mul, + v3Bias, + sv3Mul, + sv3Bias, + vOwnershipConv); +} + +/// Convert a model description from C++ to Swift +SWModelDesc modelDescToSwift(const ModelDesc* modelDesc) { + return createSWModelDesc( + modelDesc->modelVersion, + swift::String(modelDesc->name), + modelDesc->numInputChannels, + modelDesc->numInputGlobalChannels, + modelDesc->numInputMetaChannels, + modelDesc->numValueChannels, + modelDesc->numScoreValueChannels, + modelDesc->numOwnershipChannels, + trunkDescToSwift(&modelDesc->trunk), + policyHeadDescToSwift(&modelDesc->policyHead), + valueHeadDescToSwift(&modelDesc->valueHead)); +} + +} // namespace CoreMLProcess + +//------------------------------------------------------------------------------ +// LoadedModel implementation +//------------------------------------------------------------------------------ + +LoadedModel::LoadedModel(const string& fileName, const string& expectedSha256) { + modelPath = fileName; + ModelDesc::loadFromFileMaybeGZipped(fileName, modelDesc, expectedSha256); +} + +//------------------------------------------------------------------------------ +// NeuralNet namespace - Global functions +//------------------------------------------------------------------------------ + +void NeuralNet::globalInitialize() { + // No global initialization needed for Core ML +} + +void NeuralNet::globalCleanup() { + // No global cleanup needed for Core ML +} + +LoadedModel* NeuralNet::loadModelFile(const string& file, const string& expectedSha256) { + LoadedModel* loadedModel = new LoadedModel(file, expectedSha256); + return loadedModel; +} + +void NeuralNet::freeLoadedModel(LoadedModel* loadedModel) { + delete loadedModel; +} + +const ModelDesc& NeuralNet::getModelDesc(const LoadedModel* loadedModel) { + return loadedModel->modelDesc; +} + +//------------------------------------------------------------------------------ +// ComputeContext implementation +//------------------------------------------------------------------------------ + +ComputeContext::ComputeContext(int nnX, int nnY, enabled_t useFP16Mode, enabled_t useNHWCMode): +coremlContext(createCoreMLComputeContext(nnX, nnY, useFP16Mode != enabled_t::False)) { + this->useFP16Mode = useFP16Mode; + this->nnXLen = nnX; + this->nnYLen = nnY; + (void)useNHWCMode; +} + +ComputeContext::~ComputeContext() { +} + +ComputeContext* NeuralNet::createComputeContext( + const vector& gpuIdxs, + Logger* logger, + int nnXLen, + int nnYLen, + const string& openCLTunerFile, + const string& homeDataDirOverride, + bool openCLReTunePerBoardSize, + enabled_t useFP16Mode, + enabled_t useNHWCMode, + const LoadedModel* loadedModel) { + + (void)gpuIdxs; + (void)logger; + (void)openCLTunerFile; + (void)homeDataDirOverride; + (void)openCLReTunePerBoardSize; + (void)loadedModel; + + return new ComputeContext(nnXLen, nnYLen, useFP16Mode, useNHWCMode); +} + +void NeuralNet::freeComputeContext(ComputeContext* computeContext) { + delete computeContext; +} + +//------------------------------------------------------------------------------ +// ComputeHandle implementation +//------------------------------------------------------------------------------ + +static mutex computeHandleMutex; + +// Helper function to convert model and create hybrid compute handle +// This is needed because Swift Optional doesn't support assignment in C++ +static swift::Optional convertAndCreateHybridHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + bool requireExactNNLen, + int serverThreadIdx +) { + auto coremlContext = context->coremlContext; + int nnXLen = coremlContext.getNnXLen(); + int nnYLen = coremlContext.getNnYLen(); + bool useFP16 = (context->useFP16Mode != enabled_t::False); + bool optimizeMask = requireExactNNLen; + + // Convert model to CoreML format using native katagocoreml library + string coremlModelPath = CoreMLConversion::ensureModelConverted( + loadedModel->modelPath, + loadedModel->modelDesc.sha256, + nnXLen, + nnYLen, + useFP16, + optimizeMask, + serverThreadIdx + ); + + // Convert model descriptor to Swift format for MPSGraph path + SWModelDesc swModelDesc = CoreMLProcess::modelDescToSwift(&loadedModel->modelDesc); + + // Create hybrid compute handle (CoreML on CPU+ANE, MPSGraph on GPU) + return createHybridComputeHandle( + swift::String(coremlModelPath), + swModelDesc, + serverThreadIdx, + requireExactNNLen, + loadedModel->modelDesc.numInputChannels, + loadedModel->modelDesc.numInputGlobalChannels, + loadedModel->modelDesc.numInputMetaChannels, + loadedModel->modelDesc.policyHead.p2Conv.outChannels, + loadedModel->modelDesc.numValueChannels, + loadedModel->modelDesc.numScoreValueChannels, + loadedModel->modelDesc.numOwnershipChannels, + coremlContext + ); +} + +ComputeHandle::ComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + bool inputsUseNHWC, + int gpuIdx, + int serverThreadIdx, + bool requireExactNNLen): +hybridHandle(convertAndCreateHybridHandle(context, loadedModel, requireExactNNLen, serverThreadIdx)) { + const ModelDesc* modelDesc = &loadedModel->modelDesc; + auto coremlContext = context->coremlContext; + + nnXLen = coremlContext.getNnXLen(); + nnYLen = coremlContext.getNnYLen(); + gpuIndex = gpuIdx; + version = modelDesc->modelVersion; + metaEncoderVersion = modelDesc->metaEncoderVersion; + this->inputsUseNHWC = inputsUseNHWC; + this->requireExactNNLen = requireExactNNLen; + useFP16 = (context->useFP16Mode != enabled_t::False); +} + +ComputeHandle::~ComputeHandle() { +} + +ComputeHandle* NeuralNet::createComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + Logger* logger, + int maxBatchSize, + bool requireExactNNLen, + bool inputsUseNHWC, + int gpuIdxForThisThread, + int serverThreadIdx) { + + (void)logger; + (void)maxBatchSize; + + int gpuIdx = (gpuIdxForThisThread == -1) ? 0 : gpuIdxForThisThread; + ComputeHandle* handle = nullptr; + + { + lock_guard lock(computeHandleMutex); + handle = new ComputeHandle(context, loadedModel, inputsUseNHWC, gpuIdx, serverThreadIdx, requireExactNNLen); + } + + return handle; +} + +void NeuralNet::freeComputeHandle(ComputeHandle* handle) { + delete handle; +} + +bool NeuralNet::isUsingFP16(const ComputeHandle* handle) { + return handle->useFP16; +} + +//------------------------------------------------------------------------------ +// Device information +//------------------------------------------------------------------------------ + +void NeuralNet::printDevices() { + printCoreMLDevices(); +} + +//------------------------------------------------------------------------------ +// InputBuffers implementation +//------------------------------------------------------------------------------ + +InputBuffers::InputBuffers(const LoadedModel* loadedModel, int maxBatchSz, int nnXLen, int nnYLen) { + const ModelDesc& m = loadedModel->modelDesc; + + maxBatchSize = maxBatchSz; + policyResultChannels = m.policyHead.p2Conv.outChannels; + + assert(((m.modelVersion < 16) || (policyResultChannels == 4)) && + ((m.modelVersion >= 16) || (m.modelVersion < 12) || (policyResultChannels == 2)) && + ((m.modelVersion >= 12) || (policyResultChannels == 1))); + + singleSpatialElts = (size_t)m.numInputChannels * nnXLen * nnYLen; + singleInputElts = (size_t)m.numInputChannels * nnXLen * nnYLen; + singleInputGlobalElts = (size_t)m.numInputGlobalChannels; + singleInputMetaElts = (size_t)m.numInputMetaChannels; + singlePolicyResultElts = (size_t)(nnXLen * nnYLen); + singlePolicyPassResultElts = 1; + singlePolicyProbsElts = (size_t)((nnXLen * nnYLen) + 1); + singleValueResultElts = (size_t)m.numValueChannels; + singleOwnershipResultElts = (size_t)m.numOwnershipChannels * nnXLen * nnYLen; + singleOwnerMapElts = (size_t)m.numOwnershipChannels * nnXLen * nnYLen; + singleScoreValuesResultElts = (size_t)m.numScoreValueChannels; + singleMaskElts = (size_t)nnXLen * nnYLen; + + assert(NNModelVersion::getNumSpatialFeatures(m.modelVersion) == m.numInputChannels); + assert(NNModelVersion::getNumGlobalFeatures(m.modelVersion) == m.numInputGlobalChannels); + assert(singleValueResultElts == 3); + + rowSpatialBufferElts = (size_t)maxBatchSz * singleSpatialElts; + userInputBufferElts = (size_t)maxBatchSize * singleInputElts; + userInputGlobalBufferElts = (size_t)maxBatchSize * singleInputGlobalElts; + userInputMetaBufferElts = (size_t)maxBatchSize * singleInputMetaElts; + policyResultBufferElts = (size_t)maxBatchSize * singlePolicyResultElts * policyResultChannels; + policyPassResultBufferElts = (size_t)maxBatchSize * singlePolicyPassResultElts * policyResultChannels; + policyProbsBufferElts = (size_t)maxBatchSize * singlePolicyProbsElts * policyResultChannels; + valueResultBufferElts = (size_t)maxBatchSize * singleValueResultElts; + ownershipResultBufferElts = (size_t)maxBatchSize * singleOwnershipResultElts; + ownerMapBufferElts = (size_t)maxBatchSz * singleOwnerMapElts; + scoreValuesResultBufferElts = (size_t)maxBatchSize * singleScoreValuesResultElts; + userInputMaskBufferElts = (size_t)maxBatchSize * singleMaskElts; + + rowSpatialBuffer = new float[rowSpatialBufferElts]; + userInputBuffer = new float[userInputBufferElts]; + memset(&userInputBuffer[0], 0, userInputBufferElts * sizeof(userInputBuffer[0])); + + userInputGlobalBuffer = new float[userInputGlobalBufferElts]; + userInputMetaBuffer = new float[userInputMetaBufferElts]; + policyResults = new float[policyResultBufferElts]; + policyPassResults = new float[policyPassResultBufferElts]; + policyProbsBuffer = new float[policyProbsBufferElts]; + valueResults = new float[valueResultBufferElts]; + ownershipResults = new float[ownershipResultBufferElts]; + ownerMapBuffer = new float[ownerMapBufferElts]; + scoreValuesResults = new float[scoreValuesResultBufferElts]; + userInputMaskBuffer = new float[userInputMaskBufferElts]; + memset(&userInputMaskBuffer[0], 0, userInputMaskBufferElts * sizeof(userInputMaskBuffer[0])); +} + +InputBuffers::~InputBuffers() { + delete[] rowSpatialBuffer; + delete[] userInputBuffer; + delete[] userInputGlobalBuffer; + delete[] userInputMetaBuffer; + delete[] policyResults; + delete[] policyPassResults; + delete[] policyProbsBuffer; + delete[] valueResults; + delete[] ownershipResults; + delete[] ownerMapBuffer; + delete[] scoreValuesResults; + delete[] userInputMaskBuffer; +} + +InputBuffers* NeuralNet::createInputBuffers(const LoadedModel* loadedModel, int maxBatchSize, int nnXLen, int nnYLen) { + return new InputBuffers(loadedModel, maxBatchSize, nnXLen, nnYLen); +} + +void NeuralNet::freeInputBuffers(InputBuffers* inputBuffers) { + delete inputBuffers; +} + +//------------------------------------------------------------------------------ +// CoreMLProcess namespace - Helper functions +//------------------------------------------------------------------------------ + +void CoreMLProcess::copyRowData(float* dest, const float* src, size_t numElements) { + copy(src, src + numElements, dest); +} + +void CoreMLProcess::convertNCHW( + float* rowSpatialInput, + const int C, + const int H, + const int W, + const bool inputsUseNHWC) { + + if ((!inputsUseNHWC) || (C <= 0) || (H <= 0) || (W <= 0)) { + return; + } + + const int totalSize = H * W * C; + + if (totalSize <= 1) + return; + + const int HW = H * W; + + auto get_nchw_target_index = [C, W, HW](int nhwc_index) -> int { + int c = nhwc_index % C; + int temp = nhwc_index / C; + int x = temp % W; + int y = temp / W; + return (c * HW) + (y * W) + x; + }; + + std::vector processed(totalSize, false); + + for (int i = 0; i < totalSize; ++i) { + if (processed[i]) + continue; + + int target_i = get_nchw_target_index(i); + + if (target_i == i) { + processed[i] = true; + continue; + } + + int current_idx = i; + float value_in_hand = rowSpatialInput[i]; + + while (true) { + int target_idx = get_nchw_target_index(current_idx); + float value_at_target = rowSpatialInput[target_idx]; + rowSpatialInput[target_idx] = value_in_hand; + processed[target_idx] = true; + value_in_hand = value_at_target; + current_idx = target_idx; + + if (current_idx == i) + break; + } + } +} + +void CoreMLProcess::processRowData(size_t row, ComputeHandle* gpuHandle, InputBuffers* inputBuffers, NNResultBuf** inputBufs) { + int nnXLen = gpuHandle->nnXLen; + int nnYLen = gpuHandle->nnYLen; + int numSpatialFeatures = NNModelVersion::getNumSpatialFeatures(gpuHandle->version); + + float* rowSpatialInput = &inputBuffers->userInputBuffer[inputBuffers->singleSpatialElts * row]; + float* rowGlobalInput = &inputBuffers->userInputGlobalBuffer[inputBuffers->singleInputGlobalElts * row]; + float* rowMetaInput = &inputBuffers->userInputMetaBuffer[inputBuffers->singleInputMetaElts * row]; + const float* rowGlobal = inputBufs[row]->rowGlobalBuf.data(); + const float* rowSpatial = inputBufs[row]->rowSpatialBuf.data(); + const float* rowMeta = inputBufs[row]->rowMetaBuf.data(); + + CoreMLProcess::copyRowData(rowGlobalInput, rowGlobal, inputBuffers->singleInputGlobalElts); + CoreMLProcess::copyRowData(rowMetaInput, rowMeta, inputBuffers->singleInputMetaElts); + + SymmetryHelpers::copyInputsWithSymmetry( + rowSpatial, + rowSpatialInput, + 1, + nnYLen, + nnXLen, + numSpatialFeatures, + gpuHandle->inputsUseNHWC, + inputBufs[row]->symmetry); + + CoreMLProcess::convertNCHW( + rowSpatialInput, + numSpatialFeatures, + nnYLen, + nnXLen, + gpuHandle->inputsUseNHWC); + + // Copy first channel of spatial input (mask) to dedicated mask buffer + // After NCHW conversion, the first nnXLen*nnYLen elements are the mask channel + float* rowMaskInput = &inputBuffers->userInputMaskBuffer[inputBuffers->singleMaskElts * row]; + copy(rowSpatialInput, rowSpatialInput + inputBuffers->singleMaskElts, rowMaskInput); +} + +float CoreMLProcess::policyOptimismCalc(const double policyOptimism, const float p, const float pOpt) { + return p + ((pOpt - p) * policyOptimism); +} + +void CoreMLProcess::processOptimism( + InputBuffers* inputBuffers, + NNOutput* currentOutput, + const double policyOptimism, + size_t row) { + auto& buffers = *inputBuffers; + const auto singlePolicyResultElts = buffers.singlePolicyResultElts; + float* targetBuffer = &buffers.policyProbsBuffer[row * singlePolicyResultElts]; + float* policyOutputBuf = &buffers.policyResults[row * singlePolicyResultElts * buffers.policyResultChannels]; + + for(size_t i = 0; i < singlePolicyResultElts; ++i) { + const float p = policyOutputBuf[i]; + const float pOpt = policyOutputBuf[i + singlePolicyResultElts]; + targetBuffer[i] = CoreMLProcess::policyOptimismCalc(policyOptimism, p, pOpt); + } + + const auto p = buffers.policyPassResults[row * buffers.policyResultChannels]; + const auto pOpt = buffers.policyPassResults[row * buffers.policyResultChannels + 1]; + currentOutput->policyProbs[buffers.singlePolicyProbsElts - 1] = CoreMLProcess::policyOptimismCalc(policyOptimism, p, pOpt); +} + +void CoreMLProcess::processPolicy( + InputBuffers* inputBuffers, + NNOutput* currentOutput, + const ComputeHandle* gpuHandle, + NNResultBuf* inputBuf, + size_t row) { + auto& buffers = *inputBuffers; + float* targetBuffer = &buffers.policyResults[row * buffers.singlePolicyResultElts * buffers.policyResultChannels]; + const auto policyOptimism = inputBuf->policyOptimism; + + if(buffers.policyResultChannels == 1) { + currentOutput->policyProbs[buffers.singlePolicyProbsElts - 1] = + buffers.policyPassResults[row * buffers.policyResultChannels]; + } else { + CoreMLProcess::processOptimism(inputBuffers, currentOutput, policyOptimism, row); + targetBuffer = &buffers.policyProbsBuffer[row * buffers.singlePolicyResultElts]; + } + + SymmetryHelpers::copyOutputsWithSymmetry( + targetBuffer, currentOutput->policyProbs, 1, gpuHandle->nnYLen, gpuHandle->nnXLen, inputBuf->symmetry); +} + +void CoreMLProcess::processValue( + const InputBuffers* inputBuffers, + NNOutput* currentOutput, + const size_t row) { + const size_t singleValueResultElts = inputBuffers->singleValueResultElts; + assert(singleValueResultElts == 3); + const float* valueOutputBuf = &inputBuffers->valueResults[row * singleValueResultElts]; + currentOutput->whiteWinProb = valueOutputBuf[0]; + currentOutput->whiteLossProb = valueOutputBuf[1]; + currentOutput->whiteNoResultProb = valueOutputBuf[2]; +} + +void CoreMLProcess::processOwnership( + const InputBuffers* inputBuffers, + NNOutput* currentOutput, + const ComputeHandle* gpuHandle, + const int symmetry, + const size_t row) { + const int nnXLen = gpuHandle->nnXLen; + const int nnYLen = gpuHandle->nnYLen; + const size_t singleOwnershipResultElts = inputBuffers->singleOwnershipResultElts; + const size_t ownershipOutputBufOffset = row * singleOwnershipResultElts; + + if(currentOutput->whiteOwnerMap != nullptr) { + const float* ownershipOutputBuf = &inputBuffers->ownershipResults[ownershipOutputBufOffset]; + SymmetryHelpers::copyOutputsWithSymmetry( + ownershipOutputBuf, currentOutput->whiteOwnerMap, 1, nnYLen, nnXLen, symmetry); + } +} + +void CoreMLProcess::processScoreValues( + const InputBuffers* inputBuffers, + NNOutput* currentOutput, + const int modelVersion, + const size_t row) { + const size_t offset = row * inputBuffers->singleScoreValuesResultElts; + const float* currentScoreValueData = &inputBuffers->scoreValuesResults[offset]; + + if(modelVersion >= 9) { + size_t numScoreValueChannels = inputBuffers->singleScoreValuesResultElts; + assert(numScoreValueChannels == 6); + currentOutput->whiteScoreMean = currentScoreValueData[0]; + currentOutput->whiteScoreMeanSq = currentScoreValueData[1]; + currentOutput->whiteLead = currentScoreValueData[2]; + currentOutput->varTimeLeft = currentScoreValueData[3]; + currentOutput->shorttermWinlossError = currentScoreValueData[4]; + currentOutput->shorttermScoreError = currentScoreValueData[5]; + } + else if(modelVersion >= 8) { + size_t numScoreValueChannels = inputBuffers->singleScoreValuesResultElts; + assert(numScoreValueChannels == 4); + currentOutput->whiteScoreMean = currentScoreValueData[0]; + currentOutput->whiteScoreMeanSq = currentScoreValueData[1]; + currentOutput->whiteLead = currentScoreValueData[2]; + currentOutput->varTimeLeft = currentScoreValueData[3]; + currentOutput->shorttermWinlossError = 0; + currentOutput->shorttermScoreError = 0; + } + else if(modelVersion >= 4) { + size_t numScoreValueChannels = inputBuffers->singleScoreValuesResultElts; + assert(numScoreValueChannels == 2); + currentOutput->whiteScoreMean = currentScoreValueData[0]; + currentOutput->whiteScoreMeanSq = currentScoreValueData[1]; + currentOutput->whiteLead = currentOutput->whiteScoreMean; + currentOutput->varTimeLeft = 0; + currentOutput->shorttermWinlossError = 0; + currentOutput->shorttermScoreError = 0; + } + else { + assert(modelVersion >= 3); + size_t numScoreValueChannels = inputBuffers->singleScoreValuesResultElts; + assert(numScoreValueChannels == 1); + currentOutput->whiteScoreMean = currentScoreValueData[0]; + currentOutput->whiteScoreMeanSq = currentOutput->whiteScoreMean * currentOutput->whiteScoreMean; + currentOutput->whiteLead = currentOutput->whiteScoreMean; + currentOutput->varTimeLeft = 0; + currentOutput->shorttermWinlossError = 0; + currentOutput->shorttermScoreError = 0; + } +} + +void CoreMLProcess::processRow( + size_t row, + const ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + NNResultBuf** inputBufs, + vector& outputs) { + NNOutput* currentOutput = outputs[row]; + assert(currentOutput->nnXLen == gpuHandle->nnXLen); + assert(currentOutput->nnYLen == gpuHandle->nnYLen); + CoreMLProcess::processPolicy(inputBuffers, currentOutput, gpuHandle, inputBufs[row], row); + CoreMLProcess::processValue(inputBuffers, currentOutput, row); + CoreMLProcess::processOwnership(inputBuffers, currentOutput, gpuHandle, inputBufs[row]->symmetry, row); + CoreMLProcess::processScoreValues(inputBuffers, currentOutput, gpuHandle->version, row); +} + +void CoreMLProcess::getCoreMLOutput( + ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + int numBatchEltsFilled, + NNResultBuf** inputBufs, + vector& outputs) { + assert(numBatchEltsFilled > 0); + + int batchSize = numBatchEltsFilled; + + assert(batchSize <= inputBuffers->maxBatchSize); + assert((NNModelVersion::getNumSpatialFeatures(gpuHandle->version) * gpuHandle->nnXLen * gpuHandle->nnYLen) <= (int)inputBuffers->singleInputElts); + assert(NNModelVersion::getNumGlobalFeatures(gpuHandle->version) == (int)inputBuffers->singleInputGlobalElts); + + if(gpuHandle->metaEncoderVersion > 0) { + assert(SGFMetadata::METADATA_INPUT_NUM_CHANNELS == (int)inputBuffers->singleInputMetaElts); + } + + assert(inputBuffers->singleValueResultElts == 3); + + for(int row = 0; row < batchSize; row++) { + CoreMLProcess::processRowData(row, gpuHandle, inputBuffers, inputBufs); + } + + auto hybridHandle = gpuHandle->hybridHandle; + assert(hybridHandle); + + // Call hybrid inference (CoreML on CPU+ANE, MPSGraph on GPU) + // Mask buffer has correct stride (singleMaskElts = H*W per batch element) + // When requireExactNNLen is true, mask operations can be optimized (optimize_identity_mask) + hybridHandle.get().apply( + inputBuffers->userInputBuffer, + inputBuffers->userInputGlobalBuffer, + inputBuffers->userInputMetaBuffer, + inputBuffers->userInputMaskBuffer, // Dedicated mask buffer with correct stride + inputBuffers->policyResults, + inputBuffers->policyPassResults, + inputBuffers->valueResults, + inputBuffers->scoreValuesResults, + inputBuffers->ownershipResults, + batchSize); + + for(int row = 0; row < batchSize; row++) { + CoreMLProcess::processRow(row, gpuHandle, inputBuffers, inputBufs, outputs); + } +} + +void NeuralNet::getOutput( + ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + int numBatchEltsFilled, + NNResultBuf** inputBufs, + vector& outputs) { + + CoreMLProcess::getCoreMLOutput(gpuHandle, inputBuffers, numBatchEltsFilled, inputBufs, outputs); +} + +//------------------------------------------------------------------------------ +// Test functions - not supported for Core ML backend +//------------------------------------------------------------------------------ + +bool NeuralNet::testEvaluateConv( + const ConvLayerDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)outputBuffer; + return false; +} + +bool NeuralNet::testEvaluateBatchNorm( + const BatchNormLayerDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +bool NeuralNet::testEvaluateResidualBlock( + const ResidualBlockDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +bool NeuralNet::testEvaluateGlobalPoolingResidualBlock( + const GlobalPoolingResidualBlockDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +#endif // USE_COREML_BACKEND diff --git a/cpp/neuralnet/coremlbackend.h b/cpp/neuralnet/coremlbackend.h new file mode 100644 index 000000000..af52762cb --- /dev/null +++ b/cpp/neuralnet/coremlbackend.h @@ -0,0 +1,251 @@ +#pragma once + +#ifdef USE_COREML_BACKEND + +#include +#include "desc.h" +#include "../core/commontypes.h" +#include "../neuralnet/modelversion.h" +#include "../neuralnet/nneval.h" +#include "../neuralnet/nninputs.h" +#include "../neuralnet/nninterface.h" +#include + +using namespace std; +using namespace KataGoCoreML; + +namespace CoreMLProcess { + +void copyRowData(float* dest, const float* src, size_t numElements); +void convertNCHW(float* rowSpatialInput, int C, int H, int W, bool inputsUseNHWC); +void processRowData(size_t row, ComputeHandle* gpuHandle, InputBuffers* inputBuffers, NNResultBuf** inputBufs); +float policyOptimismCalc(const double policyOptimism, const float p, const float pOpt); +void processOptimism(InputBuffers* inputBuffers, NNOutput* currentOutput, const double policyOptimism, size_t row); + +void processPolicy(InputBuffers* inputBuffers, + NNOutput* currentOutput, + const ComputeHandle* gpuHandle, + NNResultBuf* inputBuf, + size_t row); + +void processValue(const InputBuffers* inputBuffers, NNOutput* currentOutput, const size_t row); + +void processOwnership(const InputBuffers* inputBuffers, + NNOutput* currentOutput, + const ComputeHandle* gpuHandle, + const int symmetry, + const size_t row); + +void +processScoreValues(const InputBuffers* inputBuffers, NNOutput* currentOutput, const int modelVersion, const size_t row); + +void processRow(size_t row, + const ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + NNResultBuf** inputBufs, + vector& outputs); + +void getCoreMLOutput(ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + int numBatchEltsFilled, + NNResultBuf** inputBufs, + vector& outputs); +}; + +/** + * @brief Represents a loaded neural network model. + * A LoadedModel object contains a ModelDesc object that describes the characteristics of the loaded model. + * For Core ML backend, we also store the model path for on-demand conversion. + */ +struct LoadedModel { + /** + * @brief The description of the loaded model. + */ + ModelDesc modelDesc; + + /** + * @brief Path to the original .bin.gz model file for conversion. + */ + string modelPath; + + /** + * @brief Construct a new Loaded Model object + * This constructor loads a machine learning model from a file and sets the modelDesc field. + * @param fileName The name of the file containing the machine learning model. + * @param expectedSha256 The expected SHA-256 hash of the model file. + */ + LoadedModel(const string& fileName, const string& expectedSha256); + + LoadedModel() = delete; + LoadedModel(const LoadedModel&) = delete; + LoadedModel& operator=(const LoadedModel&) = delete; +}; + +/** + * @brief Context for computing neural network operations using Core ML. + * Contains global configuration settings for neural network computations. + */ +struct ComputeContext { + /** + * @brief Whether to use FP16 mode for computations. + */ + enabled_t useFP16Mode; + + /** + * @brief The width of the neural network input. + */ + int nnXLen; + + /** + * @brief The height of the neural network input. + */ + int nnYLen; + + /** + * @brief Core ML compute context instance from Swift. + */ + CoreMLComputeContext coremlContext; + + /** + * @brief Constructs a ComputeContext object. + * @param nnX The width of the input tensor. + * @param nnY The height of the input tensor. + * @param useFP16Mode Whether to use half-precision floating-point (FP16) mode. + * @param useNHWCMode Whether to use the NHWC format for input tensors. + */ + ComputeContext(int nnX, int nnY, enabled_t useFP16Mode, enabled_t useNHWCMode); + + ~ComputeContext(); + ComputeContext() = delete; + ComputeContext(const ComputeContext&) = delete; + ComputeContext& operator=(const ComputeContext&) = delete; +}; + +/** + * @brief A handle for performing neural network computations using Core ML. + * This struct represents a per-thread handle for computing neural network operations. + */ +struct ComputeHandle { + /** + * @brief The x length of the neural network. + */ + int nnXLen; + + /** + * @brief The y length of the neural network. + */ + int nnYLen; + + /** + * @brief The index of the GPU to use for computation. + */ + int gpuIndex; + + /** + * @brief The version of the loaded model. + */ + int version; + + /** + * @brief The version of the metadata encoder. + */ + int metaEncoderVersion; + + /** + * @brief Whether the input data uses NHWC format. + */ + bool inputsUseNHWC; + + /** + * @brief Whether to use 16-bit floating-point precision. + */ + bool useFP16; + + /** + * @brief Whether exact neural net length is required (enables mask optimization). + */ + bool requireExactNNLen; + + /** + * @brief The hybrid compute handle instance from Swift. + * This handle dispatches work to both CoreML (CPU+ANE) and MPSGraph (GPU). + */ + swift::Optional hybridHandle; + + /** + * @brief Construct a new ComputeHandle object. + * @param context The ComputeContext object to use for computation. + * @param loadedModel A pointer to the LoadedModel object. + * @param inputsUseNHWC Whether the input data uses NHWC format. + * @param gpuIdx The index of the GPU to use. + * @param serverThreadIdx The index of the server thread. + * @param requireExactNNLen Whether exact NN length is required. + */ + ComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + bool inputsUseNHWC, + int gpuIdx, + int serverThreadIdx, + bool requireExactNNLen); + + ~ComputeHandle(); + ComputeHandle() = delete; + ComputeHandle(const ComputeHandle&) = delete; + ComputeHandle& operator=(const ComputeHandle&) = delete; +}; + +/** + * @brief Input and output buffers for neural network inference. + */ +struct InputBuffers { + int maxBatchSize; + size_t policyResultChannels; + + size_t singleSpatialElts; + size_t singleInputElts; + size_t singleInputGlobalElts; + size_t singleInputMetaElts; + size_t singlePolicyResultElts; + size_t singlePolicyPassResultElts; + size_t singlePolicyProbsElts; + size_t singleValueResultElts; + size_t singleOwnershipResultElts; + size_t singleOwnerMapElts; + size_t singleScoreValuesResultElts; + size_t singleMaskElts; + + size_t rowSpatialBufferElts; + size_t userInputBufferElts; + size_t userInputGlobalBufferElts; + size_t userInputMetaBufferElts; + size_t policyResultBufferElts; + size_t policyPassResultBufferElts; + size_t policyProbsBufferElts; + size_t valueResultBufferElts; + size_t ownershipResultBufferElts; + size_t ownerMapBufferElts; + size_t scoreValuesResultBufferElts; + size_t userInputMaskBufferElts; + + float* rowSpatialBuffer; + float* userInputBuffer; + float* userInputGlobalBuffer; + float* userInputMetaBuffer; + float* policyResults; + float* policyPassResults; + float* policyProbsBuffer; + float* valueResults; + float* ownershipResults; + float* ownerMapBuffer; + float* scoreValuesResults; + float* userInputMaskBuffer; + + InputBuffers(const LoadedModel* loadedModel, int maxBatchSz, int nnXLen, int nnYLen); + ~InputBuffers(); + InputBuffers() = delete; + InputBuffers(const InputBuffers&) = delete; + InputBuffers& operator=(const InputBuffers&) = delete; +}; + +#endif // USE_COREML_BACKEND diff --git a/cpp/neuralnet/coremlbackend.swift b/cpp/neuralnet/coremlbackend.swift new file mode 100644 index 000000000..9b17b610d --- /dev/null +++ b/cpp/neuralnet/coremlbackend.swift @@ -0,0 +1,871 @@ +import Foundation +import CoreML +import MetalPerformanceShaders +import MetalPerformanceShadersGraph + +/// A class that handles output to standard error. +class StandardError: TextOutputStream { + func write(_ string: String) { + try? FileHandle.standardError.write(contentsOf: Data(string.utf8)) + } +} + +/// Print to standard error +func printError(_ item: Any) { + var instance = StandardError() + print(item, to: &instance) +} + +// NOTE: Model caching and conversion are now handled in C++ using the native katagocoreml library. +// The Python-based CoreMLConverter and ModelCacheManager have been removed to eliminate Python dependency. + +/// Context storing board dimensions and settings +public class CoreMLComputeContext { + public let nnXLen: Int32 + public let nnYLen: Int32 + public let useFP16: Bool + + init(nnXLen: Int32, nnYLen: Int32, useFP16: Bool) { + self.nnXLen = nnXLen + self.nnYLen = nnYLen + self.useFP16 = useFP16 + } +} + +/// Create a Core ML compute context +public func createCoreMLComputeContext( + nnXLen: Int32, + nnYLen: Int32, + useFP16: Bool +) -> CoreMLComputeContext { + return CoreMLComputeContext(nnXLen: nnXLen, nnYLen: nnYLen, useFP16: useFP16) +} + +/// Handle that wraps the loaded MLModel for inference +public class CoreMLComputeHandle { + let model: MLModel + let nnXLen: Int32 + let nnYLen: Int32 + let optimizeIdentityMask: Bool + let numInputChannels: Int + let numInputGlobalChannels: Int + let numInputMetaChannels: Int + let numPolicyChannels: Int + let numValueChannels: Int + let numScoreValueChannels: Int + let numOwnershipChannels: Int + + /// Model input/output names matching KataGoCoremltools output + struct IONames { + static let spatialInput = "spatial_input" + static let globalInput = "global_input" + static let inputMask = "input_mask" + static let metaInput = "meta_input" + + static let policyOutput = "policy_p2_conv" + static let policyPassOutput = "policy_pass" + static let valueOutput = "value_v3_bias" + static let ownershipOutput = "value_ownership_conv" + static let scoreValueOutput = "value_sv3_bias" + } + + init(model: MLModel, nnXLen: Int32, nnYLen: Int32, + optimizeIdentityMask: Bool, + numInputChannels: Int, + numInputGlobalChannels: Int, + numInputMetaChannels: Int, + numPolicyChannels: Int, + numValueChannels: Int, + numScoreValueChannels: Int, + numOwnershipChannels: Int) { + self.model = model + self.nnXLen = nnXLen + self.nnYLen = nnYLen + self.optimizeIdentityMask = optimizeIdentityMask + self.numInputChannels = numInputChannels + self.numInputGlobalChannels = numInputGlobalChannels + self.numInputMetaChannels = numInputMetaChannels + self.numPolicyChannels = numPolicyChannels + self.numValueChannels = numValueChannels + self.numScoreValueChannels = numScoreValueChannels + self.numOwnershipChannels = numOwnershipChannels + } + + /// Run inference on a batch of inputs + public func apply( + spatialInput: UnsafeMutablePointer, + globalInput: UnsafeMutablePointer, + metaInput: UnsafeMutablePointer, + maskInput: UnsafeMutablePointer, + policy: UnsafeMutablePointer, + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer, + batchSize: Int + ) { + // Process batch elements in parallel using Grand Central Dispatch + // Each inference is independent, reading/writing to different buffer offsets + DispatchQueue.concurrentPerform(iterations: batchSize) { b in + autoreleasepool { + do { + try runSingleInference( + batchIndex: b, + spatialInput: spatialInput, + globalInput: globalInput, + metaInput: metaInput, + maskInput: maskInput, + policy: policy, + policyPass: policyPass, + value: value, + scoreValue: scoreValue, + ownership: ownership + ) + } catch { + printError("Core ML inference error: \(error)") + } + } + } + } + + private func runSingleInference( + batchIndex: Int, + spatialInput: UnsafeMutablePointer, + globalInput: UnsafeMutablePointer, + metaInput: UnsafeMutablePointer, + maskInput: UnsafeMutablePointer, + policy: UnsafeMutablePointer, + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer + ) throws { + let spatialSize = Int(nnXLen) * Int(nnYLen) * numInputChannels + let spatialOffset = batchIndex * spatialSize + + // Create MLMultiArray for spatial input (1, C, H, W) + let spatialArray = try MLMultiArray( + shape: [1, NSNumber(value: numInputChannels), + NSNumber(value: nnYLen), NSNumber(value: nnXLen)], + dataType: .float32) + + // Copy spatial data using fast memcpy + let spatialPtr = spatialArray.dataPointer.assumingMemoryBound(to: Float32.self) + memcpy(spatialPtr, spatialInput.advanced(by: spatialOffset), spatialSize * MemoryLayout.size) + + // Create global input array (1, C) - rank 2 as expected by converter + let globalArray = try MLMultiArray( + shape: [1, NSNumber(value: numInputGlobalChannels)], + dataType: .float32) + let globalPtr = globalArray.dataPointer.assumingMemoryBound(to: Float32.self) + let globalOffset = batchIndex * numInputGlobalChannels + memcpy(globalPtr, globalInput.advanced(by: globalOffset), numInputGlobalChannels * MemoryLayout.size) + + // Build feature provider dictionary + var inputDict: [String: MLFeatureValue] = [ + IONames.spatialInput: MLFeatureValue(multiArray: spatialArray), + IONames.globalInput: MLFeatureValue(multiArray: globalArray) + ] + + // Add mask input (always required, even with optimize_identity_mask=True) + // When optimize_identity_mask=True, the mask is still required as input but + // internal mask operations are optimized away for ~6.5% speedup + let maskArray = try MLMultiArray( + shape: [1, 1, NSNumber(value: nnYLen), NSNumber(value: nnXLen)], + dataType: .float32) + let maskPtr = maskArray.dataPointer.assumingMemoryBound(to: Float32.self) + let maskSize = Int(nnXLen) * Int(nnYLen) + let maskOffset = batchIndex * maskSize + memcpy(maskPtr, maskInput.advanced(by: maskOffset), maskSize * MemoryLayout.size) + inputDict[IONames.inputMask] = MLFeatureValue(multiArray: maskArray) + + // Add meta input if model has it + if numInputMetaChannels > 0 { + let metaArray = try MLMultiArray( + shape: [1, NSNumber(value: numInputMetaChannels)], + dataType: .float32) + let metaPtr = metaArray.dataPointer.assumingMemoryBound(to: Float32.self) + let metaOffset = batchIndex * numInputMetaChannels + memcpy(metaPtr, metaInput.advanced(by: metaOffset), numInputMetaChannels * MemoryLayout.size) + inputDict[IONames.metaInput] = MLFeatureValue(multiArray: metaArray) + } + + // Run prediction + let featureProvider = try MLDictionaryFeatureProvider(dictionary: inputDict) + let prediction = try model.prediction(from: featureProvider) + + // Extract outputs and copy to output buffers + extractOutputs( + prediction: prediction, + batchIndex: batchIndex, + policy: policy, + policyPass: policyPass, + value: value, + scoreValue: scoreValue, + ownership: ownership + ) + } + + /// Copy MLMultiArray data to destination buffer, respecting strides. + /// Core ML may return non-contiguous arrays, especially for spatial outputs after GPU computation. + private func copyMultiArray( + _ array: MLMultiArray, + to dest: UnsafeMutablePointer, + destOffset: Int + ) { + let shape = array.shape.map { $0.intValue } + let strides = array.strides.map { $0.intValue } + let ptr = array.dataPointer.assumingMemoryBound(to: Float32.self) + let totalElements = shape.reduce(1, *) + + // Check if contiguous (strides match expected for row-major C-order) + var isContiguous = true + var expectedStride = 1 + for i in (0...size) + } else { + // Slow path: copy with strides (handles non-contiguous layouts) + copyWithStrides( + from: ptr, + to: dest, + destOffset: destOffset, + shape: shape, + strides: strides, + dim: 0, + srcOffset: 0, + destIdx: 0 + ) + } + } + + /// Recursively copy array elements respecting strides (NCHW order) + @discardableResult + private func copyWithStrides( + from src: UnsafePointer, + to dest: UnsafeMutablePointer, + destOffset: Int, + shape: [Int], + strides: [Int], + dim: Int, + srcOffset: Int, + destIdx: Int + ) -> Int { + var currentDestIdx = destIdx + + if dim == shape.count - 1 { + // Innermost dimension: copy elements + for i in 0.., + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer + ) { + // Extract policy output (1, policyChannels, H, W) + // Must use stride-aware copy as Core ML may return non-contiguous arrays + if let policyArray = prediction.featureValue(for: IONames.policyOutput)?.multiArrayValue { + let policyOffset = batchIndex * Int(nnXLen) * Int(nnYLen) * numPolicyChannels + copyMultiArray(policyArray, to: policy, destOffset: policyOffset) + } + + // Extract policy pass output (1, numPolicyChannels) + if let passArray = prediction.featureValue(for: IONames.policyPassOutput)?.multiArrayValue { + let passOffset = batchIndex * numPolicyChannels + copyMultiArray(passArray, to: policyPass, destOffset: passOffset) + } + + // Extract value output (1, 3) + if let valueArray = prediction.featureValue(for: IONames.valueOutput)?.multiArrayValue { + let valueOffset = batchIndex * numValueChannels + copyMultiArray(valueArray, to: value, destOffset: valueOffset) + } + + // Extract score value output (1, numScoreValueChannels) + if let svArray = prediction.featureValue(for: IONames.scoreValueOutput)?.multiArrayValue { + let svOffset = batchIndex * numScoreValueChannels + copyMultiArray(svArray, to: scoreValue, destOffset: svOffset) + } + + // Extract ownership output (1, 1, H, W) + // Must use stride-aware copy as Core ML may return non-contiguous arrays + if let ownArray = prediction.featureValue(for: IONames.ownershipOutput)?.multiArrayValue { + let ownOffset = batchIndex * Int(nnXLen) * Int(nnYLen) * numOwnershipChannels + copyMultiArray(ownArray, to: ownership, destOffset: ownOffset) + } + } +} + +/// Create compute handle - loads pre-converted Core ML model +/// Model conversion is now handled in C++ using the native katagocoreml library +public func createCoreMLComputeHandle( + coremlModelPath: String, + serverThreadIdx: Int, + requireExactNNLen: Bool, + numInputChannels: Int32, + numInputGlobalChannels: Int32, + numInputMetaChannels: Int32, + numPolicyChannels: Int32, + numValueChannels: Int32, + numScoreValueChannels: Int32, + numOwnershipChannels: Int32, + context: CoreMLComputeContext +) -> CoreMLComputeHandle? { + + let optimizeMask = requireExactNNLen // When true: skips internal mask operations (~6.5% speedup) + let mlpackagePath = URL(fileURLWithPath: coremlModelPath) + + // Load Core ML model (already converted by C++ katagocoreml library) + do { + let config = MLModelConfiguration() + config.computeUnits = .cpuAndNeuralEngine // Exclude GPU for hybrid mode + + printError("Core ML backend \(serverThreadIdx): Compiling model...") + let compiledURL = try MLModel.compileModel(at: mlpackagePath) + + printError("Core ML backend \(serverThreadIdx): Loading compiled model...") + let model = try MLModel(contentsOf: compiledURL, configuration: config) + + printError("Core ML backend \(serverThreadIdx): Model loaded successfully, \(context.nnXLen)x\(context.nnYLen)") + + return CoreMLComputeHandle( + model: model, + nnXLen: context.nnXLen, + nnYLen: context.nnYLen, + optimizeIdentityMask: optimizeMask, + numInputChannels: Int(numInputChannels), + numInputGlobalChannels: Int(numInputGlobalChannels), + numInputMetaChannels: Int(numInputMetaChannels), + numPolicyChannels: Int(numPolicyChannels), + numValueChannels: Int(numValueChannels), + numScoreValueChannels: Int(numScoreValueChannels), + numOwnershipChannels: Int(numOwnershipChannels) + ) + } catch { + printError("Core ML backend: Failed to load model: \(error)") + return nil + } +} + +/// Print available Core ML compute units +public func printCoreMLDevices() { + printError("Core ML backend: Hybrid mode - CoreML (CPU+ANE) + MPSGraph (GPU)") +} + +// MARK: - Throughput Tracker for Adaptive Batch Sizing + +/// Tracks throughput for CoreML and MPSGraph paths to adaptively adjust batch split ratio +public class ThroughputTracker { + private var coreMLSamplesPerSec: Double = 1.0 + private var mpsGraphSamplesPerSec: Double = 1.0 + private let alpha: Double = 0.3 // EMA smoothing factor (higher = faster adaptation) + private let lock = NSLock() + + /// Update CoreML throughput measurement + public func updateCoreML(samples: Int, duration: TimeInterval) { + guard duration > 0, samples > 0 else { return } + let newRate = Double(samples) / duration + lock.lock() + coreMLSamplesPerSec = alpha * newRate + (1 - alpha) * coreMLSamplesPerSec + lock.unlock() + } + + /// Update MPSGraph throughput measurement + public func updateMPSGraph(samples: Int, duration: TimeInterval) { + guard duration > 0, samples > 0 else { return } + let newRate = Double(samples) / duration + lock.lock() + mpsGraphSamplesPerSec = alpha * newRate + (1 - alpha) * mpsGraphSamplesPerSec + lock.unlock() + } + + /// Get optimal CoreML ratio (0.0 to 1.0) based on measured throughput + public func getOptimalCoreMLRatio() -> Float { + lock.lock() + let total = coreMLSamplesPerSec + mpsGraphSamplesPerSec + let ratio = total > 0 ? Float(coreMLSamplesPerSec / total) : 0.5 + lock.unlock() + return ratio + } + + /// Get current throughput stats for logging + public func getStats() -> (coreML: Double, mpsGraph: Double, ratio: Float) { + lock.lock() + let stats = (coreMLSamplesPerSec, mpsGraphSamplesPerSec, getOptimalCoreMLRatio()) + lock.unlock() + return stats + } +} + +// MARK: - MPSGraph-based Model for GPU Inference + +/// GPU-based model using MPSGraph for inference +public class MPSGraphModelHandle { + let device: MTLDevice + let commandQueue: MTLCommandQueue + let graph: MPSGraph + let nnXLen: Int32 + let nnYLen: Int32 + let numInputChannels: Int + let numInputGlobalChannels: Int + let numInputMetaChannels: Int + let numPolicyChannels: Int + let numValueChannels: Int + let numScoreValueChannels: Int + let numOwnershipChannels: Int + + // Layers + let input: InputLayer + let inputGlobal: InputGlobalLayer + let inputMeta: InputMetaLayer + let mask: MaskLayer + let trunk: Trunk + let policyHead: PolicyHead + let valueHead: ValueHead + let targetTensors: [MPSGraphTensor] + + public init?( + modelDesc: SWModelDesc, + nnXLen: Int32, + nnYLen: Int32, + optimizeIdentityMask: Bool = false + ) { + guard let device = MTLCreateSystemDefaultDevice() else { + printError("MPSGraph backend: Failed to create Metal device") + return nil + } + + self.device = device + guard let queue = device.makeCommandQueue() else { + printError("MPSGraph backend: Failed to create command queue") + return nil + } + self.commandQueue = queue + self.graph = MPSGraph() + self.nnXLen = nnXLen + self.nnYLen = nnYLen + self.numInputChannels = modelDesc.numInputChannels.intValue + self.numInputGlobalChannels = modelDesc.numInputGlobalChannels.intValue + self.numInputMetaChannels = modelDesc.numInputMetaChannels.intValue + self.numPolicyChannels = 2 // Policy has 2 channels + self.numValueChannels = modelDesc.numValueChannels.intValue + self.numScoreValueChannels = modelDesc.numScoreValueChannels.intValue + self.numOwnershipChannels = modelDesc.numOwnershipChannels.intValue + + let nnXLenNS = nnXLen as NSNumber + let nnYLenNS = nnYLen as NSNumber + + input = InputLayer( + graph: graph, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS, + numChannels: modelDesc.numInputChannels) + + inputGlobal = InputGlobalLayer( + graph: graph, + numGlobalFeatures: modelDesc.numInputGlobalChannels) + + inputMeta = InputMetaLayer( + graph: graph, + numMetaFeatures: modelDesc.numInputMetaChannels) + + mask = MaskLayer( + graph: graph, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS) + + // Use constant tensors when mask is all 1s (requireExactNNLen=true) + let maskSum: MaskSumLayer + let maskSumSqrtS14M01: MaskSumSqrtS14M01Layer + let maskSumSqrtS14M01SquareS01: MaskSumSqrtS14M01SquareS01Layer + + if optimizeIdentityMask { + maskSum = MaskSumLayer( + graph: graph, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS) + maskSumSqrtS14M01 = MaskSumSqrtS14M01Layer( + graph: graph, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS) + maskSumSqrtS14M01SquareS01 = MaskSumSqrtS14M01SquareS01Layer( + graph: graph, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS) + } else { + maskSum = MaskSumLayer( + graph: graph, + maskTensor: mask.tensor) + maskSumSqrtS14M01 = MaskSumSqrtS14M01Layer( + graph: graph, + maskSum: maskSum) + maskSumSqrtS14M01SquareS01 = MaskSumSqrtS14M01SquareS01Layer( + graph: graph, + maskSumSqrtS14M01: maskSumSqrtS14M01) + } + + trunk = Trunk( + graph: graph, + descriptor: modelDesc.trunk, + inputTensor: input.tensor, + inputGlobalTensor: inputGlobal.tensor, + inputMetaTensor: inputMeta.tensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS, + optimizeIdentityMask: optimizeIdentityMask) + + policyHead = PolicyHead( + graph: graph, + descriptor: modelDesc.policyHead, + sourceTensor: trunk.resultTensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS, + optimizeIdentityMask: optimizeIdentityMask) + + valueHead = ValueHead( + graph: graph, + descriptor: modelDesc.valueHead, + sourceTensor: trunk.resultTensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + maskSumSqrtS14M01SquareS01Tensor: maskSumSqrtS14M01SquareS01.tensor, + nnXLen: nnXLenNS, + nnYLen: nnYLenNS, + optimizeIdentityMask: optimizeIdentityMask) + + targetTensors = [ + policyHead.policyTensor, + policyHead.policyPassTensor, + valueHead.valueTensor, + valueHead.scoreValueTensor, + valueHead.ownershipTensor, + ] + + printError("MPSGraph backend: Initialized on \(device.name)\(optimizeIdentityMask ? " (mask optimized)" : "")") + } + + /// Run inference on a batch using MPSGraph (GPU) + public func apply( + input inputPointer: UnsafeMutablePointer, + inputGlobal inputGlobalPointer: UnsafeMutablePointer, + inputMeta inputMetaPointer: UnsafeMutablePointer, + policy: UnsafeMutablePointer, + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer, + batchSize: Int + ) { + let channelAxis = InputShape.getChannelAxis() + let numInputChannels = input.shape[channelAxis] + let nnXLenNS = nnXLen as NSNumber + let nnYLenNS = nnYLen as NSNumber + + let inputShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputChannels, + nnYLen: nnYLenNS, + nnXLen: nnXLenNS) + + let inputDescriptor = MPSNDArrayDescriptor( + dataType: input.tensor.dataType, + shape: inputShape) + + let inputArray = MPSNDArray( + device: device, + descriptor: inputDescriptor) + + inputArray.writeBytes(inputPointer) + + let numInputGlobalChannels = inputGlobal.shape[channelAxis] + + let inputGlobalShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputGlobalChannels, + nnYLen: 1, + nnXLen: 1) + + let inputGlobalDescriptor = MPSNDArrayDescriptor( + dataType: inputGlobal.tensor.dataType, + shape: inputGlobalShape) + + let inputGlobalArray = MPSNDArray( + device: device, + descriptor: inputGlobalDescriptor) + + inputGlobalArray.writeBytes(inputGlobalPointer) + + let numInputMetaChannels = inputMeta.shape[channelAxis] + + let inputMetaShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputMetaChannels, + nnYLen: 1, + nnXLen: 1) + + let inputMetaDescriptor = MPSNDArrayDescriptor( + dataType: inputMeta.tensor.dataType, + shape: inputMetaShape) + + let inputMetaArray = MPSNDArray( + device: device, + descriptor: inputMetaDescriptor) + + inputMetaArray.writeBytes(inputMetaPointer) + + let maskShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: 1, + nnYLen: nnYLenNS, + nnXLen: nnXLenNS) + + let maskDescriptor = MPSNDArrayDescriptor( + dataType: mask.tensor.dataType, + shape: maskShape) + + let maskArray = MPSNDArray( + device: device, + descriptor: maskDescriptor) + + // Extract mask from first channel of spatial input + var maskStrideArray = [ + MemoryLayout.size, + Int(nnXLen) * MemoryLayout.size, + Int(nnYLen) * Int(nnXLen) * MemoryLayout.size, + numInputChannels.intValue * Int(nnYLen) * Int(nnXLen) * MemoryLayout.size, + ] + + maskArray.writeBytes(inputPointer, strideBytes: &maskStrideArray) + + let feeds = [ + input.tensor: MPSGraphTensorData(inputArray), + inputGlobal.tensor: MPSGraphTensorData(inputGlobalArray), + inputMeta.tensor: MPSGraphTensorData(inputMetaArray), + mask.tensor: MPSGraphTensorData(maskArray), + ] + + let fetch = graph.run( + with: commandQueue, + feeds: feeds, + targetTensors: targetTensors, + targetOperations: nil) + + fetch[policyHead.policyTensor]?.mpsndarray().readBytes(policy) + fetch[policyHead.policyPassTensor]?.mpsndarray().readBytes(policyPass) + fetch[valueHead.valueTensor]?.mpsndarray().readBytes(value) + fetch[valueHead.scoreValueTensor]?.mpsndarray().readBytes(scoreValue) + fetch[valueHead.ownershipTensor]?.mpsndarray().readBytes(ownership) + } +} + +// MARK: - Hybrid Compute Handle + +/// Hybrid compute handle that dispatches to both CoreML (CPU+ANE) and MPSGraph (GPU) +public class HybridComputeHandle { + let coremlHandle: CoreMLComputeHandle + let mpsGraphHandle: MPSGraphModelHandle + let throughputTracker: ThroughputTracker + let coremlQueue: DispatchQueue + let mpsGraphQueue: DispatchQueue + let nnXLen: Int32 + let nnYLen: Int32 + + public init( + coremlHandle: CoreMLComputeHandle, + mpsGraphHandle: MPSGraphModelHandle + ) { + self.coremlHandle = coremlHandle + self.mpsGraphHandle = mpsGraphHandle + self.throughputTracker = ThroughputTracker() + self.coremlQueue = DispatchQueue(label: "com.katago.coreml", qos: .userInitiated) + self.mpsGraphQueue = DispatchQueue(label: "com.katago.mpsgraph", qos: .userInitiated) + self.nnXLen = coremlHandle.nnXLen + self.nnYLen = coremlHandle.nnYLen + } + + /// Run hybrid inference - splits batch between CoreML and MPSGraph + public func apply( + spatialInput: UnsafeMutablePointer, + globalInput: UnsafeMutablePointer, + metaInput: UnsafeMutablePointer, + maskInput: UnsafeMutablePointer, + policy: UnsafeMutablePointer, + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer, + batchSize: Int + ) { + // Get optimal split ratio based on throughput + let ratio = throughputTracker.getOptimalCoreMLRatio() + let coreMLBatchSize = max(1, min(batchSize - 1, Int(Float(batchSize) * ratio))) + let mpsGraphBatchSize = batchSize - coreMLBatchSize + + // Calculate buffer offsets + let spatialSize = Int(nnXLen) * Int(nnYLen) * coremlHandle.numInputChannels + let globalSize = coremlHandle.numInputGlobalChannels + let metaSize = coremlHandle.numInputMetaChannels + let policySize = Int(nnXLen) * Int(nnYLen) * coremlHandle.numPolicyChannels + let valueSize = coremlHandle.numValueChannels + let scoreValueSize = coremlHandle.numScoreValueChannels + let ownershipSize = Int(nnXLen) * Int(nnYLen) * coremlHandle.numOwnershipChannels + + let group = DispatchGroup() + + // CoreML path (CPU + ANE) + if coreMLBatchSize > 0 { + group.enter() + coremlQueue.async { [self] in + let start = CFAbsoluteTimeGetCurrent() + + coremlHandle.apply( + spatialInput: spatialInput, + globalInput: globalInput, + metaInput: metaInput, + maskInput: maskInput, + policy: policy, + policyPass: policyPass, + value: value, + scoreValue: scoreValue, + ownership: ownership, + batchSize: coreMLBatchSize + ) + + let duration = CFAbsoluteTimeGetCurrent() - start + throughputTracker.updateCoreML(samples: coreMLBatchSize, duration: duration) + group.leave() + } + } + + // MPSGraph path (GPU) + if mpsGraphBatchSize > 0 { + group.enter() + mpsGraphQueue.async { [self] in + let start = CFAbsoluteTimeGetCurrent() + + // Offset pointers for MPSGraph batch portion + let spatialOffset = coreMLBatchSize * spatialSize + let globalOffset = coreMLBatchSize * globalSize + let metaOffset = coreMLBatchSize * metaSize + let policyOffset = coreMLBatchSize * policySize + let valueOffset = coreMLBatchSize * valueSize + let scoreValueOffset = coreMLBatchSize * scoreValueSize + let ownershipOffset = coreMLBatchSize * ownershipSize + + autoreleasepool { + mpsGraphHandle.apply( + input: spatialInput.advanced(by: spatialOffset), + inputGlobal: globalInput.advanced(by: globalOffset), + inputMeta: metaInput.advanced(by: metaOffset), + policy: policy.advanced(by: policyOffset), + policyPass: policyPass.advanced(by: coreMLBatchSize * coremlHandle.numPolicyChannels), + value: value.advanced(by: valueOffset), + scoreValue: scoreValue.advanced(by: scoreValueOffset), + ownership: ownership.advanced(by: ownershipOffset), + batchSize: mpsGraphBatchSize + ) + } + + let duration = CFAbsoluteTimeGetCurrent() - start + throughputTracker.updateMPSGraph(samples: mpsGraphBatchSize, duration: duration) + group.leave() + } + } + + // Wait for both paths to complete + group.wait() + } +} + +/// Create a hybrid compute handle +public func createHybridComputeHandle( + coremlModelPath: String, + modelDesc: SWModelDesc, + serverThreadIdx: Int, + requireExactNNLen: Bool, + numInputChannels: Int32, + numInputGlobalChannels: Int32, + numInputMetaChannels: Int32, + numPolicyChannels: Int32, + numValueChannels: Int32, + numScoreValueChannels: Int32, + numOwnershipChannels: Int32, + context: CoreMLComputeContext +) -> HybridComputeHandle? { + + // Create CoreML handle (CPU + ANE) + guard let coremlHandle = createCoreMLComputeHandle( + coremlModelPath: coremlModelPath, + serverThreadIdx: serverThreadIdx, + requireExactNNLen: requireExactNNLen, + numInputChannels: numInputChannels, + numInputGlobalChannels: numInputGlobalChannels, + numInputMetaChannels: numInputMetaChannels, + numPolicyChannels: numPolicyChannels, + numValueChannels: numValueChannels, + numScoreValueChannels: numScoreValueChannels, + numOwnershipChannels: numOwnershipChannels, + context: context + ) else { + printError("Hybrid backend \(serverThreadIdx): Failed to create CoreML handle") + return nil + } + + // Create MPSGraph handle (GPU) + guard let mpsGraphHandle = MPSGraphModelHandle( + modelDesc: modelDesc, + nnXLen: context.nnXLen, + nnYLen: context.nnYLen, + optimizeIdentityMask: requireExactNNLen + ) else { + printError("Hybrid backend \(serverThreadIdx): Failed to create MPSGraph handle") + return nil + } + + printError("Hybrid backend \(serverThreadIdx): Initialized CoreML (CPU+ANE) + MPSGraph (GPU)") + + return HybridComputeHandle( + coremlHandle: coremlHandle, + mpsGraphHandle: mpsGraphHandle + ) +} diff --git a/cpp/neuralnet/mpsgraphlayers.swift b/cpp/neuralnet/mpsgraphlayers.swift new file mode 100644 index 000000000..fce460678 --- /dev/null +++ b/cpp/neuralnet/mpsgraphlayers.swift @@ -0,0 +1,2360 @@ +// MPSGraph layer implementations shared between Metal and CoreML backends +// Extracted from metalbackend.swift to enable hybrid CoreML + MPSGraph execution + +import Foundation +import MetalPerformanceShaders +import MetalPerformanceShadersGraph + +// MARK: - Helper Extensions + +/// An extension to the Data struct for handling float data with optional FP16 conversion. +extension Data { + /// Initializes a new Data instance using an UnsafeMutablePointer, with optional conversion to FP16 format. + init( + floatsNoCopy: UnsafeMutablePointer, + shape: [NSNumber] + ) { + self.init( + bytesNoCopy: floatsNoCopy, + count: shape.countBytesOfFloat32(), + deallocator: .none) + } +} + +/// Extension to MPSNDArray to convert from MPSGraphTensor, and to read/write bytes from/to UnsafeMutableRawPointer +extension MPSNDArray { + /// Read bytes from the buffer + func readBytes(_ buffer: UnsafeMutableRawPointer) { + self.readBytes(buffer, strideBytes: nil) + } + + /// Write bytes to the buffer + func writeBytes(_ buffer: UnsafeMutableRawPointer) { + self.writeBytes(buffer, strideBytes: nil) + } +} + +/// Extension to Array to count number of elements and bytes +extension Array where Element == NSNumber { + /// Count number of elements + func countElements() -> Int { + return reduce(1, { $0 * $1.intValue }) + } + + /// Count number of bytes + func countBytesOfFloat32() -> Int { + return countElements() * MemoryLayout.size + } +} + +/// Extension to MPSGraph to the mish activation function +extension MPSGraph { + /// Mish activation: x * tanh(softplus(x)) + func mish(tensor: MPSGraphTensor) -> MPSGraphTensor { + assert(tensor.dataType == .float32) + + let one = 1.0 + let threshold = 20.0 + let thresholdTensor = constant(threshold, dataType: tensor.dataType) + let minimumTensor = minimum(tensor, thresholdTensor, name: nil) + let expTensor = exponent(with: minimumTensor, name: nil) + let oneTensor = constant(one, dataType: tensor.dataType) + let addTensor = addition(expTensor, oneTensor, name: nil) + let logTensor = logarithm(with: addTensor, name: nil) + let lessTensor = lessThan(tensor, thresholdTensor, name: nil) + let selectTensor = select( + predicate: lessTensor, trueTensor: logTensor, falseTensor: tensor, name: nil) + let tanhTensor = tanh(with: selectTensor, name: nil) + let mulTensor = multiplication(tensor, tanhTensor, name: nil) + + return mulTensor + } +} + +// MARK: - Input Shape Utilities + +/// A structure that represents the input shape (internal - not exposed to C++) +struct InputShape { + /// Create a shape for the input tensor + static func create( + batchSize: NSNumber, + numChannels: NSNumber, + nnYLen: NSNumber, + nnXLen: NSNumber + ) -> [NSNumber] { + return [batchSize, numChannels, nnYLen, nnXLen] + } + + /// Get the channel axis + static func getChannelAxis() -> Int { + return 1 + } + + /// Get the HW axes + static func getHWAxes() -> [NSNumber] { + return [2, 3] as [NSNumber] + } +} + +// MARK: - Input Layers + +/// A structure that represents the input layer +struct InputLayer { + let tensor: MPSGraphTensor + let shape: [NSNumber] + + init( + graph: MPSGraph, + nnXLen: NSNumber, + nnYLen: NSNumber, + numChannels: NSNumber, + dataType: MPSDataType = .float32 + ) { + shape = InputShape.create( + batchSize: -1, + numChannels: numChannels, + nnYLen: nnYLen, + nnXLen: nnXLen) + + self.tensor = graph.placeholder( + shape: shape, + dataType: dataType, + name: nil) + + assert(self.tensor.shape?.count == 4) + } +} + +/// A structure that represents an input global layer for a neural network model. +struct InputGlobalLayer { + let tensor: MPSGraphTensor + let shape: [NSNumber] + + init( + graph: MPSGraph, + numGlobalFeatures: NSNumber, + dataType: MPSDataType = .float32 + ) { + shape = InputShape.create( + batchSize: -1, + numChannels: numGlobalFeatures, + nnYLen: 1, + nnXLen: 1) + + self.tensor = graph.placeholder( + shape: shape, + dataType: dataType, + name: nil) + + assert(self.tensor.shape?.count == 4) + } +} + +/// A structure representing the input meta layer for a neural network graph. +struct InputMetaLayer { + let tensor: MPSGraphTensor + let shape: [NSNumber] + + init( + graph: MPSGraph, + numMetaFeatures: NSNumber, + dataType: MPSDataType = .float32 + ) { + shape = InputShape.create( + batchSize: -1, + numChannels: numMetaFeatures, + nnYLen: 1, + nnXLen: 1) + + self.tensor = graph.placeholder( + shape: shape, + dataType: dataType, + name: nil) + } +} + +/// A structure that represents a mask layer for a neural network model. +struct MaskLayer { + let tensor: MPSGraphTensor + let shape: [NSNumber] + + init( + graph: MPSGraph, + nnXLen: NSNumber, + nnYLen: NSNumber, + dataType: MPSDataType = .float32 + ) { + shape = InputShape.create( + batchSize: -1, + numChannels: 1, + nnYLen: nnYLen, + nnXLen: nnXLen) + + self.tensor = graph.placeholder( + shape: shape, + dataType: dataType, + name: nil) + + assert(self.tensor.shape?.count == 4) + } +} + +// MARK: - Mask Processing Layers + +/// A structure that represents a layer which performs the summation operation on a mask layer. +struct MaskSumLayer { + let tensor: MPSGraphTensor + + init(tensor: MPSGraphTensor) { + self.tensor = tensor + assert(self.tensor.shape?.count == 4) + } + + init( + graph: MPSGraph, + maskTensor: MPSGraphTensor + ) { + let hwAxes = InputShape.getHWAxes() + + self.tensor = graph.reductionSum( + with: maskTensor, + axes: hwAxes, + name: nil) + + assert(self.tensor.shape?.count == 4) + } + + /// Optimized init for when mask is all 1s (requireExactNNLen=true) + /// Returns constant tensor with boardSize value + init( + graph: MPSGraph, + nnXLen: NSNumber, + nnYLen: NSNumber, + dataType: MPSDataType = .float32 + ) { + let boardSize = Double(nnXLen.intValue * nnYLen.intValue) + self.tensor = graph.constant( + boardSize, + shape: [1, 1, 1, 1], + dataType: dataType) + + assert(self.tensor.shape?.count == 4) + } +} + +/// A structure that represents sqrt(maskSum) * 0.1 - 1.4 +struct MaskSumSqrtS14M01Layer { + let tensor: MPSGraphTensor + + init(tensor: MPSGraphTensor) { + self.tensor = tensor + assert(self.tensor.shape?.count == 4) + } + + init( + graph: MPSGraph, + maskSum: MaskSumLayer + ) { + let sqrtMaskSum = graph.squareRoot(with: maskSum.tensor, name: nil) + + let fourTeen = graph.constant( + 14.0, + shape: [1], + dataType: maskSum.tensor.dataType) + + let subtracted = graph.subtraction(sqrtMaskSum, fourTeen, name: nil) + + let zeroPointone = graph.constant( + 0.1, + shape: [1], + dataType: maskSum.tensor.dataType) + + self.tensor = graph.multiplication( + subtracted, + zeroPointone, + name: nil) + + assert(self.tensor.shape?.count == 4) + } + + /// Optimized init for when mask is all 1s (requireExactNNLen=true) + /// Returns constant tensor: (sqrt(boardSize) - 14) * 0.1 + init( + graph: MPSGraph, + nnXLen: NSNumber, + nnYLen: NSNumber, + dataType: MPSDataType = .float32 + ) { + let boardSize = Double(nnXLen.intValue * nnYLen.intValue) + let value = (sqrt(boardSize) - 14.0) * 0.1 + self.tensor = graph.constant( + value, + shape: [1, 1, 1, 1], + dataType: dataType) + + assert(self.tensor.shape?.count == 4) + } +} + +/// A structure for (sqrt(maskSum) * 0.1 - 1.4)^2 - 0.1 +struct MaskSumSqrtS14M01SquareS01Layer { + let tensor: MPSGraphTensor + + init(tensor: MPSGraphTensor) { + self.tensor = tensor + assert(self.tensor.shape?.count == 4) + } + + init( + graph: MPSGraph, + maskSumSqrtS14M01: MaskSumSqrtS14M01Layer + ) { + let squared = graph.square(with: maskSumSqrtS14M01.tensor, name: nil) + + let zeroPointone = graph.constant( + 0.1, + shape: [1], + dataType: maskSumSqrtS14M01.tensor.dataType) + + self.tensor = graph.subtraction( + squared, + zeroPointone, + name: nil) + + assert(self.tensor.shape?.count == 4) + } + + /// Optimized init for when mask is all 1s (requireExactNNLen=true) + /// Returns constant tensor: ((sqrt(boardSize) - 14) * 0.1)^2 - 0.1 + init( + graph: MPSGraph, + nnXLen: NSNumber, + nnYLen: NSNumber, + dataType: MPSDataType = .float32 + ) { + let boardSize = Double(nnXLen.intValue * nnYLen.intValue) + let sqrtS14M01 = (sqrt(boardSize) - 14.0) * 0.1 + let value = sqrtS14M01 * sqrtS14M01 - 0.1 + self.tensor = graph.constant( + value, + shape: [1, 1, 1, 1], + dataType: dataType) + + assert(self.tensor.shape?.count == 4) + } +} + +// MARK: - Layer Descriptors + +/// An enumeration of the different kinds of activation function. +public enum ActivationKind { + case identity + case relu + case mish +} + +/// A struct that represents a description of convolutional layer. +public struct SWConvLayerDesc { + let convYSize: NSNumber + let convXSize: NSNumber + let inChannels: NSNumber + let outChannels: NSNumber + let dilationY: Int + let dilationX: Int + let weights: UnsafeMutablePointer + + init( + convYSize: NSNumber, + convXSize: NSNumber, + inChannels: NSNumber, + outChannels: NSNumber, + dilationY: Int, + dilationX: Int, + weights: UnsafeMutablePointer + ) { + self.convYSize = convYSize + self.convXSize = convXSize + self.inChannels = inChannels + self.outChannels = outChannels + self.dilationY = dilationY + self.dilationX = dilationX + self.weights = weights + } +} + +public func createSWConvLayerDesc( + convYSize: Int32, + convXSize: Int32, + inChannels: Int32, + outChannels: Int32, + dilationY: Int32, + dilationX: Int32, + weights: UnsafeMutablePointer +) -> SWConvLayerDesc { + return SWConvLayerDesc( + convYSize: convYSize as NSNumber, + convXSize: convXSize as NSNumber, + inChannels: inChannels as NSNumber, + outChannels: outChannels as NSNumber, + dilationY: Int(dilationY), + dilationX: Int(dilationX), + weights: weights) +} + +/// A struct that represents a description of a batch normalization layer. +public struct SWBatchNormLayerDesc { + let numChannels: NSNumber + let mergedScale: UnsafeMutablePointer + let mergedBias: UnsafeMutablePointer + + init( + numChannels: NSNumber, + mergedScale: UnsafeMutablePointer, + mergedBias: UnsafeMutablePointer + ) { + self.numChannels = numChannels + self.mergedScale = mergedScale + self.mergedBias = mergedBias + } +} + +public func createSWBatchNormLayerDesc( + numChannels: Int32, + mergedScale: UnsafeMutablePointer, + mergedBias: UnsafeMutablePointer +) -> SWBatchNormLayerDesc { + return SWBatchNormLayerDesc( + numChannels: numChannels as NSNumber, + mergedScale: mergedScale, + mergedBias: mergedBias) +} + +/// A struct that represents a matrix multiplication layer descriptor +public struct SWMatMulLayerDesc { + let inChannels: NSNumber + let outChannels: NSNumber + let weights: UnsafeMutablePointer + + init( + inChannels: NSNumber, + outChannels: NSNumber, + weights: UnsafeMutablePointer + ) { + self.inChannels = inChannels + self.outChannels = outChannels + self.weights = weights + } +} + +public func createSWMatMulLayerDesc( + inChannels: Int32, + outChannels: Int32, + weights: UnsafeMutablePointer +) -> SWMatMulLayerDesc { + return SWMatMulLayerDesc( + inChannels: inChannels as NSNumber, + outChannels: outChannels as NSNumber, + weights: weights) +} + +/// A struct that represents the bias layer description. +public struct SWMatBiasLayerDesc { + let numChannels: NSNumber + let weights: UnsafeMutablePointer + + init( + numChannels: NSNumber, + weights: UnsafeMutablePointer + ) { + self.numChannels = numChannels + self.weights = weights + } +} + +public func createSWMatBiasLayerDesc( + numChannels: Int32, + weights: UnsafeMutablePointer +) -> SWMatBiasLayerDesc { + return SWMatBiasLayerDesc( + numChannels: numChannels as NSNumber, + weights: weights) +} + +// MARK: - Core Layers + +/// A class that represents a convolutional layer using MPSGraph +class ConvLayer { + let resultTensor: MPSGraphTensor + let convDescriptor = MPSGraphConvolution2DOpDescriptor( + strideInX: 1, + strideInY: 1, + dilationRateInX: 1, + dilationRateInY: 1, + groups: 1, + paddingStyle: .TF_SAME, + dataLayout: .NCHW, + weightsLayout: .OIHW)! + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + descriptor: SWConvLayerDesc, + nnXLen: NSNumber, + nnYLen: NSNumber + ) { + let weightsShape = [ + descriptor.outChannels, + descriptor.inChannels, + descriptor.convYSize, + descriptor.convXSize, + ] + + let weightsData = Data( + floatsNoCopy: descriptor.weights, + shape: weightsShape) + + let weightsTensor = graph.constant( + weightsData, + shape: weightsShape, + dataType: sourceTensor.dataType) + + resultTensor = graph.convolution2D( + sourceTensor, + weights: weightsTensor, + descriptor: convDescriptor, + name: nil) + + assert(resultTensor.shape?.count == 4) + } +} + +/// A class that represents a batch normalization layer. +class BatchNormLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + descriptor: SWBatchNormLayerDesc, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let scaleBiasShape = InputShape.create( + batchSize: 1, + numChannels: descriptor.numChannels, + nnYLen: 1, + nnXLen: 1) + + let mergedScaleData = Data( + floatsNoCopy: descriptor.mergedScale, + shape: scaleBiasShape) + + let mergedBiasData = Data( + floatsNoCopy: descriptor.mergedBias, + shape: scaleBiasShape) + + let scaleTensor = graph.constant( + mergedScaleData, + shape: scaleBiasShape, + dataType: sourceTensor.dataType) + + let biasTensor = graph.constant( + mergedBiasData, + shape: scaleBiasShape, + dataType: sourceTensor.dataType) + + let scaled = graph.multiplication( + sourceTensor, + scaleTensor, + name: nil) + + let normalized = graph.addition( + scaled, + biasTensor, + name: nil) + + // Skip mask multiplication when all mask values are 1 + if optimizeIdentityMask { + resultTensor = normalized + } else { + resultTensor = graph.multiplication( + normalized, + maskTensor, + name: nil) + } + + assert(resultTensor.shape?.count == 4) + } +} + +/// A structure that represents an activation layer +struct ActivationLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + activationKind: ActivationKind + ) { + switch activationKind { + case .relu: + resultTensor = graph.reLU(with: sourceTensor, name: nil) + case .mish: + resultTensor = graph.mish(tensor: sourceTensor) + default: + resultTensor = sourceTensor + } + + assert(resultTensor.shape == sourceTensor.shape) + } +} + +/// A structure representing a matrix multiplication layer. +struct MatMulLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + descriptor: SWMatMulLayerDesc, + sourceTensor: MPSGraphTensor + ) { + assert( + (sourceTensor.shape?.count == 4) || (sourceTensor.shape?[1] == descriptor.inChannels)) + assert( + (sourceTensor.shape?.count == 2) || (sourceTensor.shape?[1] == descriptor.inChannels)) + + let weightsShape = [ + descriptor.inChannels, + descriptor.outChannels, + ] + + let weightsData = Data( + floatsNoCopy: descriptor.weights, + shape: weightsShape) + + let weightsTensor = graph.constant( + weightsData, + shape: weightsShape, + dataType: sourceTensor.dataType) + + let shape = [-1, descriptor.inChannels] + + let reshapedSource = graph.reshape( + sourceTensor, + shape: shape, + name: nil) + + resultTensor = graph.matrixMultiplication( + primary: reshapedSource, + secondary: weightsTensor, + name: nil) + + assert(resultTensor.shape?.count == 2) + } +} + +/// A structure that performs matrix bias operations +struct MatBiasLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + descriptor: SWMatBiasLayerDesc, + sourceTensor: MPSGraphTensor + ) { + assert( + (sourceTensor.shape?.count == 2) && (sourceTensor.shape?[1] == descriptor.numChannels)) + + let weightsShape = [1, descriptor.numChannels] + + let weightsData = Data( + floatsNoCopy: descriptor.weights, + shape: weightsShape) + + let weightsTensor = graph.constant( + weightsData, + shape: weightsShape, + dataType: sourceTensor.dataType) + + resultTensor = graph.addition( + sourceTensor, + weightsTensor, + name: nil) + } +} + +/// A structure that performs bias operations in NC coordinates. +struct AddNCBiasLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + biasTensor: MPSGraphTensor, + nnXLen: NSNumber, + nnYLen: NSNumber, + numChannels: NSNumber + ) { + let shape = InputShape.create( + batchSize: -1, + numChannels: numChannels, + nnYLen: 1, + nnXLen: 1) + + assert(biasTensor.shape?[1] == shape[1]) + + let reshaped = graph.reshape(biasTensor, shape: shape, name: nil) + resultTensor = graph.addition(sourceTensor, reshaped, name: nil) + + assert(resultTensor.shape?.count == 4) + assert(resultTensor.shape?[2] == nnYLen) + assert(resultTensor.shape?[3] == nnXLen) + } +} + +// MARK: - Pooling Layers + +/// A structure that represents a global pooling layer +struct GlobalPoolingLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + optimizeIdentityMask: Bool = false + ) { + let hwAxes = InputShape.getHWAxes() + let channelAxis = InputShape.getChannelAxis() + + let sumTensor = graph.reductionSum( + with: sourceTensor, + axes: hwAxes, + name: nil) + + let meanTensor = graph.division(sumTensor, maskSumTensor, name: nil) + + let meanMaskTensor = graph.multiplication( + meanTensor, + maskSumSqrtS14M01Tensor, + name: nil) + + let maxTensor: MPSGraphTensor + if optimizeIdentityMask { + // When all mask values are 1, directly compute max without mask adjustment + maxTensor = graph.reductionMaximum( + with: sourceTensor, + axes: hwAxes, + name: nil) + } else { + // Mask out invalid positions by subtracting 1 (making them very negative) + let oneTensor = graph.constant(1.0, dataType: sourceTensor.dataType) + let maskM1Tensor = graph.subtraction(maskTensor, oneTensor, name: nil) + let addition = graph.addition(sourceTensor, maskM1Tensor, name: nil) + + maxTensor = graph.reductionMaximum( + with: addition, + axes: hwAxes, + name: nil) + } + + resultTensor = graph.concatTensors( + [ + meanTensor, + meanMaskTensor, + maxTensor, + ], + dimension: channelAxis, + name: nil) + + assert(resultTensor.shape?.count == 4) + assert(resultTensor.shape?[2] == 1) + assert(resultTensor.shape?[3] == 1) + } +} + +/// A structure that represents a layer that performs global pooling on the input tensor +struct GlobalPoolingValueLayer { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + maskSumSqrtS14M01SquareS01Tensor: MPSGraphTensor + ) { + let hwAxes = InputShape.getHWAxes() + let channelAxis = InputShape.getChannelAxis() + + let sumTensor = graph.reductionSum( + with: sourceTensor, + axes: hwAxes, + name: nil) + + let meanTensor = graph.division(sumTensor, maskSumTensor, name: nil) + + let meanMaskTensor = graph.multiplication( + meanTensor, + maskSumSqrtS14M01Tensor, + name: nil) + + let meanMaskSquareTensor = graph.multiplication( + meanTensor, + maskSumSqrtS14M01SquareS01Tensor, + name: nil) + + resultTensor = graph.concatTensors( + [ + meanTensor, + meanMaskTensor, + meanMaskSquareTensor, + ], + dimension: channelAxis, + name: nil) + + assert(resultTensor.shape?.count == 4) + assert(resultTensor.shape?[2] == 1) + assert(resultTensor.shape?[3] == 1) + } +} + +// MARK: - Block Descriptors + +/// Base class for block descriptors +public class BlockDescriptor { +} + +/// A class that represents a residual block. +public class SWResidualBlockDesc: BlockDescriptor { + let preBN: SWBatchNormLayerDesc + let preActivation: ActivationKind + let regularConv: SWConvLayerDesc + let midBN: SWBatchNormLayerDesc + let midActivation: ActivationKind + let finalConv: SWConvLayerDesc + + init( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + regularConv: SWConvLayerDesc, + midBN: SWBatchNormLayerDesc, + midActivation: ActivationKind, + finalConv: SWConvLayerDesc + ) { + self.preBN = preBN + self.preActivation = preActivation + self.regularConv = regularConv + self.midBN = midBN + self.midActivation = midActivation + self.finalConv = finalConv + } +} + +public func createSWResidualBlockDesc( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + regularConv: SWConvLayerDesc, + midBN: SWBatchNormLayerDesc, + midActivation: ActivationKind, + finalConv: SWConvLayerDesc +) -> SWResidualBlockDesc { + return SWResidualBlockDesc( + preBN: preBN, + preActivation: preActivation, + regularConv: regularConv, + midBN: midBN, + midActivation: midActivation, + finalConv: finalConv) +} + +/// A class that represents a residual block with global pooling. +public class SWGlobalPoolingResidualBlockDesc: BlockDescriptor { + let preBN: SWBatchNormLayerDesc + let preActivation: ActivationKind + let regularConv: SWConvLayerDesc + let gpoolConv: SWConvLayerDesc + let gpoolBN: SWBatchNormLayerDesc + let gpoolActivation: ActivationKind + let gpoolToBiasMul: SWMatMulLayerDesc + let midBN: SWBatchNormLayerDesc + let midActivation: ActivationKind + let finalConv: SWConvLayerDesc + + init( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + regularConv: SWConvLayerDesc, + gpoolConv: SWConvLayerDesc, + gpoolBN: SWBatchNormLayerDesc, + gpoolActivation: ActivationKind, + gpoolToBiasMul: SWMatMulLayerDesc, + midBN: SWBatchNormLayerDesc, + midActivation: ActivationKind, + finalConv: SWConvLayerDesc + ) { + self.preBN = preBN + self.preActivation = preActivation + self.regularConv = regularConv + self.gpoolConv = gpoolConv + self.gpoolBN = gpoolBN + self.gpoolActivation = gpoolActivation + self.gpoolToBiasMul = gpoolToBiasMul + self.midBN = midBN + self.midActivation = midActivation + self.finalConv = finalConv + } +} + +public func createSWGlobalPoolingResidualBlockDesc( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + regularConv: SWConvLayerDesc, + gpoolConv: SWConvLayerDesc, + gpoolBN: SWBatchNormLayerDesc, + gpoolActivation: ActivationKind, + gpoolToBiasMul: SWMatMulLayerDesc, + midBN: SWBatchNormLayerDesc, + midActivation: ActivationKind, + finalConv: SWConvLayerDesc +) -> SWGlobalPoolingResidualBlockDesc { + return SWGlobalPoolingResidualBlockDesc( + preBN: preBN, + preActivation: preActivation, + regularConv: regularConv, + gpoolConv: gpoolConv, + gpoolBN: gpoolBN, + gpoolActivation: gpoolActivation, + gpoolToBiasMul: gpoolToBiasMul, + midBN: midBN, + midActivation: midActivation, + finalConv: finalConv) +} + +/// A class that represents a nested bottleneck residual block +public class SWNestedBottleneckResidualBlockDesc: BlockDescriptor { + let preBN: SWBatchNormLayerDesc + let preActivation: ActivationKind + let preConv: SWConvLayerDesc + let blockDescriptors: [BlockDescriptor] + let postBN: SWBatchNormLayerDesc + let postActivation: ActivationKind + let postConv: SWConvLayerDesc + + init( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + preConv: SWConvLayerDesc, + blockDescriptors: [BlockDescriptor], + postBN: SWBatchNormLayerDesc, + postActivation: ActivationKind, + postConv: SWConvLayerDesc + ) { + self.preBN = preBN + self.preActivation = preActivation + self.preConv = preConv + self.blockDescriptors = blockDescriptors + self.postBN = postBN + self.postActivation = postActivation + self.postConv = postConv + } +} + +public func createSWNestedBottleneckResidualBlockDesc( + preBN: SWBatchNormLayerDesc, + preActivation: ActivationKind, + preConv: SWConvLayerDesc, + blockDescriptors: [BlockDescriptor], + postBN: SWBatchNormLayerDesc, + postActivation: ActivationKind, + postConv: SWConvLayerDesc +) -> SWNestedBottleneckResidualBlockDesc { + return SWNestedBottleneckResidualBlockDesc( + preBN: preBN, + preActivation: preActivation, + preConv: preConv, + blockDescriptors: blockDescriptors, + postBN: postBN, + postActivation: postActivation, + postConv: postConv) +} + +public class BlockDescriptorBuilder { + public var blockDescriptors: [BlockDescriptor] = [] + + init() {} + + public func enque(with descriptor: BlockDescriptor) { + blockDescriptors.append(descriptor) + } +} + +public func createBlockDescriptorBuilder() -> BlockDescriptorBuilder { + return BlockDescriptorBuilder() +} + +// MARK: - Block Implementations + +/// A class that represents a Residual Block layer +class ResidualBlock { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + descriptor: SWResidualBlockDesc, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let preBN = BatchNormLayer( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + descriptor: descriptor.preBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let preActivation = ActivationLayer( + graph: graph, + sourceTensor: preBN.resultTensor, + activationKind: descriptor.preActivation) + + let regularConv = ConvLayer( + graph: graph, + sourceTensor: preActivation.resultTensor, + descriptor: descriptor.regularConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let midBN = BatchNormLayer( + graph: graph, + sourceTensor: regularConv.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.midBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let midActivation = ActivationLayer( + graph: graph, + sourceTensor: midBN.resultTensor, + activationKind: descriptor.midActivation) + + let finalConv = ConvLayer( + graph: graph, + sourceTensor: midActivation.resultTensor, + descriptor: descriptor.finalConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + resultTensor = graph.addition( + sourceTensor, + finalConv.resultTensor, + name: nil) + + assert(resultTensor.shape?.count == 4) + } +} + +/// A class representing a residual block with global pooling +class GlobalPoolingResidualBlock { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + descriptor: SWGlobalPoolingResidualBlockDesc, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let maskSum = MaskSumLayer(tensor: maskSumTensor) + let maskSumSqrtS14M01 = MaskSumSqrtS14M01Layer(tensor: maskSumSqrtS14M01Tensor) + + let preBN = BatchNormLayer( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + descriptor: descriptor.preBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let preActivation = ActivationLayer( + graph: graph, + sourceTensor: preBN.resultTensor, + activationKind: descriptor.preActivation) + + let regularConv = ConvLayer( + graph: graph, + sourceTensor: preActivation.resultTensor, + descriptor: descriptor.regularConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let gpoolConv = ConvLayer( + graph: graph, + sourceTensor: preActivation.resultTensor, + descriptor: descriptor.gpoolConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let gpoolBN = BatchNormLayer( + graph: graph, + sourceTensor: gpoolConv.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.gpoolBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let gpoolActivation = ActivationLayer( + graph: graph, + sourceTensor: gpoolBN.resultTensor, + activationKind: descriptor.gpoolActivation) + + let gpoolConcat = GlobalPoolingLayer( + graph: graph, + sourceTensor: gpoolActivation.resultTensor, + maskTensor: maskTensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + optimizeIdentityMask: optimizeIdentityMask) + + assert(gpoolConcat.resultTensor.shape?[1] == descriptor.gpoolToBiasMul.inChannels) + + let gpoolToBiasMul = MatMulLayer( + graph: graph, + descriptor: descriptor.gpoolToBiasMul, + sourceTensor: gpoolConcat.resultTensor) + + let added = AddNCBiasLayer( + graph: graph, + sourceTensor: regularConv.resultTensor, + biasTensor: gpoolToBiasMul.resultTensor, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: descriptor.gpoolToBiasMul.outChannels) + + let midBN = BatchNormLayer( + graph: graph, + sourceTensor: added.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.midBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let midActivation = ActivationLayer( + graph: graph, + sourceTensor: midBN.resultTensor, + activationKind: descriptor.midActivation) + + let finalConv = ConvLayer( + graph: graph, + sourceTensor: midActivation.resultTensor, + descriptor: descriptor.finalConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + resultTensor = graph.addition( + sourceTensor, + finalConv.resultTensor, + name: nil) + + assert(resultTensor.shape?.count == 4) + } +} + +/// A structure that represents a block stack +struct BlockStack { + let resultTensor: MPSGraphTensor + + static func processBlockDescriptors( + _ graph: MPSGraph, + _ sourceTensor: MPSGraphTensor, + _ maskTensor: MPSGraphTensor, + _ maskSumTensor: MPSGraphTensor, + _ maskSumSqrtS14M01Tensor: MPSGraphTensor, + _ blockDescriptors: [BlockDescriptor], + _ index: Int, + _ nnXLen: NSNumber, + _ nnYLen: NSNumber, + _ optimizeIdentityMask: Bool + ) -> MPSGraphTensor { + guard index < blockDescriptors.count else { + return sourceTensor + } + + let blockDescriptor = blockDescriptors[index] + let blockInput: MPSGraphTensor + + switch blockDescriptor { + case let globalPoolingDescriptor as SWGlobalPoolingResidualBlockDesc: + let globalPooling = GlobalPoolingResidualBlock( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + descriptor: globalPoolingDescriptor, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + blockInput = globalPooling.resultTensor + case let nestedBottleneckDescriptor as SWNestedBottleneckResidualBlockDesc: + let nestedBottleneck = NestedBottleneckResidualBlock( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + descriptor: nestedBottleneckDescriptor, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + blockInput = nestedBottleneck.resultTensor + case let residualBlockDescriptor as SWResidualBlockDesc: + let ordinary = ResidualBlock( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + descriptor: residualBlockDescriptor, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + blockInput = ordinary.resultTensor + default: + blockInput = sourceTensor + } + + return processBlockDescriptors( + graph, + blockInput, + maskTensor, + maskSumTensor, + maskSumSqrtS14M01Tensor, + blockDescriptors, + index + 1, + nnXLen, + nnYLen, + optimizeIdentityMask) + } + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + blockDescriptors: [BlockDescriptor], + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + resultTensor = BlockStack.processBlockDescriptors( + graph, + sourceTensor, + maskTensor, + maskSumTensor, + maskSumSqrtS14M01Tensor, + blockDescriptors, + 0, + nnXLen, + nnYLen, + optimizeIdentityMask) + } +} + +/// A structure that represents a nested bottleneck residual block +struct NestedBottleneckResidualBlock { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + descriptor: SWNestedBottleneckResidualBlockDesc, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let preBN = BatchNormLayer( + graph: graph, + sourceTensor: sourceTensor, + maskTensor: maskTensor, + descriptor: descriptor.preBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let preActivation = ActivationLayer( + graph: graph, + sourceTensor: preBN.resultTensor, + activationKind: descriptor.preActivation) + + let preConv = ConvLayer( + graph: graph, + sourceTensor: preActivation.resultTensor, + descriptor: descriptor.preConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let blocks = BlockStack( + graph: graph, + sourceTensor: preConv.resultTensor, + maskTensor: maskTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + blockDescriptors: descriptor.blockDescriptors, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let postBN = BatchNormLayer( + graph: graph, + sourceTensor: blocks.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.postBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let postActivation = ActivationLayer( + graph: graph, + sourceTensor: postBN.resultTensor, + activationKind: descriptor.postActivation) + + let postConv = ConvLayer( + graph: graph, + sourceTensor: postActivation.resultTensor, + descriptor: descriptor.postConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + resultTensor = graph.addition( + sourceTensor, + postConv.resultTensor, + name: nil) + + assert(resultTensor.shape?.count == 4) + } +} + +// MARK: - SGF Metadata Encoder + +/// Class representing the description of the SGF Metadata Encoder. +public class SWSGFMetadataEncoderDesc { + let version: Int + let numInputMetaChannels: Int + let mul1: SWMatMulLayerDesc + let bias1: SWMatBiasLayerDesc + let act1: ActivationKind + let mul2: SWMatMulLayerDesc + let bias2: SWMatBiasLayerDesc + let act2: ActivationKind + let mul3: SWMatMulLayerDesc + + init( + version: Int, + numInputMetaChannels: Int, + mul1: SWMatMulLayerDesc, + bias1: SWMatBiasLayerDesc, + act1: ActivationKind, + mul2: SWMatMulLayerDesc, + bias2: SWMatBiasLayerDesc, + act2: ActivationKind, + mul3: SWMatMulLayerDesc + ) { + self.version = version + self.numInputMetaChannels = numInputMetaChannels + self.mul1 = mul1 + self.bias1 = bias1 + self.act1 = act1 + self.mul2 = mul2 + self.bias2 = bias2 + self.act2 = act2 + self.mul3 = mul3 + } +} + +public func createSWSGFMetadataEncoderDesc( + version: Int32, + numInputMetaChannels: Int32, + mul1: SWMatMulLayerDesc, + bias1: SWMatBiasLayerDesc, + act1: ActivationKind, + mul2: SWMatMulLayerDesc, + bias2: SWMatBiasLayerDesc, + act2: ActivationKind, + mul3: SWMatMulLayerDesc +) -> SWSGFMetadataEncoderDesc? { + return SWSGFMetadataEncoderDesc( + version: Int(version), + numInputMetaChannels: Int(numInputMetaChannels), + mul1: mul1, + bias1: bias1, + act1: act1, + mul2: mul2, + bias2: bias2, + act2: act2, + mul3: mul3) +} + +/// A class that encodes SGF metadata. +class SGFMetadataEncoder { + let resultTensor: MPSGraphTensor + + init( + graph: MPSGraph, + descriptor: SWSGFMetadataEncoderDesc, + sourceTensor: MPSGraphTensor + ) { + let mul1 = MatMulLayer( + graph: graph, + descriptor: descriptor.mul1, + sourceTensor: sourceTensor) + + let bias1 = MatBiasLayer( + graph: graph, + descriptor: descriptor.bias1, + sourceTensor: mul1.resultTensor) + + let act1 = ActivationLayer( + graph: graph, + sourceTensor: bias1.resultTensor, + activationKind: descriptor.act1) + + let mul2 = MatMulLayer( + graph: graph, + descriptor: descriptor.mul2, + sourceTensor: act1.resultTensor) + + let bias2 = MatBiasLayer( + graph: graph, + descriptor: descriptor.bias2, + sourceTensor: mul2.resultTensor) + + let act2 = ActivationLayer( + graph: graph, + sourceTensor: bias2.resultTensor, + activationKind: descriptor.act2) + + let mul3 = MatMulLayer( + graph: graph, + descriptor: descriptor.mul3, + sourceTensor: act2.resultTensor) + + resultTensor = mul3.resultTensor + + assert(resultTensor.shape?.count == 2) + } +} + +// MARK: - Trunk + +/// A class that describes a trunk for a neural network +public class SWTrunkDesc { + let version: Int + let trunkNumChannels: NSNumber + let midNumChannels: NSNumber + let regularNumChannels: NSNumber + let gpoolNumChannels: NSNumber + let initialConv: SWConvLayerDesc + let initialMatMul: SWMatMulLayerDesc + let sgfMetadataEncoder: SWSGFMetadataEncoderDesc? + let blockDescriptors: [BlockDescriptor] + let trunkTipBN: SWBatchNormLayerDesc + let trunkTipActivation: ActivationKind + + init( + version: Int, + trunkNumChannels: NSNumber, + midNumChannels: NSNumber, + regularNumChannels: NSNumber, + gpoolNumChannels: NSNumber, + initialConv: SWConvLayerDesc, + initialMatMul: SWMatMulLayerDesc, + sgfMetadataEncoder: SWSGFMetadataEncoderDesc?, + blockDescriptors: [BlockDescriptor], + trunkTipBN: SWBatchNormLayerDesc, + trunkTipActivation: ActivationKind + ) { + self.version = version + self.trunkNumChannels = trunkNumChannels + self.midNumChannels = midNumChannels + self.regularNumChannels = regularNumChannels + self.gpoolNumChannels = gpoolNumChannels + self.initialConv = initialConv + self.initialMatMul = initialMatMul + self.sgfMetadataEncoder = sgfMetadataEncoder + self.blockDescriptors = blockDescriptors + self.trunkTipBN = trunkTipBN + self.trunkTipActivation = trunkTipActivation + } +} + +public func createSWTrunkDesc( + version: Int32, + trunkNumChannels: Int32, + midNumChannels: Int32, + regularNumChannels: Int32, + gpoolNumChannels: Int32, + initialConv: SWConvLayerDesc, + initialMatMul: SWMatMulLayerDesc, + sgfMetadataEncoder: SWSGFMetadataEncoderDesc?, + blockDescriptors: [BlockDescriptor], + trunkTipBN: SWBatchNormLayerDesc, + trunkTipActivation: ActivationKind +) -> SWTrunkDesc { + return SWTrunkDesc( + version: Int(version), + trunkNumChannels: trunkNumChannels as NSNumber, + midNumChannels: midNumChannels as NSNumber, + regularNumChannels: regularNumChannels as NSNumber, + gpoolNumChannels: gpoolNumChannels as NSNumber, + initialConv: initialConv, + initialMatMul: initialMatMul, + sgfMetadataEncoder: sgfMetadataEncoder, + blockDescriptors: blockDescriptors, + trunkTipBN: trunkTipBN, + trunkTipActivation: trunkTipActivation) +} + +/// A structure representing a ResNet trunk for a neural network +struct Trunk { + let resultTensor: MPSGraphTensor + + static func getBlockSourceTensor( + graph: MPSGraph, + descriptor: SWSGFMetadataEncoderDesc?, + initialAdd: AddNCBiasLayer, + inputMetaTensor: MPSGraphTensor?, + nnXLen: NSNumber, + nnYLen: NSNumber, + numChannels: NSNumber + ) -> MPSGraphTensor { + var blockSourceTensor: MPSGraphTensor + + if let inputMetaTensor, + let descriptor, descriptor.numInputMetaChannels > 0 + { + let encoded = SGFMetadataEncoder( + graph: graph, + descriptor: descriptor, + sourceTensor: inputMetaTensor) + + let encodedAdd = AddNCBiasLayer( + graph: graph, + sourceTensor: initialAdd.resultTensor, + biasTensor: encoded.resultTensor, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: numChannels) + + blockSourceTensor = encodedAdd.resultTensor + } else { + blockSourceTensor = initialAdd.resultTensor + } + + return blockSourceTensor + } + + init( + graph: MPSGraph, + descriptor: SWTrunkDesc, + inputTensor: MPSGraphTensor, + inputGlobalTensor: MPSGraphTensor, + inputMetaTensor: MPSGraphTensor?, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let initialConv = ConvLayer( + graph: graph, + sourceTensor: inputTensor, + descriptor: descriptor.initialConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let initialMatMul = MatMulLayer( + graph: graph, + descriptor: descriptor.initialMatMul, + sourceTensor: inputGlobalTensor) + + let initialAdd = AddNCBiasLayer( + graph: graph, + sourceTensor: initialConv.resultTensor, + biasTensor: initialMatMul.resultTensor, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: descriptor.initialMatMul.outChannels) + + let blockSourceTensor = Trunk.getBlockSourceTensor( + graph: graph, + descriptor: descriptor.sgfMetadataEncoder, + initialAdd: initialAdd, + inputMetaTensor: inputMetaTensor, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: descriptor.initialMatMul.outChannels) + + let blocks = BlockStack( + graph: graph, + sourceTensor: blockSourceTensor, + maskTensor: maskTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + blockDescriptors: descriptor.blockDescriptors, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let trunkTipBN = BatchNormLayer( + graph: graph, + sourceTensor: blocks.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.trunkTipBN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let trunkTipActivation = ActivationLayer( + graph: graph, + sourceTensor: trunkTipBN.resultTensor, + activationKind: descriptor.trunkTipActivation) + + resultTensor = trunkTipActivation.resultTensor + + assert(resultTensor.shape?.count == 4) + } +} + +// MARK: - Policy Head + +/// A class that describes a policy head for a neural network +public struct SWPolicyHeadDesc { + let version: Int + let p1Conv: SWConvLayerDesc + let g1Conv: SWConvLayerDesc + let g1BN: SWBatchNormLayerDesc + let g1Activation: ActivationKind + let gpoolToBiasMul: SWMatMulLayerDesc + let p1BN: SWBatchNormLayerDesc + let p1Activation: ActivationKind + let p2Conv: SWConvLayerDesc + let gpoolToPassMul: SWMatMulLayerDesc + let gpoolToPassBias: SWMatBiasLayerDesc? + let passActivation: ActivationKind? + let gpoolToPassMul2: SWMatMulLayerDesc? + + init( + version: Int, + p1Conv: SWConvLayerDesc, + g1Conv: SWConvLayerDesc, + g1BN: SWBatchNormLayerDesc, + g1Activation: ActivationKind, + gpoolToBiasMul: SWMatMulLayerDesc, + p1BN: SWBatchNormLayerDesc, + p1Activation: ActivationKind, + p2Conv: SWConvLayerDesc, + gpoolToPassMul: SWMatMulLayerDesc, + gpoolToPassBias: SWMatBiasLayerDesc?, + passActivation: ActivationKind?, + gpoolToPassMul2: SWMatMulLayerDesc? + ) { + self.version = version + self.p1Conv = p1Conv + self.g1Conv = g1Conv + self.g1BN = g1BN + self.g1Activation = g1Activation + self.gpoolToBiasMul = gpoolToBiasMul + self.p1BN = p1BN + self.p1Activation = p1Activation + self.p2Conv = p2Conv + self.gpoolToPassMul = gpoolToPassMul + self.gpoolToPassBias = gpoolToPassBias + self.passActivation = passActivation + self.gpoolToPassMul2 = gpoolToPassMul2 + + assert( + (version >= 15) + || ((gpoolToPassBias == nil) && (passActivation == nil) && (gpoolToPassMul2 == nil)) + ) + assert( + (version < 15) + || ((gpoolToPassBias != nil) && (passActivation != nil) && (gpoolToPassMul2 != nil)) + ) + } +} + +public func createSWPolicyHeadDesc( + version: Int32, + p1Conv: SWConvLayerDesc, + g1Conv: SWConvLayerDesc, + g1BN: SWBatchNormLayerDesc, + g1Activation: ActivationKind, + gpoolToBiasMul: SWMatMulLayerDesc, + p1BN: SWBatchNormLayerDesc, + p1Activation: ActivationKind, + p2Conv: SWConvLayerDesc, + gpoolToPassMul: SWMatMulLayerDesc, + gpoolToPassBias: SWMatBiasLayerDesc, + passActivation: ActivationKind, + gpoolToPassMul2: SWMatMulLayerDesc +) -> SWPolicyHeadDesc { + if version >= 15 { + return SWPolicyHeadDesc( + version: Int(version), + p1Conv: p1Conv, + g1Conv: g1Conv, + g1BN: g1BN, + g1Activation: g1Activation, + gpoolToBiasMul: gpoolToBiasMul, + p1BN: p1BN, + p1Activation: p1Activation, + p2Conv: p2Conv, + gpoolToPassMul: gpoolToPassMul, + gpoolToPassBias: gpoolToPassBias, + passActivation: passActivation, + gpoolToPassMul2: gpoolToPassMul2) + } else { + return SWPolicyHeadDesc( + version: Int(version), + p1Conv: p1Conv, + g1Conv: g1Conv, + g1BN: g1BN, + g1Activation: g1Activation, + gpoolToBiasMul: gpoolToBiasMul, + p1BN: p1BN, + p1Activation: p1Activation, + p2Conv: p2Conv, + gpoolToPassMul: gpoolToPassMul, + gpoolToPassBias: nil, + passActivation: nil, + gpoolToPassMul2: nil) + } +} + +/// A structure that represents a policy head of a neural network. +struct PolicyHead { + let policyTensor: MPSGraphTensor + let policyPassTensor: MPSGraphTensor + + init( + graph: MPSGraph, + descriptor: SWPolicyHeadDesc, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let p1Conv = ConvLayer( + graph: graph, + sourceTensor: sourceTensor, + descriptor: descriptor.p1Conv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let g1Conv = ConvLayer( + graph: graph, + sourceTensor: sourceTensor, + descriptor: descriptor.g1Conv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let g1BN = BatchNormLayer( + graph: graph, + sourceTensor: g1Conv.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.g1BN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let g1Activation = ActivationLayer( + graph: graph, + sourceTensor: g1BN.resultTensor, + activationKind: descriptor.g1Activation) + + let g1Concat = GlobalPoolingLayer( + graph: graph, + sourceTensor: g1Activation.resultTensor, + maskTensor: maskTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + optimizeIdentityMask: optimizeIdentityMask) + + assert(g1Concat.resultTensor.shape?[1] == descriptor.gpoolToBiasMul.inChannels) + + let gpoolToBiasMul = MatMulLayer( + graph: graph, + descriptor: descriptor.gpoolToBiasMul, + sourceTensor: g1Concat.resultTensor) + + let added = AddNCBiasLayer( + graph: graph, + sourceTensor: p1Conv.resultTensor, + biasTensor: gpoolToBiasMul.resultTensor, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: descriptor.gpoolToBiasMul.outChannels) + + let p1BN = BatchNormLayer( + graph: graph, + sourceTensor: added.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.p1BN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let p1Activation = ActivationLayer( + graph: graph, + sourceTensor: p1BN.resultTensor, + activationKind: descriptor.p1Activation) + + let p2Conv = ConvLayer( + graph: graph, + sourceTensor: p1Activation.resultTensor, + descriptor: descriptor.p2Conv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + policyTensor = p2Conv.resultTensor + + assert(g1Concat.resultTensor.shape?[1] == descriptor.gpoolToPassMul.inChannels) + + let gpoolToPassMul = MatMulLayer( + graph: graph, + descriptor: descriptor.gpoolToPassMul, + sourceTensor: g1Concat.resultTensor) + + if let gpoolToPassBias = descriptor.gpoolToPassBias, + let passActivation = descriptor.passActivation, + let gpoolToPassMul2 = descriptor.gpoolToPassMul2 + { + assert(descriptor.version >= 15) + + let gpoolToPassBiasLayer = MatBiasLayer( + graph: graph, + descriptor: gpoolToPassBias, + sourceTensor: gpoolToPassMul.resultTensor) + + let passActivationLayer = ActivationLayer( + graph: graph, + sourceTensor: gpoolToPassBiasLayer.resultTensor, + activationKind: passActivation) + + let gpoolToPassMul2Layer = MatMulLayer( + graph: graph, + descriptor: gpoolToPassMul2, + sourceTensor: passActivationLayer.resultTensor) + + policyPassTensor = gpoolToPassMul2Layer.resultTensor + } else { + assert(descriptor.version < 15) + policyPassTensor = gpoolToPassMul.resultTensor + } + + assert(policyTensor.shape?.count == 4) + assert(policyPassTensor.shape?.count == 2) + } +} + +// MARK: - Value Head + +/// A struct that describes the value head of a neural network +public struct SWValueHeadDesc { + let version: Int + let v1Conv: SWConvLayerDesc + let v1BN: SWBatchNormLayerDesc + let v1Activation: ActivationKind + let v2Mul: SWMatMulLayerDesc + let v2Bias: SWMatBiasLayerDesc + let v2Activation: ActivationKind + let v3Mul: SWMatMulLayerDesc + let v3Bias: SWMatBiasLayerDesc + let sv3Mul: SWMatMulLayerDesc + let sv3Bias: SWMatBiasLayerDesc + let vOwnershipConv: SWConvLayerDesc + + init( + version: Int, + v1Conv: SWConvLayerDesc, + v1BN: SWBatchNormLayerDesc, + v1Activation: ActivationKind, + v2Mul: SWMatMulLayerDesc, + v2Bias: SWMatBiasLayerDesc, + v2Activation: ActivationKind, + v3Mul: SWMatMulLayerDesc, + v3Bias: SWMatBiasLayerDesc, + sv3Mul: SWMatMulLayerDesc, + sv3Bias: SWMatBiasLayerDesc, + vOwnershipConv: SWConvLayerDesc + ) { + self.version = version + self.v1Conv = v1Conv + self.v1BN = v1BN + self.v1Activation = v1Activation + self.v2Mul = v2Mul + self.v2Bias = v2Bias + self.v2Activation = v2Activation + self.v3Mul = v3Mul + self.v3Bias = v3Bias + self.sv3Mul = sv3Mul + self.sv3Bias = sv3Bias + self.vOwnershipConv = vOwnershipConv + } +} + +public func createSWValueHeadDesc( + version: Int32, + v1Conv: SWConvLayerDesc, + v1BN: SWBatchNormLayerDesc, + v1Activation: ActivationKind, + v2Mul: SWMatMulLayerDesc, + v2Bias: SWMatBiasLayerDesc, + v2Activation: ActivationKind, + v3Mul: SWMatMulLayerDesc, + v3Bias: SWMatBiasLayerDesc, + sv3Mul: SWMatMulLayerDesc, + sv3Bias: SWMatBiasLayerDesc, + vOwnershipConv: SWConvLayerDesc +) -> SWValueHeadDesc { + return SWValueHeadDesc( + version: Int(version), + v1Conv: v1Conv, + v1BN: v1BN, + v1Activation: v1Activation, + v2Mul: v2Mul, + v2Bias: v2Bias, + v2Activation: v2Activation, + v3Mul: v3Mul, + v3Bias: v3Bias, + sv3Mul: sv3Mul, + sv3Bias: sv3Bias, + vOwnershipConv: vOwnershipConv) +} + +/// A structure that creates a value head for the neural network +struct ValueHead { + let valueTensor: MPSGraphTensor + let scoreValueTensor: MPSGraphTensor + let ownershipTensor: MPSGraphTensor + + init( + graph: MPSGraph, + descriptor: SWValueHeadDesc, + sourceTensor: MPSGraphTensor, + maskTensor: MPSGraphTensor, + maskSumTensor: MPSGraphTensor, + maskSumSqrtS14M01Tensor: MPSGraphTensor, + maskSumSqrtS14M01SquareS01Tensor: MPSGraphTensor, + nnXLen: NSNumber, + nnYLen: NSNumber, + optimizeIdentityMask: Bool = false + ) { + let v1Conv = ConvLayer( + graph: graph, + sourceTensor: sourceTensor, + descriptor: descriptor.v1Conv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let v1BN = BatchNormLayer( + graph: graph, + sourceTensor: v1Conv.resultTensor, + maskTensor: maskTensor, + descriptor: descriptor.v1BN, + nnXLen: nnXLen, + nnYLen: nnYLen, + optimizeIdentityMask: optimizeIdentityMask) + + let v1Activation = ActivationLayer( + graph: graph, + sourceTensor: v1BN.resultTensor, + activationKind: descriptor.v1Activation) + + let v1Mean = + GlobalPoolingValueLayer( + graph: graph, + sourceTensor: v1Activation.resultTensor, + maskSumTensor: maskSumTensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01Tensor, + maskSumSqrtS14M01SquareS01Tensor: maskSumSqrtS14M01SquareS01Tensor) + + assert(v1Mean.resultTensor.shape?[1] == descriptor.v2Mul.inChannels) + + let v2Mul = MatMulLayer( + graph: graph, + descriptor: descriptor.v2Mul, + sourceTensor: v1Mean.resultTensor) + + let v2Bias = MatBiasLayer( + graph: graph, + descriptor: descriptor.v2Bias, + sourceTensor: v2Mul.resultTensor) + + let v2Activation = ActivationLayer( + graph: graph, + sourceTensor: v2Bias.resultTensor, + activationKind: descriptor.v2Activation) + + let v3Mul = MatMulLayer( + graph: graph, + descriptor: descriptor.v3Mul, + sourceTensor: v2Activation.resultTensor) + + let v3Bias = MatBiasLayer( + graph: graph, + descriptor: descriptor.v3Bias, + sourceTensor: v3Mul.resultTensor) + + let sv3Mul = MatMulLayer( + graph: graph, + descriptor: descriptor.sv3Mul, + sourceTensor: v2Activation.resultTensor) + + let sv3Bias = MatBiasLayer( + graph: graph, + descriptor: descriptor.sv3Bias, + sourceTensor: sv3Mul.resultTensor) + + let vOwnershipConv = ConvLayer( + graph: graph, + sourceTensor: v1Activation.resultTensor, + descriptor: descriptor.vOwnershipConv, + nnXLen: nnXLen, + nnYLen: nnYLen) + + valueTensor = v3Bias.resultTensor + scoreValueTensor = sv3Bias.resultTensor + ownershipTensor = vOwnershipConv.resultTensor + + assert(valueTensor.shape?.count == 2) + assert(scoreValueTensor.shape?.count == 2) + assert(ownershipTensor.shape?.count == 4) + } +} + +// MARK: - Model Descriptor + +/// A struct that describes a neural network model used for playing the game of Go. +public struct SWModelDesc { + let version: Int + let name: String + let numInputChannels: NSNumber + let numInputGlobalChannels: NSNumber + let numInputMetaChannels: NSNumber + let numValueChannels: NSNumber + let numScoreValueChannels: NSNumber + let numOwnershipChannels: NSNumber + let trunk: SWTrunkDesc + let policyHead: SWPolicyHeadDesc + let valueHead: SWValueHeadDesc + + init( + version: Int, + name: String, + numInputChannels: NSNumber, + numInputGlobalChannels: NSNumber, + numInputMetaChannels: NSNumber, + numValueChannels: NSNumber, + numScoreValueChannels: NSNumber, + numOwnershipChannels: NSNumber, + trunk: SWTrunkDesc, + policyHead: SWPolicyHeadDesc, + valueHead: SWValueHeadDesc + ) { + self.version = version + self.name = name + self.numInputChannels = numInputChannels + self.numInputGlobalChannels = numInputGlobalChannels + self.numInputMetaChannels = numInputMetaChannels + self.numValueChannels = numValueChannels + self.numScoreValueChannels = numScoreValueChannels + self.numOwnershipChannels = numOwnershipChannels + self.trunk = trunk + self.policyHead = policyHead + self.valueHead = valueHead + } +} + +public func createSWModelDesc( + version: Int32, + name: String, + numInputChannels: Int32, + numInputGlobalChannels: Int32, + numInputMetaChannels: Int32, + numValueChannels: Int32, + numScoreValueChannels: Int32, + numOwnershipChannels: Int32, + trunk: SWTrunkDesc, + policyHead: SWPolicyHeadDesc, + valueHead: SWValueHeadDesc +) -> SWModelDesc { + return SWModelDesc( + version: Int(version), + name: name, + numInputChannels: numInputChannels as NSNumber, + numInputGlobalChannels: numInputGlobalChannels as NSNumber, + numInputMetaChannels: numInputMetaChannels as NSNumber, + numValueChannels: numValueChannels as NSNumber, + numScoreValueChannels: numScoreValueChannels as NSNumber, + numOwnershipChannels: numOwnershipChannels as NSNumber, + trunk: trunk, + policyHead: policyHead, + valueHead: valueHead) +} + +// MARK: - MPSGraph Model (for GPU inference) + +/// A structure representing a neural network model for processing Go game states using MPSGraph. +struct MPSGraphModel { + let device: MTLDevice + let commandQueue: MTLCommandQueue + let graph: MPSGraph + let nnXLen: NSNumber + let nnYLen: NSNumber + let version: Int + let numValueChannels: NSNumber + let numScoreValueChannels: NSNumber + let numOwnershipChannels: NSNumber + let input: InputLayer + let inputGlobal: InputGlobalLayer + let inputMeta: InputMetaLayer + let mask: MaskLayer + let trunk: Trunk + let policyHead: PolicyHead + let valueHead: ValueHead + let targetTensors: [MPSGraphTensor] + + init( + device: MTLDevice, + graph: MPSGraph, + descriptor: SWModelDesc, + nnXLen: NSNumber, + nnYLen: NSNumber + ) { + self.device = device + self.commandQueue = device.makeCommandQueue()! + self.graph = graph + self.nnXLen = nnXLen + self.nnYLen = nnYLen + self.version = descriptor.version + self.numValueChannels = descriptor.numValueChannels + self.numScoreValueChannels = descriptor.numScoreValueChannels + self.numOwnershipChannels = descriptor.numOwnershipChannels + + input = InputLayer( + graph: graph, + nnXLen: nnXLen, + nnYLen: nnYLen, + numChannels: descriptor.numInputChannels) + + inputGlobal = InputGlobalLayer( + graph: graph, + numGlobalFeatures: descriptor.numInputGlobalChannels) + + inputMeta = InputMetaLayer( + graph: graph, + numMetaFeatures: descriptor.numInputMetaChannels) + + mask = MaskLayer( + graph: graph, + nnXLen: nnXLen, + nnYLen: nnYLen) + + let maskSum = MaskSumLayer( + graph: graph, + maskTensor: mask.tensor) + + let maskSumSqrtS14M01 = MaskSumSqrtS14M01Layer( + graph: graph, + maskSum: maskSum) + + let maskSumSqrtS14M01SquareS01 = MaskSumSqrtS14M01SquareS01Layer( + graph: graph, + maskSumSqrtS14M01: maskSumSqrtS14M01) + + trunk = Trunk( + graph: graph, + descriptor: descriptor.trunk, + inputTensor: input.tensor, + inputGlobalTensor: inputGlobal.tensor, + inputMetaTensor: inputMeta.tensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + nnXLen: nnXLen, + nnYLen: nnYLen) + + policyHead = PolicyHead( + graph: graph, + descriptor: descriptor.policyHead, + sourceTensor: trunk.resultTensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + nnXLen: nnXLen, + nnYLen: nnYLen) + + valueHead = ValueHead( + graph: graph, + descriptor: descriptor.valueHead, + sourceTensor: trunk.resultTensor, + maskTensor: mask.tensor, + maskSumTensor: maskSum.tensor, + maskSumSqrtS14M01Tensor: maskSumSqrtS14M01.tensor, + maskSumSqrtS14M01SquareS01Tensor: maskSumSqrtS14M01SquareS01.tensor, + nnXLen: nnXLen, + nnYLen: nnYLen) + + targetTensors = [ + policyHead.policyTensor, + policyHead.policyPassTensor, + valueHead.valueTensor, + valueHead.scoreValueTensor, + valueHead.ownershipTensor, + ] + } + + /// Applies the model to the given input data + public func apply( + input inputPointer: UnsafeMutablePointer, + inputGlobal inputGlobalPointer: UnsafeMutablePointer, + inputMeta inputMetaPointer: UnsafeMutablePointer, + policy: UnsafeMutablePointer, + policyPass: UnsafeMutablePointer, + value: UnsafeMutablePointer, + scoreValue: UnsafeMutablePointer, + ownership: UnsafeMutablePointer, + batchSize: Int + ) { + let channelAxis = InputShape.getChannelAxis() + let numInputChannels = input.shape[channelAxis] + + let inputShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputChannels, + nnYLen: nnYLen, + nnXLen: nnXLen) + + let inputDescriptor = MPSNDArrayDescriptor( + dataType: input.tensor.dataType, + shape: inputShape) + + let inputArray = MPSNDArray( + device: device, + descriptor: inputDescriptor) + + inputArray.writeBytes(inputPointer) + + let numInputGlobalChannels = inputGlobal.shape[channelAxis] + + let inputGlobalShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputGlobalChannels, + nnYLen: 1, + nnXLen: 1) + + let inputGlobalDescriptor = MPSNDArrayDescriptor( + dataType: inputGlobal.tensor.dataType, + shape: inputGlobalShape) + + let inputGlobalArray = MPSNDArray( + device: device, + descriptor: inputGlobalDescriptor) + + inputGlobalArray.writeBytes(inputGlobalPointer) + + let numInputMetaChannels = inputMeta.shape[channelAxis] + + let inputMetaShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: numInputMetaChannels, + nnYLen: 1, + nnXLen: 1) + + let inputMetaDescriptor = MPSNDArrayDescriptor( + dataType: inputMeta.tensor.dataType, + shape: inputMetaShape) + + let inputMetaArray = MPSNDArray( + device: device, + descriptor: inputMetaDescriptor) + + inputMetaArray.writeBytes(inputMetaPointer) + + let maskShape = InputShape.create( + batchSize: batchSize as NSNumber, + numChannels: 1, + nnYLen: nnYLen, + nnXLen: nnXLen) + + let maskDescriptor = MPSNDArrayDescriptor( + dataType: mask.tensor.dataType, + shape: maskShape) + + let maskArray = MPSNDArray( + device: device, + descriptor: maskDescriptor) + + var maskStrideArray = [ + MemoryLayout.size, + nnXLen.intValue * MemoryLayout.size, + nnYLen.intValue * nnXLen.intValue * MemoryLayout.size, + numInputChannels.intValue * nnYLen.intValue * nnXLen.intValue + * MemoryLayout.size, + ] + + maskArray.writeBytes(inputPointer, strideBytes: &maskStrideArray) + + let feeds = [ + input.tensor: MPSGraphTensorData(inputArray), + inputGlobal.tensor: MPSGraphTensorData(inputGlobalArray), + inputMeta.tensor: MPSGraphTensorData(inputMetaArray), + mask.tensor: MPSGraphTensorData(maskArray), + ] + + let fetch = graph.run( + with: commandQueue, + feeds: feeds, + targetTensors: targetTensors, + targetOperations: nil) + + assert(fetch[policyHead.policyTensor] != nil) + assert(fetch[policyHead.policyPassTensor] != nil) + assert(fetch[valueHead.valueTensor] != nil) + assert(fetch[valueHead.scoreValueTensor] != nil) + assert(fetch[valueHead.ownershipTensor] != nil) + + fetch[policyHead.policyTensor]?.mpsndarray().readBytes(policy) + fetch[policyHead.policyPassTensor]?.mpsndarray().readBytes(policyPass) + fetch[valueHead.valueTensor]?.mpsndarray().readBytes(value) + fetch[valueHead.scoreValueTensor]?.mpsndarray().readBytes(scoreValue) + fetch[valueHead.ownershipTensor]?.mpsndarray().readBytes(ownership) + } +} diff --git a/cpp/program/setup.cpp b/cpp/program/setup.cpp index a31295145..e271b3321 100644 --- a/cpp/program/setup.cpp +++ b/cpp/program/setup.cpp @@ -18,6 +18,7 @@ std::vector Setup::getBackendPrefixes() { prefixes.push_back("cuda"); prefixes.push_back("trt"); prefixes.push_back("metal"); + prefixes.push_back("coreml"); prefixes.push_back("opencl"); prefixes.push_back("eigen"); prefixes.push_back("dummybackend"); @@ -84,6 +85,8 @@ vector Setup::initializeNNEvaluators( string backendPrefix = "trt"; #elif defined(USE_METAL_BACKEND) string backendPrefix = "metal"; + #elif defined(USE_COREML_BACKEND) + string backendPrefix = "coreml"; #elif defined(USE_OPENCL_BACKEND) string backendPrefix = "opencl"; #elif defined(USE_EIGEN_BACKEND)