From ee11f00642afe00cfc14346d5c4791efa3405802 Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Wed, 1 Nov 2017 05:24:04 +0800
Subject: [PATCH] add shareLod (#5259)

* add shareLod

* fix sequence_conv grad infershape
---
 paddle/framework/op_desc.cc          | 16 ++++++++++++++++
 paddle/framework/operator.cc         | 14 ++++++++++++++
 paddle/framework/shape_inference.cc  |  3 ---
 paddle/framework/shape_inference.h   |  5 ++---
 paddle/operators/sequence_conv_op.cc |  2 +-
 5 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/paddle/framework/op_desc.cc b/paddle/framework/op_desc.cc
index c2d6f124ad..a4747e7c7c 100644
--- a/paddle/framework/op_desc.cc
+++ b/paddle/framework/op_desc.cc
@@ -52,6 +52,22 @@ class CompileTimeInferShapeContext : public InferShapeContext {
   const std::vector<std::string> &Outputs(
       const std::string &name) const override;
 
+  void ShareLoD(const std::string &in, const std::string &out, size_t i = 0,
+                size_t j = 0) const override {
+    PADDLE_ENFORCE_LT(i, Inputs(in).size());
+    PADDLE_ENFORCE_LT(j, Outputs(out).size());
+    auto *in_var = block_.FindVarRecursive(Inputs(in)[i]);
+    auto *out_var = block_.FindVarRecursive(Outputs(out)[j]);
+    if (in_var->GetType() != VarDesc::LOD_TENSOR) {
+      VLOG(3) << "input " << in << "is not LodTensor";
+      return;
+    }
+    PADDLE_ENFORCE_EQ(in_var->GetType(), VarDesc::LOD_TENSOR,
+                      "The %d-th output of Output(%s) must be LoDTensor.", j,
+                      out);
+    in_var->SetLoDLevel(out_var->GetLodLevel());
+  }
+
  private:
   DDim GetDim(const std::string &name) const override;
 
diff --git a/paddle/framework/operator.cc b/paddle/framework/operator.cc
index 222a252dc4..aa46829fdd 100644
--- a/paddle/framework/operator.cc
+++ b/paddle/framework/operator.cc
@@ -351,6 +351,20 @@ class RuntimeInferShapeContext : public InferShapeContext {
     return op_.Outputs(name);
   }
 
+  void ShareLoD(const std::string& in, const std::string& out, size_t i = 0,
+                size_t j = 0) const override {
+    PADDLE_ENFORCE_LT(i, Inputs(in).size());
+    PADDLE_ENFORCE_LT(j, Outputs(out).size());
+    Variable* in_var = scope_.FindVar(Inputs(in)[i]);
+    Variable* out_var = scope_.FindVar(Outputs(out)[j]);
+    if (!in_var->IsType<LoDTensor>()) return;
+    PADDLE_ENFORCE(out_var->IsType<LoDTensor>(),
+                   "The %d-th output of Output(%s) must be LoDTensor.", j, out);
+    auto in_tensor = in_var->Get<LoDTensor>();
+    auto* out_tensor = out_var->GetMutable<LoDTensor>();
+    out_tensor->set_lod(in_tensor.lod());
+  }
+
  private:
   DDim GetDim(const std::string& name) const override {
     Variable* var = scope_.FindVar(name);
diff --git a/paddle/framework/shape_inference.cc b/paddle/framework/shape_inference.cc
index 33a1d0b9b2..8169df8e46 100644
--- a/paddle/framework/shape_inference.cc
+++ b/paddle/framework/shape_inference.cc
@@ -28,9 +28,6 @@ void InferShapeContext::SetOutputsDim(
   SetDims(names, dims);
 }
 
-void InferShapeContext::ShareLoD(const std::string &in, const std::string &out,
-                                 size_t i, size_t j) const {}
-
 std::vector<framework::DDim> InferShapeContext::GetDims(
     const std::vector<std::string> &names) const {
   std::vector<framework::DDim> ret;
diff --git a/paddle/framework/shape_inference.h b/paddle/framework/shape_inference.h
index f1f1e44bcc..6f19900ef1 100644
--- a/paddle/framework/shape_inference.h
+++ b/paddle/framework/shape_inference.h
@@ -43,9 +43,8 @@ class InferShapeContext {
   virtual const std::vector<std::string> &Outputs(
       const std::string &name) const = 0;
 
-  // TODO(qiao) implement this function
-  void ShareLoD(const std::string &in, const std::string &out, size_t i = 0,
-                size_t j = 0) const;
+  virtual void ShareLoD(const std::string &in, const std::string &out,
+                        size_t i = 0, size_t j = 0) const = 0;
 
  protected:
   virtual framework::DDim GetDim(const std::string &name) const = 0;
diff --git a/paddle/operators/sequence_conv_op.cc b/paddle/operators/sequence_conv_op.cc
index bdb52265a5..a3f2ed1443 100644
--- a/paddle/operators/sequence_conv_op.cc
+++ b/paddle/operators/sequence_conv_op.cc
@@ -89,7 +89,7 @@ class SequenceConvGradOp : public framework::OperatorWithKernel {
     }
     if (ctx->HasOutput(framework::GradVarName("X"))) {
       ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
-      ctx->ShareLoD(framework::GradVarName("X"), "X");
+      ctx->ShareLoD("X", framework::GradVarName("X"));
     }
     if (ctx->HasOutput(framework::GradVarName("Filter"))) {
       ctx->SetOutputDim(framework::GradVarName("Filter"),
-- 
GitLab