利用 iMac Mini M4 晶片撰寫 GPU 程式

蘋果的 M 系列晶片採用所謂的「統一記憶體架構 (Unified Memory Architecture, UMA)」，與一般我們在NVIDIA CUDA 常搭配使用的獨立顯示卡記憶體架構之間有所不同。主要的差異在於記憶體的物理位置 以及 CPU 與 GPU 如何存取這些記憶體。蘋果是將 CPU、GPU、神經網路引擎 (Neural Engine) 等處理單元都整合在同一個 SoC (System on a Chip) 晶片上。它們共享同一塊物理記憶體池。記憶體存取: CPU 和 GPU 可以直接存取相同的物理記憶體位址，不需要將資料從一個記憶體區域複製到另一個區域；這樣的作法在單晶片中 CPU 以及 GPU 的資料傳輸簡單而且快速；方便在小規模的人工智慧程式開發上意外地有效。在程式上，也可以透過 pyTorch 來撰寫蘋果的 GPU 程式。

• 建立測試執行緒

首先先定義兩個執行緒，分別各自在 CPU 以及 CPU 上來執行。

import torch
import numpy as np
import time
import threading
import warnings

#-----------------------
# 定義在 CPU 執行的執行緒
class cpuThread(threading.Thread):
    def __init__(self, x,y,count):
        threading.Thread.__init__(self)
        self.x=x
        self.y=y
        self.ans=x
        self.count=count
    def run(self):
        for i in range(self.count):
            self.ans=np.matmul(self.x,self.y)
# 定義在 GPU 執行的執行緒
class gpuThread(threading.Thread):
    def __init__(self, x,y,count):
        threading.Thread.__init__(self)
        self.x=x
        self.y=y
        self.count=count
        self.ans=x
    def run(self):
        for i in range(self.count):
            self.ans=torch.matmul(self.x,self.y)
#-----------------------

• 指定執行的 CPU 及 GPU

接下來，先建立要用來測試的資料，分別放在 CPU 以及透過 PyTorch 的「device('mps')」指令，指定到 GPU。

device = torch.device('mps')
print('GPU count:',torch.mps.device_count())

# 創建執行緒物件
# 起始在 CPU 及 GPU 的亂數矩陣
np.random.seed(13)
matrixA=np.random.rand(10000,10000).astype('float32')
matrixB=np.random.rand(10000,10000).astype('float32')
tensorA=torch.tensor(matrixA).to(device)
tensorB=torch.tensor(matrixB).to(device)
#-----------------------

• 進行測試

然後，我們就可以來測試看看，在蘋果的 iMac Mini M4 晶片上，透過「pyTorch」執行 CPU 及以 GPU 的效果。

#-----------------------
# 單獨 CPU 計算時間
beginTime=time.time()
runCPU=cpuThread(matrixA,matrixB,10)
runCPU.start()
runCPU.join()
print('CPU execution time:',time.time()-beginTime)
#-----------------------
# 單獨 GPU 計算時間
beginTime=time.time()
runGPU=gpuThread(tensorA,tensorB,10)
runGPU.start()
runGPU.join()
torch.mps.synchronize() # 等所有 GPU 都計算完畢
print('GPU execution time:',time.time()-beginTime)

iMac mini M4 的執行結果

執行的結果，單純的矩陣相乘運算；GPU 大概可以比 CPU 快了四倍；當然，這中間還有「pyTorch」程式庫可以再優化的空間。不過相較於價格相近同樣是單晶片架構的 nVidia Jetson Orin 來說，這樣的速度真的已經非常驚人了，大約相差了五倍；在教學和工程研究上，真的是太值得了。

nVidia Jetson Orin 的執行結果