From 5140cbe4acf15bfa4ceae264f748b04ec69bb985 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Thu, 31 Jul 2025 03:35:41 -0700
Subject: [PATCH] Add more nvtx loggings

---
 mlx/backend/cuda/device.cpp              | 1 +
 mlx/backend/cuda/quantized/quantized.cpp | 3 +++
 mlx/backend/cuda/ternary.cu              | 2 +-
 mlx/backend/gpu/primitives.cpp           | 1 +
 4 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
index ddd235d58..79eb55406 100644
--- a/mlx/backend/cuda/device.cpp
+++ b/mlx/backend/cuda/device.cpp
@@ -269,6 +269,7 @@ void CommandEncoder::add_graph_node(cudaGraph_t child) {
 }
 
 void CommandEncoder::commit() {
+  nvtx3::scoped_range r("CommandEncoder::commit");
   if (!temporaries_.empty()) {
     add_completed_handler([temporaries = std::move(temporaries_)]() {});
   }
diff --git a/mlx/backend/cuda/quantized/quantized.cpp b/mlx/backend/cuda/quantized/quantized.cpp
index f495af53b..256f2c7d5 100644
--- a/mlx/backend/cuda/quantized/quantized.cpp
+++ b/mlx/backend/cuda/quantized/quantized.cpp
@@ -5,6 +5,8 @@
 #include "mlx/backend/gpu/copy.h"
 #include "mlx/fast_primitives.h"
 
+#include <nvtx3/nvtx3.hpp>
+
 namespace mlx::core {
 
 namespace {
@@ -42,6 +44,7 @@ inline array ensure_row_contiguous_matrix(
 void fast::AffineQuantize::eval_gpu(
     const std::vector<array>& inputs,
     std::vector<array>& outputs) {
+  nvtx3::scoped_range r("AffineQuantize::eval_gpu");
   auto& s = stream();
   auto& d = cu::device(s.device);
   auto& enc = d.get_command_encoder(s);
diff --git a/mlx/backend/cuda/ternary.cu b/mlx/backend/cuda/ternary.cu
index 58d3fa119..bc4097d99 100644
--- a/mlx/backend/cuda/ternary.cu
+++ b/mlx/backend/cuda/ternary.cu
@@ -192,7 +192,7 @@ void ternary_op_gpu(
 }
 
 void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
-  nvtx3::scoped_range r("select::eval_gpu");
+  nvtx3::scoped_range r("Select::eval_gpu");
   auto& s = out.primitive().stream();
   ternary_op_gpu<cu::Select>(inputs, out, s);
 }
diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp
index 1adb85918..56d389b4f 100644
--- a/mlx/backend/gpu/primitives.cpp
+++ b/mlx/backend/gpu/primitives.cpp
@@ -133,6 +133,7 @@ void NumberOfElements::eval_gpu(const std::vector<array>& inputs, array& out) {
 }
 
 void Pad::eval_gpu(const std::vector<array>& inputs, array& out) {
+  MLX_PROFILER_RANGE("Pad::eval_gpu");
   // Inputs must be base input array and scalar val array
   assert(inputs.size() == 2);
   auto& in = inputs[0];