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

CUTLASS 3.5.1 (#1623)

Browse Source 
* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/axpby.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
namespace cute
include/cute/algorithm/clear.hpp
+1 -3
View File
@@ -31,9 +31,7 @@
#pragma once
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/fill.hpp>
namespace cute
include/cute/algorithm/cooperative_copy.hpp
+236 -119
View File
@@ -1,51 +1,117 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
* 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.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 src to dst with element vectorization up to MaxVecBits.
// 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,
@@ -56,121 +122,171 @@ cooperative_copy(uint32_t const& tid,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
// Assumes the shapes are static, can generalize
// 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
static_assert(sizeof_bits_v<typename SrcEngine::value_type> == sizeof_bits_v<typename DstEngine::value_type>);
// 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.");
"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.");
"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())));
// Fallback - slow path, naive copy, vectorization disabled
if constexpr(size(SrcLayout{}) % NumThreads != 0) {
int index = static_cast<int>(tid);
CUTE_UNROLL
for(int i = 0; i < ceil_div(size(SrcLayout{}), NumThreads); i++) {
if(index < size(SrcLayout{})) {
dst[index] = src[index];
#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");
}
index += NumThreads;
#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 {
// Fast path with vectorization
// If the tensors can be equally partitioned by the threads,
// compute vectorization widths in elements and threads.
// Precondition on pointer alignment in DEBUG
assert(is_byte_aligned<max(MaxVecBits/8, 1u)>(raw_pointer_cast(src.data())));
assert(is_byte_aligned<max(MaxVecBits/8, 1u)>(raw_pointer_cast(dst.data())));
constexpr int elem_bits = sizeof_bits_v<typename SrcEngine::value_type>;
//
// Determine val+thr vectorization based on src/dst size and number of threads
// NOTE: This heuristic promotes parallelization over vectorization
//
// The number of elements that can be vectorized in values
constexpr int common_elem = decltype(max_common_vector(src, dst))::value;
constexpr int common_bits = common_elem * elem_bits;
constexpr int total_elem = decltype(size(src))::value;
// If there are too many threads to allow a full vectorized copy, trunc the vectorization
constexpr int total_bits = total_elem * elem_bits;
static_assert(total_bits % NumThreads == 0);
constexpr int total_bits_per_thr = total_bits / NumThreads;
// If there are too many threads to allow a full elem copy, trunc the thrs and use elem_bits
constexpr int max_vec_bits_by_thr = cute::max(elem_bits, total_bits_per_thr);
// Cap the vectorization to the common bits, the max_vec_bits_by_thr, and the MaxVecBits
constexpr int vec_bits = cute::min(common_bits, max_vec_bits_by_thr, static_cast<int>(MaxVecBits));
// Convert back to number of elements, safe_div
static_assert((vec_bits % elem_bits) == 0);
constexpr int vec_elem = vec_bits / elem_bits;
// Use only part of threads if there's not enough work for all threads
constexpr int vec_thrs = (total_elem % (vec_elem * NumThreads) == 0)
? NumThreads
: (total_elem / vec_elem);
static_assert(vec_thrs <= NumThreads);
// The common layout of the two tensors that can be vectorized over threads
// vidx -> coord
auto common_layout = max_common_layout(get_nonswizzle_portion(src.layout()),
get_nonswizzle_portion(dst.layout()));
// Scale up the common_layout to cover the entire tensors
// vidx -> coord
auto full_perm = tile_to_shape(make_layout(common_layout), size(src));
// Create the Tiler
// ((vid,tid),iter)
auto layout_vt = logical_divide(full_perm, Layout<Shape<Int<vec_elem>, Int<vec_thrs>>>{});
// Apply and slice
Tensor src_v = src.compose(layout_vt)(make_coord(_,tid),_);
Tensor dst_v = dst.compose(layout_vt)(make_coord(_,tid),_);
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 >= 8), "No support for subbyte copying");
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("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");
print(" "); print("common_layout: "); print(common_layout); print("\n");
print(" "); print("full_perm: "); print(full_perm); print("\n");
print(" "); print("Used vector: "); print(vec_elem); print("\n");
print(" "); print("Used threads: "); print(vec_thrs); print("\n");
print(" "); print("layout_vt: "); print(layout_vt); print("\n");
print(" "); print("src.compose(layout_vt): "); print(src.compose(layout_vt)); print("\n");
print(" "); print("dst.compose(layout_vt): "); print(dst.compose(layout_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 const>(dst_v): "); print(recast<VecType const>(dst_v)); print("\n");
}
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();
__syncthreads();
#endif
#endif
// If we're using all threads (static) or the tid is in in-range (dynamic)
if (vec_thrs >= NumThreads or tid < vec_thrs) {
// 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>
@@ -184,7 +300,10 @@ cooperative_copy(uint32_t const& tid,
return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst);
}
//
// Accept mutable temporaries
//
template <uint32_t NumThreads,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
@@ -197,9 +316,7 @@ cooperative_copy(uint32_t const& tid,
return cooperative_copy<NumThreads>(tid, src, dst);
}
// Accept mutable temporaries
template <uint32_t NumThreads,
uint32_t MaxVecBits,
template <uint32_t NumThreads, uint32_t MaxVecBits,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
include/cute/algorithm/cooperative_gemm.hpp
+8 -25
View File
@@ -39,7 +39,7 @@
#include <cute/algorithm/functional.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -76,29 +76,15 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
using TypeB = typename TB::value_type;
using TypeC = typename TC::value_type;
// Original, static size of the problem
auto M = size<0>(sC);
auto N = size<1>(sC);
auto K = size<1>(sA);
// Block size of the compute tile
auto BLK_M = tile_size<0>(thr_mma);
auto BLK_N = tile_size<1>(thr_mma);
auto BLK_K = tile_size<2>(thr_mma);
//
// MMA Partitioning
//
// Round the layout extents up to BLK_X to satisfy MMA partitioning safety
Tensor rounded_sA = sA.compose(make_shape(round_up(M, BLK_M), round_up(K, BLK_K)));
Tensor rounded_sB = sB.compose(make_shape(round_up(N, BLK_N), round_up(K, BLK_K)));
Tensor rounded_sC = sC.compose(make_shape(round_up(M, BLK_M), round_up(N, BLK_N)));
// Partition the sA, sB, and sC tiles across the threads for the MMA
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K)
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K)
Tensor tCsC = thr_mma.partition_C(sC); // (MMA,MMA_M,MMA_N)
// Partition the sA and sB tiles across the threads for the MMA
Tensor tCsA = thr_mma.partition_A(rounded_sA); // (MMA,MMA_M,MMA_K)
Tensor tCsB = thr_mma.partition_B(rounded_sB); // (MMA,MMA_N,MMA_K)
Tensor tCsC = thr_mma.partition_C(rounded_sC); // (MMA,MMA_M,MMA_N)
// Create register tensors for the MMA to operate on
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K)
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K)
@@ -109,9 +95,6 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
print(" sA: "); print( sA); print("\n");
print(" sB: "); print( sB); print("\n");
print(" sC: "); print( sC); print("\n");
print("r_sA: "); print(rounded_sA); print("\n");
print("r_sB: "); print(rounded_sB); print("\n");
print("r_sC: "); print(rounded_sC); print("\n");
print(thr_mma);
print("tCsA: "); print(tCsA); print("\n");
print("tCsB: "); print(tCsB); print("\n");
@@ -127,8 +110,8 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
//
// Create coordinate tensors for the problem
Tensor cA = make_identity_tensor(shape(rounded_sA)); // (M,K) -> (m,k)
Tensor cB = make_identity_tensor(shape(rounded_sB)); // (N,K) -> (n,k)
Tensor cA = make_identity_tensor(shape(sA)); // (M,K) -> (m,k)
Tensor cB = make_identity_tensor(shape(sB)); // (N,K) -> (n,k)
// Repeat partitioning with thr_mma
Tensor tCcA = thr_mma.partition_A(cA); // (MMA,MMA_M,MMA_K) -> (m,k)
@@ -222,7 +205,7 @@ cooperative_gemm_predication(ThrMMA<Args...> const& thr_mma,
//
// Create coordinate tensors for the problem
Tensor cC = make_identity_tensor(shape(rounded_sC)); // (M,N) -> (m,n)
Tensor cC = make_identity_tensor(shape(sC)); // (M,N) -> (m,n)
// Repeat partitioning with thr_mma
Tensor tCcC = thr_mma.partition_C(cC); // (MMA,MMA_M,MMA_N) -> (m,n)
include/cute/algorithm/copy.hpp
+28 -35
View File
@@ -34,7 +34,7 @@
#include <cute/container/alignment.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
#include <cute/atom/copy_atom.hpp>
@@ -199,14 +199,14 @@ copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
{
static_assert(sizeof_bits_v<VecType> >= 8 && sizeof_bits_v<VecType> % 8 == 0,
"Expected a vectorization type of at least a byte.");
using SrcType = typename SrcEngine::element_type;
using DstType = typename DstEngine::element_type;
if constexpr (sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
if constexpr (cute::is_same<SrcType, DstType>::value &&
sizeof_bits_v<VecType> > sizeof_bits_v<DstType>)
{
// Preserve volatility of Src/Dst types.
using SrcVecType = conditional_t<is_volatile_v<SrcType>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<DstType>, VecType volatile, VecType >;
using SrcVecType = conditional_t<is_volatile_v<typename SrcEngine::element_type>, VecType const volatile, VecType const>;
using DstVecType = conditional_t<is_volatile_v<typename DstEngine::element_type>, VecType volatile, VecType >;
Tensor src_v = recast<SrcVecType>(src);
Tensor dst_v = recast<DstVecType>(dst);
@@ -264,22 +264,22 @@ copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
{
constexpr int vec_elem = decltype(max_common_vector(src, dst))::value;
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
// When layouts are static, accept vec_bits up to 128
// When layouts are dynamic, accept vec_bits up to MaxVecBits
constexpr int vec_bits = (is_static<SrcLayout>::value && is_static<DstLayout>::value) ?
cute::min(vec_elem * src_bits, 128) :
cute::min(vec_elem * src_bits, MaxVecBits);
constexpr int max_align_src = decltype(max_alignment(src.layout()))::value;
constexpr int max_align_dst = decltype(max_alignment(dst.layout()))::value;
constexpr int max_align = gcd(vec_elem, max_align_src, max_align_dst);
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
constexpr int vec_bits = gcd(src_bits * max_align, MaxVecBits);
if constexpr (vec_elem > 1 && vec_bits >= 8) {
// If more than one element vectorizes to 8bits or more, then copy_vec
#if 0
if (thread0()) {
print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
print(" "); print(src); print("\n");
print(" "); print(dst); print("\n");
}
#endif
return copy_vec<uint_bit_t<vec_bits>>(src, dst);
} else {
return copy_if(TrivialPredTensor{}, src, dst);
@@ -294,10 +294,16 @@ void
copy(Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
if constexpr (is_static<SrcLayout>::value && is_static<DstLayout>::value) {
// Assume Tensors with static layouts (e.g. registers) have pointers that are 128b aligned
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
} else {
// Do not assume that dynamic layouts are aligned.
return copy(AutoVectorizingCopyWithAssumedAlignment<8>{}, src, dst);
}
}
// Auto-vectorizing copy with assumed alignment of dynamic layout strides up to 128bit.
// Auto-vectorizing copy with assumed alignment up to 128bit.
template <class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
@@ -308,19 +314,6 @@ copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopy
template <class... Args,
class SrcEngine, class SrcLayout,
class DstEngine, class DstLayout>
CUTE_HOST_DEVICE
void
copy(Copy_Atom<AutoVectorizingCopy, Args...> const&,
Tensor<SrcEngine, SrcLayout> const& src,
Tensor<DstEngine, DstLayout> & dst)
{
return copy(AutoVectorizingCopy{}, src, dst);
}
// Specializaton for Atom AutoVectorizingCopyAssumedAlignment
template <int MaxVecBits, class... Args,
class SrcEngine, class SrcLayout,
@@ -346,7 +339,7 @@ copy(Copy_Traits<SM90_BULK_COPY_AUTO, CT_Args...> const& atom, // Copy_Traits m
{
using SrcType = typename SrcEngine::value_type;
using DstType = typename DstEngine::value_type;
static_assert(sizeof_bits<SrcType>::value == sizeof_bits<DstType>::value);
static_assert(cute::is_same<SrcType, 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.");
include/cute/algorithm/fill.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/prefer.hpp>
namespace cute
include/cute/algorithm/gemm.hpp
+1 -1
View File
@@ -35,7 +35,7 @@
#include <cute/util/type_traits.hpp>
#include <cute/algorithm/functional.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/atom/mma_atom.hpp>
include/cute/algorithm/prefetch.hpp
+1 -7
View File
@@ -32,7 +32,7 @@
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/atom/copy_atom.hpp>
@@ -90,12 +90,6 @@ constexpr bool has_prefetch = false;
template <class CopyOp>
constexpr bool has_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = true;
template <class CopyOp, class = void>
constexpr bool is_prefetch = false;
template <class CopyOp>
constexpr bool is_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = is_same_v<CopyOp, typename CopyOp::PREFETCH>;
} // end namespace detail
template <class CopyOp, class... CT_Args, class... CA_Args,
include/cute/algorithm/tensor_algorithms.hpp
+13 -14
View File
@@ -33,8 +33,7 @@
#pragma once
#include <cute/config.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -100,13 +99,13 @@ transform(Tensor<Engine,Layout>&& tensor, UnaryOp&& op)
}
// Similar to std::transform transforms one tensors and assigns it to another
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> & tensor_out,
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> & tensor_out,
UnaryOp&& op)
{
CUTE_UNROLL
@@ -117,30 +116,30 @@ transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
// Accept mutable temporaries
template <class EngineIn, class LayoutIn,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class UnaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> && tensor_out,
transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
Tensor<EngineOut,LayoutOut> && tensor_out,
UnaryOp&& op)
{
return transform(tensor_in, tensor_out, op);
}
// Similar to std::transform with a binary operation
// Takes two tensors as input and one tensor as output.
// Takes two tensors as input and one tensor as output.
// Applies the binary_op to tensor_in1 and tensor_in2 and
// assigns it to tensor_out
template <class EngineIn1, class LayoutIn1,
class EngineIn2, class LayoutIn2,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class BinaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
Tensor<EngineOut,LayoutOut> & tensor_out,
Tensor<EngineOut,LayoutOut> & tensor_out,
BinaryOp&& op)
{
CUTE_UNROLL
@@ -152,11 +151,11 @@ transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
// Accept mutable temporaries
template <class EngineIn1, class LayoutIn1,
class EngineIn2, class LayoutIn2,
class EngineOut, class LayoutOut,
class EngineOut, class LayoutOut,
class BinaryOp>
CUTE_HOST_DEVICE constexpr
void
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
Tensor<EngineOut,LayoutOut> && tensor_out,
BinaryOp&& op)
include/cute/algorithm/tuple_algorithms.hpp
+53 -21
View File
@@ -404,29 +404,54 @@ namespace detail {
// This impl compiles much faster than cute::apply and variadic args
template <class T, class V, class F>
CUTE_HOST_DEVICE constexpr
decltype(auto)
fold(T&& t, V&& v, F&& f, seq<>)
auto
fold(T&&, V&& v, F&&, seq<>)
{
return static_cast<V&&>(v);
return v;
}
template <class T, class V, class F, int I, int... Is>
template <class T, class V, class F, int I0>
CUTE_HOST_DEVICE constexpr
decltype(auto)
fold(T&& t, V&& v, F&& f, seq<I,Is...>)
auto
fold(T&& t, V&& v, F&& f, seq<I0>)
{
if constexpr (sizeof...(Is) == 0) {
return f(static_cast<V&&>(v), get<I>(static_cast<T&&>(t)));
} else {
return fold(static_cast<T&&>(t),
f(static_cast<V&&>(v), get<I>(static_cast<T&&>(t))),
f,
seq<Is...>{});
}
CUTE_GCC_UNREACHABLE;
return f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1>)
{
return f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2>)
{
return f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2, int I3>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2,I3>)
{
return f(f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t))), get<I3>(static_cast<T&&>(t)));
}
template <class T, class V, class F, int I0, int I1, int I2, int I3, int... Is>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0,I1,I2,I3,Is...>)
{
return fold(static_cast<T&&>(t),
f(f(f(f(static_cast<V&&>(v), get<I0>(static_cast<T&&>(t))), get<I1>(static_cast<T&&>(t))), get<I2>(static_cast<T&&>(t))), get<I3>(static_cast<T&&>(t))),
f,
seq<Is...>{});
}
} // end namespace detail
template <class T, class V, class F>
@@ -448,7 +473,7 @@ fold(T&& t, V&& v, F&& f)
template <class T, class F>
CUTE_HOST_DEVICE constexpr
decltype(auto)
auto
fold_first(T&& t, F&& f)
{
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
@@ -457,7 +482,7 @@ fold_first(T&& t, F&& f)
f,
make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
} else {
return static_cast<T&&>(t);
return t;
}
CUTE_GCC_UNREACHABLE;
@@ -701,7 +726,14 @@ CUTE_HOST_DEVICE constexpr
auto
replace(T const& t, X const& x)
{
return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N+1,tuple_size<T>::value>{});
if constexpr (is_tuple<T>::value) {
return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N+1,tuple_size<T>::value>{});
} else {
static_assert(N == 0);
return x;
}
CUTE_GCC_UNREACHABLE;
}
// Replace the first element of the tuple with x
@@ -1077,9 +1109,9 @@ zip2_by(T const& t, TG const& guide)
/// @return A tuple of the elements of @c t in reverse order.
template <class T>
CUTE_HOST_DEVICE constexpr
CUTE_HOST_DEVICE constexpr
auto
reverse(T const& t)
reverse(T const& t)
{
if constexpr (is_tuple<T>::value) {
return detail::apply(t, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_rseq<T>{});
include/cute/arch/copy.hpp
+8 -6
View File
@@ -68,7 +68,7 @@ struct UniversalCopy
//
// Placeholder for the copy algorithm's stronger auto-vectorizing behavior
// that assumes alignment of dynamic layouts up to MaxVecBits
// that assumes alignment of pointers and dynamic layouts up to MaxVecBits
//
template <int MaxVecBits = 128>
@@ -80,15 +80,17 @@ struct AutoVectorizingCopyWithAssumedAlignment
};
//
// Placeholder for the copy algorithm's default auto-vectorizing behavior
// that does not assume alignment of dynamic layouts
// AutoVectorizingCopy alias assumes maximal alignment of pointers and dynamic strides.
// If this is not the case then AutoVectorizingCopyWithAssumedAlignment should be used instead
//
using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<128>;
// Alias
using DefaultCopy = AutoVectorizingCopy;
//
// DefaultCopy alias does not assume alignment of pointers or dynamic strides.
//
using DefaultCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
//
// Global memory prefetch into L2
include/cute/arch/copy_sm90_desc.hpp
+91 -18
View File
@@ -95,8 +95,8 @@ wait_barrier(uint64_t& smem_barrier, // 64 bits user-mange
".reg .pred P1;\n"
"LAB_WAIT:\n"
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
"@P1 bra.uni DONE;\n"
"bra.uni LAB_WAIT;\n"
"@P1 bra DONE;\n"
"bra LAB_WAIT;\n"
"DONE:\n"
"}\n"
:: "r"(smem_int_ptr),
@@ -134,6 +134,48 @@ enum class SmemSwizzleBits : uint8_t {
B128 = 3,
};
enum class OOBFill : uint8_t {
ZERO = 0,
CONSTANT = 1,
};
CUTE_HOST_DEVICE char const* to_string(OOBFill const& t) {
switch (t) {
case OOBFill::ZERO: return "ZERO";
case OOBFill::CONSTANT: return "CONSTANT";
}
return nullptr;
}
enum class L2Promotion : uint8_t {
DISABLE = 0,
B64 = 1,
B128 = 2,
B256 = 3,
};
CUTE_HOST_DEVICE char const* to_string(L2Promotion const& t) {
switch (t) {
case L2Promotion::DISABLE: return "DISABLE";
case L2Promotion::B64: return "B64";
case L2Promotion::B128: return "B128";
case L2Promotion::B256: return "B256";
}
return nullptr;
}
// Aux parameters which are independent with the problem size
struct DescriptorAuxParams {
OOBFill oobfill_ = OOBFill::ZERO;
L2Promotion l2promo_ = L2Promotion::DISABLE;
};
enum class CacheHintSm90 : uint64_t {
EVICT_NORMAL = 0x1000000000000000,
EVICT_FIRST = 0x12F0000000000000,
EVICT_LAST = 0x14F0000000000000,
};
#if (__CUDACC_VER_MAJOR__ >= 12)
#if !defined(__CUDACC_RTC__)
@@ -168,6 +210,27 @@ to_CUtensorMapSwizzle(SmemSwizzleBits const& t) {
case SmemSwizzleBits::B128: return CU_TENSOR_MAP_SWIZZLE_128B;
}
}
inline CUtensorMapFloatOOBfill
to_CUtensorMapFloatOOBfill(OOBFill const& t) {
switch(t) {
default: assert(false && "Unknown OOBFill!");
case OOBFill::ZERO: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
case OOBFill::CONSTANT: return CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA;
}
}
inline CUtensorMapL2promotion
to_CUtensorMapL2promotion(L2Promotion const& t) {
switch(t) {
default: assert(false && "Unknown L2Promotion!");
case L2Promotion::DISABLE: return CU_TENSOR_MAP_L2_PROMOTION_NONE;
case L2Promotion::B64: return CU_TENSOR_MAP_L2_PROMOTION_L2_64B;
case L2Promotion::B128: return CU_TENSOR_MAP_L2_PROMOTION_L2_128B;
case L2Promotion::B256: return CU_TENSOR_MAP_L2_PROMOTION_L2_256B;
}
}
#endif // !defined(__CUDACC_RTC__)
#endif // (__CUDACC_VER_MAJOR__ >= 12)
@@ -257,22 +320,32 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor
asm volatile (
"cvt.u64.u32 %0, %1;"
:: "l"(smem_int64_desc), "r"(smem_int_desc));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[0]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[1]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 2, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[2]));
// Strides must be a multiple of 16. Also, stride for the intermost dimension is implicitly 1
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1] >> 4));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[0]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[1]));
asm volatile (
"tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 2, %1;"
:: "l"(smem_int64_desc), "r"(prob_shape[2]));
// Strides must be a multiple of 16. Also, stride for the intermost dimension is implicitly 1
#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 5)))
// 4 LSBs are not included
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1]));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2]));
#else
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[1] >> 4));
asm volatile (
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2] >> 4));
#endif
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
#endif
include/cute/arch/copy_sm90_tma.hpp
+30 -30
View File
@@ -44,7 +44,7 @@ namespace cute
struct SM90_TMA_LOAD_1D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
@@ -53,11 +53,11 @@ struct SM90_TMA_LOAD_1D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3}], [%2];"
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3}], [%2], %4;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0)
"r"(crd0), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -89,7 +89,7 @@ struct SM90_TMA_LOAD_1D
struct SM90_TMA_LOAD_2D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
@@ -98,11 +98,11 @@ struct SM90_TMA_LOAD_2D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4}], [%2];"
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4}], [%2], %5;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1)
"r"(crd0), "r"(crd1), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -134,7 +134,7 @@ struct SM90_TMA_LOAD_2D
struct SM90_TMA_LOAD_3D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
@@ -143,11 +143,11 @@ struct SM90_TMA_LOAD_3D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5}], [%2];"
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5}], [%2], %6;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2)
"r"(crd0), "r"(crd1), "r"(crd2), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -179,7 +179,7 @@ struct SM90_TMA_LOAD_3D
struct SM90_TMA_LOAD_4D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
@@ -188,11 +188,11 @@ struct SM90_TMA_LOAD_4D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6}], [%2];"
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3)
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -224,7 +224,7 @@ struct SM90_TMA_LOAD_4D
struct SM90_TMA_LOAD_5D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4)
{
@@ -233,11 +233,11 @@ struct SM90_TMA_LOAD_5D
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2];"
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4)
"r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -269,39 +269,39 @@ struct SM90_TMA_LOAD_5D
struct SM90_TMA_LOAD
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
return SM90_TMA_LOAD_1D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0);
return SM90_TMA_LOAD_1D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
return SM90_TMA_LOAD_2D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1);
return SM90_TMA_LOAD_2D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
return SM90_TMA_LOAD_3D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2);
return SM90_TMA_LOAD_3D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
return SM90_TMA_LOAD_4D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2, crd3);
return SM90_TMA_LOAD_4D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2, crd3);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4)
{
return SM90_TMA_LOAD_5D::copy(desc_ptr, mbar_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4);
return SM90_TMA_LOAD_5D::copy(desc_ptr, mbar_ptr, cache_hint, smem_ptr, crd0, crd1, crd2, crd3, crd4);
}
struct PREFETCH
include/cute/arch/mma_sm90_desc.hpp
+14 -14
View File
@@ -85,7 +85,6 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
union GmmaDescriptor
{
CUTE_HOST_DEVICE constexpr
GmmaDescriptor() noexcept : desc_(0) {}
CUTE_HOST_DEVICE constexpr
@@ -135,21 +134,22 @@ union GmmaDescriptor
// Decay to a uint64_t
CUTE_HOST_DEVICE constexpr
operator uint64_t() const noexcept { return desc_; }
// Printer
CUTE_HOST_DEVICE friend void print(GmmaDescriptor const& t)
{
#if !defined(__CUDACC_RTC__)
printf("GmmaDescriptor: 0x%016llx\n", static_cast<unsigned long long>(t.desc_));
printf(" start_addr : 0x%04x\n", t.bitfield.start_address_);
printf(" leading_off: 0x%04x (%d)\n", t.bitfield.leading_byte_offset_, t.bitfield.leading_byte_offset_);
printf(" stride_off : 0x%04x (%d)\n", t.bitfield.stride_byte_offset_, t.bitfield.stride_byte_offset_);
printf(" base_offset: 0x%01x\n", t.bitfield.base_offset_);
printf(" layout_type: 0x%01x (%s)\n", t.bitfield.layout_type_, to_string(static_cast<GMMA::LayoutType>(t.bitfield.layout_type_)));
#endif
}
};
// Printer
CUTE_HOST_DEVICE void
print(GmmaDescriptor const& t)
{
#if !defined(__CUDACC_RTC__)
printf("GmmaDescriptor: 0x%016llx\n", static_cast<unsigned long long>(t.desc_));
printf(" start_addr : 0x%04x\n", t.bitfield.start_address_);
printf(" leading_off: 0x%04x (%d)\n", t.bitfield.leading_byte_offset_, t.bitfield.leading_byte_offset_);
printf(" stride_off : 0x%04x (%d)\n", t.bitfield.stride_byte_offset_, t.bitfield.stride_byte_offset_);
printf(" base_offset: 0x%01x\n", t.bitfield.base_offset_);
printf(" layout_type: 0x%01x (%s)\n", t.bitfield.layout_type_, to_string(static_cast<GMMA::LayoutType>(t.bitfield.layout_type_)));
#endif // !defined(__CUDACC_RTC__)
}
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cute
include/cute/arch/util.hpp
+14 -13
View File
@@ -235,24 +235,25 @@ explode(Fn fn,
}
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>
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class PtrE, int... Ie,
class PtrF, int... If>
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...>)
PtrD&& d, int_sequence<Id...>,
PtrA&& a, int_sequence<Ia...>,
PtrB&& b, int_sequence<Ib...>,
PtrC&& c, int_sequence<Ic...>,
PtrE&& e, int_sequence<Ie...>,
PtrF&& f, int_sequence<If...>)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., sfa[Isfa]..., sfb[Isfb]...);
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]..., f[If]...);
}
//
// Utility for exploding tuples into functions
//
include/cute/atom/copy_atom.hpp
+1 -1
View File
@@ -39,7 +39,7 @@
#include <cute/util/type_traits.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
include/cute/atom/copy_traits.hpp
+12 -1
View File
@@ -32,7 +32,7 @@
#include <cute/arch/copy.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
@@ -145,4 +145,15 @@ copy_unpack(Copy_Traits<CopyOp,Args...> const& traits,
copy_unpack(traits, src, dst);
}
namespace detail {
template <class CopyOp, class = void>
constexpr bool is_prefetch = false;
template <class CopyOp>
constexpr bool is_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = is_same_v<CopyOp, typename CopyOp::PREFETCH>;
} // end namespace detail
} // end namespace cute
include/cute/atom/copy_traits_sm90_im2col.hpp
+49 -41
View File
@@ -39,7 +39,7 @@
#include "cute/tensor.hpp"
#include "cute/algorithm/prefetch.hpp"
#include "cutlass/fast_math.h"
namespace cute
{
@@ -388,18 +388,19 @@ template <class EngineA, class LayoutA,
CUTE_HOST
auto
make_im2col_tma_copy_desc(
Tensor<EngineA, LayoutA> const& tensor_cwhdn, // (C,W,H,D,N)
uint32_t range_c, // TILE_C
uint32_t range_whdn, // TILE_WHDN
SmemSwizzle const& smem_swizzle, // Swizzle
TMALayout const& tma_layout_vt, // TMA layout
LowerCornerStride const& lower_corner_whd, // WHD offset of the "base pointer"
UpperCornerStride const& upper_corner_whd, // WHD upper corner
LowerPaddingStride const& lower_padding_whd, // WHD lower padding
UpperPaddingStride const& upper_padding_whd, // WHD upper padding
TraversalStride const& stride_whd, // WHD traversal stride
LowerSRTStride const& lower_srt, // SRT offset of the "base pointer"
DilationStride const& stride_srt) // SRT stride - dilation
Tensor<EngineA, LayoutA> const& tensor_cwhdn, // (C,W,H,D,N)
uint32_t range_c, // TILE_C
uint32_t range_whdn, // TILE_WHDN
SmemSwizzle const& smem_swizzle, // Swizzle
TMALayout const& tma_layout_vt, // TMA layout
LowerCornerStride const& lower_corner_whd, // WHD offset of the "base pointer"
UpperCornerStride const& upper_corner_whd, // WHD upper corner
LowerPaddingStride const& lower_padding_whd, // WHD lower padding
UpperPaddingStride const& upper_padding_whd, // WHD upper padding
TraversalStride const& stride_whd, // WHD traversal stride
LowerSRTStride const& lower_srt, // SRT offset of the "base pointer"
DilationStride const& stride_srt, // SRT stride - dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
static_assert(is_gmem<EngineA>::value, "Tensor must point to GPU global memory.");
using value_type = typename EngineA::value_type;
@@ -445,8 +446,8 @@ make_im2col_tma_copy_desc(
CUtensorMapDataType tma_format = TMA::to_CUtensorMapDataType<value_type>();
CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE;
CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE;
CUtensorMapFloatOOBfill tma_oob_fill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
CUtensorMapL2promotion tma_l2Promotion = to_CUtensorMapL2promotion(aux_params.l2promo_);
CUtensorMapFloatOOBfill tma_oob_fill = to_CUtensorMapFloatOOBfill(aux_params.oobfill_);
CUtensorMapSwizzle tma_swizzle = TMA::to_CUtensorMapSwizzle(detail::get_tma_swizzle_bits(smem_swizzle));
CUresult encode_result = cuTensorMapEncodeIm2col(
@@ -498,7 +499,11 @@ make_im2col_tma_copy_desc(
// For fprop/dgrad kernel, gemm_shapes is ((q, p, z, n), (c, s, r, t))
// For wgrad kernel, gemm_shapes is ((c, s, r, t), (q, p, z, n))
auto gemm_shapes_common = make_shape(gemm_mn, gemm_k);
auto gemm_shapes_common = make_shape(
transform_leaf(gemm_mn, [](auto s) {
return conditional_return(cute::is_static<decltype(s)>{}, s, cutlass::FastDivmod(s));
}),
gemm_k);
auto gemm_shapes = make_shape(
basis_get(stride<0,1>(tma_layout_vt), gemm_shapes_common),
basis_get(stride<0,0>(tma_layout_vt), gemm_shapes_common));
@@ -554,17 +559,18 @@ template <class CopyOp,
CUTE_HOST_RTC
auto
make_tma_atom_im2col(CopyOp,
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor: ((w, h, d, n), c)
SLayout const& slayout, // CTA Tile of SMEM, potentially swizzled
int32_t const& num_multicast, // The number of CTAs involved in multicasting
Layout<VShape,VStride> const& cta_v_map, // V: CTA val idx -> gmem mode
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt) // dilation
Tensor<GEngine,GLayout> const& gtensor, // Full GMEM Tensor: ((w, h, d, n), c)
SLayout const& slayout, // CTA Tile of SMEM, potentially swizzled
int32_t const& num_multicast, // The number of CTAs involved in multicasting
Layout<VShape,VStride> const& cta_v_map, // V: CTA val idx -> gmem mode
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt, // dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
//
// TMA parameter checking
@@ -645,7 +651,8 @@ make_tma_atom_im2col(CopyOp,
upper_padding_whd,
stride_whd,
lower_srt,
stride_srt);
stride_srt,
aux_params);
//
// Construct the Copy_Traits
@@ -697,18 +704,19 @@ template <class CopyOp,
class DilationStride>
CUTE_HOST_RTC
auto
make_tma_copy_im2col(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
Layout<TShape,TStride> const& cta_t_map, // CTA tid -> logical TMA tid
Layout<VShape,VStride> const& cta_v_map, // CTA vid -> gmem coord
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt) // dilation
make_tma_copy_im2col(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
Layout<TShape,TStride> const& cta_t_map, // CTA tid -> logical TMA tid
Layout<VShape,VStride> const& cta_v_map, // CTA vid -> gmem coord
LowerCornerStride const& lower_corner_whd,
UpperCornerStride const& upper_corner_whd,
LowerPaddingStride const& lower_padding_whd,
UpperPaddingStride const& upper_padding_whd,
TraversalStride const& stride_whd, // traversal stride
LowerSRTStride const& lower_srt,
DilationStride const& stride_srt, // dilation
TMA::DescriptorAuxParams const& aux_params = {})
{
//
// TMA parameter checking
@@ -719,7 +727,7 @@ make_tma_copy_im2col(CopyOp const& copy_op,
Copy_Atom atom = make_tma_atom_im2col(copy_op, gtensor, slayout, cosize(cta_t_map), cta_v_map,
lower_corner_whd, upper_corner_whd, lower_padding_whd,
upper_padding_whd, stride_whd, lower_srt, stride_srt);
upper_padding_whd, stride_whd, lower_srt, stride_srt, aux_params);
//
// Construct the TiledCopy
include/cute/atom/copy_traits_sm90_tma.hpp
+186 -10
View File
@@ -124,17 +124,24 @@ struct Copy_Traits<SM90_TMA_LOAD, NumBitsPerTMA, AuxParams_>
// Construct an executable SM90_TMA_LOAD with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
with(uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const {
with(
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {&tma_desc_, &tma_mbar}};
return {{}, {&tma_desc_, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
}
// Construct an executable SM90_TMA_LOAD with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc, uint64_t& tma_mbar, [[maybe_unused]] uint16_t const& multicast_mask = 0) const {
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_mbar,
[[maybe_unused]] uint16_t const& multicast_mask = 0,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
// We accept multicast_mask here to keep the API for both atoms consistent
return {{}, {new_tma_desc, &tma_mbar}};
return {{}, {new_tma_desc, &tma_mbar, static_cast<uint64_t>(cache_hint)}};
}
// Generate the TMA coord tensor
@@ -171,7 +178,8 @@ struct Copy_Traits<SM90_TMA_LOAD_OP, NumBitsPerTMA>
// SM90_TMA_LOAD arguments
tuple<
TmaDescriptor const*,
uint64_t* // smem mbarrier
uint64_t*, // smem mbarrier
uint64_t // cache hint
> const opargs_;
};
@@ -286,6 +294,38 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
///////////////////////////// TMA_STORE //////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////
// Utility for unpacking TMA_STORE arguments into a CopyOp
template <class CopyOp>
struct TMA_STORE_Unpack
{
template <class... Args,
class TS, class SLayout,
class TD, class DLayout>
CUTE_HOST_DEVICE friend constexpr void
copy_unpack(Copy_Traits<CopyOp, Args...> const& traits,
Tensor<TS,SLayout> const& src,
Tensor<TD,DLayout> & dst)
{
static_assert(is_smem<TS>::value, "Expected smem src for SM90_TMA_STORE");
void const* const desc_ptr = traits.tma_desc_;
void const* const src_ptr = cute::raw_pointer_cast(src.data());
auto dst_coord = dst.data().coord_;
#if 0
auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
printf("THR (%d,%d,%d) BLK (%d,%d,%d) TMACRD (%d,%d,%d,%d,%d) SMEMADDR (%p)\n",
threadIdx.x, threadIdx.y, threadIdx.z,
blockIdx.x, blockIdx.y, blockIdx.z,
int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), src_ptr);
#endif
return detail::explode_tuple(detail::CallCOPY<SM90_TMA_STORE>{},
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
};
struct SM90_TMA_STORE_OP : SM90_TMA_STORE {};
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA, class AuxParams_>
struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
@@ -343,6 +383,30 @@ struct Copy_Traits<SM90_TMA_STORE, NumBitsPerTMA, AuxParams_>
make_tuple(desc_ptr, src_ptr), seq<0,1>{},
dst_coord, tuple_seq<decltype(dst_coord)>{});
}
// Construct Copy_Traits executable (w/ swapped out TMA descriptor) for SM90_TMA_STORE (for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc) const {
return {{}, new_tma_desc};
}
};
// The executable SM90_TMA_STORE with tma_desc
template <class NumBitsPerTMA>
struct Copy_Traits<SM90_TMA_STORE_OP, NumBitsPerTMA>
: TMA_STORE_Unpack<SM90_TMA_STORE_OP>
{
using ThrID = Layout<_1>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape<_1,NumBitsPerTMA>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
// SM90_TMA_STORE arguments
TmaDescriptor const* tma_desc_;
};
//////////////////////////////////////////////////////////////////////////////
@@ -1240,14 +1304,14 @@ template <class TmaInternalType = void,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
class Cluster_Size = Int<1>>
CUTE_HOST_RTC
auto
make_tma_atom(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
Cluster_Size const& cluster_size = {})
{
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
@@ -1283,8 +1347,8 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
auto layout_V = make_tile(logical_divide(layout_v, tma_layout_v));
// Append with _ until we cover all Rest... modes
auto glayout_V = append<rank_v<decltype(gtensor)>>(layout_V, _);
auto slayout_V = append<rank_v<decltype(stensor)>>(layout_V, _);
auto glayout_V = append<GLayout::rank>(layout_V, _);
auto slayout_V = append<SLayout::rank>(layout_V, _);
// Transform tile mode and coalesce
Tensor gtensor_v = coalesce(gtensor.compose(glayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
Tensor stensor_v = coalesce(stensor.compose(slayout_V), Shape<Shape<_1,_1>>{}); // ((TMA,TMA_Iter), Rest...)
@@ -1304,8 +1368,8 @@ tma_partition(Copy_Atom<Args...> const& copy_atom,
// Offset inside the TMA-mode for the multicast
auto multicast_offset = cta_layout(cta_coord) * (size(tma_layout_v) / cosize(cta_layout));
auto multicast_coord = make_coord(make_coord(multicast_offset, Int<0>{}));
auto scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
auto gcoord = append<GLayout::rank>(multicast_coord, Int<0>{});
auto scoord = append<SLayout::rank>(multicast_coord, Int<0>{});
Tensor gresult = domain_offset(gcoord, gtensor_v);
Tensor sresult = domain_offset(scoord, stensor_v);
@@ -1332,4 +1396,116 @@ create_tma_multicast_mask(CtaLayout const& cta_layout_vmnk,
return mcast_mask;
}
////////////////////////////////////
// Make TMA copy A/B/C
///////////////////////////////////
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
CUTE_HOST_RTC
auto
make_tma_copy_A_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
// Keep only MK modes from MNK
auto cta_tiler_mk = remove<1>(cta_tiler);
// mcast along N mode for this M load, if any
auto cluster_size_n = size<1>(cluster_size);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_mk,
cluster_size_n);
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_mk);
auto cta_t_tile = make_layout(cluster_size_n);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_tile, cta_v_tile);
return tma_copy;
}
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler,
class Cluster_Size>
CUTE_HOST_RTC
auto
make_tma_copy_B_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler,
Cluster_Size const& cluster_size)
{
// Keep only NK modes from MNK
auto cta_tiler_nk = remove<0>(cta_tiler);
// mcast along M mode for this N load, if any
auto cluster_size_m = size<0>(cluster_size);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_nk,
cluster_size_m);
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_nk);
auto cta_t_tile = make_layout(cluster_size_m);
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_tile, cta_v_tile);
return tma_copy;
}
}
template <class TmaInternalType = void,
class CopyOp,
class GEngine, class GLayout,
class SLayout,
class CTA_Tiler>
CUTE_HOST_RTC
auto
make_tma_copy_C_sm90(CopyOp const& copy_op,
Tensor<GEngine,GLayout> const& gtensor,
SLayout const& slayout,
CTA_Tiler const& cta_tiler)
{
// Keep only MN modes from MNK
auto cta_tiler_mn = remove<2>(cta_tiler);
if constexpr (cute::is_same_v<CopyOp, SM90_TMA_LOAD_IM2COL> ||
cute::is_same_v<CopyOp, SM90_TMA_STORE_IM2COL>) {
return make_im2col_tma_copy(copy_op,
gtensor,
slayout,
cta_tiler_mn,
_1{});
} else {
auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler_mn);
// No multicast, so only 1 CTA involved
auto cta_t_map = Layout<_1,_0>{};
// Prefer TmaInternalType if specified. Fallback to GEngine::value_type
using TmaType = conditional_t<is_same<void, TmaInternalType>::value, typename GEngine::value_type, TmaInternalType>;
auto tma_copy = detail::make_tma_copy_tiled<TmaType>(copy_op, gtensor, slayout, cta_t_map, cta_v_tile);
return tma_copy;
}
}
} // end namespace cute
include/cute/atom/mma_atom.hpp
+19 -28
View File
@@ -31,11 +31,9 @@
#pragma once
#include <cute/config.hpp>
#include <cute/arch/mma.hpp>
#include <cute/atom/mma_traits.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/util/type_traits.hpp>
namespace cute {
@@ -102,7 +100,7 @@ struct MMA_Atom<MMA_Traits<Args...>>
static_assert(BLayout::rank == 1, "Expected rank-1 B tensor");
static_assert(CLayout::rank == 1, "Expected rank-1 C tensor");
return mma_unpack(*this, D, A, B, C);
return mma_unpack(static_cast<Traits const&>(*this), D, A, B, C);
}
// Three arguments reproduces C
@@ -245,12 +243,9 @@ struct TiledMMA : MMA_Atom
thrfrg_C(CTensor&& ctensor) const
{
CUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(ctensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(ctensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<0>(PermutationMNK{}),
get<1>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<0>(),
permutation_mnk<1>());
auto t_tensor = logical_divide(ctensor, t_tile); // (PermM,PermN)
// Tile the tensor for the Atom
@@ -287,12 +282,9 @@ struct TiledMMA : MMA_Atom
thrfrg_A(ATensor&& atensor) const
{
CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<0>(PermutationMNK{}),
get<2>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<0>(),
permutation_mnk<2>());
auto t_tensor = logical_divide(atensor, t_tile); // (PermM,PermK)
// Tile the tensor for the Atom
@@ -329,12 +321,9 @@ struct TiledMMA : MMA_Atom
thrfrg_B(BTensor&& btensor) const
{
CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{});
//CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
//CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
// Reorder the tensor for the TiledAtom
auto t_tile = make_tile(get<1>(PermutationMNK{}),
get<2>(PermutationMNK{}));
auto t_tile = make_tile(permutation_mnk<1>(),
permutation_mnk<2>());
auto t_tensor = logical_divide(btensor, t_tile); // (PermN,PermK)
// Tile the tensor for the Atom
@@ -377,21 +366,23 @@ struct TiledMMA : MMA_Atom
// Utility for printing and visualization
//
// The permutation applied to the MNK-mode data
template <int I>
CUTE_HOST_DEVICE constexpr
auto
permutation_mnk() const {
static_assert(0 <= I && I < 3);
auto perm = get<I>(PermutationMNK{});
return conditional_return(is_underscore<decltype(perm)>{}, size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk()), perm);
}
// The size of the MNK-mode
template <int I>
CUTE_HOST_DEVICE constexpr
auto
tile_size_mnk() const {
static_assert(0 <= I && I < 3);
auto core_size = size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk());
[[maybe_unused]] auto perm_size = size<I>(PermutationMNK{});
if constexpr (is_underscore<decltype(perm_size)>::value) {
return core_size;
} else {
return cute::max(core_size, perm_size);
}
CUTE_GCC_UNREACHABLE;
return size(permutation_mnk<I>());
}
CUTE_HOST_DEVICE constexpr
include/cute/atom/mma_traits.hpp
+1 -1
View File
@@ -32,7 +32,7 @@
#include <cute/arch/mma.hpp>
#include <cute/tensor.hpp>
#include <cute/tensor_impl.hpp>
namespace cute
{
include/cute/atom/mma_traits_sm90_gmma.hpp
+5 -3
View File
@@ -332,9 +332,6 @@ struct DescriptorIterator
{
return { GmmaDescriptor{desc_ + uint64_t(offset)} };
}
CUTE_HOST_DEVICE friend void
print(DescriptorIterator) { printf("GMMA::DescriptorIterator"); }
};
template <class T>
@@ -353,6 +350,11 @@ recast_ptr(DescriptorIterator const& iter) {
return iter; // Do nothing, it will still dereference to GmmaDescriptor and decay to uint64_t
}
CUTE_HOST_DEVICE void
print(DescriptorIterator) {
printf("GMMA::DescriptorIterator");
}
// The GMMA Traits below have custom fragment type flags for their smem desc tensors.
// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired.
template <GMMA::Major>
include/cute/container/alignment.hpp
+1 -1
View File
@@ -44,7 +44,7 @@ CUTE_HOST_DEVICE constexpr
bool
is_byte_aligned(void const* const ptr)
{
static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of 2 in alignment check");
static_assert(has_single_bit(N), "N must be a power of 2 in alignment check");
return (reinterpret_cast<uintptr_t>(ptr) & (N-1)) == 0;
}
include/cute/container/array_subbyte.hpp
+46 -57
View File
@@ -205,18 +205,22 @@ struct subbyte_iterator
private:
template <class, class> friend struct swizzle_ptr;
template <class U> friend CUTE_HOST_DEVICE constexpr U* raw_pointer_cast(subbyte_iterator<U> const&);
template <class N, class U> friend CUTE_HOST_DEVICE constexpr auto recast_ptr(subbyte_iterator<U> const&);
template <class U> friend CUTE_HOST_DEVICE void print(subbyte_iterator<U> const&);
// Pointer to storage element
storage_type* ptr_ = nullptr;
storage_type* ptr_;
// Bit index of value_type starting position within storage_type element.
// RI: 0 <= idx_ < sizeof_bit<storage_type>
uint8_t idx_ = 0;
uint8_t idx_;
public:
// Ctor
subbyte_iterator() = default;
// Default Ctor
CUTE_HOST_DEVICE constexpr
subbyte_iterator() : ptr_{nullptr}, idx_{0} {};
// Ctor
template <class PointerType>
@@ -286,43 +290,48 @@ public:
return x.ptr_ == y.ptr_ && x.idx_ == y.idx_;
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator< (subbyte_iterator const& x, subbyte_iterator const& y) {
return x.ptr_ < y.ptr_ || (x.ptr_ == y.ptr_ && x.idx_ < y.idx_);
}
CUTE_HOST_DEVICE constexpr friend
bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
CUTE_HOST_DEVICE constexpr friend
bool operator<=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator> (subbyte_iterator const& x, subbyte_iterator const& y) { return (y < x); }
CUTE_HOST_DEVICE constexpr friend
bool operator>=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x < y); }
// Conversion to raw pointer with loss of subbyte index
CUTE_HOST_DEVICE constexpr friend
T* raw_pointer_cast(subbyte_iterator const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_>
CUTE_HOST_DEVICE constexpr friend
auto recast_ptr(subbyte_iterator const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (cute::is_subbyte_v<NewT>) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
CUTE_HOST_DEVICE friend void print(subbyte_iterator x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits_v<T>), x.ptr_, x.idx_);
}
};
// Conversion to raw pointer with loss of subbyte index
template <class T>
CUTE_HOST_DEVICE constexpr
T*
raw_pointer_cast(subbyte_iterator<T> const& x) {
assert(x.idx_ == 0);
return reinterpret_cast<T*>(x.ptr_);
}
// Conversion to NewT_ with possible loss of subbyte index
template <class NewT_, class T>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(subbyte_iterator<T> const& x) {
using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
if constexpr (cute::is_subbyte_v<NewT>) { // Making subbyte_iter, preserve the subbyte idx
return subbyte_iterator<NewT>(x.ptr_, x.idx_);
} else { // Not subbyte, assume/assert subbyte idx 0
return reinterpret_cast<NewT*>(raw_pointer_cast(x));
}
CUTE_GCC_UNREACHABLE;
}
template <class T>
CUTE_HOST_DEVICE void
print(subbyte_iterator<T> const& x) {
printf("subptr[%db](%p.%u)", int(sizeof_bits_v<T>), x.ptr_, x.idx_);
}
//
// array_subbyte
// Statically sized array for non-byte-aligned data types
@@ -365,17 +374,6 @@ private:
public:
constexpr
array_subbyte() = default;
CUTE_HOST_DEVICE constexpr
array_subbyte(array_subbyte const& x) {
CUTE_UNROLL
for (size_type i = 0; i < StorageElements; ++i) {
storage[i] = x.storage[i];
}
}
CUTE_HOST_DEVICE constexpr
size_type size() const {
return N;
@@ -448,25 +446,16 @@ public:
return at(N-1);
}
// In analogy to std::vector<bool>::data(), these functions are deleted to prevent bugs.
// Instead, prefer
// auto* data = raw_pointer_cast(my_subbyte_array.begin());
// where the type of auto* is implementation-defined and
// with the knowledge that [data, data + my_subbyte_array.size()) may not be a valid range.
CUTE_HOST_DEVICE constexpr
pointer data() {
return reinterpret_cast<pointer>(storage);
}
pointer data() = delete;
CUTE_HOST_DEVICE constexpr
const_pointer data() const {
return reinterpret_cast<const_pointer>(storage);
}
CUTE_HOST_DEVICE constexpr
storage_type* raw_data() {
return storage;
}
CUTE_HOST_DEVICE constexpr
storage_type const* raw_data() const {
return storage;
}
const_pointer data() const = delete;
CUTE_HOST_DEVICE constexpr
iterator begin() {
include/cute/container/packed_tuple.hpp
+254
View File
@@ -0,0 +1,254 @@
/***************************************************************************************************
* 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/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/type_list.hpp>
namespace cute {
namespace detail {
// Empty Structure Optimization
template <bool IsFirstEmpty, bool IsRestEmpty, class... T>
struct ESO;
template <class First, class... Rest>
static constexpr bool is_first_empty_v = cute::is_empty<First>::value;
template <class First, class... Rest>
static constexpr bool is_rest_empty_v = (cute::is_empty<Rest>::value && ...);
template <class... T>
using ESO_t = ESO<is_first_empty_v<T...>, is_rest_empty_v<T...>, T...>;
// Empty First and Empty Rest...
template <class First, class... Rest>
struct ESO<true, true, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() {}
CUTE_HOST_DEVICE constexpr
ESO(First const&, Rest const&...) {}
};
// NonEmpty First and Empty Rest...
template <class First, class... Rest>
struct ESO<false, true, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : first_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const& first, Rest const&...) : first_{first} {}
First first_;
};
// Empty First and NonEmpty Rest...
template <class First, class... Rest>
struct ESO<true, false, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : rest_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const&, Rest const&... rest) : rest_{rest...} {}
ESO_t<Rest...> rest_;
};
// NonEmpty T and NonEmpty Rest...
template <class First, class... Rest>
struct ESO<false, false, First, Rest...> {
CUTE_HOST_DEVICE constexpr
ESO() : first_{}, rest_{} {}
CUTE_HOST_DEVICE constexpr
ESO(First const& first, Rest const&... rest) : first_{first}, rest_{rest...} {}
First first_;
ESO_t<Rest...> rest_;
};
// Get Nth value from ESO
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...> const& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T const&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(s.rest_); }
}
}
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...>& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(s.rest_); }
}
}
template <size_t N, class T, class... Rest, bool F, bool R>
CUTE_HOST_DEVICE constexpr decltype(auto) getv(ESO<F, R, T, Rest...>&& s) {
if constexpr (N == 0) {
if constexpr (F) { return T{}; }
else { return static_cast<T&&>(s.first_); }
} else {
if constexpr (R) { return cute::tuple_element_t<N-1, cute::type_list<Rest...>>{}; }
else { return getv<N-1>(static_cast<ESO_t<Rest...>&&>(s.rest_)); }
}
}
// findt: Implementation detail of cute::find.
// If X is the first template argument of the tuple, findt returns C<N>.
template <class X, size_t N,
bool IsFirstEmpty, bool IsRestEmpty, class First, class... Rest>
CUTE_HOST_DEVICE constexpr
auto
findt(ESO<IsFirstEmpty, IsRestEmpty, First, Rest...> const& t) noexcept
{
if constexpr (cute::is_same_v<X, First>) {
return C<N>{};
}
else {
static_assert(sizeof...(Rest) != 0,
"The type does not appear in the argument list of the tuple.");
if constexpr (IsRestEmpty) {
// The rest is empty, so creating an instance of it is cheap.
return cute::detail::findt<X, N+1>(ESO_t<Rest...>{});
}
else {
return cute::detail::findt<X, N+1>(t.rest_);
}
}
}
} // end namespace detail
// packed_tuple<T...> is a tuple type that is a standard-layout type
// whenever all of its template arguments are standard layout types:
// (cute::is_standard_layout_v<T> && ...) implies (cute::is_standard_layout_v<packed_tuple<T...>>)
template <class... T>
struct packed_tuple : detail::ESO_t<T...>
{
CUTE_HOST_DEVICE constexpr
packed_tuple() {}
CUTE_HOST_DEVICE constexpr
packed_tuple(T const&... ts)
: detail::ESO_t<T...>(ts...)
{}
};
template <>
struct packed_tuple<> {};
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...> const& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(t);
}
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...>& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(t);
}
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
decltype(auto)
get(packed_tuple<T...>&& t) {
static_assert(I < sizeof...(T), "Index out of range");
return detail::getv<I>(static_cast<detail::ESO_t<T...>&&>(t));
}
template <class... T>
CUTE_HOST_DEVICE constexpr
packed_tuple<T...>
make_packed_tuple(T const&... t)
{
return {t...};
}
// Returns the position of type X (as a static integer) in the tuple
// type's argument list. X must be unique in the argument list.
template <class X, class... T>
CUTE_HOST_DEVICE constexpr
auto
find(packed_tuple<T...> const& t) noexcept
{
return detail::findt<X, 0>(t);
}
} // end namespace cute
namespace CUTE_STL_NAMESPACE
{
template <class... T>
struct tuple_size<cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
{};
template <size_t I, class... T>
struct tuple_element<I, cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>
{};
} // end namespace CUTE_STL_NAMESPACE
#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD
namespace std {
template <class ... T>
struct tuple_size<cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
{};
template <size_t I, class ... T>
struct tuple_element<I, cute::packed_tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, cute::packed_tuple<T...>>
{};
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
include/cute/container/tuple.hpp
+58 -43
View File
@@ -36,10 +36,13 @@
#include <cute/numeric/integer_sequence.hpp>
#include <cute/container/cuda_types.hpp>
#include <cute/container/type_list.hpp>
#if defined(CUTLASS_USE_PACKED_TUPLE)
# include <cute/container/packed_tuple.hpp>
#endif
//#include <cute/container/array.hpp> // Advanced optimizations
//
// cute::tuple is like std::tuple, with two differences.
//
// 1. It works on both host and device.
@@ -50,19 +53,30 @@
// but do _not_ include references like int& or float&.
// (See std::tie for an example of a tuple of references.)
//
// This is simplified over the implementations in std::, cuda::std::, and thrust:: by ignoring much of
// the conversion SFINAE, special overloading, and avoiding cvref template types.
// Furthermore, the empty base optimization (EBO) is MORE aggressive by avoiding
// construction calls, and ignoring any need for unique element addresses.
//
// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x.
// If the template arguments of cute::tuple are all empty types (in
// the sense of std::is_empty_v), then the cute::tuple is also an
// empty type. Furthermore, if CUTLASS_USE_PACKED_TUPLE is defined,
// cute::tuple is always a standard-layout type if all of its template
// arguments are standard-layout types.
namespace cute
{
#if defined(CUTLASS_USE_PACKED_TUPLE)
template<class... T>
using tuple = packed_tuple<T...>;
#else
namespace detail
{
// This is simplified over the implementations in std::, cuda::std::, and thrust:: by ignoring much of
// the conversion SFINAE, special overloading, and avoiding cvref template types.
//
// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x.
// EBO stands for "empty base optimization."
// We use this technique to ensure that cute::tuple
// doesn't need to waste space storing any template arguments
@@ -70,6 +84,12 @@ namespace detail
// Otherwise, cute::tuple would need to spend at least 1 byte
// for each of its template arguments.
//
// This is one way in which cute::tuple differs from std::tuple.
// Empty types in the template argument list are not even constructed,
// and do not have unique element addresses. In fact, they are not
// even members of the tuple or stored in any way. Calling `get`
// constructs and returns an instance of an empty type on demand.
//
// EBO always "holds" a single value of type T.
// N is like an array index that TupleBase uses
// to access the desired tuple element.
@@ -109,9 +129,8 @@ struct EBO<N, T, false>
CUTE_HOST_DEVICE constexpr
EBO() : t_{} {}
template <class U>
CUTE_HOST_DEVICE constexpr
EBO(U const& u) : t_{u} {}
EBO(T const& t) : t_{t} {}
T t_;
};
@@ -141,15 +160,8 @@ struct TupleBase<index_sequence<I...>, T...>
CUTE_HOST_DEVICE constexpr
TupleBase() {}
template <class... U>
CUTE_HOST_DEVICE constexpr explicit
TupleBase(U const&... u)
: EBO<I,T>(u)... {}
template <class... U>
CUTE_HOST_DEVICE constexpr
TupleBase(TupleBase<index_sequence<I...>, U...> const& u)
: EBO<I,T>(getv(static_cast<EBO<I,U> const&>(u)))... {}
TupleBase(T const&... t) : EBO<I,T>(t)... {}
};
} // end namespace detail
@@ -172,16 +184,14 @@ struct tuple : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>
CUTE_HOST_DEVICE constexpr
tuple() {}
template <class... U>
CUTE_HOST_DEVICE constexpr
tuple(U const&... u) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(u...) {}
template <class... U>
CUTE_HOST_DEVICE constexpr
tuple(tuple<U...> const& u)
: detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(static_cast<detail::TupleBase<make_index_sequence<sizeof...(U)>, U...> const&>(u)) {}
tuple(T const&... t) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(t...) {}
};
template <>
struct tuple<>
{};
//
// get for cute::tuple (just like std::get for std::tuple)
//
@@ -227,6 +237,8 @@ find(tuple<T...> const& t) noexcept
return detail::findt<X>(t);
}
#endif // CUTLASS_USE_PACKED_TUPLE
//
// Custom is_tuple trait simply checks the existence of tuple_size
// and assumes std::get<I>(.), std::tuple_element<I,.>
@@ -242,6 +254,9 @@ auto has_tuple_size(...) -> false_type;
template <class T>
struct is_tuple : decltype(detail::has_tuple_size((T*)0)) {};
template<typename T>
constexpr bool is_tuple_v = cute::is_tuple<T>::value;
//
// make_tuple (value-based implementation)
//
@@ -540,20 +555,12 @@ tuple_cat(Tuples const&... ts)
namespace detail {
template <size_t I, class TupleA, class TupleB>
template <class TupleA, class TupleB, size_t... I>
CUTE_HOST_DEVICE constexpr
auto
equal_impl(TupleA const& a, TupleB const& b)
equal_impl(TupleA const& a, TupleB const& b, index_sequence<I...>)
{
if constexpr (I == tuple_size<TupleA>::value) {
return cute::true_type{}; // Terminal: TupleA is exhausted
} else if constexpr (I == tuple_size<TupleB>::value) {
return cute::false_type{}; // Terminal: TupleA is not exhausted, TupleB is exhausted
} else {
return (get<I>(a) == get<I>(b)) && equal_impl<I+1>(a,b);
}
CUTE_GCC_UNREACHABLE;
return (cute::true_type{} && ... && (get<I>(a) == get<I>(b)));
}
} // end namespace detail
@@ -564,7 +571,13 @@ CUTE_HOST_DEVICE constexpr
auto
operator==(TupleT const& t, TupleU const& u)
{
return detail::equal_impl<0>(t, u);
if constexpr (tuple_size<TupleT>::value == tuple_size<TupleU>::value) {
return detail::equal_impl(t, u, make_index_sequence<tuple_size<TupleT>::value>{});
} else {
return cute::false_type{};
}
CUTE_GCC_UNREACHABLE;
}
template <class TupleT, class TupleU,
@@ -618,19 +631,17 @@ operator!=(TupleT const& t, TupleU const& u)
namespace detail {
template <class Tuple, size_t... Is>
CUTE_HOST_DEVICE void print_tuple(Tuple const& t,
index_sequence<Is...>, char s = '(', char e = ')')
CUTE_HOST_DEVICE void print_tuple(Tuple const& t, index_sequence<Is...>, char s = '(', char e = ')')
{
using cute::print;
((void(print(Is == 0 ? s : ',')), void(print(get<Is>(t)))), ...); print(e);
print(s); ((void(print(Is == 0 ? '\0' : ',')), void(print(get<Is>(t)))), ...); print(e);
}
#if !defined(__CUDACC_RTC__)
template <class Tuple, std::size_t... Is>
CUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t,
index_sequence<Is...>, char s = '(', char e = ')')
CUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t, index_sequence<Is...>, char s = '(', char e = ')')
{
(void(os << (Is == 0 ? s : ',') << get<Is>(t)), ...);
os << s; (void(os << (Is == 0 ? '\0' : ',') << get<Is>(t)), ...);
return os << e;
}
#endif // !defined(__CUDACC_RTC__)
@@ -655,6 +666,8 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, Tuple const& t)
} // end namespace cute
#if ! defined(CUTLASS_USE_PACKED_TUPLE)
namespace CUTE_STL_NAMESPACE
{
@@ -716,5 +729,7 @@ struct tuple_element<I, const cute::tuple<T...>>
: CUTE_STL_NAMESPACE::tuple_element<I, const CUTE_STL_NAMESPACE::tuple<T...>>
{};
} // end namepsace std
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
#endif // CUTLASS_USE_PACKED_TUPLE
include/cute/container/type_list.hpp
+23 -34
View File
@@ -30,19 +30,24 @@
**************************************************************************************************/
#pragma once
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
namespace cute
{
template <class T>
struct type_c {
using type = T;
};
template <class... T>
struct type_list {};
// get<I> for type_list<T...>
// requires tuple_element_t<I,type_list<T...>> to have std::is_default_constructible
template <size_t I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...> const& t) noexcept {
return {};
}
} // end namespace cute
//
@@ -55,26 +60,6 @@ struct type_list {};
#include <tuple>
#endif
#include <cute/container/tuple.hpp>
namespace cute
{
template <int I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...>&) noexcept {
return {};
}
template <int I, class... T>
CUTE_HOST_DEVICE constexpr
CUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
get(type_list<T...> const& t) noexcept {
return {};
}
} // end namespace cute
namespace CUTE_STL_NAMESPACE
{
@@ -85,8 +70,9 @@ struct tuple_size<cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
template <class... T>
struct tuple_size<const cute::type_list<T...>>
@@ -95,8 +81,9 @@ struct tuple_size<const cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, const cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
} // end namespace std
@@ -119,8 +106,9 @@ struct tuple_size<cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
template <class... T>
struct tuple_size<const cute::type_list<T...>>
@@ -129,8 +117,9 @@ struct tuple_size<const cute::type_list<T...>>
template <size_t I, class... T>
struct tuple_element<I, const cute::type_list<T...>>
: cute::type_c<typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type>
{};
{
using type = typename CUTE_STL_NAMESPACE::tuple_element<I, CUTE_STL_NAMESPACE::tuple<T...>>::type;
};
} // end namespace std
#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD
include/cute/int_tuple.hpp
+33
View File
@@ -493,6 +493,7 @@ using is_weakly_congruent = decltype(weakly_congruent(declval<A>(), declval<B>()
/** Test if Shape A is compatible with Shape B:
* the size of A and B are the same, and
* any coordinate into A can also be used as a coordinate into B
* Equivalently, the size of Shape B is the same as Shape A at each terminal of Shape A.
* compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
@@ -523,6 +524,7 @@ using is_compatible = decltype(compatible(declval<A>(), declval<B>()));
/** Test if Shape A is weakly compatible with Shape B:
* there exists a Shape C congruent to A such that compatible(elem_scale(A,C), B)
* Equivalently, the size of Shape B is a multiple of Shape A at each terminal of Shape A.
* weakly_compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
@@ -551,6 +553,37 @@ weakly_compatible(IntTupleA const& a, IntTupleB const& b)
template <class A, class B>
using is_weakly_compatible = decltype(weakly_compatible(declval<A>(), declval<B>()));
/** Test if Shape A is softly compatible with Shape B:
* there exists a Shape C congruent to A such that compatible(shape_div(A,C), B)
* Equivalently, the size of Shape B divides Shape A at each terminal of Shape A.
* softly_compatible is a partial order on A and B: A <= B
*/
template <class IntTupleA, class IntTupleB>
CUTE_HOST_DEVICE constexpr
auto
softly_compatible(IntTupleA const& a, IntTupleB const& b)
{
if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
if constexpr (tuple_size<IntTupleA>::value != tuple_size<IntTupleB>::value) {
return false_type{};
} else {
return transform_apply(a, b, [](auto const& x, auto const& y) { return softly_compatible(x,y); },
[](auto const&... z) { return (true_type{} && ... && z); });
}
} else if constexpr (is_integral<IntTupleA>::value) {
return a % size(b) == Int<0>{};
} else if constexpr (is_integral<IntTupleB>::value) {
return false_type{};
} else {
return softly_compatible(shape(a), shape(b));
}
CUTE_GCC_UNREACHABLE;
}
template <class A, class B>
using is_softly_compatible = decltype(softly_compatible(declval<A>(), declval<B>()));
/** Replace the elements of Tuple B that are paired with an Int<0> with an Int<1>
*/
template <class IntTupleA, class IntTupleB>
include/cute/layout.hpp
+119 -26
View File
@@ -329,34 +329,23 @@ struct is_layout<Layout<Shape,Stride>> : true_type {};
// Layout construction
//
template <class Shape, class Stride,
__CUTE_REQUIRES((is_tuple<Shape >::value || is_integral<Shape >::value) &&
(is_tuple<Stride>::value || is_integral<Stride>::value))>
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, Stride const& stride)
{
static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
return Layout<Shape,Stride>(shape, stride);
}
template <class Shape,
__CUTE_REQUIRES(is_tuple<Shape>::value || is_integral<Shape>::value)>
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape)
{
return make_layout(shape, compact_col_major(shape));
}
// Construct a layout from multiple layouts by
// concatenating each layout as an independent mode
template <class... Shapes, class... Strides>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shapes,Strides> const&... layouts)
{
return make_layout(make_shape (layouts.shape()...),
make_stride(layouts.stride()...));
static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
return make_layout(shape, compact_major<LayoutLeft>(shape));
}
//
@@ -366,17 +355,57 @@ make_layout(Layout<Shapes,Strides> const&... layouts)
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, GenColMajor)
make_layout(Shape const& shape, LayoutLeft)
{
return make_layout(shape, compact_col_major(shape));
return make_layout(shape, compact_major<LayoutLeft>(shape));
}
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Shape const& shape, GenRowMajor)
make_layout(Shape const& shape, LayoutRight)
{
return make_layout(shape, compact_row_major(shape));
return make_layout(shape, compact_major<LayoutRight>(shape));
}
//
// Construct a layout from multiple layouts by concatenation
//
// One argument overload
template <class Shape0, class Stride0>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0)
{
return make_layout(make_shape (layout0.shape() ),
make_stride(layout0.stride()));
}
// Two argument overload
template <class Shape0, class Stride0,
class Shape1, class Stride1>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0,
Layout<Shape1,Stride1> const& layout1)
{
return make_layout(make_shape (layout0.shape() , layout1.shape() ),
make_stride(layout0.stride(), layout1.stride()));
}
// Var argument overload
template <class Shape0, class Stride0,
class Shape1, class Stride1,
class... Shapes, class... Strides>
CUTE_HOST_DEVICE constexpr
auto
make_layout(Layout<Shape0,Stride0> const& layout0,
Layout<Shape1,Stride1> const& layout1,
Layout<Shapes,Strides> const&... layouts)
{
return make_layout(make_shape (layout0.shape() , layout1.shape() , layouts.shape()... ),
make_stride(layout0.stride(), layout1.stride(), layouts.stride()...));
}
//
@@ -428,7 +457,7 @@ make_fragment_like(Layout<Shape,Stride> const& layout)
constexpr int R = Layout<Shape,Stride>::rank;
if constexpr (R > 1 && is_static<Shape>::value) {
return tiled_product(make_layout(get<0>(layout.shape()),
compact_col_major(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
compact_major<LayoutLeft>(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride())));
} else {
return make_layout(layout.shape());
@@ -1131,7 +1160,7 @@ complement(Shape const& shape, Stride const& stride, CoTarget const& cotarget)
// Compute the rest_shape and rest_stride
auto new_stride = get<0>(stride_) * get<0>(shape_); // new stride = min_stride * curr_shape
auto rest_shape = coalesce(ceil_div(cotarget, new_stride));
auto rest_stride = compact_col_major(rest_shape, new_stride);
auto rest_stride = compact_major<LayoutLeft>(rest_shape, new_stride);
// Coalesce and append (rest_shape, rest_stride)
return coalesce(make_layout(make_shape (result_shape , rest_shape ),
@@ -1220,7 +1249,7 @@ right_inverse(Layout<Shape,Stride> const& layout)
return Layout<_1,_0>{}; // Empty case, nothing found
} else {
// Generate the corresponding new strides and construct
auto rstride = compact_col_major(flat_layout.shape());
auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
unwrap(transform(iseq, [&](auto i) { return signum(stride<i>(flat_layout)) * get<i>(rstride); })));
}
@@ -1318,6 +1347,50 @@ max_common_vector(Layout<ShapeA,StrideA> const& a,
CUTE_GCC_UNREACHABLE;
}
/* Return a layout that distributes ShapeB over ShapeA.
*
* @returns Layout result
* @post softly_compatible(@a b, @a result)
* @post For all i,j in [0,size(@a result)) with i < j, @a result(i) < @a result(j). Surjective and Ordered.
* @post composition(make_layout(shape(@a a)), @a result) is admissible
* \code
* // Note that 6 does not divide this shape
* Layout layoutA = Layout<Shape<Int<15>,Int<14>>>{};
*
* // Want to tile any 6 elements and don't care where they come from
* Layout dist = domain_distribute(layoutA, Int<6>{}); // (_3,_2):(_1,_15)
*
* // Not guaranteed to find all 6 though...
* CUTE_STATIC_ASSERT_V(Int<6>{} == size(dist));
*
* Layout result = zipped_divide(layoutA, dist); // (_6,Rest)
* \endcode
*/
template <class ShapeA, class ShapeB>
CUTE_HOST_DEVICE constexpr
auto
domain_distribute(ShapeA const& a, ShapeB const& b)
{
static_assert(is_integral<ShapeB>::value);
static_assert(is_static<ShapeB>::value);
auto flat_shape_a = flatten(shape(a));
static_assert(is_static<decltype(flat_shape_a)>::value);
// Compute the shape of the result
auto [result_shape, b_rest] = cute::fold(flat_shape_a, cute::make_tuple(cute::tuple<>{}, size(b)), [](auto init, auto a_) {
auto [result, b_] = init;
auto gcd_ = gcd(a_, b_);
return cute::make_tuple(append(result, gcd_), b_ / gcd_);
});
// Compute the stride of the result
auto result_stride = compact_major<LayoutLeft>(flat_shape_a);
return coalesce(make_layout(result_shape, result_stride));
}
//
// Kernel (Nullspace) of a Layout
//
@@ -1363,7 +1436,7 @@ nullspace(Layout<Shape,Stride> const& layout)
return Layout<_1,_0>{}; // Empty case, nothing found
} else {
// Generate the corresponding new strides and construct
auto rstride = compact_col_major(flat_layout.shape());
auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
unwrap(transform(iseq, [&](auto i) { return get<i>(rstride); })));
}
@@ -1458,7 +1531,7 @@ auto
ceil_div(Target const& target,
Layout<TShape,TStride> const& tiler)
{
return complement(tiler, size(target));
return shape(complement(tiler, shape(target)));
}
//
@@ -1753,6 +1826,26 @@ recast_layout(Layout<Shape,Stride> const& layout)
CUTE_GCC_UNREACHABLE;
}
// Determine the maximum alignment of a Layout.
// The maximum alignment is the largest N for which upcast<N>(layout) will compile.
// upcast<N>(layout) compiles when the static shapes and strides pass divisibility checks.
// Therefore, upcast<M>(layout) will also compile for all divisors M of N.
// Note that this only considers the static shapes and strides of the Layout
// in symmetry with upcast<N> only checking against static shapes and strides and assuming all
// dynamic shapes and strides are large and multiples of N.
template <class Shape, class Stride>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(Layout<Shape,Stride> const& layout)
{
auto flat_layout = coalesce(layout);
auto static_shape = transform( shape(flat_layout), [](auto s){ return conditional_return<is_static<decltype(s)>::value>(s, Int<1>{}); });
auto static_stride = transform(stride(flat_layout), [](auto d){ return conditional_return<is_static<decltype(d)>::value>(d, Int<0>{}); });
auto filter_layout = make_layout(static_shape, static_stride);
auto permuted = logical_divide(filter_layout, right_inverse(filter_layout));
return gcd(size<0>(permuted), stride<1>(permuted));
}
//
// Display utilities
//
include/cute/layout_composed.hpp
+12
View File
@@ -577,6 +577,7 @@ coalesce(ComposedLayout<A,O,B> const& layout, Shape const& trg_profile)
return composition(layout.layout_a(), layout.offset(), coalesce(layout.layout_b(), trg_profile));
}
//
// Upcast and Downcast
//
@@ -597,6 +598,7 @@ downcast(ComposedLayout<A,O,B> const& layout)
return composition(downcast<N>(layout.layout_a()), downcast<N>(layout.offset()), downcast<N>(layout.layout_b()));
}
template <class OldType, class NewType,
class A, class O, class B>
CUTE_HOST_DEVICE constexpr
@@ -619,6 +621,16 @@ recast_layout(ComposedLayout<A,O,B> const& layout)
CUTE_GCC_UNREACHABLE;
}
template <class A, class O, class B>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(ComposedLayout<A,O,B> const& layout)
{
// Do not attempt for general ComposedLayouts
//return gcd(max_alignment(layout.layout_a()), max_alignment(layout.offset()), max_alignment(layout.layout_b()));
return Int<1>{};
}
//
// Display utilities
//
include/cute/numeric/complex.hpp
+9 -9
View File
@@ -48,13 +48,13 @@ template <class T>
static constexpr auto is_complex_v = is_complex<T>::value;
/// Fused multiply-add for complex numbers
template <class T>
template <class D, class A, class B, class C>
CUTE_HOST_DEVICE constexpr
void
fma(complex<T> & d,
complex<T> const& a,
complex<T> const& b,
complex<T> const& c)
fma(complex<D> & d,
complex<A> const& a,
complex<B> const& b,
complex<C> const& c)
{
fma(d.real(), a.real(), b.real(), c.real());
fma(d.imag(), a.real(), b.imag(), c.imag());
@@ -63,12 +63,12 @@ fma(complex<T> & d,
}
/// Fused multiply-add for triplets
template <class T>
template <class A, class B, class C>
CUTE_HOST_DEVICE constexpr
void
fma(complex<T> const& a,
complex<T> const& b,
complex<T> & c)
fma(complex<A> const& a,
complex<B> const& b,
complex<C> & c)
{
return fma(c, a, b, c);
}
include/cute/numeric/integral_constant.hpp
+5 -1
View File
@@ -33,6 +33,7 @@
#include "cute/util/print.hpp"
#include "cute/util/type_traits.hpp"
#include "cute/numeric/math.hpp"
#include "cutlass/fast_math.h"
namespace cute
{
@@ -82,8 +83,11 @@ struct is_integral<C<v> > : true_type {};
template <class T, T v>
struct is_integral<integral_constant<T,v>> : true_type {};
// is_static detects if an (abstract) value is defined completely by it's type (no members)
// Register FastDivmod as the integral type
template<>
struct is_integral<cutlass::FastDivmod> : true_type {};
// is_static detects if an (abstract) value is defined completely by its type (no members)
template <class T>
struct is_static : bool_constant<is_empty<remove_cvref_t<T>>::value> {};
include/cute/numeric/math.hpp
+30
View File
@@ -33,6 +33,7 @@
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cutlass/fast_math.h>
namespace cute
{
@@ -323,4 +324,33 @@ log_2(T x) {
return static_cast<int32_t>(bit_width(x)) - 1;
}
template <class IntDiv, class IntMod>
struct DivModReturnType {
IntDiv div_;
IntMod mod_;
CUTE_HOST_DEVICE constexpr
DivModReturnType(IntDiv const& div, IntMod const& mod) : div_(div), mod_(mod) {}
};
// General divmod
template <class CInt0, class CInt1>
CUTE_HOST_DEVICE constexpr
auto
divmod(CInt0 const& a, CInt1 const& b) {
return DivModReturnType{a / b, a % b};
}
// Specialized function with fastDivmod input
template <class CInt>
CUTE_HOST_DEVICE constexpr
auto
divmod(CInt const& a, cutlass::FastDivmod const& b) {
using val_div_type = typename cutlass::FastDivmod::value_div_type;
using val_mod_type = typename cutlass::FastDivmod::value_mod_type;
val_div_type div = 0;
val_mod_type mod = 0;
b(div, mod, a);
return DivModReturnType{div, mod};
}
} // namespace cute
include/cute/stride.hpp
+6 -4
View File
@@ -31,8 +31,9 @@
#pragma once
#include <cute/config.hpp>
#include <cute/int_tuple.hpp>
#include <cute/numeric/int.hpp>
#include <cute/numeric/math.hpp>
namespace cute
{
@@ -79,8 +80,9 @@ crd2idx_itt(CInt const& coord,
return crd2idx(_0{}, get<I0>(shape), get<I0>(stride))
+ (_0{} + ... + crd2idx(_0{}, get<Is>(shape), get<Is>(stride)));
} else { // General case
return crd2idx(coord % product(get<I0>(shape)), get<I0>(shape), get<I0>(stride))
+ crd2idx_itt(coord / product(get<I0>(shape)), shape, stride, seq<Is...>{});
auto [div, mod] = divmod(coord, product(get<I0>(shape)));
return crd2idx(mod, get<I0>(shape), get<I0>(stride))
+ crd2idx_itt(div, shape, stride, seq<Is...>{});
}
CUTE_GCC_UNREACHABLE;
@@ -229,7 +231,7 @@ idx2crd(Index const& idx,
}
} else {
if constexpr (is_tuple<Shape>::value) { // "int" tuple
return idx2crd(idx, shape, compact_col_major(shape));
return transform_leaf(as_arithmetic_tuple(crd2idx(idx, shape, make_basis_like(shape))), identity{});
} else { // "int" "int"
return idx;
}
include/cute/swizzle.hpp
+17 -1
View File
@@ -360,7 +360,7 @@ shiftr(MixedBits<S0,F0> const& m, C<S1> s)
}
//
// upcast and downcast
// Upcast and Downcast
//
template <uint32_t S0, uint32_t F0, auto S1>
@@ -410,6 +410,22 @@ downcast(T const& m)
return m * C<N>{};
}
template <uint32_t S0, uint32_t F0>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(MixedBits<S0,F0> const&)
{
return C<uint32_t(1) << countr_zero(S0 | F0)>{};
}
template <auto v>
CUTE_HOST_DEVICE constexpr
C<v>
max_alignment(C<v> const& c)
{
return c;
}
//
// Convert a Pow2Layout+Coord to a MixedBits
//
include/cute/swizzle_layout.hpp
+32 -7
View File
@@ -36,6 +36,7 @@
#include <cute/layout_composed.hpp>
#include <cute/swizzle.hpp>
#include <cute/pointer_swizzle.hpp> // get_swizzle
/* Specialized functionality for a ComposedLayout of the form
* InvolutionFn o Offset o LayoutB
@@ -56,6 +57,9 @@
namespace cute
{
template <int B, int M, int S, class Offset, class LayoutB>
struct get_swizzle<ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB>> { using type = Swizzle<B,M,S>; };
//
// Constructors
//
@@ -117,7 +121,7 @@ CUTE_HOST_DEVICE constexpr
auto
make_fragment_like(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
{
return detail::transfer_swizzle<B,M,S>(layout.layout_b(), make_fragment_like(layout.layout_b()));
return make_fragment_like(layout.layout_b());
}
//
@@ -441,7 +445,7 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
else if constexpr (scale::num == 1) {
return downcast<scale::den>(swizzle);
}
else if constexpr (scale::den == 1) {
else if constexpr (scale::den == 1) {
return upcast<scale::num>(swizzle);
}
else {
@@ -450,6 +454,24 @@ recast_layout(Swizzle<B,M,S> const& swizzle)
CUTE_GCC_UNREACHABLE;
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(Swizzle<B,M,S> const&)
{
return Int<M>{};
}
template <int B, int M, int S, class Offset, class LayoutB>
CUTE_HOST_DEVICE constexpr
auto
max_alignment(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& layout)
{
return gcd(max_alignment(layout.layout_a()),
max_alignment(layout.offset()),
max_alignment(layout.layout_b()));
}
//
// Other operations
//
@@ -485,7 +507,7 @@ max_common_vector(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& a,
Layout<Shape,Stride> const& b)
{
// This assumes that Offset is in the YZ domain of the Swizzle...
return cute::min(Int<(1 << M)>{}, max_common_vector(a.layout_b(), b));
return cute::min(max_common_vector(a.layout_b(), b), Int<(1 << M)>{});
}
template <class Shape, class Stride, int B, int M, int S, class Offset, class LayoutB>
@@ -504,12 +526,15 @@ auto
max_common_vector(ComposedLayout<Swizzle<B0,M0,S0>,Offset0,LayoutB0> const& a,
ComposedLayout<Swizzle<B1,M1,S1>,Offset1,LayoutB1> const& b)
{
auto result = coalesce(composition(a, right_inverse(b)));
// Typical impl is composition(a, right_inverse(b))
// so this is Sw0 o B0 o rinv(Sw1 o B1) = Sw0 o B0 o rinv(B1) o Sw1
auto vec = max_common_vector(a.layout_b(), b.layout_b());
if constexpr (is_constant<1, decltype(stride<0>(result.layout_b()))>::value) {
return shape<0>(result);
// This assumes that Offset is in the YZ domain of the Swizzle...
if constexpr (Swizzle<B0,M0,S0>{} == Swizzle<B1,M1,S1>{}) {
return vec;
} else {
return Int<1>{};
return cute::min(vec, Int<(1 << M0)>{}, Int<(1 << M1)>{});
}
CUTE_GCC_UNREACHABLE;
include/cute/tensor.hpp
+2 -1047
View File
File diff suppressed because it is too large Load Diff
include/cute/tensor_impl.hpp
+1153
View File
File diff suppressed because it is too large Load Diff
include/cute/util/debug.hpp
+2 -2
View File
@@ -120,7 +120,7 @@ print_type(T&&...) {
CUTE_HOST_DEVICE
bool
block(int bid)
block([[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == bid;
@@ -131,7 +131,7 @@ block(int bid)
CUTE_HOST_DEVICE
bool
thread(int tid, int bid)
thread([[maybe_unused]] int tid, [[maybe_unused]] int bid)
{
#if defined(__CUDA_ARCH__)
return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == tid) && block(bid);
include/cute/util/type_traits.hpp
+13
View File
@@ -85,6 +85,8 @@ using CUTE_STL_NAMESPACE::is_volatile_v;
using CUTE_STL_NAMESPACE::conditional;
using CUTE_STL_NAMESPACE::conditional_t;
using CUTE_STL_NAMESPACE::add_const_t;
using CUTE_STL_NAMESPACE::remove_const_t;
using CUTE_STL_NAMESPACE::remove_cv_t;
using CUTE_STL_NAMESPACE::remove_reference_t;
@@ -107,6 +109,13 @@ using CUTE_STL_NAMESPACE::is_convertible_v;
using CUTE_STL_NAMESPACE::is_same;
using CUTE_STL_NAMESPACE::is_same_v;
using CUTE_STL_NAMESPACE::is_constructible;
using CUTE_STL_NAMESPACE::is_constructible_v;
using CUTE_STL_NAMESPACE::is_default_constructible;
using CUTE_STL_NAMESPACE::is_default_constructible_v;
using CUTE_STL_NAMESPACE::is_standard_layout;
using CUTE_STL_NAMESPACE::is_standard_layout_v;
using CUTE_STL_NAMESPACE::is_arithmetic;
using CUTE_STL_NAMESPACE::is_unsigned;
using CUTE_STL_NAMESPACE::is_unsigned_v;
@@ -131,6 +140,9 @@ using CUTE_STL_NAMESPACE::common_type_t;
using CUTE_STL_NAMESPACE::remove_pointer;
using CUTE_STL_NAMESPACE::remove_pointer_t;
using CUTE_STL_NAMESPACE::alignment_of;
using CUTE_STL_NAMESPACE::alignment_of_v;
// <utility>
using CUTE_STL_NAMESPACE::declval;
@@ -261,4 +273,5 @@ struct conditional_template<false, True, False> {
template <class... U>
using type = False<U...>;
};
} // end namespace cute