CUDA (.cu文件) 通过runtime API加载PTX binary的过程分析

背景

我们经常会用nvcc直接编译一个.cu文件到binary来执行。执行时，nvcc生成的其他代码帮我们做了很多事情来实现加载PTX ELF binary等事情。现在我们需要拦截nvcc生成的代码所做的这些事情，所以写一篇文章分析一下。

细节

#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <ostream>
#include <stdio.h>
#include <stdlib.h>
#include <cuda_runtime.h>
#include <vector>

__global__ void sumArray(uint64_t *input, uint64_t *output, size_t size)
{
	for (size_t i = 0; i < size; i++)
		*output += input[i];
	printf("From device side: sum = %lu\n", (unsigned long)*output);
}

constexpr size_t ARR_SIZE = 1000;

int main()
{
	std::vector<uint64_t> arr;
	for (size_t i = 1; i <= ARR_SIZE; i++) {
		arr.push_back(i);
	}
	auto data_size = sizeof(arr[0]) * arr.size();

	uint64_t *d_input, *d_output;
	cudaMalloc(&d_input, data_size);
	cudaMalloc(&d_output, sizeof(arr[0]));
	cudaMemcpy(d_input, arr.data(), data_size, cudaMemcpyHostToDevice);
	sumArray<<<1, 1, 1>>>(d_input, d_output, arr.size());
	uint64_t host_sum;
	cudaMemcpy(&host_sum, d_output, sizeof(arr[0]), cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();
	std::cout << "Sum is " << host_sum << std::endl;
	cudaFree(d_input);
	cudaFree(d_output);

	return 0;
}

这是一个很简单的CUDA程序。传一个数组到device上，device求和并写回内存，而后host读取结果。

使用 nvcc -cuda victim.cu -o victim.cpp 来生成对应展开后的cpp文件。生成后的代码过大，这里就不粘贴了。

初始化过程

nvcc生成的代码里调用了大量内部API来实现加载。这些API的工作逻辑和driverAPI并不一致。


static void __sti____cudaRegisterAll(void) __attribute__((__constructor__));
static void __sti____cudaRegisterAll(void)
{
	__cudaFatCubinHandle =
		__cudaRegisterFatBinary((void *)&__fatDeviceText);
	{
		void (*callback_fp)(void **) =
			(void (*)(void **))(__nv_cudaEntityRegisterCallback);
		(*callback_fp)(__cudaFatCubinHandle);
		__cudaRegisterFatBinaryEnd(__cudaFatCubinHandle);
	}
	atexit(__cudaUnregisterBinaryUtil);
}

可以看到，nvcc生成的代码里给__sti____cudaRegisterAll打了constructor标记，也就是会先执行这个初始化binary的函数。我们可以看到依次调用了这些函数来初始化binary

__cudaRegisterFatBinary：加载一份binary。binary的内容是用内联汇编嵌入在源代码里的ELF
__nv_cudaEntityRegisterCallback：注册使用到的kernel函数。这个函数接下来会解释。
__cudaRegisterFatBinaryEnd：结束函数注册？（我猜的，这是没有文档的东西

同时注册了一个atexit函数__cudaUnregisterBinaryUtil

核函数注册

让我们看__nv_cudaEntityRegisterCallback。

static void __nv_cudaEntityRegisterCallback(void **__T5)
{
	{
		volatile static void **__ref __attribute__((unused));
		__ref = (volatile void **)__T5;
	};
	__nv_save_fatbinhandle_for_managed_rt(__T5);
	__cudaRegisterFunction(__T5,
			       (const char *)((void (*)(uint64_t *, uint64_t *,
							size_t))sumArray),
			       (char *)"_Z8sumArrayPmS_m", "_Z8sumArrayPmS_m",
			       -1, (uint3 *)0, (uint3 *)0, (dim3 *)0, (dim3 *)0,
			       (int *)0);
}

注意到这里会调用__cudaRegisterFunction来将kernel函数的符号名（定义在PTX ELF binary中）和host上的包装函数的地址关联起来。

host上的包装函数

void __device_stub__Z8sumArrayPmS_m(uint64_t *__par0, uint64_t *__par1,
				    size_t __par2)
{
	void *__args_arr[3];
	int __args_idx = 0;
	__args_arr[__args_idx] = (void *)(char *)&__par0;
	++__args_idx;
	__args_arr[__args_idx] = (void *)(char *)&__par1;
	++__args_idx;
	__args_arr[__args_idx] = (void *)(char *)&__par2;
	++__args_idx;
	{
		volatile static char *__f __attribute__((unused));
		__f = ((char *)((
			void (*)(uint64_t *, uint64_t *, size_t))sumArray));
		dim3 __gridDim, __blockDim;
		size_t __sharedMem;
		cudaStream_t __stream;
		if (__cudaPopCallConfiguration(&__gridDim, &__blockDim,
					       &__sharedMem,
					       &__stream) != cudaSuccess)
			return;
		if (__args_idx == 0) {
			(void)cudaLaunchKernel(
				((char *)((void (*)(uint64_t *, uint64_t *,
						    size_t))sumArray)),
				__gridDim, __blockDim, &__args_arr[__args_idx],
				__sharedMem, __stream);
		} else {
			(void)cudaLaunchKernel(
				((char *)((void (*)(uint64_t *, uint64_t *,
						    size_t))sumArray)),
				__gridDim, __blockDim, &__args_arr[0],
				__sharedMem, __stream);
		}
	};
}

void sumArray(uint64_t *__cuda_0, uint64_t *__cuda_1, size_t __cuda_2)
{
	__device_stub__Z8sumArrayPmS_m(__cuda_0, __cuda_1, __cuda_2);
}

包装函数内的核心是cudaLaunchKernel，直接传入了host上的函数地址来执行kernel。由于先前已经关联过了host上的包装函数的地址和kernel的符号名，所以只需要提供host上的包装函数的地址即可。

main函数

int main()
{
	std::vector<unsigned long> arr;
	for (size_t i = (1); i <= ARR_SIZE; i++) {
		arr.push_back(i);
	}
	auto data_size = sizeof arr[0] * arr.size();
	uint64_t *d_input, *d_output;
	cudaMalloc(&d_input, data_size);
	cudaMalloc(&d_output, sizeof arr[0]);
	cudaMemcpy(d_input, arr.data(), data_size, cudaMemcpyHostToDevice);
	(__cudaPushCallConfiguration(1, 1, 1)) ?
		(void)0 :
		sumArray(d_input, d_output, arr.size());
	uint64_t host_sum;
	cudaMemcpy(&host_sum, d_output, sizeof arr[0], cudaMemcpyDeviceToHost);
	cudaDeviceSynchronize();
	(((((std::cout << ("Sum is "))) << host_sum)) << (std::endl));
	cudaFree(d_input);
	cudaFree(d_output);
	return 0;
}

可以看到，main函数直接执行host上的包装函数即可完成核函数运行。

反注册binary

注意到我们之前提到的atexit里面的退出函数。

static void __cudaUnregisterBinaryUtil(void)
{
	____nv_dummy_param_ref((void *)&__cudaFatCubinHandle);
	__cudaUnregisterFatBinary(__cudaFatCubinHandle);
}

这里调用一个内部函数来销毁binary handle。