CUTLASS: default_epilogue_complex_tensor_op.h Source File - Cutlass

| | CUTLASS

CUDA Templates for Linear Algebra Subroutines and Solvers |

default_epilogue_complex_tensor_op.h

[Go to the documentation of this file.](default epilogue complex tensor op_8h.html)

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

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

35 #include "cutlass/cutlass.h"

36 #include "cutlass/numeric_types.h"

37 #include "cutlass/array.h"

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

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

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

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

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

47 #include "[cutlass/epilogue/warp/fragment_iterator_complex_tensor_op.h](fragment iterator complex 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)"

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

57 namespace cutlass {

58 namespace epilogue {

59 namespace threadblock {

64 template <

65typename Shape_,

66typename WarpMmaTensorOp_,

67int PartitionsK,

68typename OutputOp_,

69int ElementsPerAccess

70 >

71 struct DefaultEpilogueComplexTensorOp {

73using Shape = Shape_;

74using WarpMmaTensorOp = WarpMmaTensorOp_;

75static int const kPartitionsK = PartitionsK;

76using OutputOp = OutputOp_;

77static int const kElementsPerAccess = ElementsPerAccess;

79using ElementOutput = typename OutputOp::ElementOutput;

80using LayoutC = typename WarpMmaTensorOp::LayoutC;

81using ElementAccumulator = typename WarpMmaTensorOp::ElementC;

83//

84// Thread map

85//

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

88Shape,

89typename WarpMmaTensorOp::Shape,

90kPartitionsK,

91ElementOutput,

92 kElementsPerAccess

93 >::Type;

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

96OutputTileThreadMap,

97 ElementOutput

98 >;

100using AccumulatorFragmentIterator = cutlass::epilogue::warp::FragmentIteratorComplexTensorOp<

101typename WarpMmaTensorOp::Shape,

102typename WarpMmaTensorOp::Policy::Operator::Shape,

103typename WarpMmaTensorOp::Policy::Operator::ElementC,

104typename WarpMmaTensorOp::Policy::Operator::FragmentC,

105LayoutC

106 >;

107

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

109typename WarpMmaTensorOp::Shape,

110typename WarpMmaTensorOp::Policy::Operator::Shape,

111ElementAccumulator,

112LayoutC

113 >;

114

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

116typename OutputTileThreadMap::CompactedThreadMap,

117 ElementAccumulator

118 >;

119

121using Padding = cutlass::MatrixShape<0, 0>;

122

123//

124// Define the epilogue

125//

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

127Shape,

128WarpMmaTensorOp,

129kPartitionsK,

130OutputTileIterator,

131AccumulatorFragmentIterator,

132WarpTileIterator,

133SharedLoadIterator,

134OutputOp,

135Padding

136 >;

137 };

138

140

141 } // namespace threadblock

142 } // namespace epilogue

143 } // namespace cutlass

144

[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::epilogue::threadblock::DefaultEpilogueComplexTensorOp::ElementOutput

typename OutputOp::ElementOutput ElementOutput

Definition: default_epilogue_complex_tensor_op.h:79

cutlass

Definition: aligned_buffer.h:35

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::kPartitionsK

static int const kPartitionsK

Definition: default_epilogue_complex_tensor_op.h:75

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::LayoutC

typename WarpMmaTensorOp::LayoutC LayoutC

Definition: default_epilogue_complex_tensor_op.h:80

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::OutputTileThreadMap

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

Definition: default_epilogue_complex_tensor_op.h:93

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

Epilogue for threadblock scoped GEMMs using Tensor Ops.

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::WarpTileIterator

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

Definition: default_epilogue_complex_tensor_op.h:113

gemm.h

Defines common types used for all GEMM-like operators.

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::kElementsPerAccess

static int const kElementsPerAccess

Definition: default_epilogue_complex_tensor_op.h:77

conversion_op.h

Functor performing conversion operations used by epilogues.

[fragment_iterator_complex_tensor_op.h](fragment iterator complex tensor op_8h.html)

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

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::SharedLoadIterator

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

Definition: default_epilogue_complex_tensor_op.h:118

array.h

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

linear_combination.h

Functor performing linear combination operations used by epilogues.

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp

Defines sensible defaults for epilogues for TensorOps.

Definition: default_epilogue_complex_tensor_op.h:71

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::Shape

Shape_ Shape

Definition: default_epilogue_complex_tensor_op.h:73

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::WarpMmaTensorOp

WarpMmaTensorOp_ WarpMmaTensorOp

Definition: default_epilogue_complex_tensor_op.h:74

cutlass::epilogue::warp::FragmentIteratorComplexTensorOp

Definition: fragment_iterator_complex_tensor_op.h:61

cutlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp::AccumulatorFragmentIterator

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

Definition: default_epilogue_complex_tensor_op.h:106

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::DefaultEpilogueComplexTensorOp::OutputTileIterator

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

Definition: default_epilogue_complex_tensor_op.h:98