CUTLASS 3.1 (#915)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 20:19:34 -07:00
committed by GitHub
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37184 additions and 16419 deletions

View File

@@ -303,7 +303,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % AccessType::kElements) {
if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -388,7 +388,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % AccessType::kElements) {
if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -290,7 +290,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.K % AccessType::kElements) {
if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -288,7 +288,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % AccessType::kElements) {
if ((problem_size.C / problem_size.groups) % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -268,7 +268,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.K % AccessType::kElements) {
if (problem_size.C % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -304,7 +304,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.K % AccessType::kElements) {
if (problem_size.C % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -243,7 +243,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % AccessType::kElements) {
if (problem_size.K % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -293,7 +293,7 @@ public:
static Status can_implement(Conv2dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % AccessType::kElements) {
if (problem_size.K % AccessType::kElements) {
return Status::kErrorInvalidProblem;
}

View File

@@ -270,7 +270,7 @@ public:
static Status can_implement(Conv3dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.K % (128/sizeof_bits<Element>::value)) {
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
return Status::kErrorInvalidProblem;
}

View File

@@ -300,7 +300,7 @@ public:
static Status can_implement(Conv3dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.K % (128/sizeof_bits<Element>::value)) {
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
return Status::kErrorInvalidProblem;
}

View File

@@ -248,7 +248,7 @@ public:
static Status can_implement(Conv3dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
if (problem_size.K % (128/sizeof_bits<Element>::value)) {
return Status::kErrorInvalidProblem;
}

View File

@@ -291,7 +291,7 @@ public:
static Status can_implement(Conv3dProblemSize const &problem_size) {
// check alignment constraint on iterator's contiguous dimension
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
if (problem_size.K % (128/sizeof_bits<Element>::value)) {
return Status::kErrorInvalidProblem;
}

View File

@@ -134,6 +134,12 @@ public:
/// Number of cp.async instructions to load on group of operand B
static int const kAccessesPerGroupB =
(AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations;
// Optional staged-accumulation (e.g., tf32x3 kernels) for improved numerical
// accuracy, where each mainloop iteration first accumulates into a temporary
// set of freshly-cleared accumulators, which are subsequently added to the
// final accumulator set.
static bool const kStagedAccumulation = arch::UseStagedAccumulation<typename Operator::MathOperator>::value;
};
private:
@@ -387,10 +393,7 @@ public:
FragmentC tmp_accum;
if (platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddFastF32>::value
|| platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddComplexFastF32>::value) {
if (Detail::kStagedAccumulation) {
tmp_accum.clear();
}
@@ -444,10 +447,7 @@ public:
copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A,
group_start_iteration_B);
if (platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddFastF32>::value
|| platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddComplexFastF32>::value) {
if (Detail::kStagedAccumulation) {
warp_mma(
tmp_accum,
warp_transformed_frag_A[warp_mma_k % 2],
@@ -518,10 +518,7 @@ public:
}
if (platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddFastF32>::value
|| platform::is_same<typename Operator::MathOperator,
arch::OpMultiplyAddComplexFastF32>::value) {
if (Detail::kStagedAccumulation) {
accum = plus_accum(accum, tmp_accum);
}

View File

@@ -107,7 +107,7 @@ struct StridedDgradHorizontalThreadblockSwizzle :
// compute number of tiles in m dimension
int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m());
// compute number of tiles in n dimension
// compute number of tiles in n dimension
int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
return gemm::GemmCoord(
@@ -148,7 +148,7 @@ struct StridedDgradIdentityThreadblockSwizzle :
// compute number of tiles in m dimension
int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m());
// compute number of tiles in n dimension
// compute number of tiles in n dimension
int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
return gemm::GemmCoord(