openxla · copybara-service · Jan 13, 2025 · Jan 13, 2025
diff --git a/xla/backends/gpu/codegen/transforms/BUILD b/xla/backends/gpu/codegen/transforms/BUILD
@@ -65,13 +65,13 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/codegen/ir:xla_gpu",
+        "//xla/codegen/emitters:elemental_hlo_to_mlir",
         "//xla/codegen/ir:xla",
         "//xla/hlo/analysis:indexing_analysis",
         "//xla/mlir_hlo",
         "//xla/mlir_hlo:map_mhlo_to_scalar_op",
         "//xla/service/gpu:gpu_fusible",
         "//xla/service/gpu:ir_emission_utils",
-        "//xla/service/gpu/fusions/mlir:elemental_hlo_to_mlir",
         "//xla/service/gpu/llvm_gpu_backend",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:semantic_version",

diff --git a/xla/backends/gpu/codegen/transforms/lower_xla_gpu_to_scf.cc b/xla/backends/gpu/codegen/transforms/lower_xla_gpu_to_scf.cc
@@ -44,8 +44,8 @@ limitations under the License.
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "xla/backends/gpu/codegen/ir/xla_gpu_ops.h"
 #include "xla/backends/gpu/codegen/transforms/passes.h"
+#include "xla/codegen/emitters/elemental_hlo_to_mlir.h"
 #include "xla/hlo/analysis/indexing_map.h"
-#include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/util.h"
 
@@ -213,25 +213,23 @@ struct RewriteXlaGpuLoop : mlir::OpRewritePattern<LoopOp> {
 
     IndexingMap indexing_map = op.getIndexingMap();
     SmallVector<Value, 4> lbs, ubs, steps;
-    mlir_converter::GetLoopBoundsFromIndexingMap(b, indexing_map, &lbs, &ubs,
-                                                 &steps);
+    emitters::GetLoopBoundsFromIndexingMap(b, indexing_map, &lbs, &ubs, &steps);
     mlir::scf::LoopNest loop_nest = mlir::scf::buildLoopNest(
         b, loc, lbs, ubs, steps, op.getInits(),
         [&](OpBuilder& nested_builder, Location loc, ValueRange symbol_values,
             ValueRange iter_args) -> mlir::scf::ValueVector {
           mlir::ImplicitLocOpBuilder nested_b(loc, nested_builder);
-          auto is_in_bounds = mlir_converter::CheckConstraints(
+          auto is_in_bounds = emitters::CheckConstraints(
               indexing_map, op.getDims(), symbol_values, nested_b);
           auto if_op = nested_b.create<mlir::scf::IfOp>(
               is_in_bounds,
               [&](OpBuilder& then_builder, Location then_loc) -> void {
                 ImplicitLocOpBuilder then_b(then_loc, then_builder);
                 mlir::IRMapping mapping;
                 mapping.map(op.getInductionVars(), symbol_values);
-                mapping.map(
-                    op.getIndexingMapResults(),
-                    mlir_converter::ApplyIndexing(indexing_map, op.getDims(),
-                                                  symbol_values, then_b));
+                mapping.map(op.getIndexingMapResults(),
+                            emitters::ApplyIndexing(indexing_map, op.getDims(),
+                                                    symbol_values, then_b));
                 mapping.map(op.getRegionIterArgs(), iter_args);
                 mlir::Block* old_block = op.getBody();
                 for (auto& old_op : old_block->without_terminator()) {

diff --git a/xla/codegen/emitters/BUILD b/xla/codegen/emitters/BUILD
@@ -0,0 +1,171 @@
+load("//xla:xla.bzl", "xla_cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+cc_library(
+    name = "computation_partitioner",
+    srcs = ["computation_partitioner.cc"],
+    hdrs = ["computation_partitioner.h"],
+    deps = [
+        ":type_util",
+        "//xla:shape_util",
+        "//xla:util",
+        "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/ir:hlo",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:node_hash_map",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:DataLayoutInterfaces",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+xla_cc_test(
+    name = "computation_partitioner_test",
+    srcs = ["computation_partitioner_test.cc"],
+    deps = [
+        ":computation_partitioner",
+        "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/ir:hlo",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "elemental_hlo_to_mlir",
+    srcs = ["elemental_hlo_to_mlir.cc"],
+    hdrs = ["elemental_hlo_to_mlir.h"],
+    deps = [
+        ":computation_partitioner",
+        ":type_util",
+        "//xla:comparison_util",
+        "//xla:shape_util",
+        "//xla:status_macros",
+        "//xla:xla_data_proto_cc",
+        "//xla/codegen/ir:xla",
+        "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/translate/hlo_to_mhlo:hlo_utils",
+        "//xla/hlo/utils:hlo_traversal",
+        "//xla/mlir_hlo",
+        "//xla/mlir_hlo:map_mhlo_to_scalar_op",
+        "//xla/service:algorithm_util",
+        "//xla/stream_executor:device_description",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:node_hash_map",
+        "@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/strings:str_format",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AffineUtils",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:DataLayoutInterfaces",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:VectorDialect",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+xla_cc_test(
+    name = "elemental_hlo_to_mlir_test",
+    srcs = ["elemental_hlo_to_mlir_test.cc"],
+    deps = [
+        ":computation_partitioner",
+        ":elemental_hlo_to_mlir",
+        "//xla:status_macros",
+        "//xla/backends/gpu/codegen/ir:xla_gpu",
+        "//xla/codegen/ir:xla",
+        "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/parser:hlo_parser",
+        "//xla/hlo/testlib:filecheck",
+        "//xla/mlir_hlo",
+        "//xla/service/llvm_ir:llvm_util",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:xla_internal_test_main",
+        "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:AsmParser",
+        "@llvm-project//mlir:DLTIDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:MathDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:Transforms",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "type_util",
+    srcs = ["type_util.cc"],
+    hdrs = ["type_util.h"],
+    deps = [
+        "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/translate/hlo_to_mhlo:hlo_utils",
+        "//xla/mlir/utils:type_util",
+        "@com_google_absl//absl/log:check",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+xla_cc_test(
+    name = "type_util_test",
+    srcs = ["type_util_test.cc"],
+    deps = [
+        ":type_util",
+        "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+    ],
+)
diff --git a/...u/fusions/mlir/computation_partitioner.cc → ...degen/emitters/computation_partitioner.cc b/...u/fusions/mlir/computation_partitioner.cc → ...degen/emitters/computation_partitioner.cc
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-#include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
+#include "xla/codegen/emitters/computation_partitioner.h"
 
 #include <cstdint>
 #include <functional>
@@ -44,19 +44,18 @@ limitations under the License.
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Support/LLVM.h"
+#include "xla/codegen/emitters/type_util.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/service/gpu/fusions/mlir/type_util.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 
 namespace xla {
-namespace gpu {
-namespace mlir_converter {
+namespace emitters {
 namespace {
 
 int Arity(const Shape& shape) {
@@ -443,6 +442,5 @@ mlir::func::FuncOp CreateSubgraphMlirFunction(
   return func_op;
 }
 
-}  // namespace mlir_converter
-}  // namespace gpu
+}  // namespace emitters
 }  // namespace xla
diff --git a/...pu/fusions/mlir/computation_partitioner.h → ...odegen/emitters/computation_partitioner.h b/...pu/fusions/mlir/computation_partitioner.h → ...odegen/emitters/computation_partitioner.h
@@ -12,8 +12,8 @@ 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_SERVICE_GPU_FUSIONS_MLIR_COMPUTATION_PARTITIONER_H_
-#define XLA_SERVICE_GPU_FUSIONS_MLIR_COMPUTATION_PARTITIONER_H_
+#ifndef XLA_CODEGEN_EMITTERS_COMPUTATION_PARTITIONER_H_
+#define XLA_CODEGEN_EMITTERS_COMPUTATION_PARTITIONER_H_
 
 #include <cstdint>
 #include <functional>
@@ -33,8 +33,7 @@ limitations under the License.
 #include "xla/util.h"
 
 namespace xla {
-namespace gpu {
-namespace mlir_converter {
+namespace emitters {
 
 struct EpilogueSpecification {
   // Creates an epilogue with output indices matching the given root's shape.
@@ -206,8 +205,7 @@ mlir::func::FuncOp CreateSubgraphMlirFunction(
     const PartitionedComputation::Subgraph& subgraph,
     mlir::ImplicitLocOpBuilder& b);
 
-}  // namespace mlir_converter
-}  // namespace gpu
+}  // namespace emitters
 }  // namespace xla
 
-#endif  // XLA_SERVICE_GPU_FUSIONS_MLIR_COMPUTATION_PARTITIONER_H_
+#endif  // XLA_CODEGEN_EMITTERS_COMPUTATION_PARTITIONER_H_
diff --git a/...ions/mlir/computation_partitioner_test.cc → .../emitters/computation_partitioner_test.cc b/...ions/mlir/computation_partitioner_test.cc → .../emitters/computation_partitioner_test.cc
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-#include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
+#include "xla/codegen/emitters/computation_partitioner.h"
 
 #include <string>
 
@@ -31,8 +31,7 @@ limitations under the License.
 #include "xla/tests/hlo_test_base.h"
 
 namespace xla {
-namespace gpu {
-namespace mlir_converter {
+namespace emitters {
 namespace {
 
 using ::testing::ElementsAre;
@@ -333,6 +332,5 @@ TEST_F(ComputationPartitionerTest, SubgraphSignatures) {
 }
 
 }  // namespace
-}  // namespace mlir_converter
-}  // namespace gpu
+}  // namespace emitters
 }  // namespace xla
diff --git a/...gpu/fusions/mlir/elemental_hlo_to_mlir.cc → ...codegen/emitters/elemental_hlo_to_mlir.cc b/...gpu/fusions/mlir/elemental_hlo_to_mlir.cc → ...codegen/emitters/elemental_hlo_to_mlir.cc
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-#include "xla/service/gpu/fusions/mlir/elemental_hlo_to_mlir.h"
+#include "xla/codegen/emitters/elemental_hlo_to_mlir.h"
 
 #include <cstddef>
 #include <cstdint>
@@ -61,6 +61,8 @@ limitations under the License.
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Support/LLVM.h"
+#include "xla/codegen/emitters/computation_partitioner.h"
+#include "xla/codegen/emitters/type_util.h"
 #include "xla/codegen/ir/xla_ops.h"
 #include "xla/comparison_util.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
@@ -75,17 +77,14 @@ limitations under the License.
 #include "xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h"
 #include "xla/primitive_util.h"
 #include "xla/service/algorithm_util.h"
-#include "xla/service/gpu/fusions/mlir/computation_partitioner.h"
-#include "xla/service/gpu/fusions/mlir/type_util.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
 #include "tsl/platform/statusor.h"
 
 namespace xla {
-namespace gpu {
-namespace mlir_converter {
+namespace emitters {
 namespace {
 
 using llvm::SmallVector;
@@ -1481,8 +1480,8 @@ ValueRange EmitLoopNestImpl(
                 ValueRange symbol_values,
                 ValueRange iter_args) -> scf::ValueVector {
     ImplicitLocOpBuilder nested_b(loc, nested_builder);
-    auto is_in_bounds = mlir_converter::CheckConstraints(
-        indexing_map, dim_values, symbol_values, nested_b);
+    auto is_in_bounds =
+        CheckConstraints(indexing_map, dim_values, symbol_values, nested_b);
     auto if_op = nested_b.create<scf::IfOp>(
         is_in_bounds,
         [&](OpBuilder& then_builder, Location then_loc) -> void {
@@ -1701,6 +1700,5 @@ SmallVector<Value, 2> InlineBlock(OpBuilder& builder, Block& src_block,
   return mapped_results;
 }
 
-}  // namespace mlir_converter
-}  // namespace gpu
+}  // namespace emitters
 }  // namespace xla