Fix ROCm 5.2.3 pipeline (#15073)

Make CK optional again.

onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh

@@ -4,6 +4,7 @@
 
 #pragma once
 
+#ifdef USE_COMPOSABLE_KERNEL
 #include <cstdlib>
 
 #include "ck/ck.hpp"
@@ -75,7 +76,6 @@ using device_batched_gemm_softmax_gemm_permute_instances =
         // clang-format on
         >;
 
-
 struct PreSoftmaxAttentionScoreOp {
   PreSoftmaxAttentionScoreOp(float scale) : scale_(scale) {}
 
@@ -97,7 +97,6 @@ struct PreSoftmaxAttentionScoreOp {
   float scale_;
 };
 
-
 // Use this function to gat implementation
 template <typename DT, typename D0sDT, typename AccDT, typename D0Op, MaskingSpecialization MaskingSpec>
 std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
@@ -140,7 +139,8 @@ std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
 GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
     F16, ck::Tuple<F16, F16>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>();
 
-}  // namespace ck
+}  // namespace internal
 }  // namespace rocm
 }  // namespace contrib
 }  // namespace onnxruntime
+#endif  // USE_COMPOSABLE_KERNEL

onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#ifdef USE_COMPOSABLE_KERNEL
 #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
 #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
@@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
 }  // namespace rocm
 }  // namespace contrib
 }  // namespace onnxruntime
+#endif  // USE_COMPOSABLE_KERNEL

onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#ifdef USE_COMPOSABLE_KERNEL
 #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
 #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
@@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
 }  // namespace rocm
 }  // namespace contrib
 }  // namespace onnxruntime
+#endif  // USE_COMPOSABLE_KERNEL

onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#ifdef USE_COMPOSABLE_KERNEL
 #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
 #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
 #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
@@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
 }  // namespace rocm
 }  // namespace contrib
 }  // namespace onnxruntime
+#endif  // USE_COMPOSABLE_KERNEL

onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh

@@ -63,11 +63,13 @@ are in composable kernels. The scale and add logic is performed via Acc0ElementO
 #include "contrib_ops/cpu/bert/attention_base.h"
 #include "contrib_ops/rocm/bert/attention_impl.h"
 #include "contrib_ops/rocm/bert/attention_softmax.h"
+#ifdef USE_COMPOSABLE_KERNEL
 #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
 
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+#endif  // USE_COMPOSABLE_KERNEL
 
 #include <array>
 #include <vector>
@@ -284,24 +286,6 @@ struct GemmSoftmaxGemmPermuteGenericPipeline {
   }
 };
 
-namespace {
-template <typename T>
-struct DataTypeAdaptor {
-  using type = T;
-};
-
-template <>
-struct DataTypeAdaptor<half> {
-  using type = ck::half_t;
-};
-
-template <>
-struct DataTypeAdaptor<BFloat16> {
-  using type = ck::bhalf16_t;
-};
-
-}  // namespace
-
 template <typename T>
 class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOp<GemmSoftmaxGemmPermuteParams<T>> {
  public:
@@ -391,6 +375,24 @@ class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOp<GemmSoftmaxGem
   }
 };
 
+#ifdef USE_COMPOSABLE_KERNEL
+namespace {
+template <typename T>
+struct DataTypeAdaptor {
+  using type = T;
+};
+
+template <>
+struct DataTypeAdaptor<half> {
+  using type = ck::half_t;
+};
+
+template <>
+struct DataTypeAdaptor<BFloat16> {
+  using type = ck::bhalf16_t;
+};
+}  // namespace
+
 template <typename T, bool USE_BIAS, bool USE_MASK>
 auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() {
   constexpr const int kNumBiasBuffer = static_cast<int>(USE_BIAS) + static_cast<int>(USE_MASK);
@@ -501,6 +503,7 @@ auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() {
   }
   return ret;
 }
+#endif  // USE_COMPOSABLE_KERNEL
 
 template <typename T>
 GemmSoftmaxGemmPermuteTunableOp<T>::GemmSoftmaxGemmPermuteTunableOp() {
@@ -508,6 +511,7 @@ GemmSoftmaxGemmPermuteTunableOp<T>::GemmSoftmaxGemmPermuteTunableOp() {
     return GemmSoftmaxGemmPermuteGenericPipeline<T>::Run(params, false);
   });
 
+#ifdef USE_COMPOSABLE_KERNEL
   for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/false, /*USE_MASK=*/false>()) {
     this->RegisterOp(std::move(op));
   }
@@ -523,6 +527,7 @@ GemmSoftmaxGemmPermuteTunableOp<T>::GemmSoftmaxGemmPermuteTunableOp() {
   for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/true, /*USE_MASK=*/true>()) {
     this->RegisterOp(std::move(op));
   }
+#endif
 }
 
 }  // namespace rocm

onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py

@@ -184,6 +184,7 @@ def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, n
     _test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
 
 
+@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled")
 @pytest.mark.parametrize("mask_dim", mask_dims, ids=get_mask_dim_id)
 @pytest.mark.parametrize("biased", biaseds, ids=get_biased_id)
 @pytest.mark.parametrize("head_size", head_sizes)
@@ -329,7 +330,10 @@ def profile_with_args(dtype, batch, seqlen, total_seqlen, num_heads, head_size,
         profile_gemm_softmax_gemm_permute_func(
             getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)), *args
         )
-        profile_gemm_softmax_gemm_permute_func(getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args)
+        if ke.is_composable_kernel_available():
+            profile_gemm_softmax_gemm_permute_func(
+                getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args
+            )
         profile_gemm_softmax_gemm_permute_func(
             getattr(ke, "GemmSoftmaxGemmPermuteTunable_" + dtype_to_suffix(dtype)), *args
         )

onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu

@@ -158,6 +158,7 @@ class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplor
   }
 };
 
+#ifdef USE_COMPOSABLE_KERNEL
 template <typename T, bool USE_BIAS, bool USE_MASK>
 class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer<T> {
  public:
@@ -215,6 +216,7 @@ class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer<T>
   std::vector<std::string> type_strings_;
   size_t selected_op_{};
 };
+#endif  // USE_COMPOSABLE_KERNEL
 
 // The pipeline composed from rocblas api calls and kernel launches.
 template <typename T>
@@ -287,10 +289,12 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor
 void InitGemmSoftmaxGemmPermute(py::module m) {
   REGISTER_GENERIC(half);
 
+#ifdef USE_COMPOSABLE_KERNEL
   REGISTER_CK(half, false, false, "");
   REGISTER_CK(half, true, false, "Biased");
   REGISTER_CK(half, false, true, "Masked");
   REGISTER_CK(half, true, true, "BiasedMasked");
+#endif
 
   REGISTER_TUNABLE(half);
 }