[XLA:CPU] Re-use the same llvm::GlobalVariable for identical literals

This isn't necessary today, but it will be after an optimization change I'm
about to make.

LLVM has a constant merging pass too, but one of the motivations here is to
avoid the LLVM compile time overhead of having many large arrays in the IR.

PiperOrigin-RevId: 195032900

1 files changed, 16 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
index 0f2f3d1817..5a04076080 100644
--- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h
+++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h
@@ -530,6 +530,8 @@ class IrEmitter : public DfsHloVisitorWithDefault {
   Status EmitXfeedTransfer(XfeedKind kind, const Shape& shape,
                            llvm::Value* program_buffer_address);
 
+  llvm::GlobalVariable* EmitGlobalForLiteral(const Literal& literal);
+
   const HloModuleConfig& hlo_module_config_;
 
   bool is_top_level_computation_;
@@ -539,6 +541,20 @@ class IrEmitter : public DfsHloVisitorWithDefault {
   int64 external_global_constant_counter_ = 0;
   ExternalConstantPool* external_constant_pool_;
 
+  struct LiteralPtrHashFunctor {
+    size_t operator()(const Literal* literal) const { return literal->Hash(); }
+  };
+
+  struct LiteralPtrEqualityFunctor {
+    bool operator()(const Literal* lhs, const Literal* rhs) const {
+      return *lhs == *rhs;
+    }
+  };
+
+  tensorflow::gtl::FlatMap<const Literal*, llvm::GlobalVariable*,
+                           LiteralPtrHashFunctor, LiteralPtrEqualityFunctor>
+      emitted_literals_;
+
   TF_DISALLOW_COPY_AND_ASSIGN(IrEmitter);
 };