Add peak memory bw to XPlane to enable Roofline Analysis and Hlo Stats

PiperOrigin-RevId: 695949288
openxla · Nov 13, 2024 · a0b99e5 · a0b99e5
1 parent 3750ac9
commit a0b99e5
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/xla/tsl/profiler/utils/xplane_schema.cc b/xla/tsl/profiler/utils/xplane_schema.cc
@@ -301,6 +301,14 @@ const StatTypeMap& GetStatTypeMap() {
        {"compute_cap_minor", kDevCapComputeCapMinor},
        {"peak_teraflops_per_second", kDevCapPeakTeraflopsPerSecond},
        {"peak_hbm_bw_gigabytes_per_second", kDevCapPeakHbmBwGigabytesPerSecond},
+       {"peak_cmem_rd_bw_gigabytes_per_second",
+        kDevCapPeakCmemRdBwGigabytesPerSecond},
+       {"peak_cmem_wr_bw_gigabytes_per_second",
+        kDevCapPeakCmemWrBwGigabytesPerSecond},
+       {"peak_vmem_rd_bw_gigabytes_per_second",
+        kDevCapPeakVmemRdBwGigabytesPerSecond},
+       {"peak_vmem_wr_bw_gigabytes_per_second",
+        kDevCapPeakVmemWrBwGigabytesPerSecond},
        {"peak_sram_rd_bw_gigabytes_per_second",
         kDevCapPeakSramRdBwGigabytesPerSecond},
        {"peak_sram_wr_bw_gigabytes_per_second",

diff --git a/xla/tsl/profiler/utils/xplane_schema.h b/xla/tsl/profiler/utils/xplane_schema.h
@@ -287,6 +287,10 @@ enum StatType {
   kDevCapComputeCapMinor,
   kDevCapPeakTeraflopsPerSecond,
   kDevCapPeakHbmBwGigabytesPerSecond,
+  kDevCapPeakCmemRdBwGigabytesPerSecond,
+  kDevCapPeakCmemWrBwGigabytesPerSecond,
+  kDevCapPeakVmemRdBwGigabytesPerSecond,
+  kDevCapPeakVmemWrBwGigabytesPerSecond,
   kDevCapPeakSramRdBwGigabytesPerSecond,
   kDevCapPeakSramWrBwGigabytesPerSecond,
   kDevVendor,