[xla:cpu][roll forward] Improve compilation time by not fusing large constants into LLVM modules

Fix for the breaking of a large model without thunks.
Add tests to make sure this doesn't happen again.

Reverts 1971267

PiperOrigin-RevId: 707524044

Author: Google-ML-Automation

xla/service/cpu/BUILD

@@ -456,6 +456,7 @@ xla_test(
     ],
     tags = [
         "test_migrated_to_hlo_runner_pjrt",
+        "test_xla_cpu_no_thunks",
     ],
     deps = [
         "//xla/hlo/testlib:verified_hlo_module",
@@ -686,18 +687,39 @@ xla_cc_test(
     name = "ir_emitter_test",
     srcs = ["ir_emitter_test.cc"],
     deps = [
+        ":cpu_compiler",
+        ":cpu_executable",
+        ":cpu_options",
         ":ir_emitter",
         ":ir_function",
+        ":runtime_symbol_generator",
         ":target_machine_features_stub",
+        "//xla:cpu_function_runtime",
+        "//xla/backends/cpu/codegen:cpu_features",
+        "//xla/backends/cpu/codegen:ir_compiler",
+        "//xla/backends/cpu/codegen:jit_compiler",
+        "//xla/backends/cpu/codegen:target_machine_features",
         "//xla/hlo/analysis:hlo_ordering",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
+        "//xla/hlo/transforms:hlo_memory_scheduler",
         "//xla/service:buffer_assignment",
+        "//xla/service:buffer_value",
         "//xla/service:hlo_module_config",
         "//xla/service:logical_buffer",
+        "//xla/service/llvm_ir:llvm_util",
         "//xla/tests:hlo_test_base",
+        "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
         "@llvm-project//llvm:Core",
         "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:Target",
+        "@llvm-project//mlir:IR",
+        "@tsl//tsl/platform:env",
+        "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:statusor",
         "@tsl//tsl/platform:test",
         "@tsl//tsl/platform:test_main",
@@ -742,6 +764,7 @@ cc_library(
     copts = tsl_copts(),
     deps = [
         ":backend_config_proto_cc",
+        ":cpu_instruction_fusion",
         ":cpu_options",
         ":cpu_runtime",
         ":dot_op_emitter",

xla/service/cpu/cpu_compiler.cc

@@ -1498,17 +1498,15 @@ CpuCompiler::CompileLegacyCpuExecutable(std::unique_ptr<HloModule> module) {
 #endif
   );
 
-  // Emit global variables for constants.
-  //
-  // TODO(ezhulenev): Figure out how to emit constants that are only needed for
-  // thread local computations as with Thunks runtime we keep constants outside
-  // of the LLVM module. Currently we end up doubling memory for constants.
-  TF_RETURN_IF_ERROR(nested_ir_emitter.EmitConstantGlobals());
 
   // If we use Thunk runtime then instead of emitting LLVM function for the
   // entry computation we emit a sequence of thunks that implement the
   // computation as a sequence of interpreted commands.
   if (module->config().debug_options().xla_cpu_use_thunk_runtime()) {
+    // The thunk runtime manages large constants, therefore we only emit
+    // small ones.
+    TF_RETURN_IF_ERROR(nested_ir_emitter.EmitSmallConstantGlobals());
+
     // IR emitter is responsible for building LLVM module with host kernels for
     // corresponding HLO instructions (fusions, elemental instructions, etc.).
     IrEmitter2 ir_emitter2(*module, llvm_module.get(), &nested_ir_emitter);
@@ -1642,6 +1640,8 @@ CpuCompiler::CompileLegacyCpuExecutable(std::unique_ptr<HloModule> module) {
     return with_hlo_proto(std::move(cpu_executable));
   }
 
+  TF_RETURN_IF_ERROR(nested_ir_emitter.EmitAllConstantGlobals());
+
   // Each computation is a single function.  Emit all embedded computations
   // before the entry computation. The order of computations returned from
   // SubcomputationEmissionOrder guarantees that a called computation occurs
@@ -1899,7 +1899,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr<HloModuleGroup> module_group,
           // TODO(b/66051036): Run full msan for AOT.
           /*emit_code_for_msan=*/false);
 
-      TF_RETURN_IF_ERROR(ir_emitter.EmitConstantGlobals());
+      TF_RETURN_IF_ERROR(ir_emitter.EmitAllConstantGlobals());
 
       for (ComputationToEmit subcomputation :
            SubcomputationEmissionOrder(computation)) {

xla/service/cpu/cpu_compiler_test.cc

@@ -10,6 +10,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
+
 #include <memory>
 #include <string>
 #include <utility>
@@ -55,6 +56,25 @@ TEST_F(CpuCompilerTest, RecordsStreamzStackTrace) {
   EXPECT_GT(it->second->points.size(), 0);
 }
 
+TEST_F(CpuCompilerTest, CompilationWithLargeConstants) {
+  absl::string_view module_string = R"(
+HloModule module
+
+ENTRY main {
+  a = f32[1000,1000]{1,0} parameter(0)
+  b = f32[1000,1000]{1,0} constant({...})
+  a_plus_b = f32[1000,1000]{1,0} add(a, b)
+  c = f32[1000,1000]{1,0} constant({...})
+  ROOT result = f32[1000,1000]{1,0} add(a_plus_b, c)
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                          ParseAndReturnVerifiedModule(module_string));
+
+  EXPECT_TRUE(Run(std::move(module), /*run_hlo_passes=*/true));
+}
+
 }  // namespace
 }  // namespace cpu
 }  // namespace xla

xla/service/cpu/cpu_instruction_fusion.cc

@@ -19,6 +19,9 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/log/log.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/fusion_node_indexing_evaluation.h"
 #include "xla/service/instruction_fusion.h"
@@ -81,6 +84,10 @@ FusionDecision CpuInstructionFusion::ShouldFuse(HloInstruction* consumer,
 
   constexpr int kFusionThresholdBytes = 16 * 1024;
 
+  if (IsLargeConstant(producer)) {
+    return FusionDecision::Forbid("Don't fuse large constants.");
+  }
+
   if (CanBeOutputFused(producer, consumer)) {
     VLOG(2) << "Fusion OK: Can create output fusion.";
     return FusionDecision::Allow();
@@ -219,5 +226,12 @@ HloInstruction* CpuInstructionFusion::FuseInstruction(
   evaluation->second.UpdateEvaluationCache(new_producer, indexing_users);
   return new_producer;
 }
+
+bool CpuInstructionFusion::IsLargeConstant(
+    const HloInstruction* constant) const {
+  return constant->IsConstant() &&
+         Cast<HloConstantInstruction>(constant)->literal().size_bytes() >
+             GetLargeConstantThresholdBytes();
+}
 }  // namespace cpu
 }  // namespace xla

xla/service/cpu/cpu_instruction_fusion.h

@@ -43,6 +43,12 @@ class CpuInstructionFusion : public InstructionFusion {
     return InstructionFusion::Run(module, execution_threads);
   }
 
+  // Returns the threshold for a constant to be considered a large constant.
+  static constexpr int64_t GetLargeConstantThresholdBytes() {
+    constexpr int64_t kLargeConstantThresholdBytes = 10000;
+    return kLargeConstantThresholdBytes;
+  }
+
  protected:
   FusionDecision ShouldFuse(HloInstruction* consumer,
                             int64_t operand_index) override;
@@ -53,6 +59,9 @@ class CpuInstructionFusion : public InstructionFusion {
   HloInstruction* FuseInstruction(HloInstruction* fusion_instruction,
                                   HloInstruction* producer) override;
 
+  // Returns if a constant is large enough to be considered a large constant.
+  bool IsLargeConstant(const HloInstruction* constant) const;
+
   // Keep track of the number of times each instruction inside a fusion node is
   // indexed with different index vectors.
   absl::flat_hash_map<const HloInstruction*, FusionNodeIndexingEvaluation>

xla/service/cpu/cpu_instruction_fusion_test.cc

@@ -935,5 +935,45 @@ ENTRY main {
   EXPECT_THAT(module->entry_computation()->root_instruction(), op::Fusion());
 }
 
+TEST_F(OpcodeFusionTest, BigConstantNotInFusion) {
+  absl::string_view module_string = R"(
+HloModule module
+
+ENTRY main {
+  a = f32[1000,1000]{1,0} parameter(0)
+  b = f32[1000,1000]{1,0} constant({...})
+  a_plus_b = f32[1000,1000]{1,0} add(a, b)
+  c = f32[1000,1000]{1,0} constant({...})
+  ROOT result = f32[1000,1000]{1,0} add(a_plus_b, c)
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(module_string));
+  RunFusionAndCheckOpcodesWereFused(
+      module.get(), {HloOpcode::kParameter, HloOpcode::kParameter,
+                     HloOpcode::kParameter, HloOpcode::kAdd, HloOpcode::kAdd});
+}
+
+TEST_F(OpcodeFusionTest, SmallConstantInFusion) {
+  absl::string_view module_string = R"(
+HloModule module
+
+ENTRY main {
+  a = f32[10,10]{1,0} parameter(0)
+  b = f32[10,10]{1,0} constant({...})
+  a_plus_b = f32[10,10]{1,0} add(a, b)
+  c = f32[10,10]{1,0} constant({...})
+  ROOT result = f32[10,10]{1,0} add(a_plus_b, c)
+}
+)";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(module_string));
+  RunFusionAndCheckOpcodesWereFused(
+      module.get(), {HloOpcode::kParameter, HloOpcode::kConstant,
+                     HloOpcode::kConstant, HloOpcode::kAdd, HloOpcode::kAdd});
+}
+
 }  // namespace
 }  // namespace xla::cpu

xla/service/cpu/ir_emitter.cc

@@ -67,6 +67,7 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/collective_ops_utils.h"
 #include "xla/service/cpu/backend_config.pb.h"
+#include "xla/service/cpu/cpu_instruction_fusion.h"
 #include "xla/service/cpu/cpu_options.h"
 #include "xla/service/cpu/cpu_runtime.h"
 #include "xla/service/cpu/dot_op_emitter.h"
@@ -330,9 +331,24 @@ llvm::Constant* IrEmitter::EmitGlobalForLiteral(const Literal& literal) {
   return result_global;
 }
 
-absl::Status IrEmitter::EmitConstantGlobals() {
+absl::Status IrEmitter::EmitSmallConstantGlobals() {
+  return EmitConstantGlobals(/*max_size_bytes=*/CpuInstructionFusion::
+                                 GetLargeConstantThresholdBytes());
+}
+
+absl::Status IrEmitter::EmitAllConstantGlobals() {
+  return EmitConstantGlobals(/*max_size_bytes=*/std::nullopt);
+}
+
+absl::Status IrEmitter::EmitConstantGlobals(
+    std::optional<size_t> max_size_bytes) {
   for (const BufferAllocation& allocation : assignment_.Allocations()) {
-    if (!allocation.is_constant()) {
+    // Large constants don't get fused with other instructions, so we don't
+    // need to emit them as globals.
+    if (!allocation.is_constant() ||
+        (max_size_bytes &&
+         llvm_ir::LiteralForConstantAllocation(allocation).size_bytes() >
+             *max_size_bytes)) {
       continue;
     }
 

xla/service/cpu/ir_emitter.h

@@ -177,8 +177,11 @@ class IrEmitter : public DfsHloVisitorWithDefault,
     compute_function_.pop();
   }
 
-  // Emit an LLVM global variable for every constant buffer allocation.
-  absl::Status EmitConstantGlobals();
+  // Emit LLVM global variable for a small constant buffer allocation.
+  absl::Status EmitSmallConstantGlobals();
+
+  // Emit LLVM global variables for all constant buffer allocations.
+  absl::Status EmitAllConstantGlobals();
 
   // Emits a call to a thread local function (e.g. to the computation nested
   // within a reduce or a map).  Thread local callees (by definition) only write
@@ -239,6 +242,9 @@ class IrEmitter : public DfsHloVisitorWithDefault,
  protected:
   friend class IrEmitter2;
 
+  // Emit an LLVM global variable for every constant buffer allocation.
+  absl::Status EmitConstantGlobals(std::optional<size_t> max_size_bytes);
+
   //
   // The following methods implement the DfsHloVisitor interface.
   //