adding the cutlass examples

1 files changed, 899 insertions, 0 deletions

diff --git a/cutlass-example/cutlass/tile_iterator.h b/cutlass-example/cutlass/tile_iterator.h
new file mode 100644
index 0000000..5d39c4f
--- /dev/null
+++ b/cutlass-example/cutlass/tile_iterator.h
@@ -0,0 +1,899 @@
+/***************************************************************************************************
+ * 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);
+  }
+};
+}