통합 메모리 — Apple은 왜 새 배열 라이브러리를 만들었나

# Don't run this in the mlx env — it's a contrast, not a recipe.
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

x = torch.randn(1024, 1024).to(device)   # ← copy bytes from RAM to VRAM
y = torch.randn(1024, 1024).to(device)   # ← copy bytes from RAM to VRAM
z = (x @ y).cpu()                        # ← copy result back from VRAM to RAM

# Forget any one of these .to() / .cpu() calls and you get a runtime error
# or a silent slowdown.

import mlx.core as mx

x = mx.random.normal((1024, 1024))
y = mx.random.normal((1024, 1024))
z = x @ y

mx.eval(z)   # materialize the lazy graph (more on this in core.lesson4)
print("shape :", z.shape)
print("dtype :", z.dtype)
print("device:", mx.default_device())   # → Device(gpu, 0) on Apple Silicon

# No .to('cuda'). No .cpu(). The unified-memory pool is the storage,
# always, and 'device' is just a label about which compute unit ran it.

import mlx.core as mx

mx.set_default_device(mx.cpu)
a_cpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device())   # Device(cpu, 0)
print("a_cpu          :", a_cpu)                 # array([1, 2, 3, 4], dtype=float32)

mx.set_default_device(mx.gpu)
a_gpu = mx.array([1.0, 2.0, 3.0, 4.0])
print("default device:", mx.default_device())   # Device(gpu, 0)
print("a_gpu          :", a_gpu)                 # array([1, 2, 3, 4], dtype=float32)

# In PyTorch on CUDA, this next line is a runtime error (cross-device op).
# In MLX, it just works. There is one pool of bytes; the labels were a hint
# about scheduling, not a partition of memory.
b = a_cpu + a_gpu
mx.eval(b)
print("a_cpu + a_gpu  :", b)                    # array([2, 4, 6, 8], dtype=float32)