332 lines
14 KiB
C++
332 lines
14 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2017 - 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/atom/copy_atom.hpp>
|
|
#include <cute/algorithm/copy.hpp>
|
|
|
|
#include <cute/tensor_impl.hpp>
|
|
#include <cute/tensor_predicate.hpp>
|
|
|
|
namespace cute
|
|
{
|
|
|
|
template <uint32_t NumThreads,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE void
|
|
naive_cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> & dst)
|
|
{
|
|
auto N = size(src);
|
|
if (tid < N) {
|
|
uint32_t upper_bound = (N / NumThreads) * NumThreads;
|
|
CUTE_UNROLL
|
|
for (uint32_t i = 0; i < upper_bound; i += NumThreads) { // All in-bounds
|
|
dst[tid + i] = src[tid + i];
|
|
}
|
|
if (N % NumThreads != 0) { // Likely static condition
|
|
uint32_t final_idx = tid + upper_bound;
|
|
if (final_idx < N) { // Final in-bounds
|
|
dst[final_idx] = src[final_idx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Accept mutable temporaries
|
|
template <uint32_t NumThreads,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE void
|
|
naive_cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> && dst)
|
|
{
|
|
return naive_cooperative_copy(tid, src, dst);
|
|
}
|
|
|
|
// A heuristic to determine a "good" permutation of two tensors for later vectorization and thr-assignment
|
|
template <class AEngine, class ALayout,
|
|
class BEngine, class BLayout>
|
|
CUTE_HOST_DEVICE constexpr
|
|
auto
|
|
heuristic_permutation(Tensor<AEngine, ALayout> const& a,
|
|
Tensor<BEngine, BLayout> const& b)
|
|
{
|
|
constexpr bool swizzleA = get_swizzle_t<AEngine>::num_bits != 0 or
|
|
get_swizzle_t<ALayout>::num_bits != 0;
|
|
constexpr bool swizzleB = get_swizzle_t<BEngine>::num_bits != 0 or
|
|
get_swizzle_t<BLayout>::num_bits != 0;
|
|
auto a_inv = right_inverse(get_nonswizzle_portion(a.layout()));
|
|
auto b_inv = right_inverse(get_nonswizzle_portion(b.layout()));
|
|
|
|
constexpr uint8_t scoreA = (uint8_t(swizzleA) << 2) |
|
|
(uint8_t(is_smem<AEngine>::value) << 1) |
|
|
(uint8_t(size(a_inv) > size(b_inv)) << 0);
|
|
|
|
constexpr uint8_t scoreB = (uint8_t(swizzleB) << 2) |
|
|
(uint8_t(is_smem<BEngine>::value) << 1) |
|
|
(uint8_t(size(b_inv) > size(a_inv)) << 0);
|
|
|
|
if constexpr (scoreA >= scoreB) {
|
|
return a_inv;
|
|
} else {
|
|
return b_inv;
|
|
}
|
|
}
|
|
|
|
// cooperative_copy<NumThreads, MaxVecBits>(thr_idx, src, dst)
|
|
// Use NumThreads to copy Tensor src to Tensor dst with element-wise vectorization up to MaxVecBits.
|
|
// @pre 0 <= @a tid < NumThreads
|
|
// @pre Tensors @a src and @a dst are aligned up to MaxVecBits.
|
|
// That is, pointers and dynamic strides are assumed to be aligned up to MaxVecBits.
|
|
//
|
|
template <uint32_t NumThreads, uint32_t MaxVecBits,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE
|
|
void
|
|
cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> & dst)
|
|
{
|
|
// Assumes the shapes are static, can generalize/fallback
|
|
CUTE_STATIC_ASSERT_V(is_static<decltype(shape(src))>{} && is_static<decltype(shape(dst))>{});
|
|
CUTE_STATIC_ASSERT_V(size(src) == size(dst));
|
|
// Assumes the types are the same, can generalize/fallback
|
|
static_assert(cute::is_same<typename SrcEngine::value_type, typename DstEngine::value_type>::value);
|
|
static_assert(MaxVecBits == sizeof_bits_v<typename SrcEngine::value_type> ||
|
|
MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128,
|
|
"Expected MaxVecBits to be value size or 8 or 16 or 32 or 64 or 128 for alignment and performance.");
|
|
// Check that the tensors are likely shared across threads: either gmem or smem
|
|
static_assert((is_gmem<SrcEngine>::value || is_smem<SrcEngine>::value),
|
|
"cooperative_copy expects shared gmem or smem source tensor.");
|
|
static_assert((is_gmem<DstEngine>::value || is_smem<DstEngine>::value),
|
|
"cooperative_copy expects shared gmem or smem destination tensor.");
|
|
// Precondition on tid in DEBUG
|
|
assert(tid < NumThreads);
|
|
// Precondition on pointer alignment in DEBUG
|
|
assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(src.data())));
|
|
assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(dst.data())));
|
|
|
|
#if 0
|
|
if (thread0()) {
|
|
print(" "); print("cooperative_copy\n");
|
|
print(" "); print("NumThreads: "); print(NumThreads); print("\n");
|
|
print(" "); print("MaxVecBits: "); print(MaxVecBits); print("\n");
|
|
print(" "); print("src: "); print(src); print("\n");
|
|
print(" "); print("dst: "); print(dst); print("\n");
|
|
}
|
|
#ifdef __CUDA_ARCH__
|
|
__syncthreads();
|
|
#endif
|
|
#endif
|
|
|
|
// The common layout of the two tensors that can be vectorized over elements and threads
|
|
// vidx -> coord
|
|
auto common_layout = heuristic_permutation(src, dst);
|
|
|
|
// Apply
|
|
// (V, rest)
|
|
Tensor src_a = coalesce(logical_divide(src, common_layout), Shape<_1,_1>{});
|
|
Tensor dst_a = coalesce(logical_divide(dst, common_layout), Shape<_1,_1>{});
|
|
|
|
//
|
|
// Determine vectorization of elems and thrs based on src/dst size and number of threads
|
|
// NOTE: This heuristic promotes parallelization over vectorization
|
|
//
|
|
|
|
// The number of elements and number of bits
|
|
constexpr int elem_bits = sizeof_bits_v<typename SrcEngine::value_type>;
|
|
constexpr int total_elem = size(SrcLayout{});
|
|
|
|
// The number of elements that can be vectorized in values
|
|
constexpr int common_elem = decltype(max_common_vector(src_a, dst_a))::value;
|
|
|
|
#if 0
|
|
if (thread0()) {
|
|
print(" "); print("common_layout: "); print(common_layout); print("\n");
|
|
print(" "); print("src_a: "); print(src_a); print("\n");
|
|
print(" "); print("dst_a: "); print(dst_a); print("\n");
|
|
}
|
|
#ifdef __CUDA_ARCH__
|
|
__syncthreads();
|
|
#endif
|
|
#endif
|
|
|
|
//
|
|
if constexpr (total_elem % NumThreads != 0) {
|
|
// Not attempting to find a partitioning pattern, fallback to dynamically indexed slowpath
|
|
|
|
if constexpr (common_elem > 1 && MaxVecBits > elem_bits) {
|
|
// If the vectorization is non-trivial and divides the maximum vectorizations, then vectorize
|
|
constexpr auto max_align_src = elem_bits * decltype(max_alignment(src_a.layout()))::value;
|
|
constexpr auto max_align_dst = elem_bits * decltype(max_alignment(dst_a.layout()))::value;
|
|
constexpr auto vec_bits = gcd(max_align_src, max_align_dst, MaxVecBits);
|
|
using VecType = uint_bit_t<vec_bits>;
|
|
|
|
static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
|
|
static_assert((vec_bits >= 8), "No support for subbyte copying");
|
|
|
|
Tensor src_v = recast<VecType const>(src_a);
|
|
Tensor dst_v = recast<VecType >(dst_a);
|
|
|
|
#if 0
|
|
if (thread0()) {
|
|
print(" "); print("cooperative_copy -- naive\n");
|
|
print(" "); print("src_v: "); print(src_v); print("\n");
|
|
print(" "); print("dst_v: "); print(dst_v); print("\n");
|
|
}
|
|
#ifdef __CUDA_ARCH__
|
|
__syncthreads();
|
|
#endif
|
|
#endif
|
|
|
|
naive_cooperative_copy<NumThreads>(tid, src_v, dst_v);
|
|
} else {
|
|
naive_cooperative_copy<NumThreads>(tid, src_a, dst_a);
|
|
}
|
|
} else {
|
|
// If the tensors can be equally partitioned by the threads,
|
|
// compute vectorization widths in elements and threads.
|
|
|
|
// If there are too many threads to allow a full vectorized copy, trunc the vectorization
|
|
constexpr int total_bits = total_elem * elem_bits;
|
|
constexpr int max_bits_per_thr = total_bits / NumThreads;
|
|
// At least elem_bits, at most common_bits
|
|
constexpr int common_bits = common_elem * elem_bits;
|
|
constexpr int vec_bits = cute::max(elem_bits, cute::gcd(common_bits, int(MaxVecBits), max_bits_per_thr));
|
|
|
|
// Should account for vec_bits < 8 and/or vec_elem <= 1
|
|
// And also account for subbyte types, which could cause race conditions
|
|
// Want to ENFORCE sufficient vectorization in those cases
|
|
static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
|
|
static_assert(vec_bits >= 8, "No support for subbyte copying");
|
|
|
|
using VecType = uint_bit_t<vec_bits>;
|
|
constexpr int vec_elem = vec_bits / elem_bits;
|
|
|
|
constexpr int vec_thrs = cute::min(int(NumThreads), total_elem / vec_elem);
|
|
|
|
//
|
|
// Determine the partitioning patterns for the vec_elems and vec_thrs
|
|
//
|
|
|
|
// Distribute the rest of the V*T to some consistent portion outside of the common_layout, if needed
|
|
auto common_domain_src = domain_distribute(shape(src_a), Int<vec_elem*vec_thrs>{});
|
|
auto common_domain_dst = domain_distribute(shape(dst_a), Int<vec_elem*vec_thrs>{});
|
|
|
|
// Make sure for now, could fall back here instead
|
|
CUTE_STATIC_ASSERT_V(size(common_domain_src) == Int<vec_elem*vec_thrs>{});
|
|
CUTE_STATIC_ASSERT_V(compatible(common_domain_src, common_domain_dst) ||
|
|
compatible(common_domain_dst, common_domain_src));
|
|
// Use the "more specific" domain for the extra elements of V*T
|
|
auto common_domain = conditional_return(compatible(common_domain_src, common_domain_dst),
|
|
common_domain_dst, common_domain_src);
|
|
|
|
// Construct the tiler
|
|
auto tiler_vt = common_domain.with_shape(Int<vec_elem>{}, Int<vec_thrs>{});
|
|
|
|
// Apply and slice
|
|
Tensor src_v = logical_divide(src_a, tiler_vt)(make_coord(_,tid),_);
|
|
Tensor dst_v = logical_divide(dst_a, tiler_vt)(make_coord(_,tid),_);
|
|
|
|
#if 0
|
|
if (thread0()) {
|
|
print(" "); print("cooperative_copy -- vec\n");
|
|
print(" "); print("Used vector: "); print(vec_elem); print("\n");
|
|
print(" "); print("Used threads: "); print(vec_thrs); print("\n");
|
|
print(" "); print("tiler_vt: "); print(tiler_vt); print("\n");
|
|
print(" "); print("src_v: "); print(src_v); print("\n");
|
|
print(" "); print("dst_v: "); print(dst_v); print("\n");
|
|
print(" "); print("recast<VecType const>(src_v): "); print(recast<VecType const>(src_v)); print("\n");
|
|
print(" "); print("recast<VecType >(dst_v): "); print(recast<VecType >(dst_v)); print("\n");
|
|
}
|
|
#ifdef __CUDA_ARCH__
|
|
__syncthreads();
|
|
#endif
|
|
#endif
|
|
|
|
// If we're using all threads (static) or the tid is in-range (dynamic)
|
|
if (vec_thrs == NumThreads or tid < vec_thrs) {
|
|
return copy_if(TrivialPredTensor{}, recast<VecType const>(src_v), recast<VecType>(dst_v));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Default max-vectorization size to value_type size
|
|
template <uint32_t NumThreads,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE
|
|
void
|
|
cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> & dst)
|
|
{
|
|
constexpr uint32_t MaxVecBits = sizeof_bits_v<typename SrcEngine::value_type>;
|
|
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
|
|
}
|
|
|
|
//
|
|
// Accept mutable temporaries
|
|
//
|
|
|
|
template <uint32_t NumThreads,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE
|
|
void
|
|
cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> && dst)
|
|
{
|
|
return cooperative_copy<NumThreads>(tid, src, dst);
|
|
}
|
|
|
|
template <uint32_t NumThreads, uint32_t MaxVecBits,
|
|
class SrcEngine, class SrcLayout,
|
|
class DstEngine, class DstLayout>
|
|
CUTE_HOST_DEVICE
|
|
void
|
|
cooperative_copy(uint32_t const& tid,
|
|
Tensor<SrcEngine, SrcLayout> const& src,
|
|
Tensor<DstEngine, DstLayout> && dst)
|
|
{
|
|
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
|
|
}
|
|
|
|
} // end namespace cute
|