leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

CUTLASS 2.2 (#96)

Browse Source 
Adds support for NVIDIA Ampere Architecture features. CUDA 11 Toolkit recommended.
This commit is contained in:
Andrew Kerr
2020-06-08 16:17:35 -07:00
committed by GitHub
parent e33d90b361
commit 86931fef85
584 changed files with 51080 additions and 3373 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/device/default_gemm_configuration.h
+337 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
@@ -422,6 +422,342 @@ struct DefaultGemmConfiguration<
using Operator = arch::OpMultiplyAddSaturate;
};
////////////////////////////////////////////////////////////////////////////////
template <
typename ElementC>
struct DefaultGemmConfiguration<
arch::OpClassTensorOp,
arch::Sm75,
uint1b_t,
uint1b_t,
ElementC,
int32_t> {
static int const kAlignmentA = 128 / sizeof_bits<uint1b_t>::value;
static int const kAlignmentB = 128 / sizeof_bits<uint1b_t>::value;
using ThreadblockShape = GemmShape<128, 256, 512>;
using WarpShape = GemmShape<64, 64, 512>;
using InstructionShape = GemmShape<8, 8, 128>;
static int const kStages = 2;
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
using Operator = arch::OpXorPopc;
};
////////////////////////////////////////////////////////////////////////////////
template <typename ElementA, typename ElementB, typename ElementC,
typename ElementAccumulator>
struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm80, 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, 64>;
using WarpShape = GemmShape<64, 64, 64>;
using InstructionShape = GemmShape<16, 8, 16>;
static int const kStages = 3;
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,
typename ElementAccumulator>
struct DefaultGemmConfiguration<arch::OpClassTensorOp, arch::Sm80, double,
double, ElementC, ElementAccumulator> {
static int const kAlignmentA = 1;
static int const kAlignmentB = 1;
using ThreadblockShape = GemmShape<128, 256, 64>;
using WarpShape = GemmShape<64, 64, 64>;
using InstructionShape = GemmShape<16, 8, 16>;
static int const kStages = 3;
using EpilogueOutputOp = epilogue::thread::LinearCombination<
ElementC, 128 / sizeof_bits<ElementC>::value, ElementAccumulator,
ElementAccumulator>;
using Operator = arch::OpMultiplyAdd;
};
template <>
struct DefaultGemmConfiguration<
arch::OpClassTensorOp,
arch::Sm80,
complex<double>,
complex<double>,
complex<double>,
complex<double>
> {
static int const kAlignmentA = 1;
static int const kAlignmentB = 1;
using ThreadblockShape = GemmShape<64, 64, 16>;
using WarpShape = GemmShape<32, 32, 16>;
using InstructionShape = GemmShape<8, 8, 4>;
static int const kStages = 3;
using EpilogueOutputOp = epilogue::thread::LinearCombination<
complex<double>, 1, complex<double>,
complex<double>>;
using Operator = arch::OpMultiplyAddComplex;
};
////////////////////////////////////////////////////////////////////////////////
template <
typename ElementC>
struct DefaultGemmConfiguration<
arch::OpClassTensorOp,
arch::Sm80,
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<16, 8, 32>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 32>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 32>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 32>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 64>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 64>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 64>;
static int const kStages = 3;
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::Sm80,
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<16, 8, 64>;
static int const kStages = 3;
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::Sm80,
uint1b_t,
uint1b_t,
ElementC,
int32_t> {
static int const kAlignmentA = 128 / sizeof_bits<uint1b_t>::value;
static int const kAlignmentB = 128 / sizeof_bits<uint1b_t>::value;
using ThreadblockShape = GemmShape<128, 256, 512>;
using WarpShape = GemmShape<64, 64, 512>;
using InstructionShape = GemmShape<16, 8, 256>;
static int const kStages = 3;
using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp<
ElementC, 128 / sizeof_bits<ElementC>::value, int32_t, float>;
using Operator = arch::OpMultiplyAdd;
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
} // namespace device
} // namespace gemm
include/cutlass/gemm/device/gemm.h
+2 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
@@ -193,7 +193,7 @@ template <
ElementAccumulator_>::EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_ =
typename threadblock::GemmCohortThreadblockSwizzle<LayoutA_, LayoutB_>,
typename threadblock::GemmIdentityThreadblockSwizzle<>,
/// Number of stages used in the pipelined mainloop
int Stages =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
include/cutlass/gemm/device/gemm_array.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
include/cutlass/gemm/device/gemm_batched.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
include/cutlass/gemm/device/gemm_complex.h
+5 -3
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
@@ -192,7 +192,7 @@ template <
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle,
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
/// Number of stages used in the pipelined mainloop
int Stages =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
@@ -202,6 +202,7 @@ template <
/// Complex elementwise transformation on B operand
ComplexTransform TransformB = ComplexTransform::kNone,
/// Multiply-add operator
// (selects complex or gaussian complex)
typename Operator_ = arch::OpMultiplyAddComplex,
/// If true, kernel supports split-K with serial reduction
bool SplitKSerial = false
@@ -506,6 +507,7 @@ template <
/// Complex elementwise transformation on B operand
ComplexTransform TransformB,
/// Multiply-add operator
// (selects complex or gaussian complex)
typename Operator_,
/// If true, kernel supports split-K as a serial reduction
bool SplitKSerial
@@ -571,8 +573,8 @@ public:
EpilogueOutputOp,
ThreadblockSwizzle,
Stages,
TransformA,
TransformB,
TransformA,
Operator,
SplitKSerial
>;
include/cutlass/gemm/device/gemm_splitk_parallel.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
include/cutlass/gemm/device/gemm_universal.h
+2 -2
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
@@ -89,7 +89,7 @@ template <
OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_,
ElementAccumulator_>::EpilogueOutputOp,
/// Threadblock-level swizzling operator
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle,
typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
/// Number of stages used in the pipelined mainloop
int Stages =
DefaultGemmConfiguration<OperatorClass_, ArchTag_, ElementA_, ElementB_,
include/cutlass/gemm/device/gemm_universal_adapter.h
+99 -17
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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:
@@ -41,14 +41,78 @@ namespace device {
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
template <
typename ElementA_,
typename LayoutA_,
ComplexTransform TransformA,
int AlignmentA,
typename ElementB_,
typename LayoutB_,
ComplexTransform TransformB,
int AlignmentB,
typename LayoutC_,
bool Transpose
>
struct MapArguments {
using ElementA = ElementA_;
using LayoutA = LayoutA_;
static ComplexTransform const kTransformA = TransformA;
static int const kAlignmentA = AlignmentA;
using ElementB = ElementB_;
using LayoutB = LayoutB_;
static ComplexTransform const kTransformB = TransformB;
static int const kAlignmentB = AlignmentB;
using LayoutC = LayoutC_;
};
template <
typename ElementA_,
typename LayoutA_,
ComplexTransform TransformA,
int AlignmentA,
typename ElementB_,
typename LayoutB_,
ComplexTransform TransformB,
int AlignmentB,
typename LayoutC_
>
struct MapArguments<
ElementA_,
LayoutA_,
TransformA,
AlignmentA,
ElementB_,
LayoutB_,
TransformB,
AlignmentB,
LayoutC_,
true
> {
using ElementA = ElementB_;
using LayoutA = typename layout::LayoutTranspose<LayoutB_>::type;
static ComplexTransform const kTransformA = TransformB;
static int const kAlignmentA = AlignmentB;
using ElementB = ElementA_;
using LayoutB = typename layout::LayoutTranspose<LayoutA_>::type;
static ComplexTransform const kTransformB = TransformA;
static int const kAlignmentB = AlignmentA;
using LayoutC = typename layout::LayoutTranspose<LayoutC_>::type;
};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename GemmKernel_>
class GemmUniversalAdapter {
public:
using GemmKernel = GemmKernel_;
static_assert(std::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value,
"Universal adapter expects the kernel to be row-major and transposes its arguments.");
static bool const kInternalTranspose =
std::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
using ThreadblockShape = typename GemmKernel::Mma::Shape;
using WarpShape = typename GemmKernel::WarpShape;
@@ -56,26 +120,39 @@ public:
using OperatorClass = typename GemmKernel::OperatorClass;
using ArchTag = typename GemmKernel::ArchTag;
// Type, layout, and complex transform deliberately exchanged with B
using ElementA = typename GemmKernel::ElementB;
using LayoutA = typename layout::LayoutTranspose<typename GemmKernel::LayoutB>::type;
using TensorRefA = TensorRef<ElementA const, LayoutA>;
static ComplexTransform const kTransformA = GemmKernel::kTransformB;
using MapArguments = detail::MapArguments<
typename GemmKernel::ElementA,
typename GemmKernel::LayoutA,
GemmKernel::kTransformA,
GemmKernel::kAlignmentA,
typename GemmKernel::ElementB,
typename GemmKernel::LayoutB,
GemmKernel::kTransformB,
GemmKernel::kAlignmentB,
typename GemmKernel::LayoutC,
kInternalTranspose
>;
using ElementA = typename MapArguments::ElementA;
using LayoutA = typename MapArguments::LayoutA;
static ComplexTransform const kTransformA = MapArguments::kTransformA;
static int const kAlignmentA = GemmKernel::kAlignmentA;
// Type, layout, and complex transform deliberately exchanged with A
using ElementB = typename GemmKernel::ElementA;
using LayoutB = typename layout::LayoutTranspose<typename GemmKernel::LayoutA>::type;
using TensorRefB = TensorRef<ElementB const, LayoutB>;
static ComplexTransform const kTransformB = GemmKernel::kTransformA;
using ElementB = typename MapArguments::ElementB;
using LayoutB = typename MapArguments::LayoutB;
static ComplexTransform const kTransformB = MapArguments::kTransformB;
static int const kAlignmentB = GemmKernel::kAlignmentB;
using ElementC = typename GemmKernel::ElementC;
using LayoutC = cutlass::layout::ColumnMajor;
using LayoutC = typename MapArguments::LayoutC;
static int const kAlignmentC = GemmKernel::kAlignmentC;
using TensorRefA = TensorRef<ElementA const, LayoutA>;
using TensorRefB = TensorRef<ElementB const, LayoutB>;
using TensorRefC = TensorRef<ElementC const, LayoutC>;
using TensorRefD = TensorRef<ElementC, LayoutC>;
static int const kAlignmentC = GemmKernel::kAlignmentC;
using ElementAccumulator = typename GemmKernel::Mma::Policy::Operator::ElementC;
@@ -99,7 +176,12 @@ public:
/// Helper to construct a transposed equivalent for the underying GEMM operator
static Arguments to_underlying_arguments(Arguments const &args) {
return args.transposed_problem();
if (kInternalTranspose) {
return args.transposed_problem();
}
else {
return args;
}
}
/// Determines whether the GEMM can execute the given problem.
include/cutlass/gemm/device/gemm_universal_base.h
+1 -1
View File
@@ -1,5 +1,5 @@
/***************************************************************************************************
* Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.
* Copyright (c) 2017-2020, 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: