mirror of
				https://github.com/rust-lang/rust.git
				synced 2025-10-31 21:16:44 +00:00 
			
		
		
		
	 c6bf3a01ef
			
		
	
	
		c6bf3a01ef
		
			
		
	
	
	
	
		
			
			Autodiff batching
Enzyme supports batching, which is especially known from the ML side when training neural networks.
There we would normally have a training loop, where in each iteration we would pass in some data (e.g. an image), and a target vector. Based on how close we are with our prediction we compute our loss, and then use backpropagation to compute the gradients and update our weights.
That's quite inefficient, so what you normally do is passing in a batch of 8/16/.. images and targets, and compute the gradients for those all at once, allowing better optimizations.
Enzyme supports batching in two ways, the first one (which I implemented here) just accepts a Batch size,
and then each Dual/Duplicated argument has not one, but N shadow arguments.  So instead of
```rs
for i in 0..100 {
   df(x[i], y[i], 1234);
}
```
You can now do
```rs
for i in 0..100.step_by(4) {
   df(x[i+0],x[i+1],x[i+2],x[i+3], y[i+0], y[i+1], y[i+2], y[i+3], 1234);
}
```
which will give the same results, but allows better compiler optimizations. See the testcase for details.
There is a second variant, where we can mark certain arguments and instead of having to pass in N shadow arguments, Enzyme assumes that the argument is N times longer. I.e. instead of accepting 4 slices with 12 floats each, we would accept one slice with 48 floats. I'll implement this over the next days.
I will also add more tests for both modes.
For any one preferring some more interactive explanation, here's a video of Tim's llvm dev talk, where he presents his work. https://www.youtube.com/watch?v=edvaLAL5RqU
I'll also add some other docs to the dev guide and user docs in another PR.
r? ghost
Tracking:
- https://github.com/rust-lang/rust/issues/124509
- https://github.com/rust-lang/rust/issues/135283
		
	
			
		
			
				
	
	
		
			135 lines
		
	
	
		
			4.2 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
			
		
		
	
	
			135 lines
		
	
	
		
			4.2 KiB
		
	
	
	
		
			Rust
		
	
	
	
	
	
| #![allow(non_camel_case_types)]
 | |
| #![expect(dead_code)]
 | |
| 
 | |
| use libc::{c_char, c_uint};
 | |
| 
 | |
| use super::MetadataKindId;
 | |
| use super::ffi::{AttributeKind, BasicBlock, Metadata, Module, Type, Value};
 | |
| use crate::llvm::Bool;
 | |
| 
 | |
| #[link(name = "llvm-wrapper", kind = "static")]
 | |
| unsafe extern "C" {
 | |
|     // Enzyme
 | |
|     pub(crate) safe fn LLVMRustHasMetadata(I: &Value, KindID: MetadataKindId) -> bool;
 | |
|     pub(crate) fn LLVMRustEraseInstUntilInclusive(BB: &BasicBlock, I: &Value);
 | |
|     pub(crate) fn LLVMRustGetLastInstruction<'a>(BB: &BasicBlock) -> Option<&'a Value>;
 | |
|     pub(crate) fn LLVMRustDIGetInstMetadata(I: &Value) -> Option<&Metadata>;
 | |
|     pub(crate) fn LLVMRustEraseInstFromParent(V: &Value);
 | |
|     pub(crate) fn LLVMRustGetTerminator<'a>(B: &BasicBlock) -> &'a Value;
 | |
|     pub(crate) fn LLVMRustVerifyFunction(V: &Value, action: LLVMRustVerifierFailureAction) -> Bool;
 | |
|     pub(crate) fn LLVMRustHasAttributeAtIndex(V: &Value, i: c_uint, Kind: AttributeKind) -> bool;
 | |
|     pub(crate) fn LLVMRustGetArrayNumElements(Ty: &Type) -> u64;
 | |
| }
 | |
| 
 | |
| unsafe extern "C" {
 | |
|     // Enzyme
 | |
|     pub(crate) fn LLVMDumpModule(M: &Module);
 | |
|     pub(crate) fn LLVMDumpValue(V: &Value);
 | |
|     pub(crate) fn LLVMGetFunctionCallConv(F: &Value) -> c_uint;
 | |
|     pub(crate) fn LLVMGetReturnType(T: &Type) -> &Type;
 | |
|     pub(crate) fn LLVMGetParams(Fnc: &Value, parms: *mut &Value);
 | |
|     pub(crate) fn LLVMGetNamedFunction(M: &Module, Name: *const c_char) -> Option<&Value>;
 | |
| }
 | |
| 
 | |
| #[repr(C)]
 | |
| #[derive(Copy, Clone, PartialEq)]
 | |
| pub(crate) enum LLVMRustVerifierFailureAction {
 | |
|     LLVMAbortProcessAction = 0,
 | |
|     LLVMPrintMessageAction = 1,
 | |
|     LLVMReturnStatusAction = 2,
 | |
| }
 | |
| 
 | |
| #[cfg(llvm_enzyme)]
 | |
| pub(crate) use self::Enzyme_AD::*;
 | |
| 
 | |
| #[cfg(llvm_enzyme)]
 | |
| pub(crate) mod Enzyme_AD {
 | |
|     use libc::c_void;
 | |
|     unsafe extern "C" {
 | |
|         pub(crate) fn EnzymeSetCLBool(arg1: *mut ::std::os::raw::c_void, arg2: u8);
 | |
|     }
 | |
|     unsafe extern "C" {
 | |
|         static mut EnzymePrintPerf: c_void;
 | |
|         static mut EnzymePrintActivity: c_void;
 | |
|         static mut EnzymePrintType: c_void;
 | |
|         static mut EnzymePrint: c_void;
 | |
|         static mut EnzymeStrictAliasing: c_void;
 | |
|         static mut looseTypeAnalysis: c_void;
 | |
|         static mut EnzymeInline: c_void;
 | |
|         static mut RustTypeRules: c_void;
 | |
|     }
 | |
|     pub(crate) fn set_print_perf(print: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintPerf), print as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_print_activity(print: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintActivity), print as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_print_type(print: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrintType), print as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_print(print: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymePrint), print as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_strict_aliasing(strict: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeStrictAliasing), strict as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_loose_types(loose: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(looseTypeAnalysis), loose as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_inline(val: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(EnzymeInline), val as u8);
 | |
|         }
 | |
|     }
 | |
|     pub(crate) fn set_rust_rules(val: bool) {
 | |
|         unsafe {
 | |
|             EnzymeSetCLBool(std::ptr::addr_of_mut!(RustTypeRules), val as u8);
 | |
|         }
 | |
|     }
 | |
| }
 | |
| 
 | |
| #[cfg(not(llvm_enzyme))]
 | |
| pub(crate) use self::Fallback_AD::*;
 | |
| 
 | |
| #[cfg(not(llvm_enzyme))]
 | |
| pub(crate) mod Fallback_AD {
 | |
|     #![allow(unused_variables)]
 | |
| 
 | |
|     pub(crate) fn set_inline(val: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_print_perf(print: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_print_activity(print: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_print_type(print: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_print(print: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_strict_aliasing(strict: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_loose_types(loose: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
|     pub(crate) fn set_rust_rules(val: bool) {
 | |
|         unimplemented!()
 | |
|     }
 | |
| }
 |