如何使用钩子劫持cudaSetDevice并修改设备ID？

Question

我想用钩子（LD_PRELOAD）劫持cudaSetDevice并修改设备ID。劫持成功后，gpu任务提交将是错误的。

我试图劫持驱动程序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。这将导致无限递归。