Metal logging (#2904)

Co-authored-by: Awni Hannun <[email protected]>

docs/src/dev/metal_logging.rst

@@ -0,0 +1,40 @@
+Metal Logging
+=============
+
+In debug builds, MLX compiles Metal kernels with ``os_log`` enabled so shader
+warnings and debug messages are visible during development.
+
+.. note::
+    Metal logging is only available with Metal 3.2 or higher (macOS 15 and up,
+    iOS 18 and up).
+
+To enable logging from kernels, first make sure to build in debug mode:
+
+.. code-block:: bash
+
+    DEBUG=1 python -m pip install -e .
+
+Then, in the kernel source code include MLX's logging shim and use
+``mlx::os_log``:
+
+.. code-block::
+
+    #include "mlx/backend/metal/kernels/logging.h"
+
+    constant mlx::os_log logger("mlx", "my_kernel");
+
+    kernel void my_kernel(/* ... */) {
+    // ...
+      logger.log_debug("unexpected state: idx=%u", idx);
+    }
+
+When you run the program, set the Metal log level to your desired level and
+forward logs to ``stderr``:
+
+.. code-block:: bash
+
+    MTL_LOG_LEVEL=MTLLogLevelDebug MTL_LOG_TO_STDERR=1 python script.py
+
+See the `Metal logging guide`_ for more details.
+
+.. _`Metal logging guide`: https://developer.apple.com/documentation/metal/logging-shader-debug-messages

docs/src/index.rst

@@ -89,5 +89,6 @@ are the CPU and GPU.
 
    dev/extensions
    dev/metal_debugger
+   dev/metal_logging
    dev/custom_metal_kernels
    dev/mlx_in_cpp

mlx/backend/metal/CMakeLists.txt

@@ -22,7 +22,7 @@ function(make_jit_source SRC_FILE)
 endfunction(make_jit_source)
 
 make_jit_source(utils kernels/bf16.h kernels/bf16_math.h kernels/complex.h
-                kernels/defines.h)
+                kernels/defines.h kernels/logging.h)
 make_jit_source(unary_ops kernels/erf.h kernels/expm1f.h kernels/fp8.h)
 make_jit_source(binary_ops)
 make_jit_source(ternary_ops)

mlx/backend/metal/device.cpp

@@ -224,7 +224,7 @@ MTL::Library* load_library(
   std::ostringstream msg;
   msg << "Failed to load the metallib " << lib_name << ".metallib. "
       << "We attempted to load it from <" << current_binary_dir() << "/"
-      << lib_name << ".metallib" << ">";
+      << lib_name << ".metallib>";
 #ifdef SWIFTPM_BUNDLE
   msg << " and from the Swift PM bundle.";
 #endif
@@ -529,6 +529,11 @@ MTL::Library* Device::build_library_(const std::string& source_string) {
   auto options = MTL::CompileOptions::alloc()->init();
   options->setFastMathEnabled(false);
   options->setLanguageVersion(get_metal_version());
+#ifndef NDEBUG
+  if (options->languageVersion() >= MTL::LanguageVersion3_2) {
+    options->setEnableLogging(true);
+  }
+#endif
   auto mtl_lib = device_->newLibrary(ns_code, options, &error);
   options->release();
 

mlx/backend/metal/kernels/CMakeLists.txt

@@ -6,6 +6,7 @@ set(BASE_HEADERS
     erf.h
     expm1f.h
     fp8.h
+    logging.h
     utils.h)
 
 function(build_kernel_base TARGET SRCFILE DEPS)
@@ -20,6 +21,9 @@ function(build_kernel_base TARGET SRCFILE DEPS)
   if(MLX_METAL_DEBUG)
     set(METAL_FLAGS ${METAL_FLAGS} -gline-tables-only -frecord-sources)
   endif()
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug" AND MLX_METAL_VERSION GREATER_EQUAL 320)
+    set(METAL_FLAGS ${METAL_FLAGS} -fmetal-enable-logging)
+  endif()
   if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
     set(METAL_FLAGS ${METAL_FLAGS}
                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")

mlx/backend/metal/kernels/binary_ops.h

@@ -5,6 +5,8 @@
 #include <metal_integer>
 #include <metal_math>
 
+constant mlx::os_log logger("mlx", "binary_ops");
+
 struct Add {
   template <typename T>
   T operator()(T x, T y) {
@@ -225,6 +227,8 @@ struct Power {
     T res = 1;
     // Undefined to raise integer to negative power
     if (exp < 0) {
+      logger.log_debug(
+          "int pow exp<0 (base=%ld exp=%ld)", (long)base, (long)exp);
       return 0;
     }
 

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

@@ -2,6 +2,8 @@
 
 #pragma once
 
+constant mlx::os_log logger("mlx", "masked_assign");
+
 template <typename T, bool src_contiguous>
 [[kernel]] void masked_assign_impl(
     const device bool* mask [[buffer(0)]],
@@ -21,6 +23,7 @@ template <typename T, bool src_contiguous>
 
   const uint src_index = scatter_offsets[idx];
   if (src_index >= src_batch_size) {
+    logger.log_debug("Out of bound read from src");
     return;
   }
 

mlx/backend/metal/kernels/logging.h

@@ -0,0 +1,26 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320)
+#include <metal_logging>
+
+namespace mlx {
+using os_log = metal::os_log;
+} // namespace mlx
+
+#else
+
+namespace mlx {
+struct os_log {
+  constexpr os_log(constant char*, constant char*) constant {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const {}
+
+  template <typename... Args>
+  void log_debug(constant char*, Args...) const constant {}
+};
+} // namespace mlx
+
+#endif

mlx/backend/metal/kernels/utils.h

@@ -8,6 +8,7 @@
 #include "mlx/backend/metal/kernels/bf16_math.h"
 #include "mlx/backend/metal/kernels/complex.h"
 #include "mlx/backend/metal/kernels/defines.h"
+#include "mlx/backend/metal/kernels/logging.h"
 
 typedef half float16_t;