From 5e290dc906a3cef29356792f356d7a70274daf25 Mon Sep 17 00:00:00 2001
From: Yin Zhang <yinzz@google.com>
Date: Mon, 16 Dec 2024 20:13:37 -0800
Subject: [PATCH] Create raw_bytes_accessed Stat in XPlane Schema

PiperOrigin-RevId: 706927552
---
 xla/tsl/profiler/utils/xplane_schema.cc | 1 +
 xla/tsl/profiler/utils/xplane_schema.h  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/xla/tsl/profiler/utils/xplane_schema.cc b/xla/tsl/profiler/utils/xplane_schema.cc
index 615d7ace55147..f7320fafdac04 100644
--- a/xla/tsl/profiler/utils/xplane_schema.cc
+++ b/xla/tsl/profiler/utils/xplane_schema.cc
@@ -276,6 +276,7 @@ const StatTypeMap& GetStatTypeMap() {
        {"flops", kFlops},
        {"model_flops", kModelFlops},
        {"bytes_accessed", kBytesAccessed},
+       {"raw_bytes_accessed", kRawBytesAccessed},
        {"memory_access_breakdown", kMemoryAccessBreakdown},
        {"shape_with_layout", kShapeWithLayout},
        {"source", kSourceInfo},
diff --git a/xla/tsl/profiler/utils/xplane_schema.h b/xla/tsl/profiler/utils/xplane_schema.h
index 6cdd81c8342ca..580cfd06adc09 100644
--- a/xla/tsl/profiler/utils/xplane_schema.h
+++ b/xla/tsl/profiler/utils/xplane_schema.h
@@ -264,6 +264,7 @@ enum StatType {
   kFlops,
   kModelFlops,
   kBytesAccessed,
+  kRawBytesAccessed,
   kMemoryAccessBreakdown,
   kShapeWithLayout,
   kSourceInfo,