 namespace reduce_detail {
 
 static constexpr int kMaxReductionBlockSize = 512;
+static constexpr int64_t kMaxGridSize = int64_t{0x7fffffffUL};
 
 int64_t RoundUpToPowerOf2(int64_t x) {
     --x;

            switch (arg.out_strides().ndim()) {
                case 0:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 4, 0, 4>(arg), reduce_block_size, impl);
                    return;
                case 1:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 4, 1, 3>(arg), reduce_block_size, impl);
                    return;
                case 2:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 4, 2, 2>(arg), reduce_block_size, impl);
                    return;
                case 3:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 4, 3, 1>(arg), reduce_block_size, impl);
                    return;
                case 4:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 4, 4, 0>(arg), reduce_block_size, impl);
                    return;
            }
            XCHAINER_NEVER_REACH();

            switch (arg.out_strides().ndim()) {
                case 0:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 2, 0, 2>(arg), out_block_size, reduce_block_size, impl);
                    return;
                case 1:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 2, 1, 1>(arg), out_block_size, reduce_block_size, impl);
                    return;
                case 2:
                    reduce_detail::ReductionKernel<<<grid_size, block_size, shared_mem_size>>>(
                            MakeReductionKernelArg<In, Out, 2, 2, 0>(arg), out_block_size, reduce_block_size, impl);
                    return;
            }
            XCHAINER_NEVER_REACH();
