diff options
Diffstat (limited to 'cutlass-example/cutlass/tile_iterator.h')
| -rw-r--r-- | cutlass-example/cutlass/tile_iterator.h | 899 |
1 files changed, 0 insertions, 899 deletions
diff --git a/cutlass-example/cutlass/tile_iterator.h b/cutlass-example/cutlass/tile_iterator.h deleted file mode 100644 index 5d39c4f..0000000 --- a/cutlass-example/cutlass/tile_iterator.h +++ /dev/null @@ -1,899 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017-2018, 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 Defines the Tile Traits concept and iterators for loading and storing to tiles - efficiently. -*/ -#pragma once - -#include <cutlass/fragment.h> -#include <cutlass/load_store.h> -#include <cutlass/predicate_vector.h> -#include <cutlass/vector.h> - -namespace cutlass { - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/*!@defgroup tile_traits_concept Tile Traits Concept -@{ - -\ref tile_traits_concept is a type definining the shape of a tile and the distribution of accesses -by individual entities, either threads or other. - -@par Tile Traits Concept - Types satisfying \ref tile_traits_concept define the following members - - <b>Tile</b> - a type satisfying \ref layout_concept describing the dimensions of the tile - - <b>Delta</b> - a type satisfying \ref layout_concept describing the increments between accesses -along each dimension - - <b>Iterations</b> - a type satisfying \ref layout_concept describing the number of accesses -along each dimension - - <b>Offset</b> - the type of a <i>functor</i> computing the offset of each participating entity -as a Coord<4>. -@} -*/ - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/// Specifies dimension in which post-increment accesses advance -struct IteratorAdvance { - enum Kind { kD, kH, kW }; -}; - -/// Specifies whether iterator storage fragment consists of Scalar values or WMMA matrix -struct IteratorFragment { - enum Kind { kScalar, kWmmaMatrix }; -}; - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/** -* @brief A template defining \ref tile_traits_concept -* @concept{tile_traits_concept} -*/ -template <typename Tile_, - typename Delta_, - typename Iterations_, - typename ThreadOffset_, - int kAccessSize> -struct TileTraits { - /// Shape of the tile - typedef Tile_ Tile; - - /// Number of steps between accesses along each dimension - typedef Delta_ Delta; - - /// Number of accesses performed - typedef Iterations_ Iterations; - - /// Functor that returns the logical coordinate of each entity's initial offset in the tile - typedef ThreadOffset_ ThreadOffset; -}; - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/// Iterator for accessing a stripmined tile in memory -template <typename Traits_, - typename Scalar_, - IteratorAdvance::Kind Advance_ = IteratorAdvance::kH, - MemorySpace::Kind MemorySpace = MemorySpace::kGeneric, - typename Index_ = int, - typename FragmentElement_ = Scalar_, - IteratorFragment::Kind IteratorFragment_ = IteratorFragment::kScalar, - typename Skew_ = Shape<0, 0, 0, 0> > -struct TileIteratorBase { - /// concept TileTraits - typedef Traits_ Traits; - - /// Scalar element - typedef Scalar_ Scalar; - - /// Fragment element - typedef FragmentElement_ FragmentElement; - - /// Specifies dimension in which post-increment accesses advance. - static IteratorAdvance::Kind const kAdvance = Advance_; - - /// Specifies iterator storage fragment type (Scalar or WmmaMatrix) - static IteratorFragment::Kind const kIteratorFragment = IteratorFragment_; - - /// Source or destination memory space - static MemorySpace::Kind const kMemorySpace = MemorySpace; - - /// Index type - typedef Index_ Index; - - /// Skew quantity - typedef Skew_ Skew; - - /// Tile shape - typedef typename Traits::Tile Tile; - - /// Distance along each dimension - typedef typename Traits::Delta Delta; - - /// The strides in each dimension between different loads/stores. - typedef typename Traits::ImmediateOffsetStrides ImmediateOffsetStrides; - - /// Iterations - typedef typename Traits::Iterations Iterations; - - /// Thread offset - typedef typename Traits::ThreadOffset ThreadOffset; - - /// The number of scalars accessed per load/store. - static int const kAccessSize = Tile::kC; - - /// The elements loaded/store by one instruction. - typedef typename Vectorize<FragmentElement, kAccessSize>::Type AccessType; - - /// The size of storage needed per fragment - static int const kFragmentSize = - (kIteratorFragment == IteratorFragment::kWmmaMatrix ? 16 : sizeof(AccessType)); - /// The storage. - typedef Fragment<Scalar, ShapeCount<Tile>::kCount, kFragmentSize> Storage; - /// The fragment. - typedef Fragment<FragmentElement, ShapeCount<Iterations>::kCount * kAccessSize> Fragment; - /// The fragment iterator. - typedef FragmentIterator<Fragment, Iterations, AccessType> FragmentIterator; - /// The fragment const iterator. - typedef FragmentConstIterator<Fragment, Iterations, AccessType> FragmentConstIterator; - /// The shape of the fragment. - typedef typename FragmentIterator::FragmentShape FragmentShape; - - /// Default predicate mask type - typedef PredicateVector<ShapeCount<Iterations>::kCount> PredicateVector; - - // - // Params struct - // - - /// Parameters to the iterator - struct Params { - Index stride_d; - Index stride_h; - Index stride_w; - - Index inc_d; - Index inc_h; - Index inc_w; - - Index inc_advance; - - /// Initializes params - CUTLASS_HOST_DEVICE - int initialize(Index _stride_d, - Index _stride_h, - Index _stride_w, - Index _inc_d, - Index _inc_h, - Index _inc_w, - Index _inc_advance) { - stride_d = _stride_d; - stride_h = _stride_h; - stride_w = _stride_w; - - inc_d = _inc_d; - inc_h = _inc_h; - inc_w = _inc_w; - inc_advance = _inc_advance; - - return 0; - } - - CUTLASS_HOST_DEVICE - int initialize(Index _stride_d, Index _stride_h, Index _stride_w) { - stride_d = _stride_d; - stride_h = _stride_h; - stride_w = _stride_w; - - inc_w = stride_w * Delta::kW; - inc_h = stride_h * Delta::kH - stride_w * Delta::kW * (Iterations::kW - 1); - - if (kAdvance == IteratorAdvance::kH) { - // Advance in the H dimension. - inc_d = 0; - } else if (kAdvance == IteratorAdvance::kW) { - // Advance in the W dimension. - inc_d = stride_w * Tile::kW - stride_h * Tile::kH; - } else { - // Advance in the D dimension. - inc_d = stride_d; - } - - inc_advance = 0; - - return 0; - } - - CUTLASS_HOST_DEVICE int initialize() { - stride_d = 0; - stride_h = 0; - stride_w = 1; - - inc_d = inc_h = inc_w = inc_advance = 0; - - return 0; - } - }; - - /// Is the iterator valid? - CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { return true; } - - // - // Static function members - // - - /// Initializes a predicate vector - template <typename PredicateIterator> - CUTLASS_DEVICE static void initialize_predicates(PredicateIterator predicate_it, - Coord<3> const &bounds, - Coord<3> const &offset = make_Coord(0, 0, 0)) { - for (int d = 0; d < Iterations::kD; ++d) { - bool enable_d = (d * Delta::kD + offset[0] < bounds[0]); - for (int h = 0; h < Iterations::kH; ++h) { - bool enable_h = (h * Delta::kH + offset[1] < bounds[1]); - for (int w = 0; w < Iterations::kW; ++w) { - bool enable_w = (w * Tile::kC * Delta::kW + offset[2] < bounds[2]); - predicate_it.set(d, h, w, 0, enable_d && enable_h && enable_w); - } - } - } - } -}; - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/*!@defgroup tile_load_iterator_concept Tile Load Iterator Concept -@{ - -\ref tile_load_iterator_concept enables loading a tile from addressable memory into a fragment - -@par Tile Load Iterator Concept - Types satisfying \ref tile_load_iterator_concept define the following members - - <b>PredicateVector</b> - a \ref predicate_vector_concept with sufficient predicate storage for -each access implied by the tile traits - - <b>Fragment</b> - the destination fragment type satisfying \ref fragment_concept - - <b>initialize_predicates(pred_it, bounds, block_offset)</b> - function initializing a predicate -vector according to externally specified bounds - - <b>load_post_increment(fragment, pred_it)</b> - a method that loads a fragment and increments -the iterator to the next tile, guarded by a \ref predicate_iterator_concept - - <b>load_post_increment(fragment)</b> - a method that loads a fragment and increments the -iterator to the next tile - - <b>load(fragment, pred_it)</b> - a const method that loads a fragment, guarded by a \ref -predicate_iterator_concept - - <b>load(fragment)</b> - a method that loads a fragment - -@} -*/ - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/** -* @brief An iterator implementing \ref tile_load_iterator_concept for loading a tile from memory -* @concept{tile_load_iterator_concept} -*/ -template <typename Traits_, - typename Scalar_, - IteratorAdvance::Kind Advance_ = IteratorAdvance::kH, - MemorySpace::Kind MemorySpace = MemorySpace::kGeneric, - typename Index_ = int, - typename FragmentElement_ = Scalar_, - IteratorFragment::Kind IteratorFragment_ = IteratorFragment::kScalar, - typename Skew_ = Shape<0, 0, 0, 0> > -struct TileLoadIterator : public TileIteratorBase<Traits_, - Scalar_, - Advance_, - MemorySpace, - Index_, - FragmentElement_, - IteratorFragment_, - Skew_> { - /// Base class - typedef TileIteratorBase<Traits_, - Scalar_, - Advance_, - MemorySpace, - Index_, - FragmentElement_, - IteratorFragment_, - Skew_> - Base; - - /// concept TileTraits - typedef typename Base::Traits Traits; - - /// Scalar element - typedef typename Base::Scalar Scalar; - - /// Fragment element - typedef typename Base::FragmentElement FragmentElement; - - /// Specifies in which dimension post-increment accesses advance. - static IteratorAdvance::Kind const kAdvance = Base::kAdvance; - - /// Specifies type of iterator fragment storage (Salar or WmmaMatrix) - static IteratorFragment::Kind const kIteratorFragment = Base::kIteratorFragment; - - /// Source or destination memory space - static MemorySpace::Kind const kMemorySpace = Base::kMemorySpace; - - /// Index type - typedef typename Base::Index Index; - - /// Skew quantity - typedef typename Base::Skew Skew; - - /// Tile shape - typedef typename Base::Tile Tile; - - /// Delta - typedef typename Base::Delta Delta; - - /// Iterations - typedef typename Base::Iterations Iterations; - - /// ThreadOffset functor - typedef typename Base::ThreadOffset ThreadOffset; - - /// Fragment type - typedef typename Base::FragmentShape FragmentShape; - - /// Memory access type - typedef typename Base::AccessType AccessType; - - /// Fragment definition - typedef typename Base::Fragment Fragment; - - /// Fragment iterator definition - typedef typename Base::FragmentIterator FragmentIterator; - - /// Fragment const iterator definition - typedef typename Base::FragmentConstIterator FragmentConstIterator; - - /// Default predicate mask type - typedef typename Base::PredicateVector PredicateVector; - - /// Storage object that may be loaded from - typedef typename Base::Storage SharedStorage; - - /// IteratorBase parameters - typedef typename Base::Params BaseParams; - - /// Do we require a fence? - enum { kRequiresLoadFence = Tile::kD == 1 }; - - /// The pointer type - typedef Scalar const *Pointer; - - /// Parameters - struct Params : public BaseParams { - /// Pointer to memory - Scalar const *pointer; - - /// Initialize params to access storage object - CUTLASS_HOST_DEVICE - int initialize(SharedStorage const &storage) { - pointer = &storage[0]; - return 0; - } - - /// Initializes params to access a raw pointer - CUTLASS_HOST_DEVICE - int initialize(Scalar const *ptr, Index stride_d, Index stride_h, Index stride_w) { - Base::Params::initialize(stride_d, stride_h, stride_w); - pointer = ptr; - return 0; - } - - /// Initializes params - CUTLASS_HOST_DEVICE - int initialize(Scalar const *ptr, - Index _stride_d, - Index _stride_h, - Index _stride_w, - Index _inc_d, - Index _inc_h, - Index _inc_w, - Index _inc_advance) { - pointer = ptr; - Base::Params::initialize( - _stride_d, _stride_h, _stride_w, _inc_d, _inc_h, _inc_w, _inc_advance); - return 0; - } - - // Initializes params to default values - CUTLASS_HOST_DEVICE - int initialize() { return Base::Params::initialize(); } - }; - - // - // Data members - // - - /// Parameters structure - Params params; - - /// Offset of an individual lane from the start of the tile - Coord<4> thread_offset; - - /// Stage argument enables wrapping after some number of tiles have been loaded. - int stage; - - // - // Static member functions - // - - /// Initializes a predicate vector - template <typename PredicateIterator> - CUTLASS_HOST_DEVICE void initialize_predicates(PredicateIterator predicate_it, - Coord<3> const &bounds, - Coord<3> const &block_offset = make_Coord(0, - 0, - 0)) { - Base::initialize_predicates( - predicate_it, - bounds, - block_offset + make_Coord(0, thread_offset[1], thread_offset[2] * Tile::kC)); - } - - // - // Methods - // - - /// Default constructor - CUTLASS_HOST_DEVICE - TileLoadIterator() {} - - /// Constructs a tile load iterator - CUTLASS_HOST_DEVICE - TileLoadIterator(Params const &_params, - Coord<3> const &block_offset = make_Coord(0, 0, 0), - ThreadOffset thread_offset_func = ThreadOffset()) - : params(_params), stage(0) { - thread_offset = thread_offset_func(); - - Index block_offset_h = 0; - Index block_offset_w = 0; - if (kAdvance == IteratorAdvance::kH) { - block_offset_h = block_offset[1]; - block_offset_w = block_offset[2]; - } else { - block_offset_h = block_offset[2]; - block_offset_w = block_offset[1]; - } - - params.pointer += block_offset[0] * params.stride_d + - (block_offset_h + thread_offset[1]) * params.stride_h + - (block_offset_w + thread_offset[2] * Tile::kC) / Tile::kC * params.stride_w; - } - - /// Constructs a tile load iterator - CUTLASS_HOST_DEVICE - TileLoadIterator(Params const &, - SharedStorage &shared_storage, - Coord<3> const &block_offset = make_Coord(0, 0, 0), - ThreadOffset thread_offset_func = ThreadOffset()) - : stage(0) { - int const offset = thread_offset_func()[2]; - params.pointer = &shared_storage[offset]; - } - - /// Returns the current pointer - CUTLASS_HOST_DEVICE - Scalar const *data() const { return params.pointer; } - - /// The accessor. - CUTLASS_DEVICE void get(AccessType &value, int d, int h, int w, int c) const { - int const imm = - ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, c); - Load<Scalar, Base::kAccessSize, kMemorySpace>::load(value, params.pointer, imm); - } - - /// Increment in the D dimension - CUTLASS_HOST_DEVICE void inc_d() { params.pointer += params.inc_d; } - - /// Increment in the H dimension - CUTLASS_HOST_DEVICE void inc_h() { params.pointer += params.inc_h; } - - /// Increment in the W dimension - CUTLASS_HOST_DEVICE void inc_w() { params.pointer += params.inc_w; } - - /// Increment in the next dimension - CUTLASS_HOST_DEVICE void inc_advance() { params.pointer += params.inc_advance; } - - /// Increment the stage. - CUTLASS_DEVICE void inc_stage() { - if (Tile::kD > 1) { - int const kStageSize = Tile::kH * Tile::kW * Tile::kC; - if (stage == Tile::kD - 1) { - params.pointer -= (Tile::kD - 1) * kStageSize; - stage = 0; - } else { - params.pointer += kStageSize; - stage = stage + 1; - } - } - } - - public: - /// Loads a fragment and advances the iterator to the next tile. - template <typename Fragment, typename PredicateIterator> - CUTLASS_HOST_DEVICE void load_post_increment(Fragment &fragment, PredicateIterator pred_it) { - FragmentIterator frag_iterator(fragment); - - for (int d = 0; d < Iterations::kD; ++d) { - for (int h = 0; h < Iterations::kH; ++h) { - for (int w = 0; w < Iterations::kW; ++w, ++pred_it) { - if (*pred_it) { - Load<typename Fragment::Element, Tile::kC, kMemorySpace>::load( - reinterpret_cast<AccessType &>(frag_iterator.at(d, h, w, 0)), data(), 0); - } - - if (w < Iterations::kW - 1) { - inc_w(); - } - } - if (h < Iterations::kH - 1) { - inc_h(); - } - } - if (d < Iterations::kD - 1) { - inc_d(); - } - } - inc_advance(); - } - - /// Loads a fragment and advances the iterator to the next tile. - template <typename Fragment> - CUTLASS_HOST_DEVICE void load_post_increment(Fragment &fragment) { - typename PredicateVector::TrivialIterator pred_it; - load_post_increment(fragment, pred_it); - } - - /// Loads a fragment without advancing the iterator.. - template <typename Fragment, typename PredicateIterator> - CUTLASS_HOST_DEVICE void load(Fragment &fragment, PredicateIterator pred_it) const { - TileLoadIterator _load_it(*this); - _load_it.load_post_increment(fragment, pred_it); - } - - /// Loads a fragment without advancing the iterator.. - template <typename Fragment> - CUTLASS_HOST_DEVICE void load(Fragment &fragment) const { - typename PredicateVector::TrivialIterator pred_it; - load(fragment, pred_it); - } -}; - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/*!@defgroup tile_store_iterator_concept Tile Store Iterator Concept -@{ - -\ref tile_store_iterator_concept enables storing a tile to addressable memory - -@par Tile Store Iterator Concept - Types satisfying \ref tile_load_iterator_concept define the following members - - <b>PredicateVector</b> - a \ref predicate_vector_concept with sufficient predicate storage for -each access implied by the tile traits - - <b>Fragment</b> - the destination fragment type satisfying \ref fragment_concept - - <b>initialize_predicates(pred_it, bounds, block_offset)</b> - function initializing a predicate -vector according to externally specified bounds - - <b>store_post_increment(fragment, pred_it)</b> - a method that stores a fragment and increments -the iterator to the next tile, guarded by a \ref predicate_iterator_concept - - <b>store_post_increment(fragment)</b> - a method that stores a fragment and increments the -iterator to the next tile - - <b>store(fragment, pred_it)</b> - a const method that stores a fragment, guarded by a \ref -predicate_iterator_concept - - <b>store(fragment)</b> - a method that loads a fragment - -@} -*/ - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -/** -* @brief An iterator implementing \ref tile_store_iterator_concept for storing a tile to memory -* @concept{tile_store_iterator_concept} -*/ -template <typename Traits_, - typename Scalar_, - IteratorAdvance::Kind Advance_ = IteratorAdvance::kH, - MemorySpace::Kind MemorySpace = MemorySpace::kGeneric, - typename Index_ = int, - typename FragmentElement_ = Scalar_, - IteratorFragment::Kind IteratorFragment_ = IteratorFragment::kScalar, - typename Skew_ = Shape<0, 0, 0, 0> > -struct TileStoreIterator : public TileIteratorBase<Traits_, - Scalar_, - Advance_, - MemorySpace, - Index_, - FragmentElement_, - IteratorFragment_, - Skew_> { - /// Base class - typedef TileIteratorBase<Traits_, - Scalar_, - Advance_, - MemorySpace, - Index_, - FragmentElement_, - IteratorFragment_, - Skew_> - Base; - - /// concept TileTraits - typedef typename Base::Traits Traits; - - /// Scalar element - typedef typename Base::Scalar Scalar; - - /// Fragment element - typedef typename Base::FragmentElement FragmentElement; - - /// Specifies in which dimension post-increment accesses advance. - static IteratorAdvance::Kind const kAdvance = Base::kAdvance; - - /// Specifies type of iterator fragment storage (Salar or WmmaMatrix) - static IteratorFragment::Kind const kIteratorFragment = Base::kIteratorFragment; - - /// Source or destination memory space - static MemorySpace::Kind const kMemorySpace = Base::kMemorySpace; - - /// Index type - typedef typename Base::Index Index; - - /// Skew quantity - typedef typename Base::Skew Skew; - - /// Tile shape - typedef typename Base::Tile Tile; - - /// Delta - typedef typename Base::Delta Delta; - - /// Iterations - typedef typename Base::Iterations Iterations; - - /// ThreadOffset functor - typedef typename Base::ThreadOffset ThreadOffset; - - /// Fragment type - typedef typename Base::FragmentShape FragmentShape; - - /// Memory access type - typedef typename Base::AccessType AccessType; - - /// Fragment definition - typedef typename Base::Fragment Fragment; - - /// Fragment iterator definition - typedef typename Base::FragmentIterator FragmentIterator; - - /// Fragment const iterator definition - typedef typename Base::FragmentConstIterator FragmentConstIterator; - - /// Default predicate mask type - typedef typename Base::PredicateVector PredicateVector; - - /// Storage object which may be stored to - typedef typename Base::Storage SharedStorage; - - /// IteratorBase parameters - typedef typename Base::Params BaseParams; - - /// Parameters - struct Params : public BaseParams { - /// Pointer to memory - Scalar *pointer; - - /// Initialize params to access storage object - CUTLASS_HOST_DEVICE - int initialize(SharedStorage &storage) { - pointer = &storage[0]; - return 0; - } - - /// Initializes params to access a raw pointer - CUTLASS_HOST_DEVICE - int initialize(Scalar *ptr, Index stride_d, Index stride_h, Index stride_w) { - Base::Params::initialize(stride_d, stride_h, stride_w); - pointer = ptr; - return 0; - } - - /// Initializes params - CUTLASS_HOST_DEVICE - int initialize(Scalar *ptr, - Index _stride_d, - Index _stride_h, - Index _stride_w, - Index _inc_d, - Index _inc_h, - Index _inc_w, - Index _inc_advance) { - pointer = ptr; - Base::Params::initialize( - _stride_d, _stride_h, _stride_w, _inc_d, _inc_h, _inc_w, _inc_advance); - return 0; - } - - /// Initializes params to default values - CUTLASS_HOST_DEVICE - int initialize() { return Base::Params::initialize(); } - }; - - // - // Data members - // - - /// Parameters structure - Params params; - - /// Offset of an individual lane from the start of the tile - Coord<4> thread_offset; - - /// The stage. - int stage; - - // - // Static member functions - // - - /// Initializes a predicate vector - template <typename PredicateIterator> - CUTLASS_HOST_DEVICE void initialize_predicates(PredicateIterator predicate_it, - Coord<3> const &bounds, - Coord<3> const &block_offset = make_Coord(0, - 0, - 0)) { - Base::initialize_predicates( - predicate_it, - bounds, - block_offset + make_Coord(0, thread_offset[1], thread_offset[2] * Tile::kC)); - } - - // - // Methods - // - - /// Default constructor - CUTLASS_HOST_DEVICE - TileStoreIterator() {} - - /// Constructs a tile store iterator - CUTLASS_HOST_DEVICE - TileStoreIterator(Params const &_params, - Coord<3> const &block_offset = make_Coord(0, 0, 0), - ThreadOffset thread_offset_func = ThreadOffset()) - : params(_params), stage(0) { - thread_offset = thread_offset_func(); - - params.pointer += block_offset[0] * params.stride_d + - (block_offset[1] + thread_offset[1]) * params.stride_h + - (block_offset[2] + thread_offset[2] * Tile::kC) / Tile::kC * params.stride_w; - } - - /// Constructs a tile store iterator - CUTLASS_HOST_DEVICE - TileStoreIterator(Params const &, - SharedStorage &shared_storage, - Coord<3> const &block_offset = make_Coord(0, 0, 0), - ThreadOffset thread_offset_func = ThreadOffset()) - : stage(0) { - int const offset = thread_offset_func()[2]; - params.pointer = &shared_storage[offset]; - } - - /// Returns the current pointer - CUTLASS_HOST_DEVICE - Scalar *data() const { return params.pointer; } - - /// Increment in the D dimension - CUTLASS_HOST_DEVICE void inc_d() { params.pointer += params.inc_d; } - - /// Increment in the H dimension - CUTLASS_HOST_DEVICE void inc_h() { params.pointer += params.inc_h; } - - /// Increment in the W dimension - CUTLASS_HOST_DEVICE void inc_w() { params.pointer += params.inc_w; } - - /// Increment in the next dimension - CUTLASS_HOST_DEVICE void inc_advance() {} - - /// Increment the stage. - CUTLASS_DEVICE void inc_stage() { - if (Tile::kD > 1) { - int const kStageSize = Tile::kH * Tile::kW * Tile::kC; - if (stage == Tile::kD - 1) { - params.pointer -= (Tile::kD - 1) * kStageSize; - stage = 0; - } else { - params.pointer += kStageSize; - stage = stage + 1; - } - } - } - - /// The accessor. - CUTLASS_DEVICE void set(AccessType const &value, int d, int h, int w, int c) { - int const imm = - ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, c); - Store<Scalar, Base::kAccessSize, kMemorySpace>::store(value, params.pointer, imm); - } - - public: - /// Stores a fragment and advances to the next tile. - template <typename Fragment, typename PredicateIterator> - CUTLASS_HOST_DEVICE void store_post_increment(Fragment &fragment, PredicateIterator pred_it) { - FragmentIterator frag_iterator(fragment); - - for (int d = 0; d < Iterations::kD; ++d) { - for (int h = 0; h < Iterations::kH; ++h) { - for (int w = 0; w < Iterations::kW; ++w, ++pred_it) { - if (*pred_it) { - Store<typename Fragment::Element, Tile::kC, kMemorySpace>::store( - reinterpret_cast<AccessType &>(frag_iterator.at(d, h, w, 0)), data(), 0); - } - if (w < Iterations::kW - 1) { - inc_w(); - } - } - if (h < Iterations::kH - 1) { - inc_h(); - } - } - if (d < Iterations::kD - 1) { - inc_d(); - } - } - inc_advance(); - } - - /// Stores a fragment and advances to the next tile. - template <typename Fragment> - CUTLASS_HOST_DEVICE void store_post_increment(Fragment &fragment) { - typename PredicateVector::TrivialIterator pred_it; - store_post_increment(fragment, pred_it); - } - - /// Stores a fragment without advancing the iterator. - template <typename Fragment, typename PredicateIterator> - CUTLASS_HOST_DEVICE void store(Fragment &fragment, PredicateIterator pred_it) const { - TileStoreIterator _store_it(*this); - _store_it.store_post_increment(fragment, pred_it); - } - - /// Stores a fragment without advancing the iterator. - template <typename Fragment> - CUTLASS_HOST_DEVICE void store(Fragment &fragment) const { - typename PredicateVector::TrivialIterator pred_it; - store(fragment, pred_it); - } -}; -} |