Move triton codegen to xla/backends/gpu/codegen/triton

PiperOrigin-RevId: 715389092

Author: akuegel

xla/service/gpu/fusions/triton/BUILD renamed to xla/backends/gpu/codegen/triton/BUILD

@@ -105,18 +105,18 @@ cc_library(
 )
 
 cc_library(
-    name = "triton_fusion_emitter",
+    name = "fusion_emitter",
     srcs = if_gpu_is_configured(
-        ["triton_fusion_emitter.cc"],
-        ["triton_fusion_emitter_stub.cc"],
+        ["fusion_emitter.cc"],
+        ["fusion_emitter_stub.cc"],
     ),
-    hdrs = ["triton_fusion_emitter.h"],
+    hdrs = ["fusion_emitter.h"],
     deps = [
         ":compilation_pipeline",
         ":emitter_helpers",
+        ":fusion_emitter_legacy_matmul",
         ":passes",
-        ":triton_fusion_emitter_legacy_matmul",
-        ":triton_support",
+        ":support",
         ":xla_triton",
         ":xla_triton_passes",
         "//xla:autotuning_proto_cc",
@@ -209,12 +209,12 @@ cc_library(
 )
 
 cc_library(
-    name = "triton_fusion_emitter_legacy_matmul",
+    name = "fusion_emitter_legacy_matmul",
     srcs = if_gpu_is_configured(
-        ["triton_fusion_emitter_legacy_matmul.cc"],
-        ["triton_fusion_emitter_legacy_matmul_stub.cc"],
+        ["fusion_emitter_legacy_matmul.cc"],
+        ["fusion_emitter_legacy_matmul_stub.cc"],
     ),
-    hdrs = ["triton_fusion_emitter_legacy_matmul.h"],
+    hdrs = ["fusion_emitter_legacy_matmul.h"],
     deps = [
         ":emitter_helpers",
         ":xla_triton",
@@ -271,16 +271,16 @@ cc_library(
 )
 
 cc_library(
-    name = "triton_fusion_emitter_stub_for_testing",
+    name = "fusion_emitter_stub_for_testing",
     srcs = [
         "compilation_pipeline_stub.cc",
-        "triton_fusion_emitter_legacy_matmul_stub.cc",
-        "triton_fusion_emitter_stub.cc",
+        "fusion_emitter_legacy_matmul_stub.cc",
+        "fusion_emitter_stub.cc",
     ],
     hdrs = [
         "compilation_pipeline.h",
-        "triton_fusion_emitter.h",
-        "triton_fusion_emitter_legacy_matmul.h",
+        "fusion_emitter.h",
+        "fusion_emitter_legacy_matmul.h",
     ],
     deps = [
         "//xla:autotuning_proto_cc",
@@ -307,10 +307,10 @@ cc_library(
 )
 
 xla_cc_test(
-    name = "triton_fusion_emitter_stub_test",
-    srcs = ["triton_fusion_emitter_stub_test.cc"],
+    name = "fusion_emitter_stub_test",
+    srcs = ["fusion_emitter_stub_test.cc"],
     deps = [
-        ":triton_fusion_emitter_stub_for_testing",
+        ":fusion_emitter_stub_for_testing",
         "//xla:literal",
         "//xla:literal_util",
         "//xla/hlo/ir:hlo",
@@ -417,7 +417,7 @@ cc_library(
 )
 
 td_library(
-    name = "xla_triton_td_files",
+    name = "xla_td_files",
     srcs = glob(["*.td"]),
     includes = ["."],
     deps = [
@@ -441,7 +441,7 @@ gentbl_cc_library(
     ],
     tblgen = "@llvm-project//mlir:mlir-tblgen",
     td_file = "xla_triton_dialect.td",
-    deps = [":xla_triton_td_files"],
+    deps = [":xla_td_files"],
 )
 
 gentbl_cc_library(
@@ -460,7 +460,7 @@ gentbl_cc_library(
     tblgen = "@llvm-project//mlir:mlir-tblgen",
     td_file = "xla_triton_ops.td",
     deps = [
-        ":xla_triton_td_files",
+        ":xla_td_files",
         "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:SideEffectInterfacesTdFiles",
@@ -490,7 +490,7 @@ gentbl_cc_library(
     tblgen = "@llvm-project//mlir:mlir-tblgen",
     td_file = "xla_triton_attrs.td",
     deps = [
-        ":xla_triton_td_files",
+        ":xla_td_files",
         "@triton//:td_files",
     ],
 )
@@ -516,11 +516,11 @@ cc_library(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_deviceless_test",
-    srcs = ["triton_fusion_emitter_deviceless_test.cc"],
+    name = "fusion_emitter_deviceless_test",
+    srcs = ["fusion_emitter_deviceless_test.cc"],
     backends = ["gpu"],
     deps = [
-        ":triton_fusion_emitter",
+        ":fusion_emitter",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -539,8 +539,8 @@ xla_test(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_device_legacy_test",
-    srcs = if_gpu_is_configured(["triton_fusion_emitter_device_legacy_test.cc"]),
+    name = "fusion_emitter_device_legacy_test",
+    srcs = if_gpu_is_configured(["fusion_emitter_device_legacy_test.cc"]),
     # TODO(b/372714955): Fix the memory leak!
     backend_args = if_google(
         {
@@ -560,8 +560,8 @@ xla_test(
         "no_mac",
     ],
     deps = [
-        ":triton_fusion_emitter",
-        ":triton_test_utils",
+        ":fusion_emitter",
+        ":test_utils",
         "//xla:autotuning_proto_cc",
         "//xla:error_spec",
         "//xla:xla_proto_cc",
@@ -592,8 +592,8 @@ xla_test(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_int4_device_test",
-    srcs = if_gpu_is_configured(["triton_fusion_emitter_int4_device_test.cc"]),
+    name = "fusion_emitter_int4_device_test",
+    srcs = if_gpu_is_configured(["fusion_emitter_int4_device_test.cc"]),
     # TODO(b/372714955): Fix the memory leak!
     backend_args = if_google(
         {
@@ -651,7 +651,7 @@ xla_test(
     ],
     deps = [
         ":kernel_name_tracer",
-        ":triton_test_utils",
+        ":test_utils",
         "//xla:autotuning_proto_cc",
         "//xla:error_spec",
         "//xla:literal",
@@ -682,8 +682,8 @@ xla_test(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_device_test",
-    srcs = if_gpu_is_configured(["triton_fusion_emitter_device_test.cc"]),
+    name = "fusion_emitter_device_test",
+    srcs = if_gpu_is_configured(["fusion_emitter_device_test.cc"]),
     backends = [
         "gpu_a100",
         "gpu_h100",
@@ -694,8 +694,8 @@ xla_test(
         "no_mac",
     ],
     deps = [
-        ":triton_fusion_emitter",
-        ":triton_test_utils",
+        ":fusion_emitter",
+        ":test_utils",
         "//xla:autotuning_proto_cc",
         "//xla:error_spec",
         "//xla:shape_util",
@@ -757,12 +757,12 @@ cc_library(
 )
 
 cc_library(
-    name = "triton_test_utils",
+    name = "test_utils",
     testonly = True,
-    srcs = ["triton_test_utils.cc"],
-    hdrs = ["triton_test_utils.h"],
+    srcs = ["test_utils.cc"],
+    hdrs = ["test_utils.h"],
     deps = [
-        ":triton_fusion_emitter",
+        ":fusion_emitter",
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla/hlo/ir:hlo",
@@ -795,10 +795,10 @@ cc_library(
 )
 
 xla_cc_test(
-    name = "triton_fusion_emitter_mem_utils_test",
-    srcs = if_cuda_is_configured(["triton_fusion_emitter_mem_utils_test.cc"]),
+    name = "fusion_emitter_mem_utils_test",
+    srcs = if_cuda_is_configured(["fusion_emitter_mem_utils_test.cc"]),
     deps = [
-        ":triton_fusion_emitter",
+        ":fusion_emitter",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_traversal",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -825,8 +825,8 @@ xla_cc_test(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_large_test",
-    srcs = if_gpu_is_configured(["triton_fusion_emitter_large_test.cc"]),
+    name = "fusion_emitter_large_test",
+    srcs = if_gpu_is_configured(["fusion_emitter_large_test.cc"]),
     backends = [
         "gpu_a100",
         "gpu_h100",
@@ -853,8 +853,8 @@ xla_test(
 )
 
 xla_test(
-    name = "triton_fusion_emitter_parametrized_test",
-    srcs = if_gpu_is_configured(["triton_fusion_emitter_parametrized_test.cc"]),
+    name = "fusion_emitter_parametrized_test",
+    srcs = if_gpu_is_configured(["fusion_emitter_parametrized_test.cc"]),
     backends = [
         "gpu_a100",
         "gpu_h100",
@@ -864,8 +864,8 @@ xla_test(
     shard_count = 10,
     tags = ["no_mac"],
     deps = [
-        ":triton_support",
-        ":triton_test_utils",
+        ":support",
+        ":test_utils",
         "//xla:comparison_util",
         "//xla:error_spec",
         "//xla:xla_data_proto_cc",
@@ -881,14 +881,14 @@ xla_test(
 )
 
 cc_library(
-    name = "triton_support",
+    name = "support",
     srcs = [
-        "triton_support.cc",
-        "triton_support_legacy.cc",
+        "support.cc",
+        "support_legacy.cc",
     ],
     hdrs = [
-        "triton_support.h",
-        "triton_support_legacy.h",
+        "support.h",
+        "support_legacy.h",
     ],
     deps = [
         "//xla:shape_util",
@@ -911,16 +911,16 @@ cc_library(
 )
 
 xla_cc_test(
-    name = "triton_support_test",
-    srcs = ["triton_support_test.cc"],
+    name = "support_test",
+    srcs = ["support_test.cc"],
     shard_count = 25,
     # TODO(b/353912594): this test does not need to run on GPU, but it is broken on CPU in OSS.
     # Force it to run on GPU temporarily in order to get important OSS coverage.
     tags = ["gpu"],
     deps = [
-        ":triton_fusion_emitter",
-        ":triton_support",
-        ":triton_test_utils",
+        ":fusion_emitter",
+        ":support",
+        ":test_utils",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
@@ -941,8 +941,8 @@ xla_cc_test(
 )
 
 xla_test(
-    name = "triton_support_legacy_test",
-    srcs = if_gpu_is_configured(["triton_support_legacy_test.cc"]),
+    name = "support_legacy_test",
+    srcs = if_gpu_is_configured(["support_legacy_test.cc"]),
     backends = [
         "gpu_a100",
         "gpu_h100",
@@ -951,10 +951,10 @@ xla_test(
     ],
     tags = ["no_mac"],
     deps = [
+        ":fusion_emitter",
         ":kernel_name_tracer",
-        ":triton_fusion_emitter",
-        ":triton_support",
-        ":triton_test_utils",
+        ":support",
+        ":test_utils",
         "//xla:error_spec",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",

xla/service/gpu/fusions/triton/compilation_pipeline.h renamed to xla/backends/gpu/codegen/triton/compilation_pipeline.h

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef XLA_SERVICE_GPU_FUSIONS_TRITON_COMPILATION_PIPELINE_H_
-#define XLA_SERVICE_GPU_FUSIONS_TRITON_COMPILATION_PIPELINE_H_
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TRITON_COMPILATION_PIPELINE_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TRITON_COMPILATION_PIPELINE_H_
 
 #include <string>
 
@@ -48,4 +48,4 @@ absl::Status CreateTritonPipeline(
 }  // namespace gpu
 }  // namespace xla
 
-#endif  // XLA_SERVICE_GPU_FUSIONS_TRITON_COMPILATION_PIPELINE_H_
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TRITON_COMPILATION_PIPELINE_H_

xla/service/gpu/fusions/triton/compilation_pipeline_cuda.cc renamed to xla/backends/gpu/codegen/triton/compilation_pipeline_cuda.cc

@@ -25,7 +25,7 @@ limitations under the License.
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
-#include "xla/service/gpu/fusions/triton/xla_triton_passes.h"
+#include "xla/backends/gpu/codegen/triton/xla_triton_passes.h"
 #include "xla/service/gpu/llvm_gpu_backend/nvptx_libdevice_path.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/stream_executor/device_description.h"

xla/service/gpu/fusions/triton/compilation_pipeline_rocm.cc renamed to xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc

@@ -17,7 +17,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "mlir/Pass/PassManager.h"
-#include "xla/service/gpu/fusions/triton/compilation_pipeline.h"
+#include "xla/backends/gpu/codegen/triton/compilation_pipeline.h"
 
 namespace xla {
 namespace gpu {

xla/service/gpu/fusions/triton/compilation_pipeline_stub.cc renamed to xla/backends/gpu/codegen/triton/compilation_pipeline_stub.cc

@@ -38,6 +38,8 @@ limitations under the License.
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "xla/autotuning.pb.h"
+#include "xla/backends/gpu/codegen/triton/kernel_name_tracer.h"
+#include "xla/backends/gpu/codegen/triton/test_utils.h"
 #include "xla/error_spec.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -47,8 +49,6 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/service/dump.h"
 #include "xla/service/gpu/backend_configs.pb.h"
-#include "xla/service/gpu/fusions/triton/kernel_name_tracer.h"
-#include "xla/service/gpu/fusions/triton/triton_test_utils.h"
 #include "xla/service/gpu/tests/gpu_codegen_test.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/stream_executor/device_description.h"

xla/service/gpu/fusions/triton/dot_algorithms_test.cc renamed to xla/backends/gpu/codegen/triton/dot_algorithms_test.cc

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/service/gpu/fusions/triton/emitter_helpers.h"
+#include "xla/backends/gpu/codegen/triton/emitter_helpers.h"
 
 #include <cstdint>
 #include <variant>