[QST] uncoalesced shared accesses via Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementA>

Hi, I got a problem with sm80 cutlass 3.0 style kernel.
My kernel lies below.

template <
    typename ElementA_,
    typename ElementB_,
    typename ElementC_,
    typename ElementAccum_,
    int TileM_, int TileN_, int TileK_,
    int Pipe_>
struct AmepereCutlass3xGemm
{
public:
    using ElementA = ElementA_;
    using ElementB = ElementB_;
    using ElementC = ElementC_;
    using ElementAccum = ElementAccum_;

    using TileM = Int<TileM_>;
    using TileN = Int<TileN_>;
    using TileK = Int<TileK_>;

    using Pipe = Int<Pipe_>;

    using TilerA = Shape<TileM, TileK>;
    using TilerB = Shape<TileN, TileK>;
    using TilerC = Shape<TileM, TileN>;

    using LayoutA = cutlass::layout::RowMajor;
    using LayoutB = cutlass::layout::ColumnMajor;
    using LayoutC = cutlass::layout::RowMajor;

    using StrideA = cutlass::detail::TagToStrideA_t<LayoutA>;
    using StrideB = cutlass::detail::TagToStrideB_t<LayoutB>;
    using StrideC = cutlass::detail::TagToStrideC_t<LayoutC>;

    using MmaInstruction = typename MmaInstruction<ElementA>::mma;
    using TiledMma = TiledMMA<
        MMA_Atom<MmaInstruction>,
        Layout<Shape<_4, _2, _1>>,
        Tile<_64, _64, X>>;

    static constexpr int MaxThreadsPerBlock = size(TiledMma{});
    static constexpr int MinBlocksPerMultiprocessor = 1;

    static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
    static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
    static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;

    using G2STheadLayout = Layout<Shape<_64, _4>, Stride<_4, _1>>;
    using EpilogueThreadLayout = Layout<Shape<_64, _4>, Stride<_4, _1>>;

    CUTE_STATIC_ASSERT(size(G2STheadLayout{}) == size(EpilogueThreadLayout{}));
    CUTE_STATIC_ASSERT(size(G2STheadLayout{}) == size(TiledMma{}));

    using GmemTiledCopyA = decltype(make_tiled_copy(
        Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementA>{},
        G2STheadLayout{},
        Layout<Shape<_1, Int<AlignmentA>>>{}));

    using SmemLayoutAtomA = decltype(composition(Swizzle<3, 3, 3>{},
                                                 Layout<Shape<_32, _32>,
                                                        Stride<_32, _1>>{}));

    using SmemCopyAtomA = Copy_Atom<SM75_U32x4_LDSM_N, ElementA>;

    using GmemTiledCopyB = decltype(make_tiled_copy(
        Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementB>{},
        G2STheadLayout{},
        Layout<Shape<_1, Int<AlignmentB>>>{}));
    using SmemLayoutAtomB = decltype(composition(Swizzle<3, 3, 3>{},
                                                 Layout<Shape<_32, _32>,
                                                        Stride<_32, _1>>{}));
    using SmemCopyAtomB = Copy_Atom<SM75_U32x4_LDSM_N, ElementB>;

    using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveMma<
        cutlass::gemm::MainloopSm80CpAsync<Pipe::value>,
        Shape<TileM, TileN, TileK>,
        ElementA,
        StrideA,
        ElementB,
        StrideB,
        TiledMma,
        GmemTiledCopyA,
        SmemLayoutAtomA,
        SmemCopyAtomA,
        cute::identity, // TransformA_
        GmemTiledCopyB,
        SmemLayoutAtomB,
        SmemCopyAtomB,
        cute::identity>; // TransformB_

    using EpilogeOp = cutlass::epilogue::thread::LinearCombination<
        ElementC,
        128 / cutlass::sizeof_bits<ElementC>::value,
        ElementAccum,
        ElementAccum,
        cutlass::epilogue::thread::ScaleType::Nothing>;

    // using CollectiveEpilogue = cutlass::epilogue::collective::DefaultEpilogue<
    //     ElementC,
    //     StrideC,
    //     StrideC,
    //     EpilogeOp,
    //     cutlass::gemm::EpilogueDefault>;

    using SwizzledSmemLayout = decltype(composition(Swizzle<3, 3, 3>{}, Layout<Shape<TileM, TileN>, Stride<TileN, _1>>{}));

    using CopyAtomR2S = Copy_Atom<AutoVectorizingCopy, ElementAccum>;
    using TiledCopyS2R = decltype(make_tiled_copy(
        Copy_Atom<UniversalCopy<uint128_t>, ElementAccum>{},
        EpilogueThreadLayout{},
        Layout<Shape<_1, _8>>{}));

    using CopyAtomR2G = Copy_Atom<UniversalCopy<uint128_t>, ElementC>;

    using CollectiveEpilogue = cutlass::epilogue::collective::Epilogue<
        StrideC,
        StrideC,
        EpilogeOp,
        SwizzledSmemLayout,
        CopyAtomR2S,
        TiledCopyS2R,
        CopyAtomR2G,
        cutlass::epilogue::EpilogueSimtVectorized>;
    using GemmKernel = cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int>,
                                                            CollectiveMainloop,
                                                            CollectiveEpilogue>;

    using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;

    void run(const ElementA __restrict__ *A, const ElementB __restrict__ *B, ElementC *C, const int M, const int N, const int K)
    {
        Gemm gemm;

        auto ldA = K;
        auto ldB = K;
        auto ldC = N;
        auto dA = typename CollectiveMainloop::StrideA(ldA, cute::Int<1>{}, _);
        auto dB = typename CollectiveMainloop::StrideB(ldB, cute::Int<1>{}, _);
        auto dC = typename CollectiveEpilogue::StrideC(ldC, cute::Int<1>{}, _);
        typename Gemm::Arguments args{
            cutlass::gemm::GemmUniversalMode::kGemm,
            {M, N, K},
            {A, dA, B, dB},
            {{1.0f, 0.f}, C, dC, C, dC}};
        auto status = gemm(args);
        if (status != cutlass::Status::kSuccess)
        {
            std::cout << "Fail" << std::endl;
        }
    }
};

I test with AmepereCutlass3xGemm<cutlass::bfloat16_t, cutlass::bfloat16_t, cutlass::bfloat16_t, float, 64, 64, 64, 4>{};. After profiling with ncu, there're uncoalesced shared accesses. (reference code in cutlass)

And I finally find out that it is caused by Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementA> . The tAgAk and tAsA are 8 alignment and ncu says shared wavefronts are excessive.

Start with include/cutlass/gemm/collective/sm80_mma_multistage.hpp around the referenced line 558, then reproduce the kernel with the provided template and inspect its Nsight Compute shared-memory metrics. Compare the tAgAk and tAsA layouts for Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, ElementA>. Done means determining the cause of the excessive shared wavefronts and validating the behavior against the reference code.