Added access to the unerlying raw data of a tnsor slice/chip whenever possible

4 files changed, 126 insertions, 5 deletions

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorChipping.h b/unsupported/Eigen/CXX11/src/Tensor/TensorChipping.h
index 3aa3eba24..b862a8fd3 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorChipping.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorChipping.h
@@ -157,7 +157,14 @@ struct TensorEvaluator<const TensorChippingOp<DimId, ArgType>, Device>
 
   }*/
 
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar* data() const { return NULL; }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar* data() const {
+    Scalar* result = m_impl.data();
+    if (DimId == NumDims && result) {
+      return result + m_inputOffset;
+    } else {
+      return NULL;
+    }
+  }
 
  protected:
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const
diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorMorphing.h b/unsupported/Eigen/CXX11/src/Tensor/TensorMorphing.h
index 686bf5c24..3447592eb 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorMorphing.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorMorphing.h
@@ -366,7 +366,26 @@ struct TensorEvaluator<const TensorSlicingOp<StartIndices, Sizes, ArgType>, Devi
     }
   }
 
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar* data() const { return NULL; }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar* data() const {
+    Scalar* result = m_impl.data();
+    if (result) {
+      Index offset = 0;
+      for (int i = 0; i < NumDims; ++i) {
+        if (m_dimensions[i] != m_impl.dimensions()[i]) {
+          offset += m_offsets[i] * m_inputStrides[i];
+          for (int j = i+1; j < NumDims; ++j) {
+            if (m_dimensions[j] > 1) {
+              return NULL;
+            }
+            offset += m_offsets[j] * m_inputStrides[j];
+          }
+          break;
+        }
+      }
+      return result + offset;
+    }
+    return NULL;
+  }
 
  protected:
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index srcCoeff(Index index) const
diff --git a/unsupported/test/cxx11_tensor_chipping.cpp b/unsupported/test/cxx11_tensor_chipping.cpp
index 8c8a0cec2..0027b2888 100644
--- a/unsupported/test/cxx11_tensor_chipping.cpp
+++ b/unsupported/test/cxx11_tensor_chipping.cpp
@@ -236,9 +236,46 @@ static void test_chip_as_lvalue()
 }
 
 
+static void test_chip_raw_data()
+{
+  Tensor<float, 5> tensor(2,3,5,7,11);
+  tensor.setRandom();
+
+  typedef TensorEvaluator<decltype(tensor.chip<4>(3)), DefaultDevice> Evaluator4;
+  auto chip = Evaluator4(tensor.chip<4>(3), DefaultDevice());
+  for (int i = 0; i < 2; ++i) {
+    for (int j = 0; j < 3; ++j) {
+      for (int k = 0; k < 5; ++k) {
+        for (int l = 0; l < 7; ++l) {
+          int chip_index = i + 2 * (j + 3 * (k + 5 * l));
+          VERIFY_IS_EQUAL(chip.data()[chip_index], tensor(i,j,k,l,3));
+        }
+      }
+    }
+  }
+
+  typedef TensorEvaluator<decltype(tensor.chip<0>(0)), DefaultDevice> Evaluator0;
+  auto chip0 = Evaluator0(tensor.chip<0>(0), DefaultDevice());
+  VERIFY_IS_EQUAL(chip0.data(), static_cast<float*>(0));
+
+  typedef TensorEvaluator<decltype(tensor.chip<1>(0)), DefaultDevice> Evaluator1;
+  auto chip1 = Evaluator1(tensor.chip<1>(0), DefaultDevice());
+  VERIFY_IS_EQUAL(chip1.data(), static_cast<float*>(0));
+
+  typedef TensorEvaluator<decltype(tensor.chip<2>(0)), DefaultDevice> Evaluator2;
+  auto chip2 = Evaluator2(tensor.chip<2>(0), DefaultDevice());
+  VERIFY_IS_EQUAL(chip2.data(), static_cast<float*>(0));
+
+  typedef TensorEvaluator<decltype(tensor.chip<3>(0)), DefaultDevice> Evaluator3;
+  auto chip3 = Evaluator3(tensor.chip<3>(0), DefaultDevice());
+  VERIFY_IS_EQUAL(chip3.data(), static_cast<float*>(0));
+}
+
+
 void test_cxx11_tensor_chipping()
 {
   CALL_SUBTEST(test_simple_chip());
   CALL_SUBTEST(test_chip_in_expr());
   CALL_SUBTEST(test_chip_as_lvalue());
+  CALL_SUBTEST(test_chip_raw_data());
 }
diff --git a/unsupported/test/cxx11_tensor_morphing.cpp b/unsupported/test/cxx11_tensor_morphing.cpp
index fd1b1fa32..78b0dade0 100644
--- a/unsupported/test/cxx11_tensor_morphing.cpp
+++ b/unsupported/test/cxx11_tensor_morphing.cpp
@@ -12,7 +12,6 @@
 #include <Eigen/CXX11/Tensor>
 
 using Eigen::Tensor;
-using Eigen::IndexPair;
 
 static void test_simple_reshape()
 {
@@ -53,7 +52,8 @@ static void test_reshape_in_expr() {
   TensorMap<Tensor<float, 5>> tensor2(m2.data(), 3,5,7,11,13);
   Tensor<float, 2>::Dimensions newDims1{{2,3*5*7*11}};
   Tensor<float, 2>::Dimensions newDims2{{3*5*7*11,13}};
-  Eigen::array<IndexPair<DenseIndex>, 1> contract_along{{IndexPair<DenseIndex>(1, 0)}};
+  typedef Tensor<float, 1>::DimensionPair DimPair;
+  array<DimPair, 1> contract_along{{DimPair(1, 0)}};
   Tensor<float, 2> tensor3(2,13);
   tensor3 = tensor1.reshape(newDims1).contract(tensor2.reshape(newDims2), contract_along);
 
@@ -126,7 +126,8 @@ static void test_slice_in_expr() {
   TensorMap<Tensor<float, 2>> tensor1(m1.data(), 7, 7);
   TensorMap<Tensor<float, 2>> tensor2(m2.data(), 3, 3);
   Tensor<float, 2> tensor3(3,1);
-  array<IndexPair<DenseIndex>, 1> contract_along{{IndexPair<DenseIndex>(1, 0)}};
+  typedef Tensor<float, 1>::DimensionPair DimPair;
+  array<DimPair, 1> contract_along{{DimPair(1, 0)}};
 
   Eigen::DSizes<ptrdiff_t, 2> indices1(1,2);
   Eigen::DSizes<ptrdiff_t, 2> sizes1(3,3);
@@ -190,6 +191,62 @@ static void test_slice_as_lvalue()
 }
 
 
+static void test_slice_raw_data()
+{
+  Tensor<float, 4> tensor(3,5,7,11);
+  tensor.setRandom();
+
+  Eigen::DSizes<ptrdiff_t, 4> offsets(1,2,3,4);
+  Eigen::DSizes<ptrdiff_t, 4> extents(1,1,1,1);
+  typedef TensorEvaluator<decltype(tensor.slice(offsets, extents)), DefaultDevice> SliceEvaluator;
+  auto slice1 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice1.dimensions().TotalSize(), 1ul);
+  VERIFY_IS_EQUAL(slice1.data()[0], tensor(1,2,3,4));
+
+  extents = Eigen::DSizes<ptrdiff_t, 4>(2,1,1,1);
+  auto slice2 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice2.dimensions().TotalSize(), 2ul);
+  VERIFY_IS_EQUAL(slice2.data()[0], tensor(1,2,3,4));
+  VERIFY_IS_EQUAL(slice2.data()[1], tensor(2,2,3,4));
+
+  extents = Eigen::DSizes<ptrdiff_t, 4>(1,2,1,1);
+  auto slice3 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice3.dimensions().TotalSize(), 2ul);
+  VERIFY_IS_EQUAL(slice3.data(), static_cast<float*>(0));
+
+  offsets = Eigen::DSizes<ptrdiff_t, 4>(0,2,3,4);
+  extents = Eigen::DSizes<ptrdiff_t, 4>(3,2,1,1);
+  auto slice4 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice4.dimensions().TotalSize(), 6ul);
+  for (int i = 0; i < 3; ++i) {
+    for (int j = 0; j < 2; ++j) {
+      VERIFY_IS_EQUAL(slice4.data()[i+3*j], tensor(i,2+j,3,4));
+    }
+  }
+
+  offsets = Eigen::DSizes<ptrdiff_t, 4>(0,0,0,4);
+  extents = Eigen::DSizes<ptrdiff_t, 4>(3,5,7,2);
+  auto slice5 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice5.dimensions().TotalSize(), 210ul);
+  for (int i = 0; i < 3; ++i) {
+    for (int j = 0; j < 5; ++j) {
+      for (int k = 0; k < 7; ++k) {
+        for (int l = 0; l < 2; ++l) {
+          int slice_index = i + 3 * (j + 5 * (k + 7 * l));
+          VERIFY_IS_EQUAL(slice5.data()[slice_index], tensor(i,j,k,l+4));
+        }
+      }
+    }
+  }
+
+  offsets = Eigen::DSizes<ptrdiff_t, 4>(0,0,0,0);
+  extents = Eigen::DSizes<ptrdiff_t, 4>(3,5,7,11);
+  auto slice6 = SliceEvaluator(tensor.slice(offsets, extents), DefaultDevice());
+  VERIFY_IS_EQUAL(slice6.dimensions().TotalSize(), 3ul*5*7*11);
+  VERIFY_IS_EQUAL(slice6.data(), tensor.data());
+}
+
+
 void test_cxx11_tensor_morphing()
 {
   CALL_SUBTEST(test_simple_reshape());
@@ -199,4 +256,5 @@ void test_cxx11_tensor_morphing()
   CALL_SUBTEST(test_simple_slice());
   CALL_SUBTEST(test_slice_in_expr());
   CALL_SUBTEST(test_slice_as_lvalue());
+  CALL_SUBTEST(test_slice_raw_data());
 }