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

CUTLASS 3.6.0 (#1850)

Browse Source 
* v3.6

* update changelog

* update readme

* fix typo

* fixing typos

* hopper gemm with weight prefetch

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-10-09 15:33:27 -04:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent 0837a2a00a
commit cc3c29a81a
354 changed files with 105937 additions and 8197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cute/algorithm/clear.hpp
+3 -3
View File
@@ -30,9 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/algorithm/fill.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/algorithm/fill.hpp> // cute::fill
namespace cute
{
include/cute/algorithm/cooperative_copy.hpp
+7 -5
View File
@@ -31,12 +31,14 @@
#pragma once
#include <cute/config.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/algorithm/copy.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/layout.hpp>
#include <cute/layout_composed.hpp> // cute::logical_divide
#include <cute/swizzle.hpp> // cute::Swizzle
#include <cute/swizzle_layout.hpp> // cute::get_nonswizzle_portion
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/tensor_predicate.hpp>
#include <cute/algorithm/copy.hpp>
#include <cute/atom/copy_atom.hpp>
namespace cute
{
include/cute/algorithm/cooperative_gemm.hpp
+2 -2
View File
@@ -434,8 +434,8 @@ cooperative_gemm(uint32_t thread_idx,
static_assert(is_convertible_v<decay_t<invoke_result_t<CStoreTransformOp, TypeC>>, TypeC>,
"CStoreTransformOp functor must accept value of type TC::value_type and return value convertible to type TC::value_type");
static constexpr bool compat = weakly_compatible(tile_shape(TiledMMA<Args...>{}),
make_shape(size<0>(sA), size<0>(sB), size<1>(sA)));
static constexpr bool compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
tile_shape(TiledMMA<Args...>{}));
if constexpr (compat) {
detail::cooperative_gemm_no_predication<SmemCopyOpA, SmemCopyOpB, SmemCopyOpC>(
thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
include/cute/algorithm/copy.hpp
+4 -8
View File
@@ -30,14 +30,10 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/container/alignment.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_predicate.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/tensor_predicate.hpp> // cute::TrivialPredTensor
#include <cute/atom/copy_atom.hpp> // cute::Copy_Atom
namespace cute
{
include/cute/algorithm/functional.hpp
+3 -4
View File
@@ -30,10 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/complex.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/math.hpp> // cute::max, cute::min
#include <cute/numeric/complex.hpp> // cute::conj
/** C++14 <functional> extensions */
include/cute/algorithm/prefetch.hpp
+3 -5
View File
@@ -30,11 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/atom/copy_atom.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/atom/copy_atom.hpp> // cute::Copy_Atom
namespace cute
{
include/cute/algorithm/tuple_algorithms.hpp
+41 -91
View File
@@ -44,7 +44,7 @@
/// Code guidelines and style preferences:
///
/// For perfect forwarding, don't use std::forward, because it may not
/// be defined in device code when compiling with NVRTC. Instead, use
/// be defined in device code when compiling with NVRTC. Instead, use
/// `static_cast<ParameterType&&>(parameter_name)`.
///
/// CuTe generally does not bother forwarding functions, as
@@ -52,24 +52,9 @@
///
/// Throughout CUTLASS, cute::make_tuple always needs to be called
/// namespace-qualified, EVEN If inside the cute namespace and/or in
/// scope of a "using namespace cute" declaration. Otherwise, the
/// scope of a "using namespace cute" declaration. Otherwise, the
/// compiler may select std::make_tuple instead of cute::make_tuple,
/// due to argument-dependent lookup. Two problems may result from
/// that.
///
/// 1. Functions have an unexpected return type (std::tuple instead of
/// cute::tuple), so functions that take cute::tuple parameters
/// fail to compile (generally inside functions that have template
/// parameters expected to be cute::tuple).
///
/// 2. std::tuple does not have the required __host__ __device__
/// markings, so the CUDA compiler complains if you use it in
/// device code.
///
/// cute::make_tuple will occur more often than std::make_tuple would
/// in modern C++ code, because cute::tuple's design deprioritizes
/// correct operation of CTAD (constructor template argument
/// deduction) in favor of implementation simplicity.
/// due to argument-dependent lookup.
namespace cute
{
@@ -145,6 +130,8 @@ transform_apply(T&& t, F&& f, G&& g)
} else {
return g(f(static_cast<T&&>(t)));
}
CUTE_GCC_UNREACHABLE;
}
template <class T0, class T1, class F, class G>
@@ -157,6 +144,8 @@ transform_apply(T0&& t0, T1&& t1, F&& f, G&& g)
} else {
return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1)));
}
CUTE_GCC_UNREACHABLE;
}
template <class T0, class T1, class T2, class F, class G>
@@ -169,6 +158,8 @@ transform_apply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g)
} else {
return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2)));
}
CUTE_GCC_UNREACHABLE;
}
//
@@ -401,71 +392,36 @@ filter_tuple(T0 const& t0, T1 const& t1, T2 const& t2, F&& f)
namespace detail {
// This impl compiles much faster than cute::apply and variadic args
template <class T, class V, class F>
template <class Fn, class Val>
struct FoldAdaptor {
template <class X>
CUTE_HOST_DEVICE constexpr auto operator|(X&& x) {
auto r = fn_(val_, static_cast<X&&>(x));
return FoldAdaptor<Fn, decltype(r)>{fn_, r};
}
Fn fn_;
Val val_;
};
template <class T, class V, class F, int... Is>
CUTE_HOST_DEVICE constexpr
auto
fold(T&&, V&& v, F&&, seq<>)
fold(T&& t, V const& v, F&& f, seq<Is...>)
{
return v;
return (FoldAdaptor<F,V>{f,v} | ... | get<Is>(static_cast<T&&>(t))).val_;
}
template <class T, class V, class F, int I0>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f, seq<I0>)
{
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>
CUTE_HOST_DEVICE constexpr
auto
fold(T&& t, V&& v, F&& f)
fold(T&& t, V const& v, F&& f)
{
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::fold(static_cast<T&&>(t),
static_cast<V&&>(v),
f,
tuple_seq<T>{});
return detail::fold(static_cast<T&&>(t), v, f, tuple_seq<T>{});
} else {
return f(static_cast<V&&>(v), static_cast<T&&>(t));
return f(v, static_cast<T&&>(t));
}
CUTE_GCC_UNREACHABLE;
@@ -477,10 +433,7 @@ auto
fold_first(T&& t, F&& f)
{
if constexpr (is_tuple<remove_cvref_t<T>>::value) {
return detail::fold(static_cast<T&&>(t),
get<0>(static_cast<T&&>(t)),
f,
make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
return detail::fold(static_cast<T&&>(t), get<0>(t), f, make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
} else {
return t;
}
@@ -536,13 +489,23 @@ CUTE_HOST_DEVICE constexpr
auto
take(T const& t)
{
return detail::apply(t, [](auto const&... a) { return cute::make_tuple(a...); }, make_range<B,E>{});
if constexpr (E == -1) {
if constexpr (is_tuple<T>::value) {
return take<B,tuple_size<T>::value>(t);
} else {
return take<B,1>(t);
}
} else
if constexpr (B <= E) {
return detail::apply(t, [](auto const&... a) { return cute::make_tuple(a...); }, make_range<B,E>{});
} else {
static_assert(B <= E);
}
CUTE_GCC_UNREACHABLE;
}
//
// Select tuple elements with given indices.
//
template <int... I, class T>
CUTE_HOST_DEVICE constexpr
auto
@@ -551,19 +514,6 @@ select(T const& t)
return cute::make_tuple(get<I>(t)...);
}
template <class T, class Indices>
CUTE_HOST_DEVICE constexpr
auto
select(T const& t, Indices const& indices)
{
if constexpr (is_tuple<Indices>::value) {
return cute::transform(indices, [&t](auto i) { return select(t, i); });
} else {
static_assert(is_static<Indices>::value, "Order must be static");
return get<Indices::value>(t);
}
}
// Wrap non-tuples into rank-1 tuples or forward
template <class T>
CUTE_HOST_DEVICE constexpr
include/cute/arch/cluster_sm90.hpp
+2 -2
View File
@@ -150,7 +150,7 @@ CUTE_DEVICE dim3 cluster_shape()
}
// Get 1D ctaid in a cluster.
CUTLASS_DEVICE uint32_t block_rank_in_cluster()
CUTE_DEVICE uint32_t block_rank_in_cluster()
{
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
uint32_t rank;
@@ -162,7 +162,7 @@ CUTLASS_DEVICE uint32_t block_rank_in_cluster()
}
// Set the destination block-ID in cluster for a given SMEM Address
CUTLASS_DEVICE uint32_t set_block_rank(uint32_t smemAddr, uint32_t rank)
CUTE_DEVICE uint32_t set_block_rank(uint32_t smemAddr, uint32_t rank)
{
#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED)
uint32_t result;
include/cute/arch/config.hpp
+50
View File
@@ -0,0 +1,50 @@
/***************************************************************************************************
* Copyright (c) 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 <cutlass/arch/config.h> // CUTLASS_ARCH_MMA_SMxx_ENABLED
// TMA instructions
#if defined(CUTLASS_ARCH_MMA_SM90_ENABLED)
# define CUTE_ARCH_TMA_SM90_ENABLED
#endif
#if defined(CUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_ENABLED)
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
#endif
// STSM
#if defined(CUTLASS_ARCH_MMA_SM90_ENABLED)
# define CUTE_ARCH_STSM_SM90_ENABLED
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
include/cute/arch/copy_sm50.hpp
+28 -2
View File
@@ -40,8 +40,8 @@
namespace cute
{
struct SM50_Shuffle_U32_2x2Trans
// Shuffle data between thread pair (0, 1), (2, 3), etc.
struct SM50_Shuffle_U32_2x2Trans_XOR1
{
using SRegisters = uint32_t[2];
using DRegisters = uint32_t[2];
@@ -68,5 +68,31 @@ struct SM50_Shuffle_U32_2x2Trans
}
};
// Shuffle data between thread pair (0, 4), (1, 5), etc.
struct SM50_Shuffle_U32_2x2Trans_XOR4
{
using SRegisters = uint32_t[2];
using DRegisters = uint32_t[2];
CUTE_HOST_DEVICE static void
copy(uint32_t const& src0, uint32_t const& src1, uint32_t& dst0, uint32_t& dst1)
{
#if defined(CUTE_ARCH_WARP_SHUFFLE_ENABLED)
uint32_t x0 = threadIdx.x & 4 ? src0 : src1;
uint32_t y0 = __shfl_xor_sync(0xffffffff, x0, 4);
// Replace detination register with shuffle result.
if (threadIdx.x & 0x4) {
dst0 = y0;
}
else {
dst1 = y0;
}
#else
CUTE_INVALID_CONTROL_PATH("Trying to use __shfl_xor_sync without CUTE_ARCH_WARP_SHUFFLE_ENABLED.");
#endif
}
};
} // end namespace cute
include/cute/arch/copy_sm90.hpp
+2 -13
View File
@@ -30,21 +30,10 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/arch/config.hpp> // CUTE_ARCH_TMA_SMxx_ENABLED
#include <cute/arch/copy.hpp>
// Config
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12))
# define CUTE_ARCH_STSM_SM90_ENABLED
# define CUTE_ARCH_TMA_SM90_ENABLED
#endif
#if defined(CUTE_ARCH_TMA_SM90_ENABLED) && \
((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 3)))
# define CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED
#endif
namespace cute
{
include/cute/arch/copy_sm90_desc.hpp
+30 -23
View File
@@ -30,6 +30,8 @@
**************************************************************************************************/
#pragma once
#include "cutlass/numeric_types.h"
#if !defined(__CUDACC_RTC__)
#include <cuda.h>
#include <cinttypes>
@@ -37,6 +39,8 @@
#include <cute/config.hpp>
#include <cute/arch/util.hpp> // cute::cast_smem_ptr_to_uint
#include <cute/arch/config.hpp> // CUTE_ARCH_TMA_SMxx_ENABLED
#include <cute/arch/copy.hpp>
#include <cute/arch/copy_sm90.hpp>
@@ -134,6 +138,10 @@ enum class SmemSwizzleBits : uint8_t {
B128 = 3,
};
enum class SmemSwizzleBase : uint8_t {
SWIZZLE_BASE_16B = 0,
};
enum class OOBFill : uint8_t {
ZERO = 0,
CONSTANT = 1,
@@ -201,13 +209,21 @@ to_CUtensorMapDataType() {
}
inline CUtensorMapSwizzle
to_CUtensorMapSwizzle(SmemSwizzleBits const& t) {
to_CUtensorMapSwizzle(SmemSwizzleBits const& t, SmemSwizzleBase const& b) {
switch (t) {
default: assert(false && "Unknown SmemSwizzleBits!");
case SmemSwizzleBits::DISABLE: return CU_TENSOR_MAP_SWIZZLE_NONE;
case SmemSwizzleBits::B32: return CU_TENSOR_MAP_SWIZZLE_32B;
case SmemSwizzleBits::B64: return CU_TENSOR_MAP_SWIZZLE_64B;
case SmemSwizzleBits::B128: return CU_TENSOR_MAP_SWIZZLE_128B;
default: assert(false && "Unsupported pair of SmemSwizzleBits and SmemSwizzleBase!");
case SmemSwizzleBits::DISABLE:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 0B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_NONE;
case SmemSwizzleBits::B32:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 32B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_32B;
case SmemSwizzleBits::B64:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 64B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_64B;
case SmemSwizzleBits::B128:
assert((b == SmemSwizzleBase::SWIZZLE_BASE_16B) && "Expected 16B swizzle base for 128B swizzle bits.");
return CU_TENSOR_MAP_SWIZZLE_128B;
}
}
@@ -282,7 +298,7 @@ tma_descriptor_replace_addr_in_global_mem(TmaDescriptor const* desc_ptr,
"tensormap.replace.tile.global_address.global.b1024.b64 [%0], %1;"
:: "l"(gmem_int_desc), "l"(new_desc_addr));
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
@@ -295,15 +311,11 @@ tma_descriptor_replace_addr_in_shared_mem(TmaDescriptor& smem_desc,
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
uint32_t smem_int_desc = cast_smem_ptr_to_uint(&smem_desc);
uint64_t const new_desc_addr = reinterpret_cast<uint64_t>(new_tensor_ptr);
uint64_t const smem_int64_desc = 0;
asm volatile (
"cvt.u64.u32 %0, %1;"
:: "l"(smem_int64_desc), "r"(smem_int_desc));
asm volatile (
"tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;"
:: "l"(smem_int64_desc), "l"(new_desc_addr));
:: "r"(smem_int_desc), "l"(new_desc_addr));
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
@@ -331,7 +343,6 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor
:: "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]));
@@ -339,6 +350,7 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor
"tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 1, %1;"
:: "l"(smem_int64_desc), "l"(prob_stride[2]));
#else
// 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] >> 4));
@@ -347,7 +359,7 @@ tma_descriptor_replace_dims_strides_in_shared_mem(TmaDescriptor
:: "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");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
@@ -366,7 +378,7 @@ tma_descriptor_cp_fence_release(TmaDescriptor const* gmem_desc_ptr, TmaDescripto
"tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;"
:: "l"(gmem_int_desc), "r"(smem_int_desc));
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
@@ -381,7 +393,7 @@ tma_descriptor_fence_release()
#if defined(CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED)
asm volatile ("fence.proxy.tensormap::generic.release.gpu;");
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
@@ -400,13 +412,8 @@ tma_descriptor_fence_acquire(TmaDescriptor const* desc_ptr)
:
: "l"(gmem_int_desc)
: "memory");
asm volatile (
"cvta.global.u64 %0, %0;"
:
: "l"(gmem_int_desc), "l"(gmem_int_desc)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_TMA_SM90_ENABLED and CUDA 12.3");
CUTE_INVALID_CONTROL_PATH("Using TMA Descriptor modification without CUTE_ARCH_DEVICE_MODIFIABLE_TMA_SM90_ENABLED and CUDA 12.3");
#endif
}
include/cute/arch/copy_sm90_tma.hpp
+65 -30
View File
@@ -32,8 +32,11 @@
#include <cute/config.hpp>
#include <cute/arch/config.hpp> // CUTE_ARCH_TMA_SMxx_ENABLED
#include <cute/arch/copy.hpp>
#include <cute/arch/copy_sm90.hpp>
#include "cutlass/arch/synclog.hpp"
namespace cute
{
@@ -52,6 +55,7 @@ struct SM90_TMA_LOAD_1D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3}], [%2], %4;"
@@ -97,6 +101,7 @@ struct SM90_TMA_LOAD_2D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4}], [%2], %5;"
@@ -142,6 +147,7 @@ struct SM90_TMA_LOAD_3D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5}], [%2], %6;"
@@ -187,6 +193,7 @@ struct SM90_TMA_LOAD_4D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6}], [%2], %7;"
@@ -232,6 +239,7 @@ struct SM90_TMA_LOAD_5D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint"
" [%0], [%1, {%3, %4, %5, %6, %7}], [%2], %8;"
@@ -355,6 +363,7 @@ struct SM90_TMA_LOAD_IM2COL_3D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
@@ -405,6 +414,7 @@ struct SM90_TMA_LOAD_IM2COL_4D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
@@ -455,6 +465,7 @@ struct SM90_TMA_LOAD_IM2COL_5D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes"
@@ -565,7 +576,7 @@ struct SM90_TMA_LOAD_IM2COL
struct SM90_TMA_LOAD_MULTICAST_1D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
@@ -573,13 +584,14 @@ struct SM90_TMA_LOAD_MULTICAST_1D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%4}], [%2], %3;"
"cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4}], [%2], %3, %5;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"r"(crd0)
"r"(crd0), "l"(cache_hint)
: "memory");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -590,7 +602,7 @@ struct SM90_TMA_LOAD_MULTICAST_1D
struct SM90_TMA_LOAD_MULTICAST_2D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
@@ -598,13 +610,14 @@ struct SM90_TMA_LOAD_MULTICAST_2D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%4, %5}], [%2], %3;"
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4, %5}], [%2], %3, %6;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"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.");
@@ -615,7 +628,7 @@ struct SM90_TMA_LOAD_MULTICAST_2D
struct SM90_TMA_LOAD_MULTICAST_3D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
@@ -623,13 +636,14 @@ struct SM90_TMA_LOAD_MULTICAST_3D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%4, %5, %6}], [%2], %3;"
"cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4, %5, %6}], [%2], %3, %7;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"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.");
@@ -640,7 +654,7 @@ struct SM90_TMA_LOAD_MULTICAST_3D
struct SM90_TMA_LOAD_MULTICAST_4D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3)
{
@@ -648,13 +662,14 @@ struct SM90_TMA_LOAD_MULTICAST_4D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%4, %5, %6, %7}], [%2], %3;"
"cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4, %5, %6, %7}], [%2], %3, %8;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"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.");
@@ -665,7 +680,7 @@ struct SM90_TMA_LOAD_MULTICAST_4D
struct SM90_TMA_LOAD_MULTICAST_5D
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, 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)
{
@@ -673,13 +688,14 @@ struct SM90_TMA_LOAD_MULTICAST_5D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster"
" [%0], [%1, {%4, %5, %6, %7, %8}], [%2], %3;"
"cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster.L2::cache_hint"
" [%0], [%1, {%4, %5, %6, %7, %8}], [%2], %3, %9;"
:
: "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar),
"h"(multicast_mask),
"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.");
@@ -690,39 +706,39 @@ struct SM90_TMA_LOAD_MULTICAST_5D
struct SM90_TMA_LOAD_MULTICAST
{
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0)
{
return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0);
return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, mbar_ptr, multicast_mask, cache_hint, smem_ptr, crd0);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1)
{
return SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0, crd1);
return SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, mbar_ptr, multicast_mask, cache_hint, smem_ptr, crd0, crd1);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, uint64_t cache_hint,
void * smem_ptr,
int32_t const& crd0, int32_t const& crd1, int32_t const& crd2)
{
return SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0, crd1, crd2);
return SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, mbar_ptr, multicast_mask, cache_hint, smem_ptr, crd0, crd1, crd2);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, 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_MULTICAST_4D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3);
return SM90_TMA_LOAD_MULTICAST_4D::copy(desc_ptr, mbar_ptr, multicast_mask, cache_hint, smem_ptr, crd0, crd1, crd2, crd3);
}
CUTE_HOST_DEVICE static void
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask,
copy(void const* desc_ptr, uint64_t* mbar_ptr, uint16_t multicast_mask, 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_MULTICAST_5D::copy(desc_ptr, mbar_ptr, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3, crd4);
return SM90_TMA_LOAD_MULTICAST_5D::copy(desc_ptr, mbar_ptr, multicast_mask, cache_hint, smem_ptr, crd0, crd1, crd2, crd3, crd4);
}
using PREFETCH = typename SM90_TMA_LOAD::PREFETCH;
@@ -744,6 +760,7 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_3D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes.multicast::cluster"
@@ -772,6 +789,7 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_4D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes.multicast::cluster"
@@ -800,6 +818,7 @@ struct SM90_TMA_LOAD_IM2COL_MULTICAST_5D
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_mbar = cast_smem_ptr_to_uint(mbar_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_load(__LINE__, gmem_int_desc, smem_int_mbar, smem_int_ptr);
// Copy from global to shared::cluster.
asm volatile (
"cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes.multicast::cluster"
@@ -871,6 +890,7 @@ struct SM90_TMA_STORE_1D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [%0, {%2}], [%1];"
:
@@ -893,6 +913,7 @@ struct SM90_TMA_STORE_2D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%2, %3}], [%1];"
:
@@ -915,6 +936,7 @@ struct SM90_TMA_STORE_3D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [%0, {%2, %3, %4}], [%1];"
:
@@ -937,6 +959,7 @@ struct SM90_TMA_STORE_4D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.global.shared::cta.bulk_group [%0, {%2, %3, %4, %5}], [%1];"
:
@@ -959,6 +982,7 @@ struct SM90_TMA_STORE_5D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.bulk_group [%0, {%2, %3, %4, %5, %6}], [%1];"
:
@@ -1024,6 +1048,7 @@ struct SM90_TMA_STORE_IM2COL_3D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.3d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4}], [%1];"
@@ -1047,6 +1072,7 @@ struct SM90_TMA_STORE_IM2COL_4D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.4d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4, %5}], [%1];"
@@ -1070,6 +1096,7 @@ struct SM90_TMA_STORE_IM2COL_5D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.async.bulk.tensor.5d.global.shared::cta.im2col_no_offs.bulk_group"
" [%0, {%2, %3, %4, %5, %6}], [%1];"
@@ -1112,6 +1139,7 @@ struct SM90_TMA_STORE_IM2COL
CUTE_HOST_DEVICE static void
tma_store_fence() {
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
cutlass::arch::synclog_emit_fence_view_async_shared(__LINE__);
asm volatile ("fence.proxy.async.shared::cta;");
#elif defined(__CUDA_ARCH__)
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -1122,6 +1150,7 @@ tma_store_fence() {
CUTE_HOST_DEVICE static void
tma_store_arrive() {
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
cutlass::arch::synclog_emit_tma_store_arrive(__LINE__);
asm volatile("cp.async.bulk.commit_group;");
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
@@ -1138,6 +1167,7 @@ tma_store_wait() {
:
: "n"(Count)
: "memory");
cutlass::arch::synclog_emit_tma_store_wait(__LINE__, Count);
#else
CUTE_INVALID_CONTROL_PATH("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED.");
#endif
@@ -1157,6 +1187,7 @@ struct SM90_TMA_REDUCE_ADD_1D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.reduce.async.bulk.tensor.1d.global.shared::cta.add.bulk_group [%0, {%2}], [%1];"
:
@@ -1179,6 +1210,7 @@ struct SM90_TMA_REDUCE_ADD_2D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.reduce.async.bulk.tensor.2d.global.shared::cta.add.bulk_group [%0, {%2, %3}], [%1];"
:
@@ -1201,6 +1233,7 @@ struct SM90_TMA_REDUCE_ADD_3D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.reduce.async.bulk.tensor.3d.global.shared::cta.add.bulk_group [%0, {%2, %3, %4}], [%1];"
:
@@ -1223,6 +1256,7 @@ struct SM90_TMA_REDUCE_ADD_4D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.reduce.async.bulk.tensor.4d.global.shared::cta.add.bulk_group [%0, {%2, %3, %4, %5}], [%1];"
:
@@ -1245,6 +1279,7 @@ struct SM90_TMA_REDUCE_ADD_5D
#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
cutlass::arch::synclog_emit_tma_store(__LINE__, gmem_int_desc, smem_int_ptr);
asm volatile (
"cp.reduce.async.bulk.tensor.5d.global.shared::cta.add.bulk_group [%0, {%2, %3, %4, %5, %6}], [%1];"
:
include/cute/arch/mma.hpp
+3 -3
View File
@@ -30,9 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/util.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/complex.hpp> // cute::fma
#include <cute/numeric/real.hpp> // cute::fma
namespace cute
{
include/cute/arch/mma_sm90.hpp
+3223 -568
View File
File diff suppressed because it is too large Load Diff
include/cute/arch/mma_sm90_desc.hpp
+3 -4
View File
@@ -48,8 +48,7 @@ namespace cute {
// GMMA Descriptor and utilities
// GMMA enums and utilities
namespace GMMA
{
namespace SM90::GMMA {
enum class LayoutType : uint8_t {
INTERLEAVE = 0,
@@ -81,7 +80,7 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) {
}
#endif // !defined(__CUDACC_RTC__)
} // end namespace GMMA
} // end namespace SM90::GMMA
union GmmaDescriptor
{
@@ -146,7 +145,7 @@ print(GmmaDescriptor const& t)
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_)));
printf(" layout_type: 0x%01x (%s)\n", t.bitfield.layout_type_, to_string(static_cast<SM90::GMMA::LayoutType>(t.bitfield.layout_type_)));
#endif // !defined(__CUDACC_RTC__)
}
include/cute/arch/mma_sm90_gmma.hpp
+2056 -1364
View File
File diff suppressed because it is too large Load Diff
include/cute/arch/mma_sm90_gmma_sparse.hpp
+53789
View File
File diff suppressed because it is too large Load Diff
include/cute/arch/util.hpp
+22 -1
View File
@@ -31,7 +31,6 @@
#pragma once
#include <cute/config.hpp>
#include <cute/numeric/integer_sequence.hpp>
#if defined(__clang__) && defined(__CUDA__)
@@ -254,6 +253,28 @@ explode(Fn fn,
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]..., f[If]...);
}
template <class Fn,
class PtrD, int... Id,
class PtrA, int... Ia,
class PtrB, int... Ib,
class PtrC, int... Ic,
class PtrE, int... Ie,
class PtrF, int... If,
class PtrG, int... Ig>
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...>,
PtrE&& e, int_sequence<Ie...>,
PtrF&& f, int_sequence<If...>,
PtrG&& g, int_sequence<Ig...>)
{
return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]..., f[If]..., g[Ig]...);
}
//
// Utility for exploding tuples into functions
//
include/cute/atom/copy_atom.hpp
+43 -52
View File
@@ -30,16 +30,13 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/copy.hpp>
#include <cute/atom/copy_traits.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/util/type_traits.hpp> // cute::__CUTE_REQUIRES
#include <cute/container/tuple.hpp> // cute::is_tuple
#include <cute/numeric/integral_constant.hpp> // cute::is_constant, cute::is_integral
#include <cute/atom/copy_traits.hpp> // cute::Copy_Traits
#include <cute/atom/mma_atom.hpp> // cute::TiledMMA
namespace cute
{
@@ -651,10 +648,12 @@ print(ThrCopy<TiledCopy, ThrIdx> const& thr_copy)
print(TiledCopy{});
}
template <class... Args>
// TiledCopy to LaTeX TikZ
template <class... Args, class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
auto
print_latex(TiledCopy<Args...> const& copy)
print_latex(TiledCopy<Args...> const& copy,
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
auto [layoutS_MN, thrID_S] = copy.get_layoutS_MN();
auto [layoutD_MN, thrID_D] = copy.get_layoutD_MN();
@@ -663,13 +662,15 @@ print_latex(TiledCopy<Args...> const& copy)
layoutD_MN, thrID_D);
}
// MNK Copy Layout to Latex TIKZ -- 8-value color coded by thread
// MNK Copy Layout to LaTeX TikZ
template <class LayoutS, class ThrIDS,
class LayoutD, class ThrIDD>
class LayoutD, class ThrIDD,
class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
void
print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and tid -> thr_idx
LayoutD const& D, ThrIDD const& TD) // (m,n) -> (tid,vid) and tid -> thr_idx
LayoutD const& D, ThrIDD const& TD, // (m,n) -> (tid,vid) and tid -> thr_idx
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
CUTE_STATIC_ASSERT_V(rank(S) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(D) == Int<2>{});
@@ -677,33 +678,17 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
assert(size<0>(S) == size<0>(D));
assert(size<1>(S) == size<1>(D));
char const* latex_header =
"\\documentclass{standalone}\n"
"\\usepackage{tikz}\n"
"\\usetikzlibrary{external}\n"
"\\tikzexternalize\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n";
char const* latex_footer =
"\\end{tikzpicture}\n"
"\\end{document}\n";
char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
"{rgb,255:red,175;green,255;blue,175}",
"{rgb,255:red,255;green,255;blue,175}",
"{rgb,255:red,255;green,175;blue,175}",
"{rgb,255:red,210;green,210;blue,255}",
"{rgb,255:red,210;green,255;blue,210}",
"{rgb,255:red,255;green,255;blue,210}",
"{rgb,255:red,255;green,210;blue,210}",};
// Header
// Commented prints
printf("%% LayoutS: "); print(S); printf("\n");
printf("%% ThrIDS : "); print(TS); printf("\n");
printf("%% LayoutD: "); print(D); printf("\n");
printf("%% ThrIDD : "); print(TD); printf("\n\n");
printf(latex_header);
// Header
printf("\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
// S starting at 0,0
for (int i = 0; i < size<0>(S); ++i) {
@@ -712,12 +697,22 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
int val_idx = S(i,j) / size(TS);
int thr_idx = TS(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
i, j,
thr_idx, val_idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
0, 0, int(size<0>(S)), int(size<1>(S)));
// S Labels
for (int i = 0, j = -1; i < size<0>(S); ++i) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
}
for (int i = -1, j = 0; j < size<1>(S); ++j) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
}
// D starting at 0,size<1>(S)+3
for (int i = 0; i < size<0>(D); ++i) {
@@ -726,30 +721,26 @@ print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and
int val_idx = D(i,j) / size(TD);
int thr_idx = TD(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
i, j + size<1>(S) + 3,
thr_idx, val_idx);
}
}
// S Labels
for (int i = 0, j = -1; i < size<0>(S); ++i) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
}
for (int j = 0, i = -1; j < size<1>(S); ++j) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
0, int(size<1>(S)+3), int(size<0>(D)), int(size<1>(D)+size<1>(S)+3));
// D Labels
for (int i = 0, j = size<1>(D); i < size<0>(S); ++i) {
for (int i = 0, j = size<1>(D); i < size<0>(D); ++i) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, i);
}
for (int j = 0, i = -1; j < size<1>(D); ++j) {
for (int i = -1, j = 0; j < size<1>(D); ++j) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, j);
}
// Footer
printf(latex_footer);
printf("\\end{tikzpicture}\n"
"\\end{document}\n");
}
} // end namespace cute
include/cute/atom/copy_traits_sm50.hpp
+18 -1
View File
@@ -39,7 +39,7 @@ namespace cute
{
template <>
struct Copy_Traits<SM50_Shuffle_U32_2x2Trans>
struct Copy_Traits<SM50_Shuffle_U32_2x2Trans_XOR1>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_32>;
@@ -55,4 +55,21 @@ struct Copy_Traits<SM50_Shuffle_U32_2x2Trans>
using RefLayout = SrcLayout;
};
template <>
struct Copy_Traits<SM50_Shuffle_U32_2x2Trans_XOR4>
{
// Logical thread id to thread idx (one-thread)
using ThrID = Layout<_32>;
// Map from (src-thr,src-val) to bit
using SrcLayout = Layout<Shape <_32,_64>,
Stride<_64, _1>>;
// Map from (dst-thr,dst-val) to bit
using DstLayout = Layout<Shape <Shape < _4, _2, _4>, Shape<_32, _2>>,
Stride<Stride<_64, _32, _512>,Stride< _1, _256>>>;
// Reference map from (thr,val) to bit
using RefLayout = SrcLayout;
};
} // end namespace cute
include/cute/atom/copy_traits_sm90_im2col.hpp
+7 -5
View File
@@ -450,7 +450,9 @@ make_im2col_tma_copy_desc(
CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_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));
TMA::SmemSwizzleBits swizzle_bits = detail::get_tma_swizzle_bits(smem_swizzle);
TMA::SmemSwizzleBase swizzle_base = detail::get_tma_swizzle_base(smem_swizzle);
CUtensorMapSwizzle tma_swizzle = TMA::to_CUtensorMapSwizzle(swizzle_bits, swizzle_base);
CUresult encode_result = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeIm2col)(
&tma_desc,
@@ -636,11 +638,11 @@ make_tma_atom_im2col(CopyOp,
auto range_c = size<0,0>(tma_layout_vt);
auto range_whdn = size<0,1>(tma_layout_vt);
Tensor gtensor_cwhdn = make_tensor(gtensor.data(),
flatten(make_layout(basis_get(stride<0,0>(tma_layout_vt), gtensor.layout()),
basis_get(stride<0,1>(tma_layout_vt), gtensor.layout()))));
flatten(make_layout(make_layout(basis_get(stride<0,0>(tma_layout_vt), gtensor.shape()),
basis_get(stride<0,0>(tma_layout_vt), gtensor.stride())),
make_layout(basis_get(stride<0,1>(tma_layout_vt), gtensor.shape()),
basis_get(stride<0,1>(tma_layout_vt), gtensor.stride())))));
auto [tma_desc, tma_tensor] = make_im2col_tma_copy_desc(
gtensor_cwhdn,
range_c,
include/cute/atom/copy_traits_sm90_tma.hpp
+21 -8
View File
@@ -41,6 +41,7 @@
#include <cute/algorithm/prefetch.hpp>
#include <cute/numeric/integral_ratio.hpp>
#include <cutlass/cuda_host_adapter.hpp>
namespace cute
@@ -241,15 +242,22 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST, NumBitsPerTMA, AuxParams_>
// Construct an executable SM90_TMA_LOAD_MULTICAST with tma_mbar
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
with(uint64_t& tma_load_mbar, uint16_t const& multicast_mask) const {
return {{}, {&tma_desc_, &tma_load_mbar, multicast_mask}};
with(
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
return {{}, {&tma_desc_, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
}
// Construct an executable SM90_TMA_LOAD_MULTICAST_OP with tma_mbar (temp. overloaded for grouped gemm/ptr array gemm)
CUTE_HOST_DEVICE constexpr
Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
with(TmaDescriptor const* new_tma_desc, uint64_t& tma_load_mbar, uint16_t const& multicast_mask) const {
return {{}, {new_tma_desc, &tma_load_mbar, multicast_mask}};
with(
TmaDescriptor const* new_tma_desc,
uint64_t& tma_load_mbar,
uint16_t const& multicast_mask,
TMA::CacheHintSm90 const& cache_hint = TMA::CacheHintSm90::EVICT_NORMAL) const {
return {{}, {new_tma_desc, &tma_load_mbar, multicast_mask, static_cast<uint64_t>(cache_hint)}};
}
// Generate the TMA coord tensor
@@ -287,7 +295,8 @@ struct Copy_Traits<SM90_TMA_LOAD_MULTICAST_OP, NumBitsPerTMA>
tuple<
TmaDescriptor const*,
uint64_t*, // smem mbarrier
uint16_t // multicast mask
uint16_t, // multicast mask
uint64_t // cache hint
> const opargs_;
};
@@ -684,8 +693,10 @@ construct_tma_gbasis(Tensor<GEngine,GLayout> const& gtensor, // The origin
// TMA parameter checking
//
CUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)),
"TMA requires CTA_Tile and SLayout top-level shape equivalence.");
// CUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)),
// "TMA requires CTA_Tile and SLayout top-level shape equivalence.");
CUTE_STATIC_ASSERT_V(size(slayout) == size(cta_v_map),
"TMA requires CTA_Tile and SLayout top-level size equivalence.");
#if 0
print("gtensor : "); print(gtensor); print("\n");
@@ -983,7 +994,9 @@ make_tma_copy_desc(Tensor<GEngine,GLayout> const& gtensor, // The origin
CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE;
// TMA smem swizzle type
CUtensorMapSwizzle smem_swizzle = TMA::to_CUtensorMapSwizzle(get_tma_swizzle_bits(swizzle));
TMA::SmemSwizzleBits swizzle_bits = get_tma_swizzle_bits(swizzle);
TMA::SmemSwizzleBase swizzle_base = get_tma_swizzle_base(swizzle);
CUtensorMapSwizzle smem_swizzle = TMA::to_CUtensorMapSwizzle(swizzle_bits, swizzle_base);
CUresult result = CUTLASS_CUDA_DRIVER_WRAPPER_CALL(cuTensorMapEncodeTiled)(
&tma_desc,
tma_format,
include/cute/atom/copy_traits_sm90_tma_swizzle.hpp
+22
View File
@@ -68,4 +68,26 @@ get_tma_swizzle_bits(Layout const& layout)
return get_tma_swizzle_bits(get_swizzle_portion(layout));
}
template <int B, int M, int S>
CUTE_HOST_DEVICE constexpr
TMA::SmemSwizzleBase
get_tma_swizzle_base(Swizzle<B,M,S>)
{
if constexpr (M == 4) {
static_assert(0 <= B && B <= 3, "Expected B = 0,1,2, or 3 when M == 4. Unsupported layout swizzle.");
static_assert(S == 3, "Expected S = 3 when M == 4. Unsupported layout swizzle.");
return TMA::SmemSwizzleBase::SWIZZLE_BASE_16B;
}
else {
static_assert(M == 4, "Expected 128b=16B=(2^4)B base swizzle.");
}
}
template <class Layout>
TMA::SmemSwizzleBase
get_tma_swizzle_base(Layout const& layout)
{
return get_tma_swizzle_base(get_swizzle_portion(layout));
}
} // namespace cute::detail
include/cute/atom/mma_atom.hpp
+117 -118
View File
@@ -45,11 +45,12 @@ template <class MMAOperation>
struct MMA_Atom<MMAOperation> : MMA_Atom<MMA_Traits<MMAOperation>>
{};
template <class... Args>
struct MMA_Atom<MMA_Traits<Args...>>
: MMA_Traits<Args...>
template <class MMAOperation, class... Args>
struct MMA_Atom<MMA_Traits<MMAOperation, Args...>>
: MMA_Traits<MMAOperation, Args...>
{
using Traits = MMA_Traits<Args...>;
using MMA_Op = MMAOperation;
using Traits = MMA_Traits<MMAOperation, Args...>;
// Element value types from the MMA_Traits
using ValTypeD = typename Traits::ValTypeD;
@@ -331,7 +332,7 @@ struct TiledMMA : MMA_Atom
make_layout(size<2>(AtomShape_MNK{})));
auto b_tensor = zipped_divide(t_tensor, b_tile); // ((AtomN,AtomK),(RestN,RestK))
// Transform the Atom mode from (N,K) to (Thr,Val)
// Transform the Atom mode from (M,K) to (Thr,Val)
auto tv_tensor = b_tensor.compose(AtomLayoutB_TV{},_); // ((ThrV,FrgV),(RestN,RestK))
// Tile the tensor for the Thread
@@ -733,18 +734,22 @@ print(ThrMMA<TiledMMA, ThrVMNK> const& thr_mma)
print(static_cast<TiledMMA>(thr_mma));
}
template <class... Args>
// MMA Atom to LaTeX TikZ
template <class... Args, class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
void
print_latex(MMA_Atom<Args...> const& mma_atom)
print_latex(MMA_Atom<Args...> const& mma_atom,
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
print_latex(make_tiled_mma(mma_atom));
}
template <class... Args>
// TiledMMA to LaTeX TikZ
template <class... Args, class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
void
print_latex(TiledMMA<Args...> const& mma)
print_latex(TiledMMA<Args...> const& mma,
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
auto layout_and_thrid_C = mma.get_layoutC_MN();
auto layoutC_MN = get<0>(layout_and_thrid_C);
@@ -763,6 +768,109 @@ print_latex(TiledMMA<Args...> const& mma)
layoutB_NK, thrID_B);
}
// MNK MMA Layout to LaTeX TikZ
template <class LayoutC, class ThrIDC,
class LayoutA, class ThrIDA,
class LayoutB, class ThrIDB,
class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
void
print_latex_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx
LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx
LayoutB const& B, ThrIDB const& TB, // (n,k) -> (tid,vid) and tid -> thr_idx
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
assert(size<0>(A) == size<0>(C));
assert(size<0>(B) == size<1>(C));
assert(size<1>(A) == size<1>(B));
// Commented prints
printf("%% LayoutC: "); print(C); printf("\n");
printf("%% ThrIDC : "); print(TC); printf("\n");
printf("%% LayoutA: "); print(A); printf("\n");
printf("%% ThrIDA : "); print(TA); printf("\n");
printf("%% LayoutB: "); print(B); printf("\n");
printf("%% ThrIDB : "); print(TB); printf("\n\n");
// Header
printf("\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
// C starting at 0,0
for (int m = 0; m < size<0>(C); ++m) {
for (int n = 0; n < size<1>(C); ++n) {
int thrid = C(m,n) % size(TC);
int val_idx = C(m,n) / size(TC);
int thr_idx = TC(thrid);
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
m, n,
thr_idx, val_idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
0, 0, int(size<0>(C)), int(size<1>(C)));
// A starting at 0,-size<1>(A)-1
for (int m = 0; m < size<0>(A); ++m) {
for (int k = 0; k < size<1>(A); ++k) {
int thrid = A(m,k) % size(TA);
int val_idx = A(m,k) / size(TA);
int thr_idx = TA(thrid);
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
m, k-1-size<1>(A),
thr_idx, val_idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
0, int(-size<1>(A)-1), int(size<0>(A)), -1);
// A labels
for (int m = 0, k = -1; m < size<0>(A); ++m) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m);
}
for (int m = -1, k = 0; k < size<1>(A); ++k) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k);
}
// B starting at -size<1>(B)-1,0
for (int n = 0; n < size<0>(B); ++n) {
for (int k = 0; k < size<1>(B); ++k) {
int thrid = B(n,k) % size(TB);
int val_idx = B(n,k) / size(TB);
int thr_idx = TB(thrid);
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
k-1-size<1>(B), n,
thr_idx, val_idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
int(-size<1>(B)-1), 0, -1, int(size<0>(B)));
// B labels
for (int n = 0, k = -1; n < size<0>(B); ++n) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n);
}
for (int n = -1, k = 0; k < size<1>(B); ++k) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k);
}
// Footer
printf("\\end{tikzpicture}\n"
"\\end{document}\n");
}
// MNK MMA Layout to console printer
template <class LayoutC, class ThrIDC,
class LayoutA, class ThrIDA,
@@ -819,115 +927,6 @@ print_layout_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and
printf("+\n");
}
// MNK MMA Layout to Latex TIKZ -- 8-value color coded by thread
template <class LayoutC, class ThrIDC,
class LayoutA, class ThrIDA,
class LayoutB, class ThrIDB>
CUTE_HOST_DEVICE
void
print_latex_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx
LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx
LayoutB const& B, ThrIDB const& TB) // (n,k) -> (tid,vid) and tid -> thr_idx
{
CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
assert(size<0>(A) == size<0>(C));
assert(size<0>(B) == size<1>(C));
assert(size<1>(A) == size<1>(B));
char const* latex_header =
"\\documentclass{standalone}\n"
"\\usepackage{tikz}\n"
"\\usetikzlibrary{external}\n"
"\\tikzexternalize\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n";
char const* latex_footer =
"\\end{tikzpicture}\n"
"\\end{document}\n";
char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
"{rgb,255:red,175;green,255;blue,175}",
"{rgb,255:red,255;green,255;blue,175}",
"{rgb,255:red,255;green,175;blue,175}",
"{rgb,255:red,210;green,210;blue,255}",
"{rgb,255:red,210;green,255;blue,210}",
"{rgb,255:red,255;green,255;blue,210}",
"{rgb,255:red,255;green,210;blue,210}"};
// Header
printf("%% LayoutC: "); print(C); printf("\n");
printf("%% ThrIDC : "); print(TC); printf("\n");
printf("%% LayoutA: "); print(A); printf("\n");
printf("%% ThrIDA : "); print(TA); printf("\n");
printf("%% LayoutB: "); print(B); printf("\n");
printf("%% ThrIDB : "); print(TB); printf("\n\n");
printf(latex_header);
// C starting at 0,0
for (int m = 0; m < size<0>(C); ++m) {
for (int n = 0; n < size<1>(C); ++n) {
int thrid = C(m,n) % size(TC);
int val_idx = C(m,n) / size(TC);
int thr_idx = TC(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
m, n,
thr_idx, val_idx);
}
}
// A starting at 0,-size<1>(A)-1
for (int m = 0; m < size<0>(A); ++m) {
for (int k = 0; k < size<1>(A); ++k) {
int thrid = A(m,k) % size(TA);
int val_idx = A(m,k) / size(TA);
int thr_idx = TA(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
m, k-1-size<1>(A),
thr_idx, val_idx);
}
}
// B starting at -size<1>(B)-1,0
for (int n = 0; n < size<0>(B); ++n) {
for (int k = 0; k < size<1>(B); ++k) {
int thrid = B(n,k) % size(TB);
int val_idx = B(n,k) / size(TB);
int thr_idx = TB(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
k-1-size<1>(B), n,
thr_idx, val_idx);
}
}
// A labels
for (int m = 0, k = -1; m < size<0>(A); ++m) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m);
}
for (int k = 0, m = -1; k < size<1>(A); ++k) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k);
}
// B labels
for (int n = 0, k = -1; n < size<0>(B); ++n) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n);
}
for (int k = 0, n = -1; k < size<1>(B); ++k) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k);
}
// Footer
printf(latex_footer);
}
// MNK MMA Layout to SVG -- 8-value color coded by thread
template <class LayoutC, class ThrIDC,
class LayoutA, class ThrIDA,
include/cute/atom/mma_traits.hpp
+36 -75
View File
@@ -30,23 +30,14 @@
**************************************************************************************************/
#pragma once
#include <cute/arch/mma.hpp>
#include <cute/tensor_impl.hpp>
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/pointer.hpp> // cute::is_rmem
#include <cute/arch/mma.hpp> // cute::UniversalFMA
#include <cute/arch/util.hpp> // cute::detail::explode
namespace cute
{
namespace detail {
template <class X, class = void>
struct supports_output_scaling { static constexpr bool value = false; };
template <class X>
struct supports_output_scaling<X, void_t<decltype(declval<X>().accumulate_)>> { static constexpr bool value = true; };
} // end namespace detail
/**
* concept MMA_Traits
* {
@@ -99,17 +90,27 @@ struct MMA_Traits<UniversalFMA<D,A,B,C>>
using CLayout = Layout<Shape<_1,_1>>;
};
// Extract an MMA_Op from an MMA_Traits
template <class MMA_Traits>
struct MMA_Op {};
template <class MMA_Op_Arg, class... Args>
struct MMA_Op<MMA_Traits<MMA_Op_Arg, Args...>> {
using type = MMA_Op_Arg;
};
//
// Generic mma_unpack for any MMA_Traits
//
template <class MMA_Op, class... MMA_Args,
template <class AnyMMATraits,
class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr
void
mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
mma_unpack(AnyMMATraits const& traits,
Tensor<TD, DLayout> & D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
@@ -121,87 +122,47 @@ mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
static_assert(is_rmem<TC>::value, "Expected registers in MMA_Atom::call");
// Register value types from the MMA_Operation register arrays
using MMA_Op = typename MMA_Op<AnyMMATraits>::type;
using RegTypeD = typename remove_extent<typename MMA_Op::DRegisters>::type;
using RegTypeA = typename remove_extent<typename MMA_Op::ARegisters>::type;
using RegTypeB = typename remove_extent<typename MMA_Op::BRegisters>::type;
using RegTypeC = typename remove_extent<typename MMA_Op::CRegisters>::type;
using MMATraits = MMA_Traits<MMA_Op, MMA_Args...>;
[[maybe_unused]] constexpr int RegNumD = extent<typename MMA_Op::DRegisters>::value;
Tensor rA = recast<RegTypeA>(A);
Tensor rB = recast<RegTypeB>(B);
Tensor rD = recast<RegTypeD>(D);
Tensor rC = recast<RegTypeC>(C);
constexpr int RegNumD = extent<typename MMA_Op::DRegisters>::value;
constexpr int RegNumA = extent<typename MMA_Op::ARegisters>::value;
constexpr int RegNumB = extent<typename MMA_Op::BRegisters>::value;
constexpr int RegNumC = extent<typename MMA_Op::CRegisters>::value;
Tensor rA = recast<RegTypeA>(A);
Tensor rB = recast<RegTypeB>(B);
CUTE_STATIC_ASSERT_V(size(rA) == Int<RegNumA>{});
CUTE_STATIC_ASSERT_V(size(rB) == Int<RegNumB>{});
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumD>{});
CUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
if constexpr (is_same<RegTypeD, void>::value)
{
static_assert(is_same<typename TD::value_type, typename TC::value_type>::value, "GMMA C and D value_type must match.");
static_assert(is_same<DLayout, CLayout>::value, "GMMA C and D layouts must match.");
// assert((void*)&C == (void*)&D);
Tensor rC = recast<RegTypeC>(D); // NOTE: D and C are same, so use mutable D
//CUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
if constexpr (detail::supports_output_scaling<MMATraits>::value) {
detail::explode(MMA_Op::fma,
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{},
&(traits.accumulate_), seq<0>{});
}
else {
detail::explode(MMA_Op::fma,
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
}
}
else {
Tensor rD = recast<RegTypeD>(D);
Tensor rC = recast<RegTypeC>(C);
CUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumD>{});
CUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
if constexpr (detail::supports_output_scaling<MMATraits>::value) {
detail::explode(MMA_Op::fma,
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{},
&(traits.accumulate_), seq<0>{});
}
else {
detail::explode(MMA_Op::fma,
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
}
}
detail::explode(MMA_Op::fma,
rD, make_int_sequence<RegNumD>{},
rA, make_int_sequence<RegNumA>{},
rB, make_int_sequence<RegNumB>{},
rC, make_int_sequence<RegNumC>{});
}
//
// Accept mutable temporaries
//
template <class MMA_Op, class... MMA_Args,
template <class AnyMMATraits,
class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout>
CUTE_HOST_DEVICE constexpr
void
mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
Tensor<TD, DLayout> && D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
mma_unpack(AnyMMATraits const& traits,
Tensor<TD, DLayout> && D,
Tensor<TA, ALayout> const& A,
Tensor<TB, BLayout> const& B,
Tensor<TC, CLayout> const& C)
{
mma_unpack(traits, D, A, B, C);
}
include/cute/atom/mma_traits_sm90.hpp
+15 -3
View File
@@ -41,6 +41,8 @@ namespace cute {
//////////////////////// fp64 = fp64 * fp64 + fp64 ////////////////////////////
///////////////////////////////////////////////////////////////////////////////
using SM90_16x8x4_F64F64F64F64_TN = SM90::MMA_16x8x4_F64F64F64F64_TN;
template <>
struct MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
{
@@ -59,6 +61,8 @@ struct MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
};
using SM90_16x8x8_F64F64F64F64_TN = SM90::MMA_16x8x8_F64F64F64F64_TN;
template <>
struct MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
{
@@ -77,6 +81,8 @@ struct MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
};
using SM90_16x8x16_F64F64F64F64_TN = SM90::MMA_16x8x16_F64F64F64F64_TN;
template <>
struct MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
{
@@ -99,9 +105,11 @@ struct MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
//////////////////////// cfp64 = cfp64 * cfp64 + cfp64 ////////////////////////////
///////////////////////////////////////////////////////////////////////////////////
using SM90_16x8x4_C64C64C64C64_TN = SM90::MMA_16x8x4_C64C64C64C64_TN;
template <>
struct MMA_Traits<SM90_16x8x4_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
: MMA_Traits<SM90_16x8x4_F64F64F64F64_TN>
{
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
@@ -109,9 +117,11 @@ struct MMA_Traits<SM90_16x8x4_C64C64C64C64_TN>
using ValTypeC = complex<double>;
};
using SM90_16x8x8_C64C64C64C64_TN = SM90::MMA_16x8x8_C64C64C64C64_TN;
template <>
struct MMA_Traits<SM90_16x8x8_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
: MMA_Traits<SM90_16x8x8_F64F64F64F64_TN>
{
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
@@ -119,9 +129,11 @@ struct MMA_Traits<SM90_16x8x8_C64C64C64C64_TN>
using ValTypeC = complex<double>;
};
using SM90_16x8x16_C64C64C64C64_TN = SM90::MMA_16x8x16_C64C64C64C64_TN;
template <>
struct MMA_Traits<SM90_16x8x16_C64C64C64C64_TN>
: MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
: MMA_Traits<SM90_16x8x16_F64F64F64F64_TN>
{
using ValTypeD = complex<double>;
using ValTypeA = complex<double>;
include/cute/atom/mma_traits_sm90_gmma.hpp
+4262 -50
View File
File diff suppressed because it is too large Load Diff
include/cute/atom/mma_traits_sm90_gmma_sparse.hpp
+16915
View File
File diff suppressed because it is too large Load Diff
include/cute/config.hpp
-13
View File
@@ -142,21 +142,8 @@
# include <iomanip>
#endif
//
// Support
//
#include <cute/util/type_traits.hpp>
//
// Basic types
//
#include <cute/numeric/numeric_types.hpp>
//
// Debugging utilities
//
#include <cute/util/print.hpp>
#include <cute/util/debug.hpp>
include/cute/container/alignment.hpp
+10 -10
View File
@@ -54,17 +54,17 @@ is_byte_aligned(void const* const ptr)
# define CUTE_ALIGNAS(n) alignas(n)
#endif
template <size_t Alignment>
template <size_t Alignment, class Child = void>
struct aligned_struct {};
template <> struct CUTE_ALIGNAS( 1) aligned_struct< 1> {};
template <> struct CUTE_ALIGNAS( 2) aligned_struct< 2> {};
template <> struct CUTE_ALIGNAS( 4) aligned_struct< 4> {};
template <> struct CUTE_ALIGNAS( 8) aligned_struct< 8> {};
template <> struct CUTE_ALIGNAS( 16) aligned_struct< 16> {};
template <> struct CUTE_ALIGNAS( 32) aligned_struct< 32> {};
template <> struct CUTE_ALIGNAS( 64) aligned_struct< 64> {};
template <> struct CUTE_ALIGNAS(128) aligned_struct<128> {};
template <> struct CUTE_ALIGNAS(256) aligned_struct<256> {};
template <class Child> struct CUTE_ALIGNAS( 1) aligned_struct< 1, Child> {};
template <class Child> struct CUTE_ALIGNAS( 2) aligned_struct< 2, Child> {};
template <class Child> struct CUTE_ALIGNAS( 4) aligned_struct< 4, Child> {};
template <class Child> struct CUTE_ALIGNAS( 8) aligned_struct< 8, Child> {};
template <class Child> struct CUTE_ALIGNAS( 16) aligned_struct< 16, Child> {};
template <class Child> struct CUTE_ALIGNAS( 32) aligned_struct< 32, Child> {};
template <class Child> struct CUTE_ALIGNAS( 64) aligned_struct< 64, Child> {};
template <class Child> struct CUTE_ALIGNAS(128) aligned_struct<128, Child> {};
template <class Child> struct CUTE_ALIGNAS(256) aligned_struct<256, Child> {};
} // end namespace cute
include/cute/container/array_aligned.hpp
+2 -2
View File
@@ -30,8 +30,8 @@
**************************************************************************************************/
#pragma once
#include <cute/container/array.hpp>
#include <cute/container/alignment.hpp>
#include <cute/container/alignment.hpp> // CUTE_ALIGNAS
#include <cute/container/array.hpp> // cute::array
namespace cute
{
include/cute/container/array_subbyte.hpp
+14
View File
@@ -181,6 +181,20 @@ public:
}
};
template <class T>
CUTE_HOST_DEVICE
void
print(subbyte_reference<T> ref) {
cute::print(ref.get());
}
template <class T>
CUTE_HOST_DEVICE
void
pretty_print(subbyte_reference<T> ref) {
cute::pretty_print(ref.get());
}
//
// subbyte_iterator
// Random-access iterator over subbyte references
include/cute/container/bit_field.hpp
+2 -2
View File
@@ -35,9 +35,9 @@
#pragma once
#include <cute/config.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/numeric_types.hpp> // uint_bit_t
#include <cute/util/type_traits.hpp> // cute::is_same
namespace cute
{
include/cute/container/cuda_types.hpp
+2 -6
View File
@@ -30,12 +30,8 @@
**************************************************************************************************/
#pragma once
#include <vector_types.h>
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE, CUTE_GCC_UNREACHABLE
#include <cute/numeric/integral_constant.hpp> // cute::integral_constant
namespace cute
{
include/cute/container/tuple.hpp
+11 -2
View File
@@ -634,14 +634,23 @@ template <class Tuple, size_t... Is>
CUTE_HOST_DEVICE void print_tuple(Tuple const& t, index_sequence<Is...>, char s = '(', char e = ')')
{
using cute::print;
print(s); ((void(print(Is == 0 ? '\0' : ',')), void(print(get<Is>(t)))), ...); print(e);
if (sizeof...(Is) == 0) {
print(s);
} else {
((void(print(Is == 0 ? s : ',')), 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 = ')')
{
os << s; (void(os << (Is == 0 ? '\0' : ',') << get<Is>(t)), ...);
if (sizeof...(Is) == 0) {
os << s;
} else {
(void(os << (Is == 0 ? s : ',') << get<Is>(t)), ...);
}
return os << e;
}
#endif // !defined(__CUDACC_RTC__)
include/cute/container/type_list.hpp
+1 -2
View File
@@ -30,8 +30,7 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE, CUTE_STL_NAMESPACE
namespace cute
{
include/cute/int_tuple.hpp
+20 -146
View File
@@ -30,12 +30,11 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/container/tuple.hpp>
#include <cute/container/array.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/container/array.hpp> // cute::array
#include <cute/container/tuple.hpp> // cute::is_tuple
#include <cute/numeric/integral_constant.hpp> // cute::Int
#include <cute/algorithm/tuple_algorithms.hpp> // cute::transform
/** IntTuple is an integer or a tuple of IntTuples.
* This file holds utilities for working with IntTuples,
@@ -92,7 +91,7 @@ template <class IntTuple>
using rank_t = decltype(rank(declval<IntTuple>()));
template <class IntTuple>
static constexpr int rank_v = rank_t<IntTuple>::value;
static constexpr auto rank_v = rank_t<IntTuple>::value;
//
// shape
@@ -212,7 +211,7 @@ template <class Tuple>
using depth_t = decltype(depth(declval<Tuple>()));
template <class Tuple>
static constexpr int depth_v = depth_t<Tuple>::value;
static constexpr auto depth_v = depth_t<Tuple>::value;
//
// product
@@ -276,7 +275,7 @@ size(IntTuple const& a)
}
template <class IntTuple>
static constexpr int size_v = decltype(size(declval<IntTuple>()))::value;
static constexpr auto size_v = decltype(size(declval<IntTuple>()))::value;
//
// sum
@@ -522,68 +521,31 @@ compatible(IntTupleA const& a, IntTupleB const& b)
template <class A, class B>
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
/** Test if Shape A is evenly divided by Tiler B
* @returns Static or dynamic boolean
* @post if result is true_type, then
* size(a) == logical_divide(make_layout(shape(a)),b) will always compile
* and result in true_type.
*/
template <class IntTupleA, class IntTupleB>
template <class Shape, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
weakly_compatible(IntTupleA const& a, IntTupleB const& b)
evenly_divides(Shape const& a, Tiler const& b)
{
if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
if constexpr (tuple_size<IntTupleA>::value != tuple_size<IntTupleB>::value) {
if constexpr (is_tuple<Tiler>::value) {
if constexpr (rank_v<Tiler> > rank_v<Shape>) {
return false_type{};
} else {
return transform_apply(a, b, [](auto const& x, auto const& y) { return weakly_compatible(x,y); },
return transform_apply(b, a, [](auto const& x, auto const& y) { return evenly_divides(y,x); },
[](auto const&... z) { return (true_type{} && ... && z); });
}
} else if constexpr (is_integral<IntTupleA>::value) {
return size(b) % a == Int<0>{};
} else if constexpr (is_integral<IntTupleB>::value) {
return false_type{};
} else {
return weakly_compatible(shape(a), shape(b));
return size(a) == size(b) * size(ceil_div(shape(a), b));
}
CUTE_GCC_UNREACHABLE;
}
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>
@@ -594,7 +556,7 @@ filter_zeros(IntTupleA const& a, IntTupleB const& b)
if constexpr (is_tuple<IntTupleA>::value) {
return transform(a, b, [](auto const& x, auto const& y) { return filter_zeros(x,y); });
} else if constexpr (is_constant<0, IntTupleA>::value) {
return Int<1>{};
return repeat_like(b, Int<1>{});
} else {
return b;
}
@@ -899,92 +861,4 @@ elem_geq(T const& t, U const& u) {
return !elem_less(t, u);
}
namespace detail {
/** Increment a (dynamic) coord lexicographically within a shape
* @pre is_congruent<Coord,Shape>::value
* \code
* auto shape = make_shape(1,2,make_shape(2,3),3);
*
* int i = 0;
* for (auto coord = repeat_like(shape, 0); back(coord) != back(shape); increment(coord, shape)) {
* std::cout << i++ << ": " << coord << std::endl;
* }
* assert(i == size(shape));
* \endcode
*/
template <int I = 0, class Coord, class Shape>
CUTE_HOST_DEVICE constexpr
void
increment(Coord& coord, Shape const& shape)
{
if constexpr (is_integral<Coord>::value) {
++coord;
} else {
increment(get<I>(coord), get<I>(shape));
if constexpr (I+1 < tuple_size<Coord>::value) {
if (back(get<I>(coord)) == back(get<I>(shape))) {
back(get<I>(coord)) = 0;
increment<I+1>(coord, shape);
}
}
}
}
} // end namespace detail
struct ForwardCoordIteratorSentinal
{};
// A forward iterator for a starting coordinate in a shape's domain, and a shape.
// The starting coordinate may be zero but need not necessarily be.
template <class Coord, class Shape>
struct ForwardCoordIterator
{
static_assert(is_congruent<Coord, Shape>::value);
CUTE_HOST_DEVICE constexpr
Coord const& operator*() const { return coord; }
CUTE_HOST_DEVICE constexpr
ForwardCoordIterator& operator++() { detail::increment(coord, shape); return *this; }
// Sentinel for the end of the implied range
CUTE_HOST_DEVICE constexpr
bool operator< (ForwardCoordIteratorSentinal const&) const { return back(coord) < back(shape); }
CUTE_HOST_DEVICE constexpr
bool operator==(ForwardCoordIteratorSentinal const&) const { return back(coord) == back(shape); }
CUTE_HOST_DEVICE constexpr
bool operator!=(ForwardCoordIteratorSentinal const&) const { return back(coord) != back(shape); }
// NOTE: These are expensive, avoid use
CUTE_HOST_DEVICE constexpr
bool operator< (ForwardCoordIterator const& other) const { return colex_less(coord, other.coord); }
CUTE_HOST_DEVICE constexpr
bool operator==(ForwardCoordIterator const& other) const { return coord == other.coord; }
CUTE_HOST_DEVICE constexpr
bool operator!=(ForwardCoordIterator const& other) const { return coord != other.coord; }
Coord coord;
Shape const& shape;
};
// A forward iterator for a coordinate that starts from a provided coordinate
template <class Shape, class Coord>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Coord const& coord, Shape const& shape)
{
return ForwardCoordIterator<Coord,Shape>{coord,shape};
}
// A forward iterator for a coordinate that starts from zero
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Shape const& shape)
{
auto coord = repeat_like(shape, int(0));
return make_coord_iterator(coord, shape);
}
} // end namespace cute
include/cute/layout.hpp
+86 -74
View File
@@ -31,13 +31,13 @@
#pragma once
#include <cute/config.hpp>
#include <cute/underscore.hpp>
#include <cute/int_tuple.hpp>
#include <cute/stride.hpp>
#include <cute/underscore.hpp>
#include <cute/numeric/arithmetic_tuple.hpp>
#include <cute/numeric/integral_ratio.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/integral_ratio.hpp>
#include <cute/numeric/numeric_types.hpp> // cute::sizeof_bits
namespace cute
{
@@ -660,7 +660,7 @@ template <class Layout>
using cosize_t = decltype(cosize(declval<Layout>()));
template <class Layout>
static constexpr int cosize_v = cosize_t<Layout>::value;
static constexpr auto cosize_v = cosize_t<Layout>::value;
// With crd2idx(coord, shape), makes sense to have crd2idx(coord, Layout) as well
template <class Coord, class Shape, class Stride>
@@ -905,6 +905,15 @@ filter_zeros(Layout<Shape,Stride> const& layout)
return make_layout(filter_zeros(layout.stride(), layout.shape()), layout.stride());
}
// Replace the modes in layout that correspond to a 0 at the terminals of trg_profile with a 1-size
template <class Shape, class Stride, class IntTuple>
CUTE_HOST_DEVICE constexpr
auto
filter_zeros(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
{
return make_layout(filter_zeros(trg_profile, layout.shape()), layout.stride());
}
// Remove all of the 0-strides and 1-sizes
// Return 1-shape if empty
template <class Shape, class Stride>
@@ -1350,7 +1359,8 @@ max_common_vector(Layout<ShapeA,StrideA> const& a,
/* Return a layout that distributes ShapeB over ShapeA.
*
* @returns Layout result
* @post softly_compatible(@a b, @a result)
* @post evenly_divides(@a b, size(@a result))
* @post evenly_divides(@a a, @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
@@ -1726,8 +1736,8 @@ tile_to_shape(Layout<Shape,Stride> const& block,
// Assert proper division
if constexpr (is_static<decltype(target_shape)>::value) {
CUTE_STATIC_ASSERT_V(weakly_compatible(block_shape, target_shape),
"tile_to_shape: block shape does not divide the target shape.");
CUTE_STATIC_ASSERT_V(evenly_divides(target_shape, block_shape),
"tile_to_shape: block shape does not divide the target shape.");
}
auto product_shape = ceil_div(target_shape, block_shape);
@@ -1924,92 +1934,97 @@ print_layout(Layout const& layout, ThrID const& thrid) // (m,n) -> (tid,vid) a
printf("+\n");
}
// Generic 2D Layout to Latex printer -- B&W 8-value color coding
template <class LayoutA>
struct TikzColor_White {
CUTE_HOST_DEVICE char const*
operator()(int idx) const {
return "white";
}
};
struct TikzColor_BWx8 {
CUTE_HOST_DEVICE char const*
operator()(int idx) const {
static char const* color_map[8] = {"black!00", "black!40", "black!20", "black!60",
"black!10", "black!50", "black!30", "black!70"};
return color_map[idx % 8];
}
};
struct TikzColor_TV {
CUTE_HOST_DEVICE char const*
operator()(int tid, int vid) const {
static char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
"{rgb,255:red,175;green,255;blue,175}",
"{rgb,255:red,255;green,255;blue,175}",
"{rgb,255:red,255;green,175;blue,175}",
"{rgb,255:red,210;green,210;blue,255}",
"{rgb,255:red,210;green,255;blue,210}",
"{rgb,255:red,255;green,255;blue,210}",
"{rgb,255:red,255;green,210;blue,210}"};
return color_map[tid % 8];
}
};
// Generic 2D Layout to LaTeX printer
template <class LayoutA, class TikzColorFn = TikzColor_BWx8>
CUTE_HOST_DEVICE
void
print_latex(LayoutA const& layout_a)
print_latex(LayoutA const& layout_a, // (m,n) -> idx
TikzColorFn color = {}) // lambda(idx) -> tikz color string
{
CUTE_STATIC_ASSERT_V(rank(layout_a) <= Int<2>{});
auto layout = append<2>(layout_a, Layout<_1,_0>{});
char const* latex_header =
"\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center,font=\\Large}]\n\n";
char const* latex_footer =
"\\end{tikzpicture}\n"
"\\end{document}\n";
char const* color_map[8] = {"black!00",
"black!40",
"black!20",
"black!60",
"black!10",
"black!50",
"black!30",
"black!70"};
// Header
// Commented print(layout)
printf("%% Layout: "); print(layout); printf("\n");
printf(latex_header);
// Header
printf("\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
// Layout
for (int i = 0; i < size<0>(layout); ++i) {
for (int j = 0; j < size<1>(layout); ++j) {
int idx = layout(i,j);
printf("\\node[box,fill=%s] at (%d,%d) {%d};\n",
color_map[idx % 8],
i, j,
idx);
printf("\\node[fill=%s] at (%d,%d) {%d};\n",
color(idx), i, j, idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n",
int(size<0>(layout)), int(size<1>(layout)));
// Labels
for (int i = 0, j = -1; i < size<0>(layout); ++i) {
for (int i = 0, j = -1; i < size<0>(layout); ++i) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
}
for (int j = 0, i = -1; j < size<1>(layout); ++j) {
for (int i = -1, j = 0; j < size<1>(layout); ++j) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
}
// Footer
printf(latex_footer);
printf("\\end{tikzpicture}\n"
"\\end{document}\n");
}
// Generic ThrVal 2D Layout to Latex TIKZ -- 8-value color coded by thread
template <class Layout, class ThrID>
// Generic ThrVal 2D Layout to LaTeX TikZ
template <class Layout, class ThrID, class TikzColorFn = TikzColor_TV>
CUTE_HOST_DEVICE
void
print_latex(Layout const& layout, ThrID const& thr) // (m,n) -> (tid,vid) and tid -> thr_idx
print_latex(Layout const& layout, // (m,n) -> (tid,vid)
ThrID const& thr, // tid -> thr_idx
TikzColorFn color = {}) // lambda(thr_idx,val_idx) -> tikz color string
{
CUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{});
char const* latex_header =
"\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n";
char const* latex_footer =
"\\end{tikzpicture}\n"
"\\end{document}\n";
char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
"{rgb,255:red,175;green,255;blue,175}",
"{rgb,255:red,255;green,255;blue,175}",
"{rgb,255:red,255;green,175;blue,175}",
"{rgb,255:red,210;green,210;blue,255}",
"{rgb,255:red,210;green,255;blue,210}",
"{rgb,255:red,255;green,255;blue,210}",
"{rgb,255:red,255;green,210;blue,210}"};
// Commented prints
printf("%% Layout: "); print(layout); printf("\n");
printf("%% ThrID : "); print(thr); printf("\n");
// Header
printf("%% layout: "); print(layout); printf("\n");
printf("%% thrid: "); print(thr); printf("\n\n");
printf(latex_header);
printf("\\documentclass[convert]{standalone}\n"
"\\usepackage{tikz}\n\n"
"\\begin{document}\n"
"\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
// Layout
for (int i = 0; i < size<0>(layout); ++i) {
@@ -2018,13 +2033,15 @@ print_latex(Layout const& layout, ThrID const& thr) // (m,n) -> (tid,vid) and
int val_idx = layout(i,j) / size(thr);
int thr_idx = thr(thrid);
printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color_map[thr_idx % 8],
printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
color(thr_idx, val_idx),
i, j,
thr_idx, val_idx);
}
}
// Grid
printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n",
int(size<0>(layout)), int(size<1>(layout)));
// Labels
for (int i = 0, j = -1; i < size<0>(layout); ++i) {
printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
@@ -2034,13 +2051,8 @@ print_latex(Layout const& layout, ThrID const& thr) // (m,n) -> (tid,vid) and
}
// Footer
printf(latex_footer);
printf("\\end{tikzpicture}\n"
"\\end{document}\n");
}
} // end namespace cute
//
// Extended Layouts
//
#include <cute/swizzle_layout.hpp>
include/cute/layout_composed.hpp
+3 -3
View File
@@ -30,9 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/layout.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE, CUTE_GCC_UNREACHABLE
#include <cute/layout.hpp> // cute::tuple
#include <cute/numeric/integral_constant.hpp> // cute::true_type, cute::false_type, cute::Int
/* This implements a ComposedLayout of the form
* LayoutA o Offset o LayoutB
include/cute/numeric/arithmetic_tuple.hpp
+6 -6
View File
@@ -197,7 +197,7 @@ struct ArithmeticTupleIterator
ArithmeticTupleIterator(ArithTuple const& coord = {}) : coord_(coord) {}
CUTE_HOST_DEVICE constexpr
ArithTuple const& operator*() const { return coord_; }
ArithTuple operator*() const { return coord_; }
template <class Coord>
CUTE_HOST_DEVICE constexpr
@@ -206,7 +206,7 @@ struct ArithmeticTupleIterator
template <class Coord>
CUTE_HOST_DEVICE constexpr
auto operator+(Coord const& c) const {
return ArithmeticTupleIterator<decltype(coord_ + c)>(coord_ + c);
return ArithmeticTupleIterator<remove_cvref_t<decltype(coord_ + c)>>(coord_ + c);
}
};
@@ -268,13 +268,13 @@ basis_value(SB const& e)
// Apply the N... pack to another Tuple
template <class SB, class Tuple>
CUTE_HOST_DEVICE constexpr auto
basis_get(SB const& e, Tuple const& t)
CUTE_HOST_DEVICE decltype(auto)
basis_get(SB const& e, Tuple&& t)
{
if constexpr (is_scaled_basis<SB>::value) {
return basis_get(e.value(), get<SB::mode()>(t));
return basis_get(e.value(), get<SB::mode()>(static_cast<Tuple&&>(t)));
} else {
return t;
return static_cast<Tuple&&>(t);
}
CUTE_GCC_UNREACHABLE;
}
include/cute/numeric/complex.hpp
+3 -3
View File
@@ -30,9 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cutlass/complex.h>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/numeric_types.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cutlass/complex.h> // cutlass::complexm, cutlass::real, cutlass::imag, cutlass::is_complex
namespace cute
{
include/cute/numeric/int.hpp
+8 -6
View File
@@ -36,7 +36,9 @@
#include <cstdint>
#endif
#include <cutlass/numeric_types.h>
#include <cute/config.hpp> // CUTE_STL_NAMESPACE
#include <cutlass/numeric_types.h> // cutlass::int2b_t, cutlass::int4b_t
namespace cute
{
@@ -53,8 +55,8 @@ using CUTE_STL_NAMESPACE::int32_t;
using CUTE_STL_NAMESPACE::int64_t;
template <int N> struct int_bit;
template <> struct int_bit< 2> { using type = cutlass::int2b_t; };
template <> struct int_bit< 4> { using type = cutlass::int4b_t; };
template <> struct int_bit< 2> { using type = int2_t; };
template <> struct int_bit< 4> { using type = int4_t; };
template <> struct int_bit< 8> { using type = int8_t; };
template <> struct int_bit< 16> { using type = int16_t; };
template <> struct int_bit< 32> { using type = int32_t; };
@@ -83,9 +85,9 @@ using CUTE_STL_NAMESPACE::uint64_t;
using cutlass::uint128_t;
template <int N> struct uint_bit;
template <> struct uint_bit< 1> { using type = cutlass::uint1b_t; };
template <> struct uint_bit< 2> { using type = cutlass::uint2b_t; };
template <> struct uint_bit< 4> { using type = cutlass::uint4b_t; };
template <> struct uint_bit< 1> { using type = uint1_t; };
template <> struct uint_bit< 2> { using type = uint2_t; };
template <> struct uint_bit< 4> { using type = uint4_t; };
template <> struct uint_bit< 8> { using type = uint8_t; };
template <> struct uint_bit< 16> { using type = uint16_t; };
template <> struct uint_bit< 32> { using type = uint32_t; };
include/cute/numeric/integral_constant.hpp
+19 -6
View File
@@ -30,10 +30,9 @@
**************************************************************************************************/
#pragma once
#include "cute/util/print.hpp"
#include "cute/util/type_traits.hpp"
#include "cute/numeric/math.hpp"
#include "cutlass/fast_math.h"
#include <cute/numeric/math.hpp> // cute::max, etc
#include <cute/util/print.hpp> // cute::print
#include <cute/util/type_traits.hpp> // __CUTE_REQUIRES, cute::is_std_integral
namespace cute
{
@@ -65,7 +64,7 @@ struct integral_constant : C<v> {
static constexpr T value = v;
using value_type = T;
// Disambiguate C<v>::operator value_type()
//CUTE_HOST_DEVICE constexpr operator value_type() const noexcept { return value; }
//CUTE_HOST_DEVICE constexpr operator value_type() const noexcept { return value; }
CUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
};
@@ -406,6 +405,20 @@ conditional_return(false_type, TrueType&&, FalseType&& f) {
return static_cast<FalseType&&>(f);
}
template <auto v>
CUTE_HOST_DEVICE constexpr
auto
conditional_return(bool b, C<v> const&, C<v> const&) {
return C<v>{};
}
template <auto v, auto u>
CUTE_HOST_DEVICE constexpr
auto
conditional_return(bool b, C<v> const&, C<u> const&) {
return b ? v : u;
}
// TrueType and FalseType must have a common type
template <class TrueType, class FalseType>
CUTE_HOST_DEVICE constexpr
@@ -435,7 +448,7 @@ static_value()
return Int<Trait::value>{};
} else {
return Trait::value;
}
}
CUTE_GCC_UNREACHABLE;
}
include/cute/numeric/integral_ratio.hpp
+4 -5
View File
@@ -30,11 +30,10 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/math.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/integral_constant.hpp> // cute::false_type, cute::true_type
#include <cute/numeric/math.hpp> // cute::signum
#include <cute/util/type_traits.hpp> // __CUTE_REQUIRES
namespace cute
{
include/cute/numeric/math.hpp
+9 -9
View File
@@ -30,9 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/util/type_traits.hpp> // __CUTE_REQUIRES
#include <cute/util/type_traits.hpp>
#include <cutlass/fast_math.h>
namespace cute
@@ -143,7 +143,7 @@ has_single_bit(T x) {
// bit_width( 0b0111 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr
T
int
bit_width(T x) {
static_assert(is_unsigned<T>::value, "Only to be used for unsigned types.");
constexpr int N = (numeric_limits<T>::digits == 64 ? 6 :
@@ -224,7 +224,7 @@ rotr(T x, int s) {
// countl_zero( 0b00011100 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr
T
int
countl_zero(T x) {
return numeric_limits<T>::digits - bit_width(x);
}
@@ -235,7 +235,7 @@ countl_zero(T x) {
// countl_one( 0b11100011 ) = 3
template <class T>
CUTE_HOST_DEVICE constexpr
T
int
countl_one(T x) {
return countl_zero(~x);
}
@@ -246,7 +246,7 @@ countl_one(T x) {
// countr_zero( 0b00011100 ) = 2
template <class T>
CUTE_HOST_DEVICE constexpr
T
int
countr_zero(T x) {
return x == 0 ? numeric_limits<T>::digits : bit_width(T(x & T(-x))) - 1; // bit_width of the LSB
}
@@ -257,7 +257,7 @@ countr_zero(T x) {
// countr_one( 0b11100011 ) = 2
template <class T>
CUTE_HOST_DEVICE constexpr
T
int
countr_one(T x) {
return countr_zero(~x);
}
@@ -285,7 +285,7 @@ popcount(T x) {
// Computes the result of bitwise left-shift
template <class T>
CUTE_HOST_DEVICE constexpr
T
auto
shiftl(T x, int s) {
return s >= 0 ? (x << s) : (x >> -s);
}
@@ -293,7 +293,7 @@ shiftl(T x, int s) {
// Computes the result of bitwise right-shift
template <class T>
CUTE_HOST_DEVICE constexpr
T
auto
shiftr(T x, int s) {
return s >= 0 ? (x >> s) : (x << -s);
}
include/cute/numeric/numeric_types.hpp
+66 -6
View File
@@ -30,12 +30,11 @@
**************************************************************************************************/
#pragma once
#include <vector_types.h>
#include <cutlass/numeric_types.h>
#include <cutlass/numeric_size.h>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/int.hpp> // cute::int2_t, cute::int4_t, etc
#include <cute/numeric/int.hpp>
#include <cute/numeric/real.hpp>
#include <cutlass/numeric_size.h> // cutlass::sizeof_bits
#include <cutlass/numeric_types.h> // cutlass::float_e4m3_t, cutlass::float_e5m2_t, etc
namespace cute {
@@ -72,4 +71,65 @@ using cutlass::int4b_t;
using cutlass::uint4b_t;
using cutlass::bin1_t;
} // end namespace cute
//
// Print utility
//
CUTE_HOST_DEVICE
void
print(half_t a) {
printf("%f", static_cast<float>(a));
}
CUTE_HOST_DEVICE
void
print(bfloat16_t a) {
printf("%f", static_cast<float>(a));
}
CUTE_HOST_DEVICE
void
print(tfloat32_t a) {
printf("%f", static_cast<float>(a));
}
CUTE_HOST_DEVICE
void
print(float_e4m3_t a) {
printf("%f", static_cast<float>(a));
}
CUTE_HOST_DEVICE
void
print(float_e5m2_t a) {
printf("%f", static_cast<float>(a));
}
CUTE_HOST_DEVICE void
pretty_print(bfloat16_t v) {
printf("%*.2f", 8, float(v));
}
CUTE_HOST_DEVICE void
pretty_print(half_t v) {
printf("%*.2f", 8, float(v));
}
CUTE_HOST_DEVICE void
pretty_print(tfloat32_t v) {
printf("%*.2e", 10, static_cast<float>(v));
}
CUTE_HOST_DEVICE void
pretty_print(float_e4m3_t t) {
printf("%*.2f", 8, static_cast<float>(t));
}
CUTE_HOST_DEVICE void
pretty_print(float_e5m2_t t) {
printf("%*.2f", 8, static_cast<float>(t));
}
} // namespace cute
include/cute/numeric/real.hpp
+18
View File
@@ -35,6 +35,24 @@
namespace cute
{
/// Generic add
template <class C, class A, class B>
CUTE_HOST_DEVICE constexpr
void
add(C& c, A const& a, B const& b)
{
c = a + b;
}
/// Generic multiply
template <class C, class A, class B>
CUTE_HOST_DEVICE constexpr
void
mul(C& c, A const& a, B const& b)
{
c = a * b;
}
/// Generic fused multiply-add
template <class D, class A, class B, class C>
CUTE_HOST_DEVICE constexpr
include/cute/pointer.hpp
+19 -10
View File
@@ -30,17 +30,13 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/pointer_base.hpp> // cute::iter_adaptor
#include <cute/pointer_sparse.hpp>
#include <cute/container/array_subbyte.hpp> // cute::subbyte_iterator
#include <cute/numeric/integral_constant.hpp> // cute::true_type, cute::false_type
#include <cute/numeric/numeric_types.hpp> // sizeof_bits
#include <cute/util/type_traits.hpp>
#include <cute/numeric/numeric_types.hpp> // sizeof_bits
#include <cute/numeric/math.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer_base.hpp>
#include <cute/pointer_swizzle.hpp>
namespace cute
{
@@ -50,6 +46,9 @@ namespace cute
// Subbyte Types: uint2_t, uint4_t, etc
// Requires construction of a subbyte_iterator<T> in order to properly
// resolve each element in byte-addressed memory.
// Sparse Types: sparse_elem<int S, class T>
// A type that holds one physical element meant to represent S number of logical elements.
// Requires construction of a sparse_ptr that emulates access to the S logical elements.
//
template <class NewT>
@@ -57,6 +56,11 @@ CUTE_HOST_DEVICE constexpr
auto
recast_ptr(void* ptr)
{
if constexpr (is_sparse<NewT>::value) {
constexpr int sparsity = NewT::sparsity;
NewT* p = reinterpret_cast<NewT*>(ptr);
return make_sparse_ptr<sparsity>(p);
} else
if constexpr (cute::is_subbyte_v<NewT>) {
return subbyte_iterator<NewT>(ptr);
} else {
@@ -70,6 +74,11 @@ CUTE_HOST_DEVICE constexpr
auto
recast_ptr(void const* ptr)
{
if constexpr (is_sparse<NewT>::value) {
constexpr int sparsity = NewT::sparsity;
NewT const* p = reinterpret_cast<NewT const*>(ptr);
return make_sparse_ptr<sparsity>(p);
} else
if constexpr (cute::is_subbyte_v<NewT>) {
return subbyte_iterator<NewT const>(ptr);
} else {
include/cute/pointer_base.hpp
+3 -4
View File
@@ -30,10 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/numeric_types.hpp> // sizeof_bits
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/numeric_types.hpp> // cute::sizeof_bits
#include <cute/util/type_traits.hpp> // cute::declval, cute::void_t, etc
namespace cute
{
include/cute/pointer_flagged.hpp
+54 -9
View File
@@ -30,15 +30,13 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/arch/util.hpp> // cast_smem_ptr_to_uint
#include <cute/pointer.hpp>
#include <cute/pointer_swizzle.hpp>
#include <cute/swizzle_layout.hpp>
#include <cute/tensor.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/layout_composed.hpp> // cute::ComposedLayout
#include <cute/pointer.hpp> // cute::make_smem_ptr
#include <cute/pointer_sparse.hpp> // cute::is_sparse
#include <cute/pointer_swizzle.hpp> // cute::make_swizzle_ptr
#include <cute/arch/util.hpp> // cute::cast_smem_ptr_to_uint
#include <cute/numeric/integral_constant.hpp> // cute::Int
namespace cute
{
@@ -124,6 +122,47 @@ as_position_independent_swizzle_tensor(Tensor&& tensor)
CUTE_GCC_UNREACHABLE;
}
// A model of a nullptr sparse_ptr<S, smem_ptr<T>> with B == sizeof_bits<T>::value
// That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr
template <int Sparsity, int Bits>
struct smem_sparse_ptr_flag_bits : Int<0> {};
template <int Sparsity>
using smem_sparse_ptr_flag = smem_sparse_ptr_flag_bits<Sparsity, 1>;
// A flagged construction method to transform ComposedLayout
// Make a swizzle pointer tensor and check that the intended type size matches
template <class Iterator, class SwizzleFn, int S, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
make_tensor(Iterator const& ptr,
ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
{
static_assert(is_smem<Iterator>::value, "Expected smem.");
static_assert(is_sparse_ptr<Iterator>::value, "Expected sparse iter");
static_assert(is_sparse<iter_value_t<Iterator>>::value, "Expected sparse elem");
static_assert(S == iter_value_t<Iterator>::sparsity, "Expected sparsity S");
static_assert(B == sizeof_bits<typename iter_value_t<Iterator>::raw_type>::value, "Expected B-bit pointer type");
return make_tensor(make_swizzle_ptr(ptr, layout.layout_a()), layout.layout_b());
}
// NOTE: To preserve smem_ptr_flag_bits under recast ops
template <int N, class SwizzleFn, int S, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
upcast(ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
{
static_assert(dependent_false<SwizzleFn>, "Not implemented for safety");
}
template <int N, class SwizzleFn, int S, int B, class Layout>
CUTE_HOST_DEVICE constexpr
auto
downcast(ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
{
static_assert(dependent_false<SwizzleFn>, "Not implemented for safety");
}
//
// Display utilities
//
@@ -151,4 +190,10 @@ CUTE_HOST_DEVICE void print(smem_ptr_flag_bits<B> ptr)
printf("smem_ptr[%db](unset)", B);
}
template <int S, int B>
CUTE_HOST_DEVICE void print(smem_sparse_ptr_flag_bits<S,B>)
{
printf("smem_sparse<%d>_ptr[%db](unset)", S, B);
}
} // end namespace cute
include/cute/pointer_sparse.hpp
+172
View File
@@ -0,0 +1,172 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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> // CUTE_HOST_DEVICE
#include <cute/pointer_base.hpp> // cute::iter_adaptor
#include <cute/numeric/integral_constant.hpp> // cute::false_type, cute::true_type
#include <cute/numeric/integral_ratio.hpp> // cute::ratio
namespace cute
{
// A data type that holds one physical element meant to represent Sparsity number of logical elements
// This class is purposely not compatible with anything -- know what you're doing if you attempt to use it
template <int Sparsity, class T>
struct sparse_elem
{
static constexpr int sparsity = Sparsity;
using raw_type = T;
T elem_;
CUTE_HOST_DEVICE constexpr
explicit sparse_elem(T const& elem = {}) : elem_(elem) {}
CUTE_HOST_DEVICE constexpr friend bool operator==(sparse_elem const& a, sparse_elem const& b) { return a.elem_ == b.elem_; }
CUTE_HOST_DEVICE constexpr friend bool operator!=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ != b.elem_; }
CUTE_HOST_DEVICE constexpr friend bool operator< (sparse_elem const& a, sparse_elem const& b) { return a.elem_ < b.elem_; }
CUTE_HOST_DEVICE constexpr friend bool operator<=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ <= b.elem_; }
CUTE_HOST_DEVICE constexpr friend bool operator> (sparse_elem const& a, sparse_elem const& b) { return a.elem_ > b.elem_; }
CUTE_HOST_DEVICE constexpr friend bool operator>=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ >= b.elem_; }
};
template <class T>
struct is_sparse : false_type {};
template <class T>
struct is_sparse<T const> : is_sparse<T> {};
template <int S, class T>
struct is_sparse<sparse_elem<S,T>> : true_type {};
template<class T>
static constexpr auto is_sparse_v = is_sparse<T>::value;
// Overload sizeof_bits for sparse_elem.
// Much like subbyte element types, this is the effective number of bits in a sparse_elem
// rather than actual physical bits that may be used in storing one. Also like subbyte element
// types, modified iterators are required to properly index and access sparse_elems.
//
// Defining sizeof_bits like this makes reasonable expressions like N * sizeof_bits_v<E> meaningful
// even when E is subbyte or sparse. However, this also means that sparse_elem can rather easily be
// confused with subbyte elements and special care should be taken with each.
template <int S, class T>
struct sizeof_bits<sparse_elem<S,T>> {
// Simple implementation that conforms to sizeof_bits
//static constexpr auto value = sizeof_bits<T>::value / S;
//static_assert(value != 0, "sizeof_bits=0 detected. Sparsity is larger than width.");
//static_assert((sizeof_bits<T>::value % S) == 0, "Width needs to be a multiple of sparsity.")
// Interesting experiment that allows any sparsity level to be used by potentially presenting
// an integral_ratio rather than size_t. This is valid in most integer expressions as well.
static constexpr auto value = cute::ratio(cute::Int<cute::sizeof_bits_v<T>>{}, cute::Int<S>{});
};
//
// sparse_ptr
//
template <class T, class = void>
struct is_sparse_ptr : false_type {};
template <class T>
struct is_sparse_ptr<T, void_t<typename T::iterator>> : is_sparse_ptr<typename T::iterator> {};
template <int Sparsity, class Iterator>
struct sparse_ptr : iter_adaptor<Iterator, sparse_ptr<Sparsity, Iterator>>
{
using reference = typename iterator_traits<Iterator>::reference;
using element_type = typename iterator_traits<Iterator>::element_type;
using value_type = typename iterator_traits<Iterator>::value_type;
// Sanity, for now
static_assert(is_sparse<value_type>::value, "Enforce sparse value-type");
static_assert(Sparsity == iter_value_t<Iterator>::sparsity, "Enforce sparsity S");
static_assert(not is_sparse_ptr<Iterator>::value, "Enforce sparse singleton");
template <class Index>
CUTE_HOST_DEVICE constexpr
sparse_ptr operator+(Index const& i) const {
// Only allow offset by multiples of the sparsity factor,
// else the misalignments become a bug. E.g. (sparse_ptr<8,I>{} + 7) + 7
// Motivation for subsparse_iterator or generalization of subbyte_iterator?
assert(i % Sparsity == 0);
return {this->get() + i / Sparsity};
}
template <class Index>
CUTE_HOST_DEVICE constexpr
reference operator[](Index const& i) const {
// Allow offset by any value and dereference.
// Not implemented in terms of sparse_ptr::op+()
return *(this->get() + i / Sparsity);
}
};
template <int S, class I>
struct is_sparse_ptr<sparse_ptr<S,I>> : true_type {};
template <int Sparsity, class Iter>
CUTE_HOST_DEVICE constexpr
auto
make_sparse_ptr(Iter const& iter) {
if constexpr (Sparsity == 1) {
return iter;
} else {
return sparse_ptr<Sparsity, Iter>{iter};
}
CUTE_GCC_UNREACHABLE;
}
template <class NewT, int S, class Iter>
CUTE_HOST_DEVICE constexpr
auto
recast_ptr(sparse_ptr<S,Iter> const& ptr) {
static_assert(not is_sparse<NewT>::value);
return recast_ptr<NewT>(ptr.get());
}
//
// Display utilities
//
template <int S, class Iter>
CUTE_HOST_DEVICE void print(sparse_ptr<S,Iter> ptr)
{
printf("sparse<%d>_", S); print(ptr.get());
}
#if !defined(__CUDACC_RTC__)
template <int S, class Iter>
CUTE_HOST std::ostream& operator<<(std::ostream& os, sparse_ptr<S,Iter> ptr)
{
return os << "sparse<" << S << ">_" << ptr.get();
}
#endif
} // end namespace cute
include/cute/pointer_swizzle.hpp
+8 -12
View File
@@ -30,13 +30,11 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp> // iterator_traits
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer_base.hpp>
#include <cute/swizzle.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/pointer_base.hpp> // cute::iter_adaptor
#include <cute/swizzle.hpp> // cute::Swizzle, cute::get_swizzle primary template
#include <cute/util/type_traits.hpp> // cute::iterator_traits
#include <cute/container/array_subbyte.hpp> // cute::subbyte_iterator
/* This implements a swizzle pointer of the form
* InvolutionFn o PtrAdd
@@ -107,16 +105,14 @@ struct swizzle_ptr : iter_adaptor<Iterator,swizzle_ptr<SwizzleFn,Iterator>>
}
};
template <class T, class = void> // Default No-Swizzle
struct get_swizzle { using type = Swizzle<0,4,3>; };
//
// Helper Function
//
template <class SwizzleFn, class P> // Found the SwizzleFn
struct get_swizzle<swizzle_ptr<SwizzleFn,P>> { using type = SwizzleFn; };
template <class T> // Recurse into anything with a ::iterator
struct get_swizzle<T, void_t<typename T::iterator>> : get_swizzle<typename T::iterator> {};
template <class Iter>
using get_swizzle_t = typename get_swizzle<Iter>::type;
template <class Iterator, class SwizzleFn>
CUTE_HOST_DEVICE constexpr
swizzle_ptr<SwizzleFn,Iterator>
include/cute/stride.hpp
+126 -5
View File
@@ -30,10 +30,16 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/int_tuple.hpp>
#include <cute/numeric/int.hpp>
#include <cute/numeric/math.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/util/type_traits.hpp> // cute::__CUTE_REQUIRES
#include <cute/container/tuple.hpp> // cute::is_tuple
#include <cute/numeric/integral_constant.hpp> // cute::is_integral
#include <cute/numeric/integer_sequence.hpp> // cute::seq
#include <cute/numeric/math.hpp> // cute::divmod
#include <cute/numeric/arithmetic_tuple.hpp> // cute::basis_get
#include <cute/algorithm/functional.hpp> // cute::identity
#include <cute/algorithm/tuple_algorithms.hpp> // cute::fold
#include <cute/int_tuple.hpp> // cute::is_congruent
namespace cute
{
@@ -433,7 +439,7 @@ compact_order(Shape const& shape, Order const& order)
auto flat_order = flatten_to_tuple(order);
// Find the largest static element of order
auto max_order = cute::fold(flat_order, Int<0>{}, [](auto v, auto order) {
if constexpr (is_constant<true, decltype(v < order)>::value) {
if constexpr (is_constant<true, decltype(v < order)>::value) {
return order;
} else {
return v;
@@ -474,4 +480,119 @@ compact_order(Shape const& shape, GenRowMajor const& major)
return compact_major<LayoutRight>(shape);
}
//
// Coordinate iterator
//
namespace detail {
template <class Coord, class Shape, class Order>
CUTE_HOST_DEVICE constexpr
void
increment(Coord& coord, Shape const& shape, Order const& order)
{
++basis_get(get<0>(order), coord);
cute::for_each(make_range<1, tuple_size<Order>::value>{}, [&](auto i){
if (basis_get(get<i-1>(order), coord) == basis_get(get<i-1>(order), shape)) {
basis_get(get<i-1>(order), coord) = 0;
++basis_get(get<i>(order), coord);
}
});
}
/** Increment a (dynamic) coord colexicographically within a shape
* @pre is_congruent<Coord,Shape>::value
* \code
* auto shape = make_shape(1,2,make_shape(2,3),3);
* auto coord = repeat_like(shape, 0);
*
* for (int i = 0; i < size(shape); ++i) {
* std::cout << i << ": " << coord << std::endl;
* increment(coord, shape);
* }
* \endcode
*/
template <class Coord, class Shape>
CUTE_HOST_DEVICE constexpr
void
increment(Coord& coord, Shape const& shape)
{
increment(coord, shape, flatten_to_tuple(make_basis_like(shape)));
}
} // end namespace detail
struct ForwardCoordIteratorSentinel
{};
// A forward iterator for a starting coordinate in a shape's domain, and a shape.
// The starting coordinate may be zero but need not necessarily be.
template <class Coord, class Shape, class Order>
struct ForwardCoordIterator
{
static_assert(is_congruent<Coord, Shape>::value);
CUTE_HOST_DEVICE constexpr
Coord const& operator*() const { return coord; }
CUTE_HOST_DEVICE constexpr
ForwardCoordIterator& operator++() { detail::increment(coord, shape, Order{}); return *this; }
// Sentinel for the end of the implied range
CUTE_HOST_DEVICE constexpr
bool operator==(ForwardCoordIteratorSentinel const&) const { return basis_get(back(Order{}), coord) == basis_get(back(Order{}), shape); }
CUTE_HOST_DEVICE constexpr
bool operator!=(ForwardCoordIteratorSentinel const&) const { return basis_get(back(Order{}), coord) != basis_get(back(Order{}), shape); }
// NOTE: These are expensive, avoid use
CUTE_HOST_DEVICE constexpr
bool operator==(ForwardCoordIterator const& other) const { return coord == other.coord; }
CUTE_HOST_DEVICE constexpr
bool operator!=(ForwardCoordIterator const& other) const { return coord != other.coord; }
Coord coord;
Shape const& shape;
};
// A forward iterator for a coordinate that starts from a provided coordinate and increments in a prescribed order
template <class Order, class Shape, class Coord>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Coord const& coord, Shape const& shape)
{
static_assert(is_congruent<Coord, Shape>::value);
static_assert(is_congruent<Order, Coord>::value);
static_assert(is_congruent<Order, Shape>::value);
auto flat_order = flatten_to_tuple(Order{});
auto inv_order = transform(make_seq<rank(flat_order)>{}, [&](auto i){ return find(flat_order, i); });
auto basis_order = transform_leaf(inv_order, [&](auto i) { return get<i>(flatten_to_tuple(make_basis_like(shape))); });
return ForwardCoordIterator<Coord,Shape,decltype(basis_order)>{coord,shape};
}
// A forward iterator for a coordinate that starts from a provided coordinate and increments colex
template <class Shape, class Coord>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Coord const& coord, Shape const& shape)
{
static_assert(is_congruent<Coord, Shape>::value);
auto basis_order = flatten_to_tuple(make_basis_like(shape));
return ForwardCoordIterator<Coord,Shape,decltype(basis_order)>{coord,shape};
}
// A forward iterator for a coordinate that starts from zero and increments in a prescribed order
template <class Order, class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Shape const& shape)
{
return make_coord_iterator<Order>(repeat_like(shape, int(0)), shape);
}
// A forward iterator for a coordinate that starts from zero and increments colex
template <class Shape>
CUTE_HOST_DEVICE constexpr
auto
make_coord_iterator(Shape const& shape)
{
return make_coord_iterator(repeat_like(shape, int(0)), shape);
}
} // end namespace cute
include/cute/swizzle.hpp
+14 -7
View File
@@ -30,13 +30,11 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/container/tuple.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
#include <cute/numeric/integer_sequence.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/math.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/container/tuple.hpp> // cute::is_tuple
#include <cute/numeric/integral_constant.hpp> // cute::constant
#include <cute/numeric/math.hpp> // cute::max, cute::min
#include <cute/algorithm/tuple_algorithms.hpp> // cute::transform_apply
namespace cute
{
@@ -488,4 +486,13 @@ CUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits<S,F> const& m)
}
#endif // !defined(__CUDACC_RTC__)
//
// Helper Function
//
template <class T, class = void> // Default No-Swizzle
struct get_swizzle { using type = Swizzle<0,4,3>; };
template <class T>
using get_swizzle_t = typename get_swizzle<T>::type;
} // end namespace cute
include/cute/swizzle_layout.hpp
+20 -22
View File
@@ -30,13 +30,10 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/layout.hpp>
#include <cute/layout_composed.hpp>
#include <cute/swizzle.hpp>
#include <cute/pointer_swizzle.hpp> // get_swizzle
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/layout.hpp> // cute::Layout
#include <cute/layout_composed.hpp> // cute::ComposedLayout
#include <cute/swizzle.hpp> // cute::Swizzle, cute::get_swizzle primary template
/* Specialized functionality for a ComposedLayout of the form
* InvolutionFn o Offset o LayoutB
@@ -57,6 +54,9 @@
namespace cute
{
//
// Helper Function
//
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>; };
@@ -193,7 +193,7 @@ make_swizzle_strides(true_type,
// 0 Z DC
// 1 -Z DC
return cute::make_tuple(conditional_return((offset & (Y << Int<I>{})) == Int<0>{}, Z << Int<I>{}, -(Z << Int<I>{}))...);
return cute::make_tuple(conditional_return((offset & (Y << Int<I>{})) == Int<0>{}, Z * Int<(1 << I)>{}, -Z * Int<(1 << I)>{})...);
}
template <class IntZ, class IntY, class Offset, int... I>
@@ -214,7 +214,7 @@ make_swizzle_strides(false_type,
// 0 Y+Z Y-Z
// 1 DC DC
return cute::make_tuple(conditional_return((offset & (Z << Int<I>{})) == Int<0>{}, (Y+Z) << Int<I>{}, (Y-Z) << Int<I>{})...);
return cute::make_tuple(conditional_return((offset & (Z << Int<I>{})) == Int<0>{}, (Y+Z) * Int<(1 << I)>{}, (Y-Z) * Int<(1 << I)>{})...);
}
} // end namespace detail
@@ -240,16 +240,6 @@ slice_and_offset(Coord const& coord, ComposedLayout<Swizzle<B,M,S>,Offset,Layout
// The portion of the layout that is not yet consumed
auto sliced_layout = slice(coord, layout.layout_b());
// If the sliced_layout hits two bits that are swizzled together, then don't attempt to decay
// Compose with the layout to get the swizzle projection, P o L [The Z and Y contributing portions of L]
// (this also tests that shape/stride of layout compose with swizzle)
auto sliced_layout_only_zy = composition(swizzle_only_zy, sliced_layout);
// Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*)
auto swizzle_active_bits = sliced_layout_only_zy(size(sliced_layout_only_zy)-Int<1>{});
// Determine if any active bits collide under the swizzle
auto hit_ZandY = !(swizzle_active_bits & ~layout.layout_a()(swizzle_active_bits));
// The portion of the layout that we are consuming now
auto diced_layout = dice(coord, layout.layout_b());
auto diced_coord = dice(coord, coord);
@@ -269,8 +259,16 @@ slice_and_offset(Coord const& coord, ComposedLayout<Swizzle<B,M,S>,Offset,Layout
// If Layout's codomain hits on Y XOR Z, then it's dynamic-normal
// If Layout's codomain hits on neither Y NOR Z, then it's static-normal
// Test the sliced layout for hit_X & hit_Y for potential decay
if constexpr (is_constant<false, decltype(hit_ZandY)>::value)
// If the sliced_layout hits two bits that are swizzled together, then don't attempt to decay
// Compose with the layout to get the swizzle projection, P o L [The Z and Y contributing portions of L]
// (this also tests that shape/stride of layout compose with swizzle)
auto sliced_layout_only_zy = composition(swizzle_only_zy, sliced_layout);
// Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*)
[[maybe_unused]] auto swizzle_active_bits = sliced_layout_only_zy(size(sliced_layout_only_zy)-Int<1>{});
// Determine if any active bits collide under the swizzle for potential decay
if constexpr (is_constant<0, decltype(not (swizzle_active_bits & ~swizzle(swizzle_active_bits)))>::value)
{ // Hits on Y AND Z, so it's not reducible
return cute::make_tuple(composition(swizzle, offset_only_zy, sliced_layout), offset_anti_zy);
} else
@@ -459,7 +457,7 @@ CUTE_HOST_DEVICE constexpr
auto
max_alignment(Swizzle<B,M,S> const&)
{
return Int<M>{};
return Int<1 << M>{};
}
template <int B, int M, int S, class Offset, class LayoutB>
include/cute/tensor.hpp
+3
View File
@@ -37,7 +37,10 @@
//
#include <cute/pointer_swizzle.hpp>
#include <cute/pointer_sparse.hpp>
#include <cute/pointer_flagged.hpp>
#include <cute/tensor_zip.hpp>
//
// Tensor Algorithms
//
include/cute/tensor_impl.hpp
+61 -21
View File
@@ -41,18 +41,16 @@
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/integer_sequence.hpp>
#include <cute/container/tuple.hpp>
#include <cute/container/array_aligned.hpp>
#include <cute/container/array_subbyte.hpp>
#include <cute/pointer.hpp>
#include <cute/layout.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/layout.hpp> // cute::Shape
#include <cute/layout_composed.hpp> // cute::is_composed_layout
#include <cute/pointer.hpp> // cute::recast_ptr
#include <cute/pointer_base.hpp> // cute::iterator_traits
#include <cute/container/array_aligned.hpp> // cute::array_aligned
#include <cute/container/array_subbyte.hpp> // cute::array_subbyte
#include <cute/container/tuple.hpp> // cute::tuple
#include <cute/numeric/integral_constant.hpp> // cute::is_integral
#include <cute/util/type_traits.hpp> // __CUTE_REQUIRES
namespace cute
{
@@ -69,7 +67,7 @@ namespace cute
// iterator begin();
// };
template <class T, int N>
template <class T, size_t N>
struct ArrayEngine
{
using Storage = typename conditional<(sizeof_bits<T>::value % 8 == 0),
@@ -85,6 +83,24 @@ struct ArrayEngine
CUTE_HOST_DEVICE constexpr auto begin() { return storage_.begin(); }
};
// Specialization for sparse_elem<S,T> tensor allocation/iteration
template <int S, class T, size_t N>
struct ArrayEngine<sparse_elem<S,T>, N>
{
static_assert(N % S == 0, "Expected a multiple of the sparsity.");
using value_type = sparse_elem<S,T>;
using Storage = typename conditional<(sizeof_bits<T>::value % 8 == 0),
array_aligned<T,N/S>,
array_subbyte<T,N/S>>::type;
using iterator = sparse_ptr<S,sparse_elem<S,T>*>;
using reference = typename iterator_traits<iterator>::reference;
using element_type = typename iterator_traits<iterator>::element_type;
Storage storage_;
CUTE_HOST_DEVICE constexpr auto begin() const { return recast_ptr<value_type>(storage_.begin()); }
CUTE_HOST_DEVICE constexpr auto begin() { return recast_ptr<value_type>(storage_.begin()); }
};
template <class Iterator>
struct ViewEngine
{
@@ -622,6 +638,30 @@ filter_zeros(Tensor<Engine,Layout>&& tensor) {
return make_tensor(tensor.data(), filter_zeros(tensor.layout()));
}
template <class Engine, class Layout, class Profile>
CUTE_HOST_DEVICE constexpr
auto
filter_zeros(Tensor<Engine,Layout> const& tensor, Profile const& profile)
{
return make_tensor(tensor.data(), filter_zeros(tensor.layout(), profile));
}
template <class Engine, class Layout, class Profile>
CUTE_HOST_DEVICE constexpr
auto
filter_zeros(Tensor<Engine,Layout>& tensor, Profile const& profile)
{
return make_tensor(tensor.data(), filter_zeros(tensor.layout(), profile));
}
template <class Engine, class Layout, class Profile>
CUTE_HOST_DEVICE constexpr
auto
filter_zeros(Tensor<Engine,Layout>&& tensor, Profile const& profile)
{
return make_tensor(tensor.data(), filter_zeros(tensor.layout(), profile));
}
// Remove all of the 0-strides and 1-sizes
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
@@ -755,10 +795,10 @@ auto
max_common_vector(Tensor<SrcEngine,SrcLayout> const& a,
Tensor<DstEngine,DstLayout> const& b)
{
using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
using SrcRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using DstRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using SrcType = typename SrcEngine::value_type;
using SrcRef = typename SrcEngine::reference;
using DstType = typename DstEngine::value_type;
using DstRef = typename DstEngine::reference;
// Determine if vectorization candidates at all
if constexpr (// Should be the same value_types, else the copy is also performing a cast
@@ -795,10 +835,10 @@ auto
max_common_layout(Tensor<SrcEngine,SrcLayout> const& a,
Tensor<DstEngine,DstLayout> const& b)
{
using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
using SrcRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using DstRef = typename Tensor<SrcEngine,SrcLayout>::reference;
using SrcType = typename SrcEngine::value_type;
using SrcRef = typename SrcEngine::reference;
using DstType = typename DstEngine::value_type;
using DstRef = typename DstEngine::reference;
// Determine if vectorization candidates at all
if constexpr (// Should be the same value_types, else the copy is also performing a cast
include/cute/tensor_predicate.hpp
+2 -3
View File
@@ -30,9 +30,8 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/numeric/integral_constant.hpp> // cute::true_type
namespace cute
{
include/cute/tensor_zip.hpp
+243
View File
@@ -0,0 +1,243 @@
/***************************************************************************************************
* Copyright (c) 2023 - 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> // CUTE_HOST_DEVICE
#include <cute/tensor_impl.hpp> // cute::Tensor
#include <cute/container/tuple.hpp> // cute::tuple
namespace cute
{
// A tuple of Iterators that can be offset asymmetrically
// Note that this only accepts op+(tuple<Index...>) and op[tuple<Index...>]
// where each iterator will be offset by its respective index only.
// READ-ONLY for now until cute::tuple can be constructed with references.
template <class... Iters>
struct ZipIterator
{
using value_type = cute::tuple<iter_value_t<Iters>...>;
using element_type = cute::tuple<iter_element_t<Iters>...>;
// NOTE: cute::tuple does not support constructions with references at the moment.
// Consider fixes and/or an implementation of std::forward_as_tuple.
// For now, use a cute::tuple of value_types instead, which makes this Iterator READ-ONLY.
//using reference = cute::tuple<iter_reference_t<Iters>...>;
using reference = value_type;
ZipIterator() = delete;
CUTE_HOST_DEVICE constexpr
ZipIterator(Iters... iters)
: iters_(iters...)
{}
CUTE_HOST_DEVICE constexpr
ZipIterator(cute::tuple<Iters...> const& iters)
: iters_(iters)
{}
CUTE_HOST_DEVICE constexpr
reference operator*() const {
return cute::apply(iters_, [](auto&&... args) { return reference(*args...); });
}
template <class... Index>
CUTE_HOST_DEVICE constexpr
ZipIterator operator+(cute::tuple<Index...> const& idxs) const {
static_assert(sizeof...(Index) == sizeof...(Iters), "Expect same number of offsets as iterators.");
return cute::transform(iters_, idxs, [](auto&& iter, auto&& idx) { return iter + idx; });
}
template <class... Index>
CUTE_HOST_DEVICE constexpr
reference operator[](cute::tuple<Index...> const& idxs) const {
return *(*this + idxs);
}
cute::tuple<Iters...> iters_;
};
//------------------------------------------------------------------------------
// type traits
template <class... Iters>
struct is_rmem<ZipIterator<Iters...>> : conjunction<is_rmem<Iters>...> {};
template <class... Iters>
struct is_smem<ZipIterator<Iters...>> : conjunction<is_smem<Iters>...> {};
template <class... Iters>
struct is_gmem<ZipIterator<Iters...>> : conjunction<is_gmem<Iters>...> {};
// A tuple of Layouts that operates on each Layout symmetrically
// The Layouts need to have compatible shapes and ranks.
// The ZipLayout presents the intersection of the domain of its component Layouts.
// E.g. all Layouts accept 1D coords and ZipLayout does as well.
// The ZipLayout returns the union of the codomain of its component Layouts.
// E.g. all Layouts return an integer so ZipLayout returns a tuple of integers.
template <class... Layouts>
struct ZipLayout
{
static constexpr int rank = (int(0) | ... | Layouts::rank);
static_assert((is_layout<Layouts>::value && ...), "All template parameters must be layouts");
static_assert(((Layouts::rank == rank) && ...), "All layouts must have the same rank");
CUTE_HOST_DEVICE constexpr
ZipLayout(Layouts const&... layouts)
: layouts_(layouts...)
{}
CUTE_HOST_DEVICE constexpr
ZipLayout(cute::tuple<Layouts...> const& layouts)
: layouts_(layouts)
{}
template <class Coord>
CUTE_HOST_DEVICE constexpr
auto
operator()(Coord const& coord) const {
if constexpr (has_underscore<Coord>::value) {
return ZipLayout(cute::transform(layouts_, [&] (auto layout) { return layout(coord); }));
} else {
return cute::transform(layouts_, [&] (auto layout) { return layout(coord); });
}
CUTE_GCC_UNREACHABLE;
}
// op() convenience function for multi-dimensional coordinates
template <class Coord0, class Coord1, class... Coords>
CUTE_HOST_DEVICE constexpr
decltype(auto)
operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const {
return operator()(make_coord(c0,c1,cs...));
}
cute::tuple<Layouts...> layouts_;
};
template <class... Layouts>
struct is_layout<ZipLayout<Layouts...>> : true_type {};
//
// make_zip_tensor and unzip_tensor
//
template <class... Engines, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
make_zip_tensor(Tensor<Engines,Layouts> const&... tensors)
{
return make_tensor(ZipIterator(tensors.data()...),
ZipLayout(tensors.layout()...));
}
template <class Engine, class Layout>
CUTE_HOST_DEVICE constexpr
auto
unzip_tensor(Tensor<Engine,Layout> const& tensor)
{
return cute::transform(tensor.data().iters_, tensor.layout().layouts_,
[](auto iter, auto layout) { return make_tensor(iter, layout); });
}
//
// Utilities
//
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
rank(ZipLayout<Layouts...> const& layouts)
{
return rank<Is...>(get<0>(layouts.layouts_));
}
template <int... Is, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
size(ZipLayout<Layouts...> const& layouts)
{
return size<Is...>(get<0>(layouts.layouts_));
}
//
// Manipulation
//
// Extend each component layout to rank-N by appending Layout @a x.
template <int N, class... Layouts, class ShapeX = _1, class StrideX = _0>
CUTE_HOST_DEVICE constexpr
auto
append(ZipLayout<Layouts...> const& layouts,
Layout<ShapeX,StrideX> const& x = {})
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return append<N>(t, x); }));
}
// Extend each component layout to rank-N by prepending Layout @a x.
template <int N, class... Layouts, class ShapeX = _1, class StrideX = _0>
CUTE_HOST_DEVICE constexpr
auto
prepend(ZipLayout<Layouts...> const& layouts,
Layout<ShapeX,StrideX> const& x = {})
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return prepend<N>(t, x); }));
}
template <class... Layouts, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
logical_divide(ZipLayout<Layouts...> const& layouts,
Tiler const& tiler)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return logical_divide(t, tiler); }));
}
template <class... Layouts, class Tiler>
CUTE_HOST_DEVICE constexpr
auto
zipped_divide(ZipLayout<Layouts...> const& layouts,
Tiler const& tiler)
{
return ZipLayout(cute::transform(layouts.layouts_, [&](auto t){ return zipped_divide(t, tiler); }));
}
// Return <SlicedZipLayout, ZipOffsets> by calling slice_and_offset and all component layouts.
template <class Coord, class... Layouts>
CUTE_HOST_DEVICE constexpr
auto
slice_and_offset(Coord const& c, ZipLayout<Layouts...> const& layouts)
{
auto result = cute::zip(cute::transform(layouts.layouts_, [&c](auto const& layout) { return slice_and_offset(c, layout); }));
return cute::make_tuple(ZipLayout(get<0>(result)), get<1>(result));
}
} // end namespace cute
include/cute/underscore.hpp
+3 -6
View File
@@ -30,12 +30,9 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/container/tuple.hpp>
#include <cute/algorithm/tuple_algorithms.hpp>
#include <cute/numeric/integral_constant.hpp>
#include <cute/numeric/integer_sequence.hpp>
#include <cute/config.hpp> // CUTE_INLINE_CONSTANT, CUTE_HOST_DEVICE
#include <cute/container/tuple.hpp> // cute::is_tuple
#include <cute/numeric/integral_constant.hpp> // cute::false_type, cute::true_type
namespace cute
{
include/cute/util/print.hpp
+15 -21
View File
@@ -30,9 +30,8 @@
**************************************************************************************************/
#pragma once
#include <cute/config.hpp>
#include <cute/util/type_traits.hpp>
#include <cute/config.hpp> // CUTE_HOST_DEVICE
#include <cute/util/type_traits.hpp> // cute::is_valid
//
// CUDA compatible print and printf
@@ -156,50 +155,45 @@ print(char const* format) {
// pretty printing
//
template <class T>
CUTE_HOST_DEVICE void
pretty_print(T const& v) {
printf(" "); print(v);
}
CUTE_HOST_DEVICE void
pretty_print(bool const& v) {
pretty_print(bool v) {
printf("%*d", 3, int(v));
}
CUTE_HOST_DEVICE void
pretty_print(int32_t const& v) {
pretty_print(int32_t v) {
printf("%*d", 5, v);
}
CUTE_HOST_DEVICE void
pretty_print(uint32_t const& v) {
pretty_print(uint32_t v) {
printf("%*d", 5, v);
}
CUTE_HOST_DEVICE void
pretty_print(int64_t const& v) {
pretty_print(int64_t v) {
printf("%*lld", 5, static_cast<long long>(v));
}
CUTE_HOST_DEVICE void
pretty_print(uint64_t const& v) {
pretty_print(uint64_t v) {
printf("%*llu", 5, static_cast<unsigned long long>(v));
}
CUTE_HOST_DEVICE void
pretty_print(half_t const& v) {
printf("%*.2f", 8, float(v));
}
CUTE_HOST_DEVICE void
pretty_print(float const& v) {
pretty_print(float v) {
printf("%*.2e", 10, v);
}
CUTE_HOST_DEVICE void
pretty_print(double const& v) {
pretty_print(double v) {
printf("%*.3e", 11, v);
}
template <class T>
CUTE_HOST_DEVICE void
pretty_print(T t) {
printf(" "); print(t);
}
} // end namespace cute
include/cute/util/type_traits.hpp
+6 -5
View File
@@ -44,7 +44,7 @@
#include <limits> // numeric_limits
#endif
#include <cute/config.hpp>
#include <cute/config.hpp> // CUTE_STL_NAMESPACE
namespace cute
{
@@ -79,6 +79,7 @@ using CUTE_STL_NAMESPACE::is_const_v;
using CUTE_STL_NAMESPACE::is_volatile;
using CUTE_STL_NAMESPACE::is_volatile_v;
// Defined in cute/numeric/integral_constant.hpp
// using CUTE_STL_NAMESPACE::true_type;
// using CUTE_STL_NAMESPACE::false_type;
@@ -278,14 +279,14 @@ struct conditional_template<false, True, False> {
// is_any_of
//
/// Member `value` is true if and only if T is same as (is_same_v) at least one of the types in Us
template <typename T, typename... Us>
// Member `value` is true if and only if T is same as (is_same_v) at least one of the types in Us
template <class T, class... Us>
struct is_any_of {
constexpr static bool value = (... || CUTE_STL_NAMESPACE::is_same_v<T, Us>);
};
/// Is true if and only if T is same as (is_same_v) at least one of the types in Us
template <typename T, typename... Us>
// Is true if and only if T is same as (is_same_v) at least one of the types in Us
template <class T, class... Us>
inline constexpr bool is_any_of_v = is_any_of<T, Us...>::value;
} // end namespace cute