776 lines
30 KiB
C++
776 lines
30 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: BSD-3-Clause
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are met:
|
|
*
|
|
* 1. Redistributions of source code must retain the above copyright notice, this
|
|
* list of conditions and the following disclaimer.
|
|
*
|
|
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
* this list of conditions and the following disclaimer in the documentation
|
|
* and/or other materials provided with the distribution.
|
|
*
|
|
* 3. Neither the name of the copyright holder nor the names of its
|
|
* contributors may be used to endorse or promote products derived from
|
|
* this software without specific prior written permission.
|
|
*
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
*
|
|
**************************************************************************************************/
|
|
|
|
#pragma once
|
|
|
|
#include "cutlass/relatively_equal.h"
|
|
#include "cutlass_unit_test.h"
|
|
#include "cutlass/util/reference/host/tensor_compare.h"
|
|
|
|
#include <iostream>
|
|
|
|
#include <thrust/host_vector.h>
|
|
#include <thrust/device_vector.h>
|
|
|
|
#include <cute/tensor.hpp>
|
|
|
|
using namespace cute;
|
|
|
|
template<typename T>
|
|
struct fp64_tester {
|
|
using value_type = double;
|
|
};
|
|
|
|
template<typename T>
|
|
struct fp64_tester<complex<T>> {
|
|
using value_type = complex<double>;
|
|
};
|
|
|
|
template<class TA,
|
|
class TB,
|
|
class TC,
|
|
class ALayout, // logical shape (M, K)
|
|
class BLayout, // logical shape (N, K)
|
|
class CLayout> // logical shape (M, N)
|
|
auto host_generate_gemm_inputs(
|
|
ALayout a_layout,
|
|
BLayout b_layout,
|
|
CLayout c_layout
|
|
) {
|
|
thrust::host_vector<TA> h_a(cosize(a_layout));
|
|
thrust::host_vector<TB> h_b(cosize(b_layout));
|
|
thrust::host_vector<TC> h_c(cosize(c_layout));
|
|
thrust::host_vector<TC> h_c_out(cosize(c_layout));
|
|
|
|
auto h_a_tensor = make_tensor(h_a.data(), a_layout);
|
|
auto h_b_tensor = make_tensor(h_b.data(), b_layout);
|
|
auto h_c_tensor = make_tensor(h_c.data(), c_layout);
|
|
size_t max_size = std::max<size_t>({static_cast<size_t>(size(a_layout)),
|
|
static_cast<size_t>(size(b_layout)),
|
|
static_cast<size_t>(size(c_layout))});
|
|
for (size_t i = 0; i < max_size; ++i) {
|
|
double di = static_cast<double>(i);
|
|
if(i < size(a_layout)) {
|
|
h_a_tensor(i) = static_cast<TA>(di / size(a_layout));
|
|
}
|
|
if(i < size(b_layout)) {
|
|
h_b_tensor(i) = static_cast<TB>(di / size(a_layout));
|
|
}
|
|
if(i < size(c_layout)) {
|
|
h_c_tensor(i) = static_cast<TC>((di*di) / size(a_layout));
|
|
}
|
|
}
|
|
|
|
return std::make_tuple(h_a, h_b, h_c, h_c_out);
|
|
}
|
|
|
|
template<class Alpha, class EngineA, class ALayout,
|
|
class EngineB, class BLayout,
|
|
class Beta, class EngineC, class CLayout,
|
|
class ALoadTransform = cute::identity,
|
|
class BLoadTransform = cute::identity,
|
|
class CLoadTransform = cute::identity,
|
|
class CStoreTransform = cute::identity>
|
|
thrust::host_vector<typename EngineC::value_type>
|
|
host_reference_gemm(Alpha alpha,
|
|
Tensor<EngineA, ALayout> const& h_a_tensor,
|
|
Tensor<EngineB, BLayout> const& h_b_tensor,
|
|
Beta beta,
|
|
Tensor<EngineC, CLayout> const& h_c_tensor,
|
|
ALoadTransform const& a_load_transform = {},
|
|
BLoadTransform const& b_load_transform = {},
|
|
CLoadTransform const& c_load_transform = {},
|
|
CStoreTransform const& c_store_transform = {})
|
|
{
|
|
// Cannot use ::value_type because it propagates to complex::value_type,
|
|
// so ViewEngine<complex<double>>::value_type == double
|
|
using TA = remove_cv_t<typename EngineA::element_type>;
|
|
using TB = remove_cv_t<typename EngineB::element_type>;
|
|
using TC = remove_cv_t<typename EngineC::element_type>;
|
|
|
|
using tester = fp64_tester<TC>;
|
|
using ABC_64 = typename tester::value_type;
|
|
|
|
static_assert(std::is_same_v<typename fp64_tester<TA>::value_type, typename fp64_tester<TB>::value_type>);
|
|
static_assert(std::is_same_v<typename fp64_tester<TB>::value_type, typename fp64_tester<TC>::value_type>);
|
|
|
|
thrust::host_vector<TC> h_c_ref(cosize(h_c_tensor.layout()), static_cast<TC>(0.0));
|
|
auto h_c_ref_tensor = make_tensor(h_c_ref.data(), h_c_tensor.layout());
|
|
// A * B
|
|
for (int k = 0; k < size<1>(h_a_tensor); k++) {
|
|
for (int m = 0; m < size<0>(h_a_tensor); m++) {
|
|
for (int n = 0; n < size<0>(h_b_tensor); n++) {
|
|
const auto a_value = a_load_transform(h_a_tensor(m, k));
|
|
const auto b_value = b_load_transform(h_b_tensor(n, k));
|
|
const auto a_value_fp64 = static_cast<ABC_64>(a_value);
|
|
const auto b_value_fp64 = static_cast<ABC_64>(b_value);
|
|
h_c_ref_tensor(m, n) += static_cast<TC>(a_value_fp64 * b_value_fp64);
|
|
}
|
|
}
|
|
}
|
|
// C = A*B + C
|
|
for (int i = 0; i < size(h_c_ref_tensor); i++) {
|
|
const auto ab_value_fp64 = static_cast<ABC_64>(h_c_ref_tensor(i));
|
|
const auto c_value_fp64 = static_cast<ABC_64>(c_load_transform(h_c_tensor(i)));
|
|
h_c_ref_tensor(i) = c_store_transform(static_cast<TC>(alpha * ab_value_fp64 + beta * c_value_fp64));
|
|
}
|
|
|
|
return h_c_ref;
|
|
}
|
|
|
|
template<class EngineC, class CLayout>
|
|
void verify_gemm_correctness(cute::Tensor<EngineC, CLayout> const& h_c_out_tensor,
|
|
cute::Tensor<EngineC, CLayout> const& h_c_ref_tensor)
|
|
{
|
|
// Cannot use ::value_type because it propagates to complex::value_type,
|
|
// so ViewEngine<complex<double>>::value_type == double
|
|
using TC = remove_cv_t<typename EngineC::element_type>;
|
|
|
|
using tester = fp64_tester<TC>;
|
|
using ABC_64 = typename tester::value_type;
|
|
|
|
for (int i = 0; i < size(h_c_ref_tensor); i++) {
|
|
ABC_64 h_c_ref_i = h_c_ref_tensor(i);
|
|
ABC_64 h_c_out_i = h_c_out_tensor(i);
|
|
double epsilon(0.1f);
|
|
double nonzero_floor(std::numeric_limits<double>::min());
|
|
bool passed = cutlass::relatively_equal(h_c_out_i, h_c_ref_i, epsilon, nonzero_floor);
|
|
ASSERT_TRUE(passed) << i << " - result:" << h_c_out_i << " expected:" << h_c_ref_i;
|
|
}
|
|
}
|
|
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class GMemALayout,
|
|
class GMemBLayout,
|
|
class GMemCLayout,
|
|
class SMemALayout,
|
|
class SMemBLayout,
|
|
class SMemCLayout,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class Alpha,
|
|
class Beta,
|
|
class TiledMma,
|
|
class ALoadTransform,
|
|
class BLoadTransform,
|
|
class CLoadTransform,
|
|
class CStoreTransform,
|
|
class SMemCopyOpA,
|
|
class SMemCopyOpB,
|
|
class SMemCopyLdOpC,
|
|
class SMemCopyStOpC>
|
|
__launch_bounds__(ThreadBlockSize) __global__ void
|
|
cooperative_gemm_kernel(GMemALayout gmem_a_layout,
|
|
GMemBLayout gmem_b_layout,
|
|
GMemCLayout gmem_c_layout,
|
|
SMemALayout smem_a_layout,
|
|
SMemBLayout smem_b_layout,
|
|
SMemCLayout smem_c_layout,
|
|
TA const* a,
|
|
TB const* b,
|
|
TC const* c,
|
|
TC * c_out,
|
|
Alpha const alpha,
|
|
Beta const beta,
|
|
TiledMma tiled_mma,
|
|
ALoadTransform a_load_transform,
|
|
BLoadTransform b_load_transform,
|
|
CLoadTransform c_load_transform,
|
|
CStoreTransform c_store_transform,
|
|
SMemCopyOpA a_copy_op,
|
|
SMemCopyOpB b_copy_op,
|
|
SMemCopyLdOpC c_copy_ld_op,
|
|
SMemCopyStOpC c_copy_st_op)
|
|
{
|
|
using namespace cute;
|
|
|
|
Tensor g_a_tensor = make_tensor(make_gmem_ptr(a), gmem_a_layout);
|
|
Tensor g_b_tensor = make_tensor(make_gmem_ptr(b), gmem_b_layout);
|
|
Tensor g_c_tensor = make_tensor(make_gmem_ptr(c), gmem_c_layout);
|
|
Tensor g_c_out_tensor = make_tensor(make_gmem_ptr(c_out), gmem_c_layout);
|
|
|
|
constexpr uint32_t copy_max_vec_bytes = CopyMaxVecBits / 8;
|
|
|
|
extern __shared__ float4 smem_buf[];
|
|
auto* smem_ptr = reinterpret_cast<unsigned char*>(smem_buf);
|
|
auto* smem_ptr_a = smem_ptr;
|
|
auto* smem_ptr_b = smem_ptr_a + round_up((sizeof(TA) * cosize(smem_a_layout)), copy_max_vec_bytes);
|
|
auto* smem_ptr_c = smem_ptr_b + round_up((sizeof(TB) * cosize(smem_b_layout)), copy_max_vec_bytes);
|
|
|
|
Tensor s_a_tensor = make_tensor(make_smem_ptr<TA>(smem_ptr_a), smem_a_layout);
|
|
Tensor s_b_tensor = make_tensor(make_smem_ptr<TB>(smem_ptr_b), smem_b_layout);
|
|
Tensor s_c_tensor = make_tensor(make_smem_ptr<TC>(smem_ptr_c), smem_c_layout);
|
|
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, g_a_tensor, s_a_tensor);
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, g_b_tensor, s_b_tensor);
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, g_c_tensor, s_c_tensor);
|
|
|
|
cp_async_fence();
|
|
cp_async_wait<0>();
|
|
__syncthreads();
|
|
|
|
cooperative_gemm(
|
|
threadIdx.x, tiled_mma,
|
|
alpha, s_a_tensor, s_b_tensor, beta, s_c_tensor,
|
|
a_load_transform, b_load_transform, c_load_transform, c_store_transform,
|
|
a_copy_op, b_copy_op, c_copy_ld_op, c_copy_st_op
|
|
);
|
|
__syncthreads();
|
|
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, s_c_tensor, g_c_out_tensor);
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class GMemALayout,
|
|
class GMemBLayout,
|
|
class GMemCLayout,
|
|
class SMemALayout,
|
|
class SMemBLayout,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class TiledMma,
|
|
class ALoadTransform,
|
|
class BLoadTransform,
|
|
class CLoadTransform,
|
|
class CStoreTransform,
|
|
class SMemCopyOpA,
|
|
class SMemCopyOpB>
|
|
__launch_bounds__(ThreadBlockSize) __global__ void
|
|
cooperative_gemm_kernel_rmem_c(GMemALayout gmem_a_layout,
|
|
GMemBLayout gmem_b_layout,
|
|
GMemCLayout gmem_c_layout,
|
|
SMemALayout smem_a_layout,
|
|
SMemBLayout smem_b_layout,
|
|
TA const* a,
|
|
TB const* b,
|
|
TC const* c,
|
|
TC * c_out,
|
|
TiledMma tiled_mma,
|
|
ALoadTransform a_load_transform,
|
|
BLoadTransform b_load_transform,
|
|
CLoadTransform c_load_transform,
|
|
CStoreTransform c_store_transform,
|
|
SMemCopyOpA a_copy_op,
|
|
SMemCopyOpB b_copy_op)
|
|
{
|
|
using namespace cute;
|
|
|
|
Tensor g_a_tensor = make_tensor(make_gmem_ptr(a), gmem_a_layout);
|
|
Tensor g_b_tensor = make_tensor(make_gmem_ptr(b), gmem_b_layout);
|
|
Tensor g_c_tensor = make_tensor(make_gmem_ptr(c), gmem_c_layout);
|
|
Tensor g_c_out_tensor = make_tensor(make_gmem_ptr(c_out), gmem_c_layout);
|
|
|
|
constexpr uint32_t copy_max_vec_bytes = CopyMaxVecBits / 8;
|
|
|
|
extern __shared__ float4 smem_buf[];
|
|
auto* smem_ptr = reinterpret_cast<unsigned char*>(smem_buf);
|
|
auto* smem_ptr_a = smem_ptr;
|
|
auto* smem_ptr_b = smem_ptr_a + round_up((sizeof(TA) * cosize(smem_a_layout)), copy_max_vec_bytes);
|
|
|
|
Tensor s_a_tensor = make_tensor(make_smem_ptr<TA>(smem_ptr_a), smem_a_layout);
|
|
Tensor s_b_tensor = make_tensor(make_smem_ptr<TB>(smem_ptr_b), smem_b_layout);
|
|
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, g_a_tensor, s_a_tensor);
|
|
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, g_b_tensor, s_b_tensor);
|
|
|
|
cp_async_fence();
|
|
cp_async_wait<0>();
|
|
__syncthreads();
|
|
|
|
// Create C fragment for storing intermediate results
|
|
auto thr_mma = TiledMma().get_thread_slice(threadIdx.x);
|
|
Tensor g_c_partition = thr_mma.partition_C(g_c_tensor);
|
|
Tensor g_c_out_partition = thr_mma.partition_C(g_c_out_tensor);
|
|
Tensor r_c_partition = thr_mma.make_fragment_C(g_c_partition);
|
|
|
|
// Create indexing help for predicated GEMMs
|
|
Tensor cC = make_identity_tensor(shape(gmem_c_layout));
|
|
Tensor tCcC = thr_mma.partition_C(cC);
|
|
|
|
// Load C from global
|
|
// (always loading in predicated way)
|
|
CUTE_UNROLL
|
|
for (int i = 0; i < size(r_c_partition); ++i)
|
|
{
|
|
if (elem_less(tCcC(i), shape(g_c_tensor)))
|
|
{
|
|
r_c_partition(i) = c_load_transform(g_c_partition(i));
|
|
}
|
|
}
|
|
|
|
cooperative_gemm(
|
|
threadIdx.x, tiled_mma, s_a_tensor, s_b_tensor, r_c_partition,
|
|
a_load_transform, b_load_transform, a_copy_op, b_copy_op
|
|
);
|
|
|
|
__syncthreads();
|
|
|
|
// Store C to global
|
|
// (always storing in predicated way)
|
|
CUTE_UNROLL
|
|
for (int i = 0; i < size(r_c_partition); ++i)
|
|
{
|
|
if (elem_less(tCcC(i), shape(g_c_tensor)))
|
|
{
|
|
g_c_out_partition(i) = c_store_transform(r_c_partition(i));
|
|
}
|
|
}
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class GMemALayout, // logical shape (M, K)
|
|
class GMemBLayout, // logical shape (N, K)
|
|
class GMemCLayout, // logical shape (M, N)
|
|
class SMemALayout, // logical shape (M, K)
|
|
class SMemBLayout, // logical shape (N, K)
|
|
class SMemCLayout, // logical shape (M, N)
|
|
class TiledMma,
|
|
class ALoadTransform = cute::identity,
|
|
class BLoadTransform = cute::identity,
|
|
class CLoadTransform = cute::identity,
|
|
class CStoreTransform = cute::identity,
|
|
class ASMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
|
|
class BSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
|
|
class CSMemCopyLdOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
|
|
class CSMemCopyStOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
|
|
void test_cooperative_gemm(GMemALayout gmem_a_layout,
|
|
GMemBLayout gmem_b_layout,
|
|
GMemCLayout gmem_c_layout,
|
|
SMemALayout smem_a_layout,
|
|
SMemBLayout smem_b_layout,
|
|
SMemCLayout smem_c_layout,
|
|
TiledMma tiled_mma,
|
|
ALoadTransform a_load_transform = {},
|
|
BLoadTransform b_load_transform = {},
|
|
CLoadTransform c_load_transform = {},
|
|
CStoreTransform c_store_transform = {},
|
|
ASMemCopyOp a_smem_copy_op = {},
|
|
BSMemCopyOp b_smem_copy_op = {},
|
|
CSMemCopyLdOp c_smem_copy_ld_op = {},
|
|
CSMemCopyStOp c_smem_copy_st_op = {})
|
|
{
|
|
static_assert(std::is_same_v<typename fp64_tester<TA>::value_type, typename fp64_tester<TB>::value_type>);
|
|
static_assert(std::is_same_v<typename fp64_tester<TB>::value_type, typename fp64_tester<TC>::value_type>);
|
|
|
|
static_assert(size<0>(gmem_a_layout) == size<0>(gmem_c_layout)); // AM == CM
|
|
static_assert(size<0>(gmem_b_layout) == size<1>(gmem_c_layout)); // BN == CN
|
|
static_assert(size<1>(gmem_a_layout) == size<1>(gmem_b_layout)); // AK == BK
|
|
|
|
static_assert(size<0>(smem_a_layout) == size<0>(smem_c_layout)); // AM == CM
|
|
static_assert(size<0>(smem_b_layout) == size<1>(smem_c_layout)); // BN == CN
|
|
static_assert(size<1>(smem_a_layout) == size<1>(smem_b_layout)); // AK == BK
|
|
|
|
static_assert(cute::size(gmem_a_layout) == cute::size(smem_a_layout));
|
|
static_assert(cute::size(gmem_b_layout) == cute::size(smem_b_layout));
|
|
static_assert(cute::size(gmem_c_layout) == cute::size(smem_c_layout));
|
|
|
|
#if 0
|
|
print(" "); print("gmem: "); print(gmem_layout); print("\n");
|
|
print(" "); print("smem: "); print(smem_layout); print("\n");
|
|
print(" "); print("threads: "); print(ThreadBlockSize); print("\n");
|
|
#endif
|
|
|
|
const auto alpha = static_cast<TC>(1.1);
|
|
const auto beta = static_cast<TC>(1.2);
|
|
|
|
// Generate inputs
|
|
auto [h_a, h_b, h_c, h_c_out] = host_generate_gemm_inputs<TA, TB, TC>(gmem_a_layout, gmem_b_layout, gmem_c_layout);
|
|
|
|
thrust::device_vector<TA> d_a(h_a);
|
|
thrust::device_vector<TB> d_b(h_b);
|
|
thrust::device_vector<TC> d_c(h_c);
|
|
thrust::device_vector<TC> d_c_out(h_c_out.size(), TC(float(-1)));
|
|
|
|
constexpr uint32_t copy_max_vec_bytes = CopyMaxVecBits / 8;
|
|
|
|
const size_t shared_memory_size = round_up(sizeof(TA) * h_a.size(), copy_max_vec_bytes) +
|
|
round_up(sizeof(TB) * h_b.size(), copy_max_vec_bytes) +
|
|
sizeof(TC) * h_c.size();
|
|
|
|
|
|
auto kernel = cooperative_gemm_kernel<
|
|
ThreadBlockSize, CopyMaxVecBits,
|
|
GMemALayout, GMemBLayout, GMemCLayout,
|
|
SMemALayout, SMemBLayout, SMemCLayout,
|
|
TA, TB, TC, decltype(alpha), decltype(beta),
|
|
TiledMma,
|
|
ALoadTransform, BLoadTransform, CLoadTransform, CStoreTransform,
|
|
ASMemCopyOp, BSMemCopyOp, CSMemCopyLdOp, CSMemCopyStOp
|
|
>;
|
|
|
|
ASSERT_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_size)), 0);
|
|
|
|
kernel<<<1, ThreadBlockSize, shared_memory_size>>>(
|
|
gmem_a_layout,
|
|
gmem_b_layout,
|
|
gmem_c_layout,
|
|
smem_a_layout,
|
|
smem_b_layout,
|
|
smem_c_layout,
|
|
thrust::raw_pointer_cast(d_a.data()),
|
|
thrust::raw_pointer_cast(d_b.data()),
|
|
thrust::raw_pointer_cast(d_c.data()),
|
|
thrust::raw_pointer_cast(d_c_out.data()),
|
|
alpha,
|
|
beta,
|
|
tiled_mma,
|
|
a_load_transform,
|
|
b_load_transform,
|
|
c_load_transform,
|
|
c_store_transform,
|
|
a_smem_copy_op,
|
|
b_smem_copy_op,
|
|
c_smem_copy_ld_op,
|
|
c_smem_copy_st_op
|
|
);
|
|
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
if (result != cudaSuccess) {
|
|
cudaError_t error = cudaGetLastError();
|
|
FAIL() << "Error at kernel sync: " << cudaGetErrorString(error) << "\n";
|
|
}
|
|
|
|
// Reference gemm
|
|
auto h_c_ref = host_reference_gemm(alpha,
|
|
make_tensor(h_a.data(), gmem_a_layout),
|
|
make_tensor(h_b.data(), gmem_b_layout),
|
|
beta,
|
|
make_tensor(h_c.data(), gmem_c_layout),
|
|
a_load_transform,
|
|
b_load_transform,
|
|
c_load_transform,
|
|
c_store_transform);
|
|
|
|
// Copy result data
|
|
h_c_out = d_c_out;
|
|
|
|
// Verify correctness
|
|
verify_gemm_correctness(make_tensor(h_c_out.data(), gmem_c_layout),
|
|
make_tensor(h_c_ref.data(), gmem_c_layout));
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class GMemALayout, // logical shape (M, K)
|
|
class GMemBLayout, // logical shape (N, K)
|
|
class GMemCLayout, // logical shape (M, N)
|
|
class SMemALayout, // logical shape (M, K)
|
|
class SMemBLayout, // logical shape (N, K)
|
|
class TiledMma,
|
|
class ALoadTransform = cute::identity,
|
|
class BLoadTransform = cute::identity,
|
|
class CLoadTransform = cute::identity,
|
|
class CStoreTransform = cute::identity,
|
|
class ASMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>,
|
|
class BSMemCopyOp = AutoVectorizingCopyWithAssumedAlignment<CopyMaxVecBits>>
|
|
void test_cooperative_gemm_rmem_c(GMemALayout gmem_a_layout,
|
|
GMemBLayout gmem_b_layout,
|
|
GMemCLayout gmem_c_layout,
|
|
SMemALayout smem_a_layout,
|
|
SMemBLayout smem_b_layout,
|
|
TiledMma tiled_mma,
|
|
ALoadTransform a_load_transform = {},
|
|
BLoadTransform b_load_transform = {},
|
|
CLoadTransform c_load_transform = {},
|
|
CStoreTransform c_store_transform = {},
|
|
ASMemCopyOp a_smem_copy_op = {},
|
|
BSMemCopyOp b_smem_copy_op = {})
|
|
{
|
|
static_assert(size<0>(gmem_a_layout) == size<0>(gmem_c_layout)); // AM == CM
|
|
static_assert(size<0>(gmem_b_layout) == size<1>(gmem_c_layout)); // BN == CN
|
|
static_assert(size<1>(gmem_a_layout) == size<1>(gmem_b_layout)); // AK == BK
|
|
|
|
static_assert(size<1>(smem_a_layout) == size<1>(smem_b_layout)); // AK == BK
|
|
|
|
static_assert(cute::size(gmem_a_layout) == cute::size(smem_a_layout));
|
|
static_assert(cute::size(gmem_b_layout) == cute::size(smem_b_layout));
|
|
|
|
#if 0
|
|
print(" "); print("gmem: "); print(gmem_layout); print("\n");
|
|
print(" "); print("smem: "); print(smem_layout); print("\n");
|
|
print(" "); print("threads: "); print(ThreadBlockSize); print("\n");
|
|
#endif
|
|
|
|
const auto alpha = static_cast<TC>(1.0);
|
|
const auto beta = static_cast<TC>(1.0);
|
|
|
|
// Generate inputs
|
|
auto [h_a, h_b, h_c, h_c_out] =
|
|
host_generate_gemm_inputs<TA, TB, TC>(gmem_a_layout, gmem_b_layout, gmem_c_layout);
|
|
|
|
thrust::device_vector<TA> d_a(h_a);
|
|
thrust::device_vector<TB> d_b(h_b);
|
|
thrust::device_vector<TC> d_c(h_c);
|
|
thrust::device_vector<TC> d_c_out(h_c_out.size(), static_cast<TC>(-1));
|
|
|
|
constexpr uint32_t copy_max_vec_bytes = CopyMaxVecBits / 8;
|
|
|
|
const size_t shared_memory_size = round_up(sizeof(TA) * h_a.size(), copy_max_vec_bytes) +
|
|
round_up(sizeof(TB) * h_b.size(), copy_max_vec_bytes);
|
|
|
|
|
|
auto kernel = cooperative_gemm_kernel_rmem_c<
|
|
ThreadBlockSize, CopyMaxVecBits,
|
|
GMemALayout, GMemBLayout, GMemCLayout,
|
|
SMemALayout, SMemBLayout,
|
|
TA, TB, TC,
|
|
TiledMma,
|
|
ALoadTransform, BLoadTransform, CLoadTransform, CStoreTransform,
|
|
ASMemCopyOp, BSMemCopyOp
|
|
>;
|
|
|
|
ASSERT_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_size)), 0);
|
|
|
|
kernel<<<1, ThreadBlockSize, shared_memory_size>>>(
|
|
gmem_a_layout,
|
|
gmem_b_layout,
|
|
gmem_c_layout,
|
|
smem_a_layout,
|
|
smem_b_layout,
|
|
thrust::raw_pointer_cast(d_a.data()),
|
|
thrust::raw_pointer_cast(d_b.data()),
|
|
thrust::raw_pointer_cast(d_c.data()),
|
|
thrust::raw_pointer_cast(d_c_out.data()),
|
|
tiled_mma,
|
|
a_load_transform, b_load_transform, c_load_transform, c_store_transform,
|
|
a_smem_copy_op, b_smem_copy_op
|
|
);
|
|
|
|
cudaError_t result = cudaDeviceSynchronize();
|
|
if (result != cudaSuccess) {
|
|
cudaError_t error = cudaGetLastError();
|
|
FAIL() << "Error at kernel sync: " << cudaGetErrorString(error) << "\n";
|
|
}
|
|
|
|
// Copy result data
|
|
h_c_out = d_c_out;
|
|
|
|
// Reference gemm
|
|
auto h_c_ref = host_reference_gemm(alpha,
|
|
make_tensor(h_a.data(), gmem_a_layout),
|
|
make_tensor(h_b.data(), gmem_b_layout),
|
|
beta,
|
|
make_tensor(h_c.data(), gmem_c_layout),
|
|
a_load_transform,
|
|
b_load_transform,
|
|
c_load_transform,
|
|
c_store_transform);
|
|
|
|
// Verify correctness
|
|
verify_gemm_correctness(make_tensor(h_c_out.data(), gmem_c_layout),
|
|
make_tensor(h_c_ref.data(), gmem_c_layout));
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class ShapeMNK,
|
|
class TiledMma,
|
|
class ... Ops>
|
|
void test_cooperative_gemm_col_major_layout(ShapeMNK shape_mnk,
|
|
TiledMma tiled_mma,
|
|
Ops ... ops)
|
|
{
|
|
auto a_layout = make_layout(select<0, 2>(shape_mnk));
|
|
auto b_layout = make_layout(select<1, 2>(shape_mnk), GenRowMajor{});
|
|
auto c_layout = make_layout(select<0, 1>(shape_mnk));
|
|
|
|
test_cooperative_gemm<ThreadBlockSize,
|
|
CopyMaxVecBits,
|
|
TA, TB, TC>
|
|
(a_layout,
|
|
b_layout,
|
|
c_layout,
|
|
a_layout,
|
|
b_layout,
|
|
c_layout,
|
|
tiled_mma,
|
|
ops...);
|
|
}
|
|
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class SMemAtomLayoutA,
|
|
class SMemAtomLayoutB,
|
|
class SMemAtomLayoutC,
|
|
class ShapeMNK,
|
|
class TiledMma,
|
|
class ... Ops>
|
|
std::enable_if_t<std::conjunction_v<cute::is_layout<SMemAtomLayoutA>,
|
|
cute::is_layout<SMemAtomLayoutB>,
|
|
cute::is_layout<SMemAtomLayoutC>>>
|
|
test_cooperative_gemm_col_major_layout(SMemAtomLayoutA smem_atom_layout_a,
|
|
SMemAtomLayoutB smem_atom_layout_b,
|
|
SMemAtomLayoutC smem_atom_layout_c,
|
|
ShapeMNK shape_mnk,
|
|
TiledMma tiled_mma,
|
|
Ops&& ... ops)
|
|
{
|
|
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk));
|
|
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenRowMajor{});
|
|
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk));
|
|
|
|
auto smem_a_layout = tile_to_shape(
|
|
smem_atom_layout_a,
|
|
make_shape(shape<0>(gmem_a_layout), shape<1>(gmem_a_layout)));
|
|
|
|
auto smem_b_layout = tile_to_shape(
|
|
smem_atom_layout_b,
|
|
make_shape(shape<0>(gmem_b_layout), shape<1>(gmem_b_layout)));
|
|
|
|
auto smem_c_layout = tile_to_shape(
|
|
smem_atom_layout_c,
|
|
make_shape(shape<0>(gmem_c_layout), shape<1>(gmem_c_layout)));
|
|
|
|
test_cooperative_gemm<ThreadBlockSize,
|
|
CopyMaxVecBits,
|
|
TA, TB, TC>
|
|
(gmem_a_layout,
|
|
gmem_b_layout,
|
|
gmem_c_layout,
|
|
smem_a_layout,
|
|
smem_b_layout,
|
|
smem_c_layout,
|
|
tiled_mma,
|
|
ops...);
|
|
}
|
|
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class ShapeMNK,
|
|
class TiledMma,
|
|
class ... Ops>
|
|
void test_cooperative_gemm_col_major_layout_rmem_c(ShapeMNK shape_mnk,
|
|
TiledMma tiled_mma,
|
|
Ops ... ops)
|
|
{
|
|
auto a_layout = make_layout(select<0, 2>(shape_mnk));
|
|
auto b_layout = make_layout(select<1, 2>(shape_mnk), GenRowMajor{});
|
|
auto c_layout = make_layout(select<0, 1>(shape_mnk));
|
|
|
|
|
|
test_cooperative_gemm_rmem_c<ThreadBlockSize,
|
|
CopyMaxVecBits,
|
|
TA, TB,TC>
|
|
(a_layout,
|
|
b_layout,
|
|
c_layout,
|
|
a_layout,
|
|
b_layout,
|
|
tiled_mma,
|
|
ops...);
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
uint32_t CopyMaxVecBits,
|
|
class TA,
|
|
class TB,
|
|
class TC,
|
|
class SMemAtomLayoutA,
|
|
class SMemAtomLayoutB,
|
|
class ShapeMNK,
|
|
class TiledMma,
|
|
class ... Ops>
|
|
std::enable_if_t<std::conjunction_v<cute::is_layout<SMemAtomLayoutA>,
|
|
cute::is_layout<SMemAtomLayoutB>>>
|
|
test_cooperative_gemm_col_major_layout_rmem_c(SMemAtomLayoutA smem_atom_layout_a,
|
|
SMemAtomLayoutB smem_atom_layout_b,
|
|
ShapeMNK shape_mnk,
|
|
TiledMma tiled_mma,
|
|
Ops ... ops)
|
|
{
|
|
auto gmem_a_layout = make_layout(select<0, 2>(shape_mnk));
|
|
auto gmem_b_layout = make_layout(select<1, 2>(shape_mnk), GenRowMajor{});
|
|
auto gmem_c_layout = make_layout(select<0, 1>(shape_mnk));
|
|
|
|
auto smem_a_layout = tile_to_shape(
|
|
smem_atom_layout_a,
|
|
make_shape(shape<0>(gmem_a_layout), shape<1>(gmem_a_layout)));
|
|
|
|
auto smem_b_layout = tile_to_shape(
|
|
smem_atom_layout_b,
|
|
make_shape(shape<0>(gmem_b_layout), shape<1>(gmem_b_layout)));
|
|
|
|
test_cooperative_gemm_rmem_c<ThreadBlockSize, CopyMaxVecBits,
|
|
TA, TB, TC>
|
|
(gmem_a_layout,
|
|
gmem_b_layout,
|
|
gmem_c_layout,
|
|
smem_a_layout,
|
|
smem_b_layout,
|
|
tiled_mma,
|
|
ops...);
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
typename T,
|
|
class ... Args>
|
|
void test_cooperative_gemm_col_major_layout_rmem_c(Args&& ... args)
|
|
{
|
|
test_cooperative_gemm_col_major_layout_rmem_c<ThreadBlockSize,
|
|
cute::sizeof_bits_v<T>,
|
|
T, T, T>
|
|
(static_cast<Args&&>(args)...);
|
|
}
|
|
|
|
template<uint32_t ThreadBlockSize,
|
|
class T,
|
|
class ... Args>
|
|
void test_cooperative_gemm_col_major_layout(Args&& ... args)
|
|
{
|
|
test_cooperative_gemm_col_major_layout<ThreadBlockSize,
|
|
cute::sizeof_bits_v<T>,
|
|
T, T, T>
|
|
(static_cast<Args&&>(args)...);
|
|
}
|