Skip to content

Commit 867134e

Browse files
committed
add custom enzyme markers to target methods
1 parent 44a460c commit 867134e

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,11 @@ fn generate_enzyme_call<'ll>(
334334
let attr = llvm::AttributeKind::NoInline.create_attr(cx.llcx);
335335
attributes::apply_to_llfn(ad_fn, Function, &[attr]);
336336

337+
// We add a made-up attribute just such that we can recognize it after AD to update
338+
// (no)-inline attributes. We'll then also remove this attribute.
339+
let enzyme_marker_attr = llvm::CreateAttrString(cx.llcx, "enzyme_marker");
340+
attributes::apply_to_llfn(outer_fn, Function, &[enzyme_marker_attr]);
341+
337342
// first, remove all calls from fnc
338343
let entry = llvm::LLVMGetFirstBasicBlock(outer_fn);
339344
let br = llvm::LLVMRustGetTerminator(entry);

0 commit comments

Comments
 (0)