tile_iterator.h

Go to the documentation of this file.

/***************************************************************************************************
 * 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.
 *
 **************************************************************************************************/
#pragma once

#include "cutlass/coord.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/fragment.h"
#include "cutlass/load_store.h"
#include "cutlass/predicate_vector.h"
#include "cutlass/vector.h"
#include <cstdio>

namespace cutlass {



struct IteratorAdvance {
  enum Kind { kD, kH, kW };
};

template <typename Tile_,
          typename Delta_,
          typename Iterations_,
          typename ThreadOffset_,
          int AccessSize>
struct TileTraits {
  typedef Tile_ Tile;

  typedef Delta_ Delta;

  typedef Iterations_ Iterations;

  //
  // ThreadOffset should be a functor defined like:
  //
  // struct ThreadOffsetExample {
  //   CUTLASS_DEVICE
  //   Coord<4> operator()() const {
  //      return make_Coord(0, threadIdx.y, threadIdx.x, 0);
  //   }
  // };
  //
  typedef ThreadOffset_ ThreadOffset;

  typedef Shape<0, 0, 0, 0> ImmediateOffsetStrides;

  static int const kAccessSize = AccessSize;
};


template <typename Delta_>
struct RegularTilePredicateFunctor {
  typedef Delta_ Delta;

  Coord<3> bounds;

  CUTLASS_HOST_DEVICE
  RegularTilePredicateFunctor(Coord<3> _bounds) : bounds(_bounds) {}

  CUTLASS_HOST_DEVICE
  bool operator()(Coord<3> iteration, Coord<3> offset) const {
    return (iteration[0] * Delta::kD + offset[0] < bounds[0]) &&
           (iteration[1] * Delta::kH + offset[1] < bounds[1]) &&
           (iteration[2] * Delta::kW + offset[2] < bounds[2]);
  }
};


template <typename T>
struct DumpType {};
template <typename Traits_,
          typename Scalar_,
          IteratorAdvance::Kind Advance_ = IteratorAdvance::kH,
          MemorySpace::Kind MemorySpace = MemorySpace::kGeneric,
          typename Index_ = int,
          typename FragmentElement_ = Scalar_,
          FragmentElementType::Kind FragmentElementType_ = FragmentElementType::kScalar,
          typename Skew_ = Shape<0, 0, 0, 0> >
struct TileIteratorBase {
  typedef Traits_ Traits;

  typedef Scalar_ Scalar;

  typedef FragmentElement_ FragmentElement;

  static IteratorAdvance::Kind const kAdvance = Advance_;

  static FragmentElementType::Kind const kFragmentElementType = FragmentElementType_;

  static MemorySpace::Kind const kMemorySpace = MemorySpace;

  typedef Index_ Index;

  typedef Skew_ Skew;

  typedef typename Traits::Tile Tile;

  typedef typename Traits::Delta Delta;

  typedef typename Traits::ImmediateOffsetStrides ImmediateOffsetStrides;

  typedef typename Traits::Iterations Iterations;

  typedef typename Traits::ThreadOffset ThreadOffset;

  static int const kAccessSize = Traits::kAccessSize;

  typedef typename Vectorize<FragmentElement, kAccessSize>::Type AccessType;

  static int const kFragmentSize =
      (kFragmentElementType == FragmentElementType::kWmmaMatrix ? 16 : sizeof(AccessType));
  typedef Fragment<Scalar, ShapeCount<Tile>::kCount, kFragmentSize> Storage;
  typedef Fragment<FragmentElement, ShapeCount<Iterations>::kCount * kAccessSize> Fragment;

  typedef FragmentIterator<Fragment, Iterations, AccessType> FragmentIterator;
  typedef FragmentConstIterator<Fragment, Iterations, AccessType> FragmentConstIterator;
  typedef typename FragmentIterator::FragmentShape FragmentShape;

  typedef PredicateVector<ShapeCount<Iterations>::kCount> PredicateVector;

  //
  // Params struct
  //

  struct Params {

    //
    // Dat members
    //

    long long stride_d;
    Index stride_h;
    Index stride_w;

    long long inc_d;
    Index inc_h;
    Index inc_w;

    long long inc_advance;

    //
    // Methods
    //

    CUTLASS_HOST_DEVICE
    Params() : stride_d(0), stride_h(0), stride_w(0), inc_d(0), inc_h(0), inc_w(0) {}

    CUTLASS_HOST_DEVICE
    Params(long long _stride_d,
           Index _stride_h,
           Index _stride_w,
           long long _inc_d,
           Index _inc_h,
           Index _inc_w,
           long long _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) {}

    CUTLASS_HOST_DEVICE
    Params(Coord<4> const &stride) {
      initialize(stride);
    }

    CUTLASS_HOST_DEVICE
    int initialize(long long _stride_d,
                   Index _stride_h,
                   Index _stride_w,
                   long long _inc_d,
                   Index _inc_h,
                   Index _inc_w,
                   long long _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(Coord<4> const &stride) {
      return initialize(stride[0], stride[1], stride[2]);
    }

    CUTLASS_HOST_DEVICE
    int initialize(long long _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);
      inc_d = stride_d * Delta::kD - stride_h * Delta::kH * (Iterations::kH - 1) -
              stride_w * Delta::kW * (Iterations::kW - 1);

      inc_advance = 0;

      if (kAdvance == IteratorAdvance::kH) {
        // Advance in the H dimension.
        inc_advance = Tile::kH * stride_h;
      } else if (kAdvance == IteratorAdvance::kW) {
        // Advance in the W dimension.
        inc_advance = Tile::kW * stride_w;

      } else {
        // Advance in the D dimension.
        inc_advance = Tile::kD * stride_d;
      }

      inc_advance -= stride_d * Delta::kD * (Iterations::kD - 1) +
                     stride_h * Delta::kH * (Iterations::kH - 1) +
                     stride_w * Delta::kW * (Iterations::kW - 1);

      return 0;
    }

    CUTLASS_HOST_DEVICE int initialize() {
      stride_d = 0;
      stride_h = 0;
      stride_w = 1;

      inc_advance = 0;
      inc_d = inc_h = inc_w = 0;

      return 0;
    }
  };

  CUTLASS_HOST_DEVICE bool valid(int d, int h, int w, int c) const { return true; }

  //
  // Static function members
  //

  template <typename PredicateIterator, typename PredicateFunctor>
  CUTLASS_HOST_DEVICE static void initialize_predicates(PredicateIterator predicate_it,
                                                   PredicateFunctor const &predicate_func,
                                                   Coord<3> const &offset) {
    CUTLASS_PRAGMA_UNROLL
    for (int d = 0; d < Iterations::kD; ++d) {
      CUTLASS_PRAGMA_UNROLL
      for (int h = 0; h < Iterations::kH; ++h) {
        CUTLASS_PRAGMA_UNROLL
        for (int w = 0; w < Iterations::kW; ++w) {
          bool enable = predicate_func(make_Coord(d, h, w), offset);
          predicate_it.set(enable);
          ++predicate_it;
        }
      }
    }
  }
};



template <typename Traits_,
          typename Scalar_,
          IteratorAdvance::Kind Advance_ = IteratorAdvance::kH,
          MemorySpace::Kind MemorySpace = MemorySpace::kGeneric,
          typename Index_ = int,
          typename FragmentElement_ = Scalar_,
          FragmentElementType::Kind FragmentElementType_ = FragmentElementType::kScalar,
          typename Skew_ = Shape<0, 0, 0, 0> >
struct TileLoadIterator : public TileIteratorBase<Traits_,
                                                  Scalar_,
                                                  Advance_,
                                                  MemorySpace,
                                                  Index_,
                                                  FragmentElement_,
                                                  FragmentElementType_,
                                                  Skew_> {
  typedef TileIteratorBase<Traits_,
                           Scalar_,
                           Advance_,
                           MemorySpace,
                           Index_,
                           FragmentElement_,
                           FragmentElementType_,
                           Skew_>
      Base;

  typedef typename Base::Traits Traits;

  typedef typename Base::Scalar Scalar;

  typedef FragmentElement_ FragmentElement;

  static IteratorAdvance::Kind const kAdvance = Base::kAdvance;

  static FragmentElementType::Kind const kFragmentElementType = FragmentElementType_;

  static MemorySpace::Kind const kMemorySpace = Base::kMemorySpace;

  typedef typename Base::Index Index;

  typedef typename Base::Skew Skew;

  typedef typename Base::Tile Tile;

  typedef typename Base::Delta Delta;

  typedef typename Base::Iterations Iterations;

  typedef typename Base::ThreadOffset ThreadOffset;

  typedef typename Base::FragmentShape FragmentShape;

  typedef typename Base::AccessType AccessType;

  static int const kAccessSize = Base::kAccessSize;

  typedef typename Base::Fragment Fragment;

  typedef typename Base::FragmentIterator FragmentIterator;

  typedef typename Base::FragmentConstIterator FragmentConstIterator;

  typedef typename Base::PredicateVector PredicateVector;

  typedef typename Base::Storage SharedStorage;

  typedef typename Base::Params BaseParams;

  enum { kRequiresLoadFence = Tile::kD == 1 };

  typedef Scalar const *Pointer;

  typedef TensorRef<Scalar const, 4> TensorRef;

  struct Params : public BaseParams {
    Scalar const *pointer;

    //
    // Methods
    //

    CUTLASS_HOST_DEVICE
    Params() : pointer(0){ Base::Params::initialize(); }

    CUTLASS_HOST_DEVICE
    Params(Scalar const *ptr) : pointer(ptr) { Base::Params::initialize(); }

    CUTLASS_HOST_DEVICE
    Params(TensorRef const &ref): pointer(ref.data()) {
      Base::Params::initialize(ref.stride());
    }

    CUTLASS_HOST_DEVICE
    Params(Scalar const *ptr,
           long long _stride_d,
           Index _stride_h,
           Index _stride_w,
           long long _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);
    }

    CUTLASS_HOST_DEVICE
    Params(Scalar const *ptr, long long stride_d, Index stride_h, Index stride_w)
        : pointer(ptr) {
      Base::Params::initialize(stride_d, stride_h, stride_w);
    }

    CUTLASS_HOST_DEVICE
    int initialize(TensorRef const &ref) {
      pointer = ref.data();
      return Base::Params::initialize(ref.stride());
    }

    CUTLASS_HOST_DEVICE
    int initialize(SharedStorage const &storage) {
      pointer = &storage[0];
      Base::Params::initialize();
      return 0;
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar const *ptr) {
      pointer = ptr;
      Base::Params::initialize();
      return 0;
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar const *ptr, long long stride_d, Index stride_h, Index stride_w) {
      Base::Params::initialize(stride_d, stride_h, stride_w);
      pointer = ptr;
      return 0;
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar const *ptr,
                   long long _stride_d,
                   Index _stride_h,
                   Index _stride_w,
                   long long _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
  //

  Params params;

  Coord<4> thread_offset;

  int stage;

  //
  // Predicate initialization
  //

  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,
        RegularTilePredicateFunctor<typename Traits::Delta>(bounds),
        block_offset + make_Coord(thread_offset[0], thread_offset[1], thread_offset[2]));
  }

  template <
      typename PredicateIterator,
      typename PredicateFunctor>
  CUTLASS_HOST_DEVICE void initialize_predicates(PredicateIterator predicate_it,
                                                 PredicateFunctor const &functor,
                                                 Coord<3> const &block_offset) {
    Base::initialize_predicates(
        predicate_it,
        functor,
        block_offset + make_Coord(thread_offset[0], thread_offset[1], thread_offset[2]));
  }

  //
  // Methods
  //

  CUTLASS_HOST_DEVICE
  TileLoadIterator() {}

  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 pointer_offset = Index((block_offset[0] + thread_offset[0]) * params.stride_d) +
                           Index((block_offset[1] + thread_offset[1]) * params.stride_h) +
                           Index((block_offset[2] + thread_offset[2]) * params.stride_w);

    params.pointer += pointer_offset;
  }

  CUTLASS_HOST_DEVICE
  TileLoadIterator(Params const &,
                   Scalar const *ptr,
                   Coord<3> const &block_offset = make_Coord(0, 0, 0),
                   ThreadOffset thread_offset_func = ThreadOffset())
      : stage(0) {
    params.pointer = ptr + thread_offset_func()[2];

    params.stride_d = 0;
    params.stride_h = 0;
    params.stride_w = 1;

    params.inc_d = params.inc_h = params.inc_w = params.inc_advance = 0;
  }

  CUTLASS_HOST_DEVICE void inc_d() { params.pointer += params.inc_d; }

  CUTLASS_HOST_DEVICE void inc_h() { params.pointer += params.inc_h; }

  CUTLASS_HOST_DEVICE void inc_w() { params.pointer += params.inc_w; }

  CUTLASS_HOST_DEVICE void inc_advance() { params.pointer += params.inc_advance; }

  CUTLASS_HOST_DEVICE void load_element(AccessType &value, int d, int h, int w, int c) const {
    int const offset =
        ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, c);
    Load<Scalar,
         kAccessSize,
         kMemorySpace,
         kFragmentElementType,
         FragmentElement,
         Tile::kW,
         sizeof(FragmentElement) * kAccessSize>::load(value, params.pointer, offset);
  }

  CUTLASS_HOST_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;
      }
    }
  }

  CUTLASS_HOST_DEVICE TileLoadIterator & operator+=(Coord<3> const &offset) {
    long long _offset = offset.template dot<long long>(
      make_Coord(params.stride_d, params.stride_h, params.stride_w)
    );

    params.pointer += _offset;
    return *this;
  }

  CUTLASS_HOST_DEVICE void add_pointer_offset(Index offset) { params.pointer += offset; }

  CUTLASS_HOST_DEVICE Index stride_advance(void) {
    Index stride = params.stride_h;
    if (kAdvance == IteratorAdvance::kW) {
      stride = params.stride_w;
    }
    return stride;
  }

  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) {
          for (int c = 0; c < Iterations::kC; ++c) {
            if (*pred_it) {
              load_element(
                  reinterpret_cast<AccessType &>(frag_iterator.at(d, h, w, c)), d, h, w, c);
            }
          }
          if (w < Iterations::kW - 1) {
            inc_w();
          }
        }
        if (h < Iterations::kH - 1) {
          inc_h();
        }
      }
      if (d < Iterations::kD - 1) {
        inc_d();
      }
    }
    inc_advance();
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load_post_increment(Fragment &fragment) {
    typename PredicateVector::TrivialIterator pred_it;
    load_post_increment(fragment, pred_it);
  }

  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);
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load(Fragment &fragment) const {
    typename PredicateVector::TrivialIterator pred_it;
    load(fragment, pred_it);
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load(Fragment &fragment, int d) {
    FragmentIterator frag_iterator(fragment);
    for (int h = 0; h < Iterations::kH; ++h) {
      for (int w = 0; w < Iterations::kW; ++w) {
        for (int c = 0; c < Iterations::kC; ++c) {
          load_element(reinterpret_cast<AccessType &>(frag_iterator.at(0, h, w, c)), d, h, w, c);
        }
      }
    }
  }
};



template <typename Traits_,
          typename Scalar_,
          IteratorAdvance::Kind Advance_ = IteratorAdvance::kH,
          MemorySpace::Kind MemorySpace = MemorySpace::kGeneric,
          typename Index_ = int,
          typename FragmentElement_ = Scalar_,
          FragmentElementType::Kind FragmentElementType_ = FragmentElementType::kScalar,
          typename Skew_ = Shape<0, 0, 0, 0> >
struct TileStoreIterator : public TileIteratorBase<Traits_,
                                                   Scalar_,
                                                   Advance_,
                                                   MemorySpace,
                                                   Index_,
                                                   FragmentElement_,
                                                   FragmentElementType_,
                                                   Skew_> {
  typedef TileIteratorBase<Traits_,
                           Scalar_,
                           Advance_,
                           MemorySpace,
                           Index_,
                           FragmentElement_,
                           FragmentElementType_,
                           Skew_>
      Base;

  typedef typename Base::Traits Traits;

  typedef typename Base::Scalar Scalar;

  typedef typename Base::FragmentElement FragmentElement;

  static IteratorAdvance::Kind const kAdvance = Base::kAdvance;

  static FragmentElementType::Kind const kFragmentElementType = Base::kFragmentElementType;

  static MemorySpace::Kind const kMemorySpace = Base::kMemorySpace;

  static int const kAccessSize = Base::kAccessSize;

  typedef typename Base::Index Index;

  typedef typename Base::Skew Skew;

  typedef typename Base::Tile Tile;

  typedef typename Base::Delta Delta;

  typedef typename Base::Iterations Iterations;

  typedef typename Base::ThreadOffset ThreadOffset;

  typedef typename Base::FragmentShape FragmentShape;

  typedef typename Base::AccessType AccessType;

  typedef typename Base::Fragment Fragment;

  typedef typename Base::FragmentIterator FragmentIterator;

  typedef typename Base::FragmentConstIterator FragmentConstIterator;

  typedef typename Base::PredicateVector PredicateVector;

  typedef typename Base::Storage SharedStorage;

  typedef typename Base::Params BaseParams;

  typedef Scalar *Pointer;

  typedef TensorRef<Scalar, 4> TensorRef;

  struct Params : public BaseParams {
    Scalar *pointer;

    //
    // Methods
    //

    // Default constructor
    CUTLASS_HOST_DEVICE
    Params() : pointer(0) {}

    // Default constructor
    CUTLASS_HOST_DEVICE
    Params(Scalar *ptr) : pointer(ptr) { Base::Params::initialize(); }

    CUTLASS_HOST_DEVICE
    Params(TensorRef const &ref): pointer(ref.data()) {
      Base::Params::initialize(ref.stride());
    }

    // Default constructor
    CUTLASS_HOST_DEVICE
    Params(Scalar *ptr, long long stride_d, Index stride_h, Index stride_w) {
      initialize(ptr, stride_d, stride_h, stride_w);
    }

    // Default constructor
    CUTLASS_HOST_DEVICE
    Params(Scalar *ptr,
           long long _stride_d,
           Index _stride_h,
           Index _stride_w,
           long long _inc_d,
           Index _inc_h,
           Index _inc_w,
           Index _inc_advance) {
      initialize(ptr, _stride_d, _stride_h, _stride_w, _inc_d, _inc_h, _inc_w, _inc_advance);
    }

    CUTLASS_HOST_DEVICE
    int initialize(SharedStorage &storage) {
      pointer = &storage[0];
      return Base::Params::initialize();
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar *ptr) {
      pointer = ptr;
      return Base::Params::initialize();
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar *ptr, long long stride_d, Index stride_h, Index stride_w) {
      Base::Params::initialize(stride_d, stride_h, stride_w);
      pointer = ptr;
      return 0;
    }

    CUTLASS_HOST_DEVICE
    int initialize(Scalar *ptr,
                   long long _stride_d,
                   Index _stride_h,
                   Index _stride_w,
                   long long _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;
    }

    CUTLASS_HOST_DEVICE
    int initialize() { return Base::Params::initialize(); }
  };

  //
  // Data members
  //

  Params params;

  Coord<4> thread_offset;

  int stage;

  //
  // Predicate initialization
  //

  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,
        RegularTilePredicateFunctor<typename Traits::Delta>(bounds),
        block_offset + make_Coord(thread_offset[0], thread_offset[1], thread_offset[2]));
  }

  template <
      typename PredicateIterator,
      typename PredicateFunctor>
  CUTLASS_HOST_DEVICE void initialize_predicates(PredicateIterator predicate_it,
                                                 PredicateFunctor const &functor,
                                                 Coord<3> const &block_offset) {
    Base::initialize_predicates(
        predicate_it,
        functor,
        block_offset + make_Coord(thread_offset[0], thread_offset[1], thread_offset[2]));
  }

  //
  // Methods
  //

  CUTLASS_HOST_DEVICE
  TileStoreIterator() {}

  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] + thread_offset[0]) * params.stride_d +
                      (block_offset[1] + thread_offset[1]) * params.stride_h +
                      (block_offset[2] + thread_offset[2]) * params.stride_w;
  }

  CUTLASS_HOST_DEVICE
  TileStoreIterator(Params const &, Scalar *ptr, ThreadOffset thread_offset_func = ThreadOffset())
      : stage(0) {
    params.pointer = ptr + thread_offset_func()[2];
    params.stride_d = 0;
    params.stride_h = 0;
    params.stride_w = 1;

    params.inc_d = params.inc_h = params.inc_w = params.inc_advance = 0;
  }

  CUTLASS_HOST_DEVICE void inc_d() { params.pointer += params.inc_d; }

  CUTLASS_HOST_DEVICE void inc_h() { params.pointer += params.inc_h; }

  CUTLASS_HOST_DEVICE void inc_w() { params.pointer += params.inc_w; }

  CUTLASS_HOST_DEVICE void inc_advance() {}

  CUTLASS_HOST_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;
      }
    }
  }

  CUTLASS_HOST_DEVICE TileStoreIterator & operator+=(Coord<3> const &offset) {
    params.pointer += offset.template dot<long long>(
      make_Coord(params.stride_d, params.stride_h, params.stride_w)
    );
    return *this;
  }

  CUTLASS_HOST_DEVICE void add_pointer_offset(Index offset) { params.pointer += offset; }

  CUTLASS_HOST_DEVICE void store_element(AccessType const &value, int d, int h, int w, int c) {
    int const offset =
        ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, c);
    Store<Scalar,
          kAccessSize,
          kMemorySpace,
          kFragmentElementType,
          FragmentElement,
          Tile::kW,
          sizeof(FragmentElement) * kAccessSize>::store(value, params.pointer, offset);
  }

  template <typename Fragment, typename PredicateIterator>
  CUTLASS_HOST_DEVICE void store_post_increment(Fragment const &fragment, PredicateIterator pred_it) {
    FragmentConstIterator 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) {
          for (int c = 0; c < Iterations::kC; ++c) {
            if (*pred_it) {
              store_element(
                  reinterpret_cast<AccessType const &>(frag_iterator.at(d, h, w, c)), d, h, w, c);
            }
          }
          if (w < Iterations::kW - 1) {
            inc_w();
          }
        }
        if (h < Iterations::kH - 1) {
          inc_h();
        }
      }
      if (d < Iterations::kD - 1) {
        inc_d();
      }
    }
    inc_advance();
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void store_post_increment(Fragment const &fragment) {
    typename PredicateVector::TrivialIterator pred_it;
    store_post_increment(fragment, pred_it);
  }

  template <typename Fragment, typename PredicateIterator>
  CUTLASS_HOST_DEVICE void store(Fragment const &fragment, PredicateIterator pred_it) const {
    TileStoreIterator _store_it(*this);
    _store_it.store_post_increment(fragment, pred_it);
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void store(Fragment const &fragment) const {
    typename PredicateVector::TrivialIterator pred_it;
    store(fragment, pred_it);
  }

  CUTLASS_HOST_DEVICE void load_element(AccessType &value, int d, int h, int w, int c) const {
    int const offset =
        ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, c);

    Load<Scalar,
         kAccessSize,
         kMemorySpace,
         kFragmentElementType,
         FragmentElement,
         Tile::kW,
         sizeof(FragmentElement) * kAccessSize>::load(value, params.pointer, offset);
  }

  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) {
          for (int c = 0; c < Iterations::kC; ++c) {
            if (*pred_it) {
              load_element(
                  reinterpret_cast<AccessType &>(frag_iterator.at(d, h, w, c)), d, h, w, c);
            }
          }
          if (w < Iterations::kW - 1) {
            inc_w();
          }
        }
        if (h < Iterations::kH - 1) {
          inc_h();
        }
      }
      if (d < Iterations::kD - 1) {
        inc_d();
      }
    }
    inc_advance();
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load_post_increment(Fragment &fragment) {
    typename PredicateVector::TrivialIterator pred_it;
    load_post_increment(fragment, pred_it);
  }

  template <typename Fragment, typename PredicateIterator>
  CUTLASS_HOST_DEVICE void load(Fragment &fragment, PredicateIterator pred_it) const {
    TileStoreIterator _load_it(*this);
    _load_it.load_post_increment(fragment, pred_it);
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load(Fragment &fragment) const {
    typename PredicateVector::TrivialIterator pred_it;
    load(fragment, pred_it);
  }

  template <typename Fragment>
  CUTLASS_HOST_DEVICE void load(Fragment &fragment, int d) {
    FragmentIterator frag_iterator(fragment);
    for (int h = 0; h < Iterations::kH; ++h) {
      for (int w = 0; w < Iterations::kW; ++w) {
        for (int c = 0; c < Iterations::kC; ++c) {
          load_element(reinterpret_cast<AccessType &>(frag_iterator.at(0, h, w, c)), d, h, w, c);
        }
      }
    }
  }
};


} // namespace cutlass