Fix exit handling in codegen

This commit is contained in:
Andrzej Janik 2024-05-16 19:29:37 +02:00
parent 922692d2fa
commit 85a0e600fc

View file

@ -544,7 +544,7 @@ fn emit_method<'a, 'input>(
emit_statement(ctx, is_kernel, statement)?;
}
// happens if there is a post-ret trailing label
terminate_current_block_if_needed(ctx, None);
terminate_current_block_if_not_terminated(ctx, None);
unsafe { LLVMPositionBuilderAtEnd(ctx.builder.get(), bb_with_variables) };
unsafe { LLVMBuildBr(ctx.builder.get(), starting_bb) };
Ok(())
@ -625,6 +625,9 @@ fn emit_statement(
is_kernel: bool,
statement: crate::translate::ExpandedStatement,
) -> Result<(), TranslateError> {
if !matches!(statement, crate::translate::Statement::Label(..)) {
start_next_block_if_terminated(ctx);
}
Ok(match statement {
crate::translate::Statement::Label(label) => emit_label(ctx, label)?,
crate::translate::Statement::Variable(var) => emit_function_variable(ctx, var)?,
@ -1249,7 +1252,6 @@ fn emit_inst_isspacep_impl(
)
}
fn emit_inst_sad(
ctx: &mut EmitContext,
type_: ast::ScalarType,
@ -3583,12 +3585,15 @@ fn emit_store_var(
fn emit_label(ctx: &mut EmitContext, label: Id) -> Result<(), TranslateError> {
let new_block = unsafe { LLVMValueAsBasicBlock(ctx.names.value(label)?) };
terminate_current_block_if_needed(ctx, Some(new_block));
terminate_current_block_if_not_terminated(ctx, None);
unsafe { LLVMPositionBuilderAtEnd(ctx.builder.get(), new_block) };
Ok(())
}
fn terminate_current_block_if_needed(ctx: &mut EmitContext, new_block: Option<LLVMBasicBlockRef>) {
fn terminate_current_block_if_not_terminated(
ctx: &mut EmitContext,
new_block: Option<LLVMBasicBlockRef>,
) {
let current_block = unsafe { LLVMGetInsertBlock(ctx.builder.get()) };
if current_block == ptr::null_mut() {
return;
@ -3603,6 +3608,21 @@ fn terminate_current_block_if_needed(ctx: &mut EmitContext, new_block: Option<LL
};
}
// Some PTX instructions (exit) that are LLVM terminators
// can be emitted in the middle of a basic block
// Happens in eg. Meshroom
fn start_next_block_if_terminated(ctx: &mut EmitContext) {
let current_block = unsafe { LLVMGetInsertBlock(ctx.builder.get()) };
let terminator = unsafe { LLVMGetBasicBlockTerminator(current_block) };
if terminator == ptr::null_mut() {
return;
}
let next_block = unsafe { LLVMCreateBasicBlockInContext(ctx.context.get(), LLVM_UNNAMED) };
unsafe { LLVMBuildBr(ctx.builder.get(), next_block) };
unsafe { LLVMMoveBasicBlockAfter(next_block, current_block) };
unsafe { LLVMPositionBuilderAtEnd(ctx.builder.get(), next_block) };
}
fn emit_method_declaration<'input>(
ctx: &mut EmitContext,
method: &crate::translate::Function<'input>,