如何使用钩子劫持cudaSetDevice和修改设备id？

Question

我想用钩子(LD_PRELOAD)劫持cudaSetDevice并修改设备id。劫持成功后，gpu任务提交错误。

我试图劫持driver api中的cuCtxCreate函数，但我无法劫持它。同样的方法可以劫持cuDevicePrimaryCtxRetain函数，该函数可以被劫持，并且在修改设备id时会出现错误。

// cuda 9.0 cuda runtime api   
typedef int(*cuda_set_device_fp)(int);

// define dynamic library same name function
int cudaSetDevice(int device)
{
  static void *handle = NULL;
  static cuda_set_device_fp orig_cuda_set_device = NULL;

  if( !handle )
  {
    handle = dlopen("libcuda.so", RTLD_LAZY);
    orig_cuda_set_device = (cuda_set_device_fp)dlsym(handle, "cudaSetDevice");
  }
  device = 1;
  printf("oops!!! hack function invoked. device = %d\n", device);
  return cudaSetDevice(device);
}

劫持成功，将用户映射到设备0上的gpu任务，并重新映射到设备1。

Answer 1

在提供的源代码中，在函数的末尾再次递归调用函数，而不是使用修改后的设备id调用orig_cuda_set_device。这会导致无限递归。