CUTLASS: mma_complex_tensor_op.h Source File

| CUTLASS CUDA Templates for Linear Algebra Subroutines and Solvers

- Main Page
- Modules
- Namespaces
- Classes
- Files

- File List
- File Members

mma_complex_tensor_op.h
Go to the documentation of this file.
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 **************************************************************************************************/
30 #pragma once
31
32 #include "cutlass/cutlass.h"
33
34 #include "cutlass/array.h"
35 #include "cutlass/complex.h"
36 #include "cutlass/numeric_types.h"
37 #include "cutlass/matrix_shape.h"
38
39 #include "cutlass/arch/memory_sm75.h"
40 #include "cutlass/arch/mma_sm75.h"
41 #include "cutlass/gemm/gemm.h"
42 #include "cutlass/gemm/warp/mma.h"
43
44 #include "cutlass/gemm/warp/mma_tensor_op_policy.h"
45 #include "cutlass/gemm/warp/mma_tensor_op.h"
46
47 #include "cutlass/gemm/warp/mma_tensor_op_tile_iterator.h"
49
50 namespace cutlass {
51 namespace gemm {
52 namespace warp {
53
55
56 template <
58 typename Shape_,
60 typename RealElementA,
62 typename LayoutA_,
64 typename RealElementB,
66 typename LayoutB_,
68 typename RealElementC,
70 typename LayoutC_,
72 typename Policy_,
74 ComplexTransform TransformA = ComplexTransform::kNone,
76 ComplexTransform TransformB = ComplexTransform::kNone,
78 typename Enable = bool
79 >
80 class MmaComplexTensorOp;
81
83
85 template <
87 typename Shape_,
89 typename RealElementA,
91 typename LayoutA_,
93 typename RealElementB,
95 typename LayoutB_,
97 typename RealElementC,
99 typename LayoutC_,
101 typename Policy_,
103 ComplexTransform TransformA,
105 ComplexTransform TransformB,
107 typename Enable
108 >
109 class MmaComplexTensorOp<
110 Shape_,
111 complex<RealElementA>,
112 LayoutA_,
113 complex<RealElementB>,
114 LayoutB_,
115 complex<RealElementC>,
116 LayoutC_,
117 Policy_,
118 TransformA,
119 TransformB,
120 Enable> {
121 public:
123 using Shape = Shape_;
124
126 using ElementA = complex<RealElementA>;
127
129 using LayoutA = LayoutA_;
130
132 using ElementB = complex<RealElementB>;
133
135 using LayoutB = LayoutB_;
136
138 using ElementC = complex<RealElementC>;
139
141 using LayoutC = LayoutC_;
142
144 using Policy = Policy_;
145
147 static ComplexTransform const kTransformA = TransformA;
148
150 static ComplexTransform const kTransformB = TransformB;
151
153 using OperatorClass = arch::OpClassTensorOp;
154
156 static int const kThreadCount = 32;
157
158 public:
159
161 using IteratorA = MmaTensorOpMultiplicandTileIterator<
162 MatrixShape<Shape::kM, Shape::kK>,
163 Operand::kA,
164 ElementA,
165 LayoutA,
166 MatrixShape<Policy::Operator::Shape::kM, Policy::Operator::Shape::kK>,
167 Policy::OpDelta::kRow,
168 32,
169 1
170 >;
171
173 using FragmentA = typename IteratorA::Fragment;
174
176 using IteratorB = MmaTensorOpMultiplicandTileIterator<
177 MatrixShape<Shape::kK, Shape::kN>,
178 Operand::kB,
179 ElementB,
180 LayoutB,
181 MatrixShape<Policy::Operator::Shape::kK, Policy::Operator::Shape::kN>,
182 Policy::OpDelta::kColumn,
183 32,
184 1
185 >;
186
188 using FragmentB = typename IteratorB::Fragment;
189
190
191 static_assert(
192 !(Shape::kM % Policy::Operator::Shape::kM) &&
193 !(Shape::kN % Policy::Operator::Shape::kN),
194 "Shape of warp-level Mma must be divisible by operator shape.");
195
197 using MmaIterations = MatrixShape<
198 Shape::kM / Policy::Operator::Shape::kM,
199 Shape::kN / Policy::Operator::Shape::kN
200 >;
201
203 using IteratorC = MmaTensorOpAccumulatorTileIterator<
204 MatrixShape<Shape::kM, Shape::kN>,
205 ElementC,
206 LayoutC,
207 typename Policy::Operator::Shape,
208 typename Policy::OpDelta>;
209
214 using FragmentC = typename IteratorC::Fragment;
215
216 static_assert(
217 FragmentC::kElements == 2 * MmaIterations::kCount * Policy::Operator::FragmentC::kElements,
218 "Unexpected planar complex fragment length.");
219
220 private:
221
222 //
223 // Data members
224 //
225
227 typename Policy::Operator mma;
228
229 public:
230
231 //
232 // Methods
233 //
234
236 CUTLASS_DEVICE
237 MmaComplexTensorOp() {}
238
240 CUTLASS_DEVICE
241 void operator()(
242 FragmentC &D,
243 FragmentA const &A,
244 FragmentB const &B,
245 FragmentC const &C) const {
246
247 // Alias types for underlying real-valued matrix multiply operator
248 using MmaOperandA = typename Policy::Operator::FragmentA;
249 using MmaOperandB = typename Policy::Operator::FragmentB;
250 using MmaOperandC = typename Policy::Operator::FragmentC;
251
252 static_assert(MmaOperandA::kElements == 1,
253 "This implementation only supports math instructions in which exactly one element is needed for the A operand."
254 "We can geneneralize later.");
255
256 static_assert(MmaOperandB::kElements == 1,
257 "This implementation only supports math instructions in which exactly one element is needed for the A operand."
258 "We can geneneralize later.");
259
260 D = C;
261
262 CUTLASS_PRAGMA_UNROLL
263 for (int m = 0; m < MmaIterations::kRow; ++m) {
264
265 // mma(accum.real(), a.real(), b.real(), accum.real());
266 CUTLASS_PRAGMA_UNROLL
267 for (int n = 0; n < MmaIterations::kColumn; ++n) {
268
269 // Pack operands together. This may result in actual MOVs
270 MmaOperandA operand_A;
271 MmaOperandB operand_B;
272
273 operand_A[0] = A[m].real();
274 operand_B[0] = B[n].real();
275
276 // Real-valued accumulator part
277 MmaOperandC *accum = reinterpret_cast<MmaOperandC *>(&D) +
278 (m + n * MmaIterations::kRow);
279
280 mma(*accum, operand_A, operand_B, *accum);
281 }
282
283 // mma(accum.imag(), a.real(), b.imag(), accum.imag());
284 CUTLASS_PRAGMA_UNROLL
285 for (int n = MmaIterations::kColumn - 1; n >= 0; --n) {
286
287 // Pack operands together. This may result in actual MOVs
288 MmaOperandA operand_A;
289 MmaOperandB operand_B;
290
291 operand_A[0] = A[m].real();
292 operand_B[0] = (kTransformB == ComplexTransform::kConjugate ? -B[n].imag() : B[n].imag());
293
294 // Complex-valued accumulator part
295 MmaOperandC *accum = reinterpret_cast<MmaOperandC *>(&D) +
296 (m + n * MmaIterations::kRow) + MmaIterations::kCount;
297
298 mma(*accum, operand_A, operand_B, *accum);
299 }
300
301 // mma(accum.real(), -a.imag(), b.imag(), accum.real())
302 CUTLASS_PRAGMA_UNROLL
303 for (int n = 0; n < MmaIterations::kColumn; ++n) {
304
305 // Pack operands together. This may result in actual MOVs
306 MmaOperandA operand_A;
307 MmaOperandB operand_B;
308
309 // A imaginary part is intentionally negated
310 operand_A[0] = (kTransformA == ComplexTransform::kConjugate ? A[m].imag() : -A[m].imag());
311 operand_B[0] = (kTransformB == ComplexTransform::kConjugate ? -B[n].imag() : B[n].imag());
312
313 // Complex-valued accumulator part
314 MmaOperandC *accum = reinterpret_cast<MmaOperandC *>(&D) +
315 (m + n * MmaIterations::kRow);
316
317 mma(*accum, operand_A, operand_B, *accum);
318 }
319
320 // mma(accum.imag(), a.imag(), b.real(), accum.imag())
321 CUTLASS_PRAGMA_UNROLL
322 for (int n = MmaIterations::kColumn - 1; n >= 0; --n) {
323
324 // Pack operands together. This may result in actual MOVs
325 MmaOperandA operand_A;
326 MmaOperandB operand_B;
327
328 operand_A[0] = (kTransformA == ComplexTransform::kConjugate ? -A[m].imag() : A[m].imag());
329 operand_B[0] = B[n].real();
330
331 // Real-valued accumulator part
332 MmaOperandC *accum = reinterpret_cast<MmaOperandC *>(&D) +
333 (m + n * MmaIterations::kRow) + MmaIterations::kCount;
334
335 mma(*accum, operand_A, operand_B, *accum);
336 }
337 }
338 }
339 };
340
342
343 // TODO - partial specializations of real*complex and complex*real
344
346
347 } // namespace warp
348 } // namespace gemm
349 } // namespace cutlass
350
cutlass::MatrixShape
Describes the size of a matrix tile.
Definition: matrix_shape.h:42
cutlass
Definition: aligned_buffer.h:35
cutlass::ComplexTransform
ComplexTransform
Enumeraed type describing a transformation on a complex value.
Definition: complex.h:43
complex.h
memory_sm75.h
Architecture-specific operators on memory added for SM75.
mma_tensor_op_tile_iterator.h
Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores.
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::FragmentA
typename IteratorA::Fragment FragmentA
Storage for A tile.
Definition: mma_complex_tensor_op.h:173
gemm.h
Defines common types used for all GEMM-like operators.
cutlass::ComplexTransform::kNone
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::LayoutB
LayoutB_ LayoutB
Layout of multiplicand B.
Definition: mma_complex_tensor_op.h:135
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::FragmentB
typename IteratorB::Fragment FragmentB
Storage for B tile.
Definition: mma_complex_tensor_op.h:188
cutlass::gemm::warp::MmaComplexTensorOp
Definition: mma_complex_tensor_op.h:80
array.h
Statically sized array of elements that accommodates all CUTLASS-supported numeric types and is safe ...
CUTLASS_PRAGMA_UNROLL
#define CUTLASS_PRAGMA_UNROLL
Definition: cutlass.h:110
mma.h
Templates exposing architecture support for warp-level multiply-add operations.
cutlass::gemm::Operand::kA
matrix_shape.h
Defines a Shape template for matrix tiles.
cutlass::gemm::warp::MmaTensorOpAccumulatorTileIterator
Definition: mma_tensor_op_tile_iterator.h:1794
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::FragmentC
typename IteratorC::Fragment FragmentC
Definition: mma_complex_tensor_op.h:214
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::Shape
Shape_ Shape
Shape of warp-level matrix operation (concept: GemmShape)
Definition: mma_complex_tensor_op.h:123
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::MmaComplexTensorOp
CUTLASS_DEVICE MmaComplexTensorOp()
Ctor.
Definition: mma_complex_tensor_op.h:237
numeric_types.h
Top-level include for all CUTLASS numeric types.
cutlass::gemm::warp::MmaTensorOpMultiplicandTileIterator
Definition: mma_tensor_op_tile_iterator.h:75
static_assert
#define static_assert(__e, __m)
Definition: platform.h:153
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::OperatorClass
arch::OpClassTensorOp OperatorClass
Indicates class of matrix operator.
Definition: mma_complex_tensor_op.h:153
cutlass::complex
Definition: complex.h:92
cutlass::ComplexTransform::kConjugate
mma_sm75.h
Matrix multiply for SM75.
cutlass::gemm::Operand::kB
A multiplicand.
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::Policy
Policy_ Policy
Shape of the warp in units of thread (concept: MmaLanePolicySimt)
Definition: mma_complex_tensor_op.h:144
mma_tensor_op.h
Templates implementing warp-level matrix multiply-accumulate operations targeting Tensor Cores...
cutlass.h
Basic include for CUTLASS.
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::operator()
CUTLASS_DEVICE void operator()(FragmentC &D, FragmentA const &A, FragmentB const &B, FragmentC const &C) const
Performs a warp-level matrix multiply-accumulate operation.
Definition: mma_complex_tensor_op.h:241
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::LayoutC
LayoutC_ LayoutC
Layout of accumulator matrix C.
Definition: mma_complex_tensor_op.h:141
mma_tensor_op_policy.h
Policy describing implementation details of warp-level GEMM targeting Tensor Cores.
cutlass::gemm::warp::MmaComplexTensorOp< Shape_, complex< RealElementA >, LayoutA_, complex< RealElementB >, LayoutB_, complex< RealElementC >, LayoutC_, Policy_, TransformA, TransformB, Enable >::LayoutA
LayoutA_ LayoutA
Layout of multiplicand A.
Definition: mma_complex_tensor_op.h:129