From cde73a7bbff1ced34d8b0a19bc2553407ccbf347 Mon Sep 17 00:00:00 2001
From: Yang Zhang <yangzhang@live.com>
Date: Thu, 26 Sep 2019 15:34:28 +0800
Subject: [PATCH] Expose `mutable_data` as python binding (#19932)

* Expose `mutable_data` as python binding

test=develop

* Add test for device pointer binding

test=develop

* Make test compatible with python 2
---
 paddle/fluid/pybind/pybind.cc                 | 15 +++++++++++++
 .../fluid/tests/unittests/test_tensor.py      | 22 ++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
index 64413685c6..4c6174f25b 100644
--- a/paddle/fluid/pybind/pybind.cc
+++ b/paddle/fluid/pybind/pybind.cc
@@ -239,6 +239,21 @@ PYBIND11_MODULE(core_noavx, m) {
            [](Tensor &self, paddle::platform::CUDAPinnedPlace &place) {
              self.mutable_data<float>(place);
            })
+      .def("_mutable_data",
+           [](Tensor &self, paddle::platform::CPUPlace &place,
+              paddle::framework::proto::VarType::Type type) {
+             return reinterpret_cast<uintptr_t>(self.mutable_data(place, type));
+           })
+      .def("_mutable_data",
+           [](Tensor &self, paddle::platform::CUDAPlace &place,
+              paddle::framework::proto::VarType::Type type) {
+             return reinterpret_cast<uintptr_t>(self.mutable_data(place, type));
+           })
+      .def("_mutable_data",
+           [](Tensor &self, paddle::platform::CUDAPinnedPlace &place,
+              paddle::framework::proto::VarType::Type type) {
+             return reinterpret_cast<uintptr_t>(self.mutable_data(place, type));
+           })
       .def("_clear", &Tensor::clear)
       .def("set", PyCPUTensorSetFromArray<float>)
       .def("set", PyCPUTensorSetFromArray<int>)
diff --git a/python/paddle/fluid/tests/unittests/test_tensor.py b/python/paddle/fluid/tests/unittests/test_tensor.py
index 4615511ed8..ec180456ac 100644
--- a/python/paddle/fluid/tests/unittests/test_tensor.py
+++ b/python/paddle/fluid/tests/unittests/test_tensor.py
@@ -18,6 +18,7 @@ import paddle.fluid as fluid
 import paddle.fluid.core as core
 import unittest
 import numpy
+import numbers
 
 
 class TestTensor(unittest.TestCase):
@@ -171,7 +172,6 @@ class TestTensor(unittest.TestCase):
         var = scope.var("test_tensor")
 
         tensor = var.get_tensor()
-
         tensor._set_dims([0, 1])
         tensor._alloc_float(place)
 
@@ -256,6 +256,26 @@ class TestTensor(unittest.TestCase):
             print(tensor)
             self.assertTrue(isinstance(str(tensor), str))
 
+    def test_tensor_poiter(self):
+        place = core.CPUPlace()
+        scope = core.Scope()
+        var = scope.var("test_tensor")
+        place = core.CPUPlace()
+        tensor = var.get_tensor()
+        dtype = core.VarDesc.VarType.FP32
+        self.assertTrue(
+            isinstance(tensor._mutable_data(place, dtype), numbers.Integral))
+
+        if core.is_compiled_with_cuda():
+            place = core.CUDAPlace(0)
+            self.assertTrue(
+                isinstance(
+                    tensor._mutable_data(place, dtype), numbers.Integral))
+            place = core.CUDAPinnedPlace()
+            self.assertTrue(
+                isinstance(
+                    tensor._mutable_data(place, dtype), numbers.Integral))
+
 
 if __name__ == '__main__':
     unittest.main()