From 2ad9dd652f072266ba9ece977119a625536fb5e5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 8 Jan 2020 10:05:44 -0800
Subject: [PATCH] Add more StatType.

PiperOrigin-RevId: 288720224
Change-Id: Ieeed4dfcacd5ed435068aa17808d29f3300470bb
---
 tensorflow/core/profiler/utils/xplane_schema.cc | 12 +++++++-----
 tensorflow/core/profiler/utils/xplane_schema.h  |  7 ++++++-
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index 4af32c76457..a816add48bd 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -68,11 +68,13 @@ static const absl::string_view kStatTypeStrMap[] = {
     "device_ordinal",  "chip_ordinal",
     "node_ordinal",    "model_id",
     "queue_addr",      "request_id",
-    "run_id",          "correlation_id",
-    "graph_type",      "step_num",
-    "iter_num",        "index_on_host",
-    "bytes_reserved",  "bytes_allocated",
-    "bytes_available", "fragmentation",
+    "run_id",          "graph_type",
+    "step_num",        "iter_num",
+    "index_on_host",   "bytes_reserved",
+    "bytes_allocated", "bytes_available",
+    "fragmentation",   "device_id",
+    "context_id",      "correlation_id",
+    "memcpy_details",  "memalloc_details",
     "kernel_details",  "group_id",
     "step_name",       "level 0",
     "tf_op",           "hlo_op",
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 4216450d653..35c874a796e 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -74,7 +74,6 @@ enum StatType {
   kQueueAddr,
   kRequestId,
   kRunId,
-  kCorrelationId,
   kGraphType,
   kStepNum,
   kIterNum,
@@ -83,6 +82,12 @@ enum StatType {
   kBytesAllocated,
   kBytesAvailable,
   kFragmentation,
+  // Device trace arguments.
+  kDeviceId,
+  kContextId,
+  kCorrelationId,
+  kMemcpyDetails,
+  kMemallocDetails,
   kKernelDetails,
   // Stats added when processing traces.
   kGroupId,