diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/gemm.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/gemm.h | 344 |
1 files changed, 0 insertions, 344 deletions
diff --git a/cutlass-example/cutlass/gemm/gemm.h b/cutlass-example/cutlass/gemm/gemm.h deleted file mode 100644 index c50a3f0..0000000 --- a/cutlass-example/cutlass/gemm/gemm.h +++ /dev/null @@ -1,344 +0,0 @@ -/*************************************************************************************************** - * 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 |