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, 320 insertions, 0 deletions
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 |