Add AotCompilationMetadata field to variant of CompileAheadOfTime.

Add CompileAheadOfTime parameter that can optionally be populated during
compilation process.

This change is to allow populating metadata even if the CompileAheadOfTime
fails.

PiperOrigin-RevId: 200407917

1 files changed, 20 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h
index 6c52ffd800..d1144f97bb 100644
--- a/tensorflow/compiler/xla/service/compiler.h
+++ b/tensorflow/compiler/xla/service/compiler.h
@@ -94,6 +94,19 @@ class AotCompilationOptions {
   DebugOptions debug_options_;
 };
 
+// Abstract superclass describing metadata produced during ahead-of-time
+// compilation.
+class AotCompilationMetadata {
+ public:
+  AotCompilationMetadata(const AotCompilationMetadata&) = delete;
+  AotCompilationMetadata& operator=(AotCompilationMetadata const&) = delete;
+
+  virtual ~AotCompilationMetadata() = default;
+
+ protected:
+  AotCompilationMetadata() = default;
+};
+
 // Abstract compiler interface that is subclassed for compilation on a
 // particular platform.
 //
@@ -172,6 +185,13 @@ class Compiler {
   CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules,
                      const AotCompilationOptions& options) = 0;
 
+  // Similar to CompileAheadOfTime above but AotCompilationMetadata
+  // has an argument that can be populated during compilation.
+  virtual StatusOr<std::vector<std::unique_ptr<AotCompilationResult>>>
+  CompileAheadOfTime(std::vector<std::unique_ptr<HloModule>> modules,
+                     const AotCompilationOptions& options,
+                     std::unique_ptr<AotCompilationMetadata>* metadata);
+
   /////
   // The Compiler class also serves as a point to register compiler objects
   // for the various platforms.