diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm')
29 files changed, 6849 insertions, 0 deletions
diff --git a/cutlass-example/cutlass/gemm/clear_accumulators.h b/cutlass-example/cutlass/gemm/clear_accumulators.h new file mode 100644 index 0000000..441370f --- /dev/null +++ b/cutlass-example/cutlass/gemm/clear_accumulators.h @@ -0,0 +1,57 @@ +/*************************************************************************************************** + * 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 Defines abstractions for efficiently clearing accumulator tiles. +*/ +#pragma once + +#include <cutlass/vector.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, int kLanes_ = 1> +struct ClearAccumulators { + /// The shared storage. + struct SharedStorage {}; + + /// Ctor. + CUTLASS_DEVICE ClearAccumulators() {} + /// Ctor. + CUTLASS_DEVICE ClearAccumulators(SharedStorage& shared_storage) {} + + /// Clear the fragment. + template <typename Fragment_> + CUTLASS_DEVICE void clear(Fragment_& fragment) { + fragment.clear(); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/dgemm_traits.h b/cutlass-example/cutlass/gemm/dgemm_traits.h new file mode 100644 index 0000000..0bbc221 --- /dev/null +++ b/cutlass-example/cutlass/gemm/dgemm_traits.h @@ -0,0 +1,127 @@ +/*************************************************************************************************** + * 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 Defines structural traits of double-precision GEMM. +*/ +#pragma once + +#include <cutlass/gemm/gemm.h> +#include <cutlass/gemm/gemm_epilogue.h> +#include <cutlass/gemm/gemm_epilogue_traits.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/gemm/gemm_shared_tile.h> +#include <cutlass/gemm/gemm_traits.h> +#include <cutlass/gemm/thread_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The tile size for the GEMM KxNxM. + typename OutputTile_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_ = 1, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_ = 1> +struct DgemmConfig + : public GemmConfig< + /// The scalar type for A. + double, + /// The scalar type for B. + double, + /// The scalar type for C. + double, + /// The scalar type for D. + double, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, double, double, double>, + /// The number of scalars per LDG for A. + kScalarsPerLdgA_, + /// The number of scalars per STS for A. + kScalarsPerLdgA_, + /// The number of scalars per LDS for A. + 2, + /// The number of scalars per LDG for B. + kScalarsPerLdgB_, + /// The number of scalars per STS for B. + kScalarsPerLdgB_, + /// The number of scalars per LDS for B. + 2, + /// The number of scalars per LDG for C and STG for D. + 1, + /// The number of scalars per STS for D. + 2, + /// The number of scalars per LDS for D. + 1, + /// The number of stages in shared memory. + 2> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_ = Shape<8, 64, 128>, + /// The functor to use in the epilogue. + typename EpilogueFunctor_ = LinearScaling<double>, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<8, 8, 8>, + /// The number of doubles loaded in one LDG for A. + int kScalarsPerLdgA_ = 1, + /// The number of doubles loaded in one LDG for B. + int kScalarsPerLdgB_ = 1, + /// The index. + typename Index_ = int, + /// The DGEMM config. + typename GemmConfig_ = + DgemmConfig<OutputTile_, AccumulatorsPerThread_, kScalarsPerLdgA_, kScalarsPerLdgB_>, + /// The traits class for the epilogue. + typename GemmEpilogueTraits_ = + SimplifiedGemmEpilogueTraits<GemmConfig_, EpilogueFunctor_, Index_> > +struct DgemmTraits : public SimplifiedGemmTraits< + // The layout for A. + kLayoutA_, + // The layout for B. + kLayoutB_, + // The config. + GemmConfig_, + // The epilogue. + GemmEpilogue<GemmEpilogueTraits_>, + // The index. + Index_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass 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 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 diff --git a/cutlass-example/cutlass/gemm/gemm_epilogue_traits.h b/cutlass-example/cutlass/gemm/gemm_epilogue_traits.h new file mode 100644 index 0000000..c06fc25 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_epilogue_traits.h @@ -0,0 +1,331 @@ +/*************************************************************************************************** + * 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 Defines structural properties of the GEMM epilogue. +*/ +#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 < + /// The output tile. + typename OutputTile_, + /// The accumulators. + typename Accumulators_, + /// The iterator to load C from global memory. + typename GlobalLoadIteratorC_, + /// The transformer for C. + typename GlobalTransformerC_, + /// The transformer for D. + typename GlobalTransformerD_, + /// The iterator to store D to global memory. + typename GlobalStoreIteratorD_, + /// The iterator to store D to shared memory. + typename SharedStoreIteratorD_, + /// The shared store transformer for D. + typename SharedStoreTransformerD_, + /// The iterator to load D from shared memory. + typename SharedLoadIteratorD_, + /// The number of iterations in the epilogue. + typename Iterations_, + /// The iterations strides. + typename Delta_, + /// The functor to be used in the epilogue. + typename Functor_, + /// The index. + typename Index_ = int> +struct GemmEpilogueTraits { + // + /// The output tile. + typedef OutputTile_ OutputTile; + /// The number of iterations. + /// The accumulators. + typedef Accumulators_ Accumulators; + /// The iterator for C in global memory. + typedef GlobalLoadIteratorC_ GlobalLoadIteratorC; + /// The transformer for C. + typedef GlobalTransformerC_ GlobalTransformerC; + /// The transformer for D. + typedef GlobalTransformerD_ GlobalTransformerD; + /// The iterator for D in global memory. + typedef GlobalStoreIteratorD_ GlobalStoreIteratorD; + /// The iterator to store D in shared memory. + typedef SharedStoreIteratorD_ SharedStoreIteratorD; + /// The shared store transformer for D. + typedef SharedStoreTransformerD_ SharedStoreTransformerD; + /// The iterator to store D in shared memory. + typedef SharedLoadIteratorD_ SharedLoadIteratorD; + /// typedef typename GemmConfig::EpilogueIterations Iterations; + typedef Iterations_ Iterations; + /// The iterations strides. + typedef Delta_ Delta; + + /// The functor in charge of the math. + typedef Functor_ Functor; + /// The index. + typedef Index_ Index; + + /// We do not support 3D or 4D shapes. + static_assert(Iterations::kD == 1 && Iterations::kC == 1, "Unsupported 3D/4D shapes"); + + /// The scalar. + typedef typename Functor::Scalar Scalar; + /// The scalar for C. + typedef typename GlobalLoadIteratorC::Scalar ScalarC; + /// The scalar for D. + typedef typename GlobalStoreIteratorD::Scalar ScalarD; + + /// The params. + struct Params { + /// The strides for H and W in the different iterations of the epilogue. + Index stride_h, stride_w; + /// The params for the C iterator. + typename GlobalLoadIteratorC::Params iterator_c; + /// The params for the D global iterator. + typename GlobalStoreIteratorD::Params iterator_d; + /// The params for the D shared store iterator. + typename SharedStoreIteratorD::Params shared_store_iterator_d; + /// The params for the D shared load iterator. + typename SharedLoadIteratorD::Params shared_load_iterator_d; + /// The functor params. + typename Functor::Params functor; + + /// Setup the params. + 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); + } + }; + + /// The shared memory storage to exchange data. + union StreamSharedStorage { + // The storage for the store iterator. + typename SharedStoreIteratorD::SharedStorage store; + // The storage for the store iterator. + typename SharedLoadIteratorD::SharedStorage load; + }; + + /// The shared memory to swizzle the data in the epilogue. + struct SharedStorage { + // The storage for the shared stream D. + StreamSharedStorage shared_stream; + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename EpilogueFunctor_, typename Index_ = int> +struct GemmEpilogueTraitsHelper { + /// The scalar. + typedef typename EpilogueFunctor_::Scalar Scalar; + /// The output tile. + typedef typename GemmConfig_::OutputTile OutputTile; + + /// The number of iterations in the epilogue. + 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; + /// The functor to do the math in the epilogue. + typedef EpilogueFunctor_ Functor; + + /// The traits class to build the iterator to store to shared memory for D. + 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; + + /// The iterator to store D to shared memory. + typedef TileStoreIterator<SharedStoreTileTraits, + typename SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorD; + + /// The shared store transformer for D. + typedef Copy<typename SharedStoreIteratorD::Fragment> SharedStoreTransformerD; + + /// The traits class to build the iterator to load from shared memory for D. + 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; + + /// The iterator to load D from shared memory. + typedef TileLoadIterator<SharedLoadTileTraits, + typename SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorD; + + /// The traits class to build the iterator to load data from global memory for C^N. + 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; + + /// The iterator to load C. + typedef GemmGlobalIteratorCd<GlobalLoadTileTraits, Index_> GlobalLoadIteratorC; + /// The transformer for C. + typedef Copy<typename GlobalLoadIteratorC::Fragment> GlobalTransformerC; + + /// The traits class to build the iterator to store data to global memory for D^N. + 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; + + /// The iterator to store D. + typedef GemmGlobalIteratorCd<GlobalStoreTileTraits, Index_> GlobalStoreIteratorD; + /// The transformer for D. + typedef Copy<typename GlobalStoreIteratorD::Fragment> GlobalTransformerD; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The GEMM config. + typename GemmConfig_, + /// The epilogue functor to do the math in the epilogue. + typename EpilogueFunctor_, + /// The index. + typename Index_ = int, + /// The helper to create the traits class. + 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 diff --git a/cutlass-example/cutlass/gemm/gemm_global_stream.h b/cutlass-example/cutlass/gemm/gemm_global_stream.h new file mode 100644 index 0000000..ec675a3 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_global_stream.h @@ -0,0 +1,182 @@ +/*************************************************************************************************** + * 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 efficient loading of the thread block-level tile from global memory and + storing + to shared memory. +*/ +#pragma once + +#include <cutlass/convert.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/iterator_access.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The load iterator. + typename LoadIterator_, + /// The store iterator to copy to shared memory. + typename StoreIterator_, + /// The transformer to be applied after the data has been copied from global memory. + typename Transformer_> + +struct GlobalLoadStreamBase { + /// The load iterator. + typedef LoadIterator_ LoadIterator; + /// The transformer. + typedef Transformer_ Transformer; + /// The store iterator to write to shared memory. + typedef StoreIterator_ StoreIterator; + + /// The fragment that is copied from shared memory. + typedef typename LoadIterator::Fragment FetchedFragment; + /// The fragment that is obtained after the transformation by the transformer. + typedef typename Transformer::OutputFragment TransformedFragment; + /// Make sure the fragments match. + static_assert((platform::is_same<FetchedFragment, typename Transformer::InputFragment>::value), + ""); + /// The output fragment. + typedef TransformedFragment Fragment; + /// Make sure the transformed fragment is the same as the store fragment. + static_assert((platform::is_same<TransformedFragment, typename StoreIterator::Fragment>::value), + ""); + + /// The layout. + static MatrixLayout::Kind const kLayout = LoadIterator::kLayout; + /// The scalar type of the iterator. + typedef typename LoadIterator::Scalar Scalar; + /// The pointer. + typedef typename LoadIterator::Pointer Pointer; + /// The index. + typedef typename LoadIterator::Index Index; + + /// The params. + struct Params { + // The load iterator. + typename LoadIterator::Params load_iterator; + // The store iterator. + typename StoreIterator::Params store_iterator; + + /// Setup the params. + template <typename GemmDesc_> + CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc, Pointer pointer, Index ld) { + int error_code = load_iterator.initialize(desc, pointer, ld); + if (error_code) { + return error_code; + } + + return store_iterator.initialize(); + } + }; + + /// The amount of storage in shared memory needed to store the tile. + typedef typename StoreIterator::SharedStorage SharedStoreStorage; + + /// The storage in shared memory needed by that stream. + union SharedStorage { + // The load iterator. + typename LoadIterator::SharedStorage load_iterator; + // The store iterator. + SharedStoreStorage store_iterator; + }; + + /// Ctor. + CUTLASS_DEVICE GlobalLoadStreamBase(Params const& params, + SharedStorage& shared_storage, + Coord<3> const bounds, + Coord<3> const& block) + : load_iterator(params.load_iterator, bounds, block), + transformer(), + store_iterator(params.store_iterator, shared_storage.store_iterator) + + { + fetched_fragment.clear(); + } + + /// Load the data from shared memory to the fetch fragment. + CUTLASS_DEVICE void copy() { iterator_load(load_iterator, fetched_fragment); } + + /// Commit the data. + CUTLASS_DEVICE void commit() { + transformer.transform(fetched_fragment, transformed_fragment); + iterator_store(store_iterator, transformed_fragment); + store_iterator.inc_stage(); + } + + /// Move to the beginning of the residue code. That's a new code path in CUTLASS 1.0.1. + CUTLASS_DEVICE void move_to_residue(Index k) { load_iterator.move_to_residue(k); } + + /// Execute the residue code. + CUTLASS_DEVICE void residue(Index k, bool skip_clear = false) { + load_iterator.residue(k); + if (!skip_clear) { + fetched_fragment.clear(); + } + } + + /// Rollback to the beginning of the GEMM-k dimension. + CUTLASS_DEVICE void rollback() { load_iterator.rollback(); } + + /// The iterator. + LoadIterator load_iterator; + /// The fragment to fetch from shared memory. + FetchedFragment fetched_fragment; + /// The transformer. + Transformer transformer; + /// The fragment to convert the data after it has been fetched from shared memory. + TransformedFragment transformed_fragment; + /// The store iterator. + StoreIterator store_iterator; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The load iterator. + typename LoadIterator_, + /// The store iterator to copy to shared memory. + typename StoreIterator_, + /// The transformer to be applied after the data has been copied from global memory. + typename Transformer_ = Copy<typename LoadIterator_::Fragment> > + +struct GlobalLoadStream : public GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> { + /// The base class. + typedef GlobalLoadStreamBase<LoadIterator_, StoreIterator_, Transformer_> Base; + + /// Ctor. + CUTLASS_DEVICE GlobalLoadStream(typename Base::Params const& params, + typename Base::SharedStorage& shared_storage, + Coord<3> const& bounds, + Coord<3> const& block) + : Base(params, shared_storage, bounds, block) {} +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/gemm_global_tile.h b/cutlass-example/cutlass/gemm/gemm_global_tile.h new file mode 100644 index 0000000..1cc3b33 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_global_tile.h @@ -0,0 +1,541 @@ +/*************************************************************************************************** + * 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 Defines iterators for efficiently loading and storing to global memory. +*/ +#pragma once + +#include <cutlass/coord.h> +#include <cutlass/util/platform.h> + +#include <cutlass/gemm/gemm_operand.h> +#include <cutlass/matrix_traits.h> +#include <cutlass/predicate_vector.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/tile_iterator.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// The following functor reshapes a tile of threads to match a tile of data. The idea is that when +// the user wants to build the iterator traits, he/she may want to specify the tile independently +// from the number of scalars loaded/stored per instruction. For example, in the row-major version +// with a tile of size 128x8 - the user may want to that the iterator works with 32x8 threads if +// each thread loads 1 scalar per LDG. If the user changes to 4 scalars per LDG, then the tile of +// threads has to change. The code below detects that and correct the code automatically - it is +// a helper when the user does not specify the right configuration. + +template <typename Tile_, typename Threads_, bool = (Tile_::kW < Threads_::kW)> +struct ReshapeThreads { + typedef Threads_ Threads; +}; + +template <typename Tile_, typename Threads_> +struct ReshapeThreads<Tile_, Threads_, true> { + typedef Shape<Threads_::kD, Threads_::kH * Threads_::kW / Tile_::kW, Tile_::kW, 1> Threads; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct GemmGlobalTileTraits { + /// Identity of the operand + static GemmOperand::Kind const kOperand = kOperand_; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kAccessSize_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kGlobal; + + /// The tile shape + typedef typename ReshapeTile<Tile_, kAccessSize_>::Tile Tile; + /// The threads shape + typedef typename ReshapeThreads<Tile, Threads_>::Threads Threads; + /// The relative offset between two elements in the H/W dimension in adjacent threads. + typedef Shape<1, 1, Tile::kC> ThreadsDelta; + + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH, Threads::kW * kAccessSize> Delta; + /// Strides for immediate offset computation + typedef Shape<0, 0, Threads::kW * ThreadsDelta::kW, kAccessSize> ImmediateOffsetStrides; + /// The number of iterations needed to load/store the tile. + typedef Shape<1, Tile::kH / Threads::kH, Tile::kW / Threads::kW, Tile::kC / kAccessSize> + Iterations; + + typedef GemmMultiplicandTraits<Tile, kOperand, kLayout> MultiplicandTraits; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Threads::kW * ThreadsDelta::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kStrideH_, int kAccessSize_> +struct GemmGlobalTileCdTraits : public GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> + Base; + + /// The stride in the H dimension. + static int const kStrideH = kStrideH_; + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta; + + typedef typename Base::Iterations Iterations; + + typedef typename Base::Threads Threads; + + typedef typename Base::ThreadsDelta ThreadsDelta; + + typedef typename Base::ImmediateOffsetStrides ImmediateOffsetStrides; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Threads::kW * kStrideH * Iterations::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct GemmGlobalIteratorAb + : public TileLoadIterator<TileTraits_, + typename TileTraits_::Scalar, + TileTraits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH + : IteratorAdvance::kW, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef GemmGlobalIteratorAb<TileTraits_, Index_> This_; /// The base class. + + typedef TileLoadIterator<TileTraits_, + typename TileTraits_::Scalar, + TileTraits_::MultiplicandTraits::kKstrided ? IteratorAdvance::kH + : IteratorAdvance::kW, + MemorySpace::kGlobal, + Index_> + Base; + /// The layout. + static MatrixLayout::Kind const kLayout = TileTraits_::kLayout; + /// Fragment type loaded by the iterator + typedef typename Base::Fragment Fragment; + /// The scalar. + typedef typename TileTraits_::Scalar Scalar; + /// The threads. + typedef typename TileTraits_::Threads Threads; + /// The index. + typedef Index_ Index; + /// The thread offset + typedef typename TileTraits_::ThreadOffset ThreadOffset; + /// Specifies in which dimension post-increment accesses advance. + static IteratorAdvance::Kind const kAdvance = Base::kAdvance; + + typedef cutlass::PredicateVector<ShapeCount<typename Base::Iterations>::kCount> PredicateVector; + + /// Iterator parameters type + typedef typename Base::Params BaseParams; + + struct Params : public BaseParams { + /// Initializes params to load a strip-mined tile, given pointer and stride_h. + template <typename GemmDesc_> + CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc, Scalar const* ptr, Index stride_h) { + Index inc_d = 0; + Index inc_advance = 0; + // Move by some columns for each iteration in the H dimension. + Index inc_h = Base::Delta::kH * stride_h; + + // Move by some more columns in the number of iterations if the D dimension is > 1. + if (Base::Delta::kD > 0) { + inc_d = Base::Delta::kD * stride_h - (Base::Iterations::kH - 1) * inc_h; + } + + // Move to the beginning of the next iteration. + if (kAdvance == IteratorAdvance::kH && Base::Delta::kD > 0) { + inc_advance = inc_d; + } else if (kAdvance == IteratorAdvance::kH) { + inc_advance = inc_h; + } else if (Base::Delta::kD > 0) { + inc_advance = (Base::Iterations::kW + 0) * ShapeCount<typename Base::Delta>::kWc - + (Base::Iterations::kH - 1) * inc_h - + (Base::Iterations::kD - 1) * Base::Delta::kD * stride_h; + } else { + inc_advance = (Base::Iterations::kW + 0) * ShapeCount<typename Base::Delta>::kWc - + (Base::Iterations::kH - 1) * inc_h; + } + + // The dimensions of the tile. + int const kH = TileTraits_::Tile::kH; + int const kW = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + + // Move to the residue. + Index const kBlock = kAdvance == IteratorAdvance::kH ? kH : kW; + // The jump in the gemm-k dimension. + Index const stride = kAdvance == IteratorAdvance::kH ? stride_h : 1; + + // Compute the offset to the residue and how to "come" back. + Index const kResidue = desc.k % kBlock; + if (kResidue > 0) { + move_to_residue_offset = (desc.k - kResidue) * stride; + } else { + move_to_residue_offset = (desc.k - kBlock) * stride; + } + + Base::Params::initialize(ptr, 0, stride_h, 1, inc_d, inc_h, 0, inc_advance); + return 0; + } + + // The extra offset to control moving to the residue. + Index move_to_residue_offset; + }; + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorAb(Params const& _params, + const Coord<3>& bounds, + const Coord<3>& block, + ThreadOffset thread_offset_func = ThreadOffset()) + : params(_params) { + thread_offset = thread_offset_func(); + // The column. + Index block_h = thread_offset[1]; + // The contiguous dimension. + Index block_w = thread_offset[2]; + + // Add the blocks indices. + if (kAdvance == IteratorAdvance::kH) { + block_h += block[1]; + block_w += block[2]; + + } else { + block_h += block[2]; + block_w += block[1]; + } + + // Setup the pointer. + params.pointer += (block_h * params.stride_h + block_w); + + // Initialize predicates + initialize_predicates(bounds, make_Coord(0, block_h, block_w)); + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm); + } + + /// Increment the pointer in the H dimension. + CUTLASS_DEVICE void inc_h() { params.pointer += params.inc_h; } + /// Increment the pointer in the D dimension. + CUTLASS_DEVICE void inc_d() { params.pointer += params.inc_d; } + /// Increment the pointer to move to the next iteration. + CUTLASS_DEVICE void inc_advance() { params.pointer += params.inc_advance; } + + /// Initialize the predicates. + CUTLASS_DEVICE void initialize_predicates(const Coord<3>& bounds, const Coord<3>& block) { + // Setup the masks to control loads. + predicates.fill(0); + + int bounds_h, bounds_w; + if (kAdvance == IteratorAdvance::kH) { + bounds_w = bounds[2] - block[2]; + bounds_h = bounds[1]; + + } else { + bounds_w = bounds[1]; + bounds_h = bounds[2] - block[1]; + } + + // Fill in the bits of the predicate vector. + for (int d = 0; d < Base::Iterations::kD; ++d) { + for (int h = 0; h < Base::Iterations::kH; ++h) { + for (int w = 0; w < Base::Iterations::kW; ++w) { + for (int c = 0; c < Base::Iterations::kC; ++c) { + bool flag = w * Base::Delta::kW < bounds_w; + if (kAdvance == IteratorAdvance::kH) { + flag = flag && (h * Base::Delta::kH + d * Base::Delta::kD) < bounds_h; + } else { + flag = flag && (h * Base::Delta::kH) < bounds_h; + } + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + predicates.set(bit, flag); + } + } + } + } + } + + /// Move to residue portion. + CUTLASS_DEVICE void move_to_residue(Index k) { + // Store the pointer and the predicates. + stored_pointer = params.pointer; + stored_predicates = predicates; + + // Move the pointer to the residue. + params.pointer += params.move_to_residue_offset; + + // The dimensions of the tile. + int const kH = TileTraits_::Tile::kH; + int const kW = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + + // The unrolling factor. + int const kUnroll = kAdvance == IteratorAdvance::kH ? kH : kW; + + // Clear the predicates for the residue. TODO: We can do something smarter. + int const kResidue = (int)(k % (Index)kUnroll); + if (kResidue > 0) { + residue(kResidue); + } + } + + /// That's the residue! Update the predicates. + CUTLASS_DEVICE void residue(Index k) { + // The coordinates of the thread. + Index block_h = thread_offset[1]; + // The contiguous dimension. + Index block_w = thread_offset[2]; + + // Update the predicate vector. + for (int d = 0; d < Base::Iterations::kD; ++d) { + for (int h = 0; h < Base::Iterations::kH; ++h) { + for (int w = 0; w < Base::Iterations::kW; ++w) { + for (int c = 0; c < Base::Iterations::kC; ++c) { + Index offset = 0; + if (kAdvance == IteratorAdvance::kH) { + offset += block_h + h * Base::Delta::kH + d * Base::Delta::kD; + } else { + offset += block_w + w * Base::Delta::kW; + } + + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + if (offset >= k) { + predicates.set(bit, false); + } + } + } + } + } + } + + /// Rollback to beginning of first tile and initialize predicates. + CUTLASS_DEVICE void rollback() { + params.pointer = stored_pointer; + predicates = stored_predicates; + } + + /// Is the iterator valid? + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { + int const bit = ComputeOffsetFromShape<typename Base::Iterations>::get(d, h, w, c); + return predicates[bit]; + } + + /// Offset of an individual lane from the start of the tile + Coord<4> thread_offset; + /// The parameters + Params params; + /// The pointer. + typename Base::Scalar const* stored_pointer; + /// The predicates. + PredicateVector predicates, stored_predicates; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct GemmGlobalIteratorCd : public TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef GemmGlobalIteratorCd<TileTraits_, Index_> This_; + /// The base class. + typedef TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> + Base; + + /// The layout. + static MatrixLayout::Kind const kLayout = TileTraits_::kLayout; + + /// The scalar. + typedef typename TileTraits_::Scalar Scalar; + /// The pointer. + typedef typename TileTraits_::Pointer Pointer; + /// The threads. + typedef typename TileTraits_::Threads Threads; + /// The index. + typedef Index_ Index; + /// The thread offset + typedef typename TileTraits_::ThreadOffset ThreadOffset; + + /// The params. + struct Params { + /// The pointer. + Pointer pointer; + /// The stride in the H dimension to setup the thread in the block. + Index stride_h; + /// The strides to increment the pointer. + Index inc_advance, inc_h; + /// The strides to increment the predicate offset + Index predicate_inc_advance, predicate_inc_h; + /// The column offset to compute the predicate for the columns. + Index predicate_offset; + + /// Setup the params. + CUTLASS_HOST_DEVICE int initialize( + Pointer pointer, Index ld, Index bound, Index epilogue_stride_w, Index epilogue_delta_w) { + // The pointer. + this->pointer = pointer; + // Each column of the matrix. + stride_h = TileTraits_::ThreadsDelta::kH * ld; + // Each thread output 1 column per iteration. The stride between columns is given by the + // number of scalars that are loaded per LDS for B. + inc_h = ld * TileTraits_::kStrideH; + inc_advance = + (ld - ld * TileTraits_::kStrideH * (Base::Iterations::kH - 1)) + epilogue_stride_w; + + predicate_offset = bound; + predicate_inc_h = TileTraits_::kStrideH; + predicate_inc_advance = + -((TileTraits_::kStrideH * (Base::Iterations::kH - 1) - 1) + epilogue_delta_w); + + return 0; + } + }; + + Params params; + /// Offset of an individual lane from the start of the tile + Coord<4> thread_offset; + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorCd() {} + + /// Ctor. + CUTLASS_DEVICE GemmGlobalIteratorCd(Params const& params, + const Coord<3>& bounds, + const Coord<3>& block, + int offset = 0, + int pred_offset = 0, + ThreadOffset thread_offset_func = ThreadOffset()) + : params(params) { + thread_offset = thread_offset_func(); + // Each warp works on a different column of the tile. + int const h = thread_offset[1] + block[1]; + // Each lane writes a different element. + int const w = thread_offset[2] + block[2]; + // Setup the pointer. + this->params.pointer += ((h * params.stride_h + w) + offset); + + // Prepare the vector of predicates. + for (int i = 0; i < Base::Iterations::kW; ++i) { + predicates.set(i, w + i * Base::Delta::kW < bounds[2]); + } + this->params.predicate_offset -= (h + pred_offset); + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm); + } + + /// Increment the pointer in the C dimension. + CUTLASS_DEVICE void inc_c() {} + /// Increment the pointer in the W dimension. + CUTLASS_DEVICE void inc_w() {} + /// Increment the pointer in the H dimension. + CUTLASS_DEVICE void inc_h() { + params.pointer += params.inc_h; + params.predicate_offset -= params.predicate_inc_h; + } + /// Increment the pointer in the D dimension. + CUTLASS_DEVICE void inc_d() {} + /// Increment the pointer to move to the next iteration. + CUTLASS_DEVICE void inc_advance() { + params.pointer += params.inc_advance; + this->params.predicate_offset -= params.predicate_inc_advance; + } + + /// The accessor. + CUTLASS_DEVICE void set(typename Base::AccessType const& value, int d, int h, int w, int c) { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store( + value, params.pointer, imm); + } + + /// Test the validity of the iterator. + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { + return predicates.at(w) && params.predicate_offset > 0; + } + + /// The predicates for the row. + cutlass::PredicateVector<Base::Iterations::kW> predicates; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/gemm_operand.h b/cutlass-example/cutlass/gemm/gemm_operand.h new file mode 100644 index 0000000..737f993 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_operand.h @@ -0,0 +1,141 @@ +/*************************************************************************************************** + * 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 Defines constant expressions for mapping GEMM problem size and strides onto pitch-linear + memory. +*/ +#pragma once + +#include <cutlass/matrix_traits.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/util/platform.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to describe attributes of GEMM matrix operands +template <GemmOperand::Kind kOperand_, MatrixLayout::Kind kLayout_> +struct GemmOperandTraitsAb { + static const bool Congruous = + (kOperand_ == GemmOperand::kA ^ kLayout_ == MatrixLayout::kRowMajor); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmOperand::Kind kOperand_, typename Tile_> +struct GetExtent; + +template <typename Tile_> +struct GetExtent<GemmOperand::kA, Tile_> { + static const int kExtent = Tile_::kW; +}; + +template <typename Tile_> +struct GetExtent<GemmOperand::kB, Tile_> { + static const int kExtent = Tile_::kH; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Determines the shape of a multiplicand tile in terms of strided (H) and contiguous (W) +/// dimensions +template <typename ThreadBlockTile_, GemmOperand::Kind Usage, MatrixLayout::Kind Layout> +struct GemmMultiplicandTraits { + // Only defined for A or B + static_assert(Usage == GemmOperand::kA || Usage == GemmOperand::kB, + "MultiplicandTileShape defined only for A or B operands."); + + /// Shape of GEMM thread block tile (K, N, M) + typedef ThreadBlockTile_ ThreadBlockTile; + + /// Identifies multiplicand + static GemmOperand::Kind const kUsage = Usage; + + /// Layout of tile + static MatrixLayout::Kind const kLayout = Layout; + + // True if K is the strided dimension + static bool const kKstrided = (kUsage == GemmOperand::kA ^ kLayout == MatrixLayout::kRowMajor); + + /// Map the ThreadBlockShape onto (kH, kW) dimensions for A and B operand + typedef typename platform::conditional< + kKstrided, + Shape<1, ThreadBlockTile::kD, GetExtent<Usage, ThreadBlockTile>::kExtent>, + Shape<1, GetExtent<Usage, ThreadBlockTile>::kExtent, ThreadBlockTile::kD> >::type Shape; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Project's a coordinate (K, N, M) onto inner and outer dimensions defined for each +/// operand. +template <GemmOperand::Kind operand, bool Kstrided = true> +struct ProjectOperand; + +/// Project A operand - (0, K, M) +template <bool Kstrided> +struct ProjectOperand<GemmOperand::kA, Kstrided> { + CUTLASS_HOST_DEVICE + static Coord<3> project(Coord<3> const &coord) { + if (Kstrided) { + return make_Coord(0, coord[0], coord[2]); + } else { + return make_Coord(0, coord[2], coord[0]); + } + } +}; + +/// Project B operand - (0, K, N) +template <bool Kstrided> +struct ProjectOperand<GemmOperand::kB, Kstrided> { + CUTLASS_HOST_DEVICE + static Coord<3> project(Coord<3> const &coord) { + if (Kstrided) { + return make_Coord(0, coord[0], coord[1]); + } else { + return make_Coord(0, coord[1], coord[0]); + } + } +}; + +/// Project C operand - (0, N, M) +template <> +struct ProjectOperand<GemmOperand::kC, true> { + CUTLASS_HOST_DEVICE + static Coord<3> project(Coord<3> const &coord) { return make_Coord(0, coord[1], coord[2]); } +}; + +/// Project D operand - (0, N, M) +template <> +struct ProjectOperand<GemmOperand::kD, true> { + CUTLASS_HOST_DEVICE + static Coord<3> project(Coord<3> const &coord) { return make_Coord(0, coord[1], coord[2]); } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/gemm_shared_stream.h b/cutlass-example/cutlass/gemm/gemm_shared_stream.h new file mode 100644 index 0000000..c6ff7bd --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_shared_stream.h @@ -0,0 +1,113 @@ +/*************************************************************************************************** + * 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 Defines abstractions for managing loading and storing fragments to shared memory in the + efficient GEMM pipeline. +*/ +#pragma once + +#include <cutlass/gemm/gemm_shared_tile.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The load iterator. + typename Iterator_, + /// The transformer to be applied after the data has been copied from shared memory. + typename Transformer_ = Copy<typename Iterator_::Fragment> > + +struct SharedLoadStream { + /// The load iterator. + typedef Iterator_ Iterator; + /// The transformer. + typedef Transformer_ Transformer; + + /// The fragment that is copied from shared memory. + typedef typename Iterator::Fragment FetchedFragment; + /// The fragment that is obtained after the transformation by the transformer. + typedef typename Transformer::OutputFragment TransformedFragment; + /// Make sure the fragments match. + static_assert((platform::is_same<FetchedFragment, typename Transformer::InputFragment>::value), + ""); + /// The output fragment. + typedef TransformedFragment Fragment; + + /// The params. + struct Params { + /// The iterator params. + typename Iterator::Params iterator; + + /// Setup the params. + CUTLASS_HOST_DEVICE int initialize() { return iterator.initialize(); } + }; + + /// The storage in shared memory needed by that stream. + typedef typename Iterator::Storage SharedStorage; + + /// Ctor. + CUTLASS_DEVICE SharedLoadStream() {} + + /// Ctor. + CUTLASS_DEVICE SharedLoadStream(Params const ¶ms, SharedStorage &shared_storage) { + this->initialize(params, shared_storage); + } + + /// Initialize the stream. + CUTLASS_DEVICE void initialize(Params const ¶ms, SharedStorage &shared_storage) { + // The iterator. + iterator = Iterator(params.iterator, shared_storage); + // The transformer. + transformer = Transformer(); + } + + /// Load the data from shared memory to the fetch fragment. + CUTLASS_DEVICE void copy(FetchedFragment &fetched) { shared_iterator_load(iterator, fetched); } + + /// Load the data from shared memory to the fetch fragment. + CUTLASS_DEVICE void copy(int d, FetchedFragment &fetched) { + shared_iterator_load(iterator, fetched, d); + } + + /// Commit the data. + CUTLASS_DEVICE void commit(FetchedFragment &fetched, TransformedFragment &transformed) { + transformer.transform(fetched, transformed); + } + + /// Increment the stage. + CUTLASS_DEVICE void inc_stage() { iterator.inc_stage(); } + + /// The iterator. + Iterator iterator; + /// The transformer. + Transformer transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/gemm_shared_tile.h b/cutlass-example/cutlass/gemm/gemm_shared_tile.h new file mode 100644 index 0000000..7c61e02 --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_shared_tile.h @@ -0,0 +1,417 @@ +/*************************************************************************************************** + * 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 Defines iterators for efficiently loading and storing tiles to and from shared memory. +*/ +#pragma once + +#include <cutlass/gemm/gemm_operand.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kScalarsPerSts_> +struct GemmSharedStoreTileAbTraits { + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The tile. + typedef typename ReshapeTile<Tile_, kScalarsPerSts_>::Tile Tile; + /// The threads. + typedef Threads_ Threads; + /// The strides to compute the base position of the thread. + typedef Shape<0, ShapeCount<Tile>::kWc, Tile::kC, kScalarsPerSts_> ThreadsStrides; + /// The skew. + static int const kSkew = 0; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kScalarsPerSts_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of iterations needed to load/store the tile. + typedef Shape<1, + Tile::kH / Threads::kH, + Tile::kW / Threads::kW, + Tile::kC / Threads::kC / kAccessSize> + Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kAccessSize> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kAccessSize> + ImmediateOffsetStrides; + + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int offset = ComputeThreadOffsetFromStrides<Threads, ThreadsStrides>::get(); + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kScalarsPerSts_, int kSkew_> +struct GemmSharedStoreWithSkewTileAbTraits { + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The tile without skews. + typedef typename ReshapeTile<Tile_, kScalarsPerSts_>::Tile TileWithoutSkew; + /// The tile. + typedef typename ReshapeTile<Shape<Tile_::kD, Tile_::kH, Tile_::kW + kSkew_>, + kScalarsPerSts_>::Tile Tile; + /// The threads. + typedef Threads_ Threads; + /// The skew. + static int const kSkew = kSkew_; + /// The number of scalars per STS. + static int const kAccessSize = kScalarsPerSts_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of iterations needed to load/store the tile. + typedef Shape<1, TileWithoutSkew::kH / Threads::kW, TileWithoutSkew::kW / Threads::kH> Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, ShapeCount<Tile>::kWc, Threads::kH * kAccessSize> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, ShapeCount<Tile>::kWc, Threads::kH * kAccessSize> ImmediateOffsetStrides; + + struct ThreadOffset { + CUTLASS_HOST_DEVICE Coord<4> operator()() const { + int offset = ComputeThreadOffsetFromStrides<Threads, ThreadsStrides>::get(); + return make_Coord(0, 0, offset, 0); + } + }; + + protected: + /// The strides to compute the base position of the thread. + typedef Shape<0, kScalarsPerSts_, ShapeCount<Tile>::kHwc / Threads::kW> ThreadsStrides; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, + typename OutputTile_, + typename Warps_, + typename ThreadsPerWarp_, + typename InstructionShape_, + int kStages_, + int kScalarsPerLds_, + int kSkew_ = 0> +struct GemmSharedLoadTileATraits { + static GemmOperand::Kind const kOperand = GemmOperand::kA; + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The tile without skew. + typedef Shape<kStages_, + OutputTile_::kD / InstructionShape_::kD, + GetExtent<kOperand, OutputTile_>::kExtent * InstructionShape_::kD> + TileWithoutSkew_; + /// The tile with skew. + typedef Shape<kStages_, TileWithoutSkew_::kH, TileWithoutSkew_::kW + kSkew_> TileWithSkew; + /// The tile without skew after reshaping. + typedef typename ReshapeTile<TileWithoutSkew_, kScalarsPerLds_>::Tile TileWithoutSkew; + /// The tile. + typedef typename ReshapeTile<TileWithSkew, kScalarsPerLds_>::Tile Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The threads in a warp. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of scalars per LDG/STG. + // static int const kScalarsPerLds = kScalarsPerLds_; + static int const kAccessSize = kScalarsPerLds_; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of warps. + static int const kWarps = GetExtent<kOperand, Warps>::kExtent; + /// The number of threads in one dimension of the warp. + static int const kThreadsPerWarp = GetExtent<kOperand, ThreadsPerWarp>::kExtent; + + /// The number of iterations needed to load/store the tile. + typedef Shape<1, 1, TileWithoutSkew::kW / kWarps / kThreadsPerWarp /* / kScalarsPerLds*/> + Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<TileWithSkew::kW * Warps::kD, 0, kWarps * kThreadsPerWarp * kAccessSize, 0> + ImmediateOffsetStrides; + typedef Shape<TileWithSkew::kW * Warps::kD, 0, kWarps * kThreadsPerWarp * kAccessSize, 0> Delta; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE Coord<4> operator()() const { + // Extract the warp. + int const warp = threadIdx.x / kWarpSize; + // Extract the slice. + int const slice = warp / (Warps::kH * Warps::kW); + // Compute the row offset for each warp. + int const warp_row = warp % Warps::kW; + // Compute the row offset for each thread. + int const lane_row = (threadIdx.x & 0x0e) / 2; + // The offset. + int const offset = + slice * Tile::kW * Tile::kC + (warp_row * ThreadsPerWarp::kW + lane_row) * kAccessSize; + // Embed the offset in a 4D coordinate vector. + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, + typename OutputTile_, + typename Warps_, + typename ThreadsPerWarp_, + typename InstructionShape_, + int kStages_, + int kScalarsPerLds_, + int kSkew_ = 0> +struct GemmSharedLoadTileBTraits { + static GemmOperand::Kind const kOperand = GemmOperand::kB; + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The tile without skew. + typedef Shape<kStages_, + OutputTile_::kD / InstructionShape_::kD, + GetExtent<kOperand, OutputTile_>::kExtent * InstructionShape_::kD> + TileWithoutSkew_; + /// The tile with skew. + typedef Shape<kStages_, TileWithoutSkew_::kH, TileWithoutSkew_::kW + kSkew_> TileWithSkew; + /// The tile without skew after reshaping. + typedef typename ReshapeTile<TileWithoutSkew_, kScalarsPerLds_>::Tile TileWithoutSkew; + /// The tile. + typedef typename ReshapeTile<TileWithSkew, kScalarsPerLds_>::Tile Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The threads in a warp. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kScalarsPerLds_; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of warps. + static int const kWarps = GetExtent<kOperand, Warps>::kExtent; + /// The number of threads in one dimension of the warp. + static int const kThreadsPerWarp = GetExtent<kOperand, ThreadsPerWarp>::kExtent; + + /// The number of iterations needed to load/store the tile. + typedef Shape<1, 1, TileWithoutSkew::kW / kWarps / kThreadsPerWarp /* / kAccessSize*/> Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<TileWithSkew::kW * Warps::kD, 0, kWarps * kThreadsPerWarp * kAccessSize, 0> + ImmediateOffsetStrides; + typedef Shape<TileWithSkew::kW * Warps::kD, 0, kWarps * kThreadsPerWarp * kAccessSize, 0> Delta; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE Coord<4> operator()() const { + // Extract the warp. + int const warp = threadIdx.x / kWarpSize; + // Extract the slice. + int const slice = warp / (Warps::kH * Warps::kW); + // The warp in the slice. + int const warp_in_slice = warp % (Warps::kH * Warps::kW); + // Compute the row offset for each warp. + int const warp_col = warp_in_slice / Warps::kW; + // Compute the row offset for each thread. + int const lane_col = (threadIdx.x & 0x10) / 8 + (threadIdx.x & 0x01); + // The offset. + int const offset = + slice * Tile::kW * Tile::kC + (warp_col * ThreadsPerWarp::kH + lane_col) * kAccessSize; + // Embed the offset in a 4D coordinate. + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, + typename OutputTile_, + typename Warps_, + typename ThreadsPerWarp_, + int kScalarsPerSts_, + int kSkew_ = 0> +struct GemmSharedStoreTileDTraits { + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The dimension of the output tile. + typedef OutputTile_ OutputTile; + /// The warps in the tile. + typedef Warps_ Warps; + /// The threads in the warps. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kScalarsPerSts_; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of scalars per thread. + static int const kScalarsPerThread = OutputTile_::kW / Warps::kW / ThreadsPerWarp::kW; + /// The number of threads. + static int const kThreads = ShapeCount<Warps>::kCount * kWarpSize; + /// The number of scalars per row. We build a tile with 2 rows (to avoid bank conflicts). + static int const kScalarsPerRow = kThreads / 2 * kScalarsPerThread + kSkew; + + /// The tile. + typedef Shape<1, 2, kScalarsPerRow / kAccessSize, kAccessSize> Tile; + /// The number of iterations needed to store the tile. + typedef Shape<1, 1, kScalarsPerThread / kAccessSize> Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * ThreadsPerWarp::kW * kAccessSize> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * ThreadsPerWarp::kW * kAccessSize> ImmediateOffsetStrides; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE Coord<4> operator()() const { + // The warp. + int const warp = threadIdx.x / kWarpSize; + + // The position of the warp in the 2D tile. + int const warp_row = warp % Warps::kW; + int const warp_col = warp / Warps::kW; + + // We assume that the elements are distributed in a warps as 4 columns of 8 elements. The + // columns are stored in threads col0=[0, 2, 4, 6, 8, 10, 12, 14], col1=[1, 3, 5, 7, .., 15], + // col2=[16, 18, 20, ..., 30] and col3=[17, 19, ..., 31]. + int hi_halfwarp_offset = ((threadIdx.x >> 4) & 0x1) * OutputTile::kW; + int lo_halfwarp_offset = ((threadIdx.x >> 1) & 0x7) + ThreadsPerWarp::kW * warp_row; + + // Odd threads go to the second half of shared memory. + int const row = threadIdx.x & 0x01; + int col = warp_col * (ThreadsPerWarp::kH / 2) * OutputTile::kW + + lo_halfwarp_offset * kAccessSize + hi_halfwarp_offset; + // Embed the offset in a 4D coords. + return make_Coord(0, 0, row * kScalarsPerRow + col, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, + typename OutputTile_, + typename Warps_, + typename ThreadsPerWarp_, + int kTileH_, + int kScalarsPerLds_, + int kSkew_ = 0> +struct GemmSharedLoadTileDTraits { + /// The scalar. + typedef typename platform::remove_const<Scalar_>::type Scalar; + /// The pointer. + typedef Scalar_* Pointer; + /// The dimension of the output tile. + typedef OutputTile_ OutputTile; + /// The warps in the tile. + typedef Warps_ Warps; + /// The threads in the warps. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of scalars per LDG/STG. + static int const kAccessSize = kScalarsPerLds_; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The number of scalars per thread. + static int const kScalarsPerThread = OutputTile_::kW / Warps::kW / ThreadsPerWarp::kW; + /// The number of threads. + static int const kThreads = ShapeCount<Warps>::kCount * kWarpSize; + /// The number of scalars per row. We build a tile with 2 rows (to avoid bank conflicts). + static int const kScalarsPerRow = kThreads / 2 * kScalarsPerThread + kSkew; + + /// The tile. We have 2 rows of scalars. We use those two rows to make sure we do not have bank + /// conflicts in the epilogue. + typedef Shape<1, 2, kScalarsPerRow / kAccessSize, kAccessSize> Tile; + + // Compute the number of iterations per warp in the Tile::kH dimension. + static int const kIterationsInHPerWarp = kTileH_ / ShapeCount<Warps>::kCount; + + // As explained above, the shared memory tile is composed of 2 rows and each rows is made of + // kScalarsPerRow. A warp is expected to read from the 1st row, then move to the 2nd row and go + // back to the 1st row. To model that scheme we define the Iterations shape as Shape<X, 2, ...>. + // However, in some cases, we have only 1 iteration per warp. In that case, we must define the + // shape as Shape<1, 1, ...>. The following code does that except that we hijack the kH dimension + // to keep the number of elements to reduce for split-K. + static int const kIterationsH = kIterationsInHPerWarp == 1 ? 1 : 2; + // As soon as we know kIterationsH, it is trivial to compute kIterationsD: + static int const kIterationsD = kIterationsInHPerWarp / kIterationsH; + + // If we have split-K enabled, we have to jump over the elements from the "odd/even" column of + // threads to grab the other elements. + static int const kSplitK = OutputTile::kW * ThreadsPerWarp::kH / 2 * Warps::kH; + + /// The number of iterations needed to store the tile. + typedef Shape<kIterationsD, kIterationsH, OutputTile::kW / kWarpSize / kAccessSize, Warps::kD> + Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<OutputTile::kW, kScalarsPerRow, kWarpSize * kAccessSize, kSplitK> + ImmediateOffsetStrides; + /// The strides in each dimension between different loads/stores. + typedef Shape<OutputTile::kW, kScalarsPerRow, kWarpSize * kAccessSize, kSplitK> Delta; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE Coord<4> operator()() const { + // Each warp works on a different column. + int const h = threadIdx.x / kWarpSize; + // Compute the row. + int const w = (threadIdx.x & (kWarpSize - 1)) * kAccessSize; + int offset = 0; + if (Iterations::kH == 1) { + int const row = h & 0x1; + int const col = h / 2; + offset = row * ShapeCount<Tile>::kWc + col * OutputTile::kW * Iterations::kD + w; + } else { + offset = h * OutputTile::kW * Iterations::kD + w; + } + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/gemm_traits.h b/cutlass-example/cutlass/gemm/gemm_traits.h new file mode 100644 index 0000000..cb57c4d --- /dev/null +++ b/cutlass-example/cutlass/gemm/gemm_traits.h @@ -0,0 +1,818 @@ +/*************************************************************************************************** + * 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 Defines structural properties of complete GEMM computation. +*/ +#pragma once + +#include <cutlass/convert.h> +#include <cutlass/gemm/clear_accumulators.h> +#include <cutlass/gemm/gemm_global_stream.h> +#include <cutlass/gemm/gemm_operand.h> +#include <cutlass/gemm/gemm_shared_stream.h> +#include <cutlass/gemm/identity_block_swizzle.h> +#include <cutlass/matrix_traits.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/tile_iterator.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The scalar type for A. + typename ScalarA_, + /// The scalar type for B. + typename ScalarB_, + /// The scalar type for C. + typename ScalarC_, + /// The scalar type for D. + typename ScalarD_, + /// The output tile size for the GEMM KxNxM. + typename OutputTile_, + /// The functor to do the math. + typename MultiplyAdd_, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_, + /// The number of scalars per STS for A. + int kScalarsPerStsA_, + /// The number of scalars per LDG for A. + int kScalarsPerLdsA_, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_, + /// The number of scalars per STS for B. + int kScalarsPerStsB_, + /// The number of scalars per LDS for B. + int kScalarsPerLdsB_, + /// The number of scalars per LDG for C and STG for D. + int kScalarsPerLdgCAndStgD_, + /// The number of scalars per STS for D. + int kScalarsPerStsD_, + /// The number of scalars per LDS for D. + int kScalarsPerLdsD_, + /// The number of stages in shared memory to do single/double/triple-buffering. + int kStages_, + /// Do we do the residue in the prologue? + bool kResidueInPrologue_ = false> + +struct GemmConfig { + // + /// The scalar for A. + typedef ScalarA_ ScalarA; + /// The scalar for B. + typedef ScalarB_ ScalarB; + /// The scalar for C. + typedef ScalarC_ ScalarC; + /// The scalar for D. + typedef ScalarD_ ScalarD; + + /// The tile. + typedef OutputTile_ OutputTile; + /// The functor to do D = A*B + C. + typedef MultiplyAdd_ MultiplyAdd; + /// The shape of the instruction. + typedef typename MultiplyAdd::InstructionShape InstructionShape; + /// The number of accumulators per warp. + typedef typename MultiplyAdd::AccumulatorsPerWarp AccumulatorsPerWarp; + /// The accumulators. + typedef typename MultiplyAdd::Accumulators Accumulators; + + /// The number of warps. + typedef typename ShapeDiv<OutputTile, AccumulatorsPerWarp>::Shape Warps; + /// The default warp size (32 threads per warp). + static int const kWarpSize = cutlass::kWarpSize; + /// The numnber of threads. + static int const kThreads = ShapeCount<Warps>::kCount * kWarpSize; + + /// The number of scalars per LDG/STS/LDS for A. + static int const kScalarsPerLdgA = kScalarsPerLdgA_; + static int const kScalarsPerStsA = kScalarsPerStsA_; + static int const kScalarsPerLdsA = kScalarsPerLdsA_; + + /// The number of scalars per LDG/STS/LDS for B. + static int const kScalarsPerLdgB = kScalarsPerLdgB_; + static int const kScalarsPerStsB = kScalarsPerStsB_; + static int const kScalarsPerLdsB = kScalarsPerLdsB_; + + /// The number of scalars per LDG for C. + static int const kScalarsPerLdgC = kScalarsPerLdgCAndStgD_; + + /// The number of scalars per STS/LDS/STG for D. + static int const kScalarsPerStgD = kScalarsPerLdgCAndStgD_; + static int const kScalarsPerStsD = kScalarsPerStsD_; + static int const kScalarsPerLdsD = kScalarsPerLdsD_; + + /// The number of accumulators that are going to be fed from one LDS A/B. + static int const kAccumulatorsPerLdsA = kScalarsPerLdsA / InstructionShape::kD; + static int const kAccumulatorsPerLdsB = kScalarsPerLdsB / InstructionShape::kD; + + /// The number of stages in shared memory to implement double, triple, more-buffering. + static int const kStages = kStages_; + + /// Do we do the residue in the prologue? + static bool const kResidueInPrologue = kResidueInPrologue_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind, typename GemmConfig_> +struct GemmTileTraitsHelperA {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarA Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar; + + /// The traits class to build the iterator to load data from global memory for A^N. + typedef GemmGlobalTileTraits< + // That's A. + GemmOperand::kA, + // A is column-major. + MatrixLayout::kColumnMajor, + // The pointer is float const. + Scalar const, + // The tile has size KxM in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kW>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + /// The traits class to build the iterator to store data to shared memory for A^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer is float. + MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kW * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsA> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for A^N. + typedef GemmSharedLoadTileATraits< + // The pointer is float const. + MultiplyAddScalar const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + GemmConfig_::kScalarsPerLdsA, + // The skew. + 0> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct GemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarA Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar; + + /// The traits class to build the iterator to load data from global memory for A^T. + typedef GemmGlobalTileTraits< + // That's A. + GemmOperand::kA, + // A is row-major. + MatrixLayout::kRowMajor, + // The pointer is float const. + Scalar const, + // The tile has size MxK in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + /// The number of scalars in 4B. + static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar); + /// The skew for A. + static int const kSkewA = 128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsA / + GlobalTileTraits::Threads::kW * kScalarsIn4B; + + /// The traits class to build the iterator to store data to shared memory for A^T. + typedef GemmSharedStoreWithSkewTileAbTraits < + // The pointer is float. + MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kW * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS. + GemmConfig_::kScalarsPerStsA, + // The skew to avoid bank conflicts added in the tile W dimension. + kSkewA<GemmConfig_::kScalarsPerLdsA ? GemmConfig_::kScalarsPerLdsA : kSkewA> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for A^T. + typedef GemmSharedLoadTileATraits< + // The pointer is float const. + MultiplyAddScalar const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + GemmConfig_::kScalarsPerLdsA, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind, typename GemmConfig_> +struct GemmTileTraitsHelperB {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct GemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarB Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar; + + /// The traits class to build the iterator to load data from global memory for B^N. + typedef GemmGlobalTileTraits< + // That's B. + GemmOperand::kB, + // B is column-major. + MatrixLayout::kColumnMajor, + // The pointer is float const. + Scalar const, + // The tile has size MxK in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + /// The number of scalars in 4B. + static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar); + /// The skew for B. + static int const kSkewB = 128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsB / + GlobalTileTraits::Threads::kW * kScalarsIn4B; + + /// The traits class to build the iterator to store data to shared memory for B^N. + typedef GemmSharedStoreWithSkewTileAbTraits < + // The pointer is float. + MultiplyAddScalar, + // The tile has size KxN in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kH * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS. + GemmConfig_::kScalarsPerStsB, + // The skew to avoid bank conflicts added in the tile W dimension. + kSkewB<GemmConfig_::kScalarsPerLdsB ? GemmConfig_::kScalarsPerLdsB : kSkewB> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for B^N. + typedef GemmSharedLoadTileBTraits< + // The pointer is float const. + MultiplyAddScalar const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + GemmConfig_::kScalarsPerLdsB, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarB Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar; + + /// The traits class to build the iterator to load data from global memory for B^T. + typedef GemmGlobalTileTraits< + // That's B. + GemmOperand::kB, + // B is row-major. + MatrixLayout::kRowMajor, + // The pointer is float const. + Scalar const, + // The tile has size KxN in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kH>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + /// The traits class to build the iterator to store data to shared memory for B^T. + typedef GemmSharedStoreTileAbTraits< + // The pointer is float. + MultiplyAddScalar, + // The tile has size KxN in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kH * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsB> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for B^T. + typedef GemmSharedLoadTileBTraits< + // The pointer is float const. + MultiplyAddScalar const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + GemmConfig_::kScalarsPerLdsB, + // The skew. + 0> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmTraits_, bool kResidueInPrologue_ = GemmTraits_::kResidueInPrologue> +struct GemmResidue { + /// Move to residue portion. + template <bool kIsPrologue> + static CUTLASS_DEVICE void move_to_residue(typename GemmTraits_::GlobalLoadStreamA& stream_a, + typename GemmTraits_::GlobalLoadStreamB& stream_b, + typename GemmTraits_::Index k) { + // The new code path in CUTLASS 1.0.1: We treat the residue in the prologue so we can have + // complete main loops after that. It helps simplify the logic in the main loop. + if (kIsPrologue) { + stream_a.move_to_residue(k); + stream_b.move_to_residue(k); + } + } + + /// Rollback to beginning of first tile and initialize predicates. + static CUTLASS_DEVICE void rollback(typename GemmTraits_::GlobalLoadStreamA& stream_a, + typename GemmTraits_::GlobalLoadStreamB& stream_b) { + stream_a.rollback(); + stream_b.rollback(); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmTraits_> +struct GemmResidue<GemmTraits_, false> { + /// Move to residue portion. + template <bool kIsPrologue> + static CUTLASS_DEVICE void move_to_residue(typename GemmTraits_::GlobalLoadStreamA& stream_a, + typename GemmTraits_::GlobalLoadStreamB& stream_b, + typename GemmTraits_::Index k) { + // The index. + typedef typename GemmTraits_::Index Index; + // By how much we unroll the main loop. + Index const kUnroll = static_cast<Index>(GemmTraits_::OutputTile::kD); + + // Call the residue code. That's the same path as CUTLASS 1.0.0. + if (kIsPrologue && k < kUnroll) { + stream_a.residue(k, true); + stream_b.residue(k, true); + } else if (k <= kUnroll) { + stream_a.residue(k, false); + stream_b.residue(k, false); + } + } + + /// Rollback to beginning of first tile and initialize predicates. + static CUTLASS_DEVICE void rollback(typename GemmTraits_::GlobalLoadStreamA& stream_a, + typename GemmTraits_::GlobalLoadStreamB& stream_b) {} +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The GEMM configuration. + typename GemmConfig_, + /// The stream to load A from global memory to shared memory. + typename GlobalLoadStreamA_, + /// The stream to load B from global memory to shared memory. + typename GlobalLoadStreamB_, + /// The stream to load A from shared memory. + typename SharedLoadStreamA_, + /// The stream to load B from shared memory. + typename SharedLoadStreamB_, + /// The epilogue. + typename Epilogue_, + /// The block swizzle to reorganize the grid. + typename BlockSwizzle_ = IdentityBlockSwizzle, + /// The index. + typename Index_ = int, + /// The tool used to clear accumulators. + typename ClearAccumulators_ = ClearAccumulators<typename GemmConfig_::Accumulators::Scalar> > + +struct GemmTraits { + /// This class. + typedef GemmTraits<GemmConfig_, + GlobalLoadStreamA_, + GlobalLoadStreamB_, + SharedLoadStreamA_, + SharedLoadStreamB_, + Epilogue_, + BlockSwizzle_, + Index_, + ClearAccumulators_> + This_; + + /// The configuration. + typedef GemmConfig_ GemmConfig; + /// The output tile. + typedef typename GemmConfig::OutputTile OutputTile; + /// Is the residue treated in the prologue? + static bool const kResidueInPrologue = GemmConfig::kResidueInPrologue; + + /// The stream to load A from global memory to shared memory. + typedef GlobalLoadStreamA_ GlobalLoadStreamA; + /// The layout of A. + static MatrixLayout::Kind const kLayoutA = GlobalLoadStreamA::kLayout; + /// The scalar for A. + typedef typename GlobalLoadStreamA_::Scalar ScalarA; + + /// The stream to load B from global memory to shared memory. + typedef GlobalLoadStreamB_ GlobalLoadStreamB; + /// The layout of B. + static MatrixLayout::Kind const kLayoutB = GlobalLoadStreamB::kLayout; + /// The scalar for B. + typedef typename GlobalLoadStreamB_::Scalar ScalarB; + + /// The iterator for A to load from shared memory. + typedef SharedLoadStreamA_ SharedLoadStreamA; + /// The iterator for B to load from shared memory. + typedef SharedLoadStreamB_ SharedLoadStreamB; + + /// The multiply-add functor. + typedef typename GemmConfig::MultiplyAdd MultiplyAdd; + /// The epilogue. + typedef Epilogue_ Epilogue; + /// The scalars in the epilogue. + typedef typename Epilogue::ScalarC ScalarC; + typedef typename Epilogue::ScalarD ScalarD; + + /// The block swizzle to reorganize the grid. + typedef BlockSwizzle_ BlockSwizzle; + /// The index. + typedef Index_ Index; + /// Clear the accumulators. + typedef ClearAccumulators_ ClearAccumulators; + + /// The params. + struct Params { + /// The dimensions of the GEMM. + Index m, n, k; + /// The params for the A stream. + typename GlobalLoadStreamA::Params global_stream_a; + /// The params for the B stream. + typename GlobalLoadStreamB::Params global_stream_b; + /// The params for the A stream from shared memory. + typename SharedLoadStreamA::Params shared_stream_a; + /// The params for the B stream from shared memory. + typename SharedLoadStreamB::Params shared_stream_b; + /// The params for the epilogue. + typename Epilogue::Params epilogue; + + /// Initialize the parameters. + template <typename GemmDesc_> + CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc) { + // Set the problem size. + this->m = desc.m; + this->n = desc.n; + this->k = desc.k; + + // Initialize the iterator for A. + int error_code = + global_stream_a.initialize(desc, reinterpret_cast<ScalarA const*>(desc.d_a), desc.lda); + + if (error_code) { + return error_code; + } + + // Initialize the iterator for B. + error_code = + global_stream_b.initialize(desc, reinterpret_cast<ScalarB const*>(desc.d_b), desc.ldb); + + if (error_code) { + return error_code; + } + + // The epilogue. + return epilogue.initialize(desc); + } + }; + + // The storage for A. + template <typename GlobalLoadStream_, typename SharedLoadStream_> + union StreamSharedStorage { + // The storage needed by the global stream. + typename GlobalLoadStream_::SharedStorage global; + // The storage needed by the shared stream. + typename SharedLoadStream_::SharedStorage shared; + }; + + // The storage for the main loop + prologue. + struct MainLoopSharedStorage { + // The storage to shuffle the A matrix in shared memory. + StreamSharedStorage<GlobalLoadStreamA, SharedLoadStreamA> stream_a; + // The storage to shuffle the B matrix in shared memory. + StreamSharedStorage<GlobalLoadStreamB, SharedLoadStreamB> stream_b; + // The storage to clear the accumulators if needed. + typename ClearAccumulators::SharedStorage clear; + }; + + /// The storage in shared memory. + union SharedStorage { + // The storage for the main loop. + MainLoopSharedStorage main_loop; + // The storage for the epilogue. + typename Epilogue::SharedStorage epilogue; + }; + + /// Assemble the global load streams for A/B. + struct GlobalLoadStream { + /// Ctor. + CUTLASS_DEVICE GlobalLoadStream(Params const& params, + SharedStorage& shared_storage, + dim3 const& block) + : stream_a(params.global_stream_a, + shared_storage.main_loop.stream_a.global, + cutlass::make_Coord(0, params.k, params.m), + cutlass::make_Coord(0, 0, block.x)), + stream_b(params.global_stream_b, + shared_storage.main_loop.stream_b.global, + cutlass::make_Coord(0, params.k, params.n), + make_Coord(0, 0, block.y)) {} + + /// Trigger the copies from shared memory to registers. + CUTLASS_DEVICE void copy() { + stream_a.copy(); + stream_b.copy(); + } + + /// Commit the data. + CUTLASS_DEVICE void commit() { + stream_a.commit(); + stream_b.commit(); + } + + /// Move to residue portion. + template <bool kIsPrologue> + CUTLASS_DEVICE void move_to_residue(Index k) { + GemmResidue<This_>::move_to_residue<kIsPrologue>(stream_a, stream_b, k); + } + + /// Rollback to beginning of first tile and initialize predicates. + CUTLASS_DEVICE void rollback() { GemmResidue<This_>::rollback(stream_a, stream_b); } + + /// The stream for A. + GlobalLoadStreamA stream_a; + /// The stream for B. + GlobalLoadStreamB stream_b; + }; + + /// Assemble the shared load stream for A/B. + struct SharedLoadStream { + /// Ctor. + CUTLASS_DEVICE SharedLoadStream(Params const& params, SharedStorage& shared_storage) { + stream_a.initialize(params.shared_stream_a, shared_storage.main_loop.stream_a.shared); + stream_b.initialize(params.shared_stream_b, shared_storage.main_loop.stream_b.shared); + } + + /// Trigger the copies from shared memory to registers. + CUTLASS_DEVICE void copy(int step) { + stream_a.copy(step, fetched_a[step % 2]); + stream_b.copy(step, fetched_b[step % 2]); + } + + /// Commit the data. + CUTLASS_DEVICE void commit(int step) { + stream_a.commit(fetched_a[step % 2], transformed_a[step % 2]); + stream_b.commit(fetched_b[step % 2], transformed_b[step % 2]); + } + + /// The fragment A. + CUTLASS_DEVICE typename SharedLoadStreamA::Fragment const& fragment_a(int step) const { + return transformed_a[step % 2]; + } + + /// The fragment B. + CUTLASS_DEVICE typename SharedLoadStreamB::Fragment const& fragment_b(int step) const { + return transformed_b[step % 2]; + } + + /// Increment the stage. + CUTLASS_DEVICE void inc_stage() { + stream_a.inc_stage(); + stream_b.inc_stage(); + } + + /// The stream for A. + SharedLoadStreamA stream_a; + /// The fragments to fetch A. + typename SharedLoadStreamA::FetchedFragment fetched_a[2]; + /// The fragments to transform A. + typename SharedLoadStreamA::TransformedFragment transformed_a[2]; + /// The stream for B. + SharedLoadStreamB stream_b; + /// The fragments to fetch B. + typename SharedLoadStreamB::FetchedFragment fetched_b[2]; + /// The fragments to transform B. + typename SharedLoadStreamB::TransformedFragment transformed_b[2]; + }; + + /// The memory fence for shared loads. + static CUTLASS_DEVICE void shared_load_fence(bool in_loop) { + if (SharedLoadStreamA::Iterator::kRequiresLoadFence || + SharedLoadStreamB::Iterator::kRequiresLoadFence) { + __syncthreads(); + } + } + + /// The memory fence for shared stores. + static CUTLASS_DEVICE void shared_store_fence(bool in_loop) { __syncthreads(); } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmTileTraitsHelperA_, typename GemmTileTraitsHelperB_, typename Index_> +struct SimplifiedGemmTraitsHelper { + /// The global iterator to load A from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA_::GlobalTileTraits, Index_> + GlobalLoadIteratorA; + /// The data converter for A before storing to shared memory. + typedef Copy<typename GlobalLoadIteratorA::Fragment> GlobalTransformerA; + /// The iterator to store A to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperA_::SharedStoreTileTraits, + typename GemmTileTraitsHelperA_::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorA; + /// The stream to load A from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA> + GlobalLoadStreamA; + + /// The global iterator to load B from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperB_::GlobalTileTraits, Index_> + GlobalLoadIteratorB; + /// The data converter for B before storing to shared memory. + typedef Copy<typename GlobalLoadIteratorB::Fragment> GlobalTransformerB; + /// The iterator to store B to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperB_::SharedStoreTileTraits, + typename GemmTileTraitsHelperB_::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorB; + /// The stream to load B from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB> + GlobalLoadStreamB; + + /// The iterator to load A from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperA_::SharedLoadTileTraits, + typename GemmTileTraitsHelperA_::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorA; + /// The stream to load A from shared memory. + typedef SharedLoadStream<SharedLoadIteratorA> SharedLoadStreamA; + /// The iterator to load B from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperB_::SharedLoadTileTraits, + typename GemmTileTraitsHelperB_::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorB; + /// The stream to load B from shared memory. + typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The config for the GEMM. + typename GemmConfig_, + /// The epilogue. + typename Epilogue_, + /// The index. + typename Index_ = int, + // The configuration for the A matrix. + typename GemmTileTraitsHelperA_ = GemmTileTraitsHelperA<kLayoutA_, GemmConfig_>, + // The configuration for the B matrix. + typename GemmTileTraitsHelperB_ = GemmTileTraitsHelperB<kLayoutB_, GemmConfig_>, + // The helper class to create the streams and iterators. + typename Helper_ = + SimplifiedGemmTraitsHelper<GemmTileTraitsHelperA_, GemmTileTraitsHelperB_, Index_> > +struct SimplifiedGemmTraits : public GemmTraits< + // The config. + GemmConfig_, + // The stream to load A from global memory to shared memory. + typename Helper_::GlobalLoadStreamA, + // The stream to load B from global memory to shared memory. + typename Helper_::GlobalLoadStreamB, + // The stream to load A from shared memory. + typename Helper_::SharedLoadStreamA, + // The stream to load B from shared memory. + typename Helper_::SharedLoadStreamB, + // The epilogue. + Epilogue_, + // The block swizzle to reorganize the grid. + IdentityBlockSwizzle, + // The index. + Index_, + // The tool used to clear accumulators. + ClearAccumulators<typename GemmConfig_::Accumulators::Element> > { +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/hgemm_global_tile.h b/cutlass-example/cutlass/gemm/hgemm_global_tile.h new file mode 100644 index 0000000..f14dbb3 --- /dev/null +++ b/cutlass-example/cutlass/gemm/hgemm_global_tile.h @@ -0,0 +1,90 @@ +/*************************************************************************************************** + * 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 Tile traits used to construct global tile iterator for HGEMM. This is intended to + partition the thread block-level tile into 2D subtiles loaded by the threads and facilitate + memory accesses larger than 16 bits. +*/ +#pragma once + +#include <cutlass/coord.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/matrix_traits.h> +#include <cutlass/reshape_tile.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct HgemmCrosswiseGlobalTileTraits : public GemmGlobalTileTraits< + // Which GEMM operand? + kOperand_, + // The layout. + kLayout_, + // The scalar. + Scalar_, + // The tile. + Tile_, + // The threads. + Threads_, + // The number of scalars per LDG/STG. + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<kOperand_, kLayout_, Scalar_, Tile_, Threads_, kAccessSize_> Base; + /// The threads. + typedef typename Base::Threads Threads; + /// The threads strides. + typedef Shape<1, 2, Base::Tile::kC> ThreadsDelta; + /// The strides in each dimension between different loads/stores. + typedef Shape<Base::Threads::kH * 2, 1, Base::Threads::kW, Base::kAccessSize> Delta; + /// The number of iterations needed to load/store the tile. + typedef Shape<Base::Tile::kH / Base::Threads::kH / 2, + 2, + Base::Tile::kW / Base::Threads::kW, + Base::Tile::kC / Base::kAccessSize> + Iterations; + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Threads::kW * ThreadsDelta::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/hgemm_multiply_add.h b/cutlass-example/cutlass/gemm/hgemm_multiply_add.h new file mode 100644 index 0000000..ebbdd06 --- /dev/null +++ b/cutlass-example/cutlass/gemm/hgemm_multiply_add.h @@ -0,0 +1,104 @@ +/*************************************************************************************************** + * 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 Specialization implementing multiply-add operation on half-precision floating point + fragments. +*/ +#pragma once + +#include <cutlass/fragment.h> + +#include <cutlass/gemm/thread_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template performing matrix multiply-add operation within a thread +template <typename AccumulatorsPerThread_, typename ThreadsPerWarp_> +struct ThreadMultiplyAdd<AccumulatorsPerThread_, ThreadsPerWarp_, half, half, half> { + /// The shape of the instruction. + typedef Shape<1, 1, 2, 1> InstructionShape; + /// The number of accumulators per thread. + typedef AccumulatorsPerThread_ AccumulatorsPerThread; + /// The number of threads per warp. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of accumulators per warp. + typedef typename ShapeMul<AccumulatorsPerThread, ThreadsPerWarp>::Shape AccumulatorsPerWarp; + /// The type for A. + typedef half ScalarA; + /// The fragment for A. + typedef Fragment<ScalarA, AccumulatorsPerThread::kW> FragmentA; + /// The type for B. + typedef half ScalarB; + /// The fragment for B. + typedef Fragment<ScalarB, AccumulatorsPerThread::kH> FragmentB; + /// The type for C and D. + typedef half ScalarC; + /// The accumulators. + typedef Fragment<half, AccumulatorsPerThread::kH * AccumulatorsPerThread::kW> Accumulators; + + /// Make sure there's an even number of elements in both dimensions. + static_assert(AccumulatorsPerThread::kH % 2 == 0, "Invalid size"); + static_assert(AccumulatorsPerThread::kW % 2 == 0, "Invalid size"); + + /// Ctor. + CUTLASS_DEVICE ThreadMultiplyAdd() {} + + /// Multiply : d = a*b + c. + CUTLASS_DEVICE void multiply_add(FragmentA const& a, + FragmentB const& b, + Accumulators const& c, + Accumulators& d) { +#if defined(__CUDACC__) && __CUDA_ARCH__ >= 530 + // The inputs. + __half2 const* a_half2 = reinterpret_cast<__half2 const*>(&a[0]); + __half2 const* b_half2 = reinterpret_cast<__half2 const*>(&b[0]); + __half2 const* c_half2 = reinterpret_cast<__half2 const*>(&c[0]); + + // The output. + __half2* d_half2 = reinterpret_cast<__half2*>(&d[0]); + + for (int j = 0; j < AccumulatorsPerThread::kH / 2; ++j) { + for (int i = 0; i < AccumulatorsPerThread::kW / 2; ++i) { + // The offsets in the output fragment. + int const k0 = (2 * j + 0) * (AccumulatorsPerThread::kW / 2) + i; + int const k1 = (2 * j + 1) * (AccumulatorsPerThread::kW / 2) + i; + + // Compute the product a[i] * b[j].H0_H0. + d_half2[k0] = __hfma2(a_half2[i], __low2half2(b_half2[j]), c_half2[k0]); + // Compute the product a[i] * b[j].H1_H1. + d_half2[k1] = __hfma2(a_half2[i], __high2half2(b_half2[j]), c_half2[k1]); + } + } +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/hgemm_swizzle.h b/cutlass-example/cutlass/gemm/hgemm_swizzle.h new file mode 100644 index 0000000..ebec0d4 --- /dev/null +++ b/cutlass-example/cutlass/gemm/hgemm_swizzle.h @@ -0,0 +1,94 @@ +/*************************************************************************************************** + * 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 Transposes a tile of 16b elements. Used by HGEMM to construct a K-strided layout in + shared memory for multiplicands. +*/ +#pragma once + +#include <cuda_fp16.h> +#include <cutlass/fragment.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GlobalIterator_> +struct HgemmSwizzle { + /// The global iterator. + typedef GlobalIterator_ GlobalIterator; + /// The source fragment. + typedef typename GlobalIterator::Fragment Fragment; + /// The shape of the source fragment. + typedef typename GlobalIterator::FragmentShape FragmentShape; + + /// The input fragment. + typedef Fragment InputFragment; + /// The output fragment. + typedef Fragment OutputFragment; + + /// The src/dst must be half fragments. + static_assert((platform::is_same<typename Fragment::Element, half>::value), "Works on half"); + + /// The number of elements must be a multiple of 2. + static_assert(FragmentShape::kH == 2 && ShapeCount<FragmentShape>::kWc == 2, "Not multiple of 2"); + + /// Ctor. + CUTLASS_DEVICE HgemmSwizzle() {} + + /// Transform a fragment. + CUTLASS_DEVICE void transform(Fragment const& src, Fragment& dst) { + // Expose src/dst as int arrays. + int const* src_int = reinterpret_cast<int const*>(&src[0]); + int* dst_int = reinterpret_cast<int*>(&dst[0]); + + // Transpose the data. + for (int d = 0; d < FragmentShape::kD; ++d) { + // The indices to read two consecutive "rows". + int const i0 = 2 * d + 0; + int const i1 = 2 * d + 1; + + int a0 = src_int[i0]; + int a1 = src_int[i1]; + + int b0, b1; + asm volatile("prmt.b32 %0, %1, %2, 0x5410;" : "=r"(b0) : "r"(a0), "r"(a1)); + asm volatile("prmt.b32 %0, %1, %2, 0x7632;" : "=r"(b1) : "r"(a0), "r"(a1)); + + // The indices to store with "strides". + int const j0 = 0 * (ShapeCount<FragmentShape>::kDhw / 2) + d; + int const j1 = 1 * (ShapeCount<FragmentShape>::kDhw / 2) + d; + + dst_int[j0] = b0; + dst_int[j1] = b1; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/hgemm_traits.h b/cutlass-example/cutlass/gemm/hgemm_traits.h new file mode 100644 index 0000000..b08645b --- /dev/null +++ b/cutlass-example/cutlass/gemm/hgemm_traits.h @@ -0,0 +1,397 @@ +/*************************************************************************************************** + * 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 Defies structural properties of half-precision GEMM computation. +*/ +#pragma once + +#include <cutlass/convert.h> +#include <cutlass/reshape_tile.h> + +#include <cutlass/gemm/gemm.h> +#include <cutlass/gemm/gemm_epilogue.h> +#include <cutlass/gemm/gemm_epilogue_traits.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/gemm/gemm_shared_tile.h> +#include <cutlass/gemm/gemm_traits.h> +#include <cutlass/gemm/hgemm_global_tile.h> +#include <cutlass/gemm/hgemm_multiply_add.h> +#include <cutlass/gemm/hgemm_swizzle.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The tile size for the GEMM KxNxM. + typename OutputTile_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_ = 2, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_ = 2> +struct HgemmConfig + : public GemmConfig< + /// The scalar type for A. + half, + /// The scalar type for B. + half, + /// The scalar type for C. + half, + /// The scalar type for D. + half, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, half, half, half>, + /// The number of scalars per LDG for A. + kScalarsPerLdgA_, + /// The number of scalars per STS for A. + kScalarsPerLdgA_, + /// The number of scalars per LDS for A. + 8, + /// The number of scalars per LDG for B. + kScalarsPerLdgB_, + /// The number of scalars per STS for B. + kScalarsPerLdgB_, + /// The number of scalars per LDS for B. + 8, + /// The number of scalars per LDG for C and STG for D. + 2, + /// The number of scalars per STS for D. + 8, + /// The number of scalars per LDS for D. + 2, + /// The number of stages in shared memory. + 2> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename Iterator_> +struct HgemmTransformerA {}; + +template <typename Iterator_> +struct HgemmTransformerA<MatrixLayout::kColumnMajor, Iterator_> { + typedef Convert<typename Iterator_::Fragment, typename Iterator_::Fragment> Transformer; +}; + +template <typename Iterator_> +struct HgemmTransformerA<MatrixLayout::kRowMajor, Iterator_> { + typedef HgemmSwizzle<Iterator_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename Iterator_> +struct HgemmTransformerB {}; + +template <typename Iterator_> +struct HgemmTransformerB<MatrixLayout::kRowMajor, Iterator_> { + typedef Convert<typename Iterator_::Fragment, typename Iterator_::Fragment> Transformer; +}; + +template <typename Iterator_> +struct HgemmTransformerB<MatrixLayout::kColumnMajor, Iterator_> { + typedef HgemmSwizzle<Iterator_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_> +struct HgemmTileTraitsHelperA : public GemmTileTraitsHelperA<kLayout_, GemmConfig_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct HgemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> + : public GemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> Base; + + /// The traits class to build the iterator to load data from global memory for A^T. + typedef HgemmCrosswiseGlobalTileTraits< + GemmOperand::kA, + // The layout. + MatrixLayout::kRowMajor, + // The pointer. + half const, + // The tile has size MxK in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD>, + // The threads are distributed as (threads / K ) x K (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc) + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + /// The skew. + static int const kSkewA = 128 / sizeof(half) / GlobalTileTraits::Threads::kW / 2; + + /// The traits class to build the iterator to store data to shared memory for A^T. + typedef GemmSharedStoreWithSkewTileAbTraits < + // The pointer. + half, + // The tile has size KxM in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kW * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as warps x 32(the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + 2, + // The skew to avoid bank conflicts added in the tile W dimension. + kSkewA<GemmConfig_::kScalarsPerLdsA ? GemmConfig_::kScalarsPerLdsA : kSkewA> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for A^T. + typedef GemmSharedLoadTileATraits< + // The pointer. + half const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + 8, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_> +struct HgemmTileTraitsHelperB : public GemmTileTraitsHelperB<kLayout_, GemmConfig_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct HgemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> + : public GemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> Base; + + /// The traits class to build the iterator to load data from global memory for B^N. + typedef HgemmCrosswiseGlobalTileTraits< + GemmOperand::kB, + // The layout. + MatrixLayout::kColumnMajor, + // The pointer. + half const, + // The tile has size KxN in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc) + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + /// The skew for B. + static int const kSkewB = 128 / sizeof(half) / GlobalTileTraits::Threads::kW / 2; + + /// The traits class to build the iterator to store data to shared memory for B^N. + typedef GemmSharedStoreWithSkewTileAbTraits < + // The pointer. + half, + // The tile has size KxN in GEMM's terminology. + Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD, + GemmConfig_::OutputTile::kH * GemmConfig_::InstructionShape::kD>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + 2, + // The skew to avoid bank conflicts added in the tile W dimension. + kSkewB<GemmConfig_::kScalarsPerLdsB ? GemmConfig_::kScalarsPerLdsB : kSkewB> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for B^N. + typedef GemmSharedLoadTileBTraits< + // The pointer. + half const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + 8, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<8, 8, 16>, + /// The number of halfs loaded in one LDG for A. + int kScalarsPerLdgA_ = 2, + /// The number of halfs loaded in one LDG for B. + int kScalarsPerLdgB_ = 2, + /// The index. + typename Index_ = int> +struct HgemmTraitsHelper { + /// The HGEMM config. + typedef HgemmConfig<OutputTile_, AccumulatorsPerThread_, kScalarsPerLdgA_, kScalarsPerLdgB_> + GemmConfig; + /// The GEMM config for A. + typedef HgemmTileTraitsHelperA<kLayoutA_, GemmConfig> GemmTileTraitsHelperA; + /// The GEMM config for B. + typedef HgemmTileTraitsHelperB<kLayoutB_, GemmConfig> GemmTileTraitsHelperB; + + /// The iterator to load A from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA::GlobalTileTraits, Index_> + GlobalLoadIteratorA; + /// The default transformer for A. + typedef typename HgemmTransformerA<GemmTileTraitsHelperA::kLayout, + GlobalLoadIteratorA>::Transformer GlobalTransformerA; + /// The iterator to store A to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperA::SharedStoreTileTraits, + typename GemmTileTraitsHelperA::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorA; + /// The stream to load A from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA> + GlobalLoadStreamA; + + /// The iterator to load B from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperB::GlobalTileTraits, Index_> + GlobalLoadIteratorB; + // The default transformer for B. + typedef typename HgemmTransformerB<GemmTileTraitsHelperB::kLayout, + GlobalLoadIteratorB>::Transformer GlobalTransformerB; + /// The iterator to store B to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperB::SharedStoreTileTraits, + typename GemmTileTraitsHelperB::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorB; + /// The stream to load B from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB> + GlobalLoadStreamB; + + /// The iterator to load A from shared memory + typedef TileLoadIterator<typename GemmTileTraitsHelperA::SharedLoadTileTraits, + typename GemmTileTraitsHelperA::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorA; + /// The stream to load A from shared memory. + typedef SharedLoadStream<SharedLoadIteratorA> SharedLoadStreamA; + /// The iterator to load B from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperB::SharedLoadTileTraits, + typename GemmTileTraitsHelperB::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorB; + /// The stream to load B from shared memory. + typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB; + + /// The functor to do the multiply-add in the main loop. + typedef typename GemmConfig::MultiplyAdd MultiplyAdd; + /// The object to clear accumulators. + typedef ClearAccumulators<typename MultiplyAdd::ScalarC> ClearAccumulators; + + /// The traits class for the epilogue. + typedef SimplifiedGemmEpilogueTraits<GemmConfig, EpilogueFunctor_, Index_> GemmEpilogueTraits; + /// The epilogue. + typedef GemmEpilogue<GemmEpilogueTraits> Epilogue; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_ = Shape<8, 128, 128>, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_ = LinearScaling<half>, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<8, 8, 16>, + /// The number of halfs loaded in one LDG for A. + int kScalarsPerLdgA_ = 2, + /// The number of halfs loaded in one LDG for B. + int kScalarsPerLdgB_ = 2, + /// The index. + typename Index_ = int, + /// The helper class. + typename Helper_ = HgemmTraitsHelper<kLayoutA_, + kLayoutB_, + OutputTile_, + EpilogueFunctor_, + AccumulatorsPerThread_, + kScalarsPerLdgA_, + kScalarsPerLdgB_, + Index_> > +struct HgemmTraits : public GemmTraits< + // The config. + typename Helper_::GemmConfig, + // The stream to load A from global memory to shared memory. + typename Helper_::GlobalLoadStreamA, + // The stream to load B from global memory to shared memory. + typename Helper_::GlobalLoadStreamB, + // The stream to load A from shared memory. + typename Helper_::SharedLoadStreamA, + // The stream to load B from shared memory. + typename Helper_::SharedLoadStreamB, + // The epilogue. + typename Helper_::Epilogue, + // The block swizzle to reorganize the grid. + IdentityBlockSwizzle, + // The index. + Index_, + // The tool used to clear accumulators. + typename Helper_::ClearAccumulators> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/identity_block_swizzle.h b/cutlass-example/cutlass/gemm/identity_block_swizzle.h new file mode 100644 index 0000000..e1bdb2e --- /dev/null +++ b/cutlass-example/cutlass/gemm/identity_block_swizzle.h @@ -0,0 +1,48 @@ +/*************************************************************************************************** + * 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 Defies functors for mapping blockIdx to partitions of the GEMM computation. + + Currently, we only implement an identity mapping. +*/ +#pragma once + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct IdentityBlockSwizzle { + /// Ctor. + CUTLASS_DEVICE IdentityBlockSwizzle() {} + + /// Swizzle the block index. + CUTLASS_DEVICE dim3 swizzle() { return blockIdx; } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/igemm_epilogue.h b/cutlass-example/cutlass/gemm/igemm_epilogue.h new file mode 100644 index 0000000..0d69980 --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_epilogue.h @@ -0,0 +1,320 @@ +/*************************************************************************************************** + * 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 Defines the epilogue phase of the GEMM computation for IGEMM, supporting integer and + floating-point output matrix formats. +*/ +#pragma once + +#include <cutlass/convert.h> +#include <cutlass/fragment.h> +#include <cutlass/gemm/gemm_global_stream.h> +#include <cutlass/gemm/gemm_shared_stream.h> +#include <cutlass/gemm/igemm_global_tile.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/tile_iterator.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <int kElements_> +struct IgemmFloatToInt8Converter { + /// The input fragment. + typedef Fragment<float, kElements_> InputFragment; + /// The output fragment. + typedef Fragment<int8_t, kElements_> OutputFragment; + + // We are packing 4 floats into int32 registers so we need kElements to be multiple of 4. + static_assert(kElements_ % 4 == 0, "kElements must be multiple of 4"); + + /// Ctor. + CUTLASS_DEVICE IgemmFloatToInt8Converter() {} + + /// Transform a fragment. + CUTLASS_DEVICE void transform(InputFragment const& src, OutputFragment& dst) { + transform(src, 0, dst); + } + + /// Transform a fragment. + template <typename Fragment_> + CUTLASS_DEVICE void transform(Fragment_ const& src, int offset, OutputFragment& dst) { + // The inputs. + float4 const* src_f4 = reinterpret_cast<float4 const*>(&src[0]); + // The outputs. + int* dst_int = reinterpret_cast<int*>(&dst[0]); + + // Iterate over the floats and pack them together to produce ints. + for (int i = 0; i < kElements_ / 4; ++i) { + // Read the float4. + float4 f4 = src_f4[i]; + + // Clamp the 4 elements of the floats to the [-128, +127] range. + float x = fmaxf(-128.f, fminf(127.f, f4.x)); + float y = fmaxf(-128.f, fminf(127.f, f4.y)); + float z = fmaxf(-128.f, fminf(127.f, f4.z)); + float w = fmaxf(-128.f, fminf(127.f, f4.w)); + + // Convert to integers. + int ix = (int)x; + int iy = (int)y; + int iz = (int)z; + int iw = (int)w; + + // Extract the lower bytes to build an int32 with 4 int8. + asm volatile("prmt.b32 %0, %0, %1, 0x1140;" : "+r"(ix) : "r"(iy)); + asm volatile("prmt.b32 %0, %0, %1, 0x1140;" : "+r"(iz) : "r"(iw)); + asm volatile("prmt.b32 %0, %0, %1, 0x5410;" : "+r"(ix) : "r"(iz)); + + // Store the int. + dst_int[i] = ix; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename InputScalar_, typename OutputFragment_> +struct IgemmGlobalStoreTransformer { + typedef Convert<Fragment<InputScalar_, OutputFragment_::kElements>, OutputFragment_> Transformer; +}; + +template <int kElements_> +struct IgemmGlobalStoreTransformer<float, Fragment<int8_t, kElements_> > { + typedef IgemmFloatToInt8Converter<kElements_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <int kElements_> +struct IgemmInt8ToFloatConverter { + /// The input fragment. + typedef Fragment<int8_t, kElements_> InputFragment; + /// The output fragment. + typedef Fragment<float, kElements_> OutputFragment; + + // We are unpacking 4 int8s from int32. + static_assert(kElements_ % 4 == 0, "kElements must be multiple of 4"); + + /// Ctor. + CUTLASS_DEVICE IgemmInt8ToFloatConverter() {} + + /// Transform a fragment. + CUTLASS_DEVICE void transform(InputFragment const& src, OutputFragment& dst) { + transform(src, 0, dst); + } + + /// Transform a fragment. + template <typename Fragment_> + CUTLASS_DEVICE void transform(Fragment_ const& src, int offset, OutputFragment& dst) { + // The inputs. + int const* src_int = reinterpret_cast<int const*>(&src[0]); + // The outputs. + float4* dst_f4 = reinterpret_cast<float4*>(&dst[0]); + + // Iterate over the int8 and unpack them together to produce floats. + for (int i = 0; i < kElements_ / 4; ++i) { + // Read the int. + int ix, iy, iz, iw = src_int[i]; + + // Extract the 4 bytes. + asm volatile("prmt.b32 %0, 0x0, %1, 0x4440;" : "=r"(ix) : "r"(iw)); + asm volatile("prmt.b32 %0, 0x0, %1, 0x4441;" : "=r"(iy) : "r"(iw)); + asm volatile("prmt.b32 %0, 0x0, %1, 0x4442;" : "=r"(iz) : "r"(iw)); + asm volatile("prmt.b32 %0, 0x0, %1, 0x4443;" : "=r"(iw) : "r"(iw)); + + // The floats. + float fx, fy, fz, fw; + + // Convert to floats (make sure we generate I2F.F32.S8). + asm volatile("cvt.rn.f32.s8 %0, %1;" : "=f"(fx) : "r"(ix)); + asm volatile("cvt.rn.f32.s8 %0, %1;" : "=f"(fy) : "r"(iy)); + asm volatile("cvt.rn.f32.s8 %0, %1;" : "=f"(fz) : "r"(iz)); + asm volatile("cvt.rn.f32.s8 %0, %1;" : "=f"(fw) : "r"(iw)); + + // Store the float4. + dst_f4[i] = make_float4(fx, fy, fz, fw); + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename InputFragment_, typename OutputScalar_> +struct IgemmGlobalLoadTransformer { + typedef Convert<InputFragment_, Fragment<OutputScalar_, InputFragment_::kElements> > Transformer; +}; + +template <int kElements_> +struct IgemmGlobalLoadTransformer<Fragment<int8_t, kElements_>, float> { + typedef IgemmInt8ToFloatConverter<kElements_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename InputScalar_, typename OutputFragment_> +struct IgemmSharedStoreTransformer { + typedef Convert<Fragment<InputScalar_, OutputFragment_::kElements>, OutputFragment_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename IgemmConfig_, typename EpilogueFunctor_, typename Index_> +struct IgemmEpilogueTraitsHelper + : public GemmEpilogueTraitsHelper<IgemmConfig_, EpilogueFunctor_, Index_> { + /// The base class. + typedef GemmEpilogueTraitsHelper<IgemmConfig_, EpilogueFunctor_, Index_> Base; + /// The config. + typedef IgemmConfig_ IgemmConfig; + + /// The scalar type of the epilogue. + typedef typename Base::Scalar Scalar; + /// The iterations. + typedef typename Base::Iterations Iterations; + /// The iterations strides. + typedef typename Base::Delta Delta; + + /// The traits class for the iterator. + typedef typename Base::GlobalLoadTileTraits GlobalLoadTileTraits; + /// The iterator to store to shared memory. + typedef GemmGlobalIteratorCd<GlobalLoadTileTraits> GlobalLoadIteratorC; + /// The fragment that needs to be produced by the load iterator. + typedef typename GlobalLoadIteratorC::Fragment GlobalFragmentC; + /// The transformer from loaded data to math fragment. + typedef + typename IgemmGlobalLoadTransformer<GlobalFragmentC, Scalar>::Transformer GlobalTransformerC; + + /// The traits class for the iterator. + typedef typename Base::GlobalStoreTileTraits GlobalStoreTileTraits; + /// The iterator to store to shared memory. + typedef GemmGlobalIteratorCd<GlobalStoreTileTraits> GlobalStoreIteratorD; + /// The fragment that needs to be passed to that store iterator. + typedef typename GlobalStoreIteratorD::Fragment GlobalFragmentD; + /// The transformer from accumulators to shared memory fragments. + typedef + typename IgemmGlobalStoreTransformer<Scalar, GlobalFragmentD>::Transformer GlobalTransformerD; + + /// The traits class for the shared iterator to store D to shared memory. + typedef typename Base::SharedStoreTileTraits SharedStoreTileTraits; + /// The shared iterator to store D to shared memory. + typedef TileStoreIterator<SharedStoreTileTraits, + typename SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal> + SharedStoreIteratorD; + /// The fragment that needs to be passed to that store iterator. + typedef typename SharedStoreIteratorD::Fragment SharedStoreFragmentD; + /// The transformer from accumulators to shared memory fragments. + typedef typename IgemmSharedStoreTransformer<typename IgemmConfig::Accumulators::Element, + SharedStoreFragmentD>::Transformer + SharedStoreTransformerD; + /// The traits class for the shared iterator to load D from shared memory. + typedef typename Base::SharedLoadTileTraits SharedLoadTileTraits; + /// The shared iterator to load D from shared memory. + typedef TileLoadIterator<SharedLoadTileTraits, + typename SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorD; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The config. + typename IgemmConfig_, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_, + /// The index. + typename Index_ = int, + /// The helper class to assemble the traits. + typename Helper_ = IgemmEpilogueTraitsHelper<IgemmConfig_, EpilogueFunctor_, Index_> > +struct IgemmEpilogueTraits : public GemmEpilogueTraits< + // The output tile. + typename IgemmConfig_::OutputTile, + // The accumulators. + typename IgemmConfig_::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 iterations. + typename Helper_::Iterations, + // The strides between iterations. + typename Helper_::Delta, + // The functor to be used in the epilogue. + EpilogueFunctor_, + // The index. + Index_> { + /// Do we output in int8? + static bool const kInt8Output = + platform::is_same<typename IgemmConfig_::ScalarC, int8_t>::value != 0; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmEpilogueTraits_, bool = GemmEpilogueTraits_::kInt8Output> +struct IgemmEpilogue : public GemmEpilogue<GemmEpilogueTraits_> { + /// The base class. + typedef GemmEpilogue<GemmEpilogueTraits_> Base; + + /// Ctor. + CUTLASS_DEVICE IgemmEpilogue(typename Base::Params const& params_, + typename Base::SharedStorage& shared_storage_, + typename Base::Index m_, + typename Base::Index n_) + : Base(params_, shared_storage_, m_, n_) {} +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmEpilogueTraits_> +struct IgemmEpilogue<GemmEpilogueTraits_, true> : public GemmEpilogue<GemmEpilogueTraits_> { + /// The base class. + typedef GemmEpilogue<GemmEpilogueTraits_> Base; + + /// Ctor. + CUTLASS_DEVICE IgemmEpilogue(typename Base::Params const& params_, + typename Base::SharedStorage& shared_storage_, + typename Base::Index m_, + typename Base::Index n_) + : Base(params_, shared_storage_, m_, n_) {} +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/igemm_global_tile.h b/cutlass-example/cutlass/gemm/igemm_global_tile.h new file mode 100644 index 0000000..3f594ac --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_global_tile.h @@ -0,0 +1,161 @@ +/*************************************************************************************************** + * 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 tile iterators to partition the thread block tile into 2D subtiles and + efficiently load each. Applies permute transformation to construct 'interleaved K-strided' + data layout in which 4-element dot products from the same K index are arranged in consecutive + locations within shared memory. + + Supports efficient loads from shared memory to target the DP4A instruction. +*/ +#pragma once + +#include <cutlass/coord.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/matrix_traits.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct IgemmGlobalTileTraits : public GemmGlobalTileTraits< + // Which GEMM operand? + kOperand_, + // The layout. + kLayout_, + // The scalar. + Scalar_, + // The tile. + Tile_, + // The threads. + Threads_, + // The number of scalars per LDG/STG. + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<kOperand_, kLayout_, Scalar_, Tile_, Threads_, kAccessSize_> Base; + /// The threads. + typedef typename Base::Threads Threads; + /// The strides in each dimension between different loads/stores. + typedef Shape<Base::Threads::kH * 4, 1, Base::Threads::kW, Base::kAccessSize> Delta; + /// The number of iterations needed to load/store the tile. + typedef Shape<Base::Tile::kH / Base::Threads::kH / 4, + 4, + Base::Tile::kW / Base::Threads::kW, + Base::Tile::kC / Base::kAccessSize> + Iterations; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Threads::kW * ThreadsDelta::kH; + int thread_offset_w = threadIdx.x % Threads::kW * ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; + + public: + /// The threads strides. + typedef Shape<1, 4, Base::Tile::kC> ThreadsDelta; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Deprecated. Please use IgemmGlobalTileTraits instead. + +template <GemmOperand::Kind kOperand_, + MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Threads_, + int kAccessSize_> +struct IgemmContiguousGlobalTileTraits + : public IgemmGlobalTileTraits<kOperand_, kLayout_, Scalar_, Tile_, Threads_, kAccessSize_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct IgemmGlobalIteratorAb : public GemmGlobalIteratorAb<TileTraits_, Index_> { + /// The base class. + typedef GemmGlobalIteratorAb<TileTraits_, Index_> Base; + /// The functor to compute the thread offset. + typedef typename TileTraits_::ThreadOffset ThreadOffset; + + /// Constructor. + CUTLASS_DEVICE IgemmGlobalIteratorAb(typename Base::Params const& _params, + const Coord<3>& bounds, + const Coord<3>& block, + ThreadOffset thread_offset_func = ThreadOffset()) + : Base(_params, bounds, block, thread_offset_func), in_residue_(false), mask_(0xffffffff) { + // The number of elements read in a single iteration. + int const kBlock = TileTraits_::Tile::kW * TileTraits_::kAccessSize; + // The residue. + int const kResidue = (int)(bounds[1] % kBlock); + + // Compute the number of elements that are valid. + int const left = kResidue - Base::thread_offset[2]; + if (left > 0 && left < 4) { + mask_ = (1u << (8 * left)) - 1u; + } + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + Base::get(value, d, h, w, c); + if (in_residue_) { + reinterpret_cast<uint32_t&>(value) &= mask_; + } + } + + /// Move to residue portion. + CUTLASS_DEVICE void move_to_residue(typename Base::Index k) { + Base::move_to_residue(k); + in_residue_ = true; + } + + /// Move back to the beginning of the first tile. + CUTLASS_DEVICE void rollback() { + Base::rollback(); + in_residue_ = false; + } + + /// Are we in the residue? + bool in_residue_; + /// The mask to clean up the values. + uint32_t mask_; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/igemm_multiply_add.h b/cutlass-example/cutlass/gemm/igemm_multiply_add.h new file mode 100644 index 0000000..5a8baec --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_multiply_add.h @@ -0,0 +1,89 @@ +/*************************************************************************************************** + * 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 matrix multiply accumulate operation of 8-bit integer data using DP4A + instruction. +*/ +#pragma once + +#include <cutlass/fragment.h> + +#include <cutlass/gemm/thread_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template performing matrix multiply-add operation within a thread +template <typename AccumulatorsPerThread_, typename ThreadsPerWarp_> +struct ThreadMultiplyAdd<AccumulatorsPerThread_, ThreadsPerWarp_, int8_t, int8_t, int> { + /// The shape of the instruction. + typedef Shape<4, 1, 1> InstructionShape; + /// The number of accumulators per thread. + typedef AccumulatorsPerThread_ AccumulatorsPerThread; + /// The number of threads per warp. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of accumulators per warp. + typedef typename ShapeMul<AccumulatorsPerThread, ThreadsPerWarp>::Shape AccumulatorsPerWarp; + /// The type for A. + typedef int8_t ScalarA; + /// The fragment for A. + typedef Fragment<ScalarA, AccumulatorsPerThread::kW * 4> FragmentA; + /// The type for B. + typedef int8_t ScalarB; + /// The fragment for B. + typedef Fragment<ScalarB, AccumulatorsPerThread::kH * 4> FragmentB; + /// The type for C and D. + typedef int ScalarC; + /// The accumulators. + typedef Fragment<ScalarC, AccumulatorsPerThread::kH * AccumulatorsPerThread::kW> Accumulators; + + /// Ctor. + CUTLASS_DEVICE ThreadMultiplyAdd() {} + + /// Multiply : d = a*b + c. + CUTLASS_DEVICE void multiply_add(FragmentA const& a, + FragmentB const& b, + Accumulators const& c, + Accumulators& d) { + // The inputs. + int const* a_int = reinterpret_cast<int const*>(&a[0]); + int const* b_int = reinterpret_cast<int const*>(&b[0]); + + for (int j = 0; j < AccumulatorsPerThread::kH; ++j) { + for (int i = 0; i < AccumulatorsPerThread::kW; ++i) { + asm volatile("dp4a.s32.s32 %0, %1, %2, %3;" + : "=r"(d[j * AccumulatorsPerThread::kW + i]) + : "r"(a_int[i]), "r"(b_int[j]), "r"(c[j * AccumulatorsPerThread::kW + i])); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/igemm_swizzle.h b/cutlass-example/cutlass/gemm/igemm_swizzle.h new file mode 100644 index 0000000..77cf711 --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_swizzle.h @@ -0,0 +1,115 @@ +/*************************************************************************************************** + * 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 Transposes a fragment of data containing packed 8-bit integer elements. +*/ +#pragma once + +#include <cutlass/fragment.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GlobalIterator_> +struct IgemmSwizzle { + /// The global iterator. + typedef GlobalIterator_ GlobalIterator; + /// The source fragment. + typedef typename GlobalIterator::Fragment Fragment; + /// The shape of the source fragment. + typedef typename GlobalIterator::FragmentShape FragmentShape; + + /// The source fragment. + typedef Fragment InputFragment; + /// The destination fragment. + typedef Fragment OutputFragment; + + /// The src/dst must be int8 fragments. + static_assert((platform::is_same<typename Fragment::Element, int8_t>::value), "Works on int8"); + + /// The number of elements must be a multiple of 4. + static_assert(FragmentShape::kH % 4 == 0 && ShapeCount<FragmentShape>::kWc % 4 == 0, + "Not multiple of 4"); + + /// Ctor. + CUTLASS_DEVICE IgemmSwizzle() {} + + /// Transform a fragment. + CUTLASS_DEVICE void transform(Fragment const& src, Fragment& dst) { + // Expose src/dst as int arrays. + int const* src_int = reinterpret_cast<int const*>(&src[0]); + int* dst_int = reinterpret_cast<int*>(&dst[0]); + + // Transpose the data. + for (int d = 0; d < FragmentShape::kD; ++d) { + for (int h = 0; h < FragmentShape::kH / 4; ++h) { + for (int w = 0; w < ShapeCount<FragmentShape>::kWc / 4; ++w) { + int const i0 = d * (ShapeCount<FragmentShape>::kHwc / 4) + + (4 * h + 0) * (ShapeCount<FragmentShape>::kWc / 4) + w; + int const i1 = d * (ShapeCount<FragmentShape>::kHwc / 4) + + (4 * h + 1) * (ShapeCount<FragmentShape>::kWc / 4) + w; + int const i2 = d * (ShapeCount<FragmentShape>::kHwc / 4) + + (4 * h + 2) * (ShapeCount<FragmentShape>::kWc / 4) + w; + int const i3 = d * (ShapeCount<FragmentShape>::kHwc / 4) + + (4 * h + 3) * (ShapeCount<FragmentShape>::kWc / 4) + w; + + int a0 = src_int[i0]; + int a1 = src_int[i1]; + int a2 = src_int[i2]; + int a3 = src_int[i3]; + + int b0, b1, b2, b3, c0; + asm volatile("prmt.b32 %0, %1, %2, 0x0040;" : "=r"(b0) : "r"(a0), "r"(a1)); + asm volatile("prmt.b32 %0, %1, %2, 0x0040;" : "=r"(c0) : "r"(a2), "r"(a3)); + asm volatile("prmt.b32 %0, %1, %2, 0x5410;" : "=r"(b0) : "r"(b0), "r"(c0)); + + asm volatile("prmt.b32 %0, %1, %2, 0x0051;" : "=r"(b1) : "r"(a0), "r"(a1)); + asm volatile("prmt.b32 %0, %1, %2, 0x0051;" : "=r"(c0) : "r"(a2), "r"(a3)); + asm volatile("prmt.b32 %0, %1, %2, 0x5410;" : "=r"(b1) : "r"(b1), "r"(c0)); + + asm volatile("prmt.b32 %0, %1, %2, 0x0062;" : "=r"(b2) : "r"(a0), "r"(a1)); + asm volatile("prmt.b32 %0, %1, %2, 0x0062;" : "=r"(c0) : "r"(a2), "r"(a3)); + asm volatile("prmt.b32 %0, %1, %2, 0x5410;" : "=r"(b2) : "r"(b2), "r"(c0)); + + asm volatile("prmt.b32 %0, %1, %2, 0x0073;" : "=r"(b3) : "r"(a0), "r"(a1)); + asm volatile("prmt.b32 %0, %1, %2, 0x0073;" : "=r"(c0) : "r"(a2), "r"(a3)); + asm volatile("prmt.b32 %0, %1, %2, 0x5410;" : "=r"(b3) : "r"(b3), "r"(c0)); + + dst_int[i0] = b0; + dst_int[i1] = b1; + dst_int[i2] = b2; + dst_int[i3] = b3; + } + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/igemm_traits.h b/cutlass-example/cutlass/gemm/igemm_traits.h new file mode 100644 index 0000000..82f8de5 --- /dev/null +++ b/cutlass-example/cutlass/gemm/igemm_traits.h @@ -0,0 +1,539 @@ +/*************************************************************************************************** + * 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 Defies structural properties of mixed-precision integer GEMM. Multiplicands are assumed + to be packed 8bit integers, accumulators are assumed to be 32b signed integers, and output + formats vary. +*/ +#pragma once + +#include <cutlass/convert.h> +#include <cutlass/gemm/gemm.h> +#include <cutlass/gemm/gemm_epilogue.h> +#include <cutlass/gemm/gemm_epilogue_traits.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/gemm/gemm_shared_tile.h> +#include <cutlass/gemm/gemm_traits.h> +#include <cutlass/gemm/igemm_epilogue.h> +#include <cutlass/gemm/igemm_global_tile.h> +#include <cutlass/gemm/igemm_multiply_add.h> +#include <cutlass/gemm/igemm_swizzle.h> +#include <cutlass/reshape_tile.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The tile size for the GEMM KxNxM. + typename OutputTile_, + /// The output type. + typename ScalarD_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_> +struct IgemmConfig + : public GemmConfig< + /// The scalar type for A. + int8_t, + /// The scalar type for B. + int8_t, + /// The scalar type for C. + ScalarD_, + /// The scalar type for D. + ScalarD_, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, int8_t, int8_t, int>, + /// The number of scalars per LDG for A. + 4, + /// The number of scalars per STS for A. + 4, + /// The number of scalars per LDS for A. + 16, + /// The number of scalars per LDG for B. + 4, + /// The number of scalars per STS for B. + 4, + /// The number of scalars per LDS for B. + 16, + /// The number of scalars per LDG for C and STG for D. + 1, + /// The number of scalars per STS for D. + 4, + /// The number of scalars per LDS for D. + 1, + /// The number of stages in shared memory. + 2, + /// Enable the code path that deals with the residue in epilogue. + true> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename OutputTile_, typename AccumulatorsPerThread_> +struct IgemmConfig<OutputTile_, int8_t, AccumulatorsPerThread_> + : public GemmConfig< + /// The scalar type for A. + int8_t, + /// The scalar type for B. + int8_t, + /// The scalar type for C. + int8_t, + /// The scalar type for D. + int8_t, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, int8_t, int8_t, int>, + /// The number of scalars per LDG for A. + 4, + /// The number of scalars per STS for A. + 4, + /// The number of scalars per LDS for A. + 16, + /// The number of scalars per LDG for B. + 4, + /// The number of scalars per STS for B. + 4, + /// The number of scalars per LDS for B. + 16, + /// The number of scalars per LDG for C and STG for D. + 4, + /// The number of scalars per STS for D. + 4, + /// The number of scalars per LDS for D. + 4, + /// The number of stages in shared memory. + 2, + /// Enable the code path that deals with the residue in epilogue. + true> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperA : public GemmTileTraitsHelperA<kLayout_, GemmConfig_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_, Index_> + : public GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> Base; + + /// The number of scalars per LDG/STS/LDS for A. + static int const kScalarsPerStsA = 16; + + /// The traits class to build the iterator to load data from global memory for A^N. + typedef IgemmGlobalTileTraits< + GemmOperand::kA, + // The layout. + MatrixLayout::kColumnMajor, + // The pointer is float const. + int8_t const, + // The tile has size KxM in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kW>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + // The iterator. + typedef GemmGlobalIteratorAb<GlobalTileTraits, Index_> GlobalLoadIterator; + + /// The traits class to build the iterator to store data to shared memory for A^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer is float. + int8_t, + // The tile has size KxM in GEMM's terminology. + Shape<GemmConfig_::kStages, GemmConfig_::OutputTile::kD / 4, GemmConfig_::OutputTile::kW * 4>, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + kScalarsPerStsA> + SharedStoreTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_, Index_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor; + + /// The input scalar. + typedef int8_t Scalar; + /// The scalar stored in shared memory. + typedef int8_t MultiplyAddScalar; + + /// The number of scalars per LDG/STS/LDS for A. + static int const kScalarsPerStsA = 16; + + /// The traits class to build the iterator to load data from global memory for A^T. + typedef IgemmGlobalTileTraits< + GemmOperand::kA, + // The layout. + MatrixLayout::kRowMajor, + // The pointer is float const. + int8_t const, + // The tile has size NxK in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + // The iterator. + typedef IgemmGlobalIteratorAb<GlobalTileTraits, Index_> GlobalLoadIterator; + + /// The traits class to build the iterator to store data to shared memory for A^N. + typedef GemmSharedStoreWithSkewTileAbTraits< + // The pointer is int8. + int8_t, + // The tile has size KxN in GEMM's terminology. + Shape<GemmConfig_::kStages, GemmConfig_::OutputTile::kD / 4, GemmConfig_::OutputTile::kW * 4>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS. + kScalarsPerStsA, + // The skew to avoid bank conflicts added in the tile W dimension. + 16> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for A^N. + typedef GemmSharedLoadTileATraits< + // The pointer is float const. + int8_t const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + 16, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperB : public GemmTileTraitsHelperB<kLayout_, GemmConfig_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_, Index_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor; + + /// The input scalar. + typedef int8_t Scalar; + /// The scalar stored in shared memory. + typedef int8_t MultiplyAddScalar; + + /// The number of scalars per LDG/STS/LDS for B. + static int const kScalarsPerStsB = 16; + + /// The traits class to build the iterator to load data from global memory for B^T. + typedef IgemmGlobalTileTraits< + GemmOperand::kB, + // The layout. + MatrixLayout::kColumnMajor, + // The pointer is float const. + int8_t const, + // The tile has size NxK in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + // The iterator. + typedef IgemmGlobalIteratorAb<GlobalTileTraits, Index_> GlobalLoadIterator; + + /// The traits class to build the iterator to store data to shared memory for B^N. + typedef GemmSharedStoreWithSkewTileAbTraits< + // The pointer is int8. + int8_t, + // The tile has size KxN in GEMM's terminology. + Shape<GemmConfig_::kStages, GemmConfig_::OutputTile::kD / 4, GemmConfig_::OutputTile::kH * 4>, + // The threads are distributed as (threads / K) x K (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS. + kScalarsPerStsB, + // The skew to avoid bank conflicts added in the tile W dimension. + 16> + SharedStoreTileTraits; + + /// The traits class to build the iterator to load from shared memory for B^N. + typedef GemmSharedLoadTileBTraits< + // The pointer is float const. + int8_t const, + // 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 shape of the FMA instruction. + typename GemmConfig_::InstructionShape, + // The number of stages. + GemmConfig_::kStages, + // The number of scalars per LDS. + 16, + // The skew. + SharedStoreTileTraits::kSkew> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename Index_> +struct IgemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_, Index_> + : public GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> Base; + + /// The number of scalars per LDG/STS/LDS for B. + static int const kScalarsPerStsB = 16; + + /// The traits class to build the iterator to load data from global memory for B^T. + typedef IgemmGlobalTileTraits< + GemmOperand::kB, + // The layout. + MatrixLayout::kRowMajor, + // The pointer is float const. + int8_t const, + // The tile has size KxM in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kH>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + // The iterator. + typedef GemmGlobalIteratorAb<GlobalTileTraits, Index_> GlobalLoadIterator; + + /// The traits class to build the iterator to store data to shared memory for B^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer is float. + int8_t, + // The tile has size KxM in GEMM's terminology. + Shape<GemmConfig_::kStages, GemmConfig_::OutputTile::kD / 4, GemmConfig_::OutputTile::kH * 4>, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + kScalarsPerStsB> + SharedStoreTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename Iterator_> +struct IgemmTransformerA {}; + +template <typename Iterator_> +struct IgemmTransformerA<MatrixLayout::kRowMajor, Iterator_> { + typedef Copy<typename Iterator_::Fragment> Transformer; +}; + +template <typename Iterator_> +struct IgemmTransformerA<MatrixLayout::kColumnMajor, Iterator_> { + typedef IgemmSwizzle<Iterator_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename Iterator_> +struct IgemmTransformerB {}; + +template <typename Iterator_> +struct IgemmTransformerB<MatrixLayout::kColumnMajor, Iterator_> { + typedef Copy<typename Iterator_::Fragment> Transformer; +}; + +template <typename Iterator_> +struct IgemmTransformerB<MatrixLayout::kRowMajor, Iterator_> { + typedef IgemmSwizzle<Iterator_> Transformer; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_, + /// The output type. + typename ScalarD_, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<32, 8, 8>, + /// The index. + typename Index_ = int> +struct IgemmTraitsHelper { + /// The IGEMM config. + typedef IgemmConfig<OutputTile_, ScalarD_, AccumulatorsPerThread_> GemmConfig; + /// The GEMM config for A. + typedef IgemmTileTraitsHelperA<kLayoutA_, GemmConfig, Index_> GemmTileTraitsHelperA; + /// The GEMM config for B. + typedef IgemmTileTraitsHelperB<kLayoutB_, GemmConfig, Index_> GemmTileTraitsHelperB; + + /// The iterator to load A from global memory. + typedef typename GemmTileTraitsHelperA::GlobalLoadIterator GlobalLoadIteratorA; + + /// The default transformer for A. + typedef typename IgemmTransformerA<GemmTileTraitsHelperA::kLayout, + GlobalLoadIteratorA>::Transformer GlobalTransformerA; + /// The iterator to store A to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperA::SharedStoreTileTraits, + typename GemmTileTraitsHelperA::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorA; + /// The stream to load A from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA> + GlobalLoadStreamA; + + /// The iterator to load B from global memory. + typedef typename GemmTileTraitsHelperB::GlobalLoadIterator GlobalLoadIteratorB; + + // The default transformer for B. + typedef typename IgemmTransformerB<GemmTileTraitsHelperB::kLayout, + GlobalLoadIteratorB>::Transformer GlobalTransformerB; + /// The iterator to store B to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperB::SharedStoreTileTraits, + typename GemmTileTraitsHelperB::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorB; + /// The stream to load B from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB> + GlobalLoadStreamB; + + /// The iterator to load A from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperA::SharedLoadTileTraits, + typename GemmTileTraitsHelperA::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorA; + /// The stream to load A from shared memory. + typedef SharedLoadStream<SharedLoadIteratorA, Copy<typename SharedLoadIteratorA::Fragment> > + SharedLoadStreamA; + /// The iterator to load B from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperB::SharedLoadTileTraits, + typename GemmTileTraitsHelperB::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorB; + /// The stream to load B from shared memory. + typedef SharedLoadStream<SharedLoadIteratorB, Copy<typename SharedLoadIteratorB::Fragment> > + SharedLoadStreamB; + + /// The multiply-add functor. + typedef typename GemmConfig::MultiplyAdd MultiplyAdd; + /// The object to clear accumulators. + typedef ClearAccumulators<typename MultiplyAdd::ScalarC> ClearAccumulators; + + /// The epilogue. + typedef IgemmEpilogue<IgemmEpilogueTraits<GemmConfig, EpilogueFunctor_> > Epilogue; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename ScalarD_> +struct IgemmEpilogueScalar { + typedef float Scalar; +}; + +template <> +struct IgemmEpilogueScalar<int> { + typedef int Scalar; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_ = Shape<32, 128, 128>, + /// The output type. + typename ScalarD_ = int, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_ = LinearScaling<typename IgemmEpilogueScalar<ScalarD_>::Scalar>, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<32, 8, 8>, + /// The index. + typename Index_ = int, + /// The helper class. + typename Helper_ = IgemmTraitsHelper<kLayoutA_, + kLayoutB_, + OutputTile_, + ScalarD_, + EpilogueFunctor_, + AccumulatorsPerThread_, + Index_> > +struct IgemmTraits : public GemmTraits< + // The config. + typename Helper_::GemmConfig, + // The stream to load A from global memory to shared memory. + typename Helper_::GlobalLoadStreamA, + // The stream to load B from global memory to shared memory. + typename Helper_::GlobalLoadStreamB, + // The stream to load A from shared memory. + typename Helper_::SharedLoadStreamA, + // The stream to load B from shared memory. + typename Helper_::SharedLoadStreamB, + // The epilogue. + typename Helper_::Epilogue, + // The block swizzle to reorganize the grid. + IdentityBlockSwizzle, + // The index. + Index_, + // The tool used to clear accumulators. + typename Helper_::ClearAccumulators> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/linear_scaling.h b/cutlass-example/cutlass/gemm/linear_scaling.h new file mode 100644 index 0000000..979c93f --- /dev/null +++ b/cutlass-example/cutlass/gemm/linear_scaling.h @@ -0,0 +1,85 @@ +/*************************************************************************************************** + * 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 BLAS linear scaling function alpha*AB + beta*C +*/ +#pragma once + +#include <cutlass/fragment_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Functor to compute linear combination of fragments +template <typename Scalar_, typename FragmentMultiplyAdd_ = FragmentMultiplyAdd<Scalar_> > +struct LinearScaling { + // The scalar. + typedef Scalar_ Scalar; + // The adapater. + typedef FragmentMultiplyAdd_ FragmentMultiplyAdd; + + /// The parameters. + struct Params { + /// The alpha/beta scaling params. + Scalar alpha, beta; + + /// Initialize the parameters. + template <typename GemmDesc_> + CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc) { + alpha = desc.alpha; + beta = desc.beta; + return 0; + } + }; + + /// Ctor. + CUTLASS_DEVICE LinearScaling(Params const& params) : alpha(params.alpha), beta(params.beta) {} + + /// Evaluate the functor. + template <typename FragmentA_, typename FragmentB_> + CUTLASS_DEVICE void evaluate(FragmentA_ const& accum, FragmentB_& output) { + FragmentMultiplyAdd mad; + mad.multiply(alpha, accum, output); + } + + /// Evaluate the functor. + template <typename FragmentA_, typename FragmentB_> + CUTLASS_DEVICE void evaluate(FragmentA_ const& accum, FragmentB_ const& old, FragmentB_& output) { + FragmentMultiplyAdd mad; + FragmentB_ tmp; + mad.multiply(beta, old, tmp); + mad.multiply_add(alpha, accum, tmp, output); + } + + /// The alpha/beta scaling factors. + Scalar alpha, beta; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/sgemm_traits.h b/cutlass-example/cutlass/gemm/sgemm_traits.h new file mode 100644 index 0000000..66b7677 --- /dev/null +++ b/cutlass-example/cutlass/gemm/sgemm_traits.h @@ -0,0 +1,127 @@ +/*************************************************************************************************** + * 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 Defies structural properties of single-precision GEMM. +*/ +#pragma once + +#include <cutlass/gemm/gemm.h> +#include <cutlass/gemm/gemm_epilogue.h> +#include <cutlass/gemm/gemm_epilogue_traits.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/gemm/gemm_shared_tile.h> +#include <cutlass/gemm/gemm_traits.h> +#include <cutlass/gemm/thread_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The tile size for the GEMM KxNxM. + typename OutputTile_, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_ = 1, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_ = 1> +struct SgemmConfig + : public GemmConfig< + /// The scalar type for A. + float, + /// The scalar type for B. + float, + /// The scalar type for C. + float, + /// The scalar type for D. + float, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, float, float, float>, + /// The number of scalars per LDG for A. + kScalarsPerLdgA_, + /// The number of scalars per STS for A. + kScalarsPerLdgA_, + /// The number of scalars per LDS for A. + 4, + /// The number of scalars per LDG for B. + kScalarsPerLdgB_, + /// The number of scalars per STS for B. + kScalarsPerLdgB_, + /// The number of scalars per LDS for B. + 4, + /// The number of scalars per LDG for C and STG for D. + 1, + /// The number of scalars per STS for D. + 4, + /// The number of scalars per LDS for D. + 1, + /// The number of stages in shared memory. + 2> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_ = Shape<8, 128, 128>, + /// The functor to use in the epilogue. + typename EpilogueFunctor_ = LinearScaling<float>, + /// The number of accumulators per thread. + typename AccumulatorsPerThread_ = Shape<8, 8, 8>, + /// The number of floats loaded in one LDG for A. + int kScalarsPerLdgA_ = 1, + /// The number of floats loaded in one LDG for B. + int kScalarsPerLdgB_ = 1, + /// The index. + typename Index_ = int, + /// The SGEMM config. + typename GemmConfig_ = + SgemmConfig<OutputTile_, AccumulatorsPerThread_, kScalarsPerLdgA_, kScalarsPerLdgB_>, + /// The traits class for the epilogue. + typename GemmEpilogueTraits_ = + SimplifiedGemmEpilogueTraits<GemmConfig_, EpilogueFunctor_, Index_> > +struct SgemmTraits : public SimplifiedGemmTraits< + // The layout for A. + kLayoutA_, + // The layout for B. + kLayoutB_, + // The config. + GemmConfig_, + // The epilogue. + GemmEpilogue<GemmEpilogueTraits_>, + // The index. + Index_> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/thread_multiply_add.h b/cutlass-example/cutlass/gemm/thread_multiply_add.h new file mode 100644 index 0000000..20dca15 --- /dev/null +++ b/cutlass-example/cutlass/gemm/thread_multiply_add.h @@ -0,0 +1,84 @@ +/*************************************************************************************************** + * 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 Template implementing matrix multiply-add operations on fragments. +*/ +#pragma once + +#include <cutlass/fragment.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template performing matrix multiply-add operation within a thread +template <typename AccumulatorsPerThread_, + typename ThreadsPerWarp_, + typename ScalarA_, + typename ScalarB_, + typename ScalarC_> +struct ThreadMultiplyAdd { + /// The shape of the instruction. + typedef Shape<1, 1, 1, 1> InstructionShape; + /// The number of accumulators per thread. + typedef AccumulatorsPerThread_ AccumulatorsPerThread; + /// The number of threads per warp. + typedef ThreadsPerWarp_ ThreadsPerWarp; + /// The number of accumulators per warp. + typedef typename ShapeMul<AccumulatorsPerThread, ThreadsPerWarp>::Shape AccumulatorsPerWarp; + /// The type for A. + typedef ScalarA_ ScalarA; + /// The fragment for A. + typedef Fragment<ScalarA, AccumulatorsPerThread::kW> FragmentA; + /// The type for B. + typedef ScalarB_ ScalarB; + /// The fragment for B. + typedef Fragment<ScalarB, AccumulatorsPerThread::kH> FragmentB; + /// The type for C and D. + typedef ScalarC_ ScalarC; + /// The accumulators. + typedef Fragment<ScalarC, AccumulatorsPerThread::kH * AccumulatorsPerThread::kW, 16> Accumulators; + + /// Ctor. + CUTLASS_DEVICE ThreadMultiplyAdd() {} + + /// Multiply : d = a*b + c. + CUTLASS_DEVICE void multiply_add(FragmentA const& a, + FragmentB const& b, + Accumulators const& c, + Accumulators& d) { + for (int j = 0; j < AccumulatorsPerThread::kH; ++j) { + for (int i = 0; i < AccumulatorsPerThread::kW; ++i) { + d[j * AccumulatorsPerThread::kW + i] = a[i] * b[j] + c[j * AccumulatorsPerThread::kW + i]; + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_epilogue_traits.h b/cutlass-example/cutlass/gemm/wmma_gemm_epilogue_traits.h new file mode 100644 index 0000000..0fafacf --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_epilogue_traits.h @@ -0,0 +1,161 @@ +/*************************************************************************************************** + * 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 Defines structural properties of WMMA GEMM's epilogue phase. +*/ +#pragma once + +#include <cutlass/wmma_matrix.h> +#ifdef CUTLASS_USE_WMMA_API + +#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/gemm/wmma_gemm_global_tile.h> +#include <cutlass/gemm/wmma_gemm_shared_tile.h> +#include <cutlass/reshape_tile.h> +#include <cutlass/tile_iterator.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_, typename EpilogueFunctor_, typename Index_ = int> +struct WmmaGemmEpilogueTraitsHelper { + /// The scalar. + typedef typename EpilogueFunctor_::Scalar Scalar; + /// The output tile. + typedef typename GemmConfig_::OutputTile OutputTile; + + /// The number of WMMAs in the H dimension. + static int const kWmmasPerH = + GemmConfig_::AccumulatorsPerWarp::kH / GemmConfig_::InstructionShape::kH; + /// The number of iterations in the epilogue. That's the number of "horizontal" WMMAs. + typedef Shape<1, 1, kWmmasPerH> Iterations; + // The iteration strides in the H/W dimension. + typedef Shape<0, 0, 0> Delta; + /// The functor to do the math in the epilogue. + typedef EpilogueFunctor_ Functor; + + /// The traits class to build the iterator to store to shared memory for D. + typedef WmmaGemmSharedStoreTileDTraits< + // The output layout. + MatrixLayout::kColumnMajor, + // The pointer is float. + typename Functor::Scalar, + // The output tile size. + typename GemmConfig_::OutputTile, + // The number of warps. + typename GemmConfig_::Warps, + // The shape of the instruction. + typename GemmConfig_::InstructionShape> + SharedStoreTileTraits; + + typedef WmmaMatrix<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar, + typename GemmConfig_::InstructionShape> + WmmaMatrix; + + /// The iterator to store D to shared memory. + typedef TileStoreIterator<SharedStoreTileTraits, + typename SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared, + Index_, + WmmaMatrix, + IteratorFragment::kWmmaMatrix> + SharedStoreIteratorD; + + /// The shared store transformer for D. + typedef Copy<typename SharedStoreIteratorD::Fragment> SharedStoreTransformerD; + + /// The traits class to build the iterator to load from shared memory for D. + typedef WmmaGemmSharedLoadTileDTraits< + // The pointer. + typename Functor::Scalar, + // The tile size. + typename SharedStoreIteratorD::Tile, + // The number of threads. + Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>, + // The number of scalars per LDS. + GemmConfig_::kScalarsPerLdsD> + SharedLoadTileTraits; + + /// The iterator to load D from shared memory. + typedef TileLoadIterator<SharedLoadTileTraits, + typename SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedLoadIteratorD; + + /// The traits class to build the iterator to load data from global memory for C^N. + typedef WmmaGemmGlobalIteratorCdTraits< + // 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>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgC> + GlobalLoadTileTraits; + + /// The iterator to load C. + typedef WmmaGemmGlobalIteratorCd<GlobalLoadTileTraits, Index_> GlobalLoadIteratorC; + /// The transformer for C. + typedef Copy<typename GlobalLoadIteratorC::Fragment> GlobalTransformerC; + + /// The traits class to build the iterator to store data to global memory for D^N. + typedef WmmaGemmGlobalIteratorCdTraits< + // 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>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerStgD> + GlobalStoreTileTraits; + + /// The iterator to store D. + typedef WmmaGemmGlobalIteratorCd<GlobalStoreTileTraits, Index_> GlobalStoreIteratorD; + /// The transformer for D. + typedef Copy<typename GlobalStoreIteratorD::Fragment> GlobalTransformerD; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass + +#endif // defined CUTLASS_USE_WMMA_API diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h b/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h new file mode 100644 index 0000000..dbd57f6 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_global_tile.h @@ -0,0 +1,211 @@ +/*************************************************************************************************** + * 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 Defines tile iterator traits for loading thread block-level tile from global memory. +*/ +#pragma once + +#include <cutlass/gemm/gemm_global_tile.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kAccessSize_> +struct WmmaGemmGlobalIteratorCdTraits : public GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> { + /// The base class. + typedef GemmGlobalTileTraits<GemmOperand::kC, + MatrixLayout::kColumnMajor, + Scalar_, + Tile_, + Threads_, + kAccessSize_> + Base; + + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> Delta; + + /// Computes the thread offset in (H, W) based on thread ID + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + int thread_offset_h = threadIdx.x / Base::Threads::kW; + int thread_offset_w = threadIdx.x % Base::Threads::kW * Base::ThreadsDelta::kW; + + return make_Coord(0, thread_offset_h, thread_offset_w, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename TileTraits_, typename Index_ = int> +struct WmmaGemmGlobalIteratorCd : public TileIteratorBase<TileTraits_, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> { + /// This class. + typedef WmmaGemmGlobalIteratorCd<TileTraits_, Index_> This_; + /// The traits. + typedef TileTraits_ Traits; + /// The base class. + typedef TileIteratorBase<Traits, + typename TileTraits_::Scalar, + IteratorAdvance::kH, + MemorySpace::kGlobal, + Index_> + Base; + /// Override the strides in each dimension between different loads/stores. + typedef Shape<0, 0, Base::Delta::kW, Base::Delta::kC> ImmediateOffsetStrides; + /// The layout. + static MatrixLayout::Kind const kLayout = TileTraits_::kLayout; + + /// The scalar. + typedef typename TileTraits_::Scalar Scalar; + /// The pointer. + typedef typename TileTraits_::Pointer Pointer; + /// The threads. + typedef typename TileTraits_::Threads Threads; + /// The index. + typedef Index_ Index; + /// The thread offset functor. + typedef typename TileTraits_::ThreadOffset ThreadOffset; + + /// The params. + struct Params { + /// The pointer. + Pointer pointer; + /// The stride in the H dimension to setup the thread in the block. + Index stride_h; + /// The strides to increment the pointer. + Index inc_h, inc_advance; + /// The column offset to compute the predicate for the columns. + Index predicate_offset; + /// The strides to increment the predicate offset. + Index predicate_inc_h, predicate_inc_advance; + + /// Setup the params. + CUTLASS_HOST_DEVICE int initialize( + Pointer pointer, Index ld, Index n, Index epilogue_stride_w, Index epilogue_delta_w) { + // The pointer. + this->pointer = pointer; + // Setup the base stride. One "group of threads" per column. + stride_h = ld; + // Each thread output 1 column per iteration. . + inc_h = ld * TileTraits_::Threads::kH; + inc_advance = inc_h + epilogue_stride_w; + + predicate_offset = n; + predicate_inc_h = TileTraits_::Threads::kH; + predicate_inc_advance = predicate_inc_h + epilogue_delta_w; + + // It worked. + return 0; + } + }; + + Params params; + + Coord<4> thread_offset; + + /// Ctor. + CUTLASS_DEVICE WmmaGemmGlobalIteratorCd() {} + + /// Ctor. + CUTLASS_DEVICE WmmaGemmGlobalIteratorCd(Params const& params, + const Coord<3>& bounds, + const Coord<3>& block, + int const pointer_offset = 0, + int const pred_offset = 0, + ThreadOffset thread_offset_func = ThreadOffset()) + + : params(params) { + thread_offset = thread_offset_func(); + // Each warp works on a different column of the tile. + int const h = thread_offset[1] + block[1]; + // Each lane writes a different element. + int const w = thread_offset[2] + block[2]; + // Setup the pointer. + this->params.pointer += ((h * params.stride_h + w) + pointer_offset); + + // Prepare the vector of predicates. + for (int i = 0; i < Base::Iterations::kW; ++i) { + predicates.set(i, w + i * Base::Delta::kW < bounds[2]); + } + this->params.predicate_offset -= (h + pred_offset); + } + + /// The accessor. + CUTLASS_DEVICE void get(typename Base::AccessType& value, int d, int h, int w, int c) const { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(0, 0, w, c); + Load<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::load(value, params.pointer, imm); + } + + /// Increment the pointer in the C dimension. + CUTLASS_DEVICE void inc_c() {} + /// Increment the pointer in the W dimension. + CUTLASS_DEVICE void inc_w() {} + /// Increment the pointer in the H dimension. + CUTLASS_DEVICE void inc_h() { + params.pointer += params.inc_h; + params.predicate_offset -= params.predicate_inc_h; + } + /// Increment the pointer in the D dimension. + CUTLASS_DEVICE void inc_d() {} + /// Increment the pointer to move to the next iteration. + CUTLASS_DEVICE void inc_advance() { + params.pointer += params.inc_advance; + params.predicate_offset -= params.predicate_inc_advance; + } + + /// The accessor. + CUTLASS_DEVICE void set(typename Base::AccessType const& value, int d, int h, int w, int c) { + int const imm = + ComputeOffsetFromStrides<typename Base::ImmediateOffsetStrides>::get(d, h, w, 0); + Store<Scalar, TileTraits_::kAccessSize, MemorySpace::kGlobal>::store( + value, params.pointer, imm); + } + + /// Test the predicate. + CUTLASS_DEVICE bool valid(int d, int h, int w, int c) const { + return predicates.at(w) && params.predicate_offset > 0; + } + + /// The predicates for the row. + cutlass::PredicateVector<Base::Iterations::kW> predicates; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_multiply_add.h b/cutlass-example/cutlass/gemm/wmma_gemm_multiply_add.h new file mode 100644 index 0000000..5968350 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_multiply_add.h @@ -0,0 +1,108 @@ +/*************************************************************************************************** + * 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 warp-level matrix multiply-accumulate operation using CUDA WMMA API. +*/ +#pragma once + +#include <cutlass/wmma_matrix.h> +#ifdef CUTLASS_USE_WMMA_API +#include <cutlass/fragment.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayoutA_, + typename ScalarA_, + MatrixLayout::Kind kLayoutB_, + typename ScalarB_, + MatrixLayout::Kind kLayoutC_, + typename ScalarC_, + typename AccumulatorsPerWarp_, + typename InstructionShape_> +struct WmmaGemmMultiplyAdd { + /// The shape of the instruction. + typedef InstructionShape_ InstructionShape; + /// The number of threads per warp. That's a dummy configuration. + typedef Shape<1, InstructionShape_::kH, InstructionShape_::kW> ThreadsPerWarp; + /// The dimensions. + typedef AccumulatorsPerWarp_ AccumulatorsPerWarp; + /// The type for A. + typedef ScalarA_ ScalarA; + /// The type for B. + typedef ScalarB_ ScalarB; + /// The type for C and D. + typedef ScalarC_ ScalarC; + /// The number of iterations. + typedef typename ShapeDiv<AccumulatorsPerWarp, InstructionShape>::Shape Iterations; + + /// The element for A. + typedef WmmaMatrix<GemmOperand::kA, kLayoutA_, ScalarA, InstructionShape> ElementA; + /// The fragment for A. + typedef Fragment<ElementA, Iterations::kW> FragmentA; + + /// The element for B. + typedef WmmaMatrix<GemmOperand::kB, kLayoutB_, ScalarB, InstructionShape> ElementB; + /// The fragment for B. + typedef Fragment<ElementB, Iterations::kH> FragmentB; + + /// The element for C. + typedef WmmaMatrix<GemmOperand::kC, kLayoutC_, ScalarC, InstructionShape> ElementC; + /// The fragment for C. + typedef Fragment<ElementC, Iterations::kH * Iterations::kW> Accumulators; + + /// Ctor. + CUTLASS_DEVICE WmmaGemmMultiplyAdd() {} + + /// Multiply : d = a*b. + CUTLASS_DEVICE void multiply_add(FragmentA const& a, + FragmentB const& b, + Accumulators const& c, + Accumulators& d) { + for (int j = 0; j < Iterations::kH; ++j) { + for (int i = 0; i < Iterations::kW; ++i) { + // The input elements. + ElementA const& elt_a = a[i]; + ElementB const& elt_b = b[j]; + ElementC const& elt_c = c[j * Iterations::kW + i]; + + // The output element. + ElementC& elt_d = d[j * Iterations::kW + i]; + + // The wmma instruction. + nvcuda::wmma::mma_sync(elt_d, elt_a, elt_b, elt_c); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass + +#endif // defined CUTLASS_USE_WMMA_API diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h b/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h new file mode 100644 index 0000000..7d15b26 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_shared_tile.h @@ -0,0 +1,240 @@ +/*************************************************************************************************** + * 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 Defines iterator traits for efficiently loading and storing fragment to and from shared + memory, specialized for WMMA GEMM. +*/ +#pragma once + +#include <cutlass/wmma_matrix.h> +#ifdef CUTLASS_USE_WMMA_API + +#include <cutlass/gemm/gemm_operand.h> +#include <cutlass/reshape_tile.h> + +namespace cutlass { +namespace gemm { + +template <class> +struct Debug {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Warps_, + int kWarpStride_, + typename Iterations_, + typename Delta_, + typename WmmaShape_> +struct WmmaGemmSharedLoadTileATraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kA; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = 1; + /// The tile with skew. + typedef Tile_ Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The warps strides. + static int const kWarpStride = kWarpStride_; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The strides between iterations. + typedef Delta_ Delta; + /// The strides between iterations. + typedef Delta_ ImmediateOffsetStrides; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The offset. + int const offset = warp % Warps::kW * kWarpStride; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename Tile_, + typename Warps_, + int kWarpStride_, + typename Iterations_, + typename Delta_, + typename WmmaShape_> +struct WmmaGemmSharedLoadTileBTraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kB; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = 1; + /// The tile with skew. + typedef Tile_ Tile; + /// The number of warps. + typedef Warps_ Warps; + /// The warps strides. + static int const kWarpStride = kWarpStride_; + /// The number of iterations. + typedef Iterations_ Iterations; + /// The strides between iterations. + typedef Delta_ Delta; + /// The strides between iterations. + typedef Delta_ ImmediateOffsetStrides; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The offset. + int const offset = warp / Warps::kW * kWarpStride; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <MatrixLayout::Kind kLayout_, + typename Scalar_, + typename OutputTile_, + typename Warps_, + typename WmmaShape_, + int kSkew_ = 0> +struct WmmaGemmSharedStoreTileDTraits { + /// The operand. + static GemmOperand::Kind const kOperand = GemmOperand::kC; + /// The layout. + static MatrixLayout::Kind const kLayout = kLayout_; + /// The scalar. + typedef Scalar_ Scalar; + // The access size + static int const kAccessSize = 1; + /// The pointer. + typedef Scalar* Pointer; + /// The number of warps. + typedef Warps_ Warps; + /// The shape of the WMMA instruction. + typedef WmmaShape_ WmmaShape; + /// The skew. + static int const kSkew = kSkew_; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + /// The tile with skew. + typedef Shape<1, Warps_::kH * WmmaShape_::kH, OutputTile_::kW + kSkew_> Tile; + /// The number of iterations needed to store the tile. + typedef Shape<1, 1, OutputTile_::kW / Warps::kW / WmmaShape_::kW> Iterations; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * WmmaShape_::kW, 0> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, 0, Warps::kW * WmmaShape_::kW, 0> ImmediateOffsetStrides; + + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The warp id. + int const warp = threadIdx.x / kWarpSize; + // The starting column. + int const h = warp / Warps::kW * WmmaShape::kH; + // The w. + int const w = warp % Warps::kW * WmmaShape::kW; + // The offset. + int const offset = h * Tile::kW + w; + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename Scalar_, typename Tile_, typename Threads_, int kScalarsPerLds_> +struct WmmaGemmSharedLoadTileDTraits { + /// The scalar. + typedef Scalar_ Scalar; + /// The pointer. + typedef Scalar const* Pointer; + /// The access size + static int const kAccessSize = kScalarsPerLds_; + /// The tile. + typedef typename ReshapeTile<Tile_, kScalarsPerLds_>::Tile Tile; + /// The threads. + typedef typename ReshapeThreads<Tile, Threads_>::Threads Threads; + /// The threads strides. + typedef Shape<1, Tile::kW * Tile::kC, Tile::kC> ThreadsStrides; + /// The memory space. + static MemorySpace::Kind const kMemorySpace = MemorySpace::kShared; + + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kScalarsPerLds_> Delta; + /// The strides in each dimension between different loads/stores. + typedef Shape<0, Threads::kH * ShapeCount<Tile>::kWc, Threads::kW * kScalarsPerLds_> + ImmediateOffsetStrides; + /// The number of iterations needed to load/store the tile. + typedef Shape<1, Tile::kH / Threads::kH, Tile::kW / Threads::kW, Tile::kC / kScalarsPerLds_> + Iterations; + + /// ThreadOffset + struct ThreadOffset { + CUTLASS_HOST_DEVICE + Coord<4> operator()() const { + // The offset. + int const offset = ComputeThreadOffsetFromStrides<Threads, ThreadsStrides>::get(); + return make_Coord(0, 0, offset, 0); + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass + +#endif // defined CUTLASS_USE_WMMA_API diff --git a/cutlass-example/cutlass/gemm/wmma_gemm_traits.h b/cutlass-example/cutlass/gemm/wmma_gemm_traits.h new file mode 100644 index 0000000..7901201 --- /dev/null +++ b/cutlass-example/cutlass/gemm/wmma_gemm_traits.h @@ -0,0 +1,574 @@ +/*************************************************************************************************** + * 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 Defies structural properties of GEMM targeting WMMA API in CUDA. +*/ +#pragma once + +#include <cutlass/wmma_matrix.h> +#ifdef CUTLASS_USE_WMMA_API + +#include <cutlass/convert.h> +#include <cutlass/gemm/gemm.h> +#include <cutlass/gemm/gemm_epilogue.h> +#include <cutlass/gemm/gemm_epilogue_traits.h> +#include <cutlass/gemm/gemm_global_tile.h> +#include <cutlass/gemm/gemm_shared_tile.h> +#include <cutlass/gemm/gemm_traits.h> +#include <cutlass/gemm/wmma_gemm_epilogue_traits.h> +#include <cutlass/gemm/wmma_gemm_global_tile.h> +#include <cutlass/gemm/wmma_gemm_multiply_add.h> + +namespace cutlass { +namespace gemm { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The tile size for the GEMM KxNxM. + typename OutputTile_, + /// The output type. + typename ScalarC_, + /// The accumulator type. + typename Accumulator_, + /// The number of accumulators per warp. + typename AccumulatorsPerWarp_, + /// The shape of the WMMA instruction. + typename InstructionShape_, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_> +struct WmmaGemmConfig : public GemmConfig< + /// The scalar type for A. + half, + /// The scalar type for B. + half, + /// The scalar type for C. + ScalarC_, + /// The scalar type for D. + ScalarC_, + /// The tile size for the GEMM KxNxM. + OutputTile_, + /// The functor to do the math in the main loop. + WmmaGemmMultiplyAdd<kLayoutA_, + half, + kLayoutB_, + half, + MatrixLayout::kColumnMajor, + Accumulator_, + AccumulatorsPerWarp_, + InstructionShape_>, + /// The number of scalars per LDG for A. + kScalarsPerLdgA_, + /// The number of scalars per STS for A. + kScalarsPerLdgA_, + /// The number of scalars per LDS for A. + 8, + /// The number of scalars per LDG for B. + kScalarsPerLdgB_, + /// The number of scalars per STS for B. + kScalarsPerLdgB_, + /// The number of scalars per LDS for B. + 8, + /// The number of scalars per LDG for C and STG for D. + 16 / sizeof(ScalarC_), + /// The number of scalars per STS for D. + 16 / sizeof(ScalarC_), + /// The number of scalars per LDS for D. + 16 / sizeof(ScalarC_), + /// The number of stages in shared memory. + 1> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_> +struct WmmaGemmTileTraitsHelperA {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct WmmaGemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> + : public GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> Base; + + /// The skew. + static int const kSkew = 16 / sizeof(typename Base::MultiplyAddScalar); + /// The shared tile size. + typedef Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD, + GemmConfig_::OutputTile::kW + kSkew> + Tile; + + /// WMMA matrix + typedef WmmaMatrix<GemmOperand::kA, + MatrixLayout::kColumnMajor, + typename Base::MultiplyAddScalar, + typename GemmConfig_::InstructionShape> + WmmaMatrix; + + /// The traits class to build the iterator to store data to shared memory for A^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer. + typename Base::MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Tile, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename Base::GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsA> + SharedStoreTileTraits; + + /// The number of elements loaded in one LDG. + static int const kScalarsPerW = GemmConfig_::InstructionShape::kW * GemmConfig_::Warps::kW; + /// The number of scalars loaded per iteration. + static int const kScalarsPerIteration = Tile::kW * GemmConfig_::InstructionShape::kD; + /// The traits class to build the iterator to load from shared memory for A. + typedef WmmaGemmSharedLoadTileATraits< + // The layout of the matrix. + MatrixLayout::kColumnMajor, + // The pointer. + typename Base::MultiplyAddScalar, + // The output tile size. + Tile, + // The number of warps. + typename GemmConfig_::Warps, + // The strides between warps. + GemmConfig_::InstructionShape::kW, + // The number of iterations to load the data. + Shape<1, 1, GemmConfig_::OutputTile::kW / kScalarsPerW>, + // The stride between iterations. + Shape<kScalarsPerIteration, 0, kScalarsPerW, 0>, + // The shape of the instruction. + typename GemmConfig_::InstructionShape> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct WmmaGemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarA Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar; + + /// WMMA matrix + typedef WmmaMatrix<GemmOperand::kA, + MatrixLayout::kRowMajor, + MultiplyAddScalar, + typename GemmConfig_::InstructionShape> + WmmaMatrix; + + /// The traits class to build the iterator to load data from global memory for A^T. + typedef GemmGlobalTileTraits< + // That's A. + GemmOperand::kA, + // A is row-major. + MatrixLayout::kRowMajor, + // The pointer is float const. + Scalar const, + // The tile has size KxM in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgA> + GlobalTileTraits; + + /// The skew. + static int const kSkew = 16 / sizeof(MultiplyAddScalar); + /// The tile. + typedef Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kW, + GemmConfig_::OutputTile::kD + kSkew> + Tile; + + /// The traits class to build the iterator to store data to shared memory for A^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer. + MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Tile, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsA> + SharedStoreTileTraits; + + /// The number of elements loaded in one LDG. + static int const kScalarsPerW = GemmConfig_::InstructionShape::kW * GemmConfig_::Warps::kW; + /// The traits class to build the iterator to load from shared memory for A. + typedef WmmaGemmSharedLoadTileATraits< + // The layout of the matrix. + MatrixLayout::kRowMajor, + // The pointer. + MultiplyAddScalar, + // The tile in shared memory. + Tile, + // The number of warps. + typename GemmConfig_::Warps, + // The strides between warps. + GemmConfig_::InstructionShape::kW * Tile::kW, + // The number of iterations to load the data. + Shape<1, 1, GemmConfig_::OutputTile::kW / kScalarsPerW>, + // The stride between iterations. + Shape<GemmConfig_::InstructionShape::kD, 0, kScalarsPerW * Tile::kW>, + // The shape of the instruction. + typename GemmConfig_::InstructionShape> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_> +struct WmmaGemmTileTraitsHelperB {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct WmmaGemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> + : public GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> { + /// The base config. + typedef GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> Base; + + /// The skew. + static int const kSkew = 16 / sizeof(typename Base::MultiplyAddScalar); + /// The shared tile size. + typedef Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kD, + GemmConfig_::OutputTile::kH + kSkew> + Tile; + + /// WMMA matrix + typedef WmmaMatrix<GemmOperand::kB, + MatrixLayout::kRowMajor, + typename Base::MultiplyAddScalar, + typename GemmConfig_::InstructionShape> + WmmaMatrix; + + /// The traits class to build the iterator to store data to shared memory for B^T. + typedef GemmSharedStoreTileAbTraits< + // The pointer. + typename Base::MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Tile, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename Base::GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsB> + SharedStoreTileTraits; + + /// The number of elements loaded in one LDG. + static int const kScalarsPerW = GemmConfig_::InstructionShape::kH * GemmConfig_::Warps::kH; + /// The number of scalars loaded per iteration. + static int const kScalarsPerIteration = Tile::kW * GemmConfig_::InstructionShape::kD; + /// The traits class to build the iterator to load from shared memory for B. + typedef WmmaGemmSharedLoadTileBTraits< + // The layout of the matrix. + MatrixLayout::kRowMajor, + // The pointer. + typename Base::MultiplyAddScalar, + // The output tile size. + Tile, + // The number of warps. + typename GemmConfig_::Warps, + // The strides between warps. + GemmConfig_::InstructionShape::kH, + // The number of iterations to load the data. + Shape<1, 1, GemmConfig_::OutputTile::kH / kScalarsPerW>, + // The stride between iterations. + Shape<kScalarsPerIteration, 0, kScalarsPerW, 0>, + // The shape of the instruction. + typename GemmConfig_::InstructionShape> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename GemmConfig_> +struct WmmaGemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> { + /// The layout. + static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor; + + /// The input scalar. + typedef typename GemmConfig_::ScalarB Scalar; + /// The scalar stored in shared memory. + typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar; + + /// WMMA matrix + typedef WmmaMatrix<GemmOperand::kB, + MatrixLayout::kColumnMajor, + MultiplyAddScalar, + typename GemmConfig_::InstructionShape> + WmmaMatrix; + + /// The traits class to build the iterator to load data from global memory for B^N. + typedef GemmGlobalTileTraits< + // That's B. + GemmOperand::kB, + // A is row-major. + MatrixLayout::kColumnMajor, + // The pointer is float const. + Scalar const, + // The tile has size KxM in GEMM's terminology. + Shape<1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD>, + // The threads are distributed as warps x 32 (the traits may reorganize). + Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>, + // The number of scalars per LDG (LDG.32 or LDG.128, etc). + GemmConfig_::kScalarsPerLdgB> + GlobalTileTraits; + + /// The skew. + static int const kSkew = 16 / sizeof(MultiplyAddScalar); + /// The tile. + typedef Shape<GemmConfig_::kStages, + GemmConfig_::OutputTile::kH, + GemmConfig_::OutputTile::kD + kSkew> + Tile; + + /// The traits class to build the iterator to store data to shared memory for B^N. + typedef GemmSharedStoreTileAbTraits< + // The pointer. + MultiplyAddScalar, + // The tile has size KxM in GEMM's terminology. + Tile, + // The threads are distributed as warps x 32 (the traits may reorganize). + typename GlobalTileTraits::Threads, + // The number of scalars per STS (STS.32 or STS.128, etc). + GemmConfig_::kScalarsPerStsB> + SharedStoreTileTraits; + + /// The number of elements loaded in one LDG. + static int const kScalarsPerW = GemmConfig_::InstructionShape::kH * GemmConfig_::Warps::kH; + /// The traits class to build the iterator to load from shared memory for B. + typedef WmmaGemmSharedLoadTileBTraits< + // The layout of the matrix. + MatrixLayout::kColumnMajor, + // The pointer. + MultiplyAddScalar, + // The tile in shared memory. + Tile, + // The number of warps. + typename GemmConfig_::Warps, + // The strides between warps. + GemmConfig_::InstructionShape::kH * Tile::kW, + // The number of iterations to load the data. + Shape<1, 1, GemmConfig_::OutputTile::kH / kScalarsPerW>, + // The stride between iterations. + Shape<GemmConfig_::InstructionShape::kD, 0, kScalarsPerW * Tile::kW>, + // The shape of the instruction. + typename GemmConfig_::InstructionShape> + SharedLoadTileTraits; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The output tile. + typename OutputTile_, + /// The output type. + typename ScalarC_, + /// The accumulator type. + typename Accumulator_, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_, + /// The number of accumulators per warp. + typename AccumulatorsPerWarp_, + /// The shape of the WMMA instruction. + typename InstructionShape_, + /// The number of halfs loaded in one LDG for A. + int kScalarsPerLdgA_, + /// The number of halfs loaded in one LDG for B. + int kScalarsPerLdgB_, + /// The index. + typename Index_> +struct WmmaGemmTraitsHelper { + /// The WMMA GEMM config. + typedef WmmaGemmConfig<kLayoutA_, + kLayoutB_, + OutputTile_, + ScalarC_, + Accumulator_, + AccumulatorsPerWarp_, + InstructionShape_, + kScalarsPerLdgA_, + kScalarsPerLdgB_> + GemmConfig; + + /// The GEMM config for A. + typedef WmmaGemmTileTraitsHelperA<kLayoutA_, GemmConfig> GemmTileTraitsHelperA; + /// The GEMM config for B. + typedef WmmaGemmTileTraitsHelperB<kLayoutB_, GemmConfig> GemmTileTraitsHelperB; + + /// The iterator to load A from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA::GlobalTileTraits, Index_> + GlobalLoadIteratorA; + /// The default transformer for A. + typedef Copy<typename GlobalLoadIteratorA::Fragment> GlobalTransformerA; + /// The iterator to store A to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperA::SharedStoreTileTraits, + typename GemmTileTraitsHelperA::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorA; + /// The stream to load A from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA> + GlobalLoadStreamA; + + /// The iterator to load B from global memory. + typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperB::GlobalTileTraits, Index_> + GlobalLoadIteratorB; + // The default transformer for B. + typedef Copy<typename GlobalLoadIteratorB::Fragment> GlobalTransformerB; + /// The iterator to store B to shared memory. + typedef TileStoreIterator<typename GemmTileTraitsHelperB::SharedStoreTileTraits, + typename GemmTileTraitsHelperB::SharedStoreTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared> + SharedStoreIteratorB; + /// The stream to load B from global memory to shared memory. + typedef GlobalLoadStream<GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB> + GlobalLoadStreamB; + + /// The iterator to load A from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperA::SharedLoadTileTraits, + typename GemmTileTraitsHelperA::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared, + Index_, + typename GemmTileTraitsHelperA::WmmaMatrix, + IteratorFragment::kWmmaMatrix> + SharedLoadIteratorA; + /// The stream to load A from shared memory. + typedef SharedLoadStream<SharedLoadIteratorA> SharedLoadStreamA; + /// The iterator to load B from shared memory. + typedef TileLoadIterator<typename GemmTileTraitsHelperB::SharedLoadTileTraits, + typename GemmTileTraitsHelperB::SharedLoadTileTraits::Scalar, + IteratorAdvance::kH, + MemorySpace::kShared, + Index_, + typename GemmTileTraitsHelperB::WmmaMatrix, + IteratorFragment::kWmmaMatrix> + SharedLoadIteratorB; + /// The stream to load B from shared memory. + typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB; + + /// The functor to do the multiply-add in the main loop. + typedef typename GemmConfig::MultiplyAdd MultiplyAdd; + /// The object to clear accumulators. + typedef ClearAccumulators<typename MultiplyAdd::ScalarC> ClearAccumulators; + + /// The helper to create the epilogue traits. + typedef WmmaGemmEpilogueTraitsHelper<GemmConfig, EpilogueFunctor_, Index_> EpilogueTraitsHelper; + /// The traits class for the epilogue. + typedef SimplifiedGemmEpilogueTraits<GemmConfig, EpilogueFunctor_, Index_, EpilogueTraitsHelper> + GemmEpilogueTraits; + /// The epilogue. + typedef GemmEpilogue<GemmEpilogueTraits> Epilogue; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <typename OutputTile_, typename DefaultShape_ = Shape<64, 32, 64> > +struct WmmaGemmAccumulatorsPerWarp { + typedef typename ShapeMin<OutputTile_, DefaultShape_>::Shape Shape; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// The layout for A. + MatrixLayout::Kind kLayoutA_, + /// The layout for B. + MatrixLayout::Kind kLayoutB_, + /// The tile size for the GEMM KxNxM. + typename OutputTile_ = Shape<64, 128, 128>, + /// The output type. + typename ScalarC_ = float, + /// The functor to do the math in the epilogue. + typename EpilogueFunctor_ = LinearScaling<ScalarC_>, + /// The accumulator type. + typename Accumulator_ = ScalarC_, + /// The number of accumulators per warp. + typename AccumulatorsPerWarp_ = typename WmmaGemmAccumulatorsPerWarp<OutputTile_>::Shape, + /// The shape of the WMMA instruction. + typename InstructionShape_ = Shape<16, 16, 16>, + /// The number of scalars per LDG for A. + int kScalarsPerLdgA_ = 8, + /// The number of scalars per LDG for B. + int kScalarsPerLdgB_ = 8, + /// The index. + typename Index_ = int, + /// The helper class. + typename Helper_ = WmmaGemmTraitsHelper<kLayoutA_, + kLayoutB_, + OutputTile_, + ScalarC_, + Accumulator_, + EpilogueFunctor_, + AccumulatorsPerWarp_, + InstructionShape_, + kScalarsPerLdgA_, + kScalarsPerLdgB_, + Index_> > +struct WmmaGemmTraits : public GemmTraits< + // The config. + typename Helper_::GemmConfig, + // The stream to load A from global memory to shared memory. + typename Helper_::GlobalLoadStreamA, + // The stream to load B from global memory to shared memory. + typename Helper_::GlobalLoadStreamB, + // The stream to load A from shared memory. + typename Helper_::SharedLoadStreamA, + // The stream to load B from shared memory. + typename Helper_::SharedLoadStreamB, + // The epilogue. + typename Helper_::Epilogue, + // The block swizzle to reorganize the grid. + IdentityBlockSwizzle, + // The index. + Index_, + // The tool used to clear accumulators. + typename Helper_::ClearAccumulators> {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace gemm +} // namespace cutlass + +#endif // defined CUTLASS_USE_WMMA_API |