Intel Mac 에서 Intel GPU 로 Pytorch MPS 사용하기

if command -v protoc &> /dev/null
then
    echo "protoc command exists, please remove it first"
    echo "\tbrew uninstall protobuf; pip uninstall protobuf"
    exit
fi

USE_NATIVE_ARCH=1 \
    USE_DISTRIBUTED=0 \
    USE_MKLDNN=OFF \
    USE_FBGEMM=OFF \
    USE_NNPACK=OFF \
    USE_QNNPACK=OFF \
    USE_XNNPACK=OFF \
    USE_CUDA=0 \
    BUILD_TEST=0 \
    BUILD_CAFFE2=0 \
    BUILD_CAFFE2_OPS=0 \
    CPLUS_INCLUDE_PATH=$(python -c "import numpy as np; print(np.get_include())") \
    CMAKE_C_COMPILER_LAUNCHER=ccache \
    CMAKE_CXX_COMPILER_LAUNCHER=ccache \
    CMAKE_CUDA_COMPILER_LAUNCHER=ccache \
    CCACHE_COMPILERTYPE=clang \
    CC="clang" \
    CXX="clang++" \
    MAX_JOBS=8 \
    MACOSX_DEPLOYMENT_TARGET=14.0 \
    USE_MPS=1 \
    CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"} \
    python setup.py bdist_wheel -p macosx_14_0_x86_64 --cmake

#import <Metal/Metal.h>
#import <MetalPerformanceShaders/MetalPerformanceShaders.h>
#import <MetalPerformanceShadersGraph/MetalPerformanceShadersGraph.h>

// Objective-C wrapper for MPSGraph matrix multiplication operation

#import <MetalPerformanceShadersGraph/MPSGraph.h>

@interface MPSGraph (MatrixMultiplication)

- (MPSGraphTensor *)matrixMultiplicationWithPrimaryTensor2:(MPSGraphTensor *)selfTensor secondaryTensor:(MPSGraphTensor *)otherTensor name:(NSString * _Nullable)name;

@end

@implementation MPSGraph (MatrixMultiplication)

- (MPSGraphTensor *)_matrixMultiplicationWithPrimaryTensor_impl_3_3:(MPSGraphTensor *)selfTensor secondaryTensor:(MPSGraphTensor *)otherTensor name:(NSString * _Nullable)name {
    @autoreleasepool {
        // replace matmul with group conv2d

        // check length of selfTensor.shape == 3
        if (selfTensor.shape.count != 3) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) selfTensor.shape.count (== %lu) != 3", selfTensor.shape.count);
            return NULL;
        }

        // check length of otherTensor.shape == 3
        if (otherTensor.shape.count != 3) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) otherTensor.shape.count (== %lu) != 3", otherTensor.shape.count);
            return NULL;
        }

        // assert selfTensor.shape[2] == otherTensor.shape[1]
        if (![selfTensor.shape[2] isEqualToNumber:otherTensor.shape[1]]) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) selfTensor.shape[2] (== %@) != otherTensor.shape[1] (== %@)", selfTensor.shape[2], otherTensor.shape[1]);
            return NULL;
        }

        NSNumber *numBatch = selfTensor.shape[0];
        NSNumber *selfTensorM = selfTensor.shape[1];
        NSNumber *selfTensorK = selfTensor.shape[2];
        NSNumber *otherTensorK = otherTensor.shape[1];
        NSNumber *otherTensorN = otherTensor.shape[2];
        NSInteger axis = 0;

        // repeat _matrixMultiplicationWithPrimaryTensor_impl_2_2 numBatch times
        NSMutableArray<MPSGraphTensor *> *resultTensorArray = [NSMutableArray array];
        NSArray<MPSGraphTensor *> *splitSelfTensorArray = [self splitTensor:selfTensor numSplits:numBatch.intValue axis:axis name:NULL];
        NSArray<MPSGraphTensor *> *splitOtherTensorArray = [self splitTensor:otherTensor numSplits:numBatch.intValue axis:axis name:NULL];
        for (int i = 0; i < [numBatch intValue]; i++) {
            MPSGraphTensor *_a = [self reshapeTensor:splitSelfTensorArray[i] withShape:@[selfTensorM, selfTensorK] name:NULL];
            MPSGraphTensor *_b = [self reshapeTensor:splitOtherTensorArray[i] withShape:@[otherTensorK, otherTensorN] name:NULL];
            MPSGraphTensor *_resultTensor = [self _matrixMultiplicationWithPrimaryTensor_impl_2_2:_a secondaryTensor:_b name:NULL];
            _resultTensor = [self reshapeTensor:_resultTensor withShape:@[@1, selfTensorM, otherTensorN] name:NULL];
            [resultTensorArray addObject:_resultTensor];
        }

        // concat result tensors
        MPSGraphTensor *resultTensor = [self concatTensors:resultTensorArray dimension:axis name:name];

        return resultTensor;
    }
}

- (MPSGraphTensor *)_matrixMultiplicationWithPrimaryTensor_impl_2_2:(MPSGraphTensor *)selfTensor secondaryTensor:(MPSGraphTensor *)otherTensor name:(NSString * _Nullable)name {
    @autoreleasepool {
        // replace matmul with conv2d

        // check length of selfTensor.shape == 2
        if (selfTensor.shape.count != 2) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) selfTensor.shape.count (== %lu) != 2", selfTensor.shape.count);
            return NULL;
        }
        // assert length of otherTensor.shape == 2
        if (otherTensor.shape.count != 2) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) otherTensor.shape.count (== %lu) != 2", otherTensor.shape.count);
            return NULL;
        }
        // assert selfTensor.shape[1] == otherTensor.shape[0]
        if (![selfTensor.shape[1] isEqualToNumber:otherTensor.shape[0]]) {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) selfTensor.shape[1] (== %@) != otherTensor.shape[0] (== %@)", selfTensor.shape[1], otherTensor.shape[0]);
            return NULL;
        }

        NSNumber *selfTensorM = selfTensor.shape[0];
        NSNumber *selfTensorK = selfTensor.shape[1];
        NSNumber *otherTensorK = otherTensor.shape[0];
        NSNumber *otherTensorN = otherTensor.shape[1];

        // 1. create MPSGraphConvolution2DOpDescriptor
        MPSGraphConvolution2DOpDescriptor *_mpsGraphConv2dOpDescriptor = [MPSGraphConvolution2DOpDescriptor descriptorWithStrideInX:1
                                                                        strideInY:1
                                                                        dilationRateInX:1
                                                                        dilationRateInY:1
                                                                        groups:1
                                                                        paddingStyle:MPSGraphPaddingStyleTF_VALID
                                                                        dataLayout:MPSGraphTensorNamedDataLayoutNHWC
                                                                        weightsLayout:MPSGraphTensorNamedDataLayoutHWIO];

        // 2. reshape tensors
        MPSGraphTensor *selfTensorReshaped = [self reshapeTensor:selfTensor withShape:@[selfTensorM, @1, @1, selfTensorK] name:NULL];
        MPSGraphTensor *otherTensorReshaped = [self reshapeTensor:otherTensor withShape:@[@1, @1, otherTensorK, otherTensorN] name:NULL];

        // 3. create MPSGraphConvolution2DOp
        MPSGraphTensor *resultTensor = [self convolution2DWithSourceTensor:selfTensorReshaped
                                                                    weightsTensor:otherTensorReshaped
                                                                    descriptor:_mpsGraphConv2dOpDescriptor
                                                                    name:NULL];

        // 4. reshape result tensor
        resultTensor = [self reshapeTensor:resultTensor withShape:@[selfTensorM, otherTensorN] name:name];

        return resultTensor;
    }
}

- (MPSGraphTensor *)matrixMultiplicationWithPrimaryTensor2:(MPSGraphTensor *)selfTensor secondaryTensor:(MPSGraphTensor *)otherTensor name:(NSString * _Nullable)name {
    @autoreleasepool {
        if (selfTensor.shape.count == 2 && otherTensor.shape.count == 2) {
            return [self _matrixMultiplicationWithPrimaryTensor_impl_2_2:selfTensor secondaryTensor:otherTensor name:name];
        } else if (selfTensor.shape.count == 3 && otherTensor.shape.count == 3) {
            return [self _matrixMultiplicationWithPrimaryTensor_impl_3_3:selfTensor secondaryTensor:otherTensor name:name];
        } else {
            NSLog(@"(matrixMultiplicationWithPrimaryTensor2) selfTensor.shape.count (== %lu) != 2 or 3", selfTensor.shape.count);
            return NULL;
        }
    }
}

@end