Replace gpu_asm_extra_flags string option by individual flags

`gpu_asm_extra_flags` allowed the user to pass arbitrary command line flags
to `ptxas`. This was working well enough when calling ptxas was our only
PTX compilation option.

With library compilation and support for parallel compilation a wrong flag can easily mess up the PTX compilation with weird error. It also doesn't support flags for `nvlink` which might be needed when using parallel compilation.

So I'm replacing the list of opaque flags by individual boolean options for the two main use cases (debug compile and preserving line info). The new compilation providers will take those booleans options and generate the right flags or API calls.

As a temporary shim the function `PtxOptsFromDebugOptions` will still generate the command line flags for users of the legacy PTX compilation functions.

PiperOrigin-RevId: 701850432

Author: beckerhe

xla/debug_options_flags.cc

@@ -58,7 +58,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_autotune_max_solutions(0);
   opts.set_xla_cpu_multi_thread_eigen(true);
   opts.set_xla_gpu_cuda_data_dir("./cuda_sdk_lib");
-  opts.set_xla_gpu_asm_extra_flags("");
+  opts.set_xla_gpu_generate_debug_info(false);
+  opts.set_xla_gpu_generate_line_info(false);
 
   // As of cudnn 8.9.0, runtime fusion creates convolutions that take about 7s
   // seconds to run the first time we call them, at least on Ampere.  In
@@ -968,11 +969,16 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       bool_setter_for(&DebugOptions::set_xla_gpu_disable_gpuasm_optimizations),
       debug_options->xla_gpu_disable_gpuasm_optimizations(),
       "In XLA:GPU run ptxas in -O0 (default is -O3)."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_asm_extra_flags",
-      string_setter_for(&DebugOptions::set_xla_gpu_asm_extra_flags), "",
-      "Pass extra parameters to the GPU assembler tool (i.e., ptxas for CUDA). "
-      "If multiple parameters, separate them by comma."));
+  flag_list->push_back(
+      tsl::Flag("xla_gpu_generate_debug_info",
+                bool_setter_for(&DebugOptions::set_xla_gpu_generate_debug_info),
+                debug_options->xla_gpu_generate_debug_info(),
+                "Generate debug info for codegened CUDA kernels."));
+  flag_list->push_back(
+      tsl::Flag("xla_gpu_generate_line_info",
+                bool_setter_for(&DebugOptions::set_xla_gpu_generate_line_info),
+                debug_options->xla_gpu_generate_line_info(),
+                "Generate line info for codegened CUDA kernels."));
   flag_list->push_back(tsl::Flag(
       "xla_fuel", setter_for_xla_fuel, /*default_value_for_display=*/"",
       "Sets compiler fuel, useful for bisecting bugs in passes. Format "

xla/service/gpu/BUILD

@@ -2406,11 +2406,21 @@ cc_library(
     srcs = ["gpu_asm_opts_util.cc"],
     hdrs = ["gpu_asm_opts_util.h"],
     compatible_with = get_compatible_with_portable(),
-    copts = tsl_copts(),
     deps = [
         "//xla:xla_proto_cc",
         "//xla/stream_executor/gpu:gpu_asm_opts",
-        "@com_google_absl//absl/strings",
+    ],
+)
+
+xla_cc_test(
+    name = "gpu_asm_opts_util_test",
+    srcs = ["gpu_asm_opts_util_test.cc"],
+    deps = [
+        ":gpu_asm_opts_util",
+        "//xla:xla_proto_cc",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:test",
     ],
 )
 

xla/service/gpu/gpu_asm_opts_util.cc

@@ -16,9 +16,9 @@ limitations under the License.
 #include "xla/service/gpu/gpu_asm_opts_util.h"
 
 #include <string>
+#include <utility>
 #include <vector>
 
-#include "absl/strings/str_split.h"
 #include "xla/stream_executor/gpu/gpu_asm_opts.h"
 #include "xla/xla.pb.h"
 
@@ -27,11 +27,16 @@ namespace gpu {
 
 stream_executor::GpuAsmOpts PtxOptsFromDebugOptions(
     const DebugOptions& debug_options) {
-  std::vector<std::string> extra_flags = absl::StrSplit(
-      debug_options.xla_gpu_asm_extra_flags(), ',', absl::SkipEmpty());
+  std::vector<std::string> extra_flags;
+  if (debug_options.xla_gpu_generate_line_info()) {
+    extra_flags.emplace_back("--generate-line-info");
+  }
+  if (debug_options.xla_gpu_generate_debug_info()) {
+    extra_flags.emplace_back("--device-debug");
+  }
   return stream_executor::GpuAsmOpts(
       debug_options.xla_gpu_disable_gpuasm_optimizations(),
-      debug_options.xla_gpu_cuda_data_dir(), extra_flags);
+      debug_options.xla_gpu_cuda_data_dir(), std::move(extra_flags));
 }
 
 }  // namespace gpu

xla/service/gpu/gpu_asm_opts_util_test.cc

@@ -0,0 +1,44 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "xla/service/gpu/gpu_asm_opts_util.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/xla.pb.h"
+#include "tsl/platform/test.h"
+
+namespace xla::gpu {
+namespace {
+using ::testing::Contains;
+
+TEST(PtxOptsFromDebugOptionsTest, GenerateLineInfo) {
+  xla::DebugOptions debug_options;
+  debug_options.set_xla_gpu_generate_line_info(true);
+
+  EXPECT_THAT(PtxOptsFromDebugOptions(debug_options).extra_flags,
+              Contains("--generate-line-info"));
+}
+
+TEST(PtxOptsFromDebugOptionsTest, GenerateDebugInfo) {
+  xla::DebugOptions debug_options;
+  debug_options.set_xla_gpu_generate_debug_info(true);
+
+  EXPECT_THAT(PtxOptsFromDebugOptions(debug_options).extra_flags,
+              Contains("--device-debug"));
+}
+
+}  // namespace
+}  // namespace xla::gpu

xla/xla.proto

@@ -314,6 +314,14 @@ message DebugOptions {
   // If set to true XLA:GPU invokes `ptxas` with -O0 (default is -O3).
   bool xla_gpu_disable_gpuasm_optimizations = 103;
 
+  // If true, we generate debug info when compiling PTX. This is useful for
+  // profiling and debugging.
+  bool xla_gpu_generate_debug_info = 348;
+
+  // If true, we generate line info when compiling PTX. This is useful for
+  // profiling and debugging.
+  bool xla_gpu_generate_line_info = 349;
+
   enum ShapeChecks {
     // Do not insert any shape checks for dynamically shaped operations; output
     // buffers might contain garbage data if shapes don't match.
@@ -455,8 +463,7 @@ message DebugOptions {
   // memory, or have bugs.
   bool xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found = 138;
 
-  // Extra parameters to pass the GPU assembler.
-  string xla_gpu_asm_extra_flags = 141;
+  reserved 141;  // was xla_gpu_asm_extra_flags
 
   // Per-heap size constraint. New heaps will be created if per-heap max size is
   // reached.
@@ -1062,7 +1069,7 @@ message DebugOptions {
   // be deterministic, although with additional overhead.
   bool xla_gpu_enable_scatter_determinism_expander = 345;
 
-  // Next id: 348
+  // Next id: 350
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.