Extract register spill information from ptxas log

xla/stream_executor/BUILD

@@ -962,6 +962,12 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "kernel_stats",
+    hdrs = ["kernel_stats.h"],
+    deps = ["@com_google_absl//absl/container:flat_hash_map"],
+)
+
 xla_cc_test(
     name = "kernel_args_test",
     srcs = ["kernel_args_test.cc"],

xla/stream_executor/cuda/BUILD

@@ -798,13 +798,15 @@ cc_library(
     hdrs = ["ptx_compiler_helpers.h"],
     deps = [
         ":cuda_compute_capability",
+        "//xla/stream_executor:kernel_stats",
         "//xla/stream_executor:semantic_version",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_googlesource_code_re2//:re2",
     ],
 )
 
@@ -813,6 +815,7 @@ xla_cc_test(
     srcs = ["ptx_compiler_helpers_test.cc"],
     deps = [
         ":ptx_compiler_helpers",
+        "//xla/stream_executor:kernel_stats",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",

xla/stream_executor/cuda/ptx_compiler_helpers.cc

@@ -27,7 +27,9 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "re2/re2.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/kernel_stats.h"
 #include "xla/stream_executor/semantic_version.h"
 
 namespace stream_executor {
@@ -91,6 +93,26 @@ absl::StatusOr<int> GetLatestPtxIsaVersionFromUnsupportedVersionErrorLog(
   return major * 10 + minor;
 }
 
+ModuleStats ExtractModuleStatsFromLog(absl::string_view log) {
+  // Reads the log for registers spilled and adds up the count. An example of
+  // this line is:
+  // "Registers are spilled to local memory in function 'rr', 1080 bytes spill
+  // stores, 968 bytes spill loads"
+  static constexpr LazyRE2 kSpillRegex{
+      R"(function\s+'(\w+)',\s*(\d+)\s+bytes\s+spill\s+stores,\s*(\d+)\s+bytes\s+spill\s+loads)"};
+  ModuleStats kernel_stats_map;
+
+  int spill_stores = 0;
+  int spill_loads = 0;
+  absl::string_view function_name;
+  absl::string_view search_log = log;
+  while (RE2::FindAndConsume(&search_log, *kSpillRegex, &function_name,
+                             &spill_stores, &spill_loads)) {
+    kernel_stats_map[function_name] = {spill_stores, spill_loads};
+  }
+  return kernel_stats_map;
+}
+
 absl::Status CreateErrorFromPTXASLog(absl::string_view log,
                                      absl::string_view architecture,
                                      bool cancel_if_reg_spill) {

xla/stream_executor/cuda/ptx_compiler_helpers.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/kernel_stats.h"
 #include "xla/stream_executor/semantic_version.h"
 
 namespace stream_executor {
@@ -56,6 +57,13 @@ void WarnIfBadPtxasVersion(absl::string_view method,
 absl::StatusOr<int> GetLatestPtxIsaVersionFromUnsupportedVersionErrorLog(
     absl::string_view error_log);
 
+// Extracts the module stats from the ptxas log.
+//
+// Example: "Registers are spilled to local memory in function 'rr', 1080 bytes
+// spill stores, 968 bytes spill loads" will return:
+// ModuleStats{ KernelStats{"rr", {1080, 968}} }
+ModuleStats ExtractModuleStatsFromLog(absl::string_view log);
+
 }  // namespace stream_executor
 
 #endif  // XLA_STREAM_EXECUTOR_CUDA_PTX_COMPILER_HELPERS_H_

xla/stream_executor/cuda/ptx_compiler_helpers_test.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
+#include "xla/stream_executor/kernel_stats.h"
 
 namespace stream_executor {
 namespace {
@@ -49,7 +50,15 @@ ptxas fatal   : Ptx assembly aborted due to errors
 
 constexpr absl::string_view kPtxasLogRegisterSpillWarning = R"(
 // Something in the log before the warning.
-ptxas warning : Registers are spilled to local memory in function '__kernel', 8 bytes spill stores, 8 bytes spill loads
+ptxas warning : Registers are spilled to local memory in function '__kernel', 18 bytes spill stores, 8 bytes spill loads
+// Something in the log after the warning.
+)";
+
+constexpr absl::string_view kPtxasLogMultipleKernelsSpillingRegisters = R"(
+// Something in the log before the warning.
+ptxas warning : Registers are spilled to local memory in function '__kernel', 18 bytes spill stores, 8 bytes spill loads
+// Log log log.
+ptxas warning : Registers are spilled to local memory in function '__kernel2', 1024 bytes spill stores, 1099 bytes spill loads
 // Something in the log after the warning.
 )";
 
@@ -106,5 +115,22 @@ TEST(PtxCompilerHelpersTest, IsPtxRegisterAllocationErrorStatus) {
   EXPECT_FALSE(IsPtxRegisterAllocationError(absl::OkStatus()));
 }
 
+TEST(PtxCompilerHelpersTest, ModuleStatsAreCorrectlyExtractedFromLog) {
+  ModuleStats kernel_stats_map =
+      ExtractModuleStatsFromLog(kPtxasLogRegisterSpillWarning);
+  EXPECT_EQ(kernel_stats_map.size(), 1);
+  EXPECT_EQ(kernel_stats_map["__kernel"].store_bytes_spilled, 18);
+  EXPECT_EQ(kernel_stats_map["__kernel"].load_bytes_spilled, 8);
+  kernel_stats_map =
+      ExtractModuleStatsFromLog(kPtxasLogMultipleKernelsSpillingRegisters);
+  EXPECT_EQ(kernel_stats_map.size(), 2);
+  EXPECT_EQ(kernel_stats_map["__kernel"].store_bytes_spilled, 18);
+  EXPECT_EQ(kernel_stats_map["__kernel"].load_bytes_spilled, 8);
+  EXPECT_EQ(kernel_stats_map["__kernel2"].store_bytes_spilled, 1024);
+  EXPECT_EQ(kernel_stats_map["__kernel2"].load_bytes_spilled, 1099);
+  kernel_stats_map = ExtractModuleStatsFromLog(kPtxasLogSuccessfulCompilation);
+  EXPECT_EQ(kernel_stats_map.size(), 0);
+}
+
 }  // namespace
 }  // namespace stream_executor

xla/stream_executor/kernel_stats.h

@@ -0,0 +1,34 @@
+/* Copyright 2025 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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_KERNEL_STATS_H_
+#define XLA_STREAM_EXECUTOR_KERNEL_STATS_H_
+
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+
+// Stats about a single kernel.
+struct KernelStats {
+  // The number of spilled register bytes in the kernel for stores.
+  int store_bytes_spilled = 0;
+  // The number of spilled register bytes in the kernel for loads.
+  int load_bytes_spilled = 0;
+};
+
+// Map from a function/kernel name to its kernel stats.
+using ModuleStats = absl::flat_hash_map<std::string, KernelStats>;
+
+#endif  // XLA_STREAM_EXECUTOR_KERNEL_STATS_H_