faster contiguous gather for indices in the first axis

mlx/backend/metal/CMakeLists.txt

@@ -33,10 +33,11 @@ make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h)
 make_jit_source(binary_ops)
 make_jit_source(ternary_ops)
 make_jit_source(reduce_utils kernels/atomic.h kernels/reduction/ops.h)
-make_jit_source(scatter kernels/indexing.h)
-make_jit_source(gather kernels/indexing.h)
-make_jit_source(gather_axis)
-make_jit_source(scatter_axis)
+make_jit_source(indexing/scatter kernels/indexing/indexing.h)
+make_jit_source(indexing/gather kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_front kernels/indexing/indexing.h)
+make_jit_source(indexing/gather_axis)
+make_jit_source(indexing/scatter_axis)
 make_jit_source(hadamard)
 
 if(MLX_METAL_JIT)

mlx/backend/metal/indexing.cpp

@@ -52,8 +52,10 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto& s = stream();
   auto& d = metal::device(s.device);
 
-  int idx_ndim = nidx ? inputs[1].ndim() : 0;
-  size_t ndim = src.ndim();
+  size_t slice_size = 1;
+  for (auto s : slice_sizes_) {
+    slice_size *= s;
+  }
 
   bool large_index = nidx && inputs[1].size() > INT32_MAX;
   bool large_src = src.size() > INT32_MAX;
@@ -61,6 +63,52 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   bool large = large_index || large_src || large_out;
 
   std::string idx_type_name = nidx ? type_to_name(inputs[1]) : "";
+
+  if (src.flags().row_contiguous && nidx == 1 && axes_[0] == 0 &&
+      inputs[1].flags().row_contiguous && slice_size == src.strides()[0]) {
+    auto& indices = inputs[1];
+    std::string kernel_name = fmt::format(
+        "gather_front{0}_{1}_{2}",
+        type_to_name(out),
+        idx_type_name,
+        large ? "int64_t" : "int");
+    std::string lib_name = kernel_name;
+
+    auto lib = d.get_library(lib_name, [&]() {
+      std::string kernel_source = metal::utils();
+      kernel_source += metal::gather_front();
+      kernel_source += get_template_definition(
+          kernel_name,
+          "gather_front",
+          get_type_string(out.dtype()),
+          get_type_string(indices.dtype()),
+          large ? "int64_t" : "int");
+
+      return kernel_source;
+    });
+
+    auto& compute_encoder = d.get_command_encoder(s.index);
+    auto kernel = d.get_kernel(kernel_name, lib);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    size_t dim_x = slice_size;
+    size_t dim_y = indices.size();
+    auto group_dims = get_block_dims(dim_x, dim_y, 1);
+    MTL::Size grid_dims = MTL::Size(dim_x, dim_y, 1);
+
+    compute_encoder.set_input_array(src, 0);
+    compute_encoder.set_input_array(indices, 1);
+    compute_encoder.set_output_array(out, 2);
+    compute_encoder.set_bytes(slice_size, 3);
+    compute_encoder.set_bytes(src.shape(0), 4);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+
+    return;
+  }
+
+  int idx_ndim = nidx ? inputs[1].ndim() : 0;
+  size_t ndim = src.ndim();
+
   std::string kernel_name = fmt::format(
       "gather{0}{1}_{2}_{3}_{4}",
       type_to_name(out),
@@ -96,11 +144,6 @@ void Gather::eval_gpu(const std::vector<array>& inputs, array& out) {
   auto kernel = d.get_kernel(kernel_name, lib);
   compute_encoder.set_compute_pipeline_state(kernel);
 
-  size_t slice_size = 1;
-  for (auto s : slice_sizes_) {
-    slice_size *= s;
-  }
-
   // Launch 3D grid of threads
   // First two dimensions for the indices, the last one for the slice
   size_t dim0 = 1;

mlx/backend/metal/jit/includes.h

@@ -19,6 +19,7 @@ const char* binary_two();
 const char* copy();
 const char* fft();
 const char* gather_axis();
+const char* gather_front();
 const char* hadamard();
 const char* logsumexp();
 const char* quantized_utils();

mlx/backend/metal/kernels/indexing/gather.h

@@ -2,7 +2,7 @@
 
 #pragma once
 
-#include "mlx/backend/metal/kernels/indexing.h"
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
 
 template <typename T, typename IdxT, int NIDX, int IDX_NDIM, typename LocT>
 METAL_FUNC void gather_impl(

mlx/backend/metal/kernels/indexing/gather_front.h

@@ -0,0 +1,21 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
+
+template <typename T, typename IdxT, typename LocT>
+[[kernel]] void gather_front(
+    const device T* src,
+    const device IdxT* indices,
+    device T* out,
+    const constant int64_t& stride,
+    const constant int& size,
+    uint2 index [[thread_position_in_grid]],
+    uint2 grid_dim [[threads_per_grid]]) {
+  auto idx = offset_neg_idx(indices[index.y], size);
+  LocT src_idx = static_cast<LocT>(stride) * idx;
+  LocT out_idx = static_cast<LocT>(stride) * index.y;
+
+  out[out_idx + index.x] = src[src_idx + index.x];
+}

mlx/backend/metal/kernels/indexing/scatter.h

@@ -2,7 +2,7 @@
 
 #pragma once
 
-#include "mlx/backend/metal/kernels/indexing.h"
+#include "mlx/backend/metal/kernels/indexing/indexing.h"
 
 template <
     typename T,