Fix caching multi kernel

EnzymeAD · Jan 7, 2025 · d601991 · d601991
1 parent 9207422
commit d601991
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/enzyme_ad/jax/Passes/LowerKernel.cpp b/src/enzyme_ad/jax/Passes/LowerKernel.cpp
@@ -675,6 +675,7 @@ CallInfo CompileKernel(SymbolTableCollection &symbolTable, mlir::Location loc,
         return {};
 
       ptr = CompileHostModule(ss.str(), submod, run_init, &cuLaunchKernelPtr);
+      kernels[ss.str()] = ptr;
 
       submod.erase();
     }