From fcf40649feb6046fb9b1ed984fb9b19422cd5463 Mon Sep 17 00:00:00 2001 From: aamir Date: Sat, 21 Jul 2018 19:30:40 -0700 Subject: adding the cutlass examples --- cutlass-example/cutlass/iterator_access.h | 318 ++++++++++++++++++++++++++++++ 1 file changed, 318 insertions(+) create mode 100644 cutlass-example/cutlass/iterator_access.h (limited to 'cutlass-example/cutlass/iterator_access.h') diff --git a/cutlass-example/cutlass/iterator_access.h b/cutlass-example/cutlass/iterator_access.h new file mode 100644 index 0000000..e94beb7 --- /dev/null +++ b/cutlass-example/cutlass/iterator_access.h @@ -0,0 +1,318 @@ +/*************************************************************************************************** + * 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 +#include +#include +#include + +namespace cutlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Loads a fragment from an input iterator +template +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)) { + iterator.get(reinterpret_cast( + frag_iterator.at(d, h, w, c)), + d, + h, + w, + c); + } + } + 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 +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::get( + d, h, w, c); + + FragmentLoad::load(frag_iterator.at(d, h, w, c), + iterator.data(), + offset); + } + } + } + } +} + +/// Loads a fragment from a shared memory input iterator +template +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::get( + d, h, w, c); + + FragmentLoad::load(frag_iterator.at(0, h, w, c), + iterator.data(), + offset); + } + } + } +} + +/// Loads a fragment from an input iterator, masked by a predicate iterator +template +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:: + load(reinterpret_cast(fragment[idx]), + iterator.data(), + offset); + } + } + } + } +} + +/// Loads a fragment from an input iterator +template +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 +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 +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 +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 +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 +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)) { + iterator.set(reinterpret_cast( + frag_iterator.at(d, h, w, 0)), + d, + h, + w, + 0); + } + 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 +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::get( + d, h, w, c); + + FragmentStore::store(frag_iterator.at(d, h, w, c), + iterator.data(), + offset); + } + } + } + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Stores a fragment to an output iterator, masked by a predicate iterator +template +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:: + store(reinterpret_cast(fragment[idx]), + iterator.data(), + offset); + } + } + } + } +} + +/// Stores a fragment to an output iterator +template +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 +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 +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 +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 +CUTLASS_HOST_DEVICE void iterator_store(OutputIterator const &iterator, + Fragment const &fragment, + ConstPredicateAdapter pred_it) { + iterator_store(iterator, fragment, 0, pred_it); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cutlass -- cgit v1.3