Add extra_flags to ptxas for more control over the behavior.

PiperOrigin-RevId: 301317954
Change-Id: Ia7b2b52f108020cee625f5825ecd52de49bdbaaf

tensorflow/stream_executor/gpu/BUILD

@@ -207,7 +207,10 @@ cc_library(
         "//tensorflow/core/kernels:__subpackages__",
         "//tensorflow/stream_executor:__subpackages__",
     ],
-    deps = ["@com_google_absl//absl/strings"],
+    deps = [
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+    ],
 )
 
 cc_library(

tensorflow/stream_executor/gpu/asm_compiler.cc

@@ -212,6 +212,8 @@ port::StatusOr<std::vector<uint8>> CompileGpuAsm(int cc_major, int cc_minor,
   if (options.disable_gpuasm_optimizations) {
     ptxas_args.push_back("-O0");
   }
+  ptxas_args.insert(ptxas_args.end(), options.extra_flags.begin(),
+                    options.extra_flags.end());
   ptxas_info_dumper.SetProgram(ptxas_path, ptxas_args);
   ptxas_info_dumper.SetChannelAction(tensorflow::CHAN_STDERR,
                                      tensorflow::ACTION_PIPE);

tensorflow/stream_executor/gpu/gpu_asm_opts.h

@@ -18,8 +18,10 @@ limitations under the License.
 
 #include <string>
 #include <tuple>
+#include <vector>
 
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 
 namespace stream_executor {
 // Compilation options for compiling ptxas.
@@ -30,15 +32,21 @@ struct GpuAsmOpts {
   // Cuda directory which would be searched first.
   std::string preferred_cuda_dir;
 
-  explicit GpuAsmOpts(bool disable_gpuasm_optimizations = false,
-                      absl::string_view preferred_cuda_dir = "")
-      : disable_gpuasm_optimizations(disable_gpuasm_optimizations),
-        preferred_cuda_dir(preferred_cuda_dir) {}
+  std::vector<std::string> extra_flags;
 
-  using PtxOptionsTuple = std::tuple<bool, std::string>;
+  explicit GpuAsmOpts(bool disable_gpuasm_optimizations = false,
+                      absl::string_view preferred_cuda_dir = "",
+                      absl::Span<const std::string> extra_flags = {})
+      : disable_gpuasm_optimizations(disable_gpuasm_optimizations),
+        preferred_cuda_dir(preferred_cuda_dir),
+        extra_flags(extra_flags.begin(), extra_flags.end()) {}
+
+  using PtxOptionsTuple =
+      std::tuple<bool, std::string, std::vector<std::string>>;
 
   PtxOptionsTuple ToTuple() {
-    return std::make_tuple(disable_gpuasm_optimizations, preferred_cuda_dir);
+    return std::make_tuple(disable_gpuasm_optimizations, preferred_cuda_dir,
+                           extra_flags);
   }
 };
 }  // namespace stream_executor