Optimize division size
@@ -251,9 +251,24 @@ public:
}
__device__ void set(ptrdiff_t i) {
- ptrdiff_t a = i;
+ if (ndim == 0) {
+ index_[0] = i;
+ return;
+ }
+ if (size_ > 1LL << 31) {
+ size_t a = i;
+ for (int dim = ndim; --dim > 0; ) {
+ size_t s = shape_[dim];
+ index_[dim] = (a % s);
+ a /= s;
+ }
+ if (ndim > 0) {
+ index_[0] = a;
+ }
+ } else {
+ unsigned int a = i;
for (int dim = ndim; --dim > 0; ) {
- ptrdiff_t s = shape_[dim];
+ unsigned s = shape_[dim];
index_[dim] = (a % s);
a /= s;
}
@@ -261,6 +276,7 @@ public:
index_[0] = a;
}
}
+ }
__device__ const int* get() const {
return index_;
