more iffy improvements

Author: ZuseZ4

compiler/rustc_codegen_llvm/src/lib.rs

@@ -237,13 +237,33 @@ impl WriteBackendMethods for LlvmCodegenBackend {
         diff_fncs: Vec<AutoDiffItem>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError> {
-        if cgcx.lto != Lto::Fat {
-            let dcx = cgcx.create_dcx();
-            return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
-        }
+        //if cgcx.lto != Lto::Fat {
+        //    let dcx = cgcx.create_dcx();
+        //    return Err(dcx.handle().emit_almost_fatal(AutoDiffWithoutLTO));
+        //}
         let module_llvm = &module.module_llvm;
         builder::autodiff::differentiate(module_llvm, cgcx, diff_fncs, config)
     }
+    fn autodiff_thin(
+        cgcx: &CodegenContext<Self>,
+        thin_module: &ThinModule<Self>,
+        diff_fncs: Vec<AutoDiffItem>,
+        config: &ModuleConfig,
+    ) -> Result<(), FatalError> {
+        let dcx = cgcx.create_dcx();
+        let dcx = dcx.handle();
+
+        let module_name = &thin_module.shared.module_names[thin_module.idx];
+
+        // Right now the implementation we've got only works over serialized
+        // modules, so we create a fresh new LLVM context and parse the module
+        // into that context. One day, however, we may do this for upstream
+        // crates but for locally codegened modules we may be able to reuse
+        // that LLVM Context and Module.
+        let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx)?;
+
+        builder::autodiff::differentiate(&module_llvm, cgcx, diff_fncs, config)
+    }
 }
 
 impl LlvmCodegenBackend {

compiler/rustc_codegen_ssa/src/back/lto.rs

@@ -94,10 +94,11 @@ impl<B: WriteBackendMethods> LtoModuleCodegen<B> {
         match &self {
             LtoModuleCodegen::Fat(module) => {
                 B::autodiff(cgcx, &module, diff_fncs, config)?;
+            },
+            LtoModuleCodegen::Thin(module) => {
+                B::autodiff_thin(cgcx, module, diff_fncs, config)?;
             }
-            _ => panic!("autodiff called with non-fat LTO module"),
         }
-
         Ok(self)
     }
 }

compiler/rustc_codegen_ssa/src/back/write.rs

@@ -412,15 +412,18 @@ fn generate_lto_work<B: ExtraBackendMethods>(
         vec![(WorkItem::LTO(module), 0)]
     } else {
         if !autodiff.is_empty() {
-            let dcx = cgcx.create_dcx();
-            dcx.handle().emit_fatal(AutodiffWithoutLto {});
+            //let dcx = cgcx.create_dcx();
+            //dcx.handle().emit_fatal(AutodiffWithoutLto {});
         }
+        let config = cgcx.config(ModuleKind::Regular);
         assert!(needs_fat_lto.is_empty());
         let (lto_modules, copy_jobs) = B::run_thin_lto(cgcx, needs_thin_lto, import_only_modules)
             .unwrap_or_else(|e| e.raise());
         lto_modules
             .into_iter()
             .map(|module| {
+                let mut module =
+                    unsafe { module.autodiff(cgcx, autodiff.clone(), config).unwrap_or_else(|e| e.raise()) };
                 let cost = module.cost();
                 (WorkItem::LTO(module), cost)
             })
@@ -1476,6 +1479,7 @@ fn start_executing_work<B: ExtraBackendMethods>(
                     let autodiff_items = mem::take(&mut autodiff_items);
                     let needs_fat_lto = mem::take(&mut needs_fat_lto);
                     let needs_thin_lto = mem::take(&mut needs_thin_lto);
+                    let autodiff_items = mem::take(&mut autodiff_items);
                     let import_only_modules = mem::take(&mut lto_import_only_modules);
 
                     for (work, cost) in generate_lto_work(

compiler/rustc_codegen_ssa/src/traits/write.rs

@@ -68,6 +68,12 @@ pub trait WriteBackendMethods: 'static + Sized + Clone {
         diff_fncs: Vec<AutoDiffItem>,
         config: &ModuleConfig,
     ) -> Result<(), FatalError>;
+    fn autodiff_thin(
+        cgcx: &CodegenContext<Self>,
+        thin: &ThinModule<Self>,
+        diff_fncs: Vec<AutoDiffItem>,
+        config: &ModuleConfig,
+    ) -> Result<(), FatalError>;
 }
 
 pub trait ThinBufferMethods: Send + Sync {