diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/gemm.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/gemm.h | 344 |
1 files changed, 344 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm.h b/cutlass-example/cutlass/gemm/gemm.h new file mode 100644 index 0000000..c50a3f0 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm.h @@ -0,0 +1,344 @@ +/*************************************************************************************************** + * 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 a software-pipelined efficient GEMM. +*/ +#pragma once + +#if !defined(__CUDACC_RTC__) +#include <cuda.h> +#endif + +#include <cutlass/coord.h> +#include <cutlass/util/platform.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Gemm_> +__global__ /*__launch_bounds__(Gemm_::kThreads)*/ void gemm_kernel(typename Gemm_::Params params) { + // Declare shared memory. + __shared__ typename Gemm_::SharedStorage shared_storage; + + // Construct the GEMM object. + Gemm_ gemm(params, shared_storage); + // Run GEMM. + gemm.multiply_add(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Index_ = int> +struct GemmDesc { + /// The dimensions of the GEMM. + Index_ m, n, k; + /// The alpha/beta scaling values. + Scalar_ alpha, beta; + /// The source matrix A. + void const* d_a; + /// The stride for A. + Index_ lda; + /// The source matrix B. + void const* d_b; + /// The stride for B. + Index_ ldb; + /// The source matrix C. + void const* d_c; + /// The stride for C. + Index_ ldc; + /// The destination matrix D. + void* d_d; + /// The stride for D. + Index_ ldd; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmTraits_> +struct Gemm { + /// This class. + typedef Gemm<GemmTraits_> This_; + /// The traits. + typedef GemmTraits_ Traits; + /// The shared storage. + typedef typename Traits::SharedStorage SharedStorage; + + /// The scalar for A. + typedef typename Traits::ScalarA ScalarA; + /// The scalar for B. + typedef typename Traits::ScalarB ScalarB; + /// The scalar in the epilogue. + typedef typename Traits::Epilogue::Scalar ScalarEpilogue; + /// The scalar for C. + typedef typename Traits::Epilogue::ScalarC ScalarC; + /// The scalar for D. + typedef typename Traits::Epilogue::ScalarD ScalarD; + /// The index. + typedef typename Traits::Index Index; + + /// The number of threads. + static int const kThreads = Traits::GemmConfig::kThreads; + + /// The params. + struct Params : public Traits::Params { + CUTLASS_HOST_DEVICE int initialize(Index m, + Index n, + Index k, + ScalarEpilogue alpha, + ScalarA const* d_a, + Index lda, + ScalarB const* d_b, + Index ldb, + ScalarEpilogue beta, + ScalarC const* d_c, + Index ldc, + ScalarD* d_d, + Index ldd) { + GemmDesc<ScalarEpilogue, Index> desc; + desc.m = m; + desc.n = n; + desc.k = k; + desc.alpha = alpha; + desc.beta = beta; + desc.d_a = reinterpret_cast<void const*>(d_a); + desc.lda = lda; + desc.d_b = reinterpret_cast<void const*>(d_b); + desc.ldb = ldb; + desc.d_c = reinterpret_cast<void const*>(d_c); + desc.ldc = ldc; + desc.d_d = reinterpret_cast<void*>(d_d); + desc.ldd = ldd; + return Traits::Params::initialize(desc); + } + }; + +#if !defined(__CUDACC_RTC__) + /// Launch the kernel. + static __host__ cudaError_t launch(Params const& params, + cudaStream_t stream = cudaStreamDefault) { + // Setup the grid. + dim3 grid; + grid.x = (params.m + Traits::OutputTile::kW - 1) / Traits::OutputTile::kW; + grid.y = (params.n + Traits::OutputTile::kH - 1) / Traits::OutputTile::kH; + + // The number of threads. + dim3 block; + block.x = kThreads; + + // Launch the kernel. + void const* params_ = reinterpret_cast<void const*>(¶ms); + + return cudaLaunchKernel(reinterpret_cast<void*>(&gemm_kernel<This_>), + grid, + block, + const_cast<void**>(¶ms_), + 0, + stream); + } + + /// Launch the kernel. + static __host__ cudaError_t launch(CUfunction kernel, + Params const& params, + CUstream stream = CU_STREAM_LEGACY) { + // Setup the grid. + dim3 grid; + grid.x = (params.m + Traits::OutputTile::kW - 1) / Traits::OutputTile::kW; + grid.y = (params.n + Traits::OutputTile::kH - 1) / Traits::OutputTile::kH; + + // The number of threads. + dim3 block; + block.x = kThreads; + + // Launch the kernel. + void* params_[] = {const_cast<void*>(reinterpret_cast<void const*>(¶ms))}; + + // return cudaLaunchKernel(reinterpret_cast<void*>(&gemm_kernel<This_>), grid, block, + // const_cast<void**>(¶ms_), 0, stream); + CUresult result = cuLaunchKernel( + kernel, grid.x, grid.y, grid.z, block.x, block.y, block.z, 0, stream, params_, 0); + + if (result != CUDA_SUCCESS) { + return cudaErrorLaunchFailure; + } + return cudaSuccess; + } + +#endif + + /// Ctor. + CUTLASS_DEVICE Gemm(Params const& params_, SharedStorage& shared_storage_) + : params(params_), shared_storage(shared_storage_) {} + + /// Consume a single iteration of the loop. + template <bool kIsLastIteration> + CUTLASS_DEVICE void consume_tile(typename Traits::GlobalLoadStream& global_stream, + typename Traits::SharedLoadStream& shared_load_stream, + typename Traits::MultiplyAdd::Accumulators& accumulators, + Index outer_k) { + // If that's the last "load iteration" update the predicates. + if (!kIsLastIteration) { + global_stream.move_to_residue<false>(outer_k); + } + + // Load data for the next iteration of the main loop. + if (!kIsLastIteration) { + global_stream.copy(); + } + + // The unrolling steps for the main loop. + int const kUnrollingSteps = + Traits::MultiplyAdd::AccumulatorsPerWarp::kD / Traits::MultiplyAdd::InstructionShape::kD; + + CUTLASS_PRAGMA_UNROLL + for (int step = 0; step < kUnrollingSteps - 1; ++step) { + // Trigger the copy from shared memory for the next A/B values. + shared_load_stream.copy(step + 1); + // Make sure the values are available for the current iteration to do the multiply-add. + shared_load_stream.commit(step); + + // Do the math on the fragments of the current iteration. + typename Traits::MultiplyAdd multiply_add; + multiply_add.multiply_add(shared_load_stream.fragment_a(step), + shared_load_stream.fragment_b(step), + accumulators, + accumulators); + } + + // Make sure the data from shared memory has been entirely consumed. + Traits::shared_load_fence(true); + + // Commit the data in shared memory for A/B. + if (!kIsLastIteration) { + global_stream.commit(); + } + + // Make sure the data is in shared memory. + Traits::shared_store_fence(true); + + // Trigger the loads for the next iteration (if needed). + if (!kIsLastIteration) { + // Move to the next stage for the load (if it makes sense). + shared_load_stream.inc_stage(); + // Trigger the copy from shared memory for the next loop iteration. + shared_load_stream.copy(0); + } + + // Make sure the values are available for the current iteration to do the multiply-add. + shared_load_stream.commit(kUnrollingSteps - 1); + + // Do the math on the fragments of the current iteration. + typename Traits::MultiplyAdd multiply_add; + multiply_add.multiply_add(shared_load_stream.fragment_a(kUnrollingSteps - 1), + shared_load_stream.fragment_b(kUnrollingSteps - 1), + accumulators, + accumulators); + } + + /// Do the GEMM. + CUTLASS_DEVICE void multiply_add() { + // Swizzle the IDs of the block (to enable better cache behavior). + typename Traits::BlockSwizzle block_swizzle; + dim3 block = block_swizzle.swizzle(); + + // Scale the id. + block.x *= Traits::OutputTile::kW; + block.y *= Traits::OutputTile::kH; + + // We may want to use shared memory to clear the registers. + typedef typename Traits::ClearAccumulators ClearAccumulators; + + // The streams to read A/B from global memory to shared memory. + typename Traits::GlobalLoadStream global_stream(params, shared_storage, block); + + // Create the accumulator clear. + ClearAccumulators clear(shared_storage.main_loop.clear); + + // By how much we unroll the main loop. + Index const kUnroll = static_cast<Index>(Traits::OutputTile::kD); + + // If we do not have enough steps in the main loop, trigger the residue code. + global_stream.move_to_residue<true>(params.k); + + // Fetch the fragments for A and B from global memory. + global_stream.copy(); + + // Copy the elements to shared memory (after transformation if needed). + global_stream.commit(); + + // Make sure the data is in shared memory. + Traits::shared_store_fence(false); + + // Rollback to the beginning of the GEMM-K dimension. It may have no impact. + global_stream.rollback(); + + // The unrolling steps for the main loop. + int const kUnrollingSteps = + Traits::MultiplyAdd::AccumulatorsPerWarp::kD / Traits::MultiplyAdd::InstructionShape::kD; + + // Make sure we have at least 2 unrolling steps or our pipeling is not going to work. + static_assert(kUnrollingSteps >= 2, "The pipelining assumes at least two steps"); + + // The stream of data from shared memory to fragments. + typename Traits::SharedLoadStream shared_load_stream(params, shared_storage); + + // Trigger the copy from shared memory for the 1st stream. + shared_load_stream.copy(0); + + // Allocate the accumulators. + typename Traits::MultiplyAdd::Accumulators accumulators; + // Clear the accumulators. + clear.clear(accumulators); + + // The loop index. + Index outer_k = params.k - kUnroll; + + // Enter the main loop and iterate. + for (; outer_k > 0; outer_k -= kUnroll) { + consume_tile<false>(global_stream, shared_load_stream, accumulators, outer_k); + } + + // Residual loop. + for (; outer_k > -kUnroll; outer_k -= kUnroll) { + consume_tile<true>(global_stream, shared_load_stream, accumulators, outer_k); + } + + // Epilogue. + typedef typename Traits::Epilogue Epilogue; + Epilogue epilogue(params.epilogue, shared_storage.epilogue, params.m, params.n); + epilogue.epilogue(cutlass::make_Coord(0, block.y, block.x), accumulators); + } + + /// The params. + Params const& params; + /// The shared storage. + SharedStorage& shared_storage; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass |