C++ API parity: at::Tensor::data

pbelevich · facebook-github-bot · commit 33221b19acc3 · 2019-09-12T23:33:34.000-07:00
Summary: Pull Request resolved: #26008 Test Plan: Imported from OSS Differential Revision: D17343488 Pulled By: pbelevich fbshipit-source-id: b9ba5e26cad621a428a14292446d7fb5a6e5535d
diff --git a/aten/src/ATen/core/TensorBody.h b/aten/src/ATen/core/TensorBody.h
@@ -394,6 +394,7 @@ class CAFFE2_API Tensor {
   //Tensor * add(Tensor & b);
   void backward(const Tensor & gradient={}, bool keep_graph=false, bool create_graph=false) const;
   void set_data(const Tensor & new_data) const;
+  Tensor data() const;
   #ifdef BUILD_NAMEDTENSOR
   Tensor & names_(c10::optional<DimnameList> names) const;
   #endif
diff --git a/aten/src/ATen/core/TensorMethods.h b/aten/src/ATen/core/TensorMethods.h
@@ -73,6 +73,14 @@ inline void Tensor::set_data(const Tensor & new_data) const {
     return table->getOp<void (const Tensor &, const Tensor &)>(type_set())(const_cast<Tensor&>(*this), new_data);
 #endif
 }
+inline Tensor Tensor::data() const {
+#ifdef USE_STATIC_DISPATCH
+    return TypeDefault::data(const_cast<Tensor&>(*this));
+#else
+    static auto table = globalATenDispatch().getOpTable("aten::data(Tensor self) -> Tensor");
+    return table->getOp<Tensor (const Tensor &)>(type_set())(const_cast<Tensor&>(*this));
+#endif
+}
 #ifdef BUILD_NAMEDTENSOR
 inline Tensor & Tensor::names_(c10::optional<DimnameList> names) const {
 #ifdef USE_STATIC_DISPATCH
diff --git a/aten/src/ATen/native/VariableMethodStubs.cpp b/aten/src/ATen/native/VariableMethodStubs.cpp
@@ -12,5 +12,9 @@ void set_data(const Tensor& self, const Tensor& new_data) {
   AT_ERROR("set_data is not implemented for Tensor");
 }
 
+Tensor data(const Tensor& self) {
+  AT_ERROR("data is not implemented for Tensor");
+}
+
 } // namespace native
 } // namespace at
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -35,6 +35,9 @@
 - func: set_data(Tensor(a!) self, Tensor new_data) -> void
   variants: method
 
+- func: data(Tensor self) -> Tensor
+  variants: method
+
 - func: names_(Tensor(a!) self, Dimname[]? names) -> Tensor(a!)
   variants: method
   named_guard: False
diff --git a/aten/src/ATen/native/quantized/cpu/qconv.cpp b/aten/src/ATen/native/quantized/cpu/qconv.cpp
@@ -235,8 +235,8 @@ class QConv2dInt8 final : public c10::OperatorKernel {
           conv_p,
           act_ptr,
           *packB,
-          reinterpret_cast<uint8_t*>(output.data<c10::quint8>()),
-          buffer.data<int32_t>(),
+          reinterpret_cast<uint8_t*>(output.data_ptr<c10::quint8>()),
+          buffer.data_ptr<int32_t>(),
           outputProcObj,
           0 /* thread_id*/,
           1 /* num_threads */);
diff --git a/aten/src/ATen/test/quantized_test.cpp b/aten/src/ATen/test/quantized_test.cpp
@@ -131,14 +131,14 @@ TEST(TestQTensor, EmptyPerchannelQuantized) {
       {ch_axis},
       at::device(at::kCPU).dtype(kQUInt8));
   // Assigning to QTensor
-  auto* q_data = q.data<quint8>();
+  auto* q_data = q.data_ptr<quint8>();
   for (int i = 0; i < numel; ++i) {
     q_data[i].val_ = val;
   }
 
   // dequantize
   auto r = q.dequantize();
-  auto* r_data = r.data<float>();
+  auto* r_data = r.data_ptr<float>();
   for (int i = 0; i < numel; ++i) {
     ASSERT_EQ(
         r_data[i],
diff --git a/test/cpp/api/tensor.cpp b/test/cpp/api/tensor.cpp
@@ -311,3 +311,11 @@ TEST(TensorTest, DataPtr) {
   ASSERT_EQ(tensor_not_copy.data_ptr<float>(), tensor.data_ptr<float>());
   ASSERT_EQ(tensor_not_copy.data_ptr(), tensor.data_ptr());
 }
+
+TEST(TensorTest, Data) {
+  const auto tensor = torch::empty({3, 3});
+  ASSERT_TRUE(torch::equal(tensor, tensor.data()));
+
+  const auto tensor2 = at::empty({3, 3});
+  ASSERT_THROW(tensor2.data(), c10::Error);
+}
diff --git a/tools/autograd/gen_python_functions.py b/tools/autograd/gen_python_functions.py
@@ -36,6 +36,7 @@
     'set_quantizer_',
     'set_data',
     '.*_overrideable',  # overrideable functions for backend extension
+    'data'
 ]
 
 # These function signatures are not exposed to Python. Note that this signature
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -29,7 +29,7 @@
 
 # These functions are written manually in templates/VariableType.cpp
 MANUAL_IMPLEMENTATIONS = {
-    'resize_', 'resize_as_', 'detach', 'detach_', 'copy_', 'backward', 'set_data'
+    'resize_', 'resize_as_', 'detach', 'detach_', 'copy_', 'backward', 'set_data', 'data'
 }
 
 # These functions we don't want to record for tracing, because we always want
diff --git a/torch/csrc/autograd/VariableTypeManual.cpp b/torch/csrc/autograd/VariableTypeManual.cpp
@@ -93,6 +93,10 @@ void VariableType::set_data(const Tensor & self, const Tensor & new_data) {
   as_variable_ref(self).set_data(new_data);
 }
 
+Tensor VariableType::data(const Tensor & self) {
+  return as_variable_ref(self).variable_data();
+}
+
 // We don't have an outplace copy, so this can't be generated automatically
 Tensor & VariableType::copy_(Tensor & self, const Tensor & src, bool non_blocking) {
   jit::Value* output = nullptr;

Original file line number	Diff line number	Diff line change
`@@ -12,5 +12,9 @@ void set_data(const Tensor& self, const Tensor& new_data) {`
`12`	`12`	`AT_ERROR("set_data is not implemented for Tensor");`
`13`	`13`	`}`
`14`	`14`
	`15`	`+Tensor data(const Tensor& self) {`
	`16`	`+ AT_ERROR("data is not implemented for Tensor");`
	`17`	`+}`
	`18`	`+`
`15`	`19`	`} // namespace native`
`16`	`20`	`} // namespace at`
Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,7 @@`
`36`	`36`	`'set_quantizer_',`
`37`	`37`	`'set_data',`
`38`	`38`	`'.*_overrideable', # overrideable functions for backend extension`
	`39`	`+ 'data'`
`39`	`40`	`]`
`40`	`41`
`41`	`42`	`# These function signatures are not exposed to Python. Note that this signature`
Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@`
`29`	`29`
`30`	`30`	`# These functions are written manually in templates/VariableType.cpp`
`31`	`31`	`MANUAL_IMPLEMENTATIONS = {`
`32`		`- 'resize_', 'resize_as_', 'detach', 'detach_', 'copy_', 'backward', 'set_data'`
	`32`	`+ 'resize_', 'resize_as_', 'detach', 'detach_', 'copy_', 'backward', 'set_data', 'data'`
`33`	`33`	`}`
`34`	`34`
`35`	`35`	`# These functions we don't want to record for tracing, because we always want`