optimize elementwise_mul_grad using new interfaces

paddle/fluid/operators/elementwise/elementwise_functor.h

@@ -194,5 +194,47 @@ struct FMinFunctor<paddle::platform::float16> {
  }
 };
 
+template <typename T>
+struct MulGradFunctor {
+ inline HOSTDEVICE T operator()(const T& a, const T& b) const { return a * b; }
+};
+template <typename T>
+struct MulGradFunctor<Complex<T>> {
+ inline HOSTDEVICE Complex<T> operator()(const Complex<T>& a,
+ const Complex<T>& b) const {
+ Complex<T> b_conj(b.real, -b.imag);
+ return a * b_conj;
+ }
+};
+
+template <typename InT, typename OutT>
+struct MulGradXYFunctor {
+ inline HOSTDEVICE paddle::framework::Array<OutT, 2> operator()(const InT& a,
+ const InT& b,
+ const InT& c) {
+ paddle::framework::Array<OutT, 2> outs;
+ // dx = dout * y
+ outs[0] = a * b;
+ // dy = dout * x
+ outs[1] = a * c;
+ return outs;
+ }
+};
+
+template <typename InT, typename OutT>
+struct MulGradXYFunctor<Complex<InT>, Complex<OutT>> {
+ inline HOSTDEVICE paddle::framework::Array<Complex<OutT>, 2> operator()(
+ const Complex<InT>& a, const Complex<InT>& b, const Complex<InT>& c) {
+ paddle::framework::Array<Complex<OutT>, 2> outs;
+ // dx = dout * y
+ Complex<InT> b_conj(b.real, -b.imag);
+ outs[0] = a * b_conj;
+ // dy = dout * x
+ Complex<InT> c_conj(c.real, -c.imag);
+ outs[1] = a * c_conj;
+ return outs;
+ }
+};
+
 } // namespace operators
 } // namespace paddle

paddle/fluid/operators/elementwise/elementwise_mul_op.cu

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/operators/elementwise/elementwise_mul_op.h"
 #include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h"
+#include "paddle/fluid/operators/reduce_ops/reduce_op.cu.h"
 #include "paddle/fluid/platform/complex.h"
 #include "paddle/fluid/platform/float16.h"
 
@@ -68,69 +69,41 @@ class ElementwiseMulKernel<platform::CUDADeviceContext, T>
  }
 };
 
-template <typename T>
-static __global__ void SimpleElemwiseMulGradCUDAKernel(const T* x, const T* y,
- const T* out,
- const T* dout,
- int64_t size, T* dx,
- T* dy) {
- int col = blockIdx.x * blockDim.x + threadIdx.x;
-
- while (col < size) {
- T o = dout[col];
- dx[col] = y[col] * o;
- dy[col] = x[col] * o;
- col += blockDim.x * gridDim.x;
- }
-}
-
-template <>
-__global__ void SimpleElemwiseMulGradCUDAKernel<plat::complex<float>>(
- const plat::complex<float>* x, const plat::complex<float>* y,
- const plat::complex<float>* out, const plat::complex<float>* dout,
- int64_t size, plat::complex<float>* dx, plat::complex<float>* dy) {
- int col = blockIdx.x * blockDim.x + threadIdx.x;
-
- while (col < size) {
- plat::complex<float> o = dout[col];
- dx[col] = plat::complex<float>(y[col].real, -y[col].imag) * o;
- dy[col] = plat::complex<float>(x[col].real, -x[col].imag) * o;
- col += blockDim.x * gridDim.x;
- }
-}
-
-template <>
-__global__ void SimpleElemwiseMulGradCUDAKernel<plat::complex<double>>(
- const plat::complex<double>* x, const plat::complex<double>* y,
- const plat::complex<double>* out, const plat::complex<double>* dout,
- int64_t size, plat::complex<double>* dx, plat::complex<double>* dy) {
- int col = blockIdx.x * blockDim.x + threadIdx.x;
-
- while (col < size) {
- plat::complex<double> o = dout[col];
- dx[col] = plat::complex<double>(y[col].real, -y[col].imag) * o;
- dy[col] = plat::complex<double>(x[col].real, -x[col].imag) * o;
- col += blockDim.x * gridDim.x;
- }
-}
-
 template <typename DeviceContext, typename T>
 typename std::enable_if<
- std::is_same<DeviceContext, plat::CUDADeviceContext>::value>::type
-elementwise_mul_grad(const framework::ExecutionContext& ctx,
- const framework::Tensor* x, const framework::Tensor* y,
- const framework::Tensor* out,
- const framework::Tensor* dout, framework::Tensor* dx,
- framework::Tensor* dy) {
- dim3 block_size = dim3(ELEMENTWISE_BLOCK_SIZE, 1);
- auto size = x->numel();
- dim3 grid_size =
- dim3((size + ELEMENTWISE_BLOCK_SIZE - 1) / ELEMENTWISE_BLOCK_SIZE, 1);
- SimpleElemwiseMulGradCUDAKernel<
- T><<<grid_size, block_size, 0,
- ctx.template device_context<plat::CUDADeviceContext>().stream()>>>(
- x->data<T>(), y->data<T>(), out->data<T>(), dout->data<T>(), size,
- dx->mutable_data<T>(ctx.GetPlace()), dy->mutable_data<T>(ctx.GetPlace()));
+ std::is_same<DeviceContext, platform::CUDADeviceContext>::value>::type
+ElementwiseMulGrad(const framework::ExecutionContext& ctx,
+ const framework::Tensor* x, const framework::Tensor* y,
+ const framework::Tensor* out, const framework::Tensor* dout,
+ framework::Tensor* dx, framework::Tensor* dy) {
+ int axis = ctx.Attr<int>("axis");
+ const auto& dev_ctx =
+ ctx.template device_context<platform::CUDADeviceContext>();
+ const auto place = ctx.GetPlace();
+
+ if (dx != nullptr && dy != nullptr) {
+ dx->mutable_data<T>(place);
+ if (dx->IsSharedBufferWith(*dout)) {
+ dx->clear();
+ dx->mutable_data<T>(x->dims(), place);
+ }
+ std::vector<const framework::Tensor*> ins = {dout, y, x};
+ GetGradXAndYOut<ElementwiseType::kBinary, T>(
+ dev_ctx, place, axis, ins, dout, dx, dy, MulGradXYFunctor<T, T>());
+ } else if (dx != nullptr && dy == nullptr) {
+ dx->mutable_data<T>(place);
+ if (dx->IsSharedBufferWith(*dout)) {
+ dx->clear();
+ dx->mutable_data<T>(x->dims(), place);
+ }
+ std::vector<const framework::Tensor*> ins = {dout, y};
+ GetGradXOrYOut<ElementwiseType::kBinary, T>(dev_ctx, place, axis, ins, dout,
+ dx, MulGradFunctor<T>());
+ } else if (dx == nullptr && dy != nullptr) {
+ std::vector<const framework::Tensor*> ins = {dout, x};
+ GetGradXOrYOut<ElementwiseType::kBinary, T>(dev_ctx, place, axis, ins, dout,
+ dy, MulGradFunctor<T>());
+ }
 }
 
 } // namespace operators

paddle/fluid/operators/elementwise/elementwise_mul_op.h

@@ -174,26 +174,23 @@ struct MulGradDY<paddle::platform::complex<T>> {
 template <typename DeviceContext, typename T>
 typename std::enable_if<
  std::is_same<DeviceContext, platform::CPUDeviceContext>::value>::type
-elementwise_mul_grad(const framework::ExecutionContext& ctx,
- const framework::Tensor* x, const framework::Tensor* y,
- const framework::Tensor* out,
- const framework::Tensor* dout, framework::Tensor* dx,
- framework::Tensor* dy) {
+ElementwiseMulGrad(const framework::ExecutionContext& ctx,
+ const framework::Tensor* x, const framework::Tensor* y,
+ const framework::Tensor* out, const framework::Tensor* dout,
+ framework::Tensor* dx, framework::Tensor* dy) {
  int axis = ctx.Attr<int>("axis");
  ElemwiseGradCompute<DeviceContext, T, MulGradDX<T>, MulGradDY<T>>(
  ctx, *x, *y, *out, *dout, axis, dx, dy, MulGradDX<T>(), MulGradDY<T>());
 }
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-// cuda definition
 template <typename DeviceContext, typename T>
 typename std::enable_if<
  std::is_same<DeviceContext, platform::CUDADeviceContext>::value>::type
-elementwise_mul_grad(const framework::ExecutionContext& ctx,
- const framework::Tensor* x, const framework::Tensor* y,
- const framework::Tensor* out,
- const framework::Tensor* dout, framework::Tensor* dx,
- framework::Tensor* dy);
+ElementwiseMulGrad(const framework::ExecutionContext& ctx,
+ const framework::Tensor* x, const framework::Tensor* y,
+ const framework::Tensor* out, const framework::Tensor* dout,
+ framework::Tensor* dx, framework::Tensor* dy);
 #endif
 
 template <typename DeviceContext, typename T>
@@ -209,14 +206,8 @@ class ElementwiseMulGradKernel : public ElemwiseGradKernel<T> {
  auto* out = dout; // out is not necessary
  auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
  auto* dy = ctx.Output<Tensor>(framework::GradVarName("Y"));
- int axis = ctx.Attr<int>("axis");
- if (dx != nullptr && dy != nullptr && (dx->dims() == dy->dims())) {
- elementwise_mul_grad<DeviceContext, T>(ctx, x, y, out, dout, dx, dy);
- } else {
- ElemwiseGradCompute<DeviceContext, T, MulGradDX<T>, MulGradDY<T>>(
- ctx, *x, *y, *out, *dout, axis, dx, dy, MulGradDX<T>(),
- MulGradDY<T>());
- }
+
+ ElementwiseMulGrad<DeviceContext, T>(ctx, x, y, out, dout, dx, dy);
  }
 };