Updates and Bug fixes to CUTLASS 3.3 (#1232)

This commit is contained in:
Pradeep Ramani
2023-12-05 06:50:49 -08:00
committed by GitHub
parent 4a1709e17e
commit e9e30c2304
31 changed files with 534 additions and 227 deletions

View File

@@ -42,8 +42,8 @@
#if defined(__CUDACC_RTC__)
#include <cuda/std/type_traits>
#else
#include <cstdio>
#include <type_traits>
#include <cstdio>
#endif
#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)))

View File

@@ -154,7 +154,7 @@ sm90_get_smem_store_op_for_accumulator() {
}
else {
// auto-vectorizing store
return DefaultCopy{};
return AutoVectorizingCopyWithAssumedAlignment{};
}
}
@@ -175,7 +175,7 @@ sm90_get_smem_load_op_for_source() {
}
else {
// auto-vectorizing load
return DefaultCopy{};
return AutoVectorizingCopyWithAssumedAlignment{};
}
}

View File

@@ -438,7 +438,7 @@ struct Sm90ReLUAuxStore {
using VecType = uint_bit_t<V>;
Tensor tC_rAux_vec = recast<VecType>(tC_rAux);
Tensor tC_gAux_vec = recast<VecType>(tC_gAux);
Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{}));
Tensor tC_cAux_vec = tC_cAux.compose(make_layout(Int<size(tC_rAux_vec)>{}, Int<V>{})); // only works if vector is logically sequential
auto predicate_fn = [&] (auto&&... coords) { return elem_less(tC_cAux_vec(coords...), residue_mn); };
copy_if(FunctionPredTensor(predicate_fn), tC_rAux_vec, tC_gAux_vec);
}
@@ -662,7 +662,7 @@ struct Sm90AuxLoad<
}
if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) { // (partially) in-bounds CTA tile
copy(tC_gAux, tC_rAux);
copy_aligned(tC_gAux, tC_rAux);
}
}
}
@@ -677,7 +677,7 @@ struct Sm90AuxLoad<
}
if (elem_less(repeat_like(residue_mn, _0{}), residue_mn)) {
copy(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
copy_aligned(tC_gAux(_,_,_,epi_m,epi_n), tC_rAux);
}
}
}

View File

@@ -641,8 +641,9 @@ struct Sm90RowBroadcast {
if (epi_m == 0) { // Assumes M-major subtile loop
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
int bcast_pipe_index = (load_iteration / EpiTiles) % Stages;
copy(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
copy_aligned(filter(tCsRow(_,_,_,epi_m,epi_n,bcast_pipe_index)), filter(tCrRow));
}
}
@@ -774,7 +775,8 @@ struct Sm90ColBroadcast {
}
// Filter so we don't issue redundant copies over stride-0 modes
copy(filter(tCgCol), filter(tCrCol));
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCgCol), filter(tCrCol));
}
template <typename ElementAccumulator, int FragmentSize>

View File

@@ -915,8 +915,9 @@ public:
using ElementGmem = conditional_t<FinalReduction, ElementCompute volatile, ElementCompute>;
Tensor tCgBuf = sm90_partition_for_epilogue<ReferenceSrc>(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
copy(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
// Filter so we don't issue redundant copies over stride-0 modes
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
}
sync_fn();
}
@@ -934,7 +935,8 @@ public:
Tensor tCsBuf = sm90_partition_for_epilogue<ReferenceSrc>(sBuf(_,_,get<1>(warp_mn)), epi_tile, tiled_copy, thread_idx);
if (is_reduced_lane) {
// Filter so we don't issue redunant copies over stride-0 modes
copy(filter(tCrCol), filter(tCsBuf));
// (only works if 0-strides are in same location, which is by construction)
copy_aligned(filter(tCrCol), filter(tCsBuf));
}
sync_fn();

View File

@@ -231,7 +231,7 @@ int ceil_div(int a, int b) {
* log2_up/down codes?
*/
template <typename value_t>
CUTLASS_HOST_DEVICE int clz(value_t x) {
CUTLASS_HOST_DEVICE value_t clz(value_t x) {
for (int i = 31; i >= 0; --i) {
if ((1 << i) & x)
return value_t(31 - i);

View File

@@ -1,4 +1,4 @@
/**************************************************************************************************
/***************************************************************************************************
* Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
@@ -217,7 +217,7 @@ struct alignas(1) float8_base {
// Extract the bits in the FP32 type
uint8_t sign = uint8_t((s >> 24 & 0x80));
int32_t exp = int32_t((s >> FP32_NUM_MANTISSA_BITS) & 0xff) - FP32_EXPONENT_BIAS;
uint32_t mantissa = s & 0x7fffff;
int mantissa = s & 0x7fffff;
uint8_t u = 0;
uint8_t const kF8_NaN = 0x7f;

View File

@@ -375,6 +375,7 @@ public:
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Static initializers
/////////////////////////////////////////////////////////////////////////////////////////////////

View File

@@ -347,7 +347,7 @@ public:
// The number of tiles for which reduction is required is either:
// (a) the total number of output tiles (in the case of split-K)
// (b) the number of stream-K tiles
// To calcualte the total number of output tiles in the split-K case, we
// To calculate the total number of output tiles in the split-K case, we
// note that, in the split-K case, the units_per_problem_ member of Params will be
// the total number of output tiles.
auto reduction_tiles = params.splits_ > 1 ? params.units_per_problem_ : params.sk_tiles_;