diff options
Diffstat (limited to 'cutlass-example/cutlass/gemm/igemm_epilogue.h')
| -rw-r--r-- | cutlass-example/cutlass/gemm/igemm_epilogue.h | 320 |
1 files changed, 0 insertions, 320 deletions
diff --git a/cutlass-example/cutlass/gemm/igemm_epilogue.h b/cutlass-example/cutlass/gemm/igemm_epilogue.h deleted file mode 100644 index 0d69980..0000000 --- a/cutlass-example/cutlass/gemm/igemm_epilogue.h +++ /dev/null @@ -1,320 +0,0 @@ -/*************************************************************************************************** - * Copyright (c) 2017-2018, NVIDIA CORPORATION. All rights reserved. - * - * Redistribution and use in source and binary forms, with or without modification, are permitted - * provided that the following conditions are met: - * * Redistributions of source code must retain the above copyright notice, this list of - * conditions and the following disclaimer. - * * Redistributions in binary form must reproduce the above copyright notice, this list of - * conditions and the following disclaimer in the documentation and/or other materials - * provided with the distribution. - * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used - * to endorse or promote products derived from this software without specific prior written - * permission. - * - * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR - * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND - * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE - * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, - * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; - * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, - * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE - * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - * - **************************************************************************************************/ -/*! \file - \brief 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 |
