From 7ff5c41e061a27265e0fe793dfc5dda3f4b55e46 Mon Sep 17 00:00:00 2001 From: Jack Wind Date: Fri, 16 May 2025 03:28:03 -0400 Subject: [PATCH] Add set_threadgroup_memory_length to CommandEncoder (#2183) --- mlx/backend/metal/device.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/mlx/backend/metal/device.h b/mlx/backend/metal/device.h index 26c9a0a28..660ba65e2 100644 --- a/mlx/backend/metal/device.h +++ b/mlx/backend/metal/device.h @@ -95,6 +95,10 @@ struct CommandEncoder { return enc_->setBytes(&v, sizeof(T), idx); } + void set_threadgroup_memory_length(size_t length, int idx) { + enc_->setThreadgroupMemoryLength(length, idx); + } + ConcurrentContext start_concurrent() { return ConcurrentContext(*this); }