通过动态加载库和修改LLVM IR 来实现--动态修改生成的GPU代码

参考：https://github.com/DependableSystemsLab/LLFI-GPU/blob/master/bamboo_lib/src/profile/README.md

背景介绍

在进行CUDA编程时，我们希望能动态修改生成的GPU代码，以便插入故障注入代码或做其他调整。这篇文章介绍了一种方法，通过动态加载库和修改LLVM IR（中间表示）来实现这一目标。

什么是LLVM IR？

LLVM IR（Intermediate Representation，中间表示）是一种抽象的低级编程语言，用于描述程序的行为。它介于高级语言（如C++）和机器代码之间，是编译过程中的一个中间步骤。LLVM IR可以被优化，然后转换为目标机器代码。

什么是CUDA？

CUDA（Compute Unified Device Architecture）是NVIDIA推出的一个并行计算平台和编程模型，允许开发者利用NVIDIA GPU的强大计算能力。使用CUDA编写的代码会被编译成PTX（Parallel Thread Execution）代码，然后在GPU上执行。

问题和挑战

1. 如何定制CUDA源代码的编译：由于没有开源的CUDA前端工具，难以直接修改CUDA源代码的编译过程。
2. GPU代码优化：生成高效的GPU代码需要特定的LLVM优化集，NVIDIA的libNVVM库应用了一些专有的优化，标准的LLVM工具无法完全再现这些优化。

提出的方法

为了解决这些问题，我们创建了一个特殊的动态库，这个库附加到NVIDIA CUDA编译器中，使我们能够获取并修改未优化和优化后的LLVM IR代码。

技术细节

获取未优化的LLVM IR

1. 截获libNVVM的函数调用：我们通过截获nvvmAddModuleToProgram函数的调用来获取未优化的LLVM IR。这个函数的第二个参数是LLVM位码字符串。在CUDA编译过程中，nvcc（NVIDIA的CUDA编译器驱动程序）会调用这个函数，将生成的LLVM位码传递给NVVM库以生成PTX代码。
2. 解析位码字符串：使用LLVM库函数将位码字符串解析为LLVM Module实例，然后可以打印出IR代码。

string source = "";
source.reserve(size);
source.assign(bitcode, bitcode + size);
MemoryBuffer *input = MemoryBuffer::getMemBuffer(source);
string err;
LLVMContext &context = getGlobalContext();
initial_module = ParseBitcodeFile(input, context, &err);
if (!initial_module)
	cerr << "Error parsing module bitcode : " << err;

outs() << *initial_module;

3. 修改LLVM IR：可以通过将修改后的LLVM Module导出为位码字符串来实现对未优化LLVM的动态修改。

SmallVector<char, 128> output;
raw_svector_ostream outputStream(output);
WriteBitcodeToFile(initial_module, outputStream);
outputStream.flush();

// 调用实际的nvvmAddModuleToProgram
return nvvmAddModuleToProgram_real(prog, output.data(), output.size(), name);

获取优化后的LLVM IR

1. 拦截内存分配：通过拦截malloc函数，找到在nvvmCompileProgram开始时分配的Module大小的空间，并保存这个Module实例。

void* result = malloc_real(size);

if (called_compile)
{
	if (size == sizeof(Module))
		optimized_module = (Module*)result;
}

2. 适当时机拦截Module：选择合适的时机（比如在调用localtime函数时）拦截并修改这个Module的内容。

struct tm *localtime(const time_t *timep)
{
	static bool localtime_first_call = true;

	bind_lib(LIBC);
	bind_sym(libc, localtime, struct tm*, const time_t*);

	if (getenv("CICC_MODIFY_OPT_MODULE") && called_compile && localtime_first_call)
	{
		localtime_first_call = false;

		writeIrToFile(optimized_module, "opt_bamboo_before.ll");

		modifyModule(optimized_module);

		writeIrToFile(optimized_module, "opt_bamboo_after.ll");

	}

	return localtime_real(timep);
}

示例

假设有以下CUDA源代码：

extern "C" __device__ void kernel(int* result) { *result = 1; }

编译动态库，指定NVVM和LLVM 3.0头文件的包含路径：

$ make
g++ -g -D__STDC_LIMIT_MACROS -D__STDC_CONSTANT_MACROS -I/opt/llvm-3.0/include -I/opt/cuda/nvvm/include/ -fPIC cicc.cpp -shared -o libcicc.so -ldl
g++ -g -I/opt/cuda/nvvm/include/ -fPIC nvcc.cpp -shared -o libnvcc.so -ldl

设置环境变量并运行：

CICC_MODIFY_UNOPT_MODULE=1 LD_PRELOAD=./libnvcc.so nvcc -arch=sm_30 test.cu -c -keep
CICC_MODIFY_OPT_MODULE=1 LD_PRELOAD=./libnvcc.so nvcc -arch=sm_30 test.cu -c -keep

modifyModule函数示例（为函数名添加后缀）：

void modifyModule(Module* module)
{
	if (!module) return;

	// 为函数名添加后缀
	for (Module::iterator i = module->begin(), e = module->end(); i != e; i++)
		i->setName(i->getName() + "_modified");
}

运行后的PTX文件内容：

$ cat test.ptx 
//
// Generated by NVIDIA NVVM Compiler
// Compiler built on Thu Mar 13 19:31:35 2014 (1394735495)
// Cuda compilation tools, release 6.0, V6.0.1
//

.version 4.0
.target sm_30
.address_size 64

.visible .func kernel_modified(
	.param .b64 kernel_modified_param_0
)
{
	.reg .s32 	%r<2>;
	.reg .s64 	%rd<2>;


	ld.param.u64 	%rd1, [kernel_modified_param_0];
	mov.u32 	%r1, 1;
	st.u32 	[%rd1], %r1;
	ret;
}

这样就实现了对CUDA代码的LLVM IR的动态修改，并可以在生成PTX代码之前进行优化和故障注入等操作。这为开发领域特定的编译器提供了极大的灵活性和可能性。