[xla:gpu] Remove dependencies on LHLO from ir_emission_utils

PiperOrigin-RevId: 609899219
openxla · Feb 24, 2024 · 4ccfe33 · 4ccfe33
1 parent c1f306e
commit 4ccfe33
Show file tree

Hide file tree

Showing 8 changed files with 10 additions and 584 deletions.
diff --git a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD
@@ -1324,7 +1324,6 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/mlir_hlo",
-        "//xla/mlir_hlo:lhlo",
         "//xla/service:buffer_assignment",
         "//xla/service:hlo_parser",
         "//xla/service/llvm_ir:buffer_assignment_util",
@@ -1338,6 +1337,7 @@ cc_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Core",
@@ -1366,7 +1366,6 @@ xla_cc_test(
         "//xla:util",
         "//xla/hlo/ir:hlo",
         "//xla/mlir_hlo",
-        "//xla/mlir_hlo:lhlo",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",  # fixdeps: keep
         "//xla/translate/hlo_to_mhlo:hlo_utils",

diff --git a/xla/service/gpu/fusions/fusions.cc b/xla/service/gpu/fusions/fusions.cc
@@ -76,45 +76,6 @@ bool IsDynamicUpdateSliceFusion(const HloFusionAnalysis& analysis) {
 
 }  // namespace
 
-std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
-LmhloFusionInfo::GetCopyFusion() const {
-  auto params = GetHloOperands(fusion_op_);
-  auto outputs = GetHloOutputs(fusion_op_);
-  std::vector<mlir::Value> srcs;
-  srcs.reserve(outputs.size());
-
-  for (auto* root : analysis().fusion_roots()) {
-    if (root->opcode() != HloOpcode::kCopy ||
-        root->operand(0)->opcode() != HloOpcode::kParameter ||
-        !LayoutUtil::Equal(root->operand(0)->shape().layout(),
-                           root->shape().layout())) {
-      return std::nullopt;
-    }
-
-    mlir::Value src = params[root->operand(0)->parameter_number()];
-    if (!GetAllocationSlice(src, allocations_).ok()) return std::nullopt;
-
-    srcs.emplace_back(src);
-  }
-
-  auto dsts = std::vector<mlir::Value>(outputs.begin(), outputs.end());
-  DCHECK(srcs.size() == dsts.size());
-  std::vector<BufferAllocation::Slice> src_buffers;
-  std::vector<BufferAllocation::Slice> dst_buffers;
-  for (int i = 0; i < srcs.size(); ++i) {
-    TF_ASSIGN_OR_RETURN(BufferAllocation::Slice src_buffer,
-                        GetAllocationSlice(srcs[i], allocations_));
-    src_buffers.push_back(src_buffer);
-    TF_ASSIGN_OR_RETURN(BufferAllocation::Slice dst_buffer,
-                        GetAllocationSlice(dsts[i], allocations_));
-    dst_buffers.push_back(dst_buffer);
-  }
-
-  return std::make_unique<MemcpyFusion>(std::move(src_buffers),
-                                        std::move(dst_buffers), std::move(srcs),
-                                        std::move(dsts));
-}
-
 std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
 HloFusionInfo::GetCopyFusion() const {
   std::vector<BufferAllocation::Slice> src_buffers;
@@ -154,10 +115,6 @@ HloFusionInfo::GetCopyFusion() const {
                                         /*dsts=*/std::vector<mlir::Value>());
 }
 
-bool LmhloFusionInfo::CanEmitDynamicUpdateSliceInPlace() const {
-  return CanEmitFusedDynamicUpdateSliceInPlaceForGpu(fusion_op_, allocations_);
-}
-
 bool HloFusionInfo::CanEmitDynamicUpdateSliceInPlace() const {
   auto ret = CanEmitFusedDynamicUpdateSliceInPlaceForGpu(
       instr_, buffer_assignment_, analysis().fusion_roots());

diff --git a/xla/service/gpu/fusions/fusions.h b/xla/service/gpu/fusions/fusions.h
@@ -18,12 +18,10 @@ limitations under the License.
 #include <memory>
 #include <optional>
 
-#include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
-#include "xla/statusor.h"
 
 namespace xla {
 namespace gpu {
@@ -51,24 +49,6 @@ class FusionInfo {
   const HloFusionAnalysis& analysis_;
 };
 
-class LmhloFusionInfo : public FusionInfo {
- public:
-  LmhloFusionInfo(const HloFusionAnalysis& analysis,
-                  mlir::lmhlo::FusionOp fusion_op,
-                  absl::Span<const BufferAllocation* const> allocations)
-      : FusionInfo(analysis),
-        fusion_op_(fusion_op),
-        allocations_(allocations) {}
-
-  bool CanEmitDynamicUpdateSliceInPlace() const override;
-  std::optional<absl::StatusOr<std::unique_ptr<FusionInterface>>>
-  GetCopyFusion() const override;
-
- private:
-  mlir::lmhlo::FusionOp fusion_op_;
-  absl::Span<const BufferAllocation* const> allocations_;
-};
-
 class HloFusionInfo : public FusionInfo {
  public:
   HloFusionInfo(const HloFusionAnalysis& analysis,

diff --git a/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h b/xla/service/gpu/fusions/mlir/mlir_fusion_emitter.h
@@ -20,8 +20,10 @@ limitations under the License.
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/fusions/fusion_emitter.h"