[xla:cpu] Remove mlir/llvm kernel_definition and kernel_emitter libraries

PiperOrigin-RevId: 825133247
2025-12-06 00:19:58 +01:00 · 2025-10-28 11:33:46 -07:00 · 2025-10-28 11:33:46 -07:00 · 44446df9cf
commit 44446df9cf
parent 63d558e46f
48 changed files with 32 additions and 245 deletions
--- a/third_party/xla/xla/backends/cpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/BUILD
@ -534,8 +534,6 @@ cc_library(
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_emitter",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:work_group",
@ -633,7 +631,7 @@ cc_library(
        "//xla/backends/cpu:alignment",
        "//xla/codegen:hlo_fusion_spec",
        "//xla/codegen:ir_emission_utils",
-        "//xla/codegen:mlir_kernel_definition",
+        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters:concatenate_kernel_emitter",
        "//xla/codegen/emitters:dynamic_update_slice_kernel_emitter",
        "//xla/codegen/emitters:kernel_arguments",
--- a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc
@ -44,7 +44,6 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
--- a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h
@ -27,9 +27,7 @@ limitations under the License.
 #include "llvm/IR/Module.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/kernel_emitter.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/buffer_assignment.h"

--- a/third_party/xla/xla/backends/cpu/codegen/dot/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/dot/BUILD
@ -18,8 +18,6 @@ cc_library(
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_emitter",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:work_group",
--- a/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter.cc
@ -26,7 +26,6 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/runtime/work_group.h"
--- a/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/dot/dot_kernel_emitter.h
@ -22,8 +22,7 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_emitter.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/buffer_assignment.h"

--- a/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/BUILD
@ -27,8 +27,6 @@ cc_library(
        "//xla/backends/cpu/codegen:target_machine_features",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:work_group",
@ -73,8 +71,6 @@ cc_library(
        "//xla/backends/cpu/codegen:target_machine_features",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/runtime:work_group",
@ -106,7 +102,7 @@ xla_cc_test(
        ":elemental_kernel_emitter",
        "//xla:xla_data_proto_cc",
        "//xla/codegen:kernel_definition",
-        "//xla/codegen:llvm_kernel_definition",
+        "//xla/codegen:llvm_kernel_source",
        "//xla/hlo/analysis:alias_info",
        "//xla/hlo/analysis:hlo_ordering",
        "//xla/hlo/ir:hlo",
--- a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.cc
@ -33,7 +33,6 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
--- a/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h
@ -21,8 +21,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/buffer_assignment.h"

--- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.cc
@ -39,7 +39,6 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
--- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter.h
@ -24,8 +24,7 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/runtime/work_group.h"
 #include "xla/service/buffer_assignment.h"
--- a/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/elemental/elemental_kernel_emitter_test.cc
@ -24,7 +24,7 @@ limitations under the License.
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Type.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_definition.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_ordering.h"
 #include "xla/hlo/ir/hlo_instruction.h"
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD
@ -39,8 +39,6 @@ cc_library(
        "//xla/backends/cpu/codegen:symbol_name_util",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters:computation_partitioner",
        "//xla/codegen/emitters:elemental_hlo_to_mlir",
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.cc
@ -59,7 +59,6 @@ limitations under the License.
 #include "xla/codegen/emitters/kernel_api_builder.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h
@ -17,17 +17,17 @@ limitations under the License.

 #include <cstdint>
 #include <optional>
+#include <string>
 #include <vector>

 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/Value.h"
 #include "xla/codegen/emitters/computation_partitioner.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/buffer_assignment.h"
--- a/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.cc
@ -42,7 +42,6 @@ limitations under the License.
 #include "xla/codegen/emitters/loop_kernel_emitter.h"
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/ir_emission_utils.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
--- a/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/fusion_emitter.h
@ -18,7 +18,7 @@ limitations under the License.

 #include "absl/status/statusor.h"
 #include "xla/codegen/emitters/kernel_arguments.h"
-#include "xla/codegen/mlir_kernel_definition.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/model/experimental/symbolic_expr.h"
--- a/third_party/xla/xla/backends/cpu/codegen/tools/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/tools/BUILD
@ -79,7 +79,7 @@ xla_cc_binary(
    deps = [
        "//xla/backends/cpu/codegen:fusion_compiler",
        "//xla/backends/cpu/codegen:fusion_emitter",
-        "//xla/codegen:mlir_kernel_definition",
+        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/tools:test_lib",
        "//xla/hlo/ir:hlo",
        "//xla/service/gpu/model/experimental:symbolic_expr",
--- a/third_party/xla/xla/backends/cpu/codegen/tools/fusion_to_mlir.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/tools/fusion_to_mlir.cc
@ -21,7 +21,6 @@ limitations under the License.
 #include "llvm/Support/raw_ostream.h"
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
 #include "xla/backends/cpu/codegen/fusion_emitter.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/tools/test_lib.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instructions.h"
--- a/third_party/xla/xla/backends/cpu/testlib/BUILD
+++ b/third_party/xla/xla/backends/cpu/testlib/BUILD
@ -35,9 +35,7 @@ cc_library(
        "//xla/backends/cpu/runtime:kernel_c_api",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen:mlir_kernel_definition",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/testlib:kernel_runner",
        "//xla/runtime:work_group",
@ -92,8 +90,6 @@ cc_library(
        "//xla:util",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
        "//xla/runtime:buffer_use",
        "//xla/runtime:work_group",
@ -115,8 +111,6 @@ cc_library(
        "//xla/backends/cpu/codegen:fusion_compiler",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/runtime:buffer_use",
        "//xla/runtime:work_group",
@ -156,11 +150,7 @@ tsl_pybind_extension(
        "//xla/backends/cpu/codegen/elemental:elemental_kernel_emitter",
        "//xla/backends/cpu/codegen/emitters:cpu_fusion_emitters",
        "//xla/codegen:kernel_definition",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/testlib:kernel_runner",
        "//xla/hlo/ir:hlo",
--- a/third_party/xla/xla/backends/cpu/testlib/kernel_runner.cc
+++ b/third_party/xla/xla/backends/cpu/testlib/kernel_runner.cc
@ -37,9 +37,7 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/kernel_c_api.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/runtime/work_group.h"
 #include "xla/service/cpu/cpu_options.h"
--- a/third_party/xla/xla/backends/cpu/testlib/kernel_runner.h
+++ b/third_party/xla/xla/backends/cpu/testlib/kernel_runner.h
@ -25,9 +25,7 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/function_library.h"
 #include "xla/backends/cpu/runtime/kernel.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/codegen/testlib/kernel_runner.h"
 #include "xla/runtime/work_group.h"
--- a/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc
+++ b/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc
@ -42,11 +42,7 @@ limitations under the License.
 #include "xla/backends/cpu/testlib/llvm_ir_kernel_emitter.h"
 #include "xla/backends/cpu/testlib/mlir_kernel_emitter.h"
 #include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/codegen/testlib/kernel_runner.h"
 #include "xla/hlo/ir/hlo_instruction.h"
--- a/third_party/xla/xla/backends/cpu/testlib/llvm_ir_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/testlib/llvm_ir_kernel_emitter.cc
@ -29,7 +29,6 @@ limitations under the License.
 #include "llvm/Support/SourceMgr.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/work_group.h"
--- a/third_party/xla/xla/backends/cpu/testlib/llvm_ir_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/testlib/llvm_ir_kernel_emitter.h
@ -23,8 +23,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/work_group.h"
 #include "xla/service/buffer_assignment.h"
--- a/third_party/xla/xla/backends/cpu/testlib/mlir_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/testlib/mlir_kernel_emitter.cc
@ -26,7 +26,6 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/work_group.h"
--- a/third_party/xla/xla/backends/cpu/testlib/mlir_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/testlib/mlir_kernel_emitter.h
@ -23,8 +23,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/runtime/work_group.h"
 #include "xla/service/buffer_assignment.h"
--- a/third_party/xla/xla/codegen/BUILD
+++ b/third_party/xla/xla/codegen/BUILD
@ -90,6 +90,8 @@ cc_library(
    srcs = ["llvm_kernel_source.cc"],
    hdrs = ["llvm_kernel_source.h"],
    deps = [
+        ":kernel_definition",
+        ":kernel_emitter",
        ":kernel_source",
        "//xla/service/llvm_ir:llvm_util",
        "@llvm-project//llvm:Core",
@ -117,6 +119,8 @@ cc_library(
    srcs = ["mlir_kernel_source.cc"],
    hdrs = ["mlir_kernel_source.h"],
    deps = [
+        ":kernel_definition",
+        ":kernel_emitter",
        ":kernel_source",
        "//xla:util",
        "//xla/service/gpu/model/experimental:symbolic_expr",
@ -129,44 +133,6 @@ cc_library(
    ],
 )

-cc_library(
-    name = "mlir_kernel_definition",
-    hdrs = ["mlir_kernel_definition.h"],
-    deps = [
-        ":kernel_definition",
-        ":kernel_source",
-        ":kernel_spec",
-        ":mlir_kernel_source",
-    ],
-)
-
-cc_library(
-    name = "mlir_kernel_emitter",
-    hdrs = ["mlir_kernel_emitter.h"],
-    deps = [
-        ":kernel_emitter",
-        ":mlir_kernel_definition",
-    ],
-)
-
-cc_library(
-    name = "llvm_kernel_emitter",
-    hdrs = ["llvm_kernel_emitter.h"],
-    deps = [
-        ":kernel_emitter",
-        ":llvm_kernel_definition",
-    ],
-)
-
-cc_library(
-    name = "llvm_kernel_definition",
-    hdrs = ["llvm_kernel_definition.h"],
-    deps = [
-        ":kernel_definition",
-        ":llvm_kernel_source",
-    ],
-)
-
 cc_library(
    name = "ir_emission_utils",
    srcs = ["ir_emission_utils.cc"],
--- a/third_party/xla/xla/codegen/emitters/BUILD
+++ b/third_party/xla/xla/codegen/emitters/BUILD
@ -311,8 +311,6 @@ cc_library(
        "//xla/codegen:hlo_fusion_spec",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters/ir:xla",
        "//xla/hlo/analysis:indexing_analysis",
@ -366,8 +364,6 @@ cc_library(
        "//xla/codegen:hlo_fusion_spec",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters/ir:xla",
        "//xla/hlo/analysis:indexing_analysis",
@ -427,8 +423,6 @@ cc_library(
        "//xla/codegen:ir_emission_utils",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters/ir:xla",
        "//xla/hlo/analysis:indexing_analysis",
--- a/third_party/xla/xla/codegen/emitters/concatenate_kernel_emitter.cc
+++ b/third_party/xla/xla/codegen/emitters/concatenate_kernel_emitter.cc
@ -53,7 +53,6 @@ limitations under the License.
 #include "xla/codegen/emitters/utils.h"
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
--- a/third_party/xla/xla/codegen/emitters/concatenate_kernel_emitter.h
+++ b/third_party/xla/xla/codegen/emitters/concatenate_kernel_emitter.h
@ -30,8 +30,7 @@ limitations under the License.
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/runtime/work_dimensions.h"
--- a/third_party/xla/xla/codegen/emitters/dynamic_update_slice_kernel_emitter.cc
+++ b/third_party/xla/xla/codegen/emitters/dynamic_update_slice_kernel_emitter.cc
@ -51,7 +51,6 @@ limitations under the License.
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/ir_emission_utils.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
--- a/third_party/xla/xla/codegen/emitters/dynamic_update_slice_kernel_emitter.h
+++ b/third_party/xla/xla/codegen/emitters/dynamic_update_slice_kernel_emitter.h
@ -29,8 +29,7 @@ limitations under the License.
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/utils/hlo_traversal.h"
--- a/third_party/xla/xla/codegen/emitters/loop_kernel_emitter.cc
+++ b/third_party/xla/xla/codegen/emitters/loop_kernel_emitter.cc
@ -49,7 +49,6 @@ limitations under the License.
 #include "xla/codegen/emitters/kernel_arguments.h"
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
--- a/third_party/xla/xla/codegen/emitters/loop_kernel_emitter.h
+++ b/third_party/xla/xla/codegen/emitters/loop_kernel_emitter.h
@ -29,8 +29,7 @@ limitations under the License.
 #include "xla/codegen/hlo_fusion_spec.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
+#include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/runtime/work_dimensions.h"
--- a/third_party/xla/xla/codegen/llvm_kernel_definition.h
+++ b/third_party/xla/xla/codegen/llvm_kernel_definition.h
@ -1,28 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_CODEGEN_LLVM_KERNEL_DEFINITION_H_
-#define XLA_CODEGEN_LLVM_KERNEL_DEFINITION_H_
-
-#include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/llvm_kernel_source.h"
-
-namespace xla {
-
-using LlvmKernelDefinition = KernelDefinition<LlvmKernelSource>;
-
-}  // namespace xla
-
-#endif  // XLA_CODEGEN_LLVM_KERNEL_DEFINITION_H_
--- a/third_party/xla/xla/codegen/llvm_kernel_emitter.h
+++ b/third_party/xla/xla/codegen/llvm_kernel_emitter.h
@ -1,29 +0,0 @@
-
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_CODEGEN_LLVM_KERNEL_EMITTER_H_
-#define XLA_CODEGEN_LLVM_KERNEL_EMITTER_H_
-
-#include "xla/codegen/kernel_emitter.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-
-namespace xla {
-
-using LlvmKernelEmitter = KernelEmitter<LlvmKernelDefinition>;
-
-}  // namespace xla
-
-#endif  // XLA_CODEGEN_LLVM_KERNEL_EMITTER_H_
--- a/third_party/xla/xla/codegen/llvm_kernel_source.h
+++ b/third_party/xla/xla/codegen/llvm_kernel_source.h
@ -23,6 +23,8 @@ limitations under the License.
 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
+#include "xla/codegen/kernel_definition.h"
+#include "xla/codegen/kernel_emitter.h"
 #include "xla/codegen/kernel_source.h"

 namespace xla {
@ -50,6 +52,9 @@ class LlvmKernelSource final : public KernelSource {
  llvm::orc::ThreadSafeModule module_;
 };

+using LlvmKernelDefinition = KernelDefinition<LlvmKernelSource>;  // NOLINT
+using LlvmKernelEmitter = KernelEmitter<LlvmKernelDefinition>;    // NOLINT
+
 }  // namespace xla

 #endif  // XLA_CODEGEN_LLVM_KERNEL_SOURCE_H_
--- a/third_party/xla/xla/codegen/mlir_kernel_definition.h
+++ b/third_party/xla/xla/codegen/mlir_kernel_definition.h
@ -1,28 +0,0 @@
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_CODEGEN_MLIR_KERNEL_DEFINITION_H_
-#define XLA_CODEGEN_MLIR_KERNEL_DEFINITION_H_
-
-#include "xla/codegen/kernel_definition.h"
-#include "xla/codegen/mlir_kernel_source.h"
-
-namespace xla {
-
-using MlirKernelDefinition = KernelDefinition<MlirKernelSource>;
-
-}  // namespace xla
-
-#endif  // XLA_CODEGEN_MLIR_KERNEL_DEFINITION_H_
--- a/third_party/xla/xla/codegen/mlir_kernel_emitter.h
+++ b/third_party/xla/xla/codegen/mlir_kernel_emitter.h
@ -1,29 +0,0 @@
-
-/* Copyright 2024 The OpenXLA Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-==============================================================================*/
-
-#ifndef XLA_CODEGEN_MLIR_KERNEL_EMITTER_H_
-#define XLA_CODEGEN_MLIR_KERNEL_EMITTER_H_
-
-#include "xla/codegen/kernel_emitter.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-
-namespace xla {
-
-using MlirKernelEmitter = KernelEmitter<MlirKernelDefinition>;
-
-}  // namespace xla
-
-#endif  // XLA_CODEGEN_MLIR_KERNEL_EMITTER_H_
--- a/third_party/xla/xla/codegen/mlir_kernel_source.h
+++ b/third_party/xla/xla/codegen/mlir_kernel_source.h
@ -26,6 +26,8 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Support/DebugStringHelper.h"
+#include "xla/codegen/kernel_definition.h"
+#include "xla/codegen/kernel_emitter.h"
 #include "xla/codegen/kernel_source.h"
 #include "xla/service/gpu/model/experimental/symbolic_expr.h"

@ -81,6 +83,9 @@ class MlirKernelSource final : public KernelSource {
  Storage storage_;
 };

+using MlirKernelDefinition = KernelDefinition<MlirKernelSource>;  // NOLINT
+using MlirKernelEmitter = KernelEmitter<MlirKernelDefinition>;    // NOLINT
+
 }  // namespace xla

 #endif  // XLA_CODEGEN_MLIR_KERNEL_SOURCE_H_
--- a/third_party/xla/xla/codegen/testlib/BUILD
+++ b/third_party/xla/xla/codegen/testlib/BUILD
@ -57,11 +57,7 @@ tsl_pybind_extension(
        "//xla/codegen:kernel_emitter",
        "//xla/codegen:kernel_source",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
-        "//xla/codegen:llvm_kernel_emitter",
        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen:mlir_kernel_definition",
-        "//xla/codegen:mlir_kernel_emitter",
        "//xla/codegen:mlir_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/parser:hlo_parser",
--- a/third_party/xla/xla/codegen/testlib/kernel_runner_extension.cc
+++ b/third_party/xla/xla/codegen/testlib/kernel_runner_extension.cc
@ -38,11 +38,7 @@ limitations under the License.
 #include "xla/codegen/kernel_emitter.h"
 #include "xla/codegen/kernel_source.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/llvm_kernel_emitter.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_emitter.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/codegen/testlib/kernel_runner.h"
 #include "xla/comparison_util.h"
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@ -974,9 +974,7 @@ cc_library(
        "//xla/backends/cpu/runtime/xnnpack:xnn_fusion_thunk",
        "//xla/codegen:kernel_definition",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen:mlir_kernel_definition",
        "//xla/codegen:mlir_kernel_source",
        "//xla/codegen/emitters:computation_fingerprint",
        "//xla/codegen/emitters:kernel_api_builder",
@ -2050,9 +2048,7 @@ cc_library(
        "//xla/backends/cpu/codegen:fusion_compiler",
        "//xla/backends/cpu/codegen:fusion_emitter",
        "//xla/codegen:kernel_spec",
-        "//xla/codegen:llvm_kernel_definition",
        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen:mlir_kernel_definition",
        "//xla/codegen:mlir_kernel_source",
        "//xla/hlo/ir:hlo",
        "//xla/service:buffer_assignment",
@ -2079,7 +2075,6 @@ xla_test(
    deps = [
        ":parallel_fusion_emitter",
        "//xla/backends/cpu/codegen:fusion_compiler",
-        "//xla/codegen:llvm_kernel_definition",
        "//xla/hlo/ir:hlo",
        "//xla/hlo/testlib:filecheck",
        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
--- a/third_party/xla/xla/service/cpu/parallel_fusion_emitter.cc
+++ b/third_party/xla/xla/service/cpu/parallel_fusion_emitter.cc
@ -32,9 +32,7 @@ limitations under the License.
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
 #include "xla/backends/cpu/codegen/fusion_emitter.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/buffer_assignment.h"
--- a/third_party/xla/xla/service/cpu/parallel_fusion_emitter.h
+++ b/third_party/xla/xla/service/cpu/parallel_fusion_emitter.h
@ -26,8 +26,7 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
-#include "xla/codegen/mlir_kernel_definition.h"
+#include "xla/codegen/llvm_kernel_source.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/tsl/platform/threadpool.h"
--- a/third_party/xla/xla/service/cpu/parallel_fusion_emitter_test.cc
+++ b/third_party/xla/xla/service/cpu/parallel_fusion_emitter_test.cc
@ -29,7 +29,6 @@ limitations under the License.
 #include "llvm/IR/Module.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "xla/backends/cpu/codegen/fusion_compiler.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
--- a/third_party/xla/xla/service/cpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc
@ -78,9 +78,7 @@ limitations under the License.
 #include "xla/codegen/emitters/kernel_arguments.h"
 #include "xla/codegen/kernel_definition.h"
 #include "xla/codegen/kernel_spec.h"
-#include "xla/codegen/llvm_kernel_definition.h"
 #include "xla/codegen/llvm_kernel_source.h"
-#include "xla/codegen/mlir_kernel_definition.h"
 #include "xla/codegen/mlir_kernel_source.h"
 #include "xla/comparison_util.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"