Skip to content

Commit 9bc0401

Browse files
committed
add custom enzyme markers to target methods
1 parent f319dd9 commit 9bc0401

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

+5
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,11 @@ fn generate_enzyme_call<'ll>(
361361
let attr = llvm::AttributeKind::NoInline.create_attr(cx.llcx);
362362
attributes::apply_to_llfn(ad_fn, Function, &[attr]);
363363

364+
// We add a made-up attribute just such that we can recognize it after AD to update
365+
// (no)-inline attributes. We'll then also remove this attribute.
366+
let enzyme_marker_attr = llvm::CreateAttrString(cx.llcx, "enzyme_marker");
367+
attributes::apply_to_llfn(outer_fn, Function, &[enzyme_marker_attr]);
368+
364369
// first, remove all calls from fnc
365370
let entry = llvm::LLVMGetFirstBasicBlock(outer_fn);
366371
let br = llvm::LLVMRustGetTerminator(entry);

0 commit comments

Comments
 (0)