統一記憶體

目錄

統一記憶體#

Apple silicon 採用統一記憶體架構。CPU 與 GPU 可直接存取同一個記憶體池。MLX 的設計即是為了善用這項特性。

更具體地說，在 MLX 中建立陣列時不需要指定其位置：

a = mx.random.normal((100,))
b = mx.random.normal((100,))

a 和 b 皆存在於統一記憶體中。

在 MLX 中，你不需要把陣列搬移到裝置上，而是在執行運算時指定裝置。任何裝置都能在 a 與 b 上執行任何運算，而不必在不同記憶體位置間搬移。例如：

mx.add(a, b, stream=mx.cpu)
mx.add(a, b, stream=mx.gpu)

在上例中，CPU 與 GPU 都會執行相同的加法運算。因為兩者之間沒有相依性，所以這些運算可以（而且很可能會）並行執行。關於 MLX 中串流語意的更多資訊，請參考 使用串流。

在上述 add 範例中，運算之間沒有相依性，因此不會產生競態條件。若存在相依性，MLX 排程器會自動管理。例如：

c = mx.add(a, b, stream=mx.cpu)
d = mx.add(a, c, stream=mx.gpu)

在上述情況下，第二個 add 在 GPU 上執行，但它相依於第一個在 CPU 上執行的 add 的輸出。MLX 會自動在兩個串流間插入相依性，讓第二個 add 只有在第一個完成且 c 可用後才開始執行。

簡單範例#

以下是更有趣（雖然稍微刻意）的例子，說明統一記憶體如何提供幫助。假設我們有以下計算：

def fun(a, b, d1, d2):
  x = mx.matmul(a, b, stream=d1)
  for _ in range(500):
      b = mx.exp(b, stream=d2)
  return x, b

我們希望使用下列參數執行：

a = mx.random.uniform(shape=(4096, 512))
b = mx.random.uniform(shape=(512, 4))

第一個 matmul 運算計算密度高，適合在 GPU 上執行。第二段運算序列非常小，在 GPU 上可能受限於開銷，因此更適合在 CPU 上執行。

如果完全在 GPU 上計時，耗時為 2.8 毫秒。但若以 d1=mx.gpu、d2=mx.cpu 執行，時間約為 1.4 毫秒，快了約兩倍。這些數據是在 M1 Max 上測得。

目錄