[XLA:CPU] Make CpuFusionEmitterBase a pure interface class

PiperOrigin-RevId: 745944794

Author: WillFroom

xla/backends/cpu/codegen/emitters/BUILD

@@ -34,7 +34,6 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/cpu:alignment",
-        "//xla/backends/cpu/codegen:fusion_compiler",
         "//xla/backends/cpu/codegen:kernel_api_ir_builder",
         "//xla/backends/cpu/codegen/emitters/ir:xla_cpu",
         "//xla/codegen/emitters:computation_partitioner",
@@ -101,6 +100,7 @@ xla_cc_test(
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest",

xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc

@@ -23,16 +23,17 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
-#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
 #include "llvm/Linker/Linker.h"
 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
@@ -60,7 +61,6 @@ limitations under the License.
 #include "xla/backends/cpu/alignment.h"
 #include "xla/backends/cpu/codegen/emitters/ir/xla_cpu_ops.h"
 #include "xla/backends/cpu/codegen/emitters/ir/xla_cpu_types.h"
-#include "xla/backends/cpu/codegen/fusion_compiler.h"
 #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h"
 #include "xla/codegen/emitters/computation_partitioner.h"
 #include "xla/codegen/emitters/elemental_hlo_to_mlir.h"
@@ -76,7 +76,6 @@ limitations under the License.
 #include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/dump.h"
-#include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/platform/errors.h"
@@ -131,86 +130,6 @@ bool Needs64BitIndices(const HloComputation* computation) {
 }
 }  // namespace
 
-absl::StatusOr<CpuFusionEmissionResult> CpuFusionEmitterBase::Emit() const {
-  // Single-threaded for now.
-  TF_ASSIGN_OR_RETURN(auto module,
-                      CreateLLVMModule(*mlir_context_, *llvm_context_, *fusion_,
-                                       buffer_assignment_));
-
-  const HloModule* hlo_module = fusion_->GetModule();
-  if (hlo_module == nullptr) {
-    return Internal("HloModule is null");
-  }
-  // Create a Kernel API Builder and a throwaway kernel prototype in order to
-  // extract useful info from them, e.g. noalias, invariant_arguments and
-  // entry function attributes.
-  // TODO(ecg): find a way to obtain the same info without wasting work by
-  // creating a throwaway module. All of this additional info should probably be
-  // explicit in the generated MLIR, not added afterwards like we're doing here.
-  // TODO(ecg): some attributes on the final loads are missing wrt those
-  // generated via KernelApiIrBuilder, e.g. noalias. Add them.
-  KernelApiIrBuilder kernel_api_ir_builder(
-      *llvm_context_,
-      KernelApiIrBuilder::Options::FromHloModuleConfig(hlo_module->config()));
-  std::unique_ptr<llvm::Module> throwaway_llvm_module =
-      KernelApiIrBuilder::CreateModule(
-          absl::StrCat(fusion_->name(), "_throwaway_module"), *llvm_context_);
-  TF_ASSIGN_OR_RETURN(KernelApiIrBuilder::KernelPrototype kernel_prototype,
-                      kernel_api_ir_builder.EmitKernelPrototype(
-                          *throwaway_llvm_module, fusion_, &buffer_assignment_,
-                          "_throwaway_kernel_prototype"));
-  llvm::Function* kernel_function = module->getFunction(fusion_->name());
-  kernel_api_ir_builder.SetKernelFunctionAttributes(kernel_function);
-
-  CpuFusionEmissionResult result;
-  result.llvm_module = std::move(module);
-  result.invariant_arguments = std::move(kernel_prototype.invariant_arguments);
-  return result;
-}
-
-absl::StatusOr<std::unique_ptr<llvm::Module>>
-CpuFusionEmitterBase::CreateLLVMModule(
-    mlir::MLIRContext& mlir_context, llvm::LLVMContext& llvm_context,
-    const HloFusionInstruction& fusion,
-    const BufferAssignment& buffer_assignment) const {
-  TF_ASSIGN_OR_RETURN(auto module,
-                      CreateMLIRModule(mlir_context, fusion,
-                                       std::string(fusion.name()) + "_entry",
-                                       buffer_assignment));
-
-  FusionCompiler compiler(FusionCompiler::Options{});
-  return compiler.Compile(llvm_context, module.get());
-}
-
-absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
-CpuFusionEmitterBase::CreateMLIRModule(
-    mlir::MLIRContext& context, const HloFusionInstruction& fusion,
-    const std::string& entry_function_name,
-    const BufferAssignment& buffer_assignment,
-    mlir::interpreter::MlirCompilationTrace* trace) const {
-  mlir::OpBuilder builder(&context);
-  auto loc = mlir::NameLoc::get(builder.getStringAttr(fusion.name()));
-  mlir::OwningOpRef<mlir::ModuleOp> module = llvm_ir::CreateMlirModuleOp(loc);
-  SetDataLayoutAttribute(module.get(), fusion);
-
-  TF_ASSIGN_OR_RETURN(
-      mlir::func::FuncOp entry_func,
-      EmitFusionKernelApi(module.get(), fusion, entry_function_name,
-                          buffer_assignment));
-
-  std::vector<emitters::EpilogueSpecification> epilogues =
-      GetEpilogues(fusion, &context);
-  emitters::PartitionedComputations computations(
-      fusion.fused_instructions_computation(), &context, epilogues);
-  TF_ASSIGN_OR_RETURN(
-      emitters::CallTargetProvider call_targets,
-      EmitCallTargets(module.get(), fusion, computations, epilogues));
-
-  TF_RETURN_IF_ERROR(
-      EmitEntryFunction(computations, call_targets, entry_func, fusion));
-  return module;
-}
-
 using mlir::AffineExpr;
 
 IndexingMap GetDefaultIndexingMap(absl::Span<const int64_t> thread_tile_sizes,
@@ -405,6 +324,39 @@ void SetDataLayoutAttribute(mlir::ModuleOp module,
       mlir::DataLayoutSpecAttr::get(module->getContext(), {index_layout}));
 }
 
+absl::StatusOr<absl::flat_hash_set<int64_t>> SetKernelFunctionAttributes(
+    llvm::Module& module, const BufferAssignment& buffer_assignment,
+    const HloFusionInstruction* fusion) {
+  const HloModule* hlo_module = fusion->GetModule();
+  if (hlo_module == nullptr) {
+    return Internal("HloModule is null");
+  }
+
+  // Create a Kernel API Builder and a throwaway kernel prototype in order to
+  // extract useful info from them, e.g. noalias, invariant_arguments and
+  // entry function attributes.
+  // TODO(ecg): find a way to obtain the same info without wasting work by
+  // creating a throwaway module. All of this additional info should probably be
+  // explicit in the generated MLIR, not added afterwards like we're doing here.
+  // TODO(ecg): some attributes on the final loads are missing wrt those
+  // generated via KernelApiIrBuilder, e.g. noalias. Add them.
+  llvm::LLVMContext& context = module.getContext();
+  KernelApiIrBuilder kernel_api_ir_builder(
+      context,
+      KernelApiIrBuilder::Options::FromHloModuleConfig(hlo_module->config()));
+  std::unique_ptr<llvm::Module> throwaway_llvm_module =
+      KernelApiIrBuilder::CreateModule(
+          absl::StrCat(fusion->name(), "_throwaway_module"), context);
+  TF_ASSIGN_OR_RETURN(KernelApiIrBuilder::KernelPrototype kernel_prototype,
+                      kernel_api_ir_builder.EmitKernelPrototype(
+                          *throwaway_llvm_module, fusion, &buffer_assignment,
+                          "_throwaway_kernel_prototype"));
+  llvm::Function* kernel_function = module.getFunction(fusion->name());
+  kernel_api_ir_builder.SetKernelFunctionAttributes(kernel_function);
+
+  return kernel_prototype.invariant_arguments;
+}
+
 int64_t CeilDiv(int64_t a, int64_t b) { return (a + b - 1) / b; }
 }  // namespace cpu
 }  // namespace xla

xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.h

@@ -16,24 +16,20 @@ limitations under the License.
 #define XLA_BACKENDS_CPU_CODEGEN_EMITTERS_CPU_FUSION_EMITTER_H_
 
 #include <cstdint>
-#include <memory>
 #include <optional>
 #include <string>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
-#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
-#include "mlir/IR/Value.h"
 #include "mlir/Pass/PassManager.h"
 #include "xla/codegen/emitters/computation_partitioner.h"
 #include "xla/hlo/analysis/indexing_map.h"
@@ -44,11 +40,6 @@ limitations under the License.
 namespace xla {
 namespace cpu {
 
-struct CpuFusionEmissionResult {
-  std::unique_ptr<llvm::Module> llvm_module;
-  absl::flat_hash_set<int64_t> invariant_arguments;
-};
-
 IndexingMap GetDefaultIndexingMap(absl::Span<const int64_t> thread_tile_sizes,
                                   absl::Span<const int64_t> shape,
                                   mlir::MLIRContext* mlir_context);
@@ -69,17 +60,12 @@ absl::StatusOr<emitters::CallTargetProvider> EmitCallTargets(
 void SetDataLayoutAttribute(mlir::ModuleOp module,
                             const HloFusionInstruction& fusion);
 
+absl::StatusOr<absl::flat_hash_set<int64_t>> SetKernelFunctionAttributes(
+    llvm::Module& module, const BufferAssignment& buffer_assignment,
+    const HloFusionInstruction* fusion);
+
 class CpuFusionEmitterBase {
  public:
-  CpuFusionEmitterBase(mlir::MLIRContext* mlir_context,
-                       llvm::LLVMContext* llvm_context,
-                       const BufferAssignment& buffer_assignment,
-                       const HloFusionInstruction* fusion)
-      : mlir_context_(mlir_context),
-        llvm_context_(llvm_context),
-        buffer_assignment_(buffer_assignment),
-        fusion_(fusion) {}
-
   virtual ~CpuFusionEmitterBase() = default;
 
   virtual int64_t num_threads() const = 0;
@@ -92,42 +78,7 @@ class CpuFusionEmitterBase {
 
   virtual std::string BackendExtraOptions() { return {}; }
 
-  absl::StatusOr<CpuFusionEmissionResult> Emit() const;
-
-  // Visible for testing.
-  absl::StatusOr<std::unique_ptr<llvm::Module>> CreateLLVMModule(
-      mlir::MLIRContext& mlir_context, llvm::LLVMContext& llvm_context,
-      const HloFusionInstruction& fusion,
-      const BufferAssignment& buffer_assignment) const;
-
-  // Visible for testing.
-  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> CreateMLIRModule(
-      mlir::MLIRContext& context, const HloFusionInstruction& fusion,
-      const std::string& entry_function_name,
-      const BufferAssignment& buffer_assignment,
-      mlir::interpreter::MlirCompilationTrace* trace = nullptr) const;
-
- protected:
-  virtual absl::Status EmitEntryFunction(
-      const emitters::PartitionedComputations& computations,
-      const emitters::CallTargetProvider& call_targets,
-      mlir::func::FuncOp entry_function,
-      const HloFusionInstruction& fusion) const = 0;
-
-  virtual std::vector<emitters::EpilogueSpecification> GetEpilogues(
-      const HloFusionInstruction& fusion,
-      mlir::MLIRContext* mlir_context) const {
-    // We don't actually support epilogues for scatter, but this is how we tell
-    // the base class that we don't want it to generate code for the scatter.
-    return {};
-  }
-
-  mlir::Value EmitThreadId(mlir::ImplicitLocOpBuilder& builder, int dim) const;
-
-  mlir::MLIRContext* mlir_context_;
-  llvm::LLVMContext* llvm_context_;
-  const BufferAssignment& buffer_assignment_;
-  const HloFusionInstruction* fusion_;
+  virtual absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Emit() const = 0;
 };
 
 int64_t CeilDiv(int64_t a, int64_t b);

xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_test.cc

@@ -15,10 +15,12 @@ limitations under the License.
 
 #include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 
 #include <gtest/gtest.h>
+#include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "llvm/IR/LLVMContext.h"
@@ -135,11 +137,7 @@ TEST_F(CpuFusionEmitterTest, ScatterMlir) {
       hlo_module->entry_computation()->root_instruction());
   CpuScatterFusion emitter(mlir_context_.get(), &llvm_context_,
                            *buffer_assignment, fusion);
-  TF_ASSERT_OK_AND_ASSIGN(
-      auto mlir_module,
-      emitter.CreateMLIRModule(*mlir_context_, *fusion,
-                               std::string(fusion->name()) + "_entry",
-                               *buffer_assignment));
+  TF_ASSERT_OK_AND_ASSIGN(auto mlir_module, emitter.Emit());
   auto mlir_dump = MlirModuleToString(*mlir_module);
   TF_ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
                           RunFileCheck(mlir_dump, kExpected));
@@ -165,8 +163,15 @@ TEST_F(CpuFusionEmitterTest, ScatterLlvm) {
       hlo_module->entry_computation()->root_instruction());
   CpuScatterFusion emitter(mlir_context_.get(), &llvm_context_,
                            *buffer_assignment, fusion);
-  TF_ASSERT_OK_AND_ASSIGN(auto result, emitter.Emit());
-  auto llvm_dump = LlvmModuleToString(*result.llvm_module);
+  TF_ASSERT_OK_AND_ASSIGN(auto mlir_module, emitter.Emit());
+  FusionCompiler compiler(FusionCompiler::Options{});
+  llvm::LLVMContext llvm_context;
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<llvm::Module> llvm_module,
+                          compiler.Compile(llvm_context, mlir_module.get()));
+  TF_ASSERT_OK_AND_ASSIGN(
+      absl::flat_hash_set<int64_t> invariant_arguments,
+      SetKernelFunctionAttributes(*llvm_module, *buffer_assignment, fusion));
+  auto llvm_dump = LlvmModuleToString(*llvm_module);
   TF_ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
                           RunFileCheck(llvm_dump, kExpected));
   EXPECT_TRUE(filecheck_matched);

xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.cc

@@ -40,8 +40,10 @@ limitations under the License.
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/Location.h"
+#include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.h"
 #include "xla/codegen/emitters/computation_partitioner.h"
 #include "xla/codegen/emitters/elemental_hlo_to_mlir.h"
@@ -55,9 +57,12 @@ limitations under the License.
 #include "xla/primitive_util.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/cpu/backend_config.pb.h"
+#include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/service/scatter_simplifier.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -173,8 +178,10 @@ CpuScatterFusion::CpuScatterFusion(mlir::MLIRContext* mlir_context,
                                    llvm::LLVMContext* llvm_context,
                                    const BufferAssignment& buffer_assignment,
                                    const HloFusionInstruction* fusion)
-    : CpuFusionEmitterBase{mlir_context, llvm_context, buffer_assignment,
-                           fusion} {
+    : mlir_context_(mlir_context),
+      llvm_context_(llvm_context),
+      buffer_assignment_(buffer_assignment),
+      fusion_(fusion) {
   const auto* scatter = Cast<HloScatterInstruction>(
       fusion->fused_instructions_computation()->root_instruction());
   auto update_shape = scatter->scatter_updates().front()->shape();
@@ -236,6 +243,32 @@ IndexingMap GetScatterIndexingMap(
       {}, constraints);
 }
 
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> CpuScatterFusion::Emit()
+    const {
+  mlir::OpBuilder builder(mlir_context_);
+  auto loc = mlir::NameLoc::get(builder.getStringAttr(fusion_->name()));
+  mlir::OwningOpRef<mlir::ModuleOp> module = llvm_ir::CreateMlirModuleOp(loc);
+  SetDataLayoutAttribute(module.get(), *fusion_);
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::func::FuncOp entry_func,
+      EmitFusionKernelApi(module.get(), *fusion_,
+                          std::string(fusion_->name()) + "_entry",
+                          buffer_assignment_));
+
+  std::vector<emitters::EpilogueSpecification> epilogues =
+      GetEpilogues(*fusion_, mlir_context_);
+  emitters::PartitionedComputations computations(
+      fusion_->fused_instructions_computation(), mlir_context_, epilogues);
+  TF_ASSIGN_OR_RETURN(
+      emitters::CallTargetProvider call_targets,
+      EmitCallTargets(module.get(), *fusion_, computations, epilogues));
+
+  TF_RETURN_IF_ERROR(
+      EmitEntryFunction(computations, call_targets, entry_func, *fusion_));
+  return module;
+}
+
 absl::Status CpuScatterFusion::EmitEntryFunction(
     const emitters::PartitionedComputations& computations,
     const emitters::CallTargetProvider& call_targets,