From 73473ac26c03395cfaaa3087c4fa82ad318266e5 Mon Sep 17 00:00:00 2001
From: Shaojie WANG <wsjmessi@163.com>
Date: Sat, 8 Apr 2023 19:17:54 -0700
Subject: [PATCH] register bf16 for c ops (#52641)

---
 paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc   | 3 +++
 paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc   | 2 +-
 paddle/fluid/operators/collective/c_concat_op.cu.cc          | 3 +++
 paddle/fluid/operators/collective/c_embedding_op.cu          | 2 ++
 paddle/fluid/operators/collective/c_identity_op.cu.cc        | 3 +++
 paddle/fluid/operators/collective/c_split_op.cu              | 3 +++
 paddle/fluid/operators/collective/partial_allgather_op.cu.cc | 3 +++
 paddle/fluid/operators/collective/partial_recv_op.cu.cc      | 3 +++
 paddle/fluid/operators/collective/partial_send_op.cu.cc      | 3 +++
 paddle/fluid/operators/collective/recv_v2_op.cu.cc           | 2 +-
 paddle/fluid/operators/collective/send_v2_op.cu.cc           | 2 +-
 11 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc b/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc
index fb073c8173d..3726ff8ac90 100644
--- a/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc
+++ b/paddle/fluid/operators/collective/c_allreduce_max_op.cu.cc
@@ -20,6 +20,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(
     c_allreduce_max,
     ops::CAllReduceOpCUDAKernel<ops::kRedMax, float>,
+#if NCCL_VERSION_CODE >= 21000
+    ops::CAllReduceOpCUDAKernel<ops::kRedMax, plat::bfloat16>,
+#endif
     ops::CAllReduceOpCUDAKernel<ops::kRedMax, double>,
     ops::CAllReduceOpCUDAKernel<ops::kRedMax, int>,
     ops::CAllReduceOpCUDAKernel<ops::kRedMax, int64_t>,
diff --git a/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc b/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc
index 2072f4c0cb3..5e00c3bdba3 100644
--- a/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc
+++ b/paddle/fluid/operators/collective/c_allreduce_sum_op.cu.cc
@@ -20,7 +20,7 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(
     c_allreduce_sum,
     ops::CAllReduceOpCUDAKernel<ops::kRedSum, float>,
-#if CUDNN_VERSION_MIN(8, 1, 0) && NCCL_VERSION_CODE >= 21000
+#if NCCL_VERSION_CODE >= 21000
     ops::CAllReduceOpCUDAKernel<ops::kRedSum, plat::bfloat16>,
 #endif
     ops::CAllReduceOpCUDAKernel<ops::kRedSum, double>,
diff --git a/paddle/fluid/operators/collective/c_concat_op.cu.cc b/paddle/fluid/operators/collective/c_concat_op.cu.cc
index 74bdd2b63ae..1abe8a51b83 100644
--- a/paddle/fluid/operators/collective/c_concat_op.cu.cc
+++ b/paddle/fluid/operators/collective/c_concat_op.cu.cc
@@ -132,6 +132,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(c_concat,
                         ops::CConcatOpCUDAKernel<float>,
                         ops::CConcatOpCUDAKernel<double>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::CConcatOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::CConcatOpCUDAKernel<int>,
                         ops::CConcatOpCUDAKernel<int64_t>,
                         ops::CConcatOpCUDAKernel<plat::float16>);
diff --git a/paddle/fluid/operators/collective/c_embedding_op.cu b/paddle/fluid/operators/collective/c_embedding_op.cu
index cddbd162571..04839660b09 100644
--- a/paddle/fluid/operators/collective/c_embedding_op.cu
+++ b/paddle/fluid/operators/collective/c_embedding_op.cu
@@ -253,8 +253,10 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(c_embedding,
                         ops::CEmbeddingCUDAKernel<float>,
                         ops::CEmbeddingCUDAKernel<double>,
+                        ops::CEmbeddingCUDAKernel<plat::bfloat16>,
                         ops::CEmbeddingCUDAKernel<plat::float16>);
 REGISTER_OP_CUDA_KERNEL(c_embedding_grad,
                         ops::CEmbeddingGradCUDAKernel<float>,
                         ops::CEmbeddingGradCUDAKernel<double>,
+                        ops::CEmbeddingGradCUDAKernel<plat::bfloat16>,
                         ops::CEmbeddingGradCUDAKernel<plat::float16>);
diff --git a/paddle/fluid/operators/collective/c_identity_op.cu.cc b/paddle/fluid/operators/collective/c_identity_op.cu.cc
index 0b2f5b7eb1a..f8ddb6ab010 100644
--- a/paddle/fluid/operators/collective/c_identity_op.cu.cc
+++ b/paddle/fluid/operators/collective/c_identity_op.cu.cc
@@ -19,6 +19,9 @@ namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(c_identity,
                         ops::CIdentityOpKernel<float>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::CIdentityOpKernel<plat::bfloat16>,
+#endif
                         ops::CIdentityOpKernel<double>,
                         ops::CIdentityOpKernel<int>,
                         ops::CIdentityOpKernel<int64_t>,
diff --git a/paddle/fluid/operators/collective/c_split_op.cu b/paddle/fluid/operators/collective/c_split_op.cu
index 8bf887b954a..e3dfbb3a5f9 100644
--- a/paddle/fluid/operators/collective/c_split_op.cu
+++ b/paddle/fluid/operators/collective/c_split_op.cu
@@ -117,6 +117,9 @@ namespace plat = paddle::platform;
 
 REGISTER_OP_CUDA_KERNEL(c_split,
                         ops::CSplitOpCUDAKernel<float>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::CSplitOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::CSplitOpCUDAKernel<double>,
                         ops::CSplitOpCUDAKernel<int>,
                         ops::CSplitOpCUDAKernel<int64_t>,
diff --git a/paddle/fluid/operators/collective/partial_allgather_op.cu.cc b/paddle/fluid/operators/collective/partial_allgather_op.cu.cc
index 6bc18254737..d0b131c8d76 100644
--- a/paddle/fluid/operators/collective/partial_allgather_op.cu.cc
+++ b/paddle/fluid/operators/collective/partial_allgather_op.cu.cc
@@ -110,6 +110,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(partial_allgather,
                         ops::PartialAllGatherOpCUDAKernel<float>,
                         ops::PartialAllGatherOpCUDAKernel<double>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialAllGatherOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialAllGatherOpCUDAKernel<int>,
                         ops::PartialAllGatherOpCUDAKernel<int64_t>,
                         ops::PartialAllGatherOpCUDAKernel<plat::float16>);
diff --git a/paddle/fluid/operators/collective/partial_recv_op.cu.cc b/paddle/fluid/operators/collective/partial_recv_op.cu.cc
index ad44ce6a109..a9a98d3f3ee 100644
--- a/paddle/fluid/operators/collective/partial_recv_op.cu.cc
+++ b/paddle/fluid/operators/collective/partial_recv_op.cu.cc
@@ -121,6 +121,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(partial_recv,
                         ops::PartialRecvOpCUDAKernel<float>,
                         ops::PartialRecvOpCUDAKernel<double>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialRecvOpCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialRecvOpCUDAKernel<int>,
                         ops::PartialRecvOpCUDAKernel<int64_t>,
                         ops::PartialRecvOpCUDAKernel<plat::float16>);
diff --git a/paddle/fluid/operators/collective/partial_send_op.cu.cc b/paddle/fluid/operators/collective/partial_send_op.cu.cc
index fb49318c012..6d99e246c9d 100644
--- a/paddle/fluid/operators/collective/partial_send_op.cu.cc
+++ b/paddle/fluid/operators/collective/partial_send_op.cu.cc
@@ -120,6 +120,9 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(partial_send,
                         ops::PartialSendCUDAKernel<float>,
                         ops::PartialSendCUDAKernel<double>,
+#if NCCL_VERSION_CODE >= 21000
+                        ops::PartialSendCUDAKernel<plat::bfloat16>,
+#endif
                         ops::PartialSendCUDAKernel<int>,
                         ops::PartialSendCUDAKernel<int64_t>,
                         ops::PartialSendCUDAKernel<plat::float16>);
diff --git a/paddle/fluid/operators/collective/recv_v2_op.cu.cc b/paddle/fluid/operators/collective/recv_v2_op.cu.cc
index ec18a172e1f..2cbdae8b73d 100644
--- a/paddle/fluid/operators/collective/recv_v2_op.cu.cc
+++ b/paddle/fluid/operators/collective/recv_v2_op.cu.cc
@@ -236,7 +236,7 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(recv_v2,
                         ops::RecvOpV2CUDAKernel<float>,
                         ops::RecvOpV2CUDAKernel<double>,
-#if CUDNN_VERSION_MIN(8, 1, 0) && NCCL_VERSION_CODE >= 21000
+#if NCCL_VERSION_CODE >= 21000
                         ops::RecvOpV2CUDAKernel<plat::bfloat16>,
 #endif
                         ops::RecvOpV2CUDAKernel<int>,
diff --git a/paddle/fluid/operators/collective/send_v2_op.cu.cc b/paddle/fluid/operators/collective/send_v2_op.cu.cc
index 37b18703031..299f0595713 100644
--- a/paddle/fluid/operators/collective/send_v2_op.cu.cc
+++ b/paddle/fluid/operators/collective/send_v2_op.cu.cc
@@ -222,7 +222,7 @@ namespace plat = paddle::platform;
 REGISTER_OP_CUDA_KERNEL(send_v2,
                         ops::SendOpV2CUDAKernel<float>,
                         ops::SendOpV2CUDAKernel<double>,
-#if CUDNN_VERSION_MIN(8, 1, 0) && NCCL_VERSION_CODE >= 21000
+#if NCCL_VERSION_CODE >= 21000
                         ops::SendOpV2CUDAKernel<plat::bfloat16>,
 #endif
                         ops::SendOpV2CUDAKernel<int>,
-- 
GitLab