Avoid integer overflow in EigenMetaKernel indexing

- The current implementation computes `size + total_threads`, which can
  overflow and cause CUDA_ERROR_ILLEGAL_ADDRESS when size is close to
  the maximum representable value.
- The num_blocks calculation can also overflow due to the implementation
  of divup().
- This patch prevents these overflows and allows the kernel to work
  correctly for the full representable range of tensor sizes.
- Also adds relevant tests.
This commit is contained in:
Ben Barsdell
2021-10-18 20:58:14 +11:00
parent 55e3ae02ac
commit 50df8d3d6d
3 changed files with 109 additions and 7 deletions

View File

@@ -30,13 +30,15 @@ const T2& choose(Cond<false>, const T1&, const T2& second) {
template <typename T, typename X, typename Y>
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE
T divup(const X x, const Y y) {
return static_cast<T>((x + y - 1) / y);
// Note: This form is used because it cannot overflow.
return static_cast<T>(x == 0 ? 0 : (x - 1) / y + 1);
}
template <typename T>
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE
T divup(const T x, const T y) {
return static_cast<T>((x + y - 1) / y);
// Note: This form is used because it cannot overflow.
return static_cast<T>(x == 0 ? 0 : (x - 1) / y + 1);
}
template <size_t n> struct max_n_1 {