Reverts 67325a0

PiperOrigin-RevId: 675043690
openxla · Sep 16, 2024 · 1765b13 · 1765b13
1 parent e80c23e
commit 1765b13
Show file tree

Hide file tree

Showing 10 changed files with 97 additions and 362 deletions.
diff --git a/xla/autotuning.proto b/xla/autotuning.proto
@@ -83,10 +83,6 @@ message AutotuneResult {
     int64 num_ctas = 7;
   }
 
-  message CustomKernelFusionKey {
-    int64 kernel_index = 1;
-  }
-
   int64 scratch_bytes = 8;
   google.protobuf.Duration run_time = 9;
 
@@ -97,11 +93,10 @@ message AutotuneResult {
     GemmKey gemm = 6;
     TritonGemmKey triton = 17;
     CudaConvPlanKey cuda_conv_plan = 15;
-    CustomKernelFusionKey custom_kernel_fusion = 18;
     stream_executor.dnn.AlgorithmProto algorithm = 16;
   }
 
-  // Next ID: 19
+  // Next ID: 17
 }
 
 message AutotuningLog {

diff --git a/xla/service/gpu/BUILD b/xla/service/gpu/BUILD
@@ -1627,7 +1627,6 @@ xla_test(
         "//xla/tests:filecheck",
         "//xla/tests:hlo_test_base",
         "//xla/tests:literal_test_util",
-        "//xla/tests:verified_hlo_module",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/lib/core:status_test_util",
         "@com_google_absl//absl/log",

diff --git a/xla/service/gpu/autotuning/BUILD b/xla/service/gpu/autotuning/BUILD
@@ -58,15 +58,12 @@ cc_library(
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:buffer_comparator",
         "//xla/service/gpu:gpu_float_support",
+        "//xla/service/gpu:hlo_traversal",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:matmul_utils",
         "//xla/service/gpu:split_k_gemm_rewriter",
         "//xla/service/gpu:stream_executor_util",
-        "//xla/service/gpu/kernels:custom_kernel",
-        "//xla/service/gpu/kernels:custom_kernel_fusion",
-        "//xla/service/gpu/kernels:custom_kernel_fusion_pattern",
         "//xla/service/gpu/transforms:cudnn_fusion_compiler",
-        "//xla/service/gpu/transforms:custom_kernel_fusion_rewriter",
         "//xla/service/gpu/transforms:fusion_wrapper",
         "//xla/service/gpu/transforms:gemm_rewriter",
         "//xla/service/gpu/transforms:priority_fusion",
@@ -75,9 +72,11 @@ cc_library(
         "//xla/stream_executor:device_memory",
         "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_memory_allocator",
+        "//xla/stream_executor/gpu:redzone_allocator",
         "//xla/tools:hlo_decomposer_lib",
         "//xla/tsl/lib/core:bits",
         "//xla/tsl/util/proto:proto_utils",
+        "@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/log",
@@ -138,8 +137,6 @@ xla_test(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:device_description_proto_cc",
         "//xla/stream_executor:semantic_version",
-        "//xla/stream_executor:stream_executor_h",
-        "//xla/stream_executor/gpu:gpu_executor_header",
         "//xla/tests:filecheck",
         "//xla/tests:hlo_test_base",
         "//xla/tests:test_utils",

diff --git a/xla/service/gpu/autotuning/custom_kernel_fusion_autotuner.cc b/xla/service/gpu/autotuning/custom_kernel_fusion_autotuner.cc
@@ -192,7 +192,7 @@ absl::StatusOr<bool> AutotuneCustomKernelFusion(
   return previous_kernel_index != fastest_kernel_index;
 }
 
-bool IsCutlassCustomFusion(const HloComputation* computation) {
+bool IsCustomFusion(const HloComputation* computation) {
   if (!computation->IsFusionComputation()) {
     return false;
   }
@@ -212,18 +212,8 @@ bool IsCutlassCustomFusion(const HloComputation* computation) {
     return false;
   }
 
-  if (gpu_backend_config->fusion_backend_config().kind() != kCustomFusionKind) {
-    return false;
-  }
-
-  if (gpu_backend_config->fusion_backend_config()
-          .custom_fusion_config()
-          .name()
-          .rfind("cutlass", 0) != 0) {
-    return false;
-  }
-
-  return true;
+  return gpu_backend_config->fusion_backend_config().kind() ==
+         kCustomFusionKind;
 }
 }  // namespace
 
@@ -241,7 +231,7 @@ absl::StatusOr<bool> CustomKernelFusionAutotuner::Run(
 
   bool hlo_changed = false;
   for (const HloComputation* computation : module->computations()) {
-    if (IsCutlassCustomFusion(computation)) {
+    if (IsCustomFusion(computation)) {
       TF_ASSIGN_OR_RETURN(
           bool instruction_changed,
           AutotuneCustomKernelFusion(computation->FusionInstruction(), config_,