gemm_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/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 <
    typename ScalarA_,
    typename ScalarB_,
    typename ScalarC_,
    typename ScalarD_,
    typename OutputTile_,
    typename MultiplyAdd_,
    int kScalarsPerLdgA_,
    int kScalarsPerStsA_,
    int kScalarsPerLdsA_,
    int kScalarsPerLdgB_,
    int kScalarsPerStsB_,
    int kScalarsPerLdsB_,
    int kScalarsPerLdgCAndStgD_,
    int kScalarsPerStsD_,
    int kScalarsPerLdsD_,
    int kStages_>

struct GemmConfig {
  //
  typedef ScalarA_ ScalarA;
  typedef ScalarB_ ScalarB;
  typedef ScalarC_ ScalarC;
  typedef ScalarD_ ScalarD;

  typedef OutputTile_ OutputTile;
  typedef MultiplyAdd_ MultiplyAdd;
  typedef typename MultiplyAdd::InstructionShape InstructionShape;
  typedef typename MultiplyAdd::AccumulatorsPerWarp AccumulatorsPerWarp;
  typedef typename MultiplyAdd::Accumulators Accumulators;

  typedef typename ShapeDiv<OutputTile, AccumulatorsPerWarp>::Shape Warps;
  static int const kWarpSize = cutlass::kWarpSize;
  static int const kThreads = ShapeCount<Warps>::kCount * kWarpSize;

  static int const kScalarsPerLdgA = kScalarsPerLdgA_;
  static int const kScalarsPerStsA = kScalarsPerStsA_;
  static int const kScalarsPerLdsA = kScalarsPerLdsA_;

  static int const kScalarsPerLdgB = kScalarsPerLdgB_;
  static int const kScalarsPerStsB = kScalarsPerStsB_;
  static int const kScalarsPerLdsB = kScalarsPerLdsB_;

  static int const kScalarsPerLdgC = kScalarsPerLdgCAndStgD_;

  static int const kScalarsPerStgD = kScalarsPerLdgCAndStgD_;
  static int const kScalarsPerStsD = kScalarsPerStsD_;
  static int const kScalarsPerLdsD = kScalarsPerLdsD_;

  static int const kAccumulatorsPerLdsA = kScalarsPerLdsA / InstructionShape::kD;
  static int const kAccumulatorsPerLdsB = kScalarsPerLdsB / InstructionShape::kD;

  static int const kStages = kStages_;
};


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


template <typename GemmConfig_>
struct GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> {
  static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor;

  typedef typename GemmConfig_::ScalarA Scalar;
  typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar;

  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;

  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;

  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_> {
  static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor;

  typedef typename GemmConfig_::ScalarA Scalar;
  typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar;

  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;

  static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar);
  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.
      128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsA /
          GlobalTileTraits::Threads::kW * kScalarsIn4B>
      SharedStoreTileTraits;

  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_> {
  static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor;

  typedef typename GemmConfig_::ScalarB Scalar;
  typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar;

  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;

  static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar);
  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.
      128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsB /
          GlobalTileTraits::Threads::kW * kScalarsIn4B>
      SharedStoreTileTraits;

  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_> {
  static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor;

  typedef typename GemmConfig_::ScalarB Scalar;
  typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar;

  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;

  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;

  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 GemmConfig_,
    typename GlobalLoadStreamA_,
    typename GlobalLoadStreamB_,
    typename SharedLoadStreamA_,
    typename SharedLoadStreamB_,
    typename Epilogue_,
    typename BlockSwizzle_ = IdentityBlockSwizzle,
    typename Index_ = int,
    typename ClearAccumulators_ = ClearAccumulators<typename GemmConfig_::Accumulators::Scalar> >

struct GemmTraits {
  typedef GemmConfig_ GemmConfig;
  typedef typename GemmConfig::OutputTile OutputTile;

  typedef GlobalLoadStreamA_ GlobalLoadStreamA;
  static MatrixLayout::Kind const kLayoutA = GlobalLoadStreamA::kLayout;
  typedef typename GlobalLoadStreamA_::Scalar ScalarA;

  typedef GlobalLoadStreamB_ GlobalLoadStreamB;
  static MatrixLayout::Kind const kLayoutB = GlobalLoadStreamB::kLayout;
  typedef typename GlobalLoadStreamB_::Scalar ScalarB;

  typedef SharedLoadStreamA_ SharedLoadStreamA;
  typedef SharedLoadStreamB_ SharedLoadStreamB;

  typedef typename GlobalLoadStreamA::SharedStoreStorage SharedStoreStorageA;
  // Btw, make sure we did not messed up with the size of the storage.
  static_assert(sizeof(SharedStoreStorageA) == sizeof(typename SharedLoadStreamA::SharedStorage),
                "");

  typedef typename GlobalLoadStreamB::SharedStoreStorage SharedStoreStorageB;
  // Btw, make sure we did not messed up with the size of the storage.
  static_assert(sizeof(SharedStoreStorageB) == sizeof(typename SharedLoadStreamB::SharedStorage),
                "");

  typedef typename GemmConfig::MultiplyAdd MultiplyAdd;
  typedef Epilogue_ Epilogue;
  typedef typename Epilogue::ScalarC ScalarC;
  typedef typename Epilogue::ScalarD ScalarD;

  typedef BlockSwizzle_ BlockSwizzle;
  typedef Index_ Index;
  typedef ClearAccumulators_ ClearAccumulators;

  struct Params {
    Index m, n, k;
    typename GlobalLoadStreamA::Params global_stream_a;
    typename GlobalLoadStreamB::Params global_stream_b;
    typename SharedLoadStreamA::Params shared_stream_a;
    typename SharedLoadStreamB::Params shared_stream_b;
    typename Epilogue::Params epilogue;

    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(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(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;
  };

  union SharedStorage {
    // The storage for the main loop.
    MainLoopSharedStorage main_loop;
    // The storage for the epilogue.
    typename Epilogue::SharedStorage epilogue;
  };

  struct GlobalLoadStream {
    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)) {}

    CUTLASS_DEVICE void copy() {
      stream_a.copy();
      stream_b.copy();
    }

    CUTLASS_DEVICE void commit() {
      stream_a.commit();
      stream_b.commit();
    }

    CUTLASS_DEVICE void residue(Index k, bool skip_clear = false) {
      stream_a.residue(k, skip_clear);
      stream_b.residue(k, skip_clear);
    }

    GlobalLoadStreamA stream_a;
    GlobalLoadStreamB stream_b;
  };

  struct SharedLoadStream {
    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);
    }

    CUTLASS_DEVICE void copy(int step) {
      stream_a.copy(step, fetched_a[step % 2]);
      stream_b.copy(step, fetched_b[step % 2]);
    }

    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]);
    }

    CUTLASS_DEVICE typename SharedLoadStreamA::Fragment const& fragment_a(int step) const {
      return transformed_a[step % 2];
    }

    CUTLASS_DEVICE typename SharedLoadStreamB::Fragment const& fragment_b(int step) const {
      return transformed_b[step % 2];
    }

    CUTLASS_DEVICE void inc_stage() {
      stream_a.inc_stage();
      stream_b.inc_stage();
    }

    SharedLoadStreamA stream_a;
    typename SharedLoadStreamA::FetchedFragment fetched_a[2];
    typename SharedLoadStreamA::TransformedFragment transformed_a[2];
    SharedLoadStreamB stream_b;
    typename SharedLoadStreamB::FetchedFragment fetched_b[2];
    typename SharedLoadStreamB::TransformedFragment transformed_b[2];
  };

  static CUTLASS_DEVICE void shared_load_fence(bool in_loop) {
    if (SharedLoadStreamA::Iterator::kRequiresLoadFence ||
        SharedLoadStreamB::Iterator::kRequiresLoadFence) {
      __syncthreads();
    }
  }

  static CUTLASS_DEVICE void shared_store_fence(bool in_loop) { __syncthreads(); }
};


template <typename GemmTileTraitsHelperA_, typename GemmTileTraitsHelperB_, typename Index_>
struct SimplifiedGemmTraitsHelper {
  typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA_::GlobalTileTraits, Index_>
      GlobalLoadIteratorA;
  typedef Copy<typename GlobalLoadIteratorA::Fragment> 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;
  typedef Copy<typename GlobalLoadIteratorB::Fragment> 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_::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kShared>
      SharedLoadIteratorA;
  typedef SharedLoadStream<SharedLoadIteratorA> SharedLoadStreamA;
  typedef TileLoadIterator<typename GemmTileTraitsHelperB_::SharedLoadTileTraits,
                           typename GemmTileTraitsHelperB_::Scalar,
                           IteratorAdvance::kH,
                           MemorySpace::kShared>
      SharedLoadIteratorB;
  typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB;
};


template <
    MatrixLayout::Kind kLayoutA_,
    MatrixLayout::Kind kLayoutB_,
    typename GemmConfig_,
    typename Epilogue_,
    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