Dump HLO HBM usage info

pytorch · May 20, 2024 · b4d83e4 · b4d83e4
1 parent e0fb878
commit b4d83e4
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc
@@ -591,6 +591,20 @@ std::vector<ComputationClient::ComputationPtr> PjRtComputationClient::Compile(
           client_->Compile(instance.computation, compile_options).value();
     }
 
+    auto memory_stats_status_or = executable->GetCompiledMemoryStats();
+    if (memory_stats_status_or.ok()) {
+      xla::CompiledMemoryStats memory_stats = memory_stats_status_or.value();
+      TF_VLOG(3) << "memory usage detail = " << memory_stats.DebugString();
+      TF_VLOG(3)
+          << "total runtime device memory required to run this program = "
+          << ((memory_stats.output_size_in_bytes +
+               memory_stats.temp_size_in_bytes) >>
+              20)
+          << " MB";
+    } else {
+      TF_VLOG(3) << "memory usage is not availiable";
+    }
+
     const auto& hlo_modules = ConsumeValue(executable->GetHloModules());
     xla::HloComputation* hlo_computation = hlo_modules[0]->entry_computation();
     std::shared_ptr<PjRtComputation> pjrt_computation =