chore: import upstream snapshot with attribution

paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h

@@ -0,0 +1,318 @@
+/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#pragma once
+
+#include <type_traits>
+
+#include "glog/logging.h"
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/kernels/baddbmm_grad_kernel.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+
+COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
+
+namespace phi {
+
+template <typename T>
+struct BCopyOrScaleFunctor {
+  BCopyOrScaleFunctor(const float scale, const T* x, T* output, int64_t numel)
+      : scale_(scale), x_(x), output_(output), numel_(numel) {}
+
+  HOSTDEVICE void operator()(int64_t idx) const {
+    using MPType = typename MPTypeTrait<T>::Type;
+    const MPType mp_scale = static_cast<MPType>(scale_);
+    const MPType mp_x = static_cast<MPType>(x_[idx]);
+    output_[idx] = static_cast<T>(mp_scale * mp_x);
+  }
+
+ private:
+  const float scale_;
+  const T* x_;
+  T* output_;
+  int64_t numel_;
+};
+
+using Array1 = Eigen::DSizes<int64_t, 1>;
+using Array2 = Eigen::DSizes<int64_t, 2>;
+using Array3 = Eigen::DSizes<int64_t, 3>;
+
+template <typename T, typename Context>
+void BaddbmmGradKernel(const Context& dev_ctx,
+                       const DenseTensor& input,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       const DenseTensor& out_grad,
+                       float alpha,
+                       float beta,
+                       DenseTensor* input_grad,
+                       DenseTensor* x_grad,
+                       DenseTensor* y_grad) {
+  using MPType = typename MPTypeTrait<T>::Type;
+  bool is_float16_or_bfloat16 = false;
+  if (std::is_same<T, float16>::value || std::is_same<T, bfloat16>::value) {
+    is_float16_or_bfloat16 = true;
+  }
+
+  auto in_dims = input.dims();
+  if (input.dims().size() == 2) {
+    in_dims = {input.dims()[0], 1, input.dims()[1]};
+    input_grad->Resize(in_dims);
+  }
+  int64_t total_elems = 0;
+
+  VLOG(3) << "alpha: " << alpha << " beta: " << beta;
+
+  if (input_grad != nullptr) {
+    input_grad->set_lod(out_grad.lod());
+  }
+  if (x_grad != nullptr) {
+    x_grad->set_lod(x.lod());
+  }
+  if (y_grad != nullptr) {
+    y_grad->set_lod(y.lod());
+  }
+
+  auto blas = funcs::GetBlas<Context, T>(dev_ctx);
+  auto mt_blas = funcs::GetBlas<Context, MPType>(dev_ctx);
+  if (input_grad) {
+    dev_ctx.template Alloc<T>(input_grad);
+    total_elems = in_dims[0] * in_dims[1] * in_dims[2];
+    auto& place = *dev_ctx.eigen_device();
+    auto eigen_dout = EigenTensor<T, 3>::From(out_grad);
+    auto eigen_dinput = EigenTensor<T, 3>::From(*input_grad);
+
+    bool batch_compress = in_dims[0] != out_grad.dims()[0];
+    bool row_compress = in_dims[1] != out_grad.dims()[1];
+    bool col_compress = in_dims[2] != out_grad.dims()[2];
+    auto eigen_dinput_shape = Array3(
+        input_grad->dims()[0], input_grad->dims()[1], input_grad->dims()[2]);
+
+    if (batch_compress && row_compress && col_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum().eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum()
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (batch_compress && row_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array2(0, 1)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array2(0, 1))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (batch_compress && col_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array2(0, 2)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array2(0, 2))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (row_compress && col_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array2(1, 2)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array2(1, 2))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (batch_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array1(0)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array1(0))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (row_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array1(1)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array1(1))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else if (col_compress) {
+      if (!is_float16_or_bfloat16) {
+        eigen_dinput.device(place) =
+            eigen_dout.sum(Array1(2)).eval().reshape(eigen_dinput_shape);
+      } else {
+        eigen_dinput.device(place) = eigen_dout.template cast<MPType>()
+                                         .sum(Array1(2))
+                                         .eval()
+                                         .reshape(eigen_dinput_shape)
+                                         .template cast<T>();
+      }
+    } else {
+      // The VCOPY does not support the float16, bfloat16
+      if (!is_float16_or_bfloat16) {
+        mt_blas.VCOPY(
+            total_elems, out_grad.data<MPType>(), input_grad->data<MPType>());
+      } else {
+        funcs::ForRange<Context> for_range(dev_ctx, total_elems);
+        BCopyOrScaleFunctor<T> functor(
+            1, out_grad.data<T>(), input_grad->data<T>(), total_elems);
+        for_range(functor);
+      }
+    }
+
+    // The SCAL does not support the float16, bfloat16
+    if (!is_float16_or_bfloat16) {
+      mt_blas.SCAL(total_elems, beta, input_grad->data<MPType>());
+    } else {
+      funcs::ForRange<Context> for_range(dev_ctx, total_elems);
+      BCopyOrScaleFunctor<T> functor(
+          beta, input_grad->data<T>(), input_grad->data<T>(), total_elems);
+      for_range(functor);
+    }
+  }
+  if (x_grad) {
+    dev_ctx.template Alloc<T>(x_grad);
+    total_elems = x.dims()[0] * x.dims()[1] * x.dims()[2];
+    // x_grad = alpha * out_grad @ y^T
+    // out_grad: [B, M, N], y: [B, K, N], x_grad: [B, M, K]
+    int64_t B_dim = x.dims()[0];
+    int64_t M_dim = x.dims()[1];
+    int64_t K_dim = x.dims()[2];
+    int64_t N_dim = y.dims()[2];
+    if constexpr (std::is_same_v<MPType, float>) {
+      float gemm_alpha = FLAGS_use_accuracy_compatible_kernel ? 1.0f : alpha;
+      float zero = 0.0f;
+      blas.BatchedGEMM(CblasNoTrans,
+                       CblasTrans,
+                       M_dim,
+                       K_dim,
+                       N_dim,
+                       gemm_alpha,
+                       out_grad.data<T>(),
+                       y.data<T>(),
+                       zero,
+                       x_grad->data<T>(),
+                       B_dim,
+                       M_dim * N_dim,
+                       K_dim * N_dim);
+    } else {
+      T gemm_alpha = FLAGS_use_accuracy_compatible_kernel
+                         ? static_cast<T>(1)
+                         : static_cast<T>(alpha);
+      T zero = static_cast<T>(0);
+      blas.BatchedGEMM(CblasNoTrans,
+                       CblasTrans,
+                       M_dim,
+                       K_dim,
+                       N_dim,
+                       gemm_alpha,
+                       out_grad.data<T>(),
+                       y.data<T>(),
+                       zero,
+                       x_grad->data<T>(),
+                       B_dim,
+                       M_dim * N_dim,
+                       K_dim * N_dim);
+    }
+    if (FLAGS_use_accuracy_compatible_kernel) {
+      if (!is_float16_or_bfloat16) {
+        mt_blas.SCAL(total_elems, alpha, x_grad->data<MPType>());
+      } else {
+        funcs::ForRange<Context> for_range(dev_ctx, total_elems);
+        BCopyOrScaleFunctor<T> functor(
+            alpha, x_grad->data<T>(), x_grad->data<T>(), total_elems);
+        for_range(functor);
+      }
+    }
+  }
+  if (y_grad) {
+    dev_ctx.template Alloc<T>(y_grad);
+    total_elems = y.dims()[0] * y.dims()[1] * y.dims()[2];
+    // y_grad = alpha * x^T @ out_grad
+    // x: [B, M, K], out_grad: [B, M, N], y_grad: [B, K, N]
+    int64_t B_dim = x.dims()[0];
+    int64_t M_dim = x.dims()[1];
+    int64_t K_dim = x.dims()[2];
+    int64_t N_dim = y.dims()[2];
+    if constexpr (std::is_same_v<MPType, float>) {
+      float gemm_alpha = FLAGS_use_accuracy_compatible_kernel ? 1.0f : alpha;
+      float zero = 0.0f;
+      blas.BatchedGEMM(CblasTrans,
+                       CblasNoTrans,
+                       K_dim,
+                       N_dim,
+                       M_dim,
+                       gemm_alpha,
+                       x.data<T>(),
+                       out_grad.data<T>(),
+                       zero,
+                       y_grad->data<T>(),
+                       B_dim,
+                       M_dim * K_dim,
+                       M_dim * N_dim);
+    } else {
+      T gemm_alpha = FLAGS_use_accuracy_compatible_kernel
+                         ? static_cast<T>(1)
+                         : static_cast<T>(alpha);
+      T zero = static_cast<T>(0);
+      blas.BatchedGEMM(CblasTrans,
+                       CblasNoTrans,
+                       K_dim,
+                       N_dim,
+                       M_dim,
+                       gemm_alpha,
+                       x.data<T>(),
+                       out_grad.data<T>(),
+                       zero,
+                       y_grad->data<T>(),
+                       B_dim,
+                       M_dim * K_dim,
+                       M_dim * N_dim);
+    }
+    if (FLAGS_use_accuracy_compatible_kernel) {
+      if (!is_float16_or_bfloat16) {
+        mt_blas.SCAL(total_elems, alpha, y_grad->data<MPType>());
+      } else {
+        funcs::ForRange<Context> for_range(dev_ctx, total_elems);
+        BCopyOrScaleFunctor<T> functor(
+            alpha, y_grad->data<T>(), y_grad->data<T>(), total_elems);
+        for_range(functor);
+      }
+    }
+  }
+}
+
+}  // namespace phi