gemm_epilogue_traits.h
Go to the documentation of this file.
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include <cutlass/convert.h>
#include <cutlass/coord.h>
#include <cutlass/gemm/gemm_global_stream.h>
#include <cutlass/gemm/gemm_shared_stream.h>
#include <cutlass/gemm/linear_scaling.h>
#include <cutlass/reshape_tile.h>
#include <cutlass/tile_iterator.h>
namespace cutlass {
namespace gemm {
template <
typename OutputTile_,
typename Accumulators_,
typename GlobalLoadIteratorC_,
typename GlobalTransformerC_,
typename GlobalTransformerD_,
typename GlobalStoreIteratorD_,
typename SharedStoreIteratorD_,
typename SharedStoreTransformerD_,
typename SharedLoadIteratorD_,
typename Iterations_,
typename Delta_,
typename Functor_,
typename Index_ = int>
struct GemmEpilogueTraits {
//
typedef OutputTile_ OutputTile;
typedef Accumulators_ Accumulators;
typedef GlobalLoadIteratorC_ GlobalLoadIteratorC;
typedef GlobalTransformerC_ GlobalTransformerC;
typedef GlobalTransformerD_ GlobalTransformerD;
typedef GlobalStoreIteratorD_ GlobalStoreIteratorD;
typedef SharedStoreIteratorD_ SharedStoreIteratorD;
typedef SharedStoreTransformerD_ SharedStoreTransformerD;
typedef SharedLoadIteratorD_ SharedLoadIteratorD;
typedef Iterations_ Iterations;
typedef Delta_ Delta;
typedef Functor_ Functor;
typedef Index_ Index;
static_assert(Iterations::kD == 1 && Iterations::kC == 1, "Unsupported 3D/4D shapes");
typedef typename Functor::Scalar Scalar;
typedef typename GlobalLoadIteratorC::Scalar ScalarC;
typedef typename GlobalStoreIteratorD::Scalar ScalarD;
struct Params {
Index stride_h, stride_w;
typename GlobalLoadIteratorC::Params iterator_c;
typename GlobalStoreIteratorD::Params iterator_d;
typename SharedStoreIteratorD::Params shared_store_iterator_d;
typename SharedLoadIteratorD::Params shared_load_iterator_d;
typename Functor::Params functor;
template <typename GemmDesc_>
CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc) {
// The parameters for the functor.
int error_code = functor.initialize(desc);
if (error_code) {
return error_code;
}
// At the end of the H iteration, we jump over a number of columns.
this->stride_h = desc.ldd * Delta::kH;
// Nothing to do here.
this->stride_w = 0;
// Setup the params for the global memory iterator for C.
error_code = iterator_c.initialize(
reinterpret_cast<ScalarC const*>(desc.d_c), desc.ldc, desc.n, stride_w, Delta::kW);
if (error_code) {
return error_code;
}
// Setup the params for the global memory iterator for D.
return iterator_d.initialize(
reinterpret_cast<ScalarD*>(desc.d_d), desc.ldd, desc.n, stride_w, Delta::kW);
}
};
union StreamSharedStorage {
// The storage for the store iterator.
typename SharedStoreIteratorD::SharedStorage store;
// The storage for the store iterator.
typename SharedLoadIteratorD::SharedStorage load;
};
struct SharedStorage {
// The storage for the shared stream D.
StreamSharedStorage shared_stream;
};
};
template <typename GemmConfig_, typename EpilogueFunctor_, typename Index_ = int>
struct GemmEpilogueTraitsHelper {
typedef typename EpilogueFunctor_::Scalar Scalar;
typedef typename GemmConfig_::OutputTile OutputTile;
typedef Shape<1,
GemmConfig_::MultiplyAdd::AccumulatorsPerThread::kH /
GemmConfig_::kAccumulatorsPerLdsB,
GemmConfig_::kAccumulatorsPerLdsB>
Iterations;
// The iteration strides in the H/W dimension.
typedef Shape<0,
GemmConfig_::kAccumulatorsPerLdsB*(
GemmConfig_::Warps::kH* GemmConfig_::MultiplyAdd::ThreadsPerWarp::kH - 1),
0>
Delta;
typedef EpilogueFunctor_ Functor;
typedef GemmSharedStoreTileDTraits<
// The pointer is float.
typename Functor::Scalar,
// The output tile size.
typename GemmConfig_::OutputTile,
// The number of warps.
typename GemmConfig_::Warps,
// The number of threads per warp.
typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
// The number of scalars per STS.
GemmConfig_::kScalarsPerStsD,
// The skew -- 128 / sizeof(ScalarD) / kScalarsPerStsD is the number of threads involved in
// a single STS. We divide by 2 as our objective is to add a skew to the odd threads to
// avoid bank conflicts between odd and even threads.
128 / sizeof(typename GemmConfig_::ScalarD) / GemmConfig_::kScalarsPerStsD / 2 *
GemmConfig_::kScalarsPerStsD>
SharedStoreTileTraits;
typedef TileStoreIterator<SharedStoreTileTraits,
typename SharedStoreTileTraits::Scalar,
IteratorAdvance::kH,
MemorySpace::kShared>
SharedStoreIteratorD;
typedef Copy<typename SharedStoreIteratorD::Fragment> SharedStoreTransformerD;
typedef GemmSharedLoadTileDTraits<
// The pointer is float.
typename Functor::Scalar,
// The output tile size.
typename GemmConfig_::OutputTile,
// The number of warps.
typename GemmConfig_::Warps,
// The number of threads per warp.
typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
// The number of columns of the output tile written by iteration.
GemmConfig_::OutputTile::kH / ShapeCount<Iterations>::kCount,
// The number of scalars per LDS.
GemmConfig_::kScalarsPerLdsD,
// The skew.
SharedStoreTileTraits::kSkew>
SharedLoadTileTraits;
typedef TileLoadIterator<SharedLoadTileTraits,
typename SharedLoadTileTraits::Scalar,
IteratorAdvance::kH,
MemorySpace::kShared>
SharedLoadIteratorD;
typedef GemmGlobalTileCdTraits<
// The pointer is float const.
typename GemmConfig_::ScalarC const,
// The tile has size (N / Iterations)xM in GEMM's terminology.
Shape<1,
GemmConfig_::OutputTile::kH / ShapeCount<Iterations>::kCount,
GemmConfig_::OutputTile::kW>,
// The threads are distributed as warps x 32 (the traits may reorganize).
Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>,
// How many elements do we jump over at each iteration?
Iterations::kW,
// The number of scalars per LDG (LDG.32 or LDG.128, etc).
GemmConfig_::kScalarsPerLdgC>
GlobalLoadTileTraits;
typedef GemmGlobalIteratorCd<GlobalLoadTileTraits, Index_> GlobalLoadIteratorC;
typedef Copy<typename GlobalLoadIteratorC::Fragment> GlobalTransformerC;
typedef GemmGlobalTileCdTraits<
// The pointer is float.
typename GemmConfig_::ScalarD,
// The tile has size (N / Iterations)xM in GEMM's terminology.
Shape<1,
GemmConfig_::OutputTile::kH / ShapeCount<Iterations>::kCount,
GemmConfig_::OutputTile::kW>,
// The threads are distributed as warps x 32 (the traits may reorganize).
Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>,
// How many elements do we jump over at each iteration?
Iterations::kW,
// The number of scalars per LDG (LDG.32 or LDG.128, etc).
GemmConfig_::kScalarsPerStgD>
GlobalStoreTileTraits;
typedef GemmGlobalIteratorCd<GlobalStoreTileTraits, Index_> GlobalStoreIteratorD;
typedef Copy<typename GlobalStoreIteratorD::Fragment> GlobalTransformerD;
};
template <
typename GemmConfig_,
typename EpilogueFunctor_,
typename Index_ = int,
typename Helper_ = GemmEpilogueTraitsHelper<GemmConfig_, EpilogueFunctor_, Index_> >
struct SimplifiedGemmEpilogueTraits : public GemmEpilogueTraits<
// The output tile.
typename GemmConfig_::OutputTile,
// The accumulators.
typename GemmConfig_::Accumulators,
// The global iterator for C.
typename Helper_::GlobalLoadIteratorC,
// The transformer for C.
typename Helper_::GlobalTransformerC,
// The transformer for D.
typename Helper_::GlobalTransformerD,
// The global iterator for D.
typename Helper_::GlobalStoreIteratorD,
// The iterator to store D to shared memory.
typename Helper_::SharedStoreIteratorD,
// The shared store transformer for D.
typename Helper_::SharedStoreTransformerD,
// The iterator to load D from shared memory.
typename Helper_::SharedLoadIteratorD,
// The number of iterations.
typename Helper_::Iterations,
// The strides between iterations.
typename Helper_::Delta,
// The functor to be used in the epilogue.
EpilogueFunctor_,
// The index.
Index_> {};
} // namespace gemm
} // namespace cutlass
cutlass::gemm::GemmGlobalTileCdTraits
Definition: gemm_global_tile.h:116
cutlass::gemm::GemmEpilogueTraits::StreamSharedStorage::load
SharedLoadIteratorD::SharedStorage load
Definition: gemm_epilogue_traits.h:158
cutlass::gemm::GemmEpilogueTraits::Delta
Delta_ Delta
The iterations strides.
Definition: gemm_epilogue_traits.h:94
cutlass::MemorySpace::kShared
Definition: load_store.h:42
cutlass
Definition: convert.h:33
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalLoadTileTraits
GemmGlobalTileCdTraits< typename GemmConfig_::ScalarC const, Shape< 1, GemmConfig_::OutputTile::kH/ShapeCount< Iterations >::kCount, GemmConfig_::OutputTile::kW >, Shape< 1, ShapeCount< typename GemmConfig_::Warps >::kCount, GemmConfig_::kWarpSize >, Iterations::kW, GemmConfig_::kScalarsPerLdgC > GlobalLoadTileTraits
The traits class to build the iterator to load data from global memory for C^N.
Definition: gemm_epilogue_traits.h:260
cutlass::gemm::GemmEpilogueTraitsHelper
Definition: gemm_epilogue_traits.h:171
tile_iterator.h
Defines the Tile Traits concept and iterators for loading and storing to tiles efficiently.
cutlass::gemm::GemmEpilogueTraits::GlobalTransformerC
GlobalTransformerC_ GlobalTransformerC
The transformer for C.
Definition: gemm_epilogue_traits.h:80
cutlass::gemm::GemmEpilogueTraits::Params::iterator_d
GlobalStoreIteratorD::Params iterator_d
The params for the D global iterator.
Definition: gemm_epilogue_traits.h:118
linear_scaling.h
Implements the BLAS linear scaling function alpha*AB + beta*C.
cutlass::gemm::GemmEpilogueTraits::StreamSharedStorage
The shared memory storage to exchange data.
Definition: gemm_epilogue_traits.h:154
cutlass::gemm::GemmEpilogueTraits::ScalarC
GlobalLoadIteratorC::Scalar ScalarC
The scalar for C.
Definition: gemm_epilogue_traits.h:107
cutlass::gemm::GemmEpilogueTraitsHelper::Scalar
EpilogueFunctor_::Scalar Scalar
The scalar.
Definition: gemm_epilogue_traits.h:173
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::gemm::GemmEpilogueTraits::Params::stride_h
Index stride_h
The strides for H and W in the different iterations of the epilogue.
Definition: gemm_epilogue_traits.h:114
cutlass::gemm::GemmEpilogueTraits::Index
Index_ Index
The index.
Definition: gemm_epilogue_traits.h:99
cutlass::gemm::GemmEpilogueTraits::GlobalStoreIteratorD
GlobalStoreIteratorD_ GlobalStoreIteratorD
The iterator for D in global memory.
Definition: gemm_epilogue_traits.h:84
cutlass::Copy
Definition: convert.h:69
cutlass::gemm::GemmEpilogueTraits::OutputTile
OutputTile_ OutputTile
The output tile.
Definition: gemm_epilogue_traits.h:73
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalLoadIteratorC
GemmGlobalIteratorCd< GlobalLoadTileTraits, Index_ > GlobalLoadIteratorC
The iterator to load C.
Definition: gemm_epilogue_traits.h:263
cutlass::IteratorAdvance::kH
Definition: tile_iterator.h:62
cutlass::gemm::GemmEpilogueTraits::ScalarD
GlobalStoreIteratorD::Scalar ScalarD
The scalar for D.
Definition: gemm_epilogue_traits.h:109
cutlass::gemm::GemmEpilogueTraitsHelper::SharedStoreIteratorD
TileStoreIterator< SharedStoreTileTraits, typename SharedStoreTileTraits::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedStoreIteratorD
The iterator to store D to shared memory.
Definition: gemm_epilogue_traits.h:216
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalStoreIteratorD
GemmGlobalIteratorCd< GlobalStoreTileTraits, Index_ > GlobalStoreIteratorD
The iterator to store D.
Definition: gemm_epilogue_traits.h:284
cutlass::gemm::GemmEpilogueTraitsHelper::SharedStoreTransformerD
Copy< typename SharedStoreIteratorD::Fragment > SharedStoreTransformerD
The shared store transformer for D.
Definition: gemm_epilogue_traits.h:219
cutlass::gemm::GemmEpilogueTraitsHelper::Iterations
Shape< 1, GemmConfig_::MultiplyAdd::AccumulatorsPerThread::kH/GemmConfig_::kAccumulatorsPerLdsB, GemmConfig_::kAccumulatorsPerLdsB > Iterations
The number of iterations in the epilogue.
Definition: gemm_epilogue_traits.h:182
cutlass::gemm::GemmEpilogueTraits::Params::iterator_c
GlobalLoadIteratorC::Params iterator_c
The params for the C iterator.
Definition: gemm_epilogue_traits.h:116
cutlass::gemm::GemmEpilogueTraits::SharedStoreTransformerD
SharedStoreTransformerD_ SharedStoreTransformerD
The shared store transformer for D.
Definition: gemm_epilogue_traits.h:88
cutlass::gemm::GemmEpilogueTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const &desc)
Setup the params.
Definition: gemm_epilogue_traits.h:128
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalStoreTileTraits
GemmGlobalTileCdTraits< typename GemmConfig_::ScalarD, Shape< 1, GemmConfig_::OutputTile::kH/ShapeCount< Iterations >::kCount, GemmConfig_::OutputTile::kW >, Shape< 1, ShapeCount< typename GemmConfig_::Warps >::kCount, GemmConfig_::kWarpSize >, Iterations::kW, GemmConfig_::kScalarsPerStgD > GlobalStoreTileTraits
The traits class to build the iterator to store data to global memory for D^N.
Definition: gemm_epilogue_traits.h:281
cutlass::TileLoadIterator
An iterator implementing Tile Load Iterator Concept for loading a tile from memory.
Definition: tile_iterator.h:302
cutlass::gemm::GemmEpilogueTraits::SharedStoreIteratorD
SharedStoreIteratorD_ SharedStoreIteratorD
The iterator to store D in shared memory.
Definition: gemm_epilogue_traits.h:86
cutlass::gemm::GemmEpilogueTraits::StreamSharedStorage::store
SharedStoreIteratorD::SharedStorage store
Definition: gemm_epilogue_traits.h:156
cutlass::gemm::GemmEpilogueTraitsHelper::SharedStoreTileTraits
GemmSharedStoreTileDTraits< typename Functor::Scalar, typename GemmConfig_::OutputTile, typename GemmConfig_::Warps, typename GemmConfig_::MultiplyAdd::ThreadsPerWarp, GemmConfig_::kScalarsPerStsD, 128/sizeof(typename GemmConfig_::ScalarD)/GemmConfig_::kScalarsPerStsD/2 *GemmConfig_::kScalarsPerStsD > SharedStoreTileTraits
The traits class to build the iterator to store to shared memory for D.
Definition: gemm_epilogue_traits.h:209
cutlass::gemm::GemmSharedStoreTileDTraits::kSkew
static int const kSkew
The skew.
Definition: gemm_shared_tile.h:278
reshape_tile.h
Defines a type for restructuring a tile.
cutlass::gemm::GemmEpilogueTraits::Iterations
Iterations_ Iterations
typedef typename GemmConfig::EpilogueIterations Iterations;
Definition: gemm_epilogue_traits.h:92
cutlass::gemm::GemmSharedLoadTileDTraits
Definition: gemm_shared_tile.h:335
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
static_assert
#define static_assert(__e, __m)
Definition: platform.h:145
cutlass::Shape
A Shape implementing Layout Concept describing the dimensions of a cube.
Definition: shape.h:64
cutlass::gemm::SimplifiedGemmEpilogueTraits
Definition: gemm_epilogue_traits.h:300
cutlass::gemm::GemmEpilogueTraits::SharedStorage::shared_stream
StreamSharedStorage shared_stream
Definition: gemm_epilogue_traits.h:164
cutlass::gemm::GemmEpilogueTraitsHelper::Functor
EpilogueFunctor_ Functor
The functor to do the math in the epilogue.
Definition: gemm_epilogue_traits.h:190
cutlass::gemm::GemmEpilogueTraitsHelper::SharedLoadIteratorD
TileLoadIterator< SharedLoadTileTraits, typename SharedLoadTileTraits::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedLoadIteratorD
The iterator to load D from shared memory.
Definition: gemm_epilogue_traits.h:244
cutlass::gemm::GemmEpilogueTraitsHelper::OutputTile
GemmConfig_::OutputTile OutputTile
The output tile.
Definition: gemm_epilogue_traits.h:175
cutlass::gemm::GemmEpilogueTraits::GlobalLoadIteratorC
GlobalLoadIteratorC_ GlobalLoadIteratorC
The iterator for C in global memory.
Definition: gemm_epilogue_traits.h:78
cutlass::gemm::GemmEpilogueTraits::GlobalTransformerD
GlobalTransformerD_ GlobalTransformerD
The transformer for D.
Definition: gemm_epilogue_traits.h:82
cutlass::gemm::GemmEpilogueTraits
Definition: gemm_epilogue_traits.h:70
cutlass::gemm::GemmGlobalIteratorCd
Definition: gemm_global_tile.h:348
cutlass::gemm::GemmEpilogueTraits::Params::stride_w
Index stride_w
Definition: gemm_epilogue_traits.h:114
cutlass::Shape::kW
static int const kW
The width of the cube.
Definition: shape.h:70
cutlass::gemm::GemmEpilogueTraits::Scalar
Functor::Scalar Scalar
We do not support 3D or 4D shapes.
Definition: gemm_epilogue_traits.h:102
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalTransformerD
Copy< typename GlobalStoreIteratorD::Fragment > GlobalTransformerD
The transformer for D.
Definition: gemm_epilogue_traits.h:286
gemm_global_stream.h
Implements efficient loading of the thread block-level tile from global memory and storing to shared ...
cutlass::gemm::GemmEpilogueTraits::Params
The params.
Definition: gemm_epilogue_traits.h:112
cutlass::gemm::GemmEpilogueTraits::SharedStorage
The shared memory to swizzle the data in the epilogue.
Definition: gemm_epilogue_traits.h:162
cutlass::gemm::GemmEpilogueTraits::SharedLoadIteratorD
SharedLoadIteratorD_ SharedLoadIteratorD
The iterator to store D in shared memory.
Definition: gemm_epilogue_traits.h:90
cutlass::gemm::GemmEpilogueTraits::Params::functor
Functor::Params functor
The functor params.
Definition: gemm_epilogue_traits.h:124
cutlass::gemm::GemmEpilogueTraitsHelper::GlobalTransformerC
Copy< typename GlobalLoadIteratorC::Fragment > GlobalTransformerC
The transformer for C.
Definition: gemm_epilogue_traits.h:265
cutlass::gemm::GemmEpilogueTraits::Params::shared_load_iterator_d
SharedLoadIteratorD::Params shared_load_iterator_d
The params for the D shared load iterator.
Definition: gemm_epilogue_traits.h:122
cutlass::gemm::GemmEpilogueTraitsHelper::SharedLoadTileTraits
GemmSharedLoadTileDTraits< typename Functor::Scalar, typename GemmConfig_::OutputTile, typename GemmConfig_::Warps, typename GemmConfig_::MultiplyAdd::ThreadsPerWarp, GemmConfig_::OutputTile::kH/ShapeCount< Iterations >::kCount, GemmConfig_::kScalarsPerLdsD, SharedStoreTileTraits::kSkew > SharedLoadTileTraits
The traits class to build the iterator to load from shared memory for D.
Definition: gemm_epilogue_traits.h:237
cutlass::gemm::GemmEpilogueTraits::Accumulators
Accumulators_ Accumulators
Definition: gemm_epilogue_traits.h:76
cutlass::gemm::GemmSharedStoreTileDTraits::Scalar
platform::remove_const< Scalar_ >::type Scalar
The scalar.
Definition: gemm_shared_tile.h:266
gemm_shared_stream.h
Defines abstractions for managing loading and storing fragments to shared memory in the efficient GEM...
cutlass::ShapeCount
Compute derived counted of a Layout Concept based class.
Definition: shape.h:79
convert.h
Defines conversion operations among Fragments of different base type.
cutlass::gemm::GemmSharedLoadTileDTraits::Scalar
platform::remove_const< Scalar_ >::type Scalar
The scalar.
Definition: gemm_shared_tile.h:337
cutlass::gemm::GemmEpilogueTraits::Params::shared_store_iterator_d
SharedStoreIteratorD::Params shared_store_iterator_d
The params for the D shared store iterator.
Definition: gemm_epilogue_traits.h:120
cutlass::gemm::GemmEpilogueTraits::Functor
Functor_ Functor
The functor in charge of the math.
Definition: gemm_epilogue_traits.h:97
cutlass::gemm::GemmSharedStoreTileDTraits
Definition: gemm_shared_tile.h:264
cutlass::TileStoreIterator
An iterator implementing Tile Store Iterator Concept for storing a tile to memory.
Definition: tile_iterator.h:620