统一内存

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

如果我们在GPU上完全计时计算，我们得到2.8毫秒。但如果我们使用d1=mx.gpu和d2=mx.cpu运行计算，那么时间只有大约1.4毫秒，大约快了两倍。这些时间是在M1 Max上测量的。