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, 0 insertions, 541 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm_global_tile.h b/cutlass-example/cutlass/gemm/gemm_global_tile.h deleted file mode 100644 index 1cc3b33..0000000 --- a/cutlass-example/cutlass/gemm/gemm_global_tile.h +++ /dev/null @@ -1,541 +0,0 @@ -/*************************************************************************************************** - * 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 |