Rename hlo_module_id to program_id in XLA Op TraceMe

PiperOrigin-RevId: 643032304
tensorflow · Jun 13, 2024 · 4f9d2a7 · 4f9d2a7
1 parent ca319f5
commit 4f9d2a7
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.cc b/third_party/xla/xla/service/cpu/cpu_runtime.cc
@@ -546,12 +546,12 @@ ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int __xla_cpu_runtime_PrintfToStderr(
 
 ABSL_ATTRIBUTE_NO_SANITIZE_MEMORY int64_t __xla_cpu_runtime_TracingStart(
     const void* /* ExecutableRunOptions*  run_options_ptr*/, const char* name,
-    const char* hlo_module, int64_t hlo_module_id) {
+    const char* hlo_module, int64_t program_id) {
   VLOG(3) << "TracingStart " << name;
   auto trace_in =
       tsl::profiler::TraceMeEncode(name, {{"hlo_op", name},
                                           {"hlo_module", hlo_module},
-                                          {"hlo_module_id", hlo_module_id}});
+                                          {"program_id", program_id}});
   return tsl::profiler::TraceMe::ActivityStart(trace_in);
 }
 

diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h
@@ -110,7 +110,7 @@ extern int __xla_cpu_runtime_PrintfToStderr(const char* format, ...);
 
 extern int64_t __xla_cpu_runtime_TracingStart(
     const void* /* xla::ExecutableRunOptions* */ run_options_ptr,
-    const char* name, const char* hlo_module, int64_t hlo_module_id);
+    const char* name, const char* hlo_module, int64_t program_id);
 extern void __xla_cpu_runtime_TracingEnd(
     const void* /* xla::ExecutableRunOptions* */ run_options_ptr, int64_t id);
 

diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -3592,9 +3592,9 @@ void IrEmitter::TracingState::EmitTracingStart(llvm::IRBuilder<>* b,
 
   auto* hlo_name = b->CreateGlobalStringPtr(hlo->name());
   auto* hlo_module = b->CreateGlobalStringPtr(hlo->GetModule()->name());
-  auto* hlo_module_id = b->getInt64(hlo->GetModule()->unique_id());
+  auto* program_id = b->getInt64(hlo->GetModule()->unique_id());
   auto* activity_id = b->CreateCall(
-      trace_func, {run_options, hlo_name, hlo_module, hlo_module_id});
+      trace_func, {run_options, hlo_name, hlo_module, program_id});
   activity_id->setName(IrName(hlo, "activity_id"));
   activity_ids_[hlo] = activity_id;
 }

diff --git a/third_party/xla/xla/service/cpu/runtime/thunk.cc b/third_party/xla/xla/service/cpu/runtime/thunk.cc
@@ -65,7 +65,7 @@ std::string Thunk::TraceMeEncode() const {
   return tsl::profiler::TraceMeEncode(info_.op_name,
                                       {{"hlo_op", info_.op_name},
                                        {"hlo_module", info_.module_name},
-                                       {"hlo_module_id", info_.module_id}});
+                                       {"program_id", info_.module_id}});
 }
 
 std::ostream& operator<<(std::ostream& os, Thunk::Kind kind) {