统一内存

统一内存#

Apple 芯片采用统一内存架构。CPU 和 GPU 可以直接访问同一内存池。MLX 的设计旨在充分利用这一优势。

具体来说，在 MLX 中创建数组时，您无需指定其位置

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

Both a and b live in unified memory.

在 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，因为它计算密度更高。第二系列操作更适合 CPU，因为它们非常小，在 GPU 上可能会受限于开销。

如果我们完全在 GPU 上计时此计算，需要 2.8 毫秒。但如果我们在 d1=mx.gpu 和 d2=mx.cpu 的设置下运行计算，则时间约为 1.4 毫秒，快了一倍左右。这些时间是在 M1 Max 上测量的。

目录