cuda: Add support for constexpr array compilation inputs

Author: zatkins-dev

backends/cuda/ceed-cuda-compile.cpp

@@ -21,6 +21,7 @@
 
 #include <cstdlib>
 #include <fstream>
+#include <iomanip>
 #include <iostream>
 #include <sstream>
 #include <string>
@@ -69,7 +70,7 @@ using std::ofstream;
 using std::ostringstream;
 
 static int CeedCompileCore_Cuda(Ceed ceed, const char *source, const bool throw_error, bool *is_compile_good, CUmodule *module,
-                                const CeedInt num_defines, va_list args) {
+                                const CeedInt num_defines, const CeedInt num_extra, va_list args) {
   size_t                ptx_size;
   char                 *ptx;
   const int             num_opts            = 4;
@@ -106,6 +107,18 @@ static int CeedCompileCore_Cuda(Ceed ceed, const char *source, const bool throw_
   // Standard libCEED definitions for CUDA backends
   code << "#include <ceed/jit-source/cuda/cuda-jit.h>\n\n";
 
+  // Insert kernel specific extra code, e.g. array constants
+  if (num_extra > 0) {
+    char *val;
+
+    code << "// Kernel defined extra code\n";
+    for (int i = 0; i < num_extra; i++) {
+      val = va_arg(args, char *);
+      code << val << "\n";
+    }
+    code << "\n";
+  }
+
   // Non-macro options
   CeedCallBackend(CeedCalloc(num_opts, &opts));
   opts[0] = "-default-device";
@@ -404,12 +417,46 @@ static int CeedCompileCore_Cuda(Ceed ceed, const char *source, const bool throw_
   return CEED_ERROR_SUCCESS;
 }
 
+template <typename ArrayT>
+struct CeedArrayView {
+  const ArrayT *array;
+  CeedInt       size;
+
+  CeedArrayView<ArrayT>(const ArrayT *array_, CeedInt size_) : array(array_), size(size_) {}
+};
+
+template <typename OStream, typename ArrayT>
+OStream &operator<<(OStream &ostream, const CeedArrayView<ArrayT> &view) {
+  ostream << "{";
+  for (CeedInt i = 0; i < view.size; i++) ostream << std::setprecision(17) << view.array[i] << (i == view.size - 1 ? "}" : ", ");
+  return ostream;
+}
+
+int CeedBuildArrayConstantSize_Cuda(Ceed ceed, const char *name, CeedInt length, const CeedSize *array, char **line) {
+  std::ostringstream code;
+
+  code << "constexpr CeedSize " << name << "[" << length << "] = " << CeedArrayView<CeedSize>(array, length) << ";";
+  CeedCallBackend(CeedStringAllocCopy(code.str().c_str(), line));
+  return CEED_ERROR_SUCCESS;
+}
+
+int CeedCompileExtra_Cuda(Ceed ceed, const char *source, CUmodule *module, const CeedInt num_defines, const CeedInt num_extra, ...) {
+  bool    is_compile_good = true;
+  va_list args;
+
+  va_start(args, num_extra);
+  const CeedInt ierr = CeedCompileCore_Cuda(ceed, source, true, &is_compile_good, module, num_defines, num_extra, args);
+  va_end(args);
+  CeedCallBackend(ierr);
+  return CEED_ERROR_SUCCESS;
+}
+
 int CeedCompile_Cuda(Ceed ceed, const char *source, CUmodule *module, const CeedInt num_defines, ...) {
   bool    is_compile_good = true;
   va_list args;
 
   va_start(args, num_defines);
-  const CeedInt ierr = CeedCompileCore_Cuda(ceed, source, true, &is_compile_good, module, num_defines, args);
+  const CeedInt ierr = CeedCompileCore_Cuda(ceed, source, true, &is_compile_good, module, num_defines, 0, args);
 
   va_end(args);
   CeedCallBackend(ierr);
@@ -420,7 +467,7 @@ int CeedTryCompile_Cuda(Ceed ceed, const char *source, bool *is_compile_good, CU
   va_list args;
 
   va_start(args, num_defines);
-  const CeedInt ierr = CeedCompileCore_Cuda(ceed, source, false, is_compile_good, module, num_defines, args);
+  const CeedInt ierr = CeedCompileCore_Cuda(ceed, source, false, is_compile_good, module, num_defines, 0, args);
 
   va_end(args);
   CeedCallBackend(ierr);

backends/cuda/ceed-cuda-compile.h

@@ -12,6 +12,9 @@
 
 static inline CeedInt CeedDivUpInt(CeedInt numerator, CeedInt denominator) { return (numerator + denominator - 1) / denominator; }
 
+CEED_INTERN int CeedBuildArrayConstantSize_Cuda(Ceed ceed, const char *name, CeedInt length, const CeedSize *array, char **line);
+
+CEED_INTERN int CeedCompileExtra_Cuda(Ceed ceed, const char *source, CUmodule *module, const CeedInt num_defines, const CeedInt num_extra, ...);
 CEED_INTERN int CeedCompile_Cuda(Ceed ceed, const char *source, CUmodule *module, const CeedInt num_defines, ...);
 CEED_INTERN int CeedTryCompile_Cuda(Ceed ceed, const char *source, bool *is_compile_good, CUmodule *module, const CeedInt num_defines, ...);