diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/igemm_global_tile.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/igemm_global_tile.h | 161 |
1 files changed, 161 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/igemm_global_tile.h b/cutlass-example/cutlass/gemm/igemm_global_tile.h new file mode 100644 index 0000000..3f594ac --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_global_tile.h @@ -0,0 +1,161 @@ +/*************************************************************************************************** + * 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 Implements tile iterators to partition the thread block tile into 2D subtiles and + efficiently load each. Applies permute transformation to construct 'interleaved K-strided' + data layout in which 4-element dot products from the same K index are arranged in consecutive + locations within shared memory. + + Supports efficient loads from shared memory to target the DP4A instruction. +*/ +#pragma once + +#include <cutlass/coord.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/matrix_traits.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct IgemmGlobalTileTraits : public GemmGlobalTileTraits< + // Which GEMM operand? + kOperand_, + // The layout. + kLayout_, + // The scalar. + Scalar_, + // The tile. + Tile_, + // The threads. + Threads_, + // The number of scalars per LDG/STG. + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<kOperand_, kLayout_, Scalar_, Tile_, Threads_, kAccessSize_> Base; + /// The threads. + typedef typename Base::Threads Threads; + /// The strides in each dimension between different loads/stores. + typedef Shape<Base::Threads::kH * 4, 1, Base::Threads::kW, Base::kAccessSize> Delta; + /// The number of iterations needed to load/store the tile. + typedef Shape<Base::Tile::kH / Base::Threads::kH / 4, + 4, + Base::Tile::kW / Base::Threads::kW, + Base::Tile::kC / Base::kAccessSize> + Iterations; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Threads::kW * ThreadsDelta::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; + + public: + /// The threads strides. + typedef Shape<1, 4, Base::Tile::kC> ThreadsDelta; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Deprecated. Please use IgemmGlobalTileTraits instead. + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct IgemmContiguousGlobalTileTraits + : public IgemmGlobalTileTraits<kOperand_, kLayout_, Scalar_, Tile_, Threads_, kAccessSize_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct IgemmGlobalIteratorAb : public GemmGlobalIteratorAb<TileTraits_, Index_> { + /// The base class. + typedef GemmGlobalIteratorAb<TileTraits_, Index_> Base; + /// The functor to compute the thread offset. + typedef typename TileTraits_::ThreadOffset ThreadOffset; + + /// Constructor. + CUTLASS_DEVICE IgemmGlobalIteratorAb(typename Base::Params const& _params, + const Coord<3>& bounds, + const Coord<3>& block, + ThreadOffset thread_offset_func = ThreadOffset()) + : Base(_params, bounds, block, thread_offset_func), in_residue_(false), mask_(0xffffffff) { + // The number of elements read in a single iteration. + int const kBlock = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + // The residue. + int const kResidue = (int)(bounds[1] % kBlock); + + // Compute the number of elements that are valid. + int const left = kResidue - Base::thread_offset[2]; + if (left > 0 && left < 4) { + mask_ = (1u << (8 * left)) - 1u; + } + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + Base::get(value, d, h, w, c); + if (in_residue_) { + reinterpret_cast<uint32_t&>(value) &= mask_; + } + } + + /// Move to residue portion. + CUTLASS_DEVICE void move_to_residue(typename Base::Index k) { + Base::move_to_residue(k); + in_residue_ = true; + } + + /// Move back to the beginning of the first tile. + CUTLASS_DEVICE void rollback() { + Base::rollback(); + in_residue_ = false; + } + + /// Are we in the residue? + bool in_residue_; + /// The mask to clean up the values. + uint32_t mask_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass |