igemm_traits.h

Go to the documentation of this file.

/***************************************************************************************************
 * 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.
 *
 **************************************************************************************************/
#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 <
    typename OutputTile_,
    typename ScalarD_,
    typename AccumulatorsPerThread_>
struct IgemmConfig
    : public GemmConfig<
          int8_t,
          int8_t,
          ScalarD_,
          ScalarD_,
          OutputTile_,
          ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, int8_t, int8_t, int>,
          4,
          4,
          16,
          4,
          4,
          16,
          1,
          4,
          1,
          2> {};


template <typename OutputTile_, typename AccumulatorsPerThread_>
struct IgemmConfig<OutputTile_, int8_t, AccumulatorsPerThread_>
    : public GemmConfig<
          int8_t,
          int8_t,
          int8_t,
          int8_t,
          OutputTile_,
          ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, int8_t, int8_t, int>,
          4,
          4,
          16,
          4,
          4,
          16,
          4,
          4,
          4,
          2> {};


template <enum MatrixLayout::Kind kLayout_, typename GemmConfig_>
struct IgemmTileTraitsHelperA : public GemmTileTraitsHelperA<kLayout_, GemmConfig_> {};


template <typename GemmConfig_>
struct IgemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_>
    : public GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> {
  typedef GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> Base;

  static int const kScalarsPerStsA = 16;

  typedef IgemmContiguousGlobalTileTraits<
      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).
      4>
      GlobalTileTraits;

  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 <enum MatrixLayout::Kind kLayout_, typename GemmConfig_>
struct IgemmTileTraitsHelperB : public GemmTileTraitsHelperB<kLayout_, GemmConfig_> {};


template <typename GemmConfig_>
struct IgemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_>
    : public GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> {
  typedef GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> Base;

  static int const kScalarsPerStsB = 16;

  typedef IgemmContiguousGlobalTileTraits<
      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).
      4>
      GlobalTileTraits;

  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 <
    MatrixLayout::Kind kLayoutA_,
    MatrixLayout::Kind kLayoutB_,
    typename OutputTile_,
    typename ScalarD_,
    typename EpilogueFunctor_,
    typename AccumulatorsPerThread_ = Shape<32, 8, 8>,
    typename Index_ = int>
struct IgemmTraitsHelper {
  typedef IgemmConfig<OutputTile_, ScalarD_, AccumulatorsPerThread_> GemmConfig;
  typedef IgemmTileTraitsHelperA<kLayoutA_, GemmConfig> GemmTileTraitsHelperA;
  typedef IgemmTileTraitsHelperB<kLayoutB_, GemmConfig> GemmTileTraitsHelperB;

  typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA::GlobalTileTraits, Index_>
      GlobalLoadIteratorA;
  typedef typename IgemmTransformerA<GemmTileTraitsHelperA::kLayout,
                                     GlobalLoadIteratorA>::Transformer GlobalTransformerA;
  typedef TileStoreIterator<typename GemmTileTraitsHelperA::SharedStoreTileTraits,
                            typename GemmTileTraitsHelperA::SharedStoreTileTraits::Scalar,
                            IteratorAdvance::kH,
                            MemorySpace::kShared>
      SharedStoreIteratorA;
  typedef GlobalLoadStream<GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA>
      GlobalLoadStreamA;

  typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperB::GlobalTileTraits, Index_>
      GlobalLoadIteratorB;
  // The default transformer for B.
  typedef typename IgemmTransformerB<GemmTileTraitsHelperB::kLayout,
                                     GlobalLoadIteratorB>::Transformer GlobalTransformerB;
  typedef TileStoreIterator<typename GemmTileTraitsHelperB::SharedStoreTileTraits,
                            typename GemmTileTraitsHelperB::SharedStoreTileTraits::Scalar,
                            IteratorAdvance::kH,
                            MemorySpace::kShared>
      SharedStoreIteratorB;
  typedef GlobalLoadStream<GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB>
      GlobalLoadStreamB;

  typedef TileLoadIterator<typename GemmTileTraitsHelperA::SharedLoadTileTraits,
                           typename GemmTileTraitsHelperA::SharedLoadTileTraits::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kShared>
      SharedLoadIteratorA;
  typedef SharedLoadStream<SharedLoadIteratorA, Copy<typename SharedLoadIteratorA::Fragment> >
      SharedLoadStreamA;
  typedef TileLoadIterator<typename GemmTileTraitsHelperB::SharedLoadTileTraits,
                           typename GemmTileTraitsHelperB::SharedLoadTileTraits::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kShared>
      SharedLoadIteratorB;
  typedef SharedLoadStream<SharedLoadIteratorB, Copy<typename SharedLoadIteratorB::Fragment> >
      SharedLoadStreamB;

  typedef typename GemmConfig::MultiplyAdd MultiplyAdd;
  typedef ClearAccumulators<typename MultiplyAdd::ScalarC> ClearAccumulators;

  typedef IgemmEpilogue<IgemmEpilogueTraits<GemmConfig, EpilogueFunctor_> > Epilogue;
};


template <typename ScalarD_>
struct IgemmEpilogueScalar {
  typedef float Scalar;
};

template <>
struct IgemmEpilogueScalar<int> {
  typedef int Scalar;
};


template <
    MatrixLayout::Kind kLayoutA_,
    MatrixLayout::Kind kLayoutB_,
    typename OutputTile_ = Shape<32, 128, 128>,
    typename ScalarD_ = int,
    typename EpilogueFunctor_ = LinearScaling<typename IgemmEpilogueScalar<ScalarD_>::Scalar>,
    typename AccumulatorsPerThread_ = Shape<32, 8, 8>,
    typename Index_ = int,
    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