Added support for 32bit index on a per tensor/tensor expression. This enables us to use 32bit indices to evaluate expressions on GPU faster while keeping the ability to use 64 bit indices to manipulate large tensors on CPU in the same binary.

1 files changed, 27 insertions, 19 deletions

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorStorage.h b/unsupported/Eigen/CXX11/src/Tensor/TensorStorage.h
index 1b227e8c2..91aae091c 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorStorage.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorStorage.h
@@ -66,14 +66,16 @@ template<typename T, DenseIndex NumIndices_, int Options_>
 class TensorStorage<T, NumIndices_, Dynamic, Options_, void>
   : public TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_numeric_list_repeated<DenseIndex, NumIndices_, Dynamic>::type>
 {
+  typedef typename internal::compute_index_type<Options_ & Index32Bit>::type Index;
+  typedef DSizes<Index, NumIndices_> Dimensions;
   typedef TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_numeric_list_repeated<DenseIndex, NumIndices_, Dynamic>::type> Base_;
 
   public:
-    TensorStorage() { }
-    TensorStorage(const TensorStorage<T, NumIndices_, Dynamic, Options_, void>& other) : Base_(other) { }
+    EIGEN_DEVICE_FUNC TensorStorage() { }
+    EIGEN_DEVICE_FUNC TensorStorage(const TensorStorage<T, NumIndices_, Dynamic, Options_, void>& other) : Base_(other) { }
 
-    TensorStorage(internal::constructor_without_unaligned_array_assert) : Base_(internal::constructor_without_unaligned_array_assert()) {}
-    TensorStorage(DenseIndex size, const array<DenseIndex, NumIndices_>& dimensions) : Base_(size, dimensions) {}
+    EIGEN_DEVICE_FUNC TensorStorage(internal::constructor_without_unaligned_array_assert) : Base_(internal::constructor_without_unaligned_array_assert()) {}
+    EIGEN_DEVICE_FUNC TensorStorage(DenseIndex size, const array<Index, NumIndices_>& dimensions) : Base_(size, dimensions) {}
 
   //      TensorStorage<T, NumIndices_, Dynamic, Options_, void>& operator=(const TensorStorage<T, NumIndices_, Dynamic, Options_, void>&) = default;
 };
@@ -82,24 +84,26 @@ class TensorStorage<T, NumIndices_, Dynamic, Options_, void>
 template<typename T, DenseIndex NumIndices_, int Options_>
 class TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_numeric_list_repeated<DenseIndex, NumIndices_, Dynamic>::type>
 {
-    T *m_data;
-    DSizes<DenseIndex, NumIndices_> m_dimensions;
+  public:
+  typedef typename internal::compute_index_type<Options_&Index32Bit>::type Index;
+  typedef DSizes<Index, NumIndices_> Dimensions;
 
     typedef TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_numeric_list_repeated<DenseIndex, NumIndices_, Dynamic>::type> Self_;
-  public:
-    TensorStorage() : m_data(0), m_dimensions() {}
-    TensorStorage(internal::constructor_without_unaligned_array_assert)
-      : m_data(0), m_dimensions(internal::template repeat<NumIndices_, DenseIndex>(0)) {}
-    TensorStorage(DenseIndex size, const array<DenseIndex, NumIndices_>& dimensions)
+
+    EIGEN_DEVICE_FUNC TensorStorage() : m_data(0), m_dimensions() {}
+    EIGEN_DEVICE_FUNC TensorStorage(internal::constructor_without_unaligned_array_assert)
+      : m_data(0), m_dimensions(internal::template repeat<NumIndices_, Index>(0)) {}
+    EIGEN_DEVICE_FUNC TensorStorage(Index size, const array<Index, NumIndices_>& dimensions)
         : m_data(internal::conditional_aligned_new_auto<T,(Options_&DontAlign)==0>(size)), m_dimensions(dimensions)
       { EIGEN_INTERNAL_TENSOR_STORAGE_CTOR_PLUGIN }
-      TensorStorage(const Self_& other)
+
+    EIGEN_DEVICE_FUNC TensorStorage(const Self_& other)
       : m_data(internal::conditional_aligned_new_auto<T,(Options_&DontAlign)==0>(internal::array_prod(other.m_dimensions)))
       , m_dimensions(other.m_dimensions)
     {
       internal::smart_copy(other.m_data, other.m_data+internal::array_prod(other.m_dimensions), m_data);
     }
-    Self_& operator=(const Self_& other)
+    EIGEN_DEVICE_FUNC Self_& operator=(const Self_& other)
     {
       if (this != &other) {
         Self_ tmp(other);
@@ -108,15 +112,15 @@ class TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_nu
       return *this;
     }
 
-    ~TensorStorage() { internal::conditional_aligned_delete_auto<T,(Options_&DontAlign)==0>(m_data, internal::array_prod(m_dimensions)); }
-    void swap(Self_& other)
+    EIGEN_DEVICE_FUNC  ~TensorStorage() { internal::conditional_aligned_delete_auto<T,(Options_&DontAlign)==0>(m_data, internal::array_prod(m_dimensions)); }
+    EIGEN_DEVICE_FUNC  void swap(Self_& other)
     { std::swap(m_data,other.m_data); std::swap(m_dimensions,other.m_dimensions); }
 
-    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const DSizes<DenseIndex, NumIndices_>& dimensions() const {return m_dimensions;}
+    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const {return m_dimensions;}
 
-    EIGEN_DEVICE_FUNC void resize(DenseIndex size, const array<DenseIndex, NumIndices_>& nbDimensions)
+    EIGEN_DEVICE_FUNC void resize(Index size, const array<Index, NumIndices_>& nbDimensions)
     {
-      const DenseIndex currentSz = internal::array_prod(m_dimensions);
+      const Index currentSz = internal::array_prod(m_dimensions);
       if(size != currentSz)
       {
         internal::conditional_aligned_delete_auto<T,(Options_&DontAlign)==0>(m_data, currentSz);
@@ -132,7 +136,11 @@ class TensorStorage<T, NumIndices_, Dynamic, Options_, typename internal::gen_nu
     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T *data() { return m_data; }
     EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const T *data() const { return m_data; }
 
-    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE DenseIndex size() const { return m_dimensions.TotalSize(); }
+    EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index size() const { return m_dimensions.TotalSize(); }
+
+ private:
+  T *m_data;
+  Dimensions m_dimensions;
 };
 
 } // end namespace Eigen