Upcoming CUB/RDC changes

[gevtushenko]

Relocatable Device Code (RDC) doesn't work well with __launch_bounds__ in some cases:

// type.cuh
#pragma once

#include <cstdio>

#define SN 50
struct data_t {
  int storage[SN];
};

__device__ int foo(data_t* lhs, data_t* rhs);
__device__ int bar(data_t lhs, data_t rhs);

// bar.cu
#include "type.cuh"

__device__ int bar(data_t lhs, data_t rhs) { 
  return 42; 
}

// foo.cu
#include "type.cuh"

__device__ int foo(data_t* lhs, data_t* rhs) {
  for (int i = 0; i < SN; i++) {
    lhs->storage[i] += bar(*lhs, *rhs);
  }
  return false;
}

// main.cu
#include "type.cuh"

__global__ __launch_bounds__(1024) void kernel(data_t *lhs, data_t *rhs) {
  std::printf("%d\n", foo(lhs, rhs));
}

int main() {
}

When the code above is compiled we get following issues:

:nvcc -gencode arch=compute_86,code=sm_86 -rdc=true main.cu foo.cu bar.cu
nvlink error   : entry function '_Z6kernelP6data_tS0_' with max regcount of 64 calls function '_Z3fooP6data_tS0_' with regcount of 206

indicating mismatch in register allocations between kernel and device function. Extensive use of __launch_bounds__ in CUB leads to compilation issues in the presence of RDC. To address this, we are going to create an indirection level that would not specify __launch_bounds__ for CUB kernels in the RDC case. We are also going to clamp the threads-per-block at 128 when RDC is used. This is needed so as not to exceed the registers count available per thread block (see the following PR for details).

The change is expected to deteriorate performance in the RDC case. If you have alternative suggestions, please, feel free to participate in the discussion below.

[jrhemstad]

Note that this is documented in the CUDA Programming Guide

If a kernel is limited to a certain number of registers with the launch_bounds attribute or the --maxrregcount option, then all functions that the kernel calls must not use more than that number of registers; if they exceed the limit, then a link error will be given.

[gevtushenko]

Addition: unless there's another way to address this, we should make new behaviour opt-in so that it's possible to workaround compilation issues. Meanwhile, we'll tune CUB for the case of RDC without __launch_bounds__ and thread block size clamping. If, after tuning, new results are comparable, we'll switch the default to opt-out.