diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h | 211 |
1 files changed, 211 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h b/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h new file mode 100644 index 0000000..dbd57f6 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h @@ -0,0 +1,211 @@ +/*************************************************************************************************** + * 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 tile iterator traits for loading thread block-level tile from global memory. +*/ +#pragma once + +#include <cutlass/gemm/gemm_global_tile.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kAccessSize_> +struct WmmaGemmGlobalIteratorCdTraits : public GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> + Base; + + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta; + + /// 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 / Base::Threads::kW; + int thread_offset_w = threadIdx.x % Base::Threads::kW * Base::ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct WmmaGemmGlobalIteratorCd : public TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef WmmaGemmGlobalIteratorCd<TileTraits_, Index_> This_; + /// The traits. + typedef TileTraits_ Traits; + /// The base class. + typedef TileIteratorBase<Traits, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> + Base; + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> ImmediateOffsetStrides; + /// 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 functor. + 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_h, inc_advance; + /// The column offset to compute the predicate for the columns. + Index predicate_offset; + /// The strides to increment the predicate offset. + Index predicate_inc_h, predicate_inc_advance; + + /// Setup the params. + CUTLASS_HOST_DEVICE int initialize( + Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w) { + // The pointer. + this->pointer = pointer; + // Setup the base stride. One "group of threads" per column. + stride_h = ld; + // Each thread output 1 column per iteration. . + inc_h = ld * TileTraits_::Threads::kH; + inc_advance = inc_h + epilogue_stride_w; + + predicate_offset = n; + predicate_inc_h = TileTraits_::Threads::kH; + predicate_inc_advance = predicate_inc_h + epilogue_delta_w; + + // It worked. + return 0; + } + }; + + Params params; + + Coord<4> thread_offset; + + /// Ctor. + CUTLASS_DEVICE WmmaGemmGlobalIteratorCd() {} + + /// Ctor. + CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const& params, + const Coord<3>& bounds, + const Coord<3>& block, + int const pointer_offset = 0, + int const 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) + pointer_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; + 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(d, h, w, 0); + Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store( + value, params.pointer, imm); + } + + /// Test the predicate. + 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 |