Allow cost estimates to differ per backend and include the estimates into the HLO profile. Add a summary table for what categories have the most opportunity for optimization left in them.

PiperOrigin-RevId: 163780413

1 files changed, 7 insertions, 2 deletions

diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
index 7f9e60460c..a195d0b8e4 100644
--- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
+++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc
@@ -112,10 +112,11 @@ GpuExecutable::GpuExecutable(
     std::unique_ptr<HloModule> hlo_module,
     std::unique_ptr<BufferAssignment> assignment,
     HloCostAnalysis::ShapeSizeFunction shape_size_function)
-    : Executable(std::move(hlo_module), std::move(shape_size_function)),
+    : Executable(std::move(hlo_module)),
       ptx_(ptx),
       thunk_schedule_(std::move(thunk_schedule)),
-      assignment_(std::move(assignment)) {}
+      assignment_(std::move(assignment)),
+      shape_size_function_(std::move(shape_size_function)) {}
 
 Status GpuExecutable::ExecuteThunks(
     const ServiceExecutableRunOptions* run_options,
@@ -356,5 +357,9 @@ const PointsToSet& GpuExecutable::GetRootPointsToSet() const {
       module().entry_computation()->root_instruction());
 }
 
+std::unique_ptr<HloCostAnalysis> GpuExecutable::CreateCostAnalysis() const {
+  return MakeUnique<HloCostAnalysis>(shape_size_function_);
+}
+
 }  // namespace gpu
 }  // namespace xla