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

CUTLASS 3.0.0 (#786)

Browse Source 
* CUTLASS 3.0.0
This commit is contained in:
Vijay Thakkar
2023-01-23 20:55:28 -05:00
committed by GitHub
parent 66d9cddc83
commit 277bd6e537
377 changed files with 76396 additions and 1186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/cute/CMakeLists.txt
+50
View File
@@ -0,0 +1,50 @@
# Copyright (c) 2023 - 2023 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.
add_subdirectory(core)
add_subdirectory(ampere)
add_subdirectory(hopper)
add_subdirectory(layout)
add_custom_target(
cutlass_test_unit_cute
DEPENDS
cutlass_test_unit_cute_layout
cutlass_test_unit_cute_core
cutlass_test_unit_cute_ampere
cutlass_test_unit_cute_hopper
)
add_custom_target(
test_unit_cute
DEPENDS
test_unit_cute_layout
test_unit_cute_core
test_unit_cute_ampere
test_unit_cute_hopper
)
test/unit/cute/ampere/CMakeLists.txt
+33
View File
@@ -0,0 +1,33 @@
# Copyright (c) 2023 - 2023 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_ampere
cp_async.cu
ldsm.cu
)
test/unit/cute/ampere/cp_async.cu
+104
View File
@@ -0,0 +1,104 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <iomanip>
#include <utility>
#include <type_traits>
#include <vector>
#include <numeric>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
using namespace cute;
__global__ void
test(double const* g_in, double* g_out)
{
extern __shared__ double smem[];
smem[threadIdx.x] = g_in[threadIdx.x];
__syncthreads();
g_out[threadIdx.x] = 2 * smem[threadIdx.x];
}
__global__ void
test2(double const* g_in, double* g_out)
{
using namespace cute;
extern __shared__ double smem[];
auto s_tensor = make_tensor(make_smem_ptr(smem + threadIdx.x), Int<1>{});
auto g_tensor = make_tensor(make_gmem_ptr(g_in + threadIdx.x), Int<1>{});
copy(g_tensor, s_tensor);
cp_async_fence();
cp_async_wait<0>();
__syncthreads();
g_out[threadIdx.x] = 2 * smem[threadIdx.x];
}
TEST(SM80_CuTe_Ampere, CpAsync)
{
constexpr int count = 32;
thrust::host_vector<double> h_in(count);
for (int i = 0; i < count; ++i) {
h_in[i] = double(i);
}
thrust::device_vector<double> d_in(h_in);
thrust::device_vector<double> d_out(count, -1);
test<<<1, count, sizeof(double) * count>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<double> h_result = d_out;
thrust::device_vector<double> d_out_cp_async(count, -2);
test2<<<1, count, sizeof(double) * count>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out_cp_async.data()));
thrust::host_vector<double> h_result_cp_async = d_out_cp_async;
for (int i = 0; i < count; ++i) {
EXPECT_EQ(h_result[i], h_result_cp_async[i]);
}
}
test/unit/cute/ampere/ldsm.cu
+431
View File
@@ -0,0 +1,431 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/atom/copy_traits_sm75.hpp>
using namespace cute;
template <class T>
__global__ void
ldsm_test_device(uint16_t* g_in, uint16_t* g_out)
{
constexpr int count = sizeof(T) / 4;
int tid = threadIdx.x;
int stride = blockDim.x;
// load input gmem -> smem
__shared__ uint32_t smem[32 * count];
for (int i = 0; i < count; ++i) {
smem[tid + (stride * i)] = reinterpret_cast<uint32_t*>(g_in)[tid + (stride * i)];
}
__syncthreads();
uint32_t reg[count];
for (int i = 0; i < count; ++i) {
reg[i] = 0;
}
// load smem -> rmem using LDSM
uint128_t* smem_ptr = reinterpret_cast<uint128_t*>(smem) + tid;
T* rmem_ptr = reinterpret_cast<T*>(reg);
cute::copy_ldsm(smem_ptr, rmem_ptr);
// store output rmem -> gmem
for (int i = 0; i < count; ++i) {
reinterpret_cast<uint32_t*>(g_out)[tid + (stride * i)] = reg[i];
}
}
template <class TiledCopy, class SmemLayout>
__global__ void
ldsm_test_device_cute(uint16_t* g_in, uint16_t* g_out,
TiledCopy tiled_copy, SmemLayout smem_layout)
{
using namespace cute;
__shared__ uint16_t smem[size(smem_layout)];
auto t_g_in = make_tensor(make_gmem_ptr(g_in), smem_layout);
auto t_g_out = make_tensor(make_gmem_ptr(g_out), smem_layout);
auto t_smem = make_tensor(make_smem_ptr(smem), smem_layout);
int tid = threadIdx.x;
// Load input gmem -> smem
for (int i = tid; i < size(t_smem); i += size(tiled_copy)) {
t_smem(i) = t_g_in(i);
}
__syncthreads();
auto thr_copy = tiled_copy.get_thread_slice(tid);
auto tXsX = thr_copy.partition_S(t_smem); // (V,M,N)
auto tXgX = thr_copy.partition_D(t_g_out); // (V,M,N)
auto tXrX = make_tensor<uint16_t>(shape(tXgX)); // (V,M,N)
clear(tXrX); // Just to make sure
/*
if (thread0()) {
print("tXsX: " ); print(tXsX.layout()); print("\n");
print("tXgX: " ); print(tXgX.layout()); print("\n");
print("tXrX: " ); print(tXrX.layout()); print("\n");
}
*/
// Copy smem -> rmem via tiled_copy (LDSM, LDS)
copy(tiled_copy, tXsX, tXrX);
// Output rmem -> gmem
copy(tXrX, tXgX);
}
TEST(SM80_CuTe_Ampere, Ldsm)
{
constexpr int count = 1024;
thrust::host_vector<uint16_t> h_in(count);
for (int i = 0; i < count; ++i) {
h_in[i] = uint16_t(i);
}
thrust::device_vector<uint16_t> d_in = h_in;
//
// LDSM 1x (32b)
//
{
thrust::device_vector<uint16_t> d_out(count);
ldsm_test_device<uint32_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 32; ++i) {
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("LDSM 1x ldsm_test_device SUCCESS\n");
}
//
// LDSM 2x (64b)
//
{
thrust::device_vector<uint16_t> d_out(count);
ldsm_test_device<uint64_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 64; ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("LDSM 2x ldsm_test_device SUCCESS\n");
}
//
// LDSM 4x (128b)
//
{
thrust::device_vector<uint16_t> d_out(count);
ldsm_test_device<uint128_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 128; ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("LDSM 4x ldsm_test_device SUCCESS\n");
}
//
// CuTe LDSM
//
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x1_LDSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x1_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x2_LDSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x2_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x4_LDSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x4_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<UniversalCopy<uint16_t>, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i] , h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved LDS.U16 SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x1_LDSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x1_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x2_LDSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x2_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U32x4_LDSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x4_LDSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<UniversalCopy<uint16_t>, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 LDS.U16 SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U16x2_LDSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_2,_1>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x2_LDSM_T SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U16x4_LDSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_4,_1>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x4_LDSM_T SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM75_U16x8_LDSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_8,_1>>{});
ldsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x8_LDSM_T SUCCESS\n");
}
CUTLASS_TRACE_HOST("PASS");
}
test/unit/cute/core/CMakeLists.txt
+44
View File
@@ -0,0 +1,44 @@
# Copyright (c) 2023 - 2023 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_core
bitfield.cpp
coalesce.cpp
compare.cpp
complement.cpp
composition.cpp
inverse_left.cpp
inverse_right.cpp
logical_divide.cpp
logical_product.cpp
mixedbits.cpp
transform.cpp
tuple.cpp
)
test/unit/cute/core/bitfield.cpp
+84
View File
@@ -0,0 +1,84 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <iomanip>
#include <utility>
#include <type_traits>
#include <vector>
#include <numeric>
#include <cute/tensor.hpp>
#include <cute/container/bit_field.hpp>
using namespace cute;
TEST(CuTe_core, Bitfield)
{
for_each(make_int_range<1,65>{}, [&](auto NumBits) {
for_each(make_int_range<0,129>{}, [&](auto BitStart) {
using BF = bit_field<decltype(BitStart)::value, decltype(NumBits)::value>;
#if 0
printf("bit_field<%d,%d>:\n", decltype(BitStart)::value, decltype(NumBits)::value);
printf(" value_type_bits : %d\n", BF::value_type_bits);
printf(" storage_type_bits: %d\n", BF::storage_type_bits);
printf(" N : %d\n", BF::N);
printf(" idx : %d\n", BF::idx);
printf(" bit_lo : %d\n", BF::bit_lo);
printf(" bit_hi : %d\n", BF::bit_hi);
printf(" mask : 0x%lx\n", uint64_t(BF::mask));
printf(" mask_lo : 0x%lx\n", uint64_t(BF::mask_lo));
printf(" mask_hi : 0x%lx\n", uint64_t(BF::mask_hi));
#endif
// Test
uint64_t v = decltype(NumBits)::value == 64 ? uint64_t(-1) : ((uint64_t(1) << NumBits) - 1);
BF bf{};
bf = v;
EXPECT_EQ(v, uint64_t(bf));
});
});
for_each(make_int_range<0,129>{}, [&](auto BitStart) {
using BF = bit_field<decltype(BitStart)::value, 32, float>;
BF bf{};
bf = 3.14f;
EXPECT_EQ(3.14f, float(bf));
});
}
test/unit/cute/core/coalesce.cpp
+182
View File
@@ -0,0 +1,182 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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>
using namespace cute;
template <class Layout>
void
test_coalesce(Layout const& layout)
{
auto coalesce_layout = coalesce(layout);
CUTLASS_TRACE_HOST(shape (layout) << " => " << shape (coalesce_layout));
CUTLASS_TRACE_HOST(stride(layout) << " " << stride(coalesce_layout));
CUTE_STATIC_ASSERT_V(depth(coalesce_layout) <= Int<1>{});
ASSERT_EQ(size(coalesce_layout), size(layout));
for (int i = 0; i < size(layout); ++i) {
EXPECT_EQ(coalesce_layout(i), layout(i));
}
}
TEST(CuTe_core, Coalesce)
{
{
auto layout = make_layout(Int<1>{}, Int<0>{});
test_coalesce(layout);
}
{
auto layout = make_layout(Int<1>{}, Int<1>{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}, Int<6>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape (Int<2>{}, Int<1>{}, Int<6>{}),
make_stride(Int<1>{}, Int<6>{}, Int<2>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape (Int<2>{}, Int<1>{}, Int<6>{}),
make_stride(Int<1>{}, 7, Int<2>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape (Int<2>{}, Int<1>{}, Int<6>{}),
make_stride(Int<4>{}, 7, Int<8>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(2, Int<4>{}, Int<6>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, 4, Int<6>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}, 6));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}, Int<6>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(2, Int<4>{}, Int<6>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, 4, Int<6>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<4>{}, 6), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, Int<1>{}, Int<3>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, 1, Int<3>{}), GenRowMajor{});
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, 1, Int<3>{}), make_stride(Int<2>{}, 4, Int<4>{}));
test_coalesce(layout);
}
{
auto layout = make_layout(make_shape(Int<2>{}, 1, Int<3>{}), make_stride(Int<2>{}, Int<0>{}, Int<4>{}));
test_coalesce(layout);
}
{
auto layout = Layout<Shape<Shape<_2,_2>,Shape<_2, _2>>,
Stride<Stride<_1,_4>,Stride<_8,_32>>>{};
test_coalesce(layout);
}
}
test/unit/cute/core/compare.cpp
+168
View File
@@ -0,0 +1,168 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
TEST(CuTe_core, Compare_simple_2d_GenColMajor)
{
using namespace cute;
// Simple 2D layout
auto layout = make_layout(make_shape(Int<3>{}, Int<5>{}), GenColMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
CUTLASS_TRACE_HOST(" " << j << ": " << coord_j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), colex_less(coord_i,coord_j));
}
}
}
TEST(CuTe_core, Compare_simple_2d_GenRowMajor)
{
using namespace cute;
auto layout = make_layout(make_shape(Int<3>{}, Int<5>{}), GenRowMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), lex_less(coord_i,coord_j));
}
}
}
TEST(CuTe_core, Compare_simple_3d_GenColMajor)
{
using namespace cute;
auto layout = make_layout(make_shape(Int<2>{}, Int<3>{}, Int<5>{}), GenColMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), colex_less(coord_i,coord_j));
}
}
}
TEST(CuTe_core, Compare_simple_3d_GenRowMajor)
{
using namespace cute;
auto layout = make_layout(make_shape(Int<2>{}, Int<3>{}, Int<5>{}), GenRowMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), lex_less(coord_i,coord_j));
}
}
}
TEST(CuTe_core, Compare_hierarchical_3d_GenColMajor)
{
using namespace cute;
auto layout = make_layout(Shape<Shape<_3,_2>,Shape<_5,_2,_2>>{}, GenColMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), colex_less(coord_i,coord_j));
}
}
}
TEST(CuTe_core, Compare_hierarchical_3d_GenRowMajor)
{
using namespace cute;
auto layout = make_layout(Shape<Shape<_3,_2>,Shape<_5,_2,_2>>{}, GenRowMajor{});
CUTLASS_TRACE_HOST("Layout: " << layout);
for (int i = 0; i < size(layout); ++i) {
auto coord_i = layout.get_hier_coord(i);
CUTLASS_TRACE_HOST(i << ": " << coord_i);
EXPECT_TRUE(elem_less(coord_i, shape(layout)));
for (int j = 0; j < size(layout); ++j) {
auto coord_j = layout.get_hier_coord(j);
EXPECT_TRUE(elem_less(coord_j, shape(layout)));
EXPECT_EQ((i < j), lex_less(coord_i,coord_j));
}
}
}
test/unit/cute/core/complement.cpp
+273
View File
@@ -0,0 +1,273 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
template <class Layout, class CoSizeHi>
void
test_complement(Layout const& layout, CoSizeHi const& cosize_hi)
{
using namespace cute;
auto result = complement(layout, cosize_hi);
CUTLASS_TRACE_HOST("complement( " << layout << ", " << cosize_hi << ") => " << result);
// Post-condition on the domain size of the complement (1)
EXPECT_GE( size(result), cosize_hi / size(filter(layout)));
// Post-condition on the codomain size of the complement (2)
EXPECT_LE(cosize(result), cute::ceil_div(cosize_hi, cosize(layout)) * cosize(layout));
// Post-condition on the codomain of the complement
for (int i = 1; i < size(result); ++i) {
EXPECT_LT(result(i-1), result(i)); // Ordered (3)
for (int j = 0; j < size(layout); ++j) {
EXPECT_NE(result(i), layout(j)); // Complemented (4)
}
}
// Other observations
EXPECT_LE(size(result),cosize(result)); // As a result of the ordered condition (3)
EXPECT_GE(cosize(result), cosize_hi / size(filter(layout))); // As a result of (1) (2) and (5)
if constexpr (is_static<decltype(stride(make_layout(layout,result)))>::value) { // If we can apply complement again
EXPECT_EQ(size(complement(make_layout(layout,result))), 1); // There's no more codomain left over
}
}
template <class Layout>
void
test_complement(Layout const& layout)
{
return test_complement(layout, cosize(layout));
}
TEST(CuTe_core, Complement)
{
using namespace cute;
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("COMPLEMENT");
CUTLASS_TRACE_HOST("-------------------------------");
{
auto layout = Layout<_1,_0>{};
test_complement(layout);
test_complement(layout, Int<2>{});
}
{
auto layout = Layout<_1,_1>{};
test_complement(layout);
test_complement(layout, Int<2>{});
}
{
auto layout = Layout<_1,_2>{};
test_complement(layout, Int<1>{});
test_complement(layout, Int<2>{});
test_complement(layout, Int<8>{});
}
{
auto layout = Layout<_4,_0>{};
test_complement(layout, Int<1>{});
test_complement(layout, Int<2>{});
test_complement(layout, Int<8>{});
}
{
auto layout = Layout<_4,_1>{};
test_complement(layout, Int<1>{});
test_complement(layout, Int<2>{});
test_complement(layout, Int<8>{});
}
{
auto layout = Layout<_4,_2>{};
test_complement(layout, Int<1>{});
test_complement(layout);
test_complement(layout, Int<16>{});
}
{
auto layout = Layout<_4,_4>{};
test_complement(layout, Int<1>{});
test_complement(layout);
test_complement(layout, Int<17>{});
}
{
auto layout = Layout<Shape<_2,_4>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_2,_3>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_2,_4>, Stride<_1,_4>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_2,_4,_8>, Stride<_8,_1,_64>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_2,_4,_8>, Stride<_8,_1,_0>>{};
test_complement(layout);
test_complement(layout, Int<460>{});
}
{
auto layout = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
Stride<Stride<_1,_4>,Stride<_8,_32>>{});
test_complement(layout);
}
{
auto layout = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
Stride<Stride<_1,_32>,Stride<_8,_4>>{});
test_complement(layout);
}
// Fails due to non-injective input
//{
//auto layout = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
// Stride<Stride<_1,_8>,Stride<_8,_4>>{});
//test_complement(layout);
//}
{
auto layout = Layout<Shape<_4,_6>, Stride<_1,_6>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_4,_2>, Stride<_1,_10>>{};
test_complement(layout);
}
{
auto layout = Layout<Shape<_4,_2>, Stride<_1,_16>>{};
test_complement(layout);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Dynamic shapes/strides");
CUTLASS_TRACE_HOST("-------------------------------");
{
auto layout = make_layout(12);
test_complement(layout, 1);
test_complement(layout);
test_complement(layout, 53);
test_complement(layout, 128);
}
{
auto layout = make_layout(12, 1);
test_complement(layout, 1);
test_complement(layout);
test_complement(layout, 53);
test_complement(layout, 128);
}
{
auto layout = make_layout(12, Int<2>{});
test_complement(layout, 1);
test_complement(layout);
test_complement(layout, 53);
test_complement(layout, 128);
}
{
auto layout = make_layout(12, 2);
test_complement(layout, 1);
test_complement(layout);
test_complement(layout, 53);
test_complement(layout, 128);
}
{
auto layout = make_layout(make_shape(3,6),make_stride(_1{}, _3{}));
test_complement(layout);
}
{
auto layout = make_layout(make_shape(3,6),make_stride(_1{}, _9{}));
test_complement(layout);
}
{
auto layout = make_layout(make_shape(3,6),make_stride(_1{}, _10{}));
test_complement(layout);
}
{
auto layout = make_layout(make_shape(make_shape(2,2), make_shape(2,2)),
Stride<Stride<_1,_4>,Stride<_8,_32>>{});
test_complement(layout);
}
}
test/unit/cute/core/composition.cpp
+528
View File
@@ -0,0 +1,528 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <iostream>
#include <cute/tensor.hpp>
using namespace cute;
template <class LayoutA, class LayoutB>
void
test_composition(const LayoutA& layoutA,
const LayoutB& layoutB)
{
auto layoutR = composition(layoutA, layoutB);
CUTLASS_TRACE_HOST("test_composition()");
CUTLASS_TRACE_HOST(layoutA << " o " << layoutB);
CUTLASS_TRACE_HOST(" => ");
CUTLASS_TRACE_HOST(layoutR);
// Test that layout R is compatible with layout B
EXPECT_TRUE(compatible(layoutB, layoutR));
// True post-condition: Every coordinate c of layoutB with L1D(c) < size(layoutR) is a coordinate of layoutR.
// Test that R(c) = A(B(c)) for all coordinates c in layoutR
for (int i = 0; i < size(layoutR); ++i) {
EXPECT_EQ(layoutR(i), layoutA(layoutB(i)));
}
}
TEST(CuTe_core, Composition)
{
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("COMPOSITION" );
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Simple tests" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = Layout<_1,_0>{};
auto b = Layout<_1,_0>{};
test_composition(a, b);
}
{
auto a = Layout<_1,_0>{};
auto b = Layout<_1,_1>{};
test_composition(a, b);
}
{
auto a = Layout<_1,_1>{};
auto b = Layout<_1,_0>{};
test_composition(a, b);
}
{
auto a = Layout<_1,_1>{};
auto b = Layout<_1,_1>{};
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_2>{});
auto b = make_layout(Shape<_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_0>{});
auto b = make_layout(Shape<_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_4>{}, Stride<_0>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_1>{}, Stride<_0>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_2>{});
auto b = make_layout(Shape<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_2>{}, Stride<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_2>{});
auto b = make_layout(Shape<_2>{}, Stride<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_12>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_12>{});
auto b = make_layout(Shape<_4,_3>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_12>{}, Stride<_2>{});
auto b = make_layout(Shape<_4,_3>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_12>{});
auto b = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_12>{}, Stride<_2>{});
auto b = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_12>{});
auto b = make_layout(Shape<_2,_3>{}, Stride<_2,_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_4,_3>{});
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>{});
// test_composition(a, b);
//}
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_6>{}, Stride<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{});
auto b = make_layout(Shape<_6,_2>{}, Stride<_2,_1>{});
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>{});
// test_composition(a, b);
//}
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_4,_3>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_12>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_6>{}, Stride<_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_6,_2>{}, Stride<_2,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_8,_8>{});
auto b = make_layout(Shape<Shape<_2, _2,_2>, Shape<_2,_2, _2>>{},
Stride<Stride<_1,_16,_4>, Stride<_8,_2,_32>>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_8,_8>{}, Stride<_8,_1>{});
auto b = make_layout(Shape<Shape<_2, _2,_2>, Shape<_2,_2, _2>>{},
Stride<Stride<_1,_16,_4>, Stride<_8,_2,_32>>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<Shape<_4,_2>>{}, Stride<Stride<_1,_16>>{});
auto b = make_layout(Shape<_4,_2>{}, Stride<_2,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_2,_2>{}, Stride<_2,_1>{});
auto b = make_layout(Shape<_2,_2>{}, Stride<_2,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_8,_2>{});
auto b = make_layout(Shape<_2,_2,_2>{}, Stride<_2,_8,_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_8,_2>{}, Stride<_2,_8,_1>{});
auto b = make_layout(Shape<_2,_2,_2>{}, Stride<_1,_8,_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_8,_2>{}, Stride<_2,_8,_1>{});
auto b = make_layout(Shape<_4,_2,_2>{}, Stride<_2,_8,_1>{});
test_composition(a, b);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Dynamic shapes/strides" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = make_layout(12, 1);
auto b = make_layout(_4{}, _1{});
test_composition(a, b);
}
{
auto a = make_layout(12, 1);
auto b = make_layout(_4{}, 1);
test_composition(a, b);
}
{
auto a = make_layout(12, _1{});
auto b = make_layout(_4{}, 1);
test_composition(a, b);
}
{
auto a = make_layout(12, _1{});
auto b = make_layout(_4{}, _1{});
test_composition(a, b);
}
{
auto a = make_layout(make_shape(12,3), make_stride(1,24));
auto b = make_layout(Shape<_4>{}, Stride<_1>{});
test_composition(a, b);
}
{
auto a = make_layout(16, 2);
auto b = make_layout(4, 2);
test_composition(a, b);
}
{
auto a = make_layout(make_shape(128,24,5), make_stride(1,128,3072));
auto b = make_layout(64, 2);
test_composition(a, b);
}
{
auto a = make_layout(make_shape(128,24,5), make_stride(1,128,3072));
auto b = make_layout(480, Int<32>{});
test_composition(a, b);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("cosize(b) > size(a) and divisibility");
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = make_layout(Shape<_1>{}, Stride<_0>{});
auto b = make_layout(Shape<_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_1>{}, Stride<_1>{});
auto b = make_layout(Shape<_4>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{});
auto b = make_layout(Shape<_4>{}, Stride<_2>{});
test_composition(a, b);
}
// Last mode gets extended
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_24>{});
test_composition(a, b);
}
// Last mode extension even without last mode divisibility
{
auto a = make_layout(Shape<_4,_3>{}, Stride<_3,_1>{});
auto b = make_layout(Shape<_8>{});
test_composition(a, b);
}
// Capping a Layout with 1:0 forces divisibility and extends in stride-0
{
auto a = make_layout(Shape<_4,_3,_1>{}, Stride<_3,_1,_0>{});
auto b = make_layout(Shape<_24>{});
test_composition(a, b);
}
{
auto a = make_layout(3, _1{});
auto b = make_layout(_4{}, _1{});
test_composition(a, b);
}
{
auto a = make_layout(make_shape(48,24,5), make_stride(_1{},128,3072));
auto b = make_layout(32, Int<1>{});
test_composition(a, b);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Swizzle composition" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = Layout<Shape<_8,_8>, Stride<_8,_1>>{};
auto b = composition(Swizzle<2,0,-3>{}, Layout<Shape<_8,_8>, Stride<_8,_1>>{});
test_composition(a, b);
}
{
auto a = composition(Swizzle<2,0, 3>{}, Layout<Shape<_8,_8>, Stride<_8,_1>>{});
auto b = composition(Swizzle<2,0,-3>{}, Layout<Shape<_8,_8>, Stride<_8,_1>>{});
test_composition(a, b);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("BETA: Negative strides" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = make_layout(Shape<_4>{}, Stride<_m1>{});
auto b = make_layout(Shape<_4>{}, Stride<_1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_1>{});
auto b = make_layout(Shape<_4>{}, Stride<_m1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_m1>{});
auto b = make_layout(Shape<_4>{}, Stride<_m1>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4>{}, Stride<_1>{});
auto b = make_layout(Shape<_4>{}, Stride<_m2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_4>{}, Stride<_m1,_1>{});
auto b = make_layout(Shape<_2,_4,_2>{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_4,_4>{}, Stride<_m1,_1>{});
auto b = make_layout(Shape<_2,_4,_2>{}, Stride<_1,_4,_2>{});
test_composition(a, b);
}
// The SM80 fp64 MMA NT problem
{
auto a = make_layout(Shape<_1,Shape<_2,_4>>{}, Stride<_0,Stride<_m1,_512>>{});
auto b = make_layout(_2{}, _m1{});
test_composition(a, b);
}
{
auto a = make_layout(Shape<_1,Shape<_2,_4>>{}, Stride<_0,Stride<_m1,_512>>{});
auto b = make_layout(_4{}, _m1{});
test_composition(a, b);
}
}
test/unit/cute/core/inverse_left.cpp
+183
View File
@@ -0,0 +1,183 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <iostream>
#include <cute/tensor.hpp>
using namespace cute;
template <class Layout>
void
test_left_inverse(Layout const& layout)
{
auto inv_layout = left_inverse(layout);
CUTLASS_TRACE_HOST(layout << " ^ -1\n" << " => \n" << inv_layout);
for (int i = 0; i < size(layout); ++i) {
//printf("%3d: %3d %3d\n", i, int(layout(i)), int(inv_layout(layout(i))));
EXPECT_EQ(inv_layout(layout(i)), i);
}
CUTLASS_TRACE_HOST("Composition: " << coalesce(composition(inv_layout, layout)));
}
TEST(CuTe_core, Inverse_left)
{
{
auto layout = Layout<Shape <_1>,
Stride<_0>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <Shape <_1,_1>>,
Stride<Stride<_0,_0>>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_1>,
Stride<_1>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_4>,
Stride<_1>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_4>,
Stride<_2>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_8, _4>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_8, _4>,
Stride<_4, _1>>{};
test_left_inverse(filter(layout));
}
{
auto layout = Layout<Shape< _2,_4,_6>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_2,_4,_6>,
Stride<_4,_1,_8>>{};
test_left_inverse(layout);
}
{
auto layout = Layout<Shape <_4, _2>,
Stride<_1,_16>>{};
test_left_inverse(layout);
}
//
// Swizzle left_inverse
//
{
auto layout = ComposedLayout<Swizzle<1,0,2>, _0, Layout<Shape <_4, _4>,
Stride<_1, _4>>>{};
test_left_inverse(layout);
}
{
auto layout = ComposedLayout<Swizzle<1,0,2>, _0, Layout<Shape <_4, _4>,
Stride<_4, _1>>>{};
test_left_inverse(layout);
}
{
auto layout = ComposedLayout<Swizzle<1,0,1>, _0, Layout<Shape <_4, _4>,
Stride<_8, _1>>>{};
test_left_inverse(layout);
}
//
// Negative strides (beta support)
// Post-conditions/layout indexing aren't generalized enough to support these yet
// However, the composition post-condition is general enough.
{
auto layout = make_layout(Shape<_4>{}, Stride<Int<-1>>{});
test_left_inverse(layout);
}
//{
//auto layout = Layout<Shape < _2,_4>,
// Stride<_m1,_2>>{};
//test_left_inverse(layout);
//}
//{
//auto layout = Layout<Shape < _2, _4>,
// Stride< _4,_m1>>{};
//test_left_inverse(layout);
//}
//{
//auto layout = Layout<Shape < _2, _4, _6>,
// Stride<_m1,_12,_m2>>{};
//test_left_inverse(layout);
//}
}
test/unit/cute/core/inverse_right.cpp
+255
View File
@@ -0,0 +1,255 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <iostream>
#include <cute/tensor.hpp>
using namespace cute;
template <class Layout>
void
test_right_inverse(Layout const& layout)
{
auto inv_layout = right_inverse(layout);
CUTLASS_TRACE_HOST(layout << " ^ -1\n" << " => \n" << inv_layout);
CUTLASS_TRACE_HOST("Composition: " << coalesce(composition(layout, inv_layout)) << std::endl);
for (int i = 0; i < size(inv_layout); ++i) {
//printf("%3d: %3d %3d\n", i, int(inv_layout(i)), int(layout(inv_layout(i))));
EXPECT_EQ(layout(inv_layout(i)), i);
}
}
TEST(CuTe_core, Inverse_right)
{
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("RIGHT INVERSE" );
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Simple tests" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto layout = Layout<_1, _0>{};
test_right_inverse(layout);
}
{
auto layout = Layout<_1, _1>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_4>,
Stride<_0>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <Shape <_1,_1>>,
Stride<Stride<_0,_0>>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <Shape <_3,_7>>,
Stride<Stride<_0,_0>>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_1>,
Stride<_1>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_4>,
Stride<_1>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_4>,
Stride<_2>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_2,_4>,
Stride<_0,_2>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_8, _4>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_8, _4>,
Stride<_4, _1>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape< _2,_4,_6>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_2,_4,_6>,
Stride<_4,_1,_8>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_2,_4,_4,_6>,
Stride<_4,_1,_0,_8>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_4, _2>,
Stride<_1,_16>>{};
test_right_inverse(layout);
}
{
auto layout = Layout<Shape <_4, _2>,
Stride<_1, _5>>{};
test_right_inverse(layout);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Dynamic shapes/strides" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto layout = make_layout(Shape<_4, _2>{}, make_stride(Int<1>{}, 4));
test_right_inverse(layout);
}
{
auto layout = make_layout(make_shape(_4{}, 2), make_stride(Int<1>{}, 4));
test_right_inverse(layout);
}
{
auto layout = make_layout(make_shape(4, 2), make_stride(Int<1>{}, 4));
test_right_inverse(layout);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("Swizzle layouts" );
CUTLASS_TRACE_HOST("-------------------------------");
{
auto layout = ComposedLayout<Swizzle<1,0,2>, _0, Layout<Shape <_4, _4>,
Stride<_1, _4>>>{};
test_right_inverse(layout);
}
{
auto layout = ComposedLayout<Swizzle<1,0,2>, _0, Layout<Shape <_4, _4>,
Stride<_4, _1>>>{};
test_right_inverse(layout);
}
{
auto layout = ComposedLayout<Swizzle<1,0,1>, _0, Layout<Shape <_4, _4>,
Stride<_8, _1>>>{};
test_right_inverse(layout);
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("BETA: Negative strides" );
CUTLASS_TRACE_HOST("-------------------------------");
// Negative strides (beta support)
// Post-conditions/layout indexing aren't generalized enough to support these yet
// However, the composition post-condition is general enough.
{
auto layout = make_layout(Shape<_4>{}, Stride<Int<-1>>{});
test_right_inverse(layout);
}
//{
//auto layout = Layout<Shape < _2,_4>,
// Stride<_m1,_2>>{};
//test_right_inverse(layout);
//}
//{
//auto layout = Layout<Shape < _2, _4>,
// Stride< _4,_m1>>{};
//test_right_inverse(layout);
//}
//{
//auto layout = Layout<Shape < _2, _4, _6>,
// Stride<_m1,_12,_m2>>{};
//test_right_inverse(layout);
//}
}
test/unit/cute/core/logical_divide.cpp
+253
View File
@@ -0,0 +1,253 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
using namespace cute;
template <class LayoutA, class LayoutB>
void
test_logical_divide(LayoutA const& layoutA,
LayoutB const& layoutB)
{
auto layoutR = logical_divide(layoutA, layoutB);
CUTLASS_TRACE_HOST("test_logical_divide()");
CUTLASS_TRACE_HOST(shape(layoutA) << " / " << shape(layoutB) << " => " << shape(layoutR) );
CUTLASS_TRACE_HOST(stride(layoutA) << " " << stride(layoutB) << " => " << stride(layoutR));
// Test that layout R is compatible with layout B
ASSERT_EQ(rank(layoutR), 2);
ASSERT_TRUE(compatible(layoutB, layout<0>(layoutR)));
}
TEST(CuTe_core, Logical_divide)
{
{
auto layout = Layout<_1,_0>{};
auto tile = Layout<_1,_0>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_1,_0>{};
auto tile = Layout<_1,_1>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_1,_1>{};
auto tile = Layout<_1,_0>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_1,_1>{};
auto tile = Layout<_1,_1>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_1>{};
auto tile = Layout<_2,_1>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_1>{};
auto tile = Layout<_2,_3>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_1>{};
auto tile = Layout<Shape<_2,_3>,Stride<_3,_1>>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_2>{};
auto tile = Layout<_2,_1>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_2>{};
auto tile = Layout<_2,_3>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_6,_2>{};
auto tile = Layout<Shape<_2,_3>,Stride<_3,_1>>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_6,_6>,Stride<_1,_12>>{};
auto tile = Layout<Shape<_6,_3>,Stride<_3,_1>>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_6,_6>,Stride<_12,_1>>{};
auto tile = Layout<Shape<_6,_3>,Stride<_3,_1>>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<_32>{};
auto tile = Layout<_2,_8>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_4,_1>,Stride<_1,_1>>{};
auto tile = Layout<_2,_1>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_4,_1>,Stride<_1,_1>>{};
auto tile = Layout<_2,_2>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_8,_8>,Stride<_1,_8>>{};
auto tile = Layout<_32,_2>{};
test_logical_divide(layout, tile);
}
{
auto layout = Layout<Shape<_8,_8>,Stride<_8,_1>>{};
auto tile = Layout<_32,_2>{};
test_logical_divide(layout, tile);
}
//
// Dynamic
//
{
auto layout = make_layout(2);
auto tile = Layout<_32>{};
test_logical_divide(layout, tile);
// Enforcement for dynamic cases
auto result = logical_divide(layout, tile);
static_assert(decltype(shape<0>(result) == Int<32>{})::value);
static_assert(decltype(stride<0>(result) == Int<1>{})::value);
assert(shape<1>(result) == 1);
static_assert(decltype(stride<1>(result) == Int<32>{})::value);
}
{
auto layout = make_layout(48);
auto tile = Layout<_32>{};
test_logical_divide(layout, tile);
// Enforcement for dynamic cases
auto result = logical_divide(layout, tile);
static_assert(decltype(shape<0>(result) == Int<32>{})::value);
static_assert(decltype(stride<0>(result) == Int<1>{})::value);
assert(shape<1>(result) == 2);
static_assert(decltype(stride<1>(result) == Int<32>{})::value);
}
{
auto layout = make_layout(96);
auto tile = Layout<_32,_2>{};
test_logical_divide(layout, tile);
}
{
auto layout = make_layout(32);
auto tile = Layout<Int<48>>{};
test_logical_divide(layout, tile);
// Enforcement for dynamic cases
auto result = logical_divide(layout, tile);
static_assert(decltype(shape<0>(result) == Int<48>{})::value);
static_assert(decltype(stride<0>(result) == Int<1>{})::value);
assert(shape<1>(result) == 1);
static_assert(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>{};
//test_logical_divide(layout, tile);
//}
//{
//auto layout = make_layout(make_shape(128,4,3), make_stride(1,512,0));
//auto tile = Layout<_32,_2>{};
//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>{};
//CUTLASS_TRACE_HOST("complement: " << complement(tile, size(layout)));
//test_logical_divide(layout, tile);
//}
}
test/unit/cute/core/logical_product.cpp
+218
View File
@@ -0,0 +1,218 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
using namespace cute;
template <class LayoutA, class LayoutB>
void
test_logical_product(LayoutA const& layoutA,
LayoutB const& layoutB)
{
auto layoutR = logical_product(layoutA, layoutB);
CUTLASS_TRACE_HOST(shape(layoutA) << " x " << shape(layoutB) << " => " << shape(layoutR) );
CUTLASS_TRACE_HOST(stride(layoutA) << " " << stride(layoutB) << " => " << stride(layoutR));
// Test that layout R is compatible with layout B
ASSERT_EQ(rank(layoutR), 2);
//assert(compatible(layoutB, layout<0>(layoutR)));
//assert(consistent(layoutA, layout<1>(layoutR)));
// True post-condition:
}
TEST(CuTe_core, Logical_product)
{
{
auto vec = Layout<_1,_0>{};
auto tile = Layout<_1,_0>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_1,_1>{};
auto tile = Layout<_1,_0>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_1,_0>{};
auto tile = Layout<_1,_1>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_1,_1>{};
auto tile = Layout<_1,_1>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_3,_1>{};
auto tile = Layout<_4,_0>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_3,_0>{};
auto tile = Layout<_4,_1>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_3,_0>{};
auto tile = Layout<_4,_0>{};
test_logical_product(vec, tile);
}
{
auto vec = Layout<_3,_2>{};
auto tile = Layout<_4,_1>{};
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{});
auto tile = make_layout(Shape<_2,_4>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_2,_4>{});
auto tile = make_layout(Shape<_3>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_8,Shape<_2,_2>>{});
auto tile = make_layout(Shape<_4>{}, Stride<_2>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_2,_2>{});
auto tile = make_layout(Shape<_3,_3>{}, Stride<_3,_1>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{}, Stride<_32>{});
auto tile = make_layout(Shape<_32>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{}, Stride<_2>{});
auto tile = make_layout(Shape<_4>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{}, Stride<_32>{});
auto tile = make_layout(Shape<_128>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{}, Stride<_32>{});
auto tile = make_layout(Shape<_8,_8>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<_3>{}, Stride<_32>{});
auto tile = make_layout(Shape<_8,_8>{}, Stride<_8,_1>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<Shape<_4,_2>>{}, Stride<Stride<_1,_16>>{});
auto tile = make_layout(Shape<_4,_4>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<Shape<_4,_2>>{}, Stride<Stride<_1,_16>>{});
auto tile = make_layout(Shape<_4,_2>{}, Stride<_2,_1>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
Stride<Stride<_1,_4>,Stride<_8,_32>>{});
auto tile = make_layout(Shape<_2,_2>{}, Stride<_1,_2>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape<Shape<_2,_2>,Shape<_2, _2>>{},
Stride<Stride<_1,_4>,Stride<_8,_32>>{});
auto tile = make_layout(Shape<_2,_2>{},
Stride<_2,_1>{});
test_logical_product(vec, tile);
}
{
auto vec = make_layout(Shape <Shape <_4,_6>>{},
Stride<Stride<_1,_6>>{});
auto tile = make_layout(Shape <_3>{},
Stride<_1>{});
test_logical_product(vec, tile);
}
}
test/unit/cute/core/mixedbits.cpp
+70
View File
@@ -0,0 +1,70 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/swizzle.hpp>
TEST(CuTe_core, MixedBits) {
using namespace cute;
auto uzero = cute::integral_constant<uint32_t, 0>{};
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto S0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto F0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto S1) {
for_each(make_integer_sequence<uint32_t, 8>{}, [&](auto F1) {
if constexpr (decltype(S0 == uzero || S1 == uzero)::value) {
return;
} else if constexpr (decltype((S0 & F0) != uzero || (S1 & F1) != uzero)::value) {
return;
} else {
for (uint32_t d0 = 0; d0 < 8; ++d0) {
if ((d0 & F0) != d0) { continue; } // Skip repeats
for (uint32_t d1 = 0; d1 < 8; ++d1) {
if ((d1 & F1) != d1) { continue; } // Skip repeats
auto m0 = make_mixed_bits(S0, d0, F0);
auto m1 = make_mixed_bits(S1, d1, F1);
//print(m0); print(" & "); print(m1); print(" = "); print(m0 & m1); print("\n");
EXPECT_EQ(to_integral(m0 & m1), to_integral(m0) & to_integral(m1));
//print(m0); print(" | "); print(m1); print(" = "); print(m0 | m1); print("\n");
EXPECT_EQ(to_integral(m0 | m1), to_integral(m0) | to_integral(m1));
//print(m0); print(" ^ "); print(m1); print(" = "); print(m0 ^ m1); print("\n");
EXPECT_EQ(to_integral(m0 ^ m1), to_integral(m0) ^ to_integral(m1));
}
}
}
});
});
});
});
}
test/unit/cute/core/transform.cpp
+49
View File
@@ -0,0 +1,49 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
#include <cute/numeric/complex.hpp>
TEST(CuTe_core, Transform) {
using namespace cute;
complex<float> array[4] = {{0,0}, {1,0}, {0,1}, {1,1}};
complex<float> correct[4] = {{0,0}, {1,0}, {0,-1}, {1,-1}};
auto tensor = make_tensor(static_cast<complex<float>*>(array), make_layout(make_shape(4)));
conjugate conj;
transform(tensor, conj);
for (int i = 0; i < 4; ++i)
{
EXPECT_EQ(tensor(i), correct[i]);
}
}
test/unit/cute/core/tuple.cpp
+266
View File
@@ -0,0 +1,266 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <cutlass/trace.h>
#include <cute/tensor.hpp>
TEST(CuTe_core, Tuple)
{
using namespace cute;
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("SIMPLE STATIC AND DYNAMIC TUPLES");
CUTLASS_TRACE_HOST("-------------------------------");
using tuple_2d_s_type = tuple<_8, _4>; // (8,4)
using tuple_3d_s_type = tuple<_8, _4, _2>; // (8,4,2)
using tuple_3h_s_type = tuple<tuple<_1, _2>, _8, _2>; // ((1,2),8,2)
using tuple_2d_d_type = tuple<int, int>; // (8,4)
using tuple_3d_d_type = tuple<int, int, int>; // (8,4,2)
using tuple_3h_d_type = tuple<tuple<int, int>, int, int>; // ((1,2),8,2)
using tuple_2d_m_type = tuple<_8, int>; // (8,4)
using tuple_3d_m_type = tuple<int, int, _2>; // (8,4,2)
using tuple_3h_m_type = tuple<tuple<int, _2>, int, int>; // ((1,2),8,2)
tuple_2d_s_type tuple_2d_s;
tuple_3d_s_type tuple_3d_s;
tuple_3h_s_type tuple_3h_s;
tuple_2d_d_type tuple_2d_d(8,4);
tuple_3d_d_type tuple_3d_d(8,4,2);
tuple_3h_d_type tuple_3h_d(tuple<int,int>(1,2),8,2);
tuple_2d_m_type tuple_2d_m(_8{}, 4);
tuple_3d_m_type tuple_3d_m(8,4,_2{});
tuple_3h_m_type tuple_3h_m(tuple<int,_2>(1,_2{}),8,2);
CUTLASS_TRACE_HOST(tuple_2d_s << (is_static<tuple_2d_s_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_2d_s_type));
ASSERT_TRUE(is_static<tuple_2d_s_type>::value == true);
ASSERT_TRUE(sizeof(tuple_2d_s_type) == 1);
ASSERT_TRUE(std::is_empty<tuple_2d_s_type>::value);
CUTLASS_TRACE_HOST(tuple_3d_s << (is_static<tuple_3d_s_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3d_s_type));
ASSERT_TRUE(is_static<tuple_3d_s_type>::value == true);
ASSERT_TRUE(sizeof(tuple_3d_s_type) == 1);
ASSERT_TRUE(std::is_empty<tuple_3d_s_type>::value);
CUTLASS_TRACE_HOST(tuple_3h_s << (is_static<tuple_3h_s_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3h_s_type));
ASSERT_TRUE(is_static<tuple_3h_s_type>::value == true);
ASSERT_TRUE(sizeof(tuple_3h_s_type) == 1);
ASSERT_TRUE(std::is_empty<tuple_3h_s_type>::value);
CUTLASS_TRACE_HOST(tuple_2d_d << (is_static<tuple_2d_d_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_2d_d_type));
ASSERT_TRUE(is_static<tuple_2d_d_type>::value == false);
ASSERT_TRUE(sizeof(tuple_2d_d_type) == 8);
ASSERT_TRUE(!std::is_empty<tuple_2d_d_type>::value);
CUTLASS_TRACE_HOST(tuple_3d_d << (is_static<tuple_3d_d_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3d_d_type));
ASSERT_TRUE(is_static<tuple_3d_d_type>::value == false);
ASSERT_TRUE(sizeof(tuple_3d_d_type) == 12);
ASSERT_TRUE(!std::is_empty<tuple_3d_d_type>::value);
CUTLASS_TRACE_HOST(tuple_3h_d << (is_static<tuple_3h_d_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3h_d_type));
ASSERT_TRUE(is_static<tuple_3h_d_type>::value == false);
ASSERT_TRUE(sizeof(tuple_3h_d_type) == 16);
ASSERT_TRUE(!std::is_empty<tuple_3h_d_type>::value);
CUTLASS_TRACE_HOST(tuple_2d_m << (is_static<tuple_2d_m_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_2d_m_type));
ASSERT_TRUE(is_static<tuple_2d_m_type>::value == false);
ASSERT_TRUE(sizeof(tuple_2d_m_type) == 4);
ASSERT_TRUE(!std::is_empty<tuple_2d_m_type>::value);
CUTLASS_TRACE_HOST(tuple_3d_m << (is_static<tuple_3d_m_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3d_m_type));
ASSERT_TRUE(is_static<tuple_3d_m_type>::value == false);
ASSERT_TRUE(sizeof(tuple_3d_m_type) == 8);
ASSERT_TRUE(!std::is_empty<tuple_3d_m_type>::value);
CUTLASS_TRACE_HOST(tuple_3h_m << (is_static<tuple_3h_m_type>::value ? " Static " : " Dynamic ")
<< "sizeof = " << sizeof(tuple_3h_m_type));
ASSERT_TRUE(is_static<tuple_3h_m_type>::value == false);
ASSERT_TRUE(sizeof(tuple_3h_m_type) == 12);
ASSERT_TRUE(!std::is_empty<tuple_3h_m_type>::value);
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("SIMPLE TUPLE OPS");
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("product(" << tuple_2d_s << ") => " << product(tuple_2d_s));
CUTE_STATIC_ASSERT_V(product(tuple_2d_s) == _32{});
CUTLASS_TRACE_HOST("product(" << tuple_3d_s << ") => " << product(tuple_3d_s));
CUTE_STATIC_ASSERT_V(product(tuple_3d_s) == _64{});
CUTLASS_TRACE_HOST("product(" << tuple_3h_s << ") => " << product(tuple_3h_s));
CUTE_STATIC_ASSERT_V(product(tuple_3h_s) == _32{});
CUTLASS_TRACE_HOST("product(" << tuple_2d_d << ") => " << product(tuple_2d_d));
ASSERT_TRUE(product(tuple_2d_d) == 32);
CUTLASS_TRACE_HOST("product(" << tuple_3d_d << ") => " << product(tuple_3d_d));
ASSERT_TRUE(product(tuple_3d_d) == 64);
CUTLASS_TRACE_HOST("product(" << tuple_3h_d << ") => " << product(tuple_3h_d));
ASSERT_TRUE(product(tuple_3h_d) == 32);
CUTLASS_TRACE_HOST("product(" << tuple_2d_m << ") => " << product(tuple_2d_m));
ASSERT_TRUE(product(tuple_2d_m) == 32);
CUTLASS_TRACE_HOST("product(" << tuple_3d_m << ") => " << product(tuple_3d_m));
ASSERT_TRUE(product(tuple_3d_m) == 64);
CUTLASS_TRACE_HOST("product(" << tuple_3h_m << ") => " << product(tuple_3h_m));
ASSERT_TRUE(product(tuple_3h_m) == 32);
CUTLASS_TRACE_HOST("max(" << tuple_2d_s << ") => " << max(tuple_2d_s));
CUTE_STATIC_ASSERT_V(max(tuple_2d_s) == _8{});
CUTLASS_TRACE_HOST("max(" << tuple_3d_s << ") => " << max(tuple_3d_s));
CUTE_STATIC_ASSERT_V(max(tuple_3d_s) == _8{});
CUTLASS_TRACE_HOST("max(" << tuple_3h_s << ") => " << max(tuple_3h_s));
CUTE_STATIC_ASSERT_V(max(tuple_3h_s) == _8{});
CUTLASS_TRACE_HOST("max(" << tuple_2d_d << ") => " << max(tuple_2d_d));
ASSERT_TRUE(max(tuple_2d_d) == 8);
CUTLASS_TRACE_HOST("max(" << tuple_3d_d << ") => " << max(tuple_3d_d));
ASSERT_TRUE(max(tuple_3d_d) == 8);
CUTLASS_TRACE_HOST("max(" << tuple_3h_d << ") => " << max(tuple_3h_d));
ASSERT_TRUE(max(tuple_3h_d) == 8);
CUTLASS_TRACE_HOST("max(" << tuple_2d_m << ") => " << max(tuple_2d_m));
ASSERT_TRUE(max(tuple_2d_m) == 8);
CUTLASS_TRACE_HOST("max(" << tuple_3d_m << ") => " << max(tuple_3d_m));
ASSERT_TRUE(max(tuple_3d_m) == 8);
CUTLASS_TRACE_HOST("max(" << tuple_3h_m << ") => " << max(tuple_3h_m));
ASSERT_TRUE(max(tuple_3h_m) == 8);
// 2d s|d|m
CUTLASS_TRACE_HOST("inner_product(" << tuple_2d_s << ", " << tuple_2d_s << ") => "
<< inner_product(tuple_2d_s, tuple_2d_s));
CUTE_STATIC_ASSERT_V(inner_product(tuple_2d_s, tuple_2d_s) == Int<80>{});
CUTLASS_TRACE_HOST("inner_product(" << tuple_2d_d << ", " << tuple_2d_d << ") => "
<< inner_product(tuple_2d_d, tuple_2d_d));
ASSERT_TRUE(inner_product(tuple_2d_d, tuple_2d_d) == 80);
CUTLASS_TRACE_HOST("inner_product(" << tuple_2d_m << ", " << tuple_2d_m << ") => "
<< inner_product(tuple_2d_m, tuple_2d_m));
ASSERT_TRUE(inner_product(tuple_2d_m, tuple_2d_m) == 80);
// 3d s|d|m
CUTLASS_TRACE_HOST("inner_product(" << tuple_3d_s << ", " << tuple_3d_s << ") => "
<< inner_product(tuple_3d_s, tuple_3d_s));
CUTE_STATIC_ASSERT_V(inner_product(tuple_3d_s, tuple_3d_s) == Int<84>{});
CUTLASS_TRACE_HOST("inner_product(" << tuple_3d_d << ", " << tuple_3d_d << ") => "
<< inner_product(tuple_3d_d, tuple_3d_d));
ASSERT_TRUE(inner_product(tuple_3d_d, tuple_3d_d) == 84);
CUTLASS_TRACE_HOST("inner_product(" << tuple_3d_m << ", " << tuple_3d_m << ") => "
<< inner_product(tuple_3d_m, tuple_3d_m));
ASSERT_TRUE(inner_product(tuple_3d_m, tuple_3d_m) == 84);
// 3h s|d|m
CUTLASS_TRACE_HOST("inner_product(" << tuple_3h_s << ", " << tuple_3h_s << ") => "
<< inner_product(tuple_3h_s, tuple_3h_s));
CUTE_STATIC_ASSERT_V(inner_product(tuple_3h_s, tuple_3h_s) == Int<73>{});
CUTLASS_TRACE_HOST("inner_product(" << tuple_3h_d << ", " << tuple_3h_d << ") => "
<< inner_product(tuple_3h_d, tuple_3h_d));
ASSERT_TRUE(inner_product(tuple_3h_d, tuple_3h_d) == 73);
CUTLASS_TRACE_HOST("inner_product(" << tuple_3h_m << ", " << tuple_3h_m << ") => "
<< inner_product(tuple_3h_m, tuple_3h_m));
ASSERT_TRUE(inner_product(tuple_3h_m, tuple_3h_m) == 73);
CUTLASS_TRACE_HOST("col_major(" << tuple_2d_s << ") => " << compact_col_major(tuple_2d_s));
CUTE_STATIC_ASSERT_V((compact_col_major(tuple_2d_s) == make_tuple(_1{},_8{})));
CUTLASS_TRACE_HOST("col_major(" << tuple_3d_s << ") => " << compact_col_major(tuple_3d_s));
CUTE_STATIC_ASSERT_V((compact_col_major(tuple_3d_s) == make_tuple(_1{},_8{},_32{})));
CUTLASS_TRACE_HOST("col_major(" << tuple_3h_s << ") => " << compact_col_major(tuple_3h_s));
CUTE_STATIC_ASSERT_V((compact_col_major(tuple_3h_s) == make_tuple(make_tuple(_0{},_1{}),_2{},_16{})));
CUTLASS_TRACE_HOST("col_major(" << tuple_2d_d << ") => " << compact_col_major(tuple_2d_d));
ASSERT_TRUE((compact_col_major(tuple_2d_d) == make_tuple(_1{},8)));
CUTLASS_TRACE_HOST("col_major(" << tuple_3d_d << ") => " << compact_col_major(tuple_3d_d));
ASSERT_TRUE((compact_col_major(tuple_3d_d) == make_tuple(_1{},8,32)));
CUTLASS_TRACE_HOST("col_major(" << tuple_3h_d << ") => " << compact_col_major(tuple_3h_d));
ASSERT_TRUE((compact_col_major(tuple_3h_d) == make_tuple(make_tuple(_1{},1),2,16)));
CUTLASS_TRACE_HOST("col_major(" << tuple_2d_m << ") => " << compact_col_major(tuple_2d_m));
ASSERT_TRUE((compact_col_major(tuple_2d_m) == make_tuple(_1{},_8{})));
CUTLASS_TRACE_HOST("col_major(" << tuple_3d_m << ") => " << compact_col_major(tuple_3d_m));
ASSERT_TRUE((compact_col_major(tuple_3d_m) == make_tuple(_1{},8,32)));
CUTLASS_TRACE_HOST("col_major(" << tuple_3h_m << ") => " << compact_col_major(tuple_3h_m));
ASSERT_TRUE((compact_col_major(tuple_3h_m) == make_tuple(make_tuple(_1{},1),2,16)));
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("SLICING TUPLES");
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = Coord<_2,_3,_4,Coord<_5,_6>>{};
CUTLASS_TRACE_HOST("a = " << a);
CUTLASS_TRACE_HOST("a(1) = " << slice(1, a));
CUTLASS_TRACE_HOST("a(_) = " << slice(_, a));
CUTLASS_TRACE_HOST("a(_,1,_,_) = " << slice(make_coord(_,1,_,_), a));
CUTLASS_TRACE_HOST("a(_,1,_,(_,_)) = " << slice(make_coord(_,1,_,make_coord(_,_)), a));
CUTLASS_TRACE_HOST("a(_,1,_,(_,2)) = " << slice(make_coord(_,1,_,make_coord(_,2)), a));
CUTLASS_TRACE_HOST("a(_,1,_,(1,2)) = " << slice(make_coord(_,1,_,make_coord(1,2)), a));
}
CUTLASS_TRACE_HOST("-------------------------------");
CUTLASS_TRACE_HOST("DICING TUPLES");
CUTLASS_TRACE_HOST("-------------------------------");
{
auto a = Coord<_2,_3,_4,Coord<_5,_6>>{};
CUTLASS_TRACE_HOST("a = " << a);
CUTLASS_TRACE_HOST("a(1) = " << dice(1, a));
CUTLASS_TRACE_HOST("a(_) = " << dice(_, a));
CUTLASS_TRACE_HOST("a(_,1,_,_) = " << dice(make_coord(_,1,_,_), a));
CUTLASS_TRACE_HOST("a(_,1,_,(_,_)) = " << dice(make_coord(_,1,_,make_coord(_,_)), a));
CUTLASS_TRACE_HOST("a(_,1,_,(_,2)) = " << dice(make_coord(_,1,_,make_coord(_,2)), a));
CUTLASS_TRACE_HOST("a(_,1,_,(1,2)) = " << dice(make_coord(_,1,_,make_coord(1,2)), a));
}
}
test/unit/cute/hopper/CMakeLists.txt
+58
View File
@@ -0,0 +1,58 @@
# Copyright (c) 2023 - 2023 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.
add_custom_target(
cutlass_test_unit_cute_hopper
DEPENDS
cutlass_test_unit_cute_hopper_stsm
cutlass_test_unit_cute_hopper_tma_load
cutlass_test_unit_cute_hopper_tma_store
)
add_custom_target(
test_unit_cute_hopper
DEPENDS
test_unit_cute_hopper_stsm
test_unit_cute_hopper_tma_load
test_unit_cute_hopper_tma_store
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_stsm
stsm.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_load
tma_load.cu
)
cutlass_test_unit_add_executable(
cutlass_test_unit_cute_hopper_tma_store
tma_store.cu
)
test/unit/cute/hopper/stsm.cu
+426
View File
@@ -0,0 +1,426 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/arch/copy_sm90.hpp>
using namespace cute;
template<class T>
__global__ void
stsm_test_device(uint16_t* g_in, uint16_t* g_out)
{
constexpr int count = sizeof(T) / 4;
int tid = threadIdx.x;
int stride = blockDim.x;
// load input gmem -> rmem
uint32_t reg[count];
for (int i = 0; i < (sizeof(T) / 4); i++) {
reg[i] = reinterpret_cast<uint32_t*>(g_in)[tid + (stride * i)];
}
__shared__ uint32_t smem[32 * count];
// load rmem -> smem using STSM
uint128_t* smem_ptr = reinterpret_cast<uint128_t*>(smem) + tid;
T* rmem_ptr = reinterpret_cast<T*>(reg);
cute::copy_stsm(rmem_ptr, smem_ptr);
__syncthreads();
// store output smem -> gmem
for (int i = 0; i < (sizeof(T) / 4); i++) {
reinterpret_cast<uint32_t*>(g_out)[tid + (stride * i)] = smem[tid + (stride * i)];
}
}
template <class TiledCopy, class SmemLayout>
__global__ void
stsm_test_device_cute(uint16_t* g_in, uint16_t* g_out,
TiledCopy tiled_copy, SmemLayout smem_layout)
{
using namespace cute;
__shared__ uint16_t smem[size(smem_layout)];
Tensor t_g_in = make_tensor(make_gmem_ptr(g_in), smem_layout);
Tensor t_g_out = make_tensor(make_gmem_ptr(g_out), smem_layout);
Tensor t_smem = make_tensor(make_smem_ptr(smem), smem_layout);
int tid = threadIdx.x;
auto thr_copy = tiled_copy.get_thread_slice(tid);
Tensor tXgX = thr_copy.partition_S(t_g_in); // (V,M,N)
Tensor tXsX = thr_copy.partition_D(t_smem); // (V,M,N)
Tensor tXrX = make_tensor<uint16_t>(shape(tXgX)); // (V,M,N)
clear(tXrX); // Just to make sure
/*
if (thread0()) {
print("tXsX: " ); print(tXsX.layout()); print("\n");
print("tXgX: " ); print(tXgX.layout()); print("\n");
print("tXrX: " ); print(tXrX.layout()); print("\n");
}
*/
// Load input gmem -> rmem
copy(tXgX, tXrX);
// Copy rmem -> smem via tiled_copy (STSM, STS)
copy(tiled_copy, tXrX, tXsX);
// Output smem -> gmem
for (int i = tid; i < size(t_smem); i += size(tiled_copy)) {
t_g_out(i) = t_smem(i);
}
}
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
TEST(SM90_CuTe_Hopper, Stsm)
{
constexpr int count = 1024;
thrust::host_vector<uint16_t> h_in(count);
for (int i = 0; i < count; ++i) {
h_in[i] = uint16_t(i);
}
thrust::device_vector<uint16_t> d_in = h_in;
//
// STSM 1x (32b)
//
{
thrust::device_vector<uint16_t> d_out(count);
stsm_test_device<uint32_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 32; ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("STSM 1x stsm_test_device SUCCESS\n");
}
//
// STSM 2x (64b)
//
{
thrust::device_vector<uint16_t> d_out(count);
stsm_test_device<uint64_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 64; ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("STSM 2x stsm_test_device SUCCESS\n");
}
//
// STSM 4x (128b)
//
{
thrust::device_vector<uint16_t> d_out(count);
stsm_test_device<uint128_t><<<1, 32>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()));
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < 128; ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("STSM 4x stsm_test_device SUCCESS\n");
}
//
// CuTe STSM
//
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x1_STSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x1_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x2_STSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x2_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x4_STSM_N, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved U32x4_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,Shape <_2, _4>>,
Stride< _2,Stride<_1,_64>>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<UniversalCopy<uint16_t>, uint16_t>{},
Layout<Shape<_32,_1>>{},
Layout<Shape< _1,_8>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x8 interleaved STS.U16 SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x1_STSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x1_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x2_STSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x2_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U32x4_STSM_N, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U32x4_STSM_N SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride< _1,_32>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<UniversalCopy<uint16_t>, uint16_t>{},
Layout<Shape<_16,_2>>{},
Layout<Shape< _2,_4>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 STS.U16 SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U16x2_STSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_2,_1>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x2_STSM_T SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U16x4_STSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_4,_1>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x4_STSM_T SUCCESS\n");
}
{
thrust::device_vector<uint16_t> d_out(count);
auto smem_layout = Layout<Shape <_32,_32>,
Stride<_32, _1>>{};
auto tiled_copy = make_tiled_copy(Copy_Atom<SM90_U16x8_STSM_T, uint16_t>{},
Layout<Shape<_4,_8>>{},
Layout<Shape<_8,_1>>{});
stsm_test_device_cute<<<1, int(size(tiled_copy))>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tiled_copy,
smem_layout);
thrust::host_vector<uint16_t> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe 32x32 U16x8_STSM_T SUCCESS\n");
}
CUTLASS_TRACE_HOST("PASS");
}
#endif
test/unit/cute/hopper/tma_load.cu
+495
View File
@@ -0,0 +1,495 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
cute::uint64_t tma_load_mbar[1];
};
// __grid_constant__ was introduced in CUDA 11.7.
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
# define CUTE_GRID_CONSTANT_SUPPORTED
#endif
// __grid_constant__ can be enabled only on SM70+
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
# define CUTE_GRID_CONSTANT_ENABLED
#endif
#if ! defined(CUTE_GRID_CONSTANT)
# if defined(CUTE_GRID_CONSTANT_SUPPORTED) && defined(CUTE_GRID_CONSTANT_ENABLED)
# define CUTE_GRID_CONSTANT __grid_constant__
# else
# define CUTE_GRID_CONSTANT
# endif
#endif
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
assert(product_each(shape(gmem_layout)) == product_each(smem_layout.shape()));
// 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);
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout);
#if 0
//
// Read in trivially
//
Tensor gA_in = make_tensor(make_gmem_ptr(g_in), gmem_layout);
// Input gmem -> smem
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
sA(i) = gA_in(i);
}
__syncthreads();
#else
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor gA = tma.get_tma_tensor(shape(gmem_layout));
//
// Prepare the TMA_LOAD
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAgA = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N)
Tensor tAsA = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
#if 0
if (thread0()) {
print(" gA: "); print(gA.data()); print(" o "); print(gA.layout()); print("\n");
print("tAgA: "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print(" sA: "); print(sA.data()); print(" o "); print(sA.layout()); print("\n");
print("tAsA: "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
}
#endif
//
// Perform the TMA_LOAD
//
// Group the TMA_M and TMA_N modes
Tensor tAgA_2 = group_modes<1,rank(tAgA)>(tAgA); // (TMA,Rest)
Tensor tAsA_TR = group_modes<1,rank(tAsA)>(tAsA); // (TMA,Rest)
static_assert(size<1>(tAsA_TR) == 1);
Tensor tAsA_2 = tAsA_TR(_,0);
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA_2); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = size(sA) * sizeof(T);
if (threadIdx.x == 0)
{
/// Initialize shared memory barrier
tma_load_mbar[0] = 0;
cute::initialize_barrier(tma_load_mbar[0], 1 /*numThreads*/);
cute::set_barrier_transaction_bytes(tma_load_mbar[0], kTmaTransactionBytes);
copy(tma.with(tma_load_mbar[0]), tAgA_2(_,stage), tAsA_2);
}
__syncthreads();
/// Wait on the shared memory barrier until the phase bit flips from kPhaseBit value
constexpr int kPhaseBit = 0;
cute::wait_barrier(tma_load_mbar[0], kPhaseBit);
#endif
//
// Write out trivially
//
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
// Do the same slicing and grouping as sA
Tensor tAgA_out = cta_tma.partition_D(gA_out); // (TMA,TMA_M,TMA_N)
Tensor tAgA_2_out = group_modes<1,rank(tAgA_out)>(tAgA_out); // (TMA,Rest)
// Output smem -> gmem
for (int i = threadIdx.x; i < size(tAsA_2); i += blockDim.x) {
tAgA_2_out(i,stage) = tAsA_2(i);
}
__syncthreads();
}
}
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Col)
{
using T = half_t;
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
Layout gmem_layout = smem_layout;
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD 32x32 ColMajor SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Row)
{
using T = half_t;
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
Layout gmem_layout = smem_layout;
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD 32x32 RowMajor SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN)
{
using T = half_t;
auto smem_layout = GMMA::Layout_MN_SW128_Atom<T>{};
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_K)
{
using T = half_t;
auto smem_layout = GMMA::Layout_K_SW128_Atom<T>{};
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenRowMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_K_SW128_Atom<T> SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi)
{
using T = half_t;
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{});
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi2)
{
using T = half_t;
// Tile the GMMA::Layout atom in the K-mode first, then the M-mode to get a bigger box size
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_GMMA_SW128_MN_Multi_Dyn)
{
using T = half_t;
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
Layout gmem_layout = make_layout(make_shape(128, 128), GenColMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_32x32_Multimode)
{
using T = half_t;
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
//auto smem_layout = Layout<Shape<_32,_32>>{};
//Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_Tensor_blocking)
{
using T = half_t;
auto gmem_layout = make_shape(make_shape(336,40),make_shape(32,656)); // GMEM
auto cta_tile = make_shape(make_shape(_16{},_8{}),make_shape(_32{},_2{})); // GMEM Tiling:
// Take 16-elem from m0, 8-elem from m1,
// Take 32-elem from k0, 2-elem from k1
auto smem_layout = make_layout(cta_tile); // Col-Major SMEM
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout, cta_tile, Int<1>{});
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD Tensor blocking SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_load_Tensor_blocking_2)
{
using T = half_t;
auto gmem_layout = make_shape(make_shape(32,40),make_shape(make_shape(8,8),656)); // GMEM
auto cta_tile = make_shape(_128{},make_shape(_32{},_2{})); // GMEM Tiling:
// Take 128-elem from m: m0 must divide 128,
// m-last may be predicated
// Take 32-elem from k0, 2-elem from k1
auto smem_layout = make_layout(cta_tile); // Col-Major SMEM
thrust::host_vector<T> h_in(size(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_in.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_LOAD{}, gA, smem_layout, cta_tile, Int<1>{});
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_LOAD Tensor blocking 2 SUCCESS\n");
}
#endif
test/unit/cute/hopper/tma_store.cu
+384
View File
@@ -0,0 +1,384 @@
/***************************************************************************************************
* Copyright (c) 2017 - 2023 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 <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
};
// __grid_constant__ was introduced in CUDA 11.7.
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
# define CUTE_GRID_CONSTANT_SUPPORTED
#endif
// __grid_constant__ can be enabled only on SM70+
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
# define CUTE_GRID_CONSTANT_ENABLED
#endif
#if ! defined(CUTE_GRID_CONSTANT)
# if defined(CUTE_GRID_CONSTANT_SUPPORTED) && defined(CUTE_GRID_CONSTANT_ENABLED)
# define CUTE_GRID_CONSTANT __grid_constant__
# else
# define CUTE_GRID_CONSTANT
# endif
#endif
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma,
GmemLayout gmem_layout, SmemLayout 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.data()), smem_layout);
//
// Read in trivially
//
Tensor gA_in = make_tensor(make_gmem_ptr(g_in), gmem_layout);
// Input gmem -> smem
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
sA(i) = gA_in(i);
}
__syncthreads();
#if 0
//
// Write out trivially
//
Tensor gA_out = make_tensor(make_gmem_ptr(g_out), gmem_layout);
// Output smem -> gmem
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
gA_out(i) = sA(i);
}
#else
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor gA = tma.get_tma_tensor(shape(gmem_layout));
//
// Prepare the TMA_STORE
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAsA = cta_tma.partition_S(sA);
Tensor tAgA = cta_tma.partition_D(gA);
//
// Perform the TMA_STORE
//
if (threadIdx.x == 0) {
copy(tma, tAsA, tAgA);
}
#endif
}
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Col)
{
using T = half_t;
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
Layout gmem_layout = smem_layout;
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE 32x32 ColMajor SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Row)
{
using T = half_t;
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
Layout gmem_layout = smem_layout;
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE 32x32 RowMajor SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN)
{
using T = half_t;
auto smem_layout = GMMA::Layout_MN_SW128_Atom<T>{};
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_K)
{
using T = half_t;
auto smem_layout = GMMA::Layout_K_SW128_Atom<T>{};
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenRowMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_K_SW128_Atom<T> SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi)
{
using T = half_t;
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{});
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi2)
{
using T = half_t;
// Tile the GMMA::Layout atom in the K-mode first, then the M-mode to get a bigger box size
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
Layout gmem_layout = make_layout(make_shape(size<0>(smem_layout), size<1>(smem_layout)), GenColMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_GMMA_SW128_MN_Multi_Dyn)
{
using T = half_t;
auto smem_layout = tile_to_shape(GMMA::Layout_MN_SW128_Atom<T>{}, Shape<Int<128>,Int<128>>{}, Step<_2,_1>{});
Layout gmem_layout = make_layout(make_shape(128, 128), GenColMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
TEST(SM90_CuTe_Hopper, Tma_Store_32x32_Multimode)
{
using T = half_t;
auto smem_layout = Layout<Shape<_32,_32>, Stride<_32,_1>>{};
Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenRowMajor{});
//auto smem_layout = Layout<Shape<_32,_32>>{};
//Layout gmem_layout = make_layout(make_shape(make_shape(8,4), 32), GenColMajor{});
thrust::host_vector<T> h_in(size(smem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i); }
thrust::device_vector<T> d_in = h_in;
thrust::device_vector<T> d_out(h_in.size(), T(-1));
Tensor gA = make_tensor(d_out.data().get(), gmem_layout);
auto tma = make_tma_copy(SM90_TMA_STORE{}, gA, smem_layout);
//print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n");
int smem_size = int(sizeof(SharedStorage<T, decltype(smem_layout)>));
tma_test_device_cute<<<1, 128, smem_size>>>(
thrust::raw_pointer_cast(d_in.data()),
thrust::raw_pointer_cast(d_out.data()),
tma,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
for (int i = 0; i < size(smem_layout); ++i) {
//printf("%d %d\n", int(h_in[i]), int(h_out[i]));
EXPECT_EQ(h_out[i], h_in[i]);
}
CUTLASS_TRACE_HOST("CuTe TMA_STORE GMMA::Layout_MN_SW128_Atom<T> Multi SUCCESS\n");
}
#endif
test/unit/cute/layout/CMakeLists.txt
+32
View File
@@ -0,0 +1,32 @@
# Copyright (c) 2023 - 2023 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_layout
layout_operator.cu
)
test/unit/cute/layout/layout_operator.cu
+136
View File
@@ -0,0 +1,136 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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.
*
**************************************************************************************************/
/*! \file
\brief Unit tests Generic CuTe Layouts
*/
#include "../../common/cutlass_unit_test.h"
#include "cutlass/cutlass.h"
#include "cutlass/layout/layout.h"
#include "cutlass/matrix_coord.h"
// Cute includes
#include <cute/layout.hpp>
#include <cute/int_tuple.hpp>
using namespace cutlass;
using namespace cute;
namespace test {
namespace layout {
template <typename GenericLayout, typename Layout>
struct Testbed {
Testbed() {}
bool run() {
GenericLayout generic_layout;
Layout layout = Layout::packed({size<0>(generic_layout), size<1>(generic_layout)});
for (int m = 0; m < size<0>(generic_layout); m++) {
for (int n = 0; n < size<1>(generic_layout); n++) {
if (generic_layout(m, n) != layout({m, n})) return false;
}
}
return true;
}
};
}
}
//////////////////////////////////////////////////////////////////////////
// Test Generic CuTe Layouts
//////////////////////////////////////////////////////////////////////////
/// Canonical Layouts
TEST(GenericLayout, ColumnMajor) {
using GenericLayout = cute::Layout<Shape<_8, _4>, Stride<_1, _8>>;
using Layout = cutlass::layout::ColumnMajor;
test::layout::Testbed<GenericLayout, Layout> testbed;
EXPECT_TRUE(testbed.run());
}
//////////////////////////////////////////////////////////////////////////
TEST(GenericLayout, RowMajor) {
using GenericLayout = cute::Layout<Shape<_8, _4>, Stride<_4, _1>>;
using Layout = cutlass::layout::RowMajor;
test::layout::Testbed<GenericLayout, Layout> testbed;
EXPECT_TRUE(testbed.run());
}
//////////////////////////////////////////////////////////////////////////
/// Swizzle Shared Memory layouts
TEST(GenericLayout, RowMajorTensorOpMultiplicandCrosswise) {
using GenericLayout = decltype(
composition(
Swizzle<3,3,3>{},
Layout<Shape<_128, _64>, Stride<_64, _1>>{})
);
using Layout = cutlass::layout::RowMajorTensorOpMultiplicandCrosswise<
cutlass::sizeof_bits<cutlass::half_t>::value, 64>;
test::layout::Testbed<GenericLayout, Layout> testbed;
EXPECT_TRUE(testbed.run());
}
//////////////////////////////////////////////////////////////////////////
TEST(GenericLayout, ColumnMajorTensorOpMultiplicandCongruous) {
using GenericLayout = decltype(
composition(
Swizzle<3,3,4>{},
Layout<Shape<_128, _64>>{})
);
using Layout = cutlass::layout::ColumnMajorTensorOpMultiplicandCongruous<
cutlass::sizeof_bits<cutlass::half_t>::value, 64>;
test::layout::Testbed<GenericLayout, Layout> testbed;
EXPECT_TRUE(testbed.run());
}
//////////////////////////////////////////////////////////////////////////