From 91f80f792d985fa2fe36385f53eae20dcc780470 Mon Sep 17 00:00:00 2001
From: Qiao Longfei <qiaolongfei@baidu.com>
Date: Wed, 10 Jan 2018 12:59:43 +0800
Subject: [PATCH] Topk share lod (#7373)

* add lod tensor ToAbsOffset test

* add share lod to topk op and softmax op
---
 paddle/framework/lod_tensor_test.cc | 16 ++++++++++++++++
 paddle/operators/softmax_op.cc      |  1 +
 paddle/operators/top_k_op.cc        |  2 ++
 3 files changed, 19 insertions(+)

diff --git a/paddle/framework/lod_tensor_test.cc b/paddle/framework/lod_tensor_test.cc
index 52b87f48e5..baad9c6f98 100644
--- a/paddle/framework/lod_tensor_test.cc
+++ b/paddle/framework/lod_tensor_test.cc
@@ -115,5 +115,21 @@ TEST(LoD, AppendLoD) {
   EXPECT_EQ(origin, expected);
 }
 
+TEST(LoD, ToAbsOffset) {
+  LoD relative_lod;
+  relative_lod.push_back(std::vector<size_t>({0, 2}));
+  relative_lod.push_back(std::vector<size_t>({0, 1, 3}));
+  relative_lod.push_back(std::vector<size_t>({0, 2, 4, 5}));
+
+  LoD abs_lod = paddle::framework::ToAbsOffset(relative_lod);
+
+  LoD expected;
+  expected.push_back(std::vector<size_t>({0, 5}));
+  expected.push_back(std::vector<size_t>({0, 2, 5}));
+  expected.push_back(std::vector<size_t>({0, 2, 4, 5}));
+
+  EXPECT_EQ(abs_lod, expected);
+}
+
 }  // namespace framework
 }  // namespace paddle
diff --git a/paddle/operators/softmax_op.cc b/paddle/operators/softmax_op.cc
index e7306bc5f1..cef1f1fc99 100644
--- a/paddle/operators/softmax_op.cc
+++ b/paddle/operators/softmax_op.cc
@@ -31,6 +31,7 @@ class SoftmaxOp : public framework::OperatorWithKernel {
     PADDLE_ENFORCE(x_dims.size() == 2UL,
                    "The input of softmax op must be a matrix.");
     ctx->SetOutputDim("Out", x_dims);
+    ctx->ShareLoD("X", /*->*/ "Out");
   }
 };
 
diff --git a/paddle/operators/top_k_op.cc b/paddle/operators/top_k_op.cc
index bb72210bb6..a8ddd72973 100644
--- a/paddle/operators/top_k_op.cc
+++ b/paddle/operators/top_k_op.cc
@@ -41,6 +41,8 @@ class TopkOp : public framework::OperatorWithKernel {
     dims[dims.size() - 1] = k;
     ctx->SetOutputDim("Out", dims);
     ctx->SetOutputDim("Indices", dims);
+    ctx->ShareLoD("X", "Out");
+    ctx->ShareLoD("X", "Indices");
   }
 };