@@ -55,12 +55,29 @@ namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
typename ThreadMap_,
|
||||
conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity
|
||||
>
|
||||
class Conv2dDgradFilterTileAccessIteratorAnalytic;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Conv2dDgradFilterTileAccessIteratorAnalytic strided dgrad needs special handling to skip MMAs
|
||||
// on non-contributing w positions
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
typename ThreadMap_
|
||||
>
|
||||
class Conv2dDgradFilterTileAccessIteratorAnalytic {
|
||||
class Conv2dDgradFilterTileAccessIteratorAnalytic <
|
||||
Shape_,
|
||||
Element_,
|
||||
ThreadMap_,
|
||||
conv::StrideSupport::kStrided
|
||||
> {
|
||||
public:
|
||||
|
||||
//
|
||||
@@ -90,6 +107,197 @@ public:
|
||||
|
||||
using Params = Conv2dAnalyticParams<Layout>;
|
||||
|
||||
private:
|
||||
|
||||
Params const ¶ms_;
|
||||
Conv2dProblemSize const &problem_size_;
|
||||
LongIndex iteration_contiguous_;
|
||||
LongIndex iteration_strided_;
|
||||
char const *pointer_;
|
||||
|
||||
// For a fixed filter position (r,s) find and fill offset_k_, offset_c_ in strided and contiguous dimension
|
||||
int filter_r_;
|
||||
int filter_s_;
|
||||
int start_r_;
|
||||
int start_s_;
|
||||
int offset_k_[ThreadMap::Iterations::kStrided];
|
||||
int offset_c_[ThreadMap::Iterations::kContiguous];
|
||||
|
||||
public:
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv2dDgradFilterTileAccessIteratorAnalytic(
|
||||
Params const ¶ms,
|
||||
Conv2dProblemSize const &problem_size,
|
||||
Element const *ptr,
|
||||
int thread_idx,
|
||||
int start_r, int start_s,
|
||||
MatrixCoord const &threadblock_offset = MatrixCoord()
|
||||
):
|
||||
params_(params),
|
||||
problem_size_(problem_size),
|
||||
pointer_(reinterpret_cast<char const *>(ptr)),
|
||||
filter_r_(start_r),
|
||||
filter_s_(start_s),
|
||||
start_r_(start_r),
|
||||
start_s_(start_s) {
|
||||
|
||||
layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
|
||||
offset_c_[c] = threadblock_offset.column() + thread_coord.contiguous()
|
||||
+ c * ThreadMap::Delta::kContiguous;
|
||||
}
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
offset_k_[s] =
|
||||
threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided;
|
||||
}
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(Index index) {
|
||||
iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous;
|
||||
iteration_strided_ = index / ThreadMap::Iterations::kContiguous;
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
pointer_ += pointer_offset * sizeof_bits<Element>::value / 8;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void advance() {
|
||||
// Moves filter_s
|
||||
filter_s_ += problem_size_.stride_w;
|
||||
if (filter_s_ < problem_size_.S) {
|
||||
return;
|
||||
}
|
||||
// Restore filter_s
|
||||
filter_s_ = start_s_;
|
||||
|
||||
// Move filter_r
|
||||
filter_r_ += problem_size_.stride_h;
|
||||
if (filter_r_ < problem_size_.R) {
|
||||
return;
|
||||
}
|
||||
// Restore filter_r
|
||||
filter_r_ = start_r_;
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
offset_k_[s] += Shape::kRow * problem_size_.split_k_slices;
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the coordinate in the filter tensor w that is currently pointed to
|
||||
/// by the iterator.
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCoord at() const {
|
||||
|
||||
int c = offset_c_[iteration_contiguous_];
|
||||
int k = offset_k_[iteration_strided_];
|
||||
|
||||
return TensorCoord(k, filter_r_, filter_s_, c);
|
||||
}
|
||||
|
||||
/// Returns true if the current coordinate is within the filter tensor w
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool valid() const {
|
||||
|
||||
TensorCoord coord = at();
|
||||
|
||||
return coord.n() < problem_size_.K && coord.c() < problem_size_.C;
|
||||
}
|
||||
|
||||
/// Returns a pointer to the vector starting at the current coordinate
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType const *get() const {
|
||||
|
||||
TensorCoord coord = at();
|
||||
LongIndex offset = params_.layout(coord);
|
||||
|
||||
return reinterpret_cast<AccessType const *>(pointer_ + offset * sizeof_bits<Element>::value / 8);
|
||||
|
||||
}
|
||||
|
||||
/// Increments to the next memory access
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv2dDgradFilterTileAccessIteratorAnalytic &operator++() {
|
||||
++iteration_contiguous_;
|
||||
if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) {
|
||||
return *this;
|
||||
}
|
||||
iteration_contiguous_ = 0;
|
||||
++iteration_strided_;
|
||||
if (iteration_strided_ < ThreadMap::Iterations::kStrided) {
|
||||
return *this;
|
||||
}
|
||||
iteration_strided_ = 0;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Determines whether the Implicit GEMM can execute the given problem.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static Status can_implement(Conv2dProblemSize const &problem_size) {
|
||||
|
||||
// check alignment constraint on iterator's contiguous dimension
|
||||
if (problem_size.C % (128/sizeof_bits<Element>::value)) {
|
||||
return Status::kErrorInvalidProblem;
|
||||
}
|
||||
|
||||
return Status::kSuccess;
|
||||
}
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Conv2dDgradFilterTileAccessIteratorAnalytic unity strided dgrad is more performant for dgrad
|
||||
// on problem sizes with stride = {1x1}
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
typename ThreadMap_
|
||||
>
|
||||
class Conv2dDgradFilterTileAccessIteratorAnalytic <
|
||||
Shape_,
|
||||
Element_,
|
||||
ThreadMap_,
|
||||
conv::StrideSupport::kUnity
|
||||
>{
|
||||
public:
|
||||
|
||||
//
|
||||
// Types
|
||||
//
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout = layout::TensorNHWC;
|
||||
using ThreadMap = ThreadMap_;
|
||||
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess>;
|
||||
using TensorRef = cutlass::TensorRef<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic;
|
||||
static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity;
|
||||
static int const kConvDim = 2;
|
||||
using ConvProblemSize = typename conv::Conv2dProblemSize;
|
||||
|
||||
static_assert(sizeof_bits<Element>::value >= 8,
|
||||
"DGRAD requires elements of size 8b or larger.");
|
||||
|
||||
//
|
||||
// Parameters structure
|
||||
//
|
||||
|
||||
using Params = Conv2dAnalyticParams<Layout>;
|
||||
|
||||
private:
|
||||
|
||||
Params const ¶ms_;
|
||||
|
||||
@@ -62,7 +62,23 @@ template <
|
||||
typename ThreadMap_,
|
||||
conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity
|
||||
>
|
||||
class Conv2dDgradFilterTileAccessIteratorOptimized {
|
||||
class Conv2dDgradFilterTileAccessIteratorOptimized;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Conv2dDgradFilterTileAccessIteratorOptimized unity strided dgrad is more performant for dgrad
|
||||
// on problem sizes with stride = {1x1}
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
typename ThreadMap_
|
||||
>
|
||||
class Conv2dDgradFilterTileAccessIteratorOptimized <
|
||||
Shape_,
|
||||
Element_,
|
||||
ThreadMap_,
|
||||
conv::StrideSupport::kUnity
|
||||
> {
|
||||
public:
|
||||
|
||||
//
|
||||
@@ -79,7 +95,7 @@ public:
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized;
|
||||
static StrideSupport const kStrideSupport = StrideSupport_;
|
||||
static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity;
|
||||
static int const kConvDim = 2;
|
||||
using ConvProblemSize = typename conv::Conv2dProblemSize;
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/array.h"
|
||||
#include "cutlass/coord.h"
|
||||
#include "cutlass/functional.h"
|
||||
#include "cutlass/predicate_vector.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/tensor_view.h"
|
||||
@@ -109,7 +110,7 @@ public:
|
||||
// Parameters structure
|
||||
//
|
||||
|
||||
using Params = Conv2dAnalyticParams<Layout>;
|
||||
using Params = Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams;
|
||||
|
||||
private:
|
||||
|
||||
@@ -122,36 +123,13 @@ private:
|
||||
int filter_k_;
|
||||
int filter_r_;
|
||||
int filter_s_;
|
||||
int start_r_;
|
||||
int start_s_;
|
||||
|
||||
int offset_n_[ThreadMap::Iterations::kStrided];
|
||||
int offset_w_[ThreadMap::Iterations::kStrided];
|
||||
int offset_h_[ThreadMap::Iterations::kStrided];
|
||||
|
||||
private:
|
||||
int offset_p_[ThreadMap::Iterations::kStrided];
|
||||
int offset_q_[ThreadMap::Iterations::kStrided];
|
||||
|
||||
/// Returns the coordinate in the output tensor Dy that is currently pointed to
|
||||
/// by the iterator but DOES NOT scale by the convolution stride. This is needed
|
||||
/// to compute predicates in the valid() method. The return value of the public at()
|
||||
/// method is correctly scaled.
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCoord unscaled_at_() const {
|
||||
int n = offset_n_[iteration_strided_];
|
||||
int h = offset_h_[iteration_strided_];
|
||||
int w = offset_w_[iteration_strided_];
|
||||
|
||||
int r = filter_r_;
|
||||
int s = filter_s_;
|
||||
|
||||
if (problem_size_.mode == Mode::kConvolution) {
|
||||
r = (problem_size_.R - 1 - r);
|
||||
s = (problem_size_.S - 1 - s);
|
||||
}
|
||||
|
||||
int p = (h + problem_size_.pad_h - r * problem_size_.dilation_h);
|
||||
int q = (w + problem_size_.pad_w - s * problem_size_.dilation_w);
|
||||
|
||||
return TensorCoord(n, p, q, filter_k_);
|
||||
}
|
||||
|
||||
public:
|
||||
|
||||
@@ -161,34 +139,68 @@ public:
|
||||
Conv2dProblemSize const &problem_size,
|
||||
Element const *ptr,
|
||||
int thread_idx,
|
||||
int start_r, int start_s,
|
||||
MatrixCoord const &threadblock_offset = MatrixCoord() // threadblock offset - units are whole CTA tiles
|
||||
):
|
||||
params_(params),
|
||||
problem_size_(problem_size),
|
||||
pointer_(reinterpret_cast<char const *>(ptr)),
|
||||
filter_k_(0),
|
||||
filter_r_(0),
|
||||
filter_s_(0) {
|
||||
filter_k_(0),
|
||||
filter_r_(start_r),
|
||||
filter_s_(start_s),
|
||||
start_r_(start_r),
|
||||
start_s_(start_s) {
|
||||
|
||||
layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx);
|
||||
|
||||
filter_k_ = threadblock_offset.column() + thread_coord.contiguous();
|
||||
|
||||
int filter_r = filter_r_;
|
||||
int filter_s = filter_s_;
|
||||
|
||||
if (problem_size_.mode == Mode::kConvolution) {
|
||||
filter_r = (problem_size_.R - 1 - filter_r);
|
||||
filter_s = (problem_size_.S - 1 - filter_s);
|
||||
}
|
||||
|
||||
// Starting h, w positions for filter position in gemm_k=0
|
||||
int start_h = std::abs((problem_size_.pad_h - filter_r) % problem_size_.stride_h);
|
||||
int start_w = std::abs((problem_size_.pad_w - filter_s) % problem_size_.stride_w);
|
||||
|
||||
|
||||
// Effective P and Q for filter position required for remapping NHW rows
|
||||
int P = (problem_size_.H - start_h + problem_size_.stride_h - 1) / problem_size_.stride_h;
|
||||
int Q = (problem_size_.W - start_w + problem_size_.stride_w - 1) / problem_size_.stride_w;
|
||||
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
int offset_nhw = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided;
|
||||
int offset_npq = (threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided) % params_.tiled_rows_per_filter;
|
||||
|
||||
offset_n_[s] = offset_nhw / (problem_size_.H * problem_size_.W);
|
||||
int residual = offset_nhw % (problem_size_.H * problem_size_.W);
|
||||
// (STEP 1) [reorder NHW rows to start with same filter positions]
|
||||
offset_n_[s] = offset_npq / (P * Q);
|
||||
int residual = offset_npq % (P * Q);
|
||||
|
||||
offset_h_[s] = residual / problem_size_.W;
|
||||
offset_w_[s] = residual % problem_size_.W;
|
||||
int p = (residual / Q);
|
||||
int q = (residual % Q);
|
||||
|
||||
int mapped_h = (start_h + p * problem_size_.stride_h);
|
||||
int mapped_w = (start_w + q * problem_size_.stride_w);
|
||||
|
||||
// Access (p, q) coordinates for Dy tensor and a filter position in gemm_k=0
|
||||
// note that (h + pad_h - filter_r) and (w + pad_w - filter_s) are divisible
|
||||
// by stride_h and stride_w
|
||||
offset_p_[s] = (mapped_h + problem_size_.pad_h - filter_r) / problem_size_.stride_h;
|
||||
offset_q_[s] = (mapped_w + problem_size_.pad_w - filter_s) / problem_size_.stride_w;
|
||||
}
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) {
|
||||
return Params(problem_size, layout);
|
||||
return Params(problem_size,
|
||||
layout,
|
||||
sizeof_bits<Element>::value,
|
||||
{Shape::kRow, Shape::kColumn});
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
@@ -206,18 +218,26 @@ public:
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
void advance() {
|
||||
// move to the next tile
|
||||
++filter_s_;
|
||||
|
||||
// Move filter_s by stride_w
|
||||
filter_s_ += problem_size_.stride_w;
|
||||
if (filter_s_ < problem_size_.S) {
|
||||
return;
|
||||
}
|
||||
filter_s_ = 0;
|
||||
++filter_r_;
|
||||
|
||||
// Restore filter_s
|
||||
filter_s_ = start_s_;
|
||||
|
||||
// Move filter_r by stride_h
|
||||
filter_r_ += problem_size_.stride_h;
|
||||
if (filter_r_ < problem_size_.R) {
|
||||
return;
|
||||
}
|
||||
filter_r_ = 0;
|
||||
|
||||
// Restore filter_r
|
||||
filter_r_ = start_r_;
|
||||
|
||||
// Move filter_k
|
||||
filter_k_ += Shape_::kColumn * problem_size_.split_k_slices;
|
||||
}
|
||||
|
||||
@@ -225,14 +245,20 @@ public:
|
||||
/// by the iterator.
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCoord at() const {
|
||||
int n = offset_n_[iteration_strided_];
|
||||
int p = offset_p_[iteration_strided_];
|
||||
int q = offset_q_[iteration_strided_];
|
||||
|
||||
int conv_sign = (problem_size_.mode == Mode::kConvolution ? 1 : -1);
|
||||
|
||||
TensorCoord coord = unscaled_at_();
|
||||
p += (conv_sign * (filter_r_ / problem_size_.stride_h));
|
||||
q += (conv_sign * (filter_s_ / problem_size_.stride_w));
|
||||
|
||||
return TensorCoord(
|
||||
coord.n(),
|
||||
coord.h() / problem_size_.stride_h,
|
||||
coord.w() / problem_size_.stride_w,
|
||||
coord.c());
|
||||
n,
|
||||
p,
|
||||
q,
|
||||
filter_k_);
|
||||
}
|
||||
|
||||
|
||||
@@ -240,11 +266,9 @@ public:
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool valid() const {
|
||||
|
||||
TensorCoord unscaled_coord = unscaled_at_();
|
||||
TensorCoord coord = at();
|
||||
|
||||
return
|
||||
!(unscaled_coord.h() % problem_size_.stride_h) && !(unscaled_coord.w() % problem_size_.stride_w) &&
|
||||
coord.n() < problem_size_.N &&
|
||||
coord.h() >= 0 && coord.h() < problem_size_.P &&
|
||||
coord.w() >= 0 && coord.w() < problem_size_.Q &&
|
||||
|
||||
@@ -32,6 +32,7 @@
|
||||
backward data gradient (Dgrad), and backward weight gradient (Wgrad).
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
@@ -62,11 +63,26 @@ template <
|
||||
typename ThreadMap_,
|
||||
conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity
|
||||
>
|
||||
class Conv2dDgradOutputGradientTileAccessIteratorOptimized {
|
||||
public:
|
||||
class Conv2dDgradOutputGradientTileAccessIteratorOptimized;
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
static_assert(StrideSupport_ == conv::StrideSupport::kUnity,
|
||||
"Only unit-stride dgrad is supported at this time.");
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Conv2dDgradOutputGradientTileAccessIteratorOptimized unity stride dgrad is optimized for dgrad
|
||||
// with problem stride = {1x1}
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
typename Shape_,
|
||||
typename Element_,
|
||||
typename ThreadMap_
|
||||
>
|
||||
class Conv2dDgradOutputGradientTileAccessIteratorOptimized <
|
||||
Shape_,
|
||||
Element_,
|
||||
ThreadMap_,
|
||||
conv::StrideSupport::kUnity
|
||||
> {
|
||||
public:
|
||||
|
||||
//
|
||||
// Types
|
||||
@@ -417,5 +433,3 @@ public:
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
|
||||
@@ -99,7 +99,7 @@ public:
|
||||
|
||||
private:
|
||||
|
||||
Conv2dFpropActivationIteratorOptimizedParams<Layout> const ¶ms_;
|
||||
Params const ¶ms_;
|
||||
Conv2dProblemSize const &problem_size_;
|
||||
LongIndex iteration_contiguous_;
|
||||
LongIndex iteration_strided_;
|
||||
@@ -118,7 +118,7 @@ public:
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv2dFpropActivationTileAccessIteratorOptimized(
|
||||
Conv2dFpropActivationIteratorOptimizedParams<Layout> const ¶ms,
|
||||
Params const ¶ms,
|
||||
Conv2dProblemSize const &problem_size,
|
||||
Element const *ptr,
|
||||
int thread_idx,
|
||||
|
||||
@@ -77,6 +77,38 @@ struct Conv2dAnalyticParams {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Parameters structure used for Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams
|
||||
struct Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams {
|
||||
|
||||
using Layout = layout::TensorNHWC;
|
||||
|
||||
Layout layout;
|
||||
int tiled_rows_per_filter;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams(
|
||||
Conv2dProblemSize const &problem_size,
|
||||
Layout const &layout, ///< layout object
|
||||
int element_size_bits, ///< size of each element in bits
|
||||
MatrixCoord threadblock_shape
|
||||
): layout(layout) {
|
||||
|
||||
int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, threadblock_shape.row());
|
||||
|
||||
tiled_rows_per_filter = tile_m_per_filter * threadblock_shape.row();
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if TRACE_CONV_PARAMS_INITIALIZERS_ENABLED
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
@@ -199,6 +231,32 @@ struct Conv2dFpropActivationIteratorOptimizedParams<layout::TensorNHWC> {
|
||||
// logical offset added to internal channel counter - units are elements, not bytes
|
||||
filter_c_delta = threadblock_shape.column() * problem_size.split_k_slices;
|
||||
}
|
||||
|
||||
#if 0
|
||||
/// Prints internal state.
|
||||
CUTLASS_HOST_DEVICE
|
||||
void print() {
|
||||
auto stride = layout.stride();
|
||||
printf(
|
||||
"Conv2dFpropActivationIteratorOptimizedParams:\n"
|
||||
" layout(w: %d, h: %d, n: %d)\n"
|
||||
" inc_next[%ld, %ld, %ld]\n"
|
||||
" filter_c_delta(%d) - PQ(%d)\n"
|
||||
" pq_divmod(divisor: %d, multiplier: %u, shift_right: %u)\n"
|
||||
" q_divmod(divisor: %d, multiplier: %u, shift_right: %u)\n",
|
||||
stride[0], stride[1], stride[2],
|
||||
inc_next[0], inc_next[1], inc_next[2],
|
||||
filter_c_delta,
|
||||
PQ,
|
||||
pq_divmod.divisor,
|
||||
pq_divmod.multiplier,
|
||||
pq_divmod.shift_right,
|
||||
q_divmod.divisor,
|
||||
q_divmod.multiplier,
|
||||
q_divmod.shift_right
|
||||
);
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
/// Parameters structure used for Conv2dFpropActivationTileIteratorOptimized
|
||||
@@ -324,6 +382,23 @@ struct Conv2dFpropFilterIteratorOptimizedParams<layout::TensorNHWC>
|
||||
|
||||
filter_c_delta = threadblock_shape.row() * problem_size.split_k_slices;
|
||||
}
|
||||
|
||||
#if 0
|
||||
/// Prints internal state.
|
||||
CUTLASS_HOST_DEVICE
|
||||
void print() {
|
||||
auto stride = layout.stride();
|
||||
printf(
|
||||
"Conv2dFpropFilterIteratorOptimizedParams:\n"
|
||||
" layout[%d, %d, %d]\n"
|
||||
" RS(%d), filter_c_delta(%d), inc_next(k: %ld, rs: %ld, c: %ld)\n",
|
||||
stride[0], stride[1], stride[2],
|
||||
RS,
|
||||
filter_c_delta,
|
||||
inc_next_k, inc_next_rs, inc_next_c
|
||||
);
|
||||
}
|
||||
#endif
|
||||
};
|
||||
|
||||
template<int Interleaved_>
|
||||
@@ -382,6 +457,9 @@ struct Conv2dFpropFilterIteratorOptimizedParams<layout::TensorCxRSKx<Interleaved
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Dgrad Optimized Dy params (layout::TensorNHWC)
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Parameters object for Conv2d DGRAD OutputGradient (dy) iterator
|
||||
struct Conv2dDgradOutputGradientIteratorOptimizedParams {
|
||||
|
||||
@@ -449,7 +527,9 @@ struct Conv2dDgradOutputGradientIteratorOptimizedParams {
|
||||
}
|
||||
};
|
||||
|
||||
/// Parameters object for Conv2d DGRAD Filter (w) iterator
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Dgrad Optimized w params (layout::TensorNHWC)
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
struct Conv2dDgradFilterIteratorOptimizedParams {
|
||||
|
||||
using Layout = layout::TensorNHWC;
|
||||
@@ -609,6 +689,25 @@ struct Conv2dWgradActivationIteratorOptimizedParams {
|
||||
}
|
||||
};
|
||||
|
||||
struct PredicatedScaleBiasVectorAccessIteratorParams {
|
||||
public:
|
||||
/// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIteratorParams() { }
|
||||
|
||||
// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIteratorParams(
|
||||
Conv2dProblemSize const &problem_size,
|
||||
layout::PitchLinear const &layout) {}
|
||||
|
||||
// Default ctor
|
||||
CUTLASS_HOST_DEVICE
|
||||
PredicatedScaleBiasVectorAccessIteratorParams(
|
||||
Conv2dProblemSize const &problem_size,
|
||||
layout::RowMajor const &layout) {}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
|
||||
@@ -166,6 +166,125 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// Strided Dgrad Tile Iterator
|
||||
template <typename TileAccessIterator_>
|
||||
class TileIteratorStridedDgrad {
|
||||
public:
|
||||
using TileAccessIterator = TileAccessIterator_;
|
||||
|
||||
using Shape = typename TileAccessIterator::Shape;
|
||||
using Element = typename TileAccessIterator::Element;
|
||||
using Layout = typename TileAccessIterator::Layout;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
using ThreadMap = typename TileAccessIterator::ThreadMap;
|
||||
using AccessType = typename TileAccessIterator::AccessType;
|
||||
using TensorRef = typename TileAccessIterator::TensorRef;
|
||||
using Index = typename TileAccessIterator::Index;
|
||||
using LongIndex = typename TileAccessIterator::LongIndex;
|
||||
static IteratorAlgorithm const kIteratorAlgorithm = TileAccessIterator::kIteratorAlgorithm;
|
||||
static StrideSupport const kStrideSupport = TileAccessIterator::kStrideSupport;
|
||||
using Params = typename TileAccessIterator::Params;
|
||||
static int const kConvDim = TileAccessIterator::kConvDim;
|
||||
using ConvProblemSize = typename TileAccessIterator::ConvProblemSize;
|
||||
|
||||
/// Fragment object to be loaded or stored
|
||||
using Fragment = cutlass::Array<
|
||||
Element,
|
||||
ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
|
||||
/// Internal state
|
||||
TileAccessIterator tile_access_iterator_;
|
||||
|
||||
public:
|
||||
|
||||
/// Constructor
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorStridedDgrad(
|
||||
Params const ¶ms,
|
||||
ConvProblemSize const &problem_size,
|
||||
Element const *ptr,
|
||||
int thread_idx,
|
||||
int start_r, int start_s,
|
||||
MatrixCoord const &threadblock_offset = MatrixCoord()
|
||||
):
|
||||
tile_access_iterator_(params, problem_size, ptr, thread_idx, start_r, start_s, threadblock_offset) { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static Params getParams(ConvProblemSize const &problem_size, Layout const &layout) {
|
||||
return TileAccessIterator::getParams(problem_size, layout);
|
||||
}
|
||||
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
tile_access_iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorStridedDgrad &operator++() {
|
||||
tile_access_iterator_.advance();
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
TileIteratorStridedDgrad operator++(int) {
|
||||
TileIteratorStridedDgrad self(*this);
|
||||
operator++();
|
||||
return self;
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
|
||||
|
||||
frag.clear();
|
||||
AccessType *frag_ptr = reinterpret_cast<AccessType *>(&frag);
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) {
|
||||
|
||||
cutlass::arch::global_load<
|
||||
AccessType,
|
||||
sizeof(AccessType)
|
||||
>(
|
||||
frag_ptr[c + s * ThreadMap::Iterations::kContiguous],
|
||||
tile_access_iterator_.get() + pointer_offset,
|
||||
tile_access_iterator_.valid()
|
||||
);
|
||||
|
||||
++tile_access_iterator_;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads a fragment from memory
|
||||
CUTLASS_DEVICE
|
||||
void load(Fragment &frag) {
|
||||
tile_access_iterator_.set_iteration_index(0);
|
||||
load_with_pointer_offset(frag, 0);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void advance() {
|
||||
tile_access_iterator_.advance();
|
||||
}
|
||||
|
||||
/// Determines whether the Implicit GEMM can execute the given problem.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static Status can_implement(ConvProblemSize const &problem_size) {
|
||||
|
||||
// dispatch to iterator implementation
|
||||
return TileAccessIterator::can_implement(problem_size);
|
||||
}
|
||||
};
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
|
||||
@@ -243,6 +243,7 @@ public:
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
@@ -250,5 +251,3 @@ public:
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
|
||||
@@ -196,7 +196,7 @@ private:
|
||||
CUTLASS_HOST_DEVICE
|
||||
TensorCoord at_(int offset_npq, int k) const {
|
||||
|
||||
// The subseqnet fast_divmod() operations are equivalent to the following logical computation:
|
||||
// The subsequent fast_divmod() operations are equivalent to the following logical computation:
|
||||
//
|
||||
//
|
||||
// int npq = offset_npq;
|
||||
|
||||
@@ -355,6 +355,145 @@ struct Conv3dDgradFilterIteratorOptimizedParams {
|
||||
}
|
||||
};
|
||||
|
||||
/// Parameters object for Conv3d WGRAD OutputGradient iterator
|
||||
struct Conv3dWgradOutputGradientIteratorOptimizedParams {
|
||||
|
||||
using Layout = layout::TensorNDHWC;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
Layout layout;
|
||||
|
||||
int NZPQ; // precomputd product of N*Z*P*Q for clearing predicates
|
||||
int ZPQ; // product of Z*P*Q
|
||||
unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ
|
||||
unsigned zpq_shr; // in device code.
|
||||
|
||||
int PQ; // product of P*Q
|
||||
unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ
|
||||
unsigned pq_shr; // in device code.
|
||||
|
||||
unsigned q_mul; // precomputed quantities for fast computation of div/% by Q
|
||||
unsigned q_shr; // in device code.
|
||||
|
||||
LongIndex offset_next_strided; // offset in units of bytes to next nzpq coordinate within tile
|
||||
LongIndex offset_next_contiguous; // offset in units of bytes to next k coordinate within tile
|
||||
LongIndex inc_next_nzpq; // offset in units of bytes to next nzpq position in subsequent tile
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv3dWgradOutputGradientIteratorOptimizedParams() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv3dWgradOutputGradientIteratorOptimizedParams(
|
||||
Conv3dProblemSize const &problem_size,
|
||||
Layout const &layout,
|
||||
int element_size_bits,
|
||||
MatrixCoord threadblock_shape,
|
||||
int thread_count,
|
||||
int access_size,
|
||||
layout::PitchLinearCoord threadmap_iterations,
|
||||
layout::PitchLinearCoord threadmap_delta
|
||||
): layout(layout) {
|
||||
|
||||
TRACE_CONV_INITIALIZERS("conv3d_wgrad", "output_gradient",
|
||||
element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta);
|
||||
|
||||
// Incremental offsets in unites of bytes (number of elements) * element_size_bits / 8
|
||||
offset_next_strided = (threadmap_delta.strided() * layout.stride()[0])
|
||||
* element_size_bits / 8;
|
||||
|
||||
offset_next_contiguous = (threadmap_delta.contiguous())
|
||||
* element_size_bits / 8;
|
||||
|
||||
inc_next_nzpq = (threadblock_shape.column() * problem_size.split_k_slices * layout.stride()[0])
|
||||
* element_size_bits / 8;
|
||||
|
||||
// Precompute several quantities for fast modulo arithmetic.
|
||||
NZPQ = problem_size.N * problem_size.Z * problem_size.P * problem_size.Q;
|
||||
ZPQ = problem_size.Z * problem_size.P * problem_size.Q;
|
||||
find_divisor(zpq_mul, zpq_shr, ZPQ);
|
||||
|
||||
PQ = problem_size.P * problem_size.Q;
|
||||
find_divisor(pq_mul, pq_shr, PQ);
|
||||
|
||||
find_divisor(q_mul, q_shr, problem_size.Q);
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
/// Parameters object for Conv3d WGRAD Activation Tile Access Iterator
|
||||
struct Conv3dWgradActivationIteratorOptimizedParams {
|
||||
|
||||
using Layout = layout::TensorNDHWC;
|
||||
|
||||
Layout layout;
|
||||
|
||||
int RSC; // product of R*S*C
|
||||
unsigned rsc_mul; // precomputed quantities for fast computation of div/% by RSC
|
||||
unsigned rsc_shr; // in device code.
|
||||
|
||||
int SC; // product of S*C
|
||||
unsigned sc_mul; // precomputed quantities for fast computation of div/% by SC
|
||||
unsigned sc_shr; // in device code.
|
||||
|
||||
unsigned c_mul; // precomputed quantities for fast computation of div/% by C
|
||||
unsigned c_shr; // in device code.
|
||||
|
||||
int ZPQ; // product of Z*P*Q
|
||||
unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ
|
||||
unsigned zpq_shr; // in device code.
|
||||
|
||||
int PQ; // product of P*Q
|
||||
unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ
|
||||
unsigned pq_shr; // in device code.
|
||||
|
||||
unsigned q_mul; // precomputed quantities for fast computation of div/% by Q
|
||||
unsigned q_shr; // in device code.
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv3dWgradActivationIteratorOptimizedParams() { }
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Conv3dWgradActivationIteratorOptimizedParams(
|
||||
Conv3dProblemSize const &problem_size,
|
||||
Layout const &layout,
|
||||
int element_size_bits,
|
||||
MatrixCoord threadblock_shape,
|
||||
int thread_count,
|
||||
int access_size,
|
||||
layout::PitchLinearCoord threadmap_iterations,
|
||||
layout::PitchLinearCoord threadmap_delta
|
||||
): layout(layout) {
|
||||
|
||||
TRACE_CONV_INITIALIZERS("conv3d_wgrad", "activation",
|
||||
element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta);
|
||||
|
||||
// Precompute several quantities for fast modulo arithmetic.
|
||||
RSC = problem_size.R * problem_size.S * problem_size.C;
|
||||
find_divisor(rsc_mul, rsc_shr, RSC);
|
||||
|
||||
SC = problem_size.S * problem_size.C;
|
||||
find_divisor(sc_mul, sc_shr, SC);
|
||||
|
||||
find_divisor(c_mul, c_shr, problem_size.C);
|
||||
|
||||
ZPQ = problem_size.Z * problem_size.P * problem_size.Q;
|
||||
find_divisor(zpq_mul, zpq_shr, ZPQ);
|
||||
|
||||
PQ = problem_size.P * problem_size.Q;
|
||||
find_divisor(pq_mul, pq_shr, PQ);
|
||||
|
||||
find_divisor(q_mul, q_shr, problem_size.Q);
|
||||
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace conv
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -45,6 +45,7 @@
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/conv/convolution.h"
|
||||
#include "cutlass/conv/conv3d_problem_size.h"
|
||||
#include "cutlass/conv/threadblock/conv3d_params.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -86,62 +87,28 @@ public:
|
||||
// Parameters structure
|
||||
//
|
||||
|
||||
struct Params {
|
||||
|
||||
Layout layout;
|
||||
|
||||
int RSC; // product of R*S*C
|
||||
unsigned rsc_mul; // precomputed quantities for fast computation of div/% by RSC
|
||||
unsigned rsc_shr; // in device code.
|
||||
|
||||
int SC; // product of S*C
|
||||
unsigned sc_mul; // precomputed quantities for fast computation of div/% by SC
|
||||
unsigned sc_shr; // in device code.
|
||||
|
||||
unsigned c_mul; // precomputed quantities for fast computation of div/% by C
|
||||
unsigned c_shr; // in device code.
|
||||
|
||||
int ZPQ; // product of Z*P*Q
|
||||
unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ
|
||||
unsigned zpq_shr; // in device code.
|
||||
|
||||
int PQ; // product of P*Q
|
||||
unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ
|
||||
unsigned pq_shr; // in device code.
|
||||
|
||||
unsigned q_mul; // precomputed quantities for fast computation of div/% by Q
|
||||
unsigned q_shr; // in device code.
|
||||
|
||||
struct Params : Conv3dWgradActivationIteratorOptimizedParams {
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
Params() {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
Conv3dProblemSize const &problem_size,
|
||||
Layout const &layout
|
||||
): layout(layout) {
|
||||
Params(Conv3dWgradActivationIteratorOptimizedParams const &base)
|
||||
: Conv3dWgradActivationIteratorOptimizedParams(base) {}
|
||||
|
||||
// Precompute several quantities for fast modulo arithmetic.
|
||||
RSC = problem_size.R * problem_size.S * problem_size.C;
|
||||
find_divisor(rsc_mul, rsc_shr, RSC);
|
||||
|
||||
SC = problem_size.S * problem_size.C;
|
||||
find_divisor(sc_mul, sc_shr, SC);
|
||||
|
||||
find_divisor(c_mul, c_shr, problem_size.C);
|
||||
|
||||
ZPQ = problem_size.Z * problem_size.P * problem_size.Q;
|
||||
find_divisor(zpq_mul, zpq_shr, ZPQ);
|
||||
|
||||
PQ = problem_size.P * problem_size.Q;
|
||||
find_divisor(pq_mul, pq_shr, PQ);
|
||||
|
||||
find_divisor(q_mul, q_shr, problem_size.Q);
|
||||
|
||||
}
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(Conv3dProblemSize const &problem_size, Layout const &layout)
|
||||
: Conv3dWgradActivationIteratorOptimizedParams(
|
||||
problem_size,
|
||||
layout,
|
||||
sizeof_bits<Element>::value,
|
||||
{Shape::kRow, Shape::kColumn},
|
||||
ThreadMap::kThreads,
|
||||
ThreadMap::kElementsPerAccess,
|
||||
{ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided},
|
||||
{ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}) {}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
@@ -45,6 +45,7 @@
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/conv/convolution.h"
|
||||
#include "cutlass/conv/conv3d_problem_size.h"
|
||||
#include "cutlass/conv/threadblock/conv3d_params.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -86,61 +87,29 @@ public:
|
||||
// Parameters structure
|
||||
//
|
||||
|
||||
struct Params {
|
||||
|
||||
Layout layout;
|
||||
|
||||
int NZPQ; // precomputd product of N*Z*P*Q for clearing predicates
|
||||
int ZPQ; // product of Z*P*Q
|
||||
unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ
|
||||
unsigned zpq_shr; // in device code.
|
||||
|
||||
int PQ; // product of P*Q
|
||||
unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ
|
||||
unsigned pq_shr; // in device code.
|
||||
|
||||
unsigned q_mul; // precomputed quantities for fast computation of div/% by Q
|
||||
unsigned q_shr; // in device code.
|
||||
|
||||
LongIndex offset_next_strided; // offset in units of bytes to next nzpq coordinate within tile
|
||||
LongIndex offset_next_contiguous; // offset in units of bytes to next k coordinate within tile
|
||||
LongIndex inc_next_nzpq; // offset in units of bytes to next nzpq position in subsequent tile
|
||||
|
||||
struct Params : Conv3dWgradOutputGradientIteratorOptimizedParams {
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params() { }
|
||||
Params(Conv3dWgradOutputGradientIteratorOptimizedParams const &base)
|
||||
: Conv3dWgradOutputGradientIteratorOptimizedParams(base) {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Params(
|
||||
Conv3dProblemSize const &problem_size,
|
||||
Layout const &layout
|
||||
): layout(layout) {
|
||||
|
||||
// Incremental offsets in unites of bytes (number of elements) * sizeof_bits<Element>::value / 8
|
||||
offset_next_strided = (ThreadMap::Delta::kStrided * layout.stride()[0])
|
||||
* sizeof_bits<Element>::value / 8;
|
||||
|
||||
offset_next_contiguous = (ThreadMap::Delta::kContiguous)
|
||||
* sizeof_bits<Element>::value / 8;
|
||||
|
||||
inc_next_nzpq = (Shape::kColumn * problem_size.split_k_slices * layout.stride()[0])
|
||||
* sizeof_bits<Element>::value / 8;
|
||||
|
||||
// Precompute several quantities for fast modulo arithmetic.
|
||||
NZPQ = problem_size.N * problem_size.Z * problem_size.P * problem_size.Q;
|
||||
ZPQ = problem_size.Z * problem_size.P * problem_size.Q;
|
||||
find_divisor(zpq_mul, zpq_shr, ZPQ);
|
||||
|
||||
PQ = problem_size.P * problem_size.Q;
|
||||
find_divisor(pq_mul, pq_shr, PQ);
|
||||
|
||||
find_divisor(q_mul, q_shr, problem_size.Q);
|
||||
|
||||
}
|
||||
};
|
||||
Params(Conv3dProblemSize const &problem_size, Layout const &layout)
|
||||
: Conv3dWgradOutputGradientIteratorOptimizedParams(
|
||||
problem_size,
|
||||
layout,
|
||||
sizeof_bits<Element>::value,
|
||||
{Shape::kRow, Shape::kColumn},
|
||||
ThreadMap::kThreads,
|
||||
ThreadMap::kElementsPerAccess,
|
||||
{ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided},
|
||||
{ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}) {}
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
|
||||
@@ -377,7 +377,7 @@ public:
|
||||
|
||||
this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
|
||||
this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations);
|
||||
|
||||
|
||||
this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]);
|
||||
this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]);
|
||||
|
||||
|
||||
166
include/cutlass/conv/threadblock/threadblock_swizzle.h
Normal file
166
include/cutlass/conv/threadblock/threadblock_swizzle.h
Normal file
@@ -0,0 +1,166 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without modification, are permitted
|
||||
* provided that the following conditions are met:
|
||||
* * Redistributions of source code must retain the above copyright notice, this list of
|
||||
* conditions and the following disclaimer.
|
||||
* * 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.
|
||||
* * Neither the name of the NVIDIA CORPORATION 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 NVIDIA CORPORATION 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 TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
/*! \file
|
||||
\brief Implements several possible threadblock-swizzling functions mapping blockIdx to
|
||||
Convolution problems.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/matrix.h"
|
||||
#include "cutlass/platform/platform.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/conv/convolution.h"
|
||||
#include "cutlass/conv/conv2d_problem_size.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass {
|
||||
namespace conv {
|
||||
namespace threadblock {
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
CUTLASS_HOST_DEVICE
|
||||
static int get_strided_dgrad_tile_m(
|
||||
cutlass::conv::Conv2dProblemSize const &problem_size,
|
||||
int tile_size_m) {
|
||||
|
||||
// CTAs in M dimension per starting filter position
|
||||
int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, tile_size_m);
|
||||
|
||||
// Inflate number of CTAs in M dimension to cover every strating filter position even those that
|
||||
// may fall out of valid MMA (Dy * w) but are needed to apply epilogue (beta * Dx_source)
|
||||
// and point-wise fusion
|
||||
int tile_m = tile_m_per_filter * int(problem_size.stride().product());
|
||||
|
||||
// There is a possible performance optimization here that leads up to 2x speeds than the current
|
||||
// CUTLASS strided dgrad performance for stride > filter, i.e., stride={2x2} and filter={1x1})
|
||||
//
|
||||
// * Optimization *
|
||||
// Only launch CTAs in M dimenstion which contribute to a row in Dx output
|
||||
//
|
||||
//
|
||||
// * Constraints *
|
||||
// (A) stride <= filter, for example, stride={2x2} and filter={3x3}:
|
||||
// - (A.1): There are no constraints for this case and the optimization does
|
||||
// affect this case functionality or performance.
|
||||
// (B) stride > filter, for example, stride={2x2} and filter={1x1}:
|
||||
// - (B.1): Dx output tensor should be zero initialized
|
||||
// - (B.2): The kernel epilogue cannot apply beta. Thus, beta should be zero
|
||||
|
||||
return tile_m;
|
||||
}
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Threadblock swizzling function for strided dgrad convolution
|
||||
struct StridedDgradHorizontalThreadblockSwizzle :
|
||||
public gemm::threadblock::GemmHorizontalThreadblockSwizzle {
|
||||
|
||||
using Base = gemm::threadblock::GemmHorizontalThreadblockSwizzle;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
StridedDgradHorizontalThreadblockSwizzle() { }
|
||||
|
||||
/// Returns the shape of the problem in units of logical tiles
|
||||
/// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC)
|
||||
CUTLASS_HOST_DEVICE
|
||||
gemm::GemmCoord get_tiled_shape(
|
||||
cutlass::conv::Operator conv_operator,
|
||||
cutlass::conv::Conv2dProblemSize const &problem_size,
|
||||
gemm::GemmCoord tile_size,
|
||||
int split_k_slices) const {
|
||||
|
||||
gemm::GemmCoord implicit_gemm_problem_size =
|
||||
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
|
||||
|
||||
// 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 dimenstion
|
||||
int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
|
||||
|
||||
return gemm::GemmCoord(
|
||||
tile_m,
|
||||
tile_n,
|
||||
split_k_slices);
|
||||
}
|
||||
|
||||
/// Returns the shape of the problem in units of logical tiles
|
||||
/// For GEMM problem size (MxNxK) (Do not use base class get_tiled_shape())
|
||||
private:
|
||||
using Base::get_tiled_shape;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/// Threadblock swizzling function for strided dgrad convolution
|
||||
template <int N = 1>
|
||||
struct StridedDgradIdentityThreadblockSwizzle :
|
||||
public gemm::threadblock::GemmIdentityThreadblockSwizzle<N> {
|
||||
|
||||
using Base = gemm::threadblock::GemmIdentityThreadblockSwizzle<N>;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
StridedDgradIdentityThreadblockSwizzle() { }
|
||||
|
||||
/// Returns the shape of the problem in units of logical tiles
|
||||
/// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC)
|
||||
CUTLASS_HOST_DEVICE
|
||||
gemm::GemmCoord get_tiled_shape(
|
||||
cutlass::conv::Operator conv_operator,
|
||||
cutlass::conv::Conv2dProblemSize const &problem_size,
|
||||
gemm::GemmCoord tile_size,
|
||||
int split_k_slices) const {
|
||||
|
||||
gemm::GemmCoord implicit_gemm_problem_size =
|
||||
cutlass::conv::implicit_gemm_problem_size(conv_operator, problem_size);
|
||||
|
||||
// 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 dimenstion
|
||||
int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n();
|
||||
|
||||
return gemm::GemmCoord(
|
||||
tile_m,
|
||||
tile_n,
|
||||
split_k_slices);
|
||||
}
|
||||
|
||||
|
||||
/// Returns the shape of the problem in units of logical tiles
|
||||
/// For GEMM problem size (MxNxK) (Do not use base class get_tiled_shape())
|
||||
private:
|
||||
using Base::get_tiled_shape;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
Reference in New Issue
Block a user