[CUDA] Switch to CUDA graphs (#2317)

* cuda graph prototype

fix signal bug + start to add dependencies

capture more

capture more ops

remaining ops

fix reduce and rope deps

add concurrent context

try update, but not working

cosistent topology order

use node api

use node api directly to reduce overhead

fix bug

use kernels in unary

cache graph

format

fix synchronization

format

* comment

mlx/backend/cuda/layer_norm.cu

@@ -258,23 +258,23 @@ void LayerNorm::eval_gpu(
   encoder.set_input_array(w);
   encoder.set_input_array(b);
   encoder.set_output_array(out);
-  encoder.launch_kernel([&](cudaStream_t stream) {
-    dispatch_float_types(out.dtype(), "layernorm", [&](auto type_tag) {
-      constexpr uint32_t N_READS = 4;
-      dispatch_block_dim(
-          cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-            using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
-            auto kernel = cu::layer_norm<DataType, block_dim(), N_READS>;
-            kernel<<<n_rows, block_dim(), 0, stream>>>(
-                x.data<DataType>(),
-                w.data<DataType>(),
-                b.data<DataType>(),
-                out.data<DataType>(),
-                eps_,
-                axis_size,
-                w_stride,
-                b_stride);
-          });
+  dispatch_float_types(out.dtype(), "layernorm", [&](auto type_tag) {
+    constexpr uint32_t N_READS = 4;
+    dispatch_block_dim(cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+      using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+      auto kernel = cu::layer_norm<DataType, block_dim(), N_READS>;
+      encoder.add_kernel_node(
+          kernel,
+          n_rows,
+          block_dim(),
+          x.data<DataType>(),
+          w.data<DataType>(),
+          b.data<DataType>(),
+          out.data<DataType>(),
+          eps_,
+          axis_size,
+          w_stride,
+          b_stride);
     });
   });
 }
@@ -289,21 +289,25 @@ void LayerNormVJP::eval_gpu(
   // Ensure row contiguity. We could relax this step by checking that the array
   // is contiguous (no broadcasts or holes) and that the input strides are the
   // same as the cotangent strides but for now this is simpler.
-  auto check_input = [&s](const array& x) -> std::pair<array, bool> {
+  auto check_input = [&s](const array& x, bool& copied) {
     if (x.flags().row_contiguous) {
-      return {x, false};
+      copied = false;
+      return x;
     }
+    copied = true;
     array x_copy(x.shape(), x.dtype(), nullptr, {});
     copy_gpu(x, x_copy, CopyType::General, s);
-    return {x_copy, true};
+    return x_copy;
   };
   bool donate_x = inputs[0].is_donatable();
   bool donate_g = inputs[3].is_donatable();
-  auto [x, copied] = check_input(inputs[0]);
+  bool copied;
+  auto x = check_input(inputs[0], copied);
   donate_x |= copied;
   const array& w = inputs[1];
   const array& b = inputs[2];
-  auto [g, g_copied] = check_input(inputs[3]);
+  bool g_copied;
+  auto g = check_input(inputs[3], g_copied);
   donate_g |= g_copied;
   array& gx = outputs[0];
   array& gw = outputs[1];
@@ -334,8 +338,10 @@ void LayerNormVJP::eval_gpu(
   // gradient accumulators.
   array gw_temp =
       (has_w) ? array({n_rows, x.shape().back()}, gw.dtype(), nullptr, {}) : w;
+  bool g_in_gw = false;
   if (has_w) {
     if (!g_in_gx && donate_g) {
+      g_in_gw = true;
       gw_temp.copy_shared_buffer(g);
     } else {
       gw_temp.set_data(allocator::malloc(gw_temp.nbytes()));
@@ -343,41 +349,47 @@ void LayerNormVJP::eval_gpu(
     }
   }
 
-  // Finish with the gradient for b in case we had a b.
-  if (gb.ndim() == 1 && gb.size() == axis_size) {
+  // The gradient for b in case we had a b.
+  bool has_gb = (gb.ndim() == 1 && gb.size() == axis_size);
+  if (has_gb) {
     ReductionPlan plan(
         ReductionOpType::ContiguousStridedReduce, {n_rows}, {axis_size});
     col_reduce(encoder, g, gb, Reduce::ReduceType::Sum, {0}, plan);
   }
 
+  // Insert dependency if `g` was donated
+  if ((g_in_gx || g_in_gw) && has_gb) {
+    encoder.set_input_array(gb);
+  }
   encoder.set_input_array(x);
   encoder.set_input_array(w);
   encoder.set_input_array(g);
   encoder.set_output_array(gx);
   encoder.set_output_array(gw_temp);
-  encoder.launch_kernel([&, x = x, g = g](cudaStream_t stream) {
-    dispatch_float_types(gx.dtype(), "layernorm_vjp", [&](auto type_tag) {
-      dispatch_bool(has_w, [&](auto has_w_constant) {
-        constexpr int N_READS = 4;
-        dispatch_block_dim(
-            cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
-              using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
-              auto kernel = cu::layer_norm_vjp<
-                  DataType,
-                  has_w_constant.value,
-                  block_dim(),
-                  N_READS>;
-              kernel<<<n_rows, block_dim(), 0, stream>>>(
-                  x.data<DataType>(),
-                  w.data<DataType>(),
-                  g.data<DataType>(),
-                  gx.data<DataType>(),
-                  gw_temp.data<DataType>(),
-                  eps_,
-                  axis_size,
-                  w_stride);
-            });
-      });
+  dispatch_float_types(gx.dtype(), "layernorm_vjp", [&](auto type_tag) {
+    dispatch_bool(has_w, [&](auto has_w_constant) {
+      constexpr int N_READS = 4;
+      dispatch_block_dim(
+          cuda::ceil_div(axis_size, N_READS), [&](auto block_dim) {
+            using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+            auto kernel = cu::layer_norm_vjp<
+                DataType,
+                has_w_constant.value,
+                block_dim(),
+                N_READS>;
+            encoder.add_kernel_node(
+                kernel,
+                n_rows,
+                block_dim(),
+                x.data<DataType>(),
+                w.data<DataType>(),
+                g.data<DataType>(),
+                gx.data<DataType>(),
+                gw_temp.data<DataType>(),
+                eps_,
+                axis_size,
+                w_stride);
+          });
     });
   });