refine warpper

Author: JamesLim-sy

paddle/fluid/operators/elementwise/elementwise_mul_op.cu

@@ -36,13 +36,13 @@ class ElementwiseMulKernel<platform::CUDADeviceContext, T>
     : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
-    framework::Tensor x;
+    framework::Tensor x_for_selectedrows;
     std::vector<const framework::Tensor*> ins;
     std::vector<framework::Tensor*> outs;
     const auto& cuda_ctx =
         ctx.template device_context<platform::CUDADeviceContext>();
 
-    int axis = PackTensorsIntoVector<T>(ctx, &ins, &outs, &x);
+    int axis = PackTensorsIntoVector<T>(ctx, &ins, &outs, &x_for_selectedrows);
     LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T, T>(
         cuda_ctx, ins, &outs, axis, CudaMulFunctor<T>());
   }

paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h

@@ -509,15 +509,21 @@ void LaunchElementwiseCudaKernel(
     const platform::CUDADeviceContext &cuda_ctx,
     const std::vector<const framework::Tensor *> &ins,
     std::vector<framework::Tensor *> *outs, int axis, Functor func) {
+  std::vector<int> dims_size;
   bool no_broadcast_flag = true;
   for (auto *in : ins) {
     no_broadcast_flag = ins[0]->dims() == in->dims();
+    dims_size.emplace_back(in->dims().size());
   }
 
   if (no_broadcast_flag) {
     LaunchSameDimsElementwiseCudaKernel<ET, InT, OutT>(cuda_ctx, ins, outs,
                                                        func);
   } else {
+    axis = axis == -1
+               ? *std::max_element(dims_size.begin(), dims_size.end()) -
+                     *std::min_element(dims_size.begin(), dims_size.end())
+               : axis;
     LaunchBroadcastElementwiseCudaKernel<ET, InT, OutT>(cuda_ctx, ins, outs,
                                                         axis, func);
   }

paddle/fluid/operators/elementwise/elementwise_op_function.h

@@ -63,46 +63,52 @@ namespace operators {
 /*
 *  Pack input and output tensors into respective vectors with
 *  consideration of varible X`s class type.
+*  Input variable X is supported to be whether LoDTensor or
+*  SelectedRows class type in this package function, once X
+*  was SelectedRows type, a valid pointer x_for_selectedrows
+*  is excepted to be passed in from op kernel for acquisition
+*  of the valid address of LoDTensor created ahead in the function.
 */
 template <typename OutT>
 int PackTensorsIntoVector(const framework::ExecutionContext &ctx,
                           std::vector<const framework::Tensor *> *ins,
                           std::vector<framework::Tensor *> *outs,
-                          framework::Tensor *x_ptr = nullptr) {
+                          framework::Tensor *x_for_selectedrows = nullptr) {
   int axis = -1;
-  int x_dims_size = 0;
-  framework::Tensor *z;
-
   auto x_var = ctx.InputVar("X");
   PADDLE_ENFORCE_NOT_NULL(
       x_var, platform::errors::InvalidArgument(
-                 "Unable get input Variable X, Variable name is %s.\n",
+                 "Unable to get input Variable X, Variable name is %s.\n",
                  ctx.InputName("X")));
   auto *y = ctx.Input<framework::LoDTensor>("Y");
+  framework::Tensor *z;
 
-  if (x_ptr == nullptr || x_var->IsType<framework::LoDTensor>()) {
+  if (x_var->IsType<framework::LoDTensor>()) {
     auto *x = ctx.Input<framework::LoDTensor>("X");
     z = ctx.Output<framework::LoDTensor>("Out");
     ins->emplace_back(x);
-    x_dims_size = x->dims().size();
-
   } else if (x_var->IsType<framework::SelectedRows>()) {
     PADDLE_ENFORCE_EQ(y->dims().size() == 1 && y->dims()[0] == 1, true,
                       platform::errors::InvalidArgument(
                           "For elementwise_op, if X is Sparse, Y must be "
                           "scalar. But reveived the size of Y = %d.",
                           y->dims().size()));
+    PADDLE_ENFORCE_NOT_NULL(
+        x_for_selectedrows,
+        platform::errors::InvalidArgument(
+            "The parameter x_for_selectedrows is excepted to "
+            "be valid, once input varible X`s class type is "
+            "SelectedRows.\n"));
     auto &x_sele = x_var->Get<framework::SelectedRows>();
     auto out_sele = ctx.Output<framework::SelectedRows>("Out");
-    *x_ptr = x_sele.value();
+    *x_for_selectedrows = x_sele.value();
     out_sele->set_rows(x_sele.rows());
     out_sele->set_height(x_sele.height());
     out_sele->mutable_value()->Resize(x_sele.value().dims());
-    out_sele->mutable_value()->mutable_data(ctx.GetPlace(), x_ptr->type());
+    out_sele->mutable_value()->mutable_data(ctx.GetPlace(),
+                                            x_for_selectedrows->type());
     z = ctx.Output<framework::SelectedRows>("Out")->mutable_value();
-    ins->emplace_back(x_ptr);
-    x_dims_size = x_ptr->dims().size();
-
+    ins->emplace_back(x_for_selectedrows);
   } else {
     PADDLE_THROW(platform::errors::InvalidArgument(
         "X's type[%s] is not supported by elementwise_op. X's type should be "
@@ -115,7 +121,6 @@ int PackTensorsIntoVector(const framework::ExecutionContext &ctx,
   if (y != nullptr) {
     ins->emplace_back(y);
     axis = ctx.HasAttr("axis") ? ctx.Attr<int>("axis") : -1;
-    axis = axis == -1 ? std::abs(y->dims().size() - x_dims_size) : axis;
   }
   return axis;
 }