From 1fc313db9d56b9952e3c5fc0c4b414e81eb326bf Mon Sep 17 00:00:00 2001 From: CCYeh Date: Thu, 18 Dec 2025 05:48:07 +0100 Subject: [PATCH] Metal logging (#2904) Co-authored-by: Awni Hannun --- docs/src/dev/metal_logging.rst | 40 +++++++++++++++++++ docs/src/index.rst | 1 + mlx/backend/metal/CMakeLists.txt | 2 +- mlx/backend/metal/device.cpp | 7 +++- mlx/backend/metal/kernels/CMakeLists.txt | 4 ++ mlx/backend/metal/kernels/binary_ops.h | 4 ++ .../metal/kernels/indexing/masked_scatter.h | 3 ++ mlx/backend/metal/kernels/logging.h | 26 ++++++++++++ mlx/backend/metal/kernels/utils.h | 1 + 9 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 docs/src/dev/metal_logging.rst create mode 100644 mlx/backend/metal/kernels/logging.h diff --git a/docs/src/dev/metal_logging.rst b/docs/src/dev/metal_logging.rst new file mode 100644 index 00000000..29cecf7a --- /dev/null +++ b/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 diff --git a/docs/src/index.rst b/docs/src/index.rst index 76907dbe..2148d94b 100644 --- a/docs/src/index.rst +++ b/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 diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index 08ea623d..c15f963b 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/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) diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp index 65cd32d3..de61b0f2 100644 --- a/mlx/backend/metal/device.cpp +++ b/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(); diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 16e312e2..262b0495 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/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}") diff --git a/mlx/backend/metal/kernels/binary_ops.h b/mlx/backend/metal/kernels/binary_ops.h index cb3e8a37..4e3d881f 100644 --- a/mlx/backend/metal/kernels/binary_ops.h +++ b/mlx/backend/metal/kernels/binary_ops.h @@ -5,6 +5,8 @@ #include #include +constant mlx::os_log logger("mlx", "binary_ops"); + struct Add { template 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; } diff --git a/mlx/backend/metal/kernels/indexing/masked_scatter.h b/mlx/backend/metal/kernels/indexing/masked_scatter.h index 1fd19e2f..2ba54740 100644 --- a/mlx/backend/metal/kernels/indexing/masked_scatter.h +++ b/mlx/backend/metal/kernels/indexing/masked_scatter.h @@ -2,6 +2,8 @@ #pragma once +constant mlx::os_log logger("mlx", "masked_assign"); + template [[kernel]] void masked_assign_impl( const device bool* mask [[buffer(0)]], @@ -21,6 +23,7 @@ template const uint src_index = scatter_offsets[idx]; if (src_index >= src_batch_size) { + logger.log_debug("Out of bound read from src"); return; } diff --git a/mlx/backend/metal/kernels/logging.h b/mlx/backend/metal/kernels/logging.h new file mode 100644 index 00000000..7b3ee046 --- /dev/null +++ b/mlx/backend/metal/kernels/logging.h @@ -0,0 +1,26 @@ +// Copyright © 2025 Apple Inc. + +#pragma once + +#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320) +#include + +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 + void log_debug(constant char*, Args...) const {} + + template + void log_debug(constant char*, Args...) const constant {} +}; +} // namespace mlx + +#endif \ No newline at end of file diff --git a/mlx/backend/metal/kernels/utils.h b/mlx/backend/metal/kernels/utils.h index acdbc6ad..d3564501 100644 --- a/mlx/backend/metal/kernels/utils.h +++ b/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;