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

CUTLASS 3.1 (#915)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 20:19:34 -07:00
committed by GitHub
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37184 additions and 16419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
include/cute/algorithm/copy.hpp
View File

@@ -171,7 +171,7 @@ copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
if constexpr (sizeof(SrcType) == sizeof(DstType) && sizeof(VecType) > sizeof(DstType))
if constexpr (sizeof(SrcType) == sizeof(DstType) && sizeof(VecType) > sizeof(DstType))
{
/* @pre is_aligned<N>(src.data()) &&
* is_aligned<N>(dst.data())
@@ -259,4 +259,51 @@ copy(Copy_Atom<DefaultCopy, CopyArgs...> const&,
return copy(src, dst);
}
//////////////////////////////////////////
// Special Auto-Vectorizing Overloads
//////////////////////////////////////////
#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
template <class... CT_Args, class... CA_Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...>, CA_Args...> const& atom,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
static_assert(sizeof_bits<SrcType>::value == sizeof_bits<DstType>::value);
static_assert((is_gmem<SrcEngine>::value && is_smem<DstEngine>::value) ||
(is_smem<SrcEngine>::value && is_gmem<DstEngine>::value),
"Bulk Copy only supports gmem -> smem or smem -> gmem movement.");
// Do BulkCopy dispatch
using BULK_COPY_OP = conditional_t<is_gmem<SrcEngine>::value,
SM90_BULK_COPY_G2S,
SM90_BULK_COPY_S2G>;
constexpr int N = decltype(max_common_vector(src, dst))::value;
// Construct a new concrete Atom of the vector size
using N_BITS = Int<N*sizeof_bits<SrcType>::value>;
using COPY_ATOM = Copy_Atom<Copy_Traits<BULK_COPY_OP, N_BITS, CT_Args...>, SrcType>;
auto bulk_atom = apply(atom.opargs_, [&](auto const&... args) { return COPY_ATOM{args...}; });
// Tile the src and dst to the Atom
auto tiler = right_inverse(dst.layout()).compose(Int<N>{});
#if 0
if (thread0()) {
print("copy -- found a max_common_vector of %d\n", N);
print(" "); print(src.data()); print(" o "); print(layout(src)); print("\n");
print(" "); print(dst.data()); print(" o "); print(layout(dst)); print("\n");
}
#endif
return copy(bulk_atom, logical_divide(src, tiler), logical_divide(dst, tiler));
}
#endif // #if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED)
} // end namespace cute

4
include/cute/algorithm/functional.hpp
View File

@@ -30,10 +30,10 @@
**************************************************************************************************/
#pragma once
#include <utility>
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
/** C++14 <functional> extensions */
namespace cute {

70
include/cute/algorithm/gemm.hpp
View File

@@ -32,10 +32,12 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
/** The gemm algorithm takes four (or three) tensors and computes
* D += A * B + C
@@ -281,17 +283,15 @@ gemm(MMA_Atom<MMA> const& mma,
CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN
CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
// REGISTER .reuse OPTIMIZATIONS
auto M = size<1>(A);
auto N = size<1>(B);
// REGISTER .reuse OPTIMIZATIONS
// 64-bit traversal specialization -- serpentine path
if (size<0>(A) * sizeof(typename Tensor<TA,ALayout>::value_type) == 8 &&
size<0>(B) * sizeof(typename Tensor<TB,BLayout>::value_type) == 8)
if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 &&
decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8)
{
#if 1 // NOTE: Must depend on the C-matrix order... (which we can test)
// Row-major iteration
#if 1 // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test)
// Row-major serpentine iteration
CUTE_UNROLL
for (int m = 0; m < M; ++m) {
CUTE_UNROLL
@@ -301,7 +301,7 @@ gemm(MMA_Atom<MMA> const& mma,
}
}
#else
// Col-major iteration
// Col-major serpentine iteration
CUTE_UNROLL
for (int n = 0; n < N; ++n) {
CUTE_UNROLL
@@ -312,13 +312,12 @@ gemm(MMA_Atom<MMA> const& mma,
}
#endif
} else
// 32-bit traversal specialization -- kinked serpentine path
if (size<0>(A) * sizeof(typename Tensor<TA,ALayout>::value_type) == 4 &&
size<0>(B) * sizeof(typename Tensor<TB,BLayout>::value_type) == 4)
if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 &&
decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4)
{
#if 1 // NOTE: Must depend on the C-matrix order... (which we can test)
// Row-major iteration
#if 1 // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test)
// Row-major kinked serpentine iteration
CUTE_UNROLL
for (int m = 0; m < M; m += 2) {
CUTE_UNROLL
@@ -332,7 +331,7 @@ gemm(MMA_Atom<MMA> const& mma,
}
}
#else
// Col-major iteration
// Col-major kinked serpentine iteration
CUTE_UNROLL
for (int n = 0; n < N; n += 2) {
CUTE_UNROLL
@@ -347,9 +346,36 @@ gemm(MMA_Atom<MMA> const& mma,
}
}
#endif
} else {
// Fallback to serpentine loop
// Col-major iteration
} else
// 64-bit + 32-bit traversal order -- keep A (64-bit) in the outer loop and serpentine B
if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 &&
decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4) {
// Row-major serpentine iteration
CUTE_UNROLL
for (int m = 0; m < M; ++m) {
CUTE_UNROLL
for (int n = 0; n < N; ++n) {
int ns = (m & 1) ? N-1-n : n; // Serpentine coordinate
gemm(mma, D(_,m,ns), A(_,m), B(_,ns), C(_,m,ns));
}
}
} else
// 32-bit + 64-bit traversal order -- keep B (64-bit) in the outer loop and serpentine A
if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 &&
decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8) {
// Col-major serpentine iteration
CUTE_UNROLL
for (int n = 0; n < N; ++n) {
CUTE_UNROLL
for (int m = 0; m < M; ++m) {
int ms = (n & 1) ? M-1-m : m; // Serpentine coordinate
gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n));
}
}
} else
// Fallback to serpentine loop
{
// Col-major serpentine iteration
CUTE_UNROLL
for (int n = 0; n < N; ++n) {
CUTE_UNROLL
@@ -504,9 +530,9 @@ gemm(ThrMMA<Args...> const& thr_mma,
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
static_assert(std::is_same_v<std::decay_t<std::invoke_result_t<ALoadTransformOp, TypeA>>, TypeA>,
static_assert(is_same_v<decay_t<invoke_result_t<ALoadTransformOp, TypeA>>, TypeA>,
"ALoadTransformOp functor must accept and return value of type TA::value_type");
static_assert(std::is_same_v<std::decay_t<std::invoke_result_t<BLoadTransformOp, TypeB>>, TypeB>,
static_assert(is_same_v<decay_t<invoke_result_t<BLoadTransformOp, TypeB>>, TypeB>,
"BLoadTransformOp functor must accept and return value of type TB::value_type");
// Original, static size of the problem

2
include/cute/algorithm/prefer.hpp
View File

@@ -34,7 +34,7 @@ namespace cute
{
// Infinite types that inherit from each other
template <std::size_t N>
template <size_t N>
struct prefer : prefer<N-1> {};
template <>

21
include/cute/algorithm/tensor_algorithms.hpp
View File

@@ -99,4 +99,25 @@ transform(Tensor<Engine,Layout>&& tensor, UnaryOp&& op)
return transform(tensor, std::forward<UnaryOp>(op));
}
// Similar to std::transform transforms one tensors and assigns it to another
template <class EngineIn, class LayoutIn, class EngineOut, class LayoutOut, class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn,LayoutIn>& tensor_in, Tensor<EngineOut,LayoutOut>& tensor_out, UnaryOp&& op)
{
CUTE_UNROLL
for (int i = 0; i < size(tensor_in); ++i) {
tensor_out(i) = static_cast<UnaryOp&&>(op)(tensor_in(i));
}
}
// Accept mutable temporaries
template <class EngineIn, class LayoutIn, class EngineOut, class LayoutOut, class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn,LayoutIn>&& tensor_in, Tensor<EngineOut,LayoutOut>&& tensor_out, UnaryOp&& op)
{
return transform(tensor_in, tensor_out, std::forward<UnaryOp>(op));
}
} // end namespace cute

59
include/cute/algorithm/tuple_algorithms.hpp
View File

@@ -32,11 +32,11 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/container/tuple.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/numeric/integer_sequence.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/util/type_traits.hpp>
/** Common algorithms on (hierarchical) tuples */
/** Style choice:
@@ -150,7 +150,7 @@ CUTE_HOST_DEVICE constexpr
auto
for_each_leaf(T&& t, F&& f)
{
if constexpr (is_tuple<std::remove_reference_t<T>>::value) {
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::apply(static_cast<T&&>(t), [&](auto&&... a){ return (for_each_leaf(static_cast<decltype(a)&&>(a), f), ...); }, tuple_seq<T>{});
} else {
return f(static_cast<T&&>(t));
@@ -205,6 +205,20 @@ transform_leaf(T const& t, F&& f)
CUTE_GCC_UNREACHABLE;
}
template <class T0, class T1, class F>
CUTE_HOST_DEVICE constexpr
auto
transform_leaf(T0 const& t0, T1 const& t1, F&& f)
{
if constexpr (is_tuple<T0>::value) {
return transform(t0, t1, [&](auto const& a, auto const& b) { return transform_leaf(a, b, f); });
} else {
return f(t0, t1);
}
CUTE_GCC_UNREACHABLE;
}
//
// find and find_if
//
@@ -258,25 +272,40 @@ find(T const& t, X const& x)
}
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
none_of(T const& t, F&& f)
{
return cute::integral_constant<bool, decltype(find_if(t, f))::value == std::tuple_size<T>::value>{};
if constexpr (is_tuple<T>::value) {
return cute::integral_constant<bool, decltype(find_if(t, f))::value == tuple_size<T>::value>{};
} else {
return not f(t);
}
CUTE_GCC_UNREACHABLE;
}
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
all_of(T const& t, F&& f)
{
auto not_f = [&](auto const& a) { return !f(a); };
return cute::integral_constant<bool, decltype(find_if(t, not_f))::value == std::tuple_size<T>::value>{};
if constexpr (is_tuple<T>::value) {
auto not_f = [&](auto const& a) { return not f(a); };
return cute::integral_constant<bool, decltype(find_if(t, not_f))::value == tuple_size<T>::value>{};
} else {
return f(t);
}
CUTE_GCC_UNREACHABLE;
}
template <class T, class F>
CUTE_HOST_DEVICE constexpr
auto
any_of(T const& t, F&& f)
{
return cute::integral_constant<bool, !decltype(none_of(t, f))::value>{};
return not none_of(t, f);
}
//
@@ -340,7 +369,7 @@ CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f)
{
if constexpr (is_tuple<std::remove_reference_t<T>>::value) {
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::fold(static_cast<T&&>(t),
static_cast<V&&>(v),
f,
@@ -357,11 +386,11 @@ CUTE_HOST_DEVICE constexpr
decltype(auto)
fold_first(T&& t, F&& f)
{
if constexpr (is_tuple<std::remove_reference_t<T>>::value) {
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::fold(static_cast<T&&>(t),
get<0>(static_cast<T&&>(t)),
f,
make_range<1,std::tuple_size<std::remove_reference_t<T>>::value>{});
make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
} else {
return static_cast<T&&>(t);
}
@@ -753,12 +782,12 @@ escan(T const& t, V const& v, F&& f)
namespace detail {
template <int J, class T, int... Is>
template <int J, class... Ts>
CUTE_HOST_DEVICE constexpr
auto
zip_(T const& t, seq<Is...>)
zip_(Ts const&... ts)
{
return cute::make_tuple(get<J>(get<Is>(t))...);
return cute::make_tuple(get<J>(ts)...);
}
template <class T, int... Is, int... Js>
@@ -767,7 +796,7 @@ auto
zip(T const& t, seq<Is...>, seq<Js...>)
{
static_assert(conjunction<bool_constant<tuple_size<tuple_element_t<0,T>>::value == tuple_size<tuple_element_t<Is,T>>::value>...>::value, "Mismatched Ranks");
return cute::make_tuple(detail::zip_<Js>(t, seq<Is...>{})...);
return cute::make_tuple(zip_<Js>(get<Is>(t)...)...);
}
} // end namespace detail
@@ -817,8 +846,8 @@ zip2_by(T const& t, TG const& guide, seq<Is...>, seq<Js...>)
auto split = cute::make_tuple(zip2_by(get<Is>(t), get<Is>(guide))...);
// Rearrange and append missing modes from t to make ((A,B,...),(a,b,...,x,y))
return cute::make_tuple(cute::make_tuple(get<Is,0>(split)...),
cute::make_tuple(get<Is,1>(split)..., get<Js>(t)...));
return cute::make_tuple(cute::make_tuple(get<0>(get<Is>(split))...),
cute::make_tuple(get<1>(get<Is>(split))..., get<Js>(t)...));
}
} // end namespace detail