Speed up QMM & QMV

CMakeLists.txt

@@ -234,6 +234,10 @@ if(MLX_BUILD_METAL)
     mlx PUBLIC $<BUILD_INTERFACE:${metal_cpp_SOURCE_DIR}>
                $<INSTALL_INTERFACE:include/metal_cpp>)
   target_link_libraries(mlx PUBLIC ${METAL_LIB} ${FOUNDATION_LIB} ${QUARTZ_LIB})
+  if(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
+    find_library(IOKIT_LIB IOKit)
+    target_link_libraries(mlx PRIVATE ${IOKIT_LIB})
+  endif()
 endif()
 
 if(CMAKE_SYSTEM_NAME STREQUAL "Linux")

mlx/backend/metal/device.cpp

@@ -3,6 +3,11 @@
 #include <cstdlib>
 #include <sstream>
 
+#include <TargetConditionals.h>
+#if TARGET_OS_OSX
+#include <IOKit/IOKitLib.h>
+#endif
+
 #include <fmt/format.h>
 
 #define NS_PRIVATE_IMPLEMENTATION
@@ -1059,4 +1064,37 @@ bool is_nax_available() {
 #endif
 }
 
+int gpu_core_count() {
+  auto _count = []() {
+    int count = 0;
+#if TARGET_OS_OSX
+    io_iterator_t it;
+    if (IOServiceGetMatchingServices(
+            kIOMainPortDefault, IOServiceMatching("AGXAccelerator"), &it) !=
+        KERN_SUCCESS) {
+      return count;
+    }
+    while (io_object_t entry = IOIteratorNext(it)) {
+      CFTypeRef prop = IORegistryEntryCreateCFProperty(
+          entry, CFSTR("gpu-core-count"), kCFAllocatorDefault, 0);
+      if (prop) {
+        if (CFGetTypeID(prop) == CFNumberGetTypeID()) {
+          CFNumberGetValue(
+              static_cast<CFNumberRef>(prop), kCFNumberIntType, &count);
+        }
+        CFRelease(prop);
+      }
+      IOObjectRelease(entry);
+      if (count > 0) {
+        break;
+      }
+    }
+    IOObjectRelease(it);
+#endif
+    return count;
+  };
+  static int count_ = _count();
+  return count_;
+}
+
 } // namespace mlx::core::metal

mlx/backend/metal/device.h

@@ -271,4 +271,7 @@ NS::SharedPtr<NS::AutoreleasePool> new_scoped_memory_pool();
 
 bool is_nax_available();
 
+// Number of GPU cores, or 0 when it is not known.
+int gpu_core_count();
+
 } // namespace mlx::core::metal