diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h | 240 |
1 files changed, 240 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h b/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h new file mode 100644 index 0000000..7d15b26 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h @@ -0,0 +1,240 @@ +/*************************************************************************************************** + * 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 iterator traits for efficiently loading and storing fragment to and from shared + memory, specialized for WMMA GEMM. +*/ +#pragma once + +#include <cutlass/wmma_matrix.h> +#ifdef CUTLASS_USE_WMMA_API + +#include <cutlass/gemm/gemm_operand.h> +#include <cutlass/reshape_tile.h> + +namespace cutlass { +namespace gemm { + +template <class> +struct Debug {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Warps_, + int kWarpStride_, + typename Iterations_, + typename Delta_, + typename WmmaShape_> +struct WmmaGemmSharedLoadTileATraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kA; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = 1; + /// The tile with skew. + typedef Tile_ Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The warps strides. + static int const kWarpStride = kWarpStride_; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The strides between iterations. + typedef Delta_ Delta; + /// The strides between iterations. + typedef Delta_ ImmediateOffsetStrides; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The offset. + int const offset = warp % Warps::kW * kWarpStride; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Warps_, + int kWarpStride_, + typename Iterations_, + typename Delta_, + typename WmmaShape_> +struct WmmaGemmSharedLoadTileBTraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kB; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = 1; + /// The tile with skew. + typedef Tile_ Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The warps strides. + static int const kWarpStride = kWarpStride_; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The strides between iterations. + typedef Delta_ Delta; + /// The strides between iterations. + typedef Delta_ ImmediateOffsetStrides; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The offset. + int const offset = warp / Warps::kW * kWarpStride; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename OutputTile_, + typename Warps_, + typename WmmaShape_, + int kSkew_ = 0> +struct WmmaGemmSharedStoreTileDTraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kC; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + // The access size + static int const kAccessSize = 1; + /// The pointer. + typedef Scalar* Pointer; + /// The number of warps. + typedef Warps_ Warps; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// The tile with skew. + typedef Shape<1, Warps_::kH * WmmaShape_::kH, OutputTile_::kW + kSkew_> Tile; + /// The number of iterations needed to store the tile. + typedef Shape<1, 1, OutputTile_::kW / Warps::kW / WmmaShape_::kW> Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * WmmaShape_::kW, 0> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * WmmaShape_::kW, 0> ImmediateOffsetStrides; + + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The starting column. + int const h = warp / Warps::kW * WmmaShape::kH; + // The w. + int const w = warp % Warps::kW * WmmaShape::kW; + // The offset. + int const offset = h * Tile::kW + w; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kScalarsPerLds_> +struct WmmaGemmSharedLoadTileDTraits { + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = kScalarsPerLds_; + /// The tile. + typedef typename ReshapeTile<Tile_, kScalarsPerLds_>::Tile Tile; + /// The threads. + typedef typename ReshapeThreads<Tile, Threads_>::Threads Threads; + /// The threads strides. + typedef Shape<1, Tile::kW * Tile::kC, Tile::kC> ThreadsStrides; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kScalarsPerLds_> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kScalarsPerLds_> + ImmediateOffsetStrides; + /// The number of iterations needed to load/store the tile. + typedef Shape<1, Tile::kH / Threads::kH, Tile::kW / Threads::kW, Tile::kC / kScalarsPerLds_> + Iterations; + + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The offset. + int const offset = ComputeThreadOffsetFromStrides<Threads, ThreadsStrides>::get(); + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass + +#endif // defined CUTLASS_USE_WMMA_API |