bugfix for tensors with size larger than 2GB

springer13 · Jun 16, 2018 · 846df8a · 846df8a
1 parent 575b7b6
commit 846df8a
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 19 deletions.
diff --git a/examples/contraction.cpp b/examples/contraction.cpp
@@ -30,9 +30,9 @@ int main(int argc, char** argv)
    tcl::sizeType l1 = 6;
 
    float *dataA, *dataB, *dataC;
-   posix_memalign((void**) &dataA, 64, sizeof(float) * k2*m*k1*l1);
-   posix_memalign((void**) &dataB, 64, sizeof(float) * n*k2*k1*l1);
-   posix_memalign((void**) &dataC, 64, sizeof(float) * m*n*l1);
+   posix_memalign((void**) &dataA, 64, sizeof(float) * ((size_t)k2)*m*k1*l1);
+   posix_memalign((void**) &dataB, 64, sizeof(float) * ((size_t)n)*k2*k1*l1);
+   posix_memalign((void**) &dataC, 64, sizeof(float) * ((size_t)m)*n*l1);
 
    // Initialize tensors (data is not owned by the tensors)
    tcl::Tensor<float> A({k1,m,k2,l1}, dataA);

diff --git a/include/memoryBroker.h b/include/memoryBroker.h
@@ -37,8 +37,8 @@ class MemoryBroker {
    private:
 
       char *ptr;
-      uint64_t totalSize;
-      uint64_t currentOffset;
+      size_t totalSize;
+      size_t currentOffset;
 };
 
 }
diff --git a/include/tensor.h b/include/tensor.h
@@ -95,7 +95,7 @@ namespace tcl{
           * Calculate the product of the sizes of the indices specified by 'indices'
           * \return Product of the sizes of the specified indices
           */
-         sizeType getTotalSize( const indicesType &indices = {} ) const; 
+         size_t getTotalSize( const indicesType &indices = {} ) const; 
 
          /**
           * \return The stride of the specified index

diff --git a/src/contract.cpp b/src/contract.cpp
@@ -211,9 +211,9 @@ namespace tcl
          indices.emplace_back(x);
    }
 
-   void getBestTTGTCandidate(const indicesType &indicesA, const sizeType totalSizeA, 
-                                      const indicesType &indicesB, const sizeType totalSizeB,
-                                      const indicesType &indicesC, const sizeType totalSizeC,
+   void getBestTTGTCandidate(const indicesType &indicesA, const size_t totalSizeA, 
+                                      const indicesType &indicesB, const size_t totalSizeB,
+                                      const indicesType &indicesC, const size_t totalSizeC,
                                       const indicesType &loopIndices,
                                       const indicesType &mIndices,
                                       const indicesType &nIndices,
@@ -578,7 +578,7 @@ namespace tcl
          printVector(candidate.indicesA, "A");
          printVector(candidate.indicesB, "B");
          printVector(candidate.indicesC, "C");
-         printf("%d %d %d %d %d %d\n", totalSizeA, totalSizeB, totalSizeC, candidate.transA, candidate.transB, candidate.interchangeAB);
+         printf("%ld %ld %ld %d %d %d\n", totalSizeA, totalSizeB, totalSizeC, candidate.transA, candidate.transB, candidate.interchangeAB);
          exit(-1);
       }
 #endif
@@ -664,9 +664,9 @@ namespace tcl
       if( mIndices.size() <= 0 || nIndices.size() <= 0 || kIndices.size() <= 0 ) // TTGT is not applicable; use fallback
          return contract(alpha, A, B, beta, C);
 
-      sizeType totalSizeA = A->getTotalSize() * sizeof(floatType);
-      sizeType totalSizeB = B->getTotalSize() * sizeof(floatType);
-      sizeType totalSizeC = C->getTotalSize() * sizeof(floatType);
+      auto totalSizeA = A->getTotalSize() * sizeof(floatType);
+      auto totalSizeB = B->getTotalSize() * sizeof(floatType);
+      auto totalSizeC = C->getTotalSize() * sizeof(floatType);
       TTGTCandidate candidate;
 
       if( loopIndices.size() > 0 )
@@ -698,7 +698,7 @@ namespace tcl
       auto permB = getPermutation(indicesB, candidate.indicesB);
       auto permC = getPermutation(candidate.indicesC, indicesC);
 
-      sizeType requestedSize = 0;
+      size_t requestedSize = 0;
       if( !isIdentity(permA) )
          requestedSize += totalSizeA;
       if( !isIdentity(permB) )

diff --git a/src/memoryBroker.cpp b/src/memoryBroker.cpp
@@ -28,7 +28,7 @@ namespace tcl
 
    void MemoryBroker::alloc( size_t size )
    {
-      posix_memalign((void**)&(this->ptr), 4096, size);
+      int dummy = posix_memalign((void**)&(this->ptr), 4096, size);
       this->totalSize = size;
       this->currentOffset = 0;
    }

diff --git a/src/tensor.cpp b/src/tensor.cpp
@@ -59,9 +59,9 @@ namespace tcl
    }
 
    template<typename floatType>
-   sizeType Tensor<floatType>::getTotalSize( const indicesType &indices ) const
+   size_t Tensor<floatType>::getTotalSize( const indicesType &indices ) const
    {
-      sizeType product = 1;
+      size_t product = 1;
       if( indices.size() == 0 )
       {
          for(int i=0; i < _size.size(); ++i)
@@ -131,8 +131,10 @@ namespace tcl
          _offsets = std::vector<sizeType>(_size.size(),0);
 
       if( _data == nullptr ){
-         sizeType totalSize = std::accumulate(_outerSize.begin(), _outerSize.end(), 1, std::multiplies<sizeType>());
-         posix_memalign((void**) &_data, 64, sizeof(floatType) * totalSize);
+         size_t totalSize = 1;
+         for(auto s : _outerSize)
+				totalSize *= s;
+         int dummy = posix_memalign((void**) &_data, 64, sizeof(floatType) * totalSize);
       }
    }