default_gemm_configuration.h

Go to the documentation of this file.

/***************************************************************************************************
 * Copyright (c) 2017-2019, 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/cutlass.h"
#include "cutlass/numeric_types.h"
#include "cutlass/arch/arch.h"
#include "cutlass/arch/mma.h"
#include "cutlass/arch/wmma.h"

#include "cutlass/gemm/gemm.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/epilogue/thread/linear_combination_clamp.h"


namespace cutlass {
namespace gemm {
namespace device {


template <
  typename OperatorClass,
  typename ArchTag,
  typename ElementA, 
  typename ElementB, 
  typename ElementC,
  typename ElementAccumulator
>
struct DefaultGemmConfiguration;


template <
  typename ArchTag,
  typename ElementA, 
  typename ElementB, 
  typename ElementC, 
  typename ElementAccumulator>
struct DefaultGemmConfiguration<
  arch::OpClassSimt, 
  ArchTag,
  ElementA, 
  ElementB, 
  ElementC, 
  ElementAccumulator> {
  
  static int const kAlignmentA = 1;
  static int const kAlignmentB = 1;
  using ThreadblockShape = GemmShape<128, 128, 8>;
  using WarpShape = GemmShape<32, 64, 8>;
  using InstructionShape = GemmShape<1, 1, 1>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombination<
    ElementC,
    1,
    ElementAccumulator,
    ElementAccumulator
  >;

  using Operator = arch::OpMultiplyAdd;
};


template < 
  typename ArchTag,
  typename ElementC>
struct DefaultGemmConfiguration<arch::OpClassSimt, ArchTag, int8_t, int8_t, ElementC, int32_t> {
  
  static int const kAlignmentA = 4;
  static int const kAlignmentB = 4;
  using ThreadblockShape = GemmShape<128, 128, 32>;
  using WarpShape = GemmShape<32, 64, 32>;
  using InstructionShape = GemmShape<1, 1, 4>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
    ElementC,
    1,
    int32_t,
    float
  >;

  using Operator = arch::OpMultiplyAdd;
};


template <
  typename ArchTag,
  typename ElementA, 
  typename ElementB, 
  typename ElementC, 
  typename ElementAccumulator>
struct DefaultGemmConfiguration<
  arch::OpClassWmmaTensorOp, 
  ArchTag,
  ElementA, 
  ElementB, 
  ElementC, 
  ElementAccumulator> {
  
  static int const kAlignmentA = 128 / sizeof_bits<ElementA>::value;
  static int const kAlignmentB = 128 / sizeof_bits<ElementB>::value;

  static int const kStages = 2;
  
  using EpilogueOutputOp = epilogue::thread::LinearCombination<
    ElementC,
    128 / sizeof_bits<ElementC>::value,
    ElementAccumulator,
    ElementAccumulator
  >;

  using Operator = arch::OpMultiplyAdd;
};


template <
  typename ElementA, 
  typename ElementB, 
  typename ElementC, 
  typename ElementAccumulator>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm70,
  ElementA, 
  ElementB, 
  ElementC, 
  ElementAccumulator> {
  
  static int const kAlignmentA = 128 / sizeof_bits<ElementA>::value;
  static int const kAlignmentB = 128 / sizeof_bits<ElementB>::value;

  using ThreadblockShape = GemmShape<128, 256, 32>;
  using WarpShape = GemmShape<64, 64, 32>;
  using InstructionShape = GemmShape<16, 16, 4>;
  static int const kStages = 2;
  
  using EpilogueOutputOp = epilogue::thread::LinearCombination<
    ElementC,
    128 / sizeof_bits<ElementC>::value,
    ElementAccumulator,
    ElementAccumulator
  >;

  using Operator = arch::OpMultiplyAdd;
};


template <
  typename ElementA, 
  typename ElementB, 
  typename ElementC, 
  typename ElementAccumulator>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75,
  ElementA, 
  ElementB, 
  ElementC, 
  ElementAccumulator> {

  static int const kAlignmentA = 128 / sizeof_bits<ElementA>::value;
  static int const kAlignmentB = 128 / sizeof_bits<ElementA>::value;
  using ThreadblockShape = GemmShape<128, 256, 32>;
  using WarpShape = GemmShape<64, 64, 32>;
  using InstructionShape = GemmShape<16, 8, 8>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombination<
    ElementC,
    128 / sizeof_bits<ElementC>::value,
    ElementAccumulator,
    ElementAccumulator
  >;

  using Operator = typename platform::conditional<
      (platform::is_same<ElementA, int8_t>::value ||
       platform::is_same<ElementA, int4b_t>::value ||
       platform::is_same<ElementA, uint8_t>::value ||
       platform::is_same<ElementA, uint4b_t>::value),
      arch::OpMultiplyAddSaturate, arch::OpMultiplyAdd>::type;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  int8_t, 
  int8_t, 
  ElementC, 
  int32_t> {
  
  static int const kAlignmentA = 128 / sizeof_bits<int8_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<int8_t>::value;

  using ThreadblockShape = GemmShape<128, 256, 64>;
  using WarpShape = GemmShape<64, 64, 64>;
  using InstructionShape = GemmShape<8, 8, 16>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  int8_t, 
  uint8_t, 
  ElementC, 
  int32_t> {
  
  static int const kAlignmentA = 128 / sizeof_bits<int8_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<uint8_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 64>;
  using WarpShape = GemmShape<64, 64, 64>;
  using InstructionShape = GemmShape<8, 8, 16>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  uint8_t, 
  int8_t, 
  ElementC, 
  int32_t> {
  
  static int const kAlignmentA = 128 / sizeof_bits<uint8_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<int8_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 64>;
  using WarpShape = GemmShape<64, 64, 64>;
  using InstructionShape = GemmShape<8, 8, 16>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  uint8_t, 
  uint8_t, 
  ElementC, 
  int32_t> {
  
  static int const kAlignmentA = 128 / sizeof_bits<uint8_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<uint8_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 64>;
  using WarpShape = GemmShape<64, 64, 64>;
  using InstructionShape = GemmShape<8, 8, 16>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  int4b_t, 
  int4b_t, 
  ElementC, 
  int32_t> {
   
  static int const kAlignmentA = 128 / sizeof_bits<int4b_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<int4b_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 128>;
  using WarpShape = GemmShape<64, 64, 128>;
  using InstructionShape = GemmShape<8, 8, 32>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  int4b_t, 
  uint4b_t, 
  ElementC, 
  int32_t> {
    
  static int const kAlignmentA = 128 / sizeof_bits<int4b_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<uint4b_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 128>;
  using WarpShape = GemmShape<64, 64, 128>;
  using InstructionShape = GemmShape<8, 8, 32>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  uint4b_t, 
  int4b_t, 
  ElementC, 
  int32_t> {
  
  static int const kAlignmentA = 128 / sizeof_bits<uint4b_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<int4b_t>::value;

  using ThreadblockShape = GemmShape<128, 256, 128>;
  using WarpShape = GemmShape<64, 64, 128>;
  using InstructionShape = GemmShape<8, 8, 32>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};


template < 
  typename ElementC>
struct DefaultGemmConfiguration<
  arch::OpClassTensorOp, 
  arch::Sm75, 
  uint4b_t, 
  uint4b_t, 
  ElementC, 
  int32_t> {
   
  static int const kAlignmentA = 128 / sizeof_bits<uint4b_t>::value;
  static int const kAlignmentB = 128 / sizeof_bits<uint4b_t>::value;
 
  using ThreadblockShape = GemmShape<128, 256, 128>;
  using WarpShape = GemmShape<64, 64, 128>;
  using InstructionShape = GemmShape<8, 8, 32>;
  static int const kStages = 2;

  using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
      ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;

  using Operator = arch::OpMultiplyAddSaturate;
};

} // namespace device
} // namespace gemm
} // namespace cutlass