patch (#1320)

Allows opt-in to memory efficient GPU shader at proscribed sequence
length.  Otherwise, utilizes aggregate MLX primitives for best latency.

.circleci/config.yml

@@ -144,6 +144,7 @@ jobs:
           name: Install dependencies
           command: |
             brew install python@<< parameters.python_version >>
+            brew install openmpi
             python<< parameters.python_version >> -m venv env
             source env/bin/activate
             pip install --upgrade pip

.github/workflows/pull_request.yml

@@ -17,4 +17,4 @@ jobs:
           pip install pre-commit black isort clang-format
       - name: Run lint
         run: |
-          pre-commit run --all-files
+          pre-commit run --all-files

ACKNOWLEDGMENTS.md

@@ -10,13 +10,14 @@ 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. Added `clip_grad_norm` along with `tree_reduce`.
 - 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.
+- Diogo Da Cruz: Added `tri`, `tril`, `triu`, `tensordot`, `inner`, `outer`, `tile`, `StreamContext`, `stream`, safetensors support, `einsum`, and `einsum_path`.
 - 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`
 - AmirHossein Razlighi: Added chaining support for some of the ops in `nn.Module`. Comparison works for non array objects in `mlx.core.array`. Exception handling for invalid operations in `mlx.core.array`.
 - Gleb Pobudzey: Added the `where` primitive, and groups in 1D and 2D convolutions.
+- Paul Paczuski: Improved stability of BCE loss calculation
 
 <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" />

CMakeLists.txt

@@ -24,7 +24,7 @@ option(MLX_METAL_JIT "Use JIT compilation for Metal kernels" OFF)
 option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF)
 
 if(NOT MLX_VERSION)
-  set(MLX_VERSION 0.15.0)
+  set(MLX_VERSION 0.16.3)
 endif()
 
 # --------------------- Processor tests -------------------------
@@ -83,24 +83,21 @@ elseif (MLX_BUILD_METAL)
                   OUTPUT_VARIABLE MACOS_VERSION
                   COMMAND_ERROR_IS_FATAL ANY)
 
-  message(STATUS "Building with SDK for macOS version ${MACOS_VERSION}")
-
-  if (${MACOS_VERSION} GREATER_EQUAL 14.2)
-    set(METAL_CPP_PATCH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/metal.14.2.diff)
-    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14.2_iOS17.2.zip)
-    set(MLX_METAL_VERSION METAL_3_1)
-  elseif (${MACOS_VERSION} GREATER_EQUAL 14.0)
-    set(METAL_CPP_PATCH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/metal.14.0.diff)
-    set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS14_iOS17-beta.zip)
-    set(MLX_METAL_VERSION METAL_3_0)
-  else()
+  if (${MACOS_VERSION} LESS 14.0)
     message(FATAL_ERROR "MLX requires macOS SDK >= 14.0 to be built with MLX_BUILD_METAL=ON" )
   endif()
+  message(STATUS "Building with SDK for macOS version ${MACOS_VERSION}")
+
+  set(METAL_CPP_URL https://developer.apple.com/metal/cpp/files/metal-cpp_macOS15_iOS18-beta.zip)
+  # Get the metal version
+  execute_process(
+    COMMAND zsh "-c" "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal -E -x metal -P - | tail -1 | tr -d '\n'"
+    OUTPUT_VARIABLE MLX_METAL_VERSION
+    COMMAND_ERROR_IS_FATAL ANY)
 
   FetchContent_Declare(
     metal_cpp
     URL ${METAL_CPP_URL}
-    PATCH_COMMAND /usr/bin/patch -N -i ${METAL_CPP_PATCH} || true
   )
 
   FetchContent_MakeAvailable(metal_cpp)
@@ -115,7 +112,7 @@ elseif (MLX_BUILD_METAL)
     ${FOUNDATION_LIB}
     ${QUARTZ_LIB})
 
-  add_compile_definitions(${MLX_METAL_VERSION})
+  add_compile_definitions("MLX_METAL_VERSION=${MLX_METAL_VERSION}")
 endif()
 
 if (MLX_BUILD_CPU)
@@ -169,7 +166,26 @@ endif()
 
 find_package(MPI)
 if (MPI_FOUND)
+  execute_process(
+    COMMAND zsh "-c" "mpirun --version"
+    OUTPUT_VARIABLE MPI_VERSION
+    ERROR_QUIET
+  )
+  if (${MPI_VERSION} MATCHES ".*Open MPI.*")
     target_include_directories(mlx PRIVATE ${MPI_INCLUDE_PATH})
+  elseif (MPI_VERSION STREQUAL "")
+    set(MPI_FOUND FALSE)
+    message(
+      WARNING
+      "MPI found but mpirun is not available. Building without MPI."
+    )
+  else()
+    set(MPI_FOUND FALSE)
+    message(
+      WARNING
+      "MPI which is not OpenMPI found. Building without MPI."
+    )
+  endif() 
 endif()
 
 add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx)

benchmarks/python/einsum_bench.py

@@ -0,0 +1,84 @@
+# Copyright © 2024 Apple Inc.
+
+import time
+
+import mlx.core as mx
+import numpy as np
+
+
+def timeit(fn, its=100, args=[]):
+    for _ in range(5):
+        fn(*args)
+    tic = time.perf_counter()
+    for _ in range(its):
+        fn(*args)
+    toc = time.perf_counter()
+    return 1e3 * (toc - tic) / its
+
+
+def time_little_einsum_path():
+    subscripts = "ik,kj->ij"
+    x = mx.ones((32, 32))
+    y = mx.ones((32, 32))
+    mx_time = timeit(mx.einsum_path, args=(subscripts, x, y))
+
+    x = np.array(x)
+    y = np.array(y)
+    np_time = timeit(np.einsum_path, args=(subscripts, x, y))
+    print("Timing little einsum path...")
+    print(f"MLX ... {mx_time:.3f} ms")
+    print(f"NumPy... {np_time:.3f} ms")
+
+
+def time_big_einsum_path():
+    chars = list("abcdefgh")
+    char_to_dim = {c: v for v, c in enumerate(chars)}
+
+    num_inputs = 10
+    inputs = []
+    subscripts = []
+    for _ in range(num_inputs):
+        subscript = np.random.choice(chars, size=5, replace=False).tolist()
+        subscripts.append("".join(subscript))
+        inputs.append(np.ones(list(char_to_dim[c] for c in subscript)))
+    subscripts = ",".join(subscripts)
+
+    np_time = timeit(np.einsum_path, args=(subscripts, *inputs))
+
+    inputs = [mx.array(x) for x in inputs]
+    mx_time = timeit(mx.einsum_path, args=(subscripts, *inputs))
+    print("Timing big einsum path...")
+    print(f"MLX ... {mx_time:.3f} ms")
+    print(f"NumPy... {np_time:.3f} ms")
+
+
+def time_attention():
+    def regular_attention(x):
+        # shape [batch, sequence, num_heads, head_dim]
+        queries, keys, values = x, x, x
+        scores = queries.transpose(0, 2, 1, 3) @ keys.transpose(0, 2, 3, 1)
+        scores = mx.softmax(scores, axis=-1)
+        output = (scores @ values.transpose(0, 2, 1, 3)).swapaxes(1, 2)
+        mx.eval(output)
+
+    def einsum_attention(x):
+        # shape [batch, sequence, num_heads, head_dim]
+        queries, keys, values = x, x, x
+        scores = mx.einsum("itjk,iujk->ijtu", queries, keys)
+        scores = mx.softmax(scores, axis=-1)
+        output = mx.einsum("ijtu,iujk->itjk", scores, values)
+        mx.eval(output)
+
+    x = mx.random.uniform(shape=(8, 512, 32, 128))
+
+    regular_time = timeit(regular_attention, args=(x,))
+    ein_time = timeit(einsum_attention, args=(x,))
+    print("Timing einsum attention...")
+    print(f"Regular ... {regular_time:.3f} ms")
+    print(f"Einsum ... {ein_time:.3f} ms")
+
+
+if __name__ == "__main__":
+    time_little_einsum_path()
+    time_big_einsum_path()
+    time_attention()

benchmarks/python/hadamard_bench.py

@@ -0,0 +1,70 @@
+import argparse
+
+import matplotlib
+import mlx.core as mx
+import numpy as np
+from time_utils import measure_runtime
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+def had(x):
+    y = mx.hadamard_transform(x)
+    mx.eval(y)
+
+
+def copy(x):
+    y = x + 1.0
+    mx.eval(y)
+
+
+def run(dtype):
+    system_size = 2**26
+    outputs = {}
+    for test_fn in (had, copy):
+        for m in [1, 12, 20, 28]:
+            if test_fn == copy:
+                key = "copy"
+            elif m == 1:
+                key = "had_2^k"
+            else:
+                key = "had_m*2^k"
+            outputs.setdefault(key, {})
+            for k in range(7, 14):
+                n = m * 2**k
+                if n > 2**15:
+                    continue
+                x_np = np.random.normal(size=(system_size // n, n)).astype(dtype)
+                x = mx.array(x_np)
+                runtime_ms = measure_runtime(test_fn, x=x)
+                bytes_per_gb = 1e9
+                ms_per_s = 1e3
+                bytes_per_had = np.dtype(x_np.dtype).itemsize * 2
+                bandwidth_gb = (
+                    system_size * bytes_per_had / runtime_ms * ms_per_s / bytes_per_gb
+                )
+                print(n, bandwidth_gb)
+                outputs[key][n] = bandwidth_gb
+
+    colors = {
+        "copy": "black",
+        "had_2^k": "steelblue",
+        "had_m*2^k": "skyblue",
+    }
+    for key, output in outputs.items():
+        plt.scatter(output.keys(), output.values(), color=colors[key], label=key)
+    plt.title(f"MLX Hadamard Benchmark -- {dtype.__name__}")
+    plt.xlabel("N")
+    plt.ylabel("Bandwidth (GB/s)")
+    plt.legend()
+    plt.savefig(f"bench_{dtype.__name__}.png")
+    plt.clf()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fp16", action="store_true")
+    args = parser.parse_args()
+    dtype = np.float16 if args.fp16 else np.float32
+    run(dtype)

cmake/metal.14.0.diff

@@ -1,36 +0,0 @@
-diff -ur Metal/MTLEvent.hpp MetalNew/MTLEvent.hpp
---- Metal/MTLEvent.hpp	2023-06-01 12:18:26
-+++ MetalNew/MTLEvent.hpp	2024-04-15 07:36:59
-@@ -62,6 +62,7 @@
- 
-     uint64_t                 signaledValue() const;
-     void                     setSignaledValue(uint64_t signaledValue);
-+    bool                     waitUntilSignaledValue(uint64_t signaledValue, uint64_t timeoutMS);
- };
- 
- class SharedEventHandle : public NS::SecureCoding<SharedEventHandle>
-@@ -138,6 +139,11 @@
- _MTL_INLINE void MTL::SharedEvent::setSignaledValue(uint64_t signaledValue)
- {
-     Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSignaledValue_), signaledValue);
-+}
-+
-+// method: waitUntilSignaledValue
-+_MTL_INLINE bool MTL::SharedEvent::waitUntilSignaledValue(uint64_t signaledValue, uint64_t timeoutMS) {
-+    return Object::sendMessage<bool>(this, _MTL_PRIVATE_SEL(waitUntilSignaledValue_timeoutMS_), signaledValue, timeoutMS);
- }
- 
- // static method: alloc
-diff -ur Metal/MTLHeaderBridge.hpp MetalNew/MTLHeaderBridge.hpp
---- Metal/MTLHeaderBridge.hpp	2023-06-01 12:18:26
-+++ MetalNew/MTLHeaderBridge.hpp	2024-04-15 07:37:29
-@@ -1906,6 +1906,9 @@
-     "setShouldMaximizeConcurrentCompilation:");
- _MTL_PRIVATE_DEF_SEL(setSignaledValue_,
-     "setSignaledValue:");
-+_MTL_PRIVATE_DEF_SEL(
-+    waitUntilSignaledValue_timeoutMS_,
-+    "waitUntilSignaledValue:timeoutMS:");
- _MTL_PRIVATE_DEF_SEL(setSize_,
-     "setSize:");
- _MTL_PRIVATE_DEF_SEL(setSlice_,

cmake/metal.14.2.diff

@@ -1,36 +0,0 @@
-diff -ur Metal/MTLEvent.hpp MetalNew/MTLEvent.hpp
---- Metal/MTLEvent.hpp	2024-04-15 07:12:10
-+++ MetalNew/MTLEvent.hpp	2024-04-15 07:15:50
-@@ -62,6 +62,7 @@
- 
-     uint64_t                 signaledValue() const;
-     void                     setSignaledValue(uint64_t signaledValue);
-+    bool                     waitUntilSignaledValue(uint64_t signaledValue, uint64_t timeoutMS);
- };
- 
- class SharedEventHandle : public NS::SecureCoding<SharedEventHandle>
-@@ -138,6 +139,11 @@
- _MTL_INLINE void MTL::SharedEvent::setSignaledValue(uint64_t signaledValue)
- {
-     Object::sendMessage<void>(this, _MTL_PRIVATE_SEL(setSignaledValue_), signaledValue);
-+}
-+
-+// method: waitUntilSignaledValue
-+_MTL_INLINE bool MTL::SharedEvent::waitUntilSignaledValue(uint64_t signaledValue, uint64_t timeoutMS) {
-+    return Object::sendMessage<bool>(this, _MTL_PRIVATE_SEL(waitUntilSignaledValue_timeoutMS_), signaledValue, timeoutMS);
- }
- 
- // static method: alloc
-diff -ur Metal/MTLHeaderBridge.hpp MetalNew/MTLHeaderBridge.hpp
---- Metal/MTLHeaderBridge.hpp	2024-04-15 07:12:10
-+++ MetalNew/MTLHeaderBridge.hpp	2024-04-15 07:16:15
-@@ -1918,6 +1918,9 @@
-     "setShouldMaximizeConcurrentCompilation:");
- _MTL_PRIVATE_DEF_SEL(setSignaledValue_,
-     "setSignaledValue:");
-+_MTL_PRIVATE_DEF_SEL(
-+    waitUntilSignaledValue_timeoutMS_,
-+    "waitUntilSignaledValue:timeoutMS:");
- _MTL_PRIVATE_DEF_SEL(setSize_,
-     "setSize:");
- _MTL_PRIVATE_DEF_SEL(setSlice_,

docs/requirements.txt

@@ -1,3 +1,4 @@
 sphinx
 breathe
 sphinx-book-theme
+mlx

docs/src/conf.py

@@ -83,3 +83,15 @@ def setup(app):
 # -- Options for LaTeX output ------------------------------------------------
 
 latex_documents = [(main_doc, "MLX.tex", "MLX Documentation", author, "manual")]
+latex_elements = {
+    "preamble": r"""
+    \usepackage{enumitem}
+    \setlistdepth{5}
+    \setlist[itemize,1]{label=$\bullet$}
+    \setlist[itemize,2]{label=$\bullet$}
+    \setlist[itemize,3]{label=$\bullet$}
+    \setlist[itemize,4]{label=$\bullet$}
+    \setlist[itemize,5]{label=$\bullet$}
+    \renewlist{itemize}{itemize}{5}
+""",
+}

docs/src/dev/extensions.rst

@@ -486,9 +486,8 @@ below.
         std::ostringstream kname;
         kname << "axpby_" << "general_" << type_to_name(out);
 
-        // Make sure the metal library is available and look for it
-        // in the same folder as this executable if needed
-        d.register_library("mlx_ext", metal::get_colocated_mtllib_path);
+        // Make sure the metal library is available
+        d.register_library("mlx_ext");
 
         // Make a kernel from this metal library
         auto kernel = d.get_kernel(kname.str(), "mlx_ext");

docs/src/examples/llama-inference.rst

@@ -15,7 +15,7 @@ module to concisely define the model architecture.
 Attention layer
 ^^^^^^^^^^^^^^^^
 
-We will start with the llama attention layer which notably uses the RoPE
+We will start with the Llama attention layer which notably uses the RoPE
 positional encoding. [1]_ In addition, our attention layer will optionally use a
 key/value cache that will be concatenated with the provided keys and values to
 support efficient inference.

docs/src/examples/mlp.rst

@@ -64,7 +64,7 @@ set:
 Next, setup the problem parameters and load the data. To load the data, you need our
 `mnist data loader
 <https://github.com/ml-explore/mlx-examples/blob/main/mnist/mnist.py>`_, which
-we will import as `mnist`.
+we will import as ``mnist``.
 
 .. code-block:: python
 

docs/src/install.rst

@@ -70,36 +70,36 @@ To build and install the MLX python library from source, first, clone MLX from
 
    git clone git@github.com:ml-explore/mlx.git mlx && cd mlx
 
-Install `nanobind <https://nanobind.readthedocs.io/en/latest/>`_ with:
-
-.. code-block:: shell
-
-    pip install git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
-
 Then simply build and install MLX using pip:
 
 .. code-block:: shell
 
-   env CMAKE_BUILD_PARALLEL_LEVEL="" pip install .
+  CMAKE_BUILD_PARALLEL_LEVEL="" pip install .
 
-For developing use an editable install:
+For developing, install the package with development dependencies, and use an
+editable install:
 
 .. code-block:: shell
 
-  env CMAKE_BUILD_PARALLEL_LEVEL="" pip install -e .
+  CMAKE_BUILD_PARALLEL_LEVEL="" pip install -e ".[dev]"
 
-To make sure the install is working run the tests with:
+Once the development dependencies are installed, you can build faster with:
+
+.. code-block:: shell
+
+ CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py build_ext -j --inplace
+
+Run the tests with:
 
 .. code-block:: shell
 
-  pip install ".[testing]"
   python -m unittest discover python/tests
 
-Optional: Install stubs to enable auto completions and type checking from your IDE:
+Optional: Install stubs to enable auto completions and type checking from your
+IDE:
 
 .. code-block:: shell
 
-  pip install ".[dev]"
   python setup.py generate_stubs
 
 C++ API
@@ -195,7 +195,7 @@ GGUF, you can do:
 
 .. code-block:: shell
 
-  cmake ..
+  cmake .. \
     -DCMAKE_BUILD_TYPE=MinSizeRel \
     -DBUILD_SHARED_LIBS=ON \
     -DMLX_BUILD_CPU=OFF \

docs/src/python/array.rst

@@ -24,6 +24,7 @@ Array
     array.any
     array.argmax
     array.argmin
+    array.conj
     array.cos
     array.cummax
     array.cummin
@@ -57,3 +58,4 @@ Array
     array.transpose
     array.T
     array.var
+    array.view

docs/src/python/linalg.rst

@@ -9,7 +9,9 @@ Linear Algebra
    :toctree: _autosummary 
 
     inv
+    tri_inv
     norm
     cholesky
+    cholesky_inv
     qr
     svd

docs/src/python/ops.rst

@@ -57,6 +57,8 @@ Operations
    diagonal
    divide
    divmod
+   einsum
+   einsum_path
    equal
    erf
    erfinv
@@ -72,6 +74,7 @@ Operations
    gather_qmm
    greater
    greater_equal
+   hadamard_transform
    identity
    inner
    isclose
@@ -103,6 +106,7 @@ Operations
    minimum
    moveaxis
    multiply
+   nan_to_num
    negative
    not_equal
    ones

docs/src/python/optimizers.rst

@@ -31,6 +31,41 @@ model's parameters and the **optimizer state**.
             # Compute the new parameters but also the optimizer state.
             mx.eval(model.parameters(), optimizer.state)
 
+Saving and Loading
+------------------
+
+To serialize an optimizer, save its state. To load an optimizer, load and set
+the saved state. Here's a simple example:
+
+.. code-block:: python
+
+   import mlx.core as mx
+   from mlx.utils import tree_flatten, tree_unflatten
+   import mlx.optimizers as optim
+
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   # Perform some updates with the optimizer
+   model = {"w" : mx.zeros((5, 5))}
+   grads = {"w" : mx.ones((5, 5))}
+   optimizer.update(model, grads)
+
+   # Save the state
+   state = tree_flatten(optimizer.state)
+   mx.save_safetensors("optimizer.safetensors", dict(state))
+
+   # Later on, for example when loading from a checkpoint,
+   # recreate the optimizer and load the state
+   optimizer = optim.Adam(learning_rate=1e-2)
+
+   state = tree_unflatten(list(mx.load("optimizer.safetensors").items()))
+   optimizer.state = state
+
+Note, not every optimizer configuation parameter is saved in the state. For
+example, for Adam the learning rate is saved but the ``betas`` and ``eps``
+parameters are not. A good rule of thumb is if the parameter can be scheduled
+then it will be included in the optimizer state.
+
 .. toctree::
 
    optimizers/optimizer

docs/src/python/random.rst

@@ -44,3 +44,4 @@ we use a splittable version of Threefry, which is a counter-based PRNG.
    split
    truncated_normal
    uniform
+   laplace

docs/src/python/transforms.rst

@@ -10,6 +10,7 @@ Transforms
 
    eval
    compile
+   custom_function
    disable_compile
    enable_compile
    grad

examples/extensions/axpby/axpby.cpp

@@ -249,9 +249,8 @@ void Axpby::eval_gpu(
   kname << (contiguous_kernel ? "contiguous_" : "general_");
   kname << type_to_name(out);
 
-  // Make sure the metal library is available and look for it
-  // in the same folder as this executable if needed
-  d.register_library("mlx_ext", metal::get_colocated_mtllib_path);
+  // Make sure the metal library is available
+  d.register_library("mlx_ext");
 
   // Make a kernel from this metal library
   auto kernel = d.get_kernel(kname.str(), "mlx_ext");

examples/extensions/requirements.txt

@@ -1,4 +1,4 @@
 setuptools>=42
 cmake>=3.24
-mlx>=0.9.0
-nanobind@git+https://github.com/wjakob/nanobind.git@2f04eac452a6d9142dedb957701bdb20125561e4
+mlx>=0.16.2
+nanobind==2.0

mlx/CMakeLists.txt

@@ -6,6 +6,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/einsum.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fast.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/ops.cpp

mlx/array.cpp

@@ -17,6 +17,10 @@ bool in_tracing() {
   return detail::InTracing::in_tracing();
 }
 
+bool retain_graph() {
+  return detail::RetainGraph::retain_graph();
+}
+
 } // namespace
 
 array::array(const std::complex<float>& val, Dtype dtype /* = complex64 */)
@@ -102,7 +106,7 @@ void array::eval() {
 }
 
 bool array::is_tracer() const {
-  return array_desc_->is_tracer && in_tracing();
+  return array_desc_->is_tracer && in_tracing() || retain_graph();
 }
 
 void array::set_data(allocator::Buffer buffer, deleter_t d) {
@@ -171,10 +175,11 @@ array::~array() {
     return;
   }
 
-  // Ignore arrays that will be detached
-  if (status() != array::Status::unscheduled) {
+  // Ignore arrays that might be detached during eval
+  if (status() == array::Status::scheduled) {
     return;
   }
+
   // Break circular reference for non-detached arrays with siblings
   if (auto n = siblings().size(); n > 0) {
     bool do_detach = true;
@@ -206,7 +211,7 @@ void array::ArrayDesc::init() {
     strides[i] = size;
     size *= shape[i];
   }
-  for (auto& in : inputs) {
+  for (const auto& in : inputs) {
     is_tracer |= in.is_tracer();
   }
 }
@@ -231,7 +236,7 @@ array::ArrayDesc::ArrayDesc(
 
 array::ArrayDesc::~ArrayDesc() {
   // When an array description is destroyed it will delete a bunch of arrays
-  // that may also destory their corresponding descriptions and so on and so
+  // that may also destroy their corresponding descriptions and so on and so
   // forth.
   //
   // This calls recursively the destructor and can result in stack overflow, we

mlx/array.h

@@ -73,32 +73,32 @@ class array {
       this->array_desc_ = other.array_desc_;
     }
     return *this;
-  };
+  }
 
   /** The size of the array's datatype in bytes. */
   size_t itemsize() const {
     return size_of(dtype());
-  };
+  }
 
   /** The number of elements in the array. */
   size_t size() const {
     return array_desc_->size;
-  };
+  }
 
   /** The number of bytes in the array. */
   size_t nbytes() const {
     return size() * itemsize();
-  };
+  }
 
   /** The number of dimensions of the array. */
   size_t ndim() const {
     return array_desc_->shape.size();
-  };
+  }
 
   /** The shape of the array as a vector of integers. */
   const std::vector<int>& shape() const {
     return array_desc_->shape;
-  };
+  }
 
   /**
    *  Get the size of the corresponding dimension.
@@ -107,12 +107,12 @@ class array {
    *  bounds checking. */
   int shape(int dim) const {
     return shape().at(dim < 0 ? dim + ndim() : dim);
-  };
+  }
 
   /** The strides of the array. */
   const std::vector<size_t>& strides() const {
     return array_desc_->strides;
-  };
+  }
 
   /**
    *  Get the stride of the corresponding dimension.
@@ -121,12 +121,12 @@ class array {
    *  bounds checking. */
   size_t strides(int dim) const {
     return strides().at(dim < 0 ? dim + ndim() : dim);
-  };
+  }
 
   /** Get the arrays data type. */
   Dtype dtype() const {
     return array_desc_->dtype;
-  };
+  }
 
   /** Evaluate the array. */
   void eval();
@@ -160,10 +160,10 @@ class array {
 
     friend bool operator==(const ArrayIterator& a, const ArrayIterator& b) {
       return a.arr.id() == b.arr.id() && a.idx == b.idx;
-    };
+    }
     friend bool operator!=(const ArrayIterator& a, const ArrayIterator& b) {
       return !(a == b);
-    };
+    }
 
    private:
     const array& arr;
@@ -209,7 +209,7 @@ class array {
     allocator::Buffer buffer;
     deleter_t d;
     Data(allocator::Buffer buffer, deleter_t d = allocator::free)
-        : buffer(buffer), d(d) {};
+        : buffer(buffer), d(d) {}
     // Not copyable
     Data(const Data& d) = delete;
     Data& operator=(const Data& d) = delete;
@@ -230,22 +230,22 @@ class array {
   /** The array's primitive. */
   Primitive& primitive() const {
     return *(array_desc_->primitive);
-  };
+  }
 
   /** A shared pointer to the array's primitive. */
   std::shared_ptr<Primitive>& primitive_ptr() const {
     return array_desc_->primitive;
-  };
+  }
 
   /** Check if the array has an attached primitive or is a leaf node. */
   bool has_primitive() const {
     return array_desc_->primitive != nullptr;
-  };
+  }
 
   /** The array's inputs. */
   const std::vector<array>& inputs() const {
     return array_desc_->inputs;
-  };
+  }
 
   std::vector<array>& inputs() {
     return array_desc_->inputs;
@@ -259,12 +259,12 @@ class array {
   /** The array's siblings. */
   const std::vector<array>& siblings() const {
     return array_desc_->siblings;
-  };
+  }
 
   /** The array's siblings. */
   std::vector<array>& siblings() {
     return array_desc_->siblings;
-  };
+  }
 
   void set_siblings(std::vector<array> siblings, uint16_t position) {
     array_desc_->siblings = std::move(siblings);
@@ -281,7 +281,7 @@ class array {
     outputs.push_back(*this);
     outputs.insert(outputs.end(), siblings().begin() + idx, siblings().end());
     return outputs;
-  };
+  }
 
   /** Detach the array from the graph. */
   void detach();
@@ -289,19 +289,19 @@ class array {
   /** Get the Flags bit-field. */
   const Flags& flags() const {
     return array_desc_->flags;
-  };
+  }
 
   /** The size (in elements) of the underlying buffer the array points to. */
   size_t data_size() const {
     return array_desc_->data_size;
-  };
+  }
 
   allocator::Buffer& buffer() {
     return array_desc_->data->buffer;
-  };
+  }
   const allocator::Buffer& buffer() const {
     return array_desc_->data->buffer;
-  };
+  }
 
   // Return a copy of the shared pointer
   // to the array::Data struct
@@ -312,19 +312,20 @@ class array {
   template <typename T>
   T* data() {
     return static_cast<T*>(array_desc_->data_ptr);
-  };
+  }
 
   template <typename T>
   const T* data() const {
     return static_cast<T*>(array_desc_->data_ptr);
-  };
+  }
 
   enum Status { unscheduled, scheduled, available };
 
   bool is_available() const {
     return status() == Status::available;
   }
-  const Status status() const {
+
+  Status status() const {
     return array_desc_->status;
   }
 

mlx/backend/accelerate/conv.cpp

@@ -1,9 +1,9 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #include <cassert>
 
+#include <Accelerate/Accelerate.h>
 #include <simd/vector.h>
-#include <vecLib/vDSP.h>
 
 #include "mlx/backend/common/copy.h"
 #include "mlx/primitives.h"

mlx/backend/accelerate/matmul.cpp

@@ -2,8 +2,7 @@
 
 #include <cassert>
 
-#include <vecLib/BNNS/bnns.h>
-#include <vecLib/cblas_new.h>
+#include <Accelerate/Accelerate.h>
 
 #include "mlx/backend/accelerate/utils.h"
 #include "mlx/backend/common/copy.h"

mlx/backend/accelerate/primitives.cpp

@@ -3,8 +3,7 @@
 #include <cassert>
 #include <cmath>
 
-#include <vecLib/vDSP.h>
-#include <vecLib/vForce.h>
+#include <Accelerate/Accelerate.h>
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/binary.h"
@@ -37,7 +36,7 @@ DEFAULT(Ceil)
 DEFAULT(Concatenate)
 DEFAULT(Conjugate)
 DEFAULT(Copy)
-DEFAULT_MULTI(CustomVJP)
+DEFAULT_MULTI(CustomTransforms)
 DEFAULT_MULTI(Depends)
 DEFAULT_MULTI(DivMod)
 DEFAULT(NumberOfElements)
@@ -51,6 +50,7 @@ DEFAULT(GatherMM)
 DEFAULT(GatherQMM)
 DEFAULT(Greater)
 DEFAULT(GreaterEqual)
+DEFAULT(Hadamard)
 DEFAULT(Less)
 DEFAULT(LessEqual)
 DEFAULT(Load)
@@ -102,7 +102,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary(
+    binary_op<float>(
         a,
         b,
         out,
@@ -117,7 +117,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vadd((const float*)a, 1, (const float*)b, 1, (float*)o, 1, n);
         });
   } else if (a.dtype() == int32) {
-    binary(
+    binary_op<int>(
         a,
         b,
         out,
@@ -132,7 +132,7 @@ void Add::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vaddi((const int*)a, 1, (const int*)b, 1, (int*)o, 1, n);
         });
   } else {
-    binary(a, b, out, [](auto x, auto y) { return x + y; });
+    eval(inputs, out);
   }
 }
 
@@ -287,7 +287,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == int32) {
-    binary(
+    binary_op<int>(
         a,
         b,
         out,
@@ -300,7 +300,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vdivi((const int*)b, 1, (const int*)a, 1, (int*)o, 1, n);
         });
   } else if (a.dtype() == float32) {
-    binary(
+    binary_op<float>(
         a,
         b,
         out,
@@ -315,7 +315,7 @@ void Divide::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vdiv((const float*)b, 1, (const float*)a, 1, (float*)o, 1, n);
         });
   } else {
-    binary(a, b, out, [](auto x, auto y) { return x / y; });
+    eval(inputs, out);
   }
 }
 
@@ -326,12 +326,8 @@ void Exp::eval_cpu(const std::vector<array>& inputs, array& out) {
     set_unary_output_data(in, out);
     auto size = in.data_size();
     vvexpf(out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
-  } else if (issubdtype(out.dtype(), inexact)) {
-    unary_fp(in, out, [](auto x) { return std::exp(x); });
   } else {
-    throw std::invalid_argument(
-        "[exp] Cannot exponentiate elements in array"
-        " with non floating point type.");
+    eval(inputs, out);
   }
 }
 
@@ -393,12 +389,8 @@ void Log1p::eval_cpu(const std::vector<array>& inputs, array& out) {
     auto size = in.data_size();
     vvlog1pf(
         out.data<float>(), in.data<float>(), reinterpret_cast<int*>(&size));
-  } else if (issubdtype(out.dtype(), inexact)) {
-    unary_fp(in, out, [](auto x) { return std::log1p(x); });
   } else {
-    throw std::invalid_argument(
-        "[log1p] Cannot compute log of elements in array with"
-        " non floating point type.");
+    eval(inputs, out);
   }
 }
 
@@ -408,7 +400,7 @@ void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary(
+    binary_op<float>(
         a,
         b,
         out,
@@ -423,7 +415,7 @@ void Multiply::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vmul((const float*)a, 1, (const float*)b, 1, (float*)o, 1, n);
         });
   } else {
-    binary(a, b, out, [](auto x, auto y) { return x * y; });
+    eval(inputs, out);
   }
 }
 
@@ -434,7 +426,7 @@ void Negative::eval_cpu(const std::vector<array>& inputs, array& out) {
     set_unary_output_data(in, out);
     vDSP_vneg(in.data<float>(), 1, out.data<float>(), 1, in.data_size());
   } else {
-    unary(in, out, [](auto x) { return -x; });
+    eval(inputs, out);
   }
 }
 
@@ -521,7 +513,7 @@ void Square::eval_cpu(const std::vector<array>& inputs, array& out) {
     auto size = in.data_size();
     vDSP_vsq(in.data<float>(), 1, out.data<float>(), 1, size);
   } else {
-    unary(in, out, [](auto x) { return x * x; });
+    eval(inputs, out);
   }
 }
 
@@ -547,7 +539,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
   auto& b = inputs[1];
 
   if (a.dtype() == float32) {
-    binary(
+    binary_op<float>(
         a,
         b,
         out,
@@ -565,7 +557,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
           vDSP_vsub((const float*)b, 1, (const float*)a, 1, (float*)o, 1, n);
         });
   } else if (a.dtype() == int32) {
-    binary(
+    binary_op<int>(
         a,
         b,
         out,
@@ -577,7 +569,7 @@ void Subtract::eval_cpu(const std::vector<array>& inputs, array& out) {
         },
         UseDefaultBinaryOp());
   } else {
-    binary(a, b, out, [](auto x, auto y) { return x - y; });
+    eval(inputs, out);
   }
 }
 

mlx/backend/accelerate/reduce.cpp

@@ -2,8 +2,8 @@
 
 #include <cassert>
 
+#include <Accelerate/Accelerate.h>
 #include <simd/vector.h>
-#include <vecLib/vDSP.h>
 
 #include "mlx/backend/common/reduce.h"
 #include "mlx/primitives.h"

mlx/backend/accelerate/softmax.cpp

@@ -3,7 +3,10 @@
 #include <cassert>
 #include <limits>
 
+#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
 #include <arm_neon.h>
+#endif
+
 #include <simd/math.h>
 #include <simd/vector.h>
 
@@ -53,25 +56,26 @@ inline simd_float16 simd_fast_exp(simd_float16 x) {
   return (*(simd_float16*)&epart) * x;
 }
 
+#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
 /**
  * The ARM neon equivalent of the fast exp above.
  */
 inline float16x8_t neon_fast_exp(float16x8_t x) {
-  x = vmulq_f16(x, vdupq_n_f16(1.442695)); // multiply with log_2(e)
-  x = vmaxq_f16(x, vdupq_n_f16(-14)); // clamp under with -14
-  x = vminq_f16(x, vdupq_n_f16(14)); // clamp over with 14
+  x = vmulq_f16(x, vdupq_n_f16(float16_t(1.442695f))); // multiply with log_2(e)
+  x = vmaxq_f16(x, vdupq_n_f16(float16_t(-14.f))); // clamp under with -14
+  x = vminq_f16(x, vdupq_n_f16(float16_t(14.f))); // clamp over with 14
 
-  float16x8_t ipart = vrndmq_f16(vaddq_f16(x, vdupq_n_f16(0.5)));
+  float16x8_t ipart = vrndmq_f16(vaddq_f16(x, vdupq_n_f16(float16_t(0.5f))));
   float16x8_t fpart = vsubq_f16(x, ipart);
 
-  x = vdupq_n_f16(1.535336188319500e-4f);
-  x = vfmaq_f16(vdupq_n_f16(1.339887440266574e-3f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(1.339887440266574e-3f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(9.618437357674640e-3f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(5.550332471162809e-2f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(2.402264791363012e-1f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(6.931472028550421e-1f), x, fpart);
-  x = vfmaq_f16(vdupq_n_f16(1.000000000000000f), x, fpart);
+  x = vdupq_n_f16(float16_t(1.535336188319500e-4f));
+  x = vfmaq_f16(vdupq_n_f16(float16_t(1.339887440266574e-3f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(1.339887440266574e-3f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(9.618437357674640e-3f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(5.550332471162809e-2f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(2.402264791363012e-1f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(6.931472028550421e-1f)), x, fpart);
+  x = vfmaq_f16(vdupq_n_f16(float16_t(1.000000000000000f)), x, fpart);
 
   // generate 2**ipart in the floating point representation using integer
   // bitshifting
@@ -107,53 +111,6 @@ inline float16_t neon_reduce_add(float16x8_t x) {
   return vget_lane_f16(y, 0);
 }
 
-template <typename T, typename VT>
-struct AccelerateSimdOps {
-  VT init(T a) {
-    return a;
-  }
-
-  VT load(const T* a) {
-    return *(VT*)a;
-  }
-
-  void store(T* dst, VT x) {
-    *(VT*)dst = x;
-  }
-
-  VT max(VT a, VT b) {
-    return simd_max(a, b);
-  };
-
-  VT exp(VT x) {
-    return simd_fast_exp(x);
-  }
-
-  VT add(VT a, VT b) {
-    return a + b;
-  }
-
-  VT sub(VT a, T b) {
-    return a - b;
-  }
-
-  VT mul(VT a, VT b) {
-    return a * b;
-  }
-
-  VT mul(VT a, T b) {
-    return a * b;
-  }
-
-  T reduce_max(VT x) {
-    return simd_reduce_max(x);
-  }
-
-  T reduce_add(VT x) {
-    return simd_reduce_add(x);
-  }
-};
-
 template <typename T, typename VT>
 struct NeonFp16SimdOps {
   VT init(T a) {
@@ -170,7 +127,7 @@ struct NeonFp16SimdOps {
 
   VT max(VT a, VT b) {
     return vmaxq_f16(a, b);
-  };
+  }
 
   VT exp(VT x) {
     return neon_fast_exp(x);
@@ -201,6 +158,55 @@ struct NeonFp16SimdOps {
   }
 };
 
+#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+
+template <typename T, typename VT>
+struct AccelerateSimdOps {
+  VT init(T a) {
+    return a;
+  }
+
+  VT load(const T* a) {
+    return *(VT*)a;
+  }
+
+  void store(T* dst, VT x) {
+    *(VT*)dst = x;
+  }
+
+  VT max(VT a, VT b) {
+    return simd_max(a, b);
+  }
+
+  VT exp(VT x) {
+    return simd_fast_exp(x);
+  }
+
+  VT add(VT a, VT b) {
+    return a + b;
+  }
+
+  VT sub(VT a, T b) {
+    return a - b;
+  }
+
+  VT mul(VT a, VT b) {
+    return a * b;
+  }
+
+  VT mul(VT a, T b) {
+    return a * b;
+  }
+
+  T reduce_max(VT x) {
+    return simd_reduce_max(x);
+  }
+
+  T reduce_add(VT x) {
+    return simd_reduce_add(x);
+  }
+};
+
 template <typename T, typename AccT, typename VT, typename Ops, int N>
 void softmax(const array& in, array& out) {
   Ops ops;
@@ -362,12 +368,16 @@ void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
             AccelerateSimdOps<float, simd_float16>,
             16>(in, out);
       } else {
+#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
         softmax<
             float16_t,
             float16_t,
             float16x8_t,
             NeonFp16SimdOps<float16_t, float16x8_t>,
             8>(in, out);
+#else // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
+        eval(inputs, out); // Redirect to common backend for consistency
+#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
       }
       break;
     case bfloat16:

mlx/backend/accelerate/utils.h

@@ -1,8 +1,8 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2023-2024 Apple Inc.
 
 #pragma once
 
-#include <vecLib/BNNS/bnns.h>
+#include <Accelerate/Accelerate.h>
 #include "mlx/dtype.h"
 
 namespace mlx::core {

mlx/backend/common/CMakeLists.txt

@@ -42,10 +42,12 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/erf.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/masked_mm.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/quantized.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/reduce_utils.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scan.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/select.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/slicing.cpp

mlx/backend/common/binary.cpp

@@ -196,6 +196,20 @@ void LogAddExp::eval(const std::vector<array>& inputs, array& out) {
   }
 }
 
+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, 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, detail::LogicalOr());
+}
+
 void Maximum::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];

mlx/backend/common/common.cpp

@@ -66,7 +66,7 @@ void Copy::eval(const std::vector<array>& inputs, array& out) {
   out.copy_shared_buffer(inputs[0]);
 }
 
-void CustomVJP::eval(
+void CustomTransforms::eval(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
   assert(inputs.size() > outputs.size());

mlx/backend/common/compiled.cpp

@@ -205,8 +205,8 @@ void compiled_allocate_outputs(
       // - Donatable
       // - Correct size
       // - Not a constant
-      if (in.flags().row_contiguous && in.nbytes() == outputs[o].nbytes() &&
-          in.is_donatable() &&
+      if (in.flags().row_contiguous && in.size() == outputs[o].size() &&
+          in.itemsize() == outputs[o].itemsize() && in.is_donatable() &&
           constant_ids_.find(inputs_[i].id()) == constant_ids_.end()) {
         if (move_buffers) {
           outputs[o].move_shared_buffer(

mlx/backend/common/copy.cpp

@@ -4,6 +4,7 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/utils.h"
 
 namespace mlx::core {
 
@@ -142,29 +143,31 @@ void copy_general(
     const std::vector<int>& data_shape,
     const std::vector<stride_t>& i_strides,
     int64_t i_offset) {
-  switch (src.ndim()) {
+  auto [new_shape, new_strides] = collapse_contiguous_dims(
+      data_shape, std::vector<std::vector<stride_t>>{i_strides});
+  switch (new_shape.size()) {
     case 1:
       copy_general_dim1<SrcT, DstT, stride_t>(
-          src, dst, data_shape, i_strides, i_offset);
+          src, dst, new_shape, new_strides[0], i_offset);
       return;
     case 2:
       copy_general_dim2<SrcT, DstT, stride_t>(
-          src, dst, data_shape, i_strides, i_offset);
+          src, dst, new_shape, new_strides[0], i_offset);
       return;
     case 3:
       copy_general_dim3<SrcT, DstT, stride_t>(
-          src, dst, data_shape, i_strides, i_offset);
+          src, dst, new_shape, new_strides[0], i_offset);
       return;
     case 4:
       copy_general_dim4<SrcT, DstT, stride_t>(
-          src, dst, data_shape, i_strides, i_offset);
+          src, dst, new_shape, new_strides[0], i_offset);
       return;
   }
 
   auto src_ptr = src.data<SrcT>() + i_offset;
   auto dst_ptr = dst.data<DstT>();
   for (size_t i = 0; i < dst.size(); ++i) {
-    stride_t src_elem = elem_to_loc(i, data_shape, i_strides);
+    stride_t src_elem = elem_to_loc(i, new_shape, new_strides[0]);
     dst_ptr[i] = static_cast<DstT>(src_ptr[src_elem]);
   }
 }
@@ -195,10 +198,10 @@ inline void copy_general_general_dims(
     const std::vector<int>& data_shape,
     const std::vector<stride_t>& i_strides,
     const std::vector<stride_t>& o_strides,
-    stride_t i_offset,
-    stride_t o_offset) {
+    int64_t i_offset,
+    int64_t o_offset) {
   if constexpr (D > 1) {
-    int axis = src.ndim() - D;
+    int axis = data_shape.size() - D;
     auto stride_src = i_strides[axis];
     auto stride_dst = o_strides[axis];
     auto N = data_shape[axis];
@@ -209,7 +212,7 @@ inline void copy_general_general_dims(
       o_offset += stride_dst;
     }
   } else {
-    int axis = src.ndim() - 1;
+    int axis = data_shape.size() - 1;
     auto stride_src = i_strides[axis];
     auto stride_dst = o_strides[axis];
     auto N = data_shape[axis];
@@ -230,38 +233,76 @@ void copy_general_general(
     const std::vector<int>& data_shape,
     const std::vector<stride_t>& i_strides,
     const std::vector<stride_t>& o_strides,
-    stride_t i_offset,
-    stride_t o_offset) {
-  switch (src.ndim()) {
+    int64_t i_offset,
+    int64_t o_offset) {
+  auto [new_shape, new_strides] = collapse_contiguous_dims(
+      data_shape, std::vector<std::vector<stride_t>>{i_strides, o_strides});
+  switch (new_shape.size()) {
     case 1:
       copy_general_general_dims<SrcT, DstT, stride_t, 1>(
-          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+          src,
+          dst,
+          new_shape,
+          new_strides[0],
+          new_strides[1],
+          i_offset,
+          o_offset);
       return;
     case 2:
       copy_general_general_dims<SrcT, DstT, stride_t, 2>(
-          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+          src,
+          dst,
+          new_shape,
+          new_strides[0],
+          new_strides[1],
+          i_offset,
+          o_offset);
       return;
     case 3:
       copy_general_general_dims<SrcT, DstT, stride_t, 3>(
-          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+          src,
+          dst,
+          new_shape,
+          new_strides[0],
+          new_strides[1],
+          i_offset,
+          o_offset);
       return;
     case 4:
       copy_general_general_dims<SrcT, DstT, stride_t, 4>(
-          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+          src,
+          dst,
+          new_shape,
+          new_strides[0],
+          new_strides[1],
+          i_offset,
+          o_offset);
       return;
     case 5:
       copy_general_general_dims<SrcT, DstT, stride_t, 5>(
-          src, dst, data_shape, i_strides, o_strides, i_offset, o_offset);
+          src,
+          dst,
+          new_shape,
+          new_strides[0],
+          new_strides[1],
+          i_offset,
+          o_offset);
       return;
   }
 
   int size = std::accumulate(
-      data_shape.end() - 5, data_shape.end(), 1, std::multiplies<int>());
+      new_shape.end() - 5, new_shape.end(), 1, std::multiplies<int>());
   for (int i = 0; i < src.size(); i += size) {
-    stride_t src_offset = i_offset + elem_to_loc(i, data_shape, i_strides);
-    stride_t dst_offset = o_offset + elem_to_loc(i, dst.shape(), o_strides);
+    stride_t src_offset = i_offset + elem_to_loc(i, new_shape, new_strides[0]);
+    stride_t dst_offset = o_offset + elem_to_loc(i, new_shape, new_strides[1]);
     copy_general_general_dims<SrcT, DstT, stride_t, 5>(
-        src, dst, data_shape, i_strides, o_strides, src_offset, dst_offset);
+        src,
+        dst,
+        new_shape,
+        new_strides[0],
+        new_strides[1],
+        src_offset,
+        dst_offset);
   }
 }
 
@@ -444,8 +485,17 @@ void copy_inplace(
   }
 }
 
-template <>
-void copy_inplace<int64_t>(
+template void copy_inplace<size_t>(
+    const array& src,
+    array& dst,
+    const std::vector<int>& data_shape,
+    const std::vector<size_t>& i_strides,
+    const std::vector<size_t>& o_strides,
+    int64_t i_offset,
+    int64_t o_offset,
+    CopyType ctype);
+
+template void copy_inplace<int64_t>(
     const array& src,
     array& dst,
     const std::vector<int>& data_shape,
@@ -453,24 +503,6 @@ void copy_inplace<int64_t>(
     const std::vector<int64_t>& o_strides,
     int64_t i_offset,
     int64_t o_offset,
-    CopyType ctype) {
-  switch (ctype) {
-    case CopyType::General:
-    case CopyType::GeneralGeneral:
-      return copy_inplace_dispatch(
-          src,
-          dst,
-          ctype,
-          data_shape,
-          i_strides,
-          o_strides,
-          i_offset,
-          o_offset);
-
-    case CopyType::Scalar:
-    case CopyType::Vector:
-      return copy_inplace_dispatch(src, dst, ctype);
-  }
-}
+    CopyType ctype);
 
 } // namespace mlx::core

mlx/backend/common/default_primitives.cpp

@@ -52,7 +52,7 @@ DEFAULT(Convolution)
 DEFAULT(Copy)
 DEFAULT(Cos)
 DEFAULT(Cosh)
-DEFAULT_MULTI(CustomVJP)
+DEFAULT_MULTI(CustomTransforms)
 DEFAULT_MULTI(Depends)
 DEFAULT(Divide)
 DEFAULT(NumberOfElements)
@@ -68,6 +68,7 @@ DEFAULT(Full)
 DEFAULT(Gather)
 DEFAULT(Greater)
 DEFAULT(GreaterEqual)
+DEFAULT(Hadamard)
 DEFAULT(Less)
 DEFAULT(LessEqual)
 DEFAULT(Load)

mlx/backend/common/hadamard.cpp

@@ -0,0 +1,107 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+
+#include "mlx/backend/common/copy.h"
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+// n = 2^k component
+template <typename T>
+void hadamard_n(array& out, int n, int m, float scale) {
+  for (int b = 0; b < out.size() / n; b++) {
+    size_t loc = b * n;
+    T* data_ptr = out.data<T>() + loc;
+    int h = 1;
+    int n_over_2 = n / 2;
+    while (h < n) {
+      for (int i = 0; i < n / 2; i++) {
+        int k = i & (h - 1);
+        int j = ((i - k) << 1) + k;
+        float x = *(data_ptr + j);
+        float y = *(data_ptr + j + h);
+        *(data_ptr + j) = x + y;
+        *(data_ptr + j + h) = x - y;
+        if (h == n_over_2) {
+          *(data_ptr + j) *= scale;
+          *(data_ptr + j + h) *= scale;
+        }
+      }
+      h <<= 1;
+    }
+  }
+}
+
+// m component
+template <typename T>
+void hadamard_m(array& out, int n, int m, float scale) {
+  auto h_matrices = hadamard_matrices();
+  auto& matrix = h_matrices[m];
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+  std::vector<bool> hmat_vec;
+  while (end != std::string_view::npos) {
+    auto row = matrix.substr(start, end - start);
+    for (int i = 0; i < row.length(); i++) {
+      hmat_vec.push_back(row[i] == '+');
+    }
+    start = end + 1;
+    end = matrix.find('\n', start);
+  }
+
+  for (int b = 0; b < out.size() / m / n; b++) {
+    size_t loc = b * n * m;
+    T* data_ptr = out.data<T>() + loc;
+    for (int i = 0; i < n; i++) {
+      std::vector<float> out(m);
+      for (int j = 0; j < m; j++) {
+        for (int k = 0; k < m; k++) {
+          float x = *(data_ptr + i + k * n);
+          if (hmat_vec[k + j * m]) {
+            out[j] += x;
+          } else {
+            out[j] -= x;
+          }
+        }
+      }
+      for (int j = 0; j < m; j++) {
+        *(data_ptr + i + j * n) = out[j] * scale;
+      }
+    }
+  }
+}
+
+template <typename T>
+void hadamard(array& out, int n, int m, float scale) {
+  float n_scale = m > 1 ? 1.0 : scale;
+  hadamard_n<T>(out, n, m, n_scale);
+  if (m > 1) {
+    hadamard_m<T>(out, n, m, scale);
+  }
+}
+
+void Hadamard::eval(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  auto& in = inputs[0];
+
+  // Copy input to output
+  copy(in, out, CopyType::General);
+
+  int axis = out.ndim() - 1;
+  auto [n, m] = decompose_hadamard(out.shape(axis));
+
+  switch (in.dtype()) {
+    case float32:
+      return hadamard<float>(out, n, m, scale_);
+    case float16:
+      return hadamard<float16_t>(out, n, m, scale_);
+    case bfloat16:
+      return hadamard<bfloat16_t>(out, n, m, scale_);
+    default:
+      throw std::invalid_argument("[hadamard] Unsupported type.");
+  }
+}
+
+} // namespace mlx::core

mlx/backend/common/hadamard.h

@@ -0,0 +1,105 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include <map>
+
+#include "mlx/utils.h"
+
+namespace mlx::core {
+
+// From http://neilsloane.com/hadamard/
+constexpr std::string_view h12 = R"(
++-++++++++++
+--+-+-+-+-+-
++++-++----++
++---+--+-++-
++++++-++----
++-+---+--+-+
+++--+++-++--
++--++---+--+
+++----+++-++
++--+-++---+-
+++++----+++-
++-+--+-++---
+)";
+
+constexpr std::string_view h20 = R"(
++----+----++--++-++-
+-+----+---+++---+-++
+--+----+---+++-+-+-+
+---+----+---+++++-+-
+----+----++--++-++-+
+-+++++-----+--+++--+
++-+++-+---+-+--+++--
+++-++--+---+-+--+++-
++++-+---+---+-+--+++
+++++-----++--+-+--++
+--++-+-++-+-----++++
+---++-+-++-+---+-+++
++---++-+-+--+--++-++
+++---++-+----+-+++-+
+-++---++-+----+++++-
+-+--+--++-+----+----
++-+-----++-+----+---
+-+-+-+---+--+----+--
+--+-+++------+----+-
++--+--++------+----+
+)";
+
+constexpr std::string_view h28 = R"(
++------++----++-+--+-+--++--
+-+-----+++-----+-+--+-+--++-
+--+-----+++---+-+-+----+--++
+---+-----+++---+-+-+-+--+--+
+----+-----+++---+-+-+++--+--
+-----+-----++++--+-+--++--+-
+------++----++-+--+-+--++--+
+--++++-+-------++--+++-+--+-
+---++++-+-----+-++--+-+-+--+
++---+++--+----++-++--+-+-+--
+++---++---+----++-++--+-+-+-
++++---+----+----++-++--+-+-+
+++++--------+-+--++-++--+-+-
+-++++--------+++--++--+--+-+
+-+-++-++--++--+--------++++-
++-+-++--+--++--+--------++++
+-+-+-++--+--++--+----+---+++
++-+-+-++--+--+---+---++---++
+++-+-+-++--+------+--+++---+
+-++-+-+-++--+------+-++++---
++-++-+---++--+------+-++++--
+-++--++-+-++-+++----++------
++-++--++-+-++-+++-----+-----
+++-++---+-+-++-+++-----+----
+-++-++-+-+-+-+--+++-----+---
+--++-++++-+-+----+++-----+--
++--++-+-++-+-+----+++-----+-
+++--++-+-++-+-+----++------+
+)";
+
+inline const std::map<int, std::string_view> hadamard_matrices() {
+  return {{12, h12}, {20, h20}, {28, h28}};
+}
+
+inline std::pair<int, int> decompose_hadamard(int n) {
+  // n = m*2^k
+  int m = 1;
+  if (!is_power_of_2(n)) {
+    auto h_matrices = hadamard_matrices();
+    for (auto [factor, _] : h_matrices) {
+      if (n % factor == 0) {
+        m = factor;
+        n /= factor;
+        break;
+      }
+    }
+    if (m == 1) {
+      throw std::invalid_argument(
+          "[hadamard] Only supports n = m*2^k where m in (1, 12, 20, 28).");
+    }
+  }
+  return {n, m};
+}
+
+} // namespace mlx::core

mlx/backend/common/inverse.cpp

@@ -10,9 +10,106 @@
 #include <lapack.h>
 #endif
 
+// Wrapper to account for differences in
+// LAPACK implementations (basically how to pass the 'uplo' string to fortran).
+int strtri_wrapper(char uplo, char diag, float* matrix, int N) {
+  int info;
+
+#ifdef LAPACK_FORTRAN_STRLEN_END
+  strtri_(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ matrix,
+      /* lda = */ &N,
+      /* info = */ &info,
+      /* uplo_len = */ static_cast<size_t>(1),
+      /* diag_len = */ static_cast<size_t>(1));
+#else
+  strtri_(
+      /* uplo = */ &uplo,
+      /* diag = */ &diag,
+      /* N = */ &N,
+      /* a = */ matrix,
+      /* lda = */ &N,
+      /* info = */ &info);
+#endif
+
+  return info;
+}
+
 namespace mlx::core {
 
-void inverse_impl(const array& a, array& inv) {
+void general_inv(array& inv, int N, int i) {
+  int info;
+  auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
+  // Compute LU factorization.
+  sgetrf_(
+      /* m = */ &N,
+      /* n = */ &N,
+      /* a = */ inv.data<float>() + N * N * i,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "inverse_impl: LU factorization failed with error code " << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  static const int lwork_query = -1;
+  float workspace_size = 0;
+
+  // Compute workspace size.
+  sgetri_(
+      /* m = */ &N,
+      /* a = */ nullptr,
+      /* lda = */ &N,
+      /* ipiv = */ nullptr,
+      /* work = */ &workspace_size,
+      /* lwork = */ &lwork_query,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "inverse_impl: LU workspace calculation failed with error code "
+       << info;
+    throw std::runtime_error(ss.str());
+  }
+
+  const int lwork = workspace_size;
+  auto scratch = array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
+
+  // Compute inverse.
+  sgetri_(
+      /* m = */ &N,
+      /* a = */ inv.data<float>() + N * N * i,
+      /* lda = */ &N,
+      /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
+      /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
+      /* lwork = */ &lwork,
+      /* info = */ &info);
+
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "inverse_impl: inversion failed with error code " << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+void tri_inv(array& inv, int N, int i, bool upper) {
+  const char uplo = upper ? 'L' : 'U';
+  const char diag = 'N';
+  int info = strtri_wrapper(uplo, diag, inv.data<float>() + N * N * i, N);
+  if (info != 0) {
+    std::stringstream ss;
+    ss << "inverse_impl: triangular inversion failed with error code " << info;
+    throw std::runtime_error(ss.str());
+  }
+}
+
+void inverse_impl(const array& a, array& inv, bool tri, bool upper) {
   // Lapack uses the column-major convention. We take advantage of the following
   // identity to avoid transposing (see
   // https://math.stackexchange.com/a/340234):
@@ -24,63 +121,11 @@ void inverse_impl(const array& a, array& inv) {
   const int N = a.shape(-1);
   const size_t num_matrices = a.size() / (N * N);
 
-  int info;
-  auto ipiv = array::Data{allocator::malloc_or_wait(sizeof(int) * N)};
-
   for (int i = 0; i < num_matrices; i++) {
-    // Compute LU factorization.
-    sgetrf_(
-        /* m = */ &N,
-        /* n = */ &N,
-        /* a = */ inv.data<float>() + N * N * i,
-        /* lda = */ &N,
-        /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
-        /* info = */ &info);
-
-    if (info != 0) {
-      std::stringstream ss;
-      ss << "inverse_impl: LU factorization failed with error code " << info;
-      throw std::runtime_error(ss.str());
-    }
-
-    static const int lwork_query = -1;
-    float workspace_size = 0;
-
-    // Compute workspace size.
-    sgetri_(
-        /* m = */ &N,
-        /* a = */ nullptr,
-        /* lda = */ &N,
-        /* ipiv = */ nullptr,
-        /* work = */ &workspace_size,
-        /* lwork = */ &lwork_query,
-        /* info = */ &info);
-
-    if (info != 0) {
-      std::stringstream ss;
-      ss << "inverse_impl: LU workspace calculation failed with error code "
-         << info;
-      throw std::runtime_error(ss.str());
-    }
-
-    const int lwork = workspace_size;
-    auto scratch =
-        array::Data{allocator::malloc_or_wait(sizeof(float) * lwork)};
-
-    // Compute inverse.
-    sgetri_(
-        /* m = */ &N,
-        /* a = */ inv.data<float>() + N * N * i,
-        /* lda = */ &N,
-        /* ipiv = */ static_cast<int*>(ipiv.buffer.raw_ptr()),
-        /* work = */ static_cast<float*>(scratch.buffer.raw_ptr()),
-        /* lwork = */ &lwork,
-        /* info = */ &info);
-
-    if (info != 0) {
-      std::stringstream ss;
-      ss << "inverse_impl: inversion failed with error code " << info;
-      throw std::runtime_error(ss.str());
+    if (tri) {
+      tri_inv(inv, N, i, upper);
+    } else {
+      general_inv(inv, N, i);
     }
   }
 }
@@ -89,7 +134,7 @@ void Inverse::eval(const std::vector<array>& inputs, array& output) {
   if (inputs[0].dtype() != float32) {
     throw std::runtime_error("[Inverse::eval] only supports float32.");
   }
-  inverse_impl(inputs[0], output);
+  inverse_impl(inputs[0], output, tri_, upper_);
 }
 
 } // namespace mlx::core

mlx/backend/common/ops.h

@@ -108,105 +108,105 @@ 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 ArcTan2 {
   template <typename T>
   T operator()(T y, T x) {
     return std::atan2(y, 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 Conjugate {
@@ -219,35 +219,35 @@ 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);
@@ -258,83 +258,83 @@ struct Expm1 {
   template <typename T>
   T operator()(T x) {
     return expm1(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 {
@@ -379,49 +379,49 @@ 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 {
@@ -554,7 +554,7 @@ struct LogAddExp {
         ? maxval
         : static_cast<decltype(x)>(
               maxval + std::log1p(fast_exp(minval - maxval)));
-  };
+  }
 };
 
 struct Multiply {
@@ -602,14 +602,14 @@ 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 {
@@ -623,35 +623,35 @@ struct BitwiseAnd {
   template <typename T>
   T operator()(T x, T y) {
     return x & y;
-  };
+  }
 };
 
 struct BitwiseOr {
   template <typename T>
   T operator()(T x, T y) {
     return x | y;
-  };
+  }
 };
 
 struct BitwiseXor {
   template <typename T>
   T operator()(T x, T y) {
     return x ^ y;
-  };
+  }
 };
 
 struct LeftShift {
   template <typename T>
   T operator()(T x, T y) {
     return x << y;
-  };
+  }
 };
 
 struct RightShift {
   template <typename T>
   T operator()(T x, T y) {
     return x >> y;
-  };
+  }
 };
 
 } // namespace mlx::core::detail

mlx/backend/common/primitives.cpp

@@ -8,7 +8,6 @@
 
 #include "mlx/allocator.h"
 #include "mlx/backend/common/arange.h"
-#include "mlx/backend/common/binary.h"
 #include "mlx/backend/common/copy.h"
 #include "mlx/backend/common/ops.h"
 #include "mlx/backend/common/slicing.h"
@@ -314,20 +313,6 @@ void LogicalNot::eval(const std::vector<array>& inputs, array& out) {
   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, 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, detail::LogicalOr());
-}
-
 void Negative::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];
@@ -420,7 +405,17 @@ void Reshape::eval(const std::vector<array>& inputs, array& out) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
 
   if (copy_necessary) {
-    copy(in, out, in.data_size() == 1 ? CopyType::Scalar : CopyType::General);
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+    auto out_strides = make_contiguous_strides<size_t>(in.shape());
+    copy_inplace<size_t>(
+        in,
+        out,
+        in.shape(),
+        in.strides(),
+        out_strides,
+        0,
+        0,
+        CopyType::General);
   } else {
     shared_buffer_reshape(in, out_strides, out);
   }

mlx/backend/common/reduce.cpp

@@ -104,48 +104,14 @@ void reduce_dispatch_out(
     }
     case Reduce::Sum: {
       auto op = [](auto y, auto x) { (*y) = (*y) + x; };
-      switch (out.dtype()) {
-        case bool_:
-          reduction_op<InT, bool>(in, out, axes, false, op);
-          break;
-        case uint8:
-          reduction_op<InT, uint8_t>(in, out, axes, 0, op);
-          break;
-        case uint16:
-          reduction_op<InT, uint16_t>(in, out, axes, 0, op);
-          break;
-        case uint32:
-          reduction_op<InT, uint32_t>(in, out, axes, 0, op);
-          break;
-        case uint64:
-          reduction_op<InT, uint64_t>(in, out, axes, 0, op);
-          break;
-        case int8:
-          reduction_op<InT, int8_t>(in, out, axes, 0, op);
-          break;
-        case int16:
-          reduction_op<InT, int16_t>(in, out, axes, 0, op);
-          break;
-        case int32:
-          reduction_op<InT, int32_t>(in, out, axes, 0, op);
-          break;
-        case int64:
-          reduction_op<InT, int64_t>(in, out, axes, 0, op);
-          break;
-        case float16:
-          reduction_op<InT, float16_t>(in, out, axes, 0.0f, op);
-          break;
-        case float32:
-          reduction_op<InT, float>(in, out, axes, 0.0f, op);
-          break;
-        case bfloat16:
-          reduction_op<InT, bfloat16_t>(in, out, axes, 0.0f, op);
-          break;
-        case complex64:
-          reduction_op<InT, complex64_t>(in, out, axes, complex64_t{0.0f}, op);
-          break;
+      if (out.dtype() == int32) {
+        // special case since the input type can be bool
+        reduction_op<InT, int32_t>(in, out, axes, 0, op);
+      } else {
+        reduction_op<InT, InT>(in, out, axes, 0, op);
       }
-    } break;
+      break;
+    }
     case Reduce::Prod: {
       auto op = [](auto y, auto x) { (*y) *= x; };
       reduction_op<InT, InT>(in, out, axes, 1, op);
@@ -168,6 +134,29 @@ void reduce_dispatch_out(
 
 } // namespace
 
+void nd_loop(
+    std::function<void(int)> callback,
+    const std::vector<int>& shape,
+    const std::vector<size_t>& strides) {
+  std::function<void(int, int)> loop_inner;
+  loop_inner = [&](int dim, int offset) {
+    if (dim < shape.size() - 1) {
+      int size = shape[dim];
+      size_t stride = strides[dim];
+      for (int i = 0; i < size; i++) {
+        loop_inner(dim + 1, offset + i * stride);
+      }
+    } else {
+      int size = shape[dim];
+      size_t stride = strides[dim];
+      for (int i = 0; i < size; i++) {
+        callback(offset + i * stride);
+      }
+    }
+  };
+  loop_inner(0, 0);
+}
+
 void Reduce::eval(const std::vector<array>& inputs, array& out) {
   assert(inputs.size() == 1);
   auto& in = inputs[0];

mlx/backend/common/reduce.h

@@ -49,47 +49,18 @@ struct ReductionPlan {
   ReductionPlan(ReductionOpType type_) : type(type_) {}
 };
 
-namespace {
+ReductionPlan get_reduction_plan(const array& x, const std::vector<int> axes);
 
 // Helper for the ndimensional strided loop
 // Should this be in utils?
-inline void nd_loop(
+void nd_loop(
     std::function<void(int)> callback,
     const std::vector<int>& shape,
-    const std::vector<size_t>& strides) {
-  std::function<void(int, int)> loop_inner;
-  loop_inner = [&](int dim, int offset) {
-    if (dim < shape.size() - 1) {
-      int size = shape[dim];
-      size_t stride = strides[dim];
-      for (int i = 0; i < size; i++) {
-        loop_inner(dim + 1, offset + i * stride);
-      }
-    } else {
-      int size = shape[dim];
-      size_t stride = strides[dim];
-      for (int i = 0; i < size; i++) {
-        callback(offset + i * stride);
-      }
-    }
-  };
-  loop_inner(0, 0);
-}
+    const std::vector<size_t>& strides);
 
 std::pair<std::vector<int>, std::vector<size_t>> shapes_without_reduction_axes(
     const array& x,
-    const std::vector<int>& axes) {
-  std::vector<int> shape = x.shape();
-  std::vector<size_t> strides = x.strides();
-
-  for (int i = axes.size() - 1; i >= 0; i--) {
-    int a = axes[i];
-    shape.erase(shape.begin() + a);
-    strides.erase(strides.begin() + a);
-  }
-
-  return std::make_pair(shape, strides);
-}
+    const std::vector<int>& axes);
 
 template <typename T, typename U, typename Op>
 struct DefaultStridedReduce {
@@ -123,102 +94,6 @@ struct DefaultContiguousReduce {
   }
 };
 
-ReductionPlan get_reduction_plan(const array& x, const std::vector<int> axes) {
-  // The data is all there and we are reducing over everything
-  if (x.size() == x.data_size() && axes.size() == x.ndim() &&
-      x.flags().contiguous) {
-    return ContiguousAllReduce;
-  }
-
-  // Row contiguous input so the output is row contiguous
-  if (x.flags().row_contiguous) {
-    // Merge consecutive axes
-    std::vector<int> shape = {x.shape(axes[0])};
-    std::vector<size_t> strides = {x.strides()[axes[0]]};
-    for (int i = 1; i < axes.size(); i++) {
-      if (axes[i] - 1 == axes[i - 1]) {
-        shape.back() *= x.shape(axes[i]);
-        strides.back() = x.strides()[axes[i]];
-      } else {
-        shape.push_back(x.shape(axes[i]));
-        strides.push_back(x.strides()[axes[i]]);
-      }
-    }
-
-    if (strides.back() == 1) {
-      return ReductionPlan(ContiguousReduce, shape, strides);
-    } else if (strides.back() > 1) {
-      return ReductionPlan(ContiguousStridedReduce, shape, strides);
-    }
-  }
-
-  // Let's check if we can optimize our access patterns
-  //
-  // 1. We have a reduction axis with stride 1. Simply call
-  //    GeneralContiguousReduce and be done with it.
-  // 2. We have transpositions and we are not reducing over the axis with
-  //    stride 1. However, we are reducing over an axis where everything is
-  //    contiguous in memory to the right of that axis. We can call strided
-  //    reduce and be done with it.
-  // 2. We have weird transpositions and expands. Copy the strides to the
-  //    output, then call strided reduce.
-
-  // Sort reduction axes by stride in order to merge them and figure out if we
-  // have a contiguous reduction.
-  std::vector<std::pair<int, size_t>> reductions;
-  for (auto a : axes) {
-    reductions.push_back(std::make_pair(x.shape(a), x.strides()[a]));
-  }
-  std::sort(reductions.begin(), reductions.end(), [](auto a, auto b) {
-    return a.second > b.second;
-  });
-  // Extract the two smallest and try to merge them in case the contiguous
-  // reduction can be bigger than just the last axis.
-  for (int i = reductions.size() - 1; i >= 1; i--) {
-    auto a = reductions[i];
-    auto b = reductions[i - 1];
-
-    // b.stride = a.shape * a.stride then a and b are contiguous
-    if (b.second == a.first * a.second) {
-      reductions.erase(reductions.begin() + i);
-      reductions[i - 1] = std::make_pair(a.first * b.first, a.second);
-    }
-  }
-
-  std::vector<int> shape;
-  std::vector<size_t> strides;
-  for (auto r : reductions) {
-    shape.push_back(r.first);
-    strides.push_back(r.second);
-  }
-
-  // We can call the contiguous reduction op for every weird way the input is
-  // structured in the rest of the axes.
-  if (strides.back() == 1) {
-    return ReductionPlan(GeneralContiguousReduce, shape, strides);
-  }
-
-  // Delegate to the general strided reduction op if the axes after
-  // strides.back() are contiguous.
-  if (strides.back() > 1) {
-    int size = 1;
-    for (int i = x.ndim() - 1; i >= 0; i--) {
-      if (axes.back() == i) {
-        continue;
-      }
-      if (x.strides()[i] != size) {
-        break;
-      }
-      size *= x.shape(i);
-    }
-    if (size >= strides.back()) {
-      return ReductionPlan(GeneralStridedReduce, shape, strides);
-    }
-  }
-
-  return ReductionPlan(GeneralReduce, shape, strides);
-}
-
 template <typename T, typename U, typename OpS, typename OpC, typename Op>
 void reduction_op(
     const array& x,
@@ -361,6 +236,4 @@ void reduction_op(
   reduction_op<T, U>(x, out, axes, init, ops, opc, op);
 }
 
-} // namespace
-
 } // namespace mlx::core

mlx/backend/common/reduce_utils.cpp

@@ -0,0 +1,118 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/reduce.h"
+
+namespace mlx::core {
+
+std::pair<std::vector<int>, std::vector<size_t>> shapes_without_reduction_axes(
+    const array& x,
+    const std::vector<int>& axes) {
+  std::vector<int> shape = x.shape();
+  std::vector<size_t> strides = x.strides();
+
+  for (int i = axes.size() - 1; i >= 0; i--) {
+    int a = axes[i];
+    shape.erase(shape.begin() + a);
+    strides.erase(strides.begin() + a);
+  }
+
+  return std::make_pair(shape, strides);
+}
+
+ReductionPlan get_reduction_plan(const array& x, const std::vector<int> axes) {
+  // The data is all there and we are reducing over everything
+  if (x.size() == x.data_size() && axes.size() == x.ndim() &&
+      x.flags().contiguous) {
+    return ContiguousAllReduce;
+  }
+
+  // Row contiguous input so the output is row contiguous
+  if (x.flags().row_contiguous) {
+    // Merge consecutive axes
+    std::vector<int> shape = {x.shape(axes[0])};
+    std::vector<size_t> strides = {x.strides()[axes[0]]};
+    for (int i = 1; i < axes.size(); i++) {
+      if (axes[i] - 1 == axes[i - 1]) {
+        shape.back() *= x.shape(axes[i]);
+        strides.back() = x.strides()[axes[i]];
+      } else {
+        shape.push_back(x.shape(axes[i]));
+        strides.push_back(x.strides()[axes[i]]);
+      }
+    }
+
+    if (strides.back() == 1) {
+      return ReductionPlan(ContiguousReduce, shape, strides);
+    } else if (strides.back() > 1) {
+      return ReductionPlan(ContiguousStridedReduce, shape, strides);
+    }
+  }
+
+  // Let's check if we can optimize our access patterns
+  //
+  // 1. We have a reduction axis with stride 1. Simply call
+  //    GeneralContiguousReduce and be done with it.
+  // 2. We have transpositions and we are not reducing over the axis with
+  //    stride 1. However, we are reducing over an axis where everything is
+  //    contiguous in memory to the right of that axis. We can call strided
+  //    reduce and be done with it.
+  // 2. We have weird transpositions and expands. Copy the strides to the
+  //    output, then call strided reduce.
+
+  // Sort reduction axes by stride in order to merge them and figure out if we
+  // have a contiguous reduction.
+  std::vector<std::pair<int, size_t>> reductions;
+  for (auto a : axes) {
+    reductions.push_back(std::make_pair(x.shape(a), x.strides()[a]));
+  }
+  std::sort(reductions.begin(), reductions.end(), [](auto a, auto b) {
+    return a.second > b.second;
+  });
+  // Extract the two smallest and try to merge them in case the contiguous
+  // reduction can be bigger than just the last axis.
+  for (int i = reductions.size() - 1; i >= 1; i--) {
+    auto a = reductions[i];
+    auto b = reductions[i - 1];
+
+    // b.stride = a.shape * a.stride then a and b are contiguous
+    if (b.second == a.first * a.second) {
+      reductions.erase(reductions.begin() + i);
+      reductions[i - 1] = std::make_pair(a.first * b.first, a.second);
+    }
+  }
+
+  std::vector<int> shape;
+  std::vector<size_t> strides;
+  for (auto r : reductions) {
+    shape.push_back(r.first);
+    strides.push_back(r.second);
+  }
+
+  // We can call the contiguous reduction op for every weird way the input is
+  // structured in the rest of the axes.
+  if (strides.back() == 1) {
+    return ReductionPlan(GeneralContiguousReduce, shape, strides);
+  }
+
+  // Delegate to the general strided reduction op if the axes after
+  // strides.back() are contiguous.
+  if (strides.back() > 1) {
+    int size = 1;
+    for (int i = x.ndim() - 1; i >= 0; i--) {
+      if (axes.back() == i) {
+        continue;
+      }
+      if (x.strides()[i] != size) {
+        break;
+      }
+      size *= x.shape(i);
+    }
+    if (size >= strides.back()) {
+      return ReductionPlan(GeneralStridedReduce, shape, strides);
+    }
+  }
+
+  return ReductionPlan(GeneralReduce, shape, strides);
+}
+
+} // namespace mlx::core

mlx/backend/common/sort.cpp

@@ -113,14 +113,14 @@ void sort(const array& in, array& out, int axis) {
   axis = axis < 0 ? axis + in.ndim() : axis;
   size_t n_rows = in.size() / in.shape(axis);
 
-  auto remaining_shape = in.shape();
+  auto remaining_shape = out.shape();
   remaining_shape.erase(remaining_shape.begin() + axis);
 
-  auto remaining_strides = in.strides();
+  auto remaining_strides = out.strides();
   remaining_strides.erase(remaining_strides.begin() + axis);
 
-  size_t axis_stride = in.strides()[axis];
-  int axis_size = in.shape(axis);
+  size_t axis_stride = out.strides()[axis];
+  int axis_size = out.shape(axis);
 
   // Perform sorting in place
   for (int i = 0; i < n_rows; i++) {
@@ -143,34 +143,42 @@ void argsort(const array& in, array& out, int axis) {
   axis = axis < 0 ? axis + in.ndim() : axis;
   size_t n_rows = in.size() / in.shape(axis);
 
-  auto remaining_shape = in.shape();
-  remaining_shape.erase(remaining_shape.begin() + axis);
+  auto in_remaining_shape = in.shape();
+  in_remaining_shape.erase(in_remaining_shape.begin() + axis);
 
-  auto remaining_strides = in.strides();
-  remaining_strides.erase(remaining_strides.begin() + axis);
+  auto in_remaining_strides = in.strides();
+  in_remaining_strides.erase(in_remaining_strides.begin() + axis);
 
-  size_t axis_stride = in.strides()[axis];
+  auto out_remaining_shape = out.shape();
+  out_remaining_shape.erase(out_remaining_shape.begin() + axis);
+
+  auto out_remaining_strides = out.strides();
+  out_remaining_strides.erase(out_remaining_strides.begin() + axis);
+
+  size_t in_stride = in.strides()[axis];
+  size_t out_stride = out.strides()[axis];
   int axis_size = in.shape(axis);
 
   // Perform sorting
   for (int i = 0; i < n_rows; i++) {
-    size_t loc = elem_to_loc(i, remaining_shape, remaining_strides);
-    const T* data_ptr = in.data<T>() + loc;
-    IdxT* idx_ptr = out.data<IdxT>() + loc;
+    size_t in_loc = elem_to_loc(i, in_remaining_shape, in_remaining_strides);
+    size_t out_loc = elem_to_loc(i, out_remaining_shape, out_remaining_strides);
+    const T* data_ptr = in.data<T>() + in_loc;
+    IdxT* idx_ptr = out.data<IdxT>() + out_loc;
 
-    StridedIterator st_(idx_ptr, axis_stride, 0);
-    StridedIterator ed_(idx_ptr, axis_stride, axis_size);
+    StridedIterator st_(idx_ptr, out_stride, 0);
+    StridedIterator ed_(idx_ptr, out_stride, axis_size);
 
     // Initialize with iota
     std::iota(st_, ed_, IdxT(0));
 
     // Sort according to vals
-    StridedIterator st(idx_ptr, axis_stride, 0);
-    StridedIterator ed(idx_ptr, axis_stride, axis_size);
+    StridedIterator st(idx_ptr, out_stride, 0);
+    StridedIterator ed(idx_ptr, out_stride, axis_size);
 
-    std::stable_sort(st, ed, [data_ptr, axis_stride](IdxT a, IdxT b) {
-      auto v1 = data_ptr[a * axis_stride];
-      auto v2 = data_ptr[b * axis_stride];
+    std::stable_sort(st, ed, [data_ptr, in_stride](IdxT a, IdxT b) {
+      auto v1 = data_ptr[a * in_stride];
+      auto v2 = data_ptr[b * in_stride];
       return v1 < v2 || (v1 == v2 && a < b);
     });
   }

mlx/backend/common/utils.h

@@ -29,6 +29,15 @@ inline size_t elem_to_loc(int elem, const array& a) {
   return elem_to_loc(elem, a.shape(), a.strides());
 }
 
+template <typename stride_t>
+std::vector<stride_t> make_contiguous_strides(const std::vector<int>& shape) {
+  std::vector<stride_t> strides(shape.size(), 1);
+  for (int i = shape.size() - 1; i > 0; i--) {
+    strides[i - 1] = strides[i] * shape[i];
+  }
+  return strides;
+}
+
 // Collapse dims that are contiguous to possibly route to a better kernel
 // e.g. for x = transpose(array({0, 1, 2, 3, 4, 5, 6, 7}, {2, 2, 2}), {2, 0, 1})
 // should return {{2, 4}, {{1, 2}}}.

mlx/backend/metal/CMakeLists.txt

@@ -18,7 +18,7 @@ function(make_jit_source SRC_FILE)
               ${CMAKE_C_COMPILER}
               ${PROJECT_SOURCE_DIR}
               ${SRC_FILE}
-              "-D${MLX_METAL_VERSION}"
+              "-DMLX_METAL_VERSION=${MLX_METAL_VERSION}"
     DEPENDS make_compiled_preamble.sh
             kernels/${SRC_FILE}.h
             ${ARGN}
@@ -52,6 +52,7 @@ make_jit_source(
 )
 make_jit_source(scatter)
 make_jit_source(gather)
+make_jit_source(hadamard)
 
 if (MLX_METAL_JIT) 
   target_sources(
@@ -112,6 +113,8 @@ if (MLX_METAL_JIT)
     kernels/steel/defines.h
     kernels/steel/conv/loaders/loader_general.h
   )
+  make_jit_source(quantized)
+  make_jit_source(gemv_masked)
 else()
   target_sources(
     mlx
@@ -131,6 +134,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fft.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/matmul.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/scaled_dot_product_attention.cpp
@@ -146,6 +150,7 @@ target_sources(
   ${CMAKE_CURRENT_SOURCE_DIR}/reduce.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/ternary.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/unary.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp
 )
 
 if (NOT MLX_METAL_PATH)

mlx/backend/metal/allocator.cpp

@@ -242,8 +242,17 @@ void MetalAllocator::free(Buffer buffer) {
 }
 
 MetalAllocator& allocator() {
-  static MetalAllocator allocator_;
-  return allocator_;
+  // By creating the |allocator_| on heap, the destructor of MetalAllocator will
+  // not be called on exit and all the buffers will be leaked. This is necessary
+  // because releasing buffers can take more than 30sec when the program holds a
+  // lot of RAM (for example inferencing a LLM), and it would feel frozen to
+  // users when exiting.
+  // TODO(zcbenz): Consider using the `base::NoDestructor` class from Chromium
+  // when applying this pattern to more places, or when introducing sanitizers
+  // to MLX.
+  // https://source.chromium.org/chromium/chromium/src/+/main:base/no_destructor.h
+  static MetalAllocator* allocator_ = new MetalAllocator;
+  return *allocator_;
 }
 
 size_t set_cache_limit(size_t limit) {

mlx/backend/metal/binary.cpp

@@ -6,14 +6,58 @@
 #include "mlx/backend/metal/utils.h"
 #include "mlx/primitives.h"
 
+#define BINARY_GPU(func)                                              \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    binary_op_gpu(inputs, out, get_primitive_string(this));           \
+  }
+
+#define BINARY_GPU_MULTI(func)                                         \
+  void func::eval_gpu(                                                 \
+      const std::vector<array>& inputs, std::vector<array>& outputs) { \
+    binary_op_gpu(inputs, outputs, get_primitive_string(this));        \
+  }
+
 namespace mlx::core {
 
 constexpr int MAX_BINARY_SPECIALIZED_DIMS = 5;
 
+std::string get_kernel_name(
+    BinaryOpType bopt,
+    const std::string& op,
+    const array& a,
+    bool use_2d,
+    int ndim) {
+  std::ostringstream kname;
+  switch (bopt) {
+    case BinaryOpType::ScalarScalar:
+      kname << "ss";
+      break;
+    case BinaryOpType::ScalarVector:
+      kname << (use_2d ? "sv2" : "sv");
+      break;
+    case BinaryOpType::VectorScalar:
+      kname << (use_2d ? "vs2" : "vs");
+      break;
+    case BinaryOpType::VectorVector:
+      kname << (use_2d ? "vv2" : "vv");
+      break;
+    case BinaryOpType::General:
+      kname << "g";
+      if (ndim <= MAX_BINARY_SPECIALIZED_DIMS) {
+        kname << ndim;
+      } else {
+        kname << "n";
+      }
+      break;
+  }
+  kname << op << type_to_name(a);
+  return kname.str();
+}
+
 void binary_op_gpu_inplace(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
-    const std::string op,
+    const std::string& op,
     const Stream& s) {
   auto& a = inputs[0];
   auto& b = inputs[1];
@@ -30,38 +74,12 @@ void binary_op_gpu_inplace(
   auto& strides_b = strides[1];
   auto& strides_out = strides[2];
 
-  std::string kernel_name;
-  {
-    std::ostringstream kname;
-    switch (bopt) {
-      case BinaryOpType::ScalarScalar:
-        kname << "ss";
-        break;
-      case BinaryOpType::ScalarVector:
-        kname << "sv";
-        break;
-      case BinaryOpType::VectorScalar:
-        kname << "vs";
-        break;
-      case BinaryOpType::VectorVector:
-        kname << "vv";
-        break;
-      case BinaryOpType::General:
-        kname << "g";
-        if (shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
-          kname << shape.size();
-        } else {
-          kname << "n";
-        }
-        break;
-    }
-    kname << op << type_to_name(a);
-    kernel_name = kname.str();
-  }
-
+  bool use_2d = out.data_size() > UINT32_MAX;
+  std::string kernel_name = get_kernel_name(bopt, op, a, use_2d, shape.size());
   auto& d = metal::device(s.device);
 
-  auto kernel = get_binary_two_kernel(d, kernel_name, a, outputs[0]);
+  auto kernel =
+      get_binary_two_kernel(d, kernel_name, a.dtype(), outputs[0].dtype(), op);
 
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
@@ -105,9 +123,11 @@ void binary_op_gpu_inplace(
     MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
     compute_encoder.dispatchThreads(grid_dims, group_dims);
   } else {
-    // Launch a 1D grid of threads
+    // Launch a 1D or 2D grid of threads
     size_t nthreads = out.data_size();
-    MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+    MTL::Size grid_dims = use_2d
+        ? get_2d_grid_dims(outputs[0].shape(), outputs[0].strides())
+        : MTL::Size(nthreads, 1, 1);
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size > nthreads) {
       thread_group_size = nthreads;
@@ -120,7 +140,7 @@ void binary_op_gpu_inplace(
 void binary_op_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
-    const std::string op,
+    const std::string& op,
     const Stream& s) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
@@ -134,7 +154,7 @@ void binary_op_gpu(
 void binary_op_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
-    const std::string op) {
+    const std::string& op) {
   auto& s = outputs[0].primitive().stream();
   binary_op_gpu(inputs, outputs, op, s);
 }
@@ -142,7 +162,7 @@ void binary_op_gpu(
 void binary_op_gpu_inplace(
     const std::vector<array>& inputs,
     array& out,
-    const std::string op,
+    const std::string& op,
     const Stream& s) {
   auto& a = inputs[0];
   auto& b = inputs[1];
@@ -157,38 +177,11 @@ void binary_op_gpu_inplace(
   auto& strides_b = strides[1];
   auto& strides_out = strides[2];
 
-  std::string kernel_name;
-  {
-    std::ostringstream kname;
-    switch (bopt) {
-      case BinaryOpType::ScalarScalar:
-        kname << "ss";
-        break;
-      case BinaryOpType::ScalarVector:
-        kname << "sv";
-        break;
-      case BinaryOpType::VectorScalar:
-        kname << "vs";
-        break;
-      case BinaryOpType::VectorVector:
-        kname << "vv";
-        break;
-      case BinaryOpType::General:
-        kname << "g";
-        if (shape.size() <= MAX_BINARY_SPECIALIZED_DIMS) {
-          kname << shape.size();
-        } else {
-          kname << "n";
-        }
-        break;
-    }
-    kname << op << type_to_name(a);
-    kernel_name = kname.str();
-  }
-
+  bool use_2d = out.data_size() > UINT32_MAX;
+  std::string kernel_name = get_kernel_name(bopt, op, a, use_2d, shape.size());
   auto& d = metal::device(s.device);
 
-  auto kernel = get_binary_kernel(d, kernel_name, a, out);
+  auto kernel = get_binary_kernel(d, kernel_name, a.dtype(), out.dtype(), op);
   auto& compute_encoder = d.get_command_encoder(s.index);
   compute_encoder->setComputePipelineState(kernel);
   bool donate_a = a.data_shared_ptr() == nullptr;
@@ -225,10 +218,11 @@ void binary_op_gpu_inplace(
     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 =
-        bopt == BinaryOpType::General ? out.size() : out.data_size();
-    MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+    // Launch a 1D or 2D grid of threads
+
+    size_t nthreads = out.data_size();
+    MTL::Size grid_dims = use_2d ? get_2d_grid_dims(out.shape(), out.strides())
+                                 : MTL::Size(nthreads, 1, 1);
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size > nthreads) {
       thread_group_size = nthreads;
@@ -241,7 +235,7 @@ void binary_op_gpu_inplace(
 void binary_op_gpu(
     const std::vector<array>& inputs,
     array& out,
-    const std::string op,
+    const std::string& op,
     const Stream& s) {
   assert(inputs.size() == 2);
   auto& a = inputs[0];
@@ -254,107 +248,49 @@ void binary_op_gpu(
 void binary_op_gpu(
     const std::vector<array>& inputs,
     array& out,
-    const std::string op) {
+    const std::string& op) {
   auto& s = out.primitive().stream();
   binary_op_gpu(inputs, out, op, s);
 }
 
-void Add::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "add");
-}
-
-void ArcTan2::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "arctan2");
-}
+BINARY_GPU(Add)
+BINARY_GPU(ArcTan2)
+BINARY_GPU(Divide)
+BINARY_GPU_MULTI(DivMod)
+BINARY_GPU(Remainder)
+BINARY_GPU(Equal)
+BINARY_GPU(Greater)
+BINARY_GPU(GreaterEqual)
+BINARY_GPU(Less)
+BINARY_GPU(LessEqual)
+BINARY_GPU(LogicalAnd)
+BINARY_GPU(LogicalOr)
+BINARY_GPU(LogAddExp)
+BINARY_GPU(Maximum)
+BINARY_GPU(Minimum)
+BINARY_GPU(Multiply)
+BINARY_GPU(NotEqual)
+BINARY_GPU(Power)
+BINARY_GPU(Subtract)
 
 void BitwiseBinary::eval_gpu(const std::vector<array>& inputs, array& out) {
   switch (op_) {
     case BitwiseBinary::And:
-      binary_op_gpu(inputs, out, "bitwise_and");
+      binary_op_gpu(inputs, out, get_primitive_string(this));
       break;
     case BitwiseBinary::Or:
-      binary_op_gpu(inputs, out, "bitwise_or");
+      binary_op_gpu(inputs, out, get_primitive_string(this));
       break;
     case BitwiseBinary::Xor:
-      binary_op_gpu(inputs, out, "bitwise_xor");
+      binary_op_gpu(inputs, out, get_primitive_string(this));
       break;
     case BitwiseBinary::LeftShift:
-      binary_op_gpu(inputs, out, "left_shift");
+      binary_op_gpu(inputs, out, get_primitive_string(this));
       break;
     case BitwiseBinary::RightShift:
-      binary_op_gpu(inputs, out, "right_shift");
+      binary_op_gpu(inputs, out, get_primitive_string(this));
       break;
   }
 }
 
-void Divide::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "div");
-}
-
-void DivMod::eval_gpu(
-    const std::vector<array>& inputs,
-    std::vector<array>& outputs) {
-  binary_op_gpu(inputs, outputs, "divmod");
-}
-
-void Remainder::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "rem");
-}
-
-void Equal::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, equal_nan_ ? "naneq" : "eq");
-}
-
-void Greater::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "ge");
-}
-
-void GreaterEqual::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "geq");
-}
-
-void Less::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "le");
-}
-
-void LessEqual::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "leq");
-}
-
-void LogicalAnd::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "land");
-}
-
-void LogicalOr::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "lor");
-}
-
-void LogAddExp::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "lae");
-}
-
-void Maximum::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "max");
-}
-
-void Minimum::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "min");
-}
-
-void Multiply::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "mul");
-}
-
-void NotEqual::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "neq");
-}
-
-void Power::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "pow");
-}
-
-void Subtract::eval_gpu(const std::vector<array>& inputs, array& out) {
-  binary_op_gpu(inputs, out, "sub");
-}
-
 } // namespace mlx::core

mlx/backend/metal/binary.h

@@ -9,25 +9,25 @@ namespace mlx::core {
 void binary_op_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
-    const std::string op,
+    const std::string& op,
     const Stream& s);
 
 void binary_op_gpu(
     const std::vector<array>& inputs,
     array& out,
-    const std::string op,
+    const std::string& op,
     const Stream& s);
 
 void binary_op_gpu_inplace(
     const std::vector<array>& inputs,
     std::vector<array>& outputs,
-    const std::string op,
+    const std::string& op,
     const Stream& s);
 
 void binary_op_gpu_inplace(
     const std::vector<array>& inputs,
     array& out,
-    const std::string op,
+    const std::string& op,
     const Stream& s);
 
 } // namespace mlx::core

mlx/backend/metal/copy.cpp

@@ -64,16 +64,17 @@ void copy_gpu_inplace(
   auto& strides_in_ = strides[0];
   auto& strides_out_ = strides[1];
 
+  bool use_2d = out.data_size() > UINT32_MAX;
   auto& d = metal::device(s.device);
   std::string kernel_name;
   {
     std::ostringstream kname;
     switch (ctype) {
       case CopyType::Scalar:
-        kname << "s";
+        kname << (use_2d ? "s2" : "s");
         break;
       case CopyType::Vector:
-        kname << "v";
+        kname << (use_2d ? "v2" : "v");
         break;
       case CopyType::General:
         kname << "g";
@@ -139,7 +140,8 @@ void copy_gpu_inplace(
     compute_encoder.dispatchThreads(grid_dims, group_dims);
   } else {
     size_t nthreads = out.data_size();
-    MTL::Size grid_dims = MTL::Size(nthreads, 1, 1);
+    MTL::Size grid_dims = use_2d ? get_2d_grid_dims(out.shape(), out.strides())
+                                 : MTL::Size(nthreads, 1, 1);
     NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
     if (thread_group_size > nthreads) {
       thread_group_size = nthreads;

mlx/backend/metal/device.cpp

@@ -14,7 +14,6 @@
 #include "mlx/backend/metal/device.h"
 #include "mlx/backend/metal/metal.h"
 #include "mlx/backend/metal/metal_impl.h"
-#include "mlx/backend/metal/mps/gemm.h"
 #include "mlx/backend/metal/utils.h"
 
 namespace fs = std::filesystem;
@@ -30,13 +29,29 @@ constexpr int MAX_DISPATCHES_PER_ENCODER = 2;
 constexpr const char* default_mtllib_path = METAL_PATH;
 
 constexpr auto get_metal_version() {
-#if defined METAL_3_1
+#if (MLX_METAL_VERSION >= 320)
+  return MTL::LanguageVersion3_2;
+#elif (MLX_METAL_VERSION >= 310)
   return MTL::LanguageVersion3_1;
 #else
   return MTL::LanguageVersion3_0;
 #endif
 }
 
+std::string get_colocated_mtllib_path(const std::string& lib_name) {
+  Dl_info info;
+  std::string mtllib_path;
+  std::string lib_ext = lib_name + ".metallib";
+
+  int success = dladdr((void*)get_colocated_mtllib_path, &info);
+  if (success) {
+    auto mtllib = fs::path(info.dli_fname).remove_filename() / lib_ext;
+    mtllib_path = mtllib.c_str();
+  }
+
+  return mtllib_path;
+}
+
 auto load_device() {
   auto devices = MTL::CopyAllDevices();
   auto device = static_cast<MTL::Device*>(devices->object(0))
@@ -124,6 +139,49 @@ MTL::Library* load_library(
 
 } // namespace
 
+CommandEncoder::CommandEncoder(MTL::CommandBuffer* cbuf) : cbuf(cbuf) {
+  enc = cbuf->computeCommandEncoder(MTL::DispatchTypeConcurrent);
+  enc->retain();
+}
+
+CommandEncoder::~CommandEncoder() {
+  enc->endEncoding();
+  enc->release();
+}
+
+void CommandEncoder::set_input_array(
+    const array& a,
+    int idx,
+    int64_t offset /* = 0 */) {
+  auto r_buf = static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
+  if (auto it = outputs.find(r_buf); it != outputs.end()) {
+    // Insert a barrier
+    enc->memoryBarrier(&r_buf, 1);
+
+    // Remove the output
+    outputs.erase(it);
+  }
+  auto a_buf = static_cast<const MTL::Buffer*>(a.buffer().ptr());
+  auto base_offset = a.data<char>() -
+      static_cast<char*>(const_cast<MTL::Buffer*>(a_buf)->contents());
+  base_offset += offset;
+  enc->setBuffer(a_buf, base_offset, idx);
+}
+
+void CommandEncoder::set_output_array(
+    array& a,
+    int idx,
+    int64_t offset /* = 0 */) {
+  // Add barriers before adding the output to the output set
+  set_input_array(a, idx, offset);
+  auto buf = static_cast<MTL::Resource*>(a.buffer().ptr());
+  if (concurrent) {
+    concurrent_outputs.insert(buf);
+  } else {
+    outputs.insert(buf);
+  }
+}
+
 void CommandEncoder::dispatchThreadgroups(
     MTL::Size grid_dims,
     MTL::Size group_dims) {
@@ -253,13 +311,9 @@ void Device::register_library(
   }
 }
 
-void Device::register_library(
-    const std::string& lib_name,
-    const std::function<std::string(const std::string&)>& lib_path_func) {
+void Device::register_library(const std::string& lib_name) {
   if (auto it = library_map_.find(lib_name); it == library_map_.end()) {
-    std::string new_lib_path = lib_path_func(lib_name);
-    auto new_lib = load_library(device_, lib_name, new_lib_path.c_str());
-    library_map_.insert({lib_name, new_lib});
+    register_library(lib_name, get_colocated_mtllib_path(lib_name));
   }
 }
 
@@ -269,7 +323,7 @@ MTL::Library* Device::get_library_cache_(const std::string& lib_name) {
   if (auto it = library_map_.find(lib_name); it != library_map_.end()) {
     mtl_lib = it->second;
   } else { // Look for metallib alongside library
-    register_library(lib_name);
+    register_library(lib_name, get_colocated_mtllib_path(lib_name));
     mtl_lib = library_map_[lib_name];
   }
 

mlx/backend/metal/device.h

@@ -9,38 +9,16 @@
 #include <unordered_map>
 #include <unordered_set>
 
-#include <dlfcn.h>
-#include <filesystem>
-
 #include "mlx/array.h"
 #include "mlx/device.h"
 
-namespace fs = std::filesystem;
-
 namespace mlx::core::metal {
 
-inline std::string get_colocated_mtllib_path(const std::string& lib_name) {
-  Dl_info info;
-  std::string mtllib_path;
-  std::string lib_ext = lib_name + ".metallib";
-
-  int success = dladdr((void*)get_colocated_mtllib_path, &info);
-  if (success) {
-    auto mtllib = fs::path(info.dli_fname).remove_filename() / lib_ext;
-    mtllib_path = mtllib.c_str();
-  }
-
-  return mtllib_path;
-}
-
 using MTLFCList =
     std::vector<std::tuple<const void*, MTL::DataType, NS::UInteger>>;
 
 struct CommandEncoder {
-  CommandEncoder(MTL::CommandBuffer* cbuf) : cbuf(cbuf) {
-    enc = cbuf->computeCommandEncoder(MTL::DispatchTypeConcurrent);
-    enc->retain();
-  };
+  CommandEncoder(MTL::CommandBuffer* cbuf);
   CommandEncoder(const CommandEncoder&) = delete;
   CommandEncoder& operator=(const CommandEncoder&) = delete;
 
@@ -63,34 +41,8 @@ struct CommandEncoder {
     return enc;
   }
 
-  void set_input_array(const array& a, int idx, int64_t offset = 0) {
-    auto r_buf =
-        static_cast<MTL::Resource*>(const_cast<void*>(a.buffer().ptr()));
-    if (auto it = outputs.find(r_buf); it != outputs.end()) {
-      // Insert a barrier
-      enc->memoryBarrier(&r_buf, 1);
-
-      // Remove the output
-      outputs.erase(it);
-    }
-    auto a_buf = static_cast<const MTL::Buffer*>(a.buffer().ptr());
-    auto base_offset = a.data<char>() -
-        static_cast<char*>(const_cast<MTL::Buffer*>(a_buf)->contents());
-    base_offset += offset;
-    enc->setBuffer(a_buf, base_offset, idx);
-  }
-
-  void set_output_array(array& a, int idx, int64_t offset = 0) {
-    // Add barriers before adding the output to the output set
-    set_input_array(a, idx, offset);
-    auto buf = static_cast<MTL::Resource*>(a.buffer().ptr());
-    if (concurrent) {
-      concurrent_outputs.insert(buf);
-    } else {
-      outputs.insert(buf);
-    }
-  }
-
+  void set_input_array(const array& a, int idx, int64_t offset = 0);
+  void set_output_array(array& a, int idx, int64_t offset = 0);
   void dispatchThreadgroups(MTL::Size grid_dims, MTL::Size group_dims);
   void dispatchThreads(MTL::Size grid_dims, MTL::Size group_dims);
 
@@ -98,10 +50,7 @@ struct CommandEncoder {
     return ConcurrentContext(*this);
   }
 
-  ~CommandEncoder() {
-    enc->endEncoding();
-    enc->release();
-  }
+  ~CommandEncoder();
 
  private:
   void maybe_split();
@@ -136,10 +85,8 @@ class Device {
   void register_library(
       const std::string& lib_name,
       const std::string& lib_path);
-  void register_library(
-      const std::string& lib_name,
-      const std::function<std::string(const std::string&)>& lib_path_func =
-          get_colocated_mtllib_path);
+
+  void register_library(const std::string& lib_name);
 
   MTL::Library* get_library(const std::string& name);
 

mlx/backend/metal/fft.cpp

@@ -1,4 +1,4 @@
-// Copyright © 2023 Apple Inc.
+// Copyright © 2024 Apple Inc.
 #include <cassert>
 #include <complex>
 #include <map>
@@ -12,8 +12,7 @@
 #include "mlx/backend/metal/slicing.h"
 #include "mlx/backend/metal/unary.h"
 #include "mlx/backend/metal/utils.h"
-#include "mlx/mlx.h"
-#include "mlx/primitives.h"
+#include "mlx/utils.h"
 
 namespace mlx::core {
 
@@ -391,16 +390,16 @@ void multi_upload_bluestein_fft(
     std::vector<int> rstrides(in.ndim(), 1);
     rstarts[axis] = in.shape(axis) - back_offset;
     rstrides[axis] = -1;
-    unary_op_gpu({in}, conj_temp, "conj", s);
+    unary_op_gpu({in}, conj_temp, "Conjugate", s);
     slice_gpu(in, slice_temp, rstarts, rstrides, s);
     concatenate_gpu({conj_temp, slice_temp}, temp, (int)axis, s);
   } else if (inverse) {
-    unary_op_gpu({in}, temp, "conj", s);
+    unary_op_gpu({in}, temp, "Conjugate", s);
   } else {
     temp.copy_shared_buffer(in);
   }
 
-  binary_op_gpu({temp, w_k_broadcast}, temp1, "mul", s);
+  binary_op_gpu({temp, w_k_broadcast}, temp1, "Multiply", s);
 
   std::vector<std::pair<int, int>> pads;
   auto padded_shape = out.shape();
@@ -419,7 +418,7 @@ void multi_upload_bluestein_fft(
       /*inplace=*/false,
       s);
 
-  binary_op_gpu_inplace({pad_temp1, w_q_broadcast}, pad_temp, "mul", s);
+  binary_op_gpu_inplace({pad_temp1, w_q_broadcast}, pad_temp, "Multiply", s);
 
   fft_op(
       pad_temp,
@@ -437,7 +436,7 @@ void multi_upload_bluestein_fft(
   starts[axis] = plan.bluestein_n - offset - n;
   slice_gpu(pad_temp1, temp, starts, strides, s);
 
-  binary_op_gpu_inplace({temp, w_k_broadcast}, temp1, "mul", s);
+  binary_op_gpu_inplace({temp, w_k_broadcast}, temp1, "Multiply", s);
 
   if (real && !inverse) {
     std::vector<int> rstarts(in.ndim(), 0);
@@ -451,11 +450,11 @@ void multi_upload_bluestein_fft(
     copies.push_back(inv_n);
 
     copy_gpu(temp1, temp_float, CopyType::General, s);
-    binary_op_gpu({temp_float, inv_n}, out, "mul", s);
+    binary_op_gpu({temp_float, inv_n}, out, "Multiply", s);
   } else if (inverse) {
     auto inv_n = array({1.0f / n}, {1}, complex64);
-    unary_op_gpu({temp1}, temp, "conj", s);
-    binary_op_gpu({temp, inv_n}, out, "mul", s);
+    unary_op_gpu({temp1}, temp, "Conjugate", s);
+    binary_op_gpu({temp, inv_n}, out, "Multiply", s);
     copies.push_back(inv_n);
   } else {
     out.copy_shared_buffer(temp1);
@@ -661,34 +660,45 @@ void fft_op(
     std::ostringstream kname;
     std::string inv_string = inverse ? "true" : "false";
     std::string real_string = real ? "true" : "false";
+    std::string func_name;
     if (plan.bluestein_n > 0) {
       kname << "bluestein_fft_mem_" << threadgroup_mem_size << "_"
             << in_type_str << "_" << out_type_str;
+      func_name = "bluestein_fft";
     } else if (plan.rader_n > 1) {
       kname << "rader_fft_mem_" << threadgroup_mem_size << "_" << in_type_str
             << "_" << out_type_str;
+      func_name = "rader_fft";
     } else if (four_step_params.required) {
       step = four_step_params.first_step ? 0 : 1;
       kname << "four_step_mem_" << threadgroup_mem_size << "_" << in_type_str
             << "_" << out_type_str << "_" << step << "_" << real_string;
+      func_name = "four_step_fft";
     } else {
       kname << "fft_mem_" << threadgroup_mem_size << "_" << in_type_str << "_"
             << out_type_str;
+      func_name = "fft";
     }
     std::string base_name = kname.str();
     // We use a specialized kernel for each FFT size
     kname << "_n" << fft_size << "_inv_" << inverse;
     std::string hash_name = kname.str();
-    auto kernel = get_fft_kernel(
-        d,
-        base_name,
-        hash_name,
-        threadgroup_mem_size,
-        in_type_str,
-        out_type_str,
-        step,
-        real,
-        func_consts);
+    auto template_def = func_name == "four_step_fft" ? get_template_definition(
+                                                           base_name,
+                                                           func_name,
+                                                           threadgroup_mem_size,
+                                                           in_type_str,
+                                                           out_type_str,
+                                                           step,
+                                                           real)
+                                                     : get_template_definition(
+                                                           base_name,
+                                                           func_name,
+                                                           threadgroup_mem_size,
+                                                           in_type_str,
+                                                           out_type_str);
+    auto kernel =
+        get_fft_kernel(d, base_name, hash_name, func_consts, template_def);
 
     compute_encoder->setComputePipelineState(kernel);
     compute_encoder.set_input_array(in_contiguous, 0);
@@ -774,10 +784,9 @@ void nd_fft_op(
     fft_op(in_arr, out_arr, axis, inverse, step_real, inplace, s);
   }
 
-  std::vector<array> copies = {temp1, temp2};
   auto& d = metal::device(s.device);
   d.get_command_buffer(s.index)->addCompletedHandler(
-      [copies](MTL::CommandBuffer*) mutable { copies.clear(); });
+      [temp_arrs](MTL::CommandBuffer*) mutable { temp_arrs.clear(); });
 }
 
 void FFT::eval_gpu(const std::vector<array>& inputs, array& out) {

mlx/backend/metal/hadamard.cpp

@@ -0,0 +1,203 @@
+// Copyright © 2024 Apple Inc.
+
+#include <map>
+
+#include "mlx/backend/common/compiled.h"
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/backend/common/utils.h"
+#include "mlx/backend/metal/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/jit/includes.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+constexpr int MAX_HADAMARD_THREADS_PER_GROUP = 256;
+constexpr int MAX_HADAMARD_BYTES = 32768; // 32KB
+
+std::string gen_hadamard_codelet(int m) {
+  // Generate a O(m^2) hadamard codelet for a given M
+  // using the hadamard matrices above
+  //
+  // e.g. m = 2
+  // METAL_FUNC void hadamard_m(thread float *x) {
+  //   float tmp[2];
+  //   tmp[0] = + x[0] + x[1];
+  //   tmp[1] = + x[0] - x[1];
+  //   for (int i = 0; i < 2; i++) { x[i] = tmp[i]; }
+  // }
+  //
+  auto h_matrices = hadamard_matrices();
+  auto& matrix = h_matrices[m];
+
+  std::ostringstream source;
+  source << "METAL_FUNC void hadamard_radix_m(thread float *x) {" << std::endl;
+  if (m == 1) {
+    source << "}" << std::endl;
+    return source.str();
+  }
+  source << "  float tmp[" << m << "];" << std::endl;
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+
+  int index = 0;
+  while (end != std::string_view::npos) {
+    source << "  tmp[" << index << "] = ";
+    auto row = matrix.substr(start, end - start);
+    for (int i = 0; i < row.length(); i++) {
+      source << " " << row[i] << " x[" << i << "]";
+    }
+    source << ";" << std::endl;
+    start = end + 1;
+    end = matrix.find('\n', start);
+    index++;
+  }
+  source << "  for (int i = 0; i < " << m << "; i++) { x[i] = tmp[i]; }"
+         << std::endl;
+  source << "}" << std::endl;
+  return source.str();
+}
+
+void launch_hadamard(
+    const array& in,
+    array& out,
+    int batch_size,
+    int threads_per,
+    const std::string kernel_name,
+    float scale,
+    const Stream& s) {
+  auto& d = metal::device(s.device);
+
+  const auto& lib_name = kernel_name.substr(1);
+  auto lib = d.get_library(lib_name);
+  auto kernel = d.get_kernel(kernel_name, lib);
+  assert(threads_per <= kernel->maxTotalThreadsPerThreadgroup());
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder->setComputePipelineState(kernel);
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  compute_encoder->setBytes(&scale, sizeof(float), 2);
+
+  MTL::Size group_dims = MTL::Size(1, threads_per, 1);
+  MTL::Size grid_dims = MTL::Size(batch_size, threads_per, 1);
+  compute_encoder->dispatchThreads(grid_dims, group_dims);
+}
+
+void Hadamard::eval_gpu(const std::vector<array>& inputs, array& out) {
+  auto& s = stream();
+
+  auto& in = inputs[0];
+
+  std::vector<array> copies;
+  // Only support the last axis for now
+  int axis = in.ndim() - 1;
+  auto check_input = [&copies, &s](const array& x) {
+    // TODO(alexbarron) pass strides to kernel to relax this constraint
+    bool no_copy = x.flags().row_contiguous;
+    if (no_copy) {
+      return x;
+    } else {
+      copies.push_back(array(x.shape(), x.dtype(), nullptr, {}));
+      copy_gpu(x, copies.back(), CopyType::General, s);
+      return copies.back();
+    }
+  };
+  const array& in_contiguous = check_input(in);
+
+  if (in_contiguous.is_donatable()) {
+    out.move_shared_buffer(in_contiguous);
+  } else {
+    out.set_data(allocator::malloc_or_wait(out.nbytes()));
+  }
+
+  auto [n, m] = decompose_hadamard(in.shape(axis));
+
+  if (n * (int)size_of(in.dtype()) > MAX_HADAMARD_BYTES) {
+    throw std::invalid_argument(
+        "[hadamard] For n = m*2^k, 2^k > 8192 for FP32 or 2^k > 16384 for FP16/BF16 NYI");
+  }
+
+  int max_radix = std::min(n, 16);
+  // Use read_width 2 for m = 28 to avoid register spilling
+  int read_width = (n == 2 || m == 28) ? 2 : 4;
+
+  std::ostringstream kname;
+  kname << "hadamard_" << n * m << "_" << type_to_name(out);
+  auto kernel_name = kname.str();
+  auto& d = metal::device(s.device);
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name);
+  if (lib == nullptr) {
+    std::ostringstream kernel_source;
+    auto codelet = gen_hadamard_codelet(m);
+    kernel_source << metal::utils() << codelet << metal::hadamard();
+    kernel_source << get_template_definition(
+        "n" + kernel_name,
+        "hadamard_n",
+        get_type_string(in.dtype()),
+        n,
+        max_radix,
+        read_width);
+    kernel_source << get_template_definition(
+        "m" + kernel_name,
+        "hadamard_m",
+        get_type_string(in.dtype()),
+        n,
+        m,
+        read_width);
+    lib = d.get_library(lib_name, kernel_source.str());
+  }
+
+  int batch_size = in.size() / n;
+  int threads_per = n / max_radix;
+
+  if (m > 1) {
+    // When m is greater than 1, we decompose the
+    // computation into two uploads to the GPU:
+    //
+    // e.g. len(x) = 12*4 = 48, m = 12, n = 4
+    //
+    // y = h48 @ x
+    //
+    // Upload 1:
+    // tmp = a.reshape(12, 4) @ h4
+    //
+    // Upload 2:
+    // y = h12 @ tmp
+    array temp(in.shape(), in.dtype(), nullptr, {});
+    temp.set_data(allocator::malloc_or_wait(temp.nbytes()));
+    copies.push_back(temp);
+
+    launch_hadamard(
+        in_contiguous,
+        temp,
+        batch_size,
+        threads_per,
+        "n" + kernel_name,
+        1.0,
+        s);
+
+    // Metal sometimes reports 256 max threads per group for hadamard_m kernel
+    threads_per = std::min(n / read_width, MAX_HADAMARD_THREADS_PER_GROUP);
+    batch_size = in.size() / m / read_width / threads_per;
+    launch_hadamard(
+        temp, out, batch_size, threads_per, "m" + kernel_name, scale_, s);
+  } else {
+    launch_hadamard(
+        in_contiguous,
+        out,
+        batch_size,
+        threads_per,
+        "n" + kernel_name,
+        scale_,
+        s);
+  }
+
+  d.get_command_buffer(s.index)->addCompletedHandler(
+      [copies](MTL::CommandBuffer*) mutable { copies.clear(); });
+}
+
+} // namespace mlx::core

mlx/backend/metal/indexing.cpp

@@ -293,7 +293,18 @@ void Scatter::eval_gpu(const std::vector<array>& inputs, array& out) {
         out.shape().data(), out.shape().size() * sizeof(int), 3);
     compute_encoder->setBytes(
         out.strides().data(), out.strides().size() * sizeof(size_t), 4);
-    compute_encoder->setBytes(&upd_size, sizeof(size_t), 5);
+
+    size_t out_ndim = out.ndim();
+    compute_encoder->setBytes(&out_ndim, sizeof(out_ndim), 5);
+    if (upd_ndim <= 1) {
+      // Placeholder so Metal doesn't compalain
+      int shape_ = 0;
+      compute_encoder->setBytes(&shape_, sizeof(int), 6);
+    } else {
+      compute_encoder->setBytes(upd.shape().data(), upd_ndim * sizeof(int), 6);
+    }
+    compute_encoder->setBytes(&upd_ndim, sizeof(size_t), 7);
+    compute_encoder->setBytes(&upd_size, sizeof(size_t), 8);
 
     // Set index buffers
     for (int i = 0; i < nidx; ++i) {

mlx/backend/metal/jit/binary.h

@@ -1,87 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view binary_kernels = R"(
-template [[host_name("ss{0}")]] [[kernel]]
-void binary_ss<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("vs{0}")]] [[kernel]]
-void binary_vs<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("sv{0}")]] [[kernel]]
-void binary_sv<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("vv{0}")]] [[kernel]]
-void binary_vv<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("g4{0}")]] [[kernel]] void
-binary_g_nd<{1}, {2}, {3}, 4>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const int shape[4],
-    constant const size_t a_strides[4],
-    constant const size_t b_strides[4],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-template [[host_name("g5{0}")]] [[kernel]] void
-binary_g_nd<{1}, {2}, {3}, 5>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const int shape[5],
-    constant const size_t a_strides[5],
-    constant const size_t b_strides[5],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-
-template [[host_name("g1{0}")]] [[kernel]] void
-binary_g_nd1<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const size_t& a_stride,
-    constant const size_t& b_stride,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("g2{0}")]] [[kernel]] void
-binary_g_nd2<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const size_t a_strides[2],
-    constant const size_t b_strides[2],
-    uint2 index [[thread_position_in_grid]],
-    uint2 grid_dim [[threads_per_grid]]);
-template [[host_name("g3{0}")]] [[kernel]] void
-binary_g_nd3<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const size_t a_strides[3],
-    constant const size_t b_strides[3],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-
-template [[host_name("gn{0}")]] [[kernel]]
-void binary_g<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    constant const int* shape,
-    constant const size_t* a_strides,
-    constant const size_t* b_strides,
-    constant const int& ndim,
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-)";

mlx/backend/metal/jit/binary_two.h

@@ -1,98 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view binary_two_kernels = R"(
-template [[host_name("ss{0}")]] [[kernel]]
-void binary_ss<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("vs{0}")]] [[kernel]]
-void binary_vs<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("sv{0}")]] [[kernel]]
-void binary_sv<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("vv{0}")]] [[kernel]]
-void binary_vv<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    uint index [[thread_position_in_grid]]);
-
-template [[host_name("g4{0}")]] [[kernel]] void
-binary_g_nd<{1}, {2}, {3}, 4>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const int shape[4],
-    constant const size_t a_strides[4],
-    constant const size_t b_strides[4],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-template [[host_name("g5{0}")]] [[kernel]] void
-binary_g_nd<{1}, {2}, {3}, 5>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const int shape[5],
-    constant const size_t a_strides[5],
-    constant const size_t b_strides[5],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-
-template [[host_name("g1{0}")]] [[kernel]] void
-binary_g_nd1<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const size_t& a_stride,
-    constant const size_t& b_stride,
-    uint index [[thread_position_in_grid]]);
-template [[host_name("g2{0}")]] [[kernel]] void
-binary_g_nd2<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const size_t a_strides[2],
-    constant const size_t b_strides[2],
-    uint2 index [[thread_position_in_grid]],
-    uint2 grid_dim [[threads_per_grid]]);
-template [[host_name("g3{0}")]] [[kernel]] void
-binary_g_nd3<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const size_t a_strides[3],
-    constant const size_t b_strides[3],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-
-template [[host_name("gn{0}")]] [[kernel]]
-void binary_g<{1}, {2}, {3}>(
-    device const {1}* a,
-    device const {1}* b,
-    device {2}* c,
-    device {2}* d,
-    constant const int* shape,
-    constant const size_t* a_strides,
-    constant const size_t* b_strides,
-    constant const int& ndim,
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-)";

mlx/backend/metal/jit/fft.h

@@ -1,53 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view fft_kernel = R"(
-template [[host_name("{name}")]] [[kernel]] void
-fft<{tg_mem_size}, {in_T}, {out_T}>(
-    const device {in_T}* in [[buffer(0)]],
-    device {out_T}* out [[buffer(1)]],
-    constant const int& n,
-    constant const int& batch_size,
-    uint3 elem [[thread_position_in_grid]],
-    uint3 grid [[threads_per_grid]]);
-)";
-
-constexpr std::string_view rader_fft_kernel = R"(
-template [[host_name("{name}")]] [[kernel]] void
-rader_fft<{tg_mem_size}, {in_T}, {out_T}>(
-    const device {in_T}* in [[buffer(0)]],
-    device {out_T}* out [[buffer(1)]],
-    const device float2* raders_b_q [[buffer(2)]],
-    const device short* raders_g_q [[buffer(3)]],
-    const device short* raders_g_minus_q [[buffer(4)]],
-    constant const int& n,
-    constant const int& batch_size,
-    constant const int& rader_n,
-    uint3 elem [[thread_position_in_grid]],
-    uint3 grid [[threads_per_grid]]);
-)";
-
-constexpr std::string_view bluestein_fft_kernel = R"(
-template [[host_name("{name}")]] [[kernel]] void
-bluestein_fft<{tg_mem_size}, {in_T}, {out_T}>(
-    const device {in_T}* in [[buffer(0)]],
-    device {out_T}* out [[buffer(1)]],
-    const device float2* w_q [[buffer(2)]],
-    const device float2* w_k [[buffer(3)]],
-    constant const int& length,
-    constant const int& n,
-    constant const int& batch_size,
-    uint3 elem [[thread_position_in_grid]],
-    uint3 grid [[threads_per_grid]]);
-)";
-
-constexpr std::string_view four_step_fft_kernel = R"(
-template [[host_name("{name}")]] [[kernel]] void
-four_step_fft<{tg_mem_size}, {in_T}, {out_T}, {step}, {real}>(
-      const device {in_T}* in [[buffer(0)]],
-      device {out_T}* out [[buffer(1)]],
-      constant const int& n1,
-      constant const int& n2,
-      constant const int& batch_size,
-      uint3 elem [[thread_position_in_grid]],
-      uint3 grid [[threads_per_grid]]);
-)";

mlx/backend/metal/jit/gemv_masked.h

@@ -0,0 +1,25 @@
+// Copyright © 2024 Apple Inc.
+
+constexpr std::string_view gemv_masked_kernel = R"(
+template [[host_name("{name}")]] [[kernel]] void
+gemv_{trans}masked<{itype}, {outm_t}, {opm_t}, {bm}, {bn}, {sm}, {sn}, {tm}, {tn}, {nc}>(
+    const device {itype}* mat [[buffer(0)]],
+    const device {itype}* in_vec [[buffer(1)]],
+    device {itype}* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant size_t* vector_batch_stride [[buffer(11)]],
+    const constant size_t* matrix_batch_stride [[buffer(12)]],
+    const device {outm_t}* out_mask [[buffer(20)]],
+    const device {opm_t}* mat_mask [[buffer(21)]],
+    const device {opm_t}* vec_mask [[buffer(22)]],
+    const constant int* mask_strides [[buffer(23)]],
+    const constant size_t* mask_batch_strides [[buffer(24)]],
+    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]]);
+)";

mlx/backend/metal/jit/includes.h

@@ -18,6 +18,8 @@ const char* binary();
 const char* binary_two();
 const char* copy();
 const char* fft();
+const char* hadamard();
+const char* quantized();
 const char* ternary();
 const char* scan();
 const char* softmax();
@@ -31,5 +33,6 @@ const char* steel_gemm_splitk();
 const char* conv();
 const char* steel_conv();
 const char* steel_conv_general();
+const char* gemv_masked();
 
 } // namespace mlx::core::metal

mlx/backend/metal/jit/indexing.h

@@ -38,12 +38,24 @@ constexpr std::string_view scatter_kernels = R"(
     device mlx_atomic<{1}>* 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 size_t& out_ndim [[buffer(5)]],
+    const constant int* upd_shape [[buffer(6)]],
+    const constant size_t& upd_ndim [[buffer(7)]],
+    const constant size_t& upd_size [[buffer(8)]],
     {5}
     uint2 gid [[thread_position_in_grid]]) {{
   const array<const device {2}*, {4}> idx_buffers = {{ {6} }};
   return scatter_1d_index_impl<{1}, {2}, {3}, {4}>(
-      updates, out, out_shape, out_strides, upd_size, idx_buffers, gid);
+      updates,
+      out,
+      out_shape,
+      out_strides,
+      out_ndim,
+      upd_shape,
+      upd_ndim,
+      upd_size,
+      idx_buffers,
+      gid);
 }}
 
 [[kernel]] void scatter{0}_{4}(

mlx/backend/metal/jit/sort.h

@@ -1,81 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view block_sort_kernels = R"(
-template [[host_name("carg_{0}")]] [[kernel]] void
-block_sort<{1}, {2}, true, {3}, {4}>(
-    const device {1}* inp [[buffer(0)]],
-    device {2}* out [[buffer(1)]],
-    const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& stride_segment_axis [[buffer(4)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-template [[host_name("ncarg_{0}")]] [[kernel]] void
-block_sort_nc<{1}, {2}, true, {3}, {4}>(
-    const device {1}* inp [[buffer(0)]],
-    device {2}* out [[buffer(1)]],
-    const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& nc_dim [[buffer(4)]],
-    const device int* nc_shape [[buffer(5)]],
-    const device size_t* nc_strides [[buffer(6)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-template [[host_name("c_{0}")]] [[kernel]] void
-block_sort<{1}, {2}, false, {3}, {4}>(
-    const device {1}* inp [[buffer(0)]],
-    device {2}* out [[buffer(1)]],
-    const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& stride_segment_axis [[buffer(4)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-template [[host_name("nc_{0}")]] [[kernel]] void
-block_sort_nc<{1}, {2}, false, {3}, {4}>(
-    const device {1}* inp [[buffer(0)]],
-    device {2}* out [[buffer(1)]],
-    const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& nc_dim [[buffer(4)]],
-    const device int* nc_shape [[buffer(5)]],
-    const device size_t* nc_strides [[buffer(6)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-)";
-
-constexpr std::string_view multiblock_sort_kernels = R"(
-template [[host_name("sort_{0}")]] [[kernel]] void
-mb_block_sort<{1}, {2}, true, {3}, {4}>(
-    const device {1}* inp [[buffer(0)]],
-    device {1}* out_vals [[buffer(1)]],
-    device {2}* out_idxs [[buffer(2)]],
-    const constant int& size_sorted_axis [[buffer(3)]],
-    const constant int& stride_sorted_axis [[buffer(4)]],
-    const constant int& nc_dim [[buffer(5)]],
-    const device int* nc_shape [[buffer(6)]],
-    const device size_t* nc_strides [[buffer(7)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-template [[host_name("partition_{0}")]] [[kernel]] void
-mb_block_partition<{1}, {2}, true, {3}, {4}>(
-    device {2}* block_partitions [[buffer(0)]],
-    const device {1}* dev_vals [[buffer(1)]],
-    const device {2}* dev_idxs [[buffer(2)]],
-    const constant int& size_sorted_axis [[buffer(3)]],
-    const constant int& merge_tiles [[buffer(4)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]],
-    uint3 tgp_dims [[threads_per_threadgroup]]);
-template [[host_name("merge_{0}")]] [[kernel]] void
-mb_block_merge<{1}, {2}, true, {3}, {4}>(
-    const device {2}* block_partitions [[buffer(0)]],
-    const device {1}* dev_vals_in [[buffer(1)]],
-    const device {2}* dev_idxs_in [[buffer(2)]],
-    device {1}* dev_vals_out [[buffer(3)]],
-    device {2}* dev_idxs_out [[buffer(4)]],
-    const constant int& size_sorted_axis [[buffer(5)]],
-    const constant int& merge_tiles [[buffer(6)]],
-    const constant int& num_tiles [[buffer(7)]],
-    uint3 tid [[threadgroup_position_in_grid]],
-    uint3 lid [[thread_position_in_threadgroup]]);
-)";

mlx/backend/metal/jit/ternary.h

@@ -1,80 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view ternary_kernels = R"(
-template [[host_name("v_{0}")]] [[kernel]] void ternary_v<{1}, {2}>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* d,
-    uint index [[thread_position_in_grid]]);
-
-template [[host_name("g_{0}")]] [[kernel]] void ternary_g<{1}, {2}>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* 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]]);
-
-template [[host_name("g1_{0}")]] [[kernel]] void
-ternary_g_nd1<{1}, {2}>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* 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("g2_{0}")]] [[kernel]] void
-ternary_g_nd2<{1}, {2}>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* 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("g3_{0}")]] [[kernel]] void
-ternary_g_nd3<{1}, {2}>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* 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]]);
-template [[host_name("g4_{0}")]] [[kernel]] void
-ternary_g_nd<{1}, {2}, 4>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* d,
-    constant const int shape[4],
-    constant const size_t a_strides[4],
-    constant const size_t b_strides[4],
-    constant const size_t c_strides[4],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-template [[host_name("g5_{0}")]] [[kernel]] void
-ternary_g_nd<{1}, {2}, 5>(
-    device const bool* a,
-    device const {1}* b,
-    device const {1}* c,
-    device {1}* d,
-    constant const int shape[5],
-    constant const size_t a_strides[5],
-    constant const size_t b_strides[5],
-    constant const size_t c_strides[5],
-    uint3 index [[thread_position_in_grid]],
-    uint3 grid_dim [[threads_per_grid]]);
-)";

mlx/backend/metal/jit/unary.h

@@ -1,16 +0,0 @@
-// Copyright © 2024 Apple Inc.
-
-constexpr std::string_view unary_kernels = R"(
-template [[host_name("v{0}")]] [[kernel]] void unary_v<{1}, {2}>(
-    device const {1}* in,
-    device {1}* out,
-    uint index [[thread_position_in_grid]]);
-
-template [[host_name("g{0}")]] [[kernel]] void unary_g<{1}, {2}>(
-    device const {1}* in,
-    device {1}* out,
-    device const int* in_shape,
-    device const size_t* in_strides,
-    device const int& ndim,
-    uint index [[thread_position_in_grid]]);
-)";

mlx/backend/metal/jit_kernels.cpp

@@ -1,21 +1,16 @@
 // Copyright © 2024 Apple Inc.
-#include <fmt/format.h>
+#include <map>
 
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/metal/jit/arange.h"
-#include "mlx/backend/metal/jit/binary.h"
-#include "mlx/backend/metal/jit/binary_two.h"
 #include "mlx/backend/metal/jit/copy.h"
-#include "mlx/backend/metal/jit/fft.h"
+#include "mlx/backend/metal/jit/gemv_masked.h"
 #include "mlx/backend/metal/jit/includes.h"
 #include "mlx/backend/metal/jit/reduce.h"
 #include "mlx/backend/metal/jit/scan.h"
 #include "mlx/backend/metal/jit/softmax.h"
-#include "mlx/backend/metal/jit/sort.h"
 #include "mlx/backend/metal/jit/steel_conv.h"
 #include "mlx/backend/metal/jit/steel_gemm.h"
-#include "mlx/backend/metal/jit/ternary.h"
-#include "mlx/backend/metal/jit/unary.h"
 #include "mlx/backend/metal/kernels.h"
 #include "mlx/backend/metal/utils.h"
 
@@ -48,38 +43,81 @@ MTL::ComputePipelineState* get_arange_kernel(
 MTL::ComputePipelineState* get_unary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& out) {
+    Dtype out_type,
+    const std::string op) {
   std::string lib_name = kernel_name.substr(1);
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
+    auto u_def = get_template_definition(
+        "v" + lib_name, "unary_v", get_type_string(out_type), op);
+    auto u2_def = get_template_definition(
+        "v2" + lib_name, "unary_v2", get_type_string(out_type), op);
+    auto g_def = get_template_definition(
+        "g" + lib_name, "unary_g", get_type_string(out_type), op);
     kernel_source << metal::utils() << metal::unary_ops() << metal::unary()
-                  << fmt::format(
-                         unary_kernels,
-                         lib_name,
-                         get_type_string(out.dtype()),
-                         op_name(out));
+                  << u_def << u2_def << g_def;
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
 }
 
+void add_binary_kernels(
+    const std::string lib_name,
+    Dtype in_type,
+    Dtype out_type,
+    const std::string op,
+    std::ostringstream& kernel_source) {
+  const std::map<std::string, std::string> kernel_types = {
+      {"ss", "binary_ss"},
+      {"vs", "binary_vs"},
+      {"sv", "binary_sv"},
+      {"vv", "binary_vv"},
+      {"vs2", "binary_vs2"},
+      {"sv2", "binary_sv2"},
+      {"vv2", "binary_vv2"},
+      {"g1", "binary_g_nd1"},
+      {"g2", "binary_g_nd2"},
+      {"g3", "binary_g_nd3"},
+      {"g4", "binary_g_nd"},
+      {"g5", "binary_g_nd"},
+      {"gn", "binary_g"},
+  };
+  for (auto [name, func] : kernel_types) {
+    std::string template_def;
+    if (name == "g4" || name == "g5") {
+      int dim = std::stoi(name.substr(1));
+      template_def = get_template_definition(
+          name + lib_name,
+          func,
+          get_type_string(in_type),
+          get_type_string(out_type),
+          op,
+          dim);
+    } else {
+      template_def = get_template_definition(
+          name + lib_name,
+          func,
+          get_type_string(in_type),
+          get_type_string(out_type),
+          op);
+    }
+    kernel_source << template_def;
+  }
+}
+
 MTL::ComputePipelineState* get_binary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& in,
-    const array& out) {
+    Dtype in_type,
+    Dtype out_type,
+    const std::string op) {
   std::string lib_name = kernel_name.substr(2);
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
-    kernel_source << metal::utils() << metal::binary_ops() << metal::binary()
-                  << fmt::format(
-                         binary_kernels,
-                         lib_name,
-                         get_type_string(in.dtype()),
-                         get_type_string(out.dtype()),
-                         op_name(out));
+    kernel_source << metal::utils() << metal::binary_ops() << metal::binary();
+    add_binary_kernels(lib_name, in_type, out_type, op, kernel_source);
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
@@ -88,20 +126,16 @@ MTL::ComputePipelineState* get_binary_kernel(
 MTL::ComputePipelineState* get_binary_two_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& in,
-    const array& out) {
+    Dtype in_type,
+    Dtype out_type,
+    const std::string op) {
   std::string lib_name = kernel_name.substr(2);
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
     kernel_source << metal::utils() << metal::binary_ops()
-                  << metal::binary_two()
-                  << fmt::format(
-                         binary_two_kernels,
-                         lib_name,
-                         get_type_string(in.dtype()),
-                         get_type_string(out.dtype()),
-                         op_name(out));
+                  << metal::binary_two();
+    add_binary_kernels(lib_name, in_type, out_type, op, kernel_source);
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
@@ -110,17 +144,35 @@ MTL::ComputePipelineState* get_binary_two_kernel(
 MTL::ComputePipelineState* get_ternary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& out) {
+    Dtype type,
+    const std::string op) {
   std::string lib_name = kernel_name.substr(kernel_name.find("_") + 1);
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
-    kernel_source << metal::utils() << metal::ternary_ops() << metal::ternary()
-                  << fmt::format(
-                         ternary_kernels,
-                         lib_name,
-                         get_type_string(out.dtype()),
-                         op_name(out));
+    const std::map<std::string, std::string> kernel_types = {
+        {"v", "ternary_v"},
+        {"v2", "ternary_v2"},
+        {"g", "ternary_g"},
+        {"g1", "ternary_g_nd1"},
+        {"g2", "ternary_g_nd2"},
+        {"g3", "ternary_g_nd3"},
+        {"g4", "ternary_g_nd"},
+        {"g5", "ternary_g_nd"},
+    };
+    kernel_source << metal::utils() << metal::ternary_ops() << metal::ternary();
+    for (auto [name, func] : kernel_types) {
+      std::string template_def;
+      if (name == "g4" || name == "g5") {
+        int dim = std::stoi(name.substr(1));
+        template_def = get_template_definition(
+            name + "_" + lib_name, func, get_type_string(type), op, dim);
+      } else {
+        template_def = get_template_definition(
+            name + "_" + lib_name, func, get_type_string(type), op);
+      }
+      kernel_source << template_def;
+    }
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
@@ -205,14 +257,29 @@ MTL::ComputePipelineState* get_sort_kernel(
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
-    kernel_source << metal::utils() << metal::sort()
-                  << fmt::format(
-                         block_sort_kernels,
-                         lib_name,
-                         get_type_string(in.dtype()),
-                         get_type_string(out.dtype()),
-                         bn,
-                         tn);
+    auto in_type = get_type_string(in.dtype());
+    auto out_type = get_type_string(out.dtype());
+    kernel_source << metal::utils() << metal::sort();
+    for (bool is_argsort : {true, false}) {
+      std::string bool_string = is_argsort ? "true" : "false";
+      std::string func_string = is_argsort ? "carg_" : "c_";
+      kernel_source << get_template_definition(
+          func_string + lib_name,
+          "block_sort",
+          in_type,
+          out_type,
+          bool_string,
+          bn,
+          tn);
+      kernel_source << get_template_definition(
+          "n" + func_string + lib_name,
+          "block_sort_nc",
+          in_type,
+          out_type,
+          bool_string,
+          bn,
+          tn);
+    }
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
@@ -229,14 +296,21 @@ MTL::ComputePipelineState* get_mb_sort_kernel(
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
-    kernel_source << metal::utils() << metal::sort()
-                  << fmt::format(
-                         multiblock_sort_kernels,
-                         lib_name,
-                         get_type_string(in.dtype()),
-                         get_type_string(idx.dtype()),
-                         bn,
-                         tn);
+    kernel_source << metal::utils() << metal::sort();
+    std::vector<std::pair<std::string, std::string>> kernel_types = {
+        {"sort_", "mb_block_sort"},
+        {"partition_", "mb_block_partition"},
+        {"merge_", "mb_block_merge"}};
+    for (auto [name, func] : kernel_types) {
+      kernel_source << get_template_definition(
+          name + lib_name,
+          func,
+          get_type_string(in.dtype()),
+          get_type_string(idx.dtype()),
+          "true",
+          bn,
+          tn);
+    }
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib);
@@ -429,6 +503,49 @@ MTL::ComputePipelineState* get_steel_gemm_masked_kernel(
   return d.get_kernel(kernel_name, lib);
 }
 
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_mat,
+    int bm,
+    int bn,
+    int sm,
+    int sn,
+    int tm,
+    int tn,
+    bool contiguous) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name);
+  if (lib == nullptr) {
+    std::ostringstream kernel_source;
+    auto out_mask_type = mask_out.has_value()
+        ? get_type_string((*mask_out).dtype())
+        : "nomask_t";
+    auto op_mask_type =
+        mask_op.has_value() ? get_type_string((*mask_op).dtype()) : "nomask_t";
+    kernel_source << metal::utils() << metal::gemv_masked()
+                  << fmt::format(
+                         gemv_masked_kernel,
+                         "name"_a = lib_name,
+                         "itype"_a = get_type_string(out.dtype()),
+                         "outm_t"_a = out_mask_type,
+                         "opm_t"_a = op_mask_type,
+                         "bm"_a = bm,
+                         "bn"_a = bn,
+                         "sm"_a = sm,
+                         "sn"_a = sn,
+                         "tm"_a = tm,
+                         "tn"_a = tn,
+                         "trans"_a = transpose_mat ? "t_" : "",
+                         "nc"_a = contiguous ? "0" : "1");
+    lib = d.get_library(lib_name, kernel_source.str());
+  }
+  return d.get_kernel(kernel_name, lib);
+}
+
 MTL::ComputePipelineState* get_steel_conv_kernel(
     metal::Device& d,
     const std::string& kernel_name,
@@ -494,47 +611,32 @@ MTL::ComputePipelineState* get_fft_kernel(
     metal::Device& d,
     const std::string& kernel_name,
     const std::string& hash_name,
-    const int tg_mem_size,
-    const std::string& in_type,
-    const std::string& out_type,
-    int step,
-    bool real,
-    const metal::MTLFCList& func_consts) {
+    const metal::MTLFCList& func_consts,
+    const std::string& template_def) {
   const auto& lib_name = kernel_name;
   auto lib = d.get_library(lib_name);
   if (lib == nullptr) {
     std::ostringstream kernel_source;
     std::string kernel_string;
-    if (lib_name.find("bluestein") != std::string::npos) {
-      kernel_string = bluestein_fft_kernel;
-    } else if (lib_name.find("rader") != std::string::npos) {
-      kernel_string = rader_fft_kernel;
-    } else if (lib_name.find("four_step") != std::string::npos) {
-      kernel_string = four_step_fft_kernel;
-    } else {
-      kernel_string = fft_kernel;
-    }
-    kernel_source << metal::fft();
-    if (lib_name.find("four_step") != std::string::npos) {
-      kernel_source << fmt::format(
-          kernel_string,
-          "name"_a = lib_name,
-          "tg_mem_size"_a = tg_mem_size,
-          "in_T"_a = in_type,
-          "out_T"_a = out_type,
-          "step"_a = step,
-          "real"_a = real);
-    } else {
-      kernel_source << fmt::format(
-          kernel_string,
-          "name"_a = lib_name,
-          "tg_mem_size"_a = tg_mem_size,
-          "in_T"_a = in_type,
-          "out_T"_a = out_type);
-    }
+    kernel_source << metal::fft() << template_def;
     lib = d.get_library(lib_name, kernel_source.str());
   }
   return d.get_kernel(kernel_name, lib, hash_name, func_consts);
 }
 
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name);
+  if (lib == nullptr) {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm() << metal::quantized()
+                  << template_def;
+    lib = d.get_library(lib_name, kernel_source.str());
+  }
+  return d.get_kernel(kernel_name, lib);
+}
+
 } // namespace mlx::core

mlx/backend/metal/kernels.h

@@ -1,5 +1,7 @@
 // Copyright © 2024 Apple Inc.
 
+#include <fmt/format.h>
+
 #include "mlx/array.h"
 #include "mlx/backend/metal/device.h"
 
@@ -13,24 +15,28 @@ MTL::ComputePipelineState* get_arange_kernel(
 MTL::ComputePipelineState* get_unary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& out);
+    Dtype out_type,
+    const std::string op);
 
 MTL::ComputePipelineState* get_binary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& in,
-    const array& out);
+    Dtype in_type,
+    Dtype out_type,
+    const std::string op);
 
 MTL::ComputePipelineState* get_binary_two_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& in,
-    const array& out);
+    Dtype in_type,
+    Dtype out_type,
+    const std::string op);
 
 MTL::ComputePipelineState* get_ternary_kernel(
     metal::Device& d,
     const std::string& kernel_name,
-    const array& out);
+    Dtype type,
+    const std::string op);
 
 MTL::ComputePipelineState* get_copy_kernel(
     metal::Device& d,
@@ -145,6 +151,21 @@ MTL::ComputePipelineState* get_steel_conv_kernel(
     int n_channel_specialization,
     bool small_filter);
 
+MTL::ComputePipelineState* get_gemv_masked_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const array& out,
+    const std::optional<array>& mask_out,
+    const std::optional<array>& mask_op,
+    bool transpose_mat,
+    int bm,
+    int bn,
+    int sm,
+    int sn,
+    int tm,
+    int tn,
+    bool contiguous);
+
 MTL::ComputePipelineState* get_steel_conv_general_kernel(
     metal::Device& d,
     const std::string& kernel_name,
@@ -159,11 +180,34 @@ MTL::ComputePipelineState* get_fft_kernel(
     metal::Device& d,
     const std::string& kernel_name,
     const std::string& hash_name,
-    const int tg_mem_size,
-    const std::string& in_type,
-    const std::string& out_type,
-    int step,
-    bool real,
-    const metal::MTLFCList& func_consts);
+    const metal::MTLFCList& func_consts,
+    const std::string& template_def);
+
+MTL::ComputePipelineState* get_quantized_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& template_def);
+
+// Create a GPU kernel template definition for JIT compilation
+template <typename... Args>
+std::string
+get_template_definition(std::string name, std::string func, Args... args) {
+  std::ostringstream s;
+  s << func << "<";
+  bool first = true;
+  auto add_arg = [&s, &first](const auto& arg) {
+    if (!first) {
+      s << ", ";
+    }
+    first = false;
+    s << arg;
+  };
+  (add_arg(args), ...);
+  s << ">";
+  std::string base_string = R"(
+template [[host_name("{0}")]] [[kernel]] decltype({1}) {1};
+  )";
+  return fmt::format(base_string, name, s.str());
+}
 
 } // namespace mlx::core

mlx/backend/metal/kernels/CMakeLists.txt

@@ -1,69 +1,15 @@
 set(
-  HEADERS
+  BASE_HEADERS
   bf16.h
   bf16_math.h
   complex.h
   defines.h
+  expm1f.h
   utils.h
-  steel/conv/params.h
 )
 
-set(
-  KERNELS
-  "arg_reduce"
-  "conv"
-  "fft"
-  "gemv"
-  "quantized"
-  "random"
-  "rms_norm"
-  "layer_norm"
-  "rope"
-  "scaled_dot_product_attention"
-)
-
-if (NOT MLX_METAL_JIT)
-set(
-  KERNELS
-  ${KERNELS}
-  "arange"
-  "binary"
-  "binary_two"
-  "unary"
-  "ternary"
-  "copy"
-  "softmax"
-  "sort"
-  "scan"
-  "reduce"
-)
-set(
-  HEADERS 
-  ${HEADERS}
-  atomic.h
-  arange.h
-  unary_ops.h
-  unary.h
-  binary_ops.h
-  binary.h
-  ternary.h
-  copy.h
-  fft.h
-  fft/radix.h
-  fft/readwrite.h
-  softmax.h
-  sort.h
-  scan.h
-  reduction/ops.h
-  reduction/reduce_init.h
-  reduction/reduce_all.h
-  reduction/reduce_col.h
-  reduction/reduce_row.h
-)
-endif()
-
 function(build_kernel_base TARGET SRCFILE DEPS)
-  set(METAL_FLAGS -Wall -Wextra -fno-fast-math -D${MLX_METAL_VERSION})
+  set(METAL_FLAGS -Wall -Wextra -fno-fast-math)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS}
         -gline-tables-only
@@ -75,7 +21,7 @@ function(build_kernel_base TARGET SRCFILE DEPS)
                   -c ${SRCFILE}
                   -I${PROJECT_SOURCE_DIR}
                   -o ${TARGET}.air
-    DEPENDS ${SRCFILE} ${DEPS}
+    DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
     OUTPUT ${TARGET}.air
     COMMENT "Building ${TARGET}.air"
     VERBATIM
@@ -84,49 +30,100 @@ endfunction(build_kernel_base)
 
 function(build_kernel KERNEL)
   set(SRCFILE ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL}.metal)
-  build_kernel_base(${KERNEL} ${SRCFILE} "${HEADERS}")
+  cmake_path(GET KERNEL STEM TARGET)
+  build_kernel_base(${TARGET} ${SRCFILE} "${ARGN}")
+  set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR} PARENT_SCOPE)
 endfunction(build_kernel)
 
-foreach(KERNEL ${KERNELS})
-  build_kernel(${KERNEL})
-  set(KERNEL_AIR ${KERNEL}.air ${KERNEL_AIR})
-endforeach()
+build_kernel(arg_reduce)
+build_kernel(conv steel/conv/params.h)
+build_kernel(gemv steel/utils.h)
+build_kernel(layer_norm)
+build_kernel(random)
+build_kernel(rms_norm)
+build_kernel(rope)
+build_kernel(
+  scaled_dot_product_attention
+  scaled_dot_product_attention_params.h
+  steel/defines.h
+  steel/gemm/transforms.h
+  steel/utils.h
+)
+
+set(
+  STEEL_HEADERS 
+  steel/defines.h
+  steel/utils.h
+  steel/conv/conv.h
+  steel/conv/loader.h
+  steel/conv/loaders/loader_channel_l.h
+  steel/conv/loaders/loader_channel_n.h
+  steel/conv/loaders/loader_general.h
+  steel/conv/kernels/steel_conv.h
+  steel/conv/kernels/steel_conv_general.h
+  steel/gemm/gemm.h
+  steel/gemm/mma.h
+  steel/gemm/loader.h
+  steel/gemm/transforms.h
+  steel/gemm/kernels/steel_gemm_fused.h
+  steel/gemm/kernels/steel_gemm_masked.h
+  steel/gemm/kernels/steel_gemm_splitk.h
+)
 
 if (NOT MLX_METAL_JIT)
-  set(
-    STEEL_KERNELS
-    ${CMAKE_CURRENT_SOURCE_DIR}/steel/conv/kernels/steel_conv.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/steel/conv/kernels/steel_conv_general.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/steel/gemm/kernels/steel_gemm_fused.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/steel/gemm/kernels/steel_gemm_masked.metal
-    ${CMAKE_CURRENT_SOURCE_DIR}/steel/gemm/kernels/steel_gemm_splitk.metal
-  )
-  set(
-    STEEL_HEADERS 
-    steel/defines.h
-    steel/utils.h
-    steel/conv/conv.h
-    steel/conv/loader.h
-    steel/conv/loaders/loader_channel_l.h
-    steel/conv/loaders/loader_channel_n.h
-    steel/conv/loaders/loader_general.h
-    steel/conv/kernels/steel_conv.h
-    steel/conv/kernels/steel_conv_general.h
-    steel/gemm/gemm.h
-    steel/gemm/mma.h
-    steel/gemm/loader.h
-    steel/gemm/transforms.h
-    steel/gemm/kernels/steel_gemm_fused.h
-    steel/gemm/kernels/steel_gemm_masked.h
-    steel/gemm/kernels/steel_gemm_splitk.h
-  )
-  foreach(KERNEL ${STEEL_KERNELS})
-    cmake_path(GET KERNEL STEM TARGET)
-    build_kernel_base(${TARGET} ${KERNEL} "${STEEL_HEADERS}")
-    set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR})
-  endforeach()
+build_kernel(arange arange.h)
+build_kernel(binary binary.h binary_ops.h)
+build_kernel(binary_two binary_two.h)
+build_kernel(copy copy.h)
+build_kernel(
+  fft
+  fft.h
+  fft/radix.h
+  fft/readwrite.h
+)
+build_kernel(
+  reduce
+  atomic.h
+  reduction/ops.h
+  reduction/reduce_init.h
+  reduction/reduce_all.h
+  reduction/reduce_col.h
+  reduction/reduce_row.h
+)
+build_kernel(
+  quantized
+  quantized.h
+  ${STEEL_HEADERS}
+)
+build_kernel(scan scan.h)
+build_kernel(softmax softmax.h)
+build_kernel(sort sort.h)
+build_kernel(ternary ternary.h ternary_ops.h)
+build_kernel(unary unary.h unary_ops.h)
+build_kernel(
+  steel/conv/kernels/steel_conv
+  ${STEEL_HEADERS}
+)
+build_kernel(
+  steel/conv/kernels/steel_conv_general
+  ${STEEL_HEADERS}
+)
+build_kernel(
+  steel/gemm/kernels/steel_gemm_fused
+  ${STEEL_HEADERS}
+)
+build_kernel(
+  steel/gemm/kernels/steel_gemm_masked
+  ${STEEL_HEADERS}
+)
+build_kernel(
+  steel/gemm/kernels/steel_gemm_splitk
+  ${STEEL_HEADERS}
+)
+build_kernel(gemv_masked steel/utils.h)
 endif()
 
+
 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/bf16.h

@@ -6,7 +6,7 @@
 
 using namespace metal;
 
-#if defined METAL_3_1 || (__METAL_VERSION__ >= 310)
+#if (MLX_METAL_VERSION >= 310) || (__METAL_VERSION__ >= 310)
 
 typedef bfloat bfloat16_t;
 

mlx/backend/metal/kernels/bf16_math.h

@@ -369,7 +369,7 @@ instantiate_metal_math_funcs(
     return static_cast<otype>(__metal_simd_xor(static_cast<ctype>(data)));     \
   }
 
-#if defined METAL_3_1 || (__METAL_VERSION__ >= 310)
+#if (MLX_METAL_VERSION >= 310) || (__METAL_VERSION__ >= 310)
 
 #define bfloat16_to_uint16(x) as_type<uint16_t>(x)
 #define uint16_to_bfloat16(x) as_type<bfloat16_t>(x)

mlx/backend/metal/kernels/binary.h

@@ -36,6 +36,39 @@ template <typename T, typename U, typename Op>
   c[index] = Op()(a[index], b[index]);
 }
 
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_sv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  c[offset] = Op()(a[0], b[offset]);
+}
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_vs2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  c[offset] = Op()(a[offset], b[0]);
+}
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_vv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  c[offset] = Op()(a[offset], b[offset]);
+}
+
 template <typename T, typename U, typename Op>
 [[kernel]] void binary_g_nd1(
     device const T* a,

mlx/backend/metal/kernels/binary.metal

@@ -4,148 +4,94 @@
 #include <metal_math>
 
 // clang-format off
+#include "mlx/backend/metal/kernels/defines.h"
 #include "mlx/backend/metal/kernels/utils.h"
 #include "mlx/backend/metal/kernels/binary_ops.h"
 #include "mlx/backend/metal/kernels/binary.h"
 
-#define instantiate_binary(name, itype, otype, op, bopt)                      \
-  template                                                                    \
-      [[host_name(name)]] [[kernel]] void binary_##bopt<itype, otype, op>(    \
-          device const itype* a,                                              \
-          device const itype* b,                                              \
-          device otype* c,                                                    \
-          uint index [[thread_position_in_grid]]);
+#define instantiate_binary_all(op, tname, itype, otype)                 \
+  instantiate_kernel("ss" #op #tname, binary_ss, itype, otype, op)      \
+  instantiate_kernel("sv" #op #tname, binary_sv, itype, otype, op)      \
+  instantiate_kernel("vs" #op #tname, binary_vs, itype, otype, op)      \
+  instantiate_kernel("vv" #op #tname, binary_vv, itype, otype, op)      \
+  instantiate_kernel("sv2" #op #tname, binary_sv2, itype, otype, op)    \
+  instantiate_kernel("vs2" #op #tname, binary_vs2, itype, otype, op)    \
+  instantiate_kernel("vv2" #op #tname, binary_vv2, itype, otype, op)    \
+  instantiate_kernel("gn" #op #tname, binary_g, itype, otype, op)       \
+  instantiate_kernel("g1" #op #tname, binary_g_nd1, itype, otype, op)   \
+  instantiate_kernel("g2" #op #tname, binary_g_nd2, itype, otype, op)   \
+  instantiate_kernel("g3" #op #tname, binary_g_nd3, itype, otype, op)   \
+  instantiate_kernel("g4" #op #tname, binary_g_nd, itype, otype, op, 4) \
+  instantiate_kernel("g5" #op #tname, binary_g_nd,  itype, otype, op, 5)
 
-#define instantiate_binary_g_dim(name, itype, otype, op, dims) \
-  template [[host_name("g" #dims name)]] [[kernel]] void       \
-  binary_g_nd<itype, otype, op, dims>(                         \
-      device const itype* a,                                   \
-      device const itype* b,                                   \
-      device otype* c,                                         \
-      constant const int shape[dims],                          \
-      constant const size_t a_strides[dims],                   \
-      constant const size_t b_strides[dims],                   \
-      uint3 index [[thread_position_in_grid]],                 \
-      uint3 grid_dim [[threads_per_grid]]);
+#define instantiate_binary_integer(op)                   \
+  instantiate_binary_all(op, uint8, uint8_t, uint8_t)    \
+  instantiate_binary_all(op, uint16, uint16_t, uint16_t) \
+  instantiate_binary_all(op, uint32, uint32_t, uint32_t) \
+  instantiate_binary_all(op, uint64, uint64_t, uint64_t) \
+  instantiate_binary_all(op, int8, int8_t, int8_t)       \
+  instantiate_binary_all(op, int16, int16_t, int16_t)    \
+  instantiate_binary_all(op, int32, int32_t, int32_t)    \
+  instantiate_binary_all(op, int64, int64_t, int64_t)
 
-#define instantiate_binary_g_nd(name, itype, otype, op) \
-  template [[host_name("g1" name)]] [[kernel]] void     \
-  binary_g_nd1<itype, otype, op>(                       \
-      device const itype* a,                            \
-      device const itype* b,                            \
-      device otype* c,                                  \
-      constant const size_t& a_stride,                  \
-      constant const size_t& b_stride,                  \
-      uint index [[thread_position_in_grid]]);          \
-  template [[host_name("g2" name)]] [[kernel]] void     \
-  binary_g_nd2<itype, otype, op>(                       \
-      device const itype* a,                            \
-      device const itype* b,                            \
-      device otype* c,                                  \
-      constant const size_t a_strides[2],               \
-      constant const size_t b_strides[2],               \
-      uint2 index [[thread_position_in_grid]],          \
-      uint2 grid_dim [[threads_per_grid]]);             \
-  template [[host_name("g3" name)]] [[kernel]] void     \
-  binary_g_nd3<itype, otype, op>(                       \
-      device const itype* a,                            \
-      device const itype* b,                            \
-      device otype* c,                                  \
-      constant const size_t a_strides[3],               \
-      constant const size_t b_strides[3],               \
-      uint3 index [[thread_position_in_grid]],          \
-      uint3 grid_dim [[threads_per_grid]]);             \
-  instantiate_binary_g_dim(name, itype, otype, op, 4)   \
-  instantiate_binary_g_dim(name, itype, otype, op, 5)
+#define instantiate_binary_float(op)                \
+  instantiate_binary_all(op, float16, half, half)   \
+  instantiate_binary_all(op, float32, float, float) \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bfloat16_t)
 
-#define instantiate_binary_g(name, itype, otype, op)                            \
-  template [[host_name("gn" name)]] [[kernel]] void binary_g<itype, otype, op>( \
-      device const itype* a,                                                    \
-      device const itype* b,                                                    \
-      device otype* c,                                                          \
-      constant const int* shape,                                                \
-      constant const size_t* a_strides,                                         \
-      constant const size_t* b_strides,                                         \
-      constant const int& ndim,                                                 \
-      uint3 index [[thread_position_in_grid]],                                  \
-      uint3 grid_dim [[threads_per_grid]]);
+#define instantiate_binary_types(op)                              \
+  instantiate_binary_all(op, bool_, bool, bool)                   \
+  instantiate_binary_integer(op)                                  \
+  instantiate_binary_all(op, complex64, complex64_t, complex64_t) \
+  instantiate_binary_float(op)
 
-#define instantiate_binary_all(name, tname, itype, otype, op) \
-  instantiate_binary("ss" #name #tname, itype, otype, op, ss) \
-  instantiate_binary("sv" #name #tname, itype, otype, op, sv) \
-  instantiate_binary("vs" #name #tname, itype, otype, op, vs) \
-  instantiate_binary("vv" #name #tname, itype, otype, op, vv) \
-  instantiate_binary_g(#name #tname, itype, otype, op)        \
-  instantiate_binary_g_nd(#name #tname, itype, otype, op)
+#define instantiate_binary_types_bool(op)                \
+  instantiate_binary_all(op, bool_, bool, bool)          \
+  instantiate_binary_all(op, uint8, uint8_t, bool)       \
+  instantiate_binary_all(op, uint16, uint16_t, bool)     \
+  instantiate_binary_all(op, uint32, uint32_t, bool)     \
+  instantiate_binary_all(op, uint64, uint64_t, bool)     \
+  instantiate_binary_all(op, int8, int8_t, bool)         \
+  instantiate_binary_all(op, int16, int16_t, bool)       \
+  instantiate_binary_all(op, int32, int32_t, bool)       \
+  instantiate_binary_all(op, int64, int64_t, bool)       \
+  instantiate_binary_all(op, float16, half, bool)        \
+  instantiate_binary_all(op, float32, float, bool)       \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bool) \
+  instantiate_binary_all(op, complex64, complex64_t, bool)
 
-#define instantiate_binary_integer(name, op)                   \
-  instantiate_binary_all(name, uint8, uint8_t, uint8_t, op)    \
-  instantiate_binary_all(name, uint16, uint16_t, uint16_t, op) \
-  instantiate_binary_all(name, uint32, uint32_t, uint32_t, op) \
-  instantiate_binary_all(name, uint64, uint64_t, uint64_t, op) \
-  instantiate_binary_all(name, int8, int8_t, int8_t, op)       \
-  instantiate_binary_all(name, int16, int16_t, int16_t, op)    \
-  instantiate_binary_all(name, int32, int32_t, int32_t, op)    \
-  instantiate_binary_all(name, int64, int64_t, int64_t, op)
-
-#define instantiate_binary_float(name, op)                \
-  instantiate_binary_all(name, float16, half, half, op)   \
-  instantiate_binary_all(name, float32, float, float, op) \
-  instantiate_binary_all(name, bfloat16, bfloat16_t, bfloat16_t, op)
-
-#define instantiate_binary_types(name, op)                              \
-  instantiate_binary_all(name, bool_, bool, bool, op)                   \
-  instantiate_binary_integer(name, op)                                  \
-  instantiate_binary_all(name, complex64, complex64_t, complex64_t, op) \
-  instantiate_binary_float(name, op)
-
-#define instantiate_binary_types_bool(name, op)                \
-  instantiate_binary_all(name, bool_, bool, bool, op)          \
-  instantiate_binary_all(name, uint8, uint8_t, bool, op)       \
-  instantiate_binary_all(name, uint16, uint16_t, bool, op)     \
-  instantiate_binary_all(name, uint32, uint32_t, bool, op)     \
-  instantiate_binary_all(name, uint64, uint64_t, bool, op)     \
-  instantiate_binary_all(name, int8, int8_t, bool, op)         \
-  instantiate_binary_all(name, int16, int16_t, bool, op)       \
-  instantiate_binary_all(name, int32, int32_t, bool, op)       \
-  instantiate_binary_all(name, int64, int64_t, bool, op)       \
-  instantiate_binary_all(name, float16, half, bool, op)        \
-  instantiate_binary_all(name, float32, float, bool, op)       \
-  instantiate_binary_all(name, bfloat16, bfloat16_t, bool, op) \
-  instantiate_binary_all(name, complex64, complex64_t, bool, op)
-
-instantiate_binary_types(add, Add)
-instantiate_binary_types(div, Divide)
-instantiate_binary_types_bool(eq, Equal)
-instantiate_binary_types_bool(ge, Greater)
-instantiate_binary_types_bool(geq, GreaterEqual)
-instantiate_binary_types_bool(le, Less)
-instantiate_binary_types_bool(leq, LessEqual)
-instantiate_binary_types_bool(neq, NotEqual)
-instantiate_binary_float(lae, LogAddExp)
-instantiate_binary_types(max, Maximum)
-instantiate_binary_types(min, Minimum)
-instantiate_binary_types(mul, Multiply)
-instantiate_binary_types(sub, Subtract)
-instantiate_binary_types(pow, Power)
-instantiate_binary_types(rem, Remainder)
-instantiate_binary_float(arctan2, ArcTan2)
+instantiate_binary_types(Add)
+instantiate_binary_types(Divide)
+instantiate_binary_types_bool(Equal)
+instantiate_binary_types_bool(Greater)
+instantiate_binary_types_bool(GreaterEqual)
+instantiate_binary_types_bool(Less)
+instantiate_binary_types_bool(LessEqual)
+instantiate_binary_types_bool(NotEqual)
+instantiate_binary_float(LogAddExp)
+instantiate_binary_types(Maximum)
+instantiate_binary_types(Minimum)
+instantiate_binary_types(Multiply)
+instantiate_binary_types(Subtract)
+instantiate_binary_types(Power)
+instantiate_binary_types(Remainder)
+instantiate_binary_float(ArcTan2)
 
 // NaNEqual only needed for floating point types with boolean output
-instantiate_binary_all(naneq, float16, half, bool, NaNEqual)
-instantiate_binary_all(naneq, float32, float, bool, NaNEqual)
-instantiate_binary_all(naneq, bfloat16, bfloat16_t, bool, NaNEqual)
-instantiate_binary_all(naneq, complex64, complex64_t, bool, NaNEqual)
+instantiate_binary_all(NaNEqual, float16, half, bool)
+instantiate_binary_all(NaNEqual, float32, float, bool)
+instantiate_binary_all(NaNEqual, bfloat16, bfloat16_t, bool)
+instantiate_binary_all(NaNEqual, complex64, complex64_t, bool)
 
-instantiate_binary_all(lor, bool_, bool, bool, LogicalOr)
-instantiate_binary_all(land, bool_, bool, bool, LogicalAnd)
+instantiate_binary_all(LogicalOr, bool_, bool, bool)
+instantiate_binary_all(LogicalAnd, bool_, bool, bool)
 
 // Bitwise ops only need integer types and bool (except for l/r shift)
-instantiate_binary_integer(bitwise_and, BitwiseAnd)
-instantiate_binary_all(bitwise_and, bool_, bool, bool, BitwiseAnd)
-instantiate_binary_integer(bitwise_or, BitwiseOr)
-instantiate_binary_all(bitwise_or, bool_, bool, bool, BitwiseOr)
-instantiate_binary_integer(bitwise_xor, BitwiseXor)
-instantiate_binary_all(bitwise_xor, bool_, bool, bool, BitwiseXor)
-instantiate_binary_integer(left_shift, LeftShift)
-instantiate_binary_integer(right_shift, RightShift) // clang-format on
+instantiate_binary_integer(BitwiseAnd)
+instantiate_binary_all(BitwiseAnd, bool_, bool, bool)
+instantiate_binary_integer(BitwiseOr)
+instantiate_binary_all(BitwiseOr, bool_, bool, bool)
+instantiate_binary_integer(BitwiseXor)
+instantiate_binary_all(BitwiseXor, bool_, bool, bool)
+instantiate_binary_integer(LeftShift)
+instantiate_binary_integer(RightShift) // clang-format on

mlx/backend/metal/kernels/binary_two.h

@@ -48,6 +48,48 @@ template <typename T, typename U, typename Op>
   d[index] = out[1];
 }
 
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_sv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  auto out = Op()(a[0], b[offset]);
+  c[offset] = out[0];
+  d[offset] = out[1];
+}
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_vs2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  auto out = Op()(a[offset], b[0]);
+  c[offset] = out[0];
+  d[offset] = out[1];
+}
+
+template <typename T, typename U, typename Op>
+[[kernel]] void binary_vv2(
+    device const T* a,
+    device const T* b,
+    device U* c,
+    device U* d,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  auto out = Op()(a[offset], b[offset]);
+  c[offset] = out[0];
+  d[offset] = out[1];
+}
+
 template <typename T, typename U, typename Op>
 [[kernel]] void binary_g_nd1(
     device const T* a,

mlx/backend/metal/kernels/binary_two.metal

@@ -7,99 +7,37 @@
 #include "mlx/backend/metal/kernels/binary_ops.h"
 #include "mlx/backend/metal/kernels/binary_two.h"
 
-#define instantiate_binary(name, itype, otype, op, bopt)       \
-  template [[host_name(name)]] [[kernel]] void                 \
-      binary_##bopt<itype, otype, op>(                         \
-          device const itype* a,                               \
-          device const itype* b,                               \
-          device otype* c,                                     \
-          device otype* d,                                     \
-          uint index [[thread_position_in_grid]]);
+#define instantiate_binary_all(op, tname, itype, otype)                 \
+  instantiate_kernel("ss" #op #tname, binary_ss, itype, otype, op)      \
+  instantiate_kernel("sv" #op #tname, binary_sv, itype, otype, op)      \
+  instantiate_kernel("vs" #op #tname, binary_vs, itype, otype, op)      \
+  instantiate_kernel("vv" #op #tname, binary_vv, itype, otype, op)      \
+  instantiate_kernel("sv2" #op #tname, binary_sv2, itype, otype, op)    \
+  instantiate_kernel("vs2" #op #tname, binary_vs2, itype, otype, op)    \
+  instantiate_kernel("vv2" #op #tname, binary_vv2, itype, otype, op)    \
+  instantiate_kernel("gn" #op #tname, binary_g, itype, otype, op)       \
+  instantiate_kernel("g1" #op #tname, binary_g_nd1, itype, otype, op)   \
+  instantiate_kernel("g2" #op #tname, binary_g_nd2, itype, otype, op)   \
+  instantiate_kernel("g3" #op #tname, binary_g_nd3, itype, otype, op)   \
+  instantiate_kernel("g4" #op #tname, binary_g_nd, itype, otype, op, 4) \
+  instantiate_kernel("g5" #op #tname, binary_g_nd,  itype, otype, op, 5)
 
-#define instantiate_binary_g_dim(name, itype, otype, op, dims)       \
-  template [[host_name("g" #dims name)]] [[kernel]] void             \
-  binary_g_nd<itype, otype, op, dims>(                               \
-      device const itype* a,                                         \
-      device const itype* b,                                         \
-      device otype* c,                                               \
-      device otype* d,                                               \
-      constant const int shape[dims],                                \
-      constant const size_t a_strides[dims],                         \
-      constant const size_t b_strides[dims],                         \
-      uint3 index [[thread_position_in_grid]],                       \
-      uint3 grid_dim [[threads_per_grid]]);
+#define instantiate_binary_float(op)                \
+  instantiate_binary_all(op, float16, half, half)   \
+  instantiate_binary_all(op, float32, float, float) \
+  instantiate_binary_all(op, bfloat16, bfloat16_t, bfloat16_t)
 
-#define instantiate_binary_g_nd(name, itype, otype, op)       \
-  template [[host_name("g1" name)]] [[kernel]] void           \
-  binary_g_nd1<itype, otype, op>(                             \
-      device const itype* a,                                  \
-      device const itype* b,                                  \
-      device otype* c,                                        \
-      device otype* d,                                        \
-      constant const size_t& a_stride,                        \
-      constant const size_t& b_stride,                        \
-      uint index [[thread_position_in_grid]]);                \
-  template [[host_name("g2" name)]] [[kernel]] void           \
-  binary_g_nd2<itype, otype, op>(                             \
-      device const itype* a,                                  \
-      device const itype* b,                                  \
-      device otype* c,                                        \
-      device otype* d,                                        \
-      constant const size_t a_strides[2],                     \
-      constant const size_t b_strides[2],                     \
-      uint2 index [[thread_position_in_grid]],                \
-      uint2 grid_dim [[threads_per_grid]]);                   \
-  template [[host_name("g3" name)]] [[kernel]] void           \
-  binary_g_nd3<itype, otype, op>(                             \
-      device const itype* a,                                  \
-      device const itype* b,                                  \
-      device otype* c,                                        \
-      device otype* d,                                        \
-      constant const size_t a_strides[3],                     \
-      constant const size_t b_strides[3],                     \
-      uint3 index [[thread_position_in_grid]],                \
-      uint3 grid_dim [[threads_per_grid]]);                   \
-  instantiate_binary_g_dim(name, itype, otype, op, 4)         \
-  instantiate_binary_g_dim(name, itype, otype, op, 5)
+#define instantiate_binary_types(op)                              \
+  instantiate_binary_all(op, bool_, bool, bool)                   \
+  instantiate_binary_all(op, uint8, uint8_t, uint8_t)             \
+  instantiate_binary_all(op, uint16, uint16_t, uint16_t)          \
+  instantiate_binary_all(op, uint32, uint32_t, uint32_t)          \
+  instantiate_binary_all(op, uint64, uint64_t, uint64_t)          \
+  instantiate_binary_all(op, int8, int8_t, int8_t)                \
+  instantiate_binary_all(op, int16, int16_t, int16_t)             \
+  instantiate_binary_all(op, int32, int32_t, int32_t)             \
+  instantiate_binary_all(op, int64, int64_t, int64_t)             \
+  instantiate_binary_all(op, complex64, complex64_t, complex64_t) \
+  instantiate_binary_float(op)
 
-#define instantiate_binary_g(name, itype, otype, op)       \
-  template [[host_name("gn" name)]] [[kernel]] void        \
-  binary_g<itype, otype, op>(                              \
-      device const itype* a,                               \
-      device const itype* b,                               \
-      device otype* c,                                     \
-      device otype* d,                                     \
-      constant const int* shape,                           \
-      constant const size_t* a_strides,                    \
-      constant const size_t* b_strides,                    \
-      constant const int& ndim,                            \
-      uint3 index [[thread_position_in_grid]],             \
-      uint3 grid_dim [[threads_per_grid]]);
-
-#define instantiate_binary_all(name, tname, itype, otype, op) \
-  instantiate_binary("ss" #name #tname, itype, otype, op, ss) \
-  instantiate_binary("sv" #name #tname, itype, otype, op, sv) \
-  instantiate_binary("vs" #name #tname, itype, otype, op, vs) \
-  instantiate_binary("vv" #name #tname, itype, otype, op, vv) \
-  instantiate_binary_g(#name #tname, itype, otype, op)    \
-  instantiate_binary_g_nd(#name #tname, itype, otype, op)
-
-#define instantiate_binary_float(name, op)                \
-  instantiate_binary_all(name, float16, half, half, op)   \
-  instantiate_binary_all(name, float32, float, float, op) \
-  instantiate_binary_all(name, bfloat16, bfloat16_t, bfloat16_t, op)
-
-#define instantiate_binary_types(name, op)                              \
-  instantiate_binary_all(name, bool_, bool, bool, op)                   \
-  instantiate_binary_all(name, uint8, uint8_t, uint8_t, op)             \
-  instantiate_binary_all(name, uint16, uint16_t, uint16_t, op)          \
-  instantiate_binary_all(name, uint32, uint32_t, uint32_t, op)          \
-  instantiate_binary_all(name, uint64, uint64_t, uint64_t, op)          \
-  instantiate_binary_all(name, int8, int8_t, int8_t, op)                \
-  instantiate_binary_all(name, int16, int16_t, int16_t, op)             \
-  instantiate_binary_all(name, int32, int32_t, int32_t, op)             \
-  instantiate_binary_all(name, int64, int64_t, int64_t, op)             \
-  instantiate_binary_all(name, complex64, complex64_t, complex64_t, op) \
-  instantiate_binary_float(name, op)
-
-instantiate_binary_types(divmod, DivMod) // clang-format on
+instantiate_binary_types(DivMod) // clang-format on

mlx/backend/metal/kernels/conv.metal

@@ -344,12 +344,12 @@ winograd_conv_2d_weight_transform(
   const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
 
   // Initialize G matrix
-  simdgroup_matrix<T, 8, 8> G;
+  simdgroup_matrix<float, 8, 8> G;
   G.thread_elements()[0] = WGT::wt_transform[sm][sn];
   G.thread_elements()[1] = WGT::wt_transform[sm][sn + 1];
 
   // Initialize Gt matrix
-  simdgroup_matrix<T, 8, 8> Gt;
+  simdgroup_matrix<float, 8, 8> Gt;
   Gt.thread_elements()[0] = WGT::wt_transform[sn][sm];
   Gt.thread_elements()[1] = WGT::wt_transform[sn + 1][sm];
 
@@ -381,15 +381,15 @@ winograd_conv_2d_weight_transform(
     threadgroup_barrier(mem_flags::mem_threadgroup);
     // Do transform and store the result
     for (int c = 0; c < BC; ++c) {
-      simdgroup_matrix<T, 8, 8> g;
+      simdgroup_matrix<float, 8, 8> g;
       g.thread_elements()[0] =
           sm < R && sn < R ? Ws[simd_group_id][sm][sn][c] : T(0);
       g.thread_elements()[1] =
           sm < R && sn + 1 < R ? Ws[simd_group_id][sm][sn + 1][c] : T(0);
 
-      simdgroup_matrix<T, 8, 8> g_out = (G * g) * Gt;
-      wt_out_0[c * O] = g_out.thread_elements()[0];
-      wt_out_1[c * O] = g_out.thread_elements()[1];
+      simdgroup_matrix<float, 8, 8> g_out = (G * g) * Gt;
+      wt_out_0[c * O] = static_cast<T>(g_out.thread_elements()[0]);
+      wt_out_1[c * O] = static_cast<T>(g_out.thread_elements()[1]);
     }
 
     wt_in += BC;
@@ -433,12 +433,12 @@ winograd_conv_2d_input_transform(
   const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
 
   // Initialize B matrix
-  simdgroup_matrix<T, 8, 8> B;
+  simdgroup_matrix<float, 8, 8> B;
   B.thread_elements()[0] = WGT::in_transform[sm][sn];
   B.thread_elements()[1] = WGT::in_transform[sm][sn + 1];
 
   // Initialize Bt matrix
-  simdgroup_matrix<T, 8, 8> Bt;
+  simdgroup_matrix<float, 8, 8> Bt;
   Bt.thread_elements()[0] = WGT::in_transform[sn][sm];
   Bt.thread_elements()[1] = WGT::in_transform[sn + 1][sm];
 
@@ -493,13 +493,13 @@ winograd_conv_2d_input_transform(
     threadgroup_barrier(mem_flags::mem_threadgroup);
     // Do transform and store the result
     for (int c = simd_group_id; c < BC; c += N_SIMD_GROUPS) {
-      simdgroup_matrix<T, 8, 8> I;
+      simdgroup_matrix<float, 8, 8> I;
       I.thread_elements()[0] = Is[sm][sn][c];
       I.thread_elements()[1] = Is[sm][sn + 1][c];
 
-      simdgroup_matrix<T, 8, 8> I_out = (Bt * I) * B;
-      inp_out_0[c] = I_out.thread_elements()[0];
-      inp_out_1[c] = I_out.thread_elements()[1];
+      simdgroup_matrix<float, 8, 8> I_out = (Bt * I) * B;
+      inp_out_0[c] = static_cast<T>(I_out.thread_elements()[0]);
+      inp_out_1[c] = static_cast<T>(I_out.thread_elements()[1]);
     }
 
     inp_in += BC;
@@ -543,12 +543,12 @@ winograd_conv_2d_output_transform(
   const short sn = (qid & 2) * 2 + (simd_lane_id % 2) * 2;
 
   // Initialize A matrix
-  simdgroup_matrix<T, 8, 8> B;
+  simdgroup_matrix<float, 8, 8> B;
   B.thread_elements()[0] = WGT::out_transform[sm][sn];
   B.thread_elements()[1] = WGT::out_transform[sm][sn + 1];
 
   // Initialize At matrix
-  simdgroup_matrix<T, 8, 8> Bt;
+  simdgroup_matrix<float, 8, 8> Bt;
   Bt.thread_elements()[0] = WGT::out_transform[sn][sm];
   Bt.thread_elements()[1] = WGT::out_transform[sn + 1][sm];
 
@@ -597,16 +597,16 @@ winograd_conv_2d_output_transform(
     threadgroup_barrier(mem_flags::mem_threadgroup);
     // Do transform and store the result
     for (int c = simd_group_id; c < BO; c += N_SIMD_GROUPS) {
-      simdgroup_matrix<T, 8, 8> O_mat;
+      simdgroup_matrix<float, 8, 8> O_mat;
       O_mat.thread_elements()[0] = out_in_0[c];
       O_mat.thread_elements()[1] = out_in_1[c];
 
-      simdgroup_matrix<T, 8, 8> O_out = (Bt * (O_mat * B));
+      simdgroup_matrix<float, 8, 8> O_out = (Bt * (O_mat * B));
       if ((sm < M) && (sn < M)) {
-        Os[sm][sn][c] = O_out.thread_elements()[0];
+        Os[sm][sn][c] = static_cast<T>(O_out.thread_elements()[0]);
       }
       if ((sm < M) && ((sn + 1) < M)) {
-        Os[sm][sn + 1][c] = O_out.thread_elements()[1];
+        Os[sm][sn + 1][c] = static_cast<T>(O_out.thread_elements()[1]);
       }
     }
 
@@ -650,4 +650,5 @@ winograd_conv_2d_output_transform(
 
 // clang-format off
 instantiate_winograd_conv_2d(float32, float);
+instantiate_winograd_conv_2d(bfloat16, bfloat16_t);
 instantiate_winograd_conv_2d(float16, half); // clang-format on

mlx/backend/metal/kernels/copy.h

@@ -16,6 +16,26 @@ template <typename T, typename U>
   dst[index] = static_cast<U>(src[index]);
 }
 
+template <typename T, typename U>
+[[kernel]] void copy_s2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  dst[offset] = static_cast<U>(src[0]);
+}
+
+template <typename T, typename U>
+[[kernel]] void copy_v2(
+    device const T* src [[buffer(0)]],
+    device U* dst [[buffer(1)]],
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  dst[offset] = static_cast<U>(src[offset]);
+}
+
 template <typename T, typename U>
 [[kernel]] void copy_g_nd1(
     device const T* src [[buffer(0)]],

mlx/backend/metal/kernels/copy.metal

@@ -5,95 +5,23 @@
 #include "mlx/backend/metal/kernels/bf16.h"
 #include "mlx/backend/metal/kernels/copy.h"
 
-#define instantiate_copy(name, itype, otype, ctype)                        \
-  template [[host_name(name)]] [[kernel]] void copy_##ctype<itype, otype>( \
-      device const itype* src [[buffer(0)]],                               \
-      device otype* dst [[buffer(1)]],                                     \
-      uint index [[thread_position_in_grid]]);
-
-#define instantiate_copy_g_dim(name, itype, otype, dims)      \
-  template [[host_name("g" #dims "_" name)]] [[kernel]] void  \
-  copy_g_nd<itype, otype, dims>(                              \
-      device const itype* src [[buffer(0)]],                  \
-      device otype* dst [[buffer(1)]],                        \
-      constant const int* src_shape [[buffer(2)]],            \
-      constant const int64_t* src_strides [[buffer(3)]],      \
-      uint3 index [[thread_position_in_grid]],                \
-      uint3 grid_dim [[threads_per_grid]]);                   \
-  template [[host_name("gg" #dims "_" name)]] [[kernel]] void \
-  copy_gg_nd<itype, otype, dims>(                             \
-      device const itype* src [[buffer(0)]],                  \
-      device otype* dst [[buffer(1)]],                        \
-      constant const int* src_shape [[buffer(2)]],            \
-      constant const int64_t* src_strides [[buffer(3)]],      \
-      constant const int64_t* dst_strides [[buffer(4)]],      \
-      uint3 index [[thread_position_in_grid]]);
-
-#define instantiate_copy_g_nd(name, itype, otype)                              \
-  template [[host_name("g1_" name)]] [[kernel]] void copy_g_nd1<itype, otype>( \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t& src_stride [[buffer(3)]],                        \
-      uint index [[thread_position_in_grid]]);                                 \
-  template [[host_name("g2_" name)]] [[kernel]] void copy_g_nd2<itype, otype>( \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t* src_strides [[buffer(3)]],                       \
-      uint2 index [[thread_position_in_grid]],                                 \
-      uint2 grid_dim [[threads_per_grid]]);                                    \
-  template [[host_name("g3_" name)]] [[kernel]] void copy_g_nd3<itype, otype>( \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t* src_strides [[buffer(3)]],                       \
-      uint3 index [[thread_position_in_grid]],                                 \
-      uint3 grid_dim [[threads_per_grid]]);                                    \
-  template [[host_name("gg1_" name )]] [[kernel]] void                         \
-  copy_gg_nd1<itype, otype>(                                                   \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t& src_stride [[buffer(3)]],                        \
-      constant const int64_t& dst_stride [[buffer(4)]],                        \
-      uint index [[thread_position_in_grid]]);                                 \
-  template [[host_name("gg2_" name)]] [[kernel]] void                          \
-  copy_gg_nd2<itype, otype>(                                                   \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t* src_strides [[buffer(3)]],                       \
-      constant const int64_t* dst_strides [[buffer(4)]],                       \
-      uint2 index [[thread_position_in_grid]]);                                \
-  template [[host_name("gg3_" name)]] [[kernel]] void                          \
-  copy_gg_nd3<itype, otype>(                                                   \
-      device const itype* src [[buffer(0)]],                                   \
-      device otype* dst [[buffer(1)]],                                         \
-      constant const int64_t* src_strides [[buffer(3)]],                       \
-      constant const int64_t* dst_strides [[buffer(4)]],                       \
-      uint3 index [[thread_position_in_grid]]);                                \
-  instantiate_copy_g_dim(name, itype, otype, 4)                                \
-  instantiate_copy_g_dim(name, itype, otype, 5)
-
-#define instantiate_copy_g(name, itype, otype)                              \
-  template [[host_name("g_" name)]] [[kernel]] void copy_g<itype, otype>(   \
-      device const itype* src [[buffer(0)]],                                \
-      device otype* dst [[buffer(1)]],                                      \
-      constant const int* src_shape [[buffer(2)]],                          \
-      constant const int64_t* src_strides [[buffer(3)]],                    \
-      constant const int& ndim [[buffer(5)]],                               \
-      uint3 index [[thread_position_in_grid]],                              \
-      uint3 grid_dim [[threads_per_grid]]);                                 \
-  template [[host_name("gg_" name)]] [[kernel]] void copy_gg<itype, otype>( \
-      device const itype* src [[buffer(0)]],                                \
-      device otype* dst [[buffer(1)]],                                      \
-      constant const int* src_shape [[buffer(2)]],                          \
-      constant const int64_t* src_strides [[buffer(3)]],                    \
-      constant const int64_t* dst_strides [[buffer(4)]],                    \
-      constant const int& ndim [[buffer(5)]],                               \
-      uint3 index [[thread_position_in_grid]]);
-
 #define instantiate_copy_all(tname, itype, otype)    \
-  instantiate_copy("s_copy" #tname, itype, otype, s) \
-  instantiate_copy("v_copy" #tname, itype, otype, v) \
-  instantiate_copy_g("copy" #tname, itype, otype)    \
-  instantiate_copy_g_nd("copy" #tname, itype, otype)
+  instantiate_kernel("s_copy" #tname, copy_s, itype, otype) \
+  instantiate_kernel("v_copy" #tname, copy_v, itype, otype) \
+  instantiate_kernel("s2_copy" #tname, copy_s2, itype, otype) \
+  instantiate_kernel("v2_copy" #tname, copy_v2, itype, otype) \
+  instantiate_kernel("g1_copy" #tname, copy_g_nd1, itype, otype) \
+  instantiate_kernel("g2_copy" #tname, copy_g_nd2, itype, otype) \
+  instantiate_kernel("g3_copy" #tname, copy_g_nd3, itype, otype) \
+  instantiate_kernel("gg1_copy" #tname, copy_gg_nd1, itype, otype) \
+  instantiate_kernel("gg2_copy" #tname, copy_gg_nd2, itype, otype) \
+  instantiate_kernel("gg3_copy" #tname, copy_gg_nd3, itype, otype) \
+  instantiate_kernel("g4_copy" #tname, copy_g_nd, itype, otype, 4) \
+  instantiate_kernel("g5_copy" #tname, copy_g_nd, itype, otype, 5) \
+  instantiate_kernel("gg4_copy" #tname, copy_gg_nd, itype, otype, 4) \
+  instantiate_kernel("gg5_copy" #tname, copy_gg_nd, itype, otype, 5) \
+  instantiate_kernel("g_copy" #tname, copy_g, itype, otype) \
+  instantiate_kernel("gg_copy" #tname, copy_gg, itype, otype)
 
 #define instantiate_copy_itype(itname, itype)                \
   instantiate_copy_all(itname ##bool_, itype, bool)          \

mlx/backend/metal/kernels/defines.h

@@ -13,3 +13,11 @@ static MTL_CONST constexpr int REDUCE_N_READS = 16;
 static MTL_CONST constexpr int SOFTMAX_N_READS = 4;
 static MTL_CONST constexpr int RMS_N_READS = 4;
 static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096;
+
+// Instantiate a templated kernel.
+// Extra args are used as template parameters:
+// e.g. instantiate_kernel(binary_int, binary, a, b) ->
+// [[host_name(binary_int)]] [kernel] binary<a, b>
+#define instantiate_kernel(name, func, ...) \
+  template [[host_name(                     \
+      name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>;

mlx/backend/metal/kernels/expm1f.h

@@ -83,6 +83,7 @@ float expm1f(float a) {
   r = expm1f_scaled_unchecked(a, 1.0f);
   /* handle severe overflow and underflow */
   if (abs(a - 1.0f) > 88.0f) {
+    r = pow(2, a);
     r = fma(r, r, -1.0f);
   }
   return r;

mlx/backend/metal/kernels/fft.metal

@@ -1,58 +1,41 @@
 // Copyright © 2024 Apple Inc.
 
+#include "mlx/backend/metal/kernels/defines.h"
 #include "mlx/backend/metal/kernels/fft.h"
 
-#define instantiate_fft(tg_mem_size, in_T, out_T)        \
-  template [[host_name("fft_mem_" #tg_mem_size "_" #in_T \
-                       "_" #out_T)]] [[kernel]] void     \
-  fft<tg_mem_size, in_T, out_T>(                         \
-      const device in_T* in [[buffer(0)]],               \
-      device out_T* out [[buffer(1)]],                   \
-      constant const int& n,                             \
-      constant const int& batch_size,                    \
-      uint3 elem [[thread_position_in_grid]],            \
-      uint3 grid [[threads_per_grid]]);
+#define instantiate_fft(tg_mem_size, in_T, out_T)   \
+  instantiate_kernel(                               \
+      "fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      fft,                                          \
+      tg_mem_size,                                  \
+      in_T,                                         \
+      out_T)
 
-#define instantiate_rader(tg_mem_size, in_T, out_T)            \
-  template [[host_name("rader_fft_mem_" #tg_mem_size "_" #in_T \
-                       "_" #out_T)]] [[kernel]] void           \
-  rader_fft<tg_mem_size, in_T, out_T>(                         \
-      const device in_T* in [[buffer(0)]],                     \
-      device out_T* out [[buffer(1)]],                         \
-      const device float2* raders_b_q [[buffer(2)]],           \
-      const device short* raders_g_q [[buffer(3)]],            \
-      const device short* raders_g_minus_q [[buffer(4)]],      \
-      constant const int& n,                                   \
-      constant const int& batch_size,                          \
-      constant const int& rader_n,                             \
-      uint3 elem [[thread_position_in_grid]],                  \
-      uint3 grid [[threads_per_grid]]);
+#define instantiate_rader(tg_mem_size, in_T, out_T)       \
+  instantiate_kernel(                                     \
+      "rader_fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      rader_fft,                                          \
+      tg_mem_size,                                        \
+      in_T,                                               \
+      out_T)
 
-#define instantiate_bluestein(tg_mem_size, in_T, out_T)            \
-  template [[host_name("bluestein_fft_mem_" #tg_mem_size "_" #in_T \
-                       "_" #out_T)]] [[kernel]] void               \
-  bluestein_fft<tg_mem_size, in_T, out_T>(                         \
-      const device in_T* in [[buffer(0)]],                         \
-      device out_T* out [[buffer(1)]],                             \
-      const device float2* w_q [[buffer(2)]],                      \
-      const device float2* w_k [[buffer(3)]],                      \
-      constant const int& length,                                  \
-      constant const int& n,                                       \
-      constant const int& batch_size,                              \
-      uint3 elem [[thread_position_in_grid]],                      \
-      uint3 grid [[threads_per_grid]]);
+#define instantiate_bluestein(tg_mem_size, in_T, out_T)       \
+  instantiate_kernel(                                         \
+      "bluestein_fft_mem_" #tg_mem_size "_" #in_T "_" #out_T, \
+      bluestein_fft,                                          \
+      tg_mem_size,                                            \
+      in_T,                                                   \
+      out_T)
 
-#define instantiate_four_step(tg_mem_size, in_T, out_T, step, real)       \
-  template [[host_name("four_step_mem_" #tg_mem_size "_" #in_T "_" #out_T \
-                       "_" #step "_" #real)]] [[kernel]] void             \
-  four_step_fft<tg_mem_size, in_T, out_T, step, real>(                    \
-      const device in_T* in [[buffer(0)]],                                \
-      device out_T* out [[buffer(1)]],                                    \
-      constant const int& n1,                                             \
-      constant const int& n2,                                             \
-      constant const int& batch_size,                                     \
-      uint3 elem [[thread_position_in_grid]],                             \
-      uint3 grid [[threads_per_grid]]);
+#define instantiate_four_step(tg_mem_size, in_T, out_T, step, real)           \
+  instantiate_kernel(                                                         \
+      "four_step_mem_" #tg_mem_size "_" #in_T "_" #out_T "_" #step "_" #real, \
+      four_step_fft,                                                          \
+      tg_mem_size,                                                            \
+      in_T,                                                                   \
+      out_T,                                                                  \
+      step,                                                                   \
+      real)
 
 // clang-format off
 #define instantiate_ffts(tg_mem_size)                        \

mlx/backend/metal/kernels/gemv.metal

@@ -17,29 +17,250 @@ using namespace metal;
 
 #define MLX_MTL_CONST static constant constexpr const
 
-MLX_MTL_CONST int SIMD_SIZE = 32;
-
 template <
     typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
     const int TM, /* Thread rows (in elements) */
     const int TN, /* Thread cols (in elements) */
     const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
 struct GEMVKernel {
-  static_assert(BN == SIMD_SIZE, "gemv block must have a width of SIMD_SIZE");
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
 
-  // - The matrix of size (M = out_vec_size, N = in_vec_size) is divided up
-  //   into blocks of (BM * TM, BN * TN) divided among threadgroups
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  static_assert(
+      SN == 8 || SN == 16 || SN == 32,
+      "gemv block must have a width of 8, 16, or 32");
+
+  // - The matrix of size (M = out_vec_size, K = in_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
   // - Every thread works on a block of (TM, TN)
-  // - We assume each thead group is launched with (BN, BM, 1) threads
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
   //
-  // 1. A thread loads TN elements each from mat along TM contiguous rows
+  // 1. A thread loads TN elements each from mat along TM rows
   //    and the corresponding scalar from the vector
   // 2. The thread then multiplies and adds to accumulate its local result for
   //    the block
   // 3. At the end, each thread has accumulated results over all blocks across
   //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated blockM outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BN > 1 ? BN*(blockM + TM) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BN > 1;
+
+  static METAL_FUNC void
+  load_unsafe(const device T* src, thread T dst[TN], const int src_offset = 0) {
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      dst[tn] = src[src_offset + tn];
+    }
+  }
+
+  static METAL_FUNC void load_safe(
+      const device T* src,
+      thread T dst[TN],
+      const int src_offset = 0,
+      const int src_size = TN) {
+    if (src_offset + TN <= src_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src[src_offset + tn];
+      }
+    } else { // Edgecase
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src_offset + tn < src_size ? src[src_offset + tn] : 0;
+      }
+    }
+  }
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      const device T* bias [[buffer(2)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& matrix_ld [[buffer(6)]],
+      const constant float& alpha [[buffer(7)]],
+      const constant float& beta [[buffer(8)]],
+      const constant int& bias_stride [[buffer(14)]],
+      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]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    thread T result[TM] = {0};
+    thread T inter[TN];
+    thread T v_coeff[TN];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = BN != 1 ? SM * (simd_gid / BN) : int(SM * simd_gid);
+    const int simdN = BN != 1 ? SN * (simd_gid % BN) : 0;
+
+    int bm = (simdM + thrM) * TM;
+    int bn = (simdN + thrN) * TN;
+
+    // Block position
+    int out_row = tid.x * blockM + bm;
+
+    // Exit simdgroup if rows out of bound
+    if (out_row >= out_vec_size)
+      return;
+
+    // Adjust tail simdgroup to ensure in bound reads
+    out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
+
+    // Advance matrix
+    mat += out_row * matrix_ld;
+
+    constexpr const uniform<int> loop_stride = make_uniform(blockN);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Loop over in_vec in blocks of blockN
+    for (int i = 0; i < n_iter; ++i) {
+      load_unsafe(in_vec, v_coeff, bn);
+
+      // Per thread work loop
+      int mat_offset = 0;
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        // Load for the row
+        load_unsafe(mat, inter, mat_offset + bn);
+
+        // Accumulate results
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
+        }
+
+        mat_offset += matrix_ld;
+      }
+
+      bn += blockN;
+    }
+
+    if (leftover > 0) {
+      load_safe(in_vec, v_coeff, bn, in_size);
+
+      // Per thread work loop
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        // Load for the row
+        load_safe(&mat[tm * matrix_ld], inter, bn, in_size);
+
+        // Accumulate results
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
+        }
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tm = 0; tm < TM; tm++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sn = (SN / 2); sn >= 1; sn >>= 1) {
+        result[tm] += simd_shuffle_down(result[tm], sn);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup T* tgp_results = tgp_memory + sgN * (blockM + TM) + bm;
+      if (thrN == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          tgp_results[tm] = result[tm];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgn = 1; sgn < BN; sgn++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              result[tm] += tgp_results[sgn * (blockM + TM) + tm];
+            }
+          }
+        }
+      }
+    }
+
+    // Write outputs
+    if (simdN == 0 && thrN == 0) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        if (kDoAxpby) {
+          out_vec[out_row + tm] = static_cast<T>(alpha) * result[tm] +
+              static_cast<T>(beta) * bias[(out_row + tm) * bias_stride];
+        } else {
+          out_vec[out_row + tm] = result[tm];
+        }
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
+struct GEMVTKernel {
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then accumulates its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
   // 4. Each threadgroup writes its accumulated BN * TN outputs
   //
   // Edge case handling:
@@ -49,7 +270,8 @@ struct GEMVKernel {
   //   * The last thread that partially overlaps with the matrix is shifted
   //     inwards such that the thread block fits exactly in the matrix
 
-  MLX_MTL_CONST short tgp_mem_size = BN * TN * 2;
+  MLX_MTL_CONST short tgp_mem_size = BM > 1 ? BM*(blockN + TN) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BM > 1;
 
   static METAL_FUNC void run(
       const device T* mat [[buffer(0)]],
@@ -70,230 +292,113 @@ struct GEMVKernel {
     // Appease compiler
     (void)lid;
 
-    // Threadgroup in_vec cache
-    threadgroup T* in_vec_block = tgp_memory + simd_lid * TN * 2;
-
-    // Thread local accumulation results
-    thread T result[TM] = {0};
-    thread T inter[TN];
-    thread T v_coeff[TN];
-
-    // Block position
-    int out_row = (tid.x * BM + simd_gid) * TM;
-
-    // Exit simdgroup if rows out of bound
-    if (out_row >= out_vec_size)
-      return;
-
-    // Adjust tail simdgroup to ensure in bound reads
-    out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
-
-    // Advance matrix
-    mat += out_row * marix_ld;
-
-    // Loop over in_vec in blocks of BN * TN
-    for (int bn = simd_lid * TN; bn < in_vec_size; bn += BN * TN) {
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-
-      // Prefetch in_vector for threadgroup use
-      if (simd_gid == 0) {
-        // Main load loop
-        if (bn + TN <= in_vec_size) {
-          MLX_MTL_PRAGMA_UNROLL
-          for (int tn = 0; tn < TN; tn++) {
-            in_vec_block[tn] = in_vec[bn + tn];
-          }
-
-        } else { // Edgecase
-
-          MLX_MTL_PRAGMA_UNROLL
-          for (int tn = 0; tn < TN; tn++) {
-            in_vec_block[tn] = bn + tn < in_vec_size ? in_vec[bn + tn] : 0;
-          }
-        }
-      }
-
-      threadgroup_barrier(mem_flags::mem_threadgroup);
-
-      // Load for all rows
-      MLX_MTL_PRAGMA_UNROLL
-      for (int tn = 0; tn < TN; tn++) {
-        v_coeff[tn] = in_vec_block[tn];
-      }
-
-      // Per thread work loop
-      MLX_MTL_PRAGMA_UNROLL
-      for (int tm = 0; tm < TM; tm++) {
-        // Load for the row
-        if (bn + TN <= in_vec_size) {
-          MLX_MTL_PRAGMA_UNROLL
-          for (int tn = 0; tn < TN; tn++) {
-            inter[tn] = mat[tm * marix_ld + bn + tn];
-          }
-
-        } else { // Edgecase
-          MLX_MTL_PRAGMA_UNROLL
-          for (int tn = 0; tn < TN; tn++) {
-            int col_idx =
-                (bn + tn) < in_vec_size ? (bn + tn) : (in_vec_size - 1);
-            inter[tn] = mat[tm * marix_ld + col_idx];
-          }
-        }
-
-        // Accumulate results
-        MLX_MTL_PRAGMA_UNROLL
-        for (int tn = 0; tn < TN; tn++) {
-          result[tm] += inter[tn] * v_coeff[tn];
-        }
-      }
-    }
-
-    // Simdgroup accumulations
-    MLX_MTL_PRAGMA_UNROLL
-    for (int tm = 0; tm < TM; tm++) {
-      result[tm] = simd_sum(result[tm]);
-    }
-
-    // Write outputs
-    if (simd_lid == 0) {
-      MLX_MTL_PRAGMA_UNROLL
-      for (int tm = 0; tm < TM; tm++) {
-        if (kDoAxpby) {
-          out_vec[out_row + tm] = static_cast<T>(alpha) * result[tm] +
-              static_cast<T>(beta) * bias[(out_row + tm) * bias_stride];
-        } else {
-          out_vec[out_row + tm] = result[tm];
-        }
-      }
-    }
-  }
-};
-
-///////////////////////////////////////////////////////////////////////////////
-/// Vector matrix multiplication
-///////////////////////////////////////////////////////////////////////////////
-
-template <
-    typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
-    const int TM, /* Thread rows (in elements) */
-    const int TN, /* Thread cols (in elements) */
-    const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
-struct GEMVTKernel {
-  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up
-  //   into blocks of (BM * TM, BN * TN) divided among threadgroups
-  // - Every thread works on a block of (TM, TN)
-  // - We assume each thead group is launched with (BN, BM, 1) threads
-  //
-  // 1. A thread loads TN elements each from mat along TM contiguous rows
-  //    and the corresponding scalar from the vector
-  // 2. The thread then accumulates its local result for the block
-  // 3. At the end, each thread has accumulated results over all blocks across
-  //    the rows. These are then summed up across the threadgroup
-  // 4. Each threadgroup writes its accumulated BN * TN outputs
-  //
-  // Edge case handling:
-  // - The threadgroup with the largest tid has blocks that exceed the matrix
-  //   * The blocks that start outside the matrix are never read (thread results
-  //     remain zero)
-  //   * The last thread that partially overlaps with the matrix is shifted
-  //     inwards such that the thread block fits exactly in the matrix
-
-  MLX_MTL_CONST short tgp_mem_size = BN * BM * TN;
-
-  static METAL_FUNC void run(
-      const device T* mat [[buffer(0)]],
-      const device T* in_vec [[buffer(1)]],
-      const device T* bias [[buffer(2)]],
-      device T* out_vec [[buffer(3)]],
-      const constant int& in_vec_size [[buffer(4)]],
-      const constant int& out_vec_size [[buffer(5)]],
-      const constant int& marix_ld [[buffer(6)]],
-      const constant float& alpha [[buffer(7)]],
-      const constant float& beta [[buffer(8)]],
-      const constant int& bias_stride [[buffer(14)]],
-      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]]) {
-    // Appease compiler
-    (void)simd_gid;
-    (void)simd_lid;
-
     // Thread local accumulation results
     T result[TN] = {0};
     T inter[TN];
     T v_coeff[TM];
 
-    // Threadgroup accumulation results
-    threadgroup T* tgp_results = tgp_memory + lid.x * BM * TN;
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
 
-    int out_col = (tid.x * BN + lid.x) * TN;
-    int in_row = lid.y * TM;
+    const int sgM = BN != 1 ? (simd_gid / BN) : int(simd_gid);
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = SM * sgM;
+    const int simdN = SN * sgN;
+
+    int cm = (simdM + thrM);
+    int cn = (simdN + thrN);
+
+    int bm = cm * TM;
+    int bn = cn * TN;
+
+    int out_col = tid.x * blockN + bn;
+
+    constexpr const uniform<int> loop_stride = make_uniform(blockM);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
 
     // Edgecase handling
     if (out_col < out_vec_size) {
       out_col = out_col + TN < out_vec_size ? out_col : out_vec_size - TN;
 
       // Per thread accumulation main loop
-      int bm = in_row;
-      for (; bm < in_vec_size; bm += BM * TM) {
+      for (int i = 0; i < n_iter; ++i) {
         // Adding a threadgroup_barrier improves performance slightly
         // This is possibly it may help exploit cache better
         threadgroup_barrier(mem_flags::mem_none);
 
-        if (bm + TM <= in_vec_size) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          v_coeff[tm] = in_vec[bm + tm];
+        }
+
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          for (int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+          }
+          for (int tn = 0; tn < TN; tn++) {
+            result[tn] += v_coeff[tm] * inter[tn];
+          }
+        }
+
+        bm += blockM;
+      }
+
+      if (leftover > 0) {
+        for (int tm = 0; tm < TM && bm + tm < in_vec_size; tm++) {
+          v_coeff[tm] = in_vec[bm + tm];
+
           MLX_MTL_PRAGMA_UNROLL
-          for (int tm = 0; tm < TM; tm++) {
-            v_coeff[tm] = in_vec[bm + tm];
+          for (int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
           }
 
           MLX_MTL_PRAGMA_UNROLL
-          for (int tm = 0; tm < TM; tm++) {
-            for (int tn = 0; tn < TN; tn++) {
-              inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
-            }
-            for (int tn = 0; tn < TN; tn++) {
-              result[tn] += v_coeff[tm] * inter[tn];
-            }
-          }
-
-        } else { // Edgecase handling
-          for (int tm = 0; bm + tm < in_vec_size; tm++) {
-            v_coeff[tm] = in_vec[bm + tm];
-
-            for (int tn = 0; tn < TN; tn++) {
-              inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
-            }
-            for (int tn = 0; tn < TN; tn++) {
-              result[tn] += v_coeff[tm] * inter[tn];
-            }
+          for (int tn = 0; tn < TN; tn++) {
+            result[tn] += v_coeff[tm] * inter[tn];
           }
         }
       }
     }
 
-    // Threadgroup collection
-
+    // Simdgroup accumulations
     MLX_MTL_PRAGMA_UNROLL
-    for (int i = 0; i < TN; i++) {
-      tgp_results[lid.y * TN + i] = result[i];
+    for (int tn = 0; tn < TN; tn++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sm = (SM / 2); sm >= 1; sm >>= 1) {
+        result[tn] += simd_shuffle_down(result[tn], SN * sm);
+      }
     }
 
-    threadgroup_barrier(mem_flags::mem_threadgroup);
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup T* tgp_results = tgp_memory + sgM * (blockN + TN) + bn;
+      if (thrM == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          tgp_results[tn] = result[tn];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgM == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgm = 1; sgm < BM; sgm++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += tgp_results[sgm * (blockN + TN) + tn];
+            }
+          }
+        }
+      }
+    }
 
     // Threadgroup accumulation and writing out results
-    if (lid.y == 0 && out_col < out_vec_size) {
-      MLX_MTL_PRAGMA_UNROLL
-      for (int i = 1; i < BM; i++) {
-        MLX_MTL_PRAGMA_UNROLL
-        for (int j = 0; j < TN; j++) {
-          result[j] += tgp_results[i * TN + j];
-        }
-      }
-
+    if (cm == 0 && out_col < out_vec_size) {
       MLX_MTL_PRAGMA_UNROLL
       for (int j = 0; j < TN; j++) {
         if (kDoAxpby) {
@@ -313,13 +418,15 @@ struct GEMVTKernel {
 
 template <
     typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
     const int TM, /* Thread rows (in elements) */
     const int TN, /* Thread cols (in elements) */
     const bool kDoNCBatch, /* Batch ndim > 1 */
     const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
-[[kernel, max_total_threads_per_threadgroup(BM* BN)]] void gemv(
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv(
     const device T* mat [[buffer(0)]],
     const device T* in_vec [[buffer(1)]],
     const device T* bias [[buffer(2)]],
@@ -339,8 +446,9 @@ template <
     uint3 lid [[thread_position_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
-  using gemv_kernel = GEMVKernel<T, BM, BN, TM, TN, kDoAxpby>;
-  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+  using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   // Update batch offsets
   if (kDoNCBatch) {
@@ -373,17 +481,19 @@ template <
       alpha,
       beta,
       bias_stride,
-      tgp_memory,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
       tid,
       lid,
       simd_gid,
       simd_lid);
 }
 
-#define instantiate_gemv_helper(name, itype, bm, bn, tm, tn, nc, axpby)      \
-  template [[host_name("gemv_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn \
-                       "_nc" #nc "_axpby" #axpby)]] [[kernel]] void          \
-  gemv<itype, bm, bn, tm, tn, nc, axpby>(                                    \
+#define instantiate_gemv_helper(                                             \
+    name, itype, bm, bn, sm, sn, tm, tn, nc, axpby)                          \
+  template [[host_name("gemv_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn \
+                       "_tm" #tm "_tn" #tn "_nc" #nc                         \
+                       "_axpby" #axpby)]] [[kernel]] void                    \
+  gemv<itype, bm, bn, sm, sn, tm, tn, nc, axpby>(                            \
       const device itype* mat [[buffer(0)]],                                 \
       const device itype* in_vec [[buffer(1)]],                              \
       const device itype* bias [[buffer(2)]],                                \
@@ -405,11 +515,11 @@ template <
       uint simd_lid [[thread_index_in_simdgroup]]);
 
 // clang-format off
-#define instantiate_gemv(name, itype, bm, bn, tm, tn)        \
-  instantiate_gemv_helper(name, itype, bm, bn, tm, tn, 0, 0) \
-  instantiate_gemv_helper(name, itype, bm, bn, tm, tn, 0, 1) \
-  instantiate_gemv_helper(name, itype, bm, bn, tm, tn, 1, 0) \
-  instantiate_gemv_helper(name, itype, bm, bn, tm, tn, 1, 1) // clang-format on
+#define instantiate_gemv(name, itype, bm, bn, tm, tn)              \
+  instantiate_gemv_helper(name, itype, bm, 1, 1, bn, tm, tn, 0, 0) \
+  instantiate_gemv_helper(name, itype, bm, 1, 1, bn, tm, tn, 0, 1) \
+  instantiate_gemv_helper(name, itype, bm, 1, 1, bn, tm, tn, 1, 0) \
+  instantiate_gemv_helper(name, itype, bm, 1, 1, bn, tm, tn, 1, 1) // clang-format on
 
 // clang-format off
 #define instantiate_gemv_blocks(name, itype) \
@@ -423,11 +533,13 @@ instantiate_gemv_blocks(bfloat16, bfloat16_t);
 
 template <
     typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
     const int TM, /* Thread rows (in elements) */
     const int TN> /* Thread cols (in elements) */
-[[kernel, max_total_threads_per_threadgroup(BM* BN)]] void gemv_bs(
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_gather(
     const device T* mat [[buffer(0)]],
     const device T* in_vec [[buffer(1)]],
     const device T* bias [[buffer(2)]],
@@ -452,8 +564,9 @@ template <
     uint3 lid [[thread_position_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
-  using gemv_kernel = GEMVKernel<T, BM, BN, TM, TN, false>;
-  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+  using gemv_kernel = GEMVKernel<T, BM, BN, SM, SN, TM, TN, false>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   uint32_t indx_vec;
   uint32_t indx_mat;
@@ -501,47 +614,47 @@ template <
       alpha,
       beta,
       batch_ndim, // Not used
-      tgp_memory,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
       tid,
       lid,
       simd_gid,
       simd_lid);
 }
 
-#define instantiate_gemv_bs_helper(nm, itype, bm, bn, tm, tn)       \
-  template [[host_name("gemv_bs_" #nm "_bm" #bm "_bn" #bn "_tm" #tm \
-                       "_tn" #tn)]] [[kernel]] void                 \
-  gemv_bs<itype, bm, bn, tm, tn>(                                   \
-      const device itype* mat [[buffer(0)]],                        \
-      const device itype* in_vec [[buffer(1)]],                     \
-      const device itype* bias [[buffer(2)]],                       \
-      device itype* out_vec [[buffer(3)]],                          \
-      const constant int& in_vec_size [[buffer(4)]],                \
-      const constant int& out_vec_size [[buffer(5)]],               \
-      const constant int& marix_ld [[buffer(6)]],                   \
-      const constant float& alpha [[buffer(7)]],                    \
-      const constant float& beta [[buffer(8)]],                     \
-      const constant int& batch_ndim [[buffer(9)]],                 \
-      const constant int* batch_shape [[buffer(10)]],               \
-      const constant size_t* index_batch_strides [[buffer(11)]],    \
-      const constant int& vector_batch_ndim [[buffer(12)]],         \
-      const constant int* vector_batch_shape [[buffer(13)]],        \
-      const constant size_t* vector_batch_stride [[buffer(14)]],    \
-      const constant int& matrix_batch_ndim [[buffer(15)]],         \
-      const constant int* matrix_batch_shape [[buffer(16)]],        \
-      const constant size_t* matrix_batch_stride [[buffer(17)]],    \
-      const constant uint32_t* vec_indices [[buffer(18)]],          \
-      const constant uint32_t* mat_indices [[buffer(19)]],          \
-      uint3 tid [[threadgroup_position_in_grid]],                   \
-      uint3 lid [[thread_position_in_threadgroup]],                 \
-      uint simd_gid [[simdgroup_index_in_threadgroup]],             \
+#define instantiate_gemv_bs_helper(nm, itype, bm, bn, sm, sn, tm, tn)   \
+  template [[host_name("gemv_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm \
+                       "_sn" #sn "_tm" #tm "_tn" #tn)]] [[kernel]] void \
+  gemv_gather<itype, bm, bn, sm, sn, tm, tn>(                           \
+      const device itype* mat [[buffer(0)]],                            \
+      const device itype* in_vec [[buffer(1)]],                         \
+      const device itype* bias [[buffer(2)]],                           \
+      device itype* out_vec [[buffer(3)]],                              \
+      const constant int& in_vec_size [[buffer(4)]],                    \
+      const constant int& out_vec_size [[buffer(5)]],                   \
+      const constant int& marix_ld [[buffer(6)]],                       \
+      const constant float& alpha [[buffer(7)]],                        \
+      const constant float& beta [[buffer(8)]],                         \
+      const constant int& batch_ndim [[buffer(9)]],                     \
+      const constant int* batch_shape [[buffer(10)]],                   \
+      const constant size_t* index_batch_strides [[buffer(11)]],        \
+      const constant int& vector_batch_ndim [[buffer(12)]],             \
+      const constant int* vector_batch_shape [[buffer(13)]],            \
+      const constant size_t* vector_batch_stride [[buffer(14)]],        \
+      const constant int& matrix_batch_ndim [[buffer(15)]],             \
+      const constant int* matrix_batch_shape [[buffer(16)]],            \
+      const constant size_t* matrix_batch_stride [[buffer(17)]],        \
+      const constant uint32_t* vec_indices [[buffer(18)]],              \
+      const constant uint32_t* mat_indices [[buffer(19)]],              \
+      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]]);
 
 // clang-format off
 #define instantiate_gemv_bs_blocks(name, itype)        \
-  instantiate_gemv_bs_helper(name, itype, 4, 32, 1, 4) \
-  instantiate_gemv_bs_helper(name, itype, 4, 32, 4, 4) \
-  instantiate_gemv_bs_helper(name, itype, 8, 32, 4, 4) // clang-format on
+  instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 1, 4) \
+  instantiate_gemv_bs_helper(name, itype, 4, 1, 1, 32, 4, 4) \
+  instantiate_gemv_bs_helper(name, itype, 8, 1, 1, 32, 4, 4) // clang-format on
 
 instantiate_gemv_bs_blocks(float32, float);
 instantiate_gemv_bs_blocks(float16, half);
@@ -553,13 +666,15 @@ instantiate_gemv_bs_blocks(bfloat16, bfloat16_t);
 
 template <
     typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
     const int TM, /* Thread rows (in elements) */
     const int TN, /* Thread cols (in elements) */
     const bool kDoNCBatch, /* Batch ndim > 1 */
     const bool kDoAxpby> /* Do out = alpha * out + beta * bias */
-[[kernel, max_total_threads_per_threadgroup(BM* BN)]] void gemv_t(
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t(
     const device T* mat [[buffer(0)]],
     const device T* in_vec [[buffer(1)]],
     const device T* bias [[buffer(2)]],
@@ -579,8 +694,9 @@ template <
     uint3 lid [[thread_position_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
-  using gemv_kernel = GEMVTKernel<T, BM, BN, TM, TN, kDoAxpby>;
-  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, kDoAxpby>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   // Update batch offsets
   if (kDoNCBatch) {
@@ -613,17 +729,19 @@ template <
       alpha,
       beta,
       bias_stride,
-      tgp_memory,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
       tid,
       lid,
       simd_gid,
       simd_lid);
 }
 
-#define instantiate_gemv_t_helper(name, itype, bm, bn, tm, tn, nc, axpby)      \
-  template [[host_name("gemv_t_" #name "_bm" #bm "_bn" #bn "_tm" #tm "_tn" #tn \
-                       "_nc" #nc "_axpby" #axpby)]] [[kernel]] void            \
-  gemv_t<itype, bm, bn, tm, tn, nc, axpby>(                                    \
+#define instantiate_gemv_t_helper(                                             \
+    name, itype, bm, bn, sm, sn, tm, tn, nc, axpby)                            \
+  template [[host_name("gemv_t_" #name "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn \
+                       "_tm" #tm "_tn" #tn "_nc" #nc                           \
+                       "_axpby" #axpby)]] [[kernel]] void                      \
+  gemv_t<itype, bm, bn, sm, sn, tm, tn, nc, axpby>(                            \
       const device itype* mat [[buffer(0)]],                                   \
       const device itype* in_vec [[buffer(1)]],                                \
       const device itype* bias [[buffer(2)]],                                  \
@@ -645,20 +763,19 @@ template <
       uint simd_lid [[thread_index_in_simdgroup]]);
 
 // clang-format off
-#define instantiate_gemv_t(name, itype, bm, bn, tm, tn)        \
-  instantiate_gemv_t_helper(name, itype, bm, bn, tm, tn, 0, 0) \
-  instantiate_gemv_t_helper(name, itype, bm, bn, tm, tn, 0, 1) \
-  instantiate_gemv_t_helper(name, itype, bm, bn, tm, tn, 1, 0) \
-  instantiate_gemv_t_helper(name, itype, bm, bn, tm, tn, 1, 1) // clang-format on
+#define instantiate_gemv_t(name, itype, bm, bn, sm, sn, tm, tn)        \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 0) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 0, 1) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 0) \
+  instantiate_gemv_t_helper(name, itype, bm, bn, sm, sn, tm, tn, 1, 1) // clang-format on
 
 // clang-format off
 #define instantiate_gemv_t_blocks(name, itype) \
-  instantiate_gemv_t(name, itype, 8, 8, 4, 1)  \
-  instantiate_gemv_t(name, itype, 8, 8, 4, 4)  \
-  instantiate_gemv_t(name, itype, 8, 16, 4, 4) \
-  instantiate_gemv_t(name, itype, 8, 32, 4, 4) \
-  instantiate_gemv_t(name, itype, 8, 64, 4, 4) \
-  instantiate_gemv_t(name, itype, 8, 128, 4, 4) // clang-format on
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 1) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 4,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 16, 8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on
 
 // clang-format off
 instantiate_gemv_t_blocks(float32, float);
@@ -667,11 +784,13 @@ instantiate_gemv_t_blocks(bfloat16, bfloat16_t); // clang-format on
 
 template <
     typename T,
-    const int BM, /* Threadgroup rows (in threads) */
-    const int BN, /* Threadgroup cols (in threads) */
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
     const int TM, /* Thread rows (in elements) */
     const int TN> /* Thread cols (in elements) */
-[[kernel, max_total_threads_per_threadgroup(BM* BN)]] void gemv_t_bs(
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_gather(
     const device T* mat [[buffer(0)]],
     const device T* in_vec [[buffer(1)]],
     const device T* bias [[buffer(2)]],
@@ -696,8 +815,9 @@ template <
     uint3 lid [[thread_position_in_threadgroup]],
     uint simd_gid [[simdgroup_index_in_threadgroup]],
     uint simd_lid [[thread_index_in_simdgroup]]) {
-  using gemv_kernel = GEMVTKernel<T, BM, BN, TM, TN, false>;
-  threadgroup T tgp_memory[gemv_kernel::tgp_mem_size];
+  using gemv_kernel = GEMVTKernel<T, BM, BN, SM, SN, TM, TN, false>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
 
   uint32_t indx_vec;
   uint32_t indx_mat;
@@ -745,50 +865,49 @@ template <
       alpha,
       beta,
       batch_ndim, // Not used,
-      tgp_memory,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
       tid,
       lid,
       simd_gid,
       simd_lid);
 }
 
-#define instantiate_gemv_t_bs_helper(nm, itype, bm, bn, tm, tn)       \
-  template [[host_name("gemv_t_bs_" #nm "_bm" #bm "_bn" #bn "_tm" #tm \
-                       "_tn" #tn)]] [[kernel]] void                   \
-  gemv_t_bs<itype, bm, bn, tm, tn>(                                   \
-      const device itype* mat [[buffer(0)]],                          \
-      const device itype* in_vec [[buffer(1)]],                       \
-      const device itype* bias [[buffer(2)]],                         \
-      device itype* out_vec [[buffer(3)]],                            \
-      const constant int& in_vec_size [[buffer(4)]],                  \
-      const constant int& out_vec_size [[buffer(5)]],                 \
-      const constant int& marix_ld [[buffer(6)]],                     \
-      const constant float& alpha [[buffer(7)]],                      \
-      const constant float& beta [[buffer(8)]],                       \
-      const constant int& batch_ndim [[buffer(9)]],                   \
-      const constant int* batch_shape [[buffer(10)]],                 \
-      const constant size_t* index_batch_strides [[buffer(11)]],      \
-      const constant int& vector_batch_ndim [[buffer(12)]],           \
-      const constant int* vector_batch_shape [[buffer(13)]],          \
-      const constant size_t* vector_batch_stride [[buffer(14)]],      \
-      const constant int& matrix_batch_ndim [[buffer(15)]],           \
-      const constant int* matrix_batch_shape [[buffer(16)]],          \
-      const constant size_t* matrix_batch_stride [[buffer(17)]],      \
-      const constant uint32_t* vec_indices [[buffer(18)]],            \
-      const constant uint32_t* mat_indices [[buffer(19)]],            \
-      uint3 tid [[threadgroup_position_in_grid]],                     \
-      uint3 lid [[thread_position_in_threadgroup]],                   \
-      uint simd_gid [[simdgroup_index_in_threadgroup]],               \
+#define instantiate_gemv_t_bs_helper(nm, itype, bm, bn, sm, sn, tm, tn)   \
+  template [[host_name("gemv_t_gather_" #nm "_bm" #bm "_bn" #bn "_sm" #sm \
+                       "_sn" #sn "_tm" #tm "_tn" #tn)]] [[kernel]] void   \
+  gemv_t_gather<itype, bm, bn, sm, sn, tm, tn>(                           \
+      const device itype* mat [[buffer(0)]],                              \
+      const device itype* in_vec [[buffer(1)]],                           \
+      const device itype* bias [[buffer(2)]],                             \
+      device itype* out_vec [[buffer(3)]],                                \
+      const constant int& in_vec_size [[buffer(4)]],                      \
+      const constant int& out_vec_size [[buffer(5)]],                     \
+      const constant int& marix_ld [[buffer(6)]],                         \
+      const constant float& alpha [[buffer(7)]],                          \
+      const constant float& beta [[buffer(8)]],                           \
+      const constant int& batch_ndim [[buffer(9)]],                       \
+      const constant int* batch_shape [[buffer(10)]],                     \
+      const constant size_t* index_batch_strides [[buffer(11)]],          \
+      const constant int& vector_batch_ndim [[buffer(12)]],               \
+      const constant int* vector_batch_shape [[buffer(13)]],              \
+      const constant size_t* vector_batch_stride [[buffer(14)]],          \
+      const constant int& matrix_batch_ndim [[buffer(15)]],               \
+      const constant int* matrix_batch_shape [[buffer(16)]],              \
+      const constant size_t* matrix_batch_stride [[buffer(17)]],          \
+      const constant uint32_t* vec_indices [[buffer(18)]],                \
+      const constant uint32_t* mat_indices [[buffer(19)]],                \
+      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]]);
 
 // clang-format off
-#define instantiate_gemv_t_bs_blocks(name, itype) \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 8, 4, 1)  \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 8, 4, 4)  \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 16, 4, 4) \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 32, 4, 4) \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 64, 4, 4) \
-  instantiate_gemv_t_bs_helper(name, itype, 8, 128, 4, 4) // clang-format on
+#define instantiate_gemv_t_bs_blocks(name, itype)              \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  2, 8, 4, 4, 1) \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  2, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1,  4, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1, 16, 8, 4, 4, 4) \
+  instantiate_gemv_t_bs_helper(name, itype, 1, 16, 4, 8, 4, 4) // clang-format on
 
 // clang-format off
 instantiate_gemv_t_bs_blocks(float32, float);

mlx/backend/metal/kernels/gemv_masked.h

@@ -0,0 +1,819 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include "mlx/backend/metal/kernels/steel/utils.h"
+
+using namespace metal;
+
+#define MLX_MTL_CONST static constant constexpr const
+#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
+
+struct _NoMask {
+  char x;
+
+  constexpr METAL_FUNC operator bool() {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const threadgroup {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const device {
+    return true;
+  }
+  constexpr METAL_FUNC operator bool() const constant {
+    return true;
+  }
+};
+
+typedef struct _NoMask nomask_t;
+
+template <typename OutT, typename InT = OutT>
+struct ScaleOp {
+  OutT scale;
+
+  METAL_FUNC OutT apply(InT x) const {
+    return static_cast<OutT>(x) * scale;
+  }
+};
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+struct GEMVKernel {
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  static_assert(
+      SN == 8 || SN == 16 || SN == 32,
+      "gemv block must have a width of 8, 16, or 32");
+
+  static_assert(blockN >= blockM, "Masked gemv must have blockN >= blockM");
+
+  MLX_MTL_CONST bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  MLX_MTL_CONST bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  MLX_MTL_CONST bool has_mul_operand_mask =
+      has_operand_mask && !metal::is_same_v<op_mask_t, bool>;
+  MLX_MTL_CONST bool has_mul_output_mask =
+      has_output_mask && !metal::is_same_v<out_mask_t, bool>;
+
+  // - The matrix of size (M = out_vec_size, K = in_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then multiplies and adds to accumulate its local result for
+  //    the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated blockM outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BN > 1 ? BN*(blockM + TM) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BN > 1;
+
+  static METAL_FUNC void
+  load_unsafe(const device T* src, thread T dst[TN], const int src_offset = 0) {
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      dst[tn] = src[src_offset + tn];
+    }
+  }
+
+  static METAL_FUNC void load_safe(
+      const device T* src,
+      thread T dst[TN],
+      const int src_offset = 0,
+      const int src_size = TN) {
+    if (src_offset + TN <= src_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src[src_offset + tn];
+      }
+    } else { // Edgecase
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        dst[tn] = src_offset + tn < src_size ? src[src_offset + tn] : 0;
+      }
+    }
+  }
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& matrix_ld [[buffer(6)]],
+      const device out_mask_t* out_mask [[buffer(20)]],
+      const device op_mask_t* mat_mask [[buffer(21)]],
+      const device op_mask_t* vec_mask [[buffer(22)]],
+      const constant int* mask_strides [[buffer(23)]],
+      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]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    thread T result[TM] = {0};
+    thread T inter[TN];
+    thread T v_coeff[TN];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = BN != 1 ? SM * (simd_gid / BN) : int(SM * simd_gid);
+    const int simdN = BN != 1 ? SN * (simd_gid % BN) : 0;
+
+    int bm = (simdM + thrM) * TM;
+    int bn = (simdN + thrN) * TN;
+
+    // Block position
+    int out_row = tid.x * blockM + bm;
+
+    // Exit simdgroup if rows out of bound
+    if (out_row >= out_vec_size)
+      return;
+
+    // Adjust tail simdgroup to ensure in bound reads
+    out_row = out_row + TM <= out_vec_size ? out_row : out_vec_size - TM;
+
+    // Prepare mask offsets
+    const constant int* out_mask_strides = mask_strides;
+    const constant int* mat_mask_strides =
+        mask_strides + (has_output_mask ? 2 : 0);
+    const constant int* vec_mask_strides =
+        mat_mask_strides + (has_operand_mask ? 2 : 0);
+
+    const int m_block_idx = blockN > blockM ? out_row / blockN : int(tid.x);
+
+    const int out_mask_offset =
+        !has_output_mask ? 0 : m_block_idx * out_mask_strides[1];
+
+    int mat_mask_offset =
+        !has_operand_mask ? 0 : m_block_idx * mat_mask_strides[1];
+    int vec_mask_offset = 0;
+    const int mat_mask_step = !has_operand_mask ? 0 : mat_mask_strides[0];
+    const int vec_mask_step = !has_operand_mask ? 0 : vec_mask_strides[1];
+
+    T out_scale{1};
+
+    // Check output mask
+    if (has_output_mask) {
+      auto mask_out = out_mask[out_mask_offset];
+
+      // Write zeros and return if mask is 0
+      if (!mask_out) {
+        if (simdN == 0 && thrN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            out_vec[out_row + tm] = T(0.);
+          }
+        }
+
+        return;
+      }
+
+      // Store scalar if multiplicative mask
+      if (has_mul_output_mask) {
+        out_scale = T(mask_out);
+      }
+    }
+
+    // Advance matrix
+    mat += out_row * matrix_ld;
+
+    // Prepare for loop
+    constexpr const uniform<int> loop_stride = make_uniform(blockN);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Loop over in_vec in blocks of blockN
+    for (int i = 0; i < n_iter; ++i) {
+      if (!has_operand_mask ||
+          (bool(mat_mask[mat_mask_offset]) &&
+           bool(vec_mask[vec_mask_offset]))) {
+        T block_scale{1};
+        if (has_mul_operand_mask) {
+          block_scale =
+              T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+        }
+
+        load_unsafe(in_vec, v_coeff, bn);
+
+        // Apply scale
+        if (has_mul_operand_mask) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            v_coeff[tn] *= block_scale;
+          }
+        }
+
+        // Per thread work loop
+        int mat_offset = 0;
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          // Load for the row
+          load_unsafe(mat, inter, mat_offset + bn);
+
+          // Accumulate results
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            result[tm] += inter[tn] * v_coeff[tn];
+          }
+
+          mat_offset += matrix_ld;
+        }
+      }
+
+      bn += blockN;
+      mat_mask_offset += mat_mask_step;
+      vec_mask_offset += vec_mask_step;
+    }
+
+    if (leftover > 0 &&
+        (!has_operand_mask ||
+         (bool(mat_mask[mat_mask_offset]) &&
+          bool(vec_mask[vec_mask_offset])))) {
+      T block_scale{1};
+      if (has_mul_operand_mask) {
+        block_scale =
+            T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+      }
+
+      load_safe(in_vec, v_coeff, bn, in_size);
+
+      // Apply scale
+      if (has_mul_operand_mask) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          v_coeff[tn] *= block_scale;
+        }
+      }
+
+      // Per thread work loop
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        // Load for the row
+        load_safe(&mat[tm * matrix_ld], inter, bn, in_size);
+
+        // Accumulate results
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          result[tm] += inter[tn] * v_coeff[tn];
+        }
+      }
+    }
+
+    // Apply out scale
+    if (has_mul_output_mask) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        result[tm] *= out_scale;
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tm = 0; tm < TM; tm++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sn = (SN / 2); sn >= 1; sn >>= 1) {
+        result[tm] += simd_shuffle_down(result[tm], sn);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup T* tgp_results = tgp_memory + sgN * (blockM + TM) + bm;
+      if (thrN == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tm = 0; tm < TM; tm++) {
+          tgp_results[tm] = result[tm];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgN == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgn = 1; sgn < BN; sgn++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              result[tm] += tgp_results[sgn * (blockM + TM) + tm];
+            }
+          }
+        }
+      }
+    }
+
+    // Write outputs
+    if (simdN == 0 && thrN == 0) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tm = 0; tm < TM; tm++) {
+        out_vec[out_row + tm] = result[tm];
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN> /* Thread cols (in elements) */
+struct GEMVTKernel {
+  MLX_MTL_CONST int threadsM = BM * SM;
+  MLX_MTL_CONST int threadsN = BN * SN;
+
+  MLX_MTL_CONST int blockM = threadsM * TM;
+  MLX_MTL_CONST int blockN = threadsN * TN;
+
+  static_assert(SM * SN == 32, "simdgroup can only have 32 threads");
+
+  MLX_MTL_CONST bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  MLX_MTL_CONST bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  MLX_MTL_CONST bool has_mul_operand_mask =
+      has_operand_mask && !metal::is_same_v<op_mask_t, bool>;
+  MLX_MTL_CONST bool has_mul_output_mask =
+      has_output_mask && !metal::is_same_v<out_mask_t, bool>;
+
+  // - The matrix of size (M = in_vec_size, N = out_vec_size) is divided up
+  //   into blocks of (blockM, blockN) divided among threadgroups
+  // - Every thread works on a block of (TM, TN)
+  // - We assume each threadgroup has (threadsN, threadsM, 1) threads
+  //
+  // 1. A thread loads TN elements each from mat along TM contiguous rows
+  //    and the corresponding scalar from the vector
+  // 2. The thread then accumulates its local result for the block
+  // 3. At the end, each thread has accumulated results over all blocks across
+  //    the rows. These are then summed up across the threadgroup
+  // 4. Each threadgroup writes its accumulated BN * TN outputs
+  //
+  // Edge case handling:
+  // - The threadgroup with the largest tid has blocks that exceed the matrix
+  //   * The blocks that start outside the matrix are never read (thread results
+  //     remain zero)
+  //   * The last thread that partially overlaps with the matrix is shifted
+  //     inwards such that the thread block fits exactly in the matrix
+
+  MLX_MTL_CONST short tgp_mem_size = BM > 1 ? BM*(blockN + TN) : 0;
+  MLX_MTL_CONST bool needs_tgp_reduction = BM > 1;
+
+  static METAL_FUNC void run(
+      const device T* mat [[buffer(0)]],
+      const device T* in_vec [[buffer(1)]],
+      device T* out_vec [[buffer(3)]],
+      const constant int& in_vec_size [[buffer(4)]],
+      const constant int& out_vec_size [[buffer(5)]],
+      const constant int& marix_ld [[buffer(6)]],
+      const device out_mask_t* out_mask [[buffer(20)]],
+      const device op_mask_t* mat_mask [[buffer(21)]],
+      const device op_mask_t* vec_mask [[buffer(22)]],
+      const constant int* mask_strides [[buffer(23)]],
+      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]]) {
+    // Appease compiler
+    (void)lid;
+
+    // Thread local accumulation results
+    T result[TN] = {0};
+    T inter[TN];
+    T v_coeff[TM];
+
+    const int thrM = SN != 32 ? simd_lid / SN : 0;
+    const int thrN = SN != 32 ? simd_lid % SN : int(simd_lid);
+
+    const int sgM = BN != 1 ? (simd_gid / BN) : int(simd_gid);
+    const int sgN = BN != 1 ? (simd_gid % BN) : 0;
+
+    const int simdM = SM * sgM;
+    const int simdN = SN * sgN;
+
+    int cm = (simdM + thrM);
+    int cn = (simdN + thrN);
+
+    int bm = cm * TM;
+    int bn = cn * TN;
+
+    int out_col = tid.x * blockN + bn;
+
+    // Prepare mask offsets
+    const constant int* out_mask_strides = mask_strides;
+    const constant int* mat_mask_strides =
+        out_mask_strides + (has_output_mask ? 2 : 0);
+    const constant int* vec_mask_strides =
+        mat_mask_strides + (has_operand_mask ? 2 : 0);
+
+    const int n_block_idx = blockM > blockN ? out_col / blockM : int(tid.x);
+
+    const int out_mask_offset =
+        !has_output_mask ? 0 : n_block_idx; // * out_mask_strides[0];
+
+    int mat_mask_offset =
+        !has_operand_mask ? 0 : n_block_idx * mat_mask_strides[0];
+    int vec_mask_offset = 0;
+    const int mat_mask_step = !has_operand_mask ? 0 : mat_mask_strides[1];
+    const int vec_mask_step = !has_operand_mask ? 0 : vec_mask_strides[0];
+
+    T out_scale{1};
+
+    // Check output mask
+    if (has_output_mask) {
+      auto mask_out = out_mask[out_mask_offset];
+
+      // Write zeros and return if mask is 0
+      if (!mask_out) {
+        if (cm == 0 && out_col < out_vec_size) {
+          if (out_col + TN <= out_vec_size) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              out_vec[out_col + tn] = T(0.);
+            }
+          } else {
+            for (int tn = 0; tn < TN && (out_col + tn) < out_vec_size; tn++) {
+              out_vec[out_col + tn] = T(0.);
+            }
+          }
+        }
+
+        return;
+      }
+
+      // Store scalar if multiplicative mask
+      if (has_mul_output_mask) {
+        out_scale = T(mask_out);
+      }
+    }
+
+    // Prepare for loop
+    constexpr const uniform<int> loop_stride = make_uniform(blockM);
+    const uniform<int> in_size = make_uniform(in_vec_size);
+    const uniform<int> n_iter = in_size / loop_stride;
+    const uniform<int> last_iter = loop_stride * n_iter;
+    const uniform<int> leftover = in_size - last_iter;
+
+    // Edgecase handling
+    if (out_col < out_vec_size) {
+      out_col = (out_col + TN) <= out_vec_size ? out_col : out_vec_size - TN;
+
+      // Per thread accumulation main loop
+      for (int i = 0; i < n_iter; ++i) {
+        // Adding a threadgroup_barrier improves performance slightly
+        // This is possibly it may help exploit cache better
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (!has_operand_mask ||
+            (bool(mat_mask[mat_mask_offset]) &&
+             bool(vec_mask[vec_mask_offset]))) {
+          T block_scale{1};
+          if (has_mul_operand_mask) {
+            block_scale =
+                T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            v_coeff[tm] = in_vec[bm + tm];
+          }
+
+          // Apply scale
+          if (has_mul_operand_mask) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tm = 0; tm < TM; tm++) {
+              v_coeff[tm] *= block_scale;
+            }
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tm = 0; tm < TM; tm++) {
+            for (int tn = 0; tn < TN; tn++) {
+              inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+            }
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += v_coeff[tm] * inter[tn];
+            }
+          }
+        }
+
+        bm += blockM;
+        mat_mask_offset += mat_mask_step;
+        vec_mask_offset += vec_mask_step;
+      }
+
+      if (leftover > 0 &&
+          (!has_operand_mask ||
+           (bool(mat_mask[mat_mask_offset]) &&
+            bool(vec_mask[vec_mask_offset])))) {
+        T block_scale{1};
+        if (has_mul_operand_mask) {
+          block_scale =
+              T(mat_mask[mat_mask_offset]) * T(vec_mask[vec_mask_offset]);
+        }
+
+        for (int tm = 0; tm < TM && bm + tm < in_vec_size; tm++) {
+          v_coeff[tm] = in_vec[bm + tm];
+
+          if (has_mul_operand_mask) {
+            v_coeff[tm] *= block_scale;
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            inter[tn] = mat[(bm + tm) * marix_ld + out_col + tn];
+          }
+
+          MLX_MTL_PRAGMA_UNROLL
+          for (int tn = 0; tn < TN; tn++) {
+            result[tn] += v_coeff[tm] * inter[tn];
+          }
+        }
+      }
+    }
+
+    // Apply out scale
+    if (has_mul_output_mask) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int tn = 0; tn < TN; tn++) {
+        result[tn] *= out_scale;
+      }
+    }
+
+    // Simdgroup accumulations
+    MLX_MTL_PRAGMA_UNROLL
+    for (int tn = 0; tn < TN; tn++) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (ushort sm = (SM / 2); sm >= 1; sm >>= 1) {
+        result[tn] += simd_shuffle_down(result[tn], SN * sm);
+      }
+    }
+
+    // Threadgroup accumulation results
+    if (needs_tgp_reduction) {
+      threadgroup T* tgp_results = tgp_memory + sgM * (blockN + TN) + bn;
+      if (thrM == 0) {
+        MLX_MTL_PRAGMA_UNROLL
+        for (int tn = 0; tn < TN; tn++) {
+          tgp_results[tn] = result[tn];
+        }
+
+        threadgroup_barrier(mem_flags::mem_none);
+
+        if (sgM == 0) {
+          MLX_MTL_PRAGMA_UNROLL
+          for (int sgm = 1; sgm < BM; sgm++) {
+            MLX_MTL_PRAGMA_UNROLL
+            for (int tn = 0; tn < TN; tn++) {
+              result[tn] += tgp_results[sgm * (blockN + TN) + tn];
+            }
+          }
+        }
+      }
+    }
+
+    // Threadgroup accumulation and writing out results
+    if (cm == 0 && out_col < out_vec_size) {
+      MLX_MTL_PRAGMA_UNROLL
+      for (int j = 0; j < TN; j++) {
+        out_vec[out_col + j] = result[j];
+      }
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+/// Matrix vector multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch> /* Batch ndim > 1 */
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_masked(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant size_t* vector_batch_stride [[buffer(11)]],
+    const constant size_t* matrix_batch_stride [[buffer(12)]],
+    const device out_mask_t* out_mask [[buffer(20)]],
+    const device op_mask_t* mat_mask [[buffer(21)]],
+    const device op_mask_t* vec_mask [[buffer(22)]],
+    const constant int* mask_strides [[buffer(23)]],
+    const constant size_t* mask_batch_strides [[buffer(24)]],
+    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 gemv_kernel =
+      GEMVKernel<T, out_mask_t, op_mask_t, BM, BN, SM, SN, TM, TN>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  constexpr bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  constexpr bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (has_output_mask) {
+      out_mask +=
+          elem_to_loc(tid.z, batch_shape, mask_batch_strides, batch_ndim);
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      const constant size_t* mask_strides_mat = mask_batch_strides;
+      const constant size_t* mask_strides_vec = mask_strides_mat + batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z, batch_shape, mask_strides_mat, mask_strides_vec, batch_ndim);
+
+      mat_mask += batch_offsets.x;
+      vec_mask += batch_offsets.y;
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (has_output_mask) {
+      out_mask += tid.z * mask_batch_strides[0];
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      mat_mask += tid.z * mask_batch_strides[0];
+      vec_mask += tid.z * mask_batch_strides[batch_ndim];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      out_mask,
+      mat_mask,
+      vec_mask,
+      mask_strides,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// Vector matrix multiplication
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+    typename T,
+    typename out_mask_t,
+    typename op_mask_t,
+    const int BM, /* Threadgroup rows (in simdgroups) */
+    const int BN, /* Threadgroup cols (in simdgroups) */
+    const int SM, /* Simdgroup rows (in threads) */
+    const int SN, /* Simdgroup cols (in threads) */
+    const int TM, /* Thread rows (in elements) */
+    const int TN, /* Thread cols (in elements) */
+    const bool kDoNCBatch> /* Batch ndim > 1 */
+[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_masked(
+    const device T* mat [[buffer(0)]],
+    const device T* in_vec [[buffer(1)]],
+    device T* out_vec [[buffer(3)]],
+    const constant int& in_vec_size [[buffer(4)]],
+    const constant int& out_vec_size [[buffer(5)]],
+    const constant int& marix_ld [[buffer(6)]],
+    const constant int& batch_ndim [[buffer(9)]],
+    const constant int* batch_shape [[buffer(10)]],
+    const constant size_t* vector_batch_stride [[buffer(11)]],
+    const constant size_t* matrix_batch_stride [[buffer(12)]],
+    const device out_mask_t* out_mask [[buffer(20)]],
+    const device op_mask_t* mat_mask [[buffer(21)]],
+    const device op_mask_t* vec_mask [[buffer(22)]],
+    const constant int* mask_strides [[buffer(23)]],
+    const constant size_t* mask_batch_strides [[buffer(24)]],
+    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 gemv_kernel =
+      GEMVTKernel<T, out_mask_t, op_mask_t, BM, BN, SM, SN, TM, TN>;
+  threadgroup T tgp_memory
+      [gemv_kernel::tgp_mem_size == 0 ? 1 : gemv_kernel::tgp_mem_size];
+
+  constexpr bool has_operand_mask = !metal::is_same_v<op_mask_t, nomask_t>;
+  constexpr bool has_output_mask = !metal::is_same_v<out_mask_t, nomask_t>;
+
+  // Update batch offsets
+  if (kDoNCBatch) {
+    in_vec += elem_to_loc(tid.z, batch_shape, vector_batch_stride, batch_ndim);
+    mat += elem_to_loc(tid.z, batch_shape, matrix_batch_stride, batch_ndim);
+
+    if (has_output_mask) {
+      out_mask +=
+          elem_to_loc(tid.z, batch_shape, mask_batch_strides, batch_ndim);
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      const constant size_t* mask_strides_mat = mask_batch_strides;
+      const constant size_t* mask_strides_vec = mask_strides_mat + batch_ndim;
+
+      ulong2 batch_offsets = elem_to_loc_broadcast(
+          tid.z, batch_shape, mask_strides_mat, mask_strides_vec, batch_ndim);
+
+      mat_mask += batch_offsets.x;
+      vec_mask += batch_offsets.y;
+    }
+
+  } else {
+    in_vec += tid.z * vector_batch_stride[0];
+    mat += tid.z * matrix_batch_stride[0];
+
+    if (has_output_mask) {
+      out_mask += tid.z * mask_batch_strides[0];
+      mask_batch_strides += batch_ndim;
+    }
+
+    if (has_operand_mask) {
+      mat_mask += tid.z * mask_batch_strides[0];
+      vec_mask += tid.z * mask_batch_strides[batch_ndim];
+    }
+  }
+
+  out_vec += tid.z * out_vec_size;
+
+  gemv_kernel::run(
+      mat,
+      in_vec,
+      out_vec,
+      in_vec_size,
+      out_vec_size,
+      marix_ld,
+      out_mask,
+      mat_mask,
+      vec_mask,
+      mask_strides,
+      gemv_kernel::tgp_mem_size == 0 ? nullptr : tgp_memory,
+      tid,
+      lid,
+      simd_gid,
+      simd_lid);
+}

mlx/backend/metal/kernels/gemv_masked.metal

@@ -0,0 +1,114 @@
+// Copyright © 2023-2024 Apple Inc.
+
+// clang-format off
+#include <metal_simdgroup>
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/bf16.h"
+#include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/gemv_masked.h"
+
+#define instantiate_gemv_helper(                                           \
+    outm_n, outm_t, opm_n, opm_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  template [[host_name("gemv_outmask_" #outm_n "_opmask_" #opm_n "_" #name \
+                       "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm   \
+                       "_tn" #tn "_nc" #nc)]] [[kernel]] void              \
+  gemv_masked<itype, outm_t, opm_t, bm, bn, sm, sn, tm, tn, nc>(           \
+      const device itype* mat [[buffer(0)]],                               \
+      const device itype* in_vec [[buffer(1)]],                            \
+      device itype* out_vec [[buffer(3)]],                                 \
+      const constant int& in_vec_size [[buffer(4)]],                       \
+      const constant int& out_vec_size [[buffer(5)]],                      \
+      const constant int& marix_ld [[buffer(6)]],                          \
+      const constant int& batch_ndim [[buffer(9)]],                        \
+      const constant int* batch_shape [[buffer(10)]],                      \
+      const constant size_t* vector_batch_stride [[buffer(11)]],           \
+      const constant size_t* matrix_batch_stride [[buffer(12)]],           \
+      const device outm_t* out_mask [[buffer(20)]],                        \
+      const device opm_t* mat_mask [[buffer(21)]],                         \
+      const device opm_t* vec_mask [[buffer(22)]],                         \
+      const constant int* mask_strides [[buffer(23)]],                     \
+      const constant size_t* mask_batch_strides [[buffer(24)]],            \
+      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_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(bool_, bool, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(name, itype, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(bool_, bool, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(name, itype, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_helper(nomask, nomask_t, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(nomask, nomask_t, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(bool_, bool, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_helper(name, itype, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv(name, itype, bm, bn, sm, sn, tm, tn)   \
+  instantiate_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, 0) \
+  instantiate_gemv_base(name, itype, bm, bn, sm, sn, tm, tn, 1)
+
+#define instantiate_gemv_blocks(name, itype) \
+  instantiate_gemv(name, itype, 2, 1, 4,  8, 1, 4) \
+  instantiate_gemv(name, itype, 2, 1, 4,  8, 4, 4) \
+  instantiate_gemv(name, itype, 2, 1, 2, 16, 1, 4) \
+  instantiate_gemv(name, itype, 2, 1, 2, 16, 4, 4) \
+  instantiate_gemv(name, itype, 4, 1, 2, 16, 4, 4)
+
+instantiate_gemv_blocks(float32, float);
+instantiate_gemv_blocks(float16, half);
+instantiate_gemv_blocks(bfloat16, bfloat16_t);
+
+#define instantiate_gemv_t_helper(                                           \
+    outm_n, outm_t, opm_n, opm_t, name, itype, bm, bn, sm, sn, tm, tn, nc)   \
+  template [[host_name("gemv_t_outmask_" #outm_n "_opmask_" #opm_n "_" #name \
+                       "_bm" #bm "_bn" #bn "_sm" #sm "_sn" #sn "_tm" #tm     \
+                       "_tn" #tn "_nc" #nc)]] [[kernel]] void                \
+  gemv_t_masked<itype, outm_t, opm_t, bm, bn, sm, sn, tm, tn, nc>(           \
+      const device itype* mat [[buffer(0)]],                                 \
+      const device itype* in_vec [[buffer(1)]],                              \
+      device itype* out_vec [[buffer(3)]],                                   \
+      const constant int& in_vec_size [[buffer(4)]],                         \
+      const constant int& out_vec_size [[buffer(5)]],                        \
+      const constant int& marix_ld [[buffer(6)]],                            \
+      const constant int& batch_ndim [[buffer(9)]],                          \
+      const constant int* batch_shape [[buffer(10)]],                        \
+      const constant size_t* vector_batch_stride [[buffer(11)]],             \
+      const constant size_t* matrix_batch_stride [[buffer(12)]],             \
+      const device outm_t* out_mask [[buffer(20)]],                          \
+      const device opm_t* mat_mask [[buffer(21)]],                           \
+      const device opm_t* vec_mask [[buffer(22)]],                           \
+      const constant int* mask_strides [[buffer(23)]],                       \
+      const constant size_t* mask_batch_strides [[buffer(24)]],              \
+      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_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(bool_, bool, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(name, itype, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(bool_, bool, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(name, itype, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc)      \
+  instantiate_gemv_t_helper(nomask, nomask_t, name, itype, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(nomask, nomask_t, bool_, bool, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(bool_, bool, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc) \
+  instantiate_gemv_t_helper(name, itype, nomask, nomask_t, name, itype, bm, bn, sm, sn, tm, tn, nc)
+
+#define instantiate_gemv_t(name, itype, bm, bn, sm, sn, tm, tn)   \
+  instantiate_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, 0) \
+  instantiate_gemv_t_base(name, itype, bm, bn, sm, sn, tm, tn, 1)
+
+#define instantiate_gemv_t_blocks(name, itype) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 4, 1) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 4, 4) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 8, 1) \
+  instantiate_gemv_t(name, itype, 1, 1,  8, 4, 8, 4) \
+  instantiate_gemv_t(name, itype, 1, 2,  8, 4, 8, 4) \
+  instantiate_gemv_t(name, itype, 1, 4,  8, 4, 8, 4)
+
+instantiate_gemv_t_blocks(float32, float);
+instantiate_gemv_t_blocks(float16, half);
+instantiate_gemv_t_blocks(bfloat16, bfloat16_t); // clang-format on

mlx/backend/metal/kernels/hadamard.h

@@ -0,0 +1,167 @@
+// Copyright © 2024 Apple Inc.
+#include <metal_common>
+#include <metal_compute>
+
+#include "mlx/backend/metal/kernels/steel/defines.h"
+
+using namespace metal;
+
+// Thread local Hadamard transform for 2^R
+template <short R>
+METAL_FUNC void radix_func(thread float* x) {
+  constexpr short logR = __builtin_ctz(R);
+  short h = 1;
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < logR; s++) {
+    STEEL_PRAGMA_UNROLL
+    for (short i = 0; i < R / 2; i++) {
+      short k = i & (h - 1);
+      short j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width>
+[[kernel]] void hadamard_n(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size N = 2^k
+  //
+  // Equivalent to:
+  //    from scipy.linalg import hadamard
+  //    y = hadamard(len(x)) @ x
+
+  constexpr short num_threads = N / max_radix;
+  constexpr short logN = __builtin_ctz(N);
+  constexpr short logR = __builtin_ctz(max_radix);
+  constexpr short num_steps = logN / logR;
+  constexpr short logFinal = logN % logR;
+  constexpr short final_radix = 1 << (logFinal);
+
+  int batch_idx = elem.x * N;
+  short i = elem.y;
+
+  threadgroup T buf[N];
+
+  // Read values from device
+  STEEL_PRAGMA_UNROLL
+  for (short j = 0; j < max_radix / read_width; j++) {
+    short index = j * read_width * num_threads + i * read_width;
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      buf[index + r] = in[batch_idx + index + r];
+    }
+  }
+
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  float x[max_radix];
+  short h = 1;
+
+  STEEL_PRAGMA_UNROLL
+  for (short s = 0; s < num_steps; s++) {
+    short k = i & (h - 1);
+    short j = ((i - k) << logR) + k;
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      x[r] = buf[j + h * r];
+    }
+
+    radix_func<max_radix>(x);
+
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < max_radix; r++) {
+      buf[j + h * r] = T(x[r]);
+    }
+
+    h <<= logR;
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Do the final radix
+  // e.g. max_radix = 16
+  //      N = 1024 = 16 * 16 * 4
+  if (final_radix > 1) {
+    // Each thread does multiple butterflies
+    STEEL_PRAGMA_UNROLL
+    for (int t = 0; t < max_radix / final_radix; t++) {
+      short index = i + t * num_threads;
+      short k = index & (h - 1);
+      short j = ((index - k) << logFinal) + k;
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        x[r] = buf[j + h * r];
+      }
+
+      radix_func<final_radix>(x);
+
+      STEEL_PRAGMA_UNROLL
+      for (short r = 0; r < final_radix; r++) {
+        buf[j + h * r] = T(x[r]);
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+
+  // Write values to device
+  STEEL_PRAGMA_UNROLL
+  for (short j = 0; j < max_radix / read_width; j++) {
+    short index = j * read_width * num_threads + i * read_width;
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      out[batch_idx + index + r] = T(buf[index + r] * scale);
+    }
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+[[kernel]] void hadamard_m(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const float& scale,
+    uint3 elem [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Compute a Hadamard transform of size M
+  // using a naive O(M^2) codelet.
+  //
+  // This kernel is the second stage in the computation
+  // of a Hadamard transform of size M*N where N = 2^k.
+
+  int index = elem.x * grid.y + elem.y;
+  short i = index % (N / read_width);
+  int batch_idx = index / (N / read_width) * M * N;
+
+  float x[read_width][M];
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      x[r][c] = in[batch_idx + c * N + i * read_width + r];
+    }
+  }
+
+  STEEL_PRAGMA_UNROLL
+  for (short r = 0; r < read_width; r++) {
+    // This function is JIT compiled for M
+    // using the Hadamard matrix strings in `metal/hadamard.cpp`
+    hadamard_radix_m(x[r]);
+  }
+
+  // Write back to device
+  STEEL_PRAGMA_UNROLL
+  for (short c = 0; c < M; c++) {
+    STEEL_PRAGMA_UNROLL
+    for (short r = 0; r < read_width; r++) {
+      out[batch_idx + c * N + i * read_width + r] = T(x[r][c] * scale);
+    }
+  }
+}

mlx/backend/metal/kernels/layer_norm.metal

@@ -34,7 +34,7 @@ template <typename T, int N_READS = RMS_N_READS>
   threadgroup float local_mean[1];
   threadgroup float local_normalizer[1];
 
-  x += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
   b += b_stride * lid * N_READS;
 
@@ -89,7 +89,7 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer = local_normalizer[0];
 
   // Write the outputs
-  out += gid * axis_size + lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       thread_x[i] = (thread_x[i] - mean) * normalizer;
@@ -131,7 +131,7 @@ template <typename T, int N_READS = RMS_N_READS>
   threadgroup float local_mean[1];
   threadgroup float local_normalizer[1];
 
-  x += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
   b += b_stride * lid * N_READS;
 
@@ -188,7 +188,7 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer = local_normalizer[0];
 
   // Write the outputs
-  out += gid * axis_size + lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
       for (int i = 0; i < N_READS; i++) {
@@ -223,8 +223,8 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   // Advance the input pointers
-  x += gid * axis_size + lid * N_READS;
-  g += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
 
   // Allocate registers for the computation and accumulators
@@ -321,8 +321,8 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer2 = normalizer * normalizer;
 
   // Write the outputs
-  gx += gid * axis_size + lid * N_READS;
-  gw += gid * axis_size + lid * N_READS;
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       thread_x[i] = (thread_x[i] - mean) * normalizer;
@@ -360,8 +360,8 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   // Advance the input pointers
-  x += gid * axis_size + lid * N_READS;
-  g += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
 
   // Allocate registers for the accumulators
@@ -457,8 +457,8 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer2 = normalizer * normalizer;
 
   // Write the outputs
-  gx += gid * axis_size + lid * N_READS;
-  gw += gid * axis_size + lid * N_READS;
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
       for (int i = 0; i < N_READS; i++) {

mlx/backend/metal/kernels/random.metal

@@ -43,20 +43,22 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
   auto half_size = grid_dim.y - odd;
   out += index.x * bytes_per_key;
   bool drop_last = odd && (index.y == half_size);
-  auto count = uint2(index.y, drop_last ? 0 : index.y + grid_dim.y);
-  auto bits = threefry2x32_hash(key, count);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
   for (int i = 0; i < 4; ++i) {
-    out[4 * count.x + i] = bits.bytes[0][i];
+    out[idx + i] = bits.bytes[0][i];
   }
   if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
     if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
       int edge_bytes = (bytes_per_key % 4);
       for (int i = 0; i < edge_bytes; ++i) {
-        out[4 * count.y + i] = bits.bytes[1][i];
+        out[idx + i] = bits.bytes[1][i];
       }
     } else {
       for (int i = 0; i < 4; ++i) {
-        out[4 * count.y + i] = bits.bytes[1][i];
+        out[idx + i] = bits.bytes[1][i];
       }
     }
   }
@@ -77,22 +79,24 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
   auto k2_elem = elem_to_loc(kidx + 1, key_shape, key_strides, ndim);
   auto key = uint2(keys[k1_elem], keys[k2_elem]);
   auto half_size = grid_dim.y - odd;
-  out += index.x * bytes_per_key;
+  out += size_t(index.x) * bytes_per_key;
   bool drop_last = odd && (index.y == half_size);
-  auto count = uint2(index.y, drop_last ? 0 : index.y + grid_dim.y);
-  auto bits = threefry2x32_hash(key, count);
+  auto bits = threefry2x32_hash(
+      key, uint2(index.y, drop_last ? 0 : index.y + grid_dim.y));
+  size_t idx = size_t(index.y) << 2;
   for (int i = 0; i < 4; ++i) {
-    out[4 * count.x + i] = bits.bytes[0][i];
+    out[idx + i] = bits.bytes[0][i];
   }
   if (!drop_last) {
+    idx = (drop_last ? 0 : size_t(index.y) + grid_dim.y) << 2;
     if ((index.y + 1) == half_size && (bytes_per_key % 4) > 0) {
       int edge_bytes = (bytes_per_key % 4);
       for (int i = 0; i < edge_bytes; ++i) {
-        out[4 * count.y + i] = bits.bytes[1][i];
+        out[idx + i] = bits.bytes[1][i];
       }
     } else {
       for (int i = 0; i < 4; ++i) {
-        out[4 * count.y + i] = bits.bytes[1][i];
+        out[idx + i] = bits.bytes[1][i];
       }
     }
   }

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

@@ -23,7 +23,7 @@ template <typename U = bool>
 struct And {
   bool simd_reduce(bool val) {
     return simd_all(val);
-  };
+  }
 
   static constexpr constant bool init = true;
 
@@ -61,7 +61,7 @@ template <typename U = bool>
 struct Or {
   bool simd_reduce(bool val) {
     return simd_any(val);
-  };
+  }
 
   static constexpr constant bool init = false;
 
@@ -100,7 +100,7 @@ struct Sum {
   template <typename T>
   T simd_reduce(T val) {
     return simd_sum(val);
-  };
+  }
 
   static constexpr constant U init = U(0);
 
@@ -120,7 +120,7 @@ struct Prod {
   template <typename T>
   T simd_reduce(T val) {
     return simd_product(val);
-  };
+  }
 
   static constexpr constant U init = U(1);
 
@@ -140,7 +140,7 @@ struct Min {
   template <typename T>
   T simd_reduce(T val) {
     return simd_min(val);
-  };
+  }
 
   static constexpr constant U init = Limits<U>::max;
 
@@ -160,7 +160,7 @@ struct Max {
   template <typename T>
   T simd_reduce(T val) {
     return simd_max(val);
-  };
+  }
 
   static constexpr constant U init = Limits<U>::min;
 

mlx/backend/metal/kernels/rms_norm.metal

@@ -24,7 +24,7 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   float acc = 0;
-  x += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
@@ -62,7 +62,7 @@ template <typename T, int N_READS = RMS_N_READS>
   threadgroup_barrier(mem_flags::mem_threadgroup);
 
   // Write the outputs
-  out += gid * axis_size + lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       out[i] = w[w_stride * i] * static_cast<T>(x[i] * local_inv_mean[0]);
@@ -92,7 +92,7 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   float acc = 0;
-  x += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
@@ -132,7 +132,7 @@ template <typename T, int N_READS = RMS_N_READS>
   threadgroup_barrier(mem_flags::mem_threadgroup);
 
   // Write the outputs
-  out += gid * axis_size + lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
       for (int i = 0; i < N_READS; i++) {
@@ -165,8 +165,8 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   // Advance the input pointers
-  x += gid * axis_size + lid * N_READS;
-  g += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
 
   // Allocate registers for the computation and accumulators
@@ -233,8 +233,8 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer3 = normalizer * normalizer * normalizer;
 
   // Write the outputs
-  gx += gid * axis_size + lid * N_READS;
-  gw += gid * axis_size + lid * N_READS;
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       gx[i] = static_cast<T>(
@@ -270,8 +270,8 @@ template <typename T, int N_READS = RMS_N_READS>
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
   // Advance the input pointers
-  x += gid * axis_size + lid * N_READS;
-  g += gid * axis_size + lid * N_READS;
+  x += gid * size_t(axis_size) + lid * N_READS;
+  g += gid * size_t(axis_size) + lid * N_READS;
   w += w_stride * lid * N_READS;
 
   // Allocate registers for the accumulators
@@ -337,8 +337,8 @@ template <typename T, int N_READS = RMS_N_READS>
   float normalizer3 = normalizer * normalizer * normalizer;
 
   // Write the outputs
-  gx += gid * axis_size + lid * N_READS;
-  gw += gid * axis_size + lid * N_READS;
+  gx += gid * size_t(axis_size) + lid * N_READS;
+  gw += gid * size_t(axis_size) + lid * N_READS;
   for (uint r = 0; r < axis_size; r += lsize * N_READS) {
     if (r + lid * N_READS + N_READS <= axis_size) {
       for (int i = 0; i < N_READS; i++) {

mlx/backend/metal/kernels/rope.metal

@@ -6,36 +6,17 @@
 #include "mlx/backend/metal/kernels/utils.h"
 
 template <typename T, bool traditional, bool forward>
-[[kernel]] void rope(
+[[kernel]] void rope_single(
     const device T* in [[buffer(0)]],
     device T* out [[buffer(1)]],
-    constant const size_t strides[3],
-    constant const size_t out_strides[3],
     constant const int& offset,
     constant const float& base,
     constant const float& scale,
-    uint3 pos [[thread_position_in_grid]],
-    uint3 grid [[threads_per_grid]]) {
-  // Compute the input and output indices
-  uint in_index_1, in_index_2;
-  uint out_index_1, out_index_2;
-  if (traditional) {
-    out_index_1 = 2 * pos.x * out_strides[2] + pos.y * out_strides[1] +
-        pos.z * out_strides[0];
-    out_index_2 = out_index_1 + 1;
-    in_index_1 =
-        2 * pos.x * strides[2] + pos.y * strides[1] + pos.z * strides[0];
-    in_index_2 = in_index_1 + strides[2];
-  } else {
-    out_index_1 = pos.x * out_strides[2] + pos.y * out_strides[1] +
-        pos.z * out_strides[0];
-    out_index_2 = out_index_1 + grid.x * out_strides[2];
-    in_index_1 = pos.x * strides[2] + pos.y * strides[1] + pos.z * strides[0];
-    in_index_2 = in_index_1 + grid.x * strides[2];
-  }
-
+    constant const size_t& stride,
+    uint2 pos [[thread_position_in_grid]],
+    uint2 grid [[threads_per_grid]]) {
   // Figure out L and d.
-  float L = scale * static_cast<float>(pos.y + offset);
+  float L = scale * static_cast<float>(offset);
   float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
 
   // Compute costheta, sintheta
@@ -43,6 +24,21 @@ template <typename T, bool traditional, bool forward>
   float costheta = metal::fast::cos(theta);
   float sintheta = metal::fast::sin(theta);
 
+  // Compute the input and output indices
+  uint in_index_1, in_index_2;
+  uint out_index_1, out_index_2;
+  if (traditional) {
+    out_index_1 = 2 * pos.x + pos.y * stride;
+    out_index_2 = out_index_1 + 1;
+    in_index_1 = 2 * pos.x + pos.y * stride;
+    in_index_2 = in_index_1 + 1;
+  } else {
+    out_index_1 = pos.x + pos.y * stride;
+    out_index_2 = out_index_1 + grid.x;
+    in_index_1 = pos.x + pos.y * stride;
+    in_index_2 = in_index_1 + grid.x;
+  }
+
   // Read and write the output
   float x1 = static_cast<float>(in[in_index_1]);
   float x2 = static_cast<float>(in[in_index_2]);
@@ -59,19 +55,97 @@ template <typename T, bool traditional, bool forward>
   out[out_index_2] = static_cast<T>(rx2);
 }
 
-#define instantiate_rope(name, type, traditional, forward) \
-  template [[host_name("rope_" #name)]] [[kernel]] void    \
-  rope<type, traditional, forward>(                        \
-      const device type* in [[buffer(0)]],                 \
-      device type* out [[buffer(1)]],                      \
-      constant const size_t strides[3],                    \
-      constant const size_t out_strides[3],                \
-      constant const int& offset,                          \
-      constant const float& base,                          \
-      constant const float& scale,                         \
-      uint3 pos [[thread_position_in_grid]],               \
+template <typename T, bool traditional, bool forward, int N = 4>
+[[kernel]] void rope(
+    const device T* in [[buffer(0)]],
+    device T* out [[buffer(1)]],
+    constant const int& offset,
+    constant const float& base,
+    constant const float& scale,
+    constant const size_t strides[3],
+    constant const size_t out_strides[3],
+    constant const size_t& n_batch,
+    uint3 pos [[thread_position_in_grid]],
+    uint3 grid [[threads_per_grid]]) {
+  // Figure out L and d.
+  float L = scale * static_cast<float>(pos.y + offset);
+  float d = static_cast<float>(pos.x) / static_cast<float>(grid.x);
+
+  // Compute costheta, sintheta
+  float theta = L * metal::exp2(-d * base);
+  float costheta = metal::fast::cos(theta);
+  float sintheta = metal::fast::sin(theta);
+
+  // Compute the input and output indices
+  size_t in_index_1, in_index_2;
+  size_t out_index_1, out_index_2;
+  if (traditional) {
+    out_index_1 = 2 * pos.x * out_strides[2] + pos.y * out_strides[1] +
+        N * pos.z * out_strides[0];
+    out_index_2 = out_index_1 + 1;
+    in_index_1 =
+        2 * pos.x * strides[2] + pos.y * strides[1] + N * pos.z * strides[0];
+    in_index_2 = in_index_1 + strides[2];
+  } else {
+    out_index_1 = pos.x * out_strides[2] + pos.y * out_strides[1] +
+        N * pos.z * out_strides[0];
+    out_index_2 = out_index_1 + grid.x * out_strides[2];
+    in_index_1 =
+        pos.x * strides[2] + pos.y * strides[1] + N * pos.z * strides[0];
+    in_index_2 = in_index_1 + grid.x * strides[2];
+  }
+  for (int i = 0; i < N && pos.z * N + i < n_batch; ++i) {
+    // Read and write the output
+    float x1 = static_cast<float>(in[in_index_1]);
+    float x2 = static_cast<float>(in[in_index_2]);
+    float rx1;
+    float rx2;
+    if (forward) {
+      rx1 = x1 * costheta - x2 * sintheta;
+      rx2 = x1 * sintheta + x2 * costheta;
+    } else {
+      rx1 = x2 * sintheta + x1 * costheta;
+      rx2 = x2 * costheta - x1 * sintheta;
+    }
+    out[out_index_1] = static_cast<T>(rx1);
+    out[out_index_2] = static_cast<T>(rx2);
+    in_index_1 += strides[0];
+    in_index_2 += strides[0];
+    out_index_1 += out_strides[0];
+    out_index_2 += out_strides[0];
+  }
+}
+
+#define instantiate_rope_g(name, type, traditional, forward) \
+  template [[host_name("rope_" #name)]] [[kernel]] void      \
+  rope<type, traditional, forward>(                          \
+      const device type* in [[buffer(0)]],                   \
+      device type* out [[buffer(1)]],                        \
+      constant const int& offset,                            \
+      constant const float& base,                            \
+      constant const float& scale,                           \
+      constant const size_t strides[3],                      \
+      constant const size_t out_strides[3],                  \
+      constant const size_t& n_batch,                        \
+      uint3 pos [[thread_position_in_grid]],                 \
       uint3 grid [[threads_per_grid]]);
 
+#define instantiate_rope_s(name, type, traditional, forward)   \
+  template [[host_name("rope_single_" #name)]] [[kernel]] void \
+  rope_single<type, traditional, forward>(                     \
+      const device type* in [[buffer(0)]],                     \
+      device type* out [[buffer(1)]],                          \
+      constant const int& offset,                              \
+      constant const float& base,                              \
+      constant const float& scale,                             \
+      constant const size_t& stride,                           \
+      uint2 pos [[thread_position_in_grid]],                   \
+      uint2 grid [[threads_per_grid]]);
+
+#define instantiate_rope(name, type, traditional, forward) \
+  instantiate_rope_s(name, type, traditional, forward)     \
+      instantiate_rope_g(name, type, traditional, forward)
+
 // clang-format off
 instantiate_rope(traditional_float16, half, true, true)
 instantiate_rope(traditional_bfloat16, bfloat16_t, true, true)
@@ -84,4 +158,4 @@ instantiate_rope(vjp_traditional_bfloat16, bfloat16_t, true, false)
 instantiate_rope(vjp_traditional_float32, float, true, false)
 instantiate_rope(vjp_float16, half, false, false)
 instantiate_rope(vjp_bfloat16, bfloat16_t, false, false)
-instantiate_rope(vjp_float32, float, false, false) // clang-format on
+instantiate_rope(vjp_float32, float, false, false) // clang-format on

mlx/backend/metal/kernels/scatter.h

@@ -10,7 +10,10 @@ METAL_FUNC void scatter_1d_index_impl(
     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 size_t& out_ndim [[buffer(5)]],
+    const constant int* upd_shape [[buffer(6)]],
+    const constant size_t& upd_ndim [[buffer(7)]],
+    const constant size_t& upd_size [[buffer(8)]],
     const thread array<const device IdxT*, NIDX>& idx_buffers,
     uint2 gid [[thread_position_in_grid]]) {
   Op op;
@@ -21,7 +24,14 @@ METAL_FUNC void scatter_1d_index_impl(
     out_idx += idx_val * out_strides[i];
   }
 
-  op.atomic_update(out, updates[gid.y * upd_size + gid.x], out_idx + gid.x);
+  if (upd_ndim > 1) {
+    auto out_offset = elem_to_loc(gid.x, upd_shape + 1, out_strides, out_ndim);
+    out_idx += out_offset;
+  } else {
+    out_idx += gid.x;
+  }
+
+  op.atomic_update(out, updates[gid.y * upd_size + gid.x], out_idx);
 }
 
 template <typename T, typename IdxT, typename Op, int NIDX>

mlx/backend/metal/kernels/softmax.h

@@ -25,7 +25,7 @@ template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
 
   AccT ld[N_READS];
 
-  in += gid * axis_size + lid * N_READS;
+  in += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       ld[i] = AccT(in[i]);
@@ -83,7 +83,7 @@ template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
   normalizer = 1 / local_normalizer[0];
 
   // Normalize and write to the output
-  out += gid * axis_size + lid * N_READS;
+  out += gid * size_t(axis_size) + lid * N_READS;
   if (lid * N_READS + N_READS <= axis_size) {
     for (int i = 0; i < N_READS; i++) {
       out[i] = T(ld[i] * normalizer);
@@ -107,7 +107,7 @@ template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
     uint lsize [[threads_per_threadgroup]],
     uint simd_lane_id [[thread_index_in_simdgroup]],
     uint simd_group_id [[simdgroup_index_in_threadgroup]]) {
-  in += gid * axis_size;
+  in += gid * size_t(axis_size);
 
   constexpr int SIMD_SIZE = 32;
 
@@ -170,7 +170,7 @@ template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
 
   // Finally given the normalizer and max value we can directly write the
   // softmax output
-  out += gid * axis_size;
+  out += gid * size_t(axis_size);
   for (int r = 0; r < static_cast<int>(ceildiv(axis_size, N_READS * lsize));
        r++) {
     int offset = r * lsize * N_READS + lid * N_READS;

mlx/backend/metal/kernels/sort.h

@@ -235,19 +235,21 @@ struct KernelMergeSort {
       const device T* inp,
       device U* out,
       const constant int& size_sorted_axis,
-      const constant int& stride_sorted_axis,
-      const constant int& stride_segment_axis,
+      const constant int& in_stride_sorted_axis,
+      const constant int& out_stride_sorted_axis,
+      const constant int& in_stride_segment_axis,
+      const constant int& out_stride_segment_axis,
       threadgroup val_t* tgp_vals,
       threadgroup idx_t* tgp_idxs,
       uint3 tid [[threadgroup_position_in_grid]],
       uint3 lid [[thread_position_in_threadgroup]]) {
     // tid.y tells us the segment index
-    inp += tid.y * stride_segment_axis;
-    out += tid.y * stride_segment_axis;
+    inp += tid.y * in_stride_segment_axis;
+    out += tid.y * out_stride_segment_axis;
 
     // Copy into threadgroup memory
     for (short i = lid.x; i < N_PER_BLOCK; i += BLOCK_THREADS) {
-      tgp_vals[i] = i < size_sorted_axis ? inp[i * stride_sorted_axis]
+      tgp_vals[i] = i < size_sorted_axis ? inp[i * in_stride_sorted_axis]
                                          : val_t(CompareOp::init);
       if (ARG_SORT) {
         tgp_idxs[i] = i;
@@ -264,9 +266,9 @@ struct KernelMergeSort {
     // Write output
     for (int i = lid.x; i < size_sorted_axis; i += BLOCK_THREADS) {
       if (ARG_SORT) {
-        out[i * stride_sorted_axis] = tgp_idxs[i];
+        out[i * out_stride_sorted_axis] = tgp_idxs[i];
       } else {
-        out[i * stride_sorted_axis] = tgp_vals[i];
+        out[i * out_stride_sorted_axis] = tgp_vals[i];
       }
     }
   }
@@ -282,8 +284,10 @@ template <
     const device T* inp [[buffer(0)]],
     device U* out [[buffer(1)]],
     const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& stride_segment_axis [[buffer(4)]],
+    const constant int& in_stride_sorted_axis [[buffer(3)]],
+    const constant int& out_stride_sorted_axis [[buffer(4)]],
+    const constant int& in_stride_segment_axis [[buffer(5)]],
+    const constant int& out_stride_segment_axis [[buffer(6)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 lid [[thread_position_in_threadgroup]]) {
   using sort_kernel =
@@ -298,8 +302,10 @@ template <
         inp,
         out,
         size_sorted_axis,
-        stride_sorted_axis,
-        stride_segment_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
         tgp_vals,
         tgp_idxs,
         tid,
@@ -310,8 +316,10 @@ template <
         inp,
         out,
         size_sorted_axis,
-        stride_sorted_axis,
-        stride_segment_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        in_stride_segment_axis,
+        out_stride_segment_axis,
         tgp_vals,
         nullptr,
         tid,
@@ -331,10 +339,12 @@ template <
     const device T* inp [[buffer(0)]],
     device U* out [[buffer(1)]],
     const constant int& size_sorted_axis [[buffer(2)]],
-    const constant int& stride_sorted_axis [[buffer(3)]],
-    const constant int& nc_dim [[buffer(4)]],
-    const device int* nc_shape [[buffer(5)]],
-    const device size_t* nc_strides [[buffer(6)]],
+    const constant int& in_stride_sorted_axis [[buffer(3)]],
+    const constant int& out_stride_sorted_axis [[buffer(4)]],
+    const constant int& nc_dim [[buffer(5)]],
+    const device int* nc_shape [[buffer(6)]],
+    const device size_t* in_nc_strides [[buffer(7)]],
+    const device size_t* out_nc_strides [[buffer(8)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 lid [[thread_position_in_threadgroup]]) {
   using sort_kernel =
@@ -342,9 +352,10 @@ template <
   using val_t = typename sort_kernel::val_t;
   using idx_t = typename sort_kernel::idx_t;
 
-  auto block_idx = elem_to_loc(tid.y, nc_shape, nc_strides, nc_dim);
-  inp += block_idx;
-  out += block_idx;
+  auto in_block_idx = elem_to_loc(tid.y, nc_shape, in_nc_strides, nc_dim);
+  auto out_block_idx = elem_to_loc(tid.y, nc_shape, out_nc_strides, nc_dim);
+  inp += in_block_idx;
+  out += out_block_idx;
 
   if (ARG_SORT) {
     threadgroup val_t tgp_vals[sort_kernel::N_PER_BLOCK];
@@ -353,7 +364,9 @@ template <
         inp,
         out,
         size_sorted_axis,
-        stride_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        zero_helper,
         zero_helper,
         tgp_vals,
         tgp_idxs,
@@ -365,7 +378,9 @@ template <
         inp,
         out,
         size_sorted_axis,
-        stride_sorted_axis,
+        in_stride_sorted_axis,
+        out_stride_sorted_axis,
+        zero_helper,
         zero_helper,
         tgp_vals,
         nullptr,
@@ -507,13 +522,13 @@ template <
     bool ARG_SORT,
     short BLOCK_THREADS,
     short N_PER_THREAD>
-[[kernel, max_total_threads_per_threadgroup(BLOCK_THREADS)]] void
-mb_block_partition(
+[[kernel]] void mb_block_partition(
     device idx_t* block_partitions [[buffer(0)]],
     const device val_t* dev_vals [[buffer(1)]],
     const device idx_t* dev_idxs [[buffer(2)]],
     const constant int& size_sorted_axis [[buffer(3)]],
     const constant int& merge_tiles [[buffer(4)]],
+    const constant int& n_blocks [[buffer(5)]],
     uint3 tid [[threadgroup_position_in_grid]],
     uint3 lid [[thread_position_in_threadgroup]],
     uint3 tgp_dims [[threads_per_threadgroup]]) {
@@ -528,23 +543,29 @@ mb_block_partition(
   dev_vals += tid.y * size_sorted_axis;
   dev_idxs += tid.y * size_sorted_axis;
 
-  // Find location in merge step
-  int merge_group = lid.x / merge_tiles;
-  int merge_lane = lid.x % merge_tiles;
+  for (int i = lid.x; i <= n_blocks; i += tgp_dims.x) {
+    // Find location in merge step
+    int merge_group = i / merge_tiles;
+    int merge_lane = i % merge_tiles;
 
-  int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
-  int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
+    int sort_sz = sort_kernel::N_PER_BLOCK * merge_tiles;
+    int sort_st = sort_kernel::N_PER_BLOCK * merge_tiles * merge_group;
 
-  int A_st = min(size_sorted_axis, sort_st);
-  int A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
-  int B_st = A_ed;
-  int B_ed = min(size_sorted_axis, B_st + sort_sz / 2);
+    int A_st = min(size_sorted_axis, sort_st);
+    int A_ed = min(size_sorted_axis, sort_st + sort_sz / 2);
+    int B_st = A_ed;
+    int B_ed = min(size_sorted_axis, B_st + sort_sz / 2);
 
-  int partition_at = min(B_ed - A_st, sort_kernel::N_PER_BLOCK * merge_lane);
-  int partition = sort_kernel::merge_partition(
-      dev_vals + A_st, dev_vals + B_st, A_ed - A_st, B_ed - B_st, partition_at);
+    int partition_at = min(B_ed - A_st, sort_kernel::N_PER_BLOCK * merge_lane);
+    int partition = sort_kernel::merge_partition(
+        dev_vals + A_st,
+        dev_vals + B_st,
+        A_ed - A_st,
+        B_ed - B_st,
+        partition_at);
 
-  block_partitions[lid.x] = A_st + partition;
+    block_partitions[i] = A_st + partition;
+  }
 }
 
 template <

mlx/backend/metal/kernels/sort.metal

@@ -10,28 +10,10 @@
 
 #define instantiate_block_sort(                                          \
     name, itname, itype, otname, otype, arg_sort, bn, tn)                \
-  template [[host_name("c" #name "_" #itname "_" #otname "_bn" #bn           \
-                             "_tn" #tn)]] [[kernel]] void                \
-  block_sort<itype, otype, arg_sort, bn, tn>(                            \
-      const device itype* inp [[buffer(0)]],                             \
-      device otype* out [[buffer(1)]],                                   \
-      const constant int& size_sorted_axis [[buffer(2)]],                \
-      const constant int& stride_sorted_axis [[buffer(3)]],              \
-      const constant int& stride_segment_axis [[buffer(4)]],             \
-      uint3 tid [[threadgroup_position_in_grid]],                        \
-      uint3 lid [[thread_position_in_threadgroup]]);                     \
-  template [[host_name("nc" #name "_" #itname "_" #otname "_bn" #bn "_tn" #tn \
-                             )]] [[kernel]] void                         \
-  block_sort_nc<itype, otype, arg_sort, bn, tn>(                         \
-      const device itype* inp [[buffer(0)]],                             \
-      device otype* out [[buffer(1)]],                                   \
-      const constant int& size_sorted_axis [[buffer(2)]],                \
-      const constant int& stride_sorted_axis [[buffer(3)]],              \
-      const constant int& nc_dim [[buffer(4)]],                          \
-      const device int* nc_shape [[buffer(5)]],                          \
-      const device size_t* nc_strides [[buffer(6)]],                     \
-      uint3 tid [[threadgroup_position_in_grid]],                        \
-      uint3 lid [[thread_position_in_threadgroup]]);
+  instantiate_kernel("c" #name "_" #itname "_" #otname "_bn" #bn "_tn" #tn, \
+                     block_sort, itype, otype, arg_sort, bn, tn) \
+  instantiate_kernel("nc" #name "_" #itname "_" #otname "_bn" #bn "_tn" #tn, \
+                     block_sort_nc, itype, otype, arg_sort, bn, tn)
 
 #define instantiate_arg_block_sort_base(itname, itype, bn, tn) \
   instantiate_block_sort(                                      \
@@ -69,43 +51,12 @@ instantiate_block_sort_long(int64, int64_t)
 
 #define instantiate_multi_block_sort(                                      \
     vtname, vtype, itname, itype, arg_sort, bn, tn)                        \
-  template [[host_name("sort_mbsort_" #vtname "_" #itname "_bn" #bn      \
-                       "_tn" #tn)]] [[kernel]] void                        \
-  mb_block_sort<vtype, itype, arg_sort, bn, tn>(                           \
-      const device vtype* inp [[buffer(0)]],                               \
-      device vtype* out_vals [[buffer(1)]],                                \
-      device itype* out_idxs [[buffer(2)]],                                \
-      const constant int& size_sorted_axis [[buffer(3)]],                  \
-      const constant int& stride_sorted_axis [[buffer(4)]],                \
-      const constant int& nc_dim [[buffer(5)]],                            \
-      const device int* nc_shape [[buffer(6)]],                            \
-      const device size_t* nc_strides [[buffer(7)]],                       \
-      uint3 tid [[threadgroup_position_in_grid]],                          \
-      uint3 lid [[thread_position_in_threadgroup]]);                       \
-  template [[host_name("partition_mbsort_" #vtname "_" #itname "_bn" #bn \
-                       "_tn" #tn)]] [[kernel]] void                        \
-  mb_block_partition<vtype, itype, arg_sort, bn, tn>(                      \
-      device itype * block_partitions [[buffer(0)]],                       \
-      const device vtype* dev_vals [[buffer(1)]],                          \
-      const device itype* dev_idxs [[buffer(2)]],                          \
-      const constant int& size_sorted_axis [[buffer(3)]],                  \
-      const constant int& merge_tiles [[buffer(4)]],                       \
-      uint3 tid [[threadgroup_position_in_grid]],                          \
-      uint3 lid [[thread_position_in_threadgroup]],                        \
-      uint3 tgp_dims [[threads_per_threadgroup]]);                         \
-  template [[host_name("merge_mbsort_" #vtname "_" #itname "_bn" #bn     \
-                       "_tn" #tn)]] [[kernel]] void                        \
-  mb_block_merge<vtype, itype, arg_sort, bn, tn>(                          \
-      const device itype* block_partitions [[buffer(0)]],                  \
-      const device vtype* dev_vals_in [[buffer(1)]],                       \
-      const device itype* dev_idxs_in [[buffer(2)]],                       \
-      device vtype* dev_vals_out [[buffer(3)]],                            \
-      device itype* dev_idxs_out [[buffer(4)]],                            \
-      const constant int& size_sorted_axis [[buffer(5)]],                  \
-      const constant int& merge_tiles [[buffer(6)]],                       \
-      const constant int& num_tiles [[buffer(7)]],                         \
-      uint3 tid [[threadgroup_position_in_grid]],                          \
-      uint3 lid [[thread_position_in_threadgroup]]);
+  instantiate_kernel("sort_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_sort, vtype, itype, arg_sort, bn, tn) \
+  instantiate_kernel("partition_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_partition, vtype, itype, arg_sort, bn, tn) \
+  instantiate_kernel("merge_mbsort_" #vtname "_" #itname "_bn" #bn "_tn" #tn, \
+                     mb_block_merge, vtype, itype, arg_sort, bn, tn)
 
 #define instantiate_multi_block_sort_base(vtname, vtype) \
   instantiate_multi_block_sort(vtname, vtype, uint32, uint32_t, true, 512, 8)

mlx/backend/metal/kernels/ternary.h

@@ -10,6 +10,18 @@ template <typename T, typename Op>
   d[index] = Op()(a[index], b[index], c[index]);
 }
 
+template <typename T, typename Op>
+[[kernel]] void ternary_v2(
+    device const bool* a,
+    device const T* b,
+    device const T* c,
+    device T* d,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  d[offset] = Op()(a[offset], b[offset], c[offset]);
+}
+
 template <typename T, typename Op>
 [[kernel]] void ternary_g_nd1(
     device const bool* a,

mlx/backend/metal/kernels/ternary.metal

@@ -9,96 +9,29 @@
 #include "mlx/backend/metal/kernels/ternary_ops.h"
 #include "mlx/backend/metal/kernels/ternary.h"
 
-#define instantiate_ternary_v(name, type, op)                          \
-  template [[host_name("v_" name)]] [[kernel]] void ternary_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_all(op, tname, type)                  \
+  instantiate_kernel("v_" #op #tname, ternary_v, type, op)        \
+  instantiate_kernel("v2_" #op #tname, ternary_v2, type, op)      \
+  instantiate_kernel("g_" #op #tname, ternary_g, type, op)        \
+  instantiate_kernel("g1_" #op #tname, ternary_g_nd1, type, op)   \
+  instantiate_kernel("g2_" #op #tname, ternary_g_nd2, type, op)   \
+  instantiate_kernel("g3_" #op #tname, ternary_g_nd3, type, op)   \
+  instantiate_kernel("g4_" #op #tname, ternary_g_nd, type, op, 4) \
+  instantiate_kernel("g5_" #op #tname, ternary_g_nd, type, op, 5)
 
-#define instantiate_ternary_g(name, type, op)                          \
-  template [[host_name("g_" name)]] [[kernel]] void ternary_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_types(op)               \
+  instantiate_ternary_all(op, bool_, bool)          \
+  instantiate_ternary_all(op, uint8, uint8_t)       \
+  instantiate_ternary_all(op, uint16, uint16_t)     \
+  instantiate_ternary_all(op, uint32, uint32_t)     \
+  instantiate_ternary_all(op, uint64, uint64_t)     \
+  instantiate_ternary_all(op, int8, int8_t)         \
+  instantiate_ternary_all(op, int16, int16_t)       \
+  instantiate_ternary_all(op, int32, int32_t)       \
+  instantiate_ternary_all(op, int64, int64_t)       \
+  instantiate_ternary_all(op, float16, half)        \
+  instantiate_ternary_all(op, float32, float)       \
+  instantiate_ternary_all(op, bfloat16, bfloat16_t) \
+  instantiate_ternary_all(op, complex64, complex64_t) // clang-format on
 
-#define instantiate_ternary_g_dim(name, type, op, dims)  \
-  template [[host_name("g" #dims "_" name )]] [[kernel]] void \
-  ternary_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("g1_" name)]] [[kernel]] void \
-  ternary_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("g2_" name)]] [[kernel]] void \
-  ternary_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("g3_" name)]] [[kernel]] void \
-  ternary_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(#name #tname, type, op)    \
-  instantiate_ternary_g(#name #tname, type, op)    \
-  instantiate_ternary_g_nd(#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) // clang-format on
-
-instantiate_ternary_types(select, Select)
+instantiate_ternary_types(Select)

mlx/backend/metal/kernels/unary.h

@@ -8,6 +8,16 @@ template <typename T, typename Op>
   out[index] = Op()(in[index]);
 }
 
+template <typename T, typename Op>
+[[kernel]] void unary_v2(
+    device const T* in,
+    device T* out,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  size_t offset = index.x + grid_dim.x * size_t(index.y);
+  out[offset] = Op()(in[offset]);
+}
+
 template <typename T, typename Op>
 [[kernel]] void unary_g(
     device const T* in,

mlx/backend/metal/kernels/unary.metal

@@ -5,83 +5,69 @@
 #include "mlx/backend/metal/kernels/unary_ops.h"
 #include "mlx/backend/metal/kernels/unary.h"
 
-#define instantiate_unary_v(name, type, op)                       \
-  template [[host_name(name)]] [[kernel]] void unary_v<type, op>( \
-      device const type* in,                                      \
-      device type* out,                                           \
-      uint index [[thread_position_in_grid]]);
+#define instantiate_unary_all(op, tname, type)            \
+  instantiate_kernel("v" #op #tname, unary_v, type, op)   \
+  instantiate_kernel("v2" #op #tname, unary_v2, type, op) \
+  instantiate_kernel("g" #op #tname, unary_g, type, op)
 
-#define instantiate_unary_g(name, type, op)                       \
-  template [[host_name(name)]] [[kernel]] void unary_g<type, op>( \
-      device const type* in,                                      \
-      device type* out,                                           \
-      device const int* in_shape,                                 \
-      device const size_t* in_strides,                            \
-      device const int& ndim,                                     \
-      uint index [[thread_position_in_grid]]);
+#define instantiate_unary_float(op)               \
+  instantiate_unary_all(op, float16, half)        \
+  instantiate_unary_all(op, float32, float)       \
+  instantiate_unary_all(op, bfloat16, bfloat16_t)
 
-#define instantiate_unary_all(name, tname, type, op) \
-  instantiate_unary_v("v" #name #tname, type, op)    \
-  instantiate_unary_g("g" #name #tname, type, op)
+#define instantiate_unary_types(op)           \
+  instantiate_unary_all(op, bool_, bool)      \
+  instantiate_unary_all(op, uint8, uint8_t)   \
+  instantiate_unary_all(op, uint16, uint16_t) \
+  instantiate_unary_all(op, uint32, uint32_t) \
+  instantiate_unary_all(op, uint64, uint64_t) \
+  instantiate_unary_all(op, int8, int8_t)     \
+  instantiate_unary_all(op, int16, int16_t)   \
+  instantiate_unary_all(op, int32, int32_t)   \
+  instantiate_unary_all(op, int64, int64_t)   \
+  instantiate_unary_float(op)
 
-#define instantiate_unary_float(name, op)               \
-  instantiate_unary_all(name, float16, half, op)        \
-  instantiate_unary_all(name, float32, float, op)       \
-  instantiate_unary_all(name, bfloat16, bfloat16_t, op)
+instantiate_unary_types(Abs)
+instantiate_unary_float(ArcCos)
+instantiate_unary_float(ArcCosh)
+instantiate_unary_float(ArcSin)
+instantiate_unary_float(ArcSinh)
+instantiate_unary_float(ArcTan)
+instantiate_unary_float(ArcTanh)
+instantiate_unary_types(Ceil)
+instantiate_unary_float(Cos)
+instantiate_unary_float(Cosh)
+instantiate_unary_float(Exp)
+instantiate_unary_float(Expm1)
+instantiate_unary_types(Floor)
+instantiate_unary_float(Log)
+instantiate_unary_float(Log2)
+instantiate_unary_float(Log10)
+instantiate_unary_float(Log1p)
+instantiate_unary_types(Negative)
+instantiate_unary_float(Sigmoid)
+instantiate_unary_float(Erf)
+instantiate_unary_float(ErfInv)
+instantiate_unary_types(Sign)
+instantiate_unary_float(Sin)
+instantiate_unary_float(Sinh)
+instantiate_unary_types(Square)
+instantiate_unary_float(Sqrt)
+instantiate_unary_float(Rsqrt)
+instantiate_unary_float(Tan)
+instantiate_unary_float(Tanh)
+instantiate_unary_float(Round)
 
-#define instantiate_unary_types(name, op)           \
-  instantiate_unary_all(name, bool_, bool, op)      \
-  instantiate_unary_all(name, uint8, uint8_t, op)   \
-  instantiate_unary_all(name, uint16, uint16_t, op) \
-  instantiate_unary_all(name, uint32, uint32_t, op) \
-  instantiate_unary_all(name, uint64, uint64_t, op) \
-  instantiate_unary_all(name, int8, int8_t, op)     \
-  instantiate_unary_all(name, int16, int16_t, op)   \
-  instantiate_unary_all(name, int32, int32_t, op)   \
-  instantiate_unary_all(name, int64, int64_t, op)   \
-  instantiate_unary_float(name, op)
+instantiate_unary_all(Abs, complex64, complex64_t)
+instantiate_unary_all(Conjugate, complex64, complex64_t)
+instantiate_unary_all(Cos, complex64, complex64_t)
+instantiate_unary_all(Cosh, complex64, complex64_t)
+instantiate_unary_all(Exp, complex64, complex64_t)
+instantiate_unary_all(Negative, complex64, complex64_t)
+instantiate_unary_all(Sin, complex64, complex64_t)
+instantiate_unary_all(Sinh, complex64, complex64_t)
+instantiate_unary_all(Tan, complex64, complex64_t)
+instantiate_unary_all(Tanh, complex64, complex64_t)
+instantiate_unary_all(Round, complex64, complex64_t)
 
-instantiate_unary_types(abs, Abs)
-instantiate_unary_float(arccos, ArcCos)
-instantiate_unary_float(arccosh, ArcCosh)
-instantiate_unary_float(arcsin, ArcSin)
-instantiate_unary_float(arcsinh, ArcSinh)
-instantiate_unary_float(arctan, ArcTan)
-instantiate_unary_float(arctanh, ArcTanh)
-instantiate_unary_types(ceil, Ceil)
-instantiate_unary_float(cos, Cos)
-instantiate_unary_float(cosh, Cosh)
-instantiate_unary_float(exp, Exp)
-instantiate_unary_float(expm1, Expm1)
-instantiate_unary_types(floor, Floor)
-instantiate_unary_float(log, Log)
-instantiate_unary_float(log2, Log2)
-instantiate_unary_float(log10, Log10)
-instantiate_unary_float(log1p, Log1p)
-instantiate_unary_types(neg, Negative)
-instantiate_unary_float(sigmoid, Sigmoid)
-instantiate_unary_float(erf, Erf)
-instantiate_unary_float(erfinv, ErfInv)
-instantiate_unary_types(sign, Sign)
-instantiate_unary_float(sin, Sin)
-instantiate_unary_float(sinh, Sinh)
-instantiate_unary_types(square, Square)
-instantiate_unary_float(sqrt, Sqrt)
-instantiate_unary_float(rsqrt, Rsqrt)
-instantiate_unary_float(tan, Tan)
-instantiate_unary_float(tanh, Tanh)
-instantiate_unary_float(round, Round)
-
-instantiate_unary_all(abs, complex64, complex64_t, Abs)
-instantiate_unary_all(conj, complex64, complex64_t, Conjugate)
-instantiate_unary_all(cos, complex64, complex64_t, Cos)
-instantiate_unary_all(cosh, complex64, complex64_t, Cosh)
-instantiate_unary_all(exp, complex64, complex64_t, Exp)
-instantiate_unary_all(neg, complex64, complex64_t, Negative)
-instantiate_unary_all(sin, complex64, complex64_t, Sin)
-instantiate_unary_all(sinh, complex64, complex64_t, Sinh)
-instantiate_unary_all(tan, complex64, complex64_t, Tan)
-instantiate_unary_all(tanh, complex64, complex64_t, Tanh)
-instantiate_unary_all(round, complex64, complex64_t, Round)
-
-instantiate_unary_all(lnot, bool_, bool, LogicalNot) // clang-format on
+instantiate_unary_all(LogicalNot, bool_, bool) // clang-format on