[CINN] Optimize nvrtc compile time (#71410)

chen2016013 · web-flow · commit 8ad96043a1f3 · 2025-03-10T10:46:47.000+08:00
* update

* update

* update

* update

* update
diff --git a/paddle/cinn/backends/codegen_cuda_dev.cc b/paddle/cinn/backends/codegen_cuda_dev.cc
@@ -16,10 +16,10 @@
 
 namespace cinn {
 namespace backends {
-
-const std::string CodeGenCudaDev::source_header_ =  // NOLINT
-    R"(#include <cstdint>
-
+const std::string CodeGenCudaDev::general_source_header_ =  // NOLINT
+    R"(
+#pragma once
+#include <cstdint>
 #define CINN_WITH_CUDA
 #include "bfloat16.h"
 #include "float16.h"
@@ -34,11 +34,35 @@ using cinn::common::float162;
 using cinn::common::bfloat168;
 using cinn::common::bfloat164;
 using cinn::common::bfloat162;
-
 #include "cinn_cuda_runtime_source.cuh"
+)";
+const std::string CodeGenCudaDev::source_header_ =  // NOLINT
+    R"(
+#pragma once
+#include <cinn_with_cuda_h>
+
+#include <bfloat16_h>
+#include <cstdint>
+#include <float16_h>
+using cinn::common::bfloat16;
+using cinn::common::float16;
+using cinn::common::float8;
+using cinn::common::half4;
+using cinn::common::half8;
+using cinn::common::float168;
+using cinn::common::float164;
+using cinn::common::float162;
+using cinn::common::bfloat168;
+using cinn::common::bfloat164;
+using cinn::common::bfloat162;
+#include <cinn_cuda_runtime_source_h>
+
 )";
 
 const std::string &CodeGenCudaDev::GetSourceHeader() { return source_header_; }
+const std::string &CodeGenCudaDev::GetGeneralSourceHeader() {
+  return general_source_header_;
+}
 
 CodeGenCudaDev::CodeGenCudaDev(Target target) : CodeGenGpuDev(target) {}
 
diff --git a/paddle/cinn/backends/codegen_cuda_dev.h b/paddle/cinn/backends/codegen_cuda_dev.h
@@ -31,10 +31,15 @@ class CodeGenCudaDev : public CodeGenGpuDev {
  public:
   explicit CodeGenCudaDev(Target target);
   static const std::string& GetSourceHeader();
+  static const std::string& GetGeneralSourceHeader();
   void PrintIncludes() override;
 
  private:
   static const std::string source_header_;
+  // general_source_header_ is used for the more general situation, which load
+  // some header files while compiling but not set them into header files while
+  // creating the kernel function.
+  static const std::string general_source_header_;
 };
 
 }  // namespace backends
diff --git a/paddle/cinn/backends/nvrtc/header_generator.cc b/paddle/cinn/backends/nvrtc/header_generator.cc
@@ -14,9 +14,12 @@
 
 #include "paddle/cinn/backends/nvrtc/header_generator.h"
 
+#include <fstream>
 #include "glog/logging.h"
 #include "jitify.hpp"  // NOLINT
+#include "paddle/cinn/common/common.h"
 #include "paddle/common/enforce.h"
+
 namespace cinn {
 namespace backends {
 namespace nvrtc {
@@ -34,12 +37,60 @@ const size_t JitSafeHeaderGenerator::size() const {
   return include_names_.size();
 }
 
+std::string read_file_as_string(const std::string& file_path) {
+#ifdef RUNTIME_INCLUDE_DIR
+  static constexpr char* defined_runtime_include_dir = RUNTIME_INCLUDE_DIR;
+#else
+  static constexpr char* defined_runtime_include_dir = nullptr;
+#endif
+
+#ifdef CINN_WITH_CUDA
+  std::string cinn_path = defined_runtime_include_dir;
+  std::ifstream file(cinn_path + '/' + file_path);
+
+  if (!file.is_open()) {
+    LOG_FIRST_N(INFO, 1) << "Unable to open file : " << cinn_path << '/'
+                         << file_path;
+    return "";
+  }
+  std::stringstream buffer;
+  buffer << file.rdbuf();
+  file.close();
+  return buffer.str();
+#else
+  return "";
+#endif
+}
+#ifdef CINN_WITH_CUDA
+
+static const std::string cinn_float16_header =  // NOLINT
+    read_file_as_string("float16.h");
+static const std::string cinn_bfloat16_header =  // NOLINT
+    read_file_as_string("bfloat16.h");
+static const std::string cinn_with_cuda_header =  // NOLINT
+    R"(
+#pragma once
+#define CINN_WITH_CUDA
+)";
+static const std::string cinn_cuda_runtime_source_header =  // NOLINT
+    read_file_as_string("cinn_cuda_runtime_source.cuh");
+#endif
 JitSafeHeaderGenerator::JitSafeHeaderGenerator() {
   const auto& headers_map = ::jitify::detail::get_jitsafe_headers_map();
   for (auto& pair : headers_map) {
     include_names_.emplace_back(pair.first.data());
     headers_.emplace_back(pair.second.data());
   }
+#ifdef CINN_WITH_CUDA
+  include_names_.emplace_back("float16_h");
+  headers_.emplace_back(cinn_float16_header.data());
+  include_names_.emplace_back("bfloat16_h");
+  headers_.emplace_back(cinn_bfloat16_header.data());
+  include_names_.emplace_back("cinn_with_cuda_h");
+  headers_.emplace_back(cinn_with_cuda_header.data());
+  include_names_.emplace_back("cinn_cuda_runtime_source_h");
+  headers_.emplace_back(cinn_cuda_runtime_source_header.data());
+#endif
 }
 
 }  // namespace nvrtc
diff --git a/paddle/cinn/backends/nvrtc/nvrtc_util.cc b/paddle/cinn/backends/nvrtc/nvrtc_util.cc
@@ -24,6 +24,7 @@
 #include <fstream>
 #include <iostream>
 
+#include "paddle/cinn/backends/codegen_cuda_dev.h"
 #include "paddle/cinn/backends/cuda_util.h"
 #include "paddle/cinn/backends/nvrtc/header_generator.h"
 #include "paddle/cinn/common/common.h"
@@ -181,6 +182,17 @@ std::string Compiler::CompileCudaSource(const std::string& code,
   nvrtcResult compile_res =
       nvrtcCompileProgram(prog, param_cstrings.size(), param_cstrings.data());
 
+  if (compile_res != NVRTC_SUCCESS) {
+    std::string new_code = CodeGenCudaDev::GetGeneralSourceHeader() + code;
+    NVRTC_CALL(nvrtcCreateProgram(&prog,
+                                  new_code.c_str(),
+                                  nullptr,
+                                  header_gen.size(),
+                                  header_gen.headers().data(),
+                                  header_gen.include_names().data()));
+    compile_res =
+        nvrtcCompileProgram(prog, param_cstrings.size(), param_cstrings.data());
+  }
   {  // get log
     size_t log_size;
     NVRTC_CALL(nvrtcGetProgramLogSize(prog, &log_size));
