From fead0bb17f840cbe11609f60baa96c68f9dd51e1 Mon Sep 17 00:00:00 2001
From: Jeremy L Thompson <jeremy@jeremylt.org>
Date: Wed, 16 Oct 2024 14:47:41 -0600
Subject: [PATCH] jit - I include JiT source dirs set

---
 backends/cuda/ceed-cuda-compile.cpp | 13 ++++++++++-
 backends/hip/ceed-hip-compile.cpp   | 14 +++++++++++-
 include/ceed/backend.h              |  2 ++
 interface/ceed.c                    | 34 +++++++++++++++++++++++++++++
 4 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/backends/cuda/ceed-cuda-compile.cpp b/backends/cuda/ceed-cuda-compile.cpp
index fd2cdfd713..d415aec56b 100644
--- a/backends/cuda/ceed-cuda-compile.cpp
+++ b/backends/cuda/ceed-cuda-compile.cpp
@@ -93,7 +93,18 @@ int CeedCompile_Cuda(Ceed ceed, const char *source, CUmodule *module, const Ceed
       + std::to_string(prop.major) + std::to_string(prop.minor);
   opts[1] = arch_arg.c_str();
   opts[2] = "-Dint32_t=int";
-  opts[3] = "-I/home/jeremy/Dev/libCEED/include/ceed/jit-source/"
+  {
+    const char       **jit_source_dirs;
+    CeedInt            num_jit_source_dirs;
+    std::ostringstream include_dirs_arg;
+
+    CeedCallBackend(CeedGetJitSourceRoots(ceed, &num_jit_source_dirs, &jit_source_dirs));
+    for (CeedInt i = 0; i < num_jit_source_dirs; i++) {
+      include_dirs_arg << "-I" << jit_source_dirs[i] << " ";
+    }
+    CeedCallBackend(CeedRestoreJitSourceRoots(ceed, &jit_source_dirs));
+    opts[3] = include_dirs_arg.str().c_str();
+  }
 
   // Add string source argument provided in call
   code << source;
diff --git a/backends/hip/ceed-hip-compile.cpp b/backends/hip/ceed-hip-compile.cpp
index cafb79ed7f..edf2cb4b55 100644
--- a/backends/hip/ceed-hip-compile.cpp
+++ b/backends/hip/ceed-hip-compile.cpp
@@ -37,7 +37,7 @@ int CeedCompile_Hip(Ceed ceed, const char *source, hipModule_t *module, const Ce
   size_t                 ptx_size;
   char                  *jit_defs_source, *ptx;
   const char            *jit_defs_path;
-  const int              num_opts = 3;
+  const int              num_opts = 4;
   const char            *opts[num_opts];
   int                    runtime_version;
   hiprtcProgram          prog;
@@ -90,6 +90,18 @@ int CeedCompile_Hip(Ceed ceed, const char *source, hipModule_t *module, const Ce
   std::string arch_arg = "--gpu-architecture=" + std::string(prop.gcnArchName);
   opts[1]              = arch_arg.c_str();
   opts[2]              = "-munsafe-fp-atomics";
+  {
+    const char       **jit_source_dirs;
+    CeedInt            num_jit_source_dirs;
+    std::ostringstream include_dirs_arg;
+
+    CeedCallBackend(CeedGetJitSourceRoots(ceed, &num_jit_source_dirs, &jit_source_dirs));
+    for (CeedInt i = 0; i < num_jit_source_dirs; i++) {
+      include_dirs_arg << "-I" << jit_source_dirs[i] << " ";
+    }
+    CeedCallBackend(CeedRestoreJitSourceRoots(ceed, &jit_source_dirs));
+    opts[3] = include_dirs_arg.str().c_str();
+  }
 
   // Add string source argument provided in call
   code << source;
diff --git a/include/ceed/backend.h b/include/ceed/backend.h
index 05da6f8981..43f2d52d20 100644
--- a/include/ceed/backend.h
+++ b/include/ceed/backend.h
@@ -254,6 +254,8 @@ CEED_EXTERN int CeedSetData(Ceed ceed, void *data);
 CEED_EXTERN int CeedReference(Ceed ceed);
 CEED_EXTERN int CeedGetWorkVector(Ceed ceed, CeedSize len, CeedVector *vec);
 CEED_EXTERN int CeedRestoreWorkVector(Ceed ceed, CeedVector *vec);
+CEED_EXTERN int CeedGetJitSourceRoots(Ceed ceed, CeedInt *num_source_roots, const char ***jit_source_roots);
+CEED_EXTERN int CeedRestoreJitSourceRoots(Ceed ceed, const char ***jit_source_roots);
 
 CEED_EXTERN int CeedVectorHasValidArray(CeedVector vec, bool *has_valid_array);
 CEED_EXTERN int CeedVectorHasBorrowedArrayOfType(CeedVector vec, CeedMemType mem_type, bool *has_borrowed_array_of_type);
diff --git a/interface/ceed.c b/interface/ceed.c
index 1becb3de14..3e0a0569ae 100644
--- a/interface/ceed.c
+++ b/interface/ceed.c
@@ -863,6 +863,40 @@ int CeedRestoreWorkVector(Ceed ceed, CeedVector *vec) {
   // LCOV_EXCL_STOP
 }
 
+/**
+  @brief Retrieve list ofadditional JiT source roots from `Ceed` context.
+
+  Note: The caller is responsible for restoring `jit_source_roots` with @ref CeedRestoreJitSourceRoots().
+
+  @param[in]  ceed             `Ceed` context
+  @param[out] num_source_roots Number of JiT source directories
+  @param[out] jit_source_roots Absolute paths to additional JiT source directories
+
+  @return An error code: 0 - success, otherwise - failure
+
+  @ref Backend
+**/
+int CeedGetJitSourceRoots(Ceed ceed, CeedInt *num_source_roots, const char ***jit_source_roots) {
+  *num_source_roots = ceed->num_jit_source_roots;
+  *jit_source_roots = (const char **)ceed->jit_source_roots;
+  return CEED_ERROR_SUCCESS;
+}
+
+/**
+  @brief Restore list of additional JiT source roots from with @ref CeedGetJitSourceRoots()
+
+  @param[in]  ceed             `Ceed` context
+  @param[out] jit_source_roots Absolute paths to additional JiT source directories
+
+  @return An error code: 0 - success, otherwise - failure
+
+  @ref Backend
+**/
+int CeedRestoreJitSourceRoots(Ceed ceed, const char ***jit_source_roots) {
+  *jit_source_roots = NULL;
+  return CEED_ERROR_SUCCESS;
+}
+
 /// @}
 
 /// ----------------------------------------------------------------------------