added PrintTAFn flag for autodiff

Signed-off-by: Karan Janthe <karanjanthe@gmail.com>
This commit is contained in:
Karan Janthe
2025-06-25 02:06:53 +00:00
parent 066ae4cf98
commit 7b1c89f2b5
6 changed files with 42 additions and 4 deletions

View File

@@ -57,14 +57,19 @@ pub(crate) use self::Enzyme_AD::*;
#[cfg(llvm_enzyme)]
pub(crate) mod Enzyme_AD {
use std::ffi::{CString, c_char};
use libc::c_void;
unsafe extern "C" {
pub(crate) fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
pub(crate) fn EnzymeSetCLString(arg1: *mut ::std::os::raw::c_void, arg2: *const c_char);
}
unsafe extern "C" {
static mut EnzymePrintPerf: c_void;
static mut EnzymePrintActivity: c_void;
static mut EnzymePrintType: c_void;
static mut EnzymeFunctionToAnalyze: c_void;
static mut EnzymePrint: c_void;
static mut EnzymeStrictAliasing: c_void;
static mut looseTypeAnalysis: c_void;
@@ -86,6 +91,15 @@ pub(crate) mod Enzyme_AD {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintType), print as u8);
}
}
pub(crate) fn set_print_type_fun(fun_name: &str) {
let c_fun_name = CString::new(fun_name).unwrap();
unsafe {
EnzymeSetCLString(
std::ptr::addr_of_mut!(EnzymeFunctionToAnalyze),
c_fun_name.as_ptr() as *const c_char,
);
}
}
pub(crate) fn set_print(print: bool) {
unsafe {
EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrint), print as u8);
@@ -132,6 +146,9 @@ pub(crate) mod Fallback_AD {
pub(crate) fn set_print_type(print: bool) {
unimplemented!()
}
pub(crate) fn set_print_type_fun(fun_name: &str) {
unimplemented!()
}
pub(crate) fn set_print(print: bool) {
unimplemented!()
}