CUTLASS 2.4 (Implicit GEMM convolution) (#147)
CUTLASS 2.4 (Implicit GEMM Convolution) Co-authored-by: Manish Gupta <manigupta@nvidia.com>, Haicheng Wu <haichengw@nvidia.com>, Dustyn Blasig <dblasig@nvidia.com>, Andrew Kerr <akerr@nvidia.com>
This commit is contained in:
Manish Gupta
co-authored by
Manish Gupta <manigupta@nvidia.com>, Haicheng Wu <haichengw@nvidia.com>, Dustyn Blasig <dblasig@nvidia.com>, Andrew Kerr <akerr@nvidia.com>
Manish Gupta <manigupta@nvidia.com>, Haicheng Wu <haichengw@nvidia.com>, Dustyn Blasig <dblasig@nvidia.com>, Andrew Kerr <akerr@nvidia.com>
parent
c2b80ad4e4
commit
6615010cd0
@@ -340,6 +340,134 @@ struct PitchLinearWarpRakedThreadMap {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Policy defining a warp-raked arrangement in which a shape is partitioned into contiguous
|
||||
/// elements. Warps are arranged based on a stride.
|
||||
///
|
||||
/// This ThreadMap is used by tensor core kernels for NCxHWx layout.
|
||||
template <
|
||||
typename Shape_,
|
||||
int Threads,
|
||||
typename WarpThreadArrangement_,
|
||||
int ElementsPerAccess = 1
|
||||
>
|
||||
struct PitchLinearStridedWarpRakedThreadMap {
|
||||
|
||||
/// Tensor coordinate
|
||||
using TensorCoord = layout::PitchLinearCoord;
|
||||
|
||||
/// Tile shape
|
||||
using Shape = Shape_;
|
||||
|
||||
/// Number of threads total
|
||||
static int const kThreads = Threads;
|
||||
|
||||
using WarpThreadArrangement = WarpThreadArrangement_;
|
||||
|
||||
/// Extract vector length from Layout
|
||||
static int const kElementsPerAccess = ElementsPerAccess;
|
||||
|
||||
/// Base ThreadMap
|
||||
using BaseThreadMap = PitchLinearWarpRakedThreadMap<
|
||||
Shape,
|
||||
kThreads,
|
||||
WarpThreadArrangement,
|
||||
kElementsPerAccess
|
||||
>;
|
||||
|
||||
/// Shape of access by each thread
|
||||
using ThreadAccessShape = typename BaseThreadMap::ThreadAccessShape;
|
||||
|
||||
|
||||
struct Detail {
|
||||
|
||||
using WarpThreadArrangement = WarpThreadArrangement_;
|
||||
|
||||
using WarpAccessIterations = typename BaseThreadMap::Detail::WarpAccessIterations;
|
||||
|
||||
static int const kWarpSize = BaseThreadMap::Detail::kWarpSize;
|
||||
|
||||
static int const kWarpCount = BaseThreadMap::Detail::kWarpCount;
|
||||
|
||||
using ShapeInAccesses = typename BaseThreadMap::Detail::ShapeInAccesses;
|
||||
|
||||
// Divide it into the number of warps, first partitioning the contiguous dimension then the
|
||||
// stride.
|
||||
static int const kWarpsContiguous =
|
||||
(WarpAccessIterations::kContiguous >= kWarpCount
|
||||
? kWarpCount
|
||||
: WarpAccessIterations::kContiguous);
|
||||
|
||||
static int const kWarpsStrided =
|
||||
(kWarpCount > WarpAccessIterations::kContiguous
|
||||
? kWarpCount / kWarpsContiguous
|
||||
: 1);
|
||||
|
||||
/// Arrangement of warps within a threadblock-scoped tile
|
||||
using WarpArrangement = layout::PitchLinearShape<
|
||||
kWarpsContiguous, kWarpsStrided
|
||||
>;
|
||||
|
||||
};
|
||||
|
||||
///< Iterations along each dimension (concept: PitchLinearShape)
|
||||
using Iterations = layout::PitchLinearShape<
|
||||
Detail::WarpAccessIterations::kContiguous / Detail::kWarpsContiguous,
|
||||
Detail::WarpAccessIterations::kStrided / Detail::kWarpsStrided
|
||||
>;
|
||||
|
||||
static_assert(Iterations::kCount,
|
||||
"Number of iterations must be non-zero");
|
||||
|
||||
///< Delta betweeen accesses (units of elements, concept: PitchLinearShape)
|
||||
using Delta = typename BaseThreadMap::Delta;
|
||||
|
||||
/// Maps thread ID to a coordinate offset within the tensor's logical coordinate space
|
||||
CUTLASS_HOST_DEVICE
|
||||
static TensorCoord initial_offset(int thread_id) {
|
||||
|
||||
int warp_id = (thread_id / Detail::kWarpSize);
|
||||
int lane_id = (thread_id % Detail::kWarpSize);
|
||||
|
||||
//
|
||||
// compute warp-level offset
|
||||
//
|
||||
|
||||
// This is the shape of the entire area covered by a warp's memory access (in units of vectors)
|
||||
layout::PitchLinearCoord warp_footprint{
|
||||
Detail::WarpThreadArrangement::kContiguous * Iterations::kContiguous,
|
||||
Detail::WarpThreadArrangement::kStrided * Iterations::kStrided
|
||||
};
|
||||
|
||||
// This is the offset of a specific warp (in units of vectors)
|
||||
layout::PitchLinearCoord warp_offset{
|
||||
(warp_id % Detail::kWarpsContiguous),
|
||||
(warp_id / Detail::kWarpsContiguous)
|
||||
};
|
||||
|
||||
// This is the offset of a specific thread within a warp (units of vectors)
|
||||
layout::PitchLinearCoord thread_offset_in_warp{
|
||||
lane_id % Detail::WarpThreadArrangement::kContiguous,
|
||||
lane_id / Detail::WarpThreadArrangement::kContiguous
|
||||
};
|
||||
|
||||
// This is the offset of a thread within a threadblock tile (units of vectors)
|
||||
layout::PitchLinearCoord thread_offset_in_threadblock_tile_vec =
|
||||
warp_footprint * warp_offset + thread_offset_in_warp;
|
||||
|
||||
// This is the offset of a thread within a threadblock tile (units of elements)
|
||||
layout::PitchLinearCoord thread_offset_in_threadblock_tile_base{
|
||||
thread_offset_in_threadblock_tile_vec.contiguous() * kElementsPerAccess,
|
||||
thread_offset_in_threadblock_tile_vec.strided()
|
||||
};
|
||||
|
||||
return thread_offset_in_threadblock_tile_base;
|
||||
}
|
||||
|
||||
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Transpose the existing ThreadMap. For example, interleaved layout is like
|
||||
/// congruous in the global memory and crosswise in the shared memory. We need
|
||||
/// to transpose the coordinates between two.
|
||||
|
||||
@@ -500,7 +500,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::PitchLinear,
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for pitch-linear data.
|
||||
/// Specialization of PredicatedTileAccessIterator for column-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
@@ -676,7 +676,7 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::ColumnMajor,
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for pitch-linear data.
|
||||
/// Specialization of PredicatedTileAccessIterator for row-major data.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
@@ -852,8 +852,8 @@ class PredicatedTileAccessIterator<Shape_, Element_, layout::RowMajor,
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for interleaved data. It
|
||||
/// is mapped to the congruous layout.
|
||||
/// Specialization of PredicatedTileAccessIterator for column-major interleaved data.
|
||||
/// It is mapped to the congruous layout.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
@@ -1032,8 +1032,8 @@ class PredicatedTileAccessIterator<Shape_, Element_,
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Specialization of PredicatedTileAccessIterator for interleaved data. It
|
||||
/// is mapped to the congruous layout.
|
||||
/// Specialization of PredicatedTileAccessIterator for row-major interleaved data.
|
||||
// It is mapped to the congruous layout.
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
|
||||
@@ -805,268 +805,6 @@ class RegularTileAccessIterator<Shape_, Element_,
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator specialized for k interleaved arrangements for TensorOps
|
||||
///
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept
|
||||
///
|
||||
template <typename Shape_, typename Element_, int AdvanceRank, typename ThreadMap_, int InterleavedK, int Alignment>
|
||||
class RegularTileAccessIterator<
|
||||
Shape_, Element_,
|
||||
layout::TensorOpMultiplicandRowMajorInterleaved<sizeof_bits<Element_>::value,
|
||||
InterleavedK>,
|
||||
AdvanceRank, ThreadMap_, Alignment> {
|
||||
public:
|
||||
static_assert(
|
||||
AdvanceRank == 0 || AdvanceRank == 1,
|
||||
"Specialization for pitch-linear iterator may along advance along the "
|
||||
"contiguous(rank=0) or strided(rank=1) dimension.");
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout =
|
||||
layout::TensorOpMultiplicandRowMajorInterleaved<sizeof_bits<Element_>::value,
|
||||
InterleavedK>;
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
static int const kAlignment = Alignment;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
/// Internal details made public to facilitate introspection
|
||||
struct Detail {
|
||||
/// This iterator is specialized for an access size that is 128 bits in
|
||||
/// length.
|
||||
static int const kAccessSizeInBits = 128;
|
||||
|
||||
static_assert(sizeof_bits<Element_>::value * ThreadMap::kElementsPerAccess ==
|
||||
kAccessSizeInBits,
|
||||
"This iterator requires a policy whose access size is 128bs");
|
||||
};
|
||||
|
||||
private:
|
||||
|
||||
/// Element type per access
|
||||
using AccessType = Array<Element, Layout::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
//
|
||||
// Data members
|
||||
//
|
||||
|
||||
/// Internal pointer to first access of tile
|
||||
AccessType *pointer_;
|
||||
|
||||
/// Internal byte offset
|
||||
Index byte_offset_;
|
||||
|
||||
/// Iteration in the contiguous dimension
|
||||
int iteration_contiguous_;
|
||||
|
||||
/// Iteration in the strided dimension
|
||||
int iteration_strided_;
|
||||
|
||||
public:
|
||||
/// Construct a TileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: byte_offset_(0) {
|
||||
layout::PitchLinearCoord thread_offset_base =
|
||||
ThreadMap::initial_offset(thread_id);
|
||||
|
||||
// initialize pointer
|
||||
pointer_ = reinterpret_cast<AccessType *>(
|
||||
ref.data() + ref.offset(thread_offset_base));
|
||||
|
||||
set_iteration_index(0);
|
||||
}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int 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) {
|
||||
byte_offset_ += pointer_offset * sizeof(Element);
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
AccessType *access_ptr = pointer_;
|
||||
|
||||
int access_offset =
|
||||
(iteration_strided_ * ThreadMap::Delta::kStrided * Layout::kInterleavedK +
|
||||
iteration_contiguous_ * ThreadMap::Delta::kContiguous) / ThreadMap::kElementsPerAccess;
|
||||
|
||||
char *access_byte_ptr =
|
||||
reinterpret_cast<char *>(access_ptr + access_offset);
|
||||
|
||||
return reinterpret_cast<AccessType *>(access_byte_ptr + byte_offset_);
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator &operator++() {
|
||||
++iteration_contiguous_;
|
||||
|
||||
if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous)
|
||||
return *this;
|
||||
|
||||
// Enter here only if (iteration_contiguous_ ==
|
||||
// ThreadMap::Iteration::kContiguous)
|
||||
iteration_contiguous_ = 0;
|
||||
++iteration_strided_;
|
||||
|
||||
if (iteration_strided_ < ThreadMap::Iterations::kStrided) {
|
||||
return *this;
|
||||
}
|
||||
|
||||
// Enter here only if (iteration_strided_ == ThreadMap::Iteration::kStrided)
|
||||
// which means we enter the next tile.
|
||||
iteration_strided_ = 0;
|
||||
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator operator++(int) {
|
||||
RegularTileAccessIterator prev(*this);
|
||||
this->operator++();
|
||||
|
||||
return prev;
|
||||
}
|
||||
|
||||
/// Adds a tile offset
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &coord) {
|
||||
add_pointer_offset(coord.contiguous() * Shape::kCount);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
/// Tile iterator specialized for k interleaved arrangements for TensorOps
|
||||
///
|
||||
///
|
||||
/// Satisfies: ForwardTileIteratorConcept |
|
||||
/// ReadableContiguousTileIteratorConcept |
|
||||
/// WriteableContiguousTileIteratorConcept
|
||||
///
|
||||
|
||||
template <typename Shape_, typename Element_, int AdvanceRank, typename ThreadMap_, int InterleavedK, int Alignment>
|
||||
class RegularTileAccessIterator<
|
||||
Shape_, Element_,
|
||||
layout::TensorOpMultiplicandColumnMajorInterleaved<sizeof_bits<Element_>::value,
|
||||
InterleavedK>,
|
||||
AdvanceRank, ThreadMap_, Alignment> {
|
||||
|
||||
public:
|
||||
static_assert(
|
||||
AdvanceRank == 0 || AdvanceRank == 1,
|
||||
"Specialization for pitch-linear iterator may along advance along the "
|
||||
"contiguous(rank=0) or strided(rank=1) dimension.");
|
||||
|
||||
using Shape = Shape_;
|
||||
using Element = Element_;
|
||||
using Layout =
|
||||
layout::TensorOpMultiplicandColumnMajorInterleaved<sizeof_bits<Element_>::value,
|
||||
InterleavedK>;
|
||||
static int const kAdvanceRank = AdvanceRank;
|
||||
static int const kAlignment = Alignment;
|
||||
|
||||
using Index = typename Layout::Index;
|
||||
using LongIndex = typename Layout::LongIndex;
|
||||
|
||||
using TensorRef = TensorRef<Element, Layout>;
|
||||
using TensorCoord = typename Layout::TensorCoord;
|
||||
|
||||
using ThreadMap = ThreadMap_;
|
||||
|
||||
/// Underlying iterator type
|
||||
using UnderlyingIterator = RegularTileAccessIterator<
|
||||
cutlass::MatrixShape<Shape::kColumn, Shape::kRow>,
|
||||
Element,
|
||||
layout::TensorOpMultiplicandRowMajorInterleaved<sizeof_bits<Element_>::value, InterleavedK>,
|
||||
(kAdvanceRank == 1 ? 0 : 1),
|
||||
ThreadMap
|
||||
>;
|
||||
|
||||
private:
|
||||
|
||||
/// Element type per access
|
||||
using AccessType = Array<Element, Layout::kElementsPerAccess>;
|
||||
|
||||
private:
|
||||
|
||||
/// Underlying iterator
|
||||
UnderlyingIterator iterator_;
|
||||
|
||||
public:
|
||||
/// Construct a TileIterator with zero threadblock offset
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor
|
||||
int thread_id ///< ID of each participating thread
|
||||
)
|
||||
: iterator_({ref.data(), ref.stride()}, thread_id) {}
|
||||
|
||||
/// Overrides the internal iteration index
|
||||
CUTLASS_HOST_DEVICE
|
||||
void set_iteration_index(int index) {
|
||||
iterator_.set_iteration_index(index);
|
||||
}
|
||||
|
||||
/// Adds a pointer offset in units of Element
|
||||
CUTLASS_HOST_DEVICE
|
||||
void add_pointer_offset(LongIndex pointer_offset) {
|
||||
iterator_.add_pointer_offset(pointer_offset);
|
||||
}
|
||||
|
||||
/// Returns a pointer
|
||||
CUTLASS_HOST_DEVICE
|
||||
AccessType *get() const {
|
||||
return iterator_.get();
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator &operator++() {
|
||||
++iterator_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
/// Advances to the next tile in memory.
|
||||
CUTLASS_HOST_DEVICE
|
||||
RegularTileAccessIterator operator++(int) {
|
||||
RegularTileAccessIterator prev(*this);
|
||||
++iterator_;
|
||||
|
||||
return prev;
|
||||
}
|
||||
|
||||
/// Adds a tile offset
|
||||
CUTLASS_DEVICE
|
||||
void add_tile_offset(TensorCoord const &coord) {
|
||||
iterator_.add_tile_offset({coord.strided(), coord.contiguous()});
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace threadblock
|
||||
} // namespace transform
|
||||
} // namespace cutlass
|
||||
|
||||
Reference in New Issue
Block a user