devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Avoid integer overflow in EigenMetaKernel indexing

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

45
unsupported/test/cxx11_tensor_gpu.cu
View File

@@ -66,6 +66,47 @@ void test_gpu_nullary() {
gpuFree(d_in2);
}
// Tests that there are no indexing overflows when computing tensors with the
// max representable size.
template <typename IndexType,
IndexType N = (std::numeric_limits<IndexType>::max)()>
void test_gpu_nullary_max_size()
{
typedef int8_t DataType;
typedef Tensor<DataType, 1, 0, IndexType> TensorType;
typedef Eigen::array<IndexType, 1> ArrayType;
const IndexType n = N;
TensorType in1((ArrayType(n)));
in1.setZero();
std::size_t in1_bytes = in1.size() * sizeof(DataType);
DataType* d_in1;
gpuMalloc((void**)(&d_in1), in1_bytes);
gpuMemcpy(d_in1, in1.data(), in1_bytes, gpuMemcpyHostToDevice);
Eigen::GpuStreamDevice stream;
Eigen::GpuDevice gpu_device(&stream);
Eigen::TensorMap<TensorType> gpu_in1(d_in1, ArrayType(n));
gpu_in1.device(gpu_device) = gpu_in1.constant(123);
TensorType new1((ArrayType(n)));
assert(gpuMemcpyAsync(new1.data(), d_in1, in1_bytes, gpuMemcpyDeviceToHost,
gpu_device.stream()) == gpuSuccess);
assert(gpuStreamSynchronize(gpu_device.stream()) == gpuSuccess);
for (IndexType i = 0; i < n; ++i) {
VERIFY_IS_EQUAL(new1(ArrayType(i)), 123);
}
gpuFree(d_in1);
}
void test_gpu_elementwise_small() {
Tensor<float, 1> in1(Eigen::array<Eigen::DenseIndex, 1>(2));
Tensor<float, 1> in2(Eigen::array<Eigen::DenseIndex, 1>(2));
@@ -1524,6 +1565,10 @@ void test_gpu_gamma_sample_der_alpha()
EIGEN_DECLARE_TEST(cxx11_tensor_gpu)
{
CALL_SUBTEST_1(test_gpu_nullary());
CALL_SUBTEST_1(test_gpu_nullary_max_size<int16_t>());
CALL_SUBTEST_1(test_gpu_nullary_max_size<int32_t>());
CALL_SUBTEST_1((test_gpu_nullary_max_size<
int64_t, (std::numeric_limits<int32_t>::max)() + 100ll>()));
CALL_SUBTEST_1(test_gpu_elementwise_small());
CALL_SUBTEST_1(test_gpu_elementwise());
CALL_SUBTEST_1(test_gpu_props());