CUTLASS 2.6 (#298)

CUTLASS 2.6
This commit is contained in:
Manish Gupta
2021-07-22 21:40:53 -07:00
committed by GitHub
parent 6c29fe20ba
commit e5d51840e8
308 changed files with 32408 additions and 4722 deletions

View File

@@ -517,6 +517,9 @@ struct TransposePitchLinearThreadMap {
layout::PitchLinearShape<ThreadMap::Iterations::kStrided,
ThreadMap::Iterations::kContiguous>;
static_assert(Iterations::kContiguous == 1,
"Contiguous iteration has to be one to reuse the same shared store function with those that don't need transpose");
static_assert(Iterations::kCount, "Number of iterations must be non-zero");
///< Delta betweeen accesses (units of elements, concept: PitchLinearShape)
@@ -595,6 +598,9 @@ struct TransposePitchLinearThreadMapSimt {
static_assert(Iterations::kCount, "Number of iterations must be non-zero");
static_assert(Iterations::kStrided == 1,
"Strided iteration has to be one to reuse the same shared store function with those that don't need transpose");
/// Shape of access by each thread
using ThreadAccessShape = typename ThreadMap::ThreadAccessShape;

View File

@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
* 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:
@@ -26,6 +26,7 @@
/*! \file
\brief Basic copy routines for tensor views
*/
#pragma once
namespace cutlass {

View File

@@ -46,6 +46,7 @@
#include "cutlass/predicate_vector.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/tensor_view.h"
#include "cutlass/transform/threadblock/predicated_tile_access_iterator_params.h"
////////////////////////////////////////////////////////////////////////////////
@@ -86,6 +87,7 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::PitchLi
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorView = TensorView<Element, Layout>;
@@ -108,51 +110,31 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::PitchLi
/// Predicate vector stores mask to guard accesses
using Mask = Array<uint32_t, kPredicateWordCount>;
/// Parameters object is precomputed state and is host-constructible
class Params {
/// Uses a non-template class
struct Params : PredicatedTileAccessIteratorParams {
public:
friend PredicatedTileAccessIterator2dThreadTile;
private:
/// stride of pitch-linear layout (units of Element)
int stride_;
/// amount (in byte) to increment pointer to move to next access along
/// strided dimension
int inc_strided_;
/// amount (in byte) to increment pointer from last access to first access
/// of next tile
int inc_next_;
/// amount (in byte) to increment pointer from first access of current tile
/// to first access of next tile
int inc_advance_;
public:
using Base = PredicatedTileAccessIteratorParams;
// Default ctor
CUTLASS_HOST_DEVICE
Params(): stride_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) { }
Params() { }
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout) : stride_(layout.stride(0)) {
Params(Layout const &layout) :
Base(layout.stride(0),
MakePredicatedTileAccessIteratorDesc<Shape, Element, Layout, kAdvanceRank, ThreadMap>()()
) { }
inc_strided_ =
(stride_ * ThreadMap::Delta::kStrided) * int(sizeof(Element));
if (kAdvanceRank) {
// advance along strided dimension
inc_advance_ = Shape::kStrided * stride_ * int(sizeof(Element));
} else {
// advance along contiguous dimension
inc_advance_ = Shape::kContiguous * int(sizeof(Element));
}
inc_next_ = inc_advance_ - (ThreadMap::Iterations::kStrided - 1) *
ThreadMap::Delta::kStrided * stride_ *
int(sizeof(Element));
};
CUTLASS_HOST_DEVICE
Params(Base const &base) :
Base(base) { }
};
private:
/// Internal pointer type permits fast address arithmetic
using BytePointer = char *;
@@ -537,7 +519,12 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::ColumnM
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout)
: params_(layout::PitchLinear(layout.stride(0))){};
: params_(layout::PitchLinear(layout.stride(0))){}
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
private:
@@ -711,7 +698,12 @@ class PredicatedTileAccessIterator2dThreadTile<Shape_, Element_, layout::RowMajo
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout)
: params_(layout::PitchLinear(layout.stride(0))){};
: params_(layout::PitchLinear(layout.stride(0))){}
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
private:

View File

@@ -0,0 +1,276 @@
/***************************************************************************************************
* 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
*/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/layout/pitch_linear.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace transform {
namespace threadblock {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Predicated tile access iterator descriptor object containing template dependent state
struct PredicatedTileAccessIteratorDesc {
int element_size_bits;
int advance_rank;
layout::PitchLinearCoord threadblock_shape;
layout::PitchLinearCoord threadmap_iterations;
layout::PitchLinearCoord threadmap_delta;
//
// Methods
//
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc() { }
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc(
int element_size_bits_,
int advance_rank_,
layout::PitchLinearCoord threadblock_shape_,
layout::PitchLinearCoord threadmap_iterations_,
layout::PitchLinearCoord threadmap_delta_
):
element_size_bits(element_size_bits_),
advance_rank(advance_rank_),
threadblock_shape(threadblock_shape_),
threadmap_iterations(threadmap_iterations_),
threadmap_delta(threadmap_delta_) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Helper template to construct an PredicatedTileAccessIteratorDesc from a template
// dependent state
template <
typename Shape, typename Element, typename Layout,
int AdvanceRank, typename ThreadMap>
struct MakePredicatedTileAccessIteratorDesc;
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileAccessIterator for pitch-linear data.
template <
typename Shape, typename Element, int AdvanceRank,
typename ThreadMap>
struct MakePredicatedTileAccessIteratorDesc <
Shape, Element, layout::PitchLinear, AdvanceRank, ThreadMap> {
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc operator()() {
return PredicatedTileAccessIteratorDesc(
sizeof_bits<Element>::value,
AdvanceRank,
{Shape::kContiguous, Shape::kStrided},
{ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided},
{ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileAccessIterator for column-major data.
template <
typename Shape, typename Element, int AdvanceRank,
typename ThreadMap>
struct MakePredicatedTileAccessIteratorDesc <
Shape, Element, layout::ColumnMajor, AdvanceRank, ThreadMap> {
static int const kAdvanceRank = AdvanceRank;
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
layout::PitchLinearShape<Shape::kRow, Shape::kColumn>, Element,
layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>;
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc operator()() {
return UnderlyingMakeOperator()();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileAccessIterator for row-major data.
template <
typename Shape, typename Element, int AdvanceRank,
typename ThreadMap>
struct MakePredicatedTileAccessIteratorDesc <
Shape, Element, layout::RowMajor, AdvanceRank, ThreadMap> {
static int const kAdvanceRank = AdvanceRank;
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>, Element,
layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>;
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc operator()() {
return UnderlyingMakeOperator()();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileAccessIterator for column-major interleaved data.
template <
typename Shape, typename Element, int AdvanceRank,
typename ThreadMap, int InterleavedK>
struct MakePredicatedTileAccessIteratorDesc <
Shape, Element, layout::ColumnMajorInterleaved<InterleavedK>, AdvanceRank, ThreadMap> {
static int const kAdvanceRank = AdvanceRank;
static int const kInterleavedK = InterleavedK;
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
layout::PitchLinearShape<Shape::kRow * kInterleavedK, Shape::kColumn / kInterleavedK>, Element,
layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>;
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc operator()() {
return UnderlyingMakeOperator()();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileAccessIterator for roww-major interleaved data.
template <
typename Shape, typename Element, int AdvanceRank,
typename ThreadMap, int InterleavedK>
struct MakePredicatedTileAccessIteratorDesc <
Shape, Element, layout::RowMajorInterleaved<InterleavedK>, AdvanceRank, ThreadMap> {
static int const kAdvanceRank = AdvanceRank;
static int const kInterleavedK = InterleavedK;
using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc<
layout::PitchLinearShape<Shape::kColumn * kInterleavedK, Shape::kRow / kInterleavedK>, Element,
layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>;
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorDesc operator()() {
return UnderlyingMakeOperator()();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Parameters struct
//
struct PredicatedTileAccessIteratorParams {
using Index = int32_t;
using LongIndex = int64_t;
//
// Data members
//
/// stride of pitch-linear layout (units of Element)
LongIndex stride_;
/// amount (in byte) to increment pointer to move to next access along
/// strided dimension
LongIndex inc_strided_;
/// amount (in byte) to increment pointer from last access to first access
/// of next tile
LongIndex inc_next_;
/// amount (in byte) to increment pointer from first access of current tile
/// to first access of next tile
LongIndex inc_advance_;
//
// Methods
//
CUTLASS_HOST_DEVICE
Status initialize(LongIndex stride, PredicatedTileAccessIteratorDesc desc) {
stride_ = stride;
inc_strided_ = (LongIndex(stride_) * desc.threadmap_delta.strided()) *
desc.element_size_bits / 8;
if (desc.advance_rank) {
// advance along strided dimension
inc_advance_ =
desc.threadblock_shape.strided() * LongIndex(stride_) * desc.element_size_bits / 8;
} else {
// advance along contiguous dimension
inc_advance_ = desc.threadblock_shape.contiguous() * desc.element_size_bits / 8;
}
inc_next_ = inc_advance_ - LongIndex(desc.threadmap_iterations.strided() - 1) *
desc.threadmap_delta.strided() * LongIndex(stride_) *
desc.element_size_bits / 8;
return Status::kSuccess;
}
CUTLASS_HOST_DEVICE
Status initialize(Index stride, PredicatedTileAccessIteratorDesc desc) {
return initialize(LongIndex(stride), desc);
}
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorParams() {
initialize(LongIndex(0), PredicatedTileAccessIteratorDesc());
}
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorParams(Index stride, PredicatedTileAccessIteratorDesc desc) {
initialize(stride, desc);
}
CUTLASS_HOST_DEVICE
PredicatedTileAccessIteratorParams(LongIndex stride, PredicatedTileAccessIteratorDesc desc) {
initialize(stride, desc);
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace transform
} // namespace cutlass
////////////////////////////////////////////////////////////////////////////////

View File

@@ -189,6 +189,8 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
/// Parameters object is precomputed state and is host-constructible
class Params {
public:
using Base = typename TileAccessIterator::Params::Base;
friend PredicatedTileIterator;
private:
@@ -202,6 +204,10 @@ class PredicatedTileIterator<Shape_, Element_, layout::PitchLinear, AdvanceRank,
CUTLASS_HOST_DEVICE
Params() { }
CUTLASS_HOST_DEVICE
Params(Base const &base)
: params_(base) {}
};
private:
@@ -449,6 +455,10 @@ public:
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -651,9 +661,12 @@ public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
};
@@ -786,6 +799,657 @@ public:
////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileIterator for affine rank-2 data.
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept |
/// MaskedTileIteratorConcept
///
template <typename Shape_, typename Element_, int AdvanceRank,
typename ThreadMap_, int AccessSize>
class PredicatedTileIterator<Shape_, Element_, layout::AffineRankN<2>, AdvanceRank,
ThreadMap_, AccessSize> {
public:
static_assert(
AdvanceRank == 0 || AdvanceRank == 1,
"Specialization for pitch-linear iterator may advance along the "
"contiguous(rank=0) or strided(rank=1) dimension.");
using Shape = Shape_;
using Element = Element_;
using Layout = layout::AffineRankN<2>;
static int const kAdvanceRank = AdvanceRank;
using ThreadMap = ThreadMap_;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorRef = TensorRef<Element, Layout>;
using TensorView = TensorView<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using Pointer = Element *;
using NonConstPointer = typename platform::remove_const<Element>::type *;
/// Type used for internal memory accesses
using AccessType = AlignedArray<Element, AccessSize, (AccessSize * sizeof_bits<Element>::value / 8)>;
/// Underlying iterator to compute the addresses
using TileAccessIterator =
PredicatedTileAccessIterator<Shape, Element, Layout, kAdvanceRank,
ThreadMap, AccessType>;
static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector;
/// Fragment object to be loaded or stored
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount *
ThreadMap::kElementsPerAccess>;
/// Predicate vector stores mask to guard accesses
using Mask = typename TileAccessIterator::Mask;
/// Parameters object is precomputed state and is host-constructible
class Params {
public:
friend PredicatedTileIterator;
private:
/// Parameters object
typename TileAccessIterator::Params params_;
public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout) : params_(layout) { }
CUTLASS_HOST_DEVICE
Params() { }
};
private:
/// Internal pointer type permits fast address arithmetic
using BytePointer = char *;
private:
//
// Data members
//
/// Data member to the tile access iterator
TileAccessIterator address_iterator_;
public:
/// Constructs a TileIterator from its precomputed state, threadblock offset,
/// and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
/// Precomputed parameters object
Params const &params,
/// Pointer to start of tensor
Pointer pointer,
/// Extent of tensor
TensorCoord extent,
/// ID of each participating thread
int thread_id,
/// Initial offset of threadblock
TensorCoord const &threadblock_offset)
: address_iterator_(params.params_, pointer, extent, thread_id,
threadblock_offset) {}
/// Construct a PredicatedTileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
Params const &params, ///< Precomputed parameters object
Pointer pointer, ///< Pointer to start of tensor
TensorCoord extent, ///< Extent of tensor
int thread_id ///< ID of each participating thread
)
: PredicatedTileIterator(params, pointer, extent, thread_id,
make_Coord(0, 0)) {}
/// Adds a pointer offset in units of Element
CUTLASS_HOST_DEVICE
void add_pointer_offset(LongIndex pointer_offset) {
address_iterator_.add_pointer_offset(pointer_offset);
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the
/// iterator's internal pointer is reverted to the first "steady state" tile.
/// Subsequent calls are lightweight and must only update the internal
/// pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator &operator++() {
if (kAdvanceRank)
address_iterator_.add_tile_offset(make_Coord(0, 1));
else
address_iterator_.add_tile_offset(make_Coord(1, 0));
return *this;
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the
/// iterator's internal pointer is reverted to the first "steady state" tile.
/// Subsequent calls are lightweight and must only update the internal
/// pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator operator++(int) {
PredicatedTileIterator self(*this);
operator++();
return self;
}
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void clear_mask() { address_iterator_.clear_mask(); }
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void enable_mask() { address_iterator_.enable_mask(); }
/// Sets the predicate mask, overriding value stored in predicate iterator
CUTLASS_HOST_DEVICE
void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); }
/// Gets the mask
CUTLASS_HOST_DEVICE
void get_mask(Mask &mask) { address_iterator_.get_mask(mask); }
CUTLASS_DEVICE
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
load_with_byte_offset(frag, pointer_offset * sizeof_bits<Element>::value / 8);
}
CUTLASS_DEVICE
void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
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_PRAGMA_UNROLL
for (int v = 0; v < kAccessesPerVector; ++v) {
int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous);
address_iterator_.set_iteration_index(idx);
char const *byte_ptr = reinterpret_cast<char const *>(address_iterator_.get()) + byte_offset;
AccessType const *access_ptr = reinterpret_cast<AccessType const *>(byte_ptr);
cutlass::arch::global_load<AccessType,
sizeof(AccessType)
>(
frag_ptr[idx], access_ptr, address_iterator_.valid());
++address_iterator_;
}
}
}
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) { load_with_byte_offset(frag, 0); }
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
store_with_byte_offset(frag, pointer_offset * sizeof_bits<Element>::value / 8);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
address_iterator_.set_iteration_index(0);
AccessType const *frag_ptr = reinterpret_cast<AccessType const *>(&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_PRAGMA_UNROLL
for (int v = 0; v < kAccessesPerVector; ++v) {
int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous);
char *byte_ptr = reinterpret_cast<char *>(address_iterator_.get()) + byte_offset;
AccessType *access_ptr = reinterpret_cast<AccessType *>(byte_ptr);
if (address_iterator_.valid()) {
*access_ptr = frag_ptr[idx];
}
++address_iterator_;
}
}
}
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) { store_with_byte_offset(frag, 0); }
};
////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileIterator for affine rank 2 column-major data.
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept |
/// MaskedTileIteratorConcept
///
template <
typename Shape_,
typename Element_,
int AdvanceRank,
typename ThreadMap_,
int AccessSize
>
class PredicatedTileIterator<Shape_, Element_, layout::AffineRank2ColumnMajor, AdvanceRank, ThreadMap_, AccessSize> {
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::AffineRank2ColumnMajor;
static int const kAdvanceRank = AdvanceRank;
using ThreadMap = ThreadMap_;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorRef = TensorRef<Element, Layout>;
using TensorView = TensorView<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using Pointer = Element *;
using NonConstPointer = typename platform::remove_const<Element>::type *;
// Map to the underlying AffineRankN<2> layout
using UnderlyingIterator = PredicatedTileIterator<
layout::PitchLinearShape<Shape::kRow, Shape::kColumn>,
Element,
layout::AffineRankN<2>,
(kAdvanceRank == 0 ? 0 : 1),
ThreadMap,
AccessSize
>;
using AccessType = typename UnderlyingIterator::AccessType;
/// Fragment object to be loaded or stored
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
/// Predicate vector stores mask to guard accesses
using Mask = typename UnderlyingIterator::Mask;
/// Parameters object is precomputed state and is host-constructible
class Params {
private:
friend PredicatedTileIterator;
/// Parameters object
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) {
}
};
private:
//
// Data members
//
/// Underlying AffineRankN<2> tile iterator
UnderlyingIterator iterator_;
public:
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
Params const &params, ///< Precomputed parameters object
Pointer pointer, ///< Pointer to start of tensor
TensorCoord extent, ///< Extent of tensor
int thread_id, ///< ID of each participating thread
TensorCoord const &threadblock_offset ///< Initial offset of threadblock
):
iterator_(
params.params_,
pointer,
layout::PitchLinearCoord(extent.row(), extent.column()),
thread_id,
layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column())
) { }
/// Construct a PredicatedTileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
Params const &params, ///< Precomputed parameters object
Pointer pointer, ///< Pointer to start of tensor
TensorCoord extent, ///< Extent of tensor
int thread_id ///< ID of each participating thread
): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { }
/// Adds a pointer offset in units of Element
CUTLASS_HOST_DEVICE
void add_pointer_offset(LongIndex pointer_offset) {
iterator_.add_pointer_offset(pointer_offset);
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the iterator's
/// internal pointer is reverted to the first "steady state" tile. Subsequent calls
/// are lightweight and must only update the internal pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator &operator++() {
++iterator_;
return *this;
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the iterator's
/// internal pointer is reverted to the first "steady state" tile. Subsequent calls
/// are lightweight and must only update the internal pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator operator++(int) {
PredicatedTileIterator self(*this);
operator++();
return self;
}
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void clear_mask() {
iterator_.clear_mask();
}
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void enable_mask() {
iterator_.enable_mask();
}
/// Sets the predicate mask, overriding value stored in predicate iterator
CUTLASS_HOST_DEVICE
void set_mask(Mask const &mask) {
iterator_.set_mask(mask);
}
/// Gets the mask
CUTLASS_HOST_DEVICE
void get_mask(Mask &mask) {
iterator_.get_mask(mask);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
iterator_.load_with_pointer_offset(frag, pointer_offset);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
iterator_.load_with_byte_offset(frag, byte_offset);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) {
load_with_pointer_offset(frag, 0);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
iterator_.store_with_pointer_offset(frag, pointer_offset);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
iterator_.store_with_byte_offset(frag, byte_offset);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) {
store_with_pointer_offset(frag, 0);
}
};
////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileIterator for affine rank 2 row-major data.
///
/// Satisfies: ForwardTileIteratorConcept |
/// ReadableContiguousTileIteratorConcept |
/// WriteableContiguousTileIteratorConcept |
/// MaskedTileIteratorConcept
///
template <
typename Shape_,
typename Element_,
int AdvanceRank,
typename ThreadMap_,
int AccessSize
>
class PredicatedTileIterator<Shape_, Element_, layout::AffineRank2RowMajor, AdvanceRank, ThreadMap_, AccessSize> {
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::AffineRank2RowMajor;
static int const kAdvanceRank = AdvanceRank;
using ThreadMap = ThreadMap_;
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using TensorRef = TensorRef<Element, Layout>;
using TensorView = TensorView<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using Pointer = Element *;
using NonConstPointer = typename platform::remove_const<Element>::type *;
// Map to the underlying AffineRankN<2> layout
using UnderlyingIterator = PredicatedTileIterator<
layout::PitchLinearShape<Shape::kColumn, Shape::kRow>,
Element,
layout::AffineRankN<2>,
(kAdvanceRank == 0 ? 1 : 0),
ThreadMap,
AccessSize
>;
using AccessType = typename UnderlyingIterator::AccessType;
/// Fragment object to be loaded or stored
using Fragment = cutlass::Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
/// Predicate vector stores mask to guard accesses
using Mask = typename UnderlyingIterator::Mask;
/// Parameters object is precomputed state and is host-constructible
class Params {
private:
friend PredicatedTileIterator;
/// Parameters object
typename UnderlyingIterator::Params params_;
public:
CUTLASS_HOST_DEVICE
Params() { }
/// Construct the Params object given an AffineRankN<2> tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) {}
};
private:
//
// Data members
//
/// Underlying AffineRankN<2> tile iterator
UnderlyingIterator iterator_;
public:
/// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
Params const &params, ///< Precomputed parameters object
Pointer pointer, ///< Pointer to start of tensor
TensorCoord extent, ///< Extent of tensor
int thread_id, ///< ID of each participating thread
TensorCoord const &threadblock_offset ///< Initial offset of threadblock
):
iterator_(
params.params_,
pointer,
layout::PitchLinearCoord(extent.column(), extent.row()),
thread_id,
layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row())
) { }
/// Construct a PredicatedTileIterator with zero threadblock offset
CUTLASS_HOST_DEVICE
PredicatedTileIterator(
Params const &params, ///< Precomputed parameters object
Pointer pointer, ///< Pointer to start of tensor
TensorCoord extent, ///< Extent of tensor
int thread_id ///< ID of each participating thread
): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { }
/// Adds a pointer offset in units of Element
CUTLASS_HOST_DEVICE
void add_pointer_offset(LongIndex pointer_offset) {
iterator_.add_pointer_offset(pointer_offset);
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the iterator's
/// internal pointer is reverted to the first "steady state" tile. Subsequent calls
/// are lightweight and must only update the internal pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator &operator++() {
++iterator_;
return *this;
}
/// Advances to the next tile in memory.
///
/// The first time this method is called, predicates are updated, and the iterator's
/// internal pointer is reverted to the first "steady state" tile. Subsequent calls
/// are lightweight and must only update the internal pointer.
CUTLASS_HOST_DEVICE
PredicatedTileIterator operator++(int) {
PredicatedTileIterator self(*this);
operator++();
return self;
}
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void clear_mask() {
iterator_.clear_mask();
}
/// Clears the predicate set efficiently
CUTLASS_HOST_DEVICE
void enable_mask() {
iterator_.enable_mask();
}
/// Sets the predicate mask, overriding value stored in predicate iterator
CUTLASS_HOST_DEVICE
void set_mask(Mask const &mask) {
iterator_.set_mask(mask);
}
/// Gets the mask
CUTLASS_HOST_DEVICE
void get_mask(Mask &mask) {
iterator_.get_mask(mask);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load_with_pointer_offset(Fragment &frag, Index pointer_offset) {
iterator_.load_with_pointer_offset(frag, pointer_offset);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) {
iterator_.load_with_byte_offset(frag, byte_offset);
}
/// Loads a fragment from memory
CUTLASS_DEVICE
void load(Fragment &frag) {
load_with_pointer_offset(frag, 0);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) {
iterator_.store_with_pointer_offset(frag, pointer_offset);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) {
iterator_.store_with_byte_offset(frag, byte_offset);
}
/// Store a fragment to memory
CUTLASS_DEVICE
void store(Fragment const &frag) {
store_with_pointer_offset(frag, 0);
}
};
////////////////////////////////////////////////////////////////////////////////
/// Specialization of PredicatedTileIterator for interleaved data. It is mapped
/// to the congruous layout.
///
@@ -854,6 +1518,11 @@ class PredicatedTileIterator<Shape_, Element_,
CUTLASS_HOST_DEVICE
Params(Layout const &layout)
: params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
private:
@@ -1035,6 +1704,10 @@ class PredicatedTileIterator<Shape_, Element_,
CUTLASS_HOST_DEVICE
Params(Layout const &layout)
: params_(layout::PitchLinear(layout.stride(0))) {}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
private:

View File

@@ -194,6 +194,8 @@ class PredicatedTileIterator2dThreadTile<Shape_, Element_, layout::PitchLinear,
/// Parameters object is precomputed state and is host-constructible
class Params {
public:
using Base = typename TileAccessIterator::Params::Base;
friend PredicatedTileIterator2dThreadTile;
private:
@@ -207,6 +209,10 @@ class PredicatedTileIterator2dThreadTile<Shape_, Element_, layout::PitchLinear,
CUTLASS_HOST_DEVICE
Params() { }
CUTLASS_HOST_DEVICE
Params(Base const &base)
: params_(base) {}
};
private:
@@ -443,9 +449,11 @@ public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {}
}
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};
@@ -637,9 +645,11 @@ public:
/// Construct the Params object given a pitch-linear tensor's layout
CUTLASS_HOST_DEVICE
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {
Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { }
};
CUTLASS_HOST_DEVICE
Params(typename UnderlyingIterator::Params::Base const &base)
: params_(base) {}
};

View File

@@ -74,6 +74,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -89,7 +90,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;

View File

@@ -76,6 +76,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -107,7 +108,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_[Detail::kPointerCount];
@@ -445,6 +446,7 @@ class RegularTileAccessIterator<Shape_, Element_,
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -492,7 +494,7 @@ class RegularTileAccessIterator<Shape_, Element_,
int sections_per_stage_;
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_[Detail::kPointerCount];

View File

@@ -74,6 +74,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -107,7 +108,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;
@@ -437,6 +438,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -471,7 +473,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;
@@ -811,6 +813,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -844,7 +847,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;
@@ -1175,6 +1178,7 @@ class RegularTileAccessIterator<
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -1211,7 +1215,7 @@ class RegularTileAccessIterator<
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType *pointer_;

View File

@@ -70,11 +70,14 @@ public:
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
using Fragment = Array<Element, ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>;
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
"Advance rank may only be along the contiguous or strided dimensions.");
@@ -84,8 +87,6 @@ private:
//
// Types
//
using AccessType = AlignedArray<Element, ThreadMap::kElementsPerAccess, kAlignment>;
//
// Data members
@@ -95,7 +96,7 @@ private:
uint8_t *pointer_;
/// Stride quantity
Index stride_;
StrideIndex stride_;
/// Amount to increment pointer along strided dimension
Index increment_strided_;
@@ -242,6 +243,30 @@ public:
(coord.contiguous() * Shape::kContiguous + coord.strided() * Shape::kStrided * stride_) / 8;
add_pointer_offset(offset);
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) {
}
/// Returns a pointer
CUTLASS_HOST_DEVICE
AccessType *get() const {
#if 0
AccessType *access_ptr = pointer_[iteration_strided_ & 1];
int stride_idx = (iteration_strided_ & ~1);
int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ +
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_);
#endif
return reinterpret_cast<AccessType *>(pointer_);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -281,6 +306,8 @@ public:
kAlignment
>;
using AccessType = typename Underlying::AccessType;
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
"Advance rank may only be along the row or column dimensions.");
@@ -364,6 +391,17 @@ public:
iterator_.add_tile_offset({coord.column(), coord.row()});
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) {
}
/// Returns a pointer
CUTLASS_HOST_DEVICE
AccessType *get() const {
return iterator_.get();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -402,6 +440,8 @@ public:
ThreadMap
>;
using AccessType = typename Underlying::AccessType;
static_assert(kAdvanceRank == 0 || kAdvanceRank == 1,
"Advance rank may only be along the row or column dimensions.");
@@ -485,6 +525,17 @@ public:
iterator_.add_tile_offset({coord.row(), coord.column()});
}
/// Overrides the internal iteration index
CUTLASS_HOST_DEVICE
void set_iteration_index(int index) {
}
/// Returns a pointer
CUTLASS_HOST_DEVICE
AccessType *get() const {
return iterator_.get();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////

View File

@@ -79,6 +79,7 @@ public:
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -104,13 +105,13 @@ private:
uint8_t *pointer_;
/// Stride quantity
Index stride_;
StrideIndex stride_;
/// Amount to increment pointer along strided dimension
Index increment_strided_;
LongIndex increment_strided_;
/// Amount to advance pointer between tiles
Index increment_advance_;
LongIndex increment_advance_;
public:

View File

@@ -85,6 +85,7 @@ public:
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -123,7 +124,7 @@ private:
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType * pointer_[Detail::kPointerCount];
@@ -557,6 +558,7 @@ public:
using Index = typename Layout::Index;
using LongIndex = typename Layout::LongIndex;
using StrideIndex = typename Layout::Stride::Index;
using TensorRef = TensorRef<Element, Layout>;
using TensorCoord = typename Layout::TensorCoord;
@@ -595,7 +597,7 @@ private:
//
/// Stride value
Index stride_;
StrideIndex stride_;
/// Internal pointer to first access of tile
AccessType * pointer_[Detail::kPointerCount];