Debug time graph verification

mlx/CMakeLists.txt

@@ -1,6 +1,7 @@
 target_sources(
   mlx
-  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
+  PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp
+          ${CMAKE_CURRENT_SOURCE_DIR}/array.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/compile.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dtype.cpp

mlx/allocator.cpp

@@ -0,0 +1,53 @@
+// Copyright © 2026 Apple Inc.
+
+#include <algorithm>
+#include <mutex>
+
+#include "mlx/allocator.h"
+
+namespace mlx::core::allocator {
+
+#ifndef NDEBUG
+
+bool Allocator::has_invalid_memory(const std::vector<void*>& ptrs) {
+  std::shared_lock lock(allocations_mutex_);
+  return std::any_of(ptrs.begin(), ptrs.end(), [this](void* ptr) {
+    if (!ptr) {
+      return false;
+    }
+    auto it = allocations_.upper_bound(ptr);
+    if (it == allocations_.begin()) {
+      return true;
+    }
+    --it;
+    auto offset = static_cast<char*>(ptr) - static_cast<char*>(it->first);
+    return static_cast<size_t>(offset) > it->second;
+  });
+}
+
+void Allocator::track_allocation(void* data_ptr, size_t size) {
+  if (data_ptr && size > 0) {
+    std::unique_lock lock(allocations_mutex_);
+    allocations_[data_ptr] = size;
+  }
+}
+
+void Allocator::untrack_allocation(void* data_ptr) {
+  if (data_ptr) {
+    std::unique_lock lock(allocations_mutex_);
+    allocations_.erase(data_ptr);
+  }
+}
+
+void Allocator::swap_allocation(void* data_ptr_a, void* data_ptr_b) {
+  std::unique_lock lock(allocations_mutex_);
+  auto it = allocations_.find(data_ptr_a);
+  if (it != allocations_.end()) {
+    allocations_[data_ptr_b] = it->second;
+    allocations_.erase(it);
+  }
+}
+
+#endif // NDEBUG
+
+} // namespace mlx::core::allocator

mlx/allocator.h

@@ -6,6 +6,12 @@
 #include <functional>
 #include <utility>
 
+#ifndef NDEBUG
+#include <map>
+#include <shared_mutex>
+#include <vector>
+#endif
+
 #include "mlx/api.h"
 
 namespace mlx::core::allocator {
@@ -52,12 +58,29 @@ class MLX_API Allocator {
   };
   virtual void release(Buffer buffer) {}
 
+#ifndef NDEBUG
+  bool has_invalid_memory(const std::vector<void*>& ptrs);
+#endif
+
   Allocator() = default;
   Allocator(const Allocator& other) = delete;
   Allocator(Allocator&& other) = delete;
   Allocator& operator=(const Allocator& other) = delete;
   Allocator& operator=(Allocator&& other) = delete;
   virtual ~Allocator() = default;
+
+ protected:
+#ifndef NDEBUG
+  void track_allocation(void* data_ptr, size_t size);
+  void untrack_allocation(void* data_ptr);
+  void swap_allocation(void* data_ptr_a, void* data_ptr_b);
+#endif
+
+ private:
+#ifndef NDEBUG
+  std::map<void*, size_t> allocations_;
+  std::shared_mutex allocations_mutex_;
+#endif
 };
 
 MLX_API Allocator& allocator();

mlx/backend/common/common.cpp

@@ -41,7 +41,7 @@ void AsStrided::eval(const array& in, array& out) {
   flags.row_contiguous = row_contiguous;
   flags.col_contiguous = col_contiguous;
 
-  return out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
+  out.copy_shared_buffer(in, strides_, flags, data_size, offset_);
 }
 
 void Broadcast::eval(const std::vector<array>& inputs, array& out) {

mlx/backend/cpu/primitives.cpp

@@ -91,7 +91,7 @@ void AsStrided::eval_cpu(const std::vector<array>& inputs, array& out) {
     eval(inputs[0], out);
   } else {
     auto copy = contiguous_copy_cpu(inputs[0], stream());
-    cpu::get_command_encoder(stream()).add_temporary(copy);
+    // Donated to out.
     eval(copy, out);
   }
 }

mlx/backend/cuda/allocator.cpp

@@ -251,6 +251,9 @@ CudaAllocator::malloc_async(size_t size, int device, cudaStream_t stream) {
   if (buf->device >= 0 && buf->device != device) {
     move_to_unified_memory(*buf, stream);
   }
+#ifndef NDEBUG
+  track_allocation(buf->data, size);
+#endif
   return Buffer{buf};
 }
 
@@ -268,6 +271,10 @@ void CudaAllocator::free(Buffer buffer) {
     return;
   }
 
+#ifndef NDEBUG
+  untrack_allocation(buf->data);
+#endif
+
   std::unique_lock lock(mutex_);
   active_memory_ -= buf->size;
   if (get_cache_memory() < max_pool_size_) {
@@ -301,6 +308,9 @@ void CudaAllocator::move_to_unified_memory(
     CHECK_CUDA_ERROR(cudaMemcpy(data, buf.data, buf.size, kind));
     free_async(buf);
   }
+#ifndef NDEBUG
+  swap_allocation(buf.data, data);
+#endif
   buf.data = data;
   buf.device = -1;
 }

mlx/backend/cuda/conv/gemm_grouped_conv.cu

@@ -163,6 +163,7 @@ void gemm_grouped_conv_nd(
   wt_view.copy_shared_buffer(
       wt, {wt.strides(0), 1, C_per_group}, wt.flags(), wt.size());
   array wt_reshaped = contiguous_copy_gpu(wt_view, s);
+  encoder.add_temporary(wt_reshaped);
 
   // Batch with size of groups.
   Shape batch_shape{params.groups};

mlx/backend/cuda/copy.cu

@@ -109,7 +109,11 @@ void fill_gpu(const array& in, array& out, const Stream& s) {
   copy_contiguous(encoder, CopyType::Scalar, in, out, 0, 0);
 }
 
-void reshape_gpu(const array& in, array& out, Stream s) {
+void reshape_gpu(
+    const array& in,
+    array& out,
+    bool add_out_to_temporary_on_allocation,
+    Stream s) {
   auto [copy_necessary, out_strides] = prepare_reshape(in, out);
   if (copy_necessary) {
     auto& encoder = cu::get_command_encoder(s);
@@ -124,6 +128,9 @@ void reshape_gpu(const array& in, array& out, Stream s) {
         0,
         CopyType::General,
         s);
+    if (add_out_to_temporary_on_allocation) {
+      encoder.add_temporary(out);
+    }
   } else {
     shared_buffer_reshape(in, out_strides, out);
   }

mlx/backend/cuda/device.cpp

@@ -245,16 +245,25 @@ void CommandEncoder::set_input_array(const array& arr) {
   bytes_in_graph_ += arr.data_size();
   auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
   active_deps_.push_back(id);
+#ifndef NDEBUG
+  if (arr.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(gpu_ptr<void>(arr));
+  }
+#endif
 }
 
 void CommandEncoder::set_output_array(const array& arr) {
   if (!use_cuda_graphs()) {
     return;
   }
-
   auto id = reinterpret_cast<std::uintptr_t>(arr.buffer().ptr());
   active_deps_.push_back(id);
   active_outputs_.push_back(id);
+#ifndef NDEBUG
+  if (arr.data_shared_ptr()->is_owned()) {
+    registered_data_ptrs_.push_back(gpu_ptr<void>(arr));
+  }
+#endif
 }
 
 void CommandEncoder::add_kernel_node_raw(
@@ -514,6 +523,19 @@ void CommandEncoder::synchronize() {
   f.wait();
 }
 
+#ifndef NDEBUG
+void CommandEncoder::verify_graph(const char* primitive_name) {
+  if (allocator::allocator().has_invalid_memory(registered_data_ptrs_)) {
+    throw std::runtime_error(
+        fmt::format(
+            "There are invalid memory registered when calling {}::eval_gpu, "
+            "which is a bug that corrupts the graph, and usually caused by "
+            "temporary arrays not registered with \"add_temporary\" method.",
+            primitive_name));
+  }
+}
+#endif
+
 void CommandEncoder::commit_impl() {
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
@@ -607,6 +629,9 @@ void CommandEncoder::clear_graph_state() {
   active_outputs_.clear();
   concurrent_nodes_.clear();
   roots_.clear();
+#ifndef NDEBUG
+  registered_data_ptrs_.clear();
+#endif
   is_graph_updatable_ = true;
 }
 

mlx/backend/cuda/device.h

@@ -112,6 +112,10 @@ class MLX_API CommandEncoder {
   void add_completed_handler(std::function<void()> task);
   bool needs_commit();
   void commit();
+  void synchronize();
+#ifndef NDEBUG
+  void verify_graph(const char* primitive_name);
+#endif
 
   Device& device() {
     return device_;
@@ -121,9 +125,6 @@ class MLX_API CommandEncoder {
     return stream_;
   }
 
-  // Wait until kernels and completion handlers are finished
-  void synchronize();
-
  private:
   cudaGraphNode_t add_kernel_node_raw(const cudaKernelNodeParams& params);
   CUgraphNode add_kernel_node_raw(const CUDA_KERNEL_NODE_PARAMS& params);
@@ -147,7 +148,7 @@ class MLX_API CommandEncoder {
   CudaStream stream_;
   CudaGraph graph_;
   std::shared_ptr<Worker> worker_;
-  int node_count_{0};
+
   bool in_concurrent_{false};
   std::vector<cudaGraphNode_t> from_nodes_;
   std::vector<cudaGraphNode_t> to_nodes_;
@@ -160,11 +161,17 @@ class MLX_API CommandEncoder {
   std::vector<std::uintptr_t> active_outputs_;
   std::unordered_map<std::uintptr_t, GraphNode> node_map_;
   std::vector<cudaGraphNode_t> roots_;
-  size_t bytes_in_graph_{0};
+#ifndef NDEBUG
+  std::vector<void*> registered_data_ptrs_;
+#endif
   bool is_graph_updatable_{true};
-  bool force_single_root_{false};
+
+  int node_count_{0};
+  size_t bytes_in_graph_{0};
   int max_ops_per_graph_;
   int max_mb_per_graph_;
+
+  bool force_single_root_;
 };
 
 class MLX_API Device {

mlx/backend/cuda/eval.cpp

@@ -70,6 +70,10 @@ void eval(array& arr) {
     encoder.add_temporary(s);
   }
 
+#ifndef NDEBUG
+  encoder.verify_graph(arr.primitive().name());
+#endif
+
   if (encoder.needs_commit()) {
     scheduler::notify_new_task(stream);
     encoder.add_completed_handler(

mlx/backend/cuda/indexing.cpp

@@ -479,9 +479,6 @@ void MaskedScatter::eval_gpu(const std::vector<array>& inputs, array& out) {
   array mask_flat = (mask.ndim() == 1)
       ? reshape_in_eval(mask, {mask.shape(0), 1}, s)
       : flatten_in_eval(mask, 1, -1, s);
-  if (gpu_ptr<void>(mask_flat) != gpu_ptr<void>(mask)) {
-    encoder.add_temporary(mask_flat);
-  }
   if (!mask_flat.flags().row_contiguous) {
     mask_flat = contiguous_copy_gpu(mask_flat, s);
     encoder.add_temporary(mask_flat);

mlx/backend/gpu/copy.cpp

@@ -56,7 +56,7 @@ array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s) {
 
 array reshape_in_eval(const array& x, Shape shape, Stream s) {
   array out(std::move(shape), x.dtype(), nullptr, {});
-  reshape_gpu(x, out, s);
+  reshape_gpu(x, out, true, s);
   return out;
 }
 

mlx/backend/gpu/copy.h

@@ -48,7 +48,11 @@ void fill_gpu(const array& val, array& out, const Stream& s);
 array contiguous_copy_gpu(const array& arr, const Stream& s);
 
 // Copy data from |in| and transpose to |out|'s shape.
-void reshape_gpu(const array& in, array& out, Stream s);
+void reshape_gpu(
+    const array& in,
+    array& out,
+    bool add_out_to_temporary_on_allocation,
+    Stream s);
 
 // Like the normal ops but safe to call in eval_gpu.
 array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s);

mlx/backend/gpu/primitives.cpp

@@ -30,11 +30,7 @@ void AsStrided::eval_gpu(const std::vector<array>& inputs, array& out) {
     eval(inputs[0], out);
   } else {
     auto copy = contiguous_copy_gpu(inputs[0], stream());
-#if defined(MLX_USE_CUDA)
-    cu::get_command_encoder(stream()).add_temporary(copy);
-#else
-    metal::get_command_encoder(stream()).add_temporary(copy);
-#endif
+    // Donated to out.
     eval(copy, out);
   }
 }
@@ -183,7 +179,7 @@ void Full::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Flatten::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Flatten::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
@@ -209,7 +205,7 @@ void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Reshape::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Reshape::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void Split::eval_gpu(
@@ -248,7 +244,7 @@ void Transpose::eval_gpu(const std::vector<array>& inputs, array& out) {
 
 void Unflatten::eval_gpu(const std::vector<array>& inputs, array& out) {
   MLX_PROFILER_RANGE("Unflatten::eval_gpu");
-  reshape_gpu(inputs[0], out, stream());
+  reshape_gpu(inputs[0], out, false, stream());
 }
 
 void View::eval_gpu(const std::vector<array>& inputs, array& out) {