PR #10510: [ROCm] Triton in XLA for ROCm - gemm rewriter triton relat…

…ed changes.

Imported from GitHub PR openxla/xla#10510

First commit of the series for enabling Triton in XLA for ROCm .
Copybara import of the project:

--
f680e98aa6908d02021bc0c8bb3d24edfc6b2a5d by Zoran Jovanovic <zjovanov@amd.com>:

[ROCm] Triton in XLA for ROCm - gemm rewriter triton related changes.

Merging this change closes #10510

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10510 from ROCm:rocm_triton_backend f680e98aa6908d02021bc0c8bb3d24edfc6b2a5d
PiperOrigin-RevId: 615707894

third_party/xla/xla/pjrt/gpu/BUILD

@@ -243,7 +243,6 @@ cc_library(
         "//xla/service:local_service",
         "//xla/service:local_service_utils",
         "//xla/service/gpu:executable_proto_cc",
-        "//xla/stream_executor/cuda:cuda_platform_id",
         "//xla/stream_executor/platform",
         "@com_google_absl//absl/status",
         "@local_tsl//tsl/platform:casts",
@@ -253,11 +252,13 @@ cc_library(
         "//xla/service/gpu:gpu_compiler",
     ]) + if_cuda([
         "@local_config_cuda//cuda:cuda_headers",
+        "//xla/stream_executor/cuda:cuda_platform_id",
         "//xla/stream_executor/cuda:cuda_activation_header",
         "//xla/stream_executor/gpu:gpu_cudamallocasync_allocator",
         "//xla/service/gpu:nvptx_compiler_impl",
     ]) + if_rocm([
         "@local_config_rocm//rocm:rocm_headers",
+        "//xla/stream_executor/rocm:rocm_platform_id",
         "//xla/service/gpu:amdgpu_compiler_impl",
     ]),
     alwayslink = True,

third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc

@@ -42,13 +42,14 @@ limitations under the License.
 #include "xla/service/hlo_proto_util.h"
 #include "xla/service/local_service.h"
 #include "xla/service/local_service_utils.h"
-#include "xla/stream_executor/cuda/cuda_platform_id.h"
 #endif
 
 #if GOOGLE_CUDA
 #include "xla/service/gpu/nvptx_compiler.h"
+#include "xla/stream_executor/cuda/cuda_platform_id.h"
 #elif TENSORFLOW_USE_ROCM
 #include "xla/service/gpu/amdgpu_compiler.h"
+#include "xla/stream_executor/rocm/rocm_platform_id.h"
 #endif
 
 namespace xla {

third_party/xla/xla/service/gpu/conv_algorithm_picker.cc

@@ -36,7 +36,6 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
 #include "absl/types/span.h"
-#include "third_party/gpus/cudnn/cudnn_version.h"
 #include "xla/debug_options_flags.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -75,6 +74,7 @@ limitations under the License.
 
 #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA)
 #include "third_party/gpus/cudnn/cudnn.h"  // IWYU pragma: keep
+#include "third_party/gpus/cudnn/cudnn_version.h"
 #if CUDNN_VERSION >= 90000
 #include "third_party/gpus/cudnn/cudnn_ops.h"
 #else

third_party/xla/xla/service/gpu/cudnn_norm_rewriter_test.cc

@@ -16,13 +16,13 @@ limitations under the License.
 #include <string>
 
 #include <gtest/gtest.h>
-#include "third_party/gpus/cudnn/cudnn_version.h"
 #include "xla/error_spec.h"
 #include "xla/stream_executor/device_description.h"
 
 #if GOOGLE_CUDA
 #include "third_party/gpus/cuda/include/cuda.h"
 #include "third_party/gpus/cudnn/cudnn.h"  // IWYU pragma: keep
+#include "third_party/gpus/cudnn/cudnn_version.h"
 #endif
 
 #include "xla/service/gpu/tests/gpu_codegen_test.h"

third_party/xla/xla/service/gpu/gemm_rewriter_triton.cc

@@ -704,9 +704,14 @@ class GemmRewriterTritonVisitor : public DfsHloRewriteVisitor {
     // If a GEMM requiring padding for cuBLAS is encountered here this
     // happened because earlier ShouldTritonHandleGEMM() accepted it and padding
     // was skipped. Accept it ignoring profitability checks.
-    if (!CublasRequiresPadding(*Cast<HloDotInstruction>(dot), gpu_version_) &&
-        !should_fuse) {
-      return absl::OkStatus();
+    // TODO(rocm): check ROCM padding requirements.
+    if (std::holds_alternative<se::CudaComputeCapability>(gpu_version_)) {
+      if (!CublasRequiresPadding(
+              *Cast<HloDotInstruction>(dot),
+              std::get<se::CudaComputeCapability>(gpu_version_)) &&
+          !should_fuse) {
+        return OkStatus();
+      }
     }
 
     HloComputation* computation =
@@ -752,17 +757,31 @@ absl::StatusOr<bool> RunOnComputation(
   return visitor.changed();
 }
 
-bool IsSupportedByTriton(
-    PrecisionConfig::Algorithm algorithm,
-    const se::CudaComputeCapability& cuda_compute_capability) {
+bool IsSupportedByTriton(PrecisionConfig::Algorithm algorithm,
+                         const se::GpuComputeCapability& gpu_version) {
+  auto cuda_compute_capability =
+      std::get_if<se::CudaComputeCapability>(&gpu_version);
+  auto rocm_compute_capability =
+      std::get_if<se::RocmComputeCapability>(&gpu_version);
   switch (algorithm) {
     case PrecisionConfig::ALG_DOT_BF16_BF16_F32:
+      if (rocm_compute_capability) {
+        return rocm_compute_capability->has_bf16_dtype_support();
+      }
       return true;
 
     case PrecisionConfig::ALG_DOT_TF32_TF32_F32:
+      if (rocm_compute_capability) {
+        return false;
+      }
     case PrecisionConfig::ALG_DOT_BF16_BF16_F32_X3:
     case PrecisionConfig::ALG_DOT_BF16_BF16_F32_X6:
-      return cuda_compute_capability.IsAtLeastAmpere();
+      if (rocm_compute_capability) {
+        return rocm_compute_capability->has_bf16_dtype_support();
+      } else if (cuda_compute_capability) {
+        return cuda_compute_capability->IsAtLeastAmpere();
+      }
+      return false;
 
     // TODO(b/326579472): Fix the support of this algorithm and maybe allow it
     // here.
@@ -780,8 +799,12 @@ FusionDecision CanTritonHandleGEMM(
     const HloDotInstruction& dot, const se::GpuComputeCapability& gpu_version) {
   auto cuda_compute_capability =
       std::get_if<se::CudaComputeCapability>(&gpu_version);
+  auto rocm_compute_capability =
+      std::get_if<se::RocmComputeCapability>(&gpu_version);
 
-  if (!cuda_compute_capability) return "Non CUDA device.";
+  if (!cuda_compute_capability && !rocm_compute_capability) {
+    return "Non CUDA or ROCM device.";
+  }
 
   if (dot.precision_config().algorithm() == PrecisionConfig::ALG_UNSET) {
     if (!tsl::tensor_float_32_execution_enabled() ||
@@ -802,8 +825,13 @@ FusionDecision CanTritonHandleGEMM(
       case F32:
         return true;
       case BF16:
-        return cuda_compute_capability->IsAtLeast(
-            stream_executor::CudaComputeCapability::AMPERE);
+        if (cuda_compute_capability) {
+          return cuda_compute_capability->IsAtLeast(
+              stream_executor::CudaComputeCapability::AMPERE);
+        } else if (rocm_compute_capability) {
+          return rocm_compute_capability->has_bf16_dtype_support();
+        }
+        return false;
       default:
         return false;
     }

third_party/xla/xla/service/gpu/triton_tiling_propagation.cc

@@ -1037,7 +1037,9 @@ GetPropagatedDimOrdersAndRequirementsIfProfitablyFusible(
       std::move(std::get<DimOrdersAndReqs>(result_or_error));
   int fusion_level =
       hlo.GetModule()->config().debug_options().xla_gpu_triton_fusion_level();
-  if (!std::get<se::CudaComputeCapability>(gpu_version)
+  // TODO(ROCm) Check fusion level for ROCm.
+  if (std::holds_alternative<se::CudaComputeCapability>(gpu_version) &&
+      !std::get<se::CudaComputeCapability>(gpu_version)
            .IsAtLeast(se::CudaComputeCapability::AMPERE)) {
     fusion_level = std::min(fusion_level, 1);
   }