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

namespace cutlass {
namespace gemm {


template <
    typename LoadIterator_,
    typename StoreIterator_,
    typename Transformer_>

struct GlobalLoadStreamBase {
  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;

  struct Params {
    // The load iterator.
    typename LoadIterator::Params load_iterator;
    // The store iterator.
    typename StoreIterator::Params store_iterator;

    CUTLASS_HOST_DEVICE int initialize(Pointer pointer, Index ld) {
      int error_code = load_iterator.initialize(pointer, ld);
      if (error_code) {
        return error_code;
      }

      return store_iterator.initialize();
    }
  };

  typedef typename StoreIterator::SharedStorage SharedStoreStorage;

  union SharedStorage {
    // The load iterator.
    typename LoadIterator::SharedStorage load_iterator;
    // The store iterator.
    SharedStoreStorage store_iterator;
  };

  CUTLASS_DEVICE GlobalLoadStreamBase(Params const& params,
                                      SharedStorage& shared_storage,
                                      Coord<3> const bounds,
                                      Coord<3> const& block)
      : load_iterator(params.load_iterator, bounds, block),
        transformer(),
        store_iterator(params.store_iterator, shared_storage.store_iterator)

  {
    fetched_fragment.clear();
  }

  CUTLASS_DEVICE void copy() { iterator_load(load_iterator, fetched_fragment); }

  CUTLASS_DEVICE void commit() {
    transformer.transform(fetched_fragment, transformed_fragment);
    iterator_store(store_iterator, 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();
    }
  }

  LoadIterator load_iterator;
  FetchedFragment fetched_fragment;
  Transformer transformer;
  TransformedFragment transformed_fragment;
  StoreIterator store_iterator;
};


template <
    typename LoadIterator_,
    typename StoreIterator_,
    typename Transformer_ = Copy<typename LoadIterator_::Fragment> >

struct GlobalLoadStream : public GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> {
  typedef GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> Base;

  CUTLASS_DEVICE GlobalLoadStream(typename Base::Params const& params,
                                  typename Base::SharedStorage& shared_storage,
                                  Coord<3> const& bounds,
                                  Coord<3> const& block)
      : Base(params, shared_storage, bounds, block) {}
};

}  // namespace gemm
}  // namespace cutlass