Cutlass CUDA Templates for Linear Algebra Subroutines and Solvers

gemm_traits.h
Go to the documentation of this file.
1 /***************************************************************************************************
2 * Copyright (c) 2017-2018, 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 **************************************************************************************************/
28 #pragma once
29
30 #include "cutlass/convert.h"
31 #include "cutlass/matrix_traits.h"
32 #include "cutlass/reshape_tile.h"
33 #include "cutlass/tile_allocation.h"
34 #include "cutlass/tile_iterator.h"
35 #include "cutlass/kernel_launch.h"
36
37 #include "cutlass/gemm/clear_accumulators.h"
38 #include "cutlass/gemm/gemm_config.h"
39 #include "cutlass/gemm/gemm_desc.h"
40 #include "cutlass/gemm/gemm_stream_pair.h"
41 #include "cutlass/gemm/gemm_global_stream.h"
42 #include "cutlass/gemm/gemm_operand.h"
43 #include "cutlass/gemm/gemm_shared_stream.h"
44 #include "cutlass/gemm/threadblock_swizzle.h"
45 #include "cutlass/gemm/gemm.h"
46 namespace cutlass {
47 namespace gemm {
48
50
51 template <enum MatrixLayout::Kind, typename GemmConfig_>
52 struct GemmTileTraitsHelperA {};
53
55
56 template <typename GemmConfig_>
57 struct GemmTileTraitsHelperA<MatrixLayout::kColumnMajor, GemmConfig_> {
59 static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor;
60
62 typedef typename GemmConfig_::ScalarA Scalar;
64 typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar;
65
67 typedef GemmGlobalTileTraits<
68 // That's A.
69 GemmOperand::kA,
70 // A is column-major.
71 MatrixLayout::kColumnMajor,
72 // The pointer is float const.
73 Scalar const,
74 // The tile has size KxM in GEMM's terminology.
75 Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kW>,
76 // The threads are distributed as warps x 32 (the traits may reorganize).
77 Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>,
78 // The number of scalars per LDG (LDG.32 or LDG.128, etc).
79 GemmConfig_::kScalarsPerLdgA>
80 GlobalTileTraits;
81
83 typedef GemmSharedStoreTileAbTraits<
84 // The pointer is float.
85 MultiplyAddScalar,
86 // The tile has size KxM in GEMM's terminology.
87 Shape<GemmConfig_::kStages,
88 GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD,
89 GemmConfig_::OutputTile::kW * GemmConfig_::InstructionShape::kD>,
90 // The threads are distributed as warps x 32 (the traits may reorganize).
91 typename GlobalTileTraits::Threads,
92 // The number of scalars per STS (STS.32 or STS.128, etc).
93 GemmConfig_::kScalarsPerStsA>
94 SharedStoreTileTraits;
95
97 typedef GemmSharedLoadTileATraits<
98 // The pointer is float const.
99 MultiplyAddScalar const,
100 // The output tile size.
101 typename GemmConfig_::OutputTile,
102 // The number of warps.
103 typename GemmConfig_::Warps,
104 // The number of threads per warp.
105 typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
106 // The shape of the FMA instruction.
107 typename GemmConfig_::InstructionShape,
108 // The number of stages.
109 GemmConfig_::kStages,
110 // The number of scalars per LDS.
111 GemmConfig_::kScalarsPerLdsA,
112 // The skew.
113 0>
114 SharedLoadTileTraits;
115 };
116
118
119 template <typename GemmConfig_>
120 struct GemmTileTraitsHelperA<MatrixLayout::kRowMajor, GemmConfig_> {
122 static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor;
123
125 typedef typename GemmConfig_::ScalarA Scalar;
127 typedef typename GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar;
128
130 typedef GemmGlobalTileTraits<
131 // That's A.
132 GemmOperand::kA,
133 // A is row-major.
134 MatrixLayout::kRowMajor,
135 // The pointer is float const.
136 Scalar const,
137 // The tile has size MxK in GEMM's terminology.
138 Shape<1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD>,
139 // The threads are distributed as (threads / K) x K (the traits may reorganize).
140 Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>,
141 // The number of scalars per LDG (LDG.32 or LDG.128, etc).
142 GemmConfig_::kScalarsPerLdgA>
143 GlobalTileTraits;
144
146 static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar);
148 static int const kSkewA = 128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsA /
149 GlobalTileTraits::Threads::kW * kScalarsIn4B;
150
152 typedef GemmSharedStoreWithSkewTileAbTraits <
153 // The pointer is float.
154 MultiplyAddScalar,
155 // The tile has size KxM in GEMM's terminology.
156 Shape<GemmConfig_::kStages,
157 GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD,
158 GemmConfig_::OutputTile::kW * GemmConfig_::InstructionShape::kD>,
159 // The threads are distributed as (threads / K) x K (the traits may reorganize).
160 typename GlobalTileTraits::Threads,
161 // The number of scalars per STS.
162 GemmConfig_::kScalarsPerStsA,
163 // The skew to avoid bank conflicts added in the tile W dimension.
164 kSkewA<GemmConfig_::kScalarsPerLdsA ? GemmConfig_::kScalarsPerLdsA : kSkewA>
165 SharedStoreTileTraits;
166
168 typedef GemmSharedLoadTileATraits<
169 // The pointer is float const.
170 MultiplyAddScalar const,
171 // The output tile size.
172 typename GemmConfig_::OutputTile,
173 // The number of warps.
174 typename GemmConfig_::Warps,
175 // The number of threads per warp.
176 typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
177 // The shape of the FMA instruction.
178 typename GemmConfig_::InstructionShape,
179 // The number of stages.
180 GemmConfig_::kStages,
181 // The number of scalars per LDS.
182 GemmConfig_::kScalarsPerLdsA,
183 // The skew.
184 SharedStoreTileTraits::kSkew>
185 SharedLoadTileTraits;
186 };
187
189
190 template <enum MatrixLayout::Kind, typename GemmConfig_>
191 struct GemmTileTraitsHelperB {};
192
194
195 template <typename GemmConfig_>
196 struct GemmTileTraitsHelperB<MatrixLayout::kColumnMajor, GemmConfig_> {
198 static MatrixLayout::Kind const kLayout = MatrixLayout::kColumnMajor;
199
201 typedef typename GemmConfig_::ScalarB Scalar;
203 typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar;
204
206 typedef GemmGlobalTileTraits<
207 // That's B.
208 GemmOperand::kB,
209 // B is column-major.
210 MatrixLayout::kColumnMajor,
211 // The pointer is float const.
212 Scalar const,
213 // The tile has size MxK in GEMM's terminology.
214 Shape<1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD>,
215 // The threads are distributed as (threads / K) x K (the traits may reorganize).
216 Shape<1, GemmConfig_::kThreads / GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD>,
217 // The number of scalars per LDG (LDG.32 or LDG.128, etc).
218 GemmConfig_::kScalarsPerLdgB>
219 GlobalTileTraits;
220
222 static int const kScalarsIn4B = sizeof(MultiplyAddScalar) > 4 ? 1 : 4 / sizeof(MultiplyAddScalar);
224 static int const kSkewB = 128 / sizeof(MultiplyAddScalar) / GemmConfig_::kScalarsPerStsB /
225 GlobalTileTraits::Threads::kW * kScalarsIn4B;
226
228 typedef GemmSharedStoreWithSkewTileAbTraits <
229 // The pointer is float.
230 MultiplyAddScalar,
231 // The tile has size KxN in GEMM's terminology.
232 Shape<GemmConfig_::kStages,
233 GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD,
234 GemmConfig_::OutputTile::kH * GemmConfig_::InstructionShape::kD>,
235 // The threads are distributed as (threads / K) x K (the traits may reorganize).
236 typename GlobalTileTraits::Threads,
237 // The number of scalars per STS.
238 GemmConfig_::kScalarsPerStsB,
239 // The skew to avoid bank conflicts added in the tile W dimension.
240 kSkewB<GemmConfig_::kScalarsPerLdsB ? GemmConfig_::kScalarsPerLdsB : kSkewB>
241 SharedStoreTileTraits;
242
244 typedef GemmSharedLoadTileBTraits<
245 // The pointer is float const.
246 MultiplyAddScalar const,
247 // The output tile size.
248 typename GemmConfig_::OutputTile,
249 // The number of warps.
250 typename GemmConfig_::Warps,
251 // The number of threads per warp.
252 typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
253 // The shape of the FMA instruction.
254 typename GemmConfig_::InstructionShape,
255 // The number of stages.
256 GemmConfig_::kStages,
257 // The number of scalars per LDS.
258 GemmConfig_::kScalarsPerLdsB,
259 // The skew.
260 SharedStoreTileTraits::kSkew>
261 SharedLoadTileTraits;
262 };
263
265
266 template <typename GemmConfig_>
267 struct GemmTileTraitsHelperB<MatrixLayout::kRowMajor, GemmConfig_> {
269 static MatrixLayout::Kind const kLayout = MatrixLayout::kRowMajor;
270
272 typedef typename GemmConfig_::ScalarB Scalar;
274 typedef typename GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar;
275
277 typedef GemmGlobalTileTraits<
278 // That's B.
279 GemmOperand::kB,
280 // B is row-major.
281 MatrixLayout::kRowMajor,
282 // The pointer is float const.
283 Scalar const,
284 // The tile has size KxN in GEMM's terminology.
285 Shape<1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kH>,
286 // The threads are distributed as warps x 32 (the traits may reorganize).
287 Shape<1, ShapeCount<typename GemmConfig_::Warps>::kCount, GemmConfig_::kWarpSize>,
288 // The number of scalars per LDG (LDG.32 or LDG.128, etc).
289 GemmConfig_::kScalarsPerLdgB>
290 GlobalTileTraits;
291
293 typedef GemmSharedStoreTileAbTraits<
294 // The pointer is float.
295 MultiplyAddScalar,
296 // The tile has size KxN in GEMM's terminology.
297 Shape<GemmConfig_::kStages,
298 GemmConfig_::OutputTile::kD / GemmConfig_::InstructionShape::kD,
299 GemmConfig_::OutputTile::kH * GemmConfig_::InstructionShape::kD>,
300 // The threads are distributed as warps x 32 (the traits may reorganize).
301 typename GlobalTileTraits::Threads,
302 // The number of scalars per STS (STS.32 or STS.128, etc).
303 GemmConfig_::kScalarsPerStsB>
304 SharedStoreTileTraits;
305
307 typedef GemmSharedLoadTileBTraits<
308 // The pointer is float const.
309 MultiplyAddScalar const,
310 // The output tile size.
311 typename GemmConfig_::OutputTile,
312 // The number of warps.
313 typename GemmConfig_::Warps,
314 // The number of threads per warp.
315 typename GemmConfig_::MultiplyAdd::ThreadsPerWarp,
316 // The shape of the FMA instruction.
317 typename GemmConfig_::InstructionShape,
318 // The number of stages.
319 GemmConfig_::kStages,
320 // The number of scalars per LDS.
321 GemmConfig_::kScalarsPerLdsB,
322 // The skew.
323 0>
324 SharedLoadTileTraits;
325 };
326
328
329 template <
331 typename GemmConfig_,
333 typename GlobalLoadStreamA_,
335 typename GlobalLoadStreamB_,
337 typename SharedLoadStreamA_,
339 typename SharedLoadStreamB_,
341 typename Epilogue_,
343 typename BlockSwizzle_ = IdentityBlockSwizzle,
345 typename Index_ = int,
347 typename ClearAccumulators_ = ClearAccumulators<typename GemmConfig_::Accumulators::Element> >
348
349 struct GemmTraits {
351 typedef GemmTraits<GemmConfig_,
352 GlobalLoadStreamA_,
353 GlobalLoadStreamB_,
354 SharedLoadStreamA_,
355 SharedLoadStreamB_,
356 Epilogue_,
357 BlockSwizzle_,
358 Index_,
359 ClearAccumulators_> This_;
360
362 typedef typename cutlass::gemm::Gemm<This_> KernelClass;
363
365 typedef GemmConfig_ GemmConfig;
367 typedef typename GemmConfig::OutputTile OutputTile;
368
370 typedef GlobalLoadStreamA_ GlobalLoadStreamA;
372 static MatrixLayout::Kind const kLayoutA = GlobalLoadStreamA::kLayout;
374 typedef typename GlobalLoadStreamA_::Scalar ScalarA;
375
377 typedef GlobalLoadStreamB_ GlobalLoadStreamB;
379 static MatrixLayout::Kind const kLayoutB = GlobalLoadStreamB::kLayout;
381 typedef typename GlobalLoadStreamB_::Scalar ScalarB;
382
384 typedef SharedLoadStreamA_ SharedLoadStreamA;
386 typedef SharedLoadStreamB_ SharedLoadStreamB;
387
389 typedef typename GemmConfig::MultiplyAdd MultiplyAdd;
391 typedef Epilogue_ Epilogue;
393 typedef typename Epilogue::ScalarC ScalarC;
394 typedef typename Epilogue::ScalarD ScalarD;
395
397 typedef BlockSwizzle_ BlockSwizzle;
399 typedef Index_ Index;
401 typedef ClearAccumulators_ ClearAccumulators;
402
404 typedef GlobalLoadStreamPair<GlobalLoadStreamA,
405 GlobalLoadStreamB,
406 GemmConfig::kResidueInProlog>
407 GlobalLoadStream;
408
410 typedef typename GlobalLoadStream::ThreadblockTileStorage ThreadblockTileStorage;
411
413 typedef SharedStreamPair<SharedLoadStreamA, SharedLoadStreamB> SharedStream;
414
416 struct Params : public KernelLaunchConfiguration {
417
419 GemmCoord problem_size;
420
422 int partitionK_range;
423
425 typename GlobalLoadStream::Params global_to_shared_stream;
426
428 typename SharedStream::Params shared_stream;
429
431 typename Epilogue::Params epilogue;
432
434 template <typename GemmDesc_>
435 CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& desc) {
436 // Set the problem size.
437 problem_size = desc.problem_size;
438
439 // there is no partitionK in the default case
440 partitionK_range = problem_size[0];
441 // Compute grid dimensions
442 BlockSwizzle block_swizzle;
443 this->block = dim3(GemmConfig::kThreads);
444 this->grid = block_swizzle.get_grid_layout(
445 problem_size,
446 make_Coord_from_shape<OutputTile>());
447
448 // Compute offset to residue.
449 // partitionK_range <= problem_size[0]
450 Index gemm_k = problem_size[0];
451 Index offset_to_residue_last_partition = (gemm_k % OutputTile::kD) ? gemm_k - (gemm_k % OutputTile::kD) : 0;
452 Index offset_to_residue = (partitionK_range % OutputTile::kD) ? partitionK_range - (partitionK_range % OutputTile::kD) : 0;
453
454 // Initialize parameters objects for
455 int error_code = global_to_shared_stream.stream_a.initialize(
456 desc.A.data(),
457 desc.batch_stride_A,
458 desc.A.leading_dim(),
459 offset_to_residue,
460 offset_to_residue_last_partition
461 );
462 if (error_code) {
463 return error_code;
464 }
465
466 error_code = global_to_shared_stream.stream_b.initialize(
467 desc.B.data(),
468 desc.batch_stride_B,
469 desc.B.leading_dim(),
470 offset_to_residue,
471 offset_to_residue_last_partition
472 );
473
474 if (error_code) {
475 return error_code;
476 }
477
478 // The epilogue.
479 return epilogue.initialize(desc);
480 }
481
483 CUTLASS_HOST_DEVICE int initialize(Index m,
484 Index n,
485 Index k,
486 typename Epilogue::Scalar alpha,
487 ScalarA const* d_a,
488 Index lda,
489 ScalarB const* d_b,
490 Index ldb,
491 typename Epilogue::Scalar beta,
492 ScalarC const* d_c,
493 Index ldc,
494 ScalarD* d_d,
495 Index ldd) {
496 GemmDesc<ScalarA, ScalarB, ScalarC, ScalarD, typename Epilogue::Scalar> desc(
497 GemmCoord(k, n, m, 1),
498 alpha,
499 TensorRef<ScalarA const, 2>(d_a, lda),
500 TensorRef<ScalarB const, 2>(d_b, ldb),
501 beta,
502 TensorRef<ScalarC const, 2>(d_c, ldc),
503 TensorRef<ScalarD, 2>(d_d, ldd)
504 );
505
506 return this->initialize(desc);
507 }
508
510 CUTLASS_HOST_DEVICE int initialize(Index m,
511 Index n,
512 Index k,
513 typename Epilogue::Scalar alpha,
514 ScalarA const* d_a,
515 Index lda,
516 long long int batch_stride_A,
517 ScalarB const* d_b,
518 Index ldb,
519 long long int batch_stride_B,
520 typename Epilogue::Scalar beta,
521 ScalarC const* d_c,
522 Index ldc,
523 long long int batch_stride_C,
524 ScalarD* d_d,
525 Index ldd,
526 long long int batch_stride_D,
527 Index batch_count) {
528 GemmDesc<ScalarA, ScalarB, ScalarC, ScalarD, typename Epilogue::Scalar> desc(
529 GemmCoord(k, n, m, batch_count),
530 alpha,
531 TensorRef<ScalarA const, 2>(d_a, lda),
532 batch_stride_A,
533 TensorRef<ScalarB const, 2>(d_b, ldb),
534 batch_stride_B,
535 beta,
536 TensorRef<ScalarC const, 2>(d_c, ldc),
537 batch_stride_C,
538 TensorRef<ScalarD, 2>(d_d, ldd),
539 batch_stride_D
540 );
541
542 return this->initialize(desc);
543 }
544
546 template <typename GemmDesc_>
547 CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const& partitonK_desc, Index partitionK_count_) {
548 // partitionK GEMM is a specialized batched stried gemm with different K ranges per batch
549 // the problem_size of each batch is (lastK_size, n, m)
550 // add more comments here
551 // the k range for every batch excpet the last one
552 //assert(partitionK_count_ > 0);
553 partitionK_range = partitonK_desc.problem_size.k() / partitionK_count_;
554 // the k range of the last batch
555 // int lastK_range = (partitonK_desc.problem_size.k() % partitionK_range) + partitionK_range;
556 int lastK_range = partitonK_desc.problem_size.k() - partitionK_range * (partitionK_count_ - 1);
557 int k_size = lastK_range;
558 int lda = partitonK_desc.A.stride(0);
559 int ldb = partitonK_desc.B.stride(0);
560 int ldc = partitonK_desc.C.stride(0);
561 int ldd = partitonK_desc.D.stride(0);
562 int n = partitonK_desc.problem_size.n();
563
564
565 long long int batch_stride_A = (kLayoutA == cutlass::MatrixLayout::kColumnMajor) ? lda * partitionK_range : partitionK_range;
566 long long int batch_stride_B = (kLayoutB == cutlass::MatrixLayout::kColumnMajor) ? partitionK_range : partitionK_range * ldb;
567 long long int batch_stride_C = ldc * n;
568 long long int batch_stride_D = ldd * n;
569
570 GemmDesc<ScalarA, ScalarB, ScalarC, ScalarD, typename Epilogue::Scalar> desc(
571 //we pass lastK_size as per batch K. there is also a range that will match partitionK_size
572 GemmCoord(k_size, partitonK_desc.problem_size.n(), partitonK_desc.problem_size.m(), partitionK_count_),
573 partitonK_desc.alpha,
574 partitonK_desc.A,
575 batch_stride_A,
576 partitonK_desc.B,
577 batch_stride_B,
578 partitonK_desc.beta,
579 partitonK_desc.C,
580 batch_stride_C,
581 partitonK_desc.D,
582 batch_stride_D
583 );
584
585 // Set the problem size.
586 problem_size = desc.problem_size;
587
588 // Compute grid dimensions
589 BlockSwizzle block_swizzle;
590 this->block = dim3(GemmConfig::kThreads);
591 this->grid = block_swizzle.get_grid_layout(
592 problem_size,
593 make_Coord_from_shape<OutputTile>());
594
595 // Compute offset to residue.
596 // partitionK_range <= problem_size[0]
597 Index gemm_k = problem_size[0];
598 Index offset_to_residue_last_partition = (gemm_k % OutputTile::kD) ? gemm_k - (gemm_k % OutputTile::kD) : 0;
599 Index offset_to_residue = (partitionK_range % OutputTile::kD) ? partitionK_range - (partitionK_range % OutputTile::kD) : 0;
600
601 // Initialize parameters objects for
602 int error_code = global_to_shared_stream.stream_a.initialize(
603 desc.A.data(),
604 desc.batch_stride_A,
605 desc.A.leading_dim(),
606 offset_to_residue,
607 offset_to_residue_last_partition
608 );
609 if (error_code) {
610 return error_code;
611 }
612
613 error_code = global_to_shared_stream.stream_b.initialize(
614 desc.B.data(),
615 desc.batch_stride_B,
616 desc.B.leading_dim(),
617 offset_to_residue,
618 offset_to_residue_last_partition
619 );
620
621 if (error_code) {
622 return error_code;
623 }
624
625 // The epilogue.
626 return epilogue.initialize(desc);
627 }
628
629
631 CUTLASS_HOST_DEVICE int initialize(Index m,
632 Index n,
633 Index k,
634 typename Epilogue::Scalar alpha,
635 ScalarA const* d_a,
636 Index lda,
637 ScalarB const* d_b,
638 Index ldb,
639 typename Epilogue::Scalar beta,
640 ScalarC const* d_c,
641 Index ldc,
642 ScalarD* d_d,
643 Index ldd,
644 Index partitionK_count_) {
645
646 GemmDesc<ScalarA, ScalarB, ScalarC, ScalarD, typename Epilogue::Scalar> desc(
647 GemmCoord(k, n, m, 1),
648 alpha,
649 TensorRef<ScalarA const, 2>(d_a, lda),
650 TensorRef<ScalarB const, 2>(d_b, ldb),
651 beta,
652 TensorRef<ScalarC const, 2>(d_c, ldc),
653 TensorRef<ScalarD, 2>(d_d, ldd)
654 );
655
656
657 return this->initialize(desc, partitionK_count_);
658 }
659 };
660
661 // The storage for the main loop + prologue.
662 struct MainLoopSharedStorage {
664 ThreadblockTileStorage threadblock_tile;
665
667 typename GlobalLoadStream::SharedStorage global_to_shared_stream;
668
670 typename ClearAccumulators::SharedStorage clear;
671 };
672
674 union SharedStorage {
675 // The storage for the main loop.
676 MainLoopSharedStorage main_loop;
677 // The storage for the epilogue.
678 typename Epilogue::SharedStorage epilogue;
679 };
680
682 static CUTLASS_DEVICE void shared_load_fence(bool in_loop) {
683 if (SharedLoadStreamA::Iterator::kRequiresLoadFence ||
684 SharedLoadStreamB::Iterator::kRequiresLoadFence) {
685 __syncthreads();
686 }
687 }
688
690 static CUTLASS_DEVICE void shared_store_fence(bool in_loop) {
691 __syncthreads();
692 }
693 };
694
696
697 template <typename GemmTileTraitsHelperA_, typename GemmTileTraitsHelperB_, typename Index_>
698 struct SimplifiedGemmTraitsHelper {
700 typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperA_::GlobalTileTraits, Index_>
701 GlobalLoadIteratorA;
703 typedef Copy<typename GlobalLoadIteratorA::Fragment> GlobalTransformerA;
705 typedef TileStoreIterator<typename GemmTileTraitsHelperA_::SharedStoreTileTraits,
706 typename GemmTileTraitsHelperA_::SharedStoreTileTraits::Scalar,
707 IteratorAdvance::kH,
708 MemorySpace::kShared>
709 SharedStoreIteratorA;
711 typedef GlobalLoadStream<GemmOperand::kA,
712 GlobalLoadIteratorA,
713 SharedStoreIteratorA,
714 GlobalTransformerA>
715 GlobalLoadStreamA;
716
718 typedef GemmGlobalIteratorAb<typename GemmTileTraitsHelperB_::GlobalTileTraits, Index_>
719 GlobalLoadIteratorB;
721 typedef Copy<typename GlobalLoadIteratorB::Fragment> GlobalTransformerB;
723 typedef TileStoreIterator<typename GemmTileTraitsHelperB_::SharedStoreTileTraits,
724 typename GemmTileTraitsHelperB_::SharedStoreTileTraits::Scalar,
725 IteratorAdvance::kH,
726 MemorySpace::kShared>
727 SharedStoreIteratorB;
729 typedef GlobalLoadStream<GemmOperand::kB,
730 GlobalLoadIteratorB,
731 SharedStoreIteratorB,
732 GlobalTransformerB>
733 GlobalLoadStreamB;
734
736 typedef TileLoadIterator<typename GemmTileTraitsHelperA_::SharedLoadTileTraits,
737 typename GemmTileTraitsHelperA_::Scalar,
738 IteratorAdvance::kH,
739 MemorySpace::kShared>
740 SharedLoadIteratorA;
742 typedef SharedLoadStream<SharedLoadIteratorA> SharedLoadStreamA;
744 typedef TileLoadIterator<typename GemmTileTraitsHelperB_::SharedLoadTileTraits,
745 typename GemmTileTraitsHelperB_::Scalar,
746 IteratorAdvance::kH,
747 MemorySpace::kShared>
748 SharedLoadIteratorB;
750 typedef SharedLoadStream<SharedLoadIteratorB> SharedLoadStreamB;
751 };
752
754
755 template <
757 MatrixLayout::Kind kLayoutA_,
759 MatrixLayout::Kind kLayoutB_,
761 typename GemmConfig_,
763 typename Epilogue_,
765 typename Index_ = int,
766 // The configuration for the A matrix.
767 typename GemmTileTraitsHelperA_ = GemmTileTraitsHelperA<kLayoutA_, GemmConfig_>,
768 // The configuration for the B matrix.
769 typename GemmTileTraitsHelperB_ = GemmTileTraitsHelperB<kLayoutB_, GemmConfig_>,
770 // The helper class to create the streams and iterators.
771 typename Helper_ =
772 SimplifiedGemmTraitsHelper<GemmTileTraitsHelperA_, GemmTileTraitsHelperB_, Index_> >
773 struct SimplifiedGemmTraits : public GemmTraits<
774 // The config.
775 GemmConfig_,
776 // The stream to load A from global memory to shared memory.
777 typename Helper_::GlobalLoadStreamA,
778 // The stream to load B from global memory to shared memory.
779 typename Helper_::GlobalLoadStreamB,
780 // The stream to load A from shared memory.
781 typename Helper_::SharedLoadStreamA,
782 // The stream to load B from shared memory.
783 typename Helper_::SharedLoadStreamB,
784 // The epilogue.
785 Epilogue_,
786 // The block swizzle to reorganize the grid.
787 IdentityBlockSwizzle,
788 // The index.
789 Index_,
790 // The tool used to clear accumulators.
791 ClearAccumulators<typename GemmConfig_::Accumulators::Element> > {
792 };
793
795
796 } // namespace gemm
797 } // namespace cutlass
cutlass::gemm::GemmTraits::SharedStorage::epilogue
Epilogue::SharedStorage epilogue
Definition: gemm_traits.h:678
cutlass::gemm::GemmDesc
GEMM problem description.
Definition: gemm_desc.h:50
cutlass::gemm::GemmTraits::GlobalLoadStreamA
GlobalLoadStreamA_ GlobalLoadStreamA
The stream to load A from global memory to shared memory.
Definition: gemm_traits.h:370
cutlass::gemm::GemmEpilogue::ScalarD
GlobalStoreIteratorD::Scalar ScalarD
The scalar for D.
Definition: gemm_epilogue.h:85
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalLoadStreamA
GlobalLoadStream< GemmOperand::kA, GlobalLoadIteratorA, SharedStoreIteratorA, GlobalTransformerA > GlobalLoadStreamA
The stream to load A from global memory to shared memory.
Definition: gemm_traits.h:715
cutlass::MemorySpace::kShared
Definition: load_store.h:41
cutlass::gemm::GemmTraits::SharedLoadStreamA
SharedLoadStreamA_ SharedLoadStreamA
The iterator for A to load from shared memory.
Definition: gemm_traits.h:384
cutlass
Definition: convert.h:33
cutlass::gemm::GemmSharedLoadTileATraits
Definition: gemm_shared_tile.h:128
cutlass::gemm::GemmTraits::GlobalLoadStreamB
GlobalLoadStreamB_ GlobalLoadStreamB
The stream to load B from global memory to shared memory.
Definition: gemm_traits.h:377
cutlass::gemm::GemmSharedStoreWithSkewTileAbTraits
Definition: gemm_shared_tile.h:80
cutlass::gemm::GemmConfig::kThreads
static int const kThreads
The numnber of threads.
Definition: gemm_config.h:103
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedStoreIteratorA
TileStoreIterator< typename GemmTileTraitsHelperA_::SharedStoreTileTraits, typename GemmTileTraitsHelperA_::SharedStoreTileTraits::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedStoreIteratorA
The iterator to store A to shared memory.
Definition: gemm_traits.h:709
tile_iterator.h
Defines the Tile Traits concept and iterators for loading and storing to tiles efficiently.
cutlass::gemm::GemmTraits::ScalarD
Epilogue::ScalarD ScalarD
Definition: gemm_traits.h:394
cutlass::gemm::GemmTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(Index m, Index n, Index k, typename Epilogue::Scalar alpha, ScalarA const *d_a, Index lda, ScalarB const *d_b, Index ldb, typename Epilogue::Scalar beta, ScalarC const *d_c, Index ldc, ScalarD *d_d, Index ldd)
Helper to construct a GEMM params using a BLAS-like API.
Definition: gemm_traits.h:483
cutlass::gemm::GemmTraits::SharedStorage
The storage in shared memory.
Definition: gemm_traits.h:674
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedLoadStreamB
SharedLoadStream< SharedLoadIteratorB > SharedLoadStreamB
The stream to load B from shared memory.
Definition: gemm_traits.h:750
cutlass::gemm::GemmGlobalTileTraits
Definition: gemm_global_tile.h:70
cutlass::gemm::GemmDesc::A
TensorRefA A
The source matrix A.
Definition: gemm_desc.h:96
threadblock_swizzle.h
Defies functors for mapping blockIdx to partitions of the GEMM computation.
cutlass::gemm::GemmDesc::problem_size
GemmCoord problem_size
The dimensions of the GEMM.
Definition: gemm_desc.h:90
cutlass::gemm::GlobalLoadStreamPair::SharedStorage
Defines a structure containing shared storage for each pair.
Definition: gemm_stream_pair.h:91
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalLoadStreamB
GlobalLoadStream< GemmOperand::kB, GlobalLoadIteratorB, SharedStoreIteratorB, GlobalTransformerB > GlobalLoadStreamB
The stream to load B from global memory to shared memory.
Definition: gemm_traits.h:733
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kColumnMajor, GemmConfig_ >::Scalar
GemmConfig_::ScalarB Scalar
The input scalar.
Definition: gemm_traits.h:201
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kRowMajor, GemmConfig_ >::SharedStoreTileTraits
GemmSharedStoreTileAbTraits< MultiplyAddScalar, Shape< GemmConfig_::kStages, GemmConfig_::OutputTile::kD/GemmConfig_::InstructionShape::kD, GemmConfig_::OutputTile::kH *GemmConfig_::InstructionShape::kD >, typename GlobalTileTraits::Threads, GemmConfig_::kScalarsPerStsB > SharedStoreTileTraits
The traits class to build the iterator to store data to shared memory for B^T.
Definition: gemm_traits.h:304
cutlass::gemm::GemmCoord
Definition: gemm_coord.h:43
cutlass::gemm::GemmTraits::This_
GemmTraits< GemmConfig_, GlobalLoadStreamA_, GlobalLoadStreamB_, SharedLoadStreamA_, SharedLoadStreamB_, Epilogue_, BlockSwizzle_, Index_, ClearAccumulators_ > This_
This traits.
Definition: gemm_traits.h:359
cutlass::gemm::GemmTraits::SharedLoadStreamB
SharedLoadStreamB_ SharedLoadStreamB
The iterator for B to load from shared memory.
Definition: gemm_traits.h:386
cutlass::gemm::GemmTraits::Params::partitionK_range
int partitionK_range
The K range for every partition except the last one.
Definition: gemm_traits.h:422
kernel_launch.h
Defines structures and helpers to launch CUDA kernels within CUTLASS.
cutlass::MatrixLayout::Kind
Kind
Enumeration defining fundamental contiguous layouts.
Definition: matrix_traits.h:159
cutlass::gemm::GemmTraits::GlobalLoadStream
GlobalLoadStreamPair< GlobalLoadStreamA, GlobalLoadStreamB, GemmConfig::kResidueInProlog > GlobalLoadStream
Assemble the global load streams for A/B.
Definition: gemm_traits.h:407
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kColumnMajor, GemmConfig_ >::GlobalTileTraits
GemmGlobalTileTraits< GemmOperand::kB, MatrixLayout::kColumnMajor, Scalar const, Shape< 1, GemmConfig_::OutputTile::kH, GemmConfig_::OutputTile::kD >, Shape< 1, GemmConfig_::kThreads/GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD >, GemmConfig_::kScalarsPerLdgB > GlobalTileTraits
The traits class to build the iterator to load data from global memory for B^N.
Definition: gemm_traits.h:219
cutlass::Copy
Definition: convert.h:69
cutlass::gemm::GemmTraits::MainLoopSharedStorage::threadblock_tile
ThreadblockTileStorage threadblock_tile
Stores the threadblock tile.
Definition: gemm_traits.h:664
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedLoadStreamA
SharedLoadStream< SharedLoadIteratorA > SharedLoadStreamA
The stream to load A from shared memory.
Definition: gemm_traits.h:742
cutlass::gemm::GemmSharedStoreTileAbTraits
Definition: gemm_shared_tile.h:38
cutlass::gemm::GemmDesc::B
TensorRefB B
The source matrix B.
Definition: gemm_desc.h:102
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kColumnMajor, GemmConfig_ >::SharedLoadTileTraits
GemmSharedLoadTileATraits< MultiplyAddScalar const, typename GemmConfig_::OutputTile, typename GemmConfig_::Warps, typename GemmConfig_::MultiplyAdd::ThreadsPerWarp, typename GemmConfig_::InstructionShape, GemmConfig_::kStages, GemmConfig_::kScalarsPerLdsA, 0 > SharedLoadTileTraits
The traits class to build the iterator to load from shared memory for A^N.
Definition: gemm_traits.h:114
cutlass::gemm::GemmTraits::Epilogue
Epilogue_ Epilogue
The epilogue.
Definition: gemm_traits.h:391
cutlass::gemm::GemmTraits::ScalarA
GlobalLoadStreamA_::Scalar ScalarA
The scalar for A.
Definition: gemm_traits.h:374
cutlass::IteratorAdvance::kH
Definition: tile_iterator.h:65
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kColumnMajor, GemmConfig_ >::GlobalTileTraits
GemmGlobalTileTraits< GemmOperand::kA, MatrixLayout::kColumnMajor, Scalar const, Shape< 1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kW >, Shape< 1, ShapeCount< typename GemmConfig_::Warps >::kCount, GemmConfig_::kWarpSize >, GemmConfig_::kScalarsPerLdgA > GlobalTileTraits
The traits class to build the iterator to load data from global memory for A^N.
Definition: gemm_traits.h:80
cutlass::gemm::GemmTraits::Params::global_to_shared_stream
GlobalLoadStream::Params global_to_shared_stream
Parameters object for the global load stream.
Definition: gemm_traits.h:425
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kColumnMajor, GemmConfig_ >::Scalar
GemmConfig_::ScalarA Scalar
The input scalar.
Definition: gemm_traits.h:62
cutlass::gemm::GemmSharedLoadTileBTraits
Definition: gemm_shared_tile.h:200
cutlass::gemm::GemmGlobalIteratorAb
Definition: gemm_global_tile.h:163
cutlass::gemm::GemmTraits::ScalarC
Epilogue::ScalarC ScalarC
The scalars in the epilogue.
Definition: gemm_traits.h:393
cutlass::gemm::GemmTraits::MultiplyAdd
GemmConfig::MultiplyAdd MultiplyAdd
The multiply-add functor.
Definition: gemm_traits.h:389
cutlass::gemm::GemmTraits::shared_load_fence
static CUTLASS_DEVICE void shared_load_fence(bool in_loop)
The memory fence for shared loads.
Definition: gemm_traits.h:682
cutlass::gemm::GemmTraits::GemmConfig
GemmConfig_ GemmConfig
The configuration.
Definition: gemm_traits.h:365
cutlass::gemm::GlobalLoadStream
Definition: gemm_global_stream.h:52
cutlass::gemm::GemmTileTraitsHelperB
Definition: gemm_traits.h:191
cutlass::gemm::ClearAccumulators
Definition: clear_accumulators.h:38
cutlass::gemm::GemmTraits::Params
Parameters object constructable on the host.
Definition: gemm_traits.h:416
cutlass::gemm::SharedStreamPair
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:180
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalTransformerB
Copy< typename GlobalLoadIteratorB::Fragment > GlobalTransformerB
The data converter for B before storing to shared memory.
Definition: gemm_traits.h:721
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kRowMajor, GemmConfig_ >::Scalar
GemmConfig_::ScalarB Scalar
The input scalar.
Definition: gemm_traits.h:272
cutlass::gemm::GlobalLoadStreamPair::Params::stream_b
StreamB::Params stream_b
Parameters object for StreamB.
Definition: gemm_stream_pair.h:67
cutlass::gemm::GemmDesc::batch_stride_A
long long batch_stride_A
batch stride for A operand
Definition: gemm_desc.h:99
cutlass::gemm::Gemm
Definition: gemm.h:92
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalLoadIteratorB
GemmGlobalIteratorAb< typename GemmTileTraitsHelperB_::GlobalTileTraits, Index_ > GlobalLoadIteratorB
The global iterator to load B from global memory.
Definition: gemm_traits.h:719
cutlass::gemm::GemmConfig::kResidueInProlog
static bool const kResidueInProlog
If true, residue is computed in the prologue.
Definition: gemm_config.h:136
cutlass::TileLoadIterator
An iterator implementing Tile Load Iterator Concept for loading a tile from memory.
Definition: tile_iterator.h:402
cutlass::gemm::GemmTraits::MainLoopSharedStorage
Definition: gemm_traits.h:662
cutlass::gemm::GemmTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const &partitonK_desc, Index partitionK_count_)
Helper to construct a partitionedK GEMM params.
Definition: gemm_traits.h:547
cutlass::gemm::GlobalLoadStreamPair
Collect the global load streams for multiplicands.
Definition: gemm_stream_pair.h:50
cutlass::gemm::GemmConfig::MultiplyAdd
MultiplyAdd_ MultiplyAdd
The functor to do D = A*B + C.
Definition: gemm_config.h:90
tile_allocation.h
Defines a fragment based on a Shape<> template.
cutlass::KernelLaunchConfiguration
Structure containing the basic launch configuration of a CUDA kernel.
Definition: kernel_launch.h:38
cutlass::gemm::GemmTraits::ClearAccumulators
ClearAccumulators_ ClearAccumulators
Clear the accumulators.
Definition: gemm_traits.h:401
cutlass::gemm::SharedLoadStream
Definition: gemm_shared_stream.h:45
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kRowMajor, GemmConfig_ >::GlobalTileTraits
GemmGlobalTileTraits< GemmOperand::kA, MatrixLayout::kRowMajor, Scalar const, Shape< 1, GemmConfig_::OutputTile::kW, GemmConfig_::OutputTile::kD >, Shape< 1, GemmConfig_::kThreads/GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kD >, GemmConfig_::kScalarsPerLdgA > GlobalTileTraits
The traits class to build the iterator to load data from global memory for A^T.
Definition: gemm_traits.h:143
cutlass::gemm::GlobalLoadStreamPair::Params
Parameters object.
Definition: gemm_stream_pair.h:62
reshape_tile.h
Defines a type for restructuring a tile.
gemm_operand.h
Defines constant expressions for mapping GEMM problem size and strides onto pitch-linear memory...
cutlass::gemm::GemmTraits::Params::problem_size
GemmCoord problem_size
GEMM problem size.
Definition: gemm_traits.h:419
gemm_desc.h
Implements a software-pipelined efficient GEMM.
cutlass::gemm::GemmTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(Index m, Index n, Index k, typename Epilogue::Scalar alpha, ScalarA const *d_a, Index lda, long long int batch_stride_A, ScalarB const *d_b, Index ldb, long long int batch_stride_B, typename Epilogue::Scalar beta, ScalarC const *d_c, Index ldc, long long int batch_stride_C, ScalarD *d_d, Index ldd, long long int batch_stride_D, Index batch_count)
Helper to construct a batched GEMM params.
Definition: gemm_traits.h:510
clear_accumulators.h
Defines abstractions for efficiently clearing accumulator tiles.
cutlass::TensorRef
Definition: tensor_ref.h:131
cutlass::gemm::GemmTraits::SharedStream
SharedStreamPair< SharedLoadStreamA, SharedLoadStreamB > SharedStream
Assemble the shared load streams for A/B.
Definition: gemm_traits.h:413
cutlass::gemm::GemmTraits::shared_store_fence
static CUTLASS_DEVICE void shared_store_fence(bool in_loop)
The memory fence for shared stores.
Definition: gemm_traits.h:690
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kRowMajor, GemmConfig_ >::Scalar
GemmConfig_::ScalarA Scalar
The input scalar.
Definition: gemm_traits.h:125
cutlass::gemm::GemmTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(Index m, Index n, Index k, typename Epilogue::Scalar alpha, ScalarA const *d_a, Index lda, ScalarB const *d_b, Index ldb, typename Epilogue::Scalar beta, ScalarC const *d_c, Index ldc, ScalarD *d_d, Index ldd, Index partitionK_count_)
Helper to construct a partitionedK GEMM params.
Definition: gemm_traits.h:631
cutlass::ZipTileAllocation
Manages a pair of tile allocations as if they are one allocation.
Definition: tile_allocation.h:125
cutlass::gemm::GemmTileTraitsHelperA
Definition: gemm_traits.h:52
cutlass::gemm::GemmTraits::Params::initialize
CUTLASS_HOST_DEVICE int initialize(GemmDesc_ const &desc)
Initialize the parameters.
Definition: gemm_traits.h:435
cutlass::GemmOperand::kB
Definition: matrix_traits.h:357
cutlass::gemm::IdentityBlockSwizzle
Definition: gemm/threadblock_swizzle.h:65
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kColumnMajor, GemmConfig_ >::SharedStoreTileTraits
GemmSharedStoreTileAbTraits< MultiplyAddScalar, Shape< GemmConfig_::kStages, GemmConfig_::OutputTile::kD/GemmConfig_::InstructionShape::kD, GemmConfig_::OutputTile::kW *GemmConfig_::InstructionShape::kD >, typename GlobalTileTraits::Threads, GemmConfig_::kScalarsPerStsA > SharedStoreTileTraits
The traits class to build the iterator to store data to shared memory for A^N.
Definition: gemm_traits.h:94
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kRowMajor, GemmConfig_ >::MultiplyAddScalar
GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar
The scalar stored in shared memory.
Definition: gemm_traits.h:274
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kColumnMajor, GemmConfig_ >::MultiplyAddScalar
GemmConfig_::MultiplyAdd::ScalarB MultiplyAddScalar
The scalar stored in shared memory.
Definition: gemm_traits.h:203
cutlass::gemm::GemmTraits::ScalarB
GlobalLoadStreamB_::Scalar ScalarB
The scalar for B.
Definition: gemm_traits.h:381
CUTLASS_HOST_DEVICE
#define CUTLASS_HOST_DEVICE
Definition: cutlass.h:46
gemm_config.h
Defines properties of GEMM computation that impose some constraints on caller.
cutlass::gemm::GemmTraits
Definition: gemm_traits.h:349
cutlass::Shape
A Shape implementing Layout Concept describing the dimensions of a cube.
Definition: shape.h:64
cutlass::MatrixLayout::kColumnMajor
Definition: matrix_traits.h:159
cutlass::gemm::GemmTraits::KernelClass
cutlass::gemm::Gemm< This_ > KernelClass
The struct that consumes this Traits.
Definition: gemm_traits.h:362
cutlass::MatrixLayout::kRowMajor
Definition: matrix_traits.h:159
cutlass::gemm::GemmTraits::Params::shared_stream
SharedStream::Params shared_stream
Parameters object for the shared load stream.
Definition: gemm_traits.h:428
cutlass::gemm::GemmGlobalTileTraits::Threads
ReshapeThreads< VectorizedTile, Threads_ >::Threads Threads
The threads shape.
Definition: gemm_global_tile.h:88
cutlass::gemm::GemmDesc::batch_stride_B
long long batch_stride_B
batch stride for B operand
Definition: gemm_desc.h:105
cutlass::gemm::GemmTraits::BlockSwizzle
BlockSwizzle_ BlockSwizzle
The block swizzle to reorganize the grid.
Definition: gemm_traits.h:397
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedLoadIteratorA
TileLoadIterator< typename GemmTileTraitsHelperA_::SharedLoadTileTraits, typename GemmTileTraitsHelperA_::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedLoadIteratorA
The iterator to load A from shared memory.
Definition: gemm_traits.h:740
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedLoadIteratorB
TileLoadIterator< typename GemmTileTraitsHelperB_::SharedLoadTileTraits, typename GemmTileTraitsHelperB_::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedLoadIteratorB
The iterator to load B from shared memory.
Definition: gemm_traits.h:748
cutlass::gemm::GemmTraits::ThreadblockTileStorage
GlobalLoadStream::ThreadblockTileStorage ThreadblockTileStorage
Memory needed to store the threadblock-scoped GEMM tile.
Definition: gemm_traits.h:410
cutlass::KernelLaunchConfiguration::block
dim3 block
CUDA threablock dimensions.
Definition: kernel_launch.h:44
cutlass::gemm::GemmEpilogue::ScalarC
GlobalLoadIteratorC::Scalar ScalarC
The scalar for C.
Definition: gemm_epilogue.h:83
cutlass::gemm::GemmTraits::Index
Index_ Index
The index.
Definition: gemm_traits.h:399
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kColumnMajor, GemmConfig_ >::MultiplyAddScalar
GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar
The scalar stored in shared memory.
Definition: gemm_traits.h:64
cutlass::gemm::SimplifiedGemmTraitsHelper::SharedStoreIteratorB
TileStoreIterator< typename GemmTileTraitsHelperB_::SharedStoreTileTraits, typename GemmTileTraitsHelperB_::SharedStoreTileTraits::Scalar, IteratorAdvance::kH, MemorySpace::kShared > SharedStoreIteratorB
The iterator to store B to shared memory.
Definition: gemm_traits.h:727
cutlass::gemm::GemmTraits::Params::epilogue
Epilogue::Params epilogue
The params for the epilogue.
Definition: gemm_traits.h:431
gemm_stream_pair.h
Defines a pair of GEMM tile streams.
cutlass::gemm::ClearAccumulators::SharedStorage
The shared storage.
Definition: clear_accumulators.h:40
gemm_global_stream.h
Implements efficient loading of the thread block-level tile from global memory and storing to shared ...
cutlass::gemm::GemmTraits::SharedStorage::main_loop
MainLoopSharedStorage main_loop
Definition: gemm_traits.h:676
cutlass::gemm::GemmTraits::kLayoutA
static MatrixLayout::Kind const kLayoutA
The layout of A.
Definition: gemm_traits.h:372
cutlass::KernelLaunchConfiguration::grid
dim3 grid
CUDA grid dimensions.
Definition: kernel_launch.h:41
cutlass::GemmOperand::kA
Definition: matrix_traits.h:357
cutlass::gemm::GemmTraits::MainLoopSharedStorage::global_to_shared_stream
GlobalLoadStream::SharedStorage global_to_shared_stream
Storage for GEMM global stream.
Definition: gemm_traits.h:667
cutlass::gemm::SharedStreamPair::Params
Parameters object passed to load iterators.
Definition: gemm_stream_pair.h:192
cutlass::gemm::SimplifiedGemmTraitsHelper
Definition: gemm_traits.h:698
gemm.h
Implements a software-pipelined efficient GEMM.
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalLoadIteratorA
GemmGlobalIteratorAb< typename GemmTileTraitsHelperA_::GlobalTileTraits, Index_ > GlobalLoadIteratorA
The global iterator to load A from global memory.
Definition: gemm_traits.h:701
cutlass::gemm::GemmTraits::OutputTile
GemmConfig::OutputTile OutputTile
The output tile.
Definition: gemm_traits.h:367
matrix_traits.h
Defines properties of matrices used to denote layout and operands to GEMM kernels.
cutlass::gemm::SimplifiedGemmTraitsHelper::GlobalTransformerA
Copy< typename GlobalLoadIteratorA::Fragment > GlobalTransformerA
The data converter for A before storing to shared memory.
Definition: gemm_traits.h:703
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kRowMajor, GemmConfig_ >::SharedLoadTileTraits
GemmSharedLoadTileBTraits< MultiplyAddScalar const, typename GemmConfig_::OutputTile, typename GemmConfig_::Warps, typename GemmConfig_::MultiplyAdd::ThreadsPerWarp, typename GemmConfig_::InstructionShape, GemmConfig_::kStages, GemmConfig_::kScalarsPerLdsB, 0 > SharedLoadTileTraits
The traits class to build the iterator to load from shared memory for B^T.
Definition: gemm_traits.h:324
cutlass::gemm::GemmTraits::MainLoopSharedStorage::clear
ClearAccumulators::SharedStorage clear
Storage for clearing accumulators.
Definition: gemm_traits.h:670
cutlass::gemm::GlobalLoadStreamPair::Params::stream_a
StreamA::Params stream_a
Parameters object for StreamA.
Definition: gemm_stream_pair.h:64
cutlass::gemm::GemmTileTraitsHelperB< MatrixLayout::kRowMajor, GemmConfig_ >::GlobalTileTraits
GemmGlobalTileTraits< GemmOperand::kB, MatrixLayout::kRowMajor, Scalar const, Shape< 1, GemmConfig_::OutputTile::kD, GemmConfig_::OutputTile::kH >, Shape< 1, ShapeCount< typename GemmConfig_::Warps >::kCount, GemmConfig_::kWarpSize >, GemmConfig_::kScalarsPerLdgB > GlobalTileTraits
The traits class to build the iterator to load data from global memory for B^T.
Definition: gemm_traits.h:290
gemm_shared_stream.h
Defines abstractions for managing loading and storing fragments to shared memory in the efficient GEM...
convert.h
Defines conversion operations among Fragments of different base type.
cutlass::gemm::SimplifiedGemmTraits
Definition: gemm_traits.h:773
cutlass::gemm::GemmConfig::OutputTile
OutputTile_ OutputTile
The tile.
Definition: gemm_config.h:88
cutlass::gemm::GemmTraits::kLayoutB
static MatrixLayout::Kind const kLayoutB
The layout of B.
Definition: gemm_traits.h:379
cutlass::TileStoreIterator
An iterator implementing Tile Store Iterator Concept for storing a tile to memory.
Definition: tile_iterator.h:841
cutlass::gemm::GemmTileTraitsHelperA< MatrixLayout::kRowMajor, GemmConfig_ >::MultiplyAddScalar
GemmConfig_::MultiplyAdd::ScalarA MultiplyAddScalar
The scalar stored in shared memory.
Definition: gemm_traits.h:127