add deterministic embedding grad kernel (#50494)

* add deterministic embedding grad kernel

* minor change

* minor change

* Add new FLAG to enable deterministic embedding

* Update embedding deterministic kernel

paddle/phi/core/flags.cc

@@ -232,6 +232,22 @@ PADDLE_DEFINE_EXPORTED_bool(
     "operator. The autotuning algorithm may be non-deterministic. If "
     "true, the algorithm is deterministic.");
 
+/**
+ * CUDA related FLAG
+ * Name: FLAGS_embedding_deterministic
+ * Since Version: 2.5
+ * Value Range: bool, default=false
+ * Example:
+ * Note: whether to use deterministic algorithm in embedding op.
+ *       If true, it will use deterministic CUDA kernel in embedding op.
+ */
+PADDLE_DEFINE_EXPORTED_bool(
+    embedding_deterministic,
+    false,
+    "Whether allow using an deterministic algorithm for embedding "
+    "operator. The deterministic algorithm may be slower. If "
+    "true, the algorithm is deterministic.");
+
 /**
  * CUDNN related FLAG
  * Name: FLAGS_conv_workspace_size_limit

paddle/phi/kernels/gpu/embedding_grad_kernel.cu

@@ -25,10 +25,20 @@
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/embedding_util.h"
 
-DECLARE_bool(cudnn_deterministic);
+DECLARE_bool(embedding_deterministic);
 
 namespace phi {
 
+#ifdef PADDLE_WITH_HIP
+#define WARP_SIZE 64
+#define BLOCKDIMY 16
+#else
+#define WARP_SIZE 32
+#define BLOCKDIMY 32
+#endif
+
+#define MASK 0xffffffff
+
 template <typename InT, typename OutT>
 __global__ void InputTypeConvert(const InT* in_ids,
                                  const int64_t K,
@@ -63,6 +73,93 @@ __global__ void EmbeddingGrad(T* table,
   }
 }
 
+template <typename T, typename IdT>
+__global__ void EmbeddingGradDeterministic(T* table,
+                                           const T* output,
+                                           const IdT* ids,
+                                           const int64_t K,
+                                           const int64_t D) {
+  extern __shared__ char buf[];
+  T* smem = reinterpret_cast<T*>(buf);
+  T* my_s = smem + WARP_SIZE * threadIdx.y;
+  int64_t* indices_batch =
+      reinterpret_cast<int64_t*>(buf + sizeof(T) * WARP_SIZE * BLOCKDIMY);
+
+  const int stride = static_cast<int>(D);
+
+  const int feature = threadIdx.x + blockIdx.x * WARP_SIZE;
+
+  // To ensure determinism. If any other warps pulled grad data targeting
+  // dst_row, we elect the first warp in each matching group as the leader.
+  // Each leader warp serializes the accumulates targeting dst_row in shared
+  // memory, then adding the accumulated buffer to dst_row in table.
+  for (int batch_start = 0; batch_start < K;
+       batch_start += WARP_SIZE * BLOCKDIMY) {
+    int tid = threadIdx.x + threadIdx.y * WARP_SIZE;
+    if (batch_start + tid < K)
+      indices_batch[tid] = static_cast<int64_t>(ids[batch_start + tid]);
+
+    int batch_end =
+        min(static_cast<int64_t>(batch_start + WARP_SIZE * BLOCKDIMY), K);
+
+    // Loop over the batch of <= 1024 loaded indices in chunks of BLOCKDIMY
+    for (int chunk_start = batch_start; chunk_start < batch_end;
+         chunk_start += BLOCKDIMY) {
+      // This sync makes sure that indices_batch is ready and match-group
+      // leaders are done with their accumulates before other warps start
+      // loading again.
+      __syncthreads();
+
+      int n_this_chunk = min(batch_end - chunk_start, BLOCKDIMY);
+
+      int64_t src_row = static_cast<int64_t>(chunk_start + threadIdx.y);
+      int64_t dst_row = indices_batch[src_row - batch_start];
+      if (src_row < K && feature < stride)
+        my_s[threadIdx.x] = static_cast<T>(output[src_row * D + feature]);
+
+      __syncthreads();
+
+      if (src_row < K) {
+        int match_found_this_thread = 0;
+        if (threadIdx.x < n_this_chunk) {
+          match_found_this_thread =
+              (dst_row ==
+               indices_batch[chunk_start - batch_start + threadIdx.x]);
+        }
+#ifdef PADDLE_WITH_HIP
+        unsigned long long int matchmask =      // NOLINT
+            __ballot(match_found_this_thread);  // NOLINT
+        int first_remaining_peer = __ffsll(matchmask) - 1;
+#else
+        // If and only if match_found_this_thread of the Nth thread is non-zero,
+        // set the Nth bit of matchmask to 1.
+        unsigned int matchmask = __ballot_sync(MASK, match_found_this_thread);
+        // Find the position of the first bit set to 1 in matchmask.
+        int first_remaining_peer = __ffs(matchmask) - 1;
+#endif
+
+        // select lowest-indexed warp as the leader
+        if (threadIdx.y == first_remaining_peer) {
+          // Set the first bit 1 in matchmask to 0.
+          matchmask ^= (1 << first_remaining_peer);
+          while (matchmask) {
+#ifdef PADDLE_WITH_HIP
+            first_remaining_peer = __ffsll(matchmask) - 1;
+#else
+            first_remaining_peer = __ffs(matchmask) - 1;
+#endif
+            my_s[threadIdx.x] +=
+                smem[threadIdx.x + WARP_SIZE * first_remaining_peer];
+            matchmask ^= (1 << first_remaining_peer);
+          }
+          if (feature < stride)
+            table[dst_row * D + feature] += static_cast<T>(my_s[threadIdx.x]);
+        }
+      }
+    }
+  }
+}
+
 template <typename T, typename Context>
 struct EmbeddingGradCUDAFunctor {
   EmbeddingGradCUDAFunctor(const Context& dev_ctx,
@@ -102,17 +199,22 @@ struct EmbeddingGradCUDAFunctor {
           cudaMemsetAsync(d_table, 0, N * D * sizeof(T), dev_ctx_.stream()));
 #endif
 
-      const int gridx = 2 * dev_ctx_.GetSMCount();
-      dim3 threads(128, 8);
-      dim3 grids(gridx, 1);
-
-      if (FLAGS_cudnn_deterministic) {
-        VLOG(2) << "Run grad kernel of embedding with single thread.";
-        grids.x = 1;
-        threads.y = 1;
+      if (FLAGS_embedding_deterministic) {
+        dim3 threads(WARP_SIZE, BLOCKDIMY);
+        dim3 grids(static_cast<int>((D + WARP_SIZE - 1) / WARP_SIZE));
+        EmbeddingGradDeterministic<T, IdT>
+            <<<grids,
+               threads,
+               sizeof(T) * WARP_SIZE * BLOCKDIMY +
+                   sizeof(int) * WARP_SIZE * BLOCKDIMY,
+               dev_ctx_.stream()>>>(d_table, d_output, ids, K, D);
+      } else {
+        const int gridx = 2 * dev_ctx_.GetSMCount();
+        dim3 threads(128, 8);
+        dim3 grids(gridx, 1);
+        EmbeddingGrad<T, IdT><<<grids, threads, 0, dev_ctx_.stream()>>>(
+            d_table, d_output, ids, N, K, D);
       }
-      EmbeddingGrad<T, IdT><<<grids, threads, 0, dev_ctx_.stream()>>>(
-          d_table, d_output, ids, N, K, D);
     }
   }