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

git clone https://github.com/pytorch/pytorch.git ~/ws/pytorch
cd ~/ws/pytorch
git checkout v2.1.1
git submodule --init --recursive

// https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/aten/src/ATen/mps/MPSAllocator.h
// #L244
// m_scalar_pool(m_device, UsageFlags::SMALL | UsageFlags::SHARED | UsageFlags::SCALAR),
m_scalar_pool(m_device, UsageFlags::SMALL | UsageFlags::PRIVATE | UsageFlags::SCALAR),

// https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/aten/src/ATen/mps/MPSDevice.mm
// #L96
// if ([device isLowPower]) { // exclude Intel GPUs
    if (true) { // include Intel GPUs
      _mtl_device = [device retain];
      break;
    }

// https://github.com/pytorch/pytorch/blob/4c55dc50355d5e923642c59ad2a23d6ad54711e7/aten/src/ATen/native/mps/OperationUtils.mm
// #L391
// Scalar pools are only supported on devices with unified memory
  if (false) {
    scalar.buffer = getIMPSAllocator()->allocScalarBufferWithValue(&scalar.value, scalar.size);
    result = [[[MPSGraphTensorData alloc] initWithMTLBuffer:scalar.getMTLBuffer()
                                                      shape:@[ @1 ]
                                                   dataType:getMPSScalarType(scalar.type)] autorelease];

conda create -n pytorch-2.1.1-py11 python=3.11
conda activate pytorch-2.1.1-py11
pip install -r requirements.txt
conda install -y mkl mkl-include
conda install -y pkg-config libuv

CPLUS_INCLUDE_PATH=$(python -c "import numpy as np; print(np.get_include())") \
    MAX_JOBS=8 \
    MACOSX_DEPLOYMENT_TARGET=14.1 \
    USE_MPS=1 \
		python setup.py bdist_wheel -p macosx_14_0_x86_64

python -c "import torch; print(torch.backends.mps.is_available(), torch.backends.mps.is_built()); torch.randn(1, 3).to('mps')"
# True True

python -c "import platform; print(platform.platform())"
# macOS-14.1-x86_64-i386-64bit

# Mac 에서 python 설치
# 구 intel mac 에서 pytorch mps 사용하기 위해 python 3.11 직접 빌드하는 용도

mkdir -p ~/ws
git clone https://github.com/python/cpython ~/ws/cpython
cd ~/ws/cpython
git checkout 3.11

GDBM_CFLAGS="-I$(brew --prefix gdbm)/include" \
    GDBM_LIBS="-L$(brew --prefix gdbm)/lib -lgdbm" \
    ac_cv_working_openssl_ssl=yes \
    ac_cv_working_openssl_hashlib=yes \
    CPPFLAGS="-I$(brew --prefix openssl@3.0)/include" \  
    LDFLAGS="-L$(brew --prefix openssl@3.0)/lib" \
    PKG_CONFIG_PATH="$(brew --prefix openssl@3.0)/lib/pkgconfig" \
    ./configure \
                --with-openssl="$(brew --prefix openssl@3.0)" \
                --enable-framework \
                --enable-optimizations \
                --with-lto

echo "export CPLUS_INCLUDE_PATH=/Library/Frameworks/Python.framework/Versions/3.11/Headers" >> ~/.setting

make -s -j8
sudo make install
sudo ln -s -f /usr/local/bin/python3.11 /usr/local/bin/python
sudo ln -s -f /usr/local/bin/python3.11 /usr/local/bin/python3
sudo ln -s -f /usr/local/bin/pip3.11 /usr/local/bin/pip
sudo ln -s -f /usr/local/bin/pip3.11 /usr/local/bin/pip3

import torch
import timeit

x = torch.randn(1000, 1000)
print(timeit.timeit("x @ x", number=10, globals=globals()))
# 0.0907975767040625
y = x.to("mps")
print(timeit.timeit("y @ y; torch.mps.synchronize()", number=10, globals=globals()))
# 4.1007534249999935 (very slow)

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

@interface MatMul : NSObject {
    MPSGraph *_graph;
    MPSGraphTensor *_aTensor;
    MPSGraphTensor *_bTensor;
    MPSGraphTensor *_resultTensor;
}

- (instancetype)initWithDevice:(id<MTLDevice>)device;
- (void)matmul:(id<MTLBuffer>)A B:(id<MTLBuffer>)B C:(id<MTLBuffer>)C aShape:(MPSShape *)aShape bShape:(MPSShape *)bShape cShape:(MPSShape *)cShape;

@end

@implementation MatMul

- (instancetype)initWithDevice:(id<MTLDevice>)device
{
    if (self = [super init]) {
        _graph = [[MPSGraph alloc] init];
    }
    return self;
}

- (void)matmul:(id<MTLBuffer>)A B:(id<MTLBuffer>)B C:(id<MTLBuffer>)C aShape:(MPSShape *)aShape bShape:(MPSShape *)bShape cShape:(MPSShape *)cShape
{
    @autoreleasepool {
	    MPSDataType dataType = MPSDataTypeFloat32;
	
	    MPSGraphTensorData *aTensorData = [[MPSGraphTensorData alloc] initWithMTLBuffer:A shape:aShape dataType:dataType];
	    MPSGraphTensorData *bTensorData = [[MPSGraphTensorData alloc] initWithMTLBuffer:B shape:bShape dataType:dataType];
	    MPSGraphTensorData *cTensorData = [[MPSGraphTensorData alloc] initWithMTLBuffer:C shape:cShape dataType:dataType];
	
	    _aTensor = [_graph placeholderWithShape:aShape dataType:dataType name:NULL];
	    _bTensor = [_graph placeholderWithShape:bShape dataType:dataType name:NULL];
	
	    _resultTensor = [_graph matrixMultiplicationWithPrimaryTensor:_aTensor secondaryTensor:_bTensor name:NULL];
	
	    MPSGraphTensorDataDictionary *inputs = @{_aTensor: aTensorData, _bTensor: bTensorData};
	
	    MPSGraphTensorDataDictionary *result = [_graph runWithFeeds:inputs targetTensors:@[_resultTensor] targetOperations:@[[_resultTensor operation]]];
	
	    MPSGraphTensorData *resultData = result[_resultTensor];
	    [resultData.mpsndarray readBytes:C.contents strideBytes:nil];
	}
}

@end


int main(int argc, const char * argv[]) {
    @autoreleasepool {
        id<MTLDevice> device = MTLCreateSystemDefaultDevice();

        // shape
        MPSShape *aShape = @[@1, @128];
        MPSShape *bShape = @[@128, @1];
        MPSShape *cShape = @[@1, @1];

        unsigned long aSize = (unsigned long)[aShape[0] intValue] * (unsigned long)[aShape[1] intValue];
        unsigned long bSize = (unsigned long)[bShape[0] intValue] * (unsigned long)[bShape[1] intValue];
        unsigned long cSize = (unsigned long)[cShape[0] intValue] * (unsigned long)[cShape[1] intValue];

        float *A = (float *)malloc(sizeof(float) * aSize);
        float *B = (float *)malloc(sizeof(float) * bSize);
        float *C = (float *)malloc(sizeof(float) * cSize);

        for (int i = 0; i < aSize; i++) {
            A[i] = 1;
        }
        for (int i = 0; i < bSize; i++) {
            B[i] = 1;
        }

        id<MTLBuffer> A_buffer = [device newBufferWithBytes:A length:sizeof(float) * aSize options:MTLResourceStorageModeShared];
        id<MTLBuffer> B_buffer = [device newBufferWithBytes:B length:sizeof(float) * bSize options:MTLResourceStorageModeShared];
        id<MTLBuffer> C_buffer = [device newBufferWithBytes:C length:sizeof(float) * cSize options:MTLResourceStorageModeShared];

        MatMul *matmul = [[MatMul alloc] initWithDevice:device];
        [matmul matmul:A_buffer B:B_buffer C:C_buffer aShape:aShape bShape:bShape cShape:cShape];

        memcpy(A, A_buffer.contents, sizeof(float) * aSize);
        memcpy(B, B_buffer.contents, sizeof(float) * bSize);
        memcpy(C, C_buffer.contents, sizeof(float) * cSize);

        for (int i = 0; i < aSize; i++) {
            printf("%f ", A[i]);
        }
        printf("\n");
        for (int i = 0; i < bSize; i++) {
            printf("%f ", B[i]);
        }
        printf("\n");
        for (int i = 0; i < cSize; i++) {
            printf("%f ", C[i]);
        }
        printf("\n");
        free(A);
        free(B);
        free(C);
    }
    return 0;
}