Add more nvtx loggings

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_)]() {});
   }

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);

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);
 }

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];