Enceladus v0.1.0 · alpha
Framework interopAll pages
OverviewQuickstartProgramming modelMemory and synchronizationLanguage referenceDebuggingFramework interopPerformanceBenchmarksPorting from Triton

Guides

Framework interop

Kernels take each framework's arrays without copying them. Importing Enceladus doesn't import PyTorch or MLX. To allocate an output of the same kind and device as an input, call enceladus.new_empty(like). To get strides in elements for any array type, call enceladus.element_strides(x).

PyTorch

  • Tensors must be on the mps device. A CPU tensor raises TypeError.
  • When every array argument is an MPS tensor, the kernel runs on PyTorch's own stream through torch.mps.compile_shader, ordered with the surrounding PyTorch operations. This costs about 5 µs of host time per launch.
  • Mixed array kinds, and kernels that print or assert, take a synchronized path of about 100 µs and log a one-time warning.

MLX

  • Enceladus evaluates every argument with mx.eval() first, and the launch waits, at about 100 µs each.
  • Outputs must be freshly allocated buffers, such as enceladus.new_empty(like). Don't write to lazy results, views, or arrays from mx.broadcast_to.

NumPy

  • Enceladus wraps the array's pages without copying. If Metal can't wrap them, Enceladus copies the array in and out.
  • Negative strides raise an error, so call np.ascontiguousarray(x) first. float64 isn't supported.

DLPack

An enceladus.Tensor exports with the Metal device type, so other frameworks can view the same memory:

dlpack.py
t = enceladus.arange(8, dtype="float32")
tt = torch.from_dlpack(t)  # a torch tensor on the mps device
tt.mul_(10)
torch.mps.synchronize()
print(t.numpy())  # [ 0. 10. 20. 30. 40. 50. 60. 70.]
m = mx.from_dlpack(t)  # an MLX array over the same buffer

The export waits for pending Enceladus launches but orders nothing after that. Before another framework reads later results, call enceladus.synchronize().