CUTLASS 3.2.1 (#1113)

* Updates for 3.2.1 release.

* Minor fix in gemm op profiler for raster order.

* Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
2023-09-26 17:24:26 -04:00
committed by GitHub
parent e0aaa3c3b3
commit 90d3b0fb18
428 changed files with 22252 additions and 21761 deletions
@@ -42,8 +42,6 @@
#include "conv2d_testbed_interleaved.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
TEST(SM75_Device_Conv2d_Fprop_Analytic_ImplicitGemm_s4ncxhwx_s4cxrskx_s4ncxhwx_tensor_op_s32,
128x128_128x2_64x64x128) {
@@ -42,7 +42,6 @@
#include "conv2d_testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
TEST(SM75_Device_Conv2d_Fprop_Analytic_ImplicitGemm_s4nhwc_s4nhwc_s32nhwc_tensor_op_s32,
128x128_64x3_64x64x64) {
@@ -42,7 +42,6 @@
#include "conv2d_testbed_interleaved.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
TEST(SM75_Device_Conv2d_Fprop_Analytic_ImplicitGemm_s8ncxhwx_s8cxrskx_s8ncxhwx_tensor_op_s32,
128x128_64x2_64x64x64) {
@@ -42,7 +42,6 @@
#include "conv2d_testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
TEST(SM75_Device_Conv2d_Fprop_Analytic_ImplicitGemm_s8nhwc_s8nhwc_s32nhwc_tensor_op_s32,
128x128_64x2_64x64x64) {
+20
View File
@@ -43,3 +43,23 @@ cutlass_test_unit_add_executable(
numeric_conversion.cu
functional.cu
)
#
# CUTLASS 3x increases the host compiler requirements to C++17. However, there are
# certain existing integrations that will benefit from maintaining C++11 compatibility.
#
# This requirement only applies to select .h files which are explicitly annotated. It
# does not apply to any .hpp file.
#
# `cutlass_test_unit_core_cpp11` enforces the C++11 requirement.
#
set(CMAKE_CUDA_STANDARD 11)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
add_executable(
cutlass_test_unit_core_cpp11
cpp11.cu
)
+86
View File
@@ -0,0 +1,86 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
/////////////////////////////////////////////////////////////////////////////////////////////////
#include <cutlass/cutlass.h>
#include <cutlass/coord.h>
#include <cutlass/core_io.h>
#include <cutlass/array.h>
#include <cutlass/bfloat16.h>
#include <cutlass/fast_math.h>
#include <cutlass/float8.h>
#include <cutlass/half.h>
#include <cutlass/integer_subbyte.h>
#include <cutlass/kernel_hardware_info.h>
#include <cutlass/numeric_size.h>
#include <cutlass/numeric_types.h>
#include <cutlass/tfloat32.h>
#include <cutlass/workspace.h>
#include <cutlass/subbyte_reference.h>
#include <cutlass/conv/convolution.h>
#include <cutlass/conv/conv2d_problem_size.h>
#include <cutlass/epilogue/threadblock/predicated_tile_iterator_params.h>
#include <cutlass/gemm/gemm_enumerated_types.h>
#include <cutlass/gemm/kernel/tile_scheduler_params.h>
#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
#include <cutlass/layout/matrix.h>
#include <cutlass/layout/pitch_linear.h>
#include <cutlass/transform/threadblock/predicated_tile_access_iterator_params.h>
/////////////////////////////////////////////////////////////////////////////////////////////////
#if (201700L <= __cplusplus )
#error "This file and all of its includes must be compilable as C++11."
#endif
/////////////////////////////////////////////////////////////////////////////////////////////////
int main() {
return 0;
}
/////////////////////////////////////////////////////////////////////////////////////////////////
+1 -1
View File
@@ -36,7 +36,7 @@ cutlass_test_unit_add_executable(
compare.cpp
complement.cpp
composition.cpp
constant_arithmetic.cpp
constants.cpp
core_unit.cpp
inverse_left.cpp
inverse_right.cpp
-106
View File
@@ -1,106 +0,0 @@
/***************************************************************************************************
* 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, ConstantArithmetic) {
using namespace cute;
constexpr cute::integral_constant<uint32_t, 0> uzero{};
// This extra test exists historically as part of the diagnosis
// of a possible Clang 14 bug. However, it's a nice test for
// cute::integral_constant's arithmetic operators, so it's saved here.
// It also demonstrates how to work with cute::integral_constant
// and lambda captures. Microsoft Visual Studio ("MSVC") tends to
// disagree with other compilers about the meaning of decltype
// for variables captured by reference. MSVC and GCC 8.3.0
// also tend to disagree with other compilers (and other GCC versions)
// about whether expressions involving such variables
// are constant expressions.
//
// A typical CuTe idiom is to do lambda captures by reference [&].
// This test changes them to capture by value, except for
// the innermost lambda's capture of S1, which is by reference.
// The point is to show that MSVC and GCC 8 have issues with this
// that other compilers do not. For example,
//
// 1. MSVC needs remove_cvref_t around decltype(S1)
// in order to access decltype(S1)::value, and
// 2. MSVC and GCC 8.3.0 both report a build error with S1()
// (that is, calling operator() on S1, which returns the
// same thing as S1.value).
//
// The reason for (2) is that neither compiler thinks
// that S1() is a constant expression.
//
// This leaves S1.value as the most concise portable expression
// for the "value" member of a cute::integral_constant.
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero](auto S0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0](auto F0) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0,F0](auto S1) {
for_each(make_integer_sequence<uint32_t, 8>{}, [uzero,S0,F0,&S1](auto F1) {
static_assert((decltype(S0)::value & decltype(F0)::value) == decltype(S0 & F0)::value);
// Using S1.value means you don't have to use remove_cvref_t
// with a captured-by-reference variable.
static_assert((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) == decltype(S1 & F1)::value);
static_assert((S1.value & decltype(F1)::value) == decltype(S1 & F1)::value);
// S1() _should_ work, but does not with Visual Studio 2022,
// which emits C2131 ("expression did not evaluate to a constant").
// It also does not with GCC 8.3.0, which emits an error with messages
// "non-constant condition for static assertion" and
// "'this' is not a constant expression."
//
//static_assert((S1() & decltype(F1)::value) == decltype(S1 & F1)::value);
static_assert(decltype((S0 & F0) != uzero)::value == ((decltype(S0)::value & decltype(F0)::value) != 0));
static_assert(decltype((S1 & F1) != uzero)::value == ((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) != 0));
static_assert(decltype((S1 & F1) != uzero)::value == ((S1.value & decltype(F1)::value) != 0));
constexpr bool left = decltype((S0 & F0) != uzero || (S1 & F1) != uzero)::value;
constexpr bool right =
((decltype(S0)::value & decltype(F0)::value) != 0) ||
((cute::remove_cvref_t<decltype(S1)>::value & decltype(F1)::value) != 0);
constexpr bool right2 =
((decltype(S0)::value & decltype(F0)::value) != 0) ||
((S1.value & decltype(F1)::value) != 0);
static_assert(right == right2);
static_assert(left == right);
constexpr bool left2 = decltype((S0 & F0) != uzero)::value || decltype((S1 & F1) != uzero)::value;
static_assert(left == left2);
});
});
});
});
}
+60
View File
@@ -0,0 +1,60 @@
/***************************************************************************************************
* 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/numeric/integral_constant.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
TEST(CuTe_core, MakeIntegerSequence) {
cute::for_each(cute::make_integer_sequence<uint32_t, 13>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<uint32_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
static_assert(cute::is_same_v<typename c_type::value_type, uint32_t>);
static_assert(cute::is_constant<c_value, c_type>::value);
static_assert(cute::is_constant<0, decltype(c * cute::Int<0>{})>::value);
static_assert(cute::is_constant<2*c_value, decltype(c * cute::Int<2>{})>::value);
});
cute::for_each(cute::make_integer_sequence<int64_t, 17>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<int64_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
static_assert(cute::is_same_v<typename c_type::value_type, int64_t>);
static_assert(cute::is_constant<c_value, c_type>::value);
static_assert(cute::is_constant<0, decltype(c * cute::Int<0>{})>::value);
static_assert(cute::is_constant<2*c_value, decltype(c * cute::Int<2>{})>::value);
});
}
+17 -73
View File
@@ -31,87 +31,31 @@
#include "cutlass_unit_test.h"
// C<uint32_t(something)>::value_type is not uint32_t for GCC 7.5.0.
// This test is thus disabled for GCC < 8.
#if defined(__GNUC__) && (__GNUC__ < 8)
#include <cutlass/trace.h>
#include <cute/swizzle.hpp>
namespace { // (anonymous)
// This function exists to work around a Clang 14 issue, in which
// the compiler tries to instantiate code that lives inside the
// "else" branch of an "if constexpr," even when the "else" branch
// is false. That triggers a spurious static_assert in MixedBits.
// The work-around is to make the body of the "else" branch a
// function, rather than leaving it in line.
//
// Some compilers strangely deduce the first two terms of
// make_integer_sequence<uint32_t, 8> as C<false> and C<true>, and
// the remaining terms as C<2>, C<3>, etc. Making this function take
// cute::integral_constant<uint32_t, S0_value>, etc. doesn't work
// with those compilers.
template<class S0_type, S0_type S0_value,
class F0_type, F0_type F0_value,
class S1_type, S1_type S1_value,
class F1_type, F1_type F1_value>
void clang14_workaround(cute::integral_constant<S0_type, S0_value>,
cute::integral_constant<F0_type, F0_value>,
cute::integral_constant<S1_type, S1_value>,
cute::integral_constant<F1_type, F1_value>)
{
constexpr cute::C<static_cast<uint32_t>(S0_value)> S0{};
constexpr cute::C<static_cast<uint32_t>(F0_value)> F0{};
constexpr cute::C<static_cast<uint32_t>(S1_value)> S1{};
constexpr cute::C<static_cast<uint32_t>(F1_value)> F1{};
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(uint32_t(m0 & m1), uint32_t(m0) & uint32_t(m1));
//print(m0); print(" | "); print(m1); print(" = "); print(m0 | m1); print("\n");
EXPECT_EQ(uint32_t(m0 | m1), uint32_t(m0) | uint32_t(m1));
//print(m0); print(" ^ "); print(m1); print(" = "); print(m0 ^ m1); print("\n");
EXPECT_EQ(uint32_t(m0 ^ m1), uint32_t(m0) ^ uint32_t(m1));
}
}
}
} // namespace (anonymous)
TEST(CuTe_core, MixedBits) {
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 {
clang14_workaround(S0, F0, S1, F1);
for_each(make_int_sequence<8>{}, [&](auto S0) {
for_each(make_int_sequence<8>{}, [&](auto F0) {
for_each(make_int_sequence<8>{}, [&](auto S1) {
for_each(make_int_sequence<8>{}, [&](auto F1) {
for (uint32_t d0 = 0; d0 < 8; ++d0) {
for (uint32_t d1 = 0; d1 < 8; ++d1) {
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(uint32_t(m0 & m1), uint32_t(m0) & uint32_t(m1));
//print(m0); print(" | "); print(m1); print(" = "); print(m0 | m1); print("\n");
EXPECT_EQ(uint32_t(m0 | m1), uint32_t(m0) | uint32_t(m1));
//print(m0); print(" ^ "); print(m1); print(" = "); print(m0 ^ m1); print("\n");
EXPECT_EQ(uint32_t(m0 ^ m1), uint32_t(m0) ^ uint32_t(m1));
}
}
});
});
});
});
}
TEST(CuTe_core, MakeIntegerSequence) {
cute::for_each(cute::make_integer_sequence<uint32_t, 8>{}, [](auto c) {
using c_type = decltype(c);
constexpr auto c_value = c_type::value;
using expected_type = cute::integral_constant<uint32_t, c_value>;
static_assert(cute::is_same_v<c_type, expected_type>);
});
}
#endif // defined(__GNUC__) && (__GNUC__ < 8)
+159 -164
View File
@@ -31,182 +31,79 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "../hopper/tma_load_testbed.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];
};
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr(g_out), gmem_layout);
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_LOAD
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAgA_x = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
Tensor tAsA_x = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA.data()); print(" o "); print( mA.layout()); print("\n");
print(" gA : "); print( gA.data()); print(" o "); print( gA.layout()); print("\n");
print("tAgA_x: "); print(tAgA_x.data()); print(" o "); print(tAgA_x.layout()); print("\n");
print(" sA : "); print( sA.data()); print(" o "); print( sA.layout()); print("\n");
print("tAsA_x: "); print(tAsA_x.data()); print(" o "); print(tAsA_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_LOAD
//
// INPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tAgA = group_modes<1,rank(tAgA_x)>(tAgA_x); // (TMA,REST)
Tensor tAsA = group_modes<1,rank(tAsA_x)>(tAsA_x); // (TMA,REST)
static_assert(size<1>(tAsA) == 1);
// OUTPUT: Group the CTA_TILE_X modes and REST_X modes for output
Tensor tBgB = group_modes<0,R>(group_modes<R,rank(gB)>(gB)); // (CTA_TILE, REST)
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tAsA : "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = size(sA) * sizeof_bits_v<T> / 8;
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(_,stage), tAsA(_,0));
}
__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);
//
// Write out trivially smem -> gmem
//
//if (thread0()) {
// print_tensor(sA);
//}
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
tBgB(i,stage) = sA(i);
}
__syncthreads();
}
}
template <class T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
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);
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, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
for (int i = 0; i < size(gmem_layout); ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
}
using namespace cute;
return test_tma_load<T, TmaType>(SM90_TMA_LOAD{}, gmem_layout, smem_layout, cta_tile);
}
template <class T, class GMEM_Layout, class SMEM_Layout>
void
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
auto
test_tma_load(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
using namespace cute;
return test_tma_load<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_1D)
{
Layout smem_layout = Layout<_256, _1>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
Layout smem_layout = Layout<_256, _1>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
{
Layout gmem_layout = make_layout(128, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
{
Layout gmem_layout = make_layout(384, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
}
{
Layout gmem_layout = make_layout(128, GenColMajor{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
Layout smem_layout = Layout<Shape<_8,_8>, Stride<_1,_8>>{};
{
Layout gmem_layout = smem_layout;
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
// This doesn't result in a 1D TMA, even though it could/should...
{
Layout gmem_layout = tile_to_shape(smem_layout, Shape<_16,_16>{});
test_tma_load<int8_t>(gmem_layout, smem_layout);
test_tma_load<half_t>(gmem_layout, smem_layout);
test_tma_load< float>(gmem_layout, smem_layout);
test_tma_load<double>(gmem_layout, smem_layout);
}
}
}
@@ -270,18 +167,32 @@ template <class T, template <typename> typename SWIZZLE_ATOM>
void
test_tma_load_swizzle_atom_mn()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
auto smem_layout = SWIZZLE_ATOM<T>{};
{ // Static gmem
//Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
//test_tma_load<T>(gmem_layout, smem_layout);
}
{ // Dynamic gmem
Layout gmem_layout = make_layout(make_shape(2*uint32_t(size<0>(smem_layout)), 2*uint32_t(size<1>(smem_layout))),
GenColMajor{});
test_tma_load<T>(gmem_layout, smem_layout);
}
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
test_tma_load_swizzle_atom_k()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
auto smem_layout = SWIZZLE_ATOM<T>{};
{ // Static gmem
//Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
//test_tma_load<T>(gmem_layout, smem_layout);
}
{ // Dynamic gmem
Layout gmem_layout = make_layout(make_shape(2*uint32_t(size<0>(smem_layout)), 2*uint32_t(size<1>(smem_layout))),
GenRowMajor{});
test_tma_load<T>(gmem_layout, smem_layout);
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Atoms)
@@ -328,21 +239,21 @@ TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Atoms)
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
auto
test_tma_load_swizzle_tile_mn()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenColMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
return test_tma_load<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
void
auto
test_tma_load_swizzle_tile_k()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenRowMajor{});
return test_tma_load<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
return test_tma_load<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Load_Swizzle_Tiles)
@@ -431,4 +342,88 @@ TEST(SM90_CuTe_Hopper, Tma_Load_Tensor_Multimode)
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_Coalesce)
{
// Interleaved ColMajor
{
Layout gmem_layout = make_layout(make_shape ( 128, make_shape (_4{}, 128)),
make_stride( _4{}, make_stride(_1{}, 512)));
auto smem_layout = make_layout(make_shape (_32{}, make_shape (_4{}, _32{})),
make_stride( _4{}, make_stride(_1{}, _128{})));
// By default, uses cta_tile = Shape<_32,_128>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Interleaved RowMajor
{
Layout gmem_layout = make_layout(make_shape (make_shape (_4{}, 128), 128),
make_stride(make_stride(_1{}, 512), _4{}));
auto smem_layout = make_layout(make_shape (make_shape (_4{}, _32{}), _32{}),
make_stride(make_stride(_1{}, _128{}), _4{}));
// By default, uses cta_tile = Shape<_128,_32>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Account for stride-0 modes within the TMA tile
{
Layout gmem_layout = make_layout(make_shape ( 128, make_shape (_32{}, 4)),
make_stride( _1{}, make_stride( _0{}, 128)));
auto smem_layout = make_layout(make_shape (_64{}, make_shape (_32{} )),
make_stride( _1{}, make_stride( _0{} )));
// By default, uses cta_tile = Shape<_64,_32>
auto tma = test_tma_load<uint16_t>(gmem_layout, smem_layout);
// Check the TMA rank
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 2);
}
// Coalesce many modes and account for stride-0 modes within the TMA tile
{
Layout gmem_layout = make_layout(make_shape (make_shape (_32{},_4{}, 4), _32{}, make_shape (_4{}, 4)),
make_stride(make_stride(_16{},_4{}, 2048), _0{}, make_stride(_1{}, _512{})));
auto smem_layout = make_layout(make_shape (make_shape (_32{},_4{} ), _32{}, make_shape (_4{} )),
make_stride(make_stride(_16{},_4{} ), _0{}, make_stride(_1{} )));
// By default, uses cta_tile = Shape<_128,_32,_4>
auto tma = test_tma_load<int8_t>(gmem_layout, smem_layout);
// Check the TMA rank (Could be 3 instead of 4 with even better coalescing...?)
EXPECT_EQ(rank(tma.get_tma_tensor(shape(gmem_layout))(0)), 4);
}
}
TEST(SM90_CuTe_Hopper, Tma_Load_InternalType)
{
Layout smem_layout = Layout<Shape<_32,_32>, Stride<_1,_32>>{};
Layout gmem_layout = make_layout(make_shape(64, 64));
// Downcasted tensors to smaller TmaTypes
{
test_tma_load<int8_t, uint8_t>(gmem_layout, smem_layout);
test_tma_load<half_t, uint8_t>(gmem_layout, smem_layout);
test_tma_load< float, uint8_t>(gmem_layout, smem_layout);
test_tma_load<double, uint8_t>(gmem_layout, smem_layout);
}
// Upcasted tensors to larger TmaTypes
{
test_tma_load<int8_t, uint64_t>(gmem_layout, smem_layout);
test_tma_load<half_t, uint64_t>(gmem_layout, smem_layout);
test_tma_load< float, uint64_t>(gmem_layout, smem_layout);
test_tma_load<double, uint64_t>(gmem_layout, smem_layout);
}
// Complex<double> is 128bit, which the TMA has no concept of
{
test_tma_load<complex<double>, uint64_t>(gmem_layout, smem_layout);
test_tma_load<complex<double>, uint32_t>(gmem_layout, smem_layout);
}
}
#endif
+199
View File
@@ -0,0 +1,199 @@
/***************************************************************************************************
* 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>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
cute::uint64_t tma_load_mbar[1];
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// Shared memory barriers use 64bits in SMEM for synchronization
uint64_t* tma_load_mbar = shared_storage.tma_load_mbar;
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = tma.get_tma_tensor(shape(gmem_layout));
Tensor mB = make_tensor(make_gmem_ptr(g_out), gmem_layout);
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_LOAD
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tAgA_x = cta_tma.partition_S(gA); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
Tensor tAsA_x = cta_tma.partition_D(sA); // (TMA,TMA_M,TMA_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mA : "); print( mA.data()); print(" o "); print( mA.layout()); print("\n");
print(" gA : "); print( gA.data()); print(" o "); print( gA.layout()); print("\n");
print("tAgA_x: "); print(tAgA_x.data()); print(" o "); print(tAgA_x.layout()); print("\n");
print(" sA : "); print( sA.data()); print(" o "); print( sA.layout()); print("\n");
print("tAsA_x: "); print(tAsA_x.data()); print(" o "); print(tAsA_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_LOAD
//
// INPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tAgA = group_modes<1,rank(tAgA_x)>(tAgA_x); // (TMA,REST)
Tensor tAsA = group_modes<1,rank(tAsA_x)>(tAsA_x); // (TMA,REST)
static_assert(size<1>(tAsA) == 1);
// OUTPUT: Group the CTA_TILE_X modes and REST_X modes for output
Tensor tBgB = group_modes<0,R>(group_modes<R,rank(gB)>(gB)); // (CTA_TILE, REST)
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tAsA : "); print(tAsA.data()); print(" o "); print(tAsA.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tAgA); ++stage)
{
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
constexpr int kTmaTransactionBytes = size(sA) * sizeof_bits_v<T> / 8;
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(_,stage), tAsA(_,0));
}
__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);
//
// Write out trivially smem -> gmem
//
//if (thread0()) {
// print_tensor(sA);
//}
for (int i = threadIdx.x; i < size(sA); i += blockDim.x) {
tBgB(i,stage) = sA(i);
}
__syncthreads();
}
}
template <class T, class TmaType = T, class CopyOp, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
auto
test_tma_load(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
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<TmaType>(copy_op, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
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, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
// Validate the results, and tolerate the first 3 errors:
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
int count = 3;
for (int i = 0; i < cute::size(gmem_layout) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
return tma;
}
#endif
} // end namespace cutlass::test
+16 -137
View File
@@ -31,150 +31,30 @@
#include "cutlass_unit_test.h"
#include <iostream>
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include "../hopper/tma_store_testbed.hpp"
using namespace cute;
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
};
using namespace cutlass::test;
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor mB = tma.get_tma_tensor(shape(gmem_layout));
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_STORE
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tBsB_x = cta_tma.partition_S(sB); // (TMA,TMA_M,TMA_N)
Tensor tBgB_x = cta_tma.partition_D(gB); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mB : "); print( mB.data()); print(" o "); print( mB.layout()); print("\n");
print(" gB : "); print( gB.data()); print(" o "); print( gB.layout()); print("\n");
print("tBgB_x: "); print(tBgB_x.data()); print(" o "); print(tBgB_x.layout()); print("\n");
print(" sB : "); print( sB.data()); print(" o "); print( sB.layout()); print("\n");
print("tBsB_x: "); print(tBsB_x.data()); print(" o "); print(tBsB_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_STORE
//
// INPUT: Group the CTA_TILE_X modes and REST_X modes for input
Tensor tAgA = group_modes<0,R>(group_modes<R,rank(gA)>(gA)); // (CTA_TILE, REST)
// OUTPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tBgB = group_modes<1,rank(tBgB_x)>(tBgB_x); // (TMA,REST)
Tensor tBsB = group_modes<1,rank(tBsB_x)>(tBsB_x); // (TMA,REST)
static_assert(size<1>(tBsB) == 1);
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tBsB : "); print(tBsB.data()); print(" o "); print(tBsB.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tBgB); ++stage)
{
//
// Read in trivially gmem -> smem
//
for (int i = threadIdx.x; i < size(sB); i += blockDim.x) {
sB(i) = tAgA(i,stage);
}
__syncthreads();
//
// Perform the TMA_STORE
//
if (threadIdx.x == 0) {
copy(tma, tBsB(_,0), tBgB(_,stage));
}
tma_store_wait<0>();
__syncthreads();
}
}
template <class T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
test_tma_store(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
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, cta_tile, Int<1>{});
//print(tma);
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, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
for (int i = 0; i < size(gmem_layout); ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
}
using namespace cute;
return test_tma_store<T, TmaType>(SM90_TMA_STORE{}, gmem_layout, smem_layout, cta_tile);
}
template <class T, class GMEM_Layout, class SMEM_Layout>
template <class T, class TmaType = T, class GMEM_Layout, class SMEM_Layout>
void
test_tma_store(GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout)
{
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
using namespace cute;
return test_tma_store<T, TmaType>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
}
TEST(SM90_CuTe_Hopper, Tma_Load_1D)
@@ -258,8 +138,8 @@ void
test_tma_store_swizzle_atom_mn()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
@@ -267,8 +147,8 @@ void
test_tma_store_swizzle_atom_k()
{
auto smem_layout = SWIZZLE_ATOM<T>{};
Layout gmem_layout = make_layout(shape(smem_layout), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Atoms)
@@ -319,8 +199,8 @@ void
test_tma_store_swizzle_tile_mn()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenColMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
template <class T, template <typename> typename SWIZZLE_ATOM>
@@ -328,8 +208,8 @@ void
test_tma_store_swizzle_tile_k()
{
auto smem_layout = tile_to_shape(SWIZZLE_ATOM<T>{}, Shape<_128,_128>{});
Layout gmem_layout = make_layout(make_shape(int(size<0>(smem_layout)), int(size<1>(smem_layout))), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout, product_each(shape(smem_layout)));
Layout gmem_layout = make_layout(make_shape(2*size<0>(smem_layout), 2*size<1>(smem_layout)), GenRowMajor{});
return test_tma_store<T>(gmem_layout, smem_layout);
}
TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Tiles)
@@ -353,7 +233,6 @@ TEST(SM90_CuTe_Hopper, Tma_Store_Swizzle_Tiles)
test_tma_store_swizzle_tile_k<half_t, GMMA::Layout_K_INTER_Atom>();
}
// Tensor by-mode
TEST(SM90_CuTe_Hopper, Tma_Store_Tensor)
{
+184
View File
@@ -0,0 +1,184 @@
/***************************************************************************************************
* 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>
namespace cutlass::test {
template <class ElementType, class SmemLayout>
struct SharedStorage
{
cute::array_aligned<ElementType, cute::cosize_v<SmemLayout>> smem;
};
#if CUDA_12_0_SM90_FEATURES_SUPPORTED
template <class T, class TiledCopy, class CTA_Tiler, class GmemLayout, class SmemLayout>
__global__ void
tma_test_device_cute(T const* g_in, T* g_out,
CUTE_GRID_CONSTANT TiledCopy const tma, CTA_Tiler cta_tiler,
GmemLayout gmem_layout, SmemLayout smem_layout)
{
using namespace cute;
CUTE_STATIC_ASSERT_V(product_each(shape(cta_tiler)) == product_each(shape(smem_layout)));
// Use Shared Storage structure to allocate and distribute aligned SMEM addresses
extern __shared__ char shared_memory[];
using SharedStorage = SharedStorage<T, SmemLayout>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
// Construct SMEM tensor
Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem.data()), smem_layout); // (CTA_TILE_M,CTA_TILE_N,...)
// TMA requires special handling of strides to deal with coord codomain mapping
// Represent the full tensors -- get these from TMA
Tensor mA = make_tensor(make_gmem_ptr(g_in), gmem_layout);
Tensor mB = tma.get_tma_tensor(shape(gmem_layout));
constexpr int R = rank_v<CTA_Tiler>;
Tensor gA = local_tile(mA, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
Tensor gB = local_tile(mB, cta_tiler, repeat<R>(_)); // (CTA_TILE_M,CTA_TILE_N,...REST_M,REST_N,...)
//
// Prepare the TMA_STORE
//
auto cta_tma = tma.get_slice(Int<0>{}); // CTA slice
Tensor tBsB_x = cta_tma.partition_S(sB); // (TMA,TMA_M,TMA_N)
Tensor tBgB_x = cta_tma.partition_D(gB); // (TMA,TMA_M,TMA_N,REST_M,REST_N)
#if 0
if (thread0()) {
print(tma);
print("TILE : "); print(cta_tiler); print("\n");
print(" mB : "); print( mB.data()); print(" o "); print( mB.layout()); print("\n");
print(" gB : "); print( gB.data()); print(" o "); print( gB.layout()); print("\n");
print("tBgB_x: "); print(tBgB_x.data()); print(" o "); print(tBgB_x.layout()); print("\n");
print(" sB : "); print( sB.data()); print(" o "); print( sB.layout()); print("\n");
print("tBsB_x: "); print(tBsB_x.data()); print(" o "); print(tBsB_x.layout()); print("\n");
}
#endif
//
// Perform the TMA_STORE
//
// INPUT: Group the CTA_TILE_X modes and REST_X modes for input
Tensor tAgA = group_modes<0,R>(group_modes<R,rank(gA)>(gA)); // (CTA_TILE, REST)
// OUTPUT: Group the REST_X modes and the TMA_X modes to easily iterate through the tiles
Tensor tBgB = group_modes<1,rank(tBgB_x)>(tBgB_x); // (TMA,REST)
Tensor tBsB = group_modes<1,rank(tBsB_x)>(tBsB_x); // (TMA,REST)
static_assert(size<1>(tBsB) == 1);
#if 0
if (thread0()) {
print("tAgA : "); print(tAgA.data()); print(" o "); print(tAgA.layout()); print("\n");
print("tBsB : "); print(tBsB.data()); print(" o "); print(tBsB.layout()); print("\n");
print("tBgB : "); print(tBgB.data()); print(" o "); print(tBgB.layout()); print("\n");
}
#endif
// Loop over the TMA stages, using smem as our buffer
for (int stage = 0; stage < size<1>(tBgB); ++stage)
{
//
// Read in trivially gmem -> smem
//
for (int i = threadIdx.x; i < size(sB); i += blockDim.x) {
sB(i) = tAgA(i,stage);
}
__syncthreads();
//
// Perform the TMA_STORE
//
if (threadIdx.x == 0) {
copy(tma, tBsB(_,0), tBgB(_,stage));
}
tma_store_wait<0>();
__syncthreads();
}
}
template <class T, class TmaType = T, class CopyOp, class GMEM_Layout, class SMEM_Layout, class CTA_Tile>
void
test_tma_store(CopyOp const& copy_op,
GMEM_Layout const& gmem_layout,
SMEM_Layout const& smem_layout,
CTA_Tile const& cta_tile)
{
using namespace cute;
thrust::host_vector<T> h_in(cosize(gmem_layout));
for (int i = 0; i < h_in.size(); ++i) { h_in[i] = T(i % 13); }
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<TmaType>(copy_op, gA, smem_layout, cta_tile, Int<1>{});
//print(tma);
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, cta_tile,
gmem_layout,
smem_layout);
thrust::host_vector<T> h_out = d_out;
// Validate the results, and tolerate the first 3 errors:
Tensor hA_in = make_tensor(h_in.data(), gmem_layout);
Tensor hA_out = make_tensor(h_out.data(), gmem_layout);
int count = 3;
for (int i = 0; i < cute::size(gmem_layout) && count > 0; ++i) {
EXPECT_EQ(hA_in(i), hA_out(i));
if (hA_in(i) != hA_out(i)) {
--count;
}
}
}
#endif
} // end namespace cutlass::test
@@ -46,7 +46,8 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED)
#if defined(CUTLASS_ARCH_MMA_B1_AND_SM80_ENABLED)
////////////////////////////////////////////////////////////////////////////////
TEST(SM80_Device_Gemm_b1t_b1n_s32n_tensor_op_s32, 128x256x1024_64x64x1024) {
@@ -370,8 +371,12 @@ TEST(SM80_Device_Gemm_b1t_b1n_s32n_tensor_op_s32, 64x64x512_32x32x512) {
EXPECT_TRUE(test::gemm::device::TestAllGemmBasic<Gemm>());
}
#endif // defined(CUTLASS_ARCH_MMA_B1_AND_SM80_ENABLED)
////////////////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
TEST(SM80_Device_Gemm_XOR_b1t_b1n_s32n_tensor_op_s32, 128x256x1024_64x64x1024) {
using ElementOutput = int32_t;
using ElementAccumulator = int32_t;
@@ -694,6 +699,6 @@ TEST(SM80_Device_Gemm_XOR_b1t_b1n_s32n_tensor_op_s32, 64x64x512_32x32x512) {
EXPECT_TRUE(test::gemm::device::TestAllGemmBasic<Gemm>());
}
////////////////////////////////////////////////////////////////////////////////
#endif // defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
#endif // #if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED)
////////////////////////////////////////////////////////////////////////////////
@@ -47,10 +47,9 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED)
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
CUTLASS_TEST_L1(SM80_Device_Gemm_XOR_b1t_b1n_s32t_tensor_op_s32, 128x256x1024_64x64x1024, {
using ElementOutput = int32_t;
@@ -376,4 +375,4 @@ CUTLASS_TEST_L1(SM80_Device_Gemm_XOR_b1t_b1n_s32t_tensor_op_s32, 64x64x512_32x32
////////////////////////////////////////////////////////////////////////////////
#endif // #if defined(CUTLASS_ARCH_MMA_SM80_SUPPORTED)
#endif // #if defined(CUTLASS_ARCH_MMA_B1_XOR_SM80_ENABLED)
@@ -49,7 +49,6 @@
#include "testbed_interleaved.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s4n_s4t_s4n_tensor_op_s32, 64x128x128_32x64x128) {
@@ -195,5 +194,4 @@ TEST(SM75_Device_Gemm_s4n_s4t_s4n_tensor_op_s32, 128x256x128_64x64x128) {
}
////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s4t_s4n_s32n_tensor_op_s32, 128x256x128_64x64x128) {
@@ -245,5 +244,4 @@ TEST(SM75_Device_Gemm_s4t_s4n_s32n_tensor_op_s32, 64x64x128_32x32x128) {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "cutlass/util/reference/host/gemm.h"
#include "testbed.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
///////// WMMA Instruction Shape = 8x8x32, DataType/Instruction = s4 * s4 + s32 => s32 //////////
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -244,5 +243,4 @@ TEST(SM75_Device_Gemm_s4t_s4n_s32n_wmma_tensor_op_s32, 64x64x128_32x32x128_8x8x3
EXPECT_TRUE(test::gemm::device::TestAllGemm<Gemm>());
}
#endif //CUTLASS_SUBBYTE_INTEGER_MATRIX_MULTIPLY_ENABLED
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s4t_s4n_s32t_tensor_op_s32, 128x256x128_64x64x128) {
@@ -245,5 +244,4 @@ TEST(SM75_Device_Gemm_s4t_s4n_s32t_tensor_op_s32, 64x64x128_32x32x128) {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "cutlass/util/reference/host/gemm.h"
#include "testbed.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
///////// WMMA Instruction Shape = 8x8x32, DataType/Instruction = s4 * s4 + s32 => s32 //////////
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s4t_s4n_s4n_tensor_op_s32, 128x256x128_64x64x128) {
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s4t_s4n_s4t_tensor_op_s32, 128x256x128_64x64x128) {
@@ -243,6 +242,7 @@ TEST(SM75_Device_Gemm_s4t_s4n_s4t_tensor_op_s32, 256x64x128_64x64x128) {
EXPECT_TRUE(test::gemm::device::TestAllGemmBasic<Gemm>());
}
TEST(SM75_Device_Gemm_s4t_s4n_s4t_tensor_op_s32, 64x128x128_32x64x128) {
using ElementOutput = cutlass::int4b_t;
@@ -340,5 +340,4 @@ TEST(SM75_Device_Gemm_s4t_s4n_s4t_tensor_op_s32, 64x64x128_32x32x128) {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed_interleaved.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s8n_s8t_s8n_tensor_op_s32, 32x64x64_16x32x64) {
@@ -289,5 +288,4 @@ TEST(SM75_Device_Gemm_s8n_s8t_s8n_tensor_op_s32, 128x256x64_64x64x64) {
}
////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s8t_s8n_s32n_tensor_op_s32, 128x256x64_64x64x64) {
@@ -245,5 +244,4 @@ TEST(SM75_Device_Gemm_s8t_s8n_s32n_tensor_op_s32, 64x64x64_32x32x64) {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
TEST(SM75_Device_Gemm_s8t_s8n_s32t_tensor_op_s32, 128x256x64_64x64x64) {
@@ -245,5 +244,4 @@ TEST(SM75_Device_Gemm_s8t_s8n_s32t_tensor_op_s32, 64x64x64_32x32x64) {
}
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
CUTLASS_TEST_L0(SM75_Device_Gemm_s8t_s8n_s8n_tensor_op_s32, 128x256x64_64x64x64, {
@@ -212,5 +211,4 @@ CUTLASS_TEST_L0(SM75_Device_Gemm_s8t_s8n_s8n_tensor_op_s32, 64x64x64_32x32x64, {
} )
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
@@ -49,7 +49,6 @@
#include "testbed.h"
#if defined(CUTLASS_ARCH_MMA_SM75_SUPPORTED)
/////////////////////////////////////////////////////////////////////////////////////////////////
CUTLASS_TEST_L0(SM75_Device_Gemm_s8t_s8n_s8t_tensor_op_s32, 128x256x64_64x64x64, {
@@ -186,5 +185,4 @@ CUTLASS_TEST_L0(SM75_Device_Gemm_s8t_s8n_s8t_tensor_op_s32, 64x64x64_32x32x64, {
} )
/////////////////////////////////////////////////////////////////////////////////////////////////
#endif
+12 -6
View File
@@ -544,7 +544,6 @@ struct TestbedImpl {
// Initialize the GEMM operator
//
typename Gemm::Arguments arguments;
cutlass::KernelHardwareInfo hw_info;
hw_info.device_id = 0;
if (not profiling) {
@@ -557,12 +556,12 @@ struct TestbedImpl {
}
typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args;
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileScheduleTag, cutlass::gemm::StreamKScheduler>) {
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
scheduler_args = { static_cast<int>(splits) };
}
// DefaultEpilogue
arguments = typename Gemm::Arguments{
auto arguments = typename Gemm::Arguments {
cutlass::gemm::GemmUniversalMode::kGemm,
problem_size,
{
@@ -741,7 +740,7 @@ struct Testbed3xFusionOperation {
using ElementAux = non_void_t<typename FusionOp::ElementAux>;
using ElementAmax = non_void_t<typename FusionOp::ElementAmax>;
using LayoutTagAux = non_void_t<typename FusionOp::GmemLayoutTagAux, LayoutTagD>;
using ActivationFunctor = non_void_t<typename FusionOp::ActivationFn<ElementCompute>,
using ActivationFunctor = non_void_t<typename FusionOp::ActivationFn,
cutlass::epilogue::thread::Identity<ElementCompute>>;
static constexpr bool IsBiasEnabled = FusionOp::IsPerRowBiasSupported;
@@ -1152,7 +1151,7 @@ struct Testbed3xFusionOperation {
initialize(problem_size, alpha_, beta_);
typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args;
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileScheduleTag, cutlass::gemm::StreamKScheduler>) {
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
scheduler_args = { static_cast<int>(splits) };
}
@@ -1208,6 +1207,13 @@ struct Testbed3xFusionOperation {
fusion_args.bias_ptr = bias.device_data();
}
// example of how to set kernel activation arguments
if constexpr (cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::ScaledGELU_taylor<ElementCompute>>) {
// see ActivationFunctor::Arguments in activation.h for definition
// if Arguments doesn't exist then fusion_args.activation is empty
fusion_args.activation.scale = ElementCompute(1);
}
if constexpr (IsAbsMaxEnabled) {
fusion_args.amax_D_ptr = abs_max_D.device_data();
}
@@ -1297,7 +1303,7 @@ bool TestAll(double alpha = 1.0, double beta = 0.0, Testbed testbed = {}) {
std::vector<int> problem_size_k = {max_alignment, TileShapeK * (Stages + 1) - max_alignment};
std::vector<int> problem_splits = {1};
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileScheduleTag, cutlass::gemm::StreamKScheduler>) {
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
problem_splits.push_back(2);
problem_splits.push_back(3);
@@ -1316,7 +1316,7 @@ public:
}
typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args;
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileScheduleTag, cutlass::gemm::StreamKScheduler>) {
if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
scheduler_args = { splits };
}
@@ -454,57 +454,4 @@ using Sm90LinCombScalarReduce =
>;
} // namespace fusion
namespace collective {
template<
typename TileShape_MNK,
typename EpilogueTileType,
typename ElementC,
typename ElementD,
typename Schedule
>
struct EpilogueDescriptor{
using TileShape = TileShape_MNK;
using EpilogueTile =
decltype(detail::sm90_compute_tile_shape_or_override<ElementD, EpilogueTileType, Schedule>());
using DispatchPolicy =
decltype(detail::sm90_get_tma_dispatch_policy<TileShape_MNK,EpilogueTile,ElementC,ElementD, Schedule>());
constexpr static int StagesC = DispatchPolicy::StagesC;
constexpr static int StagesD = DispatchPolicy::StagesD;
};
template<
typename EpilogueDescriptor,
typename GmemLayoutTagAux,
typename ElementAux
>
struct AuxLoadDescriptor{
constexpr static int Stages = EpilogueDescriptor::StagesC;
using Element = ElementAux;
using Stride = gemm::TagToStrideC_t<GmemLayoutTagAux>;
using SmemLayoutAtom =
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile>());
using CopyOpS2R =
decltype(detail::sm90_get_smem_load_op_for_source<Stride, ElementAux>());
};
template<
typename EpilogueDescriptor,
typename GmemLayoutTagAux,
typename ElementAux
>
struct AuxStoreDescriptor{
constexpr static int Stages = EpilogueDescriptor::StagesD;
using Element = ElementAux;
using Stride = gemm::TagToStrideC_t<GmemLayoutTagAux>;
using SmemLayoutAtom =
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile>());
using CopyOpR2S =
decltype(detail::sm90_get_smem_store_op_for_accumulator<Stride, ElementAux>());
};
} // namespace collective
} // namespace cutlass::epilogue
@@ -70,10 +70,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t
>;
@@ -128,10 +128,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::ColumnMajor, cutlass::half_t
>;
@@ -185,10 +185,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 12
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::ColumnMajor, float
>;
@@ -72,10 +72,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombEVTDAG<
@@ -125,10 +125,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 12
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using AuxStoreDescriptor = cutlass::epilogue::collective::AuxStoreDescriptor<
using AuxStoreDescriptor = cutlass::epilogue::collective::detail::AuxStoreDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombDAGEVT<
@@ -71,7 +71,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombPerColumnBias<
@@ -121,7 +121,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombPerColumnBias<
@@ -71,10 +71,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t
>;
@@ -127,10 +127,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::ColumnMajor, cutlass::half_t
>;
@@ -182,10 +182,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule
>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::ColumnMajor, float
>;
@@ -72,10 +72,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using AuxLoadDescriptor = cutlass::epilogue::collective::AuxLoadDescriptor<
using AuxLoadDescriptor = cutlass::epilogue::collective::detail::AuxLoadDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombEVTDAG<
@@ -125,10 +125,10 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using AuxStoreDescriptor = cutlass::epilogue::collective::AuxStoreDescriptor<
using AuxStoreDescriptor = cutlass::epilogue::collective::detail::AuxStoreDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::half_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombDAGEVT<
@@ -71,7 +71,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombPerColumnBias<
@@ -121,7 +121,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::half_t, cutlass::half_t, EpilogueSchedule>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90LinCombPerColumnBias<
@@ -139,9 +139,9 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_bf16n_tensor_op_gmma_f32_epilogue, 64x128x128_
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::bfloat16_t, cutlass::bfloat16_t, EpilogueSchedule>;
using AuxStoreDescriptor = cutlass::epilogue::collective::AuxStoreDescriptor<
using AuxStoreDescriptor = cutlass::epilogue::collective::detail::AuxStoreDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::bfloat16_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90ScaledLinCombPerRowBiasEltActAmaxAux<
@@ -139,9 +139,9 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32t_tensor_op_gmma_f32_cooperative_epilogue,
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, float, float, EpilogueSchedule>;
using AuxStoreDescriptor = cutlass::epilogue::collective::AuxStoreDescriptor<
using AuxStoreDescriptor = cutlass::epilogue::collective::detail::AuxStoreDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, float>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90ScaledLinCombPerRowBiasEltActAmaxAux<
@@ -139,9 +139,9 @@ TEST(SM90_Device_Gemm_f8t_f8n_f8t_tensor_op_gmma_f32_persistent_epilogue, 64x128
using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using EpilogueDescriptor = cutlass::epilogue::collective::EpilogueDescriptor<
using EpilogueDescriptor = cutlass::epilogue::collective::detail::EpilogueDescriptor<
TileShape_MNK, EpilogueTileType, cutlass::float_e4m3_t, cutlass::float_e4m3_t, EpilogueSchedule>;
using AuxStoreDescriptor = cutlass::epilogue::collective::AuxStoreDescriptor<
using AuxStoreDescriptor = cutlass::epilogue::collective::detail::AuxStoreDescriptor<
EpilogueDescriptor, cutlass::layout::RowMajor, cutlass::float_e4m3_t>;
using FusionCallbacks = cutlass::epilogue::fusion::Sm90ScaledLinCombPerRowBiasEltActAmaxAux<
@@ -569,7 +569,6 @@ TEST(SM75_gemm_threadblock_crosswise,
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_gemm_threadblock_interleaved, tensor_op_32x32x64_16x16x64_8x8x16) {
using ElementA = uint8_t;
using LayoutA = cutlass::layout::ColumnMajorInterleaved<32>;
@@ -1793,6 +1792,7 @@ TEST(SM75_gemm_threadblock_interleaved,
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_gemm_threadblock_crosswise, tensor_op_64x64x512_64x64x512_8x8x128) {
using ElementA = cutlass::uint1b_t;
using LayoutA = cutlass::layout::RowMajor;
@@ -193,7 +193,6 @@ TEST(SM75_gemm_threadblock_wmma_tensor_op_col_row_row_s8, 64x64x64_64x64x64_16x1
///////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_SUBBYTE_INTEGER_MATRIX_MULTIPLY_ENABLED)
TEST(SM75_gemm_threadblock_wmma_tensor_op_row_col_row_s4, 64x64x128_64x64x128_8x8x32) {
using ElementA = cutlass::int4b_t;
using LayoutA = cutlass::layout::RowMajor;
@@ -262,6 +261,7 @@ TEST(SM75_gemm_threadblock_wmma_tensor_op_row_col_col_s4, 64x64x64_64x64x64_8x8x
problem_size.k(), alpha, beta)
.run(grid, block);
}
TEST(SM75_gemm_threadblock_wmma_tensor_op_row_col_row_b1, 64x64x512_64x64x512_8x8x128) {
using ElementA = cutlass::uint1b_t;
using LayoutA = cutlass::layout::RowMajor;
@@ -193,7 +193,6 @@ TEST(SM75_gemm_threadblock_singlestage_wmma_tensor_op_col_row_row_s8, 64x64x64_6
///////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_SUBBYTE_INTEGER_MATRIX_MULTIPLY_ENABLED)
TEST(SM75_gemm_threadblock_singlestage_wmma_tensor_op_row_col_row_s4, 64x64x128_64x64x128_8x8x32) {
using ElementA = cutlass::int4b_t;
using LayoutA = cutlass::layout::RowMajor;
@@ -262,6 +261,7 @@ TEST(SM75_gemm_threadblock_singlestage_wmma_tensor_op_row_col_col_s4, 64x64x64_6
problem_size.k(), alpha, beta)
.run(grid, block);
}
TEST(SM75_gemm_threadblock_singlestage_wmma_tensor_op_row_col_row_b1, 64x64x512_64x64x512_8x8x128) {
using ElementA = cutlass::uint1b_t;
using LayoutA = cutlass::layout::RowMajor;
+1 -1
View File
@@ -326,7 +326,6 @@ TEST(SM75_warp_gemm_tensor_op_crosswise_f16, 128x128x64_16x16x64_16x8x8) {
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_warp_gemm_tensor_op_crosswise_i8, 128x128x64_64x64x64_8x8x16) {
using Shape = cutlass::gemm::GemmShape<64, 64, 64>;
using InstructionShape = cutlass::gemm::GemmShape<8, 8, 16>;
@@ -746,6 +745,7 @@ TEST(SM75_warp_gemm_tensor_op_interleaved_i4, 128x128x128_16x16x128_8x8x32) {
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM75_warp_gemm_tensor_op_crosswise_b1, 128x128x512_64x64x512_8x8x128) {
using Shape = cutlass::gemm::GemmShape<64, 64, 512>;
using InstructionShape = cutlass::gemm::GemmShape<8, 8, 128>;
-6
View File
@@ -79,13 +79,9 @@ void pipeline_async_basic_device(uint32_t const num_iterations)
using SharedStorage = SharedStorage<NumStages>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
auto cta_layout = Layout<ClusterShape>{}; // (m,n) -> cta_id
int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
int lane_predicate = cute::elect_one_sync();
dim3 block_id_in_cluster = cute::block_id_in_cluster();
auto cluster_shape = ClusterShape{};
// This example showcases 2 producer 1 consumer example
typename MainloopPipeline::Params params;
@@ -158,8 +154,6 @@ struct PipelineTest {
cudaStream_t stream = nullptr) {
// Pipeline (multistage pipeline)
auto num_stages = Int<Stages>{};
auto cluster_shape = Shape<Int<ClusterShape::kM>, Int<ClusterShape::kN>, _1>{};
//
@@ -316,7 +316,6 @@ struct PipelineTest {
float elapsed_ms = 0.0f;
// Pipeline (multistage pipeline)
auto num_stages = Int<Stages>{};
auto cluster_shape = Shape<Int<ClusterShape::kM>, Int<ClusterShape::kN>, _1>{};
//