diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/gemm_global_tile.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/gemm_global_tile.h | 541 |
1 files changed, 541 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm_global_tile.h b/cutlass-example/cutlass/gemm/gemm_global_tile.h new file mode 100644 index 0000000..1cc3b33 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_global_tile.h @@ -0,0 +1,541 @@ +/*************************************************************************************************** + * 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 Defines iterators for efficiently loading and storing to global memory. +*/ +#pragma once + +#include <cutlass/coord.h> +#include <cutlass/util/platform.h> + +#include <cutlass/gemm/gemm_operand.h> +#include <cutlass/matrix_traits.h> +#include <cutlass/predicate_vector.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/tile_iterator.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The following functor reshapes a tile of threads to match a tile of data. The idea is that when +// the user wants to build the iterator traits, he/she may want to specify the tile independently +// from the number of scalars loaded/stored per instruction. For example, in the row-major version +// with a tile of size 128x8 - the user may want to that the iterator works with 32x8 threads if +// each thread loads 1 scalar per LDG. If the user changes to 4 scalars per LDG, then the tile of +// threads has to change. The code below detects that and correct the code automatically - it is +// a helper when the user does not specify the right configuration. + +template <typename Tile_, typename Threads_, bool = (Tile_::kW < Threads_::kW)> +struct ReshapeThreads { + typedef Threads_ Threads; +}; + +template <typename Tile_, typename Threads_> +struct ReshapeThreads<Tile_, Threads_, true> { + typedef Shape<Threads_::kD, Threads_::kH * Threads_::kW / Tile_::kW, Tile_::kW, 1> Threads; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct GemmGlobalTileTraits { + /// Identity of the operand + static GemmOperand::Kind const kOperand = kOperand_; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kAccessSize_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kGlobal; + + /// The tile shape + typedef typename ReshapeTile<Tile_, kAccessSize_>::Tile Tile; + /// The threads shape + typedef typename ReshapeThreads<Tile, Threads_>::Threads Threads; + /// The relative offset between two elements in the H/W dimension in adjacent threads. + typedef Shape<1, 1, Tile::kC> ThreadsDelta; + + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH, Threads::kW * kAccessSize> Delta; + /// Strides for immediate offset computation + typedef Shape<0, 0, Threads::kW * ThreadsDelta::kW, kAccessSize> ImmediateOffsetStrides; + /// The number of iterations needed to load/store the tile. + typedef Shape<1, Tile::kH / Threads::kH, Tile::kW / Threads::kW, Tile::kC / kAccessSize> + Iterations; + + typedef GemmMultiplicandTraits<Tile, kOperand, kLayout> MultiplicandTraits; + + /// 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); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kStrideH_, int kAccessSize_> +struct GemmGlobalTileCdTraits : public GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> + Base; + + /// The stride in the H dimension. + static int const kStrideH = kStrideH_; + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta; + + typedef typename Base::Iterations Iterations; + + typedef typename Base::Threads Threads; + + typedef typename Base::ThreadsDelta ThreadsDelta; + + typedef typename Base::ImmediateOffsetStrides ImmediateOffsetStrides; + + /// 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 * kStrideH * Iterations::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct GemmGlobalIteratorAb + : public TileLoadIterator<TileTraits_, + typename TileTraits_::Scalar, + TileTraits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH + : IteratorAdvance::kW, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef GemmGlobalIteratorAb<TileTraits_, Index_> This_; /// The base class. + + typedef TileLoadIterator<TileTraits_, + typename TileTraits_::Scalar, + TileTraits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH + : IteratorAdvance::kW, + MemorySpace::kGlobal, + Index_> + Base; + /// The layout. + static MatrixLayout::Kind const kLayout = TileTraits_::kLayout; + /// Fragment type loaded by the iterator + typedef typename Base::Fragment Fragment; + /// The scalar. + typedef typename TileTraits_::Scalar Scalar; + /// The threads. + typedef typename TileTraits_::Threads Threads; + /// The index. + typedef Index_ Index; + /// The thread offset + typedef typename TileTraits_::ThreadOffset ThreadOffset; + /// Specifies in which dimension post-increment accesses advance. + static IteratorAdvance::Kind const kAdvance = Base::kAdvance; + + typedef cutlass::PredicateVector<ShapeCount<typename Base::Iterations>::kCount> PredicateVector; + + /// Iterator parameters type + typedef typename Base::Params BaseParams; + + struct Params : public BaseParams { + /// Initializes params to load a strip-mined tile, given pointer and stride_h. + template <typename GemmDesc_> + CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc, Scalar const* ptr, Index stride_h) { + Index inc_d = 0; + Index inc_advance = 0; + // Move by some columns for each iteration in the H dimension. + Index inc_h = Base::Delta::kH * stride_h; + + // Move by some more columns in the number of iterations if the D dimension is > 1. + if (Base::Delta::kD > 0) { + inc_d = Base::Delta::kD * stride_h - (Base::Iterations::kH - 1) * inc_h; + } + + // Move to the beginning of the next iteration. + if (kAdvance == IteratorAdvance::kH && Base::Delta::kD > 0) { + inc_advance = inc_d; + } else if (kAdvance == IteratorAdvance::kH) { + inc_advance = inc_h; + } else if (Base::Delta::kD > 0) { + inc_advance = (Base::Iterations::kW + 0) * ShapeCount<typename Base::Delta>::kWc - + (Base::Iterations::kH - 1) * inc_h - + (Base::Iterations::kD - 1) * Base::Delta::kD * stride_h; + } else { + inc_advance = (Base::Iterations::kW + 0) * ShapeCount<typename Base::Delta>::kWc - + (Base::Iterations::kH - 1) * inc_h; + } + + // The dimensions of the tile. + int const kH = TileTraits_::Tile::kH; + int const kW = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + + // Move to the residue. + Index const kBlock = kAdvance == IteratorAdvance::kH ? kH : kW; + // The jump in the gemm-k dimension. + Index const stride = kAdvance == IteratorAdvance::kH ? stride_h : 1; + + // Compute the offset to the residue and how to "come" back. + Index const kResidue = desc.k % kBlock; + if (kResidue > 0) { + move_to_residue_offset = (desc.k - kResidue) * stride; + } else { + move_to_residue_offset = (desc.k - kBlock) * stride; + } + + Base::Params::initialize(ptr, 0, stride_h, 1, inc_d, inc_h, 0, inc_advance); + return 0; + } + + // The extra offset to control moving to the residue. + Index move_to_residue_offset; + }; + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorAb(Params const& _params, + const Coord<3>& bounds, + const Coord<3>& block, + ThreadOffset thread_offset_func = ThreadOffset()) + : params(_params) { + thread_offset = thread_offset_func(); + // The column. + Index block_h = thread_offset[1]; + // The contiguous dimension. + Index block_w = thread_offset[2]; + + // Add the blocks indices. + if (kAdvance == IteratorAdvance::kH) { + block_h += block[1]; + block_w += block[2]; + + } else { + block_h += block[2]; + block_w += block[1]; + } + + // Setup the pointer. + params.pointer += (block_h * params.stride_h + block_w); + + // Initialize predicates + initialize_predicates(bounds, make_Coord(0, block_h, block_w)); + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm); + } + + /// Increment the pointer in the H dimension. + CUTLASS_DEVICE void inc_h() { params.pointer += params.inc_h; } + /// Increment the pointer in the D dimension. + CUTLASS_DEVICE void inc_d() { params.pointer += params.inc_d; } + /// Increment the pointer to move to the next iteration. + CUTLASS_DEVICE void inc_advance() { params.pointer += params.inc_advance; } + + /// Initialize the predicates. + CUTLASS_DEVICE void initialize_predicates(const Coord<3>& bounds, const Coord<3>& block) { + // Setup the masks to control loads. + predicates.fill(0); + + int bounds_h, bounds_w; + if (kAdvance == IteratorAdvance::kH) { + bounds_w = bounds[2] - block[2]; + bounds_h = bounds[1]; + + } else { + bounds_w = bounds[1]; + bounds_h = bounds[2] - block[1]; + } + + // Fill in the bits of the predicate vector. + for (int d = 0; d < Base::Iterations::kD; ++d) { + for (int h = 0; h < Base::Iterations::kH; ++h) { + for (int w = 0; w < Base::Iterations::kW; ++w) { + for (int c = 0; c < Base::Iterations::kC; ++c) { + bool flag = w * Base::Delta::kW < bounds_w; + if (kAdvance == IteratorAdvance::kH) { + flag = flag && (h * Base::Delta::kH + d * Base::Delta::kD) < bounds_h; + } else { + flag = flag && (h * Base::Delta::kH) < bounds_h; + } + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + predicates.set(bit, flag); + } + } + } + } + } + + /// Move to residue portion. + CUTLASS_DEVICE void move_to_residue(Index k) { + // Store the pointer and the predicates. + stored_pointer = params.pointer; + stored_predicates = predicates; + + // Move the pointer to the residue. + params.pointer += params.move_to_residue_offset; + + // The dimensions of the tile. + int const kH = TileTraits_::Tile::kH; + int const kW = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + + // The unrolling factor. + int const kUnroll = kAdvance == IteratorAdvance::kH ? kH : kW; + + // Clear the predicates for the residue. TODO: We can do something smarter. + int const kResidue = (int)(k % (Index)kUnroll); + if (kResidue > 0) { + residue(kResidue); + } + } + + /// That's the residue! Update the predicates. + CUTLASS_DEVICE void residue(Index k) { + // The coordinates of the thread. + Index block_h = thread_offset[1]; + // The contiguous dimension. + Index block_w = thread_offset[2]; + + // Update the predicate vector. + for (int d = 0; d < Base::Iterations::kD; ++d) { + for (int h = 0; h < Base::Iterations::kH; ++h) { + for (int w = 0; w < Base::Iterations::kW; ++w) { + for (int c = 0; c < Base::Iterations::kC; ++c) { + Index offset = 0; + if (kAdvance == IteratorAdvance::kH) { + offset += block_h + h * Base::Delta::kH + d * Base::Delta::kD; + } else { + offset += block_w + w * Base::Delta::kW; + } + + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + if (offset >= k) { + predicates.set(bit, false); + } + } + } + } + } + } + + /// Rollback to beginning of first tile and initialize predicates. + CUTLASS_DEVICE void rollback() { + params.pointer = stored_pointer; + predicates = stored_predicates; + } + + /// Is the iterator valid? + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + return predicates[bit]; + } + + /// Offset of an individual lane from the start of the tile + Coord<4> thread_offset; + /// The parameters + Params params; + /// The pointer. + typename Base::Scalar const* stored_pointer; + /// The predicates. + PredicateVector predicates, stored_predicates; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct GemmGlobalIteratorCd : public TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef GemmGlobalIteratorCd<TileTraits_, Index_> This_; + /// The base class. + typedef TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> + Base; + + /// The layout. + static MatrixLayout::Kind const kLayout = TileTraits_::kLayout; + + /// The scalar. + typedef typename TileTraits_::Scalar Scalar; + /// The pointer. + typedef typename TileTraits_::Pointer Pointer; + /// The threads. + typedef typename TileTraits_::Threads Threads; + /// The index. + typedef Index_ Index; + /// The thread offset + typedef typename TileTraits_::ThreadOffset ThreadOffset; + + /// The params. + struct Params { + /// The pointer. + Pointer pointer; + /// The stride in the H dimension to setup the thread in the block. + Index stride_h; + /// The strides to increment the pointer. + Index inc_advance, inc_h; + /// The strides to increment the predicate offset + Index predicate_inc_advance, predicate_inc_h; + /// The column offset to compute the predicate for the columns. + Index predicate_offset; + + /// Setup the params. + CUTLASS_HOST_DEVICE int initialize( + Pointer pointer, Index ld, Index bound, Index epilogue_stride_w, Index epilogue_delta_w) { + // The pointer. + this->pointer = pointer; + // Each column of the matrix. + stride_h = TileTraits_::ThreadsDelta::kH * ld; + // Each thread output 1 column per iteration. The stride between columns is given by the + // number of scalars that are loaded per LDS for B. + inc_h = ld * TileTraits_::kStrideH; + inc_advance = + (ld - ld * TileTraits_::kStrideH * (Base::Iterations::kH - 1)) + epilogue_stride_w; + + predicate_offset = bound; + predicate_inc_h = TileTraits_::kStrideH; + predicate_inc_advance = + -((TileTraits_::kStrideH * (Base::Iterations::kH - 1) - 1) + epilogue_delta_w); + + return 0; + } + }; + + Params params; + /// Offset of an individual lane from the start of the tile + Coord<4> thread_offset; + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorCd() {} + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorCd(Params const& params, + const Coord<3>& bounds, + const Coord<3>& block, + int offset = 0, + int pred_offset = 0, + ThreadOffset thread_offset_func = ThreadOffset()) + : params(params) { + thread_offset = thread_offset_func(); + // Each warp works on a different column of the tile. + int const h = thread_offset[1] + block[1]; + // Each lane writes a different element. + int const w = thread_offset[2] + block[2]; + // Setup the pointer. + this->params.pointer += ((h * params.stride_h + w) + offset); + + // Prepare the vector of predicates. + for (int i = 0; i < Base::Iterations::kW; ++i) { + predicates.set(i, w + i * Base::Delta::kW < bounds[2]); + } + this->params.predicate_offset -= (h + pred_offset); + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm); + } + + /// Increment the pointer in the C dimension. + CUTLASS_DEVICE void inc_c() {} + /// Increment the pointer in the W dimension. + CUTLASS_DEVICE void inc_w() {} + /// Increment the pointer in the H dimension. + CUTLASS_DEVICE void inc_h() { + params.pointer += params.inc_h; + params.predicate_offset -= params.predicate_inc_h; + } + /// Increment the pointer in the D dimension. + CUTLASS_DEVICE void inc_d() {} + /// Increment the pointer to move to the next iteration. + CUTLASS_DEVICE void inc_advance() { + params.pointer += params.inc_advance; + this->params.predicate_offset -= params.predicate_inc_advance; + } + + /// The accessor. + CUTLASS_DEVICE void set(typename Base::AccessType const& value, int d, int h, int w, int c) { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store( + value, params.pointer, imm); + } + + /// Test the validity of the iterator. + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { + return predicates.at(w) && params.predicate_offset > 0; + } + + /// The predicates for the row. + cutlass::PredicateVector<Base::Iterations::kW> predicates; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass |