metal_kernel: build and hash each kernel variant once (gemma4-qat#189)

mlx/backend/common/metal_kernel.cpp

@@ -1,7 +1,9 @@
 // Copyright © 2024 Apple Inc.
 
 #include <iostream>
+#include <mutex>
 #include <sstream>
+#include <unordered_map>
 
 #include "mlx/backend/common/compiled.h"
 #include "mlx/backend/metal/metal.h"
@@ -217,6 +219,59 @@ std::string make_template_hash(const std::string& template_def) {
   return template_hash;
 }
 
+// The name, source and source hash of one kernel variant.
+struct KernelVariant {
+  std::string name;
+  std::shared_ptr<const std::string> source;
+  size_t source_hash;
+};
+
+// Shared by all copies of one kernel function. Holds at most
+// `metal_kernel_max_cached_variants` entries and is cleared when full. Callers
+// keep an owning pointer to their entry, so clearing is safe at any time.
+struct VariantCache {
+  std::mutex mutex;
+  std::unordered_map<std::string, std::shared_ptr<const KernelVariant>>
+      variants;
+};
+
+// Encodes each call argument that the name and source depend on. Template
+// arguments include their kind and name: `int` 1 and `bool` true give the same
+// name but a different signature.
+std::string variant_key(
+    const std::vector<array>& inputs,
+    const std::vector<Dtype>& output_dtypes,
+    const std::vector<std::pair<std::string, TemplateArg>>& template_args) {
+  std::string key;
+  auto append_bytes = [&key](auto v) {
+    key.append(reinterpret_cast<const char*>(&v), sizeof(v));
+  };
+  for (const auto& arr : inputs) {
+    key += static_cast<char>(arr.dtype().val());
+    key += arr.ndim() == 0                     ? 's'
+        : arr.size() < max_constant_array_size ? 'c'
+                                               : 'd';
+  }
+  for (const auto& dtype : output_dtypes) {
+    key += static_cast<char>(dtype.val());
+  }
+  for (const auto& [name, arg] : template_args) {
+    key += static_cast<char>(arg.index());
+    append_bytes(name.size());
+    key += name;
+    std::visit(
+        [&](auto v) {
+          if constexpr (std::is_same_v<decltype(v), Dtype>) {
+            key += static_cast<char>(v.val());
+          } else {
+            append_bytes(v);
+          }
+        },
+        arg);
+  }
+  return key;
+}
+
 } // namespace
 
 CustomKernelFunction metal_kernel(
@@ -270,7 +325,10 @@ CustomKernelFunction metal_kernel(
     }
   }
 
+  auto cache = std::make_shared<VariantCache>();
+
   return [=,
+          cache = std::move(cache),
           shape_infos = std::move(shape_infos),
           attributes = std::move(attributes)](
              const std::vector<array>& inputs,
@@ -298,60 +356,86 @@ CustomKernelFunction metal_kernel(
 
     auto s = resolve_metal_kernel_stream(s_);
 
-    std::string kernel_name = "custom_kernel_" + name;
-    std::string template_def = "";
-    if (!template_args.empty()) {
-      template_def = write_template(template_args);
-      auto template_hash = make_template_hash(template_def);
-      kernel_name += "_";
-      kernel_name += template_hash;
+    auto key = variant_key(inputs, output_dtypes, template_args);
+    std::shared_ptr<const KernelVariant> variant;
+    {
+      std::lock_guard lock(cache->mutex);
+      if (auto it = cache->variants.find(key); it != cache->variants.end()) {
+        variant = it->second;
+      }
     }
+    if (!variant) {
+      std::string kernel_name = "custom_kernel_" + name;
+      std::string template_def = "";
+      if (!template_args.empty()) {
+        template_def = write_template(template_args);
+        auto template_hash = make_template_hash(template_def);
+        kernel_name += "_";
+        kernel_name += template_hash;
+      }
 
-    // The generated source depends on the dtypes of the inputs and outputs
-    // and on how each input is passed (see `write_signature`). Include them
-    // in the kernel name so that a given name always maps to the same source.
-    for (const auto& arr : inputs) {
-      kernel_name += "_";
-      kernel_name += get_type_string(arr.dtype());
-      if (arr.ndim() == 0) {
-        kernel_name += "s";
-      } else if (arr.size() < max_constant_array_size) {
-        kernel_name += "c";
+      // The generated source depends on the dtypes of the inputs and outputs
+      // and on how each input is passed (see `write_signature`). Include them
+      // in the kernel name so that a given name always maps to the same
+      // source.
+      for (const auto& arr : inputs) {
+        kernel_name += "_";
+        kernel_name += get_type_string(arr.dtype());
+        if (arr.ndim() == 0) {
+          kernel_name += "s";
+        } else if (arr.size() < max_constant_array_size) {
+          kernel_name += "c";
+        }
+      }
+      for (const auto& dtype : output_dtypes) {
+        kernel_name += "_";
+        kernel_name += get_type_string(dtype);
       }
-    }
-    for (const auto& dtype : output_dtypes) {
-      kernel_name += "_";
-      kernel_name += get_type_string(dtype);
-    }
 
-    std::string kernel_source = write_signature(
-        kernel_name,
-        header,
-        source,
-        input_names,
-        inputs,
-        output_names,
-        output_dtypes,
-        template_args,
-        attributes,
-        shape_infos,
-        atomic_outputs);
+      std::string kernel_source = write_signature(
+          kernel_name,
+          header,
+          source,
+          input_names,
+          inputs,
+          output_names,
+          output_dtypes,
+          template_args,
+          attributes,
+          shape_infos,
+          atomic_outputs);
 
-    if (!template_args.empty()) {
-      template_def = kernel_name + template_def;
-      kernel_source += "\ntemplate [[host_name(\"";
-      kernel_source += kernel_name;
-      kernel_source += "\")]] [[kernel]] decltype(";
-      kernel_source += template_def;
-      kernel_source += ") ";
-      kernel_source += template_def;
-      kernel_source += ";\n";
+      if (!template_args.empty()) {
+        template_def = kernel_name + template_def;
+        kernel_source += "\ntemplate [[host_name(\"";
+        kernel_source += kernel_name;
+        kernel_source += "\")]] [[kernel]] decltype(";
+        kernel_source += template_def;
+        kernel_source += ") ";
+        kernel_source += template_def;
+        kernel_source += ";\n";
+      }
+
+      auto source_hash = std::hash<std::string>{}(kernel_source);
+      variant = std::make_shared<const KernelVariant>(KernelVariant{
+          std::move(kernel_name),
+          std::make_shared<const std::string>(std::move(kernel_source)),
+          source_hash});
+      std::lock_guard lock(cache->mutex);
+      if (auto it = cache->variants.find(key); it != cache->variants.end()) {
+        variant = it->second;
+      } else {
+        if (cache->variants.size() >= metal_kernel_max_cached_variants) {
+          cache->variants.clear();
+        }
+        cache->variants.emplace(std::move(key), variant);
+      }
     }
 
     if (verbose) {
       std::cout << "Generated source code for `" << name << "`:" << std::endl
                 << "```" << std::endl
-                << kernel_source << std::endl
+                << *variant->source << std::endl
                 << "```" << std::endl;
     }
 
@@ -360,8 +444,8 @@ CustomKernelFunction metal_kernel(
         std::move(output_dtypes),
         std::make_shared<CustomKernel>(
             s,
-            std::move(kernel_name),
-            std::move(kernel_source),
+            variant->name,
+            variant->source,
             grid,
             threadgroup,
             shape_infos,
@@ -370,7 +454,8 @@ CustomKernelFunction metal_kernel(
             std::vector<ScalarArg>{},
             false,
             0,
-            compile_options.serialize()),
+            compile_options.serialize(),
+            variant->source_hash),
         std::move(inputs));
   };
 }

mlx/backend/cuda/custom_kernel.cpp

@@ -310,14 +310,13 @@ void CustomKernel::eval_gpu(
   // Compile the custom kernel
   std::string kernel_name =
       (is_precompiled_) ? name_ : "mlx::core::cu::" + name_;
-  std::string module_name =
-      fmt::format("{}_{:x}", name_, std::hash<std::string>{}(source_));
+  std::string module_name = fmt::format("{}_{:x}", name_, source_hash_);
   cu::JitModule& mod = cu::get_jit_module(
       encoder.device(),
       module_name,
       [&]() {
         return std::make_tuple(
-            is_precompiled_, source_, std::vector{kernel_name});
+            is_precompiled_, *source_, std::vector{kernel_name});
       },
       false);
 

mlx/backend/metal/custom_kernel.cpp

@@ -47,10 +47,10 @@ void CustomKernel::eval_gpu(
 
   auto& d = metal::device(s.device);
 
-  std::string lib_name = fmt::format(
-      "{}_{:x}_{}", name_, std::hash<std::string>{}(source_), compile_options_);
+  std::string lib_name =
+      fmt::format("{}_{:x}_{}", name_, source_hash_, compile_options_);
   auto lib = d.get_library(
-      lib_name, compile_options_, [this] { return metal::utils() + source_; });
+      lib_name, compile_options_, [this] { return metal::utils() + *source_; });
   auto kernel = d.get_kernel(name_, lib);
   auto& compute_encoder = metal::get_command_encoder(s);
   compute_encoder.set_compute_pipeline_state(kernel);

mlx/fast_primitives.h

@@ -433,6 +433,9 @@ class Quantize : public Custom {
 
 using ScalarArg = std::variant<bool, int, float>;
 
+// The most variants that one metal_kernel function keeps built.
+inline constexpr size_t metal_kernel_max_cached_variants = 64;
+
 class CustomKernel : public Primitive {
  public:
   CustomKernel(
@@ -448,9 +451,40 @@ class CustomKernel : public Primitive {
       bool is_precompiled,
       int shared_memory,
       CompileOptions::Data compile_options = {})
+      : CustomKernel(
+            stream,
+            std::move(name),
+            std::make_shared<const std::string>(std::move(source)),
+            grid,
+            threadgroup,
+            std::move(shape_infos),
+            ensure_row_contiguous,
+            init_value,
+            std::move(scalar_arguments),
+            is_precompiled,
+            shared_memory,
+            compile_options) {}
+
+  // Kernels built from one source can share it and its hash.
+  CustomKernel(
+      Stream stream,
+      std::string name,
+      std::shared_ptr<const std::string> source,
+      std::tuple<int, int, int> grid,
+      std::tuple<int, int, int> threadgroup,
+      std::vector<std::tuple<bool, bool, bool>> shape_infos,
+      bool ensure_row_contiguous,
+      std::optional<float> init_value,
+      std::vector<ScalarArg> scalar_arguments,
+      bool is_precompiled,
+      int shared_memory,
+      CompileOptions::Data compile_options,
+      std::optional<size_t> source_hash = std::nullopt)
       : Primitive(stream),
         name_(std::move(name)),
         source_(std::move(source)),
+        source_hash_(
+            source_hash ? *source_hash : std::hash<std::string>{}(*source_)),
         grid_(grid),
         threadgroup_(threadgroup),
         shape_infos_(std::move(shape_infos)),
@@ -473,7 +507,7 @@ class CustomKernel : public Primitive {
   auto state() const {
     return std::make_tuple(
         name_,
-        source_,
+        *source_,
         grid_,
         threadgroup_,
         shape_infos_,
@@ -485,9 +519,14 @@ class CustomKernel : public Primitive {
         compile_options_);
   }
 
+  const std::shared_ptr<const std::string>& shared_source() const {
+    return source_;
+  }
+
  private:
   std::string name_;
-  std::string source_;
+  std::shared_ptr<const std::string> source_;
+  size_t source_hash_;
   std::tuple<int, int, int> grid_;
   std::tuple<int, int, int> threadgroup_;
   std::vector<std::tuple<bool, bool, bool>> shape_infos_;

python/tests/test_fast.py

@@ -1088,6 +1088,59 @@ def call_kernel(a: mx.array, source):
         out = call_kernel(a, source)
         self.assertTrue(mx.array_equal(out, mx.ones_like(out)))
 
+        # `int` 1 and `bool` True print the same template value but
+        # declare a different template parameter.
+        kernel = mx.fast.metal_kernel(
+            name="sizeof_kernel",
+            input_names=["inp"],
+            output_names=["out"],
+            source="out[0] = sizeof(V);",
+        )
+
+        def call_sizeof(value):
+            return kernel(
+                inputs=[a],
+                grid=(1, 1, 1),
+                threadgroup=(1, 1, 1),
+                output_shapes=[(1,)],
+                output_dtypes=[mx.int32],
+                template=[("V", value)],
+                stream=mx.gpu,
+            )[0].item()
+
+        for _ in range(2):
+            self.assertEqual(call_sizeof(1), 4)
+            self.assertEqual(call_sizeof(True), 1)
+
+    @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
+    def test_custom_kernel_many_template_values(self):
+        # More variants than the kernel keeps built. An array made before
+        # its variant was evicted must still evaluate correctly.
+        kernel = mx.fast.metal_kernel(
+            name="value_kernel",
+            input_names=["inp"],
+            output_names=["out"],
+            source="out[0] = V;",
+        )
+        a = mx.zeros((1,), dtype=mx.int32)
+
+        def call_value(value):
+            return kernel(
+                inputs=[a],
+                grid=(1, 1, 1),
+                threadgroup=(1, 1, 1),
+                output_shapes=[(1,)],
+                output_dtypes=[mx.int32],
+                template=[("V", value)],
+                stream=mx.gpu,
+            )[0]
+
+        pending = call_value(0)
+        for value in range(1, 200):
+            self.assertEqual(call_value(value).item(), value)
+        self.assertEqual(pending.item(), 0)
+        self.assertEqual(call_value(0).item(), 0)
+
     @unittest.skipIf(not mx.metal.is_available(), "Metal is not available")
     def test_custom_kernel_same_name_different_source_one_eval(self):
         # Regression test for #3832: two kernels sharing a name but with