diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/gemm_epilogue.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/gemm_epilogue.h | 231 |
1 files changed, 231 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm_epilogue.h b/cutlass-example/cutlass/gemm/gemm_epilogue.h new file mode 100644 index 0000000..bc25307 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_epilogue.h @@ -0,0 +1,231 @@ +/*************************************************************************************************** + * 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 Implements the epilogue phase of the GEMM kernel that efficiently updates global memory + with + the computed matrix product. +*/ +#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 { + /// The traits class. + typedef GemmEpilogueTraits_ Traits; + /// The params. + typedef typename Traits::Params Params; + /// The shared storage. + typedef typename Traits::SharedStorage SharedStorage; + + /// The output tile. + typedef typename Traits::OutputTile OutputTile; + /// The number of iterations. + typedef typename Traits::Iterations Iterations; + /// The accumulators. + typedef typename Traits::Accumulators Accumulators; + /// The scalar. + typedef typename Traits::Scalar Scalar; + /// The functor in charge of the math. + typedef typename Traits::Functor Functor; + + /// We do not support 3D or 4D shapes. + static_assert(Iterations::kD == 1 && Iterations::kC == 1, "Unsupported 3D/4D shapes"); + + /// The iterator for C in global memory. + typedef typename Traits::GlobalLoadIteratorC GlobalLoadIteratorC; + /// The transformer for C. + typedef typename Traits::GlobalTransformerC GlobalTransformerC; + /// The transformer for D. + typedef typename Traits::GlobalTransformerD GlobalTransformerD; + /// The iterator for D in global memory. + typedef typename Traits::GlobalStoreIteratorD GlobalStoreIteratorD; + /// The iterator to store D in shared memory. + typedef typename Traits::SharedStoreIteratorD SharedStoreIteratorD; + /// The shared store transformer for D. + typedef typename Traits::SharedStoreTransformerD SharedStoreTransformerD; + /// The iterator to load D in shared memory. + typedef typename Traits::SharedLoadIteratorD SharedLoadIteratorD; + /// The shared load transformer for D. + typedef Copy<typename SharedLoadIteratorD::Fragment> SharedLoadTransformerD; + + /// The index. + typedef typename Traits::Index Index; + + /// The scalar for C. + typedef typename GlobalLoadIteratorC::Scalar ScalarC; + /// The scalar for D. + typedef typename GlobalStoreIteratorD::Scalar ScalarD; + + /// Ctor. + CUTLASS_DEVICE GemmEpilogue(Params const& params_, + SharedStorage& shared_storage_, + Index m_, + Index n_) + : params(params_), shared_storage(shared_storage_), m(m_), n(n_) {} + + /// Execute the epilogue. + 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) { + + // The problem size. + 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); + + // The transformer to transform before storing to shared memory. + SharedStoreTransformerD shared_store_transformer; + typename SharedStoreTransformerD::OutputFragment shared_store_transformed_d; + + // The iterator to store to shared memory. + SharedStoreIteratorD shared_store_iterator(params.shared_store_iterator_d, + shared_storage.shared_stream.store); + + // The iterator to load from shared memory. TODO: Use a stream. + SharedLoadIteratorD shared_load_iterator(params.shared_load_iterator_d, + shared_storage.shared_stream.load); + + 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; + + shared_store_transformer.transform(accumulators, offset, shared_store_transformed_d); + 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. + 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); + } + } + } + + /// The memory fence for shared loads. + CUTLASS_DEVICE void shared_load_fence() { __syncthreads(); } + + /// The memory fence for shared stores. + CUTLASS_DEVICE void shared_store_fence() { __syncthreads(); } + + /// The params. + Params const& params; + /// The shared storage. + SharedStorage& shared_storage; + /// The dimensions of the GEMM. + Index m, n; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass |