Triton fixes porting from v0.6.0 (#389)

* rocprof-sdk addition,
upstream PR: openxla/pull/29769

Squash following commits..
Update rocprofiler-sdk (v3) along with roctracer (v1) for rocm-jaxlib-v0.6.0 (#302)

* update for integration of rocprofiler-sdk (along with roctracer as a backup based on bazel_options from CLI)

(cherry picked from commit 7775dd0)

use VLOG(2) to replace LOG(INFO), so PGLE has no verbose info (#357)

(cherry picked from commit 5950125)

update with kernel details for rocm-7.x (#364)

* update with kernel details for rocm-7.x

(cherry picked from commit 5597c0d)

update to remove previously hard-coded rocprofiler-sdk path (#369)

* update to remove previously hard-coded rocprofiler-sdk path and add skip_rocprofiler_sdk to avoid loading `rocprofiler-sdk`

(cherry picked from commit ff74b5f)

* fixed buffer comparator test

* misc fixes ported from rocm-jaxlib-v0.6.0

---------

Co-authored-by: Pavel Emeliyanenko <[email protected]>
(cherry picked from commit f013645)
(cherry picked from commit b03cd94)

Added support for waves_per_eu function attribute. (#181)

(cherry picked from commit bc1d816)
(cherry picked from commit d3f94e9)

removed two line change (revert of half of the openxla#25959 commit

(cherry picked from commit 109e138)

Fixes for jax 0.6.0 (#207)

* Add fixes for jax plugin 0.6.0

Drop NEEDED linking to unnecessary libs.
These are loaded by amdhipruntime and not us.

Fix missing NEEDED on MIOpen shared object.

* Minor rocblas related changes for rocm 70

(cherry picked from commit 0de7d49)

---------

Co-authored-by: Zoran Jovanovic <[email protected]>
(cherry picked from commit 28f10a0)

Add hipBLASLt support for gfx11. (#301)

(cherry picked from commit f814bff)

Add bf16 starting from gfx11, bugfix & optimize RocmComputeCapability (#303)

* Bugfix and improve device_description.h::RocmComputeCompatibility

* Enable ALG_DOT_BF16* on rocm with HW support

(cherry picked from commit 510ea06)

[ROCm] Use bundled bitcode files (#196)

Also trim bitcode file list to ockl.bc and ocml.bc only.

(cherry picked from commit fc9e3c3)

Add MIOPEN_FIND_ENFORCE For ROCm 7 for convolution gemms (#312)

* Add MIOPEN_FIND_ENFORCE For ROCm 7 for convolution gemms

* Exclude failing CollectiveOpsE2E tests

(cherry picked from commit fb6ddfb)

Restore RocmComputeCapability:: gfx11_rx7900() and gfx12_rx8900() methods (#333)

At least gfx11_rx7900() is still needed for TF build.

(cherry picked from commit 13c3de1)

Make device_count_ atomic (#343)

* Make device_count_ atomic

* Use relaxed memory order

* Fix build error

(cherry picked from commit 8513f2d)

fix hardcoded max registers (#345)

(cherry picked from commit f3e170a)

fix hardcoded ecc enabled (#348)

(cherry picked from commit 9cfa74a)

remove reserved memory (#349)

(cherry picked from commit 0015d0e)

Add rocm_dev config for remote caching (#353)

(cherry picked from commit c815420)

added rocm7 support to EnablePeerAccess (#347)

* added rocm7 support to EnablePeerAccess

* use wrap namespace, clang-format and add comments

(cherry picked from commit 85548a7)

[ROCm] Disable Cudnn fusions (#358)

(cherry picked from commit edab8b2)

* Ported all triton related changes from v0.6.0 to v0.7.1

(cherry picked from commit 1851bcc)

Disable softmax triton fusion if triton gemm is off (#281)

* Disable softmax rewriter triton if triton gemm is disabled

* Add specific flag to enable triton softmax fusion

* Address review comments

(cherry picked from commit 51a7f4b)

[ROCm][Triton] Disable transposed load in certain conditions

(cherry picked from commit 50860e9)

Enable unit tests that pass after fixing some Triton related issues. (#285)

* Enable unit tests that pass after fixing some Triton related issues.

* fusion_emitter_device_legacy_test still fails on MI200

(cherry picked from commit 97dd565)

Rocm jaxlib v0.6.0 triton support ut (#279)

* Fixed triton/support_test - no fmfa.

* Fix issue with rounding mode in accelerate amd matmul.

* Fixed issues with usage of mfma in support_test.

(cherry picked from commit 44f7d87)

Restore gpu_triton_custom_call_test (#262)

(cherry picked from commit 32eafa4)

Skipped CanNotEmitTritonCustomCallOnPreAmpereGpu test for ROCM.

(cherry picked from commit 56ec7ec)
(cherry picked from commit b1f3e9f)

fixed createTritonAMDGPULowerInstructionSchedHintsPass (#179)

(cherry picked from commit 8517a3a)
(cherry picked from commit c62e47d)

fixed bazel build issue

---------

Co-authored-by: Chunyu Jin <[email protected]>
Co-authored-by: zoranjovanovic-ns <[email protected]>
Co-authored-by: Alex <[email protected]>

Author: 4 people

build_tools/rocm/run_xla.sh

@@ -45,6 +45,9 @@ GPU_NAME=(`rocminfo | grep -m 1 gfx`)
 GPU_NAME=${GPU_NAME[1]}
 
 EXCLUDED_TESTS=(
+# //xla/service/gpu/tests:gpu_kernel_tiling_test_gpu_amd_any
+GpuKernelTilingTest.ColumnReductionWithLayoutChangeTiled
+GpuKernelTilingTest.ReductionInputTooLarge
 # //xla/pjrt/c:pjrt_c_api_gpu_test_gpu_amd_any
 PjrtCAPIGpuExtensionTest.TritonCompile
 # //xla/backends/gpu/codegen/triton:fusion_emitter_device_test_gpu_amd_any

third_party/triton/temporary/0001-AMD-Quick-fix-disabling-transposed-load-used-as-diff.patch

@@ -0,0 +1,46 @@
+From d539916e4d49cca93f54a5f99f7822050205432c Mon Sep 17 00:00:00 2001
+From: Jungwook Park <[email protected]>
+Date: Thu, 7 Aug 2025 06:34:49 -0500
+Subject: [PATCH] [AMD] Quick fix disabling transposed load used as different
+ type.
+
+Disabling transposedLoad if dot is using it as a different element type.
+Otherwise it's picking wrong vectorsize when lowering.
+---
+ .../lib/TritonAMDGPUToLLVM/MemoryOpToLLVM.cpp | 20 +++++++++++++++++++
+ 1 file changed, 20 insertions(+)
+
+diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/MemoryOpToLLVM.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/MemoryOpToLLVM.cpp
+index 661a17678..6bda3a818 100644
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/MemoryOpToLLVM.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/MemoryOpToLLVM.cpp
+@@ -214,6 +214,26 @@ private:
+       return false;
+     }
+ 
++    // transposed load can be used only when it's consumed by dot with the
++    // loaded data type.
++    int opIdx = 0;
++    triton::gpu::LocalLoadOp lLoad = cast<triton::gpu::LocalLoadOp>(localLoad);
++    if (auto dotEnc = lLoad.getSrc().getType().getEncoding())
++      opIdx = cast<triton::gpu::DotOperandEncodingAttr>(dotEnc).getOpIdx();
++    else
++      return false;
++
++    SetVector<Operation *> slice;
++    getForwardSlice(localLoad, &slice);
++    for (auto op : slice) {
++      if (auto dotOp = dyn_cast<triton::DotOp>(op)) {
++        auto inputMat = (opIdx == 0) ? dotOp.getA() : dotOp.getB();
++        auto bitwidthMat = inputMat.getType().getElementTypeBitWidth();
++        if (bitwidth != bitwidthMat)
++          return false;
++      }
++    }
++
+     return true;
+   }
+ 
+-- 
+2.34.1
+

third_party/triton/temporary/accelerateamdmatmul.patch

@@ -0,0 +1,14 @@
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp
+@@ -1068,7 +1068,10 @@ public:
+     if (isFloat(srcElTy) && isFloat(dstElTy)) {
+       auto rmode =
+           RoundingModeAttr::get(rewriter.getContext(), RoundingMode::RTNE);
+-      return rewriter.create<FpToFpOp>(loc, dstTy, v, rmode);
++      if (dstElTy.getIntOrFloatBitWidth() < srcElTy.getIntOrFloatBitWidth()) {
++        return rewriter.create<FpToFpOp>(loc, dstTy, v, rmode);
++      }
++      return rewriter.create<FpToFpOp>(loc, dstTy, v);
+     }
+     if (!isFloat(srcElTy) && isFloat(dstElTy))
+       return rewriter.create<arith::SIToFPOp>(loc, dstTy, v);

third_party/triton/temporary/accelerateamdmatmul2.patch

@@ -0,0 +1,24 @@
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/AccelerateAMDMatmul.cpp
+@@ -380,9 +380,17 @@ Value convertAndCastTensor(PatternRewriter &rewriter, Value value,
+     else if (oldElemType.isF32() && newElemType.isF16())
+       castedTensor =
+           rewriter.create<arith::TruncFOp>(loc, castedType, convertedTensor);
+-    else
+-      castedTensor =
+-          rewriter.create<tt::FpToFpOp>(loc, castedType, convertedTensor);
++    else {
++      if(oldElemType.getIntOrFloatBitWidth() > newElemType.getIntOrFloatBitWidth()) {
++        auto rmode =
++           RoundingModeAttr::get(rewriter.getContext(), RoundingMode::RTNE);
++        castedTensor =
++          rewriter.create<tt::FpToFpOp>(loc, castedType, convertedTensor, rmode);
++      } else {
++        castedTensor =
++            rewriter.create<tt::FpToFpOp>(loc, castedType, convertedTensor);
++      }
++    }
+   }
+   return castedTensor;
+ }
+ 

third_party/triton/temporary/series.bzl

@@ -19,5 +19,8 @@ temporary_patch_list = [
     "//third_party/triton:temporary/tutorial_fixes.patch",
     "//third_party/triton:temporary/ws_fix.patch",
     "//third_party/triton:temporary/ws_ub_fix.patch",
+    "//third_party/triton:temporary/0001-AMD-Quick-fix-disabling-transposed-load-used-as-diff.patch",
+    "//third_party/triton:temporary/accelerateamdmatmul.patch",
+    "//third_party/triton:temporary/accelerateamdmatmul2.patch",
     # Add new patches just above this line
 ]

xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc

@@ -146,10 +146,8 @@ absl::Status CreateTritonPipeline(
   pm->addPass(mlir::createCanonicalizerPass());
   pm->addPass(mlir::createCSEPass());
   pm->addPass(mlir::createSymbolDCEPass());
-  if (/*(instruction_sched_variant=="none") == */ false) {
-    pm->addPass(mt::createTritonAMDGPULowerInstructionSchedHintsPass(
-        cc.gfx_version(), num_stages));
-  }
+  pm->addPass(mt::createTritonAMDGPULowerInstructionSchedHintsPass(
+      cc.gfx_version(), num_stages));
   pm->addPass(mt::createConvertBuiltinFuncToLLVMPass(/*ftz=*/true));
   // There is no clusters in ROCm for now.
   out_cluster_info.clusterDimX = 1;

xla/backends/gpu/codegen/triton/emitter_helpers.cc

@@ -106,6 +106,12 @@ absl::StatusOr<Type> TritonType(EmitterLocOpBuilder& b, PrimitiveType t) {
       return b.getType<mlir::Float8E5M2Type>();
     case F8E4M3FN:
       return b.getType<mlir::Float8E4M3FNType>();
+    case F8E4M3B11FNUZ:
+      return b.getType<mlir::Float8E4M3B11FNUZType>();
+    case F8E5M2FNUZ:
+      return b.getType<mlir::Float8E5M2FNUZType>();
+    case F8E4M3FNUZ:
+      return b.getType<mlir::Float8E4M3FNUZType>();
     default:
       return absl::UnimplementedError(
           absl::StrCat("This type is not supported yet: ",
@@ -126,6 +132,9 @@ absl::StatusOr<PrimitiveType> GetPrimitiveType(Type t) {
   if (t.isInteger(1)) return PRED;
   if (mlir::isa<mlir::Float8E5M2Type>(t)) return F8E5M2;
   if (mlir::isa<mlir::Float8E4M3FNType>(t)) return F8E4M3FN;
+  if (mlir::isa<mlir::Float8E4M3B11FNUZType>(t)) return F8E4M3B11FNUZ;
+  if (mlir::isa<mlir::Float8E5M2FNUZType>(t)) return F8E5M2FNUZ;
+  if (mlir::isa<mlir::Float8E4M3FNUZType>(t)) return F8E4M3FNUZ;
   return absl::UnimplementedError("Unsupported type in getPrimitiveType.\n");
 }
 

xla/backends/gpu/codegen/triton/support.cc

@@ -62,8 +62,12 @@ bool IsTritonSupportedDataType(PrimitiveType type,
     case F64:
       return true;
     case F8E5M2:
-    case F8E4M3FN:
-      return std::holds_alternative<se::CudaComputeCapability>(gpu_version);
+    case F8E4M3FN:      
+      return std::holds_alternative<se::CudaComputeCapability>(gpu_version) ||
+             std::holds_alternative<se::RocmComputeCapability>(gpu_version);
+    case F8E5M2FNUZ:
+    case F8E4M3FNUZ:
+      return std::holds_alternative<se::RocmComputeCapability>(gpu_version);
     case BF16:
       return std::holds_alternative<se::CudaComputeCapability>(gpu_version) ||
              (std::holds_alternative<se::RocmComputeCapability>(gpu_version) &&
@@ -92,7 +96,10 @@ absl::flat_hash_set<HloOpcode> TritonSupportedUnaryElementwiseOps(
   absl::flat_hash_set<HloOpcode> ret{HloOpcode::kAbs, HloOpcode::kCopy};
 
   if (element_type != PrimitiveType::F8E5M2 &&
-      element_type != PrimitiveType::F8E4M3FN) {
+      element_type != PrimitiveType::F8E4M3FN &&
+      element_type != PrimitiveType::F8E4M3B11FNUZ &&
+      element_type != PrimitiveType::F8E5M2FNUZ &&
+      element_type != PrimitiveType::F8E4M3FNUZ) {
     ret.insert(HloOpcode::kNegate);
   }
 
@@ -171,7 +178,10 @@ absl::flat_hash_set<HloOpcode> TritonSupportedBinaryElementwiseOps(
     PrimitiveType element_type, const se::GpuComputeCapability& gpu_version) {
   if (element_type == PrimitiveType::S4 || element_type == PrimitiveType::U16 ||
       element_type == PrimitiveType::F8E5M2 ||
-      element_type == PrimitiveType::F8E4M3FN) {
+      element_type == PrimitiveType::F8E4M3FN || 
+      element_type == PrimitiveType::F8E4M3B11FNUZ ||
+      element_type == PrimitiveType::F8E5M2FNUZ ||
+      element_type == PrimitiveType::F8E4M3FNUZ) {
     return {};
   }
 
@@ -220,7 +230,10 @@ absl::flat_hash_set<HloOpcode> TritonSupportedTernaryElementwiseOps(
   }
 
   if (element_type == PrimitiveType::F8E5M2 ||
-      element_type == PrimitiveType::F8E4M3FN) {
+      element_type == PrimitiveType::F8E4M3FN || 
+      element_type == PrimitiveType::F8E4M3B11FNUZ ||
+      element_type == PrimitiveType::F8E5M2FNUZ ||
+      element_type == PrimitiveType::F8E4M3FNUZ) {
     return {HloOpcode::kSelect};
   }
 
@@ -248,7 +261,10 @@ CodegenDecision CanTritonHandleReduce(
     const HloReduceInstruction& reduce,
     const se::GpuComputeCapability& gpu_version) {
   if (reduce.shape().element_type() == PrimitiveType::F8E4M3FN ||
-      reduce.shape().element_type() == PrimitiveType::F8E5M2) {
+      reduce.shape().element_type() == PrimitiveType::F8E5M2 || 
+      reduce.shape().element_type() == PrimitiveType::F8E5M2FNUZ ||
+      reduce.shape().element_type() == PrimitiveType::F8E4M3FNUZ ||
+      reduce.shape().element_type() == PrimitiveType::F8E4M3B11FNUZ) {
     return CodegenDecision::Forbid(
         "F8E4M3FN and F8E5M2 are not supported for reductions.");
   }
@@ -296,7 +312,8 @@ absl::Status CheckSupportedCheckDotDimensions(const HloDotInstruction& dot) {
   return absl::OkStatus();
 }
 
-bool IsSupportedDotAlgorithm(PrecisionConfig::Algorithm algorithm) {
+bool IsSupportedDotAlgorithm(PrecisionConfig::Algorithm algorithm,
+                             const se::GpuComputeCapability& gpu_version) {
   switch (algorithm) {
     case PrecisionConfig::ALG_UNSET:
     case PrecisionConfig::ALG_DOT_F16_F16_F16:
@@ -309,8 +326,13 @@ bool IsSupportedDotAlgorithm(PrecisionConfig::Algorithm algorithm) {
     case PrecisionConfig::ALG_DOT_TF32_TF32_F32:
     case PrecisionConfig::ALG_DOT_TF32_TF32_F32_X3:
     case PrecisionConfig::ALG_DOT_BF16_BF16_F32_X9:
-      return true;
+      if (!std::holds_alternative<se::RocmComputeCapability>(gpu_version)) {
+        return true;
+      }
     case PrecisionConfig::ALG_DOT_BF16_BF16_BF16:
+      if (std::holds_alternative<se::RocmComputeCapability>(gpu_version)) {
+        return true;
+      }
     case PrecisionConfig::ALG_DOT_ANY_F8_ANY_F8_F32:
     case PrecisionConfig::ALG_DOT_ANY_F8_ANY_F8_F32_FAST_ACCUM:
     default:
@@ -336,7 +358,15 @@ CodegenDecision AreTypesSupportedByAlgUnsetDot(
     }
   }
 
-  auto supported_float_types = {BF16, F16, F32, F64, F8E5M2, F8E4M3FN};
+  if (input_type == F8E4M3B11FNUZ || result_type == F8E4M3B11FNUZ ||
+      input_type == F64) {
+    if (std::holds_alternative<se::RocmComputeCapability>(gpu_version)) {
+      return CodegenDecision::Forbid(
+          "Dot operation for F8E4M3B11FNUZ is not supported on ROCM.");
+    }
+  }
+
+  auto supported_float_types = {BF16, F16, F32, F64, F8E5M2};
   if (absl::c_linear_search(supported_float_types, input_type)) {
     return CodegenDecision::Allow();
   }
@@ -405,6 +435,11 @@ CodegenDecision AreDotAlgorithmInputAndOutputConversionsSupported(
     return forbid("Unsupported BF16 on GPUs before Blackwell");
   }
 
+  if (allowed_operands_types_or->front() == PrimitiveType::F64 &&
+      std::holds_alternative<se::RocmComputeCapability>(gpu_version)) {
+   return forbid("Unsupported result conversion");
+  }
+
   if (allowed_operands_types_or->size() != 1) {
     if (lhs_type == rhs_type &&
         absl::c_linear_search(*allowed_operands_types_or, lhs_type)) {
@@ -467,7 +502,7 @@ CodegenDecision IsTritonSupportedDot(
   const PrecisionConfig& precision_config = dot.precision_config();
   const PrecisionConfig::Algorithm algorithm = precision_config.algorithm();
 
-  if (!IsSupportedDotAlgorithm(algorithm)) {
+  if (!IsSupportedDotAlgorithm(algorithm, gpu_version)) {
     return CodegenDecision::Forbid(
         absl::StrCat("Unsupported dot algorithm: ",
                      PrecisionConfig::Algorithm_Name(algorithm)));
@@ -625,6 +660,8 @@ CodegenDecision IsTritonSupportedInstructionImpl(
     return CodegenDecision(
         element_type != PrimitiveType::F8E4M3FN &&
             element_type != PrimitiveType::F8E5M2 &&
+            element_type != PrimitiveType::F8E4M3FNUZ &&
+            element_type != PrimitiveType::F8E5M2FNUZ &&
             element_type != PrimitiveType::S4,
         "F8E4M3FN, F8E5M2 and S4 are not supported for iota.");
   }