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

Updates for CUTLASS 3.5.0 (#1468)

Browse Source
This commit is contained in:
Vijay Thakkar
2024-04-11 21:33:40 -04:00
committed by GitHub
parent a40e08e9d5
commit 7d49e6c7e2
171 changed files with 7526 additions and 1888 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/arch/copy_sm50.hpp
+72
View File
@@ -0,0 +1,72 @@
/***************************************************************************************************
* Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/copy.hpp>
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 500
#define CUTE_ARCH_WARP_SHUFFLE_ENABLED 1
#endif
namespace cute
{
struct SM50_Shuffle_U32_2x2Trans
{
using SRegisters = uint32_t[2];
using DRegisters = uint32_t[2];
CUTE_HOST_DEVICE static void
copy(uint32_t const& src0, uint32_t const& src1, uint32_t& dst0, uint32_t& dst1)
{
#if defined(CUTE_ARCH_WARP_SHUFFLE_ENABLED)
uint32_t x0 = src0;
uint32_t y0 = __shfl_xor_sync(0xffffffff, x0, 1);
uint32_t x1 = src1;
uint32_t y1 = __shfl_xor_sync(0xffffffff, x1, 1);
if (threadIdx.x % 2 == 0) {
dst1 = y0;
}
else {
dst0 = y1;
}
#else
CUTE_INVALID_CONTROL_PATH("Trying to use __shfl_xor_sync without CUTE_ARCH_WARP_SHUFFLE_ENABLED.");
#endif
}
};
} // end namespace cute
include/cute/arch/util.hpp
+102 -59
View File
@@ -117,7 +117,7 @@ cast_smem_ptr_to_uint(void const* const ptr)
uint32_t smem_ptr;
asm(
"{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\n"
"{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }\n"
: "=r"(smem_ptr) : "l"(ptr));
return smem_ptr;
@@ -132,11 +132,47 @@ cast_smem_ptr_to_uint(void const* const ptr)
#endif
}
namespace detail {
//
// Utility for pointer interfaces
// Wrapper for MMAOp::fma
//
namespace detail {
template <class MmaOp>
struct CallFMA {
template <class... Args>
CUTE_HOST_DEVICE constexpr void
operator()(Args&&... args) const {
return MmaOp::fma(static_cast<Args&&>(args)...);
}
};
//
// Wrapper for CopyOp::copy
//
template <class CopyOp>
struct CallCOPY {
template <class... Args>
CUTE_HOST_DEVICE constexpr void
operator()(Args&&... args) const {
return CopyOp::copy(static_cast<Args&&>(args)...);
}
};
//
// Utility for exploding pointers/arrays/tensors into functions
//
template <class Fn,
class PtrA, int... I>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn,
PtrA&& a, int_sequence<I...>)
{
return fn(a[I]...);
}
template <class Fn,
class PtrS, int... Is,
@@ -180,76 +216,83 @@ explode(Fn fn,
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]...);
}
template <class Fn,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class ParamType>
CUTE_HOST_DEVICE constexpr
void
explode_with_d_scaling(Fn fn,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
ParamType&& p0)
{
return fn(a[Ia]..., b[Ib]..., c[Ic]..., p0);
}
template <class Fn,
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class ParamType>
class PtrE, int... Ie>
CUTE_HOST_DEVICE constexpr
void
explode_with_d_scaling(Fn fn,
explode(Fn fn,
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
ParamType&& p0)
PtrE&& e, int_sequence<Ie...>)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., p0);
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]...);
}
template <class Fn,
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class PtrSFA, int... Isfa,
class PtrSFB, int... Isfb>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn,
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
PtrSFA&& sfa, int_sequence<Isfa...>,
PtrSFB&& sfb, int_sequence<Isfb...>)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., sfa[Isfa]..., sfb[Isfb]...);
}
//
// Utility for exploding tuples into functions
//
template <class Fn,
class TupleA, int... I>
CUTE_HOST_DEVICE constexpr
void
explode_tuple(Fn fn,
TupleA&& a, int_sequence<I...>)
{
return fn(get<I>(a)...);
}
template <class Fn,
class TupleA, int... Ia,
class TupleB, int... Ib>
CUTE_HOST_DEVICE constexpr
void
explode_tuple(Fn fn,
TupleA&& a, int_sequence<Ia...>,
TupleB&& b, int_sequence<Ib...>)
{
return fn(get<Ia>(a)..., get<Ib>(b)...);
}
template <class Fn,
class TupleA, int... Ia,
class TupleB, int... Ib,
class TupleC, int... Ic>
CUTE_HOST_DEVICE constexpr
void
explode_tuple(Fn fn,
TupleA&& a, int_sequence<Ia...>,
TupleB&& b, int_sequence<Ib...>,
TupleC&& c, int_sequence<Ic...>)
{
return fn(get<Ia>(a)..., get<Ib>(b)..., get<Ic>(c)...);
}
} // end namespace detail
template <int SRegCount, int DRegCount,
class Fn, class PtrS, class PtrD>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn, PtrS&& s, PtrD&& d)
{
return detail::explode(fn,
s, make_int_sequence<SRegCount>{},
d, make_int_sequence<DRegCount>{});
}
template <int ARegCount, int BRegCount, int CRegCount,
class Fn, class PtrA, class PtrB, class PtrC>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn, PtrA&& a, PtrB&& b, PtrC&& c)
{
return detail::explode(fn,
a, make_int_sequence<ARegCount>{},
b, make_int_sequence<BRegCount>{},
c, make_int_sequence<CRegCount>{});
}
template <int DRegCount, int ARegCount, int BRegCount, int CRegCount,
class Fn, class PtrD, class PtrA, class PtrB, class PtrC>
CUTE_HOST_DEVICE constexpr
void
explode(Fn fn, PtrD&& d, PtrA&& a, PtrB&& b, PtrC&& c)
{
return detail::explode(fn,
d, make_int_sequence<DRegCount>{},
a, make_int_sequence<ARegCount>{},
b, make_int_sequence<BRegCount>{},
c, make_int_sequence<CRegCount>{});
}
} // end namespace cute