Repository: pytorch
Issue: [ROCm] Inductor CK GEMM backend very slow
Labels: module: rocm,triaged,oncall: pt2

### 🐛 Describe the bug

When using the CK backend via `TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS="CK,ATEN,TRITON,CPP"` compilation of CK kernels is very slow. (>one minute per file in some cases)

It looks like some very long symbol names in these files are making compilation slower because LLVM uses `SmallString<128>` buffers to build up symbol names and now has to allocate a bunch in places that were otherwise allocation free. From some `perf` sampling it looks like this is causing LLVM to spend a lot more time in TargetMachine::getSymbol.

https://github.com/LunNova/llvm-project-rocm/blob/5a9ddc6f57430d5e8c5154779c647219c8e7cb99/llvm/lib/Target/TargetMachine.cpp#L283-L292

It's possible I've misdiagnosed this and the long symbol allocations here are mostly inside the CK code and the top level long names aren't important, or the slow compilation is entirely unrelated to the symbol names. In any case, it's very slow.

<details>
<summary>Long logs/code on torch 20241218 nightly</summary>

```
[torch/_inductor/codecache.py:3202] Compilation took 70.78182601928711 seconds. Compile command: /nix/store/baxdbiqlbq3xgfrycxz0l3lhgqr30gpg-rocmcxx/bin/clang -O3 -x hip -std=c++17 --offload-arch=gfx908 -fno-gpu-rdc -fPIC -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -mllvm -enable-post-misched=0 -DNDEBUG -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero -ffast-math -I/nix/store/jx8h8d85ymv8l950czxz8wpvd562mblx-unpack-composable_kernel-6.4.0a20241217/lib/python3.12/site-packages/ck4inductor/include -I/nix/store/jx8h8d85ymv8l950czxz8wpvd562mblx-unpack-composable_kernel-6.4.0a20241217/lib/python3.12/site-packages/ck4inductor/library/include -I/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/include -include __clang_hip_runtime_wrapper.h -L/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/lib -L/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/hip/lib -lamdhip64 -shared -o ~/ml-cache/torchinductor/yo/cyoxkhocdaqxywmu537oo7n6qsdsnn4hk6n75y3lg5sknrsei4ol.so ~/ml-cache/torchinductor/yo/cyoxkhocdaqxywmu537oo7n6qsdsnn4hk6n75y3lg5sknrsei4ol.cpp
$ cat ~/ml-cache/torchinductor/yo/cyoxkhocdaqxywmu537oo7n6qsdsnn4hk6n75y3lg5sknrsei4ol.cpp

/**
* Generated code for CK inductor backend
* See torch._inductor.codegen.rocm.ck_universal_gemm_template.CKGemmTemplate
*
* Template instance CKGemmOperation(a_layout='Row', b_layout='Col', ds_layouts=(), c_layout='Row', a_element_dtype='BF16', b_element_dtype='BF16', ds_element_dtypes=(), c_element_dtype='BF16', acc_dtype='F32', c_shuffle_dtype='BF16', a_elementwise_op='PassThrough', b_elementwise_op='PassThrough', c_elementwise_op='PassThrough', gemm_specialization='GemmSpecialization::NPadding', block_size=256, m_per_block=128, n_per_block=128, k_per_block=64, a_k1=8, b_k1=8, m_per_xdl=32, n_per_xdl=32, m_xdl_per_wave=2, n_xdl_per_wave=2, a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 32, 1), a_block_transfer_thread_cluster_arrange_order=(1, 0, 2), a_block_transfer_src_access_order=(1, 0, 2), a_block_transfer_src_vector_dim=2, a_block_transfer_src_scalar_per_vector=8, a_block_transfer_dst_scalar_per_vector_ak1=8, a_block_lds_extra_m=0, b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 32, 1), b_block_transfer_thread_cluster_arrange_order=(1, 0, 2), b_block_transfer_src_access_order=(1, 0, 2), b_block_transfer_src_vector_dim=2, b_block_transfer_src_scalar_per_vector=8, b_block_transfer_dst_scalar_per_vector_bk1=8, b_block_lds_extra_n=0, c_shuffle_m_xdl_per_wave_per_shuffle=1, c_shuffle_n_xdl_per_wave_per_shuffle=1, c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(1, 16, 1, 16), c_shuffle_block_transfer_scalar_per_vector_n_per_block=(4,), block_gemm_pipeline_scheduler='BlockGemmPipelineScheduler::Intrawave', block_gemm_pipeline_version='BlockGemmPipelineVersion::v3', a_compute_dtype=None, b_compute_dtype=None)
*
* torch.__version__='2.6.0a.post20241218'
* torch.version.git_version=Unknown
*/


#include <exception>
#include <iostream>
#include <memory>
#include <random>
#include <vector>

// CK headers

#ifdef DEBUG_LOG
#define DEBUG_LOG_TMP DEBUG_LOG
#undef DEBUG_LOG
#else
#define DEBUG_LOG_TMP 0
#endif
#include "ck/ck.hpp"
#undef DEBUG_LOG
#define DEBUG_LOG DEBUG_LOG_TMP

#include "ck/utility/data_type.hpp"
#include "ck/library/utility/check_err.hpp"
#include "ck/library/utility/device_memory.hpp"
#include "ck/library/utility/fill.hpp"
#include "ck/library/utility/host_tensor.hpp"
#include "ck/library/utility/host_tensor_generator.hpp"
#include "ck/library/utility/literals.hpp"

// CK GEMM header(s)

#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp"


// We compile all models with -fvisibility=hidden. Any symbols that need to be
// exposed in the final shared library must be declared with PT_EXPORT to make
// them visible.
#ifdef __GNUC__ // Applies to any compiler with GNU extensions (clang and g++)
#define PT_EXPORT __attribute__((__visibility__("default")))
#else
#ifdef _WIN32
#define PT_EXPORT __declspec(dllexport)
#else
#define PT_EXPORT
#endif
#endif

// as long as there is no custom arithmetic it's fine
using bfloat16 = uint16_t;
using float8_e4m3fnuz = uint8_t;
using float8_e5m2fnuz = uint8_t;

// CK globals

template <ck::index_t... Is>
using S = ck::Sequence<Is...>;

template<typename... Ts>
using Tuple = ck::Tuple<Ts...>;

using PassThrough = ck::tensor_operation::element_wise::PassThrough;
using Bilinear = ck::tensor_operation::element_wise::Bilinear;
using Scale = ck::tensor_operation::element_wise::Scale;
using ScaleAdd = ck::tensor_operation::element_wise::ScaleAdd;
using MultiplyMultiply = ck::tensor_operation::element_wise::MultiplyMultiply;

// see "composable_kernel/include/ck/utility/data_type.hpp"
using F8  = ck::f8_t;
using BF8 = ck::bf8_t;
using F16 = ck::half_t;
using F32 = float;
// using F64 = double;
using BF16 = ck::bhalf_t;
// using I32 = int32_t;
// using I8 = int8_t;
// using I4 = ck::int4_t;

#if DEBUG_LOG
static constexpr auto kDEBUG_LOG = 1;
#else
static constexpr auto kDEBUG_LOG = 0;
#endif

// CK GEMM globals

using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;

using BlockGemmPipelineScheduler = ck::BlockGemmPipelineScheduler;
using GemmSpecialization = ck::tensor_operation::device::GemmSpecialization;
using BlockGemmPipelineVersion = ck::BlockGemmPipelineVersion;

struct MultiplyMultiplyAdd {
    template <typename E, typename C, typename D0, typename D1, typename D2>
    __host__ __device__ constexpr void
    operator()(E& e, const C& c, const D0& d0, const D1& d1, const D2& d2) const {
        e = ck::type_convert<E>(
           ck::type_convert<float>(c)
           * ck::type_convert<float>(d0)
           * ck::type_convert<float>(d1)
           + ck::type_convert<float>(d2)
        );
    }
};


    // Gemm operator ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVRow_KblayoutVCol_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationNPadding_KblocksizeV256_KmperblockV128_KnperblockV128_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV32_KnperxdlV32_KmxdlperwaveV2_KnxdlperwaveV2_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV1x0x2_KablocktransfersrcaccessorderV1x0x2_KablocktransfersrcvectordimV2_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV1x0x2_KbblocktransfersrcaccessorderV1x0x2_KbblocktransfersrcvectordimV2_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV1_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x16x1x16_KcshuffleblocktransferscalarpervectornperblockV4_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone
    using Operation_ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVRow_KblayoutVCol_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationNPadding_KblocksizeV256_KmperblockV128_KnperblockV128_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV32_KnperxdlV32_KmxdlperwaveV2_KnxdlperwaveV2_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV1x0x2_KablocktransfersrcaccessorderV1x0x2_KablocktransfersrcvectordimV2_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV1x0x2_KbblocktransfersrcaccessorderV1x0x2_KbblocktransfersrcvectordimV2_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV1_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x16x1x16_KcshuffleblocktransferscalarpervectornperblockV4_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone =
        ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3<
            /* a_layout */ Row,
            /* b_layout */ Col,
            /* ds_layouts */ Tuple<>,
            /* c_layout */ Row,
            /* a_element_dtype */ BF16,
            /* b_element_dtype */ BF16,
            /* ds_element_dtypes */ Tuple<>,
            /* c_element_dtype */ BF16,
            /* acc_dtype */ F32,
            /* c_shuffle_dtype */ BF16,
            /* a_elementwise_op */ PassThrough,
            /* b_elementwise_op */ PassThrough,
            /* c_elementwise_op */ PassThrough,
            /* gemm_specialization */ GemmSpecialization::NPadding,
            /* block_size */ 256,
            /* m_per_block */ 128,
            /* n_per_block */ 128,
            /* k_per_block */ 64,
            /* a_k1 */ 8,
            /* b_k1 */ 8,
            /* m_per_xdl */ 32,
            /* n_per_xdl */ 32,
            /* m_xdl_per_wave */ 2,
            /* n_xdl_per_wave */ 2,
            /* a_block_transfer_thread_cluster_lengths_ak0_m_ak1 */ S<8, 32, 1>,
            /* a_block_transfer_thread_cluster_arrange_order */ S<1, 0, 2>,
            /* a_block_transfer_src_access_order */ S<1, 0, 2>,
            /* a_block_transfer_src_vector_dim */ 2,
            /* a_block_transfer_src_scalar_per_vector */ 8,
            /* a_block_transfer_dst_scalar_per_vector_ak1 */ 8,
            /* a_block_lds_extra_m */ 0,
            /* b_block_transfer_thread_cluster_lengths_bk0_n_bk1 */ S<8, 32, 1>,
            /* b_block_transfer_thread_cluster_arrange_order */ S<1, 0, 2>,
            /* b_block_transfer_src_access_order */ S<1, 0, 2>,
            /* b_block_transfer_src_vector_dim */ 2,
            /* b_block_transfer_src_scalar_per_vector */ 8,
            /* b_block_transfer_dst_scalar_per_vector_bk1 */ 8,
            /* b_block_lds_extra_n */ 0,
            /* c_shuffle_m_xdl_per_wave_per_shuffle */ 1,
            /* c_shuffle_n_xdl_per_wave_per_shuffle */ 1,
            /* c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block */ S<1, 16, 1, 16>,
            /* c_shuffle_block_transfer_scalar_per_vector_n_per_block */ S<4>,
            /* block_gemm_pipeline_scheduler */ BlockGemmPipelineScheduler::Intrawave,
            /* block_gemm_pipeline_version */ BlockGemmPipelineVersion::v3>;

    extern "C" {
    PT_EXPORT int rocm_fused_1(const bfloat16* X, const bfloat16* W, bfloat16* Y, int32_t M, int32_t N, int32_t K, int32_t LDA, int32_t LDB, int32_t LDC, int32_t LDD, size_t* workspace_size, uint8_t* workspace, hipStream_t stream) {
        auto gemm = 
    Operation_ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVRow_KblayoutVCol_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationNPadding_KblocksizeV256_KmperblockV128_KnperblockV128_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV32_KnperxdlV32_KmxdlperwaveV2_KnxdlperwaveV2_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV1x0x2_KablocktransfersrcaccessorderV1x0x2_KablocktransfersrcvectordimV2_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV1x0x2_KbblocktransfersrcaccessorderV1x0x2_KbblocktransfersrcvectordimV2_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV1_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x16x1x16_KcshuffleblocktransferscalarpervectornperblockV4_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone {};
        auto invoker = gemm.MakeInvoker();

        auto argument = gemm.MakeArgument(
            reinterpret_cast<const BF16*>(X),
            reinterpret_cast<const BF16*>(W),
            std::array<const void*, 0>{  },
            reinterpret_cast<BF16*>(Y),
            M,
            N,
            K,
            LDA,
            LDB,
            std::array<ck::index_t, 0>{  },
            LDC,
            1, // kBatch
            PassThrough {},
            PassThrough {},
            PassThrough {} // c_elementwise_op
        );

        if (!gemm.IsSupportedArgument(argument)) {
            // we do our best to statically avoid this case in `filter_op`
            std::cerr << "invalid argument for gemm instance " << gemm.GetTypeString() << std::endl;
            argument.Print();
            return -23;
        }
        if (workspace_size) {
            *workspace_size = gemm.GetWorkSpaceSize(&argument);
            return 0;
        }
        // run the kernel
        #ifdef GENERATE_CK_STANDALONE_RUNNER
        const auto stream_config = StreamConfig{
            stream,
            /* time kernel */ 1,
            /* log level */ 1,
            /* n_cold_iter */ 100,
            /* n_hot_iter */ 100,
            /* flush_l2_cache */ 1,
            /* rotate_count */ 5};
        #else
        const auto stream_config = StreamConfig{stream, /* time kernel */ false, /* log level */ 0};
        #endif

        const float elapsed_time = invoker.Run(argument, stream_config);

        #ifdef GENERATE_CK_STANDALONE_RUNNER
        std::cout << "elapsed time: " << elapsed_time << " ms" << std::endl;
        #else
        (void)elapsed_time;
        #endif
        return 0;
    } // kernel definition
    } // extern C

```

Also seeing some errors:

```
rank1]:E1220 13:17:04.168000 677380 torch/_inductor/select_algorithm.py:2003] [3/1] CUDA compilation error during autotuning: 
[rank1]:E1220 13:17:04.168000 677380 torch/_inductor/select_algorithm.py:2003] [3/1] C++ compile error
[rank1]:E1220 13:17:04.168000 677380 torch/_inductor/select_algorithm.py:2003] [3/1] 
[rank1]:E1220 13:17:04.168000 677380 torch/_inductor/select_algorithm.py:2003] [3/1] Command:
[rank1]:E1220 13:17:04.168000 677380 torch/_inductor/select_algorithm.py:2003] [3/1] /nix/store/baxdbiqlbq3xgfrycxz0l3lhgqr30gpg-rocmcxx/bin/clang -O3 -x hip -std=c++17 --offload-arch=gfx908 -fno-gpu-rdc -fPIC -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -mllvm -enable-post-misched=0 -DNDEBUG -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero -ffast-math -I/nix/store/jx8h8d85ymv8l950czxz8wpvd562mblx-unpack-composable_kernel-6.4.0a20241217/lib/python3.12/site-packages/ck4inductor/include -I/nix/store/jx8h8d85ymv8l950czxz8wpvd562mblx-unpack-composable_kernel-6.4.0a20241217/lib/python3.12/site-packages/ck4inductor/library/include -I/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/include -include __clang_hip_runtime_wrapper.h -L/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/lib -L/nix/store/qp3zdqix9pwcc505jgmlf27xcdq9v8xq-rocm-hip-libraries-meta/hip/lib -lamdhip64 -shared -o ~/ml-cache/torchinductor/q3/cq3qdsttyk3ptde6djpzxm46uboeevoyov3pe365w5q5xy36a525.so ~/ml-cache/torchinductor/q3/cq3qdsttyk3ptde6djpzxm46uboeevoyov3pe365w5q5xy36a525.cpp
$ cat ~/ml-cache/torchinductor/q3/cq3qdsttyk3ptde6djpzxm46uboeevoyov3pe365w5q5xy36a525.cpp
/**
* Generated code for CK inductor backend
* See torch._inductor.codegen.rocm.ck_universal_gemm_template.CKGemmTemplate
*
* Template instance CKGemmOperation(a_layout='Col', b_layout='Row', ds_layouts=(), c_layout='Row', a_element_dtype='BF16', b_element_dtype='BF16', ds_element_dtypes=(), c_element_dtype='BF16', acc_dtype='F32', c_shuffle_dtype='BF16', a_elementwise_op='PassThrough', b_elementwise_op='PassThrough', c_elementwise_op='PassThrough', gemm_specialization='GemmSpecialization::MNPadding', block_size=256, m_per_block=224, n_per_block=256, k_per_block=64, a_k1=8, b_k1=8, m_per_xdl=16, n_per_xdl=16, m_xdl_per_wave=7, n_xdl_per_wave=8, a_block_transfer_thread_cluster_lengths_ak0_m_ak1=(8, 32, 1), a_block_transfer_thread_cluster_arrange_order=(0, 2, 1), a_block_transfer_src_access_order=(0, 2, 1), a_block_transfer_src_vector_dim=1, a_block_transfer_src_scalar_per_vector=8, a_block_transfer_dst_scalar_per_vector_ak1=8, a_block_lds_extra_m=0, b_block_transfer_thread_cluster_lengths_bk0_n_bk1=(8, 32, 1), b_block_transfer_thread_cluster_arrange_order=(0, 2, 1), b_block_transfer_src_access_order=(0, 2, 1), b_block_transfer_src_vector_dim=1, b_block_transfer_src_scalar_per_vector=8, b_block_transfer_dst_scalar_per_vector_bk1=8, b_block_lds_extra_n=0, c_shuffle_m_xdl_per_wave_per_shuffle=1, c_shuffle_n_xdl_per_wave_per_shuffle=2, c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block=(1, 32, 1, 8), c_shuffle_block_transfer_scalar_per_vector_n_per_block=(8,), block_gemm_pipeline_scheduler='BlockGemmPipelineScheduler::Intrawave', block_gemm_pipeline_version='BlockGemmPipelineVersion::v3', a_compute_dtype=None, b_compute_dtype=None)
*
* torch.__version__='2.6.0a.post20241218'
* torch.version.git_version=Unknown
*/

    
#include <exception>
#include <iostream>
#include <memory>
#include <random>
#include <vector>

// CK headers

#ifdef DEBUG_LOG
#define DEBUG_LOG_TMP DEBUG_LOG
#undef DEBUG_LOG
#else
#define DEBUG_LOG_TMP 0
#endif
#include "ck/ck.hpp"
#undef DEBUG_LOG
#define DEBUG_LOG DEBUG_LOG_TMP

#include "ck/utility/data_type.hpp"
#include "ck/library/utility/check_err.hpp"
#include "ck/library/utility/device_memory.hpp"
#include "ck/library/utility/fill.hpp"
#include "ck/library/utility/host_tensor.hpp"
#include "ck/library/utility/host_tensor_generator.hpp"
#include "ck/library/utility/literals.hpp"

// CK GEMM header(s)

#include "ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp"

    
// We compile all models with -fvisibility=hidden. Any symbols that need to be
// exposed in the final shared library must be declared with PT_EXPORT to make
// them visible.
#ifdef __GNUC__ // Applies to any compiler with GNU extensions (clang and g++)
#define PT_EXPORT __attribute__((__visibility__("default")))
#else
#ifdef _WIN32
#define PT_EXPORT __declspec(dllexport)
#else
#define PT_EXPORT
#endif
#endif

// as long as there is no custom arithmetic it's fine
using bfloat16 = uint16_t;
using float8_e4m3fnuz = uint8_t;
using float8_e5m2fnuz = uint8_t;

// CK globals

template <ck::index_t... Is>
using S = ck::Sequence<Is...>;

template<typename... Ts>
using Tuple = ck::Tuple<Ts...>;

using PassThrough = ck::tensor_operation::element_wise::PassThrough;
using Bilinear = ck::tensor_operation::element_wise::Bilinear;
using Scale = ck::tensor_operation::element_wise::Scale;
using ScaleAdd = ck::tensor_operation::element_wise::ScaleAdd;
using MultiplyMultiply = ck::tensor_operation::element_wise::MultiplyMultiply;

// see "composable_kernel/include/ck/utility/data_type.hpp"
using F8  = ck::f8_t;
using BF8 = ck::bf8_t;
using F16 = ck::half_t;
using F32 = float;
// using F64 = double;
using BF16 = ck::bhalf_t;
// using I32 = int32_t;
// using I8 = int8_t;
// using I4 = ck::int4_t;

#if DEBUG_LOG
static constexpr auto kDEBUG_LOG = 1;
#else
static constexpr auto kDEBUG_LOG = 0;
#endif

// CK GEMM globals

using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;

using BlockGemmPipelineScheduler = ck::BlockGemmPipelineScheduler;
using GemmSpecialization = ck::tensor_operation::device::GemmSpecialization;
using BlockGemmPipelineVersion = ck::BlockGemmPipelineVersion;

struct MultiplyMultiplyAdd {
    template <typename E, typename C, typename D0, typename D1, typename D2>
    __host__ __device__ constexpr void
    operator()(E& e, const C& c, const D0& d0, const D1& d1, const D2& d2) const {
        e = ck::type_convert<E>(
           ck::type_convert<float>(c)
           * ck::type_convert<float>(d0)
           * ck::type_convert<float>(d1)
           + ck::type_convert<float>(d2)
        );
    }
};

    
    // Gemm operator ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVCol_KblayoutVRow_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationMNPadding_KblocksizeV256_KmperblockV224_KnperblockV256_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV16_KnperxdlV16_KmxdlperwaveV7_KnxdlperwaveV8_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV0x2x1_KablocktransfersrcaccessorderV0x2x1_KablocktransfersrcvectordimV1_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV0x2x1_KbblocktransfersrcaccessorderV0x2x1_KbblocktransfersrcvectordimV1_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV2_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x32x1x8_KcshuffleblocktransferscalarpervectornperblockV8_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone
    using Operation_ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVCol_KblayoutVRow_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationMNPadding_KblocksizeV256_KmperblockV224_KnperblockV256_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV16_KnperxdlV16_KmxdlperwaveV7_KnxdlperwaveV8_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV0x2x1_KablocktransfersrcaccessorderV0x2x1_KablocktransfersrcvectordimV1_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV0x2x1_KbblocktransfersrcaccessorderV0x2x1_KbblocktransfersrcvectordimV1_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV2_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x32x1x8_KcshuffleblocktransferscalarpervectornperblockV8_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone =
        ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3<
            /* a_layout */ Col,
            /* b_layout */ Row,
            /* ds_layouts */ Tuple<>,
            /* c_layout */ Row,
            /* a_element_dtype */ BF16,
            /* b_element_dtype */ BF16,
            /* ds_element_dtypes */ Tuple<>,
            /* c_element_dtype */ BF16,
            /* acc_dtype */ F32,
            /* c_shuffle_dtype */ BF16,
            /* a_elementwise_op */ PassThrough,
            /* b_elementwise_op */ PassThrough,
            /* c_elementwise_op */ PassThrough,
            /* gemm_specialization */ GemmSpecialization::MNPadding,
            /* block_size */ 256,
            /* m_per_block */ 224,
            /* n_per_block */ 256,
            /* k_per_block */ 64,
            /* a_k1 */ 8,
            /* b_k1 */ 8,
            /* m_per_xdl */ 16,
            /* n_per_xdl */ 16,
            /* m_xdl_per_wave */ 7,
            /* n_xdl_per_wave */ 8,
            /* a_block_transfer_thread_cluster_lengths_ak0_m_ak1 */ S<8, 32, 1>,
            /* a_block_transfer_thread_cluster_arrange_order */ S<0, 2, 1>,
            /* a_block_transfer_src_access_order */ S<0, 2, 1>,
            /* a_block_transfer_src_vector_dim */ 1,
            /* a_block_transfer_src_scalar_per_vector */ 8,
            /* a_block_transfer_dst_scalar_per_vector_ak1 */ 8,
            /* a_block_lds_extra_m */ 0,
            /* b_block_transfer_thread_cluster_lengths_bk0_n_bk1 */ S<8, 32, 1>,
            /* b_block_transfer_thread_cluster_arrange_order */ S<0, 2, 1>,
            /* b_block_transfer_src_access_order */ S<0, 2, 1>,
            /* b_block_transfer_src_vector_dim */ 1,
            /* b_block_transfer_src_scalar_per_vector */ 8,
            /* b_block_transfer_dst_scalar_per_vector_bk1 */ 8,
            /* b_block_lds_extra_n */ 0,
            /* c_shuffle_m_xdl_per_wave_per_shuffle */ 1,
            /* c_shuffle_n_xdl_per_wave_per_shuffle */ 2,
            /* c_shuffle_block_transfer_cluster_lengths_m_block_m_per_block_n_block_n_per_block */ S<1, 32, 1, 8>,
            /* c_shuffle_block_transfer_scalar_per_vector_n_per_block */ S<8>,
            /* block_gemm_pipeline_scheduler */ BlockGemmPipelineScheduler::Intrawave,
            /* block_gemm_pipeline_version */ BlockGemmPipelineVersion::v3>;

    extern "C" {
    PT_EXPORT int rocm_ck_gemm_template(const bfloat16* X, const bfloat16* W, bfloat16* Y, int32_t M, int32_t N, int32_t K, int32_t LDA, int32_t LDB, int32_t LDC, int32_t LDD, size_t* workspace_size, uint8_t* workspace, hipStream_t stream) {
        auto gemm = 
    Operation_ck_devicegemm_multid_xdl_shuffle_v3_KalayoutVCol_KblayoutVRow_KdslayoutsV_KclayoutVRow_KaelementdtypeVBF16_KbelementdtypeVBF16_KdselementdtypesV_KcelementdtypeVBF16_KaccdtypeVF32_KcshuffledtypeVBF16_KaelementwiseopVPassThrough_KbelementwiseopVPassThrough_KcelementwiseopVPassThrough_KgemmspecializationVGemmSpecializationMNPadding_KblocksizeV256_KmperblockV224_KnperblockV256_KkperblockV64_Kak1V8_Kbk1V8_KmperxdlV16_KnperxdlV16_KmxdlperwaveV7_KnxdlperwaveV8_Kablocktransferthreadclusterlengthsak0mak1V8x32x1_KablocktransferthreadclusterarrangeorderV0x2x1_KablocktransfersrcaccessorderV0x2x1_KablocktransfersrcvectordimV1_KablocktransfersrcscalarpervectorV8_Kablocktransferdstscalarpervectorak1V8_KablockldsextramV0_Kbblocktransferthreadclusterlengthsbk0nbk1V8x32x1_KbblocktransferthreadclusterarrangeorderV0x2x1_KbblocktransfersrcaccessorderV0x2x1_KbblocktransfersrcvectordimV1_KbblocktransfersrcscalarpervectorV8_Kbblocktransferdstscalarpervectorbk1V8_KbblockldsextranV0_KcshufflemxdlperwavepershuffleV1_KcshufflenxdlperwavepershuffleV2_KcshuffleblocktransferclusterlengthsmblockmperblocknblocknperblockV1x32x1x8_KcshuffleblocktransferscalarpervectornperblockV8_KblockgemmpipelineschedulerVBlockGemmPipelineSchedulerIntrawave_KblockgemmpipelineversionVBlockGemmPipelineVersionv3_KacomputedtypeVNone_KbcomputedtypeVNone {};
        auto invoker = gemm.MakeInvoker();
        
        auto argument = gemm.MakeArgument(
            reinterpret_cast<const BF16*>(X),
            reinterpret_cast<const BF16*>(W),
            std::array<const void*, 0>{  },
            reinterpret_cast<BF16*>(Y),
            M,
            N,
            K,
            LDA,
            LDB,
            std::array<ck::index_t, 0>{  },
            LDC,
            1, // kBatch
            PassThrough {},
            PassThrough {},
            PassThrough {} // c_elementwise_op
        );
        
        if (!gemm.IsSupportedArgument(argument)) {
            // we do our best to statically avoid this case in `filter_op`
            std::cerr << "invalid argument for gemm instance " << gemm.GetTypeString() << std::endl;
            argument.Print();
            return -23;
        }
        if (workspace_size) {
            *workspace_size = gemm.GetWorkSpaceSize(&argument);
            return 0;
        }
        // run the kernel
        #ifdef GENERATE_CK_STANDALONE_RUNNER
        const auto stream_config = StreamConfig{
            stream,
            /* time kernel */ 1,
            /* log level */ 1,
            /* n_cold_iter */ 100,
            /* n_hot_iter */ 100,
            /* flush_l2_cache */ 1,
            /* rotate_count */ 5};
        #else
        const auto stream_config = StreamConfig{stream, /* time kernel */ false, /* log level */ 0};
        #endif

        const float elapsed_time = invoker.Run(argument, stream_config);

        #ifdef GENERATE_CK_STANDALONE_RUNNER
        std::cout << "elapsed time: " << elapsed_time << " ms" << std::endl;
        #else
        (void)elapsed_time;
        #endif
        return 0;
    } // kernel definition
    } // extern C
    ⏎                   
```

</details>

### Versions

<details>
<summary>Very long version info on python3.12-torch-2.6.0a-nightly-20241218:</summary>

```
env | grep -E -i '(torch|hsa|rocm|rocr|ccl).*='
TORCHINDUCTOR_CK_DIR=/nix/store/8wrnjabmr02rhknibrjr6qya3fimml2f-python3.12-ck4inductor-6.4.0a20241217/lib/python3.12/site-packages/ck4inductor
TORCHINDUCTOR_CACHE_DIR=~/ml-cache/torchinductor
TORCHINDUCTOR_AUTOGRAD_CACHE=1
ROCM_BUILD_ID=release-nixos-60300
TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=CK,ATEN,TRITON,CPP
ROCM_LIBPATCH_VERSION=60300
ROCM_PATH=/nix/store/8zk35m6vnbcf339zi9k6jra4xs4ipd49-rocm-hip-libraries-meta
TORCHINDUCTOR_FX_GRAPH_CACHE=1
TORCH_ROCM_FA_PREFER_CK=1
acl-2.3.2
aotriton-unstable-20241122
attr-2.5.2
bash-5.2p37
binutils-2.43.1
binutils-2.43.1-lib
binutils-wrapper-2.43.1
blas-3
blas-3-dev
brotli-1.1.0-lib
bzip2-1.0.8
bzip2-1.0.8-bin
bzip2-1.0.8-dev
clang-rocm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
clang-rocm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-dev
clang-rocm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-lib
clr-6.3.0
clr-6.3.0-icd
cmake-3.30.5
compiler-rt-libc-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
compiler-rt-libc-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-dev
coreutils-9.5
curl-8.11.0
elfutils-0.191
expand-response-params
expat-2.6.4
expat-2.6.4-dev
find-xml-catalogs-hook
gcc-13.3.0
gcc-13.3.0-lib
gcc-13.3.0-libgcc
gcc-prefix
gcc-wrapper-13.3.0
gdbm-1.24
gdbm-1.24-dev
gdbm-1.24-lib
getopt-1.1.6
gfortran-13.3.0
gfortran-13.3.0-lib
gfortran-13.3.0-libgcc
gfortran-wrapper-13.3.0
glibc-2.40-36
glibc-2.40-36-bin
glibc-2.40-36-dev
gmp-6.3.0
gmp-with-cxx-6.3.0
gnugrep-3.11
hipblas-6.3.0
hipblas-common-unstable
hipblaslt-6.3.0
hipcub-6.3.0
hipfft-6.3.0
hipfort-6.3.0
hipify-6.3.0
hiprand-6.3.0
hipsolver-6.3.0
hipsparse-6.3.0
hwdata-0.388
hwloc-2.11.2-lib
isl-0.20
keyutils-1.6.3-lib
krb5-1.21.3-lib
libarchive-3.7.7-lib
libcap-2.70-lib
libdrm-2.4.123
libevent-2.1.12
libfabric-1.22.0
libffi-3.4.6
libffi-3.4.6-dev
libgcrypt-1.10.3-lib
libglvnd-1.7.0
libgpg-error-1.50
libidn2-2.3.7
libmpc-1.3.1
libnl-3.10.0
libpciaccess-0.18.1
libpfm-4.13.0
libpsl-0.21.5
libpsm2-12.0.1
libsodium-1.0.20
libssh2-1.11.1
libunistring-1.2
libunwind-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
libunwind-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-dev
libuv-1.48.0
libX11-1.8.10
libXau-1.0.11
libxcb-1.17.0
libxcrypt-4.4.36
libXdmcp-1.1.5
libXext-1.3.6
libxml2-2.13.4
libxml2-2.13.4-bin
libxml2-2.13.4-dev
libyaml-0.2.5
linux-headers-6.10
lld-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
lld-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-dev
lld-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-lib
llhttp-9.2.1
llvm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
llvm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-dev
llvm-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec-lib
llvm-binutils-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
llvm-binutils-wrapper-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
lsb_release
mailcap-2.1.54
miopen-6.3.0
miopen-gfx1030.kdb
miopen-gfx900.kdb
miopen-gfx906.kdb
miopen-gfx908.kdb
miopen-gfx90a.kdb
mpdecimal-4.0.0
mpdecimal-4.0.0-cxx
mpdecimal-4.0.0-dev
mpfr-4.2.1
mpich-4.2.3
mpich-4.2.3-doc
mpich-4.2.3-man
munge-0.5.16
ncurses-6.4.20221231
ncurses-6.4.20221231-dev
ncurses-6.4.20221231-man
nghttp2-1.64.0-lib
nss-cacert-3.104
numactl-2.0.18
openblas-0.3.28
openmp-18.0.0-d6e55e17f328a495bc32fddb7826e673ac9766ec
openmpi-5.0.6
openmpi-5.0.6-man
openssl-3.3.2
openssl-3.3.2-bin
openssl-3.3.2-dev
pcre2-10.44
perl-5.40.0
pkg-config-0.29.2
pkg-config-wrapper-0.29.2
pmix-5.0.4
prrte-3.0.7
publicsuffix-list-0-unstable-2024-10-25
python3.12-aiodns-3.2.0
python3.12-aiohappyeyeballs-2.4.2
python3.12-aiohttp-3.10.10
python3.12-aiosignal-1.3.1
python3.12-async-timeout-4.0.3
python3.12-attrs-24.2.0
python3.12-bcrypt-4.2.0
python3.12-brotli-1.1.0
python3.12-brotlicffi-1.1.0.0
python3.12-certifi-2024.08.30
python3.12-cffi-1.17.1
python3.12-charset-normalizer-3.3.2
python3.12-ck4inductor-6.4.0a20241217
python3.12-cryptography-43.0.1
python3.12-filelock-3.16.1
python3.12-frozenlist-1.4.1
python3.12-fsspec-2024.3.0
python3.12-huggingface-hub-0.26.2
python3.12-idna-3.10
python3.12-joblib-1.4.2
python3.12-lz4-4.3.3
python3.12-markdown-it-py-3.0.0
python3.12-mdurl-0.1.2
python3.12-msgpack-1.1.0
python3.12-multidict-6.1.0
python3.12-numpy-1.26.4
python3.12-orjson-3.10.7
python3.12-packaging-24.1
python3.12-pandas-2.2.3
python3.12-paramiko-3.5.0
python3.12-pip-24.0
python3.12-psutil-6.0.0
python3.12-pycares-4.4.0
python3.12-pycparser-2.22
python3.12-pygments-2.18.0
python3.12-pynacl-1.5.0
python3.12-pyspnego-0.11.1
python3.12-python-dateutil-2.9.0.post0
python3.12-pytz-2024.2
python3.12-pyyaml-6.0.2
python3.12-requests-2.32.3
python3.12-rich-13.8.1
python3.12-simplejson-3.19.3
python3.12-six-1.16.0
python3.12-smbprotocol-1.14.0
python3.12-tensile-6.3.0
python3.12-tensilelite-6.3.0
python3.12-torch-2.6.0a-nightly-20241218
python3.12-torch-2.6.0a-nightly-20241218-lib
python3.12-tqdm-4.66.5
python3.12-typing-extensions-4.12.2
python3.12-tzdata-2024.2
python3.12-ujson-5.10.0
python3.12-urllib3-2.2.3
python3.12-yarl-1.13.1
python3.12-zstd-1.5.5.1
python3-3.12.7
rccl-6.3.0
rdma-core-54.0
readline-8.2p13
readline-8.2p13-dev
rhash-1.4.4
rocalution-6.3.0
rocblas-6.3.0
rocfft-6.3.0
rocm-comgr-6.3.0
rocm-core-6.3.0
rocmcxx
rocm-device-libs-6.3.0
rocminfo-6.3.0
rocm-llvm-merge
rocm-merged
rocm-runtime-6.3.0
rocm-smi-6.3.0
rocprim-6.3.0
rocprofiler-register-6.3.0
rocrand-6.3.0
rocsolver-6.3.0
rocsparse-6.3.0
rocthrust-6.3.0
roctracer-6.3.0
shell-deps
source
sqlite-3.46.1
sqlite-3.46.1-bin
sqlite-3.46.1-dev
strip.sh
systemd-minimal-libs-256.7
tzdata-2024b
ucc-1.3.0
ucx-1.17.0
unpack-composable_kernel-6.4.0a20241217
util-linux-minimal-2.39.4-lib
xgcc-13.3.0-libgcc
xz-5.6.3
xz-5.6.3-bin
xz-5.6.3-dev
zlib-1.3.1
zlib-1.3.1-dev
zstd-1.5.6
zstd-1.5.6-bin
zstd-1.5.6-dev
Collecting environment information...
PyTorch version: 2.6.0a.post20241218
Is debug build: False
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: 6.3.42131-0

OS: NixOS 25.05 (Warbler) (x86_64)
GCC version: Could not collect
Clang version: 18.0.0git
CMake version: version 3.30.5
Libc version: glibc-2.40

Python version: 3.12.7 (main, Oct  1 2024, 02:05:46) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.12.4-x86_64-with-glibc2.40
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: AMD Instinct MI100 (gfx908:sramecc+:xnack-)
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: 6.3.42131
MIOpen runtime version: 3.3.0
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               128
On-line CPU(s) list:                  0-127
Vendor ID:                            AuthenticAMD
Model name:                           AMD Eng Sample: 100-000000425_37/24_N
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   64
Socket(s):                            1
Stepping:                             0
Frequency boost:                      enabled
CPU(s) scaling MHz:                   35%
CPU max MHz:                          5616.0000
CPU min MHz:                          400.0000
BogoMIPS:                             7400.01
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca fsrm debug_swap
Virtualization:                       AMD-V
L1d cache:                            2 MiB (64 instances)
L1i cache:                            2 MiB (64 instances)
L2 cache:                             32 MiB (64 instances)
L3 cache:                             256 MiB (8 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-127
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; Safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.6.0a0.post20241218
[pip3] triton==3.2.0
[conda] Could not collect
```

</details>


cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @naromero77amd @chauhang @penguinwu