hgemm_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/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 <
    typename OutputTile_,
    typename AccumulatorsPerThread_,
    int kScalarsPerLdgA_ = 2,
    int kScalarsPerLdgB_ = 2>
struct HgemmConfig
    : public GemmConfig<
          half,
          half,
          half,
          half,
          OutputTile_,
          ThreadMultiplyAdd<AccumulatorsPerThread_, Shape<1, 4, 8>, half, half, half>,
          kScalarsPerLdgA_,
          kScalarsPerLdgA_,
          8,
          kScalarsPerLdgB_,
          kScalarsPerLdgB_,
          8,
          2,
          8,
          2,
          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_> {
  typedef GemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> Base;

  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;

  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.
      128 / sizeof(half) / GlobalTileTraits::Threads::kW / 2>
      SharedStoreTileTraits;

  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_> {
  typedef GemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> Base;

  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;

  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.
      128 / sizeof(half) / GlobalTileTraits::Threads::kW / 2>
      SharedStoreTileTraits;

  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 <
    MatrixLayout::Kind kLayoutA_,
    MatrixLayout::Kind kLayoutB_,
    typename OutputTile_,
    typename EpilogueFunctor_,
    typename AccumulatorsPerThread_ = Shape<32, 8, 8>,
    int kScalarsPerLdgA_ = 2,
    int kScalarsPerLdgB_ = 2,
    typename Index_ = int>
struct HgemmTraitsHelper {
  typedef HgemmConfig<OutputTile_, AccumulatorsPerThread_, kScalarsPerLdgA_, kScalarsPerLdgB_>
      GemmConfig;
  typedef HgemmTileTraitsHelperA<kLayoutA_, GemmConfig> GemmTileTraitsHelperA;
  typedef HgemmTileTraitsHelperB<kLayoutB_, GemmConfig> GemmTileTraitsHelperB;

  typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA::GlobalTileTraits, Index_>
      GlobalLoadIteratorA;
  typedef typename HgemmTransformerA<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 HgemmTransformerB<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> SharedLoadStreamA;
  typedef TileLoadIterator<typename GemmTileTraitsHelperB::SharedLoadTileTraits,
                           typename GemmTileTraitsHelperB::SharedLoadTileTraits::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kShared>
      SharedLoadIteratorB;
  typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB;

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

  typedef SimplifiedGemmEpilogueTraits<GemmConfig, EpilogueFunctor_, Index_> GemmEpilogueTraits;
  typedef GemmEpilogue<GemmEpilogueTraits> Epilogue;
};


template <
    MatrixLayout::Kind kLayoutA_,
    MatrixLayout::Kind kLayoutB_,
    typename OutputTile_ = Shape<8, 128, 128>,
    typename EpilogueFunctor_ = LinearScaling<half>,
    typename AccumulatorsPerThread_ = Shape<8, 8, 16>,
    int kScalarsPerLdgA_ = 2,
    int kScalarsPerLdgB_ = 2,
    typename Index_ = int,
    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