Skip to content

Commit 4ccfe33

Browse files
ezhulenevcopybara-github
authored andcommitted
[xla:gpu] Remove dependencies on LHLO from ir_emission_utils
PiperOrigin-RevId: 609899219
1 parent c1f306e commit 4ccfe33

8 files changed

+10
-584
lines changed

xla/service/gpu/BUILD

+1-2
Original file line numberDiff line numberDiff line change
@@ -1324,7 +1324,6 @@ cc_library(
13241324
"//xla:xla_data_proto_cc",
13251325
"//xla/hlo/ir:hlo",
13261326
"//xla/mlir_hlo",
1327-
"//xla/mlir_hlo:lhlo",
13281327
"//xla/service:buffer_assignment",
13291328
"//xla/service:hlo_parser",
13301329
"//xla/service/llvm_ir:buffer_assignment_util",
@@ -1338,6 +1337,7 @@ cc_library(
13381337
"@com_google_absl//absl/log",
13391338
"@com_google_absl//absl/log:check",
13401339
"@com_google_absl//absl/status",
1340+
"@com_google_absl//absl/status:statusor",
13411341
"@com_google_absl//absl/strings",
13421342
"@com_google_absl//absl/types:span",
13431343
"@llvm-project//llvm:Core",
@@ -1366,7 +1366,6 @@ xla_cc_test(
13661366
"//xla:util",
13671367
"//xla/hlo/ir:hlo",
13681368
"//xla/mlir_hlo",
1369-
"//xla/mlir_hlo:lhlo",
13701369
"//xla/tests:hlo_test_base",
13711370
"//xla/tests:xla_internal_test_main", # fixdeps: keep
13721371
"//xla/translate/hlo_to_mhlo:hlo_utils",

xla/service/gpu/fusions/fusions.cc

-43
Original file line numberDiff line numberDiff line change
@@ -76,45 +76,6 @@ bool IsDynamicUpdateSliceFusion(const HloFusionAnalysis& analysis) {
7676

7777
} // namespace
7878

79-
std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
80-
LmhloFusionInfo::GetCopyFusion() const {
81-
auto params = GetHloOperands(fusion_op_);
82-
auto outputs = GetHloOutputs(fusion_op_);
83-
std::vector<mlir::Value> srcs;
84-
srcs.reserve(outputs.size());
85-
86-
for (auto* root : analysis().fusion_roots()) {
87-
if (root->opcode() != HloOpcode::kCopy ||
88-
root->operand(0)->opcode() != HloOpcode::kParameter ||
89-
!LayoutUtil::Equal(root->operand(0)->shape().layout(),
90-
root->shape().layout())) {
91-
return std::nullopt;
92-
}
93-
94-
mlir::Value src = params[root->operand(0)->parameter_number()];
95-
if (!GetAllocationSlice(src, allocations_).ok()) return std::nullopt;
96-
97-
srcs.emplace_back(src);
98-
}
99-
100-
auto dsts = std::vector<mlir::Value>(outputs.begin(), outputs.end());
101-
DCHECK(srcs.size() == dsts.size());
102-
std::vector<BufferAllocation::Slice> src_buffers;
103-
std::vector<BufferAllocation::Slice> dst_buffers;
104-
for (int i = 0; i < srcs.size(); ++i) {
105-
TF_ASSIGN_OR_RETURN(BufferAllocation::Slice src_buffer,
106-
GetAllocationSlice(srcs[i], allocations_));
107-
src_buffers.push_back(src_buffer);
108-
TF_ASSIGN_OR_RETURN(BufferAllocation::Slice dst_buffer,
109-
GetAllocationSlice(dsts[i], allocations_));
110-
dst_buffers.push_back(dst_buffer);
111-
}
112-
113-
return std::make_unique<MemcpyFusion>(std::move(src_buffers),
114-
std::move(dst_buffers), std::move(srcs),
115-
std::move(dsts));
116-
}
117-
11879
std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
11980
HloFusionInfo::GetCopyFusion() const {
12081
std::vector<BufferAllocation::Slice> src_buffers;
@@ -154,10 +115,6 @@ HloFusionInfo::GetCopyFusion() const {
154115
/*dsts=*/std::vector<mlir::Value>());
155116
}
156117

157-
bool LmhloFusionInfo::CanEmitDynamicUpdateSliceInPlace() const {
158-
return CanEmitFusedDynamicUpdateSliceInPlaceForGpu(fusion_op_, allocations_);
159-
}
160-
161118
bool HloFusionInfo::CanEmitDynamicUpdateSliceInPlace() const {
162119
auto ret = CanEmitFusedDynamicUpdateSliceInPlaceForGpu(
163120
instr_, buffer_assignment_, analysis().fusion_roots());

xla/service/gpu/fusions/fusions.h

-20
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,10 @@ limitations under the License.
1818
#include <memory>
1919
#include <optional>
2020

21-
#include "absl/types/span.h"
2221
#include "xla/hlo/ir/hlo_instructions.h"
2322
#include "xla/service/buffer_assignment.h"
2423
#include "xla/service/gpu/fusions/fusion_emitter.h"
2524
#include "xla/service/gpu/hlo_fusion_analysis.h"
26-
#include "xla/statusor.h"
2725

2826
namespace xla {
2927
namespace gpu {
@@ -51,24 +49,6 @@ class FusionInfo {
5149
const HloFusionAnalysis& analysis_;
5250
};
5351

54-
class LmhloFusionInfo : public FusionInfo {
55-
public:
56-
LmhloFusionInfo(const HloFusionAnalysis& analysis,
57-
mlir::lmhlo::FusionOp fusion_op,
58-
absl::Span<const BufferAllocation* const> allocations)
59-
: FusionInfo(analysis),
60-
fusion_op_(fusion_op),
61-
allocations_(allocations) {}
62-
63-
bool CanEmitDynamicUpdateSliceInPlace() const override;
64-
std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
65-
GetCopyFusion() const override;
66-
67-
private:
68-
mlir::lmhlo::FusionOp fusion_op_;
69-
absl::Span<const BufferAllocation* const> allocations_;
70-
};
71-
7252
class HloFusionInfo : public FusionInfo {
7353
public:
7454
HloFusionInfo(const HloFusionAnalysis& analysis,

xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h

+2
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@ limitations under the License.
2020
#include "llvm/IR/LLVMContext.h"
2121
#include "llvm/IR/Module.h"
2222
#include "mlir/IR/AffineMap.h" // from @llvm-project
23+
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
2324
#include "mlir/IR/ImplicitLocOpBuilder.h" // from @llvm-project
2425
#include "mlir/IR/MLIRContext.h" // from @llvm-project
26+
#include "mlir/IR/OwningOpRef.h" // from @llvm-project
2527
#include "mlir/IR/Value.h" // from @llvm-project
2628
#include "xla/hlo/ir/hlo_instructions.h"
2729
#include "xla/service/gpu/fusions/fusion_emitter.h"

0 commit comments

Comments
 (0)