[xla:gpu] NFC: Remove LMHLO op argument from EmitKernel #6224

PiperOrigin-RevId: 574213434
openxla · Oct 17, 2023 · dba73eb · dba73eb
1 parent fcdf7df
commit dba73eb
Show file tree

Hide file tree

Showing 12 changed files with 42 additions and 36 deletions.
diff --git a/xla/service/gpu/fusions/BUILD b/xla/service/gpu/fusions/BUILD
@@ -4,11 +4,13 @@ cc_library(
     hdrs = ["in_place_dynamic_update_slice.h"],
     deps = [
         ":fusion_emitter",
+        "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/llvm_ir:dynamic_update_slice_util",
         "//xla/service/llvm_ir:fused_ir_emitter",
+        "//xla/service/llvm_ir:ir_array",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
     ],
@@ -42,15 +44,13 @@ cc_library(
         "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu:target_util",
         "//xla/service/gpu:thunk",
-        "//xla/service/llvm_ir:buffer_assignment_util",
         "//xla/service/llvm_ir:ir_array",
         "//xla/service/llvm_ir:llvm_util",
-        "//xla/translate/mhlo_to_hlo:location_exporter",
         "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:MemRefDialect",
+        "@tsl//tsl/platform:errors",
     ],
 )
 
@@ -90,9 +90,12 @@ cc_library(
         "//xla/mlir_hlo:lhlo",
         "//xla/service:elemental_ir_emitter",
         "//xla/service/gpu:hlo_fusion_analysis",
+        "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:ir_emitter_context",
+        "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu:parallel_loop_emitter",
         "//xla/service/llvm_ir:fused_ir_emitter",
+        "//xla/service/llvm_ir:ir_array",
         "@llvm-project//llvm:ir_headers",
     ],
 )
@@ -189,6 +192,7 @@ cc_library(
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:ir_emitter_context",
+        "//xla/service/gpu:launch_dimensions",
         "//xla/service/gpu:target_util",
         "//xla/service/llvm_ir:fused_ir_emitter",
         "//xla/service/llvm_ir:ir_array",
@@ -203,6 +207,7 @@ cc_library(
     hdrs = ["input_slices.h"],
     deps = [
         ":fusion_emitter",
+        "//xla/hlo/ir:hlo",
         "//xla/service:elemental_ir_emitter",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",

diff --git a/xla/service/gpu/fusions/fusion_emitter.cc b/xla/service/gpu/fusions/fusion_emitter.cc
@@ -33,6 +33,7 @@ limitations under the License.
 #include "xla/service/gpu/target_util.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
+#include "tsl/platform/errors.h"
 
 namespace xla {
 namespace gpu {
@@ -199,9 +200,9 @@ StatusOr<FusionEmissionResult> KernelFusionEmitterBase::Emit(
               ir_emitter_context, suggested_kernel_name,
               kernel_arguments.args(), fusion_op.getInputBuffers().size(),
               launch_dims, builder);
-          TF_RETURN_IF_ERROR(EmitKernel(
-              ir_emitter_context, elemental_emitter, fusion_op, fusion,
-              launch_dims, std::move(inputs), std::move(outputs), builder, i));
+          TF_RETURN_IF_ERROR(EmitKernel(ir_emitter_context, elemental_emitter,
+                                        fusion, launch_dims, std::move(inputs),
+                                        std::move(outputs), builder, i));
           // TODO(jreiffers): Return shmem_bytes from EmitKernel when
           // converting the Triton emitters to this infrastructure.
           return KernelReuseCache::Entry{kernel->getName().str(), launch_dims,

diff --git a/xla/service/gpu/fusions/fusion_emitter.h b/xla/service/gpu/fusions/fusion_emitter.h
@@ -63,7 +63,6 @@ class KernelFusionEmitterBase : public FusionInterface {
  protected:
   virtual Status EmitKernel(IrEmitterContext& ir_emitter_context,
                             ElementalIrEmitter& elemental_emitter,
-                            mlir::lmhlo::FusionOp fusion_op,
                             const HloFusionInstruction& fusion,
                             const LaunchDimensions& launch_dims,
                             std::vector<llvm_ir::IrArray> inputs,

diff --git a/xla/service/gpu/fusions/in_place_dynamic_update_slice.cc b/xla/service/gpu/fusions/in_place_dynamic_update_slice.cc
@@ -19,9 +19,11 @@ limitations under the License.
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/IRBuilder.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/llvm_ir/dynamic_update_slice_util.h"
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
+#include "xla/service/llvm_ir/ir_array.h"
 
 namespace xla {
 namespace gpu {
@@ -35,10 +37,9 @@ StatusOr<LaunchDimensions> InPlaceDynamicUpdateSliceEmitter::launch_dimensions(
 
 Status InPlaceDynamicUpdateSliceEmitter::EmitKernel(
     IrEmitterContext& ir_emitter_context, ElementalIrEmitter& elemental_emitter,
-    mlir::lmhlo::FusionOp fusion_op, const HloFusionInstruction& fusion,
-    const LaunchDimensions& launch_dims, std::vector<llvm_ir::IrArray> inputs,
-    std::vector<llvm_ir::IrArray> outputs, llvm::IRBuilder<>* builder,
-    int kernel_index) const {
+    const HloFusionInstruction& fusion, const LaunchDimensions& launch_dims,
+    std::vector<llvm_ir::IrArray> inputs, std::vector<llvm_ir::IrArray> outputs,
+    llvm::IRBuilder<>* builder, int kernel_index) const {
   // In case a dynamic slice update's output is bitcasted, we need to ensure we
   // write to the output array using the shape and layout of the dynamic slice
   // update. This cast is known to be safe to do iff, in the case the output of

diff --git a/xla/service/gpu/fusions/in_place_dynamic_update_slice.h b/xla/service/gpu/fusions/in_place_dynamic_update_slice.h
@@ -60,7 +60,6 @@ class InPlaceDynamicUpdateSliceEmitter : public KernelFusionEmitterBase {
  protected:
   Status EmitKernel(IrEmitterContext& ir_emitter_context,
                     ElementalIrEmitter& elemental_emitter,
-                    mlir::lmhlo::FusionOp fusion_op,
                     const HloFusionInstruction& fusion,
                     const LaunchDimensions& launch_dims,
                     std::vector<llvm_ir::IrArray> inputs,

diff --git a/xla/service/gpu/fusions/input_slices.cc b/xla/service/gpu/fusions/input_slices.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include "xla/service/gpu/fusions/input_slices.h"
 
 #include "llvm/IR/IRBuilder.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/service/elemental_ir_emitter.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/parallel_loop_emitter.h"
@@ -158,10 +159,9 @@ StatusOr<LaunchDimensions> InputSlicesFusion::launch_dimensions(
 
 Status InputSlicesFusion::EmitKernel(
     IrEmitterContext& ir_emitter_context, ElementalIrEmitter& elemental_emitter,
-    mlir::lmhlo::FusionOp fusion_op, const HloFusionInstruction& fusion,
-    const LaunchDimensions& launch_dims, std::vector<llvm_ir::IrArray> inputs,
-    std::vector<llvm_ir::IrArray> outputs, llvm::IRBuilder<>* builder,
-    int kernel_index) const {
+    const HloFusionInstruction& fusion, const LaunchDimensions& launch_dims,
+    std::vector<llvm_ir::IrArray> inputs, std::vector<llvm_ir::IrArray> outputs,
+    llvm::IRBuilder<>* builder, int kernel_index) const {
   TF_ASSIGN_OR_RETURN(Shape element_shape,
                       GetConsistentInputShapeForRootSlices(
                           fusion.fused_instructions_computation()));
@@ -172,9 +172,9 @@ Status InputSlicesFusion::EmitKernel(
                    inputs, outputs, index, builder);
              },
              element_shape, launch_dims, builder)
-      .EmitLoop(llvm_ir::IrName(GetIrNameFromLoc(fusion_op.getLoc())),
-                GetIndexTypeForKernel(fusion_op, launch_dims.launch_bound(),
-                                      builder));
+      .EmitLoop(
+          fusion.name(),
+          GetIndexTypeForKernel(&fusion, launch_dims.launch_bound(), builder));
 }
 
 }  // namespace gpu

diff --git a/xla/service/gpu/fusions/input_slices.h b/xla/service/gpu/fusions/input_slices.h
@@ -40,7 +40,6 @@ class InputSlicesFusion : public KernelFusionEmitterBase {
  protected:
   Status EmitKernel(IrEmitterContext& ir_emitter_context,
                     ElementalIrEmitter& elemental_emitter,
-                    mlir::lmhlo::FusionOp fusion_op,
                     const HloFusionInstruction& fusion,
                     const LaunchDimensions& launch_dims,
                     std::vector<llvm_ir::IrArray> inputs,

diff --git a/xla/service/gpu/fusions/loop.cc b/xla/service/gpu/fusions/loop.cc
@@ -17,20 +17,23 @@ limitations under the License.
 #include <vector>
 
 #include "llvm/IR/IRBuilder.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/parallel_loop_emitter.h"
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
+#include "xla/service/llvm_ir/ir_array.h"
 
 namespace xla {
 namespace gpu {
 
 Status LoopFusion::EmitKernel(
     IrEmitterContext& ir_emitter_context, ElementalIrEmitter& elemental_emitter,
-    mlir::lmhlo::FusionOp fusion_op, const HloFusionInstruction& fusion,
-    const LaunchDimensions& launch_dims, std::vector<llvm_ir::IrArray> inputs,
-    std::vector<llvm_ir::IrArray> outputs, llvm::IRBuilder<>* builder,
-    int kernel_index) const {
+    const HloFusionInstruction& fusion, const LaunchDimensions& launch_dims,
+    std::vector<llvm_ir::IrArray> inputs, std::vector<llvm_ir::IrArray> outputs,
+    llvm::IRBuilder<>* builder, int kernel_index) const {
   FusedIrEmitter fused_emitter(elemental_emitter);
-  for (int i = 0; i < fusion_op.getInputBuffers().size(); i++) {
+  for (int i = 0; i < fusion.fused_parameters().size(); i++) {
     fused_emitter.BindGenerator(
         *fusion.fused_parameter(i), [&, i](llvm_ir::IrArray::Index index) {
           return inputs[i].EmitReadArrayElement(index, builder);
@@ -41,11 +44,11 @@ Status LoopFusion::EmitKernel(
       fused_emitter.GetGenerator(*fusion.fused_expression_root()));
 
   llvm::Type* index_type =
-      GetIndexTypeForKernel(fusion_op, launch_dims.launch_bound(), builder);
+      GetIndexTypeForKernel(&fusion, launch_dims.launch_bound(), builder);
 
   return ParallelLoopEmitter(element_generator, outputs, launch_dims, builder,
                              *analysis_.GetLoopFusionConfig())
-      .EmitLoop(GetIrNameFromLoc(fusion_op->getLoc()), index_type);
+      .EmitLoop(fusion.name(), index_type);
 }
 
 StatusOr<LaunchDimensions> LoopFusion::launch_dimensions(

diff --git a/xla/service/gpu/fusions/loop.h b/xla/service/gpu/fusions/loop.h
@@ -37,7 +37,6 @@ class LoopFusion : public KernelFusionEmitterBase {
  protected:
   Status EmitKernel(IrEmitterContext& ir_emitter_context,
                     ElementalIrEmitter& elemental_emitter,
-                    mlir::lmhlo::FusionOp fusion_op,
                     const HloFusionInstruction& fusion,
                     const LaunchDimensions& launch_dims,
                     std::vector<llvm_ir::IrArray> inputs,

diff --git a/xla/service/gpu/fusions/transpose.cc b/xla/service/gpu/fusions/transpose.cc
@@ -17,9 +17,11 @@ limitations under the License.
 #include <vector>
 
 #include "llvm/IR/IRBuilder.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/permutation_util.h"
 #include "xla/service/gpu/fusions/tiling_util.h"
 #include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/target_util.h"
 #include "xla/service/llvm_ir/fused_ir_emitter.h"
 #include "xla/service/llvm_ir/ir_array.h"
@@ -73,10 +75,9 @@ llvm_ir::IrArray::Index PermuteIndex(const llvm_ir::IrArray::Index& index,
 
 Status TransposeFusion::EmitKernel(
     IrEmitterContext& ir_emitter_context, ElementalIrEmitter& elemental_emitter,
-    mlir::lmhlo::FusionOp fusion_op, const HloFusionInstruction& fusion,
-    const LaunchDimensions& launch_dims, std::vector<llvm_ir::IrArray> inputs,
-    std::vector<llvm_ir::IrArray> outputs, llvm::IRBuilder<>* builder,
-    int kernel_index) const {
+    const HloFusionInstruction& fusion, const LaunchDimensions& launch_dims,
+    std::vector<llvm_ir::IrArray> inputs, std::vector<llvm_ir::IrArray> outputs,
+    llvm::IRBuilder<>* builder, int kernel_index) const {
   const auto& tiling_scheme = *analysis_.GetTransposeTilingScheme();
   const auto& hlo_roots = analysis_.fusion_roots();
   FusedIrEmitter fused_emitter(elemental_emitter);
@@ -233,7 +234,7 @@ Status TransposeFusion::EmitKernel(
       };
 
   llvm::Type* index_type =
-      GetIndexTypeForKernel(fusion_op, launch_dims.launch_bound(), builder);
+      GetIndexTypeForKernel(&fusion, launch_dims.launch_bound(), builder);
   return EmitTilingKernel(builder, tiling_scheme, index_type, tile_generator)
       .status();
 }

diff --git a/xla/service/gpu/fusions/transpose.h b/xla/service/gpu/fusions/transpose.h
@@ -61,7 +61,6 @@ class TransposeFusion : public KernelFusionEmitterBase {
  protected:
   Status EmitKernel(IrEmitterContext& ir_emitter_context,
                     ElementalIrEmitter& elemental_emitter,
-                    mlir::lmhlo::FusionOp fusion_op,
                     const HloFusionInstruction& fusion,
                     const LaunchDimensions& launch_dims,
                     std::vector<llvm_ir::IrArray> inputs,

diff --git a/xla/service/gpu/tests/fusion.hlo b/xla/service/gpu/tests/fusion.hlo
@@ -46,9 +46,9 @@ HloModule TestModule, is_scheduled=true
 // CHECK:         %[[VAL_36:.*]] = udiv i32 %[[VAL_29]], 802816
 // CHECK:         %[[VAL_37:.*]] = icmp ult i32 %[[VAL_5]], 102760448
 // CHECK:         br i1 %[[VAL_37]], label %[[VAL_38:.*]], label %[[VAL_39:.*]]
-// CHECK:       fusion_1.in_bounds-after:                         ; preds = %[[VAL_38]], %[[VAL_40:.*]]
+// CHECK:       fusion.1.in_bounds-after:                         ; preds = %[[VAL_38]], %[[VAL_40:.*]]
 // CHECK:         ret void
-// CHECK:       fusion_1.in_bounds-true:                          ; preds = %[[VAL_40]]
+// CHECK:       fusion.1.in_bounds-true:                          ; preds = %[[VAL_40]]
 // CHECK:         %[[VAL_41:.*]] = getelementptr inbounds [64 x float], ptr %[[VAL_42:.*]], i32 0, i32 %[[VAL_7]]
 // CHECK:         %[[VAL_43:.*]] = load float, ptr %[[VAL_41]], align 4, !invariant.load
 // CHECK:         %[[VAL_44:.*]] = getelementptr inbounds [64 x float], ptr %[[VAL_45:.*]], i32 0, i32 %[[VAL_7]]