Back to Cutlass

CUTLASS: default_epilogue_tensor_op.h Source File

docs/default__epilogue__tensor__op_8h_source.html

4.4.228.6 KB
Original Source

| | CUTLASS

CUDA Templates for Linear Algebra Subroutines and Solvers |

default_epilogue_tensor_op.h

[Go to the documentation of this file.](default epilogue tensor__op_8h.html)

1 /***************************************************************************************************

2 * Copyright (c) 2017-2019, NVIDIA CORPORATION. All rights reserved.

3 *

4 * Redistribution and use in source and binary forms, with or without modification, are permitted

5 * provided that the following conditions are met:

6 * * Redistributions of source code must retain the above copyright notice, this list of

7 * conditions and the following disclaimer.

8 * * Redistributions in binary form must reproduce the above copyright notice, this list of

9 * conditions and the following disclaimer in the documentation and/or other materials

10 * provided with the distribution.

11 * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used

12 * to endorse or promote products derived from this software without specific prior written

13 * permission.

14 *

15 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR

16 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND

17 * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE

18 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,

19 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;

20 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,

21 * STRICT LIABILITY, OR TOR (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

22 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

23 *

24 **************************************************************************************************/

33 #pragma once

34

35 #include "cutlass/cutlass.h"

36 #include "cutlass/numeric_types.h"

37 #include "cutlass/array.h"

38

39 #include "cutlass/gemm/gemm.h"

40

41 #include "cutlass/epilogue/thread/linear_combination.h"

42 #include "cutlass/epilogue/thread/conversion_op.h"

43 #include "cutlass/epilogue/thread/reduction_op.h"

44

45 #include "[cutlass/transform/threadblock/regular_tile_iterator_pitch_linear.h](regular tile iterator pitch linear_8h.html)"

46

47 #include "[cutlass/epilogue/warp/fragment_iterator_tensor_op.h](fragment iterator tensor__op_8h.html)"

48 #include "[cutlass/epilogue/warp/tile_iterator_tensor_op.h](tile iterator tensor__op_8h.html)"

49 #include "[cutlass/epilogue/threadblock/default_thread_map_tensor_op.h](default thread map tensor op_8h.html)"

50 #include "[cutlass/epilogue/threadblock/predicated_tile_iterator.h](epilogue_2threadblock_2predicated tile iterator_8h.html)"

51 #include "[cutlass/epilogue/threadblock/shared_load_iterator.h](shared load iterator_8h.html)"

52

53 #include "cutlass/epilogue/threadblock/epilogue.h"

54 #include "cutlass/epilogue/threadblock/interleaved_epilogue.h"

55

57

58 namespace cutlass {

59 namespace epilogue {

60 namespace threadblock {

61

63

65 template <

66typename Shape_,

67typename WarpMmaTensorOp_,

68int PartitionsK,

69typename OutputOp_,

70int ElementsPerAccess

71 >

72 struct DefaultEpilogueTensorOp {

73

74using Shape = Shape_;

75using WarpMmaTensorOp = WarpMmaTensorOp_;

76static int const kPartitionsK = PartitionsK;

77using OutputOp = OutputOp_;

78static int const kElementsPerAccess = ElementsPerAccess;

79

80using ElementOutput = typename OutputOp::ElementOutput;

81using LayoutC = typename WarpMmaTensorOp::LayoutC;

82using ElementAccumulator = typename WarpMmaTensorOp::ElementC;

83

84//

85// Thread map

86//

87

88using OutputTileThreadMap = typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp<

89Shape,

90typename WarpMmaTensorOp::Shape,

91kPartitionsK,

92ElementOutput,

93 kElementsPerAccess

94 >::Type;

95

96using OutputTileIterator = cutlass::epilogue::threadblock::PredicatedTileIterator<

97OutputTileThreadMap,

98 ElementOutput

99 >;

100

101using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorTensorOp<

102typename WarpMmaTensorOp::Shape,

103typename WarpMmaTensorOp::Policy::Operator::Shape,

104typename WarpMmaTensorOp::Policy::Operator::ElementC,

105typename WarpMmaTensorOp::Policy::Operator::FragmentC,

106LayoutC

107 >;

108

109using WarpTileIterator = cutlass::epilogue::warp::TileIteratorTensorOp<

110typename WarpMmaTensorOp::Shape,

111typename WarpMmaTensorOp::Policy::Operator::Shape,

112ElementAccumulator,

113LayoutC

114 >;

115

116using SharedLoadIterator = cutlass::epilogue::threadblock::SharedLoadIterator<

117typename OutputTileThreadMap::CompactedThreadMap,

118 ElementAccumulator

119 >;

120

122using Padding = cutlass::MatrixShape<0, 64 / sizeof_bits<ElementAccumulator>::value * 4>;

123

124//

125// Define the epilogue

126//

127using Epilogue = cutlass::epilogue::threadblock::Epilogue<

128Shape,

129WarpMmaTensorOp,

130kPartitionsK,

131OutputTileIterator,

132AccumulatorFragmentIterator,

133WarpTileIterator,

134SharedLoadIterator,

135OutputOp,

136Padding

137 >;

138 };

139

141

144 template <typename Shape_, typename WarpMmaTensorOp_, int PartitionsK,

145typename OutputOp_, int ElementsPerAccess, int InterleavedK,

146bool IsBetaZero = false, bool isSplitK = false>

147 struct DefaultInterleavedEpilogueTensorOp {

148using Shape = Shape_;

149using WarpMmaTensorOp = WarpMmaTensorOp_;

150static int const kPartitionsK = PartitionsK;

151using OutputOp = OutputOp_;

152static int const kElementsPerAccess = ElementsPerAccess;

153

154using ElementOutput = typename OutputOp::ElementOutput;

155using LayoutC = typename WarpMmaTensorOp::LayoutC;

156using ElementAccumulator = typename WarpMmaTensorOp::ElementC;

157

158//

159// Thread map

160//

161using OutputTileThreadMap = typename cutlass::epilogue::threadblock::

162 DefaultInterleavedThreadMapTensorOp<

163Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput,

164kElementsPerAccess, InterleavedK>::Type;

165

166using OutputTileIterator =

167cutlass::epilogue::threadblock::InterleavedPredicatedTileIterator<

168OutputTileThreadMap, ElementOutput, InterleavedK>;

169

170using AccumulatorFragmentIterator =

171cutlass::epilogue::warp::FragmentIteratorTensorOp<

172typename WarpMmaTensorOp::Shape,

173typename WarpMmaTensorOp::Policy::Operator::Shape,

174typename WarpMmaTensorOp::Policy::Operator::ElementC,

175typename WarpMmaTensorOp::Policy::Operator::FragmentC,

176LayoutC>;

177

178//

179// Define the epilogue

180//

181using Epilogue = cutlass::epilogue::threadblock::InterleavedEpilogue<

182Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator,

183AccumulatorFragmentIterator, OutputOp, InterleavedK, IsBetaZero>;

184 };

185

187 } // namespace threadblock

188 } // namespace epilogue

189 } // namespace cutlass

190

[default_thread_map_tensor_op.h](default thread map tensor op_8h.html)

cutlass::MatrixShape

Describes the size of a matrix tile.

Definition: matrix_shape.h:42

[regular_tile_iterator_pitch_linear.h](regular tile iterator pitch linear_8h.html)

Templates implementing loading of tiles from pitch-linear rank=2 tensors.

cutlass

Definition: aligned_buffer.h:35

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::LayoutC

typename WarpMmaTensorOp::LayoutC LayoutC

Definition: default_epilogue_tensor_op.h:81

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::ElementOutput

typename OutputOp::ElementOutput ElementOutput

Definition: default_epilogue_tensor_op.h:80

[predicated_tile_iterator.h](epilogue_2threadblock_2predicated tile iterator_8h.html)

Epilogue for threadblock scoped GEMMs using Tensor Ops.

cutlass::epilogue::threadblock::InterleavedEpilogue

Epilogue operator without splitk.

Definition: interleaved_epilogue.h:79

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::WarpMmaTensorOp

WarpMmaTensorOp_ WarpMmaTensorOp

Definition: default_epilogue_tensor_op.h:75

gemm.h

Defines common types used for all GEMM-like operators.

conversion_op.h

Functor performing conversion operations used by epilogues.

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::kPartitionsK

static int const kPartitionsK

Definition: default_epilogue_tensor_op.h:76

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::OutputOp

OutputOp_ OutputOp

Definition: default_epilogue_tensor_op.h:151

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::WarpMmaTensorOp

WarpMmaTensorOp_ WarpMmaTensorOp

Definition: default_epilogue_tensor_op.h:149

[fragment_iterator_tensor_op.h](fragment iterator tensor__op_8h.html)

This defines a "fragment" iterator for visiting the fragments of an accumulator tile that participate...

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::OutputTileIterator

cutlass::epilogue::threadblock::PredicatedTileIterator< OutputTileThreadMap, ElementOutput > OutputTileIterator

Definition: default_epilogue_tensor_op.h:99

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::ElementAccumulator

typename WarpMmaTensorOp::ElementC ElementAccumulator

Definition: default_epilogue_tensor_op.h:82

array.h

Statically sized array of elements that accommodates all CUTLASS-supported numeric types and is safe ...

cutlass::epilogue::threadblock::DefaultInterleavedThreadMapTensorOp

Defines the optimal thread map for TensorOp accumulator layouts.

Definition: default_thread_map_tensor_op.h:104

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::Shape

Shape_ Shape

Definition: default_epilogue_tensor_op.h:74

linear_combination.h

Functor performing linear combination operations used by epilogues.

cutlass::sizeof_bits

Defines the size of an element in bits.

Definition: numeric_types.h:42

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::LayoutC

typename WarpMmaTensorOp::LayoutC LayoutC

Definition: default_epilogue_tensor_op.h:155

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::ElementAccumulator

typename WarpMmaTensorOp::ElementC ElementAccumulator

Definition: default_epilogue_tensor_op.h:156

cutlass::epilogue::threadblock::DefaultThreadMapTensorOp

Defines the optimal thread map for TensorOp accumulator layouts.

Definition: default_thread_map_tensor_op.h:52

numeric_types.h

Top-level include for all CUTLASS numeric types.

cutlass::epilogue::warp::TileIteratorTensorOp

Template for reading and writing tiles of accumulators to shared memory.

Definition: tile_iterator_tensor_op.h:52

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::SharedLoadIterator

cutlass::epilogue::threadblock::SharedLoadIterator< typename OutputTileThreadMap::CompactedThreadMap, ElementAccumulator > SharedLoadIterator

Definition: default_epilogue_tensor_op.h:119

cutlass::epilogue::threadblock::InterleavedPredicatedTileIterator

Definition: epilogue/threadblock/predicated_tile_iterator.h:452

[shared_load_iterator.h](shared load iterator_8h.html)

Epilogue for threadblock scoped GEMMs using Tensor Ops.

cutlass::epilogue::warp::FragmentIteratorTensorOp

Definition: fragment_iterator_tensor_op.h:61

[tile_iterator_tensor_op.h](tile iterator tensor__op_8h.html)

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::ElementOutput

typename OutputOp::ElementOutput ElementOutput

Definition: default_epilogue_tensor_op.h:154

cutlass::epilogue::threadblock::Epilogue

Epilogue operator without splitk.

Definition: epilogue.h:74

epilogue.h

Epilogue for threadblock scoped GEMMs using Tensor Ops.

cutlass::epilogue::threadblock::PredicatedTileIterator

Definition: epilogue/threadblock/predicated_tile_iterator.h:65

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::WarpTileIterator

cutlass::epilogue::warp::TileIteratorTensorOp< typename WarpMmaTensorOp::Shape, typename WarpMmaTensorOp::Policy::Operator::Shape, ElementAccumulator, LayoutC > WarpTileIterator

Definition: default_epilogue_tensor_op.h:114

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp

Definition: default_epilogue_tensor_op.h:147

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::OutputTileThreadMap

typename cutlass::epilogue::threadblock::DefaultInterleavedThreadMapTensorOp< Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput, kElementsPerAccess, InterleavedK >::Type OutputTileThreadMap

Definition: default_epilogue_tensor_op.h:164

cutlass::epilogue::threadblock::SharedLoadIterator

Definition: shared_load_iterator.h:61

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::OutputTileThreadMap

typename cutlass::epilogue::threadblock::DefaultThreadMapTensorOp< Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput, kElementsPerAccess >::Type OutputTileThreadMap

Definition: default_epilogue_tensor_op.h:94

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp

Defines sensible defaults for epilogues for TensorOps.

Definition: default_epilogue_tensor_op.h:72

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::AccumulatorFragmentIterator

cutlass::epilogue::warp::FragmentIteratorTensorOp< typename WarpMmaTensorOp::Shape, typename WarpMmaTensorOp::Policy::Operator::Shape, typename WarpMmaTensorOp::Policy::Operator::ElementC, typename WarpMmaTensorOp::Policy::Operator::FragmentC, LayoutC > AccumulatorFragmentIterator

Definition: default_epilogue_tensor_op.h:107

reduction_op.h

Functor performing reduction operations used by epilogues.

cutlass::epilogue::threadblock::DefaultInterleavedEpilogueTensorOp::Shape

Shape_ Shape

Definition: default_epilogue_tensor_op.h:148

cutlass.h

Basic include for CUTLASS.

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::kElementsPerAccess

static int const kElementsPerAccess

Definition: default_epilogue_tensor_op.h:78

cutlass::epilogue::threadblock::DefaultEpilogueTensorOp::OutputOp

OutputOp_ OutputOp

Definition: default_epilogue_tensor_op.h:77

interleaved_epilogue.h

Epilogue for threadblock scoped GEMMs using Tensor Ops.


Generated by 1.8.11