Back to Cutlass

CUTLASS: default_mma_wmma_tensor_op.h Source File

docs/default__mma__wmma__tensor__op_8h_source.html

4.4.25.1 KB
Original Source

| | CUTLASS

CUDA Templates for Linear Algebra Subroutines and Solvers |

default_mma_wmma_tensor_op.h

[Go to the documentation of this file.](default mma wmma 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 **************************************************************************************************/

29 #pragma once

30

31 #include "cutlass/arch/wmma.h"

32

33 #if defined(CUTLASS_ARCH_WMMA_ENABLED)

34

35 #include "cutlass/cutlass.h"

36 #include "[cutlass/gemm/warp/mma_tensor_op_wmma.h](mma tensor op__wmma_8h.html)"

37

38 namespace cutlass {

39 namespace gemm {

40 namespace warp {

41

43

44 template <

46typename WarpShape_,

48typename InstructionShape_,

50typename ElementA_,

52typename LayoutA_,

54typename ElementB_,

56typename LayoutB_,

58typename ElementC_,

60typename LayoutC_,

62typename Operator_ = arch::OpMultiplyAdd,

64int PartitionsK = 1,

66int PartitionsN = 1

67 >

68 struct DefaultMmaTensorOpWmma;

69

71

73 template <

75typename WarpShape_,

77typename InstructionShape_,

79typename ElementA,

81typename LayoutA,

83typename ElementB,

85typename LayoutB,

87typename ElementC,

89typename LayoutC,

91typename Operator_,

93int PartitionsK,

95int PartitionsN>

96 struct DefaultMmaTensorOpWmma {

97using Policy = cutlass::gemm::warp::MmaTensorOpPolicy<

98 cutlass::arch::Wmma<

99 InstructionShape_,

100 ElementA,

101 LayoutA,

102 ElementB,

103 LayoutB,

104 ElementC,

105 LayoutC,

106 Operator_>,

107cutlass::MatrixShape<1, 1> >;

108

109// Define the warp-level tensor op

110using Type = cutlass::gemm::warp::MmaTensorOpWmma<

111 WarpShape_,

112 ElementA,

113 LayoutA,

114 ElementB,

115 LayoutB,

116 ElementC,

117 LayoutC,

118 Policy,

119 PartitionsK,

120 PartitionsN>;

121 };

122

124

125 } // namespace warp

126 } // namespace gemm

127 } // namespace cutlass

128

129 #endif

130

cutlass::MatrixShape

Describes the size of a matrix tile.

Definition: matrix_shape.h:42

cutlass

Definition: aligned_buffer.h:35

cutlass::gemm::warp::MmaTensorOpPolicy

Policy.

Definition: mma_tensor_op_policy.h:48

[mma_tensor_op_wmma.h](mma tensor op__wmma_8h.html)

Templates implementing warp-level matrix multiply-accumulate operations targeting Tensor Cores...

wmma.h

Templates exposing architecture support for warp matrix multiply-add (WMMA) operations.

cutlass.h

Basic include for CUTLASS.


Generated by 1.8.11