Sparse Kernel for Lookup Table Grad Operator

paddle/framework/backward.cc

@@ -179,8 +179,9 @@ static std::unique_ptr<OperatorBase> BackwardRecursive(
  // collect all the offset for each alias,
  // insert a sum operator to add all aliases to output
  insert_position.push_back(
- {dup_op.back(), OpRegistry::CreateOp("sum", {{"X", dup_outputs}},
- {{"Out", {name}}}, {})});
+ {dup_op.back(),
+ OpRegistry::CreateOp("sum", {{"X", dup_outputs}}, {{"Out", {name}}},
+ {{"is_internal", true}})});
  }
 
  // make sure the inserted `sum` ops follow the BFS order.
@@ -315,6 +316,7 @@ static void CreateGradVarInBlock(
  return false; /* not break */
  });
  if (need_infer_shape) {
+ ops[op_index]->InferVarType(block_desc);
  ops[op_index]->InferShape(*block_desc);
  }
  }
@@ -341,7 +343,6 @@ std::vector<std::unique_ptr<OpDescBind>> MakeOpGrad(
  grad_op_descs = OpInfoMap::Instance()
  .Get(op_desc->Type())
  .GradOpMaker()(*op_desc, *no_grad_vars, grad_to_var);
-
  std::list<std::unique_ptr<OpDescBind>> pending_fill_zeros_ops;
  for (auto& desc : grad_op_descs) {
  for (const std::string& in_name : desc->InputArgumentNames()) {
@@ -413,8 +414,9 @@ std::vector<std::unique_ptr<OpDescBind>> MakeBlockBackward(
  backward_descs[dup_op[i]]->Rename(out_name, new_name);
  sum_op_inputs.emplace_back(new_name);
  }
- std::unique_ptr<OpDescBind> sum_op(new OpDescBind(
- "sum", {{"X", sum_op_inputs}}, {{"Out", {out_name}}}, {}));
+ std::unique_ptr<OpDescBind> sum_op(
+ new OpDescBind("sum", {{"X", sum_op_inputs}}, {{"Out", {out_name}}},
+ {{"is_internal", true}}));
  pending_sum_ops.push_back({dup_op.back(), std::move(sum_op)});
  }
  }
@@ -442,7 +444,7 @@ ParamGradInfoMap AppendBackward(
  }
 
  const int root_block_idx = 0;
- auto root_block = program_desc.Block(root_block_idx);
+ auto* root_block = program_desc.Block(root_block_idx);
 
  // insert fill one op for target
  // TODO(qiao) add some check to the target.
@@ -457,6 +459,9 @@ ParamGradInfoMap AppendBackward(
  {{"shape", target_shape},
  {"value", static_cast<float>(1.0)},
  {"data_type", framework::DataType::FP32}}));
+ // infer var type of fill_one_op_out
+ fill_one_op->InferVarType(root_block);
+
  root_block->AppendAllocatedOp(std::move(fill_one_op));
  size_t forward_op_num = root_block->OpSize();
  size_t forward_block_num = program_desc.Size();
@@ -473,8 +478,8 @@ ParamGradInfoMap AppendBackward(
 
  // Create target gradient variable
  std::unordered_map<std::string, GradVarInfo> retv;
-
  auto var = root_block->Var(fill_one_op_out);
+
  // FIXME(qiao) infer the data type
  var->SetDataType(framework::DataType::FP32);
  var->SetShape(target.Shape());
@@ -490,6 +495,7 @@ ParamGradInfoMap AppendBackward(
  CreateGradVarInBlock(0, grad_to_var, program_desc.Block(block_index),
  &retv);
  }
+ LOG(INFO) << "AppendBackward Done";
  return retv;
 }
 

paddle/framework/executor.cc

@@ -20,6 +20,7 @@ limitations under the License. */
 #include <set>
 #include <vector>
 
+#include "paddle/framework/feed_fetch_type.h"
 #include "paddle/framework/lod_tensor.h"
 #include "paddle/framework/op_registry.h"
 #include "paddle/framework/scope.h"
@@ -56,6 +57,20 @@ Executor::~Executor() {
  }
 }
 
+void CreateTensor(Variable* var, VarDesc::VarType var_type) {
+ if (var_type == VarDesc::LOD_TENSOR) {
+ var->GetMutable<LoDTensor>();
+ } else if (var_type == VarDesc::SELECTED_ROWS) {
+ var->GetMutable<SelectedRows>();
+ } else if (var_type == VarDesc::FEED_MINIBATCH) {
+ var->GetMutable<FeedFetchList>();
+ } else if (var_type == VarDesc::FETCH_LIST) {
+ var->GetMutable<FeedFetchList>();
+ } else {
+ PADDLE_THROW("Variable type must be LoDTensor/SelectedRows.");
+ }
+}
+
 void Executor::Run(const ProgramDesc& pdesc, Scope* scope, int block_id) {
  // TODO(tonyyang-svail):
  // - only runs on the first device (i.e. no interdevice communication)
@@ -69,10 +84,12 @@ void Executor::Run(const ProgramDesc& pdesc, Scope* scope, int block_id) {
  for (auto& var : block.vars()) {
  if (var.persistable()) {
  auto* ptr = scope->Var(var.name());
+ CreateTensor(ptr, var.type());
  VLOG(3) << "Create Variable " << var.name()
  << " global, which pointer is " << ptr;
  } else {
  auto* ptr = local_scope.Var(var.name());
+ CreateTensor(ptr, var.type());
  VLOG(3) << "Create Variable " << var.name()
  << " locally, which pointer is " << ptr;
  }

paddle/framework/op_desc.h

@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include <unordered_map>
 #include <vector>
+#include "glog/logging.h"
 #include "paddle/framework/attribute.h"
 #include "paddle/framework/type_defs.h"
 #include "paddle/framework/var_desc.h"

paddle/framework/op_registry.h

@@ -147,19 +147,19 @@ class OpKernelRegistrar : public Registrar {
 /**
  * Macro to register Operator.
  */
-#define REGISTER_OP(op_type, op_class, op_maker_class, grad_op_type,  \
- grad_op_class)  \
- REGISTER_OPERATOR(grad_op_type, grad_op_class);  \
- class _GradOpDescMaker_##grad_op_type##_  \
- : public ::paddle::framework::DefaultGradOpDescMaker<true> {  \
- using ::paddle::framework::DefaultGradOpDescMaker<  \
- true>::DefaultGradOpDescMaker;  \
-  \
- protected:  \
- virtual std::string GradOpType() const { return #grad_op_type; }  \
- };  \
- REGISTER_OPERATOR(op_type, op_class, _GradOpDescMaker_##grad_op_type##_, \
- op_maker_class);
+#define REGISTER_OP(op_type, op_class, op_maker_class, grad_op_type, \
+ grad_op_class) \
+ REGISTER_OPERATOR(grad_op_type, grad_op_class); \
+ class _GradOpDescMaker_##grad_op_type##_ \
+ : public ::paddle::framework::DefaultGradOpDescMaker<true> { \
+ using ::paddle::framework::DefaultGradOpDescMaker< \
+ true>::DefaultGradOpDescMaker; \
+ \
+ protected: \
+ virtual std::string GradOpType() const { return #grad_op_type; } \
+ }; \
+ REGISTER_OPERATOR(op_type, op_class, op_maker_class,  \
+ _GradOpDescMaker_##grad_op_type##_);
 
 #define REGISTER_OP_WITHOUT_GRADIENT(op_type, op_class, op_maker_class) \
  REGISTER_OPERATOR(op_type, op_class, op_maker_class)

paddle/framework/operator.cc

@@ -33,24 +33,6 @@ ExecutionContext::GetEigenDevice<platform::GPUPlace, Eigen::GpuDevice>() const {
 }
 #endif
 
-const Tensor* GetTensorFromVar(const Variable* var) {
- if (var->IsType<LoDTensor>()) {
- return &var->Get<LoDTensor>();
- }
- PADDLE_ENFORCE(var->IsType<Tensor>(),
- "The Input must be LoDTensor or Tensor.");
- return &var->Get<Tensor>();
-}
-
-Tensor* GetTensorFromVar(Variable* var) {
- if (var->IsType<LoDTensor>()) {
- return var->GetMutable<LoDTensor>();
- }
- PADDLE_ENFORCE(var->IsType<Tensor>(),
- "The Input must be LoDTensor or Tensor.");
- return var->GetMutable<Tensor>();
-}
-
 std::string OperatorBase::Input(const std::string& name) const {
  auto& ins = Inputs(name);
  PADDLE_ENFORCE_LE(ins.size(), 1UL,
@@ -227,7 +209,7 @@ const std::vector<const Tensor*> ExecutionContext::MultiInput<Tensor>(
 template <>
 Tensor* ExecutionContext::Output<Tensor>(const std::string& name) const {
  auto var = OutputVar(name);
- return var == nullptr ? nullptr : var->GetMutable<LoDTensor>();
+ return var == nullptr ? nullptr : GetMutableTensorFromVar(var);
 }
 
 template <>
@@ -240,7 +222,7 @@ std::vector<Tensor*> ExecutionContext::MultiOutput<Tensor>(
  [&](const std::string& sub_name) {
  auto var = scope_.FindVar(sub_name);
  return var == nullptr ? nullptr
- : var->GetMutable<LoDTensor>();
+ : GetMutableTensorFromVar(var);
  });
  return res;
 }

paddle/framework/operator.h

@@ -28,6 +28,7 @@ limitations under the License. */
 #include "paddle/framework/lod_tensor.h"
 #include "paddle/framework/op_info.h"
 #include "paddle/framework/scope.h"
+#include "paddle/framework/selected_rows.h"
 #include "paddle/framework/shape_inference.h"
 #include "paddle/framework/tensor.h"
 #include "paddle/platform/device_context.h"
@@ -60,9 +61,6 @@ inline std::string GradVarName(const std::string& var_name) {
 class OperatorBase;
 class ExecutionContext;
 
-extern const Tensor* GetTensorFromVar(const Variable* var);
-extern Tensor* GetTensorFromVar(Variable* var);
-
 /**
  * OperatorBase has the basic element that Net will call to do computation.
  * Only CreateOperator from OpRegistry will new Operator directly. User
@@ -169,6 +167,30 @@ class NOP : public OperatorBase {
  }
 };
 
+static const Tensor* GetTensorFromVar(const Variable* var) {
+ const Tensor* t = nullptr;
+ if (var->IsType<LoDTensor>()) {
+ t = &(var->Get<LoDTensor>());
+ } else if (var->IsType<SelectedRows>()) {
+ t = &(var->Get<SelectedRows>().value());
+ } else {
+ PADDLE_THROW("Variable type must be LoDTensor/SelectedRows.");
+ }
+ return t;
+}
+
+static Tensor* GetMutableTensorFromVar(Variable* var) {
+ Tensor* t = nullptr;
+ if (var->IsType<LoDTensor>()) {
+ t = var->GetMutable<LoDTensor>();
+ } else if (var->IsType<SelectedRows>()) {
+ t = var->GetMutable<SelectedRows>()->mutable_value();
+ } else {
+ PADDLE_THROW("Variable type must be LoDTensor/SelectedRows.");
+ }
+ return t;
+}
+
 class ExecutionContext {
  public:
  ExecutionContext(const OperatorBase& op, const Scope& scope,
@@ -511,28 +533,14 @@ class RuntimeInferShapeContext : public InferShapeContext {
  }
 
  private:
- template <bool Allocate>
- Tensor* GetTensor(const std::string& name) const {
- Tensor* t = nullptr;
- auto* var = scope_.FindVar(name);
- if (!var->IsType<LoDTensor>() && !var->IsType<Tensor>()) {
- if (Allocate) {
- t = var->GetMutable<LoDTensor>();
- } else {
- PADDLE_THROW("Variable(%s) should be tensor", name);
- }
- } else {
- t = GetTensorFromVar(scope_.FindVar(name));
- }
- return t;
- }
-
  DDim GetDim(const std::string& name) const override {
- return GetTensor<false>(name)->dims();
+ Variable* var = scope_.FindVar(name);
+ return GetTensorFromVar(var)->dims();
  }
 
  void SetDim(const std::string& name, const DDim& dim) override {
- GetTensor<true>(name)->Resize(dim);
+ Variable* var = scope_.FindVar(name);
+ GetMutableTensorFromVar(var)->Resize(dim);
  }
 
  const OperatorBase& op_;
@@ -655,6 +663,8 @@ class OperatorWithKernel : public OperatorBase {
  t = &var->Get<Tensor>();
  } else if (var->IsType<LoDTensor>()) {
  t = &var->Get<LoDTensor>();
+ } else if (var->IsType<SelectedRows>()) {
+ t = &(var->Get<SelectedRows>().value());
  }
  if (t != nullptr) {
  int tmp = static_cast<int>(ToDataType(t->type()));

paddle/framework/selected_rows.h

@@ -37,6 +37,8 @@ class SelectedRows {
 
  const Vector<int64_t>& rows() const { return rows_; }
 
+ Vector<int64_t>& mutable_rows() { return rows_; }
+
  void set_rows(const Vector<int64_t>& rows) { rows_ = rows; }
 
  DDim GetCompleteDims() const {

paddle/operators/CMakeLists.txt

@@ -131,7 +131,7 @@ op_library(recurrent_op SRCS recurrent_op.cc rnn/recurrent_op_utils.cc
 op_library(cond_op SRCS cond_op.cc DEPS framework_proto tensor operator net_op)
 op_library(cross_entropy_op DEPS cross_entropy)
 op_library(softmax_with_cross_entropy_op DEPS cross_entropy softmax)
-op_library(sum_op DEPS net_op)
+op_library(sum_op DEPS net_op selected_rows_functor)
 op_library(pool_op DEPS pooling)
 op_library(pool_with_index_op DEPS pooling)
 op_library(lstm_op DEPS sequence2batch lstm_compute)

paddle/operators/cross_entropy_op.cu

@@ -21,7 +21,7 @@ namespace {
 
 template <typename T>
 __global__ void CrossEntropyGradientKernel(T* dX, const T* dY, const T* X,
- const int* label, const int N,
+ const int64_t* label, const int N,
  const int D) {
  // TOOD(qingqing) define CUDA_1D_KERNEL_LOOP macro in a common file.
  // CUDA_1D_KERNEL_LOOP(i, N) {
@@ -77,8 +77,8 @@ class CrossEntropyGradientOpCUDAKernel : public framework::OpKernel<T> {
  T* dx_data = dx->mutable_data<T>(ctx.GetPlace());
  const T* x_data = x->data<T>();
 
- int batch_size = x->dims()[0];
- int class_num = x->dims()[1];
+ int64_t batch_size = x->dims()[0];
+ int64_t class_num = x->dims()[1];
 
  int block = 512;
  int grid = (batch_size * class_num + block - 1) / block;
@@ -93,7 +93,7 @@ class CrossEntropyGradientOpCUDAKernel : public framework::OpKernel<T> {
  } else {
  math::SetConstant<platform::GPUPlace, T> functor;
  functor(ctx.device_context(), dx, 0);
- auto* label_data = label->data<int>();
+ auto* label_data = label->data<int64_t>();
  grid = (batch_size + block - 1) / block;
  CrossEntropyGradientKernel<T><<<
  grid, block, 0, reinterpret_cast<const platform::CUDADeviceContext&>(

paddle/operators/cross_entropy_op.h

@@ -54,28 +54,28 @@ class CrossEntropyGradientOpKernel : public framework::OpKernel<T> {
  Tensor* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
  T* dx_data = dx->mutable_data<T>(ctx.GetPlace());
 
- int class_num = x->dims()[1];
+ int64_t class_num = x->dims()[1];
  if (ctx.Attr<bool>("soft_label")) {
  auto x_mat = EigenMatrix<T>::From(*x);
  auto dy_mat = EigenMatrix<T>::From(*dy);
  auto lbl_mat = EigenMatrix<T>::From(*label);
  auto dx_mat = EigenMatrix<T>::From(*dx);
 
  dx_mat.device(ctx.GetEigenDevice<platform::CPUPlace>()) =
- -(lbl_mat * dy_mat.broadcast(Eigen::DSizes<int, 2>(1, class_num)) /
- x_mat);
+ -(lbl_mat *
+ dy_mat.broadcast(Eigen::DSizes<int64_t, 2>(1, class_num)) / x_mat);
  } else {
- int batch_size = x->dims()[0];
+ int64_t batch_size = x->dims()[0];
  const T* dy_data = dy->data<T>();
  const T* x_data = x->data<T>();
- const int* label_data = label->data<int>();
+ const int64_t* label_data = label->data<int64_t>();
 
  math::SetConstant<platform::CPUPlace, T> functor;
  functor(ctx.device_context(), dx, 0);
 
- for (int i = 0; i < batch_size; ++i) {
+ for (int64_t i = 0; i < batch_size; ++i) {
  PADDLE_ASSERT(label_data[i] >= 0 || label_data[i] < class_num);
- int index = i * class_num + label_data[i];
+ int64_t index = i * class_num + label_data[i];
  dx_data[index] = -dy_data[i] / x_data[index];
  }
  }

paddle/operators/feed_op.cc

@@ -41,7 +41,7 @@ class FeedOp : public framework::OperatorBase {
 
  auto col = Attr<int>("col");
 
- VLOG(3) << "Feed Var " << feed_var_name << "'s " << col << " column to var"
+ VLOG(3) << "Feed Var " << feed_var_name << "'s " << col << " column to var "
  << out_name;
 
  auto &feed_list = feed_var->Get<framework::FeedFetchList>();