Add Flag to Control NVRTC Cubin Fmad Option (#58794)

zhhsplendid · web-flow · commit 6aafcf27fca9 · 2023-11-08T19:06:24.000+08:00
Add flag to control NVRTC Cubin Fmad option so that we can test precision correctly.
diff --git a/paddle/cinn/backends/nvrtc/nvrtc_util.cc b/paddle/cinn/backends/nvrtc/nvrtc_util.cc
@@ -32,6 +32,7 @@
 
 PD_DECLARE_string(cinn_nvcc_cmd_path);
 PD_DECLARE_bool(nvrtc_compile_to_cubin);
+PD_DECLARE_bool(cinn_nvrtc_cubin_with_fmad);
 
 namespace cinn {
 namespace backends {
@@ -106,6 +107,9 @@ std::string Compiler::CompileCudaSource(const std::string& code,
   }
   if (compile_to_cubin_) {
     compile_options.push_back("-arch=sm_" + cc);
+    std::string enable_fmad =
+        FLAGS_cinn_nvrtc_cubin_with_fmad ? "true" : "false";
+    compile_options.push_back("--fmad=" + enable_fmad);
   } else {
     compile_options.push_back("-arch=compute_" + cc);
   }
diff --git a/paddle/cinn/runtime/flags.cc b/paddle/cinn/runtime/flags.cc
@@ -112,6 +112,14 @@ PD_DEFINE_bool(cinn_compile_with_nvrtc,
                BoolFromEnv("FLAGS_cinn_compile_with_nvrtc", true),
                "Whether nvrtc compile cuda source with nvrtc(default nvcc).");
 
+PD_DEFINE_bool(
+    cinn_nvrtc_cubin_with_fmad,
+    BoolFromEnv("FLAGS_cinn_nvrtc_cubin_with_fmad", true),
+    "Whether nvrtc enables fmad when compile to cubin. This flag only works "
+    "when FLAGS_nvrtc_compile_to_cubin=true. Fmad is the cuda speed up "
+    "technique which contract fp mulitplication and addition/subtraction into "
+    "multiply-add operation. It may result in different fp precision.");
+
 // FLAGS for performance analysis and accuracy debug
 PD_DEFINE_bool(cinn_sync_run,
                BoolFromEnv("FLAGS_cinn_sync_run", false),
