try again with checkpointed classes

* Implemented pickling and copy for Python arrays(#300 & #367)

* Fixing typos

* Pickle with NumPy arrays

* Pickle: workaround for bfloat16

* Revert "Pickle: workaround for bfloat16"

This reverts commit 25afe6bc09.

* Added an error when pickling bfloat16

* Update python/tests/test_array.py

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

* Update python/tests/test_array.py

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

* Update python/src/array.cpp

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

* Update python/src/array.cpp

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

* clang-format applied

---------

Co-authored-by: Awni Hannun <awni.hannun@gmail.com>

.circleci/config.yml

@@ -237,6 +237,14 @@ workflows:
     jobs:
       - mac_build_and_test
       - linux_build_and_test
+
+  build_pypi_release:
+    when:
+      and:
+        - not: << pipeline.parameters.nightly_build >>
+        - not: << pipeline.parameters.weekly_build >>
+        - not: << pipeline.parameters.test_release >>
+    jobs:
       - build_release:
           filters:
             tags:

ACKNOWLEDGMENTS.md

@@ -10,9 +10,11 @@ MLX was developed with contributions from the following individuals:
 - Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. Added `LogicalAnd` and `LogicalOR` ops.
 - Juarez Bochi: Fixed bug in cross attention.
 - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example.
-- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream` and safetensor support
-- Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented ``MaxPool1d``, ``MaxPool2d``, ``AvgPool1d``, ``AvgPool2d``.
-
+- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream` and safetensor support.
+- Gabrijel Boduljak: Added `mlx.core.linalg`, implemented `norm` method and `InstanceNorm` layer. Implemented pooling layers and ``Upsample``.
+- Hinrik Snær Guðmundsson: Added `atleast_1d`, `atleast_2d`, `atleast_3d` ops.
+- Luca Arnaboldi: Added `Ceil` and `Floor` ops; implemented pickling, copy and deepcopy for mlx arrays.
+- Brian Keene & Atila Orhon, with Argmax Inc.: Added `fast.scaled_dot_product_attention`
 <a href="https://github.com/ml-explore/mlx/graphs/contributors">
   <img class="dark-light" src="https://contrib.rocks/image?repo=ml-explore/mlx&anon=0&columns=20&max=100&r=true" />
 </a>
@@ -252,4 +254,4 @@ Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
-limitations under the License.
+limitations under the License.

CMakeLists.txt

@@ -18,7 +18,7 @@ option(MLX_BUILD_METAL "Build metal backend" ON)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.3.0)
+  set(MLX_VERSION 0.5.1)
 endif()
 
 # --------------------- Processor tests -------------------------
@@ -28,7 +28,6 @@ message(STATUS "Building MLX for ${CMAKE_HOST_SYSTEM_PROCESSOR} processor on ${C
 set(MLX_BUILD_ARM OFF)
 
 if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
-
   if (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64" AND ${CMAKE_HOST_APPLE})
     message(FATAL_ERROR
       "Building for x86_64 on macOS is not supported."
@@ -67,8 +66,6 @@ if (MLX_BUILD_METAL AND NOT METAL_LIB)
   set(MLX_BUILD_METAL OFF)
 elseif (MLX_BUILD_METAL)
   message(STATUS "Building METAL sources")
-  add_compile_definitions(_METAL_)
-
   # Throw an error if xcrun not found
   execute_process(COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
                   OUTPUT_VARIABLE MACOS_VERSION

README.md

@@ -11,10 +11,12 @@ brought to you by Apple machine learning research.
 
 Some key features of MLX include:
 
- - **Familiar APIs**: MLX has a Python API that closely follows NumPy.
-   MLX also has a fully featured C++ API, which closely mirrors the Python API. 
-   MLX has higher-level packages like `mlx.nn` and `mlx.optimizers` with APIs
-   that closely follow PyTorch to simplify building more complex models.
+ - **Familiar APIs**: MLX has a Python API that closely follows NumPy.  MLX
+   also has fully featured C++, [C](https://github.com/ml-explore/mlx-c), and
+   [Swift](https://github.com/ml-explore/mlx-swift/) APIs, which closely mirror
+   the Python API.  MLX has higher-level packages like `mlx.nn` and
+   `mlx.optimizers` with APIs that closely follow PyTorch to simplify building
+   more complex models.
 
  - **Composable function transformations**: MLX supports composable function
    transformations for automatic differentiation, automatic vectorization,

benchmarks/cpp/single_ops.cpp

@@ -73,6 +73,7 @@ void time_unary_ops() {
 
 void time_binary_ops() {
   int M = 1000, N = 100, K = 10;
+  auto condition = random::randint(0, 2, {M, N, K});
   auto a = random::uniform({M, N, K});
   auto b = random::uniform({M, N, K});
   auto device = default_device();
@@ -84,7 +85,9 @@ void time_binary_ops() {
   TIME(divide, a, b, device);
   TIME(maximum, a, b, device);
   TIME(minimum, a, b, device);
+  TIME(where, condition, a, b, device);
 
+  condition = array({true});
   b = random::uniform({1});
   eval(b);
   TIMEM("scalar", add, a, b, device);
@@ -93,7 +96,9 @@ void time_binary_ops() {
   TIMEM("scalar", multiply, a, b, device);
   TIMEM("vector-scalar", divide, a, b, device);
   TIMEM("scalar-vector", divide, b, a, device);
+  TIMEM("scalar-vector", where, condition, a, b, device);
 
+  condition = broadcast_to(array({true}), {1000, 100});
   a = broadcast_to(random::uniform({1}), {1000, 100});
   b = broadcast_to(random::uniform({1}), {1000, 100});
   eval(a, b);
@@ -101,6 +106,7 @@ void time_binary_ops() {
   TIMEM("scalar-scalar broadcast", subtract, a, b, device);
   TIMEM("scalar-scalar broadcast", multiply, a, b, device);
   TIMEM("scalar-scalar broadcast", divide, a, b, device);
+  TIMEM("scalar-scalar broadcast", where, condition, a, b, device);
 }
 
 void time_strided_ops() {

benchmarks/python/comparative/bench_mlx.py

@@ -380,10 +380,6 @@ if __name__ == "__main__":
     if len(args.axis) > 1:
         args.axis.pop(0)
 
-    if args.print_pid:
-        print(os.getpid())
-        input("Press enter to run")
-
     if args.cpu:
         mx.set_default_device(mx.cpu)
     else:
@@ -406,6 +402,10 @@ if __name__ == "__main__":
     x = xs[0]
     axis = args.axis[0]
 
+    if args.print_pid:
+        print(os.getpid())
+        input("Press enter to run")
+
     if args.benchmark == "matmul_square":
         print(bench(matmul_square, x))
 

benchmarks/python/comparative/bench_torch.py

@@ -331,10 +331,6 @@ if __name__ == "__main__":
     if len(args.axis) > 1:
         args.axis.pop(0)
 
-    if args.print_pid:
-        print(os.getpid())
-        input("Press enter to run")
-
     torch.set_num_threads(1)
     device = "cpu" if args.cpu else "mps"
 
@@ -354,6 +350,10 @@ if __name__ == "__main__":
     x = xs[0]
     axis = args.axis[0]
 
+    if args.print_pid:
+        print(os.getpid())
+        input("Press enter to run")
+
     if args.benchmark == "matmul_square":
         print(bench(matmul_square, x))
 

benchmarks/python/compile_bench.py

@@ -0,0 +1,109 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+import math
+import random
+
+import mlx.core as mx
+from time_utils import time_fn
+
+
+def bench_gelu():
+
+    def gelu(x):
+        return x * (1 + mx.erf(x / math.sqrt(2))) / 2
+
+    x = mx.random.uniform(shape=(1000, 1024))
+
+    def gen_fun(fun):
+        def bench_fun(x):
+            for _ in range(10):
+                x = fun(x)
+            return x
+
+        return bench_fun
+
+    time_fn(gen_fun(gelu), x, msg="fixed gelu")
+    time_fn(gen_fun(mx.compile(gelu)), x, msg="compiled fixed gelu")
+
+    def randint():
+        return random.randint(1, x.shape[0])
+
+    def gen_fun(fun):
+        def bench_fun(x, y):
+            x = x[: randint()]
+            for _ in range(10):
+                x = fun(x)
+                y = fun(y)
+            return x, y
+
+        return bench_fun
+
+    y = mx.random.uniform(shape=(1000, 1024))
+    time_fn(gen_fun(gelu), x, y, msg="variable gelu")
+    time_fn(gen_fun(mx.compile(gelu)), x, y, msg="compiled variable gelu")
+    time_fn(
+        gen_fun(mx.compile(gelu, shapeless=True)),
+        x,
+        y,
+        msg="shapeless variable gelu",
+    )
+
+
+def bench_layernorm():
+
+    weight = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    bias = mx.random.uniform(shape=(4096,)).astype(mx.float16)
+    mx.eval(weight, bias)
+
+    def layernorm(x):
+        x = x.astype(mx.float32)
+        means = mx.mean(x, axis=-1, keepdims=True)
+        var = mx.var(x, axis=-1, keepdims=True)
+        x = (x - means) * mx.rsqrt(var + 1e-4)
+        x = x.astype(mx.float16)
+        return weight * x + bias
+
+    x = mx.random.uniform(shape=(1000, 4096)).astype(mx.float16)
+
+    def gen_fun(fun):
+        def bench_fun(x):
+            for _ in range(10):
+                x = fun(x)
+            return x
+
+        return bench_fun
+
+    time_fn(gen_fun(layernorm), x, msg="fixed layernorm")
+    time_fn(gen_fun(mx.compile(layernorm)), x, msg="compiled fixed layernorm")
+
+    def randint():
+        return random.randint(1, x.shape[0])
+
+    def gen_fun(fun):
+        def bench_fun(x):
+            x = x[: randint()]
+            for _ in range(10):
+                x = fun(x)
+            return x
+
+        return bench_fun
+
+    random.seed(0)
+    time_fn(gen_fun(layernorm), x, msg="variable layernorm")
+    random.seed(0)
+    time_fn(gen_fun(mx.compile(layernorm)), x, msg="compiled variable layernorm")
+    random.seed(0)
+    time_fn(
+        gen_fun(mx.compile(layernorm, shapeless=True)),
+        x,
+        msg="shapeless variable layernorm",
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Compile benchmarks.")
+    args = parser.parse_args()
+
+    bench_gelu()
+    bench_layernorm()

benchmarks/python/conv_bench.py

@@ -0,0 +1,129 @@
+import argparse
+import math
+import os
+import subprocess
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+device_name = device_name.decode("utf-8").strip("\n")
+
+N_warmup = 10
+N_iter_bench = 100
+N_iter_func = 5
+
+
+def bench(f, a, b):
+    for i in range(N_warmup):
+        f(a, b)
+    torch.mps.synchronize()
+
+    s = time.perf_counter_ns()
+    for i in range(N_iter_bench):
+        f(a, b)
+    e = time.perf_counter_ns()
+    return (e - s) * 1e-9
+
+
+def make_mx_conv_2D(strides=(1, 1), padding=(0, 0)):
+    def mx_conv_2D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = mx.conv2d(a, b, stride=strides, padding=padding)
+            ys.append(y)
+        mx.eval(ys)
+        return ys
+
+    return mx_conv_2D
+
+
+def make_pt_conv_2D(strides=(1, 1), padding=(0, 0)):
+    @torch.no_grad()
+    def pt_conv_2D(a, b):
+        ys = []
+        for i in range(N_iter_func):
+            y = torch.conv2d(a, b, stride=strides, padding=padding)
+            ys.append(y)
+        torch.mps.synchronize()
+        return ys
+
+    return pt_conv_2D
+
+
+def bench_shape(N, H, W, C, kH, kW, O, strides, padding, np_dtype):
+
+    scale = 1.0 / math.sqrt(kH * kH * C)
+    a_np = np.random.uniform(0, 0.5, (N, H, W, C)).astype(np_dtype)
+    b_np = np.random.uniform(-scale, scale, (O, kH, kW, C)).astype(np_dtype)
+
+    a_mx = mx.array(a_np)
+    b_mx = mx.array(b_np)
+
+    a_pt = torch.from_numpy(a_np.transpose((0, 3, 1, 2))).to("mps")
+    b_pt = torch.from_numpy(b_np.transpose((0, 3, 1, 2))).to("mps")
+
+    torch.mps.synchronize()
+
+    f_mx = make_mx_conv_2D(strides, padding)
+    f_pt = make_pt_conv_2D(strides, padding)
+
+    time_torch = bench(f_pt, a_pt, b_pt)
+    time_mlx = bench(f_mx, a_mx, b_mx)
+
+    out_mx = mx.conv2d(a_mx, b_mx, stride=strides, padding=padding)
+    out_pt = torch.conv2d(
+        a_pt.to("cpu"), b_pt.to("cpu"), stride=strides, padding=padding
+    )
+    out_pt = torch.permute(out_pt, (0, 2, 3, 1))
+    out_pt = out_pt.numpy(force=True)
+
+    atol = 2e-5 if np_dtype == np.float32 else 1e-4
+
+    if not np.allclose(out_pt, out_mx, atol=atol):
+        print(
+            f"Failed at {(N, H, W, C)}, {(O, kH, kW, C)} [strides = {strides}, padding = {padding}] with max(|a - b|) = {np.max(np.abs(out_pt - out_mx))}"
+        )
+
+    return time_mlx, time_torch
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run conv benchmarks")
+
+    dtypes = ("float32",)
+    shapes = (
+        (4, 32, 32, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 32, 32, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 32, 32, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 32, 32, 256, 5, 5, 256, (1, 1), (2, 2)),
+        (4, 32, 32, 512, 5, 5, 512, (1, 1), (2, 2)),
+        (4, 64, 64, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 64, 64, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 64, 64, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 64, 64, 256, 5, 5, 256, (1, 1), (2, 2)),
+        (4, 128, 128, 32, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 128, 128, 64, 5, 5, 64, (1, 1), (2, 2)),
+        (4, 128, 128, 128, 5, 5, 128, (1, 1), (2, 2)),
+        (4, 256, 256, 32, 5, 5, 3, (1, 1), (2, 2)),
+        (4, 256, 256, 3, 5, 5, 32, (1, 1), (2, 2)),
+        (4, 128, 128, 64, 5, 5, 3, (1, 1), (2, 2)),
+        (4, 128, 128, 3, 5, 5, 64, (1, 1), (2, 2)),
+    )
+
+    for dtype in dtypes:
+        print("(N,   H,   W,   C), (  O, kH, kW,   C),   dtype, stride,   pads,  diff%")
+        for N, H, W, C, kH, kW, O, strides, padding in shapes:
+            np_dtype = getattr(np, dtype)
+            time_mlx, time_torch = bench_shape(
+                N, H, W, C, kH, kW, O, strides, padding, np_dtype
+            )
+            diff = time_torch / time_mlx - 1.0
+
+            print(
+                f"({N}, {H:3d}, {W:3d}, {C:3d}), ({O:3d}, {kH:2d}, {kW:2d}, {C:3d}), {dtype}, {strides}, {padding}, {100. * diff:+5.2f}%"
+            )
+            if time_mlx >= 2.0 * time_torch:
+                print("ATTENTION ^^^^^^^")

benchmarks/python/scatter_bench.py

@@ -7,12 +7,14 @@ import torch
 from time_utils import measure_runtime
 
 
-def benchmark_scatter_mlx(dst_shape, x_shape, idx_shape):
+def benchmark_scatter_mlx(dst_shape, x_shape, idx_shapes):
     def scatter(dst, x, idx):
-        dst[idx] = x
+        dst[*idx] = x
         mx.eval(dst)
 
-    idx = mx.random.randint(0, dst_shape[0] - 1, idx_shape)
+    idx = []
+    for idx_shape in idx_shapes:
+        idx.append(mx.random.randint(0, dst_shape[0] - 1, idx_shape))
     x = mx.random.normal(x_shape).astype(mx.float32)
     dst = mx.random.normal(dst_shape).astype(mx.float32)
 
@@ -20,13 +22,15 @@ def benchmark_scatter_mlx(dst_shape, x_shape, idx_shape):
     print(f"MLX: {runtime:.3f}ms")
 
 
-def benchmark_scatter_torch(dst_shape, x_shape, idx_shape, device):
+def benchmark_scatter_torch(dst_shape, x_shape, idx_shapes, device):
     def gather(dst, x, idx, device):
-        dst[idx] = x
+        dst[*idx] = x
         if device == torch.device("mps"):
             torch.mps.synchronize()
 
-    idx = torch.randint(0, dst_shape[0] - 1, idx_shape).to(device)
+    idx = []
+    for idx_shape in idx_shapes:
+        idx.append(torch.randint(0, dst_shape[0] - 1, idx_shape).to(device))
     x = torch.randn(x_shape, dtype=torch.float32).to(device)
     dst = torch.randn(dst_shape, dtype=torch.float32).to(device)
 
@@ -45,9 +49,45 @@ if __name__ == "__main__":
     else:
         device = torch.device("mps")
 
-    dst_shapes = [(10, 64), (100_000, 64), (1_000_000, 64)]
-    idx_shapes = [(1_000_000,), (1_000_000,), (100_000,)]
-    x_shapes = [(1_000_000, 64), (1_000_000, 64), (100_000, 64)]
+    dst_shapes = [
+        (10, 64),
+        (100_000, 64),
+        (1_000_000, 64),
+        (100_000,),
+        (2_000_00,),
+        (20_000_000,),
+        (10000, 64),
+        (100, 64),
+        (100, 10_000, 64),
+        (10, 100, 100, 21),
+        (1_000, 1_000, 10),
+    ]
+    idx_shapes = [
+        [(1_000_000,)],
+        [(1_000_000,)],
+        [(100_000,)],
+        [(1_000_000,)],
+        [(20_000_000,)],
+        [(20_000_000,)],
+        [(1000000,)],
+        [(10000000,)],
+        [(1_000,)],
+        [(10_000,)],
+        [(1_000,), (1_000,)],
+    ]
+    x_shapes = [
+        (1_000_000, 64),
+        (1_000_000, 64),
+        (100_000, 64),
+        (1_000_000,),
+        (20_000_000,),
+        (20_000_000,),
+        (1000000, 64),
+        (10000000, 64),
+        (1_000, 10_000, 64),
+        (10_000, 100, 100, 21),
+        (1_000, 10),
+    ]
 
     for dst_shape, x_shape, idx_shape in zip(dst_shapes, x_shapes, idx_shapes):
         print("=" * 20)

benchmarks/python/time_utils.py

@@ -6,7 +6,11 @@ import mlx.core as mx
 
 
 def time_fn(fn, *args, **kwargs):
-    print(f"Timing {fn.__name__} ...", end=" ")
+    msg = kwargs.pop("msg", None)
+    if msg:
+        print(f"Timing {msg} ...", end=" ")
+    else:
+        print(f"Timing {fn.__name__} ...", end=" ")
 
     # warmup
     for _ in range(5):

docs/src/conf.py

@@ -49,10 +49,12 @@ html_theme_options = {
     "repository_url": "https://github.com/ml-explore/mlx",
     "use_repository_button": True,
     "navigation_with_keys": False,
+    "logo": {
+        "image_light": "_static/mlx_logo.png",
+        "image_dark": "_static/mlx_logo_dark.png",
+    },
 }
 
-html_logo = "_static/mlx_logo.png"
-
 
 # -- Options for HTMLHelp output ---------------------------------------------
 

docs/src/index.rst

@@ -64,6 +64,7 @@ are the CPU and GPU.
    python/transforms
    python/fft
    python/linalg
+   python/metal
    python/nn
    python/optimizers
    python/tree_utils

docs/src/python/metal.rst

@@ -0,0 +1,14 @@
+Metal
+=====
+
+.. currentmodule:: mlx.core.metal
+
+.. autosummary:: 
+  :toctree: _autosummary
+
+  is_available
+  get_active_memory
+  get_peak_memory
+  get_cache_memory
+  set_memory_limit
+  set_cache_limit

docs/src/python/nn.rst

@@ -173,6 +173,7 @@ In detail:
    :toctree: _autosummary
 
    value_and_grad
+   checkpoint
 
 .. toctree::
 

docs/src/python/nn/functions.rst

@@ -12,13 +12,24 @@ simple functions.
    :toctree: _autosummary_functions
    :template: nn-module-template.rst
 
+   elu
    gelu
    gelu_approx
    gelu_fast_approx
+   glu
+   hardswish
+   leaky_relu
+   log_sigmoid
+   log_softmax
    mish
    prelu
    relu
+   relu6
    selu
-   softshrink
+   sigmoid
    silu
+   softmax
+   softplus
+   softshrink
    step
+   tanh

docs/src/python/nn/layers.rst

@@ -40,3 +40,4 @@ Layers
    Softshrink
    Step
    Transformer
+   Upsample

docs/src/python/ops.rst

@@ -25,6 +25,9 @@ Operations
    argpartition
    argsort
    array_equal
+   atleast_1d
+   atleast_2d
+   atleast_3d
    broadcast_to
    ceil
    clip
@@ -32,6 +35,7 @@ Operations
    convolve
    conv1d
    conv2d
+   conv_general
    cos
    cosh
    dequantize
@@ -53,6 +57,7 @@ Operations
    greater_equal
    identity
    inner
+   isclose
    isnan
    isposinf
    isneginf
@@ -117,6 +122,8 @@ Operations
    tan
    tanh
    tensordot
+   tile
+   topk
    transpose
    tri
    tril

docs/src/python/optimizers/schedulers.rst

@@ -8,6 +8,8 @@ Schedulers
 .. autosummary::
    :toctree: _autosummary
 
-   step_decay    
-   exponential_decay    
    cosine_decay    
+   exponential_decay    
+   join_schedules
+   linear_schedule
+   step_decay    

docs/src/python/transforms.rst

@@ -17,3 +17,4 @@ Transforms
    jvp
    vjp
    vmap
+   checkpoint

mlx/backend/accelerate/primitives.cpp

@@ -33,7 +33,6 @@ DEFAULT(ArgSort)
 DEFAULT(AsStrided)
 DEFAULT(Broadcast)
 DEFAULT(Ceil)
-DEFAULT_MULTI(Compiled)
 DEFAULT(Concatenate)
 DEFAULT(Copy)
 DEFAULT_MULTI(CustomVJP)
@@ -65,6 +64,7 @@ DEFAULT(Reshape)
 DEFAULT(Remainder)
 DEFAULT(Round)
 DEFAULT(Scatter)
+DEFAULT(Select)
 DEFAULT(Sigmoid)
 DEFAULT(Sign)
 DEFAULT(Slice)
@@ -82,11 +82,8 @@ void Abs::eval_cpu(const std::vector<array>& inputs, array& out) {
   } else if (in.dtype() == int32 && in.flags().contiguous) {
     set_unary_output_data(in, out);
     vDSP_vabsi(in.data<int>(), 1, out.data<int>(), 1, in.data_size());
-  } else if (is_unsigned(in.dtype())) {
-    // No-op for unsigned types
-    out.copy_shared_buffer(in);
   } else {
-    unary(in, out, AbsOp());
+    eval(inputs, out);
   }
 }
 

mlx/backend/accelerate/quantized.cpp

@@ -24,8 +24,6 @@ void _qmm_t_4_64(
   constexpr int bitmask = (1 << bits) - 1;
   constexpr int pack_factor = 32 / bits;
   constexpr int packs_in_group = group_size / pack_factor;
-  const int Kg = K / group_size;
-  const int Kw = K / pack_factor;
 
   for (int m = 0; m < M; m++) {
     const uint32_t* w_local = w;

mlx/backend/common/CMakeLists.txt

@@ -1,3 +1,36 @@
+
+if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+  set(COMPILER ${CMAKE_C_COMPILER})
+  set(CLANG TRUE)
+else()
+  set(COMPILER ${CMAKE_CXX_COMPILER})
+endif()
+
+add_custom_command(
+    OUTPUT  compiled_preamble.cpp
+    COMMAND /bin/bash
+              ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
+              ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
+              ${COMPILER}
+              ${PROJECT_SOURCE_DIR}
+              ${CLANG}
+
+    DEPENDS make_compiled_preamble.sh
+            compiled_preamble.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/half_types.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/fp16.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/bf16.h
+            ${PROJECT_SOURCE_DIR}/mlx/types/complex.h
+            ops.h
+)
+
+add_custom_target(
+  cpu_compiled_preamble
+  DEPENDS compiled_preamble.cpp
+)
+
+add_dependencies(mlx cpu_compiled_preamble)
+
 target_sources(
   mlx
   PRIVATE
@@ -13,10 +46,26 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/sort.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/threefry.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/load.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/qrf.cpp
+  ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
 )
+
+if (IOS)
+  target_sources(
+    mlx
+    PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/compiled_nocpu.cpp
+  )
+else()
+  target_sources(
+    mlx
+    PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/compiled_cpu.cpp
+  )
+endif()

mlx/backend/common/binary.cpp

@@ -7,6 +7,7 @@
 #include "mlx/allocator.h"
 #include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/binary_two.h"
+#include "mlx/backend/common/ops.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
@@ -73,7 +74,7 @@ void Add::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return x + y; });
+  binary(a, b, out, detail::Add());
 }
 
 void DivMod::eval(
@@ -135,106 +136,56 @@ void Divide::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return x / y; });
+  binary(a, b, out, detail::Divide());
 }
 
-struct RemainderFn {
-  template <typename T>
-  std::enable_if_t<std::is_integral_v<T> & !std::is_signed_v<T>, T> operator()(
-      T numerator,
-      T denominator) {
-    return numerator % denominator;
-  }
-
-  template <typename T>
-  std::enable_if_t<std::is_integral_v<T> & std::is_signed_v<T>, T> operator()(
-      T numerator,
-      T denominator) {
-    auto r = numerator % denominator;
-    if (r != 0 && (r < 0 != denominator < 0))
-      r += denominator;
-    return r;
-  }
-
-  template <typename T>
-  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
-      T numerator,
-      T denominator) {
-    auto r = std::fmod(numerator, denominator);
-    if (r != 0 && (r < 0 != denominator < 0)) {
-      r += denominator;
-    }
-    return r;
-  }
-
-  complex64_t operator()(complex64_t numerator, complex64_t denominator) {
-    return numerator % denominator;
-  }
-};
-
 void Remainder::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, RemainderFn{});
+  binary(a, b, out, detail::Remainder());
 }
 
 void Equal::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   if (equal_nan_) {
-    comparison_op(inputs[0], inputs[1], out, [](auto x, auto y) {
-      return x == y || (std::isnan(x) && std::isnan(y));
-    });
+    comparison_op(inputs[0], inputs[1], out, detail::NaNEqual());
   } else {
-    comparison_op(
-        inputs[0], inputs[1], out, [](auto x, auto y) { return x == y; });
+    comparison_op(inputs[0], inputs[1], out, detail::Equal());
   }
 }
 
 void Greater::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(
-      inputs[0], inputs[1], out, [](auto x, auto y) { return x > y; });
+  comparison_op(inputs[0], inputs[1], out, detail::Greater());
 }
 
 void GreaterEqual::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(
-      inputs[0], inputs[1], out, [](auto x, auto y) { return x >= y; });
+  comparison_op(inputs[0], inputs[1], out, detail::GreaterEqual());
 }
 
 void Less::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(
-      inputs[0], inputs[1], out, [](auto x, auto y) { return x < y; });
+  comparison_op(inputs[0], inputs[1], out, detail::Less());
 }
 
 void LessEqual::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(
-      inputs[0], inputs[1], out, [](auto x, auto y) { return x <= y; });
+  comparison_op(inputs[0], inputs[1], out, detail::LessEqual());
 }
 
 void LogAddExp::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  auto op = [](auto x, auto y) {
-    constexpr float inf = std::numeric_limits<float>::infinity();
-    auto maxval = (x > y) ? x : y;
-    auto minval = (x > y) ? y : x;
-    return (minval == -inf || maxval == inf)
-        ? maxval
-        : static_cast<decltype(x)>(
-              maxval + std::log1p(std::exp(minval - maxval)));
-  };
   if (is_floating_point(out.dtype())) {
     if (out.dtype() == float32) {
-      binary_op<float>(a, b, out, op);
+      binary_op<float>(a, b, out, detail::LogAddExp());
     } else if (out.dtype() == float16) {
-      binary_op<float16_t>(a, b, out, op);
+      binary_op<float16_t>(a, b, out, detail::LogAddExp());
     } else if (out.dtype() == bfloat16) {
-      binary_op<bfloat16_t>(a, b, out, op);
+      binary_op<bfloat16_t>(a, b, out, detail::LogAddExp());
     } else {
       std::ostringstream err;
       err << "[logaddexp] Does not support " << out.dtype();
@@ -251,84 +202,40 @@ void Maximum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-
-  if (is_floating_point(out.dtype())) {
-    binary(a, b, out, [](auto x, auto y) {
-      if (std::isnan(x)) {
-        return x;
-      }
-      return (x > y) ? x : y;
-    });
-  } else {
-    binary(a, b, out, [](auto x, auto y) { return (x > y) ? x : y; });
-  }
+  binary(a, b, out, detail::Maximum());
 }
 
 void Minimum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  if (is_floating_point(out.dtype())) {
-    binary(a, b, out, [](auto x, auto y) {
-      if (std::isnan(x)) {
-        return x;
-      }
-      return (x < y) ? x : y;
-    });
-  } else {
-    binary(a, b, out, [](auto x, auto y) { return (x < y) ? x : y; });
-  }
+  binary(a, b, out, detail::Minimum());
 }
 
 void Multiply::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return x * y; });
+  binary(a, b, out, detail::Multiply());
 }
 
 void NotEqual::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
-  comparison_op(
-      inputs[0], inputs[1], out, [](auto x, auto y) { return x != y; });
+  comparison_op(inputs[0], inputs[1], out, detail::NotEqual());
 }
 
-struct PowerFn {
-  template <typename T>
-  std::enable_if_t<!std::is_integral_v<T>, T> operator()(T base, T exp) {
-    return std::pow(base, exp);
-  }
-
-  template <typename T>
-  std::enable_if_t<std::is_integral_v<T>, T> operator()(T base, T exp) {
-    if (exp < 0) {
-      throw std::invalid_argument(
-          "Integers cannot be raise to negative powers");
-    }
-    T res = 1;
-    while (exp) {
-      if (exp & 1) {
-        res *= base;
-      }
-      exp >>= 1;
-      base *= base;
-    }
-    return res;
-  }
-};
-
 void Power::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, PowerFn{});
+  binary(a, b, out, detail::Power());
 }
 
 void Subtract::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
   auto& b = inputs[1];
-  binary(a, b, out, [](auto x, auto y) { return x - y; });
+  binary(a, b, out, detail::Subtract());
 }
 
 } // namespace mlx::core

mlx/backend/common/binary.h

@@ -9,7 +9,7 @@ namespace mlx::core {
 
 namespace {
 
-enum BinaryOpType {
+enum class BinaryOpType {
   ScalarScalar,
   ScalarVector,
   VectorScalar,
@@ -20,17 +20,17 @@ enum BinaryOpType {
 BinaryOpType get_binary_op_type(const array& a, const array& b) {
   BinaryOpType bopt;
   if (a.data_size() == 1 && b.data_size() == 1) {
-    bopt = ScalarScalar;
+    bopt = BinaryOpType::ScalarScalar;
   } else if (a.data_size() == 1 && b.flags().contiguous) {
-    bopt = ScalarVector;
+    bopt = BinaryOpType::ScalarVector;
   } else if (b.data_size() == 1 && a.flags().contiguous) {
-    bopt = VectorScalar;
+    bopt = BinaryOpType::VectorScalar;
   } else if (
       a.flags().row_contiguous && b.flags().row_contiguous ||
       a.flags().col_contiguous && b.flags().col_contiguous) {
-    bopt = VectorVector;
+    bopt = BinaryOpType::VectorVector;
   } else {
-    bopt = General;
+    bopt = BinaryOpType::General;
   }
   return bopt;
 }
@@ -42,11 +42,11 @@ void set_binary_op_output_data(
     BinaryOpType bopt,
     bool donate_with_move = false) {
   switch (bopt) {
-    case ScalarScalar:
+    case BinaryOpType::ScalarScalar:
       out.set_data(
           allocator::malloc_or_wait(out.itemsize()), 1, a.strides(), a.flags());
       break;
-    case ScalarVector:
+    case BinaryOpType::ScalarVector:
       if (b.is_donatable() && b.itemsize() == out.itemsize()) {
         if (donate_with_move) {
           out.move_shared_buffer(b);
@@ -61,7 +61,7 @@ void set_binary_op_output_data(
             b.flags());
       }
       break;
-    case VectorScalar:
+    case BinaryOpType::VectorScalar:
       if (a.is_donatable() && a.itemsize() == out.itemsize()) {
         if (donate_with_move) {
           out.move_shared_buffer(a);
@@ -76,7 +76,7 @@ void set_binary_op_output_data(
             a.flags());
       }
       break;
-    case VectorVector:
+    case BinaryOpType::VectorVector:
       if (a.is_donatable() && a.itemsize() == out.itemsize()) {
         if (donate_with_move) {
           out.move_shared_buffer(a);
@@ -97,7 +97,7 @@ void set_binary_op_output_data(
             a.flags());
       }
       break;
-    case General:
+    case BinaryOpType::General:
       if (a.is_donatable() && a.flags().row_contiguous &&
           a.itemsize() == out.itemsize() && a.size() == out.size()) {
         if (donate_with_move) {
@@ -424,25 +424,25 @@ void binary_op(
   set_binary_op_output_data(a, b, out, bopt);
 
   // The full computation is scalar scalar so call the base op once
-  if (bopt == ScalarScalar) {
+  if (bopt == BinaryOpType::ScalarScalar) {
     *(out.data<U>()) = op(*a.data<T>(), *b.data<T>());
     return;
   }
 
   // The full computation is scalar vector so delegate to the op
-  if (bopt == ScalarVector) {
+  if (bopt == BinaryOpType::ScalarVector) {
     opsv(a.data<T>(), b.data<T>(), out.data<U>(), b.data_size());
     return;
   }
 
   // The full computation is vector scalar so delegate to the op
-  if (bopt == VectorScalar) {
+  if (bopt == BinaryOpType::VectorScalar) {
     opvs(a.data<T>(), b.data<T>(), out.data<U>(), a.data_size());
     return;
   }
 
   // The full computation is vector vector so delegate to the op
-  if (bopt == VectorVector) {
+  if (bopt == BinaryOpType::VectorVector) {
     opvv(a.data<T>(), b.data<T>(), out.data<U>(), out.size());
     return;
   }
@@ -475,17 +475,17 @@ void binary_op(
   // Case 1: LxM and FxM where L and F are broadcastable and M is row contiguous
   int dim = ndim;
   if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
-    bopt = VectorVector;
+    bopt = BinaryOpType::VectorVector;
     dim = d;
     // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
     // contiguous
   } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
-    bopt = VectorScalar;
+    bopt = BinaryOpType::VectorScalar;
     dim = d;
     // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
     // contiguous
   } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
-    bopt = ScalarVector;
+    bopt = BinaryOpType::ScalarVector;
     dim = d;
   }
 
@@ -495,20 +495,20 @@ void binary_op(
   size_t stride;
   if (dim == 0 || strides[dim - 1] < 16) {
     stride = 1;
-    bopt = General;
+    bopt = BinaryOpType::General;
     dim = ndim;
   } else {
     stride = strides[dim - 1];
   }
 
   switch (bopt) {
-    case VectorVector:
+    case BinaryOpType::VectorVector:
       binary_op_dispatch_dims<T, U>(a, b, out, opvv, dim, stride);
       break;
-    case VectorScalar:
+    case BinaryOpType::VectorScalar:
       binary_op_dispatch_dims<T, U>(a, b, out, opvs, dim, stride);
       break;
-    case ScalarVector:
+    case BinaryOpType::ScalarVector:
       binary_op_dispatch_dims<T, U>(a, b, out, opsv, dim, stride);
       break;
     default:

mlx/backend/common/binary_two.h

@@ -260,14 +260,14 @@ void binary_op(
   set_binary_op_output_data(a, b, out_b, bopt);
 
   // The full computation is scalar scalar so call the base op once
-  if (bopt == ScalarScalar) {
+  if (bopt == BinaryOpType::ScalarScalar) {
     std::tie(*(out_a.data<U>()), *(out_b.data<U>())) =
         op(*a.data<T>(), *b.data<T>());
     return;
   }
 
   // The full computation is scalar vector so delegate to the op
-  if (bopt == ScalarVector) {
+  if (bopt == BinaryOpType::ScalarVector) {
     opsv(
         a.data<T>(),
         b.data<T>(),
@@ -278,7 +278,7 @@ void binary_op(
   }
 
   // The full computation is vector scalar so delegate to the op
-  if (bopt == VectorScalar) {
+  if (bopt == BinaryOpType::VectorScalar) {
     opvs(
         a.data<T>(),
         b.data<T>(),
@@ -289,7 +289,7 @@ void binary_op(
   }
 
   // The full computation is vector vector so delegate to the op
-  if (bopt == VectorVector) {
+  if (bopt == BinaryOpType::VectorVector) {
     opvv(
         a.data<T>(),
         b.data<T>(),
@@ -327,17 +327,17 @@ void binary_op(
   // Case 1: LxM and FxM where L and F are broadcastable and M is row contiguous
   int dim = ndim;
   if (int d = std::max(a_rc_dim, b_rc_dim); d < ndim) {
-    bopt = VectorVector;
+    bopt = BinaryOpType::VectorVector;
     dim = d;
     // Case 2: LxM and Fx1 where L and F are broadcastable and M is row
     // contiguous
   } else if (int d = std::max(a_rc_dim, b_s_dim); d < ndim) {
-    bopt = VectorScalar;
+    bopt = BinaryOpType::VectorScalar;
     dim = d;
     // Case 3: Lx1 and FxM where L and F are broadcastable and M is row
     // contiguous
   } else if (int d = std::max(a_s_dim, b_rc_dim); d < ndim) {
-    bopt = ScalarVector;
+    bopt = BinaryOpType::ScalarVector;
     dim = d;
   }
 
@@ -347,20 +347,20 @@ void binary_op(
   size_t stride;
   if (dim == 0 || strides[dim - 1] < 16) {
     stride = 1;
-    bopt = General;
+    bopt = BinaryOpType::General;
     dim = ndim;
   } else {
     stride = strides[dim - 1];
   }
 
   switch (bopt) {
-    case VectorVector:
+    case BinaryOpType::VectorVector:
       binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opvv, dim, stride);
       break;
-    case VectorScalar:
+    case BinaryOpType::VectorScalar:
       binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opvs, dim, stride);
       break;
-    case ScalarVector:
+    case BinaryOpType::ScalarVector:
       binary_op_dispatch_dims<T, U>(a, b, out_a, out_b, opsv, dim, stride);
       break;
     default:

mlx/backend/common/compiled.cpp

@@ -1,59 +1,114 @@
 // Copyright © 2023-2024 Apple Inc.
 
-#include <queue>
-
+#include "mlx/backend/common/compiled.h"
 #include "mlx/primitives.h"
+#include "mlx/utils.h"
 
 namespace mlx::core {
 
-// Build the real tape
-std::pair<std::queue<array>, std::vector<array>> trace_to_real(
-    const std::vector<array>& trace_tape,
-    const std::vector<array>& trace_inputs,
-    const std::vector<array>& trace_outputs,
-    const std::vector<array>& inputs) {
-  std::unordered_map<uintptr_t, array> trace_to_real;
-  for (int i = 0; i < inputs.size(); ++i) {
-    trace_to_real.insert({trace_inputs[i].id(), inputs[i]});
+void print_constant(std::ostream& os, const array& x) {
+  switch (x.dtype()) {
+    case float32:
+      return print_float_constant<float>(os, x);
+    case float16:
+      return print_float_constant<float16_t>(os, x);
+    case bfloat16:
+      return print_float_constant<bfloat16_t>(os, x);
+    case complex64:
+      return print_complex_constant<complex64_t>(os, x);
+    case int8:
+      return print_int_constant<int8_t>(os, x);
+    case int16:
+      return print_int_constant<int16_t>(os, x);
+    case int32:
+      return print_int_constant<int32_t>(os, x);
+    case int64:
+      return print_int_constant<int64_t>(os, x);
+    case uint8:
+      return print_int_constant<uint8_t>(os, x);
+    case uint16:
+      return print_int_constant<uint16_t>(os, x);
+    case uint32:
+      return print_int_constant<uint32_t>(os, x);
+    case uint64:
+      return print_int_constant<uint64_t>(os, x);
+    case bool_:
+      os << std::boolalpha << x.item<bool>();
+      return;
+    default:
+      throw std::runtime_error("Unsupported constant type");
   }
-  std::queue<array> tape;
-  for (auto& a : trace_tape) {
-    // Find real inputs
-    std::vector<array> real_inputs;
-    for (auto& in : a.inputs()) {
-      real_inputs.push_back(trace_to_real.at(in.id()));
-    }
-    tape.push(
-        array(a.shape(), a.dtype(), a.primitive_ptr(), std::move(real_inputs)));
-    trace_to_real.insert({a.id(), tape.back()});
-  }
-
-  std::vector<array> outputs;
-  for (auto& o : trace_outputs) {
-    outputs.push_back(trace_to_real.at(o.id()));
-  }
-  return {tape, outputs};
 }
 
-void Compiled::eval(
+std::string get_type_string(Dtype d) {
+  switch (d) {
+    case float32:
+      return "float";
+    case float16:
+      return "float16_t";
+    case bfloat16:
+      return "bfloat16_t";
+    case complex64:
+      return "complex64_t";
+    case bool_:
+      return "bool";
+    case int8:
+      return "int8_t";
+    case int16:
+      return "int16_t";
+    case int32:
+      return "int32_t";
+    case int64:
+      return "int64_t";
+    case uint8:
+      return "uint8_t";
+    case uint16:
+      return "uint16_t";
+    case uint32:
+      return "uint32_t";
+    case uint64:
+      return "uint64_t";
+    default: {
+      std::ostringstream msg;
+      msg << "Unsupported compilation type " << d;
+      throw std::runtime_error(msg.str());
+    }
+  }
+}
+
+std::string build_lib_name(
     const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  // Make the a real tape from the tracers
-  auto [tape, real_outputs] = trace_to_real(tape_, inputs_, outputs_, inputs);
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::unordered_set<uintptr_t>& constant_ids) {
+  std::ostringstream os;
+  std::ostringstream constant_hasher;
 
-  // Run the tape
-  while (!tape.empty()) {
-    auto a = std::move(tape.front());
-    tape.pop();
-    auto outputs = a.outputs();
-    a.primitive().eval_cpu(a.inputs(), outputs);
-    a.detach();
+  // The primitives describing the tape. For unary and binary primitives this
+  // must be enough to describe the full computation.
+  for (auto& a : tape) {
+    a.primitive().print(os);
   }
+  os << "_";
 
-  // Copy results into outputs
-  for (int o = 0; o < real_outputs.size(); ++o) {
-    outputs[o].copy_shared_buffer(real_outputs[o]);
+  for (auto& x : inputs) {
+    if (constant_ids.find(x.id()) != constant_ids.end()) {
+      os << "C";
+      print_constant(constant_hasher, x);
+    } else {
+      os << (is_scalar(x) ? "S" : "V");
+    }
   }
+  os << "_";
+  for (auto& x : inputs) {
+    if (constant_ids.find(x.id()) != constant_ids.end()) {
+      continue;
+    }
+    os << kindof(x.dtype()) << x.itemsize();
+  }
+  os << "_" << std::hash<std::string>{}(constant_hasher.str());
+
+  return os.str();
 }
 
 } // namespace mlx::core

mlx/backend/common/compiled.h

@@ -0,0 +1,56 @@
+// Copyright © 2023-2024 Apple Inc.
+#pragma once
+
+#include <iomanip>
+#include <sstream>
+#include <unordered_set>
+
+#include "mlx/array.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+inline bool is_static_cast(const Primitive& p) {
+  return (
+      typeid(p) == typeid(Broadcast) || typeid(p) == typeid(Copy) ||
+      typeid(p) == typeid(StopGradient) || typeid(p) == typeid(AsType));
+}
+
+std::string build_lib_name(
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::unordered_set<uintptr_t>& constant_ids);
+
+std::string get_type_string(Dtype d);
+
+template <typename T>
+void print_float_constant(std::ostream& os, const array& x) {
+  auto old_precision = os.precision();
+  os << std::setprecision(std::numeric_limits<float>::digits10 + 1)
+     << x.item<T>() << std::setprecision(old_precision);
+}
+
+template <typename T>
+void print_int_constant(std::ostream& os, const array& x) {
+  os << x.item<T>();
+}
+
+template <typename T>
+void print_complex_constant(std::ostream& os, const array& x) {
+  auto old_precision = os.precision();
+  T constant = x.item<T>();
+
+  os << get_type_string(x.dtype()) << "("
+     << std::setprecision(std::numeric_limits<float>::digits10 + 1)
+     << constant.real() << ", " << constant.imag() << ")"
+     << std::setprecision(old_precision);
+}
+
+void print_constant(std::ostream& os, const array& x);
+
+inline bool is_scalar(const array& x) {
+  return x.ndim() == 0;
+}
+
+} // namespace mlx::core

mlx/backend/common/compiled_cpu.cpp

@@ -0,0 +1,408 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <dlfcn.h>
+#include <filesystem>
+#include <list>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/compiled_preamble.h"
+#include "mlx/device.h"
+#include "mlx/graph_utils.h"
+
+namespace mlx::core {
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is also available.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return true;
+}
+} // namespace detail
+
+std::string get_temp_file(const std::string& name) {
+  return std::filesystem::temp_directory_path().append(name);
+}
+
+// Return a pointer to a compiled function
+void* compile(
+    const std::string& kernel_name,
+    const std::string& source_code = "") {
+  struct DLib {
+    DLib(const std::string& libname) {
+      lib = dlopen(libname.c_str(), RTLD_NOW);
+      if (!lib) {
+        std::ostringstream msg;
+        msg << "Could not load C++ shared library " << dlerror();
+        throw std::runtime_error(msg.str());
+      }
+    }
+
+    ~DLib() {
+      dlclose(lib);
+    }
+    void* lib;
+  };
+  // Statics to cache compiled libraries and functions
+  static std::list<DLib> libs;
+  static std::unordered_map<std::string, void*> kernels;
+  if (auto it = kernels.find(kernel_name); it != kernels.end()) {
+    return it->second;
+  }
+  if (source_code.empty()) {
+    return nullptr;
+  }
+
+  std::ostringstream shared_lib_name;
+  shared_lib_name << "lib" << kernel_name << ".so";
+  auto shared_lib_path = get_temp_file(shared_lib_name.str());
+  bool lib_exists = false;
+  {
+    std::ifstream f(shared_lib_path.c_str());
+    lib_exists = f.good();
+  }
+
+  if (!lib_exists) {
+    // Open source file and write source code to it
+    std::ostringstream source_file_name;
+    source_file_name << kernel_name << ".cpp";
+    auto source_file_path = get_temp_file(source_file_name.str());
+
+    std::ofstream source_file(source_file_path);
+    source_file << source_code;
+    source_file.close();
+
+    std::ostringstream build_command;
+    build_command << "g++ -std=c++17 -O2 -Wall -fPIC -shared "
+                  << source_file_path << " -o " << shared_lib_path;
+    std::string build_command_str = build_command.str();
+    auto return_code = system(build_command_str.c_str());
+    if (return_code) {
+      std::ostringstream msg;
+      msg << "[Compile::eval_cpu] Failed to compile function " << kernel_name
+          << " with error code " << return_code << "." << std::endl;
+      throw std::runtime_error(msg.str());
+    }
+  }
+
+  // load library
+  libs.emplace_back(shared_lib_path);
+
+  // Load function
+  void* fun = dlsym(libs.back().lib, kernel_name.c_str());
+  if (!fun) {
+    std::ostringstream msg;
+    msg << "[Compile::eval_cpu] Failed to load compiled function "
+        << kernel_name << std::endl
+        << dlerror();
+    throw std::runtime_error(msg.str());
+  }
+  kernels.insert({kernel_name, fun});
+  return fun;
+}
+
+inline void build_kernel(
+    std::ostream& os,
+    const std::string& kernel_name,
+    const std::vector<array>& inputs,
+    const std::vector<array>& outputs,
+    const std::vector<array>& tape,
+    const std::unordered_set<uintptr_t>& constant_ids,
+    bool contiguous,
+    int ndim) {
+  // All outputs should have the exact same shape and will be row contiguous
+  auto output_shape = outputs[0].shape();
+  auto output_strides = outputs[0].strides();
+
+  // Constants are scalars that are captured by value and cannot change
+  auto is_constant = [&constant_ids](const array& x) {
+    return constant_ids.find(x.id()) != constant_ids.end();
+  };
+
+  NodeNamer namer;
+
+  // Start the kernel
+  os << "void " << kernel_name << "(void** args) {" << std::endl;
+
+  // Add the input arguments
+  int cnt = 0;
+  for (auto& x : inputs) {
+    auto& xname = namer.get_name(x);
+
+    // Skip constants from the input list
+    if (is_constant(x)) {
+      continue;
+    }
+
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << xname << " = (" << tstr << "*)args[" << cnt++
+       << "];" << std::endl;
+    // Scalars and contiguous need no strides
+    if (!is_scalar(x) && !contiguous) {
+      os << "  const size_t* " << xname << "_strides = (size_t*)args[" << cnt++
+         << "];" << std::endl;
+    }
+  }
+
+  // Add the output arguments
+  for (auto& x : outputs) {
+    auto tstr = get_type_string(x.dtype());
+    os << "  " << tstr << "* " << namer.get_name(x) << " = (" << tstr
+       << "*)args[" << cnt++ << "];" << std::endl;
+  }
+  // Add output strides and shape to extract the indices.
+  if (!contiguous) {
+    os << "  const int* shape = (int*)args[" << cnt++ << "];" << std::endl;
+  } else {
+    os << "  const size_t size = (size_t)args[" << cnt++ << "];" << std::endl;
+  }
+
+  if (contiguous) {
+    os << "  for (size_t i = 0; i < size; ++i) {" << std::endl;
+  } else {
+    for (int d = 0; d < ndim; ++d) {
+      os << "  for (int i" << d << " = 0; i" << d << " < shape[" << d
+         << "]; ++i" << d << ") {" << std::endl;
+    }
+  }
+
+  // Read the inputs in tmps
+  for (auto& x : inputs) {
+    auto& xname = namer.get_name(x);
+
+    if (is_constant(x)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = ";
+      print_constant(os, x);
+      os << ";" << std::endl;
+    } else if (is_scalar(x)) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[0];" << std::endl;
+    } else if (contiguous) {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = "
+         << xname << "[i];" << std::endl;
+    } else {
+      os << "  " << get_type_string(x.dtype()) << " tmp_" << xname << " = *"
+         << xname << ";" << std::endl;
+    }
+  }
+
+  // Actually write the computation
+  for (auto& x : tape) {
+    os << "  " << get_type_string(x.dtype()) << " tmp_" << namer.get_name(x)
+       << " = ";
+    if (is_static_cast(x.primitive())) {
+      os << "static_cast<" << get_type_string(x.dtype()) << ">(tmp_"
+         << namer.get_name(x.inputs()[0]) << ");" << std::endl;
+    } else {
+      x.primitive().print(os);
+      os << "()(";
+      for (int i = 0; i < x.inputs().size() - 1; i++) {
+        os << "tmp_" << namer.get_name(x.inputs()[i]) << ", ";
+      }
+      os << "tmp_" << namer.get_name(x.inputs().back()) << ");" << std::endl;
+    }
+  }
+
+  // Write the outputs from tmps
+  for (auto& x : outputs) {
+    if (contiguous) {
+      os << "  " << namer.get_name(x) << "[i] = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    } else {
+      os << "  *" << namer.get_name(x) << "++ = tmp_" << namer.get_name(x)
+         << ";" << std::endl;
+    }
+  }
+
+  // Close loops
+  if (contiguous) {
+    os << "  }" << std::endl;
+  } else {
+    for (int d = ndim - 1; d >= 0; --d) {
+      // Update pointers
+      for (auto& x : inputs) {
+        if (is_constant(x) || is_scalar(x)) {
+          continue;
+        }
+        auto& xname = namer.get_name(x);
+        os << "  " << xname << " += " << xname << "_strides[" << d << "];"
+           << std::endl;
+        if (d < ndim - 1) {
+          os << "  " << xname << " -= " << xname << "_strides[" << d + 1 << "]"
+             << " * shape[" << d + 1 << "];" << std::endl;
+        }
+      }
+      os << "  }" << std::endl;
+    }
+  }
+
+  // Finish the kernel
+  os << "}" << std::endl;
+}
+
+void Compiled::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  if (kernel_lib_.empty()) {
+    kernel_lib_ = build_lib_name(inputs_, outputs_, tape_, constant_ids_);
+  }
+
+  // Figure out which kernel we are using
+  auto& shape = outputs[0].shape();
+  bool contiguous = true;
+  {
+    bool all_contig = true;
+    bool all_row_contig = true;
+    bool all_col_contig = true;
+    int non_scalar_inputs = 0;
+    for (auto& x : inputs) {
+      if (is_scalar(x)) {
+        continue;
+      }
+      non_scalar_inputs++;
+      bool shape_eq = x.shape() == shape;
+      all_contig &= (x.flags().contiguous && shape_eq);
+      all_row_contig &= (x.flags().row_contiguous && shape_eq);
+      all_col_contig &= (x.flags().col_contiguous && shape_eq);
+    }
+    if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) {
+      contiguous = false;
+    } else if (non_scalar_inputs == 1 && !all_contig) {
+      contiguous = false;
+    }
+  }
+
+  // Handle all broadcasting and collect function input arguments
+  std::vector<void*> args;
+  std::vector<std::vector<size_t>> strides;
+  for (int i = 0; i < inputs.size(); i++) {
+    // Skip constants.
+    if (constant_ids_.find(inputs_[i].id()) != constant_ids_.end()) {
+      continue;
+    }
+    auto& x = inputs[i];
+    args.push_back((void*)x.data<void>());
+
+    if (contiguous || is_scalar(x)) {
+      continue;
+    }
+
+    // Broadcast the input to the output shape.
+    std::vector<size_t> xstrides;
+    int j = 0;
+    for (; j < shape.size() - x.ndim(); j++) {
+      if (shape[j] == 1) {
+        xstrides.push_back(outputs[0].strides()[j]);
+      } else {
+        xstrides.push_back(0);
+      }
+    }
+    for (int i = 0; i < x.ndim(); i++, j++) {
+      if (x.shape(i) == 1) {
+        if (shape[j] == 1) {
+          xstrides.push_back(outputs[0].strides()[j]);
+        } else {
+          xstrides.push_back(0);
+        }
+      } else {
+        xstrides.push_back(x.strides()[i]);
+      }
+    }
+    strides.push_back(std::move(xstrides));
+    args.push_back(strides.back().data());
+  }
+
+  // Get the kernel name from the lib
+  int ndim = shape.size();
+  auto kernel_name = kernel_lib_ + (contiguous ? "_contiguous" : "_strided_");
+  if (!contiguous) {
+    kernel_name += std::to_string(shape.size());
+  }
+
+  // Get the function
+  auto fn_ptr = compile(kernel_name);
+
+  // If it doesn't exist, compile it
+  if (fn_ptr == nullptr) {
+    std::ostringstream kernel;
+    kernel << get_kernel_preamble() << std::endl;
+    kernel << "extern \"C\"  {" << std::endl;
+    build_kernel(
+        kernel,
+        kernel_name,
+        inputs_,
+        outputs_,
+        tape_,
+        constant_ids_,
+        contiguous,
+        ndim);
+    // Close extern "C"
+    kernel << "}" << std::endl;
+
+    // Compile and get function pointer
+    fn_ptr = compile(kernel_name, kernel.str());
+  }
+
+  // Allocate space for the outputs possibly with input donation
+  if (contiguous) {
+    int o = 0;
+    std::vector<size_t> strides;
+    size_t data_size;
+    array::Flags flags;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().contiguous && !is_scalar(in) && in.is_donatable() &&
+          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
+        outputs[o++].copy_shared_buffer(in);
+      }
+      // Get representative input flags to properly set non-donated outputs
+      if (strides.empty() && in.size() == outputs[0].size()) {
+        strides = in.strides();
+        flags = in.flags();
+        data_size = in.data_size();
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(
+          allocator::malloc_or_wait(data_size * outputs[o].itemsize()),
+          data_size,
+          strides,
+          flags);
+    }
+  } else {
+    int o = 0;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Row contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().row_contiguous && in.nbytes() == outputs[o].nbytes() &&
+          in.is_donatable() &&
+          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
+        outputs[o++].copy_shared_buffer(in);
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(allocator::malloc_or_wait(outputs[o].nbytes()));
+    }
+  }
+
+  for (auto& x : outputs) {
+    args.push_back(x.data<void>());
+  }
+  if (!contiguous) {
+    args.push_back((void*)outputs[0].shape().data());
+  } else {
+    args.push_back((void*)outputs[0].data_size());
+  }
+  auto fun = (void (*)(void**))fn_ptr;
+  fun(args.data());
+}
+
+} // namespace mlx::core

mlx/backend/common/compiled_nocpu.cpp

@@ -0,0 +1,23 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/common/compiled.h"
+
+namespace mlx::core {
+
+// GPU compile is always available if the GPU is available and since we are in
+// this file CPU compile is not available so check if the device is a GPU
+// device.
+namespace detail {
+bool compile_available_for_device(const Device& device) {
+  return device == Device::gpu;
+}
+} // namespace detail
+
+void Compiled::eval_cpu(
+    const std::vector<array>& inputs,
+    std::vector<array>& outputs) {
+  throw std::runtime_error(
+      "[Compiled::eval_cpu] CPU compialtion not supported on the platform.");
+}
+
+} // namespace mlx::core

mlx/backend/common/compiled_preamble.h

@@ -0,0 +1,11 @@
+// Copyright © 2023-24 Apple Inc.
+
+#pragma once
+
+// clang-format off
+#include "mlx/types/half_types.h"
+#include "mlx/types/complex.h"
+#include "mlx/backend/common/ops.h"
+// clang-format on
+
+const char* get_kernel_preamble();

mlx/backend/common/conv.cpp

@@ -1,6 +1,7 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cassert>
+#include <numeric>
 
 #ifdef ACCELERATE_NEW_LAPACK
 #include <Accelerate/Accelerate.h>
@@ -27,14 +28,16 @@ void slow_conv_1D(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
   const T* start_wt_ptr = wt.data<T>();
 
   const T* in_ptr = in.data<T>();
   T* out_ptr = out.data<T>();
 
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
-  const int iH = in.shape(1); // Input spatial dim
+  const int iH = 1 + in_dilation[0] * (in.shape(1) - 1); // Input spatial dim
   const int oH = out.shape(1); // Output spatial dim
   const int O = wt.shape(0); // Out channels
   const int C = wt.shape(2); // In channels
@@ -61,12 +64,15 @@ void slow_conv_1D(
         for (int wh = 0; wh < wH; ++wh) {
           const T* wt_ptr = filter_wt_ptr + wh * wt_stride_H;
 
-          int ih = oh * wt_strides[0] - padding[0] + wh * wt_dilation[0];
+          int wh_flip = flip ? (wH - wh - 1) : wh;
+          int ih = oh * wt_strides[0] - padding[0] + wh_flip * wt_dilation[0];
 
-          if (ih >= 0 && ih < iH) {
+          auto ih_div = std::div(ih, in_dilation[0]);
+
+          if (ih >= 0 && ih < iH && ih_div.rem == 0) {
             for (int c = 0; c < C; ++c) {
               r += static_cast<float>(
-                       in_ptr[ih * in_stride_H + c * in_stride_C]) *
+                       in_ptr[ih_div.quot * in_stride_H + c * in_stride_C]) *
                   static_cast<float>(wt_ptr[c * wt_stride_C]);
             } // c
 
@@ -90,14 +96,16 @@ void slow_conv_2D(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
   const T* st_wt_ptr = wt.data<T>();
   const T* st_in_ptr = in.data<T>();
   T* st_out_ptr = out.data<T>();
 
   const int N = in.shape(0); // Batch size, should be the same as out.shape(0)
-  const int iH = in.shape(1); // Input spatial dim
-  const int iW = in.shape(2); // Input spatial dim
+  const int iH = 1 + in_dilation[0] * (in.shape(1) - 1); // Input spatial dim
+  const int iW = 1 + in_dilation[1] * (in.shape(2) - 1); // Input spatial dim
   const int oH = out.shape(1); // Output spatial dim
   const int oW = out.shape(2); // Output spatial dim
   const int O = wt.shape(0); // Out channels
@@ -120,6 +128,8 @@ void slow_conv_2D(
   const size_t out_stride_W = out.strides()[2];
   const size_t out_stride_O = out.strides()[3];
 
+  bool is_idil_one = in_dilation[0] == 1 && in_dilation[1] == 1;
+
   auto pt_conv_no_checks =
       [&](const T* in_ptr, const T* wt_ptr, T* out_ptr, int oh, int ow) {
         out_ptr += oh * out_stride_H + ow * out_stride_W;
@@ -131,8 +141,10 @@ void slow_conv_2D(
 
           for (int wh = 0; wh < wH; ++wh) {
             for (int ww = 0; ww < wW; ++ww) {
-              int ih = ih_base + wh * wt_dilation[0];
-              int iw = iw_base + ww * wt_dilation[1];
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int ih = ih_base + wh_flip * wt_dilation[0];
+              int iw = iw_base + ww_flip * wt_dilation[1];
 
               const T* wt_ptr_pt = wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
               const T* in_ptr_pt = in_ptr + ih * in_stride_H + iw * in_stride_W;
@@ -153,25 +165,74 @@ void slow_conv_2D(
         } // o
       };
 
+  int jump_h = flip ? -wt_dilation[0] : wt_dilation[0];
+  int jump_w = flip ? -wt_dilation[1] : wt_dilation[1];
+
+  int init_h = (flip ? (wH - 1) * wt_dilation[0] : 0);
+  int init_w = (flip ? (wW - 1) * wt_dilation[1] : 0);
+
+  int f_wgt_jump_h = std::lcm(in_dilation[0], wt_dilation[0]) / wt_dilation[0];
+  int f_wgt_jump_w = std::lcm(in_dilation[1], wt_dilation[1]) / wt_dilation[1];
+
+  int f_out_jump_h = std::lcm(in_dilation[0], wt_strides[0]) / wt_strides[0];
+  int f_out_jump_w = std::lcm(in_dilation[1], wt_strides[1]) / wt_strides[1];
+
+  std::vector<int> base_h(f_out_jump_h);
+  std::vector<int> base_w(f_out_jump_w);
+
+  for (int i = 0; i < f_out_jump_h; ++i) {
+    int ih_loop = i * wt_strides[0] - padding[0] + init_h;
+
+    int wh_base = 0;
+    while (wh_base < wH && ih_loop % in_dilation[0] != 0) {
+      wh_base++;
+      ih_loop += jump_h;
+    }
+
+    base_h[i] = wh_base;
+  }
+
+  for (int j = 0; j < f_out_jump_w; ++j) {
+    int iw_loop = j * wt_strides[1] - padding[1] + init_w;
+
+    int ww_base = 0;
+    while (ww_base < wW && iw_loop % in_dilation[1] != 0) {
+      ww_base++;
+      iw_loop += jump_w;
+    }
+
+    base_w[j] = ww_base;
+  }
+
   auto pt_conv_all_checks =
       [&](const T* in_ptr, const T* wt_ptr, T* out_ptr, int oh, int ow) {
         out_ptr += oh * out_stride_H + ow * out_stride_W;
+
         int ih_base = oh * wt_strides[0] - padding[0];
         int iw_base = ow * wt_strides[1] - padding[1];
 
+        int wh_base = base_h[oh % f_out_jump_h];
+        int ww_base = base_w[ow % f_out_jump_w];
+
         for (int o = 0; o < O; ++o) {
           float r = 0.;
 
-          for (int wh = 0; wh < wH; ++wh) {
-            for (int ww = 0; ww < wW; ++ww) {
-              int ih = ih_base + wh * wt_dilation[0];
-              int iw = iw_base + ww * wt_dilation[1];
+          for (int wh = wh_base; wh < wH; wh += f_wgt_jump_h) {
+            for (int ww = ww_base; ww < wW; ww += f_wgt_jump_w) {
+              int wh_flip = flip ? wH - wh - 1 : wh;
+              int ww_flip = flip ? wW - ww - 1 : ww;
+              int ih = ih_base + wh_flip * wt_dilation[0];
+              int iw = iw_base + ww_flip * wt_dilation[1];
 
               if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) {
                 const T* wt_ptr_pt =
                     wt_ptr + wh * wt_stride_H + ww * wt_stride_W;
+
+                int ih_dil = !is_idil_one ? (ih / in_dilation[0]) : ih;
+                int iw_dil = !is_idil_one ? (iw / in_dilation[1]) : iw;
+
                 const T* in_ptr_pt =
-                    in_ptr + ih * in_stride_H + iw * in_stride_W;
+                    in_ptr + ih_dil * in_stride_H + iw_dil * in_stride_W;
 
                 for (int c = 0; c < C; ++c) {
                   r += static_cast<float>(in_ptr_pt[0]) *
@@ -191,13 +252,17 @@ void slow_conv_2D(
       };
 
   int oH_border_0 = 0;
-  int oH_border_1 = (padding[0] + wt_strides[0] + 1) / wt_strides[0];
-  int oH_border_2 = (iH + padding[0] - wH * wt_dilation[0]) / wt_strides[0];
+  int oH_border_1 =
+      is_idil_one ? ((padding[0] + wt_strides[0] - 1) / wt_strides[0]) : oH;
+  int oH_border_2 = std::max(
+      oH_border_1, (iH + padding[0] - wH * wt_dilation[0]) / wt_strides[0]);
   int oH_border_3 = oH;
 
   int oW_border_0 = 0;
-  int oW_border_1 = (padding[1] + wt_strides[0] + 1) / wt_strides[1];
-  int oW_border_2 = (iW + padding[1] - wW * wt_dilation[1]) / wt_strides[1];
+  int oW_border_1 =
+      is_idil_one ? ((padding[1] + wt_strides[1] - 1) / wt_strides[1]) : oW;
+  int oW_border_2 = std::max(
+      oW_border_1, (iW + padding[1] - wW * wt_dilation[1]) / wt_strides[1]);
   int oW_border_3 = oW;
 
   for (int n = 0; n < N; ++n) {
@@ -246,15 +311,18 @@ void dispatch_slow_conv_1D(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
   if (in.dtype() == float32) {
-    return slow_conv_1D<float>(in, wt, out, padding, wt_strides, wt_dilation);
+    return slow_conv_1D<float>(
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else if (in.dtype() == float16) {
     return slow_conv_1D<float16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation);
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else if (in.dtype() == bfloat16) {
     return slow_conv_1D<bfloat16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation);
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else {
     throw std::invalid_argument(
         "[Convolution::eval] got unsupported data type.");
@@ -267,15 +335,18 @@ void dispatch_slow_conv_2D(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
   if (in.dtype() == float32) {
-    return slow_conv_2D<float>(in, wt, out, padding, wt_strides, wt_dilation);
+    return slow_conv_2D<float>(
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else if (in.dtype() == float16) {
     return slow_conv_2D<float16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation);
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else if (in.dtype() == bfloat16) {
     return slow_conv_2D<bfloat16_t>(
-        in, wt, out, padding, wt_strides, wt_dilation);
+        in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
   } else {
     throw std::invalid_argument(
         "[Convolution::eval] got unsupported data type.");
@@ -493,13 +564,16 @@ void conv_1D_cpu(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
-  if (wt_dilation[0] == 1) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
+  if (wt_dilation[0] == 1 && in_dilation[0] == 1 && !flip) {
     return explicit_gemm_conv_1D_cpu(
         in, wt, out, padding, wt_strides, wt_dilation);
   }
 
-  return dispatch_slow_conv_1D(in, wt, out, padding, wt_strides, wt_dilation);
+  return dispatch_slow_conv_1D(
+      in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
 }
 
 void conv_2D_cpu(
@@ -508,8 +582,11 @@ void conv_2D_cpu(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
-  return dispatch_slow_conv_2D(in, wt, out, padding, wt_strides, wt_dilation);
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
+  return dispatch_slow_conv_2D(
+      in, wt, out, padding, wt_strides, wt_dilation, in_dilation, flip);
 }
 
 } // namespace
@@ -523,12 +600,26 @@ void Convolution::eval(const std::vector<array>& inputs, array& out) {
   // 2D convolution
   if (in.ndim() == (2 + 2)) {
     return conv_2D_cpu(
-        in, wt, out, padding_, kernel_strides_, kernel_dilation_);
+        in,
+        wt,
+        out,
+        padding_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_);
   }
   // 1D convolution
   else if (in.ndim() == (1 + 2)) {
     return conv_1D_cpu(
-        in, wt, out, padding_, kernel_strides_, kernel_dilation_);
+        in,
+        wt,
+        out,
+        padding_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_);
   }
   // Throw error
   else {

mlx/backend/common/default_primitives.cpp

@@ -43,7 +43,6 @@ DEFAULT(AsStrided)
 DEFAULT(Broadcast)
 DEFAULT_MULTI(DivMod)
 DEFAULT(Ceil)
-DEFAULT_MULTI(Compiled)
 DEFAULT(Concatenate)
 DEFAULT(Convolution)
 DEFAULT(Copy)
@@ -88,6 +87,7 @@ DEFAULT(Reshape)
 DEFAULT(Round)
 DEFAULT(Scan)
 DEFAULT(Scatter)
+DEFAULT(Select)
 DEFAULT(Sigmoid)
 DEFAULT(Sign)
 DEFAULT(Sin)

mlx/backend/common/erf.h

@@ -1,11 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-namespace mlx::core {
-
-/* Approximation to the inverse error function.
- * Based on code from:
- *   https://stackoverflow.com/questions/27229371/inverse-error-function-in-c#answer-49743348
- */
-float erfinv(float a);
-
-} // namespace mlx::core

mlx/backend/common/make_compiled_preamble.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+#
+# This script generates a C++ function that provides the CPU
+# code for use with kernel generation.
+#
+# Copyright © 2023-24 Apple Inc.
+
+
+OUTPUT_FILE=$1
+GCC=$2
+SRCDIR=$3
+CLANG=$4
+
+if [ $CLANG = "TRUE" ]; then
+  read -r -d '' INCLUDES <<- EOM
+  #include <cmath>
+  #include <complex>
+  #include <cstdint>
+  #include <vector>
+EOM
+
+fi
+
+CONTENT=$($GCC -I $SRCDIR -E $SRCDIR/mlx/backend/common/compiled_preamble.h 2>/dev/null)
+
+cat << EOF > "$OUTPUT_FILE"
+const char* get_kernel_preamble() {
+return R"preamble(
+$INCLUDES
+$CONTENT
+using namespace mlx::core::detail;
+)preamble";
+}
+EOF

mlx/backend/common/ops.h

@@ -0,0 +1,602 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+#include <stdint.h>
+#include <cmath>
+#include <complex>
+
+namespace mlx::core::detail {
+
+namespace {
+constexpr float inf = std::numeric_limits<float>::infinity();
+} // namespace
+
+typedef union {
+  int i;
+  float f;
+} IntOrFloat;
+
+inline float fast_exp(float x) {
+  if (x == -std::numeric_limits<float>::infinity()) {
+    return 0.0f;
+  } else if (x == std::numeric_limits<float>::infinity() || std::isnan(x)) {
+    return x;
+  }
+  x *= 1.442695; // multiply with log_2(e)
+  float ipart, fpart;
+  IntOrFloat epart;
+  x = std::max(-80.f, std::min(x, 80.f));
+  ipart = std::floor(x + 0.5);
+  fpart = x - ipart;
+
+  x = 1.535336188319500e-4f;
+  x = x * fpart + 1.339887440266574e-3f;
+  x = x * fpart + 9.618437357674640e-3f;
+  x = x * fpart + 5.550332471162809e-2f;
+  x = x * fpart + 2.402264791363012e-1f;
+  x = x * fpart + 6.931472028550421e-1f;
+  x = x * fpart + 1.000000000000000f;
+
+  // generate 2**ipart in the floating point representation using integer
+  // bitshifting
+  epart.i = (int(ipart) + 127) << 23;
+
+  return epart.f * x;
+}
+
+inline float fast_erf(float a) {
+  float r, s, t, u;
+  t = std::abs(a);
+  s = a * a;
+  if (t > 0.927734375f) {
+    // maximum error 0.99527 ulp
+    r = std::fma(
+        -1.72853470e-5f, t, 3.83197126e-4f); // -0x1.220000p-16,0x1.91cfb2p-12
+    u = std::fma(
+        -3.88396438e-3f, t, 2.42546219e-2f); // -0x1.fd1438p-9, 0x1.8d6342p-6
+    r = std::fma(r, s, u);
+    r = std::fma(r, t, -1.06777877e-1f); // -0x1.b55cb8p-4
+    r = std::fma(r, t, -6.34846687e-1f); // -0x1.450aa0p-1
+    r = std::fma(r, t, -1.28717512e-1f); // -0x1.079d0cp-3
+    r = std::fma(r, t, -t);
+    // TODO, replace with expm1 when implemented
+    r = 1.0f - std::exp(r);
+    r = std::copysign(r, a);
+  } else {
+    // maximum error 0.98929 ulp
+    r = -5.96761703e-4f; // -0x1.38e000p-11
+    r = std::fma(r, s, 4.99119423e-3f); //  0x1.471a58p-8
+    r = std::fma(r, s, -2.67681349e-2f); // -0x1.b691b2p-6
+    r = std::fma(r, s, 1.12819925e-1f); //  0x1.ce1c44p-4
+    r = std::fma(r, s, -3.76125336e-1f); // -0x1.812700p-2
+    r = std::fma(r, s, 1.28379166e-1f); //  0x1.06eba8p-3
+    r = std::fma(r, a, a);
+  }
+  return r;
+}
+
+inline float fast_erfinv(float a) {
+  auto t = std::fma(a, 0.0f - a, 1.0f);
+  t = std::log(t);
+  float p;
+  if (std::abs(t) > 6.125f) { // maximum ulp error = 2.35793
+    p = 3.03697567e-10f; //  0x1.4deb44p-32
+    p = std::fma(p, t, 2.93243101e-8f); //  0x1.f7c9aep-26
+    p = std::fma(p, t, 1.22150334e-6f); //  0x1.47e512p-20
+    p = std::fma(p, t, 2.84108955e-5f); //  0x1.dca7dep-16
+    p = std::fma(p, t, 3.93552968e-4f); //  0x1.9cab92p-12
+    p = std::fma(p, t, 3.02698812e-3f); //  0x1.8cc0dep-9
+    p = std::fma(p, t, 4.83185798e-3f); //  0x1.3ca920p-8
+    p = std::fma(p, t, -2.64646143e-1f); // -0x1.0eff66p-2
+    p = std::fma(p, t, 8.40016484e-1f); //  0x1.ae16a4p-1
+  } else { // maximum ulp error = 2.35002
+    p = 5.43877832e-9f; //  0x1.75c000p-28
+    p = std::fma(p, t, 1.43285448e-7f); //  0x1.33b402p-23
+    p = std::fma(p, t, 1.22774793e-6f); //  0x1.499232p-20
+    p = std::fma(p, t, 1.12963626e-7f); //  0x1.e52cd2p-24
+    p = std::fma(p, t, -5.61530760e-5f); // -0x1.d70bd0p-15
+    p = std::fma(p, t, -1.47697632e-4f); // -0x1.35be90p-13
+    p = std::fma(p, t, 2.31468678e-3f); //  0x1.2f6400p-9
+    p = std::fma(p, t, 1.15392581e-2f); //  0x1.7a1e50p-7
+    p = std::fma(p, t, -2.32015476e-1f); // -0x1.db2aeep-3
+    p = std::fma(p, t, 8.86226892e-1f); //  0x1.c5bf88p-1
+  }
+  return a * p;
+}
+
+struct Abs {
+  template <typename T>
+  T operator()(T x) {
+    return std::abs(x);
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+};
+
+struct ArcCos {
+  template <typename T>
+  T operator()(T x) {
+    return std::acos(x);
+  };
+};
+
+struct ArcCosh {
+  template <typename T>
+  T operator()(T x) {
+    return std::acosh(x);
+  };
+};
+
+struct ArcSin {
+  template <typename T>
+  T operator()(T x) {
+    return std::asin(x);
+  };
+};
+
+struct ArcSinh {
+  template <typename T>
+  T operator()(T x) {
+    return std::asinh(x);
+  };
+};
+
+struct ArcTan {
+  template <typename T>
+  T operator()(T x) {
+    return std::atan(x);
+  };
+};
+
+struct ArcTanh {
+  template <typename T>
+  T operator()(T x) {
+    return std::atanh(x);
+  };
+};
+
+struct Ceil {
+  template <typename T>
+  T operator()(T x) {
+    return std::ceil(x);
+  };
+  int8_t operator()(int8_t x) {
+    return x;
+  };
+  int16_t operator()(int16_t x) {
+    return x;
+  };
+  int32_t operator()(int32_t x) {
+    return x;
+  };
+  int64_t operator()(int64_t x) {
+    return x;
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+};
+
+struct Cos {
+  template <typename T>
+  T operator()(T x) {
+    return std::cos(x);
+  };
+};
+
+struct Cosh {
+  template <typename T>
+  T operator()(T x) {
+    return std::cosh(x);
+  };
+};
+
+struct Erf {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<T>(fast_erf(static_cast<float>(x)));
+  };
+};
+
+struct ErfInv {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<T>(fast_erfinv(static_cast<float>(x)));
+  };
+};
+
+struct Exp {
+  template <typename T>
+  T operator()(T x) {
+    return fast_exp(x);
+  };
+
+  complex64_t operator()(complex64_t x) {
+    return std::exp(x);
+  }
+};
+
+struct Floor {
+  template <typename T>
+  T operator()(T x) {
+    return std::floor(x);
+  };
+  int8_t operator()(int8_t x) {
+    return x;
+  };
+  int16_t operator()(int16_t x) {
+    return x;
+  };
+  int32_t operator()(int32_t x) {
+    return x;
+  };
+  int64_t operator()(int64_t x) {
+    return x;
+  };
+  uint8_t operator()(uint8_t x) {
+    return x;
+  };
+  uint16_t operator()(uint16_t x) {
+    return x;
+  };
+  uint32_t operator()(uint32_t x) {
+    return x;
+  };
+  uint64_t operator()(uint64_t x) {
+    return x;
+  };
+  bool operator()(bool x) {
+    return x;
+  };
+};
+
+struct Log {
+  template <typename T>
+  T operator()(T x) {
+    return std::log(x);
+  };
+};
+
+struct Log2 {
+  template <typename T>
+  T operator()(T x) {
+    return std::log2(x);
+  };
+};
+
+struct Log10 {
+  template <typename T>
+  T operator()(T x) {
+    return std::log10(x);
+  };
+};
+
+struct Log1p {
+  template <typename T>
+  T operator()(T x) {
+    return log1p(x);
+  };
+};
+
+struct LogicalNot {
+  template <typename T>
+  T operator()(T x) {
+    return !x;
+  };
+};
+
+struct Negative {
+  template <typename T>
+  T operator()(T x) {
+    return -x;
+  };
+};
+
+struct Round {
+  template <typename T>
+  T operator()(T x) {
+    return std::rint(x);
+  }
+
+  complex64_t operator()(complex64_t x) {
+    return {std::rint(x.real()), std::rint(x.imag())};
+  }
+};
+
+struct Sigmoid {
+  template <typename T>
+  T operator()(T x) {
+    auto one = static_cast<decltype(x)>(1.0);
+    return one / (one + fast_exp(-x));
+  }
+};
+
+struct Sign {
+  template <typename T>
+  T operator()(T x) {
+    return (x > T(0)) - (x < T(0));
+  }
+  uint8_t operator()(uint8_t x) {
+    return x != 0;
+  }
+  uint16_t operator()(uint16_t x) {
+    return x != 0;
+  }
+  uint32_t operator()(uint32_t x) {
+    return x != 0;
+  }
+  uint64_t operator()(uint64_t x) {
+    return x != 0;
+  }
+};
+
+struct Sin {
+  template <typename T>
+  T operator()(T x) {
+    return std::sin(x);
+  };
+};
+
+struct Sinh {
+  template <typename T>
+  T operator()(T x) {
+    return std::sinh(x);
+  };
+};
+
+struct Square {
+  template <typename T>
+  T operator()(T x) {
+    return x * x;
+  };
+};
+
+struct Sqrt {
+  template <typename T>
+  T operator()(T x) {
+    return std::sqrt(x);
+  };
+};
+
+struct Rsqrt {
+  template <typename T>
+  T operator()(T x) {
+    return static_cast<decltype(x)>(1.0) / std::sqrt(x);
+  };
+};
+
+struct Tan {
+  template <typename T>
+  T operator()(T x) {
+    return std::tan(x);
+  };
+};
+
+struct Tanh {
+  template <typename T>
+  T operator()(T x) {
+    return std::tanh(x);
+  };
+};
+
+struct Add {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x + y;
+  }
+};
+
+struct Divide {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x / y;
+  }
+};
+
+struct Remainder {
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T> & !std::is_signed_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    return numerator % denominator;
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T> & std::is_signed_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    auto r = numerator % denominator;
+    if (r != 0 && (r < 0 != denominator < 0))
+      r += denominator;
+    return r;
+  }
+
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(
+      T numerator,
+      T denominator) {
+    auto r = std::fmod(numerator, denominator);
+    if (r != 0 && (r < 0 != denominator < 0)) {
+      r += denominator;
+    }
+    return r;
+  }
+
+  complex64_t operator()(complex64_t numerator, complex64_t denominator) {
+    return numerator % denominator;
+  }
+};
+
+struct Equal {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x == y;
+  }
+};
+
+struct NaNEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x == y || (std::isnan(x) && std::isnan(y));
+  }
+};
+
+struct Greater {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x > y;
+  }
+};
+
+struct GreaterEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x >= y;
+  }
+};
+
+struct Less {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x < y;
+  }
+};
+
+struct LessEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x <= y;
+  }
+};
+
+struct Maximum {
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(T x, T y) {
+    return (x > y) ? x : y;
+  }
+
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(T x, T y) {
+    if (std::isnan(x)) {
+      return x;
+    }
+    return (x > y) ? x : y;
+  }
+};
+
+struct Minimum {
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(T x, T y) {
+    return x < y ? x : y;
+  }
+
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(T x, T y) {
+    if (std::isnan(x)) {
+      return x;
+    }
+    return x < y ? x : y;
+  }
+};
+
+struct LogAddExp {
+  template <typename T>
+  T operator()(T x, T y) {
+    constexpr float inf = std::numeric_limits<float>::infinity();
+    auto maxval = Maximum()(x, y);
+    auto minval = Minimum()(x, y);
+    return (minval == -inf || maxval == inf)
+        ? maxval
+        : static_cast<decltype(x)>(
+              maxval + std::log1p(fast_exp(minval - maxval)));
+  };
+};
+
+struct Multiply {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x * y;
+  }
+};
+
+struct NotEqual {
+  template <typename T>
+  bool operator()(T x, T y) {
+    return x != y;
+  }
+};
+
+struct Power {
+  template <typename T>
+  std::enable_if_t<!std::is_integral_v<T>, T> operator()(T base, T exp) {
+    return std::pow(base, exp);
+  }
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, T> operator()(T base, T exp) {
+    T res = 1;
+    while (exp) {
+      if (exp & 1) {
+        res *= base;
+      }
+      exp >>= 1;
+      base *= base;
+    }
+    return res;
+  }
+};
+
+struct Subtract {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x - y;
+  }
+};
+
+struct LogicalAnd {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x && y;
+  };
+};
+
+struct LogicalOr {
+  template <typename T>
+  T operator()(T x, T y) {
+    return x || y;
+  };
+};
+
+struct Select {
+  template <typename T>
+  T operator()(bool condition, T x, T y) {
+    return condition ? x : y;
+  }
+};
+
+} // namespace mlx::core::detail

mlx/backend/common/primitives.cpp

@@ -10,7 +10,7 @@
 #include "mlx/backend/common/arange.h"
 #include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/copy.h"
-#include "mlx/backend/common/erf.h"
+#include "mlx/backend/common/ops.h"
 #include "mlx/backend/common/threefry.h"
 #include "mlx/backend/common/unary.h"
 #include "mlx/backend/common/utils.h"
@@ -26,7 +26,7 @@ void Abs::eval(const std::vector<array>& inputs, array& out) {
     // No-op for unsigned types
     out.copy_shared_buffer(in);
   } else {
-    unary(in, out, AbsOp());
+    unary(in, out, detail::Abs());
   }
 }
 
@@ -38,7 +38,7 @@ void ArcCos::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::acos(x); });
+    unary_fp(in, out, detail::ArcCos());
   } else {
     throw std::invalid_argument(
         "[arccos] Cannot compute inverse cosine of elements in array"
@@ -50,7 +50,7 @@ void ArcCosh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::acosh(x); });
+    unary_fp(in, out, detail::ArcCosh());
   } else {
     throw std::invalid_argument(
         "[arccosh] Cannot compute inverse hyperbolic cosine of elements in"
@@ -62,7 +62,7 @@ void ArcSin::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::asin(x); });
+    unary_fp(in, out, detail::ArcSin());
   } else {
     throw std::invalid_argument(
         "[arcsin] Cannot compute inverse sine of elements in array"
@@ -74,7 +74,7 @@ void ArcSinh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::asinh(x); });
+    unary_fp(in, out, detail::ArcSinh());
   } else {
     throw std::invalid_argument(
         "[arcsinh] Cannot compute inverse hyperbolic sine of elements in"
@@ -86,7 +86,7 @@ void ArcTan::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::atan(x); });
+    unary_fp(in, out, detail::ArcTan());
   } else {
     throw std::invalid_argument(
         "[arctan] Cannot compute inverse tangent of elements in array"
@@ -98,7 +98,7 @@ void ArcTanh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::atanh(x); });
+    unary_fp(in, out, detail::ArcTanh());
   } else {
     throw std::invalid_argument(
         "[arctanh] Cannot compute inverse hyperbolic tangent of elements in"
@@ -172,7 +172,7 @@ void Ceil::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (not is_integral(in.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::ceil(x); });
+    unary_fp(in, out, detail::Ceil());
   } else {
     // No-op integer types
     out.copy_shared_buffer(in);
@@ -212,7 +212,7 @@ void Cos::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::cos(x); });
+    unary_fp(in, out, detail::Cos());
   } else {
     throw std::invalid_argument(
         "[cos] Cannot compute cosine of elements in array"
@@ -224,7 +224,7 @@ void Cosh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::cosh(x); });
+    unary_fp(in, out, detail::Cosh());
   } else {
     throw std::invalid_argument(
         "[cosh] Cannot compute hyperbolic cosine of elements in array"
@@ -256,17 +256,13 @@ void Erf::eval(const std::vector<array>& inputs, array& out) {
   const auto& in = inputs[0];
   switch (out.dtype()) {
     case float32:
-      unary_op<float>(in, out, [](auto x) { return std::erf(x); });
+      unary_op<float>(in, out, detail::Erf());
       break;
     case float16:
-      unary_op<float16_t>(in, out, [](auto x) {
-        return static_cast<float16_t>(std::erf(static_cast<float>(x)));
-      });
+      unary_op<float16_t>(in, out, detail::Erf());
       break;
     case bfloat16:
-      unary_op<bfloat16_t>(in, out, [](auto x) {
-        return static_cast<bfloat16_t>(std::erf(static_cast<float>(x)));
-      });
+      unary_op<bfloat16_t>(in, out, detail::Erf());
       break;
     default:
       throw std::invalid_argument(
@@ -280,17 +276,13 @@ void ErfInv::eval(const std::vector<array>& inputs, array& out) {
   const auto& in = inputs[0];
   switch (out.dtype()) {
     case float32:
-      unary_op<float>(in, out, [](auto x) { return erfinv(x); });
+      unary_op<float>(in, out, detail::ErfInv());
       break;
     case float16:
-      unary_op<float16_t>(in, out, [](auto x) {
-        return static_cast<float16_t>(erfinv(static_cast<float>(x)));
-      });
+      unary_op<float16_t>(in, out, detail::ErfInv());
       break;
     case bfloat16:
-      unary_op<bfloat16_t>(in, out, [](auto x) {
-        return static_cast<bfloat16_t>(erfinv(static_cast<float>(x)));
-      });
+      unary_op<bfloat16_t>(in, out, detail::ErfInv());
       break;
     default:
       throw std::invalid_argument(
@@ -302,9 +294,8 @@ void ErfInv::eval(const std::vector<array>& inputs, array& out) {
 void Exp::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
-
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::exp(x); });
+    unary_fp(in, out, detail::Exp());
   } else {
     throw std::invalid_argument(
         "[exp] Cannot exponentiate elements in array"
@@ -316,7 +307,7 @@ void Floor::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (not is_integral(in.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::floor(x); });
+    unary_fp(in, out, detail::Floor());
   } else {
     // No-op integer types
     out.copy_shared_buffer(in);
@@ -344,13 +335,13 @@ void Log::eval(const std::vector<array>& inputs, array& out) {
   if (is_floating_point(out.dtype())) {
     switch (base_) {
       case Base::e:
-        unary_fp(in, out, [](auto x) { return std::log(x); });
+        unary_fp(in, out, detail::Log());
         break;
       case Base::two:
-        unary_fp(in, out, [](auto x) { return std::log2(x); });
+        unary_fp(in, out, detail::Log2());
         break;
       case Base::ten:
-        unary_fp(in, out, [](auto x) { return std::log10(x); });
+        unary_fp(in, out, detail::Log10());
         break;
     }
   } else {
@@ -364,7 +355,7 @@ void Log1p::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::log1p(x); });
+    unary_fp(in, out, detail::Log1p());
   } else {
     throw std::invalid_argument(
         "[log1p] Cannot compute log of elements in array with"
@@ -375,27 +366,27 @@ void Log1p::eval(const std::vector<array>& inputs, array& out) {
 void LogicalNot::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  unary(in, out, [](auto x) { return !x; });
+  unary(in, out, detail::LogicalNot());
 }
 
 void LogicalAnd::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2); // LogicalAnd requires two input arrays
   auto& in1 = inputs[0];
   auto& in2 = inputs[1];
-  binary(in1, in2, out, [](auto x, auto y) { return x && y; });
+  binary(in1, in2, out, detail::LogicalAnd());
 }
 
 void LogicalOr::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2); // LogicalOr requires two input arrays
   auto& in1 = inputs[0];
   auto& in2 = inputs[1];
-  binary(in1, in2, out, [](auto x, auto y) { return x || y; });
+  binary(in1, in2, out, detail::LogicalOr());
 }
 
 void Negative::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  unary(in, out, [](auto x) { return -x; });
+  unary(in, out, detail::Negative());
 }
 
 void Pad::eval(const std::vector<array>& inputs, array& out) {
@@ -498,7 +489,7 @@ void Round::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (not is_integral(in.dtype())) {
-    unary_fp(in, out, RoundOp());
+    unary_fp(in, out, detail::Round());
   } else {
     // No-op integer types
     out.copy_shared_buffer(in);
@@ -509,11 +500,7 @@ void Sigmoid::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    auto sigmoid_op = [](auto x) {
-      auto one = static_cast<decltype(x)>(1.0);
-      return one / (one + std::exp(-x));
-    };
-    unary_fp(in, out, sigmoid_op);
+    unary_fp(in, out, detail::Sigmoid());
   } else {
     throw std::invalid_argument(
         "[sigmoid] Cannot sigmoid of elements in array with"
@@ -527,7 +514,7 @@ void Sign::eval(const std::vector<array>& inputs, array& out) {
   if (in.dtype() == bool_) {
     out.copy_shared_buffer(in);
   } else {
-    unary(in, out, SignOp());
+    unary(in, out, detail::Sign());
   }
 }
 
@@ -535,7 +522,7 @@ void Sin::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::sin(x); });
+    unary_fp(in, out, detail::Sin());
   } else {
     throw std::invalid_argument(
         "[sin] Cannot compute sine of elements in array"
@@ -547,7 +534,7 @@ void Sinh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::sinh(x); });
+    unary_fp(in, out, detail::Sinh());
   } else {
     throw std::invalid_argument(
         "[sinh] Cannot compute hyperbolic sine of elements in array"
@@ -656,18 +643,16 @@ void Split::eval(
 void Square::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
-  unary(in, out, [](auto x) { return x * x; });
+  unary(in, out, detail::Square());
 }
 
 void Sqrt::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
   if (recip_) {
-    unary_fp(in, out, [](auto x) {
-      return static_cast<decltype(x)>(1.0) / sqrt(x);
-    });
+    unary_fp(in, out, detail::Rsqrt());
   } else {
-    unary_fp(in, out, [](auto x) { return sqrt(x); });
+    unary_fp(in, out, detail::Sqrt());
   }
 }
 
@@ -680,7 +665,7 @@ void Tan::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::tan(x); });
+    unary_fp(in, out, detail::Tan());
   } else {
     throw std::invalid_argument(
         "[tan] Cannot compute tangent of elements in array"
@@ -692,7 +677,7 @@ void Tanh::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   const auto& in = inputs[0];
   if (is_floating_point(out.dtype())) {
-    unary_fp(in, out, [](auto x) { return std::tanh(x); });
+    unary_fp(in, out, detail::Tanh());
   } else {
     throw std::invalid_argument(
         "[tanh] Cannot compute hyperbolic tangent of elements in array"

mlx/backend/common/rope.cpp

@@ -1,7 +1,6 @@
 // Copyright © 2023-2024 Apple Inc.
 
-#include "mlx/fast.h"
-#include "mlx/primitives.h"
+#include "mlx/fast_primitives.h"
 
 namespace mlx::core::fast {
 

mlx/backend/common/select.cpp

@@ -0,0 +1,72 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/ternary.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+template <typename Op>
+void select_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  switch (out.dtype()) {
+    case bool_:
+      ternary_op<bool, bool, bool, bool>(a, b, c, out, op);
+      break;
+    case uint8:
+      ternary_op<bool, uint8_t, uint8_t, uint8_t>(a, b, c, out, op);
+      break;
+    case uint16:
+      ternary_op<bool, uint16_t, uint16_t, uint16_t>(a, b, c, out, op);
+      break;
+    case uint32:
+      ternary_op<bool, uint32_t, uint32_t, uint32_t>(a, b, c, out, op);
+      break;
+    case uint64:
+      ternary_op<bool, uint64_t, uint64_t, uint64_t>(a, b, c, out, op);
+      break;
+    case int8:
+      ternary_op<bool, int8_t, int8_t, int8_t>(a, b, c, out, op);
+      break;
+    case int16:
+      ternary_op<bool, int16_t, int16_t, int16_t>(a, b, c, out, op);
+      break;
+    case int32:
+      ternary_op<bool, int32_t, int32_t, int32_t>(a, b, c, out, op);
+      break;
+    case int64:
+      ternary_op<bool, int64_t, int64_t, int64_t>(a, b, c, out, op);
+      break;
+    case float16:
+      ternary_op<bool, float16_t, float16_t, float16_t>(a, b, c, out, op);
+      break;
+    case float32:
+      ternary_op<bool, float, float, float>(a, b, c, out, op);
+      break;
+    case bfloat16:
+      ternary_op<bool, bfloat16_t, bfloat16_t, bfloat16_t>(a, b, c, out, op);
+      break;
+    case complex64:
+      ternary_op<bool, complex64_t, complex64_t, complex64_t>(a, b, c, out, op);
+      break;
+  }
+}
+
+} // namespace
+
+void Select::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 3);
+  const auto& condition = inputs[0];
+  const auto& a = inputs[1];
+  const auto& b = inputs[2];
+  select_op(condition, a, b, out, detail::Select());
+}
+
+} // namespace mlx::core

mlx/backend/common/ternary.h

@@ -0,0 +1,226 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+#include "mlx/allocator.h"
+#include "mlx/array.h"
+#include "mlx/backend/common/ops.h"
+#include "mlx/backend/common/utils.h"
+namespace mlx::core {
+
+namespace {
+
+// TODO: Add support for more combinations of input types.
+enum class TernaryOpType {
+  ScalarScalarScalar,
+  General,
+};
+
+TernaryOpType
+get_ternary_op_type(const array& a, const array& b, const array& c) {
+  TernaryOpType topt;
+  if (a.data_size() == 1 && b.data_size() == 1 && c.data_size() == 1) {
+    topt = TernaryOpType::ScalarScalarScalar;
+  } else {
+    topt = TernaryOpType::General;
+  }
+  return topt;
+}
+
+void set_ternary_op_output_data(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    TernaryOpType topt,
+    bool donate_with_move = false) {
+  switch (topt) {
+    case TernaryOpType::ScalarScalarScalar:
+      out.set_data(
+          allocator::malloc_or_wait(out.itemsize()), 1, b.strides(), b.flags());
+      break;
+    case TernaryOpType::General:
+      out.set_data(allocator::malloc_or_wait(out.nbytes()));
+      break;
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dims1(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+
+  U* dst = out.data<U>();
+  size_t a_idx = 0;
+  size_t b_idx = 0;
+  size_t c_idx = 0;
+  for (size_t i = 0; i < out.size(); ++i) {
+    dst[i] = op(a_ptr[a_idx], b_ptr[b_idx], c_ptr[c_idx]);
+    a_idx += a.strides()[0];
+    b_idx += b.strides()[0];
+    c_idx += c.strides()[0];
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dims2(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+
+  U* dst = out.data<U>();
+  size_t a_idx = 0;
+  size_t b_idx = 0;
+  size_t c_idx = 0;
+  size_t out_idx = 0;
+  for (size_t i = 0; i < a.shape()[0]; ++i) {
+    for (size_t j = 0; j < a.shape()[1]; ++j) {
+      dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx], c_ptr[c_idx]);
+      a_idx += a.strides()[1];
+      b_idx += b.strides()[1];
+      c_idx += c.strides()[1];
+    }
+    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
+    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
+    c_idx += c.strides()[0] - c.strides()[1] * c.shape()[1];
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dims3(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+  U* dst = out.data<U>();
+  size_t a_idx = 0;
+  size_t b_idx = 0;
+  size_t c_idx = 0;
+  size_t out_idx = 0;
+  for (size_t i = 0; i < a.shape()[0]; ++i) {
+    for (size_t j = 0; j < a.shape()[1]; ++j) {
+      for (size_t k = 0; k < a.shape()[2]; ++k) {
+        dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx], c_ptr[c_idx]);
+        a_idx += a.strides()[2];
+        b_idx += b.strides()[2];
+        c_idx += c.strides()[2];
+      }
+      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
+      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
+      c_idx += c.strides()[1] - c.strides()[2] * c.shape()[2];
+    }
+    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
+    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
+    c_idx += c.strides()[0] - c.strides()[1] * c.shape()[1];
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dims4(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+
+  U* dst = out.data<U>();
+  size_t a_idx = 0;
+  size_t b_idx = 0;
+  size_t c_idx = 0;
+  size_t out_idx = 0;
+  for (size_t i = 0; i < a.shape()[0]; ++i) {
+    for (size_t j = 0; j < a.shape()[1]; ++j) {
+      for (size_t k = 0; k < a.shape()[2]; ++k) {
+        for (size_t ii = 0; ii < a.shape()[3]; ++ii) {
+          dst[out_idx++] = op(a_ptr[a_idx], b_ptr[b_idx], c_ptr[c_idx]);
+          a_idx += a.strides()[3];
+          b_idx += b.strides()[3];
+          c_idx += c.strides()[3];
+        }
+        a_idx += a.strides()[2] - a.strides()[3] * a.shape()[3];
+        b_idx += b.strides()[2] - b.strides()[3] * b.shape()[3];
+        c_idx += c.strides()[2] - c.strides()[3] * c.shape()[3];
+      }
+      a_idx += a.strides()[1] - a.strides()[2] * a.shape()[2];
+      b_idx += b.strides()[1] - b.strides()[2] * b.shape()[2];
+      c_idx += c.strides()[1] - c.strides()[2] * c.shape()[2];
+    }
+    a_idx += a.strides()[0] - a.strides()[1] * a.shape()[1];
+    b_idx += b.strides()[0] - b.strides()[1] * b.shape()[1];
+    c_idx += c.strides()[0] - c.strides()[1] * c.shape()[1];
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op_dispatch_dims(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  switch (out.ndim()) {
+    case 1:
+      ternary_op_dims1<T1, T2, T3, U, Op>(a, b, c, out, op);
+      return;
+    case 2:
+      ternary_op_dims2<T1, T2, T3, U, Op>(a, b, c, out, op);
+      return;
+    case 3:
+      ternary_op_dims3<T1, T2, T3, U, Op>(a, b, c, out, op);
+      return;
+    case 4:
+      ternary_op_dims4<T1, T2, T3, U, Op>(a, b, c, out, op);
+      return;
+  }
+
+  const T1* a_ptr = a.data<T1>();
+  const T2* b_ptr = b.data<T2>();
+  const T3* c_ptr = c.data<T3>();
+  U* dst = out.data<U>();
+  for (size_t i = 0; i < out.size(); i++) {
+    int a_idx = elem_to_loc(i, a.shape(), a.strides());
+    int b_idx = elem_to_loc(i, b.shape(), b.strides());
+    int c_idx = elem_to_loc(i, c.shape(), c.strides());
+    dst[i] = op(a_ptr[a_idx], b_ptr[b_idx], c_ptr[c_idx]);
+  }
+}
+
+template <typename T1, typename T2, typename T3, typename U, typename Op>
+void ternary_op(
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    Op op) {
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt);
+
+  // The full computation is scalar-scalar-scalar so we call the base op once.
+  if (topt == TernaryOpType::ScalarScalarScalar) {
+    *(out.data<U>()) = op(*a.data<T1>(), *b.data<T2>(), *c.data<T3>());
+    return;
+  }
+
+  ternary_op_dispatch_dims<T1, T2, T3, U>(a, b, c, out, op);
+}
+
+} // namespace
+
+} // namespace mlx::core

mlx/backend/common/unary.h

@@ -11,59 +11,6 @@ namespace mlx::core {
 
 namespace {
 
-struct AbsOp {
-  template <typename T>
-  T operator()(T x) {
-    return std::abs(x);
-  }
-  uint8_t operator()(uint8_t x) {
-    return x;
-  }
-  uint16_t operator()(uint16_t x) {
-    return x;
-  }
-  uint32_t operator()(uint32_t x) {
-    return x;
-  }
-  uint64_t operator()(uint64_t x) {
-    return x;
-  }
-  bool operator()(bool x) {
-    return x;
-  }
-};
-
-struct SignOp {
-  template <typename T>
-  T operator()(T x) {
-    return (x > T(0)) - (x < T(0));
-  }
-
-  uint8_t operator()(uint8_t x) {
-    return x != 0;
-  }
-  uint16_t operator()(uint16_t x) {
-    return x != 0;
-  }
-  uint32_t operator()(uint32_t x) {
-    return x != 0;
-  }
-  uint64_t operator()(uint64_t x) {
-    return x != 0;
-  }
-};
-
-struct RoundOp {
-  template <typename T>
-  T operator()(T x) {
-    return std::rint(x);
-  }
-
-  complex64_t operator()(complex64_t x) {
-    return {std::rint(x.real()), std::rint(x.imag())};
-  }
-};
-
 void set_unary_output_data(const array& in, array& out) {
   if (in.is_donatable() && in.itemsize() == out.itemsize()) {
     out.copy_shared_buffer(in);

mlx/backend/metal/CMakeLists.txt

@@ -4,7 +4,7 @@ add_custom_command(
               ${CMAKE_CURRENT_SOURCE_DIR}/make_compiled_preamble.sh
               ${CMAKE_CURRENT_BINARY_DIR}/compiled_preamble.cpp
               ${CMAKE_C_COMPILER}
-              ${CMAKE_SOURCE_DIR}
+              ${PROJECT_SOURCE_DIR}
     DEPENDS make_compiled_preamble.sh
             kernels/compiled_preamble.h
             kernels/unary.h
@@ -29,6 +29,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/metal.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp

mlx/backend/metal/allocator.cpp

@@ -1,5 +1,4 @@
-// Copyright © 2023 Apple Inc.
-
+// Copyright © 2023-2024 Apple Inc.
 #include "mlx/backend/metal/allocator.h"
 #include "mlx/backend/metal/metal.h"
 
@@ -23,16 +22,6 @@ void* Buffer::raw_ptr() {
 
 namespace metal {
 
-static bool cache_enabled_ = true;
-
-bool cache_enabled() {
-  return cache_enabled_;
-}
-
-void set_cache_enabled(bool enabled) {
-  cache_enabled_ = enabled;
-}
-
 namespace {
 
 BufferCache::BufferCache(MTL::Device* device)
@@ -44,7 +33,6 @@ BufferCache::~BufferCache() {
 }
 
 void BufferCache::clear() {
-  std::lock_guard<std::mutex> lk(cache_mutex_);
   for (auto& [size, holder] : buffer_pool_) {
     if (holder->buf)
       holder->buf->release();
@@ -57,12 +45,9 @@ void BufferCache::clear() {
 }
 
 MTL::Buffer* BufferCache::reuse_from_cache(size_t size) {
-  std::lock_guard<std::mutex> lk(cache_mutex_);
-
   // Find the closest buffer in pool
   MTL::Buffer* pbuf = nullptr;
 
-  // Make sure we use most of the available memory
   auto it = buffer_pool_.lower_bound(size);
 
   // Make sure we use most of the available memory
@@ -85,8 +70,6 @@ MTL::Buffer* BufferCache::reuse_from_cache(size_t size) {
 }
 
 void BufferCache::recycle_to_cache(MTL::Buffer* buf) {
-  std::lock_guard<std::mutex> lk(cache_mutex_);
-
   // Add to cache
   if (buf) {
     BufferHolder* bh = new BufferHolder(buf);
@@ -100,7 +83,6 @@ void BufferCache::release_cached_buffers(size_t min_bytes_to_free) {
   if (min_bytes_to_free >= 0.9 * pool_size_) {
     clear();
   } else {
-    std::lock_guard<std::mutex> lk(cache_mutex_);
     size_t total_bytes_freed = 0;
 
     while (tail_ && (total_bytes_freed < min_bytes_to_free)) {
@@ -158,9 +140,23 @@ void BufferCache::remove_from_list(BufferCache::BufferHolder* to_remove) {
 MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
       buffer_cache_(device_),
-      peak_allocated_size_(0),
       block_limit_(1.5 * device_->recommendedMaxWorkingSetSize()),
-      gc_limit_(0.95 * device_->recommendedMaxWorkingSetSize()) {}
+      gc_limit_(0.95 * device_->recommendedMaxWorkingSetSize()),
+      max_pool_size_(block_limit_) {}
+
+size_t MetalAllocator::set_cache_limit(size_t limit) {
+  std::swap(limit, max_pool_size_);
+  return limit;
+};
+
+size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) {
+  std::swap(limit, block_limit_);
+  relaxed_ = relaxed;
+  gc_limit_ = std::min(
+      block_limit_,
+      static_cast<size_t>(0.95 * device_->recommendedMaxWorkingSetSize()));
+  return limit;
+};
 
 Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   // Metal doesn't like empty buffers
@@ -174,41 +170,53 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   }
 
   // Try the cache
+  std::unique_lock lk(mutex_);
   MTL::Buffer* buf = buffer_cache_.reuse_from_cache(size);
-
   if (!buf) {
+    size_t mem_required = get_active_memory() + get_cache_memory() + size;
+
     // If there is too much memory pressure, fail (likely causes a wait).
-    if (!allow_swap && device_->currentAllocatedSize() + size >= block_limit_) {
+    if (!(allow_swap && relaxed_) && mem_required >= block_limit_) {
       return Buffer{nullptr};
     }
 
     auto thread_pool = metal::new_scoped_memory_pool();
 
-    // If we have a lot of memory pressure, check if we can reclaim some memory
-    // from the cache
-    if (device_->currentAllocatedSize() + size >= gc_limit_) {
-      size_t min_bytes_to_free =
-          size + device_->currentAllocatedSize() - gc_limit_;
-      buffer_cache_.release_cached_buffers(min_bytes_to_free);
+    // If we have a lot of memory pressure or are over the maximum cache size,
+    // try to reclaim memory from the cache
+    if (mem_required >= gc_limit_) {
+      buffer_cache_.release_cached_buffers(mem_required - gc_limit_);
     }
 
     // Allocate new buffer if needed
     size_t res_opt = MTL::ResourceStorageModeShared;
     res_opt |= MTL::ResourceHazardTrackingModeTracked;
+    lk.unlock();
     buf = device_->newBuffer(size, res_opt);
+    lk.lock();
   }
 
-  peak_allocated_size_ =
-      std::max(peak_allocated_size_, device_->currentAllocatedSize());
+  active_memory_ += buf->length();
+  peak_memory_ = std::max(peak_memory_, active_memory_);
+
+  // Maintain the cache below the requested limit
+  if (get_cache_memory() >= max_pool_size_) {
+    auto thread_pool = metal::new_scoped_memory_pool();
+    buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
+  }
 
   return Buffer{static_cast<void*>(buf)};
 }
 
 void MetalAllocator::free(Buffer buffer) {
   auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
-  if (cache_enabled()) {
+  std::unique_lock lk(mutex_);
+  active_memory_ -= buf->length();
+  if (get_cache_memory() < max_pool_size_) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
+    lk.unlock();
+    auto thread_pool = metal::new_scoped_memory_pool();
     buf->release();
   }
 }
@@ -218,6 +226,22 @@ MetalAllocator& allocator() {
   return allocator_;
 }
 
+size_t set_cache_limit(size_t limit) {
+  return allocator().set_cache_limit(limit);
+}
+size_t set_memory_limit(size_t limit, bool relaxed /* = true */) {
+  return allocator().set_memory_limit(limit, relaxed);
+}
+size_t get_active_memory() {
+  return allocator().get_active_memory();
+}
+size_t get_peak_memory() {
+  return allocator().get_peak_memory();
+}
+size_t get_cache_memory() {
+  return allocator().get_cache_memory();
+}
+
 } // namespace metal
 
 } // namespace mlx::core

mlx/backend/metal/allocator.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -19,11 +19,13 @@ class BufferCache {
  public:
   BufferCache(MTL::Device* device);
   ~BufferCache();
-  void clear();
 
   MTL::Buffer* reuse_from_cache(size_t size);
   void recycle_to_cache(MTL::Buffer* buf);
   void release_cached_buffers(size_t min_bytes_to_free);
+  size_t cache_size() {
+    return pool_size_;
+  }
 
  private:
   struct BufferHolder {
@@ -35,11 +37,11 @@ class BufferCache {
     MTL::Buffer* buf;
   };
 
+  void clear();
   void add_at_head(BufferHolder* to_add);
   void remove_from_list(BufferHolder* to_remove);
 
   MTL::Device* device_;
-  std::mutex cache_mutex_;
 
   std::multimap<size_t, BufferHolder*> buffer_pool_;
   BufferHolder* head_;
@@ -54,6 +56,17 @@ class MetalAllocator : public allocator::Allocator {
  public:
   virtual Buffer malloc(size_t size, bool allow_swap = false) override;
   virtual void free(Buffer buffer) override;
+  size_t get_active_memory() {
+    return active_memory_;
+  };
+  size_t get_peak_memory() {
+    return peak_memory_;
+  };
+  size_t get_cache_memory() {
+    return buffer_cache_.cache_size();
+  };
+  size_t set_cache_limit(size_t limit);
+  size_t set_memory_limit(size_t limit, bool relaxed);
 
  private:
   MTL::Device* device_;
@@ -64,9 +77,14 @@ class MetalAllocator : public allocator::Allocator {
   BufferCache buffer_cache_;
 
   // Allocation stats
-  size_t peak_allocated_size_;
   size_t block_limit_;
   size_t gc_limit_;
+  size_t active_memory_{0};
+  size_t peak_memory_{0};
+  size_t max_pool_size_;
+  bool relaxed_{true};
+
+  std::mutex mutex_;
 };
 
 MetalAllocator& allocator();

mlx/backend/metal/compiled.cpp

@@ -2,6 +2,7 @@
 
 #include <sstream>
 
+#include "mlx/backend/common/compiled.h"
 #include "mlx/backend/metal/compiled_preamble.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/utils.h"
@@ -11,125 +12,6 @@
 
 namespace mlx::core {
 
-inline bool is_static_cast(const Primitive& p) {
-  return (
-      typeid(p) == typeid(Broadcast) || typeid(p) == typeid(Copy) ||
-      typeid(p) == typeid(StopGradient) || typeid(p) == typeid(AsType));
-}
-
-inline auto get_type_string(Dtype d) {
-  switch (d) {
-    case float32:
-      return "float";
-    case float16:
-      return "half";
-    case bfloat16:
-      return "bfloat16_t";
-    case bool_:
-      return "bool";
-    case int8:
-      return "int8_t";
-    case int16:
-      return "int16_t";
-    case int32:
-      return "int32_t";
-    case int64:
-      return "int64_t";
-    case uint8:
-      return "uint8_t";
-    case uint16:
-      return "uint16_t";
-    case uint32:
-      return "uint32_t";
-    case uint64:
-      return "uint64_t";
-    default: {
-      std::ostringstream msg;
-      msg << "Unsupported compilation type " << d;
-      throw std::runtime_error(msg.str());
-    }
-  }
-}
-
-template <typename T>
-void print_float_constant(std::ostream& os, const array& x) {
-  auto old_precision = os.precision();
-  os << std::setprecision(std::numeric_limits<float>::digits10 + 1)
-     << x.item<T>() << std::setprecision(old_precision);
-}
-
-template <typename T>
-void print_int_constant(std::ostream& os, const array& x) {
-  os << x.item<T>();
-}
-
-void print_constant(std::ostream& os, const array& x) {
-  switch (x.dtype()) {
-    case float32:
-      return print_float_constant<float>(os, x);
-    case float16:
-      return print_float_constant<float16_t>(os, x);
-    case bfloat16:
-      return print_float_constant<bfloat16_t>(os, x);
-    case int8:
-      return print_int_constant<int8_t>(os, x);
-    case int16:
-      return print_int_constant<int16_t>(os, x);
-    case int32:
-      return print_int_constant<int32_t>(os, x);
-    case int64:
-      return print_int_constant<int64_t>(os, x);
-    case uint8:
-      return print_int_constant<uint8_t>(os, x);
-    case uint16:
-      return print_int_constant<uint16_t>(os, x);
-    case uint32:
-      return print_int_constant<uint32_t>(os, x);
-    case uint64:
-      return print_int_constant<uint64_t>(os, x);
-    case bool_:
-      os << std::boolalpha << x.item<bool>();
-      return;
-    default:
-      throw std::runtime_error("Unsupported constant type");
-  }
-}
-
-inline std::string build_lib_name(
-    const std::vector<array>& inputs,
-    const std::vector<array>& outputs,
-    const std::vector<array>& tape,
-    const std::unordered_set<uintptr_t>& constant_ids) {
-  std::ostringstream os;
-  std::ostringstream constant_hasher;
-
-  // The primitives describing the tape. For unary and binary primitives this
-  // must be enough to describe the full computation.
-  for (auto& a : tape) {
-    a.primitive().print(os);
-  }
-  os << "_";
-
-  for (auto& x : inputs) {
-    if (constant_ids.find(x.id()) != constant_ids.end()) {
-      os << "C";
-      print_constant(constant_hasher, x);
-    } else {
-      os << ((x.size() == 1) ? "S" : "V");
-    }
-  }
-  os << "_";
-  for (auto& x : inputs) {
-    if (constant_ids.find(x.id()) != constant_ids.end()) {
-      continue;
-    }
-    os << kindof(x.dtype()) << x.itemsize();
-  }
-  os << "_" << std::hash<std::string>{}(constant_hasher.str());
-
-  return os.str();
-}
-
 inline void build_kernel(
     std::ostream& os,
     const std::string& kernel_name,
@@ -149,9 +31,6 @@ inline void build_kernel(
     return constant_ids.find(x.id()) != constant_ids.end();
   };
 
-  // For scalar we shouldn't do the indexing things, just read at 0
-  auto is_scalar = [](const array& x) { return x.size() == 1; };
-
   NodeNamer namer;
   bool add_indices = false;
   int cnt = 0;
@@ -286,7 +165,7 @@ inline void build_kernel(
 
   if (cnt > 31) {
     std::ostringstream msg;
-    msg << "[compile] Too many inputs/outputs fused in the Metal Compile "
+    msg << "[compile] Too many inputs/outputs fused in the Metal Compiled "
         << "primitive which exhausted the available argument buffers for "
         << "the kernel. Please file an issue with the function that results "
         << "in this error. The name of the kernel is '" << kernel_name << "'";
@@ -344,13 +223,7 @@ void Compiled::eval_gpu(
         /* ndim = */ 0,
         /* dynamic_dims = */ true);
 
-    kernel_source_ = kernel.str();
-    lib = d.get_library(kernel_lib_, kernel_source_);
-  }
-
-  // Allocate space for the outputs
-  for (auto& out : outputs) {
-    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    lib = d.get_library(kernel_lib_, kernel.str());
   }
 
   // Figure out which kernel we are using
@@ -358,7 +231,7 @@ void Compiled::eval_gpu(
   bool contiguous = true;
   for (auto& x : inputs) {
     if ((!x.flags().row_contiguous || x.shape() != output_shape) &&
-        x.size() > 1) {
+        !is_scalar(x)) {
       contiguous = false;
       break;
     }
@@ -379,7 +252,7 @@ void Compiled::eval_gpu(
       auto& x = inputs[i];
 
       // Skip scalar inputs.
-      if (x.size() <= 1) {
+      if (is_scalar(x)) {
         continue;
       }
 
@@ -434,7 +307,7 @@ void Compiled::eval_gpu(
     }
     auto& x = inputs[i];
     set_array_buffer(compute_encoder, x, cnt++);
-    if (!contiguous && x.size() > 1) {
+    if (!contiguous && !is_scalar(x)) {
       compute_encoder->setBytes(
           strides[stride_idx].data(),
           strides[stride_idx].size() * sizeof(size_t),
@@ -443,6 +316,27 @@ void Compiled::eval_gpu(
     }
   }
 
+  // Allocate space for the outputs possibly with input donation
+  {
+    int o = 0;
+    for (int i = 0; i < inputs.size() && o < outputs.size(); ++i) {
+      auto& in = inputs[i];
+      // Conditions for donation
+      // - Row contiguous
+      // - Donatable
+      // - Correct size
+      // - Not a constant
+      if (in.flags().row_contiguous && in.nbytes() == outputs[o].nbytes() &&
+          in.is_donatable() &&
+          constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
+        outputs[o++].move_shared_buffer(in);
+      }
+    }
+    for (; o < outputs.size(); ++o) {
+      outputs[o].set_data(allocator::malloc_or_wait(outputs[o].nbytes()));
+    }
+  }
+
   // Put the outputs in
   for (auto& x : outputs) {
     set_array_buffer(compute_encoder, x, cnt++);

mlx/backend/metal/conv.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <algorithm>
 #include <cassert>
@@ -7,80 +7,72 @@
 
 #include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/device.h"
-#include "mlx/backend/metal/kernels/conv_params.h"
 #include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
 #include "mlx/backend/metal/matmul.h"
 #include "mlx/backend/metal/utils.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
+using namespace mlx::steel;
+
 namespace mlx::core {
 
 namespace {
 
-void explicit_gemm_conv_1D_gpu(
+template <int N>
+void explicit_gemm_conv_ND_gpu(
     const Stream& s,
     metal::Device& d,
     const array& in,
     const array& wt,
     array out,
-    const MLXConvParams<1>& conv_params) {
-  // Pad input
-  std::vector<int> padded_shape = {
-      conv_params.N, conv_params.iS[0] + 2 * conv_params.pad[0], conv_params.C};
-  array in_padded(padded_shape, in.dtype(), nullptr, {});
+    const MLXConvParams<N>& conv_params) {
+  // Prepare unfolding array
+  std::vector<int> unfolded_shape = {
+      static_cast<int>(out.size() / conv_params.O),
+      static_cast<int>(wt.size() / conv_params.O)};
+  array in_unfolded(unfolded_shape, in.dtype(), nullptr, {});
 
-  // Fill with zeros
-  copy_gpu(array(0, in.dtype()), in_padded, CopyType::Scalar, s);
+  in_unfolded.set_data(allocator::malloc_or_wait(in_unfolded.nbytes()));
 
-  // Pick input slice from padded
-  size_t data_offset = conv_params.pad[0] * in_padded.strides()[1];
-  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
-  in_padded_slice.copy_shared_buffer(
-      in_padded,
-      in_padded.strides(),
-      in_padded.flags(),
-      in_padded_slice.size(),
-      data_offset);
+  // Prepare unfolding kernel
+  std::ostringstream kname;
+  kname << "naive_unfold_nd_" << type_to_name(in_unfolded) << "_" << N;
+  auto compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname.str());
+  compute_encoder->setComputePipelineState(kernel);
 
-  // Copy input values into the slice
-  copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s);
+  set_array_buffer(compute_encoder, in, 0);
+  set_array_buffer(compute_encoder, in_unfolded, 1);
 
-  // Make strided view
-  std::vector<int> strided_shape = {
-      conv_params.N, conv_params.oS[0], conv_params.wS[0], conv_params.C};
+  compute_encoder->setBytes(&conv_params, sizeof(conv_params), 2);
 
-  std::vector<size_t> strided_strides = {
-      in_padded.strides()[0],
-      in_padded.strides()[1] * conv_params.str[0],
-      in_padded.strides()[1],
-      in_padded.strides()[2]};
-  auto flags = in_padded.flags();
+  // Launch unfolding kernel
+  int tgp_x = std::min(conv_params.C, 64);
+  tgp_x = 32 * ((tgp_x + 32 - 1) / 32);
+  int tgp_y = 256 / tgp_x;
 
-  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
-  in_strided_view.copy_shared_buffer(
-      in_padded, strided_strides, flags, in_strided_view.size(), 0);
+  MTL::Size group_dims = MTL::Size(tgp_x, tgp_y, 1);
+  MTL::Size grid_dims = MTL::Size(
+      conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]);
 
-  // Materialize strided view
-  std::vector<int> strided_reshape = {
-      conv_params.N * conv_params.oS[0], conv_params.wS[0] * conv_params.C};
-  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
-  copy_gpu(in_strided_view, in_strided, CopyType::General, s);
+  compute_encoder->dispatchThreads(grid_dims, group_dims);
 
   // Perform gemm
-  std::vector<array> copies = {in_padded, in_strided};
+  std::vector<array> copies;
   return steel_matmul(
       s,
       d,
-      /*a = */ in_strided,
+      /*a = */ in_unfolded,
       /*b = */ wt,
       /*c = */ out,
-      /*M = */ strided_reshape[0],
+      /*M = */ unfolded_shape[0],
       /*N = */ conv_params.O,
-      /*K = */ strided_reshape[1],
+      /*K = */ unfolded_shape[1],
       /*batch_size_out = */ 1,
-      /*a_cols = */ strided_reshape[1],
-      /*b_cols = */ strided_reshape[1],
+      /*a_cols = */ unfolded_shape[1],
+      /*b_cols = */ unfolded_shape[1],
       /*a_transposed = */ false,
       /*b_transposed = */ true,
       /*copies = */ copies);
@@ -94,7 +86,9 @@ void conv_1D_gpu(
     array out,
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
-    const std::vector<int>& wt_dilation) {
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip) {
   // Make conv params
   MLXConvParams<1> conv_params{
       /* const int  N = */ in.shape(0),
@@ -105,24 +99,19 @@ void conv_1D_gpu(
       /* const int oS[NDIM] = */ {out.shape(1)},
       /* const int str[NDIM] = */ {wt_strides[0]},
       /* const int pad[NDIM] = */ {padding[0]},
-      /* const int dil[NDIM] = */ {wt_dilation[0]},
+      /* const int kdil[NDIM] = */ {wt_dilation[0]},
+      /* const int idil[NDIM] = */ {in_dilation[0]},
       /* const size_t in_strides[NDIM + 2] = */
       {in.strides()[0], in.strides()[1], in.strides()[2]},
       /* const size_t wt_strides[NDIM + 2] = */
       {wt.strides()[0], wt.strides()[1], wt.strides()[2]},
       /* const size_t out_strides[NDIM + 2] = */
       {out.strides()[0], out.strides()[1], out.strides()[2]},
-  };
+      /* const int groups = */ 1,
+      /* const bool flip = */ flip};
 
   // Direct to explicit gemm conv
-  if (wt_dilation[0] == 1) {
-    explicit_gemm_conv_1D_gpu(s, d, in, wt, out, conv_params);
-  }
-
-  // Direct to fallback conv
-  else {
-    throw std::invalid_argument("[conv_1D_gpu] Dilation needs to be 1.");
-  }
+  return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params);
 }
 
 void slow_conv_2D_gpu(
@@ -168,114 +157,262 @@ void implicit_gemm_conv_2D_gpu(
     const array& wt,
     array out,
     const MLXConvParams<2>& conv_params) {
-  int bm = 32, bn = 32, bk = 16;
+  // Deduce implicit gemm size
+  int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1];
+  int implicit_N = conv_params.O;
+  int implicit_K = conv_params.wS[0] * conv_params.wS[1] * conv_params.C;
+
+  // Determine block and warp tiles
   int wm = 2, wn = 2;
 
+  int bm = implicit_M >= 8192 && conv_params.C >= 64 ? 64 : 32;
+  int bn = (bm == 64 || implicit_N >= 64) ? 64 : 32;
+  int bk = 16;
+
+  if (implicit_N <= 16) {
+    bn = 8;
+    wm = 4;
+    wn = 1;
+  }
+
+  int tn = (implicit_N + bn - 1) / bn;
+  int tm = (implicit_M + bm - 1) / bm;
+  int swizzle_log = 0;
+
+  // Fix small channel specialization
+  int n_channel_specialization = 0;
+  int channel_k_iters = ((conv_params.C + bk - 1) / bk);
+  int gemm_k_iters = conv_params.wS[0] * conv_params.wS[1] * channel_k_iters;
+
+  if (conv_params.C <= 2) {
+    gemm_k_iters = (implicit_K + bk - 1) / bk;
+    n_channel_specialization = conv_params.C;
+  } else if (conv_params.C <= 4) {
+    gemm_k_iters = ((conv_params.wS[0] * conv_params.wS[1] * 4) + bk - 1) / bk;
+    n_channel_specialization = conv_params.C;
+  }
+
+  bool small_filter = (!n_channel_specialization) &&
+      (conv_params.wS[0] <= 16 && conv_params.wS[1] <= 16);
+
+  // Fix host side helper params
+  int sign = (conv_params.flip ? -1 : 1);
+  int ijw = conv_params.in_strides[2] * conv_params.kdil[1];
+  int ijh = conv_params.in_strides[1] * conv_params.kdil[0];
+
+  int inp_jump_w = sign * ijw;
+  int inp_jump_h = sign * (ijh - (conv_params.wS[1] - 1) * ijw);
+  int inp_jump_c = bk - sign * (conv_params.wS[0] - 1) * ijh -
+      sign * (conv_params.wS[1] - 1) * ijw;
+
+  // Build implicit gemm params
+  ImplicitGemmConv2DParams gemm_params{
+      /* const int M = */ implicit_M,
+      /* const int N = */ implicit_N,
+      /* const int K = */ implicit_K,
+
+      /* const int gemm_k_iterations = */ gemm_k_iters,
+
+      /* const int inp_jump_w = */ inp_jump_w,
+      /* const int inp_jump_h = */ inp_jump_h,
+      /* const int inp_jump_c = */ inp_jump_c,
+
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int swizzle_log = */ swizzle_log};
+
+  // Determine kernel
   std::ostringstream kname;
   kname << "implicit_gemm_conv_2d_" << type_to_name(out) << "_bm" << bm << "_bn"
-        << bn << "_bk" << bk << "_wm" << wm << "_wn" << wn;
+        << bn << "_bk" << bk << "_wm" << wm << "_wn" << wn << "_channel_"
+        << (n_channel_specialization ? std::to_string(n_channel_specialization)
+                                     : "l")
+        << "_filter_" << (small_filter ? 's' : 'l');
 
   // Encode and dispatch kernel
   auto compute_encoder = d.get_command_encoder(s.index);
   auto kernel = d.get_kernel(kname.str());
   compute_encoder->setComputePipelineState(kernel);
 
-  int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1];
-  int implicit_N = conv_params.O;
-  int implicit_K = conv_params.wS[0] * conv_params.wS[1] * conv_params.C;
-
-  size_t grid_dim_x = (implicit_N + bn - 1) / bn;
-  size_t grid_dim_y = (implicit_M + bm - 1) / bm;
+  // Deduce grid launch dimensions
+  int tile = 1 << swizzle_log;
+  size_t grid_dim_y = (tm + tile - 1) / tile;
+  size_t grid_dim_x = tn * tile;
 
   MTL::Size group_dims = MTL::Size(32, wn, wm);
   MTL::Size grid_dims = MTL::Size(grid_dim_x, grid_dim_y, 1);
 
+  // Encode arrays
   set_array_buffer(compute_encoder, in, 0);
   set_array_buffer(compute_encoder, wt, 1);
   set_array_buffer(compute_encoder, out, 2);
 
+  // Encode params
   compute_encoder->setBytes(&conv_params, sizeof(MLXConvParams<2>), 3);
+  compute_encoder->setBytes(&gemm_params, sizeof(ImplicitGemmConv2DParams), 4);
+
+  // Launch kernel
   compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
 }
 
-void explicit_gemm_conv_2D_gpu(
+void implicit_gemm_conv_2D_general_gpu(
     const Stream& s,
     metal::Device& d,
     const array& in,
     const array& wt,
     array out,
     const MLXConvParams<2>& conv_params) {
-  // Pad input
-  std::vector<int> padded_shape = {
-      conv_params.N,
-      conv_params.iS[0] + 2 * conv_params.pad[0],
-      conv_params.iS[1] + 2 * conv_params.pad[1],
-      conv_params.C};
-  array in_padded(padded_shape, in.dtype(), nullptr, {});
+  // Deduce implicit gemm size
+  int implicit_M = conv_params.N * conv_params.oS[0] * conv_params.oS[1];
+  int implicit_N = conv_params.O;
+  int implicit_K = conv_params.wS[0] * conv_params.wS[1] * conv_params.C;
 
-  // Fill with zeros
-  copy_gpu(array(0, in.dtype()), in_padded, CopyType::Scalar, s);
+  // Determine block and warp tiles
+  int wm = 2, wn = 2;
 
-  // Pick input slice from padded
-  size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] +
-      conv_params.pad[1] * in_padded.strides()[2];
-  array in_padded_slice(in.shape(), in_padded.dtype(), nullptr, {});
-  in_padded_slice.copy_shared_buffer(
-      in_padded,
-      in_padded.strides(),
-      in_padded.flags(),
-      in_padded_slice.size(),
-      data_offset);
+  // Make jump params
+  int f_wgt_jump_h =
+      std::lcm(conv_params.idil[0], conv_params.kdil[0]) / conv_params.kdil[0];
+  int f_wgt_jump_w =
+      std::lcm(conv_params.idil[1], conv_params.kdil[1]) / conv_params.kdil[1];
 
-  // Copy input values into the slice
-  copy_gpu_inplace(in, in_padded_slice, CopyType::GeneralGeneral, s);
+  int f_out_jump_h =
+      std::lcm(conv_params.idil[0], conv_params.str[0]) / conv_params.str[0];
+  int f_out_jump_w =
+      std::lcm(conv_params.idil[1], conv_params.str[1]) / conv_params.str[1];
 
-  // Make strided view
-  std::vector<int> strided_shape = {
-      conv_params.N,
-      conv_params.oS[0],
-      conv_params.oS[1],
-      conv_params.wS[0],
-      conv_params.wS[1],
-      conv_params.C};
+  int adj_out_h = (conv_params.oS[0] + f_out_jump_h - 1) / f_out_jump_h;
+  int adj_out_w = (conv_params.oS[1] + f_out_jump_w - 1) / f_out_jump_w;
+  int adj_out_hw = adj_out_h * adj_out_w;
+  int adj_implicit_m = conv_params.N * adj_out_hw;
 
-  std::vector<size_t> strided_strides = {
-      in_padded.strides()[0],
-      in_padded.strides()[1] * conv_params.str[0],
-      in_padded.strides()[2] * conv_params.str[1],
-      in_padded.strides()[1],
-      in_padded.strides()[2],
-      in_padded.strides()[3]};
-  auto flags = in_padded.flags();
+  Conv2DGeneralJumpParams jump_params{
+      /* const int f_wgt_jump_h = */ f_wgt_jump_h,
+      /* const int f_wgt_jump_w = */ f_wgt_jump_w,
 
-  array in_strided_view(strided_shape, in_padded.dtype(), nullptr, {});
-  in_strided_view.copy_shared_buffer(
-      in_padded, strided_strides, flags, in_strided_view.size(), 0);
+      /* const int f_out_jump_h = */ f_out_jump_h,
+      /* const int f_out_jump_w = */ f_out_jump_w,
 
-  // Materialize strided view
-  std::vector<int> strided_reshape = {
-      conv_params.N * conv_params.oS[0] * conv_params.oS[1],
-      conv_params.wS[0] * conv_params.wS[1] * conv_params.C};
-  array in_strided(strided_reshape, in_strided_view.dtype(), nullptr, {});
-  copy_gpu(in_strided_view, in_strided, CopyType::General, s);
+      /* const int adj_out_h = */ adj_out_h,
+      /* const int adj_out_w = */ adj_out_w,
+      /* const int adj_out_hw = */ adj_out_hw,
+      /* const int adj_implicit_m = */ adj_implicit_m};
 
-  // Perform gemm
-  std::vector<array> copies = {in_padded, in_strided};
-  return steel_matmul(
-      s,
-      d,
-      /*a = */ in_strided,
-      /*b = */ wt,
-      /*c = */ out,
-      /*M = */ strided_reshape[0],
-      /*N = */ conv_params.O,
-      /*K = */ strided_reshape[1],
-      /*batch_size_out = */ 1,
-      /*a_cols = */ strided_reshape[1],
-      /*b_cols = */ strided_reshape[1],
-      /*a_transposed = */ false,
-      /*b_transposed = */ true,
-      /*copies = */ copies);
+  // Make base info
+  std::vector<Conv2DGeneralBaseInfo> base_h(f_out_jump_h);
+  std::vector<Conv2DGeneralBaseInfo> base_w(f_out_jump_w);
+
+  int jump_h = conv_params.flip ? -conv_params.kdil[0] : conv_params.kdil[0];
+  int jump_w = conv_params.flip ? -conv_params.kdil[1] : conv_params.kdil[1];
+
+  int init_h =
+      (conv_params.flip ? (conv_params.wS[0] - 1) * conv_params.kdil[0] : 0);
+  int init_w =
+      (conv_params.flip ? (conv_params.wS[1] - 1) * conv_params.kdil[1] : 0);
+
+  for (int i = 0; i < f_out_jump_h; ++i) {
+    int ih_loop = i * conv_params.str[0] - conv_params.pad[0] + init_h;
+
+    int wh_base = 0;
+    while (wh_base < conv_params.wS[0] && ih_loop % conv_params.idil[0] != 0) {
+      wh_base++;
+      ih_loop += jump_h;
+    }
+
+    int wh_size =
+        ((conv_params.wS[0] - wh_base) + f_wgt_jump_h - 1) / f_wgt_jump_h;
+    base_h[i] = {wh_base, wh_size};
+  }
+
+  for (int j = 0; j < f_out_jump_w; ++j) {
+    int iw_loop = j * conv_params.str[1] - conv_params.pad[1] + init_w;
+
+    int ww_base = 0;
+    while (ww_base < conv_params.wS[1] && iw_loop % conv_params.idil[1] != 0) {
+      ww_base++;
+      iw_loop += jump_w;
+    }
+
+    int ww_size =
+        ((conv_params.wS[1] - ww_base) + f_wgt_jump_w - 1) / f_wgt_jump_w;
+    base_w[j] = {ww_base, ww_size};
+  }
+
+  // Collect block sizes
+  int bm = adj_implicit_m >= 8192 && conv_params.C >= 64 ? 64 : 32;
+  int bn = (bm == 64 && implicit_N >= 64) ? 64 : 32;
+  int bk = 16;
+
+  int tn = (implicit_N + bn - 1) / bn;
+  int tm = (adj_implicit_m + bm - 1) / bm;
+  int swizzle_log = 0;
+
+  // Get channel iteration info
+  int channel_k_iters = ((conv_params.C + bk - 1) / bk);
+  int gemm_k_iters = channel_k_iters;
+
+  // Fix host side helper params
+  int sign = (conv_params.flip ? -1 : 1);
+  int ijw = conv_params.in_strides[2] * conv_params.kdil[1];
+  int ijh = conv_params.in_strides[1] * conv_params.kdil[0];
+
+  int inp_jump_w = sign * ijw;
+  int inp_jump_h = sign * (ijh - (conv_params.wS[1] - 1) * ijw);
+  int inp_jump_c = bk - sign * (conv_params.wS[0] - 1) * ijh -
+      sign * (conv_params.wS[1] - 1) * ijw;
+
+  // Build implicit gemm params
+  ImplicitGemmConv2DParams gemm_params{
+      /* const int M = */ implicit_M,
+      /* const int N = */ implicit_N,
+      /* const int K = */ implicit_K,
+
+      /* const int gemm_k_iterations = */ gemm_k_iters,
+
+      /* const int inp_jump_w = */ inp_jump_w,
+      /* const int inp_jump_h = */ inp_jump_h,
+      /* const int inp_jump_c = */ inp_jump_c,
+
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int swizzle_log = */ swizzle_log};
+
+  // Determine kernel
+  std::ostringstream kname;
+  kname << "implicit_gemm_conv_2d_general_" << type_to_name(out) << "_bm" << bm
+        << "_bn" << bn << "_bk" << bk << "_wm" << wm << "_wn" << wn;
+
+  // Encode and dispatch kernel
+  auto compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname.str());
+  compute_encoder->setComputePipelineState(kernel);
+
+  // Deduce grid launch dimensions
+  int tile = 1 << swizzle_log;
+  size_t grid_dim_y = (tm + tile - 1) / tile;
+  size_t grid_dim_x = tn * tile;
+  size_t grid_dim_z = f_out_jump_h * f_out_jump_w;
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  MTL::Size grid_dims = MTL::Size(grid_dim_x, grid_dim_y, grid_dim_z);
+
+  // Encode arrays
+  set_array_buffer(compute_encoder, in, 0);
+  set_array_buffer(compute_encoder, wt, 1);
+  set_array_buffer(compute_encoder, out, 2);
+
+  // Encode params
+  compute_encoder->setBytes(&conv_params, sizeof(MLXConvParams<2>), 3);
+  compute_encoder->setBytes(&gemm_params, sizeof(ImplicitGemmConv2DParams), 4);
+  compute_encoder->setBytes(&jump_params, sizeof(Conv2DGeneralJumpParams), 5);
+
+  compute_encoder->setBytes(
+      base_h.data(), sizeof(Conv2DGeneralBaseInfo) * base_h.size(), 6);
+  compute_encoder->setBytes(
+      base_w.data(), sizeof(Conv2DGeneralBaseInfo) * base_w.size(), 7);
+
+  // Launch kernel
+  compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
 }
 
 void winograd_conv_2D_gpu(
@@ -300,6 +437,7 @@ void winograd_conv_2D_gpu(
   // Fill with zeros
   array zero_arr = array(0, in.dtype());
   copy_gpu(zero_arr, in_padded, CopyType::Scalar, s);
+  copies_w.push_back(zero_arr);
 
   // Pick input slice from padded
   size_t data_offset = conv_params.pad[0] * in_padded.strides()[1] +
@@ -328,7 +466,8 @@ void winograd_conv_2D_gpu(
       /* const int oS[NDIM] = */ {out.shape(1), out.shape(2)},
       /* const int str[NDIM] = */ {1, 1},
       /* const int pad[NDIM] = */ {0, 0},
-      /* const int dil[NDIM] = */ {1, 1},
+      /* const int kdil[NDIM] = */ {1, 1},
+      /* const int idil[NDIM] = */ {1, 1},
       /* const size_t in_strides[NDIM + 2] = */
       {in_padded.strides()[0],
        in_padded.strides()[1],
@@ -338,6 +477,8 @@ void winograd_conv_2D_gpu(
       {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]},
       /* const size_t out_strides[NDIM + 2] = */
       {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]},
+      /* const int groups = */ 1,
+      /* const bool flip = */ false,
   };
 
   int O_c = conv_params.O;
@@ -461,6 +602,8 @@ void conv_2D_gpu(
     const std::vector<int>& padding,
     const std::vector<int>& wt_strides,
     const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    bool flip,
     std::vector<array>& copies) {
   // Make conv params
   MLXConvParams<2> conv_params{
@@ -472,37 +615,47 @@ void conv_2D_gpu(
       /* const int oS[NDIM] = */ {out.shape(1), out.shape(2)},
       /* const int str[NDIM] = */ {wt_strides[0], wt_strides[1]},
       /* const int pad[NDIM] = */ {padding[0], padding[1]},
-      /* const int dil[NDIM] = */ {wt_dilation[0], wt_dilation[1]},
+      /* const int kdil[NDIM] = */ {wt_dilation[0], wt_dilation[1]},
+      /* const int idil[NDIM] = */ {in_dilation[0], in_dilation[1]},
       /* const size_t in_strides[NDIM + 2] = */
       {in.strides()[0], in.strides()[1], in.strides()[2], in.strides()[3]},
       /* const size_t wt_strides[NDIM + 2] = */
       {wt.strides()[0], wt.strides()[1], wt.strides()[2], wt.strides()[3]},
       /* const size_t out_strides[NDIM + 2] = */
       {out.strides()[0], out.strides()[1], out.strides()[2], out.strides()[3]},
+      /* const int groups = */ 1,
+      /* const bool flip = */ flip,
   };
 
+  bool is_stride_one = conv_params.str[0] == 1 && conv_params.str[1] == 1;
+  bool is_kdil_one = conv_params.kdil[0] == 1 && conv_params.kdil[1] == 1;
+  bool is_idil_one = conv_params.idil[0] == 1 && conv_params.idil[1] == 1;
+
+  bool inp_large = (conv_params.in_strides[0] >= 1ul << 18);
+  bool channels_large = (conv_params.C + conv_params.O) >= 512;
+  bool channels_med = (conv_params.C + conv_params.O) >= 256;
+
   // Direct to winograd conv
-  if (conv_params.C % 32 == 0 && conv_params.O % 32 == 0 &&
-      conv_params.C >= 64 && conv_params.O >= 64 && conv_params.wS[0] == 3 &&
-      conv_params.wS[1] == 3 && conv_params.str[0] == 1 &&
-      conv_params.str[1] == 1 && conv_params.dil[0] == 1 &&
-      conv_params.dil[1] == 1) {
-    winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
+  if (!flip && is_stride_one && is_kdil_one && is_idil_one &&
+      conv_params.wS[0] == 3 && conv_params.wS[1] == 3 &&
+      conv_params.C % 32 == 0 && conv_params.O % 32 == 0 &&
+      (channels_large || (channels_med && inp_large))) {
+    return winograd_conv_2D_gpu(s, d, in, wt, out, conv_params, copies);
   }
 
   // Direct to implicit gemm conv
-  else if (conv_params.C % 32 == 0 && conv_params.O % 32 == 0) {
-    implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);
+  if (is_idil_one && (conv_params.C <= 4 || conv_params.C % 16 == 0) &&
+      (conv_params.O <= 16 || conv_params.O % 16 == 0)) {
+    return implicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);
+  }
+
+  else if (conv_params.C % 16 == 0 && conv_params.O % 16 == 0) {
+    return implicit_gemm_conv_2D_general_gpu(s, d, in, wt, out, conv_params);
   }
 
   // Direct to explicit gemm conv
-  else if (wt_dilation[0] == 1 && wt_dilation[1] == 1) {
-    explicit_gemm_conv_2D_gpu(s, d, in, wt, out, conv_params);
-  }
-
-  // Direct to fallback conv
   else {
-    slow_conv_2D_gpu(s, d, in, wt, out, conv_params);
+    return explicit_gemm_conv_ND_gpu(s, d, in, wt, out, conv_params);
   }
 }
 
@@ -533,11 +686,31 @@ void Convolution::eval_gpu(const std::vector<array>& inputs, array& out) {
   // 2D conv
   if (out.ndim() == 4) {
     conv_2D_gpu(
-        s, d, in, wt, out, padding_, kernel_strides_, kernel_dilation_, copies);
+        s,
+        d,
+        in,
+        wt,
+        out,
+        padding_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_,
+        copies);
   }
   // 1D conv
   else if (out.ndim() == 3) {
-    conv_1D_gpu(s, d, in, wt, out, padding_, kernel_strides_, kernel_dilation_);
+    conv_1D_gpu(
+        s,
+        d,
+        in,
+        wt,
+        out,
+        padding_,
+        kernel_strides_,
+        kernel_dilation_,
+        input_dilation_,
+        flip_);
   }
   // Throw error
   else {

mlx/backend/metal/indexing.cpp

@@ -142,7 +142,28 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   // Get kernel name
   std::ostringstream kname;
   std::string idx_type_name = nidx ? type_to_name(inputs[1]) : "";
-  kname << "scatter" << type_to_name(out) << idx_type_name;
+
+  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  bool index_nd1_specialization = (idx_ndim == 1);
+
+  // Bail from fast path (1d index specialization) if scatter dims aren't
+  // the outermost dims and contiguous since update access won't be raster
+  // order.
+  for (auto i = 0; i < axes_.size() && index_nd1_specialization; i++) {
+    index_nd1_specialization &= (axes_[i] == i);
+  }
+
+  // Bail from fast path (1d index specialization) if any of the dims are
+  // broadcasted, since we can't rely on linear indexing in that case.
+  for (int i = 1; i < inputs.size() && index_nd1_specialization; i++) {
+    index_nd1_specialization &= inputs[i].flags().row_contiguous;
+  }
+
+  if (index_nd1_specialization) {
+    kname << "scatter_1d_index" << type_to_name(out) << idx_type_name;
+  } else {
+    kname << "scatter" << type_to_name(out) << idx_type_name;
+  }
   switch (reduce_type_) {
     case Scatter::None:
       kname << "_none";
@@ -167,92 +188,109 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
 
   auto& upd = inputs.back();
   size_t nthreads = upd.size();
-  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
-  if (thread_group_size > nthreads) {
-    thread_group_size = nthreads;
-  }
 
   compute_encoder->setComputePipelineState(kernel);
 
-  // Collect all idx shapes and strides into one place
-  int idx_ndim = nidx ? inputs[1].ndim() : 0;
-  std::vector<int> idx_shapes;
-  std::vector<size_t> idx_strides;
-
-  for (int i = 0; i < nidx; ++i) {
-    idx_shapes.insert(
-        idx_shapes.end(),
-        inputs[i + 1].shape().begin(),
-        inputs[i + 1].shape().end());
-
-    idx_strides.insert(
-        idx_strides.end(),
-        inputs[i + 1].strides().begin(),
-        inputs[i + 1].strides().end());
-  }
-
   // Set all the buffers
   set_array_buffer(compute_encoder, upd, 1);
   set_array_buffer(compute_encoder, out, 2);
 
   // Set update info
-  size_t upd_ndim = upd.ndim();
+  uint upd_ndim = upd.ndim();
   size_t upd_size = 1;
   for (int i = idx_ndim; i < upd.ndim(); ++i) {
     upd_size *= upd.shape(i);
   }
-  if (upd_ndim == 0) {
-    // Need placeholders so Metal doesn't compalain
-    int shape_ = 0;
-    size_t stride_ = 0;
-    compute_encoder->setBytes(&shape_, sizeof(int), 3);
-    compute_encoder->setBytes(&stride_, sizeof(size_t), 4);
-  } else {
-    compute_encoder->setBytes(upd.shape().data(), upd_ndim * sizeof(int), 3);
+
+  if (index_nd1_specialization) {
+    bool upd_col_contiguous = upd.flags().col_contiguous;
     compute_encoder->setBytes(
-        upd.strides().data(), upd_ndim * sizeof(size_t), 4);
-  }
-  compute_encoder->setBytes(&upd_ndim, sizeof(size_t), 5);
-  compute_encoder->setBytes(&upd_size, sizeof(size_t), 6);
-
-  // Set output info
-  size_t out_ndim = out.ndim();
-  if (out_ndim == 0) {
-    // Need placeholders so Metal doesn't compalain
-    int shape_ = 0;
-    size_t stride_ = 0;
-    compute_encoder->setBytes(&shape_, sizeof(int), 7);
-    compute_encoder->setBytes(&stride_, sizeof(size_t), 8);
-  } else {
-    compute_encoder->setBytes(out.shape().data(), out_ndim * sizeof(int), 7);
+        out.shape().data(), out.shape().size() * sizeof(int), 3);
     compute_encoder->setBytes(
-        out.strides().data(), out_ndim * sizeof(size_t), 8);
-  }
-  compute_encoder->setBytes(&out_ndim, sizeof(size_t), 9);
-  compute_encoder->setBytes(axes_.data(), axes_.size() * sizeof(int), 10);
+        out.strides().data(), out.strides().size() * sizeof(size_t), 4);
+    compute_encoder->setBytes(&upd_size, sizeof(size_t), 5);
+    compute_encoder->setBytes(&upd_col_contiguous, sizeof(bool), 6);
 
-  // Set index info
-  if (idx_ndim == 0) {
-    // Add a 0 in idx_shapes and strides to avoid the missing buffer binding
-    // error in the metal API.
-    idx_shapes.push_back(0);
-    idx_strides.push_back(0);
-  }
-  compute_encoder->setBytes(
-      idx_shapes.data(), idx_shapes.size() * sizeof(int), 11);
-  compute_encoder->setBytes(
-      idx_strides.data(), idx_strides.size() * sizeof(size_t), 12);
-  compute_encoder->setBytes(&idx_ndim, sizeof(int), 13);
+    // Set index buffers
+    for (int i = 1; i < nidx + 1; ++i) {
+      set_array_buffer(compute_encoder, inputs[i], 20 + i);
+    }
 
-  // Set index buffers
-  for (int i = 1; i < nidx + 1; ++i) {
-    set_array_buffer(compute_encoder, inputs[i], 20 + i);
-  }
+    // Launch grid
+    MTL::Size grid_dims = MTL::Size(upd_size, nthreads / upd_size, 1);
+    MTL::Size group_dims = get_block_dims(upd_size, nthreads / upd_size, 1);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
 
-  // Launch grid
-  MTL::Size grid_dims = MTL::Size(upd_size, nthreads / upd_size, 1);
-  MTL::Size group_dims = get_block_dims(upd_size, nthreads / upd_size, 1);
-  compute_encoder->dispatchThreads(grid_dims, group_dims);
+  } else {
+    // Collect all idx shapes and strides into one place
+    std::vector<int> idx_shapes;
+    std::vector<size_t> idx_strides;
+
+    for (int i = 0; i < nidx; ++i) {
+      idx_shapes.insert(
+          idx_shapes.end(),
+          inputs[i + 1].shape().begin(),
+          inputs[i + 1].shape().end());
+
+      idx_strides.insert(
+          idx_strides.end(),
+          inputs[i + 1].strides().begin(),
+          inputs[i + 1].strides().end());
+    }
+
+    if (upd_ndim == 0) {
+      // Need placeholders so Metal doesn't compalain
+      int shape_ = 0;
+      size_t stride_ = 0;
+      compute_encoder->setBytes(&shape_, sizeof(int), 3);
+      compute_encoder->setBytes(&stride_, sizeof(size_t), 4);
+    } else {
+      compute_encoder->setBytes(upd.shape().data(), upd_ndim * sizeof(int), 3);
+      compute_encoder->setBytes(
+          upd.strides().data(), upd_ndim * sizeof(size_t), 4);
+    }
+    compute_encoder->setBytes(&upd_ndim, sizeof(size_t), 5);
+    compute_encoder->setBytes(&upd_size, sizeof(size_t), 6);
+
+    // Set output info
+    size_t out_ndim = out.ndim();
+    if (out_ndim == 0) {
+      // Need placeholders so Metal doesn't compalain
+      int shape_ = 0;
+      size_t stride_ = 0;
+      compute_encoder->setBytes(&shape_, sizeof(int), 7);
+      compute_encoder->setBytes(&stride_, sizeof(size_t), 8);
+    } else {
+      compute_encoder->setBytes(out.shape().data(), out_ndim * sizeof(int), 7);
+      compute_encoder->setBytes(
+          out.strides().data(), out_ndim * sizeof(size_t), 8);
+    }
+    compute_encoder->setBytes(&out_ndim, sizeof(size_t), 9);
+    compute_encoder->setBytes(axes_.data(), axes_.size() * sizeof(int), 10);
+
+    // Set index info
+    if (idx_ndim == 0) {
+      // Add a 0 in idx_shapes and strides to avoid the missing buffer binding
+      // error in the metal API.
+      idx_shapes.push_back(0);
+      idx_strides.push_back(0);
+    }
+    compute_encoder->setBytes(
+        idx_shapes.data(), idx_shapes.size() * sizeof(int), 11);
+    compute_encoder->setBytes(
+        idx_strides.data(), idx_strides.size() * sizeof(size_t), 12);
+    compute_encoder->setBytes(&idx_ndim, sizeof(int), 13);
+
+    // Set index buffers
+    for (int i = 1; i < nidx + 1; ++i) {
+      set_array_buffer(compute_encoder, inputs[i], 20 + i);
+    }
+
+    // Launch grid
+    MTL::Size grid_dims = MTL::Size(upd_size, nthreads / upd_size, 1);
+    MTL::Size group_dims = get_block_dims(upd_size, nthreads / upd_size, 1);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+  }
 }
 
 } // namespace mlx::core

mlx/backend/metal/kernels/CMakeLists.txt

@@ -3,11 +3,12 @@ set(
   ${CMAKE_CURRENT_SOURCE_DIR}/atomic.h
   ${CMAKE_CURRENT_SOURCE_DIR}/bf16.h
   ${CMAKE_CURRENT_SOURCE_DIR}/bf16_math.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/binary.h
   ${CMAKE_CURRENT_SOURCE_DIR}/complex.h
   ${CMAKE_CURRENT_SOURCE_DIR}/defines.h
   ${CMAKE_CURRENT_SOURCE_DIR}/erf.h
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.h
-  ${CMAKE_CURRENT_SOURCE_DIR}/reduce.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/unary.h
   ${CMAKE_CURRENT_SOURCE_DIR}/utils.h
 )
 
@@ -22,11 +23,12 @@ set(
   "gemv"
   "quantized"
   "random"
-  "reduce"
   "rope"
   "scan"
+  "scaled_dot_product_attention"
   "softmax"
   "sort"
+  "ternary"
   "unary"
   "gather"
   "scatter"
@@ -48,11 +50,7 @@ endfunction(build_kernel_base)
 
 function(build_kernel KERNEL)
   set(SRCFILE ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL}.metal)
-  set(HEADERS_PADDED ${HEADERS})
-  if(${KERNEL} STREQUAL "conv")
-    set(HEADERS_PADDED ${HEADERS_PADDED} ${CMAKE_CURRENT_SOURCE_DIR}/conv.h)
-  endif()
-  build_kernel_base(${KERNEL} ${SRCFILE} "${HEADERS_PADDED}")
+  build_kernel_base(${KERNEL} ${SRCFILE} "${HEADERS}")
 endfunction(build_kernel)
 
 foreach(KERNEL ${KERNELS})
@@ -69,6 +67,15 @@ foreach(KERNEL ${STEEL_KERNELS})
   set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR})
 endforeach()
 
+file(GLOB_RECURSE REDUCE_KERNELS ${CMAKE_CURRENT_SOURCE_DIR}/reduction/*.metal)
+file(GLOB_RECURSE REDUCE_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/reduction/*.h)
+
+foreach(KERNEL ${REDUCE_KERNELS})
+  cmake_path(GET KERNEL STEM TARGET)
+  build_kernel_base(${TARGET} ${KERNEL} "${REDUCE_HEADERS}")
+  set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR})
+endforeach()
+
 add_custom_command(
   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
   COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o ${MLX_METAL_PATH}/mlx.metallib

mlx/backend/metal/kernels/arg_reduce.metal

@@ -11,8 +11,6 @@ template <typename U>
 struct IndexValPair {
   uint32_t index;
   U val;
-
-  IndexValPair(uint32_t _index, U _val) : index(_index), val(_val) {}
 };
 
 template <typename U>
@@ -65,10 +63,10 @@ struct ArgMax {
 
 template <typename U>
 IndexValPair<U> simd_shuffle_down(IndexValPair<U> data, uint16_t delta) {
-  return IndexValPair<U>(
+  return IndexValPair<U>{
     simd_shuffle_down(data.index, delta),
     simd_shuffle_down(data.val, delta)
-  );
+  };
 }
 
 
@@ -82,7 +80,6 @@ template <typename T, typename Op, int N_READS>
     const device size_t& ndim [[buffer(5)]],
     const device size_t& axis_stride [[buffer(6)]],
     const device size_t& axis_size [[buffer(7)]],
-    threadgroup IndexValPair<T> *local_data [[threadgroup(0)]],
     uint gid [[thread_position_in_grid]],
     uint lid [[thread_position_in_threadgroup]],
     uint lsize [[threads_per_threadgroup]],
@@ -111,7 +108,9 @@ template <typename T, typename Op, int N_READS>
   auto in_idx = elem_to_loc(gid / lsize, shape, in_strides, ndim);
   auto out_idx = elem_to_loc(gid / lsize, shape, out_strides, ndim);
 
-  IndexValPair<T> best(0, Op::init);
+  IndexValPair<T> best{0, Op::init};
+
+  threadgroup IndexValPair<T> local_data[32];
 
   // Loop over the reduction axis in lsize*N_READS buckets
   for (uint r=0; r < ceildiv(axis_size, N_READS*lsize); r++) {
@@ -172,7 +171,6 @@ template <typename T, typename Op, int N_READS>
       const device size_t& ndim [[buffer(5)]], \
       const device size_t& axis_stride [[buffer(6)]], \
       const device size_t& axis_size [[buffer(7)]], \
-      threadgroup IndexValPair<itype> *local_data [[threadgroup(0)]], \
       uint gid [[thread_position_in_grid]], \
       uint lid [[thread_position_in_threadgroup]], \
       uint lsize [[threads_per_threadgroup]], \

mlx/backend/metal/kernels/binary.metal

@@ -2,16 +2,6 @@
 
 #include "mlx/backend/metal/kernels/binary.h"
 
-template <typename T, typename U, typename Op>
-[[kernel]] void binary_op_s2s(
-    device const T* a,
-    device const T* b,
-    device U* c,
-    uint index [[thread_position_in_grid]]) {
-  c[index] = Op()(a[0], b[0]);
-}
-
-
 template <typename T, typename U, typename Op>
 [[kernel]] void binary_op_ss(
     device const T* a,

mlx/backend/metal/kernels/compiled_preamble.h

@@ -1,4 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/metal/kernels/binary.h"
+#include "mlx/backend/metal/kernels/ternary.h"
 #include "mlx/backend/metal/kernels/unary.h"
+
+typedef half float16_t;

mlx/backend/metal/kernels/conv.h

@@ -1,481 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-#pragma once
-
-#include <metal_simdgroup>
-#include <metal_simdgroup_matrix>
-#include <metal_stdlib>
-
-#include "mlx/backend/metal/kernels/bf16.h"
-#include "mlx/backend/metal/kernels/conv_params.h"
-
-#define MLX_MTL_CONST static constant constexpr const
-
-using namespace metal;
-
-///////////////////////////////////////////////////////////////////////////////
-// Loading helper
-///////////////////////////////////////////////////////////////////////////////
-
-template <
-    typename T,
-    int BM,
-    int BN,
-    int BK,
-    int vec_size,
-    int tgp_size,
-    int tgp_padding = 0>
-struct Conv2DInputBlockLoader {
-  // Destination dimensions
-  MLX_MTL_CONST int dst_fd = BM;
-  MLX_MTL_CONST int dst_ld = BK + tgp_padding;
-  MLX_MTL_CONST int n_vecs = BK / vec_size;
-
-  // Stride along block row within the block
-  MLX_MTL_CONST int bstride = tgp_size / n_vecs;
-  MLX_MTL_CONST int n_rows = dst_fd / bstride;
-
-  // Thread location indices
-  const short thread_idx;
-  const short bi;
-  const short bj;
-
-  // threadgroup and device memory
-  threadgroup T* dst;
-  const device T* src;
-
-  const constant MLXConvParams<2>& params;
-
-  int weight_h;
-  int weight_w;
-
-  int offsets_n[n_rows];
-  int offsets_oh[n_rows];
-  int offsets_ow[n_rows];
-
-  /* Constructor */
-  METAL_FUNC Conv2DInputBlockLoader(
-      const device T* src_,
-      threadgroup T* dst_,
-      const constant MLXConvParams<2>& params_,
-      uint3 tid [[threadgroup_position_in_grid]],
-      uint3 lid [[thread_position_in_threadgroup]],
-      uint simd_group_id [[simdgroup_index_in_threadgroup]],
-      uint simd_lane_id [[thread_index_in_simdgroup]])
-      : thread_idx(simd_group_id * 32 + simd_lane_id),
-        bi(thread_idx / n_vecs),
-        bj(vec_size * (thread_idx % n_vecs)),
-        dst(dst_ + bi * dst_ld + bj),
-        src(src_ + bj),
-        params(params_),
-        weight_h(0),
-        weight_w(0) {
-    int out_n_pixels = params.oS[0] * params.oS[1];
-
-    for (int i = 0; i < n_rows; ++i) {
-      int offset_nhw = tid.y * BM + bi + i * bstride;
-      offsets_n[i] = offset_nhw / out_n_pixels;
-      int hw = offset_nhw % out_n_pixels;
-      offsets_oh[i] = hw / params.oS[1];
-      offsets_ow[i] = hw % params.oS[1];
-    }
-
-    (void)lid;
-  }
-
-  /* Load from device memory into threadgroup memory - without bound checking */
-  METAL_FUNC void load_unsafe() const {
-#pragma clang loop unroll(full)
-    for (short i = 0, is = 0; i < n_rows; ++i, is += bstride) {
-      int n = offsets_n[i];
-      int oh = offsets_oh[i];
-      int ow = offsets_ow[i];
-
-      int ih = oh * params.str[0] - params.pad[0] + weight_h * params.dil[0];
-      int iw = ow * params.str[1] - params.pad[1] + weight_w * params.dil[1];
-
-      // Read from input if in bounds
-      if (ih >= 0 && ih < params.iS[0] && iw >= 0 && iw < params.iS[1]) {
-        const device T* curr_src = src + n * params.in_strides[0] +
-            ih * params.in_strides[1] + iw * params.in_strides[2];
-
-#pragma clang loop unroll(full)
-        for (short j = 0; j < vec_size; ++j) {
-          dst[is * dst_ld + j] = curr_src[j];
-        }
-      }
-
-      // Zero pad otherwise
-      else {
-#pragma clang loop unroll(full)
-        for (short j = 0; j < vec_size; ++j) {
-          dst[is * dst_ld + j] = T(0);
-        }
-      }
-    }
-  }
-
-  /* Iteration helper */
-  METAL_FUNC void next() {
-    if (++weight_w < params.wS[1]) {
-      return;
-    }
-
-    weight_w = 0;
-
-    if (++weight_h < params.wS[0]) {
-      return;
-    }
-
-    weight_h = 0;
-
-    src += BK;
-  }
-};
-
-template <
-    typename T,
-    int BM,
-    int BN,
-    int BK,
-    int vec_size,
-    int tgp_size,
-    int tgp_padding = 0>
-struct Conv2DWeightBlockLoader {
-  // Destination dimensions
-  MLX_MTL_CONST int dst_fd = BN;
-  MLX_MTL_CONST int dst_ld = BK + tgp_padding;
-  MLX_MTL_CONST int n_vecs = BK / vec_size;
-
-  // Stride along block row within the block
-  MLX_MTL_CONST int bstride = tgp_size / n_vecs;
-  MLX_MTL_CONST int n_rows = dst_fd / bstride;
-
-  // Leading dimension for src
-  const int src_ld;
-
-  // Thread location indices
-  const short thread_idx;
-  const short bi;
-  const short bj;
-
-  // threadgroup and device memory
-  threadgroup T* dst;
-  const device T* src;
-
-  const constant MLXConvParams<2>& params;
-
-  int weight_h;
-  int weight_w;
-
-  /* Constructor */
-  METAL_FUNC Conv2DWeightBlockLoader(
-      const device T* src_,
-      threadgroup T* dst_,
-      const constant MLXConvParams<2>& params_,
-      uint3 tid [[threadgroup_position_in_grid]],
-      uint3 lid [[thread_position_in_threadgroup]],
-      uint simd_group_id [[simdgroup_index_in_threadgroup]],
-      uint simd_lane_id [[thread_index_in_simdgroup]])
-      : src_ld(params_.wt_strides[0]),
-        thread_idx(simd_group_id * 32 + simd_lane_id),
-        bi(thread_idx / n_vecs),
-        bj(vec_size * (thread_idx % n_vecs)),
-        dst(dst_ + bi * dst_ld + bj),
-        src(src_ + bi * src_ld + bj),
-        params(params_),
-        weight_h(0),
-        weight_w(0) {
-    (void)lid;
-    (void)tid;
-  }
-
-  /* Load from device memory into threadgroup memory - without bound checking */
-  METAL_FUNC void load_unsafe() const {
-    const device T* curr_src =
-        src + weight_h * params.wt_strides[1] + weight_w * params.wt_strides[2];
-#pragma clang loop unroll(full)
-    for (short i = 0; i < dst_fd; i += bstride) {
-#pragma clang loop unroll(full)
-      for (short j = 0; j < vec_size; j++) {
-        dst[i * dst_ld + j] = curr_src[i * src_ld + j];
-      }
-    }
-  }
-
-  /* Iteration helper */
-  METAL_FUNC void next() {
-    if (++weight_w < params.wS[1]) {
-      return;
-    }
-
-    weight_w = 0;
-
-    if (++weight_h < params.wS[0]) {
-      return;
-    }
-
-    weight_h = 0;
-
-    src += BK;
-  }
-};
-
-///////////////////////////////////////////////////////////////////////////////
-// Transforms
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename OutT, typename InT>
-struct TransformNone {
-  static METAL_FUNC OutT apply(InT x) {
-    return static_cast<OutT>(x);
-  }
-};
-
-template <typename T>
-struct AccumHelper {
-  typedef float accum_type;
-};
-
-///////////////////////////////////////////////////////////////////////////////
-// MMA helper
-///////////////////////////////////////////////////////////////////////////////
-
-template <
-    typename T,
-    int BM,
-    int BN,
-    int BK,
-    int WM,
-    int WN,
-    bool transpose_a,
-    bool transpose_b,
-    int tgp_padding_a = 0,
-    int tgp_padding_b = 0,
-    typename AccumType = typename AccumHelper<T>::accum_type,
-    typename Epilogue = TransformNone<T, AccumType>>
-struct Conv2DBlockMMA {
-  // Warp tile size along M
-  MLX_MTL_CONST int TM = BM / (WM * 8);
-  // Warp tile size along N
-  MLX_MTL_CONST int TN = BN / (WN * 8);
-
-  // Warp tile simdgroup matrix strides along M
-  MLX_MTL_CONST int TM_stride = 8 * WM;
-  // Warp tile simdgroup matrix strides along M
-  MLX_MTL_CONST int TN_stride = 8 * WN;
-
-  // Leading dimensions of threadgroup A, B blocks
-  MLX_MTL_CONST int lda_tgp = (transpose_a ? BM : BK) + tgp_padding_a;
-  MLX_MTL_CONST int ldb_tgp = (transpose_b ? BK : BN) + tgp_padding_b;
-
-  // Strides of A, B along reduction axis
-  MLX_MTL_CONST short simd_stride_a =
-      transpose_a ? TM_stride : TM_stride * lda_tgp;
-  MLX_MTL_CONST short simd_stride_b =
-      transpose_b ? TN_stride * ldb_tgp : TN_stride;
-
-  // Jump between elements
-  MLX_MTL_CONST short jump_a = transpose_a ? lda_tgp : 1;
-  MLX_MTL_CONST short jump_b = transpose_b ? ldb_tgp : 1;
-
-  // Offsets within threadgroup
-  const int tm;
-  const int tn;
-
-  // Simdgroup matrices
-  simdgroup_matrix<AccumType, 8, 8> Asimd[TM];
-  simdgroup_matrix<AccumType, 8, 8> Bsimd[TN];
-  simdgroup_matrix<AccumType, 8, 8> results[TM * TN] = {
-      simdgroup_matrix<AccumType, 8, 8>(0)};
-
-  short sm;
-  short sn;
-
-  /* Constructor */
-  METAL_FUNC Conv2DBlockMMA(
-      uint simd_group_id [[simdgroup_index_in_threadgroup]],
-      uint simd_lane_id [[thread_index_in_simdgroup]])
-      : tm(8 * (simd_group_id / WN)), tn(8 * (simd_group_id % WN)) {
-    short qid = simd_lane_id / 4;
-    sm = (qid & 4) + (simd_lane_id / 2) % 4;
-    sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
-  }
-
-  /* (BM, BK) X (BK, BN) multiply accumulate function */
-  METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) {
-// Iterate over BK in blocks of 8
-#pragma clang loop unroll(full)
-    for (short kk = 0; kk < BK; kk += 8) {
-      short2 offset_a =
-          transpose_a ? short2(tm + sm, kk + sn) : short2(kk + sn, tm + sm);
-      short2 offset_b =
-          transpose_b ? short2(kk + sm, tn + sn) : short2(tn + sn, kk + sm);
-
-      const threadgroup T* As__ = As + offset_a.y * lda_tgp + offset_a.x;
-      const threadgroup T* Bs__ = Bs + offset_b.y * ldb_tgp + offset_b.x;
-
-      simdgroup_barrier(mem_flags::mem_none);
-// Load elements from threadgroup A as simdgroup matrices
-#pragma clang loop unroll(full)
-      for (short i = 0; i < TM; i++) {
-        Asimd[i].thread_elements()[0] = static_cast<AccumType>(As__[0]);
-        Asimd[i].thread_elements()[1] = static_cast<AccumType>(As__[jump_a]);
-        As__ += simd_stride_a;
-      }
-
-      simdgroup_barrier(mem_flags::mem_none);
-// Load elements from threadgroup B as simdgroup matrices
-#pragma clang loop unroll(full)
-      for (short j = 0; j < TN; j++) {
-        Bsimd[j].thread_elements()[0] = static_cast<AccumType>(Bs__[0]);
-        Bsimd[j].thread_elements()[1] = static_cast<AccumType>(Bs__[jump_b]);
-        Bs__ += simd_stride_b;
-      }
-
-      simdgroup_barrier(mem_flags::mem_none);
-// Multiply and accumulate into result simdgroup matrices
-#pragma clang loop unroll(full)
-      for (short i = 0; i < TM; i++) {
-#pragma clang loop unroll(full)
-        for (short j = 0; j < TN; j++) {
-          simdgroup_multiply_accumulate(
-              results[i * TN + j], Asimd[i], Bsimd[j], results[i * TN + j]);
-        }
-      }
-    }
-  }
-
-  /* Store results from simdgroup_matrix results into device memory */
-  METAL_FUNC void store_result(device T* C, const int ldc) const {
-#pragma clang loop unroll(full)
-    for (int i = 0; i < TM; i++) {
-#pragma clang loop unroll(full)
-      for (int j = 0; j < TN; j++) {
-        C[(i * TM_stride + sm + tm) * ldc + j * TN_stride + tn + sn] =
-            Epilogue::apply(results[i * TN + j].thread_elements()[0]);
-        C[(i * TM_stride + sm + tm) * ldc + j * TN_stride + tn + sn + 1] =
-            Epilogue::apply(results[i * TN + j].thread_elements()[1]);
-      }
-    }
-  }
-
-  METAL_FUNC void
-  store_result_safe(device T* C, const int ldc, short2 dst_tile_dims) const {
-#pragma clang loop unroll(full)
-    for (int i = 0; i < TM; i++) {
-      if (tm + i * TM_stride + sm < dst_tile_dims.y) {
-#pragma clang loop unroll(full)
-        for (int j = 0; j < TN; j++) {
-          if (tn + j * TN_stride + sn < dst_tile_dims.x) {
-            C[(tm + i * TM_stride + sm) * ldc + tn + j * TN_stride + sn] =
-                Epilogue::apply(results[i * TN + j].thread_elements()[0]);
-          }
-
-          if (tn + j * TN_stride + sn + 1 < dst_tile_dims.x) {
-            C[(tm + i * TM_stride + sm) * ldc + tn + j * TN_stride + sn + 1] =
-                Epilogue::apply(results[i * TN + j].thread_elements()[1]);
-          }
-        }
-      }
-    }
-  }
-};
-
-///////////////////////////////////////////////////////////////////////////////
-// GEMM kernels
-///////////////////////////////////////////////////////////////////////////////
-
-template <
-    typename T,
-    int BM,
-    int BN,
-    int BK,
-    int WM,
-    int WN,
-    bool transpose_a,
-    bool transpose_b,
-    typename AccumType = typename AccumHelper<T>::accum_type,
-    typename Epilogue = TransformNone<T, AccumType>>
-struct Conv2DImplicitGEMMKernel {
-  MLX_MTL_CONST short tgp_padding_a = 16 / sizeof(T);
-  MLX_MTL_CONST short tgp_padding_b = 16 / sizeof(T);
-  MLX_MTL_CONST short tgp_mem_size_a =
-      transpose_a ? BK * (BM + tgp_padding_a) : BM * (BK + tgp_padding_a);
-  MLX_MTL_CONST short tgp_mem_size_b =
-      transpose_b ? BN * (BK + tgp_padding_b) : BK * (BN + tgp_padding_b);
-  MLX_MTL_CONST short tgp_mem_size = tgp_mem_size_a + tgp_mem_size_b;
-
-  MLX_MTL_CONST short tgp_size = WM * WN * 32;
-  MLX_MTL_CONST short vec_size = (BM == 64 && BN == 64) ? 8 : 4;
-
-  using loader_a_t =
-      Conv2DInputBlockLoader<T, BM, BN, BK, vec_size, tgp_size, tgp_padding_a>;
-  using loader_b_t =
-      Conv2DWeightBlockLoader<T, BM, BN, BK, vec_size, tgp_size, tgp_padding_b>;
-  using mma_t = Conv2DBlockMMA<
-      T,
-      BM,
-      BN,
-      BK,
-      WM,
-      WN,
-      transpose_a,
-      transpose_b,
-      tgp_padding_a,
-      tgp_padding_b,
-      AccumType,
-      Epilogue>;
-
-  /* Main kernel function */
-  static METAL_FUNC void run(
-      const device T* A [[buffer(0)]],
-      const device T* B [[buffer(1)]],
-      device T* C [[buffer(2)]],
-      const constant MLXConvParams<2>& params [[buffer(3)]],
-      threadgroup T* tgp_memory [[threadgroup(0)]],
-      uint3 tid [[threadgroup_position_in_grid]],
-      uint3 lid [[thread_position_in_threadgroup]],
-      uint simd_gid [[simdgroup_index_in_threadgroup]],
-      uint simd_lid [[thread_index_in_simdgroup]]) {
-    const int c_row = tid.y * BM;
-    const int c_col = tid.x * BN;
-    const int K = params.wt_strides[0];
-    const int N = params.O;
-
-    B += c_col * K;
-    C += c_row * N + c_col;
-
-    // Prepare threadgroup memory for loading
-    threadgroup T* As = tgp_memory;
-    threadgroup T* Bs = tgp_memory + tgp_mem_size_a;
-
-    // Prepare threadgroup loading operations
-    loader_a_t loader_a(A, As, params, tid, lid, simd_gid, simd_lid);
-    loader_b_t loader_b(B, Bs, params, tid, lid, simd_gid, simd_lid);
-
-    // Prepare threadgroup mma operation
-    mma_t mma_op(simd_gid, simd_lid);
-
-    for (int k = 0; k < K; k += BK) {
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-      // Load elements into threadgroup
-      loader_a.load_unsafe();
-      loader_b.load_unsafe();
-
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-
-      // Multiply and accumulate threadgroup elements
-      mma_op.mma(As, Bs);
-
-      // Prepare for next iteration
-      loader_a.next();
-      loader_b.next();
-    }
-
-    threadgroup_barrier(mem_flags::mem_none);
-
-    // Store results to device memory
-    mma_op.store_result(C, N);
-  }
-};

mlx/backend/metal/kernels/conv.metal

@@ -1,16 +1,102 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <metal_stdlib>
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
 
-#include "mlx/backend/metal/kernels/conv_params.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
 #include "mlx/backend/metal/kernels/bf16.h"
 
-#include "mlx/backend/metal/kernels/conv.h"
+#define MLX_MTL_CONST static constant constexpr const
 
 using namespace metal;
 
 ///////////////////////////////////////////////////////////////////////////////
-/// Slow and naive kernels
+/// Naive unfold with dilation
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, int N>
+[[kernel]] void naive_unfold_Nd(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    const constant MLXConvParams<N>* params [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+
+  int filter_size = params->C;
+  for(short i = 0; i < N; i++) filter_size *= params->wS[i];
+
+  int out_pixels = 1;
+  for(short i = 0; i < N; i++) out_pixels *= params->oS[i];
+
+  // Set out 
+  out += gid.z * filter_size + gid.y * (params->C);
+
+  // Corrdinates in input
+  int is[N] = {0};
+
+  // gid.z: N oS (Batch and row in unfolded output)
+  // gid.y: wS (Filter location to unfold input)
+  // gid.x: C (channel)
+
+  int n = (gid.z) / out_pixels;
+  int oS = (gid.z) % out_pixels;
+  int wS = gid.y;
+
+  bool valid = n < params->N;
+
+  // Unroll dimensions 
+  for (int i = N - 1; i >= 0; --i) {
+    int os_ = (oS % params->oS[i]);
+    int ws_ = (wS % params->wS[i]);
+    
+    ws_ = params->flip ? params->wS[i] - ws_ - 1 : ws_;
+
+    int is_ = os_ * params->str[i] - params->pad[i] + ws_ * params->kdil[i];
+    int is_max = 1 + params->idil[i] * (params->iS[i] - 1);
+
+    valid &= is_ >= 0 && is_ < is_max && (is_ % params->idil[i] == 0);
+
+    is[i] = is_ / params->idil[i];
+
+    oS /= params->oS[i];
+    wS /= params->wS[i];
+  }
+
+  if(valid) {
+    size_t in_offset = n * params->in_strides[0];
+
+    for(int i = 0; i < N; ++i) {
+      in_offset += is[i] * params->in_strides[i + 1];
+    }
+
+    out[gid.x] = in[in_offset + gid.x];
+  } else {
+    out[gid.x] = T(0);
+  }
+
+}
+
+#define instantiate_naive_unfold_nd(name, itype, n) \
+  template [[host_name("naive_unfold_nd_" #name "_" #n)]] \
+  [[kernel]] void naive_unfold_Nd( \
+      const device itype* in [[buffer(0)]], \
+      device itype* out [[buffer(1)]], \
+      const constant MLXConvParams<n>* params [[buffer(2)]], \
+      uint3 gid [[thread_position_in_grid]]);
+
+#define instantiate_naive_unfold_nd_dims(name, itype) \
+  instantiate_naive_unfold_nd(name, itype, 1) \
+  instantiate_naive_unfold_nd(name, itype, 2) \
+  instantiate_naive_unfold_nd(name, itype, 3)
+
+instantiate_naive_unfold_nd_dims(float32, float);
+instantiate_naive_unfold_nd_dims(float16, half);
+instantiate_naive_unfold_nd_dims(bfloat16, bfloat16_t);
+
+///////////////////////////////////////////////////////////////////////////////
+/// Slow and naive conv2d kernels
 ///////////////////////////////////////////////////////////////////////////////
 
 template <typename T, 
@@ -58,8 +144,8 @@ template <typename T,
 
         // Local in
         for(int m = 0; m < TM; m++) {
-          int i = out_h[m] * params.str[0] - params.pad[0] + h * params.dil[0];
-          int j = out_w[m] * params.str[1] - params.pad[1] + w * params.dil[1];
+          int i = out_h[m] * params.str[0] - params.pad[0] + h * params.kdil[0];
+          int j = out_w[m] * params.str[1] - params.pad[1] + w * params.kdil[1];
 
           bool valid = i >= 0 && i < params.iS[0] && j >= 0 && j < params.iS[1];
           in_local[m] = valid ? in[i * params.in_strides[1] + j * params.in_strides[2] + c] : T(0);
@@ -116,59 +202,6 @@ instantiate_naive_conv_2d_blocks(float32, float);
 instantiate_naive_conv_2d_blocks(float16, half);
 instantiate_naive_conv_2d_blocks(bfloat16, bfloat16_t);
 
-///////////////////////////////////////////////////////////////////////////////
-/// Implicit gemm kernels
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T,
-          int BM,
-          int BN,
-          int BK,
-          int WM,
-          int WN>
-[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void implicit_gemm_conv_2d(
-    const device T* in [[buffer(0)]],
-    const device T* wt [[buffer(1)]],
-    device T* out [[buffer(2)]],
-    const constant MLXConvParams<2>& params [[buffer(3)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint simd_gid [[simdgroup_index_in_threadgroup]],
-    uint simd_lid [[thread_index_in_simdgroup]]) {
-
-  using gemm_kernel = Conv2DImplicitGEMMKernel<T, BM, BN, BK, WM, WN, /*transpose_a*/ false, /*transpose_b*/ true>;
-    
-  threadgroup T tgp_memory[gemm_kernel::tgp_mem_size];
-
-  gemm_kernel::run( 
-    in, wt, out, 
-    params, tgp_memory,
-    tid, lid, simd_gid, simd_lid
-  );
-
-}
-
-#define instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn) \
-  template [[host_name("implicit_gemm_conv_2d_" #name "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn)]] \
-  [[kernel]] void implicit_gemm_conv_2d<itype, bm, bn, bk, wm, wn>( \
-      const device itype* in [[buffer(0)]], \
-      const device itype* wt [[buffer(1)]], \
-      device itype* out [[buffer(2)]], \
-      const constant MLXConvParams<2>& params [[buffer(3)]], \
-      uint3 tid [[threadgroup_position_in_grid]], \
-      uint3 lid [[thread_position_in_threadgroup]], \
-      uint simd_gid [[simdgroup_index_in_threadgroup]], \
-      uint simd_lid [[thread_index_in_simdgroup]]);
-
-#define instantiate_implicit_2d_blocks(name, itype) \
-    instantiate_implicit_conv_2d(name, itype, 32, 32, 32, 2, 2) \
-    instantiate_implicit_conv_2d(name, itype, 32, 32, 16, 2, 2) \
-    instantiate_implicit_conv_2d(name, itype, 64, 64, 16, 2, 2)
-
-instantiate_implicit_2d_blocks(float32, float);
-instantiate_implicit_2d_blocks(float16, half);
-instantiate_implicit_2d_blocks(bfloat16, bfloat16_t);
-
 ///////////////////////////////////////////////////////////////////////////////
 /// Winograd kernels
 ///////////////////////////////////////////////////////////////////////////////

mlx/backend/metal/kernels/conv_params.h

@@ -1,19 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-#pragma once
-
-template <int NDIM>
-struct MLXConvParams {
-  const int N; // Batch size
-  const int C; // In channels
-  const int O; // Out channels
-  const int iS[NDIM]; // Input spatial dim
-  const int wS[NDIM]; // Weight spatial dim
-  const int oS[NDIM]; // Output spatial dim
-  const int str[NDIM]; // Kernel strides
-  const int pad[NDIM]; // Input padding
-  const int dil[NDIM]; // Kernel dilation
-  const size_t in_strides[NDIM + 2]; // In strides
-  const size_t wt_strides[NDIM + 2]; // Wt strides
-  const size_t out_strides[NDIM + 2]; // Out strides
-};

mlx/backend/metal/kernels/quantized.metal

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <metal_stdlib>
 #include <metal_simdgroup>
@@ -23,7 +23,143 @@ template <> struct AccT<bfloat16_t> {
   typedef float acc_t;
 };
 
-template <typename T, const int BM, const int BN, const int group_size, const int bits>
+
+template <typename T, typename U, int values_per_thread, int bits>
+inline U load_vector(const device T *x, thread U *x_thread) {
+  static_assert(bits == 2 || bits == 4 || bits == 8, "Template undefined for bits not in {2, 4, 8}");
+
+  U sum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i+1] + x[i+2] + x[i+3];
+      x_thread[i] = x[i];
+      x_thread[i+1] = x[i+1] / 4.0f;
+      x_thread[i+2] = x[i+2] / 16.0f;
+      x_thread[i+3] = x[i+3] / 64.0f;
+    }
+  }
+
+  else if (bits == 4) {
+    for (int i = 0; i < values_per_thread; i += 4) {
+      sum += x[i] + x[i+1] + x[i+2] + x[i+3];
+      x_thread[i] = x[i];
+      x_thread[i+1] = x[i+1] / 16.0f;
+      x_thread[i+2] = x[i+2] / 256.0f;
+      x_thread[i+3] = x[i+3] / 4096.0f;
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      sum += x[i];
+      x_thread[i] = x[i];
+    }
+  }
+
+  return sum;
+}
+
+template <typename U, int values_per_thread, int bits>
+inline U qdot(const device uint8_t* w, const thread U *x_thread, U scale, U bias, U sum) {
+  static_assert(bits == 2 || bits == 4 || bits == 8, "Template undefined for bits not in {2, 4, 8}");
+
+  U accum = 0;
+
+  if (bits == 2) {
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum += (
+          x_thread[4*i] * (w[i] & 0x03)
+          + x_thread[4*i+1] * (w[i] & 0x0c)
+          + x_thread[4*i+2] * (w[i] & 0x30)
+          + x_thread[4*i+3] * (w[i] & 0xc0));
+    }
+  }
+
+  else if (bits == 4) {
+    const device uint16_t* ws = (const device uint16_t*)w;
+    for (int i = 0; i < (values_per_thread / 4); i++) {
+      accum += (
+          x_thread[4*i] * (ws[i] & 0x000f)
+          + x_thread[4*i+1] * (ws[i] & 0x00f0)
+          + x_thread[4*i+2] * (ws[i] & 0x0f00)
+          + x_thread[4*i+3] * (ws[i] & 0xf000));
+    }
+  }
+
+  else if (bits == 8) {
+    for (int i = 0; i < values_per_thread; i++) {
+      accum += x_thread[i] * w[i];
+    }
+  }
+
+  return scale * accum + sum * bias;
+}
+
+template <typename T, int group_size, int bits, int packs_per_thread>
+[[kernel]] void qmv_fast(
+    const device uint32_t* w [[buffer(0)]],
+    const device T* scales [[buffer(1)]],
+    const device T* biases [[buffer(2)]],
+    const device T* x [[buffer(3)]],
+    device T* y [[buffer(4)]],
+    const constant int& in_vec_size [[buffer(5)]],
+    const constant int& out_vec_size [[buffer(6)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int pack_factor = 32 / bits;
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
+
+  typedef float U;
+
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
+
+  // Adjust positions
+  const int in_vec_size_w = in_vec_size / pack_factor;
+  const int in_vec_size_g = in_vec_size / group_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup;
+  w += out_row * in_vec_size_w + simd_lid * packs_per_thread;
+  scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+  biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+  x += tid.z * in_vec_size + simd_lid * values_per_thread;
+  y += tid.z * out_vec_size + out_row;
+
+  for (int k = 0; k < in_vec_size; k += block_size) {
+    U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      const device uint8_t* wl = (const device uint8_t *)(w + row * in_vec_size_w);
+      const device T* sl = scales + row * in_vec_size_g;
+      const device T* bl = biases + row * in_vec_size_g;
+
+      U s = sl[0];
+      U b = bl[0];
+      result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+    }
+
+    w += block_size / pack_factor;
+    scales += block_size / group_size;
+    biases += block_size / group_size;
+    x += block_size;
+  }
+
+  for (int row = 0; row < results_per_simdgroup; row++) {
+    result[row] = simd_sum(result[row]);
+    if (simd_lid == 0) {
+      y[row] = static_cast<T>(result[row]);
+    }
+  }
+}
+
+
+template <typename T, const int group_size, const int bits>
 [[kernel]] void qmv(
     const device uint32_t* w [[buffer(0)]],
     const device T* scales [[buffer(1)]],
@@ -33,91 +169,101 @@ template <typename T, const int BM, const int BN, const int group_size, const in
     const constant int& in_vec_size [[buffer(5)]],
     const constant int& out_vec_size [[buffer(6)]],
     uint3 tid [[threadgroup_position_in_grid]],
-    uint lid [[thread_index_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
 
-  static_assert(BN == SIMD_SIZE, "qmv expects BN to be equal to SIMD_SIZE");
+  constexpr int num_simdgroups = 2;
+  constexpr int results_per_simdgroup = 4;
+  constexpr int packs_per_thread = 1;
+  constexpr int pack_factor = 32 / bits;
+  constexpr int values_per_thread = pack_factor * packs_per_thread;
+  constexpr int block_size = values_per_thread * SIMD_SIZE;
+  constexpr int scale_step_per_thread = group_size / values_per_thread;
 
-  (void)lid;
+  typedef float U;
 
-  constexpr int bitmask = (1 << bits) - 1;
-  constexpr int el_per_thread = 32 / bits;
-  constexpr int colgroup = BN * el_per_thread;
-  constexpr int groups_per_block = colgroup / group_size;
-
-  typedef typename AccT<T>::acc_t U;
-  threadgroup U scales_block[BM * groups_per_block];
-  threadgroup U biases_block[BM * groups_per_block];
-  threadgroup U x_block[colgroup];
-
-  thread uint32_t w_local;
-  thread U result = 0;
-  thread U scale = 1;
-  thread U bias = 0;
-  thread U x_thread[el_per_thread];
+  thread U x_thread[values_per_thread];
+  thread U result[results_per_simdgroup] = {0};
 
   // Adjust positions
-  const int in_vec_size_w = in_vec_size / el_per_thread;
+  const int in_vec_size_w = in_vec_size / pack_factor;
   const int in_vec_size_g = in_vec_size / group_size;
-  int out_row = tid.y * BM + simd_gid;
-  w += out_row * in_vec_size_w;
-  scales += out_row * in_vec_size_g;
-  biases += out_row * in_vec_size_g;
-  x += tid.z * in_vec_size;
-  y += tid.z * out_vec_size;
+  const int out_row = tid.y * (num_simdgroups * results_per_simdgroup) + simd_gid * results_per_simdgroup;
+  const int used_out_row = min(out_vec_size - results_per_simdgroup, out_row);
 
   if (out_row >= out_vec_size) {
     return;
   }
 
-  // Loop over in_vec in blocks of colgroup
-  for (int i=0; i<in_vec_size; i+=colgroup) {
-    // Load the vec to shared memory
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-    if (simd_gid == 0) {
-      #pragma clang loop unroll(full)
-      for (int j=0; j<el_per_thread; j++) {
-        x_block[simd_lid * el_per_thread + j] = x[i + simd_lid * el_per_thread + j];
-      }
-    }
-    if (simd_lid == 0) {
-      #pragma clang loop unroll(full)
-      for (int j=0; j<groups_per_block; j++) {
-        scales_block[simd_gid * groups_per_block + j] = scales[i / group_size + j];
-      }
-      #pragma clang loop unroll(full)
-      for (int j=0; j<groups_per_block; j++) {
-        biases_block[simd_gid * groups_per_block + j] = biases[i / group_size + j];
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
+  // In this case we need to properly guard all our reads because there isn't
+  // even 1 tile in the matrix
+  if (out_vec_size < (num_simdgroups * results_per_simdgroup)) {
+    w += out_row * in_vec_size_w + simd_lid * packs_per_thread;
+    scales += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    biases += out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.z * in_vec_size + simd_lid * values_per_thread;
+    y += tid.z * out_vec_size + out_row;
 
-    // Load in_vec, scale, bias to registers
-    #pragma clang loop unroll(full)
-    for (int j=0; j<el_per_thread; j++) {
-      x_thread[j] = x_block[simd_lid*el_per_thread + j];
+    for (int k = 0; k < in_vec_size; k += block_size) {
+      U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+      for (int row = 0; out_row + row < out_vec_size; row++) {
+        const device uint8_t* wl = (const device uint8_t *)(w + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+      }
+
+      w += block_size / pack_factor;
+      scales += block_size / group_size;
+      biases += block_size / group_size;
+      x += block_size;
     }
-    scale = scales_block[simd_gid * groups_per_block + simd_lid * el_per_thread / group_size];
-    bias = biases_block[simd_gid * groups_per_block + simd_lid * el_per_thread / group_size];
 
-    // Load the matrix elements
-    w_local = w[i / el_per_thread + simd_lid];
-
-    // Do all the work.
-    #pragma clang loop unroll(full)
-    for (int k=0; k<el_per_thread; k++) {
-      result += (scale * static_cast<U>(w_local & bitmask) + bias) * x_thread[k];
-      w_local >>= bits;
+    for (int row = 0; out_row + row < out_vec_size; row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
     }
   }
 
-  // Accumulate in the simdgroup
-  result = simd_sum(result);
+  // In this case the last tile is moved back to redo some output values
+  else {
+    w += used_out_row * in_vec_size_w + simd_lid * packs_per_thread;
+    scales += used_out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    biases += used_out_row * in_vec_size_g + simd_lid / scale_step_per_thread;
+    x += tid.z * in_vec_size + simd_lid * values_per_thread;
+    y += tid.z * out_vec_size + used_out_row;
 
-  // Store the result
-  if (simd_lid == 0) {
-    y[out_row] = static_cast<T>(result);
+    for (int k = 0; k < in_vec_size; k += block_size) {
+      U sum = load_vector<T, U, values_per_thread, bits>(x, x_thread);
+
+      for (int row = 0; row < results_per_simdgroup; row++) {
+        const device uint8_t* wl = (const device uint8_t *)(w + row * in_vec_size_w);
+        const device T* sl = scales + row * in_vec_size_g;
+        const device T* bl = biases + row * in_vec_size_g;
+
+        U s = sl[0];
+        U b = bl[0];
+        result[row] += qdot<U, values_per_thread, bits>(wl, x_thread, s, b, sum);
+      }
+
+      w += block_size / pack_factor;
+      scales += block_size / group_size;
+      biases += block_size / group_size;
+      x += block_size;
+    }
+
+    for (int row = 0; row < results_per_simdgroup; row++) {
+      result[row] = simd_sum(result[row]);
+      if (simd_lid == 0) {
+        y[row] = static_cast<T>(result[row]);
+      }
+    }
   }
 }
 
@@ -532,9 +678,38 @@ template <typename T, const int BM, const int BK, const int BN, const int group_
 }
 
 
-#define instantiate_qmv(name, itype, group_size, bits) \
-  template [[host_name("qmv_" #name "_gs_" #group_size "_b_" #bits)]] \
-  [[kernel]] void qmv<itype, 32, 32, group_size, bits>( \
+#define instantiate_qmv_fast(name, itype, group_size, bits, packs_per_thread) \
+  template [[host_name("qmv_" #name "_gs_" #group_size "_b_" #bits "_fast")]] \
+  [[kernel]] void qmv_fast<itype, group_size, bits, packs_per_thread>( \
+    const device uint32_t* w [[buffer(0)]], \
+    const device itype* scales [[buffer(1)]], \
+    const device itype* biases [[buffer(2)]], \
+    const device itype* x [[buffer(3)]], \
+    device itype* y [[buffer(4)]], \
+    const constant int& in_vec_size [[buffer(5)]], \
+    const constant int& out_vec_size [[buffer(6)]], \
+    uint3 tid [[threadgroup_position_in_grid]], \
+    uint simd_gid [[simdgroup_index_in_threadgroup]], \
+    uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_qmv_fast_types(group_size, bits, packs_per_thread) \
+  instantiate_qmv_fast(float32, float, group_size, bits, packs_per_thread) \
+  instantiate_qmv_fast(float16, half, group_size, bits, packs_per_thread) \
+  instantiate_qmv_fast(bfloat16, bfloat16_t, group_size, bits, packs_per_thread)
+
+instantiate_qmv_fast_types(128, 2, 1)
+instantiate_qmv_fast_types(128, 4, 2)
+instantiate_qmv_fast_types(128, 8, 2)
+instantiate_qmv_fast_types( 64, 2, 1)
+instantiate_qmv_fast_types( 64, 4, 2)
+instantiate_qmv_fast_types( 64, 8, 2)
+instantiate_qmv_fast_types( 32, 2, 1)
+instantiate_qmv_fast_types( 32, 4, 2)
+instantiate_qmv_fast_types( 32, 8, 2)
+
+#define instantiate_qmv(name, itype, group_size, bits) \
+  template [[host_name("qmv_" #name "_gs_" #group_size "_b_" #bits)]] \
+  [[kernel]] void qmv<itype, group_size, bits>( \
     const device uint32_t* w [[buffer(0)]], \
     const device itype* scales [[buffer(1)]], \
     const device itype* biases [[buffer(2)]], \
@@ -543,7 +718,6 @@ template <typename T, const int BM, const int BK, const int BN, const int group_
     const constant int& in_vec_size [[buffer(5)]], \
     const constant int& out_vec_size [[buffer(6)]], \
     uint3 tid [[threadgroup_position_in_grid]], \
-    uint lid [[thread_index_in_threadgroup]], \
     uint simd_gid [[simdgroup_index_in_threadgroup]], \
     uint simd_lid [[thread_index_in_simdgroup]]);
 

mlx/backend/metal/kernels/reduce.metal

@@ -1,619 +0,0 @@
-// Copyright © 2023 Apple Inc.
-
-#include <metal_atomic>
-#include <metal_simdgroup>
-
-#include "mlx/backend/metal/kernels/defines.h"
-#include "mlx/backend/metal/kernels/reduce.h"
-#include "mlx/backend/metal/kernels/utils.h"
-
-using namespace metal;
-
-static constant uint8_t simd_size = 32;
-
-template <typename T, typename Op>
-[[kernel]] void init_reduce(
-    device T *out [[buffer(0)]],
-    uint tid [[thread_position_in_grid]]) {
-  out[tid] = Op::init;
-}
-
-#define instantiate_init_reduce(name, otype, op) \
-  template [[host_name("i" #name)]] \
-    [[kernel]] void init_reduce<otype, op>( \
-      device otype *out [[buffer(1)]], \
-      uint tid [[thread_position_in_grid]]);
-
-///////////////////////////////////////////////////////////////////////////////
-// All reduce
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-inline U per_thread_all_reduce(
-    const device T *in,
-    const device size_t& in_size,
-    uint gid,
-    uint grid_size) {
-  Op op;
-  U total_val = Op::init;
-
-  if (gid * N_READS < in_size) {
-    in += gid * N_READS;
-
-    int r = 0;
-    for(; r < (int)ceildiv(in_size, grid_size * N_READS) - 1; r++) {
-      U vals[N_READS] = {op.init};
-
-      for(int i = 0; i < N_READS; i++) {
-        vals[i] = static_cast<U>(in[i]);
-      }
-      for(int i = 0; i < N_READS; i++) {
-        total_val = op(vals[i], total_val);
-      }
-
-      in += grid_size * N_READS;
-    }
-
-    // Separate case for the last set as we close the reduction size
-    size_t curr_idx = (gid + r * (size_t)grid_size) * N_READS;
-    if (curr_idx < in_size) {
-      int max_reads = in_size - curr_idx;
-      T vals[N_READS];
-
-      for(int i = 0, idx = 0; i < N_READS; i++, idx++) {
-        idx = idx < max_reads ? idx : max_reads - 1;
-        vals[i] = in[idx];
-      }
-      for(int i = 0; i < N_READS; i++) {
-        U val = i < max_reads ? vals[i] : Op::init;
-        total_val = op(static_cast<U>(val), total_val);
-      }
-    }
-  }
-
-  return total_val;
-}
-
-// NB: This kernel assumes threads_per_threadgroup is at most
-// 1024. This way with a simd_size of 32, we are guaranteed to
-// complete the reduction in two steps of simd-level reductions.
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-[[kernel]] void all_reduce(
-    const device T *in [[buffer(0)]],
-    device mlx_atomic<U> *out [[buffer(1)]],
-    const device size_t& in_size [[buffer(2)]],
-    uint gid [[thread_position_in_grid]],
-    uint lid [[thread_position_in_threadgroup]],
-    uint grid_size [[threads_per_grid]],
-    uint simd_per_group [[simdgroups_per_threadgroup]],
-    uint simd_lane_id [[thread_index_in_simdgroup]],
-    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-
-  Op op;
-  threadgroup U local_vals[simd_size];
-
-  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
-
-  // Reduction within simd group
-  total_val = op.simd_reduce(total_val);
-  if (simd_lane_id == 0) {
-    local_vals[simd_group_id] = total_val;
-  }
-
-  // Reduction within thread group
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  total_val = lid < simd_per_group ? local_vals[lid] : op.init;
-  total_val = op.simd_reduce(total_val);
-
-  // Reduction across threadgroups
-  if (lid == 0) {
-    op.atomic_update(out, total_val);
-  }
-}
-
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-[[kernel]] void all_reduce_no_atomics(
-    const device T *in [[buffer(0)]],
-    device U *out [[buffer(1)]],
-    const device size_t& in_size [[buffer(2)]],
-    uint gid [[thread_position_in_grid]],
-    uint lid [[thread_position_in_threadgroup]],
-    uint grid_size [[threads_per_grid]],
-    uint simd_per_group [[simdgroups_per_threadgroup]],
-    uint simd_lane_id [[thread_index_in_simdgroup]],
-    uint simd_group_id [[simdgroup_index_in_threadgroup]],
-    uint thread_group_id [[threadgroup_position_in_grid]]) {
-
-  Op op;
-  threadgroup U local_vals[simd_size];
-
-  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
-
-  // Reduction within simd group (simd_add isn't supported for uint64/int64 types)
-  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
-    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
-  }
-  // Write simd group reduction results to local memory
-  if (simd_lane_id == 0) {
-    local_vals[simd_group_id] = total_val;
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  // Reduction of simdgroup reduction results within threadgroup.
-  total_val = lid < simd_per_group ? local_vals[lid] : op.init;
-  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
-    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
-  }
-
-  // Reduction across threadgroups
-  if (lid == 0) {
-    out[thread_group_id] = total_val;
-  }
-}
-
-#define instantiate_all_reduce(name, itype, otype, op) \
-  template [[host_name("all_reduce_" #name)]] \
-  [[kernel]] void all_reduce<itype, otype, op>( \
-      const device itype *in [[buffer(0)]], \
-      device mlx_atomic<otype> *out [[buffer(1)]], \
-      const device size_t& in_size [[buffer(2)]], \
-      uint gid [[thread_position_in_grid]], \
-      uint lid [[thread_position_in_threadgroup]], \
-      uint grid_size [[threads_per_grid]], \
-      uint simd_per_group [[simdgroups_per_threadgroup]], \
-      uint simd_lane_id [[thread_index_in_simdgroup]], \
-      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
-
-#define instantiate_all_reduce_no_atomics(name, itype, otype, op) \
-  template [[host_name("all_reduce_no_atomics_" #name)]] \
-  [[kernel]] void all_reduce_no_atomics<itype, otype, op>( \
-      const device itype *in [[buffer(0)]], \
-      device otype *out [[buffer(1)]], \
-      const device size_t& in_size [[buffer(2)]], \
-      uint gid [[thread_position_in_grid]], \
-      uint lid [[thread_position_in_threadgroup]], \
-      uint grid_size [[threads_per_grid]], \
-      uint simd_per_group [[simdgroups_per_threadgroup]], \
-      uint simd_lane_id [[thread_index_in_simdgroup]], \
-      uint simd_group_id [[simdgroup_index_in_threadgroup]], \
-      uint thread_group_id [[threadgroup_position_in_grid]]);
-
-///////////////////////////////////////////////////////////////////////////////
-// Row atomics
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-inline U per_thread_row_reduce(
-    const device T *in,
-    const constant size_t& reduction_size,
-    const constant size_t& out_size,
-    const constant int* shape,
-    const constant size_t* strides,
-    const constant int& ndim,
-    uint lsize_x,
-    uint lid_x,
-    uint2 tid) {
-
-  Op op;
-
-  // Each threadgroup handles 1 reduction
-  // TODO: Specializing elem_to_loc would be slightly faster
-  int idx = tid.y * out_size + tid.x;
-  int extra_offset = elem_to_loc(idx, shape, strides, ndim);
-  in += extra_offset + lid_x * N_READS;
-  
-  // The reduction is accumulated here
-  U total_val = Op::init;
-
-  // Loop over the reduction size within thread group
-  int r = 0;
-  for (; r < (int)ceildiv(reduction_size, N_READS*lsize_x) - 1; r++) {
-    T vals[N_READS]; 
-    for(int i = 0; i < N_READS; i++) {
-      vals[i] = in[i];
-    }
-    for(int i = 0; i < N_READS; i++) {
-      total_val = op(static_cast<U>(vals[i]), total_val);
-    }
-
-    in += lsize_x * N_READS;
-  }
-
-  // Separate case for the last set as we close the reduction size   
-  size_t reduction_index = (lid_x + (size_t)lsize_x * r) * N_READS;
-  if(reduction_index < reduction_size) {
-    int max_reads = reduction_size - reduction_index;
-
-    T vals[N_READS]; 
-    for(int i = 0; i < N_READS; i++) {
-      int idx = min(i, max_reads - 1);
-      vals[i] = static_cast<U>(in[idx]);
-    }
-    for(int i = 0; i < N_READS; i++) {
-      T val = i < max_reads ? vals[i] : Op::init;
-      total_val = op(static_cast<U>(val), total_val);
-    }
-  }
-
-  return total_val;
-}
-
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-[[kernel]] void row_reduce_general(
-    const device T *in [[buffer(0)]],
-    device mlx_atomic<U> *out [[buffer(1)]],
-    const constant size_t& reduction_size [[buffer(2)]],
-    const constant size_t& out_size [[buffer(3)]],
-    const constant int* shape [[buffer(4)]],
-    const constant size_t* strides [[buffer(5)]],
-    const constant int& ndim [[buffer(6)]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 lsize [[threads_per_threadgroup]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint simd_lane_id [[thread_index_in_simdgroup]],
-    uint simd_per_group [[simdgroups_per_threadgroup]],
-    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-
-  Op op;
-  threadgroup U local_vals[simd_size];
-
-  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
-
-  total_val = op.simd_reduce(total_val);
-  
-  // Prepare next level
-  if (simd_lane_id == 0) {
-    local_vals[simd_group_id] = total_val;
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-      
-  // Reduction within thread group
-  //    Only needed if multiple simd groups
-  if(reduction_size > simd_size) {
-    total_val = lid.x < simd_per_group ? local_vals[lid.x] : op.init;
-    total_val = op.simd_reduce(total_val);
-  }
-  // Update output
-  if (lid.x == 0) {
-    op.atomic_update(out, total_val, tid.x);
-  }
-}
-
-template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
-[[kernel]] void row_reduce_general_no_atomics(
-    const device T *in [[buffer(0)]],
-    device U *out [[buffer(1)]],
-    const constant size_t& reduction_size [[buffer(2)]],
-    const constant size_t& out_size [[buffer(3)]],
-    const constant int* shape [[buffer(4)]],
-    const constant size_t* strides [[buffer(5)]],
-    const constant int& ndim [[buffer(6)]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 lsize [[threads_per_threadgroup]],
-    uint3 gsize [[threads_per_grid]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint simd_lane_id [[thread_index_in_simdgroup]],
-    uint simd_per_group [[simdgroups_per_threadgroup]],
-    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-
-  Op op;
-
-  threadgroup U local_vals[simd_size];
-  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
-
-  // Reduction within simd group - simd_add isn't supported for int64 types
-  for (uint16_t i = simd_size/2; i > 0; i /= 2) {
-    total_val = op(total_val, simd_shuffle_down(total_val, i));
-  }
-
-  // Prepare next level
-  if (simd_lane_id == 0) {
-    local_vals[simd_group_id] = total_val;
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  // Reduction within thread group
-  // Only needed if thread group has multiple simd groups
-  if(ceildiv(reduction_size, N_READS) > simd_size) {
-    total_val = lid.x < simd_per_group ? local_vals[lid.x] : op.init;
-    for (uint16_t i = simd_size/2; i > 0; i /= 2) {
-      total_val = op(total_val, simd_shuffle_down(total_val, i));
-    }
-  }
-  // Write row reduce output for threadgroup with 1st thread in thread group
-  if (lid.x == 0) {
-    out[(ceildiv(gsize.y, lsize.y) * tid.x) + tid.y] = total_val;
-  }
-}
-
-#define instantiate_row_reduce_general(name, itype, otype, op) \
-  template [[host_name("row_reduce_general_" #name)]] \
-  [[kernel]] void row_reduce_general<itype, otype, op>( \
-      const device itype *in [[buffer(0)]],  \
-      device mlx_atomic<otype> *out [[buffer(1)]],  \
-      const constant size_t& reduction_size [[buffer(2)]],  \
-      const constant size_t& out_size [[buffer(3)]],  \
-      const constant int* shape [[buffer(4)]],  \
-      const constant size_t* strides [[buffer(5)]],  \
-      const constant int& ndim [[buffer(6)]],  \
-      uint3 lid [[thread_position_in_threadgroup]],  \
-      uint3 lsize [[threads_per_threadgroup]],  \
-      uint3 tid [[threadgroup_position_in_grid]],  \
-      uint simd_lane_id [[thread_index_in_simdgroup]],  \
-      uint simd_per_group [[simdgroups_per_threadgroup]],  \
-      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
-
-#define instantiate_row_reduce_general_no_atomics(name, itype, otype, op) \
-  template [[host_name("row_reduce_general_no_atomics_" #name)]] \
-  [[kernel]] void row_reduce_general_no_atomics<itype, otype, op>( \
-      const device itype *in [[buffer(0)]],  \
-      device otype *out [[buffer(1)]],  \
-      const constant size_t& reduction_size [[buffer(2)]],  \
-      const constant size_t& out_size [[buffer(3)]],  \
-      const constant int* shape [[buffer(4)]],  \
-      const constant size_t* strides [[buffer(5)]],  \
-      const constant int& ndim [[buffer(6)]],  \
-      uint3 lid [[thread_position_in_threadgroup]],  \
-      uint3 lsize [[threads_per_threadgroup]],  \
-      uint3 gsize [[threads_per_grid]], \
-      uint3 tid [[threadgroup_position_in_grid]],  \
-      uint simd_lane_id [[thread_index_in_simdgroup]],  \
-      uint simd_per_group [[simdgroups_per_threadgroup]],  \
-      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
-
-///////////////////////////////////////////////////////////////////////////////
-// Column reduce
-///////////////////////////////////////////////////////////////////////////////
-
-template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
-inline U _contiguous_strided_reduce(
-    const device T *in,
-    threadgroup U *local_data,
-    uint in_idx,
-    uint reduction_size,
-    uint reduction_stride,
-    uint2 tid,
-    uint2 lid,
-    uint2 lsize) {
-
-  Op op;
-  U total_val = Op::init;
-
-  uint base_offset = (tid.y * lsize.y + lid.y) * N_READS;
-  for(uint r = 0; r < N_READS && (base_offset + r) < reduction_size; r++) {
-    uint offset = base_offset + r;
-    total_val = op(static_cast<U>(total_val), in[in_idx + offset * reduction_stride]);
-  }
-  local_data[lsize.y * lid.x + lid.y] = total_val;
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-
-  U val = Op::init;
-  if(lid.y == 0) {
-    // Perform reduction across columns in thread group
-    for(uint i = 0; i < lsize.y; i++) {
-      val = op(val, local_data[lsize.y * lid.x + i]);
-    }
-  }
-
-  return val;
-}
-
-template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
-[[kernel]] void col_reduce_general(
-    const device T *in [[buffer(0)]],
-    device mlx_atomic<U> *out [[buffer(1)]],
-    const constant size_t& reduction_size [[buffer(2)]],
-    const constant size_t& reduction_stride [[buffer(3)]],
-    const constant size_t& out_size [[buffer(4)]],
-    const constant int* shape [[buffer(5)]],
-    const constant size_t* strides [[buffer(6)]],
-    const constant int& ndim [[buffer(7)]],
-    threadgroup U *local_data [[threadgroup(0)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 lsize [[threads_per_threadgroup]]) {
-  auto out_idx = tid.x * lsize.x + lid.x;
-  auto in_idx = elem_to_loc(
-    out_idx + tid.z * out_size,
-    shape,
-    strides,
-    ndim
-  );
-
-  Op op;
-  if(out_idx < out_size) {
-    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
-              in,
-              local_data,
-              in_idx,
-              reduction_size,
-              reduction_stride,
-              tid.xy,
-              lid.xy,
-              lsize.xy);
-
-    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
-    if (lid.y == 0) {
-      op.atomic_update(out, val, out_idx);
-    }
-  }
-}
-
-template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
-[[kernel]] void col_reduce_general_no_atomics(
-    const device T *in [[buffer(0)]],
-    device U *out [[buffer(1)]],
-    const constant size_t& reduction_size [[buffer(2)]],
-    const constant size_t& reduction_stride [[buffer(3)]],
-    const constant size_t& out_size [[buffer(4)]],
-    const constant int* shape [[buffer(5)]],
-    const constant size_t* strides [[buffer(6)]],
-    const constant int& ndim [[buffer(7)]],
-    threadgroup U *local_data [[threadgroup(0)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 gid [[thread_position_in_grid]],
-    uint3 lsize [[threads_per_threadgroup]],
-    uint3 gsize [[threads_per_grid]]) {
-  auto out_idx = tid.x * lsize.x + lid.x;
-  auto in_idx = elem_to_loc(
-    out_idx + tid.z * out_size,
-    shape,
-    strides,
-    ndim
-  );
-
-  if(out_idx < out_size) {
-    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
-              in,
-              local_data,
-              in_idx,
-              reduction_size,
-              reduction_stride,
-              tid.xy,
-              lid.xy,
-              lsize.xy);
-
-    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
-    if (lid.y == 0) {
-      uint tgsize_y = ceildiv(gsize.y, lsize.y);
-      uint tgsize_z = ceildiv(gsize.z, lsize.z);
-      out[tgsize_y * tgsize_z * gid.x + tgsize_y * tid.z + tid.y] = val;
-    }
-  }
-}
-
-#define instantiate_col_reduce_general(name, itype, otype, op) \
-  template [[host_name("col_reduce_general_" #name)]] \
-  [[kernel]] void col_reduce_general<itype, otype, op>( \
-      const device itype *in [[buffer(0)]], \
-      device mlx_atomic<otype> *out [[buffer(1)]], \
-      const constant size_t& reduction_size [[buffer(2)]], \
-      const constant size_t& reduction_stride [[buffer(3)]], \
-      const constant size_t& out_size [[buffer(4)]], \
-      const constant int* shape [[buffer(5)]],  \
-      const constant size_t* strides [[buffer(6)]],  \
-      const constant int& ndim [[buffer(7)]],  \
-      threadgroup otype *local_data [[threadgroup(0)]], \
-      uint3 tid [[threadgroup_position_in_grid]], \
-      uint3 lid [[thread_position_in_threadgroup]], \
-      uint3 lsize [[threads_per_threadgroup]]);
-
-#define instantiate_col_reduce_general_no_atomics(name, itype, otype, op) \
-  template [[host_name("col_reduce_general_no_atomics_" #name)]] \
-  [[kernel]] void col_reduce_general_no_atomics<itype, otype, op>( \
-      const device itype *in [[buffer(0)]], \
-      device otype *out [[buffer(1)]], \
-      const constant size_t& reduction_size [[buffer(2)]], \
-      const constant size_t& reduction_stride [[buffer(3)]], \
-      const constant size_t& out_size [[buffer(4)]], \
-      const constant int* shape [[buffer(5)]],  \
-      const constant size_t* strides [[buffer(6)]],  \
-      const constant int& ndim [[buffer(7)]],  \
-      threadgroup otype *local_data [[threadgroup(0)]], \
-      uint3 tid [[threadgroup_position_in_grid]], \
-      uint3 lid [[thread_position_in_threadgroup]], \
-      uint3 gid [[thread_position_in_grid]], \
-      uint3 lsize [[threads_per_threadgroup]], \
-      uint3 gsize [[threads_per_grid]]);
-
-///////////////////////////////////////////////////////////////////////////////
-// Instantiations
-///////////////////////////////////////////////////////////////////////////////
-
-#define instantiate_reduce(name, itype, otype, op) \
-  instantiate_all_reduce(name, itype, otype, op) \
-  instantiate_row_reduce_general(name, itype, otype, op) \
-  instantiate_col_reduce_general(name, itype, otype, op)
-
-#define instantiate_reduce_no_atomics(name, itype, otype, op) \
-  instantiate_all_reduce_no_atomics(name, itype, otype, op) \
-  instantiate_row_reduce_general_no_atomics(name, itype, otype, op) \
-  instantiate_col_reduce_general_no_atomics(name, itype, otype, op)
-
-#define instantiate_same_reduce_no_atomics(name, tname, type, op) \
-  instantiate_init_reduce(name ##tname, type, op<type>) \
-  instantiate_reduce_no_atomics(name ##tname, type, type, op<type>)
-
-#define instantiate_same_reduce(name, tname, type, op) \
-  instantiate_init_reduce(name ##tname, type, op<type>) \
-  instantiate_reduce(name ##tname, type, type, op<type>)
-
-#define instantiate_reduce_from_types_helper(name, tname, itype, otype, op) \
-  instantiate_reduce(name ##tname, itype, otype, op)
-
-#define instantiate_reduce_from_types(name, otype, op) \
-  instantiate_reduce_from_types_helper(name, bool_, bool, otype, op) \
-  instantiate_reduce_from_types_helper(name, uint8, uint8_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, uint16, uint16_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, uint32, uint32_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, int8, int8_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, int16, int16_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, int32, int32_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, int64, int64_t, otype, op) \
-  instantiate_reduce_from_types_helper(name, float16, half, otype, op) \
-  instantiate_reduce_from_types_helper(name, float32, float, otype, op) \
-  instantiate_reduce_from_types_helper(name, bfloat16, bfloat16_t, otype, op)
-
-// special case bool with larger output type
-instantiate_reduce(sumbool_, bool, uint32_t, Sum<uint32_t>)
-instantiate_same_reduce(sum, uint8, uint8_t, Sum)
-instantiate_same_reduce(sum, uint16, uint16_t, Sum)
-instantiate_same_reduce(sum, uint32, uint32_t, Sum)
-instantiate_same_reduce(sum, int8, int8_t, Sum)
-instantiate_same_reduce(sum, int16, int16_t, Sum)
-instantiate_same_reduce(sum, int32, int32_t, Sum)
-instantiate_same_reduce(sum, float16, half, Sum)
-instantiate_same_reduce(sum, float32, float, Sum)
-
-instantiate_same_reduce_no_atomics(sum, int64, int64_t, Sum)
-instantiate_same_reduce_no_atomics(sum, uint64, uint64_t, Sum)
-
-instantiate_same_reduce(prod, uint8, uint8_t, Prod)
-instantiate_same_reduce(prod, uint16, uint16_t, Prod)
-instantiate_same_reduce(prod, uint32, uint32_t, Prod)
-instantiate_same_reduce(prod, int8, int8_t, Prod)
-instantiate_same_reduce(prod, int16, int16_t, Prod)
-instantiate_same_reduce(prod, int32, int32_t, Prod)
-instantiate_same_reduce(prod, float16, half, Prod)
-instantiate_same_reduce(prod, float32, float, Prod)
-
-instantiate_same_reduce_no_atomics(prod, int64, int64_t, Prod)
-instantiate_same_reduce_no_atomics(prod, uint64, uint64_t, Prod)
-
-instantiate_same_reduce(sum, bfloat16, bfloat16_t, Sum)
-instantiate_same_reduce(prod, bfloat16, bfloat16_t, Prod)
-
-instantiate_init_reduce(andbool_, bool, And)
-instantiate_reduce_from_types(and, bool, And)
-
-instantiate_init_reduce(orbool_, bool, Or)
-instantiate_reduce_from_types(or, bool, Or)
-
-// Compiler segfaulted with the names "min" or "max" ...
-instantiate_same_reduce(min_, uint8, uint8_t, Min)
-instantiate_same_reduce(min_, uint16, uint16_t, Min)
-instantiate_same_reduce(min_, uint32, uint32_t, Min)
-instantiate_same_reduce(min_, int8, int8_t, Min)
-instantiate_same_reduce(min_, int16, int16_t, Min)
-instantiate_same_reduce(min_, int32, int32_t, Min)
-instantiate_same_reduce(min_, float16, half, Min)
-instantiate_same_reduce(min_, float32, float, Min)
-
-instantiate_same_reduce_no_atomics(min_, int64, int64_t, Min)
-instantiate_same_reduce_no_atomics(min_, uint64, uint64_t, Min)
-
-instantiate_same_reduce(max_, uint8, uint8_t, Max)
-instantiate_same_reduce(max_, uint16, uint16_t, Max)
-instantiate_same_reduce(max_, uint32, uint32_t, Max)
-instantiate_same_reduce(max_, int8, int8_t, Max)
-instantiate_same_reduce(max_, int16, int16_t, Max)
-instantiate_same_reduce(max_, int32, int32_t, Max)
-instantiate_same_reduce(max_, float16, half, Max)
-instantiate_same_reduce(max_, float32, float, Max)
-
-instantiate_same_reduce_no_atomics(max_, int64, int64_t, Max)
-instantiate_same_reduce_no_atomics(max_, uint64, uint64_t, Max)
-
-instantiate_same_reduce(min_, bfloat16, bfloat16_t, Min)
-instantiate_same_reduce(max_, bfloat16, bfloat16_t, Max)

mlx/backend/metal/kernels/reduction/kernels/reduce_all.metal

@@ -0,0 +1,185 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/reduction/utils.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_inst.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// All reduce helper
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+METAL_FUNC U per_thread_all_reduce(
+    const device T* in,
+    const device size_t& in_size,
+    uint gid,
+    uint grid_size) {
+  Op op;
+  U total_val = Op::init;
+
+  if (gid * N_READS < in_size) {
+    in += gid * N_READS;
+
+    int r = 0;
+    for (; r < (int)ceildiv(in_size, grid_size * N_READS) - 1; r++) {
+      U vals[N_READS] = {op.init};
+
+      for (int i = 0; i < N_READS; i++) {
+        vals[i] = static_cast<U>(in[i]);
+      }
+      for (int i = 0; i < N_READS; i++) {
+        total_val = op(vals[i], total_val);
+      }
+
+      in += grid_size * N_READS;
+    }
+
+    // Separate case for the last set as we close the reduction size
+    size_t curr_idx = (gid + r * (size_t)grid_size) * N_READS;
+    if (curr_idx < in_size) {
+      int max_reads = in_size - curr_idx;
+      T vals[N_READS];
+
+      for (int i = 0, idx = 0; i < N_READS; i++, idx++) {
+        idx = idx < max_reads ? idx : max_reads - 1;
+        vals[i] = in[idx];
+      }
+      for (int i = 0; i < N_READS; i++) {
+        U val = i < max_reads ? vals[i] : Op::init;
+        total_val = op(static_cast<U>(val), total_val);
+      }
+    }
+  }
+
+  return total_val;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// All reduce kernel
+///////////////////////////////////////////////////////////////////////////////
+
+
+// NB: This kernel assumes threads_per_threadgroup is at most
+// 1024. This way with a simd_size of 32, we are guaranteed to
+// complete the reduction in two steps of simd-level reductions.
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void all_reduce(
+    const device T *in [[buffer(0)]],
+    device mlx_atomic<U> *out [[buffer(1)]],
+    const device size_t& in_size [[buffer(2)]],
+    uint gid [[thread_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint grid_size [[threads_per_grid]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+
+  Op op;
+  threadgroup U local_vals[simd_size];
+
+  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
+
+  // Reduction within simd group
+  total_val = op.simd_reduce(total_val);
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+
+  // Reduction within thread group
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  total_val = lid < simd_per_group ? local_vals[lid] : op.init;
+  total_val = op.simd_reduce(total_val);
+
+  // Reduction across threadgroups
+  if (lid == 0) {
+    op.atomic_update(out, total_val);
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void all_reduce_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const device size_t& in_size [[buffer(2)]],
+    uint gid [[thread_position_in_grid]],
+    uint lid [[thread_position_in_threadgroup]],
+    uint grid_size [[threads_per_grid]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint thread_group_id [[threadgroup_position_in_grid]]) {
+
+  Op op;
+  threadgroup U local_vals[simd_size];
+
+  U total_val = per_thread_all_reduce<T, U, Op, N_READS>(in, in_size, gid, grid_size);
+
+  // Reduction within simd group (simd_add isn't supported for uint64/int64 types)
+  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
+  }
+  // Write simd group reduction results to local memory
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Reduction of simdgroup reduction results within threadgroup.
+  total_val = lid < simd_per_group ? local_vals[lid] : op.init;
+  for (uint16_t lane_offset = simd_size/2; lane_offset > 0; lane_offset /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, lane_offset));
+  }
+
+  // Reduction across threadgroups
+  if (lid == 0) {
+    out[thread_group_id] = total_val;
+  }
+}
+
+#define instantiate_all_reduce(name, itype, otype, op) \
+  template [[host_name("all_reduce_" #name)]] \
+  [[kernel]] void all_reduce<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device mlx_atomic<otype> *out [[buffer(1)]], \
+      const device size_t& in_size [[buffer(2)]], \
+      uint gid [[thread_position_in_grid]], \
+      uint lid [[thread_position_in_threadgroup]], \
+      uint grid_size [[threads_per_grid]], \
+      uint simd_per_group [[simdgroups_per_threadgroup]], \
+      uint simd_lane_id [[thread_index_in_simdgroup]], \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_all_reduce_no_atomics(name, itype, otype, op) \
+  template [[host_name("all_reduce_no_atomics_" #name)]] \
+  [[kernel]] void all_reduce_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const device size_t& in_size [[buffer(2)]], \
+      uint gid [[thread_position_in_grid]], \
+      uint lid [[thread_position_in_threadgroup]], \
+      uint grid_size [[threads_per_grid]], \
+      uint simd_per_group [[simdgroups_per_threadgroup]], \
+      uint simd_lane_id [[thread_index_in_simdgroup]], \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]], \
+      uint thread_group_id [[threadgroup_position_in_grid]]);
+
+///////////////////////////////////////////////////////////////////////////////
+// Instantiations
+///////////////////////////////////////////////////////////////////////////////
+
+#define instantiate_same_all_reduce_helper(name, tname, type, op) \
+  instantiate_all_reduce(name ##tname, type, type, op<type>)
+
+#define instantiate_same_all_reduce_na_helper(name, tname, type, op) \
+  instantiate_all_reduce_no_atomics(name ##tname, type, type, op<type>)
+
+instantiate_reduce_ops(instantiate_same_all_reduce_helper, instantiate_reduce_helper_types)
+instantiate_reduce_ops(instantiate_same_all_reduce_na_helper, instantiate_reduce_helper_64b)
+
+instantiate_reduce_from_types(instantiate_all_reduce, and, bool, And)
+instantiate_reduce_from_types(instantiate_all_reduce, or, bool, Or)
+
+// special case bool with larger output type
+instantiate_all_reduce(sumbool_, bool, uint32_t, Sum<uint32_t>)

mlx/backend/metal/kernels/reduction/kernels/reduce_col.metal

@@ -0,0 +1,184 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/reduction/utils.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_inst.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// Column reduce helper
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+METAL_FUNC U _contiguous_strided_reduce(
+    const device T* in,
+    threadgroup U* local_data,
+    uint in_idx,
+    uint reduction_size,
+    uint reduction_stride,
+    uint2 tid,
+    uint2 lid,
+    uint2 lsize) {
+  Op op;
+  U total_val = Op::init;
+
+  uint base_offset = (tid.y * lsize.y + lid.y) * N_READS;
+  for (uint r = 0; r < N_READS && (base_offset + r) < reduction_size; r++) {
+    uint offset = base_offset + r;
+    total_val =
+        op(static_cast<U>(total_val), in[in_idx + offset * reduction_stride]);
+  }
+  local_data[lsize.y * lid.x + lid.y] = total_val;
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  U val = Op::init;
+  if (lid.y == 0) {
+    // Perform reduction across columns in thread group
+    for (uint i = 0; i < lsize.y; i++) {
+      val = op(val, local_data[lsize.y * lid.x + i]);
+    }
+  }
+
+  return val;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Column reduce kernel
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+[[kernel]] void col_reduce_general(
+    const device T *in [[buffer(0)]],
+    device mlx_atomic<U> *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& reduction_stride [[buffer(3)]],
+    const constant size_t& out_size [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    threadgroup U *local_data [[threadgroup(0)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]]) {
+  auto out_idx = tid.x * lsize.x + lid.x;
+  auto in_idx = elem_to_loc(
+    out_idx + tid.z * out_size,
+    shape,
+    strides,
+    ndim
+  );
+
+  Op op;
+  if(out_idx < out_size) {
+    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
+              in,
+              local_data,
+              in_idx,
+              reduction_size,
+              reduction_stride,
+              tid.xy,
+              lid.xy,
+              lsize.xy);
+
+    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
+    if (lid.y == 0) {
+      op.atomic_update(out, val, out_idx);
+    }
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+[[kernel]] void col_reduce_general_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& reduction_stride [[buffer(3)]],
+    const constant size_t& out_size [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    threadgroup U *local_data [[threadgroup(0)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 gid [[thread_position_in_grid]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint3 gsize [[threads_per_grid]]) {
+  auto out_idx = tid.x * lsize.x + lid.x;
+  auto in_idx = elem_to_loc(
+    out_idx + tid.z * out_size,
+    shape,
+    strides,
+    ndim
+  );
+
+  if(out_idx < out_size) {
+    U val = _contiguous_strided_reduce<T, U, Op, N_READS>(
+              in,
+              local_data,
+              in_idx,
+              reduction_size,
+              reduction_stride,
+              tid.xy,
+              lid.xy,
+              lsize.xy);
+
+    // Write out reduction results generated by threadgroups working on specific output element, contiguously.
+    if (lid.y == 0) {
+      uint tgsize_y = ceildiv(gsize.y, lsize.y);
+      uint tgsize_z = ceildiv(gsize.z, lsize.z);
+      out[tgsize_y * tgsize_z * gid.x + tgsize_y * tid.z + tid.y] = val;
+    }
+  }
+}
+
+#define instantiate_col_reduce_general(name, itype, otype, op) \
+  template [[host_name("col_reduce_general_" #name)]] \
+  [[kernel]] void col_reduce_general<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device mlx_atomic<otype> *out [[buffer(1)]], \
+      const constant size_t& reduction_size [[buffer(2)]], \
+      const constant size_t& reduction_stride [[buffer(3)]], \
+      const constant size_t& out_size [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]],  \
+      const constant size_t* strides [[buffer(6)]],  \
+      const constant int& ndim [[buffer(7)]],  \
+      threadgroup otype *local_data [[threadgroup(0)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint3 lid [[thread_position_in_threadgroup]], \
+      uint3 lsize [[threads_per_threadgroup]]);
+
+#define instantiate_col_reduce_general_no_atomics(name, itype, otype, op) \
+  template [[host_name("col_reduce_general_no_atomics_" #name)]] \
+  [[kernel]] void col_reduce_general_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const constant size_t& reduction_size [[buffer(2)]], \
+      const constant size_t& reduction_stride [[buffer(3)]], \
+      const constant size_t& out_size [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]],  \
+      const constant size_t* strides [[buffer(6)]],  \
+      const constant int& ndim [[buffer(7)]],  \
+      threadgroup otype *local_data [[threadgroup(0)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint3 lid [[thread_position_in_threadgroup]], \
+      uint3 gid [[thread_position_in_grid]], \
+      uint3 lsize [[threads_per_threadgroup]], \
+      uint3 gsize [[threads_per_grid]]);
+
+///////////////////////////////////////////////////////////////////////////////
+// Instantiations
+///////////////////////////////////////////////////////////////////////////////
+
+#define instantiate_same_col_reduce_helper(name, tname, type, op) \
+  instantiate_col_reduce_general(name ##tname, type, type, op<type>)
+
+#define instantiate_same_col_reduce_na_helper(name, tname, type, op) \
+  instantiate_col_reduce_general_no_atomics(name ##tname, type, type, op<type>)
+
+instantiate_reduce_ops(instantiate_same_col_reduce_helper, instantiate_reduce_helper_types)
+instantiate_reduce_ops(instantiate_same_col_reduce_na_helper, instantiate_reduce_helper_64b)
+
+instantiate_col_reduce_general(sumbool_, bool, uint32_t, Sum<uint32_t>)
+instantiate_reduce_from_types(instantiate_col_reduce_general, and, bool, And)
+instantiate_reduce_from_types(instantiate_col_reduce_general, or, bool, Or)

mlx/backend/metal/kernels/reduction/kernels/reduce_init.metal

@@ -0,0 +1,33 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/reduction/utils.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_inst.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// Reduce init
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename Op>
+[[kernel]] void init_reduce(
+    device T *out [[buffer(0)]],
+    uint tid [[thread_position_in_grid]]) {
+  out[tid] = Op::init;
+}
+
+#define instantiate_init_reduce(name, otype, op) \
+  template [[host_name("i" #name)]] \
+    [[kernel]] void init_reduce<otype, op>( \
+      device otype *out [[buffer(1)]], \
+      uint tid [[thread_position_in_grid]]);
+
+#define instantiate_init_reduce_helper(name, tname, type, op) \
+  instantiate_init_reduce(name ##tname, type, op<type>)
+
+instantiate_reduce_ops(instantiate_init_reduce_helper, instantiate_reduce_helper_types)
+instantiate_reduce_ops(instantiate_init_reduce_helper, instantiate_reduce_helper_64b)
+
+instantiate_init_reduce(andbool_, bool, And)
+instantiate_init_reduce(orbool_, bool, Or)

mlx/backend/metal/kernels/reduction/kernels/reduce_row.metal

@@ -0,0 +1,369 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/reduction/utils.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+#include "mlx/backend/metal/kernels/reduction/reduce_inst.h"
+
+using namespace metal;
+
+///////////////////////////////////////////////////////////////////////////////
+// Small row reductions
+///////////////////////////////////////////////////////////////////////////////
+
+// Each thread reduces for one output
+template <typename T, typename U, typename Op>
+[[kernel]] void row_reduce_general_small(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& out_size [[buffer(3)]],
+    const constant size_t& non_row_reductions [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    uint lid [[thread_position_in_grid]]) {
+
+  Op op;
+  
+  uint out_idx = lid;
+
+  if(out_idx >= out_size) {
+    return;
+  }
+
+  U total_val = Op::init;
+
+  for(short r = 0; r < short(non_row_reductions); r++) {
+    uint in_idx = elem_to_loc(out_idx + r * out_size, shape, strides, ndim);
+    const device T * in_row = in + in_idx;
+    
+    for(short i = 0; i < short(reduction_size); i++) {
+      total_val = op(static_cast<U>(in_row[i]), total_val);
+    }
+  }
+
+  out[out_idx] = total_val;
+}
+
+// Each simdgroup reduces for one output
+template <typename T, typename U, typename Op>
+[[kernel]] void row_reduce_general_med(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& out_size [[buffer(3)]],
+    const constant size_t& non_row_reductions [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    uint tid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[dispatch_simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+
+  Op op;
+  
+  uint out_idx = simd_per_group * tid + simd_group_id;
+
+  if(out_idx >= out_size) {
+    return;
+  }
+
+  U total_val = Op::init;
+
+  if(short(non_row_reductions) == 1) {
+    uint in_idx = elem_to_loc(out_idx, shape, strides, ndim);
+    const device T * in_row = in + in_idx;
+
+    for(short i = simd_lane_id; i < short(reduction_size); i += 32) {
+      total_val = op(static_cast<U>(in_row[i]), total_val);
+    }
+  }
+
+  else if (short(non_row_reductions) >= 32) {
+
+    for(short r = simd_lane_id; r < short(non_row_reductions); r+=32) {
+
+      uint in_idx = elem_to_loc(out_idx + r * out_size, shape, strides, ndim);
+      const device T * in_row = in + in_idx;
+
+      for(short i = 0; i < short(reduction_size); i++) {
+        total_val = op(static_cast<U>(in_row[i]), total_val);
+      }
+
+    }
+
+  }
+
+  else {
+
+    const short n_reductions = short(reduction_size) * short(non_row_reductions);
+    const short reductions_per_thread = (n_reductions + simd_size - 1) / simd_size;
+
+    const short r_st = simd_lane_id / reductions_per_thread;
+    const short r_ed = short(non_row_reductions);
+    const short r_jump = simd_size / reductions_per_thread;
+
+    const short i_st = simd_lane_id % reductions_per_thread;
+    const short i_ed = short(reduction_size);
+    const short i_jump = reductions_per_thread;
+
+    for(short r = r_st; r < r_ed; r += r_jump) {
+
+      uint in_idx = elem_to_loc(out_idx + r * out_size, shape, strides, ndim);
+      const device T * in_row = in + in_idx;
+
+      for(short i = i_st; i < i_ed; i += i_jump) {
+        total_val = op(static_cast<U>(in_row[i]), total_val);
+      }
+
+    }
+
+  }
+
+
+  total_val = op.simd_reduce(total_val);
+
+  if(simd_lane_id == 0) {
+    out[out_idx] = total_val;
+  }
+}
+
+#define instantiate_row_reduce_small(name, itype, otype, op) \
+  template[[host_name("row_reduce_general_small_" #name)]] \
+  [[kernel]] void row_reduce_general_small<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const constant size_t& reduction_size [[buffer(2)]], \
+      const constant size_t& out_size [[buffer(3)]], \
+      const constant size_t& non_row_reductions [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]], \
+      const constant size_t* strides [[buffer(6)]], \
+      const constant int& ndim [[buffer(7)]], \
+      uint lid [[thread_position_in_grid]]); \
+  template[[host_name("row_reduce_general_med_" #name)]] \
+  [[kernel]] void row_reduce_general_med<itype, otype, op>( \
+      const device itype *in [[buffer(0)]], \
+      device otype *out [[buffer(1)]], \
+      const constant size_t& reduction_size [[buffer(2)]], \
+      const constant size_t& out_size [[buffer(3)]], \
+      const constant size_t& non_row_reductions [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]], \
+      const constant size_t* strides [[buffer(6)]], \
+      const constant int& ndim [[buffer(7)]], \
+      uint tid [[threadgroup_position_in_grid]], \
+      uint simd_lane_id [[thread_index_in_simdgroup]], \
+      uint simd_per_group [[dispatch_simdgroups_per_threadgroup]], \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+///////////////////////////////////////////////////////////////////////////////
+// Large row reductions
+///////////////////////////////////////////////////////////////////////////////
+
+template <typename T, typename U, typename Op, int N_READS = REDUCE_N_READS>
+METAL_FUNC U per_thread_row_reduce(
+    const device T* in,
+    const constant size_t& reduction_size,
+    const constant size_t& out_size,
+    const constant int* shape,
+    const constant size_t* strides,
+    const constant int& ndim,
+    uint lsize_x,
+    uint lid_x,
+    uint2 tid) {
+  Op op;
+
+  // Each threadgroup handles 1 reduction
+  // TODO: Specializing elem_to_loc would be slightly faster
+  int idx = tid.y * out_size + tid.x;
+  int extra_offset = elem_to_loc(idx, shape, strides, ndim);
+  in += extra_offset + lid_x * N_READS;
+
+  // The reduction is accumulated here
+  U total_val = Op::init;
+
+  // Loop over the reduction size within thread group
+  int r = 0;
+  for (; r < (int)ceildiv(reduction_size, N_READS * lsize_x) - 1; r++) {
+    T vals[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      vals[i] = in[i];
+    }
+    for (int i = 0; i < N_READS; i++) {
+      total_val = op(static_cast<U>(vals[i]), total_val);
+    }
+
+    in += lsize_x * N_READS;
+  }
+
+  // Separate case for the last set as we close the reduction size
+  size_t reduction_index = (lid_x + (size_t)lsize_x * r) * N_READS;
+  if (reduction_index < reduction_size) {
+    int max_reads = reduction_size - reduction_index;
+
+    T vals[N_READS];
+    for (int i = 0; i < N_READS; i++) {
+      int idx = min(i, max_reads - 1);
+      vals[i] = static_cast<U>(in[idx]);
+    }
+    for (int i = 0; i < N_READS; i++) {
+      T val = i < max_reads ? vals[i] : Op::init;
+      total_val = op(static_cast<U>(val), total_val);
+    }
+  }
+
+  return total_val;
+}
+
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void row_reduce_general(
+    const device T *in [[buffer(0)]],
+    device mlx_atomic<U> *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& out_size [[buffer(3)]],
+    const constant size_t& non_row_reductions [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+
+  (void)non_row_reductions;
+
+  Op op;
+  threadgroup U local_vals[simd_size];
+
+  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
+
+  total_val = op.simd_reduce(total_val);
+  
+  // Prepare next level
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+      
+  // Reduction within thread group
+  //    Only needed if multiple simd groups
+  if(reduction_size > simd_size) {
+    total_val = lid.x < simd_per_group ? local_vals[lid.x] : op.init;
+    total_val = op.simd_reduce(total_val);
+  }
+  // Update output
+  if (lid.x == 0) {
+    op.atomic_update(out, total_val, tid.x);
+  }
+}
+
+template <typename T, typename U, typename Op, int N_READS=REDUCE_N_READS>
+[[kernel]] void row_reduce_general_no_atomics(
+    const device T *in [[buffer(0)]],
+    device U *out [[buffer(1)]],
+    const constant size_t& reduction_size [[buffer(2)]],
+    const constant size_t& out_size [[buffer(3)]],
+    const constant size_t& non_row_reductions [[buffer(4)]],
+    const constant int* shape [[buffer(5)]],
+    const constant size_t* strides [[buffer(6)]],
+    const constant int& ndim [[buffer(7)]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint3 lsize [[threads_per_threadgroup]],
+    uint3 gsize [[threads_per_grid]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint simd_lane_id [[thread_index_in_simdgroup]],
+    uint simd_per_group [[simdgroups_per_threadgroup]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
+
+  (void)non_row_reductions;
+
+  Op op;
+
+  threadgroup U local_vals[simd_size];
+  U total_val = per_thread_row_reduce<T, U, Op, N_READS>(in, reduction_size, out_size, shape, strides, ndim, lsize.x, lid.x, tid.xy);
+
+  // Reduction within simd group - simd_add isn't supported for int64 types
+  for (uint16_t i = simd_size/2; i > 0; i /= 2) {
+    total_val = op(total_val, simd_shuffle_down(total_val, i));
+  }
+
+  // Prepare next level
+  if (simd_lane_id == 0) {
+    local_vals[simd_group_id] = total_val;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  // Reduction within thread group
+  // Only needed if thread group has multiple simd groups
+  if(ceildiv(reduction_size, N_READS) > simd_size) {
+    total_val = lid.x < simd_per_group ? local_vals[lid.x] : op.init;
+    for (uint16_t i = simd_size/2; i > 0; i /= 2) {
+      total_val = op(total_val, simd_shuffle_down(total_val, i));
+    }
+  }
+  // Write row reduce output for threadgroup with 1st thread in thread group
+  if (lid.x == 0) {
+    out[(ceildiv(gsize.y, lsize.y) * tid.x) + tid.y] = total_val;
+  }
+}
+
+#define instantiate_row_reduce_general(name, itype, otype, op) \
+  instantiate_row_reduce_small(name, itype, otype, op) \
+  template [[host_name("row_reduce_general_" #name)]] \
+  [[kernel]] void row_reduce_general<itype, otype, op>( \
+      const device itype *in [[buffer(0)]],  \
+      device mlx_atomic<otype> *out [[buffer(1)]],  \
+      const constant size_t& reduction_size [[buffer(2)]],  \
+      const constant size_t& out_size [[buffer(3)]],  \
+      const constant size_t& non_row_reductions [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]], \
+      const constant size_t* strides [[buffer(6)]], \
+      const constant int& ndim [[buffer(7)]], \
+      uint3 lid [[thread_position_in_threadgroup]],  \
+      uint3 lsize [[threads_per_threadgroup]],  \
+      uint3 tid [[threadgroup_position_in_grid]],  \
+      uint simd_lane_id [[thread_index_in_simdgroup]],  \
+      uint simd_per_group [[simdgroups_per_threadgroup]],  \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+#define instantiate_row_reduce_general_no_atomics(name, itype, otype, op) \
+  instantiate_row_reduce_small(name, itype, otype, op) \
+  template [[host_name("row_reduce_general_no_atomics_" #name)]] \
+  [[kernel]] void row_reduce_general_no_atomics<itype, otype, op>( \
+      const device itype *in [[buffer(0)]],  \
+      device otype *out [[buffer(1)]],  \
+      const constant size_t& reduction_size [[buffer(2)]],  \
+      const constant size_t& out_size [[buffer(3)]],  \
+      const constant size_t& non_row_reductions [[buffer(4)]], \
+      const constant int* shape [[buffer(5)]], \
+      const constant size_t* strides [[buffer(6)]], \
+      const constant int& ndim [[buffer(7)]], \
+      uint3 lid [[thread_position_in_threadgroup]],  \
+      uint3 lsize [[threads_per_threadgroup]],  \
+      uint3 gsize [[threads_per_grid]], \
+      uint3 tid [[threadgroup_position_in_grid]],  \
+      uint simd_lane_id [[thread_index_in_simdgroup]],  \
+      uint simd_per_group [[simdgroups_per_threadgroup]],  \
+      uint simd_group_id [[simdgroup_index_in_threadgroup]]);
+
+
+///////////////////////////////////////////////////////////////////////////////
+// Instantiations
+///////////////////////////////////////////////////////////////////////////////
+
+#define instantiate_same_row_reduce_helper(name, tname, type, op) \
+  instantiate_row_reduce_general(name ##tname, type, type, op<type>)
+
+#define instantiate_same_row_reduce_na_helper(name, tname, type, op) \
+  instantiate_row_reduce_general_no_atomics(name ##tname, type, type, op<type>)
+
+instantiate_reduce_ops(instantiate_same_row_reduce_helper, instantiate_reduce_helper_types)
+instantiate_reduce_ops(instantiate_same_row_reduce_na_helper, instantiate_reduce_helper_64b)
+
+
+instantiate_reduce_from_types(instantiate_row_reduce_general, and, bool, And)
+instantiate_reduce_from_types(instantiate_row_reduce_general, or, bool, Or)
+
+instantiate_row_reduce_general(sumbool_, bool, uint32_t, Sum<uint32_t>)

mlx/backend/metal/kernels/reduction/ops.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 

mlx/backend/metal/kernels/reduction/reduce_inst.h

@@ -0,0 +1,71 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_atomic>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+
+#define instantiate_reduce_helper_floats(inst_f, name, op)         \
+  inst_f(name, float16, half, op) inst_f(name, float32, float, op) \
+      inst_f(name, bfloat16, bfloat16_t, op)
+
+#define instantiate_reduce_helper_uints(inst_f, name, op)             \
+  inst_f(name, uint8, uint8_t, op) inst_f(name, uint16, uint16_t, op) \
+      inst_f(name, uint32, uint32_t, op)
+
+#define instantiate_reduce_helper_ints(inst_f, name, op)          \
+  inst_f(name, int8, int8_t, op) inst_f(name, int16, int16_t, op) \
+      inst_f(name, int32, int32_t, op)
+
+#define instantiate_reduce_helper_64b(inst_f, name, op) \
+  inst_f(name, int64, int64_t, op) inst_f(name, uint64, uint64_t, op)
+
+#define instantiate_reduce_helper_types(inst_f, name, op) \
+  instantiate_reduce_helper_floats(inst_f, name, op)      \
+      instantiate_reduce_helper_uints(inst_f, name, op)   \
+          instantiate_reduce_helper_ints(inst_f, name, op)
+
+#define instantiate_reduce_ops(inst_f, type_f)        \
+  type_f(inst_f, sum, Sum) type_f(inst_f, prod, Prod) \
+      type_f(inst_f, min_, Min) type_f(inst_f, max_, Max)
+
+// Special case for bool reductions
+#define instantiate_reduce_from_types_helper( \
+    inst_f, name, tname, itype, otype, op)    \
+  inst_f(name##tname, itype, otype, op)
+
+#define instantiate_reduce_from_types(inst_f, name, otype, op)                  \
+  instantiate_reduce_from_types_helper(inst_f, name, bool_, bool, otype, op)    \
+      instantiate_reduce_from_types_helper(                                     \
+          inst_f, name, uint8, uint8_t, otype, op)                              \
+          instantiate_reduce_from_types_helper(                                 \
+              inst_f, name, uint16, uint16_t, otype, op)                        \
+              instantiate_reduce_from_types_helper(                             \
+                  inst_f, name, uint32, uint32_t, otype, op)                    \
+                  instantiate_reduce_from_types_helper(                         \
+                      inst_f, name, int8, int8_t, otype, op)                    \
+                      instantiate_reduce_from_types_helper(                     \
+                          inst_f, name, int16, int16_t, otype, op)              \
+                          instantiate_reduce_from_types_helper(                 \
+                              inst_f, name, int32, int32_t, otype, op)          \
+                              instantiate_reduce_from_types_helper(             \
+                                  inst_f, name, int64, int64_t, otype, op)      \
+                                  instantiate_reduce_from_types_helper(         \
+                                      inst_f, name, float16, half, otype, op)   \
+                                      instantiate_reduce_from_types_helper(     \
+                                          inst_f,                               \
+                                          name,                                 \
+                                          float32,                              \
+                                          float,                                \
+                                          otype,                                \
+                                          op)                                   \
+                                          instantiate_reduce_from_types_helper( \
+                                              inst_f,                           \
+                                              name,                             \
+                                              bfloat16,                         \
+                                              bfloat16_t,                       \
+                                              otype,                            \
+                                              op)

mlx/backend/metal/kernels/reduction/utils.h

@@ -0,0 +1,14 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <metal_atomic>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/steel/utils.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/reduction/ops.h"
+
+static constant constexpr const uint8_t simd_size = 32;

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -0,0 +1,451 @@
+#include <metal_stdlib>
+#include <metal_simdgroup>
+
+#include "mlx/backend/metal/kernels/scaled_dot_product_attention_params.h"
+using namespace metal;
+
+template<typename T, typename T2, typename T4, uint16_t TILE_SIZE_CONST, uint16_t NSIMDGROUPS>
+[[kernel]] void fast_inference_sdpa_compute_partials_template(const device T *Q [[buffer(0)]],
+                              const device T *K [[buffer(1)]],
+                              const device T *V [[buffer(2)]],
+                              const device uint64_t& L [[buffer(3)]],
+                              const device MLXScaledDotProductAttentionParams& params [[buffer(4)]],
+                              device float* O_partials [[buffer(5)]],
+                              device float* p_lse [[buffer(6)]],
+                              device float* p_maxes [[buffer(7)]],
+                              threadgroup T* threadgroup_block [[threadgroup(0)]],
+                              uint simd_lane_id [[thread_index_in_simdgroup]],
+                              uint simd_group_id [[simdgroup_index_in_threadgroup]],
+                              uint3 tid [[threadgroup_position_in_grid]]) {
+    constexpr const size_t DK = 128;
+    constexpr const ulong SIMDGROUP_MATRIX_LOAD_FACTOR = 8;
+    constexpr const size_t THREADS_PER_SIMDGROUP = 32;
+    constexpr const uint iter_offset = NSIMDGROUPS * 4;
+    const bool is_gqa = params.N_KV_HEADS != params.N_Q_HEADS;
+    uint kv_head_offset_factor = tid.x;
+    if(is_gqa) {
+        int q_kv_head_ratio = params.N_Q_HEADS / params.N_KV_HEADS;
+        kv_head_offset_factor = tid.x / q_kv_head_ratio;
+    }
+    constexpr const uint16_t P_VEC4 = TILE_SIZE_CONST / NSIMDGROUPS / 4;
+    constexpr const size_t MATRIX_LOADS_PER_SIMDGROUP = TILE_SIZE_CONST / (SIMDGROUP_MATRIX_LOAD_FACTOR * NSIMDGROUPS);
+    constexpr const size_t MATRIX_COLS = DK / SIMDGROUP_MATRIX_LOAD_FACTOR;
+    constexpr const uint totalSmemV = SIMDGROUP_MATRIX_LOAD_FACTOR * SIMDGROUP_MATRIX_LOAD_FACTOR * (MATRIX_LOADS_PER_SIMDGROUP + 1) * NSIMDGROUPS;
+
+    threadgroup T4* smemFlush = (threadgroup T4*)threadgroup_block;
+    #pragma clang loop unroll(full)
+    for(uint i = 0; i < 8; i++) {
+        smemFlush[simd_lane_id + simd_group_id * THREADS_PER_SIMDGROUP + i * NSIMDGROUPS * THREADS_PER_SIMDGROUP] = T4(0.f);
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // TODO: multiple query sequence length for speculative decoding
+    const uint tgroup_query_head_offset = tid.x * DK + tid.z * (params.N_Q_HEADS * DK);
+
+    const uint tgroup_k_head_offset = kv_head_offset_factor * DK * L;
+    const uint tgroup_k_tile_offset = tid.y * TILE_SIZE_CONST * DK;
+    const uint tgroup_k_batch_offset = tid.z * L * params.N_KV_HEADS * DK;
+
+    const device T* baseK = K + tgroup_k_batch_offset + tgroup_k_tile_offset + tgroup_k_head_offset;
+    const device T* baseQ = Q + tgroup_query_head_offset;
+
+    device T4* simdgroupQueryData = (device T4*)baseQ;
+
+    constexpr const size_t ACCUM_PER_GROUP = TILE_SIZE_CONST / NSIMDGROUPS;
+    float threadAccum[ACCUM_PER_GROUP];
+
+    #pragma clang loop unroll(full)
+    for(size_t threadAccumIndex = 0; threadAccumIndex < ACCUM_PER_GROUP; threadAccumIndex++) {
+        threadAccum[threadAccumIndex] = -INFINITY;
+    }
+
+    uint KROW_ACCUM_INDEX = 0;
+
+    const int32_t SEQUENCE_LENGTH_LESS_TILE_SIZE = L - TILE_SIZE_CONST;
+    const bool LAST_TILE = (tid.y + 1) * TILE_SIZE_CONST >= L;
+    const bool LAST_TILE_ALIGNED = (SEQUENCE_LENGTH_LESS_TILE_SIZE == int32_t(tid.y * TILE_SIZE_CONST));
+
+    T4 thread_data_x4;
+    T4 thread_data_y4;
+    if(!LAST_TILE || LAST_TILE_ALIGNED) {
+        thread_data_x4 = *(simdgroupQueryData + simd_lane_id);
+        #pragma clang loop unroll(full)
+        for(size_t KROW = simd_group_id; KROW < TILE_SIZE_CONST; KROW += NSIMDGROUPS) {
+            const uint KROW_OFFSET = KROW * DK;
+            const device T* baseKRow = baseK + KROW_OFFSET;
+            device T4* keysData = (device T4*)baseKRow;
+            thread_data_y4 = *(keysData + simd_lane_id);
+            T kq_scalar = dot(thread_data_x4, thread_data_y4);
+            threadAccum[KROW_ACCUM_INDEX] = float(kq_scalar);
+            KROW_ACCUM_INDEX++;
+        }
+    } else {
+        thread_data_x4 = *(simdgroupQueryData + simd_lane_id);
+        const uint START_ROW = tid.y * TILE_SIZE_CONST;
+        const device T* baseKThisHead = K + tgroup_k_batch_offset + tgroup_k_head_offset;
+
+        for(size_t KROW = START_ROW + simd_group_id; KROW < L; KROW += NSIMDGROUPS) {
+            const uint KROW_OFFSET = KROW * DK;
+            const device T* baseKRow = baseKThisHead + KROW_OFFSET;
+            device T4* keysData = (device T4*)baseKRow;
+            thread_data_y4 = *(keysData + simd_lane_id);
+            T kq_scalar = dot(thread_data_x4, thread_data_y4);
+            threadAccum[KROW_ACCUM_INDEX] = float(kq_scalar);
+            KROW_ACCUM_INDEX++;
+        }
+    }
+    threadgroup float* smemP = (threadgroup float*)threadgroup_block;
+
+    #pragma clang loop unroll(full)
+    for(size_t i = 0; i < P_VEC4; i++) {
+        thread_data_x4 = T4(threadAccum[4 * i], threadAccum[4 * i + 1], threadAccum[4 * i + 2], threadAccum[4 * i + 3]);
+        simdgroup_barrier(mem_flags::mem_none);
+        thread_data_y4 = simd_sum(thread_data_x4);
+        if(simd_lane_id == 0) {
+            const uint base_smem_p_offset = i * iter_offset + simd_group_id;
+            smemP[base_smem_p_offset + NSIMDGROUPS * 0] = float(thread_data_y4.x);
+            smemP[base_smem_p_offset + NSIMDGROUPS * 1] = float(thread_data_y4.y);
+            smemP[base_smem_p_offset + NSIMDGROUPS * 2] = float(thread_data_y4.z);
+            smemP[base_smem_p_offset + NSIMDGROUPS * 3] = float(thread_data_y4.w);
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    float groupMax;
+    float lse = 0.f;
+
+    constexpr const size_t THREADS_PER_THREADGROUP_TIMES_4 = 4 * 32;
+    constexpr const size_t ACCUM_ARRAY_LENGTH = TILE_SIZE_CONST / THREADS_PER_THREADGROUP_TIMES_4 + 1;
+    float4 pvals[ACCUM_ARRAY_LENGTH];
+
+    #pragma clang loop unroll(full)
+    for(uint accum_array_iter = 0; accum_array_iter < ACCUM_ARRAY_LENGTH; accum_array_iter++) {
+        pvals[accum_array_iter] = float4(-INFINITY);
+    }
+
+    if (TILE_SIZE_CONST == 64) {
+        threadgroup float2* smemPtrFlt2 = (threadgroup float2*)threadgroup_block;
+        float2 vals = smemPtrFlt2[simd_lane_id];
+        vals *= params.INV_ALPHA;
+        float maxval = max(vals.x, vals.y);
+        simdgroup_barrier(mem_flags::mem_none);
+        groupMax = simd_max(maxval);
+
+        float2 expf_shifted = exp(vals - groupMax);
+        float sumExpLocal = expf_shifted.x + expf_shifted.y;
+        simdgroup_barrier(mem_flags::mem_none);
+        float tgroupExpSum = simd_sum(sumExpLocal);
+
+        lse = log(tgroupExpSum);
+        float2 local_p_hat = expf_shifted / tgroupExpSum;
+        pvals[0].x = local_p_hat.x;
+        pvals[0].y = local_p_hat.y;
+        smemPtrFlt2[simd_lane_id] = float2(0.f);
+    }
+    constexpr const bool TILE_SIZE_LARGER_THAN_64 = TILE_SIZE_CONST > 64;
+    constexpr const int TILE_SIZE_ITERS_128 = TILE_SIZE_CONST / 128;
+
+    if (TILE_SIZE_LARGER_THAN_64) {
+        float maxval = -INFINITY;
+        threadgroup float4* smemPtrFlt4 = (threadgroup float4*)threadgroup_block;
+        #pragma clang loop unroll(full)
+        for(int i = 0; i < TILE_SIZE_ITERS_128; i++) {
+            float4 vals = smemPtrFlt4[simd_lane_id + i * THREADS_PER_SIMDGROUP];
+            vals *= params.INV_ALPHA;
+            pvals[i] = vals;
+            maxval = fmax3(vals.x, vals.y, maxval);
+            maxval = fmax3(vals.z, vals.w, maxval);
+        }
+        simdgroup_barrier(mem_flags::mem_none);
+        groupMax = simd_max(maxval);
+
+        float sumExpLocal = 0.f;
+        #pragma clang loop unroll(full)
+        for(int i = 0; i < TILE_SIZE_ITERS_128; i++) {
+            pvals[i] = exp(pvals[i] - groupMax);
+            sumExpLocal += pvals[i].x + pvals[i].y + pvals[i].z + pvals[i].w;
+        }
+        simdgroup_barrier(mem_flags::mem_none);
+        float tgroupExpSum = simd_sum(sumExpLocal);
+        lse = log(tgroupExpSum);
+        #pragma clang loop unroll(full)
+        for(int i = 0; i < TILE_SIZE_ITERS_128; i++) {
+            pvals[i] = pvals[i] / tgroupExpSum;
+            smemPtrFlt4[simd_lane_id + i * THREADS_PER_SIMDGROUP] = float4(0.f);
+        }
+    }
+
+    threadgroup T* smemV = (threadgroup T*)threadgroup_block;
+
+    const size_t v_batch_offset = tid.z * params.N_KV_HEADS * L * DK;
+    const size_t v_head_offset = kv_head_offset_factor * L * DK;
+
+    const size_t v_tile_offset = tid.y * TILE_SIZE_CONST * DK;
+    const size_t v_offset = v_batch_offset + v_head_offset + v_tile_offset;
+    device T* baseV = (device T*)V + v_offset;
+
+    threadgroup float* smemOpartial = (threadgroup float*)(smemV + totalSmemV);
+
+    if (!LAST_TILE || LAST_TILE_ALIGNED) {
+        #pragma clang loop unroll(full)
+        for(size_t col = 0; col < MATRIX_COLS; col++) {
+            uint matrix_load_loop_iter = 0;
+            constexpr const size_t TILE_SIZE_CONST_DIV_8 = TILE_SIZE_CONST / 8;
+            
+            for(size_t tile_start = simd_group_id; tile_start < TILE_SIZE_CONST_DIV_8; tile_start += NSIMDGROUPS) {
+                simdgroup_matrix<T, 8, 8> tmp;
+                ulong simdgroup_matrix_offset = matrix_load_loop_iter * NSIMDGROUPS * SIMDGROUP_MATRIX_LOAD_FACTOR + simd_group_id * SIMDGROUP_MATRIX_LOAD_FACTOR;
+                ulong2 matrixOrigin = ulong2(col * SIMDGROUP_MATRIX_LOAD_FACTOR, simdgroup_matrix_offset);
+                simdgroup_load(tmp, baseV, DK, matrixOrigin, true);
+                const ulong2 matrixOriginSmem = ulong2(simdgroup_matrix_offset, 0);
+                const ulong elemsPerRowSmem = TILE_SIZE_CONST;
+                simdgroup_store(tmp, smemV, elemsPerRowSmem, matrixOriginSmem, false);
+                matrix_load_loop_iter++;
+            };
+            
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+            
+            if (TILE_SIZE_CONST == 64) {
+                T2 local_p_hat = T2(pvals[0].x, pvals[0].y);
+                uint loop_iter = 0;
+                threadgroup float* oPartialSmem = smemOpartial + SIMDGROUP_MATRIX_LOAD_FACTOR * col;
+                
+                #pragma clang loop unroll(full)
+                for(size_t row = simd_group_id; row < SIMDGROUP_MATRIX_LOAD_FACTOR; row += NSIMDGROUPS) {
+                    threadgroup T* smemV_row = smemV + (TILE_SIZE_CONST * row);
+                    threadgroup T2* smemV2 = (threadgroup T2*)smemV_row;
+                    T2 v_local = *(smemV2 + simd_lane_id);
+                    
+                    T val = dot(local_p_hat, v_local);
+                    simdgroup_barrier(mem_flags::mem_none);
+
+                    T row_sum = simd_sum(val);
+                    oPartialSmem[simd_group_id + loop_iter * NSIMDGROUPS] = float(row_sum);
+                    loop_iter++;
+                }
+            }
+            
+            if (TILE_SIZE_CONST > 64) {
+                constexpr const size_t TILE_SIZE_CONST_DIV_128 = (TILE_SIZE_CONST + 1) / 128;
+                threadgroup float* oPartialSmem = smemOpartial + SIMDGROUP_MATRIX_LOAD_FACTOR * col;
+                uint loop_iter = 0;
+                for(size_t row = simd_group_id; row < SIMDGROUP_MATRIX_LOAD_FACTOR; row += NSIMDGROUPS) {
+                    threadgroup T* smemV_row = smemV + (TILE_SIZE_CONST * row);
+                    
+                    T row_sum = 0.f;
+                    for(size_t i = 0; i < TILE_SIZE_CONST_DIV_128; i++) {
+                        threadgroup T4* smemV2 = (threadgroup T4*)smemV_row;
+                        T4 v_local = *(smemV2 + simd_lane_id + i * THREADS_PER_SIMDGROUP);
+                        T4 p_local = T4(pvals[i]);
+                        T val = dot(p_local, v_local);
+                        row_sum += val;
+                    }
+                    simdgroup_barrier(mem_flags::mem_none);
+                    row_sum = simd_sum(row_sum);
+                    oPartialSmem[simd_group_id + loop_iter * NSIMDGROUPS] = float(row_sum);
+                    loop_iter++;
+                }
+            }
+        }
+    } else {
+        const int32_t START_ROW = tid.y * TILE_SIZE_CONST;
+        const int32_t MAX_START_ROW = L - SIMDGROUP_MATRIX_LOAD_FACTOR + 1;
+        const device T* baseVThisHead = V + v_batch_offset + v_head_offset;
+        constexpr const int ROWS_PER_ITER = 8;
+        #pragma clang loop unroll(full)
+        for(size_t col = 0; col < MATRIX_COLS; col++) {
+            uint smem_col_index = simd_group_id * SIMDGROUP_MATRIX_LOAD_FACTOR;
+            int32_t tile_start;
+            for(tile_start = START_ROW + simd_group_id * SIMDGROUP_MATRIX_LOAD_FACTOR; tile_start < MAX_START_ROW; tile_start += NSIMDGROUPS * SIMDGROUP_MATRIX_LOAD_FACTOR) {
+                simdgroup_matrix<T, 8, 8> tmp;
+                ulong2 matrixOrigin = ulong2(col * SIMDGROUP_MATRIX_LOAD_FACTOR, tile_start);
+                simdgroup_load(tmp, baseVThisHead, DK, matrixOrigin, /* transpose */ true);
+                const ulong2 matrixOriginSmem = ulong2(smem_col_index, 0);
+                constexpr const ulong elemsPerRowSmem = TILE_SIZE_CONST;
+                simdgroup_store(tmp, smemV, elemsPerRowSmem, matrixOriginSmem, /* transpose */ false);
+                smem_col_index += NSIMDGROUPS * SIMDGROUP_MATRIX_LOAD_FACTOR;
+            };
+
+            tile_start = ((L / SIMDGROUP_MATRIX_LOAD_FACTOR) * SIMDGROUP_MATRIX_LOAD_FACTOR);
+
+            const int32_t INT_L = int32_t(L);
+            for(int row_index  = tile_start + simd_group_id ; row_index < INT_L; row_index += NSIMDGROUPS) {
+                if(simd_lane_id < SIMDGROUP_MATRIX_LOAD_FACTOR) {
+                    const uint elems_per_row_gmem = DK;
+                    const uint col_index_v_gmem = col * SIMDGROUP_MATRIX_LOAD_FACTOR + simd_lane_id;
+                    const uint row_index_v_gmem = row_index;
+
+                    const uint elems_per_row_smem = TILE_SIZE_CONST;
+                    const uint col_index_v_smem = row_index % TILE_SIZE_CONST;
+                    const uint row_index_v_smem = simd_lane_id;
+
+                    const uint scalar_offset_gmem = row_index_v_gmem * elems_per_row_gmem + col_index_v_gmem;
+                    const uint scalar_offset_smem = row_index_v_smem * elems_per_row_smem + col_index_v_smem;
+                    T vdata = T(*(baseVThisHead + scalar_offset_gmem));
+                    smemV[scalar_offset_smem] = vdata;
+                    smem_col_index += NSIMDGROUPS;
+                }
+            }
+
+            threadgroup_barrier(mem_flags::mem_threadgroup);
+
+            if (TILE_SIZE_CONST == 64) {
+                T2 local_p_hat = T2(pvals[0].x, pvals[0].y);
+                threadgroup float* oPartialSmem = smemOpartial + SIMDGROUP_MATRIX_LOAD_FACTOR * col;
+                for(size_t smem_row_index = simd_group_id;
+                    smem_row_index < ROWS_PER_ITER; smem_row_index += NSIMDGROUPS) {
+                    threadgroup T* smemV_row = smemV + (TILE_SIZE_CONST * smem_row_index);
+                    threadgroup T2* smemV2 = (threadgroup T2*)smemV_row;
+                    T2 v_local = *(smemV2 + simd_lane_id);
+                    T val = dot(local_p_hat, v_local);
+                    simdgroup_barrier(mem_flags::mem_none);
+                    T row_sum = simd_sum(val);
+                    oPartialSmem[smem_row_index] = float(row_sum);
+                }
+            }
+
+            if (TILE_SIZE_CONST > 64) {
+                threadgroup float* oPartialSmem = smemOpartial + SIMDGROUP_MATRIX_LOAD_FACTOR * col;
+                uint loop_count = 0;
+                for(size_t row_index = simd_group_id;
+                    row_index < ROWS_PER_ITER; row_index += NSIMDGROUPS) {
+                    T row_sum = 0.f;
+                    for(size_t tile_iters = 0; tile_iters < TILE_SIZE_ITERS_128; tile_iters++) {
+                        threadgroup T* smemV_row = smemV + (TILE_SIZE_CONST * row_index);
+                        threadgroup T4* smemV2 = (threadgroup T4*)smemV_row;
+                        T4 v_local = *(smemV2 + simd_lane_id + tile_iters * THREADS_PER_SIMDGROUP);
+                        T4 p_local = T4(pvals[tile_iters]);
+                        row_sum += dot(p_local, v_local);
+                        
+                    }
+                    simdgroup_barrier(mem_flags::mem_none);
+                    row_sum = simd_sum(row_sum);
+                    oPartialSmem[simd_group_id + NSIMDGROUPS * loop_count] = float(row_sum);
+                    loop_count++;
+                }
+            }
+        }
+    }
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    if(simd_group_id == 0) {
+        threadgroup float4* oPartialVec4 = (threadgroup float4*)smemOpartial;
+        float4 vals = *(oPartialVec4 + simd_lane_id);
+        device float* oPartialGmem = O_partials + tid.x * DK * params.KV_TILES + tid.y * DK;
+        device float4* oPartialGmemVec4 = (device float4*)oPartialGmem;
+        oPartialGmemVec4[simd_lane_id] = vals;
+    }
+
+    if(simd_group_id == 0 && simd_lane_id == 0) {
+        const uint tileIndex = tid.y;
+        const uint gmem_partial_scalar_offset = tid.z * params.N_Q_HEADS * params.KV_TILES + tid.x * params.KV_TILES + tileIndex;
+        p_lse[gmem_partial_scalar_offset] = lse;
+        p_maxes[gmem_partial_scalar_offset] = groupMax;
+    }
+}
+
+#define instantiate_fast_inference_sdpa_to_partials_kernel(itype, itype2, itype4, tile_size, nsimdgroups) \
+template [[host_name("fast_inference_sdpa_compute_partials_" #itype "_" #tile_size "_" #nsimdgroups )]] \
+[[kernel]] void fast_inference_sdpa_compute_partials_template<itype, itype2, itype4, tile_size, nsimdgroups>( \
+    const device itype *Q [[buffer(0)]], \
+    const device itype *K [[buffer(1)]], \
+    const device itype *V [[buffer(2)]], \
+    const device uint64_t& L [[buffer(3)]], \
+    const device MLXScaledDotProductAttentionParams& params [[buffer(4)]], \
+    device float* O_partials [[buffer(5)]], \
+    device float* p_lse [[buffer(6)]], \
+    device float* p_maxes [[buffer(7)]], \
+    threadgroup itype *threadgroup_block [[threadgroup(0)]], \
+    uint simd_lane_id [[thread_index_in_simdgroup]], \
+    uint simd_group_id [[simdgroup_index_in_threadgroup]], \
+    uint3 tid [[threadgroup_position_in_grid]]);
+
+
+#define instantiate_fast_inference_sdpa_to_partials_shapes_helper(itype, itype2, itype4, tile_size) \
+    instantiate_fast_inference_sdpa_to_partials_kernel(itype, itype2, itype4, tile_size, 4) \
+    instantiate_fast_inference_sdpa_to_partials_kernel(itype, itype2, itype4, tile_size, 8) \
+
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(float, float2, float4, 64);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(float, float2, float4, 128);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(float, float2, float4, 256);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(float, float2, float4, 512);
+
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(half, half2, half4, 64);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(half, half2, half4, 128);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(half, half2, half4, 256);
+instantiate_fast_inference_sdpa_to_partials_shapes_helper(half, half2, half4, 512);
+
+
+template <typename T>
+void fast_inference_sdpa_reduce_tiles_template(
+    const device float *O_partials [[buffer(0)]],
+    const device float *p_lse[[buffer(1)]],
+    const device float *p_maxes [[buffer(2)]],
+    const device MLXScaledDotProductAttentionParams& params [[buffer(3)]],
+    device T* O [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]]) {
+
+    constexpr const int DK = 128;
+    const ulong offset_rows = tid.z * params.KV_TILES * params.N_Q_HEADS + tid.x * params.KV_TILES;
+    const device float* p_lse_row = p_lse + offset_rows;
+    const device float* p_rowmax_row = p_maxes + offset_rows;
+    // reserve some number of registers.  this constitutes an assumption on max value of KV TILES.
+    constexpr const uint8_t reserve = 128;
+    float p_lse_regs[reserve];
+    float p_rowmax_regs[reserve];
+    float weights[reserve];
+
+    float true_max = -INFINITY;
+    for(size_t i = 0; i < params.KV_TILES; i++) {
+        p_lse_regs[i] = float(*(p_lse_row + i));
+        p_rowmax_regs[i] = float(*(p_rowmax_row + i));
+        true_max = fmax(p_rowmax_regs[i], true_max);
+        weights[i] = exp(p_lse_regs[i]);
+    }
+
+    float denom = 0.f;
+    for(size_t i = 0; i < params.KV_TILES; i++) {
+        weights[i] *= exp(p_rowmax_regs[i]-true_max);
+        denom += weights[i];
+    }
+
+    const device float* O_partials_with_offset = O_partials + tid.z * params.N_Q_HEADS * DK * params.KV_TILES + tid.x * DK * params.KV_TILES;
+
+    float o_value = 0.f;
+    for(size_t i = 0; i < params.KV_TILES; i++) {
+        float val = *(O_partials_with_offset + i * DK + lid.x);
+        o_value += val * weights[i] / denom;
+    }
+    device T* O_gmem = O + tid.z * params.N_Q_HEADS * DK + tid.x * DK;
+    O_gmem[lid.x] = T(o_value);
+    return;
+}
+
+
+kernel void fast_inference_sdpa_reduce_tiles_float(
+    const device float *O_partials [[buffer(0)]],
+    const device float *p_lse[[buffer(1)]],
+    const device float *p_maxes [[buffer(2)]],
+    const device MLXScaledDotProductAttentionParams& params [[buffer(3)]],
+    device float* O [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]])
+{
+    fast_inference_sdpa_reduce_tiles_template<float>(O_partials, p_lse, p_maxes, params,
+                                     O, tid, lid);
+}
+
+kernel void fast_inference_sdpa_reduce_tiles_half(
+    const device float *O_partials [[buffer(0)]],
+    const device float *p_lse[[buffer(1)]],
+    const device float *p_maxes [[buffer(2)]],
+    const device MLXScaledDotProductAttentionParams& params [[buffer(3)]],
+    device half* O [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]])
+{
+    fast_inference_sdpa_reduce_tiles_template<half>(O_partials, p_lse, p_maxes, params,
+                                     O, tid, lid);
+}

mlx/backend/metal/kernels/scaled_dot_product_attention_params.h

@@ -0,0 +1,14 @@
+//
+//  scaled_dot_product_attention_params.h
+//  mlx
+
+#pragma once
+
+struct MLXScaledDotProductAttentionParams {
+  // Associated dimensions & transposition information
+  const uint QUERY_SEQUENCE_LENGTH = 1;
+  const uint N_Q_HEADS = 32;
+  const uint N_KV_HEADS = 32;
+  const uint KV_TILES = 1;
+  const float INV_ALPHA = 0.08838834764831843f;
+};

mlx/backend/metal/kernels/scatter.metal

@@ -4,7 +4,7 @@
 
 #include "mlx/backend/metal/kernels/bf16.h"
 #include "mlx/backend/metal/kernels/indexing.h"
-#include "mlx/backend/metal/kernels/reduce.h"
+#include "mlx/backend/metal/kernels/reduction/ops.h"
 #include "mlx/backend/metal/kernels/utils.h"
 
 using namespace metal;
@@ -13,6 +13,58 @@ using namespace metal;
 // Scatter kernel
 /////////////////////////////////////////////////////////////////////
 
+template <typename T, typename IdxT, typename Op, int NIDX> \
+METAL_FUNC void scatter_1d_index_impl(
+  const device T *updates [[buffer(1)]],
+  device mlx_atomic<T> *out [[buffer(2)]],
+  const constant int* out_shape [[buffer(3)]],
+  const constant size_t* out_strides [[buffer(4)]],
+  const constant size_t& upd_size [[buffer(5)]],
+  const constant bool& upd_col_contiguous [[buffer(6)]],
+  const thread array<const device IdxT*, NIDX>& idx_buffers,
+  uint2 gid [[thread_position_in_grid]]) {
+
+  Op op;
+
+  uint out_idx = 0;
+  for (int i = 0; i < NIDX; i++) {
+    auto idx_val = offset_neg_idx(
+        idx_buffers[i][gid.y], out_shape[i]);
+    out_idx += idx_val * out_strides[i];
+  }
+
+  if (!upd_col_contiguous) {
+    op.atomic_update(out, updates[gid.y * upd_size + gid.x], out_idx + gid.x);
+  } else {
+    op.atomic_update(out, updates[gid.x * upd_size + gid.y], out_idx + gid.x);
+  }
+}
+
+#define make_scatter_1d_index(IDX_ARG, IDX_ARR) \
+template <typename T, typename IdxT, typename Op, int NIDX> \
+[[kernel]] void scatter_1d_index( \
+  const device T *updates [[buffer(1)]], \
+  device mlx_atomic<T> *out [[buffer(2)]], \
+  const constant int* out_shape [[buffer(3)]], \
+  const constant size_t* out_strides [[buffer(4)]], \
+  const constant size_t& upd_size [[buffer(5)]], \
+  const constant bool& upd_col_contiguous [[buffer(6)]], \
+  IDX_ARG(IdxT) \
+  uint2 gid [[thread_position_in_grid]]) { \
+  \
+  const array<const device IdxT*, NIDX> idx_buffers = {IDX_ARR()}; \
+  \
+  return scatter_1d_index_impl<T, IdxT, Op, NIDX>( \
+    updates, \
+    out, \
+    out_shape, \
+    out_strides, \
+    upd_size, \
+    upd_col_contiguous, \
+    idx_buffers, \
+    gid); \
+  \
+}
 
 template <typename T, typename IdxT, typename Op, int NIDX>
 METAL_FUNC void scatter_impl(
@@ -46,10 +98,14 @@ METAL_FUNC void scatter_impl(
     out_idx += idx_val * out_strides[ax];
   }
 
-  auto out_offset = elem_to_loc(
-      ind_offset, upd_shape + indices.ndim, out_strides, out_ndim);
+  if (upd_size > 1) {
+    auto out_offset = elem_to_loc(
+        ind_offset, upd_shape + indices.ndim, out_strides, out_ndim);
+    out_idx += out_offset;
+  }
+
   auto upd_idx = elem_to_loc(gid.y * upd_size + gid.x, upd_shape, upd_strides, upd_ndim);
-  op.atomic_update(out, updates[upd_idx], out_idx + out_offset);
+  op.atomic_update(out, updates[upd_idx], out_idx);
 }
 
 #define make_scatter_impl(IDX_ARG, IDX_ARR) \
@@ -90,9 +146,11 @@ template <typename T, typename IdxT, typename Op, int NIDX>  \
       axes, \
       idxs, \
       gid); \
-} 
+}
 
-#define make_scatter(n) make_scatter_impl(IDX_ARG_ ##n, IDX_ARR_ ##n)
+#define make_scatter(n) \
+make_scatter_impl(IDX_ARG_ ##n, IDX_ARR_ ##n) \
+make_scatter_1d_index(IDX_ARG_ ##n, IDX_ARR_ ##n)
 
 make_scatter(0)
 make_scatter(1)
@@ -129,8 +187,21 @@ template [[host_name("scatter" name "_" #nidx)]] \
     IDX_ARG(idx_t) \
     uint2 gid [[thread_position_in_grid]]);
 
+#define instantiate_scatter6(name, src_t, idx_t, op_t, nidx, IDX_ARG) \
+template [[host_name("scatter_1d_index" name "_" #nidx)]] \
+[[kernel]] void scatter_1d_index<src_t, idx_t, op_t, nidx>( \
+  const device src_t *updates [[buffer(1)]], \
+  device mlx_atomic<src_t> *out [[buffer(2)]], \
+  const constant int* out_shape [[buffer(3)]], \
+  const constant size_t* out_strides [[buffer(4)]], \
+  const constant size_t& upd_size [[buffer(5)]], \
+  const constant bool& upd_col_contiguous [[buffer(6)]], \
+  IDX_ARG(idx_t) \
+  uint2 gid [[thread_position_in_grid]]);
+
 #define instantiate_scatter4(name, src_t, idx_t, op_t, nidx) \
-  instantiate_scatter5(name, src_t, idx_t, op_t, nidx, IDX_ARG_ ##nidx)
+  instantiate_scatter5(name, src_t, idx_t, op_t, nidx, IDX_ARG_ ##nidx) \
+  instantiate_scatter6(name, src_t, idx_t, op_t, nidx, IDX_ARG_ ##nidx)
 
 // Special case NINDEX=0
 #define instantiate_scatter_nd0(name, type) \

mlx/backend/metal/kernels/steel/conv/conv.h

@@ -0,0 +1,11 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/loader.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+using namespace metal;
+using namespace mlx::steel;

mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.metal

@@ -0,0 +1,189 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/conv.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/bf16.h"
+
+using namespace metal;
+
+template <typename T,
+          int BM,
+          int BN,
+          int BK,
+          int WM,
+          int WN,
+          int N_CHANNELS = 0,
+          bool SMALL_FILTER = false>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void implicit_gemm_conv_2d(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* C [[buffer(2)]],
+    const constant MLXConvParams<2>* params [[buffer(3)]],
+    const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+
+  using namespace mlx::steel;
+
+  (void)lid;
+  
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+  constexpr short tgp_padding_a = 16 / sizeof(T);
+  constexpr short tgp_padding_b = 16 / sizeof(T);
+
+  constexpr short shape_a_cols = (transpose_a ? BM : BK) + tgp_padding_a;
+  constexpr short shape_b_cols = (transpose_b ? BK : BN) + tgp_padding_b;
+  constexpr short shape_a_rows = (transpose_a ? BK : BM);
+  constexpr short shape_b_rows = (transpose_b ? BN : BK);
+  constexpr short tgp_mem_size_a = shape_a_cols * shape_a_rows;
+  constexpr short tgp_mem_size_b = shape_b_cols * shape_b_rows;
+
+  constexpr short tgp_size = WM * WN * 32;
+  
+  // Input loader 
+
+  using loader_a_t = typename metal::conditional_t<
+      // Check for small channel specialization
+      N_CHANNELS != 0 && N_CHANNELS <= 4,
+
+        // Go to small channel specialization
+        Conv2DInputBlockLoaderSmallChannels<
+            T, BM, BN, BK, tgp_size, N_CHANNELS, tgp_padding_a>,
+
+        // Else go to general loader 
+        typename metal::conditional_t< 
+            // Check if filter size is small enough
+            SMALL_FILTER,
+
+              // Go to small filter specialization
+              Conv2DInputBlockLoaderSmallFilter<
+                  T, BM, BN, BK, tgp_size, tgp_padding_a>,
+
+              // Else go to large filter generalization 
+              Conv2DInputBlockLoaderLargeFilter<
+                  T, BM, BN, BK, tgp_size, tgp_padding_a>
+        >
+  >;
+
+
+  // Weight loader
+  using loader_b_t = typename metal::conditional_t<
+      // Check for small channel specialization
+      N_CHANNELS != 0 && N_CHANNELS <= 4,
+
+        // Go to small channel specialization
+        Conv2DWeightBlockLoaderSmallChannels<
+            T, BM, BN, BK, tgp_size, N_CHANNELS, tgp_padding_b>,
+
+        // Else go to general loader 
+        Conv2DWeightBlockLoader<T, BM, BN, BK, tgp_size, tgp_padding_b>
+  >;
+  
+  using mma_t = BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      shape_a_cols,
+      shape_b_cols>;
+    
+  threadgroup T As[tgp_mem_size_a];
+  threadgroup T Bs[tgp_mem_size_b];
+
+  const int tid_y = ((tid.y) << gemm_params->swizzle_log) +
+    ((tid.x) & ((1 << gemm_params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> gemm_params->swizzle_log;
+
+  if (gemm_params->tiles_n <= tid_x || gemm_params->tiles_m <= tid_y) {
+    return;
+  }
+
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int K = gemm_params->K;
+  const int N = gemm_params->N;
+
+  B += c_col * K;
+  C += c_row * N + c_col;
+
+  const int2 offsets_a(0, c_row);
+  const int2 offsets_b(0, c_col);
+
+  // Prepare threadgroup loading operations
+  loader_a_t loader_a(A, As, offsets_a, params, gemm_params, simd_gid, simd_lid);
+  loader_b_t loader_b(B, Bs, offsets_b, params, gemm_params, simd_gid, simd_lid);
+
+  // Prepare threadgroup mma operation
+  mma_t mma_op(simd_gid, simd_lid);
+
+  int gemm_k_iterations = gemm_params->gemm_k_iterations;
+  for (int k = 0; k < gemm_k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Load elements into threadgroup
+    loader_a.load_unsafe();
+    loader_b.load_unsafe();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results to device memory
+  short tgp_bm = min(BM, gemm_params->M - c_row);
+  short tgp_bn = min(BN, gemm_params->N - c_col);
+  mma_op.store_result_safe(C, N, short2(tgp_bn, tgp_bm));
+
+}
+
+#define instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, channel_name, n_channels, filter_name, small_filter) \
+  template [[host_name("implicit_gemm_conv_2d_" #name "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn "_channel_" #channel_name "_filter_" #filter_name)]] \
+  [[kernel]] void implicit_gemm_conv_2d<itype, bm, bn, bk, wm, wn, n_channels, small_filter>( \
+      const device itype* A [[buffer(0)]], \
+      const device itype* B [[buffer(1)]], \
+      device itype* C [[buffer(2)]], \
+      const constant MLXConvParams<2>* params [[buffer(3)]], \
+      const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint3 lid [[thread_position_in_threadgroup]], \
+      uint simd_gid [[simdgroup_index_in_threadgroup]], \
+      uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_implicit_2d_filter(name, itype, bm, bn, bk, wm, wn) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, l, 0, s, true) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, l, 0, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 1, 1, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 2, 2, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 3, 3, l, false) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn, 4, 4, l, false)
+
+#define instantiate_implicit_2d_blocks(name, itype) \
+    instantiate_implicit_2d_filter(name, itype, 32,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 64,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2)
+
+instantiate_implicit_2d_blocks(float32, float);
+instantiate_implicit_2d_blocks(float16, half);
+instantiate_implicit_2d_blocks(bfloat16, bfloat16_t);

mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal

@@ -0,0 +1,209 @@
+// Copyright © 2024 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/conv.h"
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h"
+#include "mlx/backend/metal/kernels/bf16.h"
+
+using namespace metal;
+using namespace mlx::steel;
+
+template <typename T,
+          int BM,
+          int BN,
+          int BK,
+          int WM,
+          int WN,
+          typename AccumType = float,
+          typename Epilogue = TransformNone<T, AccumType>>
+[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void implicit_gemm_conv_2d_general(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device T* C [[buffer(2)]],
+    const constant MLXConvParams<2>* params [[buffer(3)]],
+    const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]],
+    const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]],
+    const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]],
+    const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 lid [[thread_position_in_threadgroup]],
+    uint simd_gid [[simdgroup_index_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+
+  (void)lid;
+  
+  constexpr bool transpose_a = false;
+  constexpr bool transpose_b = true;
+  constexpr short tgp_padding_a = 16 / sizeof(T);
+  constexpr short tgp_padding_b = 16 / sizeof(T);
+
+  constexpr short shape_a_cols = (transpose_a ? BM : BK) + tgp_padding_a;
+  constexpr short shape_b_cols = (transpose_b ? BK : BN) + tgp_padding_b;
+  constexpr short shape_a_rows = (transpose_a ? BK : BM);
+  constexpr short shape_b_rows = (transpose_b ? BN : BK);
+  constexpr short tgp_mem_size_a = shape_a_cols * shape_a_rows;
+  constexpr short tgp_mem_size_b = shape_b_cols * shape_b_rows;
+
+  constexpr short tgp_size = WM * WN * 32;
+  
+  // Input loader 
+  using loader_a_t = Conv2DInputBlockLoaderGeneral<
+      T, BM, BN, BK, tgp_size, tgp_padding_a>;
+
+  // Weight loader
+  using loader_b_t = Conv2DWeightBlockLoaderGeneral<
+      T, BM, BN, BK, tgp_size, tgp_padding_b>;
+  
+  using mma_t = BlockMMA<
+      T,
+      T,
+      BM,
+      BN,
+      BK,
+      WM,
+      WN,
+      transpose_a,
+      transpose_b,
+      shape_a_cols,
+      shape_b_cols>;
+    
+  threadgroup T As[tgp_mem_size_a];
+  threadgroup T Bs[tgp_mem_size_b];
+
+  const int tid_y = ((tid.y) << gemm_params->swizzle_log) +
+    ((tid.x) & ((1 << gemm_params->swizzle_log) - 1));
+  const int tid_x = (tid.x) >> gemm_params->swizzle_log;
+
+  if (gemm_params->tiles_n <= tid_x || gemm_params->tiles_m <= tid_y) {
+    return;
+  }
+
+  const int tid_z = tid.z;
+
+  const int base_oh = tid_z / jump_params->f_out_jump_w;
+  const int base_ow = tid_z % jump_params->f_out_jump_w;
+
+  const int base_wh = base_h[base_oh].weight_base;
+  const int base_ww = base_w[base_ow].weight_base;
+
+  const int base_wh_size = base_h[base_oh].weight_size;
+  const int base_ww_size = base_w[base_ow].weight_size;
+
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int K = gemm_params->K;
+
+  B += c_col * K;
+
+  const int4 offsets_a(0, c_row, base_oh, base_ow);
+  const int2 offsets_b(0, c_col);
+
+  // Prepare threadgroup loading operations
+  loader_a_t loader_a(A, As, offsets_a, params, jump_params, base_wh, base_ww, simd_gid, simd_lid);
+  loader_b_t loader_b(B, Bs, offsets_b, params, jump_params, base_wh, base_ww, simd_gid, simd_lid);
+
+  // Prepare threadgroup mma operation
+  mma_t mma_op(simd_gid, simd_lid);
+
+  int gemm_k_iterations = base_wh_size * base_ww_size * gemm_params->gemm_k_iterations;
+
+  for (int k = 0; k < gemm_k_iterations; k++) {
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Load elements into threadgroup
+    loader_a.load_unsafe();
+    loader_b.load_unsafe();
+
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+
+    // Multiply and accumulate threadgroup elements
+    mma_op.mma(As, Bs);
+
+    // Prepare for next iteration
+    loader_a.next();
+    loader_b.next();
+  }
+
+  threadgroup_barrier(mem_flags::mem_none);
+
+  // Store results to device memory
+  {
+    // Adjust for simdgroup and thread locatio
+    int offset_m = c_row + mma_op.sm + mma_op.tm;
+    int offset_n = c_col + mma_op.sn + mma_op.tn;
+    C += offset_n;
+
+    if (offset_n >= gemm_params->N)
+      return;
+
+    short diff = gemm_params->N - offset_n;
+
+    STEEL_PRAGMA_UNROLL
+    for (int i = 0; i < mma_t::TM; i++) {
+
+      int cm = offset_m + i * mma_t::TM_stride;
+
+      int n = cm / jump_params->adj_out_hw;
+      int hw = cm % jump_params->adj_out_hw;
+      int oh = (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + base_oh;
+      int ow = (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + base_ow;
+
+      if(n < params->N && oh < params->oS[0] && ow < params->oS[1]) {
+
+        int offset_cm = n * params->out_strides[0] + oh * params->out_strides[1] + ow * params->out_strides[2];
+
+        STEEL_PRAGMA_UNROLL
+        for (int j = 0; j < mma_t::TN; j++) {
+          // Get accumulated result and associated offset in C
+          thread const auto& accum = mma_op.results[i * mma_t::TN + j].thread_elements();
+          int offset = offset_cm + (j * mma_t::TN_stride);
+
+          // Apply epilogue and output C
+          if (j * mma_t::TN_stride < diff) {
+            C[offset] = Epilogue::apply(accum[0]);
+          }
+
+          if (j * mma_t::TN_stride + 1 < diff) {
+            C[offset + 1] = Epilogue::apply(accum[1]);
+          }
+        }
+
+      }
+    }
+  }
+
+}
+
+#define instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn) \
+  template [[host_name("implicit_gemm_conv_2d_general_" #name "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn)]] \
+  [[kernel]] void implicit_gemm_conv_2d_general<itype, bm, bn, bk, wm, wn>( \
+      const device itype* A [[buffer(0)]], \
+      const device itype* B [[buffer(1)]], \
+      device itype* C [[buffer(2)]], \
+      const constant MLXConvParams<2>* params [[buffer(3)]], \
+      const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]], \
+      const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]], \
+      const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]], \
+      const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]], \
+      uint3 tid [[threadgroup_position_in_grid]], \
+      uint3 lid [[thread_position_in_threadgroup]], \
+      uint simd_gid [[simdgroup_index_in_threadgroup]], \
+      uint simd_lid [[thread_index_in_simdgroup]]);
+
+#define instantiate_implicit_2d_filter(name, itype, bm, bn, bk, wm, wn) \
+    instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn)
+
+#define instantiate_implicit_2d_blocks(name, itype) \
+    instantiate_implicit_2d_filter(name, itype, 32,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 64,  8, 16, 4, 1) \
+    instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \
+    instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2)
+
+instantiate_implicit_2d_blocks(float32, float);
+instantiate_implicit_2d_blocks(float16, half);
+instantiate_implicit_2d_blocks(bfloat16, bfloat16_t);

mlx/backend/metal/kernels/steel/conv/loader.h

@@ -0,0 +1,6 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h"
+#include "mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h"

mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h

@@ -0,0 +1,449 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderLargeFilter {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderLargeFilter(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_h(0),
+        weight_w(0) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Adjust for flip
+      if (params->flip) {
+        ih += (params->wS[0] - 1) * params->kdil[0];
+        iw += (params->wS[1] - 1) * params->kdil[1];
+      }
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2] + bj;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+      int ih = read_ih[i] + weight_h * params->kdil[0];
+      int iw = read_iw[i] + weight_w * params->kdil[1];
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) &&
+          (iw >= 0 && iw < params->iS[1])) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = src[i][j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_w < params->wS[1]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_w;
+      }
+
+      return;
+    }
+
+    weight_w = 0;
+
+    if (++weight_h < params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_h;
+      }
+
+      return;
+    }
+
+    weight_h = 0;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += gemm_params->inp_jump_c;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderSmallFilter {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  using mask_t = short;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  mask_t mask_h[n_rows];
+  mask_t mask_w[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderSmallFilter(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_h(0),
+        weight_w(0) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    int read_n[n_rows];
+    int read_ih[n_rows];
+    int read_iw[n_rows];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Adjust for flip
+      if (params->flip) {
+        ih += (params->wS[0] - 1) * params->kdil[0];
+        iw += (params->wS[1] - 1) * params->kdil[1];
+      }
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2] + bj;
+    }
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      mask_h[i] = 0;
+      mask_w[i] = 0;
+    }
+
+    for (short kh = 0; kh < params->wS[0]; kh++) {
+      short flip_h = params->flip ? params->wS[0] - kh - 1 : kh;
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; ++i) {
+        int n = read_n[i];
+        int ih = read_ih[i] + flip_h * params->kdil[0];
+
+        bool in_bounds = n < params->N && ih >= 0 && ih < params->iS[0];
+
+        mask_h[i] |= (in_bounds << kh);
+      }
+    }
+
+    for (short kw = 0; kw < params->wS[1]; kw++) {
+      short flip_w = params->flip ? params->wS[1] - kw - 1 : kw;
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; ++i) {
+        int iw = read_iw[i] + flip_w * params->kdil[1];
+
+        bool in_bounds = iw >= 0 && iw < params->iS[1];
+
+        mask_w[i] |= (in_bounds << kw);
+      }
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    mask_t h_mask = mask_t(1) << weight_h;
+    mask_t w_mask = mask_t(1) << weight_w;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Read from input if in bounds
+      if ((mask_h[i] & h_mask) && (mask_w[i] & w_mask)) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = src[i][j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_w < params->wS[1]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_w;
+      }
+
+      return;
+    }
+
+    weight_w = 0;
+
+    if (++weight_h < params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < n_rows; i++) {
+        src[i] += gemm_params->inp_jump_h;
+      }
+
+      return;
+    }
+
+    weight_h = 0;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += gemm_params->inp_jump_c;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoader {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size =
+      (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4);
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+
+  int weight_hw;
+
+  const int read_n;
+  const bool do_read;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoader(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj),
+        params(params_),
+        weight_hw(0),
+        read_n(offsets.y + bi),
+        do_read(read_n + n_rows * TROWS <= gemm_params_->N) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (BN != 8 || do_read) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BN; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = src[i * src_ld + j];
+        }
+      }
+    } else {
+      for (short i = 0; i < BN; i += TROWS) {
+        if ((read_n + i) < params->O) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = src[i * src_ld + j];
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    if (++weight_hw < (params->wS[1] * params->wS[0])) {
+      src += params->wt_strides[2];
+      return;
+    }
+
+    weight_hw = 0;
+
+    src += BK - (params->wS[1] * params->wS[0] - 1) * params->wt_strides[2];
+  }
+};
+
+} // namespace steel
+} // namespace mlx

mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h

@@ -0,0 +1,319 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <short n_channels_>
+struct ChannelHelper {
+  STEEL_CONST short n_channels = n_channels_;
+  STEEL_CONST short vec_size = n_channels_ <= 4 ? 4 : 8;
+  STEEL_CONST short excess = vec_size - n_channels_;
+};
+
+template <>
+struct ChannelHelper<1> {
+  STEEL_CONST short n_channels = 1;
+  STEEL_CONST short vec_size = 1;
+  STEEL_CONST short excess = 0;
+};
+
+template <>
+struct ChannelHelper<2> {
+  STEEL_CONST short n_channels = 2;
+  STEEL_CONST short vec_size = 2;
+  STEEL_CONST short excess = 0;
+};
+
+template <>
+struct ChannelHelper<3> {
+  STEEL_CONST short n_channels = 3;
+  STEEL_CONST short vec_size = 4;
+  STEEL_CONST short excess = 1;
+};
+
+template <>
+struct ChannelHelper<4> {
+  STEEL_CONST short n_channels = 4;
+  STEEL_CONST short vec_size = 4;
+  STEEL_CONST short excess = 0;
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short n_channels,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderSmallChannels {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = ChannelHelper<n_channels>::vec_size;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant ImplicitGemmConv2DParams* gemm_params;
+
+  short weight_hw;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderSmallChannels(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        gemm_params(gemm_params_),
+        weight_hw(thread_idx % TCOLS) {
+    int out_n_pixels = params->oS[0] * params->oS[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / out_n_pixels;
+      int hw = offset_nhw % out_n_pixels;
+      int oh = hw / params->oS[1];
+      int ow = hw % params->oS[1];
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] +
+          iw * params->in_strides[2];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (weight_hw >= params->wS[1] * params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+      return;
+    }
+
+    int wh = (weight_hw / params->wS[1]);
+    int ww = (weight_hw % params->wS[1]);
+
+    int flip_h = params->flip ? params->wS[0] - wh - 1 : wh;
+    int flip_w = params->flip ? params->wS[1] - ww - 1 : ww;
+
+    int weight_h = flip_h * params->kdil[0];
+    int weight_w = flip_w * params->kdil[1];
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+      int ih = read_ih[i] + weight_h;
+      int iw = read_iw[i] + weight_w;
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) &&
+          (iw >= 0 && iw < params->iS[1])) {
+        const device T* curr_src = src[i] + weight_h * params->in_strides[1] +
+            weight_w * params->in_strides[2];
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < n_channels; ++j) {
+          dst[is * dst_ld + j] = curr_src[j];
+        }
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = n_channels; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_hw += TCOLS;
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short n_channels,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoaderSmallChannels {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = ChannelHelper<n_channels>::vec_size;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+
+  int weight_hw;
+
+  const int read_n;
+  const bool do_read;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoaderSmallChannels(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant ImplicitGemmConv2DParams* gemm_params_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld),
+        params(params_),
+        weight_hw(thread_idx % TCOLS),
+        read_n(offsets.y + bi),
+        do_read(read_n + BN <= gemm_params_->N) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    if (bi >= BROWS || bj >= BCOLS)
+      return;
+
+    if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+
+      return;
+    }
+
+    const device T* curr_src = src + weight_hw * params->wt_strides[2];
+
+    if (BN != 8 || do_read) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BROWS; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < n_channels; j++) {
+          dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+        }
+
+        STEEL_PRAGMA_UNROLL
+        for (short j = n_channels; j < vec_size; j++) {
+          dst[i * dst_ld + j] = T(0);
+        }
+      }
+    } else {
+      for (short i = 0; i < BROWS; i += TROWS) {
+        if (((read_n + i) < params->O)) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < n_channels; j++) {
+            dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+          }
+
+          STEEL_PRAGMA_UNROLL
+          for (short j = n_channels; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_hw += TCOLS;
+  }
+};
+
+} // namespace steel
+} // namespace mlx

mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h

@@ -0,0 +1,288 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/conv/params.h"
+
+///////////////////////////////////////////////////////////////////////////////
+// Loading helper
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mlx {
+namespace steel {
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DInputBlockLoaderGeneral {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BM;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4;
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+
+  const constant MLXConvParams<2>* params;
+  const constant Conv2DGeneralJumpParams* jump_params;
+
+  const short base_wh;
+  const short base_ww;
+
+  short weight_h;
+  short weight_w;
+
+  const device T* src[n_rows];
+
+  int read_n[n_rows];
+  int read_ih[n_rows];
+  int read_iw[n_rows];
+
+  /* Constructor */
+  METAL_FUNC Conv2DInputBlockLoaderGeneral(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int4 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant Conv2DGeneralJumpParams* jump_params_,
+      const short base_wh_,
+      const short base_ww_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        params(params_),
+        jump_params(jump_params_),
+        base_wh(base_wh_),
+        base_ww(base_ww_),
+        weight_h(base_wh_),
+        weight_w(base_ww_) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; ++i) {
+      int offset_nhw = offsets.y + bi + i * TROWS;
+      int n = offset_nhw / jump_params->adj_out_hw;
+      int hw = offset_nhw % jump_params->adj_out_hw;
+      int oh =
+          (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + offsets.z;
+      int ow =
+          (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + offsets.w;
+
+      int ih = oh * params->str[0] - params->pad[0];
+      int iw = ow * params->str[1] - params->pad[1];
+
+      read_n[i] = n;
+      read_ih[i] = ih;
+      read_iw[i] = iw;
+
+      // Read from input if in bounds
+      src[i] = src_ + n * params->in_strides[0] + bj;
+    }
+  }
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) {
+      // Find bounds
+      int n = read_n[i];
+
+      int h_flip = params->flip ? params->wS[0] - weight_h - 1 : weight_h;
+      int w_flip = params->flip ? params->wS[1] - weight_w - 1 : weight_w;
+
+      int ih_dil = read_ih[i] + h_flip * params->kdil[0];
+      int iw_dil = read_iw[i] + w_flip * params->kdil[1];
+
+      int ih = ih_dil / params->idil[0];
+      int iw = iw_dil / params->idil[1];
+
+      size_t offset = ih * params->in_strides[1] + iw * params->in_strides[2];
+
+      // Read from input if in bounds
+      if ((n < params->N) && (ih_dil >= 0 && ih < params->iS[0]) &&
+          (iw_dil >= 0 && iw < params->iS[1])) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = (src[i])[offset + j];
+        }
+      }
+
+      // Zero pad otherwise
+      else {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; ++j) {
+          dst[is * dst_ld + j] = T(0);
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_w += jump_params->f_wgt_jump_w;
+    if (weight_w < params->wS[1]) {
+      return;
+    }
+
+    weight_w = base_ww;
+
+    weight_h += jump_params->f_wgt_jump_h;
+    if (weight_h < params->wS[0]) {
+      return;
+    }
+
+    weight_h = base_wh;
+
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < n_rows; i++) {
+      src[i] += BK;
+    }
+  }
+};
+
+template <
+    typename T,
+    short BM,
+    short BN,
+    short BK,
+    short tgp_size,
+    short tgp_padding = 0>
+struct Conv2DWeightBlockLoaderGeneral {
+  // Destination dimensions
+  STEEL_CONST short BROWS = BN;
+  STEEL_CONST short BCOLS = BK;
+
+  // Read dimensions
+  STEEL_CONST short dst_ld = BCOLS + tgp_padding;
+  STEEL_CONST short vec_size =
+      (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4);
+
+  // Thread read shape
+  STEEL_CONST short TCOLS = BCOLS / vec_size;
+  STEEL_CONST short TROWS = tgp_size / TCOLS;
+
+  // Rows / strided reads within the block
+  STEEL_CONST short n_rows = BROWS / TROWS;
+
+  // Leading dimension for src
+  const int src_ld;
+
+  // Thread location indices
+  const short thread_idx;
+  const short bi;
+  const short bj;
+
+  // threadgroup and device memory
+  threadgroup T* dst;
+  const device T* src;
+
+  const constant MLXConvParams<2>* params;
+  const constant Conv2DGeneralJumpParams* jump_params;
+
+  const short base_wh;
+  const short base_ww;
+
+  short weight_h;
+  short weight_w;
+
+  const int start_row;
+
+  /* Constructor */
+  METAL_FUNC Conv2DWeightBlockLoaderGeneral(
+      const device T* src_,
+      threadgroup T* dst_,
+      const int2 offsets,
+      const constant MLXConvParams<2>* params_,
+      const constant Conv2DGeneralJumpParams* jump_params_,
+      const short base_wh_,
+      const short base_ww_,
+      uint simd_group_id [[simdgroup_index_in_threadgroup]],
+      uint simd_lane_id [[thread_index_in_simdgroup]])
+      : src_ld(params_->wt_strides[0]),
+        thread_idx(simd_group_id * 32 + simd_lane_id),
+        bi(thread_idx / TCOLS),
+        bj(vec_size * (thread_idx % TCOLS)),
+        dst(dst_ + bi * dst_ld + bj),
+        src(src_ + bi * src_ld + bj),
+        params(params_),
+        jump_params(jump_params_),
+        base_wh(base_wh_),
+        base_ww(base_ww_),
+        weight_h(base_wh_),
+        weight_w(base_ww_),
+        start_row(offsets.y + bi) {}
+
+  /* Load from device memory into threadgroup memory - without bound checking */
+  METAL_FUNC void load_unsafe() const {
+    const device T* curr_src = src + weight_h * params->wt_strides[1] +
+        weight_w * params->wt_strides[2];
+
+    if ((start_row + BN <= params->O)) {
+      STEEL_PRAGMA_UNROLL
+      for (short i = 0; i < BN; i += TROWS) {
+        STEEL_PRAGMA_UNROLL
+        for (short j = 0; j < vec_size; j++) {
+          dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+        }
+      }
+    } else {
+      for (short i = 0; i < BN; i += TROWS) {
+        if ((start_row + i) < params->O) {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = curr_src[i * src_ld + j];
+          }
+        } else {
+          STEEL_PRAGMA_UNROLL
+          for (short j = 0; j < vec_size; j++) {
+            dst[i * dst_ld + j] = T(0);
+          }
+        }
+      }
+    }
+  }
+
+  /* Iteration helper */
+  METAL_FUNC void next() {
+    weight_w += jump_params->f_wgt_jump_w;
+    if (weight_w < params->wS[1]) {
+      return;
+    }
+
+    weight_w = base_ww;
+
+    weight_h += jump_params->f_wgt_jump_h;
+    if (weight_h < params->wS[0]) {
+      return;
+    }
+
+    weight_h = base_wh;
+
+    src += BK;
+  }
+};
+
+} // namespace steel
+} // namespace mlx

mlx/backend/metal/kernels/steel/conv/params.h

@@ -0,0 +1,62 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+template <int NDIM>
+struct MLXConvParams {
+  const int N; // Batch size
+  const int C; // In channels
+  const int O; // Out channels
+  const int iS[NDIM]; // Input spatial dim
+  const int wS[NDIM]; // Weight spatial dim
+  const int oS[NDIM]; // Output spatial dim
+  const int str[NDIM]; // Kernel strides
+  const int pad[NDIM]; // Input padding
+  const int kdil[NDIM]; // Kernel dilation
+  const int idil[NDIM]; // Input dilation
+  const size_t in_strides[NDIM + 2]; // In strides
+  const size_t wt_strides[NDIM + 2]; // Wt strides
+  const size_t out_strides[NDIM + 2]; // Out strides
+  const int groups; // Input channel groups
+  const bool flip;
+};
+
+namespace mlx {
+namespace steel {
+
+struct ImplicitGemmConv2DParams {
+  const int M;
+  const int N;
+  const int K;
+
+  const int gemm_k_iterations;
+
+  const int inp_jump_w;
+  const int inp_jump_h;
+  const int inp_jump_c;
+
+  const int tiles_n;
+  const int tiles_m;
+  const int swizzle_log;
+};
+
+struct Conv2DGeneralJumpParams {
+  const int f_wgt_jump_h;
+  const int f_wgt_jump_w;
+
+  const int f_out_jump_h;
+  const int f_out_jump_w;
+
+  const int adj_out_h;
+  const int adj_out_w;
+  const int adj_out_hw;
+  const int adj_implicit_m;
+};
+
+struct Conv2DGeneralBaseInfo {
+  int weight_base;
+  int weight_size;
+};
+
+} // namespace steel
+} // namespace mlx

mlx/backend/metal/kernels/steel/gemm/gemm.h

@@ -4,6 +4,7 @@
 
 #include "mlx/backend/metal/kernels/steel/gemm/loader.h"
 #include "mlx/backend/metal/kernels/steel/gemm/mma.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
 #include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
 #include "mlx/backend/metal/kernels/steel/utils.h"
 

mlx/backend/metal/kernels/steel/gemm/mma.h

@@ -2,9 +2,15 @@
 
 #pragma once
 
+#include <metal_simdgroup>
+#include <metal_simdgroup_matrix>
+#include <metal_stdlib>
+
 #include "mlx/backend/metal/kernels/steel/gemm/transforms.h"
 #include "mlx/backend/metal/kernels/steel/utils.h"
 
+using namespace metal;
+
 ///////////////////////////////////////////////////////////////////////////////
 // MMA helper
 ///////////////////////////////////////////////////////////////////////////////
@@ -167,6 +173,9 @@ struct BlockMMA {
     C += (sm + tm) * ldc + (tn + sn);
     dst_tile_dims -= short2(tn + sn, sm + tm);
 
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
     STEEL_PRAGMA_UNROLL
     for (int i = 0; i < TM; i++) {
       if (i * TM_stride < dst_tile_dims.y) {
@@ -236,6 +245,9 @@ struct BlockMMA {
     D += (sm + tm) * ldd + tn + sn;
     dst_tile_dims -= short2(tn + sn, sm + tm);
 
+    if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0)
+      return;
+
     STEEL_PRAGMA_UNROLL
     for (int i = 0; i < TM; i++) {
       if (i * TM_stride < dst_tile_dims.y) {

mlx/backend/metal/kernels/steel/host.h

@@ -1,5 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-#pragma once
-
-#include "mlx/backend/metal/kernels/steel/gemm/params.h"

mlx/backend/metal/kernels/steel/utils.h

@@ -3,7 +3,6 @@
 #pragma once
 
 #include <metal_stdlib>
-#include "mlx/backend/metal/kernels/steel/host.h"
 
 #define STEEL_CONST static constant constexpr const
 #define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")

mlx/backend/metal/kernels/ternary.h

@@ -0,0 +1,10 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+struct Select {
+  template <typename T>
+  T operator()(bool condition, T x, T y) {
+    return condition ? x : y;
+  }
+};

mlx/backend/metal/kernels/ternary.metal

@@ -0,0 +1,201 @@
+// Copyright © 2023 Apple Inc.
+
+#include <metal_integer>
+#include <metal_math>
+
+#include "mlx/backend/metal/kernels/utils.h"
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/ternary.h"
+
+template <typename T, typename Op>
+[[kernel]] void ternary_op_v(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    uint index [[thread_position_in_grid]]) {
+  d[index] = Op()(a[index], b[index], c[index]);
+}
+
+template <typename T, typename Op>
+[[kernel]] void ternary_op_g_nd1(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const size_t& a_strides,
+    constant const size_t& b_strides,
+    constant const size_t& c_strides,
+    uint index [[thread_position_in_grid]]) {
+  auto a_idx = elem_to_loc_1(index, a_strides);
+  auto b_idx = elem_to_loc_1(index, b_strides);
+  auto c_idx = elem_to_loc_1(index, c_strides);
+  d[index] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op>
+[[kernel]] void ternary_op_g_nd2(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const size_t a_strides[2],
+    constant const size_t b_strides[2],
+    constant const size_t c_strides[2],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_2(index, a_strides);
+  auto b_idx = elem_to_loc_2(index, b_strides);
+  auto c_idx = elem_to_loc_2(index, c_strides);
+  size_t out_idx = index.x + (size_t)grid_dim.x * index.y;
+  d[out_idx] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op>
+[[kernel]] void ternary_op_g_nd3(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const size_t a_strides[3],
+    constant const size_t b_strides[3],
+    constant const size_t c_strides[3],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto a_idx = elem_to_loc_3(index, a_strides);
+  auto b_idx = elem_to_loc_3(index, b_strides);
+  auto c_idx = elem_to_loc_3(index, c_strides);
+  size_t out_idx = index.x + (size_t)grid_dim.x * (index.y + (size_t)grid_dim.y * index.z);
+  d[out_idx] = Op()(a[a_idx], b[b_idx], c[c_idx]);
+}
+
+template <typename T, typename Op, int DIM>
+[[kernel]] void ternary_op_g_nd(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int shape[DIM],
+    constant const size_t a_strides[DIM],
+    constant const size_t b_strides[DIM],
+    constant const size_t c_strides[DIM],
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc_3_nd<DIM>(index, shape, a_strides, b_strides, c_strides);
+  size_t out_idx = index.x + (size_t)grid_dim.x * (index.y + (size_t)grid_dim.y * index.z);
+  d[out_idx] = Op()(a[idx.x], b[idx.y], c[idx.z]);
+}
+
+template <typename T, typename Op>
+[[kernel]] void ternary_op_g(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    constant const int* shape,
+    constant const size_t* a_strides,
+    constant const size_t* b_strides,
+    constant const size_t* c_strides,
+    constant const int& ndim,
+    uint3 index [[thread_position_in_grid]],
+    uint3 grid_dim [[threads_per_grid]]) {
+  auto idx = elem_to_loc_3_nd(index, shape, a_strides, b_strides, c_strides, ndim);
+  size_t out_idx = index.x + grid_dim.x * (index.y + grid_dim.y * index.z);
+  d[out_idx] = Op()(a[idx.x], b[idx.y], c[idx.z]);
+}
+
+#define instantiate_ternary_v(name, type, op) \
+  template [[host_name(name)]] \
+  [[kernel]] void ternary_op_v<type, op>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      uint index [[thread_position_in_grid]]); \
+
+#define instantiate_ternary_g(name, type, op) \
+  template [[host_name(name)]] \
+  [[kernel]] void ternary_op_g<type, op>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      constant const int* shape, \
+      constant const size_t* a_strides, \
+      constant const size_t* b_strides, \
+      constant const size_t* c_strides, \
+      constant const int& ndim, \
+      uint3 index [[thread_position_in_grid]], \
+      uint3 grid_dim [[threads_per_grid]]); \
+
+#define instantiate_ternary_g_dim(name, type, op, dims) \
+  template [[host_name(name "_" #dims)]] \
+  [[kernel]] void ternary_op_g_nd<type, op, dims>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      constant const int shape[dims], \
+      constant const size_t a_strides[dims], \
+      constant const size_t b_strides[dims], \
+      constant const size_t c_strides[dims], \
+      uint3 index [[thread_position_in_grid]], \
+      uint3 grid_dim [[threads_per_grid]]); \
+
+#define instantiate_ternary_g_nd(name, type, op) \
+  template [[host_name(name "_1")]] \
+  [[kernel]] void ternary_op_g_nd1<type, op>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      constant const size_t& a_strides, \
+      constant const size_t& b_strides, \
+      constant const size_t& c_strides, \
+      uint index [[thread_position_in_grid]]); \
+  template [[host_name(name "_2")]] \
+  [[kernel]] void ternary_op_g_nd2<type, op>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      constant const size_t a_strides[2], \
+      constant const size_t b_strides[2], \
+      constant const size_t c_strides[2], \
+      uint2 index [[thread_position_in_grid]], \
+      uint2 grid_dim [[threads_per_grid]]); \
+  template [[host_name(name "_3")]] \
+  [[kernel]] void ternary_op_g_nd3<type, op>( \
+      device const bool* a, \
+      device const type* b, \
+      device const type* c, \
+      device type* d, \
+      constant const size_t a_strides[3], \
+      constant const size_t b_strides[3], \
+      constant const size_t c_strides[3], \
+      uint3 index [[thread_position_in_grid]], \
+      uint3 grid_dim [[threads_per_grid]]); \
+  instantiate_ternary_g_dim(name, type, op, 4) \
+  instantiate_ternary_g_dim(name, type, op, 5) \
+
+#define instantiate_ternary_all(name, tname, type, op) \
+  instantiate_ternary_v("v" #name #tname, type, op) \
+  instantiate_ternary_g("g" #name #tname, type, op) \
+  instantiate_ternary_g_nd("g" #name #tname, type, op) \
+
+#define instantiate_ternary_types(name, op) \
+  instantiate_ternary_all(name, bool_, bool, op) \
+  instantiate_ternary_all(name, uint8, uint8_t, op) \
+  instantiate_ternary_all(name, uint16, uint16_t, op) \
+  instantiate_ternary_all(name, uint32, uint32_t, op) \
+  instantiate_ternary_all(name, uint64, uint64_t, op) \
+  instantiate_ternary_all(name, int8, int8_t, op) \
+  instantiate_ternary_all(name, int16, int16_t, op) \
+  instantiate_ternary_all(name, int32, int32_t, op) \
+  instantiate_ternary_all(name, int64, int64_t, op) \
+  instantiate_ternary_all(name, float16, half, op) \
+  instantiate_ternary_all(name, float32, float, op) \
+  instantiate_ternary_all(name, bfloat16, bfloat16_t, op) \
+  instantiate_ternary_all(name, complex64, complex64_t, op) \
+
+instantiate_ternary_types(select, Select)

mlx/backend/metal/kernels/unary.h

@@ -9,6 +9,10 @@
 #include "mlx/backend/metal/kernels/erf.h"
 #include "mlx/backend/metal/kernels/utils.h"
 
+namespace {
+constant float inf = metal::numeric_limits<float>::infinity();
+}
+
 struct Abs {
   template <typename T>
   T operator()(T x) {

mlx/backend/metal/kernels/utils.h

@@ -91,6 +91,30 @@ inline size_t elem_to_loc(
   return loc;
 }
 
+template <int NDIM>
+inline uint3 elem_to_loc_3_nd(
+    uint3 elem,
+    constant const int shape[NDIM],
+    constant const size_t a_strides[NDIM],
+    constant const size_t b_strides[NDIM],
+    constant const size_t c_strides[NDIM]) {
+  uint3 loc = {
+      static_cast<uint>(
+          elem.x * a_strides[NDIM - 1] + elem.y * a_strides[NDIM - 2]),
+      static_cast<uint>(
+          elem.x * b_strides[NDIM - 1] + elem.y * b_strides[NDIM - 2]),
+      static_cast<uint>(
+          elem.x * c_strides[NDIM - 1] + elem.y * c_strides[NDIM - 2])};
+  for (int d = NDIM - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * a_strides[d];
+    loc.y += l * b_strides[d];
+    loc.z += l * c_strides[d];
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
 template <int NDIM>
 inline uint2 elem_to_loc_2_nd(
     uint3 elem,
@@ -150,6 +174,30 @@ inline size_t elem_to_loc(
   return loc;
 }
 
+inline uint3 elem_to_loc_3_nd(
+    uint3 elem,
+    constant const int* shape,
+    constant const size_t* a_strides,
+    constant const size_t* b_strides,
+    constant const size_t* c_strides,
+    int ndim) {
+  uint3 loc = {
+      static_cast<uint>(
+          elem.x * a_strides[ndim - 1] + elem.y * a_strides[ndim - 2]),
+      static_cast<uint>(
+          elem.x * b_strides[ndim - 1] + elem.y * b_strides[ndim - 2]),
+      static_cast<uint>(
+          elem.x * c_strides[ndim - 1] + elem.y * c_strides[ndim - 2])};
+  for (int d = ndim - 3; d >= 0; --d) {
+    uint l = elem.z % shape[d];
+    loc.x += l * a_strides[d];
+    loc.y += l * b_strides[d];
+    loc.z += l * c_strides[d];
+    elem.z /= shape[d];
+  }
+  return loc;
+}
+
 inline uint2 elem_to_loc_2_nd(
     uint3 elem,
     constant const int* shape,

mlx/backend/metal/matmul.cpp

@@ -8,7 +8,7 @@
 #include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/kernels/defines.h"
-#include "mlx/backend/metal/kernels/steel/host.h"
+#include "mlx/backend/metal/kernels/steel/gemm/params.h"
 #include "mlx/backend/metal/matmul.h"
 #include "mlx/backend/metal/mps/gemm.h"
 #include "mlx/backend/metal/utils.h"

mlx/backend/metal/metal.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cstdlib>
 #include <future>
@@ -10,6 +10,10 @@
 
 namespace mlx::core::metal {
 
+bool is_available() {
+  return true;
+}
+
 int max_ops_per_buffer() {
   auto get_val = []() {
     if (const char* buff_str = std::getenv("MLX_MAX_OPS_PER_BUFFER")) {

mlx/backend/metal/metal.h

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
@@ -11,16 +11,52 @@
 
 namespace mlx::core::metal {
 
-constexpr bool is_available() {
-#ifdef _METAL_
-  return true;
-#else
-  return false;
-#endif
-}
+bool is_available();
 
-bool cache_enabled(void);
-void set_cache_enabled(bool enabled);
+/* Get the actively used memory in bytes.
+ *
+ * Note, this will not always match memory use reported by the system because
+ * it does not include cached memory buffers.
+ * */
+size_t get_active_memory();
+
+/* Get the peak amount of used memory in bytes.
+ *
+ * The maximum memory used is recorded from the beginning of the program
+ * execution.
+ * */
+size_t get_peak_memory();
+
+/* Get the cache size in bytes.
+ *
+ * The cache includes memory not currently used that has not been returned
+ * to the system allocator.
+ * */
+size_t get_cache_memory();
+
+/* Set the memory limit.
+ * Calls to malloc will wait on scheduled tasks if the limit is exceeded.  If
+ * there are no more scheduled tasks an error will be raised if relaxed
+ * is false or memory will be allocated (including the potential for
+ * swap) if relaxed is true.
+ *
+ * The memory limit defaults to 1.5 times the maximum recommended working set
+ * size reported by the device.
+ *
+ * Returns the previous memory limit.
+ * */
+size_t set_memory_limit(size_t limit, bool relaxed = true);
+
+/* Set the free cache limit.
+ * If using more than the given limit, free memory will be reclaimed
+ * from the cache on the next allocation. To disable the cache,
+ * set the limit to 0.
+ *
+ * The cache limit defaults to the memory limit.
+ *
+ * Returns the previous cache limit.
+ * */
+size_t set_cache_limit(size_t limit);
 
 void new_stream(Stream stream);
 std::shared_ptr<void> new_scoped_memory_pool();

mlx/backend/metal/primitives.cpp

@@ -1,11 +1,11 @@
 // Copyright © 2023-2024 Apple Inc.
-
 #include <algorithm>
 #include <cassert>
 #include <numeric>
 #include <sstream>
 
 #include "mlx/backend/common/binary.h"
+#include "mlx/backend/common/ternary.h"
 #include "mlx/backend/metal/copy.h"
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/kernels/defines.h"
@@ -43,24 +43,25 @@ void binary_op(
 
   std::ostringstream kname;
   switch (bopt) {
-    case ScalarScalar:
+    case BinaryOpType::ScalarScalar:
       kname << "ss";
       break;
-    case ScalarVector:
+    case BinaryOpType::ScalarVector:
       kname << "sv";
       break;
-    case VectorScalar:
+    case BinaryOpType::VectorScalar:
       kname << "vs";
       break;
-    case VectorVector:
+    case BinaryOpType::VectorVector:
       kname << "vv";
       break;
-    case General:
+    case BinaryOpType::General:
       kname << "g";
       break;
   }
   kname << op << type_to_name(a);
-  if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
+  if (bopt == BinaryOpType::General &&
+      shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
     kname << "_" << shape.size();
   }
 
@@ -80,7 +81,7 @@ void binary_op(
   set_array_buffer(compute_encoder, outputs[0], 2);
   set_array_buffer(compute_encoder, outputs[1], 3);
 
-  if (bopt == General) {
+  if (bopt == BinaryOpType::General) {
     auto ndim = shape.size();
     if (ndim > 3) {
       compute_encoder->setBytes(shape.data(), ndim * sizeof(int), 4);
@@ -141,24 +142,25 @@ void binary_op(
 
   std::ostringstream kname;
   switch (bopt) {
-    case ScalarScalar:
+    case BinaryOpType::ScalarScalar:
       kname << "ss";
       break;
-    case ScalarVector:
+    case BinaryOpType::ScalarVector:
       kname << "sv";
       break;
-    case VectorScalar:
+    case BinaryOpType::VectorScalar:
       kname << "vs";
       break;
-    case VectorVector:
+    case BinaryOpType::VectorVector:
       kname << "vv";
       break;
-    case General:
+    case BinaryOpType::General:
       kname << "g";
       break;
   }
   kname << op << type_to_name(a);
-  if (bopt == General && shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
+  if (bopt == BinaryOpType::General &&
+      shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
     kname << "_" << shape.size();
   }
 
@@ -173,7 +175,7 @@ void binary_op(
   set_array_buffer(compute_encoder, donate_b ? out : b, 1);
   set_array_buffer(compute_encoder, out, 2);
 
-  if (bopt == General) {
+  if (bopt == BinaryOpType::General) {
     auto ndim = shape.size();
     if (ndim > 3) {
       compute_encoder->setBytes(shape.data(), ndim * sizeof(int), 3);
@@ -202,7 +204,94 @@ void binary_op(
     compute_encoder->dispatchThreads(grid_dims, group_dims);
   } else {
     // Launch a 1D grid of threads
-    size_t nthreads = bopt == General ? out.size() : out.data_size();
+    size_t nthreads =
+        bopt == BinaryOpType::General ? out.size() : out.data_size();
+    MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+    NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+  }
+}
+
+void ternary_op(
+    const std::vector<array>& inputs,
+    array& out,
+    const std::string op) {
+  assert(inputs.size() == 3);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt, true /* donate_with_move */);
+
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto [shape, strides] = collapse_contiguous_dims(a, b, c, out);
+  auto& strides_a = strides[0];
+  auto& strides_b = strides[1];
+  auto& strides_c = strides[2];
+  auto& strides_out = strides[3];
+
+  std::ostringstream kname;
+  if (topt == TernaryOpType::General) {
+    kname << "g";
+    kname << op << type_to_name(b);
+    if (shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
+      kname << "_" << shape.size();
+    }
+  } else {
+    kname << "v";
+    kname << op << type_to_name(b);
+  }
+
+  auto& s = out.primitive().stream();
+  auto& d = metal::device(s.device);
+  auto kernel = d.get_kernel(kname.str());
+  auto compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder->setComputePipelineState(kernel);
+  set_array_buffer(compute_encoder, a, 0);
+  set_array_buffer(compute_encoder, b, 1);
+  set_array_buffer(compute_encoder, c, 2);
+  set_array_buffer(compute_encoder, out, 3);
+
+  if (topt == TernaryOpType::General) {
+    auto ndim = shape.size();
+    if (ndim > 3) {
+      compute_encoder->setBytes(shape.data(), ndim * sizeof(int), 4);
+      compute_encoder->setBytes(strides_a.data(), ndim * sizeof(size_t), 5);
+      compute_encoder->setBytes(strides_b.data(), ndim * sizeof(size_t), 6);
+      compute_encoder->setBytes(strides_c.data(), ndim * sizeof(size_t), 7);
+
+      if (ndim > MAX_BINARY_SPECIALIZED_DIMS) {
+        compute_encoder->setBytes(&ndim, sizeof(int), 8);
+      }
+    } else {
+      // The shape is implicit in the grid for <= 3D
+      compute_encoder->setBytes(strides_a.data(), ndim * sizeof(size_t), 4);
+      compute_encoder->setBytes(strides_b.data(), ndim * sizeof(size_t), 5);
+      compute_encoder->setBytes(strides_c.data(), ndim * sizeof(size_t), 6);
+    }
+
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+    NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::binary] Must use 1024 sized block");
+    }
+    MTL::Size group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder->dispatchThreads(grid_dims, group_dims);
+  } else {
+    // Launch a 1D grid of threads
+    size_t nthreads = out.data_size();
     MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size > nthreads) {
@@ -430,8 +519,6 @@ void ArgReduce::eval_gpu(const std::vector<array>& inputs, array& out) {
     compute_encoder->setBytes(&ndim, sizeof(size_t), 5);
     compute_encoder->setBytes(&axis_stride, sizeof(size_t), 6);
     compute_encoder->setBytes(&axis_size, sizeof(size_t), 7);
-    compute_encoder->setThreadgroupMemoryLength(
-        simd_size * (sizeof(uint32_t) + in.itemsize()), 0);
     compute_encoder->dispatchThreads(grid_dims, group_dims);
   }
 }
@@ -621,6 +708,10 @@ void Multiply::eval_gpu(const std::vector<array>& inputs, array& out) {
   binary_op(inputs, out, "mul");
 }
 
+void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
+  ternary_op(inputs, out, "select");
+}
+
 void Negative::eval_gpu(const std::vector<array>& inputs, array& out) {
   unary_op(inputs, out, "neg");
 }
@@ -691,7 +782,6 @@ void RandomBits::eval_gpu(const std::vector<array>& inputs, array& out) {
   // organize into grid nkeys x elem_per_key
   MTL::Size grid_dims = MTL::Size(num_keys, half_size + odd, 1);
   NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
-  auto nthreads = std::min(num_keys * (half_size + odd), thread_group_size);
   MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
   auto compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);

mlx/backend/metal/quantized.cpp

@@ -41,8 +41,35 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
   int B = x.size() / D;
   int O = out.shape(-1);
   if (transpose_) {
+    // Route to the fast qmv kernel that has no bounds checking
+    if (B < 6 && O % 8 == 0 && D % 512 == 0 && D >= 512) {
+      std::ostringstream kname;
+      kname << "qmv_" << type_to_name(out) << "_gs_" << group_size_ << "_b_"
+            << bits_ << "_fast";
+
+      // Encode and dispatch kernel
+      auto compute_encoder = d.get_command_encoder(s.index);
+      auto kernel = d.get_kernel(kname.str());
+      compute_encoder->setComputePipelineState(kernel);
+
+      int bo = 8;
+      int bd = 32;
+      MTL::Size group_dims = MTL::Size(bd, 2, 1);
+      MTL::Size grid_dims = MTL::Size(1, O / bo, B);
+
+      set_array_buffer(compute_encoder, w, 0);
+      set_array_buffer(compute_encoder, scales, 1);
+      set_array_buffer(compute_encoder, biases, 2);
+      set_array_buffer(compute_encoder, x, 3);
+      set_array_buffer(compute_encoder, out, 4);
+      compute_encoder->setBytes(&D, sizeof(int), 5);
+      compute_encoder->setBytes(&O, sizeof(int), 6);
+
+      compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+    }
+
     // Route to the qmv kernel
-    if (B < 6) {
+    else if (B < 6) {
       std::ostringstream kname;
       kname << "qmv_" << type_to_name(out) << "_gs_" << group_size_ << "_b_"
             << bits_;
@@ -52,9 +79,9 @@ void QuantizedMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {
       auto kernel = d.get_kernel(kname.str());
       compute_encoder->setComputePipelineState(kernel);
 
-      int bo = std::min(32, O);
+      int bo = 8;
       int bd = 32;
-      MTL::Size group_dims = MTL::Size(bd, bo, 1);
+      MTL::Size group_dims = MTL::Size(bd, 2, 1);
       MTL::Size grid_dims = MTL::Size(1, (O + bo - 1) / bo, B);
 
       set_array_buffer(compute_encoder, w, 0);

mlx/backend/metal/reduce.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <algorithm>
 #include <cassert>
@@ -130,15 +130,8 @@ void row_reduce_general_dispatch(
     const Stream& s) {
   Dtype out_dtype = out.dtype();
   bool is_out_64b_int = is_64b_int(out_dtype);
-  auto kernel = (is_out_64b_int)
-      ? d.get_kernel(
-            "row_reduce_general_no_atomics_" + op_name + type_to_name(in))
-      : d.get_kernel("row_reduce_general_" + op_name + type_to_name(in));
-
-  compute_encoder->setComputePipelineState(kernel);
 
   // Prepare the arguments for the kernel
-  int n_reads = REDUCE_N_READS;
   size_t reduction_size = plan.shape.back();
   auto shape = plan.shape;
   auto strides = plan.strides;
@@ -160,32 +153,72 @@ void row_reduce_general_dispatch(
   }
   int ndim = shape.size();
 
-  // Each thread group is responsible for 1 output
-  NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
-  thread_group_size =
-      std::min((reduction_size + n_reads - 1) / n_reads, thread_group_size);
+  // Determine dispatch kernel
+  std::ostringstream kname;
 
-  // Align thread group size with simd_size
-  uint simd_size = kernel->threadExecutionWidth();
-  thread_group_size =
-      (thread_group_size + simd_size - 1) / simd_size * simd_size;
-  assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
+  bool is_small = non_row_reductions * reduction_size < 32;
+  bool is_med = non_row_reductions * reduction_size <= 256;
+  is_out_64b_int &= !is_small && !is_med;
 
-  // Launch enough thread groups for each output
-  size_t n_threads = out.size() * thread_group_size;
-  MTL::Size grid_dims = MTL::Size(n_threads, non_row_reductions, 1);
-  MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+  std::string small_desc = "_";
+  if (is_small) {
+    small_desc = "_small_";
+  } else if (is_med) {
+    small_desc = "_med_";
+  }
 
-  if (is_out_64b_int == false || non_row_reductions == 1) {
+  small_desc = is_out_64b_int ? "_no_atomics_" : small_desc;
+
+  kname << "row_reduce_general" << small_desc << op_name << type_to_name(in);
+
+  auto kernel = d.get_kernel(kname.str());
+  compute_encoder->setComputePipelineState(kernel);
+
+  // Get dispatch grid dims
+  MTL::Size grid_dims;
+  MTL::Size group_dims;
+
+  // Each thread handles one output
+  if (is_small) {
+    grid_dims = MTL::Size(out.size(), 1, 1);
+    group_dims = MTL::Size(std::min(1024ul, out.size()), 1, 1);
+  }
+  // Each simdgroup handles one output
+  else if (is_med) {
+    grid_dims = MTL::Size(out.size() * 32, 1, 1);
+    group_dims = MTL::Size(std::min(8ul, out.size()) * 32, 1, 1);
+  }
+  // Each theadgroup handles one output
+  else {
+    int n_reads = REDUCE_N_READS;
+    NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    thread_group_size =
+        std::min((reduction_size + n_reads - 1) / n_reads, thread_group_size);
+
+    // Align thread group size with simd_size
+    uint simd_size = kernel->threadExecutionWidth();
+    thread_group_size =
+        (thread_group_size + simd_size - 1) / simd_size * simd_size;
+    assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
+
+    // Launch enough thread groups for each output
+    size_t n_threads = out.size() * thread_group_size;
+    grid_dims = MTL::Size(n_threads, non_row_reductions, 1);
+    group_dims = MTL::Size(thread_group_size, 1, 1);
+  }
+
+  // Dispatch kernel
+  if (!is_out_64b_int || non_row_reductions == 1) {
     // Set the arguments for the kernel
     set_array_buffer(compute_encoder, in, 0);
     set_array_buffer(compute_encoder, out, 1);
     compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
     compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
-    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(&non_row_reductions, sizeof(size_t), 4);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 5);
     compute_encoder->setBytes(
-        strides.data(), strides.size() * sizeof(size_t), 5);
-    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+        strides.data(), strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
     compute_encoder->dispatchThreads(grid_dims, group_dims);
 
   } else {
@@ -203,10 +236,11 @@ void row_reduce_general_dispatch(
     set_array_buffer(compute_encoder, intermediate, 1);
     compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
     compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
-    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 4);
+    compute_encoder->setBytes(&non_row_reductions, sizeof(size_t), 4);
+    compute_encoder->setBytes(shape.data(), shape.size() * sizeof(int), 5);
     compute_encoder->setBytes(
-        strides.data(), strides.size() * sizeof(size_t), 5);
-    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+        strides.data(), strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
     compute_encoder->dispatchThreads(grid_dims, group_dims);
 
     // Set up second dispatch
@@ -230,24 +264,27 @@ void row_reduce_general_dispatch(
     set_array_buffer(compute_encoder, out, 1);
     compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
     compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(&non_row_reductions, sizeof(size_t), 4);
     compute_encoder->setBytes(
-        new_shape.data(), new_shape.size() * sizeof(int), 4);
+        new_shape.data(), new_shape.size() * sizeof(int), 5);
     compute_encoder->setBytes(
-        new_strides.data(), new_strides.size() * sizeof(size_t), 5);
-    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+        new_strides.data(), new_strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
 
     // Each thread group is responsible for 1 output
-    thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+    int n_reads = REDUCE_N_READS;
+    size_t thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     thread_group_size =
         std::min((reduction_size + n_reads - 1) / n_reads, thread_group_size);
 
     // Align thread group size with simd_size
+    uint simd_size = kernel->threadExecutionWidth();
     thread_group_size =
         (thread_group_size + simd_size - 1) / simd_size * simd_size;
     assert(thread_group_size <= kernel->maxTotalThreadsPerThreadgroup());
 
     // Launch enough thread groups for each output
-    n_threads = thread_group_size;
+    size_t n_threads = thread_group_size;
     grid_dims = MTL::Size(n_threads, out.size(), 1);
     group_dims = MTL::Size(thread_group_size, 1, 1);
 
@@ -417,11 +454,12 @@ void strided_reduce_general_dispatch(
     set_array_buffer(compute_encoder, out, 1);
     compute_encoder->setBytes(&reduction_size, sizeof(size_t), 2);
     compute_encoder->setBytes(&out_size, sizeof(size_t), 3);
+    compute_encoder->setBytes(&reduction_size, sizeof(size_t), 4);
     compute_encoder->setBytes(
-        new_shape.data(), new_shape.size() * sizeof(int), 4);
+        new_shape.data(), new_shape.size() * sizeof(int), 5);
     compute_encoder->setBytes(
-        new_strides.data(), new_strides.size() * sizeof(size_t), 5);
-    compute_encoder->setBytes(&ndim, sizeof(int), 6);
+        new_strides.data(), new_strides.size() * sizeof(size_t), 6);
+    compute_encoder->setBytes(&ndim, sizeof(int), 7);
 
     // Each thread group is responsible for 1 output
     size_t n_reads = REDUCE_N_READS;

mlx/backend/metal/rope.cpp

@@ -1,8 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/metal/utils.h"
-#include "mlx/fast.h"
-#include "mlx/primitives.h"
+#include "mlx/fast_primitives.h"
 
 namespace mlx::core::fast {
 

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -0,0 +1,222 @@
+//
+//  scaled_dot_product_attention.cpp
+//  mlx
+
+#include <algorithm>
+#include <cassert>
+#include <numeric>
+#include <sstream>
+
+#include "mlx/backend/metal/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels/scaled_dot_product_attention_params.h"
+#include "mlx/backend/metal/metal.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/primitives.h"
+#include "mlx/utils.h"
+
+namespace mlx::core::fast {
+
+namespace {
+
+void sdpa_metal(
+    const Stream& s,
+    metal::Device& d,
+    const array& q,
+    const array& k,
+    const array& v,
+    const array& p_lse,
+    const array& p_rowmaxes,
+    const array& o_partial,
+    const uint heads,
+    const uint tile_size,
+    const uint n_tiles,
+    const float alpha,
+    array& out,
+    std::vector<array>& temporaries) {
+  std::ostringstream kname_partials;
+
+  kname_partials << "fast_inference_sdpa_compute_partials_";
+
+  std::ostringstream kname_reduce;
+  std::string delimiter = "_";
+  kname_reduce << "fast_inference_sdpa_reduce_tiles" + delimiter;
+
+  for (const auto& arr : {k, v, out}) {
+    if (arr.dtype() != q.dtype()) {
+      throw std::runtime_error(
+          "[ScaledDotProductAttention::eval_gpu]: expected matching dtypes for q,k,v,o");
+    }
+  }
+
+  if (q.dtype() == float32) {
+    kname_partials << "float" + delimiter;
+    kname_reduce << "float";
+  } else if (q.dtype() == float16) {
+    kname_partials << "half" + delimiter;
+    kname_reduce << "half";
+  } else {
+    throw std::runtime_error(
+        "[ScaledDotProductAttention::eval_gpu]: unexpected dtype found for queries: expected either float32 or float16.");
+  }
+
+  std::string kname_suffix_tile_size = std::to_string(tile_size) + delimiter;
+
+  uint nsimd = 8;
+  std::string kname_suffix_nsimdgroups = std::to_string(nsimd);
+
+  // maximum number of splits == 128 at the moment (reserved tile registers in
+  // reduction kernel). this is arbitrary and could be changed in the shader.
+
+  std::string kname_suffix = kname_suffix_tile_size + kname_suffix_nsimdgroups;
+  kname_partials << kname_suffix;
+  auto compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = d.get_kernel(kname_partials.str());
+  compute_encoder->setComputePipelineState(kernel);
+
+  constexpr const uint batch = 1;
+  MTL::Size grid_dims = MTL::Size(heads, n_tiles, batch);
+  MTL::Size group_dims = MTL::Size(32, nsimd, 1);
+
+  const uint64_t KV_sequence_length = k.shape(-2);
+  const uint query_sequence_length = q.shape(-2);
+  const uint n_q_heads = q.shape(1);
+  const uint n_kv_heads = k.shape(1);
+
+  MLXScaledDotProductAttentionParams params{
+      query_sequence_length, n_q_heads, n_kv_heads, n_tiles, alpha};
+
+  set_array_buffer(compute_encoder, q, 0);
+  set_array_buffer(compute_encoder, k, 1);
+  set_array_buffer(compute_encoder, v, 2);
+  compute_encoder->setBytes(&KV_sequence_length, sizeof(KV_sequence_length), 3);
+  compute_encoder->setBytes(
+      &params, sizeof(MLXScaledDotProductAttentionParams), 4);
+  set_array_buffer(compute_encoder, o_partial, 5);
+  set_array_buffer(compute_encoder, p_lse, 6);
+  set_array_buffer(compute_encoder, p_rowmaxes, 7);
+
+  constexpr const uint tgroupMemorySize = 32768;
+  compute_encoder->setThreadgroupMemoryLength(tgroupMemorySize, 0);
+  compute_encoder->dispatchThreadgroups(grid_dims, group_dims);
+
+  {
+    auto kernel_accum = d.get_kernel(kname_reduce.str());
+    compute_encoder->setComputePipelineState(kernel_accum);
+    set_array_buffer(compute_encoder, o_partial, 0);
+    set_array_buffer(compute_encoder, p_lse, 1);
+    set_array_buffer(compute_encoder, p_rowmaxes, 2);
+    compute_encoder->setBytes(
+        &params, sizeof(MLXScaledDotProductAttentionParams), 3);
+    set_array_buffer(compute_encoder, out, 4);
+
+    MTL::Size grid_dims_reduce = MTL::Size(heads, 1, batch);
+    MTL::Size group_dims_reduce = MTL::Size(128, 1, 1);
+
+    compute_encoder->dispatchThreadgroups(grid_dims_reduce, group_dims_reduce);
+
+    d.get_command_buffer(s.index)->addCompletedHandler(
+        [temporaries](MTL::CommandBuffer*) mutable { temporaries.clear(); });
+    return;
+  }
+}
+} // namespace
+
+void ScaledDotProductAttention::eval_gpu(
+    const std::vector<array>& inputs,
+    array& out) {
+  assert(inputs.size() >= 3);
+  if (!is_floating_point(out.dtype())) {
+    throw std::runtime_error(
+        "[ScaledDotProductAttention] Does not yet support non-floating point types.");
+  }
+
+  if (inputs.size() == 4) {
+    out = fallback_(inputs)[0];
+    return;
+  }
+
+  out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+
+  auto& q_pre = inputs[0];
+  auto& k_pre = inputs[1];
+  auto& v_pre = inputs[2];
+  auto& o = out;
+  /////////////////////////////////////////////////////////////////////////////
+  // Init checks and prep
+
+  // Keep a vector with copies to be cleared in the completed buffer to release
+  // the arrays
+  std::vector<array> temporaries;
+  auto check_transpose = [&temporaries, &s](const array& arr) {
+    auto stx = arr.strides()[arr.ndim() - 2];
+    auto sty = arr.strides()[arr.ndim() - 1];
+    if (stx == arr.shape(-1) && sty == 1) {
+      return arr;
+    } else {
+      array arr_copy(arr.shape(), arr.dtype(), nullptr, {});
+      copy_gpu(arr, arr_copy, CopyType::General, s);
+      temporaries.push_back(arr_copy);
+      size_t stx = arr.shape(-1);
+      return arr_copy;
+    }
+  };
+
+  auto q = check_transpose(q_pre);
+  auto k = check_transpose(k_pre);
+  auto v = check_transpose(v_pre);
+
+  const int heads = q.shape(-3);
+  int tile_size = 64;
+  const int kv_seq_len = k.shape(-2);
+  if (kv_seq_len > 8000) {
+    tile_size = 128;
+  }
+  if (kv_seq_len > 16000) {
+    tile_size = 256;
+  }
+  if (kv_seq_len > 32000) {
+    tile_size = 512;
+  }
+
+  const int n_tiles = (kv_seq_len + tile_size - 1) / tile_size;
+
+  array o_partials(
+      {q.shape(-4), q.shape(-3), q.shape(-2), n_tiles * v.shape(-1)},
+      float32,
+      nullptr,
+      {});
+  o_partials.set_data(allocator::malloc_or_wait(o_partials.nbytes()));
+
+  array p_lse(
+      {q.shape(-4), q.shape(-3), q.shape(-2), n_tiles}, float32, nullptr, {});
+  array p_rowmaxes(
+      {q.shape(-4), q.shape(-3), q.shape(-2), n_tiles}, float32, nullptr, {});
+  p_lse.set_data(allocator::malloc_or_wait(p_lse.nbytes()));
+  p_rowmaxes.set_data(allocator::malloc_or_wait(p_rowmaxes.nbytes()));
+
+  temporaries.push_back(p_lse);
+  temporaries.push_back(p_rowmaxes);
+  temporaries.push_back(o_partials);
+
+  return sdpa_metal(
+      s,
+      d,
+      q,
+      k,
+      v,
+      p_lse,
+      p_rowmaxes,
+      o_partials,
+      heads,
+      tile_size,
+      n_tiles,
+      scale_,
+      out,
+      temporaries);
+}
+
+} // namespace mlx::core::fast

mlx/backend/no_metal/metal.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <stdexcept>
 
@@ -6,6 +6,10 @@
 
 namespace mlx::core::metal {
 
+bool is_available() {
+  return false;
+}
+
 void new_stream(Stream) {}
 std::shared_ptr<void> new_scoped_memory_pool() {
   return nullptr;
@@ -19,10 +23,21 @@ std::function<void()> make_task(
       "[metal::make_task] Cannot make GPU task without metal backend");
 }
 
-// No cache for CPU only
-bool cache_enabled(void) {
-  return false;
+// No-ops when Metal is not available.
+size_t get_active_memory() {
+  return 0;
+}
+size_t get_peak_memory() {
+  return 0;
+}
+size_t get_cache_memory() {
+  return 0;
+}
+size_t set_memory_limit(size_t, bool) {
+  return 0;
+}
+size_t set_cache_limit(size_t) {
+  return 0;
 }
-void set_cache_enabled(bool) {}
 
 } // namespace mlx::core::metal

mlx/backend/no_metal/primitives.cpp

@@ -1,7 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/primitives.h"
-#include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
 
 #define NO_GPU_MULTI(func)                                             \
   void func::eval_gpu(                                                 \
@@ -80,6 +80,7 @@ NO_GPU(Reshape)
 NO_GPU(Round)
 NO_GPU(Scan)
 NO_GPU(Scatter)
+NO_GPU(Select)
 NO_GPU(Sigmoid)
 NO_GPU(Sign)
 NO_GPU(Sin)
@@ -98,6 +99,7 @@ NO_GPU(Transpose)
 
 namespace fast {
 NO_GPU_MULTI(RoPE)
+NO_GPU(ScaledDotProductAttention)
 } // namespace fast
 
 } // namespace mlx::core

mlx/compile.cpp

@@ -7,13 +7,14 @@
 
 #include "mlx/allocator.h"
 #include "mlx/compile.h"
+#include "mlx/compile_impl.h"
 #include "mlx/primitives.h"
 #include "mlx/transforms.h"
 #include "mlx/transforms_impl.h"
 
 namespace mlx::core {
 
-constexpr int max_compile_depth = 10;
+constexpr int max_compile_depth = 11;
 
 bool is_unary(const Primitive& p) {
   return (
@@ -47,6 +48,10 @@ bool is_binary(const Primitive& p) {
       typeid(p) == typeid(Subtract));
 }
 
+bool is_ternary(const Primitive& p) {
+  return typeid(p) == typeid(Select);
+}
+
 bool is_broadcast(const Primitive& p) {
   return typeid(p) == typeid(Broadcast);
 }
@@ -55,19 +60,22 @@ bool is_noop(const Primitive& p) {
   return typeid(p) == typeid(Copy) || typeid(p) == typeid(StopGradient);
 }
 
-bool is_fusable(const Primitive& p) {
-  return is_unary(p) || is_binary(p) || is_broadcast(p) || is_noop(p);
+bool is_reduction(const Primitive& p) {
+  return typeid(p) == typeid(Reduce) || typeid(p) == typeid(ArgReduce);
 }
 
-namespace detail {
+bool is_fusable(const Primitive& p) {
+  return is_unary(p) || is_binary(p) || is_ternary(p) || is_broadcast(p) ||
+      is_noop(p);
+}
 
-std::vector<array> compile_replace(
-    const std::vector<array>& tape,
-    const std::vector<array>& trace_inputs,
-    const std::vector<array>& trace_outputs,
-    const std::vector<array>& inputs);
-
-} // namespace detail
+bool allows_shapeless(const Primitive& p) {
+  return typeid(p) == typeid(Compiled) || is_unary(p) || is_binary(p) ||
+      is_noop(p) || is_reduction(p) || typeid(p) == typeid(Softmax) ||
+      typeid(p) == typeid(Sort) || typeid(p) == typeid(ArgSort) ||
+      typeid(p) == typeid(ArgPartition) || typeid(p) == typeid(Partition) ||
+      typeid(p) == typeid(Select);
+}
 
 Compiled::Compiled(
     Stream stream,
@@ -123,6 +131,23 @@ void Compiled::print(std::ostream& os) {
   }
 }
 
+std::vector<std::vector<int>> Compiled::output_shapes(
+    const std::vector<array>& inputs) {
+  size_t nd = 0;
+  for (auto& in : inputs) {
+    nd = std::max(nd, in.ndim());
+  }
+  std::vector<int> out_shape(nd, 0);
+  for (auto& in : inputs) {
+    auto dd = nd - in.ndim();
+    for (auto i = dd; i < nd; ++i) {
+      out_shape[i] = std::max(out_shape[i], in.shape()[i - dd]);
+    }
+  }
+  // All outputs have the same shape
+  return std::vector<std::vector<int>>(outputs_.size(), out_shape);
+}
+
 namespace detail {
 
 CompileMode& compile_mode() {
@@ -180,21 +205,30 @@ struct CompilerCache {
     std::vector<array> outputs;
     std::vector<array> tape;
     bool empty{true};
+    std::vector<uint64_t> constants;
   };
 
   // Returns a reference to a CacheEntry which can be updated
   // by the caller to avoid copying large tapes / inputs / outputs
-  CacheEntry& find(size_t fun_id, const std::vector<array>& inputs) {
+  CacheEntry& find(
+      size_t fun_id,
+      const std::vector<array>& inputs,
+      bool shapeless,
+      const std::vector<uint64_t>& constants) {
     // Try to find the entry
     auto [entry_it, inserted] = cache_.insert({fun_id, {}});
     auto& entries = entry_it->second;
-    auto is_match = [](const std::vector<array>& in1,
-                       const std::vector<array>& in2) {
+    auto is_match = [shapeless](
+                        const std::vector<array>& in1,
+                        const std::vector<array>& in2) {
       if (in1.size() != in2.size()) {
         return false;
       }
       for (int i = 0; i < in1.size(); ++i) {
-        if (in1[i].shape() != in2[i].shape()) {
+        if (in1[i].ndim() != in2[i].ndim()) {
+          return false;
+        }
+        if (!shapeless && in1[i].shape() != in2[i].shape()) {
           return false;
         }
         if (in1[i].dtype() != in2[i].dtype()) {
@@ -210,7 +244,7 @@ struct CompilerCache {
     // more easily searchable structure.
     for (auto& entry : entries) {
       // Check the inputs match and return if so
-      if (is_match(inputs, entry.inputs)) {
+      if (is_match(inputs, entry.inputs) && constants == entry.constants) {
         return entry;
       }
     }
@@ -319,6 +353,9 @@ void compile_simplify(
       case 1:
         v = *a.data<uint8_t>();
         break;
+      case 2:
+        v = *a.data<uint16_t>();
+        break;
       case 4:
         v = *a.data<uint32_t>();
         break;
@@ -648,7 +685,8 @@ std::vector<array> compile_replace(
     const std::vector<array>& tape,
     const std::vector<array>& trace_inputs,
     const std::vector<array>& trace_outputs,
-    const std::vector<array>& inputs) {
+    const std::vector<array>& inputs,
+    bool shapeless) {
   std::unordered_map<uintptr_t, array> trace_to_real;
   for (int i = 0; i < inputs.size(); ++i) {
     trace_to_real.insert({trace_inputs[i].id(), inputs[i]});
@@ -666,18 +704,29 @@ std::vector<array> compile_replace(
         real_inputs.push_back(trace_to_real.at(in.id()));
       }
       if (a.siblings().empty()) {
+        auto shape =
+            shapeless ? a.primitive().output_shapes(real_inputs)[0] : a.shape();
         auto real_a = array(
-            a.shape(), a.dtype(), a.primitive_ptr(), std::move(real_inputs));
+            std::move(shape),
+            a.dtype(),
+            a.primitive_ptr(),
+            std::move(real_inputs));
         trace_to_real.insert({a.id(), std::move(real_a)});
       } else {
         // Ensure the order is correct for multi-output primitives
-        std::vector<std::vector<int>> shapes;
         std::vector<Dtype> types;
         auto trace_out = a.outputs();
         for (auto& o : trace_out) {
-          shapes.push_back(o.shape());
           types.push_back(o.dtype());
         }
+        std::vector<std::vector<int>> shapes;
+        if (shapeless) {
+          shapes = a.primitive().output_shapes(real_inputs);
+        } else {
+          for (auto& o : trace_out) {
+            shapes.push_back(o.shape());
+          }
+        }
         auto real_out =
             array::make_arrays(shapes, types, a.primitive_ptr(), real_inputs);
         for (int i = 0; i < trace_out.size(); ++i) {
@@ -694,13 +743,35 @@ std::vector<array> compile_replace(
   return outputs;
 }
 
+void compile_validate_shapeless(const std::vector<array>& tape) {
+  for (auto& t : tape) {
+    if (!t.has_primitive()) {
+      continue;
+    }
+    auto& p = t.primitive();
+    if (allows_shapeless(p)) {
+      continue;
+    }
+
+    std::ostringstream msg;
+    msg << "[compile] Cannot compile primitive ";
+    p.print(msg);
+    msg << " with shapeless enabled.";
+    throw std::invalid_argument(msg.str());
+  }
+}
+
 std::function<std::vector<array>(const std::vector<array>&)> compile(
     const std::function<std::vector<array>(const std::vector<array>&)>& fun,
-    size_t fun_id) {
-  if (compile_mode() == CompileMode::disabled) {
+    size_t fun_id,
+    bool shapeless /* = false */,
+    std::vector<uint64_t> constants /* = {} */) {
+  if (compile_mode() == CompileMode::disabled ||
+      !(compile_available_for_device(default_device()))) {
     return fun;
   }
-  return [fun, fun_id](const std::vector<array>& inputs) {
+  return [fun, fun_id, shapeless, constants = std::move(constants)](
+             const std::vector<array>& inputs) {
     // If the inputs are tracers, trace the original graph
     if (std::any_of(inputs.begin(), inputs.end(), [](auto& in) {
           return in.is_tracer();
@@ -709,12 +780,14 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
     }
 
     // Find a cache entry with the correct inputs
-    auto& entry = compiler_cache().find(fun_id, inputs);
+    auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);
 
     // No matching cache entry existed, so compile
     if (entry.empty) {
       // Mark the entry as not empty since we are about to fill it
       entry.empty = false;
+      // Set the constants
+      entry.constants = std::move(constants);
       // Trace to build the graph
       std::tie(entry.inputs, entry.outputs) = compile_trace(fun, inputs);
 
@@ -736,11 +809,16 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
       if (compile_mode() != CompileMode::no_fuse) {
         compile_fuse(entry.tape, parents_map, entry.inputs, entry.outputs);
       }
+
+      if (shapeless) {
+        compile_validate_shapeless(entry.tape);
+      }
     }
 
     // At this point we must have a tape, now replace the placeholders
     // with real arrays that can be evaluated
-    return compile_replace(entry.tape, entry.inputs, entry.outputs, inputs);
+    return compile_replace(
+        entry.tape, entry.inputs, entry.outputs, inputs, shapeless);
   };
 }
 
@@ -751,12 +829,13 @@ void compile_erase(size_t fun_id) {
 } // namespace detail
 
 std::function<std::vector<array>(const std::vector<array>&)> compile(
-    const std::function<std::vector<array>(const std::vector<array>&)>& fun) {
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    bool shapeless /* false */) {
   if (detail::compile_mode() == CompileMode::disabled) {
     return fun;
   }
   auto fun_id = detail::getAddress(fun);
-  return detail::compile(fun, fun_id);
+  return detail::compile(fun, fun_id, shapeless);
 }
 
 void disable_compile() {

mlx/compile.h

@@ -8,9 +8,10 @@ namespace mlx::core {
 
 enum class CompileMode { disabled, no_simplify, no_fuse, enabled };
 
-// Compile takes a function and returns a new function
+/** Compile takes a function and returns a compiled function. */
 std::function<std::vector<array>(const std::vector<array>&)> compile(
-    const std::function<std::vector<array>(const std::vector<array>&)>& fun);
+    const std::function<std::vector<array>(const std::vector<array>&)>& fun,
+    bool shapeless = false);
 
 /** Globally disable compilation.
  * Setting the environment variable ``MLX_DISABLE_COMPILE`` can also

mlx/compile_impl.h

@@ -0,0 +1,11 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/device.h"
+
+namespace mlx::core::detail {
+
+bool compile_available_for_device(const Device& device);
+
+}

mlx/dtype.h

@@ -46,10 +46,6 @@ struct Dtype {
   };
 };
 
-inline bool is_available(const Dtype& dtype) {
-  return true;
-}
-
 static constexpr Dtype bool_{Dtype::Val::bool_, sizeof(bool)};
 
 static constexpr Dtype uint8{Dtype::Val::uint8, sizeof(uint8_t)};

mlx/fast.cpp

@@ -1,6 +1,8 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/fast.h"
+#include "mlx/fast_primitives.h"
+#include "mlx/ops.h"
 #include "mlx/transforms.h"
 
 namespace mlx::core::fast {
@@ -125,4 +127,147 @@ bool RoPE::is_equivalent(const Primitive& other) const {
       offset_ == a_other.offset_);
 }
 
+/** Computes: O = softmax(Q @ K.T) @ V **/
+array scaled_dot_product_attention(
+    const array& queries,
+    const array& keys,
+    const array& values,
+    const float scale,
+    const std::optional<array>& mask,
+    StreamOrDevice s) {
+  for (const auto& tensor : {queries, keys, values}) {
+    if (tensor.ndim() != 4) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] input with shape "
+          << tensor.shape() << " expected to be rank 4";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  const size_t batch_dim = queries.shape(0);
+  for (const auto& tensor : {keys, values}) {
+    if (tensor.shape(0) != batch_dim) {
+      std::ostringstream msg;
+      msg << "[scaled_dot_product_attention] mismatching batch dimension for input with shape "
+          << tensor.shape() << ".";
+      throw std::invalid_argument(msg.str());
+    }
+  }
+
+  // Q, K must have matching last dims (d_k aka 'head_dim');
+  if (queries.shape(-1) != keys.shape(-1)) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] query, keys expected to have matching last dimension; found query shape "
+        << queries.shape() << " for keys shape " << keys.shape() << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // K, V must have matching number of heads (n_kv_heads);
+  size_t n_q_heads = queries.shape(-3);
+  size_t n_kv_heads = keys.shape(-3);
+
+  if (keys.shape(-3) != values.shape(-3)) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] keys, values expected to have matching n_kv_heads; found keys with n_heads "
+        << keys.shape(-3) << " for values with n_heads " << values.shape(-3)
+        << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  // n_heads % n_kv_heads == 0; n_heads >= 1, n_kv_heads >= 1.
+  if (n_q_heads % n_kv_heads != 0) {
+    std::ostringstream msg;
+    msg << "[scaled_dot_product_attention] n_heads must be a multiple of n_kv_heads, found n_heads "
+        << n_q_heads << " for n_kv_heads " << n_kv_heads << ".";
+    throw std::invalid_argument(msg.str());
+  }
+
+  auto final_type = result_type({queries, keys, values});
+
+  auto q = astype(queries, final_type, s);
+  auto k = astype(keys, final_type, s);
+  auto v = astype(values, final_type, s);
+
+  auto out_shape =
+      std::vector<int>({q.shape(0), q.shape(1), q.shape(2), v.shape(-1)});
+
+  /* generic implementation for use cases that Metal implementation does not
+   * support. For non-supported cases listed below, use MLX primitives:
+   * * CPU implementation
+   * * batch size > 1
+   * * query sequence length > 1
+   * * non-null mask
+   */
+  bool needs_mask = mask.has_value();
+  auto fallback = [scale, needs_mask, final_type, n_q_heads, n_kv_heads, &s](
+                      const std::vector<array>& inputs) {
+    auto& q_tensor = inputs[0];
+    auto& k_tensor = inputs[1];
+    auto& v_tensor = inputs[2];
+    auto q_scaled = multiply(array(scale, q_tensor.dtype()), q_tensor, s);
+
+    auto tile_if_needs_repeat =
+        [n_q_heads, n_kv_heads](const array& arr, StreamOrDevice& s) -> array {
+      if (n_q_heads == n_kv_heads)
+        return arr;
+      int n_repeats = n_q_heads / n_kv_heads;
+      constexpr const int heads_axis =
+          1; // heads axis, assumes tensors arranged  as [0, 1, 2, 3] ->
+             // [Batch, Heads, Sequence, Hidden]
+      auto ret = repeat(arr, n_repeats, heads_axis, s);
+      return ret;
+    };
+    auto k_tensor_tiled = tile_if_needs_repeat(k_tensor, s);
+    auto v_tensor_tiled = tile_if_needs_repeat(v_tensor, s);
+
+    // dim check on k, v; repeat if untiled, since naive matmul will have
+    // dim mismatch for GQA (MQA could make use of broadcast)
+    auto k_transposed = transpose(k_tensor_tiled, {0, 1, 3, 2}, s);
+    auto s_tensor = matmul(q_scaled, k_transposed, s);
+    if (needs_mask) {
+      auto mask_tensor = inputs[3];
+      s_tensor = add(s_tensor, mask_tensor, s);
+    }
+    auto p = astype(
+        softmax(astype(s_tensor, float32, s), std::vector<int>{-1}, s),
+        final_type,
+        s);
+    auto out_tensor = matmul(p, v_tensor_tiled, s);
+    return std::vector<array>{out_tensor};
+  };
+
+  auto stream = to_stream(s);
+
+  // current implementation use case: batch size 1, query sequence length 1, no
+  // mask.  Likewise, requires head_dim == 128
+  constexpr const int supported_head_dim = 128;
+  const size_t query_head_dim = q.shape(-1);
+  const size_t query_sequence_length = q.shape(2);
+  bool implementation_supports_use_case = batch_dim == 1 &&
+      query_sequence_length == 1 && !mask.has_value() &&
+      query_head_dim == supported_head_dim;
+
+  if (stream.device == Device::gpu && implementation_supports_use_case) {
+    auto out = array(
+        out_shape,
+        final_type,
+        std::make_unique<ScaledDotProductAttention>(
+            stream, fallback, scale, false),
+        {q, k, v});
+    return out;
+  }
+
+  if (mask.has_value()) {
+    return fallback({q, k, v, mask.value()})[0];
+  } else {
+    return fallback({q, k, v})[0];
+  }
+}
+
+bool ScaledDotProductAttention::is_equivalent(const Primitive& other) const {
+  const ScaledDotProductAttention& a_other =
+      static_cast<const ScaledDotProductAttention&>(other);
+  return needs_mask_ == a_other.needs_mask_ && scale_ == a_other.scale_;
+}
+
 } // namespace mlx::core::fast

mlx/fast.h

@@ -2,40 +2,12 @@
 
 #pragma once
 
-#include "mlx/ops.h"
-#include "mlx/primitives.h"
+#include <optional>
+
+#include "mlx/utils.h"
 
 namespace mlx::core::fast {
 
-// Custom primitive accepts a fallback function which it uses for
-// transformations. Transformations are virtual so that derived classes may to
-// override the default behavior
-class Custom : public Primitive {
- public:
-  explicit Custom(
-      Stream stream,
-      std::function<std::vector<array>(std::vector<array>)> fallback)
-      : Primitive(stream), fallback_(fallback){};
-
-  virtual std::pair<std::vector<array>, std::vector<int>> vmap(
-      const std::vector<array>& inputs,
-      const std::vector<int>& axes) override;
-
-  virtual std::vector<array> jvp(
-      const std::vector<array>& primals,
-      const std::vector<array>& tangents,
-      const std::vector<int>& argnums) override;
-
-  virtual std::vector<array> vjp(
-      const std::vector<array>& primals,
-      const std::vector<array>& cotangents,
-      const std::vector<int>& argnums,
-      const std::vector<array>& outputs) override;
-
- private:
-  std::function<std::vector<array>(std::vector<array>)> fallback_;
-};
-
 array rope(
     const array& x,
     int dims,
@@ -45,38 +17,13 @@ array rope(
     int offset,
     StreamOrDevice s /* = {} */);
 
-class RoPE : public Custom {
- public:
-  RoPE(
-      Stream stream,
-      std::function<std::vector<array>(std::vector<array>)> fallback,
-      int dims,
-      bool traditional,
-      float base,
-      float scale,
-      int offset)
-      : Custom(stream, fallback),
-        dims_(dims),
-        traditional_(traditional),
-        base_(base),
-        scale_(scale),
-        offset_(offset){};
-
-  void eval_cpu(const std::vector<array>& inputs, std::vector<array>& outputs)
-      override;
-  void eval_gpu(const std::vector<array>& inputs, std::vector<array>& outputs)
-      override;
-
-  DEFINE_PRINT(RoPE)
-  bool is_equivalent(const Primitive& other) const override;
-
- private:
-  std::function<std::vector<array>(std::vector<array>)> fallback_;
-  int dims_;
-  bool traditional_;
-  float base_;
-  float scale_;
-  int offset_;
-};
+/** Computes: O = softmax(Q @ K.T) @ V **/
+array scaled_dot_product_attention(
+    const array& queries,
+    const array& keys,
+    const array& values,
+    const float scale,
+    const std::optional<array>& mask = std::nullopt,
+    StreamOrDevice s = {});
 
 } // namespace mlx::core::fast