gemm_epilogue.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/fragment.h>
namespace cutlass {
namespace gemm {
template <typename T>
CUTLASS_DEVICE bool is_zero(T x) {
return x == T(0);
}
#if !defined(__CUDACC_RTC__) || defined(CUTLASS_NVRTC_HAS_FP16)
CUTLASS_DEVICE bool is_zero(half x) { return reinterpret_cast<int16_t&>(x) == int16_t(0); }
#endif
template <typename GemmEpilogueTraits_>
struct GemmEpilogue {
typedef GemmEpilogueTraits_ Traits;
typedef typename Traits::Params Params;
typedef typename Traits::SharedStorage SharedStorage;
typedef typename Traits::OutputTile OutputTile;
typedef typename Traits::Iterations Iterations;
typedef typename Traits::Accumulators Accumulators;
typedef typename Traits::Scalar Scalar;
typedef typename Traits::Functor Functor;
static_assert(Iterations::kD == 1 && Iterations::kC == 1, "Unsupported 3D/4D shapes");
typedef typename Traits::GlobalLoadIteratorC GlobalLoadIteratorC;
typedef typename Traits::GlobalTransformerC GlobalTransformerC;
typedef typename Traits::GlobalTransformerD GlobalTransformerD;
typedef typename Traits::GlobalStoreIteratorD GlobalStoreIteratorD;
typedef typename Traits::SharedStoreIteratorD SharedStoreIteratorD;
typedef typename Traits::SharedStoreTransformerD SharedStoreTransformerD;
typedef typename Traits::SharedLoadIteratorD SharedLoadIteratorD;
typedef Copy<typename SharedLoadIteratorD::Fragment> SharedLoadTransformerD;
typedef typename Traits::Index Index;
typedef typename GlobalLoadIteratorC::Scalar ScalarC;
typedef typename GlobalStoreIteratorD::Scalar ScalarD;
CUTLASS_DEVICE GemmEpilogue(Params const& params_,
SharedStorage& shared_storage_,
Index m_,
Index n_)
: params(params_), shared_storage(shared_storage_), m(m_), n(n_) {}
CUTLASS_DEVICE void epilogue(Coord<3> const& block, Accumulators& accumulators) {
if (is_zero(params.functor.beta)) {
epilogue_with_or_without_beta<true>(block, accumulators);
} else {
epilogue_with_or_without_beta<false>(block, accumulators);
}
}
template <bool kBetaIsZero_>
CUTLASS_DEVICE void epilogue_with_or_without_beta(Coord<3> const& block,
Accumulators& accumulators) {
Coord<3> const bounds = cutlass::make_Coord(0, n, m);
// The functor.
Functor functor(params.functor);
// The C fragment.
typename GlobalLoadIteratorC::Fragment fragment_c;
// The transformed C fragment.
typename GlobalTransformerC::OutputFragment transformed_c;
CUTLASS_PRAGMA_UNROLL
for (int h = 0; h < Iterations::kH; ++h) {
// Compute pointer and predicate offsets for C and D global iterators.
int const pointer_offset =
((params.iterator_d.inc_h * (GlobalStoreIteratorD::Iterations::kH - 1) +
params.iterator_d.inc_advance) *
Iterations::kW +
params.stride_h) *
h;
int const predicate_offset =
((params.iterator_d.predicate_inc_h * (GlobalStoreIteratorD::Iterations::kH - 1) +
params.iterator_d.predicate_inc_advance) *
Iterations::kW +
Traits::Delta::kH) *
h;
// The iterator to load the elements of the C matrix.
GlobalLoadIteratorC global_load_iterator(
params.iterator_c, bounds, block, pointer_offset, predicate_offset);
// The transformer for C.
GlobalTransformerC transformer_c;
// The transformer for D.
GlobalTransformerD transformer_d;
// The iterator to store into the D matrix.
GlobalStoreIteratorD global_store_iterator(
params.iterator_d, bounds, block, pointer_offset, predicate_offset);
CUTLASS_PRAGMA_UNROLL
for (int w = 0; w < Iterations::kW; ++w) {
// Load the C matrix into fragment.
if (!kBetaIsZero_) {
iterator_load(global_load_iterator, fragment_c);
}
// Make sure we can write to shared memory.
shared_load_fence();
// Copy the accumulators to shared memory.
int const offset = (h * Iterations::kW + w) * SharedStoreIteratorD::Fragment::kElements;
SharedStoreTransformerD shared_store_transformer;
typename SharedStoreTransformerD::OutputFragment shared_store_transformed_d;
shared_store_transformer.transform(accumulators, offset, shared_store_transformed_d);
SharedStoreIteratorD shared_store_iterator(params.shared_store_iterator_d,
shared_storage.shared_stream.store);
shared_iterator_store(shared_store_iterator, shared_store_transformed_d);
// Make sure the data is in shared memory.
shared_store_fence();
// Copy the accumulators back to registers from shared memory.
SharedLoadIteratorD shared_load_iterator(params.shared_load_iterator_d,
shared_storage.shared_stream.load);
typename SharedLoadIteratorD::Fragment fetched_d;
shared_iterator_load(shared_load_iterator, fetched_d);
// Do the math.
typename GlobalTransformerD::InputFragment fragment_d;
if (kBetaIsZero_) {
functor.evaluate(fetched_d, fragment_d);
} else {
// Transform C fragment.
transformer_c.transform(fragment_c, transformed_c);
// Do the math.
functor.evaluate(fetched_d, transformed_c, fragment_d);
}
// Transform D fragment.
typename GlobalTransformerD::OutputFragment transformed_d;
transformer_d.transform(fragment_d, transformed_d);
// Copy the results to global memory.
iterator_store(global_store_iterator, transformed_d);
}
}
}
CUTLASS_DEVICE void shared_load_fence() { __syncthreads(); }
CUTLASS_DEVICE void shared_store_fence() { __syncthreads(); }
Params const& params;
SharedStorage& shared_storage;
Index m, n;
};
} // namespace gemm
} // namespace cutlass
cutlass::gemm::GemmEpilogue::ScalarD
GlobalStoreIteratorD::Scalar ScalarD
The scalar for D.
Definition: gemm_epilogue.h:98
cutlass::gemm::GemmEpilogue::SharedStoreIteratorD
Traits::SharedStoreIteratorD SharedStoreIteratorD
The iterator to store D in shared memory.
Definition: gemm_epilogue.h:84
cutlass
Definition: convert.h:33
cutlass::shared_iterator_load
CUTLASS_DEVICE void shared_iterator_load(InputIterator &iterator, Fragment &fragment)
Loads a fragment from a shared memory input iterator.
Definition: iterator_access.h:75
cutlass::gemm::GemmEpilogue::Params
Traits::Params Params
The params.
Definition: gemm_epilogue.h:57
cutlass::gemm::GemmEpilogue
Definition: gemm_epilogue.h:53
cutlass::gemm::GemmEpilogue::epilogue_with_or_without_beta
CUTLASS_DEVICE void epilogue_with_or_without_beta(Coord< 3 > const &block, Accumulators &accumulators)
Definition: gemm_epilogue.h:117
cutlass::gemm::GemmEpilogue::GemmEpilogue
CUTLASS_DEVICE GemmEpilogue(Params const &params_, SharedStorage &shared_storage_, Index m_, Index n_)
Ctor.
Definition: gemm_epilogue.h:101
coord.h
A Coord is a coordinate of arbitrary rank into a tensor or matrix.
cutlass::make_Coord
CUTLASS_HOST_DEVICE Coord< 1 > make_Coord(int _0)
Helper to make a 2-element coordinate.
Definition: coord.h:241
cutlass::Copy
Definition: convert.h:69
cutlass::gemm::GemmEpilogue::SharedStorage
Traits::SharedStorage SharedStorage
The shared storage.
Definition: gemm_epilogue.h:59
cutlass::gemm::GemmEpilogue::GlobalTransformerD
Traits::GlobalTransformerD GlobalTransformerD
The transformer for D.
Definition: gemm_epilogue.h:80
cutlass::gemm::GemmEpilogue::OutputTile
Traits::OutputTile OutputTile
The output tile.
Definition: gemm_epilogue.h:62
cutlass::gemm::GemmEpilogue::Accumulators
Traits::Accumulators Accumulators
The accumulators.
Definition: gemm_epilogue.h:66
CUTLASS_PRAGMA_UNROLL
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:60
cutlass::gemm::GemmEpilogue::shared_load_fence
CUTLASS_DEVICE void shared_load_fence()
The memory fence for shared loads.
Definition: gemm_epilogue.h:209
cutlass::gemm::GemmEpilogue::shared_storage
SharedStorage & shared_storage
The shared storage.
Definition: gemm_epilogue.h:217
cutlass::gemm::GemmEpilogue::Traits
GemmEpilogueTraits_ Traits
The traits class.
Definition: gemm_epilogue.h:55
cutlass::gemm::is_zero
CUTLASS_DEVICE bool is_zero(T x)
Definition: gemm_epilogue.h:42
cutlass::gemm::GemmEpilogue::params
Params const & params
The params.
Definition: gemm_epilogue.h:215
cutlass::gemm::GemmEpilogue::SharedLoadIteratorD
Traits::SharedLoadIteratorD SharedLoadIteratorD
The iterator to load D in shared memory.
Definition: gemm_epilogue.h:88
cutlass::gemm::GemmEpilogue::Index
Traits::Index Index
The index.
Definition: gemm_epilogue.h:93
static_assert
#define static_assert(__e, __m)
Definition: platform.h:145
cutlass::gemm::GemmEpilogue::SharedStoreTransformerD
Traits::SharedStoreTransformerD SharedStoreTransformerD
The shared store transformer for D.
Definition: gemm_epilogue.h:86
cutlass::shared_iterator_store
CUTLASS_DEVICE void shared_iterator_store(OutputIterator &iterator, Fragment const &fragment)
Stores a fragment to a shared memory output iterator.
Definition: iterator_access.h:228
cutlass::gemm::GemmEpilogue::GlobalStoreIteratorD
Traits::GlobalStoreIteratorD GlobalStoreIteratorD
The iterator for D in global memory.
Definition: gemm_epilogue.h:82
cutlass::Coord
Statically-sized array specifying Coords within a tensor.
Definition: coord.h:48
cutlass::iterator_store
CUTLASS_HOST_DEVICE void iterator_store(OutputIterator &iterator, Fragment &fragment)
Stores a fragment to an output iterator.
Definition: iterator_access.h:193
cutlass::gemm::GemmEpilogue::ScalarC
GlobalLoadIteratorC::Scalar ScalarC
The scalar for C.
Definition: gemm_epilogue.h:96
cutlass::gemm::GemmEpilogue::n
Index n
Definition: gemm_epilogue.h:219
cutlass::iterator_load
CUTLASS_HOST_DEVICE void iterator_load(InputIterator &iterator, Fragment &fragment)
Loads a fragment from an input iterator.
Definition: iterator_access.h:41
cutlass::gemm::GemmEpilogue::Functor
Traits::Functor Functor
The functor in charge of the math.
Definition: gemm_epilogue.h:70
cutlass::gemm::GemmEpilogue::Iterations
Traits::Iterations Iterations
The number of iterations.
Definition: gemm_epilogue.h:64
cutlass::gemm::GemmEpilogue::epilogue
CUTLASS_DEVICE void epilogue(Coord< 3 > const &block, Accumulators &accumulators)
Execute the epilogue.
Definition: gemm_epilogue.h:108
fragment.h
Defines Fragment, a statically-sized array for storing parts of matrices within a thread&#39;s registers...
cutlass::gemm::GemmEpilogue::SharedLoadTransformerD
Copy< typename SharedLoadIteratorD::Fragment > SharedLoadTransformerD
The shared load transformer for D.
Definition: gemm_epilogue.h:90
cutlass::gemm::GemmEpilogue::Scalar
Traits::Scalar Scalar
The scalar.
Definition: gemm_epilogue.h:68
convert.h
Defines conversion operations among Fragments of different base type.
cutlass::gemm::GemmEpilogue::m
Index m
The dimensions of the GEMM.
Definition: gemm_epilogue.h:219
cutlass::gemm::GemmEpilogue::shared_store_fence
CUTLASS_DEVICE void shared_store_fence()
The memory fence for shared stores.
Definition: gemm_epilogue.h:212
cutlass::gemm::GemmEpilogue::GlobalTransformerC
Traits::GlobalTransformerC GlobalTransformerC
The transformer for C.
Definition: gemm_epilogue.h:78
cutlass::gemm::GemmEpilogue::GlobalLoadIteratorC
Traits::GlobalLoadIteratorC GlobalLoadIteratorC
We do not support 3D or 4D shapes.
Definition: gemm_epilogue.h:73