[XLA:CPU] Port scatter to kernel API

PiperOrigin-RevId: 750108264

Author: WillFroom

xla/backends/cpu/codegen/BUILD

@@ -99,6 +99,8 @@ cc_library(
         "//xla:util",
         "//xla/backends/cpu/codegen/emitters/ir:xla_cpu",
         "//xla/backends/cpu/codegen/emitters/transforms:passes",
+        "//xla/codegen:llvm_ir_kernel_source",
+        "//xla/codegen:mlir_kernel_source",
         "//xla/codegen/emitters/ir:xla",
         "//xla/codegen/emitters/ir:xla_attrs_inc_gen",
         "//xla/codegen/emitters/transforms:passes",
@@ -107,12 +109,12 @@ cc_library(
         "//xla/mlir_hlo:mhlo_passes",
         "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
         "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Core",
-        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
         "@llvm-project//mlir:AffineToStandard",
         "@llvm-project//mlir:AffineTransforms",
@@ -485,6 +487,22 @@ cc_library(
     ],
 )
 
+py_strict_test(
+    name = "scatter_kernel_emitter_test",
+    srcs = ["scatter_kernel_emitter_test.py"],
+    tags = [
+        "no_oss",
+    ],
+    deps = [
+        "//third_party/py/numpy",
+        "//xla/backends/cpu/testlib",
+        "//xla/codegen/testlib",
+        "//xla/python:xla_extension",
+        "@absl_py//absl/testing:absltest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 xla_cc_test(
     name = "object_loader_test",
     srcs = ["object_loader_test.cc"],

xla/backends/cpu/codegen/emitters/BUILD

@@ -34,8 +34,13 @@ 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:kernel_definition",
+        "//xla/codegen:kernel_emitter",
+        "//xla/codegen:kernel_spec",
+        "//xla/codegen:mlir_kernel_source",
         "//xla/codegen/emitters:computation_partitioner",
         "//xla/codegen/emitters:elemental_hlo_to_mlir",
         "//xla/codegen/emitters:type_util",
@@ -51,6 +56,7 @@ cc_library(
         "//xla/service:scatter_simplifier",
         "//xla/service/cpu:backend_config_proto_cc",
         "//xla/service/llvm_ir:llvm_util",
+        "//xla/stream_executor:launch_dim",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
@@ -93,6 +99,9 @@ xla_cc_test(
     deps = [
         ":cpu_fusion_emitters",
         "//xla/backends/cpu/codegen:fusion_compiler",
+        "//xla/codegen:kernel_definition",
+        "//xla/codegen:llvm_ir_kernel_source",
+        "//xla/codegen:mlir_kernel_source",
         "//xla/hlo/analysis:hlo_ordering",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
@@ -101,7 +110,6 @@ 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",
@@ -111,5 +119,6 @@ xla_cc_test(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LLVMToLLVMIRTranslation",
         "@llvm-project//mlir:Pass",
+        "@tsl//tsl/platform:casts",
     ],
 )

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

@@ -26,6 +26,7 @@ limitations under the License.
 #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"
@@ -39,6 +40,7 @@ limitations under the License.
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -128,6 +130,33 @@ bool Needs64BitIndices(const HloComputation* computation) {
   }
   return false;
 }
+
+absl::Status SetKernelFunctionAttributes(const HloFusionInstruction& fusion,
+                                         mlir::Builder& builder,
+                                         mlir::func::FuncOp& func) {
+  const HloModule* hlo_module = fusion.GetModule();
+  if (hlo_module == nullptr) {
+    return Internal("HloModule is null");
+  }
+
+  mlir::MLIRContext* context = func->getContext();
+
+  // This is a hack until https://github.com/llvm/llvm-project/pull/135811 is
+  // merged, the value "2" corresponds to the default enum value.
+  mlir::ArrayAttr uwtable_attr = builder.getStrArrayAttr({"uwtable", "2"});
+  int32_t vector_width =
+      hlo_module->config().debug_options().xla_cpu_prefer_vector_width();
+  mlir::ArrayAttr prefer_vector_width_attr = builder.getStrArrayAttr(
+      {"prefer-vector-width", absl::StrCat(vector_width)});
+  func->setAttr("passthrough",
+                builder.getArrayAttr({uwtable_attr, prefer_vector_width_attr}));
+  func->setAttr(
+      "frame_pointer",
+      mlir::LLVM::FramePointerKindAttr::get(
+          context, mlir::LLVM::framePointerKind::FramePointerKind::All));
+
+  return absl::OkStatus();
+}
 }  // namespace
 
 using mlir::AffineExpr;
@@ -239,6 +268,10 @@ absl::StatusOr<mlir::func::FuncOp> EmitFusionKernelApi(
       loc, fusion.name(),
       builder.getFunctionType(/*arg_types=*/{call_frame_type},
                               /*result_types=*/{error_type}));
+
+  TF_RETURN_IF_ERROR(
+      SetKernelFunctionAttributes(fusion, builder, call_frame_func));
+
   builder.setInsertionPointToStart(call_frame_func.addEntryBlock());
   mlir::Value call_frame_arg = call_frame_func.getArgument(0);
   SmallVector<mlir::Value> extracted_values;
@@ -324,39 +357,6 @@ 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

@@ -60,10 +60,6 @@ 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:
   virtual ~CpuFusionEmitterBase() = default;

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

@@ -13,14 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.h"
-
-#include <cstdint>
 #include <memory>
 #include <string>
+#include <utility>
 
 #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"
@@ -32,6 +29,9 @@ limitations under the License.
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h"
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
+#include "xla/codegen/kernel_definition.h"
+#include "xla/codegen/llvm_ir_kernel_source.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/hlo_ordering.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/service/logical_buffer.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/platform/statusor.h"
+#include "tsl/platform/casts.h"
 
 namespace xla {
 namespace cpu {
@@ -62,18 +63,13 @@ std::string MlirModuleToString(const mlir::ModuleOp& module) {
 
 class CpuFusionEmitterTest : public HloTestBase {
  protected:
-  CpuFusionEmitterTest() : mlir_context_(FusionCompiler::CreateContext()) {}
-
   absl::StatusOr<std::unique_ptr<BufferAssignment>> RunBufferAssignment(
       const HloModule& hlo) {
     return BufferAssigner::Run(
         &hlo, std::make_unique<DependencyHloOrdering>(&hlo),
         backend().compiler()->BufferSizeBytesFunction(),
         [](LogicalBuffer::Color) { return /*alignment=*/1; });
   }
-
-  std::unique_ptr<mlir::MLIRContext> mlir_context_;
-  llvm::LLVMContext llvm_context_;
 };
 
 static constexpr absl::string_view kScatterHlo = R"(
@@ -136,10 +132,12 @@ TEST_F(CpuFusionEmitterTest, ScatterMlir) {
                           RunBufferAssignment(*hlo_module));
   auto fusion = Cast<HloFusionInstruction>(
       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.Emit());
-  auto mlir_dump = MlirModuleToString(*mlir_module);
+  CpuScatterFusion emitter(*buffer_assignment, fusion);
+  TF_ASSERT_OK_AND_ASSIGN(KernelDefinition kernel_definition,
+                          emitter.EmitKernelDefinition());
+  const auto& mlir_source =
+      tsl::down_cast<const MlirKernelSource&>(kernel_definition.source());
+  auto mlir_dump = mlir_source.ToString();
   TF_ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
                           RunFileCheck(mlir_dump, kExpected));
   EXPECT_TRUE(filecheck_matched);
@@ -162,17 +160,15 @@ TEST_F(CpuFusionEmitterTest, ScatterLlvm) {
                           RunBufferAssignment(*hlo_module));
   auto fusion = Cast<HloFusionInstruction>(
       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.Emit());
+  CpuScatterFusion emitter(*buffer_assignment, fusion);
+  TF_ASSERT_OK_AND_ASSIGN(KernelDefinition kernel_definition,
+                          emitter.EmitKernelDefinition());
+  auto [spec, source] = std::move(kernel_definition).release();
+  auto& mlir_source = tsl::down_cast<MlirKernelSource&>(*source);
   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(LlvmIrKernelSource llvm_source,
+                          compiler.Compile(std::move(mlir_source)));
+  auto llvm_dump = llvm_source.ToString();
   TF_ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
                           RunFileCheck(llvm_dump, kExpected));
   EXPECT_TRUE(filecheck_matched);