leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/CMakeLists.txt
+3
View File
@@ -28,6 +28,7 @@
add_subdirectory(core)
add_subdirectory(volta)
add_subdirectory(turing)
add_subdirectory(ampere)
add_subdirectory(hopper)
add_subdirectory(layout)
@@ -39,6 +40,7 @@ add_custom_target(
cutlass_test_unit_cute_layout
cutlass_test_unit_cute_core
cutlass_test_unit_cute_volta
cutlass_test_unit_cute_turing
cutlass_test_unit_cute_ampere
cutlass_test_unit_cute_hopper
cutlass_test_unit_cute_msvc_compilation
@@ -51,6 +53,7 @@ add_custom_target(
test_unit_cute_core
test_unit_cute_volta
test_unit_cute_ampere
test_unit_cute_turing
test_unit_cute_hopper
test_unit_cute_msvc_compilation
)
test/unit/cute/ampere/CMakeLists.txt
+1
View File
@@ -30,6 +30,7 @@ cutlass_test_unit_add_executable(
cutlass_test_unit_cute_ampere
cp_async.cu
ldsm.cu
cooperative_gemm.cu
)
cutlass_test_unit_add_executable(
test/unit/cute/ampere/cooperative_gemm.cu
+300
View File
@@ -0,0 +1,300 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2024 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.
*
**************************************************************************************************/
#include "cutlass_unit_test.h"
#include <cute/tensor.hpp>
#include "../cooperative_gemm_common.hpp"
using namespace cute;
TEST(SM80_CuTe_Ampere, CooperativeGemm1_Half_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x8_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm2_Double_MMA) {
using value_type = double;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,
Layout<Shape<_2,_2,_1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm3_Half_MMA_CustomSmemLayouts) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 128;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>, // 2x2x1 thread group
Tile<_32, _32, _16> // 32x32x16 MMA for LDSM, 1x2x1 value group`
>;
using smem_a_atom_layout_t = Layout<Shape<_64, _8>, Stride< _1,_64>>;
using smem_b_atom_layout_t = Layout<Shape< _8,_32>, Stride<_32, _1>>;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
test_cooperative_gemm_col_major_layout<smem_a_atom_layout_t,
smem_b_atom_layout_t,
smem_c_atom_layout_t,
m,
n,
k,
thread_block_size,
tiled_mma_t,
128,
value_type,
value_type,
value_type>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm4_Half_MMA_SwizzledSmemLayouts) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 128;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x16_F16F16F16F16_TN>,
Layout<Shape<_2, _2, _1>>, // 2x2x1 thread group
Tile<_32, _32, _16> // 32x32x16 MMA for LDSM, 1x2x1 value group`
>;
// RowMajor
using smem_rowmajor_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape < _8,_64>,
Stride<_64, _1>>{}));
// ColMajor
using smem_colmajor_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape <_64, _8>,
Stride< _1,_64>>{}));
using smem_a_atom_layout_t = smem_rowmajor_atom_layout_t;
using smem_b_atom_layout_t = smem_colmajor_atom_layout_t;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
SM75_U32x4_LDSM_N, // A
SM75_U16x8_LDSM_T, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
value_type,
value_type,
value_type>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm5_Double_MMA_SwizzledSmemLayouts) {
using value_type = double;
constexpr uint32_t m = 128;
constexpr uint32_t n = 64;
constexpr uint32_t k = 16;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>, // Atom
Layout<Shape<_2, _2, _1>>, // Atom layout
Tile<Layout<Shape<_16, _2>, Stride<_2, _1>>, // 32x32x4 MMA with perm for load vectorization
Layout<Shape<_16, _2>, Stride<_2, _1>>,
Underscore>>;
using smem_a_atom_layout_t = decltype(
composition(Swizzle<2,2,2>{},
Layout<Shape <_16, _4>,
Stride< _1,_16>>{})); // M, K
using smem_b_atom_layout_t = decltype(
composition(Swizzle<2,2,2>{},
Layout<Shape <_16, _4>,
Stride< _1,_16>>{})); // N, K
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
value_type,
value_type,
value_type>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm6_MixedPrecisionFP16FP32_MMA) {
using TA = cutlass::half_t;
using TB = cutlass::half_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm7_MixedPrecisionBF16FP32_MMA) {
using TA = cutlass::bfloat16_t;
using TB = cutlass::bfloat16_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32BF16BF16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
}
TEST(SM80_CuTe_Ampere, CooperativeGemm8_MixedPrecisionTF32FP32_MMA) {
using TA = cutlass::tfloat32_t;
using TB = cutlass::tfloat32_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
}
test/unit/cute/cooperative_gemm_common.hpp
+414
View File
@@ -0,0 +1,414 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2024 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_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<class ALayout,
class BLayout,
class CLayout,
class SMemALayout,
class SMemBLayout,
class SMemCLayout,
class SmemCopyOpA,
class SmemCopyOpB,
class SmemCopyOpC,
uint32_t ThreadBlockSize,
class TiledMma,
uint32_t CopyMaxVecBits,
class TA,
class TB,
class TC,
class Alpha,
class Beta,
class ALoadTransform,
class BLoadTransform,
class CLoadTransform,
class CStoreTransform>
__launch_bounds__(ThreadBlockSize) __global__ void
cooperative_gemm_kernel(TA const* a,
TB const* b,
TC* c,
TC* c_out,
Alpha const alpha,
Beta const beta,
ALoadTransform a_load_transform,
BLoadTransform b_load_transform,
CLoadTransform c_load_transform,
CStoreTransform c_store_transform)
{
using namespace cute;
Tensor g_a_tensor = make_tensor(make_gmem_ptr(a), ALayout{});
Tensor g_b_tensor = make_tensor(make_gmem_ptr(b), BLayout{});
Tensor g_c_tensor = make_tensor(make_gmem_ptr(c), CLayout{});
Tensor g_c_out_tensor = make_tensor(make_gmem_ptr(c_out), CLayout{});
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(SMemALayout {})), copy_max_vec_bytes);
auto* smem_ptr_c = smem_ptr_b + round_up((sizeof(TB) * cosize(SMemBLayout {})), copy_max_vec_bytes);
Tensor s_a_tensor = make_tensor(make_smem_ptr<TA>(smem_ptr_a), SMemALayout{});
Tensor s_b_tensor = make_tensor(make_smem_ptr<TB>(smem_ptr_b), SMemBLayout{});
Tensor s_c_tensor = make_tensor(make_smem_ptr<TC>(smem_ptr_c), SMemCLayout{});
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();
TiledMma tiled_mma;
cooperative_gemm<SmemCopyOpA, SmemCopyOpB, SmemCopyOpC>(
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
);
__syncthreads();
cooperative_copy<ThreadBlockSize, CopyMaxVecBits>(threadIdx.x, s_c_tensor, g_c_out_tensor);
}
template<class ALayout, // logical shape (M, K)
class BLayout, // logical shape (N, K)
class CLayout, // logical shape (M, N)
class SMemALayout, // logical shape (M, K)
class SMemBLayout, // logical shape (N, K)
class SMemCLayout, // logical shape (M, N)
class SmemCopyOpA,
class SmemCopyOpB,
class SmemCopyOpC,
uint32_t ThreadBlockSize,
class TiledMma,
uint32_t CopyMaxVecBits,
class TA,
class TB,
class TC,
class ALoadTransform = cute::identity,
class BLoadTransform = cute::identity,
class CLoadTransform = cute::identity,
class CStoreTransform = cute::identity>
void test_cooperative_gemm(ALoadTransform const& a_load_transform = {},
BLoadTransform const& b_load_transform = {},
CLoadTransform const& c_load_transform = {},
CStoreTransform const& c_store_transform = {})
{
using gmem_a_layout_t = ALayout;
using gmem_b_layout_t = BLayout;
using gmem_c_layout_t = CLayout;
using smem_a_layout_t = SMemALayout;
using smem_b_layout_t = SMemBLayout;
using smem_c_layout_t = SMemCLayout;
static_assert(size<0>(gmem_a_layout_t{}) == size<0>(gmem_c_layout_t{})); // AM == CM
static_assert(size<0>(gmem_b_layout_t{}) == size<1>(gmem_c_layout_t{})); // BN == CN
static_assert(size<1>(gmem_a_layout_t{}) == size<1>(gmem_b_layout_t{})); // AK == BK
static_assert(size<0>(smem_a_layout_t{}) == size<0>(smem_c_layout_t{})); // AM == CM
static_assert(size<0>(smem_b_layout_t{}) == size<1>(smem_c_layout_t{})); // BN == CN
static_assert(size<1>(smem_a_layout_t{}) == size<1>(smem_b_layout_t{})); // AK == BK
static_assert(cute::size(gmem_a_layout_t {}) == cute::size(smem_a_layout_t {}));
static_assert(cute::size(gmem_b_layout_t {}) == cute::size(smem_b_layout_t {}));
static_assert(cute::size(gmem_c_layout_t {}) == cute::size(smem_c_layout_t {}));
#if 0
print(" "); print("gmem: "); print(gmem_layout_t{}); print("\n");
print(" "); print("smem: "); print(smem_layout_t{}); 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);
thrust::host_vector<TA> h_a(cosize(gmem_a_layout_t{}));
thrust::host_vector<TB> h_b(cosize(gmem_b_layout_t{}));
thrust::host_vector<TC> h_c(cosize(gmem_c_layout_t{}));
thrust::host_vector<TC> h_c_out(cosize(gmem_c_layout_t{}));
auto h_a_tensor = make_tensor(h_a.data(), gmem_a_layout_t{});
auto h_b_tensor = make_tensor(h_b.data(), gmem_b_layout_t{});
auto h_c_tensor = make_tensor(h_c.data(), gmem_c_layout_t{});
size_t max_size = std::max<size_t>({static_cast<size_t>(size(gmem_a_layout_t {})),
static_cast<size_t>(size(gmem_b_layout_t {})),
static_cast<size_t>(size(gmem_c_layout_t {}))});
for (size_t i = 0; i < max_size; ++i) {
double di = static_cast<double>(i);
if(i < size(gmem_a_layout_t{})) {
h_a_tensor(i) = static_cast<TA>(di / size(gmem_a_layout_t{}));
}
if(i < size(gmem_b_layout_t{})) {
h_b_tensor(i) = static_cast<TA>(di / size(gmem_a_layout_t{}));
}
if(i < size(gmem_c_layout_t{})) {
h_c_tensor(i) = static_cast<TC>((di*di) / size(gmem_a_layout_t{}));
}
}
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)));
const size_t shared_memory_size =
(sizeof(TA) * h_a.size()) + (sizeof(TB) * h_b.size()) + (sizeof(TC) * h_c.size());
auto kernel = cooperative_gemm_kernel<
gmem_a_layout_t, gmem_b_layout_t, gmem_c_layout_t,
smem_a_layout_t, smem_b_layout_t, smem_c_layout_t,
SmemCopyOpA, SmemCopyOpB, SmemCopyOpC,
ThreadBlockSize, TiledMma, CopyMaxVecBits,
TA, TB, TC, decltype(alpha), decltype(beta),
ALoadTransform, BLoadTransform, CLoadTransform, CStoreTransform
>;
ASSERT_EQ(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(shared_memory_size)), 0);
kernel<<<1, ThreadBlockSize, shared_memory_size>>>(
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,
a_load_transform,
b_load_transform,
c_load_transform,
c_store_transform
);
cudaError_t result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
cudaError_t error = cudaGetLastError();
FAIL() << "Error at kernel sync: " << cudaGetErrorString(error) << "\n";
}
thrust::host_vector<TC> h_c_ref(h_c.size(), static_cast<TC>(0.0));
auto h_c_ref_tensor = make_tensor(h_c_ref.data(), gmem_c_layout_t{});
// 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<double>(a_value);
const auto b_value_fp64 = static_cast<double>(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<double>(h_c_ref_tensor(i));
const auto c_value_fp64 = static_cast<double>(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));
}
h_c_out = d_c_out;
auto h_c_out_tensor = make_tensor(h_c_out.data(), gmem_c_layout_t{});
for (int i = 0; i < size(h_c_ref_tensor); i++) {
double h_c_ref_i = h_c_ref_tensor(i);
double 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 M,
uint32_t N,
uint32_t K,
uint32_t ThreadBlockSize,
class TiledMMAType,
uint32_t CopyMaxVecBits,
class TA,
class TB,
class TC,
class ALoadTransform = cute::identity,
class BLoadTransform = cute::identity,
class CLoadTransform = cute::identity,
class CStoreTransform = cute::identity>
void test_cooperative_gemm_col_major_layout(ALoadTransform const& a_load_transform = {},
BLoadTransform const& b_load_transform = {},
CLoadTransform const& c_load_transform = {},
CStoreTransform const& c_store_transform = {})
{
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<K> {})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<N> {}, Int<K> {}), GenRowMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<N> {})));
using smem_a_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<K> {})));
using smem_b_layout_t = decltype(make_layout(make_shape(Int<N> {}, Int<K> {}), GenRowMajor{}));
using smem_c_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<N> {})));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TA>>,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TB>>,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TC>>,
ThreadBlockSize,
TiledMMAType,
CopyMaxVecBits,
TA,
TB,
TC>(a_load_transform, b_load_transform, c_load_transform, c_store_transform);
}
template<uint32_t M,
uint32_t N,
uint32_t K,
uint32_t ThreadBlockSize,
class TiledMMAType,
class T,
class ALoadTransform = cute::identity,
class BLoadTransform = cute::identity,
class CLoadTransform = cute::identity,
class CStoreTransform = cute::identity>
void test_cooperative_gemm_col_major_layout(ALoadTransform const& a_load_transform = {},
BLoadTransform const& b_load_transform = {},
CLoadTransform const& c_load_transform = {},
CStoreTransform const& c_store_transform = {})
{
test_cooperative_gemm_col_major_layout<M, N, K, ThreadBlockSize, TiledMMAType, cute::sizeof_bits_v<T>, T, T, T>(
a_load_transform, b_load_transform, c_load_transform, c_store_transform);
}
template<class SMemAAtomLayout,
class SMemBAtomLayout,
class SMemCAtomLayout,
uint32_t M,
uint32_t N,
uint32_t K,
uint32_t ThreadBlockSize,
class TiledMMAType,
uint32_t CopyMaxVecBits,
class TA,
class TB,
class TC,
class ALoadTransform = cute::identity,
class BLoadTransform = cute::identity,
class CLoadTransform = cute::identity,
class CStoreTransform = cute::identity>
void test_cooperative_gemm_col_major_layout(ALoadTransform const& a_load_transform = {},
BLoadTransform const& b_load_transform = {},
CLoadTransform const& c_load_transform = {},
CStoreTransform const& c_store_transform = {})
{
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<K> {})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<N> {}, Int<K> {}), GenRowMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<M> {}, Int<N> {})));
using smem_a_atom_layout_t = SMemAAtomLayout;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
using smem_b_atom_layout_t = SMemBAtomLayout;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
using smem_c_atom_layout_t = SMemCAtomLayout;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TA>>,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TB>>,
AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<TC>>,
ThreadBlockSize,
TiledMMAType,
CopyMaxVecBits,
TA,
TB,
TC>(a_load_transform, b_load_transform, c_load_transform, c_store_transform);
}
template<class SMemAAtomLayout,
class SMemBAtomLayout,
class SMemCAtomLayout,
uint32_t M,
uint32_t N,
uint32_t K,
uint32_t ThreadBlockSize,
class TiledMMAType,
class T,
class ALoadTransform = cute::identity,
class BLoadTransform = cute::identity,
class CLoadTransform = cute::identity,
class CStoreTransform = cute::identity>
void test_cooperative_gemm_col_major_layout(ALoadTransform const& a_load_transform = {},
BLoadTransform const& b_load_transform = {},
CLoadTransform const& c_load_transform = {},
CStoreTransform const& c_store_transform = {})
{
test_cooperative_gemm_col_major_layout<SMemAAtomLayout,
SMemBAtomLayout,
SMemCAtomLayout,
M,
N,
K,
ThreadBlockSize,
TiledMMAType,
cute::sizeof_bits_v<T>,
T,
T,
T>(a_load_transform, b_load_transform, c_load_transform, c_store_transform);
}
test/unit/cute/core/complement.cpp
+25 -10
View File
@@ -35,22 +35,22 @@
#include <cute/tensor.hpp>
template <class Layout, class CoSizeHi>
template <class Layout, class CoTarget>
void
test_complement(Layout const& layout, CoSizeHi const& cosize_hi)
test_complement(Layout const& layout, CoTarget const& cotarget)
{
using namespace cute;
auto result = complement(layout, cosize_hi);
auto result = complement(layout, cotarget);
CUTLASS_TRACE_HOST("complement(" << layout << ", " << cosize_hi << ") => " << result);
CUTLASS_TRACE_HOST("complement(" << layout << ", " << cotarget << ") => " << result);
auto completed = make_layout(layout, result);
// Lower-bound on the codomain size of the layout ++ complement (1)
EXPECT_GE(cosize(completed), cosize_hi);
EXPECT_GE(cosize(completed), size(cotarget));
// Upper-bound on the codomain size of the complement (2)
EXPECT_LE(cosize(result), cute::round_up(cosize_hi, cosize(layout)));
EXPECT_LE(cosize(result), cute::round_up(size(cotarget), cosize(layout)));
// Post-condition on the codomain of the complement
for (int i = 1; i < size(result); ++i) {
@@ -62,9 +62,9 @@ test_complement(Layout const& layout, CoSizeHi const& cosize_hi)
// Other observations
EXPECT_LE(size(result), cosize(result)); // As a result of the ordered condition (3)
EXPECT_GE(size(result), cosize_hi / size(filter(layout)));
EXPECT_GE(size(result), size(cotarget) / size(filter(layout)));
EXPECT_LE(cosize(completed), cosize(result) + cosize(layout));
EXPECT_GE(cosize(result), cosize_hi / size(filter(layout)));
EXPECT_GE(cosize(result), size(cotarget) / size(filter(layout)));
if constexpr (is_static<decltype(stride(completed))>::value) { // If we can apply complement again
EXPECT_EQ(size(complement(completed)), 1); // There's no more codomain left over
}
@@ -90,6 +90,8 @@ TEST(CuTe_core, Complement)
test_complement(layout);
test_complement(layout, Int<2>{});
test_complement(layout, Int<5>{});
test_complement(layout, make_shape(Int<2>{}, 2));
}
{
@@ -97,6 +99,8 @@ TEST(CuTe_core, Complement)
test_complement(layout);
test_complement(layout, Int<2>{});
test_complement(layout, Int<5>{});
test_complement(layout, make_shape(Int<2>{}, 2));
}
{
@@ -105,6 +109,8 @@ TEST(CuTe_core, Complement)
test_complement(layout, Int<1>{});
test_complement(layout, Int<2>{});
test_complement(layout, Int<8>{});
test_complement(layout, Int<5>{});
test_complement(layout, make_shape(Int<2>{}, 2));
}
{
@@ -130,6 +136,7 @@ TEST(CuTe_core, Complement)
test_complement(layout);
test_complement(layout, Int<16>{});
test_complement(layout, Int<19>{});
test_complement(layout, make_shape(Int<2>{}, 2));
}
{
@@ -138,6 +145,7 @@ TEST(CuTe_core, Complement)
test_complement(layout, Int<1>{});
test_complement(layout);
test_complement(layout, Int<17>{});
test_complement(layout, make_shape(Int<2>{}, 2));
}
{
@@ -193,8 +201,8 @@ TEST(CuTe_core, Complement)
// Fails due to non-injective layout
// {
// auto layout = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
// Stride<Stride<_1,_8>,Stride<_8,_4>>{});
// auto layout = make_layout(Shape <Shape <_2,_2>,Shape <_2,_2>>{},
// Stride<Stride<_1,_8>,Stride<_8,_4>>{});
// test_complement(layout);
// }
@@ -289,4 +297,11 @@ TEST(CuTe_core, Complement)
test_complement(layout);
}
{
auto layout = make_layout(Int<64>{});
test_complement(layout, make_shape(Int<32>{}, Int<4>{}, Int<4>{}));
test_complement(layout, make_shape(Int<32>{}, Int<4>{}, 4));
}
}
test/unit/cute/core/composition.cpp
+27 -12
View File
@@ -212,13 +212,12 @@ TEST(CuTe_core, Composition)
test_composition(a, b);
}
// FAILS due to b not "dividing into" a properly
//{
// auto a = make_layout(Shape<_4,_3>{});
// auto b = make_layout(Shape<_6>{});
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_6>{});
// test_composition(a, b);
//}
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{});
@@ -234,13 +233,12 @@ TEST(CuTe_core, Composition)
test_composition(a, b);
}
// FAILS due to b not "dividing into" a properly
//{
// auto a = make_layout(Shape<_4,_3>{});
// auto b = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
// test_composition(a, b);
//}
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
@@ -523,4 +521,21 @@ TEST(CuTe_core, Composition)
test_composition(a, b);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("BETA: Tuple strides" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = make_layout(Shape<_4,_4>{}, Stride<_4,_1>{});
auto b = make_layout(Shape<_4,_4>{}, Stride<E<1>,E<0>>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,Shape<_2,_3>>{}, Stride<_6,Stride<_3,_1>>{});
auto b = make_layout(Shape<_2,_4>{}, Stride<E<1,1>,E<0>>{});
test_composition(a, b);
}
}
test/unit/cute/core/logical_divide.cpp
+33 -18
View File
@@ -227,27 +227,42 @@ TEST(CuTe_core, Logical_divide)
ASSERT_TRUE(decltype(stride<1>(result) == Int<48>{})::value);
}
// DISALLOWED
//{
//auto layout = make_layout(make_shape(128,4,3), make_stride(1,512,0));
//auto tile = Layout<_32>{};
{
auto layout = make_layout(make_shape(Int<32>{}, Int<4>{}, 4));
auto tile = Layout<_64>{};
//test_logical_divide(layout, tile);
//}
test_logical_divide(layout, tile);
//{
//auto layout = make_layout(make_shape(128,4,3), make_stride(1,512,0));
//auto tile = Layout<_32,_2>{};
// Enforcement of result
auto result = logical_divide(layout, tile);
ASSERT_TRUE(bool( shape(result) == make_shape (_64{}, make_shape ( _2{}, 4))));
ASSERT_TRUE(bool(stride(result) == make_stride( _1{}, make_stride(_64{},_128{}))));
}
//CUTLASS_TRACE_HOST("complement: " << complement(tile, size(layout)));
//test_logical_divide(layout, tile);
//}
//{
//auto layout = make_layout(make_shape(16,4,3), make_stride(1,512,0));
//auto tile = Layout<_32>{};
//
// ALLOWED, but dangerous due to the dynamic lhs shapes
// Consider disallowing...
//
//CUTLASS_TRACE_HOST("complement: " << complement(tile, size(layout)));
//test_logical_divide(layout, tile);
//}
{
auto layout = make_layout(make_shape(128,4,3), make_stride(1,512,0));
auto tile = Layout<_32>{};
test_logical_divide(layout, tile);
}
{
auto layout = make_layout(make_shape(128,4,3), make_stride(1,512,0));
auto tile = Layout<_32,_2>{};
test_logical_divide(layout, tile);
}
{
auto layout = make_layout(make_shape(16,4,3), make_stride(1,512,0));
auto tile = Layout<_32>{};
test_logical_divide(layout, tile);
}
}
test/unit/cute/hopper/CMakeLists.txt
+5
View File
@@ -56,6 +56,11 @@ cutlass_test_unit_add_executable(
tma_load.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_mcast_load
tma_mcast_load.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_store
tma_store.cu
test/unit/cute/hopper/tma_load.cu
-2
View File
@@ -44,7 +44,6 @@ test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
return test_tma_load<T, TmaType>(SM90_TMA_LOAD{}, gmem_layout, smem_layout, cta_tile);
}
@@ -53,7 +52,6 @@ auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
using namespace cute;
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
test/unit/cute/hopper/tma_mcast_load.cu
+76
View File
@@ -0,0 +1,76 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 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.
*
**************************************************************************************************/
#include "cutlass_unit_test.h"
#include "../hopper/tma_mcast_load_testbed.hpp"
using namespace cute;
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile, class Cluster_Size = Int<1>>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile,
Cluster_Size const& cluster_size = {})
{
return test_tma_load<T, TmaType>(SM90_TMA_LOAD_MULTICAST{}, gmem_layout, smem_layout, cta_tile, cluster_size);
}
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_32x32_Col_MCast)
{
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
{
Layout gmem_layout = make_layout(make_shape(32,32), GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<half_t>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load< float>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<double>(gmem_layout, smem_layout, shape(smem_layout), Int<2>{});
test_tma_load<int8_t>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load<half_t>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load< float>(gmem_layout, smem_layout, shape(smem_layout), 2);
test_tma_load<double>(gmem_layout, smem_layout, shape(smem_layout), 2);
}
}
#endif
test/unit/cute/hopper/tma_mcast_load_testbed.hpp
+242
View File
@@ -0,0 +1,242 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2024 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_unit_test.h"
#include <iostream>
#include <cstdint>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/arch/cluster_sm90.hpp>
#include <cutlass/cluster_launch.hpp>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::ArrayEngine<ElementType, cute::cosize_v<SmemLayout>> smem;
alignas(16) cute::uint64_t tma_load_mbar[1];
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class GmemLayout, class SmemLayout,
class CopyAtom, class CTA_Tiler, class Cluster_Size>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out, GmemLayout gmem_layout, SmemLayout smem_layout,
CUTE_GRID_CONSTANT CopyAtom const tma, CTA_Tiler cta_tiler, Cluster_Size cluster_size)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.begin()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr<T>(g_out), gmem_layout);
Tensor gA = zipped_divide(mA, cta_tiler); // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
Tensor gB = zipped_divide(mB, cta_tiler); // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
#if 1
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA); print("\n");
print(" mB : "); print( mB); print("\n");
print(" gA : "); print( gA); print("\n");
print(" gB : "); print( gB); print("\n");
print(" sA : "); print( sA); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// Prepare the TMA_LOAD
//
Tensor sA_x = make_tensor(sA.data(), make_layout(sA.layout(), Layout<_1>{})); // ((CTA_TILE_M,CTA_TILE_N,...),_1)
Tensor tBgB = gB; // ((CTA_TILE_M,CTA_TILE_N,...),(REST_M,REST_N,...))
int cta_rank_in_cluster = cute::block_rank_in_cluster();
auto [tAgA, tAsA] = tma_partition(tma, cta_rank_in_cluster, make_layout(cluster_size), sA_x, gA);
#if 1
if (thread0()) {
print("sA_x : "); print(sA_x); print("\n");
print("tBgB : "); print(tBgB); print("\n");
print("tAgA : "); print(tAgA); print("\n");
print("tAsA : "); print(tAsA); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// TMA Multicast Masks -- Get a mask of the active ctas in each TMA
//
int elected_cta_rank = 0;
bool elect_one_cta = (elected_cta_rank == cta_rank_in_cluster);
bool elect_one_thr = cute::elect_one_sync();
uint16_t tma_mcast_mask = ((uint16_t(1) << cluster_size) - 1);
#if 1
if (thread0()) {
print("tma_mcast_mask : "); print(tma_mcast_mask); print("\n");
} __syncthreads(); cute::cluster_sync();
#endif
//
// Perform the TMA_LOAD
//
if (elect_one_thr) {
// Initialize TMA barrier
cute::initialize_barrier(tma_load_mbar[0], /* num_threads */ 1);
}
int tma_phase_bit = 0;
// Ensures all CTAs in the Cluster have initialized
__syncthreads();
cute::cluster_sync();
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = sizeof(ArrayEngine<T, size(sA)>);
if (elect_one_thr)
{
cute::set_barrier_transaction_bytes(tma_load_mbar[0], kTmaTransactionBytes);
copy(tma.with(tma_load_mbar[0], tma_mcast_mask), tAgA(_,stage), tAsA(_,0));
}
__syncthreads();
/// Wait on the shared memory barrier until the phase bit flips from tma_phase_bit value
cute::wait_barrier(tma_load_mbar[0], tma_phase_bit);
tma_phase_bit ^= 1;
//
// Write out trivially smem -> gmem
//
// Subbyte elements could cause race conditions, so be even more conservative
if (elect_one_cta && elect_one_thr) {
copy(sA, tBgB(_,stage));
}
__syncthreads();
cute::cluster_sync();
}
}
template <class T, class TmaType = T, class CopyOp,
class GMEM_Layout, class SMEM_Layout,
class CTA_Tiler, class Cluster_Size>
auto
test_tma_load(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
using namespace cute;
// Allocate and initialize host test data
size_t N = ceil_div(cosize(gmem_layout) * sizeof_bits<T>::value, 8);
thrust::host_vector<uint8_t> h_in(N);
for (size_t i = 0; i < h_in.size(); ++i) {
h_in[i] = uint8_t(i % 13);
}
Tensor hA_in = make_tensor(recast_ptr<T>(h_in.data()), gmem_layout);
// Allocate and initialize device test data
thrust::device_vector<uint8_t> d_in = h_in;
thrust::device_vector<uint8_t> d_out(h_in.size(), uint8_t(-1)); // overflow uint
// Create TMA for this device Tensor
Tensor gA = make_tensor(make_gmem_ptr<T>(raw_pointer_cast(d_in.data())), gmem_layout);
auto tma = make_tma_atom<TmaType>(copy_op, gA, smem_layout, cta_tiler, cluster_size);
//print(tma);
// Launch
dim3 dimBlock(32);
dim3 dimCluster(size(cluster_size));
dim3 dimGrid = dimCluster;
int smem_size = sizeof(SharedStorage<T, SMEM_Layout>);
void* kernel_ptr = (void*) &tma_test_device_cute<T, GMEM_Layout, SMEM_Layout,
decltype(tma), CTA_Tiler, Cluster_Size>;
cutlass::launch_kernel_on_cluster({dimGrid, dimBlock, dimCluster, smem_size},
kernel_ptr,
reinterpret_cast<T const*>(raw_pointer_cast(d_in.data())),
reinterpret_cast<T *>(raw_pointer_cast(d_out.data())),
gmem_layout,
smem_layout,
tma, cta_tiler, cluster_size);
// Copy results back to host
thrust::host_vector<uint8_t> h_out = d_out;
Tensor hA_out = make_tensor(recast_ptr<T>(h_out.data()), gmem_layout);
// Validate the results. Print only the first 3 errors.
int count = 3;
for (int i = 0; i < int(size(hA_out)) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
return tma;
}
#endif
} // end namespace cutlass::test
test/unit/cute/turing/CMakeLists.txt
+32
View File
@@ -0,0 +1,32 @@
# Copyright (c) 2023 - 2024 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.
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_turing
cooperative_gemm.cu
)
test/unit/cute/turing/cooperative_gemm.cu
+58
View File
@@ -0,0 +1,58 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2024 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.
*
**************************************************************************************************/
#include "cutlass_unit_test.h"
#include <cute/tensor.hpp>
#include "../cooperative_gemm_common.hpp"
using namespace cute;
TEST(SM75_CuTe_Turing, CooperativeGemm1_MixedPrecisionFP16FP32_MMA) {
using TA = cutlass::half_t;
using TB = cutlass::half_t;
using TC = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 64;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<SM75_16x8x8_F32F16F16F32_TN>,
Layout<Shape<_2, _2, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 128, TA, TB, TC>();
}
test/unit/cute/volta/CMakeLists.txt
+1
View File
@@ -30,4 +30,5 @@ cutlass_test_unit_add_executable(
cutlass_test_unit_cute_volta
vectorization_auto.cu
cooperative_copy.cu
cooperative_gemm.cu
)
test/unit/cute/volta/cooperative_copy.cu
+48
View File
@@ -263,6 +263,21 @@ TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefault1D)
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefault1DFallback)
{
using value_type = float;
constexpr uint32_t count = 99;
using gmem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<count>{})));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2D)
{
using value_type = float;
@@ -279,6 +294,22 @@ TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2D)
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2DFallback)
{
using value_type = float;
constexpr uint32_t x = 37;
constexpr uint32_t y = 37;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{})));
constexpr uint32_t thread_block_size = 64;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2DCustomStride)
{
using value_type = float;
@@ -312,6 +343,23 @@ TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG3D)
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG3DFallback)
{
using value_type = cute::half_t;
constexpr uint32_t x = 44;
constexpr uint32_t y = 24;
constexpr uint32_t z = 14;
using gmem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
using smem_layout_t = decltype(make_layout(make_shape(Int<x>{}, Int<y>{}, Int<z>{})));
constexpr uint32_t thread_block_size = 128;
test_cooperative_copy_default<typename TestFixture::mode,
TestFixture::max_vec_bits,
gmem_layout_t,
smem_layout_t,
thread_block_size,
value_type>();
}
TYPED_TEST(SM70_CuTe_Volta, CooperativeCopyDefaultGSSG2Dto3D)
{
using value_type = double;
test/unit/cute/volta/cooperative_gemm.cu
+421
View File
@@ -0,0 +1,421 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2024 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.
*
**************************************************************************************************/
#include "cutlass_unit_test.h"
#include <cute/tensor.hpp>
#include "../cooperative_gemm_common.hpp"
using namespace cute;
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA) {
using value_type = float;
constexpr uint32_t m = 64;
constexpr uint32_t n = 32;
constexpr uint32_t k = 16;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_16, _8, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication) {
using value_type = float;
constexpr uint32_t m = 88;
constexpr uint32_t n = 20;
constexpr uint32_t k = 12;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_2, _64, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication2) {
using value_type = float;
constexpr uint32_t m = 88;
constexpr uint32_t n = 36;
constexpr uint32_t k = 24;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_4, _32, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm1_FloatFMA_Predication3) {
using value_type = float;
constexpr uint32_t m = 67;
constexpr uint32_t n = 13;
constexpr uint32_t k = 11;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_1, _128, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm2_DoubleFMA) {
using value_type = double;
constexpr uint32_t m = 16;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t =
TiledMMA<
MMA_Atom<UniversalFMA<value_type, value_type, value_type, value_type>>,
Layout<Shape<_16, _8, _1>>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm3_Float_FMA_CustomPermutationMNK) {
using value_type = float;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 256;
using tiled_mma_t = TiledMMA<
MMA_Atom<
UniversalFMA<value_type, value_type, value_type, value_type>
>,
Layout<
Shape<_16, _16, _1>
>,
Tile<
Layout<
Shape<_16,_2>, Stride<_2,_1>
>, // 32x32x1 MMA with perm for load vectorization
Layout<
Shape<_16,_2>, Stride<_2,_1>
>,
Underscore
>
>;
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm4_Half_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
using smem_a_atom_layout_t = typename tiled_mma_t::AtomLayoutB_TV;
using smem_b_atom_layout_t = typename tiled_mma_t::AtomLayoutA_TV;
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {})));
test_cooperative_gemm_col_major_layout<smem_a_atom_layout_t,
smem_b_atom_layout_t,
smem_c_atom_layout_t,
m,
n,
k,
thread_block_size,
tiled_mma_t,
value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm5_Half_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
using smem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using smem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using smem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
value_type,
value_type,
value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm5_Half_MMA_Predicated) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 31;
constexpr uint32_t n = 27;
constexpr uint32_t k = 17;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
using smem_a_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<k>{})));
using smem_b_layout_t = decltype(make_layout(make_shape(Int<n>{}, Int<k>{}), GenColMajor{}));
using smem_c_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{})));
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<16>, // A
AutoVectorizingCopyWithAssumedAlignment<16>, // B
AutoVectorizingCopyWithAssumedAlignment<16>, // C
thread_block_size,
tiled_mma_t,
16,
value_type,
value_type,
value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm6_Half_MAA_SwizzledSmemLayouts) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 128;
constexpr uint32_t n = 128;
constexpr uint32_t k = 64;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
using smem_a_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape < _8,_64>,
Stride<_64, _1>>{}));
using smem_b_atom_layout_t = decltype(
composition(Swizzle<3,3,3>{},
Layout<Shape <_64, _8>,
Stride< _1,_64>>{}));
using smem_c_atom_layout_t = decltype(make_layout(make_shape(Int<m>{}, Int<n>{}), GenRowMajor{}));
using gmem_a_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<k> {}), GenRowMajor{}));
using gmem_b_layout_t = decltype(make_layout(make_shape(Int<n> {}, Int<k> {}), GenColMajor{}));
using gmem_c_layout_t = decltype(make_layout(make_shape(Int<m> {}, Int<n> {}), GenRowMajor{}));
using smem_a_atom_layout_t = smem_a_atom_layout_t;
using smem_a_layout_t = decltype(tile_to_shape(
smem_a_atom_layout_t{},
make_shape(shape<0>(gmem_a_layout_t{}), shape<1>(gmem_a_layout_t{})))
);
// Transposed
using smem_b_atom_layout_t = smem_b_atom_layout_t;
using smem_b_layout_t = decltype(tile_to_shape(
smem_b_atom_layout_t{},
make_shape(shape<0>(gmem_b_layout_t{}), shape<1>(gmem_b_layout_t{})))
);
using smem_c_atom_layout_t = smem_c_atom_layout_t;
using smem_c_layout_t = decltype(tile_to_shape(
smem_c_atom_layout_t{},
make_shape(shape<0>(gmem_c_layout_t{}), shape<1>(gmem_c_layout_t{})))
);
test_cooperative_gemm<gmem_a_layout_t,
gmem_b_layout_t,
gmem_c_layout_t,
smem_a_layout_t,
smem_b_layout_t,
smem_c_layout_t,
AutoVectorizingCopyWithAssumedAlignment<128>, // A
AutoVectorizingCopyWithAssumedAlignment<128>, // B
AutoVectorizingCopyWithAssumedAlignment<128>, // C
thread_block_size,
tiled_mma_t,
128,
value_type,
value_type,
value_type>();
}
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformNegate_FMA) {
using TA = float;
using TB = float;
using TC = double;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<UniversalFMA<TC, TA, TB, TC>>,
Layout<Shape<_16, _8, _1>>
>;
auto aload = cute::negate {};
auto bload = cute::negate {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 64, TA, TB, TC>(
aload, bload, cload, cstore);
}
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformNegate_MMA) {
using value_type = cutlass::half_t;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<SM70_8x8x4_F16F16F16F16_TN>,
Layout<Shape<_4, _4, _1>>
>;
auto aload = cute::negate {};
auto bload = cute::negate {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, value_type>(
aload, bload, cload, cstore);
}
template<class ConstantType>
struct increment_by_x {
ConstantType x;
template <class T>
CUTE_HOST_DEVICE constexpr
T operator()(const T& arg) const {
return arg + x;
}
};
template<class From, class To>
struct convert_to {
CUTE_HOST_DEVICE constexpr
To operator()(const From& arg) const {
return static_cast<To>(arg);
}
};
TEST(SM70_CuTe_Volta, CooperativeGemm7_TransformCustomOp_FMA) {
using TA = float;
using TB = float;
using TC = double;
constexpr uint32_t m = 32;
constexpr uint32_t n = 32;
constexpr uint32_t k = 32;
constexpr uint32_t thread_block_size = 128;
using tiled_mma_t = TiledMMA<
MMA_Atom<UniversalFMA<TC, TA, TB, TC>>,
Layout<Shape<_16, _8, _1>>
>;
auto aload = increment_by_x<float>{1.111f};
auto bload = convert_to<float, double> {};
auto cload = cute::negate {};
auto cstore = cute::negate {};
test_cooperative_gemm_col_major_layout<m, n, k, thread_block_size, tiled_mma_t, 64, TA, TB, TC>(
aload, bload, cload, cstore);
}