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

3.6.0 update (#2005)

Browse Source 
* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-12-25 01:34:40 -05:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent e1cd8c7866
commit 3d261a5974
258 changed files with 10863 additions and 3883 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp
+95 -5
View File
@@ -189,7 +189,7 @@ private:
-problem_shape.dilation[NumSpatialDimensions-1-i] :
problem_shape.dilation[NumSpatialDimensions-1-i];
}
return make_im2col_tma_copy(
GmemTiledCopyA{},
tensor_a,
@@ -225,7 +225,7 @@ private:
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
auto lower_srt = detail::compute_lower_srt(problem_shape);
return make_im2col_tma_copy(
GmemTiledCopyB{},
tensor_b,
@@ -372,6 +372,96 @@ public:
return false;
}
if (is_im2col_A || is_im2col_B) {
// Check valid corner values for TMA_LOAD_IM2COL, signed int ranging from [-corner_limit, corner_limit - 1]
constexpr int32_t corner_limit = 1 << (16 / NumSpatialDimensions - 1);
auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
for (int i = 0; i < problem_shape.RankS; ++i) {
implementable = implementable && lower_corner_whd[i] >= -corner_limit && lower_corner_whd[i] <= (corner_limit - 1);
}
auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
for (int i = 0; i < problem_shape.RankS; ++i) {
implementable = implementable && upper_corner_whd[i] >= -corner_limit && upper_corner_whd[i] <= (corner_limit - 1);
}
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Padding values don't meet requirements for TMA LOAD IM2COL.\n");
return false;
}
}
// Wgrad kernels don't support non-packed output strides, non-packed tensor A stride (linearized)
if constexpr (ConvOp == conv::Operator::kWgrad) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
std::ostringstream os;
#endif
const auto & input_shape = problem_shape.shape_A;
const auto & input_stride = problem_shape.stride_A;
implementable &= input_stride[ProblemShape::RankT - 1] == 1;
int input_shape_size = 1;
for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
input_shape_size *= input_shape[i + 1];
implementable &= input_stride[i] == input_shape_size;
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
if (input_stride[i] != input_shape_size) {
os << "\n *** input_stride[" << i << "] = " << input_stride[i] << " != input_shape_size = " << input_shape_size << " ***";
}
#endif
}
if (!implementable) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
os << "\n input_shape_size: " << input_shape_size
<< "\n input_shape: " << input_shape
<< "\n input_stride: " << input_stride
<< "\n";
#endif
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed input strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
return false;
}
const auto & output_shape = problem_shape.shape_C;
const auto & output_stride = problem_shape.stride_C;
implementable &= output_stride[ProblemShape::RankT - 1] == 1;
int output_shape_size = 1;
for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
output_shape_size *= output_shape[i + 1];
implementable &= output_stride[i] == output_shape_size;
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
if (output_stride[i] != output_shape_size) {
os << "\n *** output_stride[" << i << "] = " << output_stride[i] << " != output_shape_size = " << output_shape_size << " ***";
}
#endif
}
if (!implementable) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
os << "\n output_shape_size: " << input_shape_size
<< "\n output_shape: " << input_shape
<< "\n output_stride: " << input_stride
<< "\n";
#endif
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Wgrad kernels don't support non-packed output strides.\n");
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
CUTLASS_TRACE_HOST(os.str());
#endif
return false;
}
}
// Conv kernels only support cross correlation mode currently.
implementable &= problem_shape.mode == cutlass::conv::Mode::kCrossCorrelation;
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Conv kernels only support cross correlation mode currently.\n");
return false;
}
if (problem_shape.groups > 1) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: This kernel does not support conv groups > 1.\n");
return false;
@@ -516,9 +606,9 @@ public:
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
* then would just be acquired since the phase was
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(smem_pipe_producer_state);
@@ -645,7 +735,7 @@ public:
k_tile_count -= prologue_mma_count;
smem_pipe_release.advance(k_tile_count);
// Wait on all GMMAs to complete
warpgroup_wait<0>();
include/cutlass/conv/convnd_problem_shape.hpp
+40
View File
@@ -319,6 +319,7 @@ struct ConvProblemShape {
// | ShapeB | KTRSC | KTRSC | NDHWC |
// | ShapeC | NZPQK | NDHWC | KTRSC |
//
// Input comes from calculate_xformed_act, which does NOT depend on ConvOp.
CUTLASS_HOST_DEVICE
constexpr void
set_shape_stride_ABC(
@@ -328,6 +329,31 @@ struct ConvProblemShape {
TensorStride stride_flt,
TensorExtent shape_xformed_act,
TensorStride stride_xformed_act) {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
printf("*** set_shape_stride_ABC ***");
printf("\n shape_act: ");
print(shape_act);
printf("\n stride_act: ");
print(stride_act);
printf("\n shape_flt: ");
print(shape_flt);
printf("\n stride_flt: ");
print(stride_flt);
printf("\n shape_xformed_act: ");
print(shape_xformed_act);
printf("\n stride_xformed_act: ");
print(stride_xformed_act);
if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
printf("\n ConvOp: Fprop");
}
if constexpr (ConvOp == cutlass::conv::Operator::kDgrad) {
printf("\n ConvOp: Dgrad");
}
if constexpr (ConvOp == cutlass::conv::Operator::kWgrad) {
printf("\n ConvOp: Wgrad");
}
printf("\n");
#endif
if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
shape_A = shape_act;
@@ -353,6 +379,20 @@ struct ConvProblemShape {
shape_C = shape_flt;
stride_C = stride_flt;
}
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
printf("\n shape_A: ");
print(shape_A);
printf("\n stride_A: ");
print(stride_A);
printf("\n shape_B: ");
print(shape_B);
printf("\n stride_B: ");
print(stride_B);
printf("\n shape_C: ");
print(shape_C);
printf("\n stride_C: ");
print(stride_C);
#endif
}
// Get A extents.
include/cutlass/conv/kernel/direct_convolution.h
+2 -1
View File
@@ -40,6 +40,7 @@
#include "cutlass/array.h"
#include "cutlass/numeric_types.h"
#include "cutlass/matrix_shape.h"
#include "cutlass/platform/platform.h"
#include "cutlass/semaphore.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/layout/tensor.h"
@@ -155,7 +156,7 @@ struct DirectConvolutionParams {
swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
// Dynamic SMEM usage because stride and dilation are runtime params.
smem_size_ = (max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
smem_size_ = (cutlass::platform::max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
}
CUTLASS_HOST_DEVICE