gemm_global_stream.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/convert.h"
#include "cutlass/gemm/gemm_global_tile.h"
#include "cutlass/tile_allocation.h"

namespace cutlass {
namespace gemm {


template <
    GemmOperand::Kind Operand,
    typename LoadIterator_,
    typename StoreIterator_,
    typename Transformer_>

struct GlobalLoadStream {
  static GemmOperand::Kind const kOperand = Operand;
  typedef LoadIterator_ LoadIterator;
  typedef Transformer_ Transformer;
  typedef StoreIterator_ StoreIterator;

  typedef typename LoadIterator::Fragment FetchedFragment;
  typedef typename Transformer::OutputFragment TransformedFragment;
  static_assert((platform::is_same<FetchedFragment, typename Transformer::InputFragment>::value),
                "");
  typedef TransformedFragment Fragment;
  static_assert((platform::is_same<TransformedFragment, typename StoreIterator::Fragment>::value),
                "");

  static MatrixLayout::Kind const kLayout = LoadIterator::kLayout;
  typedef typename LoadIterator::Scalar Scalar;
  typedef typename LoadIterator::Pointer Pointer;
  typedef typename LoadIterator::Index Index;
  typedef typename LoadIterator::Tile Tile;

  typedef TileAllocation<typename StoreIterator::Scalar, typename StoreIterator::Tile>
      ThreadblockTileStorage;

  typedef typename ThreadblockTileStorage::TensorRef ThreadblockTileRef;

  struct Params {
    // The load iterator.
    typename LoadIterator::Params load_iterator;
    // The store iterator.
    typename StoreIterator::Params store_iterator;
    // Offset to residue.
    Index offset_to_residue;

    CUTLASS_HOST_DEVICE int initialize(Pointer pointer,
                                       long long batch_stride,
                                       Index ldm,
                                       Index _offset_to_residue) {

      offset_to_residue = _offset_to_residue;
      int error_code = load_iterator.initialize(pointer, batch_stride, ldm);
      if (error_code) {
        return error_code;
      }
      return store_iterator.initialize();
    }
  };

  struct SharedStorage {};

  //
  // Static member functions
  //

  CUTLASS_DEVICE static Coord<3> project_coordinate(Coord<3> const& coord, Index d_offset = 0) {
    bool const kKstrided =
        GemmMultiplicandTraits<typename LoadIterator::Tile, kOperand, kLayout>::kKstrided;
    Coord<3> tile_coord = ProjectOperand<kOperand, kKstrided>::project(coord);
    return make_Coord(
        tile_coord[0] + d_offset, tile_coord[1], tile_coord[2] / LoadIterator::Tile::kC);
  }

  CUTLASS_DEVICE GlobalLoadStream(
      Params const& _params,
      SharedStorage& shared_storage,
      ThreadblockTileRef const& threadblock_tile_ref,
      Coord<3> const bounds,
      Coord<3> const& _threadblock_offset)
      : params(_params),
        multiplicand_bounds(project_coordinate(bounds, 1)),
        threadblock_offset(project_coordinate(_threadblock_offset)),
        load_iterator(params.load_iterator,
          project_coordinate(bounds, 1), /*multiplicant_bounds*/
          project_coordinate(_threadblock_offset) /*threablock_offset*/),
        transformer(),
        store_iterator(params.store_iterator, threadblock_tile_ref.data())
  {
    load_iterator.initialize_predicates(multiplicand_bounds, threadblock_offset);
    fetched_fragment.clear();
  }


  CUTLASS_DEVICE void copy() { load_iterator.load_post_increment(fetched_fragment); }

  CUTLASS_DEVICE void commit() {
    transformer.transform(fetched_fragment, transformed_fragment);
    store_iterator.store_post_increment(transformed_fragment);
    store_iterator.inc_stage();
  }

  CUTLASS_DEVICE void residue(Index k, bool skip_clear = false) {
    load_iterator.residue(k);
    if (!skip_clear) {
      fetched_fragment.clear();
    }
  }

  CUTLASS_DEVICE void move_to_residue(Index k, Index kTileK) {
    Index kResidue = k % kTileK;
    if (kResidue) {
      residue(kResidue);
    }
    load_iterator.add_pointer_offset(params.offset_to_residue * load_iterator.stride_advance());
  }

  CUTLASS_DEVICE void rollback(void) {
    load_iterator.initialize_predicates(multiplicand_bounds, threadblock_offset);

    int const kBlock = kOperand == GemmOperand::kA
                           ? (kLayout == MatrixLayout::kColumnMajor ? Tile::kH : Tile::kW)
                           : (kLayout == MatrixLayout::kRowMajor ? Tile::kH : Tile::kW);

    load_iterator.add_pointer_offset(-(params.offset_to_residue + kBlock) *
                                     load_iterator.stride_advance());
  }

  CUTLASS_DEVICE GlobalLoadStream &operator+=(Coord<3> const &offset) {
    load_iterator += offset;
    return *this;
  }

  //
  // Data members
  //

  Params params;
  Coord<3> multiplicand_bounds;
  Coord<3> threadblock_offset;
  LoadIterator load_iterator;
  FetchedFragment fetched_fragment;
  Transformer transformer;
  TransformedFragment transformed_fragment;
  StoreIterator store_iterator;
};

}  // namespace gemm
}  // namespace cutlass