326 lines
15 KiB
C++
326 lines
15 KiB
C++
/***************************************************************************************************
|
|
* 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 Free functions for loading and storing to implementations of tile iteartor concepts.
|
|
*/
|
|
#pragma once
|
|
|
|
#include <cutlass/fragment_load_store.h>
|
|
#include <cutlass/load_store.h>
|
|
#include <cutlass/predicate_vector.h>
|
|
#include <cutlass/shape.h>
|
|
|
|
namespace cutlass {
|
|
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Loads a fragment from an input iterator
|
|
template <typename InputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_load(InputIterator &iterator, Fragment &fragment) {
|
|
typename InputIterator::FragmentIterator frag_iterator(fragment);
|
|
for (int d = 0; d < InputIterator::Iterations::kD; ++d) {
|
|
for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
|
|
for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
|
|
for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
|
|
if (iterator.valid(d, h, w, c)) {
|
|
int const offset =
|
|
ComputeOffsetFromStrides<typename InputIterator::ImmediateOffsetStrides>::get(
|
|
0, 0, w, c);
|
|
Load<typename Fragment::Element, InputIterator::Tile::kC, InputIterator::kMemorySpace>::
|
|
load(reinterpret_cast<typename InputIterator::AccessType &>(
|
|
frag_iterator.at(d, h, w, c)),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
if (w < InputIterator::Iterations::kW - 1) {
|
|
iterator.inc_w();
|
|
}
|
|
}
|
|
if (h < InputIterator::Iterations::kH - 1) {
|
|
iterator.inc_h();
|
|
}
|
|
}
|
|
if (d < InputIterator::Iterations::kD - 1) {
|
|
iterator.inc_d();
|
|
}
|
|
}
|
|
iterator.inc_advance();
|
|
}
|
|
|
|
/// Loads a fragment from a shared memory input iterator
|
|
template <typename InputIterator, typename Fragment>
|
|
CUTLASS_DEVICE void shared_iterator_load(InputIterator &iterator, Fragment &fragment) {
|
|
typename InputIterator::FragmentIterator frag_iterator(fragment);
|
|
for (int d = 0; d < InputIterator::Iterations::kD; ++d) {
|
|
for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
|
|
for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
|
|
for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
|
|
int const offset =
|
|
ComputeOffsetFromStrides<typename InputIterator::ImmediateOffsetStrides>::get(
|
|
d, h, w, c);
|
|
|
|
FragmentLoad<InputIterator::kIteratorFragment,
|
|
InputIterator::Tile::kC,
|
|
typename InputIterator::Scalar,
|
|
InputIterator::kMemorySpace,
|
|
typename InputIterator::FragmentElement,
|
|
InputIterator::Tile::kW>::load(frag_iterator.at(d, h, w, c),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Loads a fragment from a shared memory input iterator
|
|
template <typename InputIterator, typename Fragment>
|
|
CUTLASS_DEVICE void shared_iterator_load(InputIterator &iterator, Fragment &fragment, int d) {
|
|
typename InputIterator::FragmentIterator frag_iterator(fragment);
|
|
for (int h = 0; h < InputIterator::Iterations::kH; ++h) {
|
|
for (int w = 0; w < InputIterator::Iterations::kW; ++w) {
|
|
for (int c = 0; c < InputIterator::Iterations::kC; ++c) {
|
|
int const offset =
|
|
ComputeOffsetFromStrides<typename InputIterator::ImmediateOffsetStrides>::get(
|
|
d, h, w, c);
|
|
|
|
FragmentLoad<InputIterator::kIteratorFragment,
|
|
InputIterator::Tile::kC,
|
|
typename InputIterator::Scalar,
|
|
InputIterator::kMemorySpace,
|
|
typename InputIterator::FragmentElement,
|
|
InputIterator::Tile::kW>::load(frag_iterator.at(0, h, w, c),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Loads a fragment from an input iterator, masked by a predicate iterator
|
|
template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
|
|
Fragment &fragment,
|
|
typename InputIterator::Index offset,
|
|
ConstPredicateAdapter predicate_adapter) {
|
|
for (int d = 0; d < InputIterator::Iterations::kD; ++d, iterator.inc_d()) {
|
|
for (int h = 0; h < InputIterator::Iterations::kH; ++h, iterator.inc_h()) {
|
|
for (int w = 0; w < InputIterator::Iterations::kW; ++w, iterator.inc_w()) {
|
|
if (predicate_adapter.at(d, h, w, 0)) {
|
|
int idx = InputIterator::Tile::kC *
|
|
(w + InputIterator::Iterations::kW * (h + InputIterator::Iterations::kH * d));
|
|
|
|
Load<typename Fragment::Element, InputIterator::Tile::kC, InputIterator::kMemorySpace>::
|
|
load(reinterpret_cast<typename InputIterator::AccessType &>(fragment[idx]),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Loads a fragment from an input iterator
|
|
template <typename InputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
|
|
Fragment &fragment,
|
|
typename InputIterator::Index offset = 0) {
|
|
TrivialPredicateTileAdapter pred;
|
|
iterator_load_post_increment(iterator, fragment, offset, pred);
|
|
}
|
|
|
|
/// Loads a fragment from an input iterator
|
|
template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_load_post_increment(InputIterator &iterator,
|
|
Fragment &fragment,
|
|
ConstPredicateAdapter pred_it) {
|
|
iterator_load_post_increment(iterator, fragment, 0, pred_it);
|
|
}
|
|
|
|
template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &_iterator,
|
|
Fragment &fragment,
|
|
typename InputIterator::Index offset,
|
|
ConstPredicateAdapter predicate_adapter) {
|
|
InputIterator iterator(_iterator);
|
|
iterator_load_post_increment(iterator, fragment, offset, predicate_adapter);
|
|
}
|
|
|
|
/// Loads a fragment from an input iterator
|
|
template <typename InputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &iterator,
|
|
Fragment &fragment,
|
|
typename InputIterator::Index offset = 0) {
|
|
TrivialPredicateTileAdapter pred;
|
|
iterator_load(iterator, fragment, offset, pred);
|
|
}
|
|
|
|
/// Loads a fragment from an input iterator
|
|
template <typename InputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_load(InputIterator const &iterator,
|
|
Fragment &fragment,
|
|
ConstPredicateAdapter pred_it) {
|
|
iterator_load(iterator, fragment, 0, pred_it);
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Stores a fragment to an output iterator
|
|
template <typename OutputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_store(OutputIterator &iterator, Fragment &fragment) {
|
|
typename OutputIterator::FragmentIterator frag_iterator(fragment);
|
|
for (int d = 0; d < OutputIterator::Iterations::kD; ++d) {
|
|
for (int h = 0; h < OutputIterator::Iterations::kH; ++h) {
|
|
for (int w = 0; w < OutputIterator::Iterations::kW; ++w) {
|
|
if (iterator.valid(d, h, w, 0)) {
|
|
int const offset =
|
|
ComputeOffsetFromStrides<typename OutputIterator::ImmediateOffsetStrides>::get(
|
|
d, h, w, 0);
|
|
|
|
Store<typename Fragment::Element,
|
|
OutputIterator::Tile::kC,
|
|
OutputIterator::kMemorySpace>::
|
|
store(reinterpret_cast<typename OutputIterator::AccessType &>(
|
|
frag_iterator.at(d, h, w, 0)),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
if (w < OutputIterator::Iterations::kW - 1) {
|
|
iterator.inc_w();
|
|
}
|
|
}
|
|
if (h < OutputIterator::Iterations::kH - 1) {
|
|
iterator.inc_h();
|
|
}
|
|
}
|
|
if (d < OutputIterator::Iterations::kD - 1) {
|
|
iterator.inc_d();
|
|
}
|
|
}
|
|
iterator.inc_advance();
|
|
}
|
|
|
|
/// Stores a fragment to a shared memory output iterator
|
|
template <typename OutputIterator, typename Fragment>
|
|
CUTLASS_DEVICE void shared_iterator_store(OutputIterator &iterator, Fragment const &fragment) {
|
|
typename OutputIterator::FragmentConstIterator frag_iterator(fragment);
|
|
for (int d = 0; d < OutputIterator::Iterations::kD; ++d) {
|
|
for (int h = 0; h < OutputIterator::Iterations::kH; ++h) {
|
|
for (int w = 0; w < OutputIterator::Iterations::kW; ++w) {
|
|
for (int c = 0; c < OutputIterator::Iterations::kC; ++c) {
|
|
int const offset =
|
|
ComputeOffsetFromStrides<typename OutputIterator::ImmediateOffsetStrides>::get(
|
|
d, h, w, c);
|
|
|
|
FragmentStore<OutputIterator::kIteratorFragment,
|
|
OutputIterator::Tile::kC,
|
|
typename OutputIterator::Scalar,
|
|
OutputIterator::kMemorySpace,
|
|
typename OutputIterator::FragmentElement,
|
|
OutputIterator::Tile::kW>::store(frag_iterator.at(d, h, w, c),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Stores a fragment to an output iterator, masked by a predicate iterator
|
|
template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
|
|
Fragment const &fragment,
|
|
typename OutputIterator::Index offset,
|
|
ConstPredicateAdapter predicate_adapter) {
|
|
for (int d = 0; d < OutputIterator::Iterations::kD; ++d, iterator.inc_d()) {
|
|
for (int h = 0; h < OutputIterator::Iterations::kH; ++h, iterator.inc_h()) {
|
|
for (int w = 0; w < OutputIterator::Iterations::kW; ++w, iterator.inc_w()) {
|
|
if (predicate_adapter.at(d, h, w, 0)) {
|
|
int idx = OutputIterator::Tile::kC *
|
|
(w + OutputIterator::Iterations::kW * (h + OutputIterator::Iterations::kH * d));
|
|
|
|
Store<typename Fragment::Element,
|
|
OutputIterator::Tile::kC,
|
|
OutputIterator::kMemorySpace>::
|
|
store(reinterpret_cast<typename OutputIterator::AccessType const &>(fragment[idx]),
|
|
iterator.data(),
|
|
offset);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Stores a fragment to an output iterator
|
|
template <typename OutputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
|
|
Fragment const &fragment,
|
|
typename OutputIterator::Index offset = 0) {
|
|
TrivialPredicateTileAdapter pred;
|
|
iterator_store_post_increment(iterator, fragment, offset, pred);
|
|
}
|
|
|
|
/// Stores a fragment to an output iterator
|
|
template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_store_post_increment(OutputIterator &iterator,
|
|
Fragment const &fragment,
|
|
ConstPredicateAdapter pred_it) {
|
|
iterator_store_post_increment(iterator, fragment, 0, pred_it);
|
|
}
|
|
|
|
/// Stores a fragment to an output iterator, masked by a predicate iterator
|
|
template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &_iterator,
|
|
Fragment const &fragment,
|
|
typename OutputIterator::Index offset,
|
|
ConstPredicateAdapter predicate_adapter) {
|
|
OutputIterator iterator(_iterator);
|
|
iterator_store_post_increment(iterator, fragment, offset, predicate_adapter);
|
|
}
|
|
|
|
/// Stores a fragment to an output iterator
|
|
template <typename OutputIterator, typename Fragment>
|
|
CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &iterator,
|
|
Fragment const &fragment,
|
|
typename OutputIterator::Index offset = 0) {
|
|
TrivialPredicateTileAdapter pred;
|
|
iterator_store(iterator, fragment, offset, pred);
|
|
}
|
|
|
|
/// Stores a fragment to an output iterator
|
|
template <typename OutputIterator, typename Fragment, typename ConstPredicateAdapter>
|
|
CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &iterator,
|
|
Fragment const &fragment,
|
|
ConstPredicateAdapter pred_it) {
|
|
iterator_store(iterator, fragment, 0, pred_it);
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
} // namespace cutlass
|