     return size_;
   }
 
+  __device__ const index_t& get() const {
+    return index_;
+  }
+
   __device__ void set(ptrdiff_t i) {
     // ndim == 0 case uses partial template specialization
     if (ndim == 1) {
       index_[0] = i;
-      return;
-    }
-    if (size_ > 1LL << 31) {
+    } else if (size_ > 1LL << 31) {
       // 64-bit division is very slow on GPU
-      size_t a = static_cast<size_t>(i);
+      this->_set(static_cast<unsigned long long int>(i));
+    } else {
+      this->_set(static_cast<unsigned int>(i));
+    }
+  }
+
+private:
+  template<typename index_t>
+  __device__ void _set(index_t i) {
       for (int dim = ndim; --dim > 0; ) {
-        size_t s = static_cast<size_t>(shape_[dim]);
+        index_t s = static_cast<index_t>(shape_[dim]);
         if (s & (s - 1)) {
-          size_t t = a / s;
-          index_[dim] = a - t * s;
-          a = t;
+          index_t t = i / s;
+          index_[dim] = i - t * s;
+          i = t;
         } else { // exp of 2
-          index_[dim] = a & (s - 1);
-          a >>= __popcll(s - 1);
+          index_[dim] = i & (s - 1);
+          i >>= _log2(s);
         }
       }
-      index_[0] = a;
-    } else {
-      unsigned int a = static_cast<unsigned int>(i);
-      for (int dim = ndim; --dim > 0; ) {
-        unsigned int s = static_cast<unsigned int>(shape_[dim]);
-        if (s & (s - 1)) {
-          unsigned int t = a / s;
-          index_[dim] = a - t * s;
-          a = t;
-        } else { // exp of 2
-          index_[dim] = a & (s - 1);
-          a >>= __popc(s - 1);
-        }
-      }
-      index_[0] = a;
-    }
-  }
-
-  __device__ const index_t& get() const {
-    return index_;
-  }
+      index_[0] = i;
+  }
+
+  // can also be implemented as __ffs(x)-1 or 31-__clz(x)
+  static unsigned int __device__ _log2(unsigned int x) { return __popc(x-1); }
+  static unsigned long long int __device__ _log2(unsigned long long int x) { return __popcll(x-1); }
 };
 
 template <>

// in NVRTC std:initializer_list is pre-defined (no need to include it)
#include <initializer_list>
#endif
#endif

// math

  static const int ndim = _ndim;
  static const bool c_contiguous = _c_contiguous;
  static const bool use_32bit_indexing = _use_32bit_indexing;
  typedef typename std::conditional<_use_32bit_indexing, int, ptrdiff_t>::type index_t;
private:
  T* data_;
  ptrdiff_t size_;
