OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 00:19:58 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

[xla:codegen] Remove MlirKernelEmitter alias

Browse Source 
PiperOrigin-RevId: 825346639
This commit is contained in:
Eugene Zhulenev 2025-10-28 21:11:41 -07:00 committed by TensorFlower Gardener
parent 1444679887
commit 27a45465a0
10 changed files with 28 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
third_party/xla/xla/backends/cpu/codegen/emitters/BUILD vendored
View File

@ -38,6 +38,7 @@ cc_library(
"//xla/backends/cpu/codegen:kernel_api_ir_builder",
"//xla/backends/cpu/codegen:symbol_name_util",
"//xla/codegen:kernel_definition",
"//xla/codegen:kernel_emitter",
"//xla/codegen:kernel_spec",
"//xla/codegen:mlir_kernel_source",
"//xla/codegen/emitters:computation_partitioner",

6
third_party/xla/xla/backends/cpu/codegen/emitters/cpu_scatter_emitter.h vendored
View File

@ -27,7 +27,7 @@ limitations under the License.
#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/kernel_emitter.h"
#include "xla/codegen/mlir_kernel_source.h"
#include "xla/hlo/analysis/indexing_map.h"
#include "xla/hlo/ir/hlo_instructions.h"
@ -38,14 +38,14 @@ namespace xla {
namespace cpu {
// Generic scatter fusion. Lowers to LLVM via MLIR.
class CpuScatterFusion final : public MlirKernelEmitter {
class CpuScatterFusion final : public KernelEmitter<MlirKernelSource> {
public:
CpuScatterFusion(const BufferAssignment& buffer_assignment,
const HloFusionInstruction* fusion,
gpu::SymbolicExprContext* symbolic_expr_context);
absl::string_view name() const final { return "cpu_scatter_fusion"; }
absl::StatusOr<MlirKernelDefinition> EmitKernelDefinition() final;
absl::StatusOr<KernelDefinition> EmitKernelDefinition() final;
private:
absl::Status EmitEntryFunction(

1
third_party/xla/xla/backends/cpu/testlib/BUILD vendored
View File

@ -112,6 +112,7 @@ cc_library(
deps = [
"//xla/backends/cpu/codegen:fusion_compiler",
"//xla/codegen:kernel_definition",
"//xla/codegen:kernel_emitter",
"//xla/codegen:kernel_spec",
"//xla/codegen:mlir_kernel_source",
"//xla/runtime:buffer_use",

12
third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc vendored
View File

@ -100,10 +100,10 @@ NB_MODULE(_extension, kernel_runner_module) {
{});
});
nb::class_<MlirTestKernelEmitter, MlirKernelEmitter>(kernel_runner_module,
"MlirTestKernelEmitter")
.def("__init__", [](MlirKernelEmitter* self, absl::string_view ir,
absl::string_view kernel_name,
nb::class_<MlirTestKernelEmitter, KernelEmitter<MlirKernelSource>>(
kernel_runner_module, "MlirTestKernelEmitter")
.def("__init__", [](KernelEmitter<MlirKernelSource>* self,
absl::string_view ir, absl::string_view kernel_name,
NbNumWorkGroups num_workgroups) {
new (self)
MlirTestKernelEmitter(ir, kernel_name,
@ -189,8 +189,8 @@ NB_MODULE(_extension, kernel_runner_module) {
nb::keep_alive<1, 2>(), nb::keep_alive<1, 3>(),
nb::keep_alive<1, 4>());
nb::class_<CpuScatterFusion, MlirKernelEmitter>(kernel_runner_module,
"ScatterKernelEmitter")
nb::class_<CpuScatterFusion, KernelEmitter<MlirKernelSource>>(
kernel_runner_module, "ScatterKernelEmitter")
.def(
"__init__",
[](CpuScatterFusion* self, const HloFusionInstruction* instruction,

5
third_party/xla/xla/backends/cpu/testlib/mlir_kernel_emitter.h vendored
View File

@ -23,6 +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/kernel_emitter.h"
#include "xla/codegen/mlir_kernel_source.h"
#include "xla/runtime/buffer_use.h"
#include "xla/runtime/work_group.h"
@ -34,7 +35,7 @@ namespace xla::cpu {
// into the dedicated MLIR context and module instance. This kernel emitter is
// intended to be used for testing purposes only: (1) load pre-compiled LLVM IR
// into the XLA kernel spec; (2) Execute it with user provided input buffers.
class MlirTestKernelEmitter : public MlirKernelEmitter {
class MlirTestKernelEmitter : public KernelEmitter<MlirKernelSource> {
public:
// When loading kernel IR into the KernelSpec we create a separate buffer
// allocation for every kernel argument. We don't use buffer assignment in
@ -49,7 +50,7 @@ class MlirTestKernelEmitter : public MlirKernelEmitter {
absl::Span<const KernelArg> args);
absl::string_view name() const override { return "mlir_test_kernel_emitter"; }
absl::StatusOr<MlirKernelDefinition> EmitKernelDefinition() final;
absl::StatusOr<KernelDefinition> EmitKernelDefinition() final;
private:
std::string mlir_;

3
third_party/xla/xla/codegen/emitters/BUILD vendored
View File

@ -310,6 +310,7 @@ cc_library(
"//xla:util",
"//xla/codegen:hlo_fusion_spec",
"//xla/codegen:kernel_definition",
"//xla/codegen:kernel_emitter",
"//xla/codegen:kernel_spec",
"//xla/codegen:mlir_kernel_source",
"//xla/codegen/emitters/ir:xla",
@ -363,6 +364,7 @@ cc_library(
"//xla:util",
"//xla/codegen:hlo_fusion_spec",
"//xla/codegen:kernel_definition",
"//xla/codegen:kernel_emitter",
"//xla/codegen:kernel_spec",
"//xla/codegen:mlir_kernel_source",
"//xla/codegen/emitters/ir:xla",
@ -422,6 +424,7 @@ cc_library(
"//xla/codegen:hlo_fusion_spec",
"//xla/codegen:ir_emission_utils",
"//xla/codegen:kernel_definition",
"//xla/codegen:kernel_emitter",
"//xla/codegen:kernel_spec",
"//xla/codegen:mlir_kernel_source",
"//xla/codegen/emitters/ir:xla",

8
third_party/xla/xla/codegen/emitters/concatenate_kernel_emitter.h vendored
View File

@ -16,7 +16,6 @@ limitations under the License.
#ifndef XLA_CODEGEN_EMITTERS_CONCATENATE_KERNEL_EMITTER_H_
#define XLA_CODEGEN_EMITTERS_CONCATENATE_KERNEL_EMITTER_H_
#include <cstdint>
#include <string>
#include <vector>
@ -28,7 +27,7 @@ limitations under the License.
#include "xla/codegen/emitters/ir/xla_ops.h"
#include "xla/codegen/emitters/kernel_arguments.h"
#include "xla/codegen/hlo_fusion_spec.h"
#include "xla/codegen/kernel_definition.h"
#include "xla/codegen/kernel_emitter.h"
#include "xla/codegen/mlir_kernel_source.h"
#include "xla/hlo/analysis/indexing_map.h"
#include "xla/hlo/ir/hlo_instructions.h"
@ -39,7 +38,8 @@ limitations under the License.
namespace xla::emitters {
class ConcatenateFusionKernelEmitter final : public MlirKernelEmitter {
class ConcatenateFusionKernelEmitter final
: public KernelEmitter<MlirKernelSource> {
public:
ConcatenateFusionKernelEmitter(
gpu::SymbolicExprContext& symbolic_expr_context,
@ -53,7 +53,7 @@ class ConcatenateFusionKernelEmitter final : public MlirKernelEmitter {
return "concatenate_fusion_kernel_emitter";
}
absl::StatusOr<MlirKernelDefinition> EmitKernelDefinition() override;
absl::StatusOr<KernelDefinition> EmitKernelDefinition() override;
static IndexingMap ComputeWorkItemIdToOutputIndexing(
const WorkDimensions& work_dimensions, const Shape& largest_shape,

7
third_party/xla/xla/codegen/emitters/dynamic_update_slice_kernel_emitter.h vendored
View File

@ -27,7 +27,7 @@ limitations under the License.
#include "xla/codegen/emitters/ir/xla_ops.h"
#include "xla/codegen/emitters/kernel_arguments.h"
#include "xla/codegen/hlo_fusion_spec.h"
#include "xla/codegen/kernel_definition.h"
#include "xla/codegen/kernel_emitter.h"
#include "xla/codegen/kernel_spec.h"
#include "xla/codegen/mlir_kernel_source.h"
#include "xla/hlo/analysis/indexing_map.h"
@ -46,7 +46,8 @@ namespace xla::emitters {
// 3. a tuple op returning the result of several dynamic-update-slice ops
// 4. a tuple op returning the result of several bitcast
// dynamic-update-slice ops
class DynamicUpdateSliceKernelEmitter final : public MlirKernelEmitter {
class DynamicUpdateSliceKernelEmitter final
: public KernelEmitter<MlirKernelSource> {
public:
DynamicUpdateSliceKernelEmitter(
gpu::SymbolicExprContext& symbolic_expr_context,
@ -60,7 +61,7 @@ class DynamicUpdateSliceKernelEmitter final : public MlirKernelEmitter {
return "dynamic_update_slice_kernel_emitter";
}
absl::StatusOr<MlirKernelDefinition> EmitKernelDefinition() override;
absl::StatusOr<KernelDefinition> EmitKernelDefinition() override;
// Get the shape that will be used for loop indexing for the given fusion
// specification.

6
third_party/xla/xla/codegen/emitters/loop_kernel_emitter.h vendored
View File

@ -26,7 +26,7 @@ limitations under the License.
#include "xla/codegen/emitters/ir/xla_ops.h"
#include "xla/codegen/emitters/kernel_arguments.h"
#include "xla/codegen/hlo_fusion_spec.h"
#include "xla/codegen/kernel_definition.h"
#include "xla/codegen/kernel_emitter.h"
#include "xla/codegen/mlir_kernel_source.h"
#include "xla/hlo/analysis/indexing_map.h"
#include "xla/hlo/ir/hlo_instructions.h"
@ -38,7 +38,7 @@ limitations under the License.
namespace xla::emitters {
// Generic loop fusion.
class LoopFusionKernelEmitter final : public MlirKernelEmitter {
class LoopFusionKernelEmitter final : public KernelEmitter<MlirKernelSource> {
public:
LoopFusionKernelEmitter(gpu::SymbolicExprContext& symbolic_expr_context,
const HloFusionInstruction& fusion,
@ -50,7 +50,7 @@ class LoopFusionKernelEmitter final : public MlirKernelEmitter {
BackendKind backend_kind);
absl::string_view name() const final { return "loop_fusion_kernel_emitter"; }
absl::StatusOr<MlirKernelDefinition> EmitKernelDefinition() override;
absl::StatusOr<KernelDefinition> EmitKernelDefinition() override;
static IndexingMap ComputeWorkItemIdToOutputIndexing(
const WorkDimensions& work_dimensions, const Shape& root_shape,

2
third_party/xla/xla/codegen/mlir_kernel_source.h vendored
View File

@ -27,7 +27,6 @@ limitations under the License.
#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"
@ -84,7 +83,6 @@ class MlirKernelSource final : public KernelSource {
};
using MlirKernelDefinition = KernelDefinition<MlirKernelSource>; // NOLINT
using MlirKernelEmitter = KernelEmitter<MlirKernelSource>; // NOLINT
} // namespace xla