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:
-2
@@ -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) {
|
||||
|
||||
|
||||
-1
@@ -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) {
|
||||
|
||||
|
||||
-1
@@ -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) {
|
||||
|
||||
|
||||
-1
@@ -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) {
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
{
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+6
-6
@@ -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
|
||||
>;
|
||||
|
||||
|
||||
+4
-4
@@ -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<
|
||||
|
||||
+2
-2
@@ -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<
|
||||
|
||||
+6
-6
@@ -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
|
||||
>;
|
||||
|
||||
|
||||
+4
-4
@@ -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<
|
||||
|
||||
+2
-2
@@ -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<
|
||||
|
||||
+2
-2
@@ -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;
|
||||
|
||||
@@ -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>;
|
||||
|
||||
@@ -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>{};
|
||||
|
||||
//
|
||||
|
||||
Reference in New Issue
Block a user