Skip to content

Commit

Permalink
[llvm] Fix broken llvm module on CUDA backend when TI_LLVM_15 is on (t…
Browse files Browse the repository at this point in the history
…aichi-dev#6458)

Issue: taichi-dev#6447

### Brief Summary
Previously we inject some cuda kernel metadata after linking, that was
problematic since when you modify metadata for a llvm::Function you'll
have to use the exactly same LLVMContext it was created in. Thus we
failed verifyModule check from LLVM15. And we happened to only TI_WARN
there so didn't notice this issue.
![Screenshot from 2022-10-26
17-10-43](https://user-images.githubusercontent.com/5248122/198422336-fd4ed5e8-aa3c-4baf-841a-95bba5c148e0.png)

This PR fixes it by moving the mark as CUDA step up to parallel
compilation so that it was guaranteed to be in the same LLVMContext.
  • Loading branch information
ailzhang authored Oct 28, 2022
1 parent 87d66a1 commit 10b4844
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 7 deletions.
6 changes: 0 additions & 6 deletions taichi/codegen/cuda/codegen_cuda.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -780,12 +780,6 @@ FunctionType CUDAModuleToFunctionConverter::convert(
auto &mod = data.module;
auto &tasks = data.tasks;
#ifdef TI_WITH_CUDA
for (const auto &task : tasks) {
llvm::Function *func = mod->getFunction(task.name);
TI_ASSERT(func);
tlctx_->mark_function_as_cuda_kernel(func, task.block_dim);
}

auto jit = tlctx_->jit.get();
auto cuda_module =
jit->add_module(std::move(mod), executor_->get_config()->gpu_max_reg);
Expand Down
9 changes: 9 additions & 0 deletions taichi/codegen/llvm/codegen_llvm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2749,6 +2749,15 @@ LLVMCompiledTask TaskCodeGenLLVM::run_compilation() {
emit_to_module();
eliminate_unused_functions();

if (config.arch == Arch::cuda) {
// CUDA specific metadata
for (const auto &task : offloaded_tasks) {
llvm::Function *func = module->getFunction(task.name);
TI_ASSERT(func);
tlctx->mark_function_as_cuda_kernel(func, task.block_dim);
}
}

return {std::move(offloaded_tasks), std::move(module),
std::move(used_tree_ids), std::move(struct_for_tls_sizes)};
}
Expand Down
2 changes: 1 addition & 1 deletion taichi/runtime/cuda/jit_cuda.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ std::string JITSessionCUDA::compile_module_to_ptx(
// Part of this function is borrowed from Halide::CodeGen_PTX_Dev.cpp
if (llvm::verifyModule(*module, &llvm::errs())) {
module->print(llvm::errs(), nullptr);
TI_WARN("Module broken");
TI_ERROR("LLVM Module broken");
}

using namespace llvm;
Expand Down

0 comments on commit 10b4844

Please sign in to comment.