From 1efee9db09143089fd8d67ca6411933a350fad6f Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Wed, 13 Mar 2024 20:34:06 -0700
Subject: [PATCH] Add types and order in kernel name (#831)

---
 mlx/backend/common/compiled.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp
index 453ffa63a..8c95eb9ca 100644
--- a/mlx/backend/common/compiled.cpp
+++ b/mlx/backend/common/compiled.cpp
@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include "mlx/backend/common/compiled.h"
+#include "mlx/graph_utils.h"
 #include "mlx/primitives.h"
 #include "mlx/utils.h"
 
@@ -81,13 +82,27 @@ std::string build_lib_name(
     const std::vector<array>& outputs,
     const std::vector<array>& tape,
     const std::unordered_set<uintptr_t>& constant_ids) {
+  NodeNamer namer;
   std::ostringstream os;
   std::ostringstream constant_hasher;
 
+  // Fill the input names. This is not really necessary, I just like having A,
+  // B, C, ... as the inputs.
+  for (auto& x : inputs) {
+    namer.get_name(x);
+  }
+
   // The primitives describing the tape. For unary and binary primitives this
   // must be enough to describe the full computation.
   for (auto& a : tape) {
+    // name and type of output
+    os << namer.get_name(a) << kindof(a.dtype()) << a.itemsize();
+    // computation performed
     a.primitive().print(os);
+    // name of inputs to the function
+    for (auto& inp : a.inputs()) {
+      os << namer.get_name(inp);
+    }
   }
   os << "_";